zrb 1.5.6__py3-none-any.whl → 1.5.8__py3-none-any.whl

zrb/builtin/llm/tool/file.py

@@ -4,8 +4,7 @@ import os
 import re
 from typing import Any, Optional
 
-from zrb.util.file import read_file as _read_file
-from zrb.util.file import write_file as _write_file
+from zrb.util.file import read_file, read_file_with_line_numbers, write_file
 
 DEFAULT_EXCLUDED_PATTERNS = [
     # Common Python artifacts
@@ -182,7 +181,7 @@ def read_from_file(
     start_line: Optional[int] = None,
     end_line: Optional[int] = None,
 ) -> str:
-    """Read file content (or specific lines) at a path.
+    """Read file content (or specific lines) at a path, including line numbers.
     Args:
         path (str): Path to read. Pass exactly as provided, including '~'.
         start_line (Optional[int]): Starting line number (1-based).
@@ -191,6 +190,7 @@ def read_from_file(
             Defaults to None (end of file).
     Returns:
         str: JSON: {"path": "...", "content": "...", "start_line": N, ...} or {"error": "..."}
+        The content includes line numbers.
     Raises:
         Exception: If an error occurs.
     """
@@ -199,7 +199,7 @@ def read_from_file(
         # Check if file exists
         if not os.path.exists(abs_path):
             return json.dumps({"error": f"File {path} does not exist"})
-        content = _read_file(abs_path)
+        content = read_file_with_line_numbers(abs_path)
         lines = content.splitlines()
         total_lines = len(lines)
         # Adjust line indices (convert from 1-based to 0-based)
@@ -259,7 +259,7 @@ def write_to_file(
         directory = os.path.dirname(abs_path)
         if directory and not os.path.exists(directory):
             os.makedirs(directory, exist_ok=True)
-        _write_file(abs_path, content)
+        write_file(abs_path, content)
         result_data = {"success": True, "path": path}
         if warning:
             result_data["warning"] = warning
@@ -391,6 +391,7 @@ def apply_diff(
         replace_marker (str): Marker for end of replacement block.
             Defaults to ">>>>>> REPLACE".
     SEARCH block must exactly match file content including whitespace/indentation.
+    SEARCH block should NOT contains line numbers
     Format example:
         [Search Marker, e.g., <<<<<< SEARCH]
         :start_line:10
@@ -414,7 +415,7 @@ def apply_diff(
             return json.dumps(
                 {"success": False, "path": path, "error": f"File not found at {path}"}
             )
-        content = _read_file(abs_path)
+        content = read_file(abs_path)
         lines = content.splitlines()
         if start_line < 1 or end_line > len(lines) or start_line > end_line:
             return json.dumps(
@@ -444,7 +445,7 @@ def apply_diff(
         new_content = "\n".join(new_lines)
         if content.endswith("\n"):
             new_content += "\n"
-        _write_file(abs_path, new_content)
+        write_file(abs_path, new_content)
         return json.dumps({"success": True, "path": path})
     except ValueError as e:
         raise ValueError(f"Error parsing diff: {e}")

zrb/builtin/llm/tool/rag.py

@@ -152,13 +152,14 @@ def create_rag_from_directory(
 
     retrieve.__name__ = tool_name
     retrieve.__doc__ = dedent(
-        f"""{tool_description}
+        f"""
+        {tool_description}
         Args:
             query (str): The user query to search for in documents.
         Returns:
             str: JSON string with search results: {{"ids": [...], "documents": [...], ...}}
-    """
-    )
+        """
+    ).strip()
     return retrieve
 
 

zrb/task/cmd_task.py

@@ -130,7 +130,8 @@ class CmdTask(BaseTask):
             partial(ctx.print, plain=True) if self._should_plain_print else ctx.print
         )
         xcom_pid_key = f"{self.name}-pid"
-        ctx.xcom[xcom_pid_key] = Xcom([])
+        if xcom_pid_key not in ctx.xcom:
+            ctx.xcom[xcom_pid_key] = Xcom([])
         cmd_result, return_code = await run_command(
             cmd=[shell, shell_flag, cmd_script],
             cwd=cwd,

zrb/task/llm/agent.py

@@ -0,0 +1,141 @@
+import inspect
+from collections.abc import Callable
+
+from openai import APIError
+from pydantic_ai import Agent, Tool
+from pydantic_ai.agent import AgentRun
+from pydantic_ai.mcp import MCPServer
+from pydantic_ai.messages import ModelMessagesTypeAdapter
+from pydantic_ai.models import Model
+from pydantic_ai.settings import ModelSettings
+
+from zrb.context.any_context import AnyContext
+from zrb.context.any_shared_context import AnySharedContext
+from zrb.task.llm.error import extract_api_error_details
+from zrb.task.llm.print_node import print_node
+from zrb.task.llm.tool_wrapper import wrap_tool
+from zrb.task.llm.typing import ListOfDict
+
+ToolOrCallable = Tool | Callable
+
+
+def create_agent_instance(
+    ctx: AnyContext,
+    model: str | Model | None,
+    system_prompt: str,
+    model_settings: ModelSettings | None,
+    tools_attr: (
+        list[ToolOrCallable] | Callable[[AnySharedContext], list[ToolOrCallable]]
+    ),
+    additional_tools: list[ToolOrCallable],
+    mcp_servers_attr: list[MCPServer] | Callable[[AnySharedContext], list[MCPServer]],
+    additional_mcp_servers: list[MCPServer],
+) -> Agent:
+    """Creates a new Agent instance with configured tools and servers."""
+    tools_or_callables = list(tools_attr(ctx) if callable(tools_attr) else tools_attr)
+    tools_or_callables.extend(additional_tools)
+    tools = []
+    for tool_or_callable in tools_or_callables:
+        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))
+
+    mcp_servers = list(
+        mcp_servers_attr(ctx) if callable(mcp_servers_attr) else mcp_servers_attr
+    )
+    mcp_servers.extend(additional_mcp_servers)
+
+    return Agent(
+        model=model,
+        system_prompt=system_prompt,
+        tools=tools,
+        mcp_servers=mcp_servers,
+        model_settings=model_settings,
+        retries=3,  # Consider making retries configurable?
+    )
+
+
+def get_agent(
+    ctx: AnyContext,
+    agent_attr: Agent | Callable[[AnySharedContext], Agent] | None,
+    model: str | Model | None,
+    system_prompt: str,
+    model_settings: ModelSettings | None,
+    tools_attr: (
+        list[ToolOrCallable] | Callable[[AnySharedContext], list[ToolOrCallable]]
+    ),
+    additional_tools: list[ToolOrCallable],
+    mcp_servers_attr: list[MCPServer] | Callable[[AnySharedContext], list[MCPServer]],
+    additional_mcp_servers: list[MCPServer],
+) -> Agent:
+    """Retrieves the configured Agent instance or creates one if necessary."""
+    if isinstance(agent_attr, Agent):
+        return agent_attr
+    if callable(agent_attr):
+        agent_instance = agent_attr(ctx)
+        if not isinstance(agent_instance, Agent):
+            err_msg = (
+                "Callable agent factory did not return an Agent instance, "
+                f"got: {type(agent_instance)}"
+            )
+            raise TypeError(err_msg)
+        return agent_instance
+    # If no agent provided, create one using the configuration
+    return create_agent_instance(
+        ctx=ctx,
+        model=model,
+        system_prompt=system_prompt,
+        model_settings=model_settings,
+        tools_attr=tools_attr,
+        additional_tools=additional_tools,
+        mcp_servers_attr=mcp_servers_attr,
+        additional_mcp_servers=additional_mcp_servers,
+    )
+
+
+async def run_agent_iteration(
+    ctx: AnyContext,
+    agent: Agent,
+    user_prompt: str,
+    history_list: ListOfDict,
+) -> AgentRun:
+    """
+    Runs a single iteration of the agent execution loop.
+
+    Args:
+        ctx: The task context.
+        agent: The Pydantic AI agent instance.
+        user_prompt: The user's input prompt.
+        history_list: The current conversation history.
+
+    Returns:
+        The agent run result object.
+
+    Raises:
+        Exception: If any error occurs during agent execution.
+    """
+    async with agent.run_mcp_servers():
+        async with agent.iter(
+            user_prompt=user_prompt,
+            message_history=ModelMessagesTypeAdapter.validate_python(history_list),
+        ) as agent_run:
+            async for node in agent_run:
+                # Each node represents a step in the agent's execution
+                # Reference: https://ai.pydantic.dev/agents/#streaming
+                try:
+                    await print_node(ctx.print, agent_run, node)
+                except APIError as e:
+                    # Extract detailed error information from the response
+                    error_details = extract_api_error_details(e)
+                    ctx.log_error(f"API Error: {error_details}")
+                    raise
+                except Exception as e:
+                    ctx.log_error(f"Error processing node: {str(e)}")
+                    ctx.log_error(f"Error type: {type(e).__name__}")
+                    raise
+            return agent_run

zrb/task/llm/config.py

@@ -0,0 +1,83 @@
+from typing import Callable
+
+from pydantic_ai.models import Model
+from pydantic_ai.settings import ModelSettings
+
+from zrb.attr.type import StrAttr, fstring
+from zrb.context.any_context import AnyContext
+from zrb.context.any_shared_context import AnySharedContext
+from zrb.llm_config import LLMConfig
+from zrb.llm_config import llm_config as default_llm_config
+from zrb.util.attr import get_attr
+
+
+def get_model_settings(
+    ctx: AnyContext,
+    model_settings_attr: (
+        ModelSettings | Callable[[AnySharedContext], ModelSettings] | None
+    ),
+) -> ModelSettings | None:
+    """Gets the model settings, resolving callables if necessary."""
+    if callable(model_settings_attr):
+        return model_settings_attr(ctx)
+    return model_settings_attr
+
+
+def get_model_base_url(
+    ctx: AnyContext,
+    model_base_url_attr: StrAttr | None,
+    render_model_base_url: bool,
+) -> str | None:
+    """Gets the model base URL, rendering if configured."""
+    base_url = get_attr(
+        ctx, model_base_url_attr, None, auto_render=render_model_base_url
+    )
+    if isinstance(base_url, str) or base_url is None:
+        return base_url
+    raise ValueError(f"Invalid model base URL: {base_url}")
+
+
+def get_model_api_key(
+    ctx: AnyContext,
+    model_api_key_attr: StrAttr | None,
+    render_model_api_key: bool,
+) -> str | None:
+    """Gets the model API key, rendering if configured."""
+    api_key = get_attr(ctx, model_api_key_attr, None, auto_render=render_model_api_key)
+    if isinstance(api_key, str) or api_key is None:
+        return api_key
+    raise ValueError(f"Invalid model API key: {api_key}")
+
+
+def get_model(
+    ctx: AnyContext,
+    model_attr: Callable[[AnySharedContext], Model | str | fstring] | Model | None,
+    render_model: bool,
+    model_base_url_attr: StrAttr | None,
+    render_model_base_url: bool,
+    model_api_key_attr: StrAttr | None,
+    render_model_api_key: bool,
+) -> str | Model | None:
+    """Gets the model instance or name, handling defaults and configuration."""
+    model = get_attr(ctx, model_attr, None, auto_render=render_model)
+    if model is None:
+        return default_llm_config.get_default_model()
+    if isinstance(model, str):
+        model_base_url = get_model_base_url(
+            ctx, model_base_url_attr, render_model_base_url
+        )
+        model_api_key = get_model_api_key(ctx, model_api_key_attr, render_model_api_key)
+        llm_config = LLMConfig(
+            default_model_name=model,
+            default_base_url=model_base_url,
+            default_api_key=model_api_key,
+        )
+        if model_base_url is None and model_api_key is None:
+            default_model_provider = default_llm_config.get_default_model_provider()
+            if default_model_provider is not None:
+                llm_config.set_default_provider(default_model_provider)
+        return llm_config.get_default_model()
+    # If it's already a Model instance, return it directly
+    if isinstance(model, Model):
+        return model
+    raise ValueError(f"Invalid model type resolved: {type(model)}, value: {model}")

zrb/task/llm/context.py

@@ -0,0 +1,95 @@
+import datetime
+import inspect
+import os
+import platform
+import re
+from collections.abc import Callable
+from typing import Any
+
+from zrb.context.any_context import AnyContext
+from zrb.context.any_shared_context import AnySharedContext
+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)
+    current_references = []
+
+    for ref in references:
+        resource_path = os.path.abspath(os.path.expanduser(ref))
+        if os.path.isfile(resource_path):
+            content = read_file_with_line_numbers(resource_path)
+            current_references.append(
+                {
+                    "reference": ref,
+                    "name": resource_path,
+                    "type": "file",
+                    "note": "line numbers are included in the content",
+                    "content": content,
+                }
+            )
+        elif os.path.isdir(resource_path):
+            content = read_dir(resource_path)
+            current_references.append(
+                {
+                    "reference": ref,
+                    "name": resource_path,
+                    "type": "directory",
+                    "content": content,
+                }
+            )
+
+    return {
+        "current_time": datetime.datetime.now().isoformat(),
+        "current_working_directory": os.getcwd(),
+        "current_os": platform.system(),
+        "os_version": platform.version(),
+        "python_version": platform.python_version(),
+        "current_references": current_references,
+    }
+
+
+def get_conversation_context(
+    ctx: AnyContext,
+    conversation_context_attr: (
+        dict[str, Any] | Callable[[AnySharedContext], dict[str, Any]] | None
+    ),
+) -> dict[str, Any]:
+    """
+    Retrieves the conversation context.
+    If a value in the context dict is callable, it executes it with ctx.
+    """
+    raw_context = get_attr(
+        ctx, conversation_context_attr, {}, auto_render=False
+    )  # Context usually shouldn't be rendered
+    if not isinstance(raw_context, dict):
+        ctx.log_warning(
+            f"Conversation context resolved to type {type(raw_context)}, "
+            "expected dict. Returning empty context."
+        )
+        return {}
+    # If conversation_context contains callable value, execute them.
+    processed_context: dict[str, Any] = {}
+    for key, value in raw_context.items():
+        if callable(value):
+            try:
+                # Check if the callable expects 'ctx'
+                sig = inspect.signature(value)
+                if "ctx" in sig.parameters:
+                    processed_context[key] = value(ctx)
+                else:
+                    processed_context[key] = value()
+            except Exception as e:
+                ctx.log_warning(
+                    f"Error executing callable for context key '{key}': {e}. "
+                    "Skipping."
+                )
+                processed_context[key] = None
+        else:
+            processed_context[key] = value
+    return processed_context
+
+
+# Context enrichment functions moved to context_enrichment.py

zrb/task/llm/{context_enricher.py → context_enrichment.py}

@@ -8,12 +8,13 @@ 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.context.any_context import AnyContext
-from zrb.task.llm.agent_runner import run_agent_iteration
-from zrb.task.llm.history import ListOfDict
+from zrb.task.llm.agent import run_agent_iteration
+from zrb.task.llm.typing import ListOfDict
+from zrb.util.attr import get_bool_attr
 
 
-# Configuration model for context enrichment
 class EnrichmentConfig(BaseModel):
     model_config = {"arbitrary_types_allowed": True}
     model: Model | str | None = None
@@ -74,13 +75,58 @@ async def enrich_context(
             response = enrichment_run.result.data.response
             if response:
                 conversation_context.update(response)
-                ctx.log_info("History summarized and added/updated in context.")
+                ctx.log_info("Context enriched based on history.")
                 ctx.log_info(
-                    f"New conversation context: {json.dumps(conversation_context)}"
+                    f"Updated conversation context: {json.dumps(conversation_context)}"
                 )
         else:
-            ctx.log_warning("Context enrichment return no data")
+            ctx.log_warning("Context enrichment returned no data")
     except Exception as e:
         ctx.log_warning(f"Error during context enrichment LLM call: {e}")
         traceback.print_exc()
     return conversation_context
+
+
+def should_enrich_context(
+    ctx: AnyContext,
+    history_list: ListOfDict,
+    should_enrich_context_attr: BoolAttr,
+    render_enrich_context: bool,
+) -> bool:
+    """Determines if context enrichment should occur based on history and config."""
+    if len(history_list) == 0:
+        return False
+    return get_bool_attr(
+        ctx,
+        should_enrich_context_attr,
+        True,  # Default to True if not specified
+        auto_render=render_enrich_context,
+    )
+
+
+async def maybe_enrich_context(
+    ctx: AnyContext,
+    history_list: ListOfDict,
+    conversation_context: dict[str, Any],
+    should_enrich_context_attr: BoolAttr,
+    render_enrich_context: bool,
+    model: str | Model | None,
+    model_settings: ModelSettings | None,
+    context_enrichment_prompt: str,
+) -> dict[str, Any]:
+    """Enriches context based on history if enabled."""
+    if should_enrich_context(
+        ctx, history_list, should_enrich_context_attr, render_enrich_context
+    ):
+        # Use the enrich_context function now defined in this file
+        return await enrich_context(
+            ctx=ctx,
+            config=EnrichmentConfig(
+                model=model,
+                settings=model_settings,
+                prompt=context_enrichment_prompt,
+            ),
+            conversation_context=conversation_context,
+            history_list=history_list,
+        )
+    return conversation_context

zrb/task/llm/history.py

@@ -5,12 +5,14 @@ from typing import Any, Optional
 
 from pydantic import BaseModel
 
+from zrb.attr.type import StrAttr
 from zrb.context.any_context import AnyContext
-from zrb.util.file import read_file
+from zrb.context.any_shared_context import AnySharedContext
+from zrb.task.llm.typing import ListOfDict
+from zrb.util.attr import get_str_attr
+from zrb.util.file import read_file, write_file
 from zrb.util.run import run_async
 
-ListOfDict = list[dict[str, Any]]
-
 
 # Define the new ConversationHistoryData model
 class ConversationHistoryData(BaseModel):
@@ -90,3 +92,151 @@ class ConversationHistoryData(BaseModel):
                 f"Error validating/parsing history data from {source}: {e}. Ignoring."
             )
             return None
+
+
+def get_history_file(
+    ctx: AnyContext,
+    conversation_history_file_attr: StrAttr | None,
+    render_history_file: bool,
+) -> str:
+    """Gets the path to the conversation history file, rendering if configured."""
+    return get_str_attr(
+        ctx,
+        conversation_history_file_attr,
+        "",
+        auto_render=render_history_file,
+    )
+
+
+async def read_conversation_history(
+    ctx: AnyContext,
+    conversation_history_reader: (
+        Callable[[AnySharedContext], ConversationHistoryData | dict | list | None]
+        | None
+    ),
+    conversation_history_file_attr: StrAttr | None,
+    render_history_file: bool,
+    conversation_history_attr: (
+        ConversationHistoryData
+        | Callable[[AnySharedContext], ConversationHistoryData | dict | list]
+        | dict
+        | list
+    ),
+) -> ConversationHistoryData:
+    """Reads conversation history from reader, file, or attribute, with validation."""
+    history_file = get_history_file(
+        ctx, conversation_history_file_attr, render_history_file
+    )
+    # Use the class method defined above
+    history_data = await ConversationHistoryData.read_from_sources(
+        ctx=ctx,
+        reader=conversation_history_reader,
+        file_path=history_file,
+    )
+    if history_data:
+        return history_data
+    # Priority 3: Callable or direct conversation_history attribute
+    raw_data_attr: Any = None
+    if callable(conversation_history_attr):
+        try:
+            raw_data_attr = await run_async(conversation_history_attr(ctx))
+        except Exception as e:
+            ctx.log_warning(
+                f"Error executing callable conversation_history attribute: {e}. "
+                "Ignoring."
+            )
+    if raw_data_attr is None:
+        raw_data_attr = conversation_history_attr
+    if raw_data_attr:
+        # Use the class method defined above
+        history_data = ConversationHistoryData.parse_and_validate(
+            ctx, raw_data_attr, "attribute"
+        )
+        if history_data:
+            return history_data
+    # Fallback: Return default value
+    return ConversationHistoryData()
+
+
+async def write_conversation_history(
+    ctx: AnyContext,
+    history_data: ConversationHistoryData,
+    conversation_history_writer: (
+        Callable[[AnySharedContext, ConversationHistoryData], None] | None
+    ),
+    conversation_history_file_attr: StrAttr | None,
+    render_history_file: bool,
+):
+    """Writes conversation history using the writer or to a file."""
+    if conversation_history_writer is not None:
+        await run_async(conversation_history_writer(ctx, history_data))
+    history_file = get_history_file(
+        ctx, conversation_history_file_attr, render_history_file
+    )
+    if history_file != "":
+        write_file(history_file, history_data.model_dump_json(indent=2))
+
+
+async def prepare_initial_state(
+    ctx: AnyContext,
+    conversation_history_reader: (
+        Callable[[AnySharedContext], ConversationHistoryData | dict | list | None]
+        | None
+    ),
+    conversation_history_file_attr: StrAttr | None,
+    render_history_file: bool,
+    conversation_history_attr: (
+        ConversationHistoryData
+        | Callable[[AnySharedContext], ConversationHistoryData | dict | list]
+        | dict
+        | list
+    ),
+    conversation_context_getter: Callable[[AnyContext], dict[str, Any]],
+) -> tuple[ListOfDict, dict[str, Any]]:
+    """Reads history and prepares the initial conversation context."""
+    history_data: ConversationHistoryData = await read_conversation_history(
+        ctx,
+        conversation_history_reader,
+        conversation_history_file_attr,
+        render_history_file,
+        conversation_history_attr,
+    )
+    # Clean the history list to remove context from historical user prompts
+    cleaned_history_list = []
+    for interaction in history_data.history:
+        cleaned_history_list.append(
+            remove_context_from_interaction_history(interaction)
+        )
+    conversation_context = conversation_context_getter(ctx)
+    # Merge history context from loaded data without overwriting existing keys
+    for key, value in history_data.context.items():
+        if key not in conversation_context:
+            conversation_context[key] = value
+    # Return the CLEANED history list
+    return cleaned_history_list, conversation_context
+
+
+def remove_context_from_interaction_history(
+    interaction: dict[str, Any],
+) -> dict[str, Any]:
+    try:
+        cleaned_interaction = json.loads(json.dumps(interaction))
+    except Exception:
+        # Fallback to shallow copy if not JSON serializable (less safe)
+        cleaned_interaction = interaction.copy()
+    if "parts" in cleaned_interaction and isinstance(
+        cleaned_interaction["parts"], list
+    ):
+        for part in cleaned_interaction["parts"]:
+            is_user_prompt = part.get("part_kind") == "user-prompt"
+            has_str_content = isinstance(part.get("content"), str)
+            if is_user_prompt and has_str_content:
+                content = part["content"]
+                user_message_marker = "# User Message\n"
+                marker_index = content.find(user_message_marker)
+                if marker_index != -1:
+                    # Extract message after the marker and strip whitespace
+                    start_index = marker_index + len(user_message_marker)
+                    part["content"] = content[start_index:].strip()
+                # else: If marker not found, leave content as is (old format/error)
+    return cleaned_interaction