zrb 1.11.0__py3-none-any.whl → 1.13.0__py3-none-any.whl

zrb/task/llm/agent.py

@@ -7,15 +7,15 @@ 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.tool_wrapper import wrap_func, wrap_tool
 from zrb.task.llm.typing import ListOfDict
 
 if TYPE_CHECKING:
     from pydantic_ai import Agent, Tool
     from pydantic_ai.agent import AgentRun
-    from pydantic_ai.mcp import MCPServer
     from pydantic_ai.models import Model
     from pydantic_ai.settings import ModelSettings
+    from pydantic_ai.toolsets import AbstractToolset
 
     ToolOrCallable = Tool | Callable
 else:
@@ -28,26 +28,43 @@ def create_agent_instance(
     system_prompt: str = "",
     model_settings: "ModelSettings | None" = None,
     tools: list[ToolOrCallable] = [],
-    mcp_servers: list["MCPServer"] = [],
+    toolsets: list["AbstractToolset[Agent]"] = [],
     retries: int = 3,
 ) -> "Agent":
     """Creates a new Agent instance with configured tools and servers."""
     from pydantic_ai import Agent, Tool
+    from pydantic_ai.tools import GenerateToolJsonSchema
 
     # Normalize tools
     tool_list = []
     for tool_or_callable in tools:
         if isinstance(tool_or_callable, Tool):
             tool_list.append(tool_or_callable)
+            # Update tool's function
+            tool = tool_or_callable
+            tool_list.append(
+                Tool(
+                    function=wrap_func(tool.function),
+                    takes_ctx=tool.takes_ctx,
+                    max_retries=tool.max_retries,
+                    name=tool.name,
+                    description=tool.description,
+                    prepare=tool.prepare,
+                    docstring_format=tool.docstring_format,
+                    require_parameter_descriptions=tool.require_parameter_descriptions,
+                    schema_generator=GenerateToolJsonSchema,
+                    strict=tool.strict,
+                )
+            )
         else:
-            # Pass ctx to wrap_tool
+            # Turn function into tool
             tool_list.append(wrap_tool(tool_or_callable, ctx))
     # Return Agent
     return Agent(
         model=model,
         system_prompt=system_prompt,
         tools=tool_list,
-        toolsets=mcp_servers,
+        toolsets=toolsets,
         model_settings=model_settings,
         retries=retries,
     )
@@ -63,8 +80,8 @@ def get_agent(
         list[ToolOrCallable] | Callable[[AnySharedContext], list[ToolOrCallable]]
     ),
     additional_tools: list[ToolOrCallable],
-    mcp_servers_attr: "list[MCPServer] | Callable[[AnySharedContext], list[MCPServer]]",
-    additional_mcp_servers: "list[MCPServer]",
+    toolsets_attr: "list[AbstractToolset[Agent]] | Callable[[AnySharedContext], list[AbstractToolset[Agent]]]",  # noqa
+    additional_toolsets: "list[AbstractToolset[Agent]]",
     retries: int = 3,
 ) -> "Agent":
     """Retrieves the configured Agent instance or creates one if necessary."""
@@ -85,18 +102,16 @@ def get_agent(
     # Get tools for agent
     tools = list(tools_attr(ctx) if callable(tools_attr) else tools_attr)
     tools.extend(additional_tools)
-    # Get MCP Servers for agent
-    mcp_servers = list(
-        mcp_servers_attr(ctx) if callable(mcp_servers_attr) else mcp_servers_attr
-    )
-    mcp_servers.extend(additional_mcp_servers)
+    # Get Toolsets for agent
+    tool_sets = list(toolsets_attr(ctx) if callable(toolsets_attr) else toolsets_attr)
+    tool_sets.extend(additional_toolsets)
     # If no agent provided, create one using the configuration
     return create_agent_instance(
         ctx=ctx,
         model=model,
         system_prompt=system_prompt,
         tools=tools,
-        mcp_servers=mcp_servers,
+        toolsets=tool_sets,
         model_settings=model_settings,
         retries=retries,
     )

zrb/task/llm/conversation_history_model.py

@@ -176,46 +176,23 @@ class ConversationHistory:
         """
         return json.dumps({"content": self._fetch_long_term_note()})
 
-    def add_long_term_info(self, new_info: str) -> str:
+    def write_long_term_note(self, content: str) -> str:
         """
-        Add new info for long-term reference.
+        Write the entire content of the long-term references.
+        This will overwrite any existing long-term notes.
 
         Args:
-            new_info (str): New info to be added into long-term references.
+            content (str): The full content of the long-term notes.
 
         Returns:
-            str: JSON with new content of the notes.
-
-        Raises:
-            Exception: If the note cannot be read.
-        """
-        llm_context_config.add_to_context(new_info, cwd="/")
-        return json.dumps({"success": True, "content": self._fetch_long_term_note()})
-
-    def remove_long_term_info(self, irrelevant_info: str) -> str:
+            str: JSON indicating success.
         """
-        Remove irrelevant info from long-term reference.
-
-        Args:
-            irrelevant_info (str): Irrelevant info to be removed from long-term references.
-
-        Returns:
-            str: JSON with new content of the notes and deletion status.
-
-        Raises:
-            Exception: If the note cannot be read.
-        """
-        was_removed = llm_context_config.remove_from_context(irrelevant_info, cwd="/")
-        return json.dumps(
-            {
-                "success": was_removed,
-                "content": self._fetch_long_term_note(),
-            }
-        )
+        llm_context_config.write_context(content, context_path="/")
+        return json.dumps({"success": True})
 
     def read_contextual_note(self) -> str:
         """
-        Read the content of the contextual references.
+        Read the content of the contextual references for the current project.
 
         This tool helps you retrieve knowledge or notes stored for contextual reference.
         If the note does not exist, you may want to create it using the write tool.
@@ -228,52 +205,25 @@ class ConversationHistory:
         """
         return json.dumps({"content": self._fetch_contextual_note()})
 
-    def add_contextual_info(self, new_info: str, context_path: str | None) -> str:
-        """
-        Add new info for contextual reference.
-
-        Args:
-            new_info (str): New info to be added into contextual references.
-            context_path (str, optional): contextual directory path for new info
-
-        Returns:
-            str: JSON with new content of the notes.
-
-        Raises:
-            Exception: If the note cannot be read.
-        """
-        if context_path is None:
-            context_path = self.project_path
-        llm_context_config.add_to_context(new_info, context_path=context_path)
-        return json.dumps({"success": True, "content": self._fetch_contextual_note()})
-
-    def remove_contextual_info(
-        self, irrelevant_info: str, context_path: str | None
+    def write_contextual_note(
+        self, content: str, context_path: str | None = None
     ) -> str:
         """
-        Remove irrelevant info from contextual reference.
+        Write the entire content of the contextual references for a specific path.
+        This will overwrite any existing contextual notes for that path.
 
         Args:
-            irrelevant_info (str): Irrelevant info to be removed from contextual references.
-            context_path (str, optional): contextual directory path of the irrelevant info
+            content (str): The full content of the contextual notes.
+            context_path (str, optional): The directory path for the context.
+                                        Defaults to the current project path.
 
         Returns:
-            str: JSON with new content of the notes and deletion status.
-
-        Raises:
-            Exception: If the note cannot be read.
+            str: JSON indicating success.
         """
         if context_path is None:
             context_path = self.project_path
-        was_removed = llm_context_config.remove_from_context(
-            irrelevant_info, context_path=context_path
-        )
-        return json.dumps(
-            {
-                "success": was_removed,
-                "content": self._fetch_contextual_note(),
-            }
-        )
+        llm_context_config.write_context(content, context_path=context_path)
+        return json.dumps({"success": True})
 
     def _fetch_long_term_note(self):
         contexts = llm_context_config.get_contexts(cwd=self.project_path)

zrb/{config/default_workflow/code.md → task/llm/default_workflow/coding.md}

@@ -1,5 +1,3 @@
-# Special Instructions for Software Engineering
-
 When the user's request involves writing or modifying code, you MUST follow these domain-specific rules in addition to your core workflow.
 
 ## 1. Critical Prohibitions

zrb/{config/default_workflow/content.md → task/llm/default_workflow/copywriting.md}

@@ -1,5 +1,3 @@
-# Special Instructions for Content Creation & Management
-
 When the user's request involves creating, refining, or organizing textual content, you MUST follow these domain-specific rules in addition to your core workflow.
 
 ## 1. Core Principles

zrb/{config/default_workflow/research.md → task/llm/default_workflow/researching.md}

@@ -1,5 +1,3 @@
-# Special Instructions for Research, Analysis, and Summarization
-
 When the user's request involves finding, synthesizing, or analyzing information, you MUST follow these domain-specific rules in addition to your core workflow.
 
 ## 1. Core Principles

zrb/task/llm/history_summarization.py

@@ -146,11 +146,9 @@ async def summarize_history(
             conversation_history.write_past_conversation_summary,
             conversation_history.write_past_conversation_transcript,
             conversation_history.read_long_term_note,
-            conversation_history.add_long_term_info,
-            conversation_history.remove_long_term_info,
+            conversation_history.write_long_term_note,
             conversation_history.read_contextual_note,
-            conversation_history.add_contextual_info,
-            conversation_history.remove_contextual_info,
+            conversation_history.write_contextual_note,
         ],
     )
     try:

zrb/task/llm/print_node.py

@@ -14,6 +14,7 @@ async def print_node(print_func: Callable, agent_run: Any, node: Any):
         PartDeltaEvent,
         PartStartEvent,
         TextPartDelta,
+        ThinkingPartDelta,
         ToolCallPartDelta,
     )
 
@@ -33,7 +34,9 @@ async def print_node(print_func: Callable, agent_run: Any, node: Any):
                     )
                     is_streaming = False
                 elif isinstance(event, PartDeltaEvent):
-                    if isinstance(event.delta, TextPartDelta):
+                    if isinstance(event.delta, TextPartDelta) or isinstance(
+                        event.delta, ThinkingPartDelta
+                    ):
                         print_func(
                             stylize_faint(f"{event.delta.content_delta}"),
                             end="",

zrb/task/llm/prompt.py

@@ -3,11 +3,12 @@ import platform
 import re
 from datetime import datetime, timezone
 
-from zrb.attr.type import StrAttr
+from zrb.attr.type import StrAttr, StrListAttr
 from zrb.config.llm_config import llm_config as llm_config
+from zrb.config.llm_context.config import llm_context_config
 from zrb.context.any_context import AnyContext
 from zrb.task.llm.conversation_history_model import ConversationHistory
-from zrb.util.attr import get_attr, get_str_attr
+from zrb.util.attr import get_attr, get_str_attr, get_str_list_attr
 from zrb.util.file import read_dir, read_file_with_line_numbers
 from zrb.util.llm.prompt import make_prompt_section
 
@@ -15,13 +16,14 @@ from zrb.util.llm.prompt import make_prompt_section
 def get_persona(
     ctx: AnyContext,
     persona_attr: StrAttr | None,
+    render_persona: bool,
 ) -> str:
     """Gets the persona, prioritizing task-specific, then default."""
     persona = get_attr(
         ctx,
         persona_attr,
         None,
-        auto_render=False,
+        auto_render=render_persona,
     )
     if persona is not None:
         return persona
@@ -31,13 +33,14 @@ def get_persona(
 def get_base_system_prompt(
     ctx: AnyContext,
     system_prompt_attr: StrAttr | None,
+    render_system_prompt: bool,
 ) -> str:
     """Gets the base system prompt, prioritizing task-specific, then default."""
     system_prompt = get_attr(
         ctx,
         system_prompt_attr,
         None,
-        auto_render=False,
+        auto_render=render_system_prompt,
     )
     if system_prompt is not None:
         return system_prompt
@@ -47,33 +50,95 @@ def get_base_system_prompt(
 def get_special_instruction_prompt(
     ctx: AnyContext,
     special_instruction_prompt_attr: StrAttr | None,
+    render_spcecial_instruction_prompt: bool,
 ) -> str:
     """Gets the special instruction prompt, prioritizing task-specific, then default."""
     special_instruction = get_attr(
         ctx,
         special_instruction_prompt_attr,
         None,
-        auto_render=False,
+        auto_render=render_spcecial_instruction_prompt,
     )
     if special_instruction is not None:
         return special_instruction
     return llm_config.default_special_instruction_prompt
 
 
+def get_modes(
+    ctx: AnyContext,
+    modes_attr: StrAttr | None,
+    render_modes: bool,
+) -> str:
+    """Gets the modes, prioritizing task-specific, then default."""
+    raw_modes = get_str_list_attr(
+        ctx,
+        modes_attr,
+        auto_render=render_modes,
+    )
+    modes = [mode.strip() for mode in raw_modes if mode.strip() != ""]
+    if len(modes) > 0:
+        return modes
+    return llm_config.default_modes or []
+
+
+def get_workflow_prompt(
+    ctx: AnyContext,
+    modes_attr: StrAttr | None,
+    render_modes: bool,
+) -> str:
+    modes = get_modes(ctx, modes_attr, render_modes)
+    # Get user-defined workflows
+    workflows = {
+        workflow_name: content
+        for workflow_name, content in llm_context_config.get_workflows().items()
+        if workflow_name in modes
+    }
+    # Get requested builtin-workflow names
+    requested_builtin_workflow_names = [
+        workflow_name
+        for workflow_name in ("coding", "copywriting", "researching")
+        if workflow_name in modes and workflow_name not in workflows
+    ]
+    # add builtin-workflows if requested
+    if len(requested_builtin_workflow_names) > 0:
+        dir_path = os.path.dirname(__file__)
+        for workflow_name in requested_builtin_workflow_names:
+            workflow_file_path = os.path.join(
+                dir_path, "default_workflow", f"{workflow_name}.md"
+            )
+            with open(workflow_file_path, "r") as f:
+                workflows[workflow_name] = f.read()
+    return "\n".join(
+        [
+            make_prompt_section(header.capitalize(), content)
+            for header, content in workflows.items()
+            if header.lower() in modes
+        ]
+    )
+
+
 def get_system_and_user_prompt(
     ctx: AnyContext,
     user_message: str,
     persona_attr: StrAttr | None = None,
+    render_persona: bool = False,
     system_prompt_attr: StrAttr | None = None,
+    render_system_prompt: bool = False,
     special_instruction_prompt_attr: StrAttr | None = None,
+    render_special_instruction_prompt: bool = False,
+    modes_attr: StrListAttr | None = None,
+    render_modes: bool = False,
     conversation_history: ConversationHistory | None = None,
 ) -> tuple[str, str]:
     """Combines persona, base system prompt, and special instructions."""
-    persona = get_persona(ctx, persona_attr)
-    base_system_prompt = get_base_system_prompt(ctx, system_prompt_attr)
-    special_instruction = get_special_instruction_prompt(
-        ctx, special_instruction_prompt_attr
+    persona = get_persona(ctx, persona_attr, render_persona)
+    base_system_prompt = get_base_system_prompt(
+        ctx, system_prompt_attr, render_system_prompt
     )
+    special_instruction_prompt = get_special_instruction_prompt(
+        ctx, special_instruction_prompt_attr, render_special_instruction_prompt
+    )
+    workflow_prompt = get_workflow_prompt(ctx, modes_attr, render_modes)
     if conversation_history is None:
         conversation_history = ConversationHistory()
     conversation_context, new_user_message = extract_conversation_context(user_message)
@@ -81,7 +146,8 @@ def get_system_and_user_prompt(
         [
             make_prompt_section("Persona", persona),
             make_prompt_section("System Prompt", base_system_prompt),
-            make_prompt_section("Special Instruction", special_instruction),
+            make_prompt_section("Special Instruction", special_instruction_prompt),
+            make_prompt_section("Special Workflows", workflow_prompt),
             make_prompt_section(
                 "Past Conversation",
                 "\n".join(
@@ -194,30 +260,15 @@ def get_user_message(
 def get_summarization_system_prompt(
     ctx: AnyContext,
     summarization_prompt_attr: StrAttr | None,
+    render_summarization_prompt: bool,
 ) -> str:
     """Gets the summarization prompt, rendering if configured and handling defaults."""
     summarization_prompt = get_attr(
         ctx,
         summarization_prompt_attr,
         None,
-        auto_render=False,
+        auto_render=render_summarization_prompt,
     )
     if summarization_prompt is not None:
         return summarization_prompt
     return llm_config.default_summarization_prompt
-
-
-def get_context_enrichment_prompt(
-    ctx: AnyContext,
-    context_enrichment_prompt_attr: StrAttr | None,
-) -> str:
-    """Gets the context enrichment prompt, rendering if configured and handling defaults."""
-    context_enrichment_prompt = get_attr(
-        ctx,
-        context_enrichment_prompt_attr,
-        None,
-        auto_render=False,
-    )
-    if context_enrichment_prompt is not None:
-        return context_enrichment_prompt
-    return llm_config.default_context_enrichment_prompt

zrb/task/llm/tool_wrapper.py

@@ -5,9 +5,12 @@ import typing
 from collections.abc import Callable
 from typing import TYPE_CHECKING
 
+from zrb.config.config import CFG
 from zrb.context.any_context import AnyContext
 from zrb.task.llm.error import ToolExecutionError
+from zrb.util.callable import get_callable_name
 from zrb.util.run import run_async
+from zrb.util.string.conversion import to_boolean
 
 if TYPE_CHECKING:
     from pydantic_ai import Tool
@@ -18,16 +21,19 @@ def wrap_tool(func: Callable, ctx: AnyContext) -> "Tool":
     from pydantic_ai import RunContext, Tool
 
     original_sig = inspect.signature(func)
-    # Use helper function for clarity
     needs_run_context_for_pydantic = _has_context_parameter(original_sig, RunContext)
+    wrapper = wrap_func(func, ctx)
+    return Tool(wrapper, takes_ctx=needs_run_context_for_pydantic)
+
+
+def wrap_func(func: Callable, ctx: AnyContext) -> Callable:
+    original_sig = inspect.signature(func)
     needs_any_context_for_injection = _has_context_parameter(original_sig, AnyContext)
     takes_no_args = len(original_sig.parameters) == 0
     # Pass individual flags to the wrapper creator
     wrapper = _create_wrapper(func, original_sig, ctx, needs_any_context_for_injection)
-    # Adjust signature - _adjust_signature determines exclusions based on type
     _adjust_signature(wrapper, original_sig, takes_no_args)
-    # takes_ctx in pydantic-ai Tool is specifically for RunContext
-    return Tool(wrapper, takes_ctx=needs_run_context_for_pydantic)
+    return wrapper
 
 
 def _has_context_parameter(original_sig: inspect.Signature, context_type: type) -> bool:
@@ -71,13 +77,11 @@ def _create_wrapper(
     async def wrapper(*args, **kwargs):
         # Identify AnyContext parameter name from the original signature if needed
         any_context_param_name = None
-
         if needs_any_context_for_injection:
             for param in original_sig.parameters.values():
                 if _is_annotated_with_context(param.annotation, AnyContext):
                     any_context_param_name = param.name
                     break  # Found it, no need to continue
-
             if any_context_param_name is None:
                 # This should not happen if needs_any_context_for_injection is True,
                 # but check for safety
@@ -87,24 +91,25 @@ def _create_wrapper(
             # Inject the captured ctx into kwargs. This will overwrite if the LLM
             # somehow provided it.
             kwargs[any_context_param_name] = ctx
-
         # If the dummy argument was added for schema generation and is present in kwargs,
         # remove it before calling the original function, unless the original function
         # actually expects a parameter named '_dummy'.
         if "_dummy" in kwargs and "_dummy" not in original_sig.parameters:
             del kwargs["_dummy"]
-
         try:
-            # Call the original function.
-            # pydantic-ai is responsible for injecting RunContext if takes_ctx is True.
-            # Our wrapper injects AnyContext if needed.
-            # The arguments received by the wrapper (*args, **kwargs) are those
-            # provided by the LLM, potentially with RunContext already injected by
-            # pydantic-ai if takes_ctx is True. We just need to ensure AnyContext
-            # is injected if required by the original function.
-            # The dummy argument handling is moved to _adjust_signature's logic
-            # for schema generation, it's not needed here before calling the actual
-            # function.
+            if not CFG.LLM_YOLO_MODE and not ctx.is_web_mode and ctx.is_tty:
+                func_name = get_callable_name(func)
+                ctx.print(f"✅ >> Allow to run tool: {func_name} (Y/n)", plain=True)
+                user_confirmation_str = await _read_line()
+                try:
+                    user_confirmation = to_boolean(user_confirmation_str)
+                except Exception:
+                    user_confirmation = False
+                if not user_confirmation:
+                    ctx.print(f"❌ >> Rejecting {func_name} call. Why?", plain=True)
+                    reason = await _read_line()
+                    ctx.print("", plain=True)
+                    raise ValueError(f"User disapproval: {reason}")
             return await run_async(func(*args, **kwargs))
         except Exception as e:
             error_model = ToolExecutionError(
@@ -118,6 +123,13 @@ def _create_wrapper(
     return wrapper
 
 
+async def _read_line():
+    from prompt_toolkit import PromptSession
+
+    reader = PromptSession()
+    return await reader.prompt_async()
+
+
 def _adjust_signature(
     wrapper: Callable, original_sig: inspect.Signature, takes_no_args: bool
 ):