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

langroid/cachedb/redis_cachedb.py

@@ -1,7 +1,8 @@
 import json
 import logging
 import os
-from typing import Any, Dict, Optional
+from contextlib import AbstractContextManager, contextmanager
+from typing import Any, Dict, List, TypeVar
 
 import fakeredis
 import redis
@@ -10,6 +11,7 @@ from pydantic import BaseModel
 
 from langroid.cachedb.base import CacheDB
 
+T = TypeVar("T", bound="RedisCache")
 logger = logging.getLogger(__name__)
 
 
@@ -33,7 +35,7 @@ class RedisCache(CacheDB):
         load_dotenv()
 
         if self.config.fake:
-            self.client = fakeredis.FakeStrictRedis()  # type: ignore
+            self.pool = fakeredis.FakeStrictRedis()  # type: ignore
         else:
             redis_password = os.getenv("REDIS_PASSWORD")
             redis_host = os.getenv("REDIS_HOST")
@@ -43,21 +45,46 @@ class RedisCache(CacheDB):
                     """REDIS_PASSWORD, REDIS_HOST, REDIS_PORT not set in .env file,
                     using fake redis client"""
                 )
-                self.client = fakeredis.FakeStrictRedis()  # type: ignore
+                self.pool = fakeredis.FakeStrictRedis()  # type: ignore
             else:
-                self.client = redis.Redis(  # type: ignore
+                self.pool = redis.ConnectionPool(  # type: ignore
                     host=redis_host,
                     port=redis_port,
                     password=redis_password,
+                    max_connections=50,
+                    socket_timeout=5,
+                    socket_keepalive=True,
+                    retry_on_timeout=True,
+                    health_check_interval=30,
                 )
 
+    @contextmanager  # type: ignore
+    def redis_client(self) -> AbstractContextManager[T]:  # type: ignore
+        """Cleanly open and close a redis client, avoids max clients exceeded error"""
+        if isinstance(self.pool, fakeredis.FakeStrictRedis):
+            yield self.pool
+        else:
+            client: T = redis.Redis(connection_pool=self.pool)
+            try:
+                yield client
+            finally:
+                client.close()
+
+    def close_all_connections(self) -> None:
+        with self.redis_client() as client:  # type: ignore
+            clients = client.client_list()
+            for c in clients:
+                client.client_kill(c["addr"])
+
     def clear(self) -> None:
         """Clear keys from current db."""
-        self.client.flushdb()
+        with self.redis_client() as client:  # type: ignore
+            client.flushdb()
 
     def clear_all(self) -> None:
         """Clear all keys from all dbs."""
-        self.client.flushall()
+        with self.redis_client() as client:  # type: ignore
+            client.flushall()
 
     def store(self, key: str, value: Any) -> None:
         """
@@ -67,9 +94,14 @@ class RedisCache(CacheDB):
             key (str): The key under which to store the value.
             value (Any): The value to store.
         """
-        self.client.set(key, json.dumps(value))
-
-    def retrieve(self, key: str) -> Optional[Dict[str, Any]]:
+        with self.redis_client() as client:  # type: ignore
+            try:
+                client.set(key, json.dumps(value))
+            except redis.exceptions.ConnectionError:
+                logger.warning("Redis connection error, not storing key/value")
+                return None
+
+    def retrieve(self, key: str) -> Dict[str, Any] | str | None:
         """
         Retrieve the value associated with a key.
 
@@ -79,5 +111,40 @@ class RedisCache(CacheDB):
         Returns:
             dict: The value associated with the key.
         """
-        value = self.client.get(key)
-        return json.loads(value) if value else None
+        with self.redis_client() as client:  # type: ignore
+            try:
+                value = client.get(key)
+            except redis.exceptions.ConnectionError:
+                logger.warning("Redis connection error, returning None")
+                return None
+            return json.loads(value) if value else None
+
+    def delete_keys(self, keys: List[str]) -> None:
+        """
+        Delete the keys from the cache.
+
+        Args:
+            keys (List[str]): The keys to delete.
+        """
+        with self.redis_client() as client:  # type: ignore
+            try:
+                client.delete(*keys)
+            except redis.exceptions.ConnectionError:
+                logger.warning("Redis connection error, not deleting keys")
+                return None
+
+    def delete_keys_pattern(self, pattern: str) -> None:
+        """
+        Delete the keys matching the pattern from the cache.
+
+        Args:
+            prefix (str): The pattern to match.
+        """
+        with self.redis_client() as client:  # type: ignore
+            try:
+                keys = client.keys(pattern)
+                if len(keys) > 0:
+                    client.delete(*keys)
+            except redis.exceptions.ConnectionError:
+                logger.warning("Redis connection error, not deleting keys")
+                return None

langroid/embedding_models/__init__.py

@@ -0,0 +1,34 @@
+from . import base
+from . import models
+from . import remote_embeds
+
+from .base import (
+    EmbeddingModel,
+    EmbeddingModelsConfig,
+)
+from .models import (
+    OpenAIEmbeddings,
+    OpenAIEmbeddingsConfig,
+    SentenceTransformerEmbeddingsConfig,
+    SentenceTransformerEmbeddings,
+    embedding_model,
+)
+from .remote_embeds import (
+    RemoteEmbeddingsConfig,
+    RemoteEmbeddings,
+)
+
+__all__ = [
+    "base",
+    "models",
+    "remote_embeds",
+    "EmbeddingModel",
+    "EmbeddingModelsConfig",
+    "OpenAIEmbeddings",
+    "OpenAIEmbeddingsConfig",
+    "SentenceTransformerEmbeddingsConfig",
+    "SentenceTransformerEmbeddings",
+    "embedding_model",
+    "RemoteEmbeddingsConfig",
+    "RemoteEmbeddings",
+]

langroid/embedding_models/base.py

@@ -1,15 +1,19 @@
 import logging
 from abc import ABC, abstractmethod
 
-from chromadb.api.types import EmbeddingFunction
+import numpy as np
 from pydantic import BaseSettings
 
+from langroid.mytypes import EmbeddingFunction
+
 logging.getLogger("openai").setLevel(logging.ERROR)
 
 
 class EmbeddingModelsConfig(BaseSettings):
     model_type: str = "openai"
     dims: int = 0
+    context_length: int = 512
+    batch_size: int = 512
 
 
 class EmbeddingModel(ABC):
@@ -25,8 +29,14 @@ class EmbeddingModel(ABC):
             SentenceTransformerEmbeddings,
             SentenceTransformerEmbeddingsConfig,
         )
+        from langroid.embedding_models.remote_embeds import (
+            RemoteEmbeddings,
+            RemoteEmbeddingsConfig,
+        )
 
-        if isinstance(config, OpenAIEmbeddingsConfig):
+        if isinstance(config, RemoteEmbeddingsConfig):
+            return RemoteEmbeddings(config)
+        elif isinstance(config, OpenAIEmbeddingsConfig):
             return OpenAIEmbeddings(config)
         elif isinstance(config, SentenceTransformerEmbeddingsConfig):
             return SentenceTransformerEmbeddings(config)
@@ -41,3 +51,12 @@ class EmbeddingModel(ABC):
     @abstractmethod
     def embedding_dims(self) -> int:
         pass
+
+    def similarity(self, text1: str, text2: str) -> float:
+        """Compute cosine similarity between two texts."""
+        [emb1, emb2] = self.embedding_fn()([text1, text2])
+        return float(
+            np.array(emb1)
+            @ np.array(emb2)
+            / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
+        )

langroid/embedding_models/models.py

@@ -1,32 +1,98 @@
+import atexit
 import os
-from typing import Callable, List
+from typing import Callable, List, Optional
 
-import openai
+import tiktoken
 from dotenv import load_dotenv
+from openai import OpenAI
 
 from langroid.embedding_models.base import EmbeddingModel, EmbeddingModelsConfig
-from langroid.language_models.utils import retry_with_exponential_backoff
 from langroid.mytypes import Embeddings
+from langroid.parsing.utils import batched
 
 
 class OpenAIEmbeddingsConfig(EmbeddingModelsConfig):
     model_type: str = "openai"
     model_name: str = "text-embedding-ada-002"
     api_key: str = ""
+    api_base: Optional[str] = None
+    organization: str = ""
     dims: int = 1536
+    context_length: int = 8192
 
 
 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 EmbeddingFunctionCallable:
+    """
+    A callable class designed to generate embeddings for a list of texts using
+    the OpenAI API, with automatic retries on failure.
+
+    Attributes:
+        model (OpenAIEmbeddings): An instance of OpenAIEmbeddings 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, model: "OpenAIEmbeddings", batch_size: int = 512):
+        """
+        Initialize the EmbeddingFunctionCallable with a specific model.
+
+        Args:
+            model (OpenAIEmbeddings): An instance of OpenAIEmbeddings to use for
+            generating embeddings.
+            batch_size (int): Batch size
+        """
+        self.model = 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.
+        """
+        tokenized_texts = self.model.truncate_texts(input)
+        embeds = []
+        for batch in batched(tokenized_texts, self.batch_size):
+            result = self.model.client.embeddings.create(
+                input=batch, model=self.model.config.model_name
+            )
+            batch_embeds = [d.embedding for d in result.data]
+            embeds.extend(batch_embeds)
+        return embeds
 
 
 class OpenAIEmbeddings(EmbeddingModel):
-    def __init__(self, config: OpenAIEmbeddingsConfig):
+    def __init__(self, config: OpenAIEmbeddingsConfig = OpenAIEmbeddingsConfig()):
         super().__init__()
         self.config = config
         load_dotenv()
         self.config.api_key = os.getenv("OPENAI_API_KEY", "")
+        self.config.organization = os.getenv("OPENAI_ORGANIZATION", "")
         if self.config.api_key == "":
             raise ValueError(
                 """OPENAI_API_KEY env variable must be set to use 
@@ -34,28 +100,38 @@ class OpenAIEmbeddings(EmbeddingModel):
                 in your .env file.
                 """
             )
-        openai.api_key = self.config.api_key
+        self.client = OpenAI(base_url=self.config.api_base, api_key=self.config.api_key)
+        self.tokenizer = tiktoken.encoding_for_model(self.config.model_name)
+
+    def truncate_texts(self, texts: 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]:
-        @retry_with_exponential_backoff
-        def fn(texts: List[str]) -> Embeddings:
-            result = openai.Embedding.create(  # type: ignore
-                input=texts, model=self.config.model_name
-            )
-            return [d["embedding"] for d in result["data"]]
-
-        return fn
+        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: SentenceTransformerEmbeddingsConfig):
+    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(
                 """
@@ -67,13 +143,39 @@ class SentenceTransformerEmbeddings(EmbeddingModel):
 
         super().__init__()
         self.config = config
-        self.model = SentenceTransformer(self.config.model_name)
+
+        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]:
         def fn(texts: List[str]) -> Embeddings:
-            return self.model.encode(  # type: ignore
-                texts, convert_to_numpy=True
-            ).tolist()
+            if self.config.data_parallel:
+                embeds: Embeddings = self.model.encode_multi_process(
+                    texts,
+                    self.pool,
+                    batch_size=self.config.batch_size,
+                ).tolist()
+            else:
+                embeds = []
+                for batch in batched(texts, self.config.batch_size):
+                    batch_embeds = self.model.encode(
+                        batch, convert_to_numpy=True
+                    ).tolist()  # type: ignore
+                    embeds.extend(batch_embeds)
+
+            return embeds
 
         return fn
 

langroid/embedding_models/protoc/embeddings.proto

@@ -0,0 +1,19 @@
+syntax = "proto3";
+
+service Embedding {
+    rpc Embed (EmbeddingRequest) returns (BatchEmbeds) {};
+}
+
+message EmbeddingRequest {
+    string model_name = 1;
+    int32 batch_size = 2;
+    repeated string strings = 3;
+}
+
+message BatchEmbeds {
+    repeated Embed embeds = 1;
+}
+
+message Embed {
+    repeated float embed = 1;
+}

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,
+        )