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

zrb/task/llm/agent.py

@@ -1,243 +0,0 @@
-import json
-from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
-
-from zrb.config.llm_rate_limitter import LLMRateLimiter, llm_rate_limitter
-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_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.messages import UserContent
-    from pydantic_ai.models import Model
-    from pydantic_ai.settings import ModelSettings
-    from pydantic_ai.toolsets import AbstractToolset
-
-    ToolOrCallable = Tool | Callable
-
-
-def create_agent_instance(
-    ctx: AnyContext,
-    model: "str | Model",
-    system_prompt: str = "",
-    model_settings: "ModelSettings | None" = None,
-    tools: "list[ToolOrCallable]" = [],
-    toolsets: list["AbstractToolset[Agent]"] = [],
-    retries: int = 3,
-    is_yolo_mode: bool | None = None,
-) -> "Agent":
-    """Creates a new Agent instance with configured tools and servers."""
-    from pydantic_ai import Agent, Tool
-    from pydantic_ai.tools import GenerateToolJsonSchema
-
-    if is_yolo_mode is None:
-        is_yolo_mode = False
-    # 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, ctx, is_yolo_mode),
-                    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:
-            # Turn function into tool
-            tool_list.append(wrap_tool(tool_or_callable, ctx, is_yolo_mode))
-    # Return Agent
-    return Agent(
-        model=model,
-        system_prompt=system_prompt,
-        tools=tool_list,
-        toolsets=toolsets,
-        model_settings=model_settings,
-        retries=retries,
-    )
-
-
-def get_agent(
-    ctx: AnyContext,
-    agent_attr: "Agent | Callable[[AnySharedContext], Agent] | None",
-    model: "str | Model",
-    system_prompt: str,
-    model_settings: "ModelSettings | None",
-    tools_attr: (
-        "list[ToolOrCallable] | Callable[[AnySharedContext], list[ToolOrCallable]]"
-    ),
-    additional_tools: "list[ToolOrCallable]",
-    toolsets_attr: "list[AbstractToolset[Agent]] | Callable[[AnySharedContext], list[AbstractToolset[Agent]]]",  # noqa
-    additional_toolsets: "list[AbstractToolset[Agent]]",
-    retries: int = 3,
-    is_yolo_mode: bool | None = None,
-) -> "Agent":
-    """Retrieves the configured Agent instance or creates one if necessary."""
-    from pydantic_ai import Agent
-
-    # Render agent instance and return if agent_attr is already an agent
-    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
-    # Get tools for agent
-    tools = list(tools_attr(ctx) if callable(tools_attr) else tools_attr)
-    tools.extend(additional_tools)
-    # 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,
-        toolsets=tool_sets,
-        model_settings=model_settings,
-        retries=retries,
-        is_yolo_mode=is_yolo_mode,
-    )
-
-
-async def run_agent_iteration(
-    ctx: AnyContext,
-    agent: "Agent",
-    user_prompt: str,
-    attachments: "list[UserContent] | None" = None,
-    history_list: ListOfDict | None = None,
-    rate_limitter: LLMRateLimiter | None = None,
-    max_retry: int = 2,
-) -> "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.
-    """
-    if max_retry < 0:
-        raise ValueError("Max retry cannot be less than 0")
-    attempt = 0
-    while attempt < max_retry:
-        try:
-            return await _run_single_agent_iteration(
-                ctx=ctx,
-                agent=agent,
-                user_prompt=user_prompt,
-                attachments=[] if attachments is None else attachments,
-                history_list=[] if history_list is None else history_list,
-                rate_limitter=(
-                    llm_rate_limitter if rate_limitter is None else rate_limitter
-                ),
-            )
-        except BaseException:
-            attempt += 1
-            if attempt == max_retry:
-                raise
-    raise Exception("Max retry exceeded")
-
-
-async def _run_single_agent_iteration(
-    ctx: AnyContext,
-    agent: "Agent",
-    user_prompt: str,
-    attachments: "list[UserContent]",
-    history_list: ListOfDict,
-    rate_limitter: LLMRateLimiter,
-) -> "AgentRun":
-    from openai import APIError
-    from pydantic_ai.messages import ModelMessagesTypeAdapter
-
-    agent_payload = _estimate_request_payload(
-        agent, user_prompt, attachments, history_list
-    )
-    if rate_limitter:
-        await rate_limitter.throttle(agent_payload)
-    else:
-        await llm_rate_limitter.throttle(agent_payload)
-
-    user_prompt_with_attachments = [user_prompt] + attachments
-    async with agent:
-        async with agent.iter(
-            user_prompt=user_prompt_with_attachments,
-            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(_get_plain_printer(ctx), 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
-
-
-def _estimate_request_payload(
-    agent: "Agent",
-    user_prompt: str,
-    attachments: "list[UserContent]",
-    history_list: ListOfDict,
-) -> str:
-    system_prompts = agent._system_prompts if hasattr(agent, "_system_prompts") else ()
-    return json.dumps(
-        [
-            {"role": "system", "content": "\n".join(system_prompts)},
-            *history_list,
-            {"role": "user", "content": user_prompt},
-            *[_estimate_attachment_payload(attachment) for attachment in attachments],
-        ]
-    )
-
-
-def _estimate_attachment_payload(attachment: "UserContent") -> Any:
-    if hasattr(attachment, "url"):
-        return {"role": "user", "content": attachment.url}
-    if hasattr(attachment, "data"):
-        return {"role": "user", "content": "x" * len(attachment.data)}
-    return ""
-
-
-def _get_plain_printer(ctx: AnyContext):
-    def printer(*args, **kwargs):
-        if "plain" not in kwargs:
-            kwargs["plain"] = True
-        return ctx.print(*args, **kwargs)
-
-    return printer

zrb/task/llm/config.py

@@ -1,103 +0,0 @@
-from typing import TYPE_CHECKING, Callable
-
-if TYPE_CHECKING:
-    from pydantic_ai.models import Model
-    from pydantic_ai.settings import ModelSettings
-
-from zrb.attr.type import BoolAttr, StrAttr, fstring
-from zrb.config.llm_config import LLMConfig, llm_config
-from zrb.context.any_context import AnyContext
-from zrb.context.any_shared_context import AnySharedContext
-from zrb.util.attr import get_attr, get_bool_attr
-
-
-def get_is_yolo_mode(
-    ctx: AnyContext,
-    is_yolo_mode_attr: BoolAttr | None = None,
-    render_yolo_mode: bool = True,
-):
-    return get_bool_attr(
-        ctx,
-        is_yolo_mode_attr,
-        llm_config.default_yolo_mode,
-        auto_render=render_yolo_mode,
-    )
-
-
-def get_model_settings(
-    ctx: AnyContext,
-    model_settings_attr: (
-        "ModelSettings | Callable[[AnySharedContext], ModelSettings] | None"
-    ) = None,
-) -> "ModelSettings | None":
-    """Gets the model settings, resolving callables if necessary."""
-    model_settings = get_attr(ctx, model_settings_attr, None, auto_render=False)
-    if model_settings is None:
-        return llm_config.default_model_settings
-    return model_settings
-
-
-def get_model_base_url(
-    ctx: AnyContext,
-    model_base_url_attr: StrAttr | None = None,
-    render_model_base_url: bool = True,
-) -> 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 base_url is None and llm_config.default_model_base_url is not None:
-        return llm_config.default_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 = None,
-    render_model_api_key: bool = True,
-) -> 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 api_key is None and llm_config.default_api_key is not None:
-        return llm_config.default_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 = None,
-    render_model_base_url: bool = True,
-    model_api_key_attr: StrAttr | None = None,
-    render_model_api_key: bool = True,
-) -> "str | Model":
-    """Gets the model instance or name, handling defaults and configuration."""
-    from pydantic_ai.models import Model
-
-    model = get_attr(ctx, model_attr, None, auto_render=render_model)
-    if model is None:
-        return llm_config.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)
-        new_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 = llm_config.default_model_provider
-            if default_model_provider is not None:
-                new_llm_config.set_default_model_provider(default_model_provider)
-        return new_llm_config.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/conversation_history.py

@@ -1,128 +0,0 @@
-import json
-from collections.abc import Callable
-from copy import deepcopy
-from typing import Any
-
-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.conversation_history_model import ConversationHistory
-from zrb.task.llm.typing import ListOfDict
-from zrb.util.attr import get_str_attr
-from zrb.util.file import write_file
-from zrb.util.run import run_async
-
-
-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], ConversationHistory | dict | list | None] | None
-    ),
-    conversation_history_file_attr: StrAttr | None,
-    render_history_file: bool,
-    conversation_history_attr: (
-        ConversationHistory
-        | Callable[[AnySharedContext], ConversationHistory | dict | list]
-        | dict
-        | list
-    ),
-) -> ConversationHistory:
-    """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 ConversationHistory.read_from_source(
-        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 = ConversationHistory.parse_and_validate(
-            ctx, raw_data_attr, "attribute"
-        )
-        if history_data:
-            return history_data
-    # Fallback: Return default value
-    return ConversationHistory()
-
-
-async def write_conversation_history(
-    ctx: AnyContext,
-    history_data: ConversationHistory,
-    conversation_history_writer: (
-        Callable[[AnySharedContext, ConversationHistory], 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, json.dumps(history_data.to_dict(), indent=2))
-
-
-def replace_system_prompt_in_history(
-    history_list: ListOfDict, replacement: str = "<main LLM system prompt>"
-) -> ListOfDict:
-    """
-    Returns a new history list where any part with part_kind 'system-prompt'
-    has its 'content' replaced with the given replacement string.
-    Args:
-        history: List of history items (each item is a dict with a 'parts' list).
-        replacement: The string to use in place of system-prompt content.
-
-    Returns:
-        A deep-copied list of history items with system-prompt content replaced.
-    """
-    new_history = deepcopy(history_list)
-    for item in new_history:
-        parts = item.get("parts", [])
-        for part in parts:
-            if part.get("part_kind") == "system-prompt":
-                part["content"] = replacement
-    return new_history
-
-
-def count_part_in_history_list(history_list: ListOfDict) -> int:
-    """Calculates the total number of 'parts' in a history list."""
-    history_part_len = 0
-    for history in history_list:
-        if "parts" in history:
-            history_part_len += len(history["parts"])
-        else:
-            history_part_len += 1
-    return history_part_len