zrb 1.5.11__py3-none-any.whl → 1.5.12__py3-none-any.whl

zrb/builtin/llm/llm_chat.py

@@ -93,7 +93,7 @@ llm_chat: LLMTask = llm_group.add_task(
         ),
         conversation_history_reader=read_chat_conversation,
         conversation_history_writer=write_chat_conversation,
-        description="Chat with LLM",
+        description="💬 Chat with LLM",
         system_prompt=lambda ctx: (
             None if ctx.input.system_prompt.strip() == "" else ctx.input.system_prompt
         ),

zrb/builtin/llm/tool/sub_agent.py

@@ -0,0 +1,125 @@
+import json
+from collections.abc import Callable
+from textwrap import dedent
+
+from pydantic_ai import Tool
+from pydantic_ai.mcp import MCPServer
+from pydantic_ai.models import Model
+from pydantic_ai.settings import ModelSettings
+
+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
+
+ToolOrCallable = Tool | Callable
+
+
+def create_sub_agent_tool(
+    tool_name: str,
+    tool_description: str,
+    sub_agent_system_prompt: str | None = None,  # Make optional
+    sub_agent_model: str | Model | None = None,
+    sub_agent_model_settings: ModelSettings | None = None,
+    sub_agent_tools: list[ToolOrCallable] = [],
+    sub_agent_mcp_servers: list[MCPServer] = [],
+) -> Callable[[AnyContext, str], str]:
+    """
+    Create an LLM "sub-agent" tool function for use by a main LLM agent.
+
+    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.
+
+    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).
+
+    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.
+    """
+
+    async def run_sub_agent(ctx: AnyContext, query: str) -> str:
+        """
+        Runs the sub-agent with the given query.
+        """
+        # Resolve parameters, falling back to llm_config defaults if None
+        resolved_model = get_model(
+            ctx=ctx,
+            model_attr=sub_agent_model,
+            render_model=True,  # Assuming we always want to render model string attributes
+            model_base_url_attr=None,
+            # Sub-agent tool doesn't have separate base_url/api_key params
+            render_model_base_url=False,
+            model_api_key_attr=None,
+            render_model_api_key=False,
+        )
+        resolved_model_settings = get_model_settings(
+            ctx=ctx,
+            model_settings_attr=sub_agent_model_settings,
+        )
+
+        if sub_agent_system_prompt is None:
+            resolved_system_prompt = get_combined_system_prompt(
+                ctx=ctx,
+                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 = sub_agent_system_prompt
+
+        # Create the sub-agent instance
+        sub_agent_agent = create_agent_instance(
+            ctx=ctx,
+            model=resolved_model,
+            system_prompt=resolved_system_prompt,
+            model_settings=resolved_model_settings,
+            tools_attr=sub_agent_tools,  # Pass tools from factory closure
+            additional_tools=[],  # No additional tools added after factory creation
+            mcp_servers_attr=sub_agent_mcp_servers,  # Pass servers from factory closure
+            additional_mcp_servers=[],  # No additional servers added after factory creation
+        )
+
+        # Run the sub-agent iteration
+        # Start with an empty history for the sub-agent
+        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
+        )
+
+        # 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."
+
+    # Set the name and docstring for the callable function
+    run_sub_agent.__name__ = tool_name
+    run_sub_agent.__doc__ = dedent(
+        f"""
+        {tool_description}
+
+        Args:
+            query (str): The query or task for the sub-agent.
+
+        Returns:
+            str: The final response or result from the sub-agent.
+        """
+    ).strip()
+
+    return run_sub_agent

zrb/builtin/llm/tool/web.py

@@ -1,8 +1,6 @@
 import json
 from collections.abc import Callable
 
-# Annotated import removed
-
 
 async def open_web_page(url: str) -> str:
     """Get parsed text content and links from a web page URL.

zrb/config.py

@@ -80,9 +80,6 @@ WEB_AUTH_REFRESH_TOKEN_EXPIRE_MINUTES = int(
 LLM_HISTORY_DIR = os.getenv(
     "ZRB_LLM_HISTORY_DIR", os.path.expanduser(os.path.join("~", ".zrb-llm-history"))
 )
-LLM_HISTORY_FILE = os.getenv(
-    "ZRB_LLM_HISTORY_FILE", os.path.join(LLM_HISTORY_DIR, "history.json")
-)
 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"))

zrb/llm_config.py

@@ -60,6 +60,7 @@ class LLMConfig:
         default_summarize_history: bool | None = None,
         default_history_summarization_threshold: int | None = None,
         default_enrich_context: bool | None = None,
+        default_context_enrichment_threshold: int | None = None,
     ):
         self._default_model_name = (
             default_model_name
@@ -104,7 +105,7 @@ class LLMConfig:
         self._default_summarize_history = (
             default_summarize_history
             if default_summarize_history is not None
-            else os.getenv("ZRB_LLM_SUMMARIZE_HISTORY", "true").lower() == "true"
+            else to_boolean(os.getenv("ZRB_LLM_SUMMARIZE_HISTORY", "true"))
         )
         self._default_history_summarization_threshold = (
             default_history_summarization_threshold
@@ -114,7 +115,12 @@ class LLMConfig:
         self._default_enrich_context = (
             default_enrich_context
             if default_enrich_context is not None
-            else to_boolean(os.getenv("ZRB_LLM_ENRICH_CONTEXT", "0"))
+            else to_boolean(os.getenv("ZRB_LLM_ENRICH_CONTEXT", "true"))
+        )
+        self._default_context_enrichment_threshold = (
+            default_context_enrichment_threshold
+            if default_context_enrichment_threshold is not None
+            else int(os.getenv("ZRB_LLM_CONTEXT_ENRICHMENT_THRESHOLD", "5"))
         )
         self._default_provider = None
         self._default_model = None
@@ -186,6 +192,9 @@ class LLMConfig:
     def get_default_enrich_context(self) -> bool:
         return self._default_enrich_context
 
+    def get_default_context_enrichment_threshold(self) -> int:
+        return self._default_context_enrichment_threshold
+
     def set_default_persona(self, persona: str):
         self._default_persona = persona
 
@@ -227,5 +236,10 @@ class LLMConfig:
     def set_default_enrich_context(self, enrich_context: bool):
         self._default_enrich_context = enrich_context
 
+    def set_default_context_enrichment_threshold(
+        self, context_enrichment_threshold: int
+    ):
+        self._default_context_enrichment_threshold = context_enrichment_threshold
+
 
 llm_config = LLMConfig()

zrb/task/base_task.py

@@ -23,6 +23,26 @@ from zrb.task.base.operators import handle_lshift, handle_rshift
 
 
 class BaseTask(AnyTask):
+    """
+    Implements a concrete task class `BaseTask` derived from the abstract base class `AnyTask`.
+
+    This class serves as a robust and flexible task implementation that can be tailored for
+    various execution scenarios within the Zrb framework. It supports functionalities such as:
+
+    - **Task Definition and Initialization:** Setting up task attributes like `name`, `color`,
+    `icon`, `description`, `cli_only`, `inputs`, `envs`, `action`, among others.
+    - **Dependency Management:** Managing task dependencies using properties and methods to
+    append upstreams, fallbacks, readiness checks, and successors, ensuring tasks are executed
+    in the correct order and conditions.
+    - **Execution Control:** Contains methods for both synchronous (`run`) and asynchronous
+    execution (`async_run`), alongside internal task lifecycle methods (`exec_root_tasks`,
+    `exec_chain`, `exec`).
+    - **Readiness and Monitoring:** Supports readiness checks, retry mechanisms, and monitoring
+    before task execution to ensure the task is executed under proper conditions.
+    - **Operator Overloading:** Implements operators to handle task chaining and dependencies
+    conveniently.
+    """
+
     def __init__(
         self,
         name: str,

zrb/task/llm/agent.py

@@ -1,4 +1,3 @@
-import inspect
 from collections.abc import Callable
 
 from openai import APIError
@@ -32,6 +31,7 @@ def create_agent_instance(
     additional_mcp_servers: list[MCPServer],
 ) -> Agent:
     """Creates a new Agent instance with configured tools and servers."""
+    # Get tools
     tools_or_callables = list(tools_attr(ctx) if callable(tools_attr) else tools_attr)
     tools_or_callables.extend(additional_tools)
     tools = []
@@ -39,17 +39,14 @@ def create_agent_instance(
         if isinstance(tool_or_callable, Tool):
             tools.append(tool_or_callable)
         else:
-            # Inspect original callable for 'ctx' parameter (pydantic-ai context)
-            original_sig = inspect.signature(tool_or_callable)
-            takes_ctx = "ctx" in original_sig.parameters
-            wrapped_tool = wrap_tool(tool_or_callable)
-            tools.append(Tool(wrapped_tool, takes_ctx=takes_ctx))
-
+            # Pass ctx to wrap_tool
+            tools.append(wrap_tool(tool_or_callable, ctx))
+    # Get MCP Servers
     mcp_servers = list(
         mcp_servers_attr(ctx) if callable(mcp_servers_attr) else mcp_servers_attr
     )
     mcp_servers.extend(additional_mcp_servers)
-
+    # Return Agent
     return Agent(
         model=model,
         system_prompt=system_prompt,

zrb/task/llm/context.py

@@ -12,13 +12,19 @@ from zrb.util.attr import get_attr
 from zrb.util.file import read_dir, read_file_with_line_numbers
 
 
-def get_default_context(user_message: str) -> dict[str, Any]:
-    """Generates default context including time, OS, and file references."""
-    references = re.findall(r"@(\S+)", user_message)
+def extract_default_context(user_message: str) -> tuple[str, dict[str, Any]]:
+    """
+    Return modified user message and default context including time, OS, and file references.
+    """
+    modified_user_message = user_message
+    # Match “@” + any non-space/comma sequence that contains at least one “/”
+    pattern = r"(?<!\w)@(?=[^,\s]*/)([^,\s]+)"
+    potential_resource_path = re.findall(pattern, user_message)
     current_references = []
 
-    for ref in references:
+    for ref in potential_resource_path:
         resource_path = os.path.abspath(os.path.expanduser(ref))
+        print("RESOURCE PATH", resource_path)
         if os.path.isfile(resource_path):
             content = read_file_with_line_numbers(resource_path)
             current_references.append(
@@ -30,6 +36,8 @@ def get_default_context(user_message: str) -> dict[str, Any]:
                     "content": content,
                 }
             )
+            # Remove the '@' from the modified user message for valid file paths
+            modified_user_message = modified_user_message.replace(f"@{ref}", ref, 1)
         elif os.path.isdir(resource_path):
             content = read_dir(resource_path)
             current_references.append(
@@ -40,8 +48,10 @@ def get_default_context(user_message: str) -> dict[str, Any]:
                     "content": content,
                 }
             )
+            # Remove the '@' from the modified user message for valid directory paths
+            modified_user_message = modified_user_message.replace(f"@{ref}", ref, 1)
 
-    return {
+    context = {
         "current_time": datetime.datetime.now().isoformat(),
         "current_working_directory": os.getcwd(),
         "current_os": platform.system(),
@@ -50,6 +60,8 @@ def get_default_context(user_message: str) -> dict[str, Any]:
         "current_references": current_references,
     }
 
+    return modified_user_message, context
+
 
 def get_conversation_context(
     ctx: AnyContext,
@@ -90,6 +102,3 @@ def get_conversation_context(
         else:
             processed_context[key] = value
     return processed_context
-
-
-# Context enrichment functions moved to context_enrichment.py

zrb/task/llm/context_enrichment.py

@@ -8,12 +8,12 @@ from pydantic_ai import Agent
 from pydantic_ai.models import Model
 from pydantic_ai.settings import ModelSettings
 
-from zrb.attr.type import BoolAttr
+from zrb.attr.type import BoolAttr, IntAttr
 from zrb.context.any_context import AnyContext
 from zrb.llm_config import llm_config
 from zrb.task.llm.agent import run_agent_iteration
 from zrb.task.llm.typing import ListOfDict
-from zrb.util.attr import get_bool_attr
+from zrb.util.attr import get_bool_attr, get_int_attr
 
 
 class EnrichmentConfig(BaseModel):
@@ -62,7 +62,7 @@ async def enrich_context(
         mcp_servers=[],
         model_settings=config.settings,
         retries=config.retries,
-        result_type=EnrichmentResult,
+        output_type=EnrichmentResult,
     )
 
     try:
@@ -72,8 +72,8 @@ async def enrich_context(
             user_prompt=user_prompt_data,  # Pass the formatted data as user prompt
             history_list=[],  # Enrichment agent doesn't need prior history itself
         )
-        if enrichment_run and enrichment_run.result.data:
-            response = enrichment_run.result.data.response
+        if enrichment_run and enrichment_run.result.output:
+            response = enrichment_run.result.output.response
             if response:
                 conversation_context.update(response)
                 ctx.log_info("Context enriched based on history.")
@@ -88,21 +88,53 @@ async def enrich_context(
     return conversation_context
 
 
+def get_context_enrichment_threshold(
+    ctx: AnyContext,
+    context_enrichment_threshold_attr: IntAttr | None,
+    render_context_enrichment_threshold: bool,
+) -> int:
+    """Gets the context enrichment threshold, handling defaults and errors."""
+    try:
+        return get_int_attr(
+            ctx,
+            context_enrichment_threshold_attr,
+            # Use llm_config default if attribute is None
+            llm_config.get_default_context_enrichment_threshold(),
+            auto_render=render_context_enrichment_threshold,
+        )
+    except ValueError as e:
+        ctx.log_warning(
+            f"Could not convert context_enrichment_threshold to int: {e}. "
+            "Defaulting to -1 (no threshold)."
+        )
+        return -1
+
+
 def should_enrich_context(
     ctx: AnyContext,
     history_list: ListOfDict,
     should_enrich_context_attr: BoolAttr | None,  # Allow None
     render_enrich_context: bool,
+    context_enrichment_threshold_attr: IntAttr | None,
+    render_context_enrichment_threshold: bool,
 ) -> bool:
-    """Determines if context enrichment should occur based on history and config."""
-    if len(history_list) == 0:
+    """
+    Determines if context enrichment should occur based on history, threshold, and config.
+    """
+    history_len = len(history_list)
+    if history_len == 0:
+        return False
+    enrichment_threshold = get_context_enrichment_threshold(
+        ctx,
+        context_enrichment_threshold_attr,
+        render_context_enrichment_threshold,
+    )
+    if enrichment_threshold == -1 or enrichment_threshold > history_len:
         return False
-    # Use llm_config default if attribute is None
-    default_value = llm_config.get_default_enrich_context()
     return get_bool_attr(
         ctx,
         should_enrich_context_attr,
-        default_value,  # Pass the default from llm_config
+        llm_config.get_default_enrich_context(),
         auto_render=render_enrich_context,
     )
 
@@ -111,15 +143,22 @@ async def maybe_enrich_context(
     ctx: AnyContext,
     history_list: ListOfDict,
     conversation_context: dict[str, Any],
-    should_enrich_context_attr: BoolAttr | None,  # Allow None
+    should_enrich_context_attr: BoolAttr | None,
     render_enrich_context: bool,
+    context_enrichment_threshold_attr: IntAttr | None,
+    render_context_enrichment_threshold: bool,
     model: str | Model | None,
     model_settings: ModelSettings | None,
     context_enrichment_prompt: str,
 ) -> dict[str, Any]:
-    """Enriches context based on history if enabled."""
+    """Enriches context based on history if enabled and threshold met."""
     if should_enrich_context(
-        ctx, history_list, should_enrich_context_attr, render_enrich_context
+        ctx,
+        history_list,
+        should_enrich_context_attr,
+        render_enrich_context,
+        context_enrichment_threshold_attr,
+        render_context_enrichment_threshold,
     ):
         # Use the enrich_context function now defined in this file
         return await enrich_context(

zrb/task/llm/history_summarization.py

@@ -64,9 +64,7 @@ def should_summarize_history(
         history_summarization_threshold_attr,
         render_history_summarization_threshold,
     )
-    if summarization_threshold == -1:  # -1 means no summarization trigger
-        return False
-    if summarization_threshold > history_part_len:
+    if summarization_threshold == -1 or summarization_threshold > history_part_len:
         return False
     return get_bool_attr(
         ctx,
@@ -122,8 +120,8 @@ async def summarize_history(
             user_prompt=summarization_user_prompt,
             history_list=[],  # Summarization agent doesn't need prior history
         )
-        if summary_run and summary_run.result.data:
-            summary_text = str(summary_run.result.data)
+        if summary_run and summary_run.result.output:
+            summary_text = str(summary_run.result.output)
             # Update context with the new summary
             conversation_context["history_summary"] = summary_text
             ctx.log_info("History summarized and added/updated in context.")

zrb/task/llm/prompt.py

@@ -5,7 +5,7 @@ from typing import Any
 from zrb.attr.type import StrAttr
 from zrb.context.any_context import AnyContext
 from zrb.llm_config import llm_config as default_llm_config
-from zrb.task.llm.context import get_default_context  # Updated import
+from zrb.task.llm.context import extract_default_context
 from zrb.util.attr import get_attr, get_str_attr
 
 
@@ -135,14 +135,17 @@ def build_user_prompt(
     conversation_context: dict[str, Any],
 ) -> str:
     """Constructs the final user prompt including context."""
-    user_message = get_user_message(ctx, message_attr)
+    original_user_message = get_user_message(ctx, message_attr)
     # Combine default context, conversation context (potentially enriched/summarized)
-    enriched_context = {**get_default_context(user_message), **conversation_context}
+    modified_user_message, default_context = extract_default_context(
+        original_user_message
+    )
+    enriched_context = {**default_context, **conversation_context}
     return dedent(
         f"""
         # Context
         {json.dumps(enriched_context)}
         # User Message
-        {user_message}
+        {modified_user_message}
         """
     ).strip()