langroid 0.37.7__tar.gz → 0.38.0__tar.gz

{langroid-0.37.7 → langroid-0.38.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langroid
-Version: 0.37.7
+Version: 0.38.0
 Summary: Harness LLMs with Multi-Agent Programming
 Author-email: Prasad Chalasani <pchalasani@gmail.com>
 License: MIT
@@ -106,6 +106,8 @@ Provides-Extra: docx
 Requires-Dist: python-docx<2.0.0,>=1.1.0; extra == 'docx'
 Provides-Extra: fastembed
 Requires-Dist: fastembed<0.4.0,>=0.3.1; extra == 'fastembed'
+Provides-Extra: google-generativeai
+Requires-Dist: google-generativeai>=0.8.4; extra == 'google-generativeai'
 Provides-Extra: hf-embeddings
 Requires-Dist: sentence-transformers<3.0.0,>=2.2.2; extra == 'hf-embeddings'
 Requires-Dist: torch<3.0.0,>=2.0.0; extra == 'hf-embeddings'

{langroid-0.37.7 → langroid-0.38.0}/langroid/embedding_models/__init__.py

@@ -13,6 +13,8 @@ from .models import (
     SentenceTransformerEmbeddingsConfig,
     LlamaCppServerEmbeddings,
     LlamaCppServerEmbeddingsConfig,
+    GeminiEmbeddings,
+    GeminiEmbeddingsConfig,
     embedding_model,
 )
 from .remote_embeds import (
@@ -33,6 +35,8 @@ __all__ = [
     "SentenceTransformerEmbeddingsConfig",
     "LlamaCppServerEmbeddings",
     "LlamaCppServerEmbeddingsConfig",
+    "GeminiEmbeddings",
+    "GeminiEmbeddingsConfig",
     "embedding_model",
     "RemoteEmbeddingsConfig",
     "RemoteEmbeddings",

{langroid-0.37.7 → langroid-0.38.0}/langroid/embedding_models/base.py

@@ -28,6 +28,8 @@ class EmbeddingModel(ABC):
             AzureOpenAIEmbeddingsConfig,
             FastEmbedEmbeddings,
             FastEmbedEmbeddingsConfig,
+            GeminiEmbeddings,
+            GeminiEmbeddingsConfig,
             LlamaCppServerEmbeddings,
             LlamaCppServerEmbeddingsConfig,
             OpenAIEmbeddings,
@@ -52,6 +54,8 @@ class EmbeddingModel(ABC):
             return FastEmbedEmbeddings(config)
         elif isinstance(config, LlamaCppServerEmbeddingsConfig):
             return LlamaCppServerEmbeddings(config)
+        elif isinstance(config, GeminiEmbeddingsConfig):
+            return GeminiEmbeddings(config)
         else:
             raise ValueError(f"Unknown embedding config: {config.__repr_name__}")
 

{langroid-0.37.7 → langroid-0.38.0}/langroid/embedding_models/models.py

@@ -77,6 +77,14 @@ class LlamaCppServerEmbeddingsConfig(EmbeddingModelsConfig):
     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
@@ -160,6 +168,8 @@ class EmbeddingFunctionCallable:
                         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
 
 
@@ -437,6 +447,54 @@ class LlamaCppServerEmbeddings(EmbeddingModel):
         return self.config.dims
 
 
+class GeminiEmbeddings(EmbeddingModel):
+    def __init__(self, config: GeminiEmbeddingsConfig = GeminiEmbeddingsConfig()):
+        try:
+            import google.generativeai as genai
+        except ImportError as e:
+            raise LangroidImportError(extra="google-generativeai", 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.
+                """
+            )
+        genai.configure(api_key=self.config.api_key)  # type: ignore[attr-defined]
+        self.client = genai
+
+    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
+        return EmbeddingFunctionCallable(self, self.config.batch_size)
+
+    def generate_embeddings(self, texts: List[str]) -> Embeddings:
+        all_embeddings = []  # More precise type hint
+        for batch in batched(texts, self.config.batch_size):
+            result = self.client.embed_content(  # type: ignore[attr-defined]
+                model=self.config.model_name,
+                content=batch,
+                task_type="RETRIEVAL_DOCUMENT",
+            )
+
+            embeddings = result["embedding"]
+            if not isinstance(embeddings, list):
+                raise ValueError("Unexpected format for embeddings: not a list")
+
+            if embeddings and isinstance(embeddings[0], list):
+                all_embeddings.extend(embeddings)
+            else:
+                all_embeddings.append(embeddings)
+
+        return all_embeddings
+
+    @property
+    def embedding_dims(self) -> int:
+        return self.config.dims
+
+
 def embedding_model(embedding_fn_type: str = "openai") -> EmbeddingModel:
     """
     Args:
@@ -457,5 +515,7 @@ def embedding_model(embedding_fn_type: str = "openai") -> EmbeddingModel:
         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-0.37.7 → langroid-0.38.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "langroid"
-version = "0.37.7"
+version = "0.38.0"
 authors = [
     {name = "Prasad Chalasani", email = "pchalasani@gmail.com"},
 ]
@@ -220,6 +220,9 @@ momento = [
 fastembed = [
     "fastembed<0.4.0,>=0.3.1",
 ]
+google-generativeai = [
+    "google-generativeai>=0.8.4",
+]
 
 
 [dependency-groups]