zrb 1.2.1__tar.gz → 1.3.0__tar.gz

{zrb-1.2.1 → zrb-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.2.1
+Version: 1.3.0
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -24,7 +24,7 @@ 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: psutil (>=6.1.1,<7.0.0)
-Requires-Dist: pydantic-ai (>=0.0.19,<0.0.20)
+Requires-Dist: pydantic-ai (>=0.0.31,<0.0.32)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: python-jose[cryptography] (>=3.4.0,<4.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)
@@ -39,29 +39,96 @@ Description-Content-Type: text/markdown
 
 # 🤖 Zrb: Your Automation Powerhouse
 
-Zrb allows you to write your automation tasks in Python and declaratively:
+Zrb allows you to write your automation tasks in Python. For example, you can define the following script in your home directory (`/home/<your-user-name>/zrb_init.py`).
 
 
 ```python
-# Filename: zrb_init.py
-from zrb import cli, Task, Group, IntInput
-
-math = cli.add_group(Group("math", description="Math tools"))
-math.add_task(Task(
-    name="add",
-    input=[
-        IntInput("a", description="First number"),
-        IntInput("b", description="Second number")
-    ],
-    action=lambda ctx: ctx.input.a + ctx.input.b
-))
+import os
+from zrb import cli, LLMTask, CmdTask, StrInput
+from zrb.builtin.llm.tool.file import read_source_code, write_text_file
+from pydantic_ai.models.openai import OpenAIModel
+
+
+CURRENT_DIR = os.getcwd()
+OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
+OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "")
+OPENROUTER_MODEL_NAME = os.getenv(
+    "AGENT_MODEL_NAME", "anthropic/claude-3.7-sonnet"
+)
+
+
+# Defining a LLM Task to create a Plantuml script based on source code in current directory.
+# User can choose the diagram type. By default it is "state diagram"
+make_uml = cli.add_task(
+    LLMTask(
+        name="make-uml",
+        description="Creating plantuml diagram based on source code in current directory",
+        input=StrInput(name="diagram", default="state diagram"),
+        model=OpenAIModel(
+            OPENROUTER_MODEL_NAME,
+            base_url=OPENROUTER_BASE_URL,
+            api_key=OPENROUTER_API_KEY,
+        ),
+        message=(
+            f"Read source code in {CURRENT_DIR}, "
+            "make a {ctx.input.diagram} in plantuml format. "
+            f"Write the script into {CURRENT_DIR}/{{ctx.input.diagram}}.uml"
+        ),
+        tools=[
+            read_source_code,
+            write_text_file,
+        ],
+    )
+)
+
+# Defining a Cmd Task to transform Plantuml script into a png image.
+make_png = cli.add_task(
+    CmdTask(
+        name="make-png",
+        description="Creating png based on source code in current directory",
+        input=StrInput(name="diagram", default="state diagram"),
+        cmd="plantuml -tpng '{ctx.input.diagram}.uml'",
+        cwd=CURRENT_DIR,
+    )
+)
+
+# Making sure that make_png has make_uml as its dependency.
+make_uml >> make_png
 ```
 
-Once defined, you will be able to access your automation tasks from the CLI, Web Interface, or via HTTP API.
+Once defined, your automation tasks are immediately accessible from the CLI. You can then invoke the tasks by invoking.
 
-For more complex scenario, you can also defined Task dependencies (upstreams) and retry mechanisms. You can also make a scheduled tasks, just like in Apache Airflow.
+```bash
+zrb make-png --diagram "state diagram"
+```
+
+Or you can invoke the tasks without parameter.
+
+```bash
+zrb make-png
+```
+
+At this point, Zrb will politely ask you to provide the diagram type.
+
+```
+diagram [state diagram]:
+```
+
+You can just press enter if you want to use the default value.
+
+Finally, you can run Zrb as a server and make your tasks available for non technical users by invoking the following command.
+
+```bash
+zrb server start
+```
+
+You will have a nice web interface running on `http://localhost:12123`
+
+![Zrb Web UI](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/zrb-web-ui.png)
+
+Now, let's see how Zrb generate the state diagram. Based on the source code in your current directory, Zrb will generate a `state diagram.uml` and transform it into `state diagram.png`.
 
-Furthermore, Zrb has some builtin tasks to manage monorepo, generate FastAPI application, or play around with LLM.
+![State Diagram](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/state-diagram.png)
 
 See the [getting started guide](https://github.com/state-alchemists/zrb/blob/main/docs/recipes/getting-started/README.md) for more information. Or just watch the demo:
 

zrb-1.3.0/README.md

@@ -0,0 +1,139 @@
+![](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/zrb/android-chrome-192x192.png)
+
+[Documentation](https://github.com/state-alchemists/zrb/blob/main/docs/README.md)
+
+# 🤖 Zrb: Your Automation Powerhouse
+
+Zrb allows you to write your automation tasks in Python. For example, you can define the following script in your home directory (`/home/<your-user-name>/zrb_init.py`).
+
+
+```python
+import os
+from zrb import cli, LLMTask, CmdTask, StrInput
+from zrb.builtin.llm.tool.file import read_source_code, write_text_file
+from pydantic_ai.models.openai import OpenAIModel
+
+
+CURRENT_DIR = os.getcwd()
+OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1"
+OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY", "")
+OPENROUTER_MODEL_NAME = os.getenv(
+    "AGENT_MODEL_NAME", "anthropic/claude-3.7-sonnet"
+)
+
+
+# Defining a LLM Task to create a Plantuml script based on source code in current directory.
+# User can choose the diagram type. By default it is "state diagram"
+make_uml = cli.add_task(
+    LLMTask(
+        name="make-uml",
+        description="Creating plantuml diagram based on source code in current directory",
+        input=StrInput(name="diagram", default="state diagram"),
+        model=OpenAIModel(
+            OPENROUTER_MODEL_NAME,
+            base_url=OPENROUTER_BASE_URL,
+            api_key=OPENROUTER_API_KEY,
+        ),
+        message=(
+            f"Read source code in {CURRENT_DIR}, "
+            "make a {ctx.input.diagram} in plantuml format. "
+            f"Write the script into {CURRENT_DIR}/{{ctx.input.diagram}}.uml"
+        ),
+        tools=[
+            read_source_code,
+            write_text_file,
+        ],
+    )
+)
+
+# Defining a Cmd Task to transform Plantuml script into a png image.
+make_png = cli.add_task(
+    CmdTask(
+        name="make-png",
+        description="Creating png based on source code in current directory",
+        input=StrInput(name="diagram", default="state diagram"),
+        cmd="plantuml -tpng '{ctx.input.diagram}.uml'",
+        cwd=CURRENT_DIR,
+    )
+)
+
+# Making sure that make_png has make_uml as its dependency.
+make_uml >> make_png
+```
+
+Once defined, your automation tasks are immediately accessible from the CLI. You can then invoke the tasks by invoking.
+
+```bash
+zrb make-png --diagram "state diagram"
+```
+
+Or you can invoke the tasks without parameter.
+
+```bash
+zrb make-png
+```
+
+At this point, Zrb will politely ask you to provide the diagram type.
+
+```
+diagram [state diagram]:
+```
+
+You can just press enter if you want to use the default value.
+
+Finally, you can run Zrb as a server and make your tasks available for non technical users by invoking the following command.
+
+```bash
+zrb server start
+```
+
+You will have a nice web interface running on `http://localhost:12123`
+
+![Zrb Web UI](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/zrb-web-ui.png)
+
+Now, let's see how Zrb generate the state diagram. Based on the source code in your current directory, Zrb will generate a `state diagram.uml` and transform it into `state diagram.png`.
+
+![State Diagram](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/state-diagram.png)
+
+See the [getting started guide](https://github.com/state-alchemists/zrb/blob/main/docs/recipes/getting-started/README.md) for more information. Or just watch the demo:
+
+[![Video Title](https://img.youtube.com/vi/W7dgk96l__o/0.jpg)](https://www.youtube.com/watch?v=W7dgk96l__o)
+
+
+# 🫰 Installing Zrb
+
+You can install Zrb as a pip package by invoking the following command:
+
+```bash
+pip install --pre zrb
+```
+
+Alternatively, you can also use our installation script to install Zrb along with some prerequisites:
+
+```bash
+bash -c "$(curl -fsSL https://raw.githubusercontent.com/state-alchemists/zrb/main/install.sh)"
+```
+
+# 🐞 Bug Report + Feature Request
+
+You can submit bug reports and feature requests by creating a new [issue](https://github.com/state-alchemists/zrb/issues) on Zrb's GitHub Repositories. When reporting a bug or requesting a feature, please be sure to:
+
+- Include the version of Zrb you are using (i.e., `zrb version`)
+- Tell us what you have tried
+- Tell us what you expect
+- Tell us what you get
+
+We will also welcome your [pull requests and contributions](https://github.com/state-alchemists/zrb/pulls).
+
+
+# ☕ Donation
+
+Help Red Skull to click the donation button:
+
+[![](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/donator.png)](https://stalchmst.com/donation)
+
+# 🎉 Fun Fact
+
+> Madou Ring Zaruba (魔導輪ザルバ, Madōrin Zaruba) is a Madougu which supports bearers of the Garo Armor. [(Garo Wiki | Fandom)](https://garo.fandom.com/wiki/Zaruba)
+
+![Madou Ring Zaruba on Kouga's Hand](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/madou-ring-zaruba.jpg)

{zrb-1.2.1 → zrb-1.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.2.1"
+version = "1.3.0"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"
@@ -50,7 +50,7 @@ beautifulsoup4 = "^4.12.3"
 fastapi = {extras = ["standard"], version = "^0.115.6"}
 python-jose = {extras = ["cryptography"], version = "^3.4.0"}
 ulid-py = "^1.1.0"
-pydantic-ai = "^0.0.19"
+pydantic-ai = "^0.0.31"
 fastembed = "^0.5.1"
 psutil = "^6.1.1"
 

{zrb-1.2.1 → zrb-1.3.0}/src/zrb/builtin/llm/llm_chat.py

@@ -2,22 +2,40 @@ import json
 import os
 from typing import Any
 
+from pydantic_ai.models import Model
+
 from zrb.builtin.group import llm_group
 from zrb.builtin.llm.tool.api import get_current_location, get_current_weather
 from zrb.builtin.llm.tool.cli import run_shell_command
-from zrb.builtin.llm.tool.web import open_web_route, query_internet
+from zrb.builtin.llm.tool.file import (
+    list_file,
+    read_source_code,
+    read_text_file,
+    write_text_file,
+)
+from zrb.builtin.llm.tool.web import (
+    create_search_internet_tool,
+    open_web_page,
+    search_arxiv,
+    search_wikipedia,
+)
 from zrb.config import (
     LLM_ALLOW_ACCESS_INTERNET,
+    LLM_ALLOW_ACCESS_LOCAL_FILE,
     LLM_ALLOW_ACCESS_SHELL,
     LLM_HISTORY_DIR,
     LLM_MODEL,
     LLM_SYSTEM_PROMPT,
+    SERP_API_KEY,
 )
+from zrb.context.any_context import AnyContext
 from zrb.context.any_shared_context import AnySharedContext
+from zrb.input.any_input import AnyInput
 from zrb.input.bool_input import BoolInput
 from zrb.input.str_input import StrInput
 from zrb.input.text_input import TextInput
 from zrb.task.llm_task import LLMTask
+from zrb.util.attr import get_attr
 from zrb.util.file import read_file, write_file
 from zrb.util.string.conversion import to_pascal_case
 
@@ -72,23 +90,53 @@ def _write_chat_conversation(
     write_file(last_session_file_path, current_session_name)
 
 
-llm_chat: LLMTask = llm_group.add_task(
-    LLMTask(
-        name="llm-chat",
-        input=[
+class _LLMChat(LLMTask):
+
+    _default_model: Model | str | None = None
+
+    def set_default_model(self, model: Model | str):
+        self._default_model = model
+
+    @property
+    def inputs(self) -> list[AnyInput]:
+        task_inputs = super().inputs
+        model_input_default = LLM_MODEL if self._default_model is None else "default"
+        return [
             StrInput(
                 "model",
                 description="LLM Model",
                 prompt="LLM Model",
-                default=LLM_MODEL,
+                default=model_input_default,
                 allow_positional_parsing=False,
+                always_prompt=False,
             ),
+            *task_inputs,
+        ]
+
+    def _get_model(self, ctx: AnyContext) -> str | Model | None:
+        if ctx.input.model == "default":
+            if self._default_model is not None:
+                return self._default_model
+            return super()._get_model(ctx)
+        model = get_attr(
+            ctx, ctx.input.model, "ollama_chat/llama3.1", auto_render=self._render_model
+        )
+        if isinstance(model, (Model, str)) or model is None:
+            return model
+        raise ValueError("Invalid model")
+
+
+llm_chat: LLMTask = llm_group.add_task(
+    _LLMChat(
+        name="llm-chat",
+        input=[
             TextInput(
                 "system-prompt",
                 description="System prompt",
                 prompt="System prompt",
                 default=LLM_SYSTEM_PROMPT,
                 allow_positional_parsing=False,
+                always_prompt=False,
             ),
             BoolInput(
                 "start-new",
@@ -96,6 +144,7 @@ llm_chat: LLMTask = llm_group.add_task(
                 prompt="Start new conversation (LLM will forget everything)",
                 default=False,
                 allow_positional_parsing=False,
+                always_prompt=False,
             ),
             TextInput("message", description="User message", prompt="Your message"),
             PreviousSessionInput(
@@ -104,12 +153,12 @@ llm_chat: LLMTask = llm_group.add_task(
                 prompt="Previous conversation session (can be empty)",
                 allow_positional_parsing=False,
                 allow_empty=True,
+                always_prompt=False,
             ),
         ],
         conversation_history_reader=_read_chat_conversation,
         conversation_history_writer=_write_chat_conversation,
         description="Chat with LLM",
-        model="{ctx.input.model}",
         system_prompt="{ctx.input['system-prompt']}",
         message="{ctx.input.message}",
         retries=0,
@@ -117,11 +166,21 @@ llm_chat: LLMTask = llm_group.add_task(
     alias="chat",
 )
 
+
+if LLM_ALLOW_ACCESS_LOCAL_FILE:
+    llm_chat.add_tool(read_source_code)
+    llm_chat.add_tool(list_file)
+    llm_chat.add_tool(read_text_file)
+    llm_chat.add_tool(write_text_file)
+
 if LLM_ALLOW_ACCESS_SHELL:
     llm_chat.add_tool(run_shell_command)
 
 if LLM_ALLOW_ACCESS_INTERNET:
-    llm_chat.add_tool(open_web_route)
-    llm_chat.add_tool(query_internet)
+    llm_chat.add_tool(open_web_page)
+    llm_chat.add_tool(search_wikipedia)
+    llm_chat.add_tool(search_arxiv)
+    if SERP_API_KEY != "":
+        llm_chat.add_tool(create_search_internet_tool(SERP_API_KEY))
     llm_chat.add_tool(get_current_location)
     llm_chat.add_tool(get_current_weather)

{zrb-1.2.1 → zrb-1.3.0}/src/zrb/builtin/llm/tool/api.py

@@ -1,13 +1,13 @@
 import json
 from typing import Annotated, Literal
 
-import requests
-
 
 def get_current_location() -> (
     Annotated[str, "JSON string representing latitude and longitude"]
 ):  # noqa
     """Get the user's current location."""
+    import requests
+
     return json.dumps(requests.get("http://ip-api.com/json?fields=lat,lon").json())
 
 
@@ -17,6 +17,8 @@ def get_current_weather(
     temperature_unit: Literal["celsius", "fahrenheit"],
 ) -> str:
     """Get the current weather in a given location."""
+    import requests
+
     resp = requests.get(
         "https://api.open-meteo.com/v1/forecast",
         params={

zrb-1.3.0/src/zrb/builtin/llm/tool/file.py

@@ -0,0 +1,39 @@
+import os
+
+from zrb.util.file import read_file, write_file
+
+
+def list_file(
+    directory: str = ".",
+    extensions: list[str] = [".py", ".go", ".js", ".ts", ".java", ".c", ".cpp"],
+) -> list[str]:
+    """List all files in a directory"""
+    all_files: list[str] = []
+    for root, _, files in os.walk(directory):
+        for filename in files:
+            for extension in extensions:
+                if filename.lower().endswith(extension):
+                    all_files.append(os.path.join(root, filename))
+    return all_files
+
+
+def read_text_file(file: str) -> str:
+    """Read a text file"""
+    return read_file(os.path.abspath(file))
+
+
+def write_text_file(file: str, content: str):
+    """Write a text file"""
+    return write_file(os.path.abspath(file), content)
+
+
+def read_source_code(
+    directory: str = ".",
+    extensions: list[str] = [".py", ".go", ".js", ".ts", ".java", ".c", ".cpp"],
+) -> list[str]:
+    """Read source code in a directory"""
+    files = list_file(directory, extensions)
+    for index, file in enumerate(files):
+        content = read_text_file(file)
+        files[index] = f"# {file}\n```\n{content}\n```"
+    return files

{zrb-1.2.1 → zrb-1.3.0}/src/zrb/builtin/llm/tool/rag.py

@@ -1,7 +1,9 @@
+import fnmatch
 import hashlib
 import json
 import os
 import sys
+from collections.abc import Callable
 
 import ulid
 
@@ -15,6 +17,20 @@ from zrb.util.cli.style import stylize_error, stylize_faint
 from zrb.util.file import read_file
 
 
+class RAGFileReader:
+    def __init__(self, glob_pattern: str, read: Callable[[str], str]):
+        self.glob_pattern = glob_pattern
+        self.read = read
+
+    def is_match(self, file_name: str):
+        if os.sep not in self.glob_pattern and (
+            os.altsep is None or os.altsep not in self.glob_pattern
+        ):
+            # Pattern like "*.txt" – match only the basename.
+            return fnmatch.fnmatch(os.path.basename(file_name), self.glob_pattern)
+        return fnmatch.fnmatch(file_name, self.glob_pattern)
+
+
 def create_rag_from_directory(
     tool_name: str,
     tool_description: str,
@@ -25,6 +41,7 @@ def create_rag_from_directory(
     chunk_size: int = RAG_CHUNK_SIZE,
     overlap: int = RAG_OVERLAP,
     max_result_count: int = RAG_MAX_RESULT_COUNT,
+    file_reader: list[RAGFileReader] = [],
 ):
     async def retrieve(query: str) -> str:
         from chromadb import PersistentClient
@@ -36,35 +53,31 @@ def create_rag_from_directory(
             path=vector_db_path, settings=Settings(allow_reset=True)
         )
         collection = 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)
         current_hashes = {}
-
+        # Get updated_files
         updated_files = []
-
         for root, _, files in os.walk(document_dir_path):
             for file in files:
                 file_path = os.path.join(root, file)
                 file_hash = _compute_file_hash(file_path)
                 relative_path = os.path.relpath(file_path, document_dir_path)
                 current_hashes[relative_path] = file_hash
-
                 if previous_hashes.get(relative_path) != file_hash:
                     updated_files.append(file_path)
-
+        # Upsert updated_files to vector db
         if updated_files:
             print(
                 stylize_faint(f"Updating {len(updated_files)} changed files"),
                 file=sys.stderr,
             )
-
             for file_path in updated_files:
                 try:
                     relative_path = os.path.relpath(file_path, document_dir_path)
                     collection.delete(where={"file_path": relative_path})
-                    content = _read_file_content(file_path)
+                    content = _read_txt_content(file_path, file_reader)
                     file_id = ulid.new().str
                     for i in range(0, len(content), chunk_size - overlap):
                         chunk = content[i : i + chunk_size]
@@ -92,14 +105,13 @@ def create_rag_from_directory(
                         stylize_error(f"Error processing {file_path}: {e}"),
                         file=sys.stderr,
                     )
-
             _save_hashes(hash_file_path, current_hashes)
         else:
             print(
                 stylize_faint("No changes detected. Skipping database update."),
                 file=sys.stderr,
             )
-
+        # 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]
@@ -123,7 +135,22 @@ def _compute_file_hash(file_path: str) -> str:
     return hash_md5.hexdigest()
 
 
-def _read_file_content(file_path: str) -> str:
+def _load_hashes(file_path: str) -> dict:
+    if os.path.exists(file_path):
+        with open(file_path, "r") as f:
+            return json.load(f)
+    return {}
+
+
+def _save_hashes(file_path: str, hashes: dict):
+    with open(file_path, "w") as f:
+        json.dump(hashes, f)
+
+
+def _read_txt_content(file_path: str, file_reader: list[RAGFileReader]):
+    for reader in file_reader:
+        if reader.is_match(file_path):
+            return reader.read(file_path)
     if file_path.lower().endswith(".pdf"):
         return _read_pdf(file_path)
     return read_file(file_path)
@@ -136,15 +163,3 @@ def _read_pdf(file_path: str) -> str:
         return "\n".join(
             page.extract_text() for page in pdf.pages if page.extract_text()
         )
-
-
-def _load_hashes(file_path: str) -> dict:
-    if os.path.exists(file_path):
-        with open(file_path, "r") as f:
-            return json.load(f)
-    return {}
-
-
-def _save_hashes(file_path: str, hashes: dict):
-    with open(file_path, "w") as f:
-        json.dump(hashes, f)