zrb 1.8.10__py3-none-any.whl → 1.21.29__py3-none-any.whl

zrb/builtin/llm/tool/note.py

@@ -0,0 +1,84 @@
+import os
+
+from zrb.config.llm_context.config import llm_context_config
+
+
+def read_long_term_note() -> str:
+    """
+    Retrieves the GLOBAL long-term memory shared across ALL sessions and projects.
+
+    CRITICAL: Consult this first for user preferences, facts, and cross-project context.
+
+    Returns:
+        str: The current global note content.
+    """
+    contexts = llm_context_config.get_notes()
+    return contexts.get("/", "")
+
+
+def write_long_term_note(content: str) -> str:
+    """
+    Persists CRITICAL facts to the GLOBAL long-term memory.
+
+    USE EAGERLY to save or update:
+    - User preferences (e.g., "I prefer Python", "No unit tests").
+    - User information (e.g., user name, user email address).
+    - Important facts (e.g., "My API key is in .env").
+    - Cross-project goals.
+    - Anything that will be useful for future interaction across projects.
+
+    WARNING: This OVERWRITES the entire global note.
+
+    Args:
+        content (str): The text to strictly memorize.
+
+    Returns:
+        str: Confirmation message.
+    """
+    llm_context_config.write_note(content, "/")
+    return "Global long-term note saved."
+
+
+def read_contextual_note(path: str | None = None) -> str:
+    """
+    Retrieves LOCAL memory specific to a file or directory path.
+
+    Use to recall project-specific architecture, code summaries, or past decisions
+    relevant to the current working location.
+
+    Args:
+        path (str | None): Target file/dir. Defaults to current working directory (CWD).
+
+    Returns:
+        str: The local note content for the path.
+    """
+    if path is None:
+        path = os.getcwd()
+    abs_path = os.path.abspath(path)
+    contexts = llm_context_config.get_notes(cwd=abs_path)
+    return contexts.get(abs_path, "")
+
+
+def write_contextual_note(content: str, path: str | None = None) -> str:
+    """
+    Persists LOCAL facts specific to a file or directory.
+
+    USE EAGERLY to save or update:
+    - Architectural patterns for this project/directory.
+    - Summaries of large files or directories.
+    - Specific guidelines for this project.
+    - Anything related to this directory that will be useful for future interaction.
+
+    WARNING: This OVERWRITES the note for the specific path.
+
+    Args:
+        content (str): The text to memorize for this location.
+        path (str | None): Target file/dir. Defaults to CWD.
+
+    Returns:
+        str: Confirmation message.
+    """
+    if path is None:
+        path = os.getcwd()
+    llm_context_config.write_note(content, path)
+    return f"Contextual note saved for: {path}"

zrb/builtin/llm/tool/rag.py

@@ -5,10 +5,11 @@ import os
 import sys
 from collections.abc import Callable
 from textwrap import dedent
+from typing import Any
 
 import ulid
 
-from zrb.config import CFG
+from zrb.config.config import CFG
 from zrb.util.cli.style import stylize_error, stylize_faint
 from zrb.util.file import read_file
 
@@ -43,13 +44,41 @@ def create_rag_from_directory(
     openai_base_url: str | None = None,
     openai_embedding_model: str | None = None,
 ):
-    """Create a RAG retrieval tool function for LLM use.
-    This factory configures and returns an async function that takes a query,
-    updates a vector database if needed, performs a similarity search,
-    and returns relevant document chunks.
+    """
+    Create a powerful RAG (Retrieval-Augmented Generation) tool for querying a local
+    knowledge base.
+
+    This factory function generates a tool that performs semantic search over a directory of
+    documents. It automatically indexes the documents into a vector database (ChromaDB) and
+    keeps it updated as files change.
+
+    The generated tool is ideal for answering questions based on a specific set of documents,
+    such as project documentation or internal wikis.
+
+    Args:
+        tool_name (str): The name for the generated RAG tool (e.g., "search_project_docs").
+        tool_description (str): A clear description of what the tool does and when to use it.
+            This is what the LLM will see.
+        document_dir_path (str, optional): The path to the directory containing the documents
+            to be indexed.
+        vector_db_path (str, optional): The path where the ChromaDB vector database will be
+            stored.
+        vector_db_collection (str, optional): The name of the collection within the vector
+            database.
+        chunk_size (int, optional): The size of text chunks for embedding.
+        overlap (int, optional): The overlap between text chunks.
+        max_result_count (int, optional): The maximum number of search results to return.
+        file_reader (list[RAGFileReader], optional): A list of custom file readers for
+            specific file types.
+        openai_api_key (str, optional): Your OpenAI API key for generating embeddings.
+        openai_base_url (str, optional): An optional base URL for the OpenAI API.
+        openai_embedding_model (str, optional): The embedding model to use.
+
+    Returns:
+        An asynchronous function that serves as the RAG tool.
     """
 
-    async def retrieve(query: str) -> str:
+    async def retrieve(query: str) -> dict[str, Any]:
         # Docstring will be set dynamically below
         from chromadb import PersistentClient
         from chromadb.config import Settings
@@ -164,7 +193,7 @@ def create_rag_from_directory(
             query_embeddings=query_vector,
             n_results=max_result_count_val,
         )
-        return json.dumps(results)
+        return dict(results)
 
     retrieve.__name__ = tool_name
     retrieve.__doc__ = dedent(
@@ -173,7 +202,8 @@ def create_rag_from_directory(
         Args:
             query (str): The user query to search for in documents.
         Returns:
-            str: JSON string with search results: {{"ids": [...], "documents": [...], ...}}
+            dict[str, Any]: dictionary with search results:
+                {{"ids": [...], "documents": [...], ...}}
         """
     ).strip()
     return retrieve

zrb/builtin/llm/tool/sub_agent.py

@@ -3,60 +3,70 @@ from collections.abc import Callable
 from textwrap import dedent
 from typing import TYPE_CHECKING, Any, Coroutine
 
+from zrb.context.any_context import AnyContext
+from zrb.task.llm.agent import create_agent_instance
+from zrb.task.llm.agent_runner import run_agent_iteration
+from zrb.task.llm.config import get_model, get_model_settings
+from zrb.task.llm.prompt import get_system_and_user_prompt
+from zrb.task.llm.subagent_conversation_history import (
+    get_ctx_subagent_history,
+    set_ctx_subagent_history,
+)
+
 if TYPE_CHECKING:
     from pydantic_ai import Tool
-    from pydantic_ai.mcp import MCPServer
+    from pydantic_ai._agent_graph import HistoryProcessor
     from pydantic_ai.models import Model
     from pydantic_ai.settings import ModelSettings
-else:
-    Tool = Any
-    MCPServer = Any
-    Model = Any
-    ModelSettings = Any
+    from pydantic_ai.toolsets import AbstractToolset
 
-from zrb.context.any_context import AnyContext
-from zrb.task.llm.agent import create_agent_instance, run_agent_iteration
-from zrb.task.llm.config import get_model, get_model_settings
-from zrb.task.llm.prompt import get_combined_system_prompt
-
-if TYPE_CHECKING:
     ToolOrCallable = Tool | Callable
-else:
-    ToolOrCallable = Any
 
 
 def create_sub_agent_tool(
     tool_name: str,
     tool_description: str,
     system_prompt: str | None = None,
-    model: str | Model | None = None,
-    model_settings: ModelSettings | None = None,
-    tools: list[ToolOrCallable] = [],
-    mcp_servers: list[MCPServer] = [],
-) -> Callable[[AnyContext, str], Coroutine[Any, Any, str]]:
+    model: "str | Model | None" = None,
+    model_settings: "ModelSettings | None" = None,
+    tools: "list[ToolOrCallable]" = [],
+    toolsets: list["AbstractToolset[None]"] = [],
+    yolo_mode: bool | list[str] | None = None,
+    history_processors: list["HistoryProcessor"] | None = None,
+    log_indent_level: int = 2,
+    agent_name: str | None = None,
+    auto_summarize: bool = True,
+    remember_history: bool = True,
+) -> Callable[[AnyContext, str], Coroutine[Any, Any, Any]]:
     """
-    Create an LLM "sub-agent" tool function for use by a main LLM agent.
+    Create a tool that is another AI agent, capable of handling complex, multi-step sub-tasks.
 
-    This factory configures and returns an async function that, when called
-    by the main agent, instantiates and runs a sub-agent (the sub-agent)
-    with a given query and returns the sub-agent's final response.
+    This factory function generates a tool that, when used, spins up a temporary, specialized
+    AI agent. This "sub-agent" has its own system prompt, tools, and context, allowing it to
+    focus on accomplishing a specific task without being distracted by the main conversation.
+
+    This is ideal for delegating complex tasks like analyzing a file or a repository.
 
     Args:
-        tool_name: The name of the tool for the main agent.
-        tool_description: The description of the tool for the main agent.
-        sub_agent_system_prompt: The system prompt for the sub-agent.
-        sub_agent_model: The model for the sub-agent (optional).
-        sub_agent_model_settings: Model settings for the sub-agent (optional).
-        sub_agent_tools: A list of tools (Tool instances or callables) for the
-            sub-agent (optional).
-        sub_agent_mcp_servers: A list of MCP servers for the sub-agent (optional).
+        tool_name (str): The name for the generated sub-agent tool.
+        tool_description (str): A clear description of the sub-agent's purpose and when to
+            use it. This is what the LLM will see.
+        system_prompt (str, optional): The system prompt that will guide the sub-agent's
+            behavior.
+        model (str | Model, optional): The language model the sub-agent will use.
+        model_settings (ModelSettings, optional): Specific settings for the sub-agent's model.
+        tools (list, optional): A list of tools that will be exclusively available to the
+            sub-agent.
+        toolsets (list, optional): A list of Toolsets for the sub-agent.
 
     Returns:
-        An async callable function that takes a context and a query string,
-        runs the sub-agent, and returns the sub-agent's final message content.
+        An asynchronous function that serves as the sub-agent tool. When called, it runs the
+        sub-agent with a given query and returns its final result.
     """
+    if agent_name is None:
+        agent_name = f"{tool_name}_agent"
 
-    async def run_sub_agent(ctx: AnyContext, query: str) -> str:
+    async def run_sub_agent(ctx: AnyContext, query: str) -> Any:
         """
         Runs the sub-agent with the given query.
         """
@@ -75,16 +85,13 @@ def create_sub_agent_tool(
             ctx=ctx,
             model_settings_attr=model_settings,
         )
-
         if system_prompt is None:
-            resolved_system_prompt = get_combined_system_prompt(
+            resolved_system_prompt, query = get_system_and_user_prompt(
                 ctx=ctx,
+                user_message=query,
                 persona_attr=None,
-                render_persona=False,
                 system_prompt_attr=None,
-                render_system_prompt=False,
                 special_instruction_prompt_attr=None,
-                render_special_instruction_prompt=False,
             )
         else:
             resolved_system_prompt = system_prompt
@@ -95,26 +102,36 @@ def create_sub_agent_tool(
             system_prompt=resolved_system_prompt,
             model_settings=resolved_model_settings,
             tools=tools,
-            mcp_servers=mcp_servers,
+            toolsets=toolsets,
+            yolo_mode=yolo_mode,
+            history_processors=history_processors,
+            auto_summarize=auto_summarize,
         )
-
         sub_agent_run = None
         # Run the sub-agent iteration
-        # Start with an empty history for the sub-agent
+        history_list = (
+            get_ctx_subagent_history(ctx, agent_name) if remember_history else []
+        )
         sub_agent_run = await run_agent_iteration(
             ctx=ctx,
             agent=sub_agent_agent,
             user_prompt=query,
-            history_list=[],  # Start with empty history for the sub-agent
+            attachments=[],
+            history_list=history_list,
+            log_indent_level=log_indent_level,
         )
-
         # Return the sub-agent's final message content
         if sub_agent_run and sub_agent_run.result:
-            # Return the final message content as a string
-            return json.dumps({"result": sub_agent_run.result.output})
-        else:
-            ctx.log_warning("Sub-agent run did not produce a result.")
-            return "Sub-agent failed to produce a result."
+            # Return the final message content
+            if remember_history:
+                set_ctx_subagent_history(
+                    ctx,
+                    agent_name,
+                    json.loads(sub_agent_run.result.all_messages_json()),
+                )
+            return sub_agent_run.result.output
+        ctx.log_warning("Sub-agent run did not produce a result.")
+        raise ValueError(f"{tool_name} not returning any result")
 
     # Set the name and docstring for the callable function
     run_sub_agent.__name__ = tool_name
@@ -126,7 +143,7 @@ def create_sub_agent_tool(
             query (str): The query or task for the sub-agent.
 
         Returns:
-            str: The final response or result from the sub-agent.
+            Any: The final response or result from the sub-agent.
         """
     ).strip()
 

zrb/builtin/llm/tool/web.py

@@ -1,147 +1,171 @@
-import json
 from collections.abc import Callable
+from typing import Any
+from urllib.parse import urljoin
 
+from zrb.config.config import CFG
+from zrb.config.llm_config import llm_config
+
+_DEFAULT_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
+
+
+async def open_web_page(url: str) -> dict[str, Any]:
+    """
+    Fetches, parses, and converts a web page to readable Markdown.
+    Preserves semantic structure, removes non-essentials, and extracts all absolute links.
+
+    Example:
+    open_web_page(url='https://www.example.com/article')
 
-async def open_web_page(url: str) -> str:
-    """Get parsed text content and links from a web page URL.
     Args:
-        url (str): The URL of the web page to open.
+        url (str): The full URL of the web page.
+
     Returns:
-        str: JSON: {"content": "parsed text content", "links_on_page": ["url1", ...]}
+        dict: Markdown content and a list of absolute links.
     """
+    html_content, links = await _fetch_page_content(url)
+    markdown_content = _convert_html_to_markdown(html_content)
+    return {"content": markdown_content, "links_on_page": links}
+
+
+def create_search_internet_tool() -> Callable:
+    if llm_config.default_search_internet_tool is not None:
+        return llm_config.default_search_internet_tool
+
+    def search_internet(query: str, page: int = 1) -> dict[str, Any]:
+        """
+        Performs an internet search using a search engine.
+        Use to find information, answer general knowledge, or research topics.
+
+        Example:
+        search_internet(query='latest AI advancements', page=1)
 
-    async def get_page_content(page_url: str):
-        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
-        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()
-                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": user_agent})
-            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]:
-    def search_internet(query: str, num_results: int = 10) -> str:
-        """Search the internet using SerpApi (Google Search) and return parsed results.
         Args:
-            query (str): Search query.
-            num_results (int): Search result count. Defaults to 10.
+            query (str): The search query.
+            page (int, optional): Search result page number. Defaults to 1.
+
         Returns:
-            str: JSON: {"content": "parsed text content", "links_on_page": ["url1", ...]}
+            dict: Summary of search results (titles, links, snippets).
         """
         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 (
+            CFG.SEARCH_INTERNET_METHOD.strip().lower() == "serpapi"
+            and CFG.SERPAPI_KEY != ""
+        ):
+            response = requests.get(
+                "https://serpapi.com/search",
+                headers={"User-Agent": _DEFAULT_USER_AGENT},
+                params={
+                    "q": query,
+                    "start": (page - 1) * 10,
+                    "hl": CFG.SERPAPI_LANG,
+                    "safe": CFG.SERPAPI_SAFE,
+                    "api_key": CFG.SERPAPI_KEY,
+                },
+            )
+        elif (
+            CFG.SEARCH_INTERNET_METHOD.strip().lower() == "brave"
+            and CFG.BRAVE_API_KEY != ""
+        ):
+            response = requests.get(
+                "https://api.search.brave.com/res/v1/web/search",
+                headers={
+                    "User-Agent": _DEFAULT_USER_AGENT,
+                    "Accept": "application/json",
+                    "x-subscription-token": CFG.BRAVE_API_KEY,
+                },
+                params={
+                    "q": query,
+                    "count": "10",
+                    "offset": (page - 1) * 10,
+                    "safesearch": CFG.BRAVE_API_SAFE,
+                    "search_lang": CFG.BRAVE_API_LANG,
+                    "summary": "true",
+                },
+            )
+        else:
+            response = requests.get(
+                url=f"{CFG.SEARXNG_BASE_URL}/search",
+                headers={"User-Agent": _DEFAULT_USER_AGENT},
+                params={
+                    "q": query,
+                    "format": "json",
+                    "pageno": page,
+                    "safesearch": CFG.SEARXNG_SAFE,
+                    "language": CFG.SEARXNG_LANG,
+                },
+            )
         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 response.json()
 
     return search_internet
 
 
-def search_wikipedia(query: str) -> str:
-    """Search Wikipedia using its API.
-    Args:
-        query (str): Search query.
-    Returns:
-        str: JSON from Wikipedia API: {"batchcomplete": ..., "query": {"search": [...]}}
-    """
-    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: str, num_results: int = 10) -> str:
-    """Search ArXiv for papers using its API.
-    Args:
-        query (str): Search query.
-        num_results (int): Search result count. Defaults to 10.
-    Returns:
-        str: XML string from ArXiv API containing search results.
-    """
-    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
-
+async def _fetch_page_content(url: str) -> tuple[str, list[str]]:
+    """Fetches the HTML content and all absolute links from a URL."""
+    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()
+            await page.set_extra_http_headers({"User-Agent": _DEFAULT_USER_AGENT})
+            try:
+                await page.goto(url, wait_until="networkidle", timeout=30000)
+                await page.wait_for_load_state("domcontentloaded")
+                content = await page.content()
+                links = await page.eval_on_selector_all(
+                    "a[href]",
+                    """
+                    (elements, baseUrl) => elements.map(el => {
+                        const href = el.getAttribute('href');
+                        if (!href || href.startsWith('#')) return null;
+                        try {
+                            return new URL(href, baseUrl).href;
+                        } catch (e) {
+                            return null;
+                        }
+                    }).filter(href => href !== null)
+                    """,
+                    url,
+                )
+                return content, links
+                # return json.dumps({"content": content, "links": links})
+            finally:
+                await browser.close()
+    except Exception:
+        import requests
+        from bs4 import BeautifulSoup
 
-def parse_html_text(html_text: str) -> dict[str, str]:
+        response = requests.get(url, headers={"User-Agent": _DEFAULT_USER_AGENT})
+        if response.status_code != 200:
+            raise Exception(
+                f"Unable to retrieve page content. Status code: {response.status_code}"
+            )
+        content = response.text
+        soup = BeautifulSoup(content, "html.parser")
+        links = [
+            urljoin(url, a["href"])
+            for a in soup.find_all("a", href=True)
+            if not a["href"].startswith("#")
+        ]
+        return content, links
+        # return json.dumps({"content": content, "links": links})
+
+
+def _convert_html_to_markdown(html_text: str) -> str:
+    """Converts HTML content to a clean Markdown representation."""
     from bs4 import BeautifulSoup
+    from markdownify import markdownify as md
 
-    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):
+    # Remove non-content tags
+    for tag in soup(
+        ["script", "link", "meta", "style", "header", "footer", "nav", "aside"]
+    ):
         tag.decompose()
-    html_text = soup.get_text(separator=" ", strip=True)
-    return {"content": html_text, "links_on_page": links}
+    # Convert the cleaned HTML to Markdown
+    return md(str(soup))