zrb 1.5.7__tar.gz → 1.5.8__tar.gz

{zrb-1.5.7 → zrb-1.5.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.5.7
+Version: 1.5.8
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later

{zrb-1.5.7 → zrb-1.5.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.5.7"
+version = "1.5.8"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"

{zrb-1.5.7 → zrb-1.5.8}/src/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-1.5.7 → zrb-1.5.8}/src/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-1.5.8/src/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-1.5.8/src/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-1.5.8/src/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-1.5.7/src/zrb/task/llm/context_enricher.py → zrb-1.5.8/src/zrb/task/llm/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-1.5.8/src/zrb/task/llm/history.py

@@ -0,0 +1,242 @@
+import json
+import os
+from collections.abc import Callable
+from typing import Any, Optional
+
+from pydantic import BaseModel
+
+from zrb.attr.type import StrAttr
+from zrb.context.any_context import AnyContext
+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
+
+
+# Define the new ConversationHistoryData model
+class ConversationHistoryData(BaseModel):
+    context: dict[str, Any] = {}
+    history: ListOfDict = []
+
+    @classmethod
+    async def read_from_sources(
+        cls,
+        ctx: AnyContext,
+        reader: Callable[[AnyContext], dict[str, Any] | list | None] | None,
+        file_path: str | None,
+    ) -> Optional["ConversationHistoryData"]:
+        """Reads conversation history from various sources with priority."""
+        # Priority 1: Reader function
+        if reader:
+            try:
+                raw_data = await run_async(reader(ctx))
+                if raw_data:
+                    instance = cls.parse_and_validate(ctx, raw_data, "reader")
+                    if instance:
+                        return instance
+            except Exception as e:
+                ctx.log_warning(
+                    f"Error executing conversation history reader: {e}. Ignoring."
+                )
+        # Priority 2: History file
+        if file_path and os.path.isfile(file_path):
+            try:
+                content = read_file(file_path)
+                raw_data = json.loads(content)
+                instance = cls.parse_and_validate(ctx, raw_data, f"file '{file_path}'")
+                if instance:
+                    return instance
+            except json.JSONDecodeError:
+                ctx.log_warning(
+                    f"Could not decode JSON from history file '{file_path}'. "
+                    "Ignoring file content."
+                )
+            except Exception as e:
+                ctx.log_warning(
+                    f"Error reading history file '{file_path}': {e}. "
+                    "Ignoring file content."
+                )
+        # If neither reader nor file provided valid data
+        return None
+
+    @classmethod
+    def parse_and_validate(
+        cls, ctx: AnyContext, data: Any, source: str
+    ) -> Optional["ConversationHistoryData"]:
+        """Parses raw data into ConversationHistoryData, handling validation & old formats."""
+        try:
+            if isinstance(data, cls):
+                return data  # Already a valid instance
+            if isinstance(data, dict) and "history" in data:
+                # Standard format {'context': ..., 'history': ...}
+                # Ensure context exists, even if empty
+                data.setdefault("context", {})
+                return cls.model_validate(data)
+            elif isinstance(data, list):
+                # Handle old format (just a list) - wrap it
+                ctx.log_warning(
+                    f"History from {source} contains old list format. "
+                    "Wrapping it into the new structure {'context': {}, 'history': [...]}. "
+                    "Consider updating the source format."
+                )
+                return cls(history=data, context={})
+            else:
+                ctx.log_warning(
+                    f"History data from {source} has unexpected format "
+                    f"(type: {type(data)}). Ignoring."
+                )
+                return None
+        except Exception as e:  # Catch validation errors too
+            ctx.log_warning(
+                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