codekingpro/portable-devtools

1from __future__ import annotations2 3import warnings4from typing import Any, Iterator, List, Optional5 6from langchain_core.embeddings import Embeddings7from pydantic import BaseModel8 9 10def _chunk(texts: List[str], size: int) -> Iterator[List[str]]:11    for i in range(0, len(texts), size):12        yield texts[i : i + size]13 14 15class MlflowAIGatewayEmbeddings(Embeddings, BaseModel):16    """MLflow AI Gateway embeddings.17 18    To use, you should have the ``mlflow[gateway]`` python package installed.19    For more information, see https://mlflow.org/docs/latest/gateway/index.html.20 21    Example:22        .. code-block:: python23 24            from langchain_community.embeddings import MlflowAIGatewayEmbeddings25 26            embeddings = MlflowAIGatewayEmbeddings(27                gateway_uri="<your-mlflow-ai-gateway-uri>",28                route="<your-mlflow-ai-gateway-embeddings-route>"29            )30    """31 32    route: str33    """The route to use for the MLflow AI Gateway API."""34    gateway_uri: Optional[str] = None35    """The URI for the MLflow AI Gateway API."""36 37    def __init__(self, **kwargs: Any):38        warnings.warn(39            "`MlflowAIGatewayEmbeddings` is deprecated. Use `MlflowEmbeddings` or "40            "`DatabricksEmbeddings` instead.",41            DeprecationWarning,42        )43        try:44            import mlflow.gateway45        except ImportError as e:46            raise ImportError(47                "Could not import `mlflow.gateway` module. "48                "Please install it with `pip install mlflow[gateway]`."49            ) from e50 51        super().__init__(**kwargs)52        if self.gateway_uri:53            mlflow.gateway.set_gateway_uri(self.gateway_uri)54 55    def _query(self, texts: List[str]) -> List[List[float]]:56        try:57            import mlflow.gateway58        except ImportError as e:59            raise ImportError(60                "Could not import `mlflow.gateway` module. "61                "Please install it with `pip install mlflow[gateway]`."62            ) from e63 64        embeddings = []65        for txt in _chunk(texts, 20):66            resp = mlflow.gateway.query(self.route, data={"text": txt})67            # response is List[List[float]]68            if isinstance(resp["embeddings"][0], List):69                embeddings.extend(resp["embeddings"])70            # response is List[float]71            else:72                embeddings.append(resp["embeddings"])73        return embeddings74 75    def embed_documents(self, texts: List[str]) -> List[List[float]]:76        return self._query(texts)77 78    def embed_query(self, text: str) -> List[float]:79        return self._query([text])[0]80