zrb 1.15.3__py3-none-any.whl → 2.0.0a4__py3-none-any.whl

zrb/task/llm/print_node.py

@@ -1,101 +0,0 @@
-from collections.abc import Callable
-from typing import Any
-
-from zrb.util.cli.style import stylize_faint
-
-
-async def print_node(print_func: Callable, agent_run: Any, node: Any):
-    """Prints the details of an agent execution node using a provided print function."""
-    from pydantic_ai import Agent
-    from pydantic_ai.messages import (
-        FinalResultEvent,
-        FunctionToolCallEvent,
-        FunctionToolResultEvent,
-        PartDeltaEvent,
-        PartStartEvent,
-        TextPartDelta,
-        ThinkingPartDelta,
-        ToolCallPartDelta,
-    )
-
-    if Agent.is_user_prompt_node(node):
-        print_func(stylize_faint(" 🔠 Receiving input..."))
-    elif Agent.is_model_request_node(node):
-        # A model request node => We can stream tokens from the model's request
-        print_func(stylize_faint(" 🧠 Processing..."))
-        # Reference: https://ai.pydantic.dev/agents/#streaming
-        async with node.stream(agent_run.ctx) as request_stream:
-            is_streaming = False
-            async for event in request_stream:
-                if isinstance(event, PartStartEvent) and event.part:
-                    if is_streaming:
-                        print_func("")
-                    content = _get_event_part_content(event)
-                    print_func(stylize_faint(f"     {content}"), end="")
-                    is_streaming = False
-                elif isinstance(event, PartDeltaEvent):
-                    if isinstance(event.delta, TextPartDelta) or isinstance(
-                        event.delta, ThinkingPartDelta
-                    ):
-                        print_func(
-                            stylize_faint(f"{event.delta.content_delta}"),
-                            end="",
-                        )
-                    elif isinstance(event.delta, ToolCallPartDelta):
-                        print_func(
-                            stylize_faint(f"{event.delta.args_delta}"),
-                            end="",
-                        )
-                    is_streaming = True
-                elif isinstance(event, FinalResultEvent) and event.tool_name:
-                    if is_streaming:
-                        print_func("")
-                    print_func(
-                        stylize_faint(f"     Result: tool_name={event.tool_name}"),
-                    )
-                    is_streaming = False
-            if is_streaming:
-                print_func("")
-    elif Agent.is_call_tools_node(node):
-        # A handle-response node => The model returned some data, potentially calls a tool
-        print_func(stylize_faint(" 🧰 Calling Tool..."))
-        async with node.stream(agent_run.ctx) as handle_stream:
-            async for event in handle_stream:
-                if isinstance(event, FunctionToolCallEvent):
-                    # Handle empty arguments across different providers
-                    if event.part.args == "" or event.part.args is None:
-                        event.part.args = {}
-                    elif isinstance(
-                        event.part.args, str
-                    ) and event.part.args.strip() in ["null", "{}"]:
-                        # Some providers might send "null" or "{}" as a string
-                        event.part.args = {}
-                    # Handle dummy property if present (from our schema sanitization)
-                    if (
-                        isinstance(event.part.args, dict)
-                        and "_dummy" in event.part.args
-                    ):
-                        del event.part.args["_dummy"]
-                    print_func(
-                        stylize_faint(
-                            f"     {event.part.tool_call_id} | "
-                            f"Call {event.part.tool_name} {event.part.args}"
-                        )
-                    )
-                elif isinstance(event, FunctionToolResultEvent):
-                    print_func(
-                        stylize_faint(
-                            f"     {event.tool_call_id} | {event.result.content}"
-                        )
-                    )
-    elif Agent.is_end_node(node):
-        # Once an End node is reached, the agent run is complete
-        print_func(stylize_faint(" ✅ Completed..."))
-
-
-def _get_event_part_content(event: Any) -> str:
-    if not hasattr(event, "part"):
-        return f"{event}"
-    if not hasattr(event.part, "content"):
-        return f"{event.part}"
-    return getattr(event.part, "content")

zrb/task/llm/prompt.py

@@ -1,325 +0,0 @@
-import os
-import platform
-import re
-from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Callable
-
-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.context.any_shared_context import AnySharedContext
-from zrb.task.llm.conversation_history_model import ConversationHistory
-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
-
-if TYPE_CHECKING:
-    from pydantic_ai.messages import UserContent
-
-
-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=render_persona,
-    )
-    if persona is not None:
-        return persona
-    return llm_config.default_persona or ""
-
-
-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=render_system_prompt,
-    )
-    if system_prompt is not None:
-        return system_prompt
-    return llm_config.default_system_prompt or ""
-
-
-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=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: StrListAttr | None,
-    render_modes: bool,
-) -> list[str]:
-    """Gets the modes, prioritizing task-specific, then default."""
-    raw_modes = get_str_list_attr(
-        ctx,
-        [] if modes_attr is None else modes_attr,
-        auto_render=render_modes,
-    )
-    if raw_modes is None:
-        raw_modes = []
-    modes = [mode.strip().lower() 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: StrListAttr | None,
-    render_modes: bool,
-) -> str:
-    builtin_workflow_dir = os.path.join(os.path.dirname(__file__), "default_workflow")
-    modes = set(get_modes(ctx, modes_attr, render_modes))
-
-    # Get user-defined workflows
-    workflows = {
-        workflow_name.strip().lower(): content
-        for workflow_name, content in llm_context_config.get_workflows().items()
-        if workflow_name.strip().lower() in modes
-    }
-
-    # Get available builtin workflow names from the file system
-    available_builtin_workflow_names = set()
-    try:
-        for filename in os.listdir(builtin_workflow_dir):
-            if filename.endswith(".md"):
-                available_builtin_workflow_names.add(filename[:-3].lower())
-    except FileNotFoundError:
-        # Handle case where the directory might not exist
-        ctx.log_error(
-            f"Warning: Default workflow directory not found at {builtin_workflow_dir}"
-        )
-    except Exception as e:
-        # Catch other potential errors during directory listing
-        ctx.log_error(f"Error listing default workflows: {e}")
-
-    # Determine which builtin workflows are requested and not already loaded
-    requested_builtin_workflow_names = [
-        workflow_name
-        for workflow_name in available_builtin_workflow_names
-        if workflow_name in modes and workflow_name not in workflows
-    ]
-
-    # Add builtin-workflows if requested
-    if len(requested_builtin_workflow_names) > 0:
-        for workflow_name in requested_builtin_workflow_names:
-            workflow_file_path = os.path.join(
-                builtin_workflow_dir, f"{workflow_name}.md"
-            )
-            try:
-                with open(workflow_file_path, "r") as f:
-                    workflows[workflow_name] = f.read()
-            except FileNotFoundError:
-                ctx.log_error(
-                    f"Warning: Builtin workflow file not found: {workflow_file_path}"
-                )
-            except Exception as e:
-                ctx.log_error(f"Error reading builtin workflow {workflow_name}: {e}")
-
-    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, 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)
-    new_system_prompt = "\n".join(
-        [
-            make_prompt_section("Persona", persona),
-            make_prompt_section("System Prompt", base_system_prompt),
-            make_prompt_section("Special Instruction", special_instruction_prompt),
-            make_prompt_section("Special Workflows", workflow_prompt),
-            make_prompt_section(
-                "Past Conversation",
-                "\n".join(
-                    [
-                        make_prompt_section(
-                            "Summary",
-                            conversation_history.past_conversation_summary,
-                            as_code=True,
-                        ),
-                        make_prompt_section(
-                            "Last Transcript",
-                            conversation_history.past_conversation_transcript,
-                            as_code=True,
-                        ),
-                    ]
-                ),
-            ),
-            make_prompt_section(
-                "Notes",
-                "\n".join(
-                    [
-                        make_prompt_section(
-                            "Long Term",
-                            conversation_history.long_term_note,
-                            as_code=True,
-                        ),
-                        make_prompt_section(
-                            "Contextual",
-                            conversation_history.contextual_note,
-                            as_code=True,
-                        ),
-                    ]
-                ),
-            ),
-            make_prompt_section("Conversation Context", conversation_context),
-        ]
-    )
-    return new_system_prompt, new_user_message
-
-
-def extract_conversation_context(user_message: str) -> tuple[str, str]:
-    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)
-    apendixes = []
-    for i, ref in enumerate(potential_resource_path):
-        resource_path = os.path.abspath(os.path.expanduser(ref))
-        content = ""
-        ref_type = ""
-        if os.path.isfile(resource_path):
-            content = read_file_with_line_numbers(resource_path)
-            ref_type = "file"
-        elif os.path.isdir(resource_path):
-            content = read_dir(resource_path)
-            ref_type = "directory"
-        if content != "":
-            # Replace the @-reference in the user message with the placeholder
-            placeholder = f"[Reference {i+1}: {os.path.basename(ref)}]"
-            modified_user_message = modified_user_message.replace(
-                f"@{ref}", placeholder, 1
-            )
-            apendixes.append(
-                make_prompt_section(
-                    f"{placeholder} ({ref_type} path: `{resource_path}`)",
-                    content,
-                    as_code=True,
-                )
-            )
-    conversation_context = "\n".join(
-        [
-            make_prompt_section("Current OS", platform.system()),
-            make_prompt_section("OS Version", platform.version()),
-            make_prompt_section("Python Version", platform.python_version()),
-        ]
-    )
-    iso_date = datetime.now(timezone.utc).astimezone().isoformat()
-    current_directory = os.getcwd()
-    modified_user_message = "\n".join(
-        [
-            make_prompt_section("User Message", modified_user_message),
-            make_prompt_section(
-                "Context",
-                "\n".join(
-                    [
-                        make_prompt_section(
-                            "Current working directory", current_directory
-                        ),
-                        make_prompt_section("Current time", iso_date),
-                        make_prompt_section("Apendixes", "\n".join(apendixes)),
-                    ]
-                ),
-            ),
-        ]
-    )
-    return conversation_context, modified_user_message
-
-
-def get_user_message(
-    ctx: AnyContext,
-    message_attr: StrAttr | None,
-    render_user_message: bool,
-) -> str:
-    """Gets the user message, rendering and providing a default."""
-    return get_str_attr(
-        ctx, message_attr, "How are you?", auto_render=render_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=render_summarization_prompt,
-    )
-    if summarization_prompt is not None:
-        return summarization_prompt
-    return llm_config.default_summarization_prompt
-
-
-def get_attachments(
-    ctx: AnyContext,
-    attachment: "UserContent | list[UserContent] | Callable[[AnySharedContext], UserContent | list[UserContent]] | None" = None,  # noqa
-) -> "list[UserContent]":
-    if attachment is None:
-        return []
-    if callable(attachment):
-        result = attachment(ctx)
-        if result is None:
-            return []
-        if isinstance(result, list):
-            return result
-        return [result]
-    if isinstance(attachment, list):
-        return attachment
-    return [attachment]

zrb/task/llm/tool_wrapper.py

@@ -1,220 +0,0 @@
-import functools
-import inspect
-import traceback
-import typing
-from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
-
-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.cli.style import (
-    stylize_blue,
-    stylize_error,
-    stylize_green,
-    stylize_yellow,
-)
-from zrb.util.run import run_async
-from zrb.util.string.conversion import to_boolean
-
-if TYPE_CHECKING:
-    from pydantic_ai import Tool
-
-
-def wrap_tool(func: Callable, ctx: AnyContext, is_yolo_mode: bool) -> "Tool":
-    """Wraps a tool function to handle exceptions and context propagation."""
-    from pydantic_ai import RunContext, Tool
-
-    original_sig = inspect.signature(func)
-    needs_run_context_for_pydantic = _has_context_parameter(original_sig, RunContext)
-    wrapper = wrap_func(func, ctx, is_yolo_mode)
-    return Tool(wrapper, takes_ctx=needs_run_context_for_pydantic)
-
-
-def wrap_func(func: Callable, ctx: AnyContext, is_yolo_mode: bool) -> 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, is_yolo_mode
-    )
-    _adjust_signature(wrapper, original_sig, takes_no_args)
-    return wrapper
-
-
-def _has_context_parameter(original_sig: inspect.Signature, context_type: type) -> bool:
-    """
-    Checks if the function signature includes a parameter with the specified
-    context type annotation.
-    """
-    return any(
-        _is_annotated_with_context(param.annotation, context_type)
-        for param in original_sig.parameters.values()
-    )
-
-
-def _is_annotated_with_context(param_annotation, context_type):
-    """
-    Checks if the parameter annotation is the specified context type
-    or a generic type containing it (e.g., Optional[ContextType]).
-    """
-    if param_annotation is inspect.Parameter.empty:
-        return False
-    if param_annotation is context_type:
-        return True
-    # Check for generic types like Optional[ContextType] or Union[ContextType, ...]
-    origin = typing.get_origin(param_annotation)
-    args = typing.get_args(param_annotation)
-    if origin is not None and args:
-        # Check if context_type is one of the arguments of the generic type
-        return any(arg is context_type for arg in args)
-    return False
-
-
-def _create_wrapper(
-    func: Callable,
-    original_sig: inspect.Signature,
-    ctx: AnyContext,
-    needs_any_context_for_injection: bool,
-    is_yolo_mode: bool,
-) -> Callable:
-    """Creates the core wrapper function."""
-
-    @functools.wraps(func)
-    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
-                raise ValueError(
-                    "AnyContext parameter name not found in function signature."
-                )
-            # 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:
-            if not is_yolo_mode and not ctx.is_web_mode and ctx.is_tty:
-                approval, reason = await _ask_for_approval(ctx, func, args, kwargs)
-                if not approval:
-                    raise ValueError(f"User disapproving: {reason}")
-            return await run_async(func(*args, **kwargs))
-        except Exception as e:
-            error_model = ToolExecutionError(
-                tool_name=func.__name__,
-                error_type=type(e).__name__,
-                message=str(e),
-                details=traceback.format_exc(),
-            )
-            return error_model.model_dump_json()
-
-    return wrapper
-
-
-async def _ask_for_approval(
-    ctx: AnyContext, func: Callable, args: list[Any], kwargs: dict[str, Any]
-) -> tuple[bool, str]:
-    func_name = get_callable_name(func)
-    normalized_args = [stylize_green(_truncate_arg(arg)) for arg in args]
-    normalized_kwargs = []
-    for key, val in kwargs.items():
-        truncated_val = _truncate_arg(f"{val}")
-        normalized_kwargs.append(
-            f"{stylize_yellow(key)}={stylize_green(truncated_val)}"
-        )
-    func_param_str = ",".join(normalized_args + normalized_kwargs)
-    func_call_str = (
-        f"{stylize_blue(func_name + '(')}{func_param_str}{stylize_blue(')')}"
-    )
-    while True:
-        ctx.print(
-            f"\n✅ >> Allow to run tool: {func_call_str} (Yes | No, <reason>) ",
-            plain=True,
-        )
-        user_input = await _read_line()
-        ctx.print("", plain=True)
-        user_responses = [val.strip() for val in user_input.split(",", maxsplit=1)]
-        while len(user_responses) < 2:
-            user_responses.append("")
-        approval_str, reason = user_responses
-        try:
-            approved = True if approval_str.strip() == "" else to_boolean(approval_str)
-            if not approved and reason == "":
-                ctx.print(
-                    stylize_error(
-                        f"You must specify rejection reason (i.e., No, <why>) for {func_call_str}"  # noqa
-                    ),
-                    plain=True,
-                )
-                continue
-            return approved, reason
-        except Exception:
-            ctx.print(
-                stylize_error(
-                    f"Invalid approval value for {func_call_str}: {approval_str}"
-                ),
-                plain=True,
-            )
-            continue
-
-
-def _truncate_arg(arg: str, length: int = 19) -> str:
-    if len(arg) > length:
-        return f"{arg[:length-4]} ..."
-    return arg
-
-
-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
-):
-    """Adjusts the wrapper function's signature for schema generation."""
-    # The wrapper's signature should represent the arguments the *LLM* needs to provide.
-    # The LLM does not provide RunContext (pydantic-ai injects it) or AnyContext
-    # (we inject it). So, the wrapper's signature should be the original signature,
-    # minus any parameters annotated with RunContext or AnyContext.
-
-    from pydantic_ai import RunContext
-
-    params_for_schema = [
-        param
-        for param in original_sig.parameters.values()
-        if not _is_annotated_with_context(param.annotation, RunContext)
-        and not _is_annotated_with_context(param.annotation, AnyContext)
-    ]
-
-    # If after removing context parameters, there are no parameters left,
-    # and the original function took no args, keep the dummy.
-    # If after removing context parameters, there are no parameters left,
-    # but the original function *did* take args (only context), then the schema
-    # should have no parameters.
-    if not params_for_schema and takes_no_args:
-        # Keep the dummy if the original function truly had no parameters
-        new_sig = inspect.Signature(
-            parameters=[
-                inspect.Parameter(
-                    "_dummy", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=None
-                )
-            ]
-        )
-    else:
-        new_sig = inspect.Signature(parameters=params_for_schema)
-
-    wrapper.__signature__ = new_sig

zrb/task/llm/typing.py

@@ -1,3 +0,0 @@
-from typing import Any
-
-ListOfDict = list[dict[str, Any]]