langroid 0.37.7__py3-none-any.whl → 0.38.0__py3-none-any.whl

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/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/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.dist-info → langroid-0.38.0.dist-info}/METADATA

@@ -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.dist-info → langroid-0.38.0.dist-info}/RECORD

@@ -55,9 +55,9 @@ langroid/cachedb/__init__.py,sha256=icAT2s7Vhf-ZGUeqpDQGNU6ob6o0aFEyjwcxxUGRFjg,
 langroid/cachedb/base.py,sha256=ztVjB1DtN6pLCujCWnR6xruHxwVj3XkYniRTYAKKqk0,1354
 langroid/cachedb/momento_cachedb.py,sha256=YEOJ62hEcV6iIeMr5aGgRYgWQqFYaej9gEDEcY0sm7M,3172
 langroid/cachedb/redis_cachedb.py,sha256=7kgnbf4b5CKsCrlL97mHWKvdvlLt8zgn7lc528jEpiE,5141
-langroid/embedding_models/__init__.py,sha256=XhVIMQJbQRpImcnhA9sJR7h6r7QgPo1SKDCvwEUD9j4,851
-langroid/embedding_models/base.py,sha256=DUhvzALoW2UMbtmLxP4eJTfPii99WjUNX7bwFpj_K-0,2395
-langroid/embedding_models/models.py,sha256=W8BG_OFgpA7lkEU0Fk4Qt7AzxuN0JjGf7Y-29EXZQ4E,16776
+langroid/embedding_models/__init__.py,sha256=KyYxR3jDFUCfYjSuCL86qjAmrq6mXXjOT4lFNOKVj6Y,955
+langroid/embedding_models/base.py,sha256=Ml7oA6PzQm0wZmIYn3fhF7dvZCi-amviWUwOeBegH3A,2562
+langroid/embedding_models/models.py,sha256=fKQBiaaG7uYoeELDbAiNxwLdn-CWN8dyiVEZcdk_bjI,18959
 langroid/embedding_models/remote_embeds.py,sha256=6_kjXByVbqhY9cGwl9R83ZcYC2km-nGieNNAo1McHaY,5151
 langroid/embedding_models/protoc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/embedding_models/protoc/embeddings.proto,sha256=_O-SgFpTaylQeOTgSpxhEJ7CUw7PeCQQJLaPqpPYKJg,321
@@ -123,7 +123,7 @@ langroid/vector_store/meilisearch.py,sha256=6frB7GFWeWmeKzRfLZIvzRjllniZ1cYj3Hmh
 langroid/vector_store/momento.py,sha256=xOaU7Hlyyn_5ihb0ARS5JHtmrKrTCt2IdRA-ioMM5ek,10307
 langroid/vector_store/qdrantdb.py,sha256=v7TAsIoj_vxeKDYS9tpwJLBZA8fuTweTYxHo0X_uawM,17949
 langroid/vector_store/weaviatedb.py,sha256=FOzgvqLqvdN5jJebVtJ-8tu2CeBzBfSP3ih4_ODEOOw,10605
-langroid-0.37.7.dist-info/METADATA,sha256=iaVoVKbhBTjGSt00U8RQAKJ1LRfLaUcGDQrrCvkIvAA,60524
-langroid-0.37.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langroid-0.37.7.dist-info/licenses/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
-langroid-0.37.7.dist-info/RECORD,,
+langroid-0.38.0.dist-info/METADATA,sha256=PMVjo4UfoNgUyP9RWp2AAEIWA9amAnTHUvzrPPjt_ms,60634
+langroid-0.38.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langroid-0.38.0.dist-info/licenses/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
+langroid-0.38.0.dist-info/RECORD,,