zrb 1.0.0b5__py3-none-any.whl → 1.0.0b6__py3-none-any.whl

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

{zrb-1.0.0b5.dist-info → zrb-1.0.0b6.dist-info}/METADATA

@@ -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.dist-info → zrb-1.0.0b6.dist-info}/RECORD

@@ -11,7 +11,7 @@ zrb/builtin/llm/llm_chat.py,sha256=vUUchYKJuHr-N_HarpKRFsV0EdQDAFZzAfbK4C0vx88,4
 zrb/builtin/llm/previous-session.js,sha256=xMKZvJoAbrwiyHS0OoPrWuaKxWYLoyR5sguePIoCjTY,816
 zrb/builtin/llm/tool/api.py,sha256=yQ3XV8O7Fx7hHssLSOcmiHDnevPhz9ktWi44HK7zTls,801
 zrb/builtin/llm/tool/cli.py,sha256=to_IjkfrMGs6eLfG0cpVN9oyADWYsJQCtyluUhUdBww,253
-zrb/builtin/llm/tool/rag.py,sha256=jJRLERW6824JeEzEQ_OqLMaaa3mjuNqsRcRWoL1wVx0,5192
+zrb/builtin/llm/tool/rag.py,sha256=PawaLZL-ThctxtBtsQuP3XsgTxQKyCGFqrudCANPJKk,5162
 zrb/builtin/llm/tool/web.py,sha256=N2HYuXbKPUpjVAq_UnQMbUrTIE8u0Ut3TeQadZ7_NJc,2217
 zrb/builtin/md5.py,sha256=0pNlrfZA0wlZlHvFHLgyqN0JZJWGKQIF5oXxO44_OJk,949
 zrb/builtin/project/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -171,7 +171,7 @@ zrb/callback/callback.py,sha256=hKefB_Jd1XGjPSLQdMKDsGLHPzEGO2dqrIArLl_EmD0,848
 zrb/cmd/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/cmd/cmd_result.py,sha256=L8bQJzWCpcYexIxHBNsXj2pT3BtLmWex0iJSMkvimOA,597
 zrb/cmd/cmd_val.py,sha256=7Doowyg6BK3ISSGBLt-PmlhzaEkBjWWm51cED6fAUOQ,1014
-zrb/config.py,sha256=vBBhSZ5xQV5S_20PQS54ShFgV-ko3Zbzxw23wDDJnbQ,4556
+zrb/config.py,sha256=RB0ikM82RnCDZaarJOEBfD7GeYPfcq6jLb6F359il1s,4677
 zrb/content_transformer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/content_transformer/any_content_transformer.py,sha256=v8ZUbcix1GGeDQwB6OKX_1TjpY__ksxWVeqibwa_iZA,850
 zrb/content_transformer/content_transformer.py,sha256=YU6Xr3G_IaCWKQGsf9z9YlCclbiwcJ7ytQv3wKpPIiI,2125
@@ -262,7 +262,7 @@ zrb/task/base_task.py,sha256=ImA0ReyB6neVUfY4nKLnL0h2EMGIJ9wvvNvIAN92-RE,21194
 zrb/task/base_trigger.py,sha256=jC722rDvodaBLeNaFghkTyv1u0QXrK6BLZUUqcmBJ7Q,4581
 zrb/task/cmd_task.py,sha256=JpClYoEmJTqKSxhuCErXd2kHLS3Hk2zXeYnl7entNeU,10378
 zrb/task/http_check.py,sha256=Gf5rOB2Se2EdizuN9rp65HpGmfZkGc-clIAlHmPVehs,2565
-zrb/task/llm_task.py,sha256=hU1waf3NFc41YxdIe0k9z1eKrcPDaBUo4imuEq737TU,7782
+zrb/task/llm_task.py,sha256=ptXC3x9Dwn7-4JrGQyEtzOXZ4dNQATDgCeowkvwAu9U,7723
 zrb/task/make_task.py,sha256=PD3b_aYazthS8LHeJsLAhwKDEgdurQZpymJDKeN60u0,2265
 zrb/task/rsync_task.py,sha256=pVVslZ46qgcpU_EKhyTQEQie8kUOMuTsVQdbQG2L-yk,6318
 zrb/task/scaffolder.py,sha256=rME18w1HJUHXgi9eTYXx_T2G4JdqDYzBoNOkdOOo5-o,6806
@@ -303,7 +303,7 @@ zrb/util/string/name.py,sha256=8picJfUBXNpdh64GNaHv3om23QHhUZux7DguFLrXHp8,1163
 zrb/util/todo.py,sha256=1nDdwPc22oFoK_1ZTXyf3638Bg6sqE2yp_U4_-frHoc,16015
 zrb/xcom/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/xcom/xcom.py,sha256=o79rxR9wphnShrcIushA0Qt71d_p3ZTxjNf7x9hJB78,1571
-zrb-1.0.0b5.dist-info/METADATA,sha256=BhCdzcbWjo2XbxLbBVo7JtQpGUAnQS1CnKyRgTkN-cM,4228
-zrb-1.0.0b5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-zrb-1.0.0b5.dist-info/entry_points.txt,sha256=-Pg3ElWPfnaSM-XvXqCxEAa-wfVI6BEgcs386s8C8v8,46
-zrb-1.0.0b5.dist-info/RECORD,,
+zrb-1.0.0b6.dist-info/METADATA,sha256=ZGOBMXjA82fN6lsJ3yNmJTGqwoZWJcNEKCxXtTouk68,4270
+zrb-1.0.0b6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+zrb-1.0.0b6.dist-info/entry_points.txt,sha256=-Pg3ElWPfnaSM-XvXqCxEAa-wfVI6BEgcs386s8C8v8,46
+zrb-1.0.0b6.dist-info/RECORD,,