from__future__importannotationsimportwarningsfromtypingimportAny,Dict,List,Mapping,Optionalfromlangchain_core.callbacksimportCallbackManagerForLLMRunfromlangchain_core.language_models.llmsimportLLMfrompydanticimportBaseModel# Ignoring type because below is valid pydantic code# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]classParams(BaseModel,extra="allow"):# type: ignore[call-arg]"""Parameters for the MLflow AI Gateway LLM."""temperature:float=0.0candidate_count:int=1"""The number of candidates to return."""stop:Optional[List[str]]=Nonemax_tokens:Optional[int]=None
[docs]classMlflowAIGateway(LLM):"""MLflow AI Gateway LLMs. To use, you should have the ``mlflow[gateway]`` python package installed. For more information, see https://mlflow.org/docs/latest/gateway/index.html. Example: .. code-block:: python from langchain_community.llms import MlflowAIGateway completions = MlflowAIGateway( gateway_uri="<your-mlflow-ai-gateway-uri>", route="<your-mlflow-ai-gateway-completions-route>", params={ "temperature": 0.1 } ) """route:strgateway_uri:Optional[str]=Noneparams:Optional[Params]=Nonedef__init__(self,**kwargs:Any):warnings.warn("`MlflowAIGateway` is deprecated. Use `Mlflow` or `Databricks` instead.",DeprecationWarning,)try:importmlflow.gatewayexceptImportErrorase:raiseImportError("Could not import `mlflow.gateway` module. ""Please install it with `pip install mlflow[gateway]`.")fromesuper().__init__(**kwargs)ifself.gateway_uri:mlflow.gateway.set_gateway_uri(self.gateway_uri)@propertydef_default_params(self)->Dict[str,Any]:params:Dict[str,Any]={"gateway_uri":self.gateway_uri,"route":self.route,**(self.params.dict()ifself.paramselse{}),}returnparams@propertydef_identifying_params(self)->Mapping[str,Any]:returnself._default_paramsdef_call(self,prompt:str,stop:Optional[List[str]]=None,run_manager:Optional[CallbackManagerForLLMRun]=None,**kwargs:Any,)->str:try:importmlflow.gatewayexceptImportErrorase:raiseImportError("Could not import `mlflow.gateway` module. ""Please install it with `pip install mlflow[gateway]`.")fromedata:Dict[str,Any]={"prompt":prompt,**(self.params.dict()ifself.paramselse{}),}ifs:=(stopor(self.params.stopifself.paramselseNone)):data["stop"]=sresp=mlflow.gateway.query(self.route,data=data)returnresp["candidates"][0]["text"]@propertydef_llm_type(self)->str:return"mlflow-ai-gateway"