zrb 1.0.0b5__tar.gz → 1.0.0b6__tar.gz

{zrb-1.0.0b5 → zrb-1.0.0b6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.0.0b5
+Version: 1.0.0b6
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -19,6 +19,7 @@ Requires-Dist: beautifulsoup4 (>=4.12.3,<5.0.0)
 Requires-Dist: black (>=24.10.0,<24.11.0)
 Requires-Dist: chromadb (>=0.5.20,<0.6.0) ; extra == "rag"
 Requires-Dist: fastapi[standard] (>=0.115.6,<0.116.0)
+Requires-Dist: fastembed (>=0.5.1,<0.6.0)
 Requires-Dist: isort (>=5.13.2,<5.14.0)
 Requires-Dist: libcst (>=1.5.0,<2.0.0)
 Requires-Dist: pdfplumber (>=0.11.4,<0.12.0) ; extra == "rag"

{zrb-1.0.0b5 → zrb-1.0.0b6}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.0.0b5"
+version = "1.0.0b6"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"
@@ -51,6 +51,7 @@ fastapi = {extras = ["standard"], version = "^0.115.6"}
 python-jose = {extras = ["cryptography"], version = "^3.3.0"}
 ulid-py = "^1.1.0"
 pydantic-ai = "^0.0.19"
+fastembed = "^0.5.1"
 
 [tool.poetry.extras]
 # poetry install -E rag

{zrb-1.0.0b5 → zrb-1.0.0b6}/src/zrb/builtin/llm/tool/rag.py

@@ -3,7 +3,6 @@ import json
 import os
 import sys
 
-import litellm
 import ulid
 
 from zrb.config import (
@@ -30,7 +29,9 @@ def create_rag_from_directory(
     async def retrieve(query: str) -> str:
         from chromadb import PersistentClient
         from chromadb.config import Settings
+        from fastembed import TextEmbedding
 
+        embedding_model = TextEmbedding(model_name=model)
         client = PersistentClient(
             path=vector_db_path, settings=Settings(allow_reset=True)
         )
@@ -75,10 +76,8 @@ def create_rag_from_directory(
                                 ),
                                 file=sys.stderr,
                             )
-                            response = await litellm.aembedding(
-                                model=model, input=[chunk]
-                            )
-                            vector = response["data"][0]["embedding"]
+                            embedding_result = list(embedding_model.embed([chunk]))
+                            vector = embedding_result[0]
                             collection.upsert(
                                 ids=[chunk_id],
                                 embeddings=[vector],
@@ -102,9 +101,8 @@ def create_rag_from_directory(
             )
 
         print(stylize_faint("Vectorizing query"), file=sys.stderr)
-        query_response = await litellm.aembedding(model=model, input=[query])
-        query_vector = query_response["data"][0]["embedding"]
-
+        embedding_result = list(embedding_model.embed([query]))
+        query_vector = embedding_result[0]
         print(stylize_faint("Searching documents"), file=sys.stderr)
         results = collection.query(
             query_embeddings=query_vector,

{zrb-1.0.0b5 → zrb-1.0.0b6}/src/zrb/config.py

@@ -95,7 +95,10 @@ LLM_HISTORY_FILE = os.getenv(
 )
 LLM_ALLOW_ACCESS_SHELL = to_boolean(os.getenv("ZRB_LLM_ACCESS_FILE", "1"))
 LLM_ALLOW_ACCESS_INTERNET = to_boolean(os.getenv("ZRB_LLM_ACCESS_INTERNET", "1"))
-RAG_EMBEDDING_MODEL = os.getenv("ZRB_RAG_EMBEDDING_MODEL", "ollama/nomic-embed-text")
+# noqa See: https://qdrant.github.io/fastembed/examples/Supported_Models/#supported-text-embedding-models
+RAG_EMBEDDING_MODEL = os.getenv(
+    "ZRB_RAG_EMBEDDING_MODEL", "nomic-ai/nomic-embed-text-v1.5-Q"
+)
 RAG_CHUNK_SIZE = int(os.getenv("ZRB_RAG_CHUNK_SIZE", "1024"))
 RAG_OVERLAP = int(os.getenv("ZRB_RAG_OVERLAP", "128"))
 RAG_MAX_RESULT_COUNT = int(os.getenv("ZRB_RAG_MAX_RESULT_COUNT", "5"))

{zrb-1.0.0b5 → zrb-1.0.0b6}/src/zrb/task/llm_task.py

@@ -154,8 +154,6 @@ class LLMTask(BaseTask):
             tool if isinstance(tool, Tool) else Tool(tool, takes_ctx=False)
             for tool in tools_or_callables
         ]
-        for tool in tools:
-            print("tool", tool)
         return Agent(
             self._get_model(ctx),
             system_prompt=self._get_system_prompt(ctx),