langroid 0.2.12__py3-none-any.whl → 0.3.0__py3-none-any.whl

langroid/embedding_models/base.py

@@ -24,6 +24,8 @@ class EmbeddingModel(ABC):
     @classmethod
     def create(cls, config: EmbeddingModelsConfig) -> "EmbeddingModel":
         from langroid.embedding_models.models import (
+            FastEmbedEmbeddings,
+            FastEmbedEmbeddingsConfig,
             OpenAIEmbeddings,
             OpenAIEmbeddingsConfig,
             SentenceTransformerEmbeddings,
@@ -40,6 +42,8 @@ class EmbeddingModel(ABC):
             return OpenAIEmbeddings(config)
         elif isinstance(config, SentenceTransformerEmbeddingsConfig):
             return SentenceTransformerEmbeddings(config)
+        elif isinstance(config, FastEmbedEmbeddingsConfig):
+            return FastEmbedEmbeddings(config)
         else:
             raise ValueError(f"Unknown embedding config: {config.__repr_name__}")
 

langroid/embedding_models/models.py

@@ -1,12 +1,14 @@
 import atexit
 import os
-from typing import Callable, List, Optional
+from functools import cached_property
+from typing import Any, Callable, Dict, List, Optional
 
 import tiktoken
 from dotenv import load_dotenv
 from openai import OpenAI
 
 from langroid.embedding_models.base import EmbeddingModel, EmbeddingModelsConfig
+from langroid.exceptions import LangroidImportError
 from langroid.mytypes import Embeddings
 from langroid.parsing.utils import batched
 
@@ -32,6 +34,20 @@ class SentenceTransformerEmbeddingsConfig(EmbeddingModelsConfig):
     devices: Optional[list[str]] = None
 
 
+class FastEmbedEmbeddingsConfig(EmbeddingModelsConfig):
+    """Config for qdrant/fastembed embeddings,
+    see here: https://github.com/qdrant/fastembed
+    """
+
+    model_type: str = "fastembed"
+    model_name: str = "BAAI/bge-small-en-v1.5"
+    batch_size: int = 256
+    cache_dir: Optional[str] = None
+    threads: Optional[int] = None
+    parallel: Optional[int] = None
+    additional_kwargs: Dict[str, Any] = {}
+
+
 class EmbeddingFunctionCallable:
     """
     A callable class designed to generate embeddings for a list of texts using
@@ -189,6 +205,41 @@ class SentenceTransformerEmbeddings(EmbeddingModel):
         return dims  # type: ignore
 
 
+class FastEmbedEmbeddings(EmbeddingModel):
+    def __init__(self, config: FastEmbedEmbeddingsConfig = FastEmbedEmbeddingsConfig()):
+        try:
+            from fastembed import TextEmbedding
+        except ImportError:
+            raise LangroidImportError("fastembed", extra="fastembed")
+
+        super().__init__()
+        self.config = config
+        self._batch_size = config.batch_size
+        self._parallel = config.parallel
+
+        self._model = TextEmbedding(
+            model_name=self.config.model_name,
+            cache_dir=self.config.cache_dir,
+            threads=self.config.threads,
+            **self.config.additional_kwargs,
+        )
+
+    def embedding_fn(self) -> Callable[[List[str]], Embeddings]:
+        def fn(texts: List[str]) -> Embeddings:
+            embeddings = self._model.embed(
+                texts, batch_size=self._batch_size, parallel=self._parallel
+            )
+
+            return [embedding.tolist() for embedding in embeddings]
+
+        return fn
+
+    @cached_property
+    def embedding_dims(self) -> int:
+        embed_func = self.embedding_fn()
+        return len(embed_func(["text"])[0])
+
+
 def embedding_model(embedding_fn_type: str = "openai") -> EmbeddingModel:
     """
     Args:
@@ -198,5 +249,7 @@ def embedding_model(embedding_fn_type: str = "openai") -> EmbeddingModel:
     """
     if embedding_fn_type == "openai":
         return OpenAIEmbeddings  # type: ignore
+    elif embedding_fn_type == "fastembed":
+        return FastEmbedEmbeddings  # type: ignore
     else:  # default sentence transformer
         return SentenceTransformerEmbeddings  # type: ignore

{langroid-0.2.12.dist-info → langroid-0.3.0.dist-info}/METADATA

@@ -1,11 +1,11 @@
 Metadata-Version: 2.1
 Name: langroid
-Version: 0.2.12
+Version: 0.3.0
 Summary: Harness LLMs with Multi-Agent Programming
 License: MIT
 Author: Prasad Chalasani
 Author-email: pchalasani@gmail.com
-Requires-Python: >=3.10,<4.0
+Requires-Python: >=3.10,<3.13
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
@@ -17,6 +17,7 @@ Provides-Extra: chromadb
 Provides-Extra: db
 Provides-Extra: doc-chat
 Provides-Extra: docx
+Provides-Extra: fastembed
 Provides-Extra: hf-embeddings
 Provides-Extra: hf-transformers
 Provides-Extra: lancedb
@@ -44,6 +45,7 @@ Requires-Dist: docstring-parser (>=0.15,<0.16)
 Requires-Dist: duckduckgo-search (>=6.0.0,<7.0.0)
 Requires-Dist: faker (>=18.9.0,<19.0.0)
 Requires-Dist: fakeredis (>=2.12.1,<3.0.0)
+Requires-Dist: fastembed (>=0.3.1,<0.4.0) ; extra == "all" or extra == "fastembed"
 Requires-Dist: fire (>=0.5.0,<0.6.0)
 Requires-Dist: google-api-python-client (>=2.95.0,<3.0.0)
 Requires-Dist: google-generativeai (>=0.5.2,<0.6.0)

{langroid-0.2.12.dist-info → langroid-0.3.0.dist-info}/RECORD

@@ -52,9 +52,9 @@ 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=lsu8qxCjfGujXGueJWU-VI3LMZYGjLSYgqUKDd4F3Qo,715
-langroid/embedding_models/base.py,sha256=MSjaTkFcfoMGY6SHPOqAsbZbKctj8-1N6zgaFYmOFTg,1830
+langroid/embedding_models/base.py,sha256=3dK0nW3XNjK3Vyh2kxhIffzDuUYumVVkCIimB3UPHeU,2009
 langroid/embedding_models/clustering.py,sha256=tZWElUqXl9Etqla0FAa7og96iDKgjqWjucZR_Egtp-A,6684
-langroid/embedding_models/models.py,sha256=-xeN0irBPc1tUgRFHGM1ki4NwOIHr6F3SKuEjD5nTOg,7144
+langroid/embedding_models/models.py,sha256=NQ1Cfw5MmeR69fzthIWjnkSAQuWcxnUecS5z_U2IPNs,8900
 langroid/embedding_models/protoc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 langroid/embedding_models/protoc/embeddings.proto,sha256=_O-SgFpTaylQeOTgSpxhEJ7CUw7PeCQQJLaPqpPYKJg,321
 langroid/embedding_models/protoc/embeddings_pb2.py,sha256=4Q57PhOunv-uZNJrxYrWBXAI0ZtfnVZXFRhRj5JuRSg,1662
@@ -129,8 +129,8 @@ langroid/vector_store/meilisearch.py,sha256=6frB7GFWeWmeKzRfLZIvzRjllniZ1cYj3Hmh
 langroid/vector_store/momento.py,sha256=QaPzUnTwlswoawGB-paLtUPyLRvckFXLfLDfvbTzjNQ,10505
 langroid/vector_store/qdrant_cloud.py,sha256=3im4Mip0QXLkR6wiqVsjV1QvhSElfxdFSuDKddBDQ-4,188
 langroid/vector_store/qdrantdb.py,sha256=wYOuu5c2vIKn9ZgvTXcAiZXMpV8AOXEWFAzI8S8UP-0,16828
-pyproject.toml,sha256=IQC1A8fpg0cKWje_Gt3YxZpRQw9ie1xlJ8TaiS4Vr3w,6960
-langroid-0.2.12.dist-info/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
-langroid-0.2.12.dist-info/METADATA,sha256=66JdF6SmaH5yisjOpT6dgYsjXAG7FbnD8OaTAmVnV7I,54101
-langroid-0.2.12.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-langroid-0.2.12.dist-info/RECORD,,
+pyproject.toml,sha256=gh7X6Y1IgIluL77xY0ejjIJMOIXBPw2-oLYcXUg01Ng,7054
+langroid-0.3.0.dist-info/LICENSE,sha256=EgVbvA6VSYgUlvC3RvPKehSg7MFaxWDsFuzLOsPPfJg,1065
+langroid-0.3.0.dist-info/METADATA,sha256=MUh8_clSRdnwAFTiu8ESCAyXKNkFbU59uXvthr6r2CI,54210
+langroid-0.3.0.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+langroid-0.3.0.dist-info/RECORD,,

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "langroid"
-version = "0.2.12"
+version = "0.3.0"
 description = "Harness LLMs with Multi-Agent Programming"
 authors = ["Prasad Chalasani <pchalasani@gmail.com>"]
 readme = "README.md"
@@ -10,7 +10,7 @@ include = ["pyproject.toml"]
 
 # =============== MAIN DEPS ==============
 [tool.poetry.dependencies]
-python = "^3.10"
+python = ">=3.10,<3.13"
 
 # =========== OPTIONALS ==============================
 chromadb = {version=">=0.4.21, <=0.4.23", optional=true}
@@ -39,6 +39,7 @@ pyarrow = {version="15.0.0", optional=true}
 pdfplumber = {version="^0.10.2", optional=true}
 python-docx = {version="^1.1.0", optional=true}
 scrapy = {version="^2.11.0", optional=true}
+fastembed = {version="^0.3.1", optional=true}
 
 # ====CORE================================
 pyyaml = "^6.0.1"
@@ -111,6 +112,7 @@ all = [
     "metaphor-python", "neo4j",
     "litellm",
     "chainlit", "python-socketio",
+    "fastembed"
 ]
 # more granular groupings
 lancedb = ["lancedb", "tantivy", "pyarrow"]
@@ -135,6 +137,7 @@ mkdocs = [
 ]
 meilisearch = ["meilisearch-python-sdk"]
 momento = ["momento"]
+fastembed = ["fastembed"]
 
 
 # ================= DEV DEPS =================