zrb 1.2.0__py3-none-any.whl → 1.2.2__py3-none-any.whl

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/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/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/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/builtin/llm/tool/web.py

@@ -1,8 +1,9 @@
 import json
+from collections.abc import Callable
 from typing import Annotated
 
 
-def open_web_route(url: str) -> str:
+def open_web_page(url: str) -> str:
     """Get content from a web page."""
     import requests
 
@@ -19,30 +20,55 @@ def open_web_route(url: str) -> str:
     return json.dumps(parse_html_text(response.text))
 
 
-def query_internet(
+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 factual information from the internet by using Google."""
+    """Search on Arxiv"""
     import requests
 
-    response = requests.get(
-        "https://google.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",
-        },
-    )
-    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))
+    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]:

zrb/builtin/project/add/fastapp/fastapp_task.py

@@ -62,7 +62,7 @@ scaffold_fastapp = Scaffolder(
         ),
         # Register fastapp's tasks to project's zrb_init (project_dir/zrb_init.py)
         ContentTransformer(
-            name="trasnform-zrb-init",
+            name="transform-zrb-init",
             match=is_project_zrb_init_file,
             transform=update_project_zrb_init_file,
         ),

zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/_zrb/column/add_column_task.py

@@ -1,12 +1,12 @@
 import os
 
 from my_app_name._zrb.column.add_column_util import (
-    update_fastapp_schema,
-    update_fastapp_test_create,
-    update_fastapp_test_delete,
-    update_fastapp_test_read,
-    update_fastapp_test_update,
-    update_fastapp_ui,
+    update_my_app_name_schema,
+    update_my_app_name_test_create,
+    update_my_app_name_test_delete,
+    update_my_app_name_test_read,
+    update_my_app_name_test_update,
+    update_my_app_name_ui,
 )
 from my_app_name._zrb.config import APP_DIR
 from my_app_name._zrb.format_task import format_my_app_name_code
@@ -23,14 +23,14 @@ from zrb import AnyContext, Task, make_task
 
 
 @make_task(
-    name="validate-add-fastapp-column",
+    name="validate-add-my-app-name-column",
     input=[
         existing_module_input,
         existing_entity_input,
     ],
     retries=0,
 )
-async def validate_add_fastapp_column(ctx: AnyContext):
+async def validate_add_my_app_name_column(ctx: AnyContext):
     module_name = ctx.input.module
     if module_name not in get_existing_module_names():
         raise ValueError(f"Module not exist: {module_name}")
@@ -39,96 +39,96 @@ async def validate_add_fastapp_column(ctx: AnyContext):
         raise ValueError(f"Schema not exist: {schema_name}")
 
 
-update_fastapp_schema_task = Task(
-    name="update-fastapp-schema",
+update_my_app_name_schema_task = Task(
+    name="update-my-app-name-schema",
     input=[
         existing_module_input,
         existing_entity_input,
         new_column_input,
         new_column_type_input,
     ],
-    action=update_fastapp_schema,
+    action=update_my_app_name_schema,
     retries=0,
-    upstream=validate_add_fastapp_column,
+    upstream=validate_add_my_app_name_column,
 )
 
-update_fastapp_ui_task = Task(
-    name="update-fastapp-ui",
+update_my_app_name_ui_task = Task(
+    name="update-my-app-name-ui",
     input=[
         existing_module_input,
         existing_entity_input,
         new_column_input,
         new_column_type_input,
     ],
-    action=update_fastapp_ui,
+    action=update_my_app_name_ui,
     retries=0,
-    upstream=validate_add_fastapp_column,
+    upstream=validate_add_my_app_name_column,
 )
 
-update_fastapp_test_create_task = Task(
-    name="update-fastapp-test-create",
+update_my_app_name_test_create_task = Task(
+    name="update-my-app-name-test-create",
     input=[
         existing_module_input,
         existing_entity_input,
         new_column_input,
         new_column_type_input,
     ],
-    action=update_fastapp_test_create,
+    action=update_my_app_name_test_create,
     retries=0,
-    upstream=validate_add_fastapp_column,
+    upstream=validate_add_my_app_name_column,
 )
 
-update_fastapp_test_read_task = Task(
-    name="update-fastapp-test-read",
+update_my_app_name_test_read_task = Task(
+    name="update-my-app-name-test-read",
     input=[
         existing_module_input,
         existing_entity_input,
         new_column_input,
         new_column_type_input,
     ],
-    action=update_fastapp_test_read,
+    action=update_my_app_name_test_read,
     retries=0,
-    upstream=validate_add_fastapp_column,
+    upstream=validate_add_my_app_name_column,
 )
 
-update_fastapp_test_update_task = Task(
-    name="update-fastapp-test-update",
+update_my_app_name_test_update_task = Task(
+    name="update-my-app-name-test-update",
     input=[
         existing_module_input,
         existing_entity_input,
         new_column_input,
         new_column_type_input,
     ],
-    action=update_fastapp_test_update,
+    action=update_my_app_name_test_update,
     retries=0,
-    upstream=validate_add_fastapp_column,
+    upstream=validate_add_my_app_name_column,
 )
 
-update_fastapp_test_delete_task = Task(
-    name="update-fastapp-test-delete",
+update_my_app_name_test_delete_task = Task(
+    name="update-my-app-name-test-delete",
     input=[
         existing_module_input,
         existing_entity_input,
         new_column_input,
         new_column_type_input,
     ],
-    action=update_fastapp_test_delete,
+    action=update_my_app_name_test_delete,
     retries=0,
-    upstream=validate_add_fastapp_column,
+    upstream=validate_add_my_app_name_column,
 )
 
 
-add_fastapp_column = app_create_group.add_task(
+add_my_app_name_column = app_create_group.add_task(
     Task(
-        name="add-fastapp-column",
+        name="add-my-app-name-column",
         description="📊 Create new column on an entity",
         upstream=[
-            update_fastapp_schema_task,
-            update_fastapp_ui_task,
-            update_fastapp_test_create_task,
-            update_fastapp_test_read_task,
-            update_fastapp_test_update_task,
-            update_fastapp_test_delete_task,
+            update_my_app_name_schema_task,
+            update_my_app_name_ui_task,
+            update_my_app_name_test_create_task,
+            update_my_app_name_test_read_task,
+            update_my_app_name_test_update_task,
+            update_my_app_name_test_delete_task,
         ],
         successor=format_my_app_name_code,
         retries=0,

zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/_zrb/column/add_column_util.py

@@ -6,11 +6,7 @@ from bs4 import BeautifulSoup, formatter
 from my_app_name._zrb.config import APP_DIR
 
 from zrb.context.any_context import AnyContext
-from zrb.util.codemod.modify_class import append_code_to_class
-from zrb.util.codemod.modify_class_parent import prepend_parent_class
 from zrb.util.codemod.modify_class_property import append_property_to_class
-from zrb.util.codemod.modify_function import append_code_to_function
-from zrb.util.codemod.modify_module import prepend_code_to_module
 from zrb.util.file import read_file, write_file
 from zrb.util.string.conversion import (
     to_human_case,
@@ -20,7 +16,7 @@ from zrb.util.string.conversion import (
 )
 
 
-def update_fastapp_schema(ctx: AnyContext):
+def update_my_app_name_schema(ctx: AnyContext):
     snake_entity_name = to_snake_case(ctx.input.entity)
     pascal_entity_name = to_pascal_case(ctx.input.entity)
     snake_column_name = to_snake_case(ctx.input.column)
@@ -64,7 +60,7 @@ def _get_default_column_value(data_type: str) -> str:
     return "None"
 
 
-def update_fastapp_ui(ctx: AnyContext):
+def update_my_app_name_ui(ctx: AnyContext):
     kebab_module_name = to_kebab_case(ctx.input.module)
     kebab_entity_name = to_kebab_case(ctx.input.entity)
     snake_column_name = to_snake_case(ctx.input.column)
@@ -237,7 +233,7 @@ def _alter_js_function_returned_array(
     return new_html
 
 
-def update_fastapp_test_create(ctx: AnyContext):
+def update_my_app_name_test_create(ctx: AnyContext):
     snake_module_name = to_snake_case(ctx.input.module)
     snake_entity_name = to_snake_case(ctx.input.entity)
     test_file_path = os.path.join(
@@ -253,7 +249,7 @@ def update_fastapp_test_create(ctx: AnyContext):
     write_file(test_file_path, new_code)
 
 
-def update_fastapp_test_read(ctx: AnyContext):
+def update_my_app_name_test_read(ctx: AnyContext):
     snake_module_name = to_snake_case(ctx.input.module)
     snake_entity_name = to_snake_case(ctx.input.entity)
     test_file_path = os.path.join(
@@ -269,7 +265,7 @@ def update_fastapp_test_read(ctx: AnyContext):
     write_file(test_file_path, new_code)
 
 
-def update_fastapp_test_update(ctx: AnyContext):
+def update_my_app_name_test_update(ctx: AnyContext):
     snake_module_name = to_snake_case(ctx.input.module)
     snake_entity_name = to_snake_case(ctx.input.entity)
     test_file_path = os.path.join(
@@ -285,7 +281,7 @@ def update_fastapp_test_update(ctx: AnyContext):
     write_file(test_file_path, new_code)
 
 
-def update_fastapp_test_delete(ctx: AnyContext):
+def update_my_app_name_test_delete(ctx: AnyContext):
     snake_module_name = to_snake_case(ctx.input.module)
     snake_entity_name = to_snake_case(ctx.input.entity)
     test_file_path = os.path.join(