zrb 1.2.1__tar.gz → 1.2.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.2.1
+Version: 1.2.2
 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)

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

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.2.1"
+version = "1.2.2"
 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.2.2}/src/zrb/builtin/llm/llm_chat.py

@@ -5,13 +5,26 @@ from typing import Any
 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_shared_context import AnySharedContext
 from zrb.input.bool_input import BoolInput
@@ -117,11 +130,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.2.2}/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.2.2/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.2.2}/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)

zrb-1.2.2/src/zrb/builtin/llm/tool/web.py

@@ -0,0 +1,100 @@
+import json
+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))
+
+
+def create_search_internet_tool(serp_api_key: str) -> Callable[[str, int], str]:
+    def search_internet(
+        query: Annotated[str, "Search query"],
+        num_results: Annotated[int, "Search result count, by default 10"] = 10,
+    ) -> str:
+        """Search factual information from the internet by using Google."""
+        import requests
+
+        response = requests.get(
+            "https://serpapi.com/search",
+            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
+            },
+            params={
+                "q": query,
+                "num": num_results,
+                "hl": "en",
+                "safe": "off",
+                "api_key": serp_api_key,
+            },
+        )
+        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))
+
+    return search_internet
+
+
+def search_wikipedia(query: Annotated[str, "Search query"]) -> str:
+    """Search on wikipedia"""
+    import requests
+
+    params = {"action": "query", "list": "search", "srsearch": query, "format": "json"}
+    response = requests.get("https://en.wikipedia.org/w/api.php", params=params)
+    return response.json()
+
+
+def search_arxiv(
+    query: Annotated[str, "Search query"],
+    num_results: Annotated[int, "Search result count, by default 10"] = 10,
+) -> str:
+    """Search on Arxiv"""
+    import requests
+
+    params = {"search_query": f"all:{query}", "start": 0, "max_results": num_results}
+    response = requests.get("http://export.arxiv.org/api/query", params=params)
+    return response.content
+
+
+def parse_html_text(html_text: str) -> dict[str, str]:
+    from bs4 import BeautifulSoup
+
+    ignored_tags = [
+        "script",
+        "link",
+        "meta",
+        "style",
+        "code",
+        "footer",
+        "nav",
+        "header",
+        "aside",
+    ]
+    soup = BeautifulSoup(html_text, "html.parser")
+    links = []
+    for anchor in soup.find_all("a"):
+        if not anchor or "href" not in anchor.attrs:
+            continue
+        link: str = anchor["href"]
+        if link.startswith("#") or link.startswith("/"):
+            continue
+        links.append(link)
+    for tag in soup(ignored_tags):
+        tag.decompose()
+    html_text = soup.get_text(separator=" ", strip=True)
+    return {"content": html_text, "links_on_page": links}

{zrb-1.2.1 → zrb-1.2.2}/src/zrb/config.py

@@ -92,7 +92,8 @@ LLM_HISTORY_DIR = os.getenv(
 LLM_HISTORY_FILE = os.getenv(
     "ZRB_LLM_HISTORY_FILE", os.path.join(LLM_HISTORY_DIR, "history.json")
 )
-LLM_ALLOW_ACCESS_SHELL = to_boolean(os.getenv("ZRB_LLM_ACCESS_FILE", "1"))
+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(
@@ -101,6 +102,7 @@ RAG_EMBEDDING_MODEL = os.getenv(
 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"))
+SERP_API_KEY = os.getenv("SERP_API_KEY", "")
 
 
 BANNER = f"""

{zrb-1.2.1 → zrb-1.2.2}/src/zrb/content_transformer/content_transformer.py

@@ -1,4 +1,5 @@
 import fnmatch
+import os
 import re
 from collections.abc import Callable
 
@@ -40,7 +41,12 @@ class ContentTransformer(AnyContentTransformer):
                     return True
             except re.error:
                 pass
-            return fnmatch.fnmatch(file_path, pattern)
+            if os.sep not in pattern and (
+                os.altsep is None or os.altsep not in pattern
+            ):
+                # Pattern like "*.txt" – match only the basename.
+                return fnmatch.fnmatch(file_path, os.path.basename(file_path))
+            return fnmatch.fnmatch(file_path, file_path)
 
     def transform_file(self, ctx: AnyContext, file_path: str):
         if callable(self._transform_file):

{zrb-1.2.1 → zrb-1.2.2}/src/zrb/context/context.py

@@ -88,7 +88,7 @@ class Context(AnyContext):
             return template
         return int(self.render(template))
 
-    def render_float(self, template: str) -> float:
+    def render_float(self, template: str | float) -> float:
         if isinstance(template, float):
             return template
         return float(self.render(template))
@@ -102,9 +102,10 @@ class Context(AnyContext):
         flush: bool = True,
         plain: bool = False,
     ):
+        sep = " " if sep is None else sep
         message = sep.join([f"{value}" for value in values])
         if plain:
-            self.append_to_shared_log(remove_style(message))
+            # self.append_to_shared_log(remove_style(message))
             print(message, sep=sep, end=end, file=file, flush=flush)
             return
         color = self._color
@@ -132,6 +133,7 @@ class Context(AnyContext):
         flush: bool = True,
     ):
         if self._shared_ctx.get_logging_level() <= logging.DEBUG:
+            sep = " " if sep is None else sep
             message = sep.join([f"{value}" for value in values])
             stylized_message = stylize_log(f"[DEBUG] {message}")
             self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
@@ -145,6 +147,7 @@ class Context(AnyContext):
         flush: bool = True,
     ):
         if self._shared_ctx.get_logging_level() <= logging.INFO:
+            sep = " " if sep is None else sep
             message = sep.join([f"{value}" for value in values])
             stylized_message = stylize_log(f"[INFO] {message}")
             self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
@@ -158,6 +161,7 @@ class Context(AnyContext):
         flush: bool = True,
     ):
         if self._shared_ctx.get_logging_level() <= logging.INFO:
+            sep = " " if sep is None else sep
             message = sep.join([f"{value}" for value in values])
             stylized_message = stylize_warning(f"[WARNING] {message}")
             self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
@@ -171,6 +175,7 @@ class Context(AnyContext):
         flush: bool = True,
     ):
         if self._shared_ctx.get_logging_level() <= logging.ERROR:
+            sep = " " if sep is None else sep
             message = sep.join([f"{value}" for value in values])
             stylized_message = stylize_error(f"[ERROR] {message}")
             self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)
@@ -184,6 +189,7 @@ class Context(AnyContext):
         flush: bool = True,
     ):
         if self._shared_ctx.get_logging_level() <= logging.CRITICAL:
+            sep = " " if sep is None else sep
             message = sep.join([f"{value}" for value in values])
             stylized_message = stylize_error(f"[CRITICAL] {message}")
             self.print(stylized_message, sep=sep, end=end, file=file, flush=flush)

{zrb-1.2.1 → zrb-1.2.2}/src/zrb/input/text_input.py

@@ -69,16 +69,17 @@ class TextInput(BaseInput):
         )
 
     def _prompt_cli_str(self, shared_ctx: AnySharedContext) -> str:
-        prompt_message = (
-            f"{self.comment_start}{super().prompt_message}{self.comment_end}"
+        prompt_message = super().prompt_message
+        comment_prompt_message = (
+            f"{self.comment_start}{prompt_message}{self.comment_end}"
         )
-        prompt_message_eol = f"{prompt_message}\n"
+        comment_prompt_message_eol = f"{comment_prompt_message}\n"
         default_value = self.get_default_str(shared_ctx)
         with tempfile.NamedTemporaryFile(
             delete=False, suffix=self._extension
         ) as temp_file:
             temp_file_name = temp_file.name
-            temp_file.write(prompt_message_eol.encode())
+            temp_file.write(comment_prompt_message_eol.encode())
             # Pre-fill with default content
             if default_value:
                 temp_file.write(default_value.encode())
@@ -87,7 +88,10 @@ class TextInput(BaseInput):
         subprocess.call([self._editor, temp_file_name])
         # Read the edited content
         edited_content = read_file(temp_file_name)
-        parts = [text.strip() for text in edited_content.split(prompt_message, 1)]
+        parts = [
+            text.strip() for text in edited_content.split(comment_prompt_message, 1)
+        ]
         edited_content = "\n".join(parts).lstrip()
         os.remove(temp_file_name)
+        print(f"{prompt_message}: {edited_content}")
         return edited_content