zrb 1.4.2__py3-none-any.whl → 1.5.0__py3-none-any.whl

zrb/builtin/llm/tool/rag.py

@@ -9,6 +9,8 @@ import ulid
 
 from zrb.config import (
     RAG_CHUNK_SIZE,
+    RAG_EMBEDDING_API_KEY,
+    RAG_EMBEDDING_BASE_URL,
     RAG_EMBEDDING_MODEL,
     RAG_MAX_RESULT_COUNT,
     RAG_OVERLAP,
@@ -35,24 +37,34 @@ def create_rag_from_directory(
     tool_name: str,
     tool_description: str,
     document_dir_path: str = "./documents",
-    model: str = RAG_EMBEDDING_MODEL,
     vector_db_path: str = "./chroma",
     vector_db_collection: str = "documents",
     chunk_size: int = RAG_CHUNK_SIZE,
     overlap: int = RAG_OVERLAP,
     max_result_count: int = RAG_MAX_RESULT_COUNT,
     file_reader: list[RAGFileReader] = [],
+    openai_api_key: str = RAG_EMBEDDING_API_KEY,
+    openai_base_url: str = RAG_EMBEDDING_BASE_URL,
+    openai_embedding_model: str = RAG_EMBEDDING_MODEL,
 ):
     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(
+        from openai import OpenAI
+
+        # Initialize OpenAI client with custom URL if provided
+        client_args = {}
+        if openai_api_key:
+            client_args["api_key"] = openai_api_key
+        if openai_base_url:
+            client_args["base_url"] = openai_base_url
+        # Initialize OpenAI client for embeddings
+        openai_client = OpenAI(**client_args)
+        # Initialize ChromaDB client
+        chroma_client = PersistentClient(
             path=vector_db_path, settings=Settings(allow_reset=True)
         )
-        collection = client.get_or_create_collection(vector_db_collection)
+        collection = chroma_client.get_or_create_collection(vector_db_collection)
         # Track file changes using a hash-based approach
         hash_file_path = os.path.join(vector_db_path, "file_hashes.json")
         previous_hashes = _load_hashes(hash_file_path)
@@ -89,8 +101,11 @@ def create_rag_from_directory(
                                 ),
                                 file=sys.stderr,
                             )
-                            embedding_result = list(embedding_model.embed([chunk]))
-                            vector = embedding_result[0]
+                            # Get embeddings using OpenAI
+                            embedding_response = openai_client.embeddings.create(
+                                input=chunk, model=openai_embedding_model
+                            )
+                            vector = embedding_response.data[0].embedding
                             collection.upsert(
                                 ids=[chunk_id],
                                 embeddings=[vector],
@@ -113,8 +128,11 @@ def create_rag_from_directory(
             )
         # Vectorize query and get related document chunks
         print(stylize_faint("Vectorizing query"), file=sys.stderr)
-        embedding_result = list(embedding_model.embed([query]))
-        query_vector = embedding_result[0]
+        # Get embeddings using OpenAI
+        embedding_response = openai_client.embeddings.create(
+            input=query, model=openai_embedding_model
+        )
+        query_vector = embedding_response.data[0].embedding
         print(stylize_faint("Searching documents"), file=sys.stderr)
         results = collection.query(
             query_embeddings=query_vector,

zrb/builtin/llm/tool/web.py

@@ -3,21 +3,53 @@ from collections.abc import Callable
 from typing import Annotated
 
 
-def open_web_page(url: str) -> str:
-    """Get content from a web page."""
-    import requests
-
-    response = requests.get(
-        url,
-        headers={
-            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"  # noqa
-        },
-    )
-    if response.status_code != 200:
-        raise Exception(
-            f"Error: Unable to retrieve search results (status code: {response.status_code})"  # noqa
-        )
-    return json.dumps(parse_html_text(response.text))
+async def open_web_page(url: str) -> str:
+    """Get content from a web page using a headless browser."""
+
+    async def get_page_content(page_url: str):
+        user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"  # noqa
+        try:
+            from playwright.async_api import async_playwright
+
+            async with async_playwright() as p:
+                browser = await p.chromium.launch(headless=True)
+                page = await browser.new_page()
+                await page.set_extra_http_headers({"User-Agent": user_agent})
+                try:
+                    # Navigate to the URL with a timeout of 30 seconds
+                    await page.goto(page_url, wait_until="networkidle", timeout=30000)
+                    # Wait for the content to load
+                    await page.wait_for_load_state("domcontentloaded")
+                    # Get the page content
+                    content = await page.content()
+                    # Extract all links from the page
+                    links = await page.eval_on_selector_all(
+                        "a[href]",
+                        """
+                        (elements) => elements.map(el => {
+                            const href = el.getAttribute('href');
+                            if (href && !href.startsWith('#') && !href.startsWith('/')) {
+                                return href;
+                            }
+                            return null;
+                        }).filter(href => href !== null)
+                    """,
+                    )
+                    return {"content": content, "links_on_page": links}
+                finally:
+                    await browser.close()
+        except ImportError:
+            import requests
+
+            response = requests.get(url, headers={"User-Agent": user_agent})
+            if response.status_code != 200:
+                msg = f"Unable to retrieve search results. Status code: {response.status_code}"
+                raise Exception(msg)
+            return {"content": response.text, "links_on_page": []}
+
+    result = await get_page_content(url)
+    # Parse the HTML content
+    return json.dumps(parse_html_text(result["content"]))
 
 
 def create_search_internet_tool(serp_api_key: str) -> Callable[[str, int], str]:

zrb/builtin/todo.py

@@ -25,6 +25,18 @@ from zrb.util.todo import (
 )
 
 
+def _get_filter_input(allow_positional_parsing: bool = False) -> StrInput:
+    return StrInput(
+        name="filter",
+        description="Visual filter",
+        prompt="Visual Filter",
+        allow_empty=True,
+        allow_positional_parsing=allow_positional_parsing,
+        always_prompt=False,
+        default=TODO_VISUAL_FILTER,
+    )
+
+
 @make_task(
     name="add-todo",
     input=[
@@ -51,6 +63,7 @@ from zrb.util.todo import (
             prompt="Task context (space separated)",
             allow_empty=True,
         ),
+        _get_filter_input(),
     ],
     description="➕ Add todo",
     group=todo_group,
@@ -82,16 +95,22 @@ def add_todo(ctx: AnyContext):
         )
     )
     save_todo_list(todo_file_path, todo_list)
-    return get_visual_todo_list(todo_list, TODO_VISUAL_FILTER)
+    return get_visual_todo_list(todo_list, filter=ctx.input.filter)
 
 
-@make_task(name="list-todo", description="📋 List todo", group=todo_group, alias="list")
+@make_task(
+    name="list-todo",
+    input=_get_filter_input(allow_positional_parsing=True),
+    description="📋 List todo",
+    group=todo_group,
+    alias="list",
+)
 def list_todo(ctx: AnyContext):
     todo_file_path = os.path.join(TODO_DIR, "todo.txt")
     todo_list: list[TodoTaskModel] = []
     if os.path.isfile(todo_file_path):
         todo_list = load_todo_list(todo_file_path)
-    return get_visual_todo_list(todo_list, TODO_VISUAL_FILTER)
+    return get_visual_todo_list(todo_list, filter=ctx.input.filter)
 
 
 @make_task(
@@ -127,7 +146,10 @@ def show_todo(ctx: AnyContext):
 
 @make_task(
     name="complete-todo",
-    input=StrInput(name="keyword", prompt="Task keyword", description="Task Keyword"),
+    input=[
+        StrInput(name="keyword", prompt="Task keyword", description="Task Keyword"),
+        _get_filter_input(),
+    ],
     description="✅ Complete todo",
     group=todo_group,
     alias="complete",
@@ -141,10 +163,10 @@ def complete_todo(ctx: AnyContext):
     todo_task = select_todo_task(todo_list, ctx.input.keyword)
     if todo_task is None:
         ctx.log_error("Task not found")
-        return get_visual_todo_list(todo_list, TODO_VISUAL_FILTER)
+        return get_visual_todo_list(todo_list, filter=ctx.input.filter)
     if todo_task.completed:
         ctx.log_error("Task already completed")
-        return get_visual_todo_list(todo_list, TODO_VISUAL_FILTER)
+        return get_visual_todo_list(todo_list, filter=ctx.input.filter)
     # Update todo task
     todo_task = cascade_todo_task(todo_task)
     if todo_task.creation_date is not None:
@@ -152,11 +174,12 @@ def complete_todo(ctx: AnyContext):
     todo_task.completed = True
     # Save todo list
     save_todo_list(todo_file_path, todo_list)
-    return get_visual_todo_list(todo_list, TODO_VISUAL_FILTER)
+    return get_visual_todo_list(todo_list, filter=ctx.input.filter)
 
 
 @make_task(
     name="archive-todo",
+    input=_get_filter_input(),
     description="📚 Archive todo",
     group=todo_group,
     alias="archive",
@@ -180,7 +203,7 @@ def archive_todo(ctx: AnyContext):
     ]
     if len(new_archived_todo_list) == 0:
         ctx.print("No completed task to archive")
-        return get_visual_todo_list(todo_list, TODO_VISUAL_FILTER)
+        return get_visual_todo_list(todo_list, filter=ctx.input.filter)
     archive_file_path = os.path.join(TODO_DIR, "archive.txt")
     if not os.path.isdir(TODO_DIR):
         os.make_dirs(TODO_DIR, exist_ok=True)
@@ -192,7 +215,7 @@ def archive_todo(ctx: AnyContext):
     # Save the new todo list and add the archived ones
     save_todo_list(archive_file_path, archived_todo_list)
     save_todo_list(todo_file_path, working_todo_list)
-    return get_visual_todo_list(todo_list, TODO_VISUAL_FILTER)
+    return get_visual_todo_list(todo_list, filter=ctx.input.filter)
 
 
 @make_task(
@@ -216,6 +239,7 @@ def archive_todo(ctx: AnyContext):
             description="Working stop time",
             default=lambda _: _get_default_stop_work_time_str(),
         ),
+        _get_filter_input(),
     ],
     description="🕒 Log work todo",
     group=todo_group,
@@ -230,7 +254,7 @@ def log_todo(ctx: AnyContext):
     todo_task = select_todo_task(todo_list, ctx.input.keyword)
     if todo_task is None:
         ctx.log_error("Task not found")
-        return get_visual_todo_list(todo_list, TODO_VISUAL_FILTER)
+        return get_visual_todo_list(todo_list, filter=ctx.input.filter)
     # Update todo task
     todo_task = cascade_todo_task(todo_task)
     current_duration_str = todo_task.keyval.get("duration", "0")
@@ -268,7 +292,7 @@ def log_todo(ctx: AnyContext):
         log_work_list = json.loads(read_file(log_work_path))
     return "\n".join(
         [
-            get_visual_todo_list(todo_list, TODO_VISUAL_FILTER),
+            get_visual_todo_list(todo_list, filter=ctx.input.filter),
             "",
             get_visual_todo_card(todo_task, log_work_list),
         ]
@@ -296,6 +320,7 @@ def _get_default_stop_work_time_str() -> str:
             default=lambda _: _get_todo_txt_content(),
             allow_positional_parsing=False,
         ),
+        _get_filter_input(),
     ],
     description="📝 Edit todo",
     group=todo_group,
@@ -311,7 +336,7 @@ def edit_todo(ctx: AnyContext):
     todo_file_path = os.path.join(TODO_DIR, "todo.txt")
     write_file(todo_file_path, new_content)
     todo_list = load_todo_list(todo_file_path)
-    return get_visual_todo_list(todo_list, TODO_VISUAL_FILTER)
+    return get_visual_todo_list(todo_list, filter=ctx.input.filter)
 
 
 def _get_todo_txt_content() -> str:

zrb/config.py

@@ -85,10 +85,10 @@ LLM_HISTORY_FILE = os.getenv(
 LLM_ALLOW_ACCESS_LOCAL_FILE = to_boolean(os.getenv("ZRB_LLM_ACCESS_LOCAL_FILE", "1"))
 LLM_ALLOW_ACCESS_SHELL = to_boolean(os.getenv("ZRB_LLM_ACCESS_SHELL", "1"))
 LLM_ALLOW_ACCESS_INTERNET = to_boolean(os.getenv("ZRB_LLM_ACCESS_INTERNET", "1"))
-# 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 Configuration
+RAG_EMBEDDING_API_KEY = os.getenv("ZRB_RAG_EMBEDDING_API_KEY", None)
+RAG_EMBEDDING_BASE_URL = os.getenv("ZRB_RAG_EMBEDDING_BASE_URL", None)
+RAG_EMBEDDING_MODEL = os.getenv("ZRB_RAG_EMBEDDING_MODEL", "text-embedding-ada-002")
 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/llm_config.py

@@ -2,20 +2,29 @@ import os
 
 from pydantic_ai.models import Model
 from pydantic_ai.models.openai import OpenAIModel
+from pydantic_ai.providers import Provider
 from pydantic_ai.providers.openai import OpenAIProvider
 
 DEFAULT_SYSTEM_PROMPT = """
 You have access to tools.
-Your goal to to answer user queries accurately.
+Your goal is to provide insightful and accurate information based on user queries.
 Follow these instructions precisely:
-1. ALWAYS use available tools to gather information BEFORE asking the user questions
-2. For tools that require arguments: provide arguments in valid JSON format
-3. For tools that require NO arguments: call with empty JSON object ({}) NOT empty string ('')
-4. NEVER pass arguments to tools that don't accept parameters
-5. NEVER ask users for information obtainable through tools
-6. Use tools in logical sequence until you have sufficient information
-7. If a tool call fails, check if you're passing arguments in the correct format
-8. Only after exhausting relevant tools should you request clarification
+1. ALWAYS use available tools to gather information BEFORE asking the user questions.
+2. For tools that require arguments: provide arguments in valid JSON format.
+3. For tools with no args: call the tool without args. Do NOT pass "" or {}.
+4. NEVER pass arguments to tools that don't accept parameters.
+5. NEVER ask users for information obtainable through tools.
+6. Use tools in a logical sequence until you have sufficient information.
+7. If a tool call fails, check if you're passing arguments in the correct format.
+   Consider alternative strategies if the issue persists.
+8. Only after exhausting relevant tools should you request clarification.
+9. Understand the context of user queries to provide relevant and accurate responses.
+10. Engage with users in a conversational manner once the necessary information is gathered.
+11. Adapt to different query types or scenarios to improve flexibility and effectiveness.
+""".strip()
+
+DEFAULT_PERSONA = """
+You are an expert in various fields including technology, science, history, and more.
 """.strip()
 
 
@@ -26,6 +35,7 @@ class LLMConfig:
         default_model_name: str | None = None,
         default_base_url: str | None = None,
         default_api_key: str | None = None,
+        default_persona: str | None = None,
         default_system_prompt: str | None = None,
     ):
         self._model_name = (
@@ -48,12 +58,20 @@ class LLMConfig:
             if default_system_prompt is not None
             else os.getenv("ZRB_LLM_SYSTEM_PROMPT", None)
         )
+        self._persona = (
+            default_persona
+            if default_persona is not None
+            else os.getenv("ZRB_LLM_PERSONA", None)
+        )
+        self._default_provider = None
         self._default_model = None
 
     def _get_model_name(self) -> str | None:
         return self._model_name if self._model_name is not None else None
 
-    def _get_model_provider(self) -> OpenAIProvider:
+    def get_default_model_provider(self) -> Provider | str:
+        if self._default_provider is not None:
+            return self._default_provider
         if self._model_base_url is None and self._model_api_key is None:
             return "openai"
         return OpenAIProvider(
@@ -61,9 +79,15 @@ class LLMConfig:
         )
 
     def get_default_system_prompt(self) -> str:
-        if self._system_prompt is not None:
-            return self._system_prompt
-        return DEFAULT_SYSTEM_PROMPT
+        system_prompt = (
+            DEFAULT_SYSTEM_PROMPT
+            if self._system_prompt is None
+            else self._system_prompt
+        )
+        persona = DEFAULT_PERSONA if self._persona is None else self._persona
+        if persona is not None:
+            return f"{persona}\n{system_prompt}"
+        return system_prompt
 
     def get_default_model(self) -> Model | str | None:
         if self._default_model is not None:
@@ -73,7 +97,7 @@ class LLMConfig:
             return None
         return OpenAIModel(
             model_name=model_name,
-            provider=self._get_model_provider(),
+            provider=self.get_default_model_provider(),
         )
 
     def set_default_system_prompt(self, system_prompt: str):
@@ -88,6 +112,9 @@ class LLMConfig:
     def set_default_model_base_url(self, model_base_url: str):
         self._model_base_url = model_base_url
 
+    def set_default_provider(self, provider: Provider | str):
+        self._default_provider = provider
+
     def set_default_model(self, model: Model | str | None):
         self._default_model = model
 

zrb/task/llm_task.py

@@ -1,5 +1,8 @@
+import functools
+import inspect
 import json
 import os
+import traceback
 from collections.abc import Callable
 from typing import Any
 
@@ -201,6 +204,9 @@ class LLMTask(BaseTask):
             async with node.stream(agent_run.ctx) as handle_stream:
                 async for event in handle_stream:
                     if isinstance(event, FunctionToolCallEvent):
+                        # Fixing anthrophic claude when call function with empty parameter
+                        if event.part.args == "":
+                            event.part.args = {}
                         ctx.print(
                             stylize_faint(
                                 f"[Tools] The LLM calls tool={event.part.tool_name!r} with args={event.part.args} (tool_call_id={event.part.tool_call_id!r})"  # noqa
@@ -240,7 +246,7 @@ class LLMTask(BaseTask):
         )
         tools_or_callables.extend(self._additional_tools)
         tools = [
-            tool if isinstance(tool, Tool) else Tool(tool, takes_ctx=False)
+            tool if isinstance(tool, Tool) else Tool(_wrap_tool(tool), takes_ctx=False)
             for tool in tools_or_callables
         ]
         return Agent(
@@ -256,21 +262,17 @@ class LLMTask(BaseTask):
         if model is None:
             return default_llm_config.get_default_model()
         if isinstance(model, str):
+            model_base_url = self._get_model_base_url(ctx)
+            model_api_key = self._get_model_api_key(ctx)
             llm_config = LLMConfig(
                 default_model_name=model,
-                default_base_url=get_attr(
-                    ctx,
-                    self._get_model_base_url(ctx),
-                    None,
-                    auto_render=self._render_model_base_url,
-                ),
-                default_api_key=get_attr(
-                    ctx,
-                    self._get_model_api_key(ctx),
-                    None,
-                    auto_render=self._render_model_api_key,
-                ),
+                default_base_url=model_base_url,
+                default_api_key=model_api_key,
             )
+            if model_base_url is None and model_api_key is None:
+                default_model_provider = default_llm_config.get_default_model_provider()
+                if default_model_provider is not None:
+                    llm_config.set_default_provider(default_model_provider)
             return llm_config.get_default_model()
         raise ValueError(f"Invalid model: {model}")
 
@@ -288,7 +290,7 @@ class LLMTask(BaseTask):
         )
         if isinstance(api_key, str) or api_key is None:
             return api_key
-        raise ValueError(f"Invalid model base URL: {api_key}")
+        raise ValueError(f"Invalid model API key: {api_key}")
 
     def _get_system_prompt(self, ctx: AnyContext) -> str:
         system_prompt = get_attr(
@@ -325,3 +327,16 @@ class LLMTask(BaseTask):
             "",
             auto_render=self._render_history_file,
         )
+
+
+def _wrap_tool(func):
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        try:
+            return await run_async(func(*args, **kwargs))
+        except Exception as e:
+            # Optionally, you can include more details from traceback if needed.
+            error_details = traceback.format_exc()
+            return f"Error: {e}\nDetails: {error_details}"
+
+    return wrapper