zrb 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

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/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_template/my_app_name/_zrb/column/add_column_util.py

@@ -2,7 +2,7 @@ import os
 import re
 import textwrap
 
-from bs4 import BeautifulSoup, formatter
+from bs4 import BeautifulSoup, Tag, formatter
 from my_app_name._zrb.config import APP_DIR
 
 from zrb.context.any_context import AnyContext
@@ -65,6 +65,7 @@ def update_my_app_name_ui(ctx: AnyContext):
     kebab_entity_name = to_kebab_case(ctx.input.entity)
     snake_column_name = to_snake_case(ctx.input.column)
     human_column_name = to_human_case(ctx.input.column).title()
+    column_type = ctx.input.type
     subroute_file_path = os.path.join(
         APP_DIR,
         "module",
@@ -85,18 +86,21 @@ def update_my_app_name_ui(ctx: AnyContext):
         form_id="crud-create-form",
         column_label=human_column_name,
         column_name=snake_column_name,
+        column_type=column_type,
     )
     new_code = _add_input_to_form(
         new_code,
         form_id="crud-update-form",
         column_label=human_column_name,
         column_name=snake_column_name,
+        column_type=column_type,
     )
     new_code = _add_input_to_form(
         new_code,
         form_id="crud-delete-form",
         column_label=human_column_name,
         column_name=snake_column_name,
+        column_type=column_type,
     )
     # JS Function
     new_code = _alter_js_function_returned_array(
@@ -140,7 +144,7 @@ def _add_th_before_last(html_str, table_id, th_content):
 
 
 def _add_input_to_form(
-    html_str: str, form_id: str, column_label: str, column_name: str
+    html_str: str, form_id: str, column_label: str, column_name: str, column_type: str
 ) -> str:
     soup = BeautifulSoup(html_str, "html.parser")
     # Find the form by id.
@@ -151,9 +155,7 @@ def _add_input_to_form(
     new_label = soup.new_tag("label")
     new_label.append(f"{column_label}: ")
     # Create a new input element with the provided column name.
-    new_input = soup.new_tag(
-        "input", attrs={"type": "text", "name": column_name, "required": "required"}
-    )
+    new_input = _get_html_input(soup, column_name, column_type)
     new_label.append(new_input)
     # Look for a footer element inside the form.
     footer = form.find("footer")
@@ -168,6 +170,26 @@ def _add_input_to_form(
     )
 
 
+def _get_html_input(soup: BeautifulSoup, column_name: str, column_type: str) -> Tag:
+    # Map your custom types to HTML input types.
+    type_mapping = {
+        "str": "text",
+        "int": "number",
+        "float": "number",
+        "bool": "checkbox",
+        "datetime": "datetime-local",
+        "date": "date",
+    }
+    # Get the corresponding HTML input type; default to "text" if not found.
+    html_input_type = type_mapping.get(column_type, "text")
+    # Create the new input tag with the appropriate attributes.
+    new_input = soup.new_tag(
+        "input",
+        attrs={"type": html_input_type, "name": column_name, "required": "required"},
+    )
+    return new_input
+
+
 def _infer_html_indent_width(html_str: str) -> int:
     """
     Infer the indentation width (number of spaces) from the HTML string.
@@ -208,7 +230,7 @@ def _alter_js_function_returned_array(
     #  3. Captures the newline and leading whitespace (indent) of the return statement.
     #  4. Captures the rest of the return line.
     pattern = (
-        r"(function\s+"
+        r"("
         + re.escape(js_function_name)
         + r"\s*\([^)]*\)\s*\{)"  # group1: function header
         r"([\s\S]*?)"  # group2: code before return