langroid 0.33.6__py3-none-any.whl → 0.33.8__py3-none-any.whl

langroid/embedding_models/remote_embeds.py

@@ -0,0 +1,153 @@
+"""
+If run as a script, starts an RPC server which handles remote
+embedding requests:
+
+For example:
+python3 -m langroid.embedding_models.remote_embeds --port `port`
+
+where `port` is the port at which the service is exposed.  Currently,
+supports insecure connections only, and this should NOT be exposed to
+the internet.
+"""
+
+import atexit
+import subprocess
+import time
+from typing import Callable, Optional
+
+import grpc
+from fire import Fire
+
+import langroid.embedding_models.models as em
+import langroid.embedding_models.protoc.embeddings_pb2 as embeddings_pb
+import langroid.embedding_models.protoc.embeddings_pb2_grpc as embeddings_grpc
+from langroid.mytypes import Embeddings
+
+
+class RemoteEmbeddingRPCs(embeddings_grpc.EmbeddingServicer):
+    def __init__(
+        self,
+        model_name: str,
+        batch_size: int,
+        data_parallel: bool,
+        device: Optional[str],
+        devices: Optional[list[str]],
+    ):
+        super().__init__()
+        self.embedding_fn = em.SentenceTransformerEmbeddings(
+            em.SentenceTransformerEmbeddingsConfig(
+                model_name=model_name,
+                batch_size=batch_size,
+                data_parallel=data_parallel,
+                device=device,
+                devices=devices,
+            )
+        ).embedding_fn()
+
+    def Embed(
+        self, request: embeddings_pb.EmbeddingRequest, _: grpc.RpcContext
+    ) -> embeddings_pb.BatchEmbeds:
+        embeds = self.embedding_fn(list(request.strings))
+
+        embeds_pb = [embeddings_pb.Embed(embed=e) for e in embeds]
+
+        return embeddings_pb.BatchEmbeds(embeds=embeds_pb)
+
+
+class RemoteEmbeddingsConfig(em.SentenceTransformerEmbeddingsConfig):
+    api_base: str = "localhost"
+    port: int = 50052
+    # The below are used only when waiting for server creation
+    poll_delay: float = 0.01
+    max_retries: int = 1000
+
+
+class RemoteEmbeddings(em.SentenceTransformerEmbeddings):
+    def __init__(self, config: RemoteEmbeddingsConfig = RemoteEmbeddingsConfig()):
+        super().__init__(config)
+        self.config: RemoteEmbeddingsConfig = config
+        self.have_started_server: bool = False
+
+    def embedding_fn(self) -> Callable[[list[str]], Embeddings]:
+        def fn(texts: list[str]) -> Embeddings:
+            url = f"{self.config.api_base}:{self.config.port}"
+            with grpc.insecure_channel(url) as channel:
+                stub = embeddings_grpc.EmbeddingStub(channel)  # type: ignore
+                response = stub.Embed(
+                    embeddings_pb.EmbeddingRequest(
+                        strings=texts,
+                    )
+                )
+
+                return [list(emb.embed) for emb in response.embeds]
+
+        def with_handling(texts: list[str]) -> Embeddings:
+            # In local mode, start the server if it has not already
+            # been started
+            if self.config.api_base == "localhost" and not self.have_started_server:
+                try:
+                    return fn(texts)
+                # Occurs when the server hasn't been started
+                except grpc.RpcError:
+                    self.have_started_server = True
+                    # Start the server
+                    proc = subprocess.Popen(
+                        [
+                            "python3",
+                            __file__,
+                            "--bind_address_base",
+                            self.config.api_base,
+                            "--port",
+                            str(self.config.port),
+                            "--batch_size",
+                            str(self.config.batch_size),
+                            "--model_name",
+                            self.config.model_name,
+                        ],
+                    )
+
+                    atexit.register(lambda: proc.terminate())
+
+                    for _ in range(self.config.max_retries - 1):
+                        try:
+                            return fn(texts)
+                        except grpc.RpcError:
+                            time.sleep(self.config.poll_delay)
+
+            # The remote is not local or we have exhausted retries
+            # We should now raise an error if the server is not accessible
+            return fn(texts)
+
+        return with_handling
+
+
+async def serve(
+    bind_address_base: str = "localhost",
+    port: int = 50052,
+    batch_size: int = 512,
+    data_parallel: bool = False,
+    device: Optional[str] = None,
+    devices: Optional[list[str]] = None,
+    model_name: str = "BAAI/bge-large-en-v1.5",
+) -> None:
+    """Starts the RPC server."""
+    server = grpc.aio.server()
+    embeddings_grpc.add_EmbeddingServicer_to_server(
+        RemoteEmbeddingRPCs(
+            model_name=model_name,
+            batch_size=batch_size,
+            data_parallel=data_parallel,
+            device=device,
+            devices=devices,
+        ),
+        server,
+    )  # type: ignore
+    url = f"{bind_address_base}:{port}"
+    server.add_insecure_port(url)
+    await server.start()
+    print(f"Embedding server started, listening on {url}")
+    await server.wait_for_termination()
+
+
+if __name__ == "__main__":
+    Fire(serve)

langroid/exceptions.py

@@ -0,0 +1,71 @@
+from typing import Optional
+
+
+class XMLException(Exception):
+    def __init__(self, message: str) -> None:
+        super().__init__(message)
+
+
+class InfiniteLoopException(Exception):
+    def __init__(self, message: str = "Infinite loop detected", *args: object) -> None:
+        super().__init__(message, *args)
+
+
+class LangroidImportError(ImportError):
+    def __init__(
+        self,
+        package: Optional[str] = None,
+        extra: Optional[str] = None,
+        error: str = "",
+        *args: object,
+    ) -> None:
+        """
+        Generate helpful warning when attempting to import package or module.
+
+        Args:
+            package (str): The name of the package to import.
+            extra (str): The name of the extras package required for this import.
+            error (str): The error message to display. Depending on context, we
+                can set this by capturing the ImportError message.
+
+        """
+        if error == "" and package is not None:
+            error = f"{package} is not installed by default with Langroid.\n"
+
+        if extra:
+            install_help = f"""
+                If you want to use it, please install langroid 
+                with the `{extra}` extra, for example:
+                
+                If you are using pip:
+                pip install "langroid[{extra}]"
+                
+                For multiple extras, you can separate them with commas:
+                pip install "langroid[{extra},another-extra]"
+                
+                If you are using Poetry:
+                poetry add langroid --extras "{extra}"
+                
+                For multiple extras with Poetry, list them with spaces:
+                poetry add langroid --extras "{extra} another-extra"
+
+                If you are using uv:
+                uv add "langroid[{extra}]"
+
+                For multiple extras with uv, you can separate them with commas: 
+                uv add "langroid[{extra},another-extra]"
+                
+                If you are working within the langroid dev env (which uses uv),
+                you can do:
+                uv sync --dev --extra "{extra}"
+                or if you want to include multiple extras:
+                uv sync --dev --extra "{extra}" --extra "another-extra"
+                """
+        else:
+            install_help = """
+                If you want to use it, please install it in the same
+                virtual environment as langroid.
+                """
+        msg = error + install_help
+
+        super().__init__(msg, *args)

langroid/language_models/__init__.py

@@ -0,0 +1,53 @@
+from . import utils
+from . import config
+from . import base
+from . import openai_gpt
+from . import azure_openai
+from . import prompt_formatter
+
+from .base import (
+    LLMConfig,
+    LLMMessage,
+    LLMFunctionCall,
+    LLMFunctionSpec,
+    Role,
+    LLMTokenUsage,
+    LLMResponse,
+)
+from .openai_gpt import (
+    OpenAIChatModel,
+    AnthropicModel,
+    GeminiModel,
+    OpenAICompletionModel,
+    OpenAIGPTConfig,
+    OpenAIGPT,
+)
+from .mock_lm import MockLM, MockLMConfig
+from .azure_openai import AzureConfig, AzureGPT
+
+
+__all__ = [
+    "utils",
+    "config",
+    "base",
+    "openai_gpt",
+    "azure_openai",
+    "prompt_formatter",
+    "LLMConfig",
+    "LLMMessage",
+    "LLMFunctionCall",
+    "LLMFunctionSpec",
+    "Role",
+    "LLMTokenUsage",
+    "LLMResponse",
+    "OpenAIChatModel",
+    "AnthropicModel",
+    "GeminiModel",
+    "OpenAICompletionModel",
+    "OpenAIGPTConfig",
+    "OpenAIGPT",
+    "AzureConfig",
+    "AzureGPT",
+    "MockLM",
+    "MockLMConfig",
+]

langroid/language_models/azure_openai.py

@@ -0,0 +1,153 @@
+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,
+)
+
+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): can be 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): can be set in the ``.env``
+            file as ``AZURE_OPENAI_MODEL_NAME``
+            and should be based on the model name chosen during setup.
+        model_version (str): can be set in the ``.env`` file as
+            ``AZURE_OPENAI_MODEL_VERSION`` and should be based on the model name
+            chosen during setup.
+    """
+
+    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 = ""
+    model_name: str = ""
+    model_version: str = ""  # is used to determine the cost of using the model
+    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=...`
+    class Config:
+        env_prefix = "AZURE_OPENAI_"
+
+
+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``
+    Attributes:
+        config (AzureConfig): AzureConfig object
+        api_key (str): Azure API key
+        api_base (str): Azure API base url
+        api_version (str): Azure API version
+        model_name (str): the name of gpt model in your deployment
+        model_version (str): the version of gpt model in your deployment
+    """
+
+    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.deployment_name == "":
+            raise ValueError(
+                """
+                AZURE_OPENAI_DEPLOYMENT_NAME not set in .env file,
+                please set it to your Azure openai deployment name."""
+            )
+        self.deployment_name = self.config.deployment_name
+
+        if self.config.model_name == "":
+            raise ValueError(
+                """
+                AZURE_OPENAI_MODEL_NAME not set in .env file,
+                please set it to chat model name in your deployment."""
+            )
+
+        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),
+            )
+
+        # set the chat model to be the same as the model_name
+        self.config.chat_model = self.config.model_name
+
+        self.supports_json_schema = (
+            self.config.api_version >= azureStructuredOutputAPIMin
+            and self.config.model_version in azureStructuredOutputList
+        )