codekingpro/portable-devtools

1from __future__ import annotations2 3from typing import Any, Dict, List, Mapping, Optional4from urllib.parse import urlparse5 6from langchain_core.callbacks import CallbackManagerForLLMRun7from langchain_core.language_models import LLM8from pydantic import Field, PrivateAttr9 10 11class Mlflow(LLM):12    """MLflow LLM service.13 14    To use, you should have the `mlflow[genai]` python package installed.15    For more information, see https://mlflow.org/docs/latest/llms/deployments.16 17    Example:18        .. code-block:: python19 20            from langchain_community.llms import Mlflow21 22            completions = Mlflow(23                target_uri="http://localhost:5000",24                endpoint="test",25                temperature=0.1,26            )27    """28 29    endpoint: str30    """The endpoint to use."""31    target_uri: str32    """The target URI to use."""33    temperature: float = 0.034    """The sampling temperature."""35    n: int = 136    """The number of completion choices to generate."""37    stop: Optional[List[str]] = None38    """The stop sequence."""39    max_tokens: Optional[int] = None40    """The maximum number of tokens to generate."""41    extra_params: Dict[str, Any] = Field(default_factory=dict)42    """Any extra parameters to pass to the endpoint."""43 44    """Extra parameters such as `temperature`."""45    _client: Any = PrivateAttr()46 47    def __init__(self, **kwargs: Any):48        super().__init__(**kwargs)49        self._validate_uri()50        try:51            from mlflow.deployments import get_deploy_client52 53            self._client = get_deploy_client(self.target_uri)54        except ImportError as e:55            raise ImportError(56                "Failed to create the client. "57                "Please run `pip install mlflow[genai]` to install "58                "required dependencies."59            ) from e60 61    def _validate_uri(self) -> None:62        if self.target_uri == "databricks":63            return64        allowed = ["http", "https", "databricks"]65        if urlparse(self.target_uri).scheme not in allowed:66            raise ValueError(67                f"Invalid target URI: {self.target_uri}. "68                f"The scheme must be one of {allowed}."69            )70 71    @property72    def _default_params(self) -> Dict[str, Any]:73        return {74            "target_uri": self.target_uri,75            "endpoint": self.endpoint,76            "temperature": self.temperature,77            "n": self.n,78            "stop": self.stop,79            "max_tokens": self.max_tokens,80            "extra_params": self.extra_params,81        }82 83    @property84    def _identifying_params(self) -> Mapping[str, Any]:85        return self._default_params86 87    def _call(88        self,89        prompt: str,90        stop: Optional[List[str]] = None,91        run_manager: Optional[CallbackManagerForLLMRun] = None,92        **kwargs: Any,93    ) -> str:94        data: Dict[str, Any] = {95            "prompt": prompt,96            "temperature": self.temperature,97            "n": self.n,98            **self.extra_params,99            **kwargs,100        }101        if stop := self.stop or stop:102            data["stop"] = stop103        if self.max_tokens is not None:104            data["max_tokens"] = self.max_tokens105 106        resp = self._client.predict(endpoint=self.endpoint, inputs=data)107        return resp["choices"][0]["text"]108 109    @property110    def _llm_type(self) -> str:111        return "mlflow"112