zrb 1.4.2__tar.gz → 1.4.3__tar.gz

{zrb-1.4.2 → zrb-1.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.4.2
+Version: 1.4.3
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -13,16 +13,20 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Provides-Extra: all
+Provides-Extra: playwright
 Provides-Extra: rag
 Requires-Dist: autopep8 (>=2.0.4,<3.0.0)
 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: chromadb (>=0.5.20,<0.6.0) ; extra == "rag" or extra == "all"
 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"
+Requires-Dist: openai (>=1.10.0,<2.0.0) ; extra == "rag" or extra == "all"
+Requires-Dist: pdfplumber (>=0.11.4,<0.12.0) ; extra == "rag" or extra == "all"
+Requires-Dist: playwright (>=1.43.0,<2.0.0) ; extra == "playwright" or extra == "all"
 Requires-Dist: psutil (>=6.1.1,<7.0.0)
 Requires-Dist: pydantic-ai (>=0.0.42,<0.0.43)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
@@ -79,40 +83,48 @@ bash -c "$(curl -fsSL https://raw.githubusercontent.com/state-alchemists/zrb/mai
 
 ```
 
-# 🍲 Quick Start
+# 🍲 Quick Start: Build Your First Automation Workflow
 
-Create a file at `/home/<your-user-name>/zrb_init.py` with the following content:
+Zrb empowers you to create custom automation tasks using Python. This guide shows you how to define two simple tasks: one to generate a PlantUML script from your source code and another to convert that script into a PNG image.
 
+## 1. Create Your Task Definition File
+
+Place a file named `zrb_init.py` in a directory that's accessible from your projects. Zrb will automatically search for this file by starting in your current directory and then moving upward (i.e., checking parent directories) until it finds one. This means if you place your `zrb_init.py` in your home directory (e.g., `/home/<your-user-name>/zrb_init.py`), the tasks defined there will be available for any project.
+
+Add the following content to your zrb_init.py:
 
 ```python
 import os
-from zrb import cli, llm_config, LLMTask, CmdTask, StrInput, Group
-from zrb.builtin.llm.tool.file import read_all_files, write_text_file
+from zrb import cli, LLMTask, CmdTask, StrInput, Group
+from zrb.builtin.llm.tool.file import (
+    read_text_file, list_files, write_text_file
+)
 
 CURRENT_DIR = os.getcwd()
 
-# Make UML group
+# Create a group for UML-related tasks
 uml_group = cli.add_group(Group(name="uml", description="UML related tasks"))
 
-# Generate UML script
+# Task 1: Generate a PlantUML script from your source code
 make_uml_script = uml_group.add_task(
     LLMTask(
         name="make-script",
         description="Creating plantuml diagram based on source code in current directory",
         input=StrInput(name="diagram", default="state diagram"),
         message=(
-            f"Read source code in {CURRENT_DIR}, "
+            f"Read all necessary files in {CURRENT_DIR}, "
             "make a {ctx.input.diagram} in plantuml format. "
             f"Write the script into {CURRENT_DIR}/{{ctx.input.diagram}}.uml"
         ),
         tools=[
-            read_all_files,
+            list_files,
+            read_text_file,
             write_text_file,
         ],
     )
 )
 
-# Defining a Cmd Task to transform Plantuml script into a png image.
+# Task 2: Convert the PlantUML script into a PNG image
 make_uml_image = uml_group.add_task(
     CmdTask(
         name="make-image",
@@ -123,63 +135,75 @@ make_uml_image = uml_group.add_task(
     )
 )
 
-# Making sure that make_png has make_uml as its dependency.
+# Set up the dependency: the image task runs after the script is created
 make_uml_script >> make_uml_image
 ```
 
-You have just define two automation tasks.
+**What This Does**
 
-The first one use LLM to read files in your current directory and create a `PlantUML script` on that directory.
+- **Task 1 – make-script**:
 
-The second task turn the PlantUML script into a `*.png` file. The second task depends on the first task and both of them are located under the same group.
+    Uses an LLM to read all files in your current directory and generate a PlantUML script (e.g., `state diagram.uml`).
 
-You can run the tasks by invoking `zrb uml make-script` or `zrb uml make-image` respectively.
+- **Task 2 – make-image**:
 
-When you run zrb, it automatically searches for a file named `zrb_init.py` starting from your current directory and moving upward through its parent directories. This design lets you set up common automation tasks in a central location—like placing a `zrb_init.py` in your home directory (`/home/<your-user>/zrb_init.py`)—so that your tasks are available across all your projects.
+    Executes a command that converts the PlantUML script into a PNG image (e.g., `state diagram.png`). This task will run only after the script has been generated.
 
-Now, go to your project and create a state diagram:
 
-```bash
-git clone git@github.com:jjinux/gotetris.git
-cd gotetris
-zrb uml make-image --diagram "state diagram"
-```
+## 2. Run Your Tasks
 
-You can also invoke the task without specifying parameter.
+After setting up your tasks, you can execute them from any project. For example:
 
-```bash
-zrb uml make-image
-```
+- Clone/Create a Project:
 
-Once you do so, Zrb will ask you to provide the diagram type.
+    ```bash
+    git clone git@github.com:jjinux/gotetris.git
+    cd gotetris
+    ```
 
-```
-diagram [state diagram]:
-```
+- Create a state diagram:
+
+    ```bash
+    zrb uml make-image --diagram "state diagram"
+    ```
+
+- Or use the interactive mode:
 
-You can just press enter if you want to use the default value (i.e., in this case `state diagram`).
+    ```bash
+    zrb uml make-image
+    ```
 
-Finally, you can also serve the tasks via a Web UI interface by invoking the following command:
+    Zrb will prompt:
+
+    ```bash
+    diagram [state diagram]:
+    ```
+
+    Press **Enter** to use the default value
+
+![State Diagram](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/state-diagram.png)
+
+
+## 3. Try Out the Web UI
+
+You can also serve your tasks through a user-friendly web interface:
 
 ```bash
 zrb server start
 ```
 
-You will have a nice web interface running on `http://localhost:12123`
+Then open your browser and visit `http://localhost:21213`
 
 ![Zrb Web UI](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/zrb-web-ui.png)
 
-Now, let's see how things work in detail. First, Zrb generates a `state diagram.uml` in your current directory, it then transform the UML script into a PNG image `state diagram.png`.
-
-![State Diagram](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/state-diagram.png)
-
 
 # 🎥 Demo & Documentation
 
 - **Step by step guide:** [Getting started with Zrb](https://github.com/state-alchemists/zrb/blob/main/docs/recipes/getting-started/README.md).
 - **Full documentation:** [Zrb Documentation](https://github.com/state-alchemists/zrb/blob/main/docs/README.md)
 - **Video demo:**
-  [![Video Title](https://img.youtube.com/vi/W7dgk96l__o/0.jpg)](https://www.youtube.com/watch?v=W7dgk96l__o)
+
+    [![Video Title](https://img.youtube.com/vi/W7dgk96l__o/0.jpg)](https://www.youtube.com/watch?v=W7dgk96l__o)
 
 
 # 🤝 Join the Community

{zrb-1.4.2 → zrb-1.4.3}/README.md

@@ -44,40 +44,48 @@ bash -c "$(curl -fsSL https://raw.githubusercontent.com/state-alchemists/zrb/mai
 
 ```
 
-# 🍲 Quick Start
+# 🍲 Quick Start: Build Your First Automation Workflow
 
-Create a file at `/home/<your-user-name>/zrb_init.py` with the following content:
+Zrb empowers you to create custom automation tasks using Python. This guide shows you how to define two simple tasks: one to generate a PlantUML script from your source code and another to convert that script into a PNG image.
 
+## 1. Create Your Task Definition File
+
+Place a file named `zrb_init.py` in a directory that's accessible from your projects. Zrb will automatically search for this file by starting in your current directory and then moving upward (i.e., checking parent directories) until it finds one. This means if you place your `zrb_init.py` in your home directory (e.g., `/home/<your-user-name>/zrb_init.py`), the tasks defined there will be available for any project.
+
+Add the following content to your zrb_init.py:
 
 ```python
 import os
-from zrb import cli, llm_config, LLMTask, CmdTask, StrInput, Group
-from zrb.builtin.llm.tool.file import read_all_files, write_text_file
+from zrb import cli, LLMTask, CmdTask, StrInput, Group
+from zrb.builtin.llm.tool.file import (
+    read_text_file, list_files, write_text_file
+)
 
 CURRENT_DIR = os.getcwd()
 
-# Make UML group
+# Create a group for UML-related tasks
 uml_group = cli.add_group(Group(name="uml", description="UML related tasks"))
 
-# Generate UML script
+# Task 1: Generate a PlantUML script from your source code
 make_uml_script = uml_group.add_task(
     LLMTask(
         name="make-script",
         description="Creating plantuml diagram based on source code in current directory",
         input=StrInput(name="diagram", default="state diagram"),
         message=(
-            f"Read source code in {CURRENT_DIR}, "
+            f"Read all necessary files in {CURRENT_DIR}, "
             "make a {ctx.input.diagram} in plantuml format. "
             f"Write the script into {CURRENT_DIR}/{{ctx.input.diagram}}.uml"
         ),
         tools=[
-            read_all_files,
+            list_files,
+            read_text_file,
             write_text_file,
         ],
     )
 )
 
-# Defining a Cmd Task to transform Plantuml script into a png image.
+# Task 2: Convert the PlantUML script into a PNG image
 make_uml_image = uml_group.add_task(
     CmdTask(
         name="make-image",
@@ -88,63 +96,75 @@ make_uml_image = uml_group.add_task(
     )
 )
 
-# Making sure that make_png has make_uml as its dependency.
+# Set up the dependency: the image task runs after the script is created
 make_uml_script >> make_uml_image
 ```
 
-You have just define two automation tasks.
+**What This Does**
 
-The first one use LLM to read files in your current directory and create a `PlantUML script` on that directory.
+- **Task 1 – make-script**:
 
-The second task turn the PlantUML script into a `*.png` file. The second task depends on the first task and both of them are located under the same group.
+    Uses an LLM to read all files in your current directory and generate a PlantUML script (e.g., `state diagram.uml`).
 
-You can run the tasks by invoking `zrb uml make-script` or `zrb uml make-image` respectively.
+- **Task 2 – make-image**:
 
-When you run zrb, it automatically searches for a file named `zrb_init.py` starting from your current directory and moving upward through its parent directories. This design lets you set up common automation tasks in a central location—like placing a `zrb_init.py` in your home directory (`/home/<your-user>/zrb_init.py`)—so that your tasks are available across all your projects.
+    Executes a command that converts the PlantUML script into a PNG image (e.g., `state diagram.png`). This task will run only after the script has been generated.
 
-Now, go to your project and create a state diagram:
 
-```bash
-git clone git@github.com:jjinux/gotetris.git
-cd gotetris
-zrb uml make-image --diagram "state diagram"
-```
+## 2. Run Your Tasks
 
-You can also invoke the task without specifying parameter.
+After setting up your tasks, you can execute them from any project. For example:
 
-```bash
-zrb uml make-image
-```
+- Clone/Create a Project:
 
-Once you do so, Zrb will ask you to provide the diagram type.
+    ```bash
+    git clone git@github.com:jjinux/gotetris.git
+    cd gotetris
+    ```
 
-```
-diagram [state diagram]:
-```
+- Create a state diagram:
+
+    ```bash
+    zrb uml make-image --diagram "state diagram"
+    ```
+
+- Or use the interactive mode:
 
-You can just press enter if you want to use the default value (i.e., in this case `state diagram`).
+    ```bash
+    zrb uml make-image
+    ```
 
-Finally, you can also serve the tasks via a Web UI interface by invoking the following command:
+    Zrb will prompt:
+
+    ```bash
+    diagram [state diagram]:
+    ```
+
+    Press **Enter** to use the default value
+
+![State Diagram](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/state-diagram.png)
+
+
+## 3. Try Out the Web UI
+
+You can also serve your tasks through a user-friendly web interface:
 
 ```bash
 zrb server start
 ```
 
-You will have a nice web interface running on `http://localhost:12123`
+Then open your browser and visit `http://localhost:21213`
 
 ![Zrb Web UI](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/zrb-web-ui.png)
 
-Now, let's see how things work in detail. First, Zrb generates a `state diagram.uml` in your current directory, it then transform the UML script into a PNG image `state diagram.png`.
-
-![State Diagram](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/state-diagram.png)
-
 
 # 🎥 Demo & Documentation
 
 - **Step by step guide:** [Getting started with Zrb](https://github.com/state-alchemists/zrb/blob/main/docs/recipes/getting-started/README.md).
 - **Full documentation:** [Zrb Documentation](https://github.com/state-alchemists/zrb/blob/main/docs/README.md)
 - **Video demo:**
-  [![Video Title](https://img.youtube.com/vi/W7dgk96l__o/0.jpg)](https://www.youtube.com/watch?v=W7dgk96l__o)
+
+    [![Video Title](https://img.youtube.com/vi/W7dgk96l__o/0.jpg)](https://www.youtube.com/watch?v=W7dgk96l__o)
 
 
 # 🤝 Join the Community

{zrb-1.4.2 → zrb-1.4.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.4.2"
+version = "1.4.3"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"
@@ -52,11 +52,15 @@ python-jose = {extras = ["cryptography"], version = "^3.4.0"}
 ulid-py = "^1.1.0"
 pydantic-ai = "^0.0.42"
 fastembed = "^0.5.1"
+openai = {version = "^1.10.0", optional = true}
+playwright = {version = "^1.43.0", optional = true}
 psutil = "^6.1.1"
 
 [tool.poetry.extras]
 # poetry install -E rag
-rag = ["chromadb", "pdfplumber"]
+rag = ["chromadb", "pdfplumber", "openai"]
+playwright = ["playwright"]
+all = ["chromadb", "pdfplumber", "openai", "playwright"]
 
 [tool.poetry.group.dev.dependencies]
 alembic = "^1.14.0" # FastApp dependencies

{zrb-1.4.2 → zrb-1.4.3}/src/zrb/builtin/llm/tool/api.py

@@ -5,7 +5,7 @@ from typing import Annotated, Literal
 def get_current_location() -> (
     Annotated[str, "JSON string representing latitude and longitude"]
 ):  # noqa
-    """Get the user's current location."""
+    """Get the user's current location. This function take no argument."""
     import requests
 
     return json.dumps(requests.get("http://ip-api.com/json?fields=lat,lon").json())

{zrb-1.4.2 → zrb-1.4.3}/src/zrb/builtin/llm/tool/file.py

@@ -93,6 +93,7 @@ def list_files(
     """List all files in a directory that match any of the included glob patterns
     and do not reside in any directory matching an excluded pattern.
     Patterns are evaluated using glob-style matching.
+    included_patterns and excluded_patterns already has sane default values.
     """
     all_files: list[str] = []
     for root, dirs, files in os.walk(directory):
@@ -144,6 +145,7 @@ def read_all_files(
     """Read all files in a directory that match any of the included glob patterns
     and do not match any of the excluded glob patterns.
     Patterns are evaluated using glob-style matching.
+    included_patterns and excluded_patterns already has sane default values.
     """
     files = list_files(directory, included_patterns, excluded_patterns)
     for index, file in enumerate(files):

{zrb-1.4.2 → zrb-1.4.3}/src/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-1.4.2 → zrb-1.4.3}/src/zrb/builtin/llm/tool/web.py

@@ -3,21 +3,61 @@ 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):
+        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()
+                # Set user agent to mimic a regular browser
+                user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+                user_agent += "AppleWebKit/537.36 (KHTML, like Gecko) "
+                user_agent += "Chrome/91.0.4472.124 Safari/537.36"
+                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": "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:
+                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-1.4.2 → zrb-1.4.3}/src/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-1.4.2 → zrb-1.4.3}/src/zrb/llm_config.py

@@ -2,18 +2,19 @@ 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 answer user queries accurately.
 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 ('')
+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 logical sequence until you have sufficient information
+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
 8. Only after exhausting relevant tools should you request clarification
 """.strip()
@@ -48,12 +49,15 @@ class LLMConfig:
             if default_system_prompt is not None
             else os.getenv("ZRB_LLM_SYSTEM_PROMPT", 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(
@@ -73,7 +77,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 +92,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