Source code for langchain.llms.mlflow_ai_gateway
from __future__ import annotations
from typing import Any, Dict, List, Mapping, Optional
from langchain.callbacks.manager import CallbackManagerForLLMRun
from langchain.llms.base import LLM
from langchain.pydantic_v1 import BaseModel, Extra
# Ignoring type because below is valid pydantic code
# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]class Params(BaseModel, extra=Extra.allow):  # type: ignore[call-arg]
    """Parameters for the MLflow AI Gateway LLM."""
    temperature: float = 0.0
    candidate_count: int = 1
    """The number of candidates to return."""
    stop: Optional[List[str]] = None
    max_tokens: Optional[int] = None
[docs]class MlflowAIGateway(LLM):
    """
    Wrapper around completions LLMs in the MLflow AI Gateway.
    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.llms import MlflowAIGateway
            completions = MlflowAIGateway(
                gateway_uri="<your-mlflow-ai-gateway-uri>",
                route="<your-mlflow-ai-gateway-completions-route>",
                params={
                    "temperature": 0.1
                }
            )
    """
    route: str
    gateway_uri: Optional[str] = None
    params: Optional[Params] = None
    def __init__(self, **kwargs: Any):
        try:
            import mlflow.gateway