langroid 0.1.139__py3-none-any.whl → 0.1.219__py3-none-any.whl

langroid/embedding_models/protoc/embeddings_pb2.py

@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: embeddings.proto
+# Protobuf Python Version: 4.25.1
+"""Generated protocol buffer code."""
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+from google.protobuf.internal import builder as _builder
+
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(
+    b'\n\x10\x65mbeddings.proto"K\n\x10\x45mbeddingRequest\x12\x12\n\nmodel_name\x18\x01 \x01(\t\x12\x12\n\nbatch_size\x18\x02 \x01(\x05\x12\x0f\n\x07strings\x18\x03 \x03(\t"%\n\x0b\x42\x61tchEmbeds\x12\x16\n\x06\x65mbeds\x18\x01 \x03(\x0b\x32\x06.Embed"\x16\n\x05\x45mbed\x12\r\n\x05\x65mbed\x18\x01 \x03(\x02\x32\x37\n\tEmbedding\x12*\n\x05\x45mbed\x12\x11.EmbeddingRequest\x1a\x0c.BatchEmbeds"\x00\x62\x06proto3'
+)
+
+_globals = globals()
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals)
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "embeddings_pb2", _globals)
+if _descriptor._USE_C_DESCRIPTORS == False:
+    DESCRIPTOR._options = None
+    _globals["_EMBEDDINGREQUEST"]._serialized_start = 20
+    _globals["_EMBEDDINGREQUEST"]._serialized_end = 95
+    _globals["_BATCHEMBEDS"]._serialized_start = 97
+    _globals["_BATCHEMBEDS"]._serialized_end = 134
+    _globals["_EMBED"]._serialized_start = 136
+    _globals["_EMBED"]._serialized_end = 158
+    _globals["_EMBEDDING"]._serialized_start = 160
+    _globals["_EMBEDDING"]._serialized_end = 215
+# @@protoc_insertion_point(module_scope)

langroid/embedding_models/protoc/embeddings_pb2.pyi

@@ -0,0 +1,50 @@
+from typing import (
+    ClassVar as _ClassVar,
+)
+from typing import (
+    Iterable as _Iterable,
+)
+from typing import (
+    Mapping as _Mapping,
+)
+from typing import (
+    Optional as _Optional,
+)
+from typing import (
+    Union as _Union,
+)
+
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from google.protobuf.internal import containers as _containers
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class EmbeddingRequest(_message.Message):
+    __slots__ = ("model_name", "batch_size", "strings")
+    MODEL_NAME_FIELD_NUMBER: _ClassVar[int]
+    BATCH_SIZE_FIELD_NUMBER: _ClassVar[int]
+    STRINGS_FIELD_NUMBER: _ClassVar[int]
+    model_name: str
+    batch_size: int
+    strings: _containers.RepeatedScalarFieldContainer[str]
+    def __init__(
+        self,
+        model_name: _Optional[str] = ...,
+        batch_size: _Optional[int] = ...,
+        strings: _Optional[_Iterable[str]] = ...,
+    ) -> None: ...
+
+class BatchEmbeds(_message.Message):
+    __slots__ = ("embeds",)
+    EMBEDS_FIELD_NUMBER: _ClassVar[int]
+    embeds: _containers.RepeatedCompositeFieldContainer[Embed]
+    def __init__(
+        self, embeds: _Optional[_Iterable[_Union[Embed, _Mapping]]] = ...
+    ) -> None: ...
+
+class Embed(_message.Message):
+    __slots__ = ("embed",)
+    EMBED_FIELD_NUMBER: _ClassVar[int]
+    embed: _containers.RepeatedScalarFieldContainer[float]
+    def __init__(self, embed: _Optional[_Iterable[float]] = ...) -> None: ...

langroid/embedding_models/protoc/embeddings_pb2_grpc.py

@@ -0,0 +1,79 @@
+# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
+"""Client and server classes corresponding to protobuf-defined services."""
+import grpc
+
+import langroid.embedding_models.protoc.embeddings_pb2 as embeddings__pb2
+
+
+class EmbeddingStub(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.Embed = channel.unary_unary(
+            "/Embedding/Embed",
+            request_serializer=embeddings__pb2.EmbeddingRequest.SerializeToString,
+            response_deserializer=embeddings__pb2.BatchEmbeds.FromString,
+        )
+
+
+class EmbeddingServicer(object):
+    """Missing associated documentation comment in .proto file."""
+
+    def Embed(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details("Method not implemented!")
+        raise NotImplementedError("Method not implemented!")
+
+
+def add_EmbeddingServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+        "Embed": grpc.unary_unary_rpc_method_handler(
+            servicer.Embed,
+            request_deserializer=embeddings__pb2.EmbeddingRequest.FromString,
+            response_serializer=embeddings__pb2.BatchEmbeds.SerializeToString,
+        ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+        "Embedding", rpc_method_handlers
+    )
+    server.add_generic_rpc_handlers((generic_handler,))
+
+
+# This class is part of an EXPERIMENTAL API.
+class Embedding(object):
+    """Missing associated documentation comment in .proto file."""
+
+    @staticmethod
+    def Embed(
+        request,
+        target,
+        options=(),
+        channel_credentials=None,
+        call_credentials=None,
+        insecure=False,
+        compression=None,
+        wait_for_ready=None,
+        timeout=None,
+        metadata=None,
+    ):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            "/Embedding/Embed",
+            embeddings__pb2.EmbeddingRequest.SerializeToString,
+            embeddings__pb2.BatchEmbeds.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+        )

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/language_models/__init__.py

@@ -21,3 +21,25 @@ from .openai_gpt import (
     OpenAIGPT,
 )
 from .azure_openai import AzureConfig, AzureGPT
+
+__all__ = [
+    "utils",
+    "config",
+    "base",
+    "openai_gpt",
+    "azure_openai",
+    "prompt_formatter",
+    "LLMConfig",
+    "LLMMessage",
+    "LLMFunctionCall",
+    "LLMFunctionSpec",
+    "Role",
+    "LLMTokenUsage",
+    "LLMResponse",
+    "OpenAIChatModel",
+    "OpenAICompletionModel",
+    "OpenAIGPTConfig",
+    "OpenAIGPT",
+    "AzureConfig",
+    "AzureGPT",
+]

langroid/language_models/azure_openai.py

@@ -22,6 +22,9 @@ class AzureConfig(OpenAIGPTConfig):
             chose for your deployment when you deployed a model.
         model_name (str): can be set in the ``.env`` file as ``AZURE_GPT_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
@@ -29,6 +32,7 @@ class AzureConfig(OpenAIGPTConfig):
     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 = ""
 
     # all of the vars above can be set via env vars,
@@ -50,6 +54,7 @@ class AzureGPT(OpenAIGPT):
         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):
@@ -89,10 +94,7 @@ class AzureGPT(OpenAIGPT):
         # set the chat model to be the same as the model_name
         # This corresponds to the gpt model you chose for your deployment
         # when you deployed a model
-        if "35-turbo" in self.config.model_name:
-            self.config.chat_model = OpenAIChatModel.GPT3_5_TURBO
-        else:
-            self.config.chat_model = OpenAIChatModel.GPT4
+        self.set_chat_model()
 
         self.client = AzureOpenAI(
             api_key=self.config.api_key,
@@ -107,3 +109,44 @@ class AzureGPT(OpenAIGPT):
             azure_deployment=self.config.deployment_name,
             timeout=Timeout(self.config.timeout),
         )
+
+    def set_chat_model(self) -> None:
+        """
+        Sets the chat model configuration based on the model name specified in the
+        ``.env``. This function checks the `model_name` in the configuration and sets
+        the appropriate chat model in the `config.chat_model`. It supports handling for
+        '35-turbo' and 'gpt-4' models. For 'gpt-4', it further delegates the handling
+        to `handle_gpt4_model` method. If the model name does not match any predefined
+        models, it defaults to `OpenAIChatModel.GPT4`.
+        """
+        MODEL_35_TURBO = "35-turbo"
+        MODEL_GPT4 = "gpt-4"
+
+        if self.config.model_name == MODEL_35_TURBO:
+            self.config.chat_model = OpenAIChatModel.GPT3_5_TURBO
+        elif self.config.model_name == MODEL_GPT4:
+            self.handle_gpt4_model()
+        else:
+            self.config.chat_model = OpenAIChatModel.GPT4
+
+    def handle_gpt4_model(self) -> None:
+        """
+        Handles the setting of the GPT-4 model in the configuration.
+        This function checks the `model_version` in the configuration.
+        If the version is not set, it raises a ValueError indicating that the model
+        version needs to be specified in the ``.env`` file.
+        It sets `OpenAIChatModel.GPT4_TURBO` if the version is
+        '1106-Preview', otherwise, it defaults to setting `OpenAIChatModel.GPT4`.
+        """
+        VERSION_1106_PREVIEW = "1106-Preview"
+
+        if self.config.model_version == "":
+            raise ValueError(
+                "AZURE_OPENAI_MODEL_VERSION not set in .env file. "
+                "Please set it to the chat model version used in your deployment."
+            )
+
+        if self.config.model_version == VERSION_1106_PREVIEW:
+            self.config.chat_model = OpenAIChatModel.GPT4_TURBO
+        else:
+            self.config.chat_model = OpenAIChatModel.GPT4

langroid/language_models/base.py

@@ -3,18 +3,18 @@ import asyncio
 import json
 import logging
 from abc import ABC, abstractmethod
+from datetime import datetime
 from enum import Enum
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 import aiohttp
-from pydantic import BaseModel, BaseSettings
+from pydantic import BaseModel, BaseSettings, Field
 
 from langroid.cachedb.momento_cachedb import MomentoCacheConfig
 from langroid.cachedb.redis_cachedb import RedisCacheConfig
-from langroid.language_models.config import Llama2FormatterConfig, PromptFormatterConfig
 from langroid.mytypes import Document
 from langroid.parsing.agent_chats import parse_message
-from langroid.parsing.json import top_level_json_field
+from langroid.parsing.parse_json import top_level_json_field
 from langroid.prompts.dialog import collate_chat_history
 from langroid.prompts.templates import (
     EXTRACTION_PROMPT_GPT4,
@@ -26,16 +26,21 @@ from langroid.utils.output.printing import show_if_debug
 logger = logging.getLogger(__name__)
 
 
+def noop_fn(*args: List[Any], **kwargs: Dict[str, Any]) -> None:
+    pass
+
+
 class LLMConfig(BaseSettings):
     type: str = "openai"
+    streamer: Optional[Callable[[Any], None]] = noop_fn
     api_base: str | None = None
-    formatter: None | PromptFormatterConfig = Llama2FormatterConfig()
+    formatter: None | str = None
     timeout: int = 20  # timeout for API requests
     chat_model: str = ""
     completion_model: str = ""
     temperature: float = 0.0
-    chat_context_length: int = 1024
-    completion_context_length: int = 1024
+    chat_context_length: int = 8000
+    completion_context_length: int = 8000
     max_output_tokens: int = 1024  # generate at most this many tokens
     # if input length + max_output_tokens > context length of model,
     # we will try shortening requested output
@@ -72,8 +77,11 @@ class LLMFunctionCall(BaseModel):
         fun_args_str = message["arguments"]
         # sometimes may be malformed with invalid indents,
         # so we try to be safe by removing newlines.
-        fun_args_str = fun_args_str.replace("\n", "").strip()
-        fun_args = ast.literal_eval(fun_args_str)
+        if fun_args_str is not None:
+            fun_args_str = fun_args_str.replace("\n", "").strip()
+            fun_args = ast.literal_eval(fun_args_str)
+        else:
+            fun_args = None
         fun_call.arguments = fun_args
 
         return fun_call
@@ -135,6 +143,7 @@ class LLMMessage(BaseModel):
     tool_id: str = ""  # used by OpenAIAssistant
     content: str
     function_call: Optional[LLMFunctionCall] = None
+    timestamp: datetime = Field(default_factory=datetime.utcnow)
 
     def api_dict(self) -> Dict[str, Any]:
         """
@@ -157,6 +166,7 @@ class LLMMessage(BaseModel):
                     dict_no_none["function_call"]["arguments"]
                 )
         dict_no_none.pop("tool_id", None)
+        dict_no_none.pop("timestamp", None)
         return dict_no_none
 
     def __str__(self) -> str:
@@ -179,6 +189,12 @@ class LLMResponse(BaseModel):
     usage: Optional[LLMTokenUsage]
     cached: bool = False
 
+    def __str__(self) -> str:
+        if self.function_call is not None:
+            return str(self.function_call)
+        else:
+            return self.message
+
     def to_LLMMessage(self) -> LLMMessage:
         content = self.message
         role = Role.ASSISTANT if self.function_call is None else Role.FUNCTION
@@ -245,7 +261,7 @@ class LanguageModel(ABC):
     # usage cost by model, accumulates here
     usage_cost_dict: Dict[str, LLMTokenUsage] = {}
 
-    def __init__(self, config: LLMConfig):
+    def __init__(self, config: LLMConfig = LLMConfig()):
         self.config = config
 
     @staticmethod

langroid/language_models/config.py

@@ -11,3 +11,8 @@ class PromptFormatterConfig(BaseSettings):
 
 class Llama2FormatterConfig(PromptFormatterConfig):
     use_bos_eos: bool = False
+
+
+class HFPromptFormatterConfig(PromptFormatterConfig):
+    type: str = "hf"
+    model_name: str