langroid 0.59.0b3__py3-none-any.whl → 0.59.2__py3-none-any.whl

langroid/embedding_models/models.py-e

@@ -1,563 +0,0 @@
-import atexit
-import os
-from functools import cached_property
-from typing import Any, Callable, Dict, List, Optional
-
-import requests
-import tiktoken
-from dotenv import load_dotenv
-from openai import AzureOpenAI, OpenAI
-
-from langroid.embedding_models.base import EmbeddingModel, EmbeddingModelsConfig
-from langroid.exceptions import LangroidImportError
-from langroid.language_models.provider_params import LangDBParams
-from langroid.mytypes import Embeddings
-from langroid.parsing.utils import batched
-from pydantic import ConfigDict
-
-AzureADTokenProvider = Callable[[], str]
-
-
-class OpenAIEmbeddingsConfig(EmbeddingModelsConfig):
-    model_type: str = "openai"
-    model_name: str = "text-embedding-3-small"
-    api_key: str = ""
-    api_base: Optional[str] = None
-    organization: str = ""
-    dims: int = 1536
-    context_length: int = 8192
-    langdb_params: LangDBParams = LangDBParams()
-
-    model_config = ConfigDict(env_prefix="OPENAI_")
-
-class AzureOpenAIEmbeddingsConfig(EmbeddingModelsConfig):
-    model_type: str = "azure-openai"
-    model_name: str = "text-embedding-3-small"
-    api_key: str = ""
-    api_base: str = ""
-    deployment_name: Optional[str] = None
-    # api_version defaulted to 2024-06-01 as per https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/embeddings?tabs=python-new
-    # change this to required  supported version
-    api_version: Optional[str] = "2024-06-01"
-    # TODO: Add auth support for Azure OpenAI via AzureADTokenProvider
-    azure_ad_token: Optional[str] = None
-    azure_ad_token_provider: Optional[AzureADTokenProvider] = None
-    dims: int = 1536
-    context_length: int = 8192
-
-    model_config = ConfigDict(env_prefix="AZURE_OPENAI_")
-
-class SentenceTransformerEmbeddingsConfig(EmbeddingModelsConfig):
-    model_type: str = "sentence-transformer"
-    model_name: str = "BAAI/bge-large-en-v1.5"
-    context_length: int = 512
-    data_parallel: bool = False
-    # Select device (e.g. "cuda", "cpu") when data parallel is disabled
-    device: Optional[str] = None
-    # Select devices when data parallel is enabled
-    devices: Optional[list[str]] = None
-
-
-class FastEmbedEmbeddingsConfig(EmbeddingModelsConfig):
-    """Config for qdrant/fastembed embeddings,
-    see here: https://github.com/qdrant/fastembed
-    """
-
-    model_type: str = "fastembed"
-    model_name: str = "BAAI/bge-small-en-v1.5"
-    batch_size: int = 256
-    cache_dir: Optional[str] = None
-    threads: Optional[int] = None
-    parallel: Optional[int] = None
-    additional_kwargs: Dict[str, Any] = {}
-
-
-class LlamaCppServerEmbeddingsConfig(EmbeddingModelsConfig):
-    api_base: str = ""
-    context_length: int = 2048
-    batch_size: int = 2048
-
-
-class GeminiEmbeddingsConfig(EmbeddingModelsConfig):
-    model_type: str = "gemini"
-    model_name: str = "models/text-embedding-004"
-    api_key: str = ""
-    dims: int = 768
-    batch_size: int = 512
-
-
-class EmbeddingFunctionCallable:
-    """
-    A callable class designed to generate embeddings for a list of texts using
-    the OpenAI or Azure OpenAI API, with automatic retries on failure.
-
-    Attributes:
-        embed_model (EmbeddingModel): An instance of EmbeddingModel that provides
-               configuration and utilities for generating embeddings.
-
-    Methods:
-        __call__(input: List[str]) -> Embeddings: Generate embeddings for
-                                a list of input texts.
-    """
-
-    def __init__(self, embed_model: EmbeddingModel, batch_size: int = 512):
-        """
-        Initialize the EmbeddingFunctionCallable with a specific model.
-
-        Args:
-            model ( OpenAIEmbeddings or AzureOpenAIEmbeddings): An instance of
-                            OpenAIEmbeddings or AzureOpenAIEmbeddings to use for
-                            generating embeddings.
-            batch_size (int): Batch size
-        """
-        self.embed_model = embed_model
-        self.batch_size = batch_size
-
-    def __call__(self, input: List[str]) -> Embeddings:
-        """
-        Generate embeddings for a given list of input texts using the OpenAI API,
-        with retries on failure.
-
-        This method:
-        - Truncates each text in the input list to the model's maximum context length.
-        - Processes the texts in batches to generate embeddings efficiently.
-        - Automatically retries the embedding generation process with exponential
-        backoff in case of failures.
-
-        Args:
-            input (List[str]): A list of input texts to generate embeddings for.
-
-        Returns:
-            Embeddings: A list of embedding vectors corresponding to the input texts.
-        """
-        embeds = []
-        if isinstance(self.embed_model, (OpenAIEmbeddings, AzureOpenAIEmbeddings)):
-            # Truncate texts to context length while preserving text format
-            truncated_texts = self.embed_model.truncate_texts(input)
-
-            # Process in batches
-            for batch in batched(truncated_texts, self.batch_size):
-                result = self.embed_model.client.embeddings.create(
-                    input=batch, model=self.embed_model.config.model_name  # type: ignore
-                )
-                batch_embeds = [d.embedding for d in result.data]
-                embeds.extend(batch_embeds)
-
-        elif isinstance(self.embed_model, SentenceTransformerEmbeddings):
-            if self.embed_model.config.data_parallel:
-                embeds = self.embed_model.model.encode_multi_process(
-                    input,
-                    self.embed_model.pool,
-                    batch_size=self.batch_size,
-                ).tolist()
-            else:
-                for str_batch in batched(input, self.batch_size):
-                    batch_embeds = self.embed_model.model.encode(
-                        str_batch, convert_to_numpy=True
-                    ).tolist()  # type: ignore
-                    embeds.extend(batch_embeds)
-
-        elif isinstance(self.embed_model, FastEmbedEmbeddings):
-            embeddings = self.embed_model.model.embed(
-                input, batch_size=self.batch_size, parallel=self.embed_model.parallel
-            )
-
-            embeds = [embedding.tolist() for embedding in embeddings]
-        elif isinstance(self.embed_model, LlamaCppServerEmbeddings):
-            for input_string in input:
-                tokenized_text = self.embed_model.tokenize_string(input_string)
-                for token_batch in batched(tokenized_text, self.batch_size):
-                    gen_embedding = self.embed_model.generate_embedding(
-                        self.embed_model.detokenize_string(list(token_batch))
-                    )
-                    embeds.append(gen_embedding)
-        elif isinstance(self.embed_model, GeminiEmbeddings):
-            embeds = self.embed_model.generate_embeddings(input)
-        return embeds
-
-
-class OpenAIEmbeddings(EmbeddingModel):
-    def __init__(self, config: OpenAIEmbeddingsConfig = OpenAIEmbeddingsConfig()):
-        super().__init__()
-        self.config = config
-        load_dotenv()
-
-        # Check if using LangDB
-        self.is_langdb = self.config.model_name.startswith("langdb/")
-
-        if self.is_langdb:
-            self.config.model_name = self.config.model_name.replace("langdb/", "")
-            self.config.api_base = self.config.langdb_params.base_url
-            project_id = self.config.langdb_params.project_id
-            if project_id:
-                self.config.api_base += "/" + project_id + "/v1"
-            self.config.api_key = self.config.langdb_params.api_key
-
-        if not self.config.api_key:
-            self.config.api_key = os.getenv("OPENAI_API_KEY", "")
-
-        self.config.organization = os.getenv("OPENAI_ORGANIZATION", "")
-
-        if self.config.api_key == "":
-            if self.is_langdb:
-                raise ValueError(
-                    """
-                    LANGDB_API_KEY must be set in .env or your environment 
-                    to use OpenAIEmbeddings via LangDB.
-                    """
-                )
-            else:
-                raise ValueError(
-                    """
-                    OPENAI_API_KEY must be set in .env or your environment 
-                    to use OpenAIEmbeddings.
-                    """
-                )
-
-        self.client = OpenAI(
-            base_url=self.config.api_base,
-            api_key=self.config.api_key,
-            organization=self.config.organization,
-        )
-        model_for_tokenizer = self.config.model_name
-        if model_for_tokenizer.startswith("openai/"):
-            self.config.model_name = model_for_tokenizer.replace("openai/", "")
-        self.tokenizer = tiktoken.encoding_for_model(self.config.model_name)
-
-    def truncate_texts(self, texts: List[str]) -> List[str] | List[List[int]]:
-        """
-        Truncate texts to the embedding model's context length.
-        TODO: Maybe we should show warning, and consider doing T5 summarization?
-        """
-        truncated_tokens = [
-            self.tokenizer.encode(text, disallowed_special=())[
-                : self.config.context_length
-            ]
-            for text in texts
-        ]
-
-        if self.is_langdb:
-            # LangDB embedding endpt only works with strings, not tokens
-            return [self.tokenizer.decode(tokens) for tokens in truncated_tokens]
-        return truncated_tokens
-
-    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
-        return EmbeddingFunctionCallable(self, self.config.batch_size)
-
-    @property
-    def embedding_dims(self) -> int:
-        return self.config.dims
-
-
-class AzureOpenAIEmbeddings(EmbeddingModel):
-    """
-    Azure OpenAI embeddings model implementation.
-    """
-
-    def __init__(
-        self, config: AzureOpenAIEmbeddingsConfig = AzureOpenAIEmbeddingsConfig()
-    ):
-        """
-        Initializes Azure OpenAI embeddings model.
-
-        Args:
-            config: Configuration for Azure OpenAI embeddings model.
-        Raises:
-            ValueError: If required Azure config values are not set.
-        """
-        super().__init__()
-        self.config = config
-        load_dotenv()
-
-        if self.config.api_key == "":
-            raise ValueError(
-                """AZURE_OPENAI_API_KEY env variable must be set to use 
-            AzureOpenAIEmbeddings. Please set the AZURE_OPENAI_API_KEY value 
-            in your .env file."""
-            )
-
-        if self.config.api_base == "":
-            raise ValueError(
-                """AZURE_OPENAI_API_BASE env variable must be set to use 
-            AzureOpenAIEmbeddings. Please set the AZURE_OPENAI_API_BASE value 
-            in your .env file."""
-            )
-        self.client = AzureOpenAI(
-            api_key=self.config.api_key,
-            api_version=self.config.api_version,
-            azure_endpoint=self.config.api_base,
-            azure_deployment=self.config.deployment_name,
-        )
-        self.tokenizer = tiktoken.encoding_for_model(self.config.model_name)
-
-    def truncate_texts(self, texts: List[str]) -> List[str] | List[List[int]]:
-        """
-        Truncate texts to the embedding model's context length.
-        TODO: Maybe we should show warning, and consider doing T5 summarization?
-        """
-        return [
-            self.tokenizer.encode(text, disallowed_special=())[
-                : self.config.context_length
-            ]
-            for text in texts
-        ]
-
-    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
-        """Get the embedding function for Azure OpenAI.
-
-        Returns:
-            Callable that generates embeddings for input texts.
-        """
-        return EmbeddingFunctionCallable(self, self.config.batch_size)
-
-    @property
-    def embedding_dims(self) -> int:
-        return self.config.dims
-
-
-STEC = SentenceTransformerEmbeddingsConfig
-
-
-class SentenceTransformerEmbeddings(EmbeddingModel):
-    def __init__(self, config: STEC = STEC()):
-        # this is an "extra" optional dependency, so we import it here
-        try:
-            from sentence_transformers import SentenceTransformer
-            from transformers import AutoTokenizer
-        except ImportError:
-            raise ImportError(
-                """
-                To use sentence_transformers embeddings, 
-                you must install langroid with the [hf-embeddings] extra, e.g.:
-                pip install "langroid[hf-embeddings]"
-                """
-            )
-
-        super().__init__()
-        self.config = config
-
-        self.model = SentenceTransformer(
-            self.config.model_name,
-            device=self.config.device,
-        )
-        if self.config.data_parallel:
-            self.pool = self.model.start_multi_process_pool(
-                self.config.devices  # type: ignore
-            )
-            atexit.register(
-                lambda: SentenceTransformer.stop_multi_process_pool(self.pool)
-            )
-
-        self.tokenizer = AutoTokenizer.from_pretrained(self.config.model_name)
-        self.config.context_length = self.tokenizer.model_max_length
-
-    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
-        return EmbeddingFunctionCallable(self, self.config.batch_size)
-
-    @property
-    def embedding_dims(self) -> int:
-        dims = self.model.get_sentence_embedding_dimension()
-        if dims is None:
-            raise ValueError(
-                f"Could not get embedding dimension for model {self.config.model_name}"
-            )
-        return dims  # type: ignore
-
-
-class FastEmbedEmbeddings(EmbeddingModel):
-    def __init__(self, config: FastEmbedEmbeddingsConfig = FastEmbedEmbeddingsConfig()):
-        try:
-            from fastembed import TextEmbedding
-        except ImportError:
-            raise LangroidImportError("fastembed", extra="fastembed")
-
-        super().__init__()
-        self.config = config
-        self.batch_size = config.batch_size
-        self.parallel = config.parallel
-
-        self.model = TextEmbedding(
-            model_name=self.config.model_name,
-            cache_dir=self.config.cache_dir,
-            threads=self.config.threads,
-            **self.config.additional_kwargs,
-        )
-
-    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
-        return EmbeddingFunctionCallable(self, self.config.batch_size)
-
-    @cached_property
-    def embedding_dims(self) -> int:
-        embed_func = self.embedding_fn()
-        return len(embed_func(["text"])[0])
-
-
-LCSEC = LlamaCppServerEmbeddingsConfig
-
-
-class LlamaCppServerEmbeddings(EmbeddingModel):
-    def __init__(self, config: LCSEC = LCSEC()):
-        super().__init__()
-        self.config = config
-
-        if self.config.api_base == "":
-            raise ValueError(
-                """Api Base MUST be set for Llama Server Embeddings.
-                """
-            )
-
-        self.tokenize_url = self.config.api_base + "/tokenize"
-        self.detokenize_url = self.config.api_base + "/detokenize"
-        self.embedding_url = self.config.api_base + "/embeddings"
-
-    def tokenize_string(self, text: str) -> List[int]:
-        data = {"content": text, "add_special": False, "with_pieces": False}
-        response = requests.post(self.tokenize_url, json=data)
-
-        if response.status_code == 200:
-            tokens = response.model_dump_json()["tokens"]
-            if not (isinstance(tokens, list) and isinstance(tokens[0], (int, float))):
-                # not all(isinstance(token, (int, float)) for token in tokens):
-                raise ValueError(
-                    """Tokenizer endpoint has not returned the correct format. 
-                   Is the URL correct?
-                """
-                )
-            return tokens
-        else:
-            raise requests.HTTPError(
-                self.tokenize_url,
-                response.status_code,
-                "Failed to connect to tokenization provider",
-            )
-
-    def detokenize_string(self, tokens: List[int]) -> str:
-        data = {"tokens": tokens}
-        response = requests.post(self.detokenize_url, json=data)
-
-        if response.status_code == 200:
-            text = response.model_dump_json()["content"]
-            if not isinstance(text, str):
-                raise ValueError(
-                    """Deokenizer endpoint has not returned the correct format. 
-                   Is the URL correct?
-                """
-                )
-            return text
-        else:
-            raise requests.HTTPError(
-                self.detokenize_url,
-                response.status_code,
-                "Failed to connect to detokenization provider",
-            )
-
-    def truncate_string_to_context_size(self, text: str) -> str:
-        tokens = self.tokenize_string(text)
-        tokens = tokens[: self.config.context_length]
-        return self.detokenize_string(tokens)
-
-    def generate_embedding(self, text: str) -> List[int | float]:
-        data = {"content": text}
-        response = requests.post(self.embedding_url, json=data)
-
-        if response.status_code == 200:
-            embeddings = response.model_dump_json()["embedding"]
-            if not (
-                isinstance(embeddings, list) and isinstance(embeddings[0], (int, float))
-            ):
-                raise ValueError(
-                    """Embedding endpoint has not returned the correct format. 
-                   Is the URL correct?
-                """
-                )
-            return embeddings
-        else:
-            raise requests.HTTPError(
-                self.embedding_url,
-                response.status_code,
-                "Failed to connect to embedding provider",
-            )
-
-    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
-        return EmbeddingFunctionCallable(self, self.config.batch_size)
-
-    @property
-    def embedding_dims(self) -> int:
-        return self.config.dims
-
-
-class GeminiEmbeddings(EmbeddingModel):
-    def __init__(self, config: GeminiEmbeddingsConfig = GeminiEmbeddingsConfig()):
-        try:
-            from google import genai
-        except ImportError as e:
-            raise LangroidImportError(extra="google-genai", error=str(e))
-        super().__init__()
-        self.config = config
-        load_dotenv()
-        self.config.api_key = os.getenv("GEMINI_API_KEY", "")
-
-        if self.config.api_key == "":
-            raise ValueError(
-                """
-                GEMINI_API_KEY env variable must be set to use GeminiEmbeddings.
-                """
-            )
-        self.client = genai.Client(api_key=self.config.api_key)
-
-    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
-        return EmbeddingFunctionCallable(self, self.config.batch_size)
-
-    def generate_embeddings(self, texts: List[str]) -> List[List[float]]:
-        """Generates embeddings for a list of input texts."""
-        all_embeddings: List[List[float]] = []
-
-        for batch in batched(texts, self.config.batch_size):
-            result = self.client.models.embed_content(  # type: ignore[attr-defined]
-                model=self.config.model_name,
-                contents=batch,  # type: ignore
-            )
-
-            if not hasattr(result, "embeddings") or not isinstance(
-                result.embeddings, list
-            ):
-                raise ValueError(
-                    "Unexpected format for embeddings: missing or incorrect type"
-                )
-
-            # Extract .values from ContentEmbedding objects
-            all_embeddings.extend(
-                [emb.values for emb in result.embeddings]  # type: ignore
-            )
-
-        return all_embeddings
-
-    @property
-    def embedding_dims(self) -> int:
-        return self.config.dims
-
-
-def embedding_model(embedding_fn_type: str = "openai") -> EmbeddingModel:
-    """
-    Args:
-        embedding_fn_type: Type of embedding model to use. Options are:
-         - "openai",
-         - "azure-openai",
-         - "sentencetransformer", or
-         - "fastembed".
-            (others may be added in the future)
-    Returns:
-        EmbeddingModel: The corresponding embedding model class.
-    """
-    if embedding_fn_type == "openai":
-        return OpenAIEmbeddings  # type: ignore
-    elif embedding_fn_type == "azure-openai":
-        return AzureOpenAIEmbeddings  # type: ignore
-    elif embedding_fn_type == "fastembed":
-        return FastEmbedEmbeddings  # type: ignore
-    elif embedding_fn_type == "llamacppserver":
-        return LlamaCppServerEmbeddings  # type: ignore
-    elif embedding_fn_type == "gemini":
-        return GeminiEmbeddings  # type: ignore
-    else:  # default sentence transformer
-        return SentenceTransformerEmbeddings  # type: ignore

langroid/language_models/azure_openai.py-e

@@ -1,134 +0,0 @@
-import logging
-from typing import Callable
-
-from dotenv import load_dotenv
-from httpx import Timeout
-from openai import AsyncAzureOpenAI, AzureOpenAI
-
-from langroid.language_models.openai_gpt import (
-    OpenAIGPT,
-    OpenAIGPTConfig,
-)
-from pydantic import ConfigDict
-
-azureStructuredOutputList = [
-    "2024-08-06",
-    "2024-11-20",
-]
-
-azureStructuredOutputAPIMin = "2024-08-01-preview"
-
-logger = logging.getLogger(__name__)
-
-
-class AzureConfig(OpenAIGPTConfig):
-    """
-    Configuration for Azure OpenAI GPT.
-
-    Attributes:
-        type (str): should be ``azure.``
-        api_version (str): can be set in the ``.env`` file as
-            ``AZURE_OPENAI_API_VERSION.``
-        deployment_name (str|None): can be optionally set in the ``.env`` file as
-            ``AZURE_OPENAI_DEPLOYMENT_NAME`` and should be based the custom name you
-            chose for your deployment when you deployed a model.
-        model_name (str): [DEPRECATED] can be set in the ``.env``
-            file as ``AZURE_OPENAI_MODEL_NAME``
-            and should be based on the model name chosen during setup.
-        chat_model (str): the chat model name to use. Can be set via
-            the env variable ``AZURE_OPENAI_CHAT_MODEL``.
-            Recommended to use this instead of ``model_name``.
-
-    """
-
-    api_key: str = ""  # CAUTION: set this ONLY via env var AZURE_OPENAI_API_KEY
-    type: str = "azure"
-    api_version: str = "2023-05-15"
-    deployment_name: str | None = None
-    model_name: str = ""
-    api_base: str = ""
-
-    # Alternatively, bring your own clients:
-    azure_openai_client_provider: Callable[[], AzureOpenAI] | None = None
-    azure_openai_async_client_provider: Callable[[], AsyncAzureOpenAI] | None = None
-
-    # all of the vars above can be set via env vars,
-    # by upper-casing the name and prefixing with `env_prefix`, e.g.
-    # AZURE_OPENAI_API_VERSION=2023-05-15
-    # This is either done in the .env file, or via an explicit
-    # `export AZURE_OPENAI_API_VERSION=...`
-    model_config = ConfigDict(env_prefix="AZURE_OPENAI_")
-
-    def __init__(self, **kwargs) -> None:  # type: ignore
-        if "model_name" in kwargs and "chat_model" not in kwargs:
-            kwargs["chat_model"] = kwargs["model_name"]
-        super().__init__(**kwargs)
-
-
-class AzureGPT(OpenAIGPT):
-    """
-    Class to access OpenAI LLMs via Azure. These env variables can be obtained from the
-    file `.azure_env`. Azure OpenAI doesn't support ``completion``
-    """
-
-    def __init__(self, config: AzureConfig):
-        # This will auto-populate config values from .env file
-        load_dotenv()
-        super().__init__(config)
-        self.config: AzureConfig = config
-
-        if (
-            self.config.azure_openai_client_provider
-            or self.config.azure_openai_async_client_provider
-        ):
-            if not self.config.azure_openai_client_provider:
-                self.client = None
-                logger.warning(
-                    "Using user-provided Azure OpenAI client, but only async "
-                    "client has been provided. Synchronous calls will fail."
-                )
-            if not self.config.azure_openai_async_client_provider:
-                self.async_client = None
-                logger.warning(
-                    "Using user-provided Azure OpenAI client, but no async "
-                    "client has been provided. Asynchronous calls will fail."
-                )
-
-            if self.config.azure_openai_client_provider:
-                self.client = self.config.azure_openai_client_provider()
-            if self.config.azure_openai_async_client_provider:
-                self.async_client = self.config.azure_openai_async_client_provider()
-                self.async_client.timeout = Timeout(self.config.timeout)
-        else:
-            if self.config.api_key == "":
-                raise ValueError(
-                    """
-                    AZURE_OPENAI_API_KEY not set in .env file,
-                    please set it to your Azure API key."""
-                )
-
-            if self.config.api_base == "":
-                raise ValueError(
-                    """
-                    AZURE_OPENAI_API_BASE not set in .env file,
-                    please set it to your Azure API key."""
-                )
-
-            self.client = AzureOpenAI(
-                api_key=self.config.api_key,
-                azure_endpoint=self.config.api_base,
-                api_version=self.config.api_version,
-                azure_deployment=self.config.deployment_name,
-            )
-            self.async_client = AsyncAzureOpenAI(
-                api_key=self.config.api_key,
-                azure_endpoint=self.config.api_base,
-                api_version=self.config.api_version,
-                azure_deployment=self.config.deployment_name,
-                timeout=Timeout(self.config.timeout),
-            )
-
-        self.supports_json_schema = (
-            self.config.api_version >= azureStructuredOutputAPIMin
-            and self.config.api_version in azureStructuredOutputList
-        )