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

zrb/config/llm_rate_limitter.py

@@ -1,190 +0,0 @@
-import asyncio
-import json
-import time
-from collections import deque
-from typing import Any, Callable
-
-from zrb.config.config import CFG
-
-
-class LLMRateLimitter:
-    """
-    Helper class to enforce LLM API rate limits and throttling.
-    Tracks requests and tokens in a rolling 60-second window.
-    """
-
-    def __init__(
-        self,
-        max_requests_per_minute: int | None = None,
-        max_tokens_per_minute: int | None = None,
-        max_tokens_per_request: int | None = None,
-        max_tokens_per_tool_call_result: int | None = None,
-        throttle_sleep: float | None = None,
-        use_tiktoken: bool | None = None,
-        tiktoken_encoding_name: str | None = None,
-    ):
-        self._max_requests_per_minute = max_requests_per_minute
-        self._max_tokens_per_minute = max_tokens_per_minute
-        self._max_tokens_per_request = max_tokens_per_request
-        self._max_tokens_per_tool_call_result = max_tokens_per_tool_call_result
-        self._throttle_sleep = throttle_sleep
-        self._use_tiktoken = use_tiktoken
-        self._tiktoken_encoding_name = tiktoken_encoding_name
-        self.request_times = deque()
-        self.token_times = deque()
-
-    @property
-    def max_requests_per_minute(self) -> int:
-        if self._max_requests_per_minute is not None:
-            return self._max_requests_per_minute
-        return CFG.LLM_MAX_REQUESTS_PER_MINUTE
-
-    @property
-    def max_tokens_per_minute(self) -> int:
-        if self._max_tokens_per_minute is not None:
-            return self._max_tokens_per_minute
-        return CFG.LLM_MAX_TOKENS_PER_MINUTE
-
-    @property
-    def max_tokens_per_request(self) -> int:
-        if self._max_tokens_per_request is not None:
-            return self._max_tokens_per_request
-        return CFG.LLM_MAX_TOKENS_PER_REQUEST
-
-    @property
-    def max_tokens_per_tool_call_result(self) -> int:
-        if self._max_tokens_per_tool_call_result is not None:
-            return self._max_tokens_per_tool_call_result
-        return CFG.LLM_MAX_TOKENS_PER_TOOL_CALL_RESULT
-
-    @property
-    def throttle_sleep(self) -> float:
-        if self._throttle_sleep is not None:
-            return self._throttle_sleep
-        return CFG.LLM_THROTTLE_SLEEP
-
-    @property
-    def use_tiktoken(self) -> bool:
-        if self._use_tiktoken is not None:
-            return self._use_tiktoken
-        return CFG.USE_TIKTOKEN
-
-    @property
-    def tiktoken_encoding_name(self) -> str:
-        if self._tiktoken_encoding_name is not None:
-            return self._tiktoken_encoding_name
-        return CFG.TIKTOKEN_ENCODING_NAME
-
-    def set_max_requests_per_minute(self, value: int):
-        self._max_requests_per_minute = value
-
-    def set_max_tokens_per_minute(self, value: int):
-        self._max_tokens_per_minute = value
-
-    def set_max_tokens_per_request(self, value: int):
-        self._max_tokens_per_request = value
-
-    def set_max_tokens_per_tool_call_result(self, value: int):
-        self._max_tokens_per_tool_call_result = value
-
-    def set_throttle_sleep(self, value: float):
-        self._throttle_sleep = value
-
-    def count_token(self, prompt: Any) -> int:
-        str_prompt = self._prompt_to_str(prompt)
-        if not self.use_tiktoken:
-            return self._fallback_count_token(str_prompt)
-        try:
-            import tiktoken
-
-            enc = tiktoken.get_encoding(self.tiktoken_encoding_name)
-            return len(enc.encode(str_prompt))
-        except Exception:
-            return self._fallback_count_token(str_prompt)
-
-    def _fallback_count_token(self, str_prompt: str) -> int:
-        return len(str_prompt) // 4
-
-    def clip_prompt(self, prompt: Any, limit: int) -> str:
-        str_prompt = self._prompt_to_str(prompt)
-        if not self.use_tiktoken:
-            return self._fallback_clip_prompt(str_prompt, limit)
-        try:
-            import tiktoken
-
-            enc = tiktoken.get_encoding(self.tiktoken_encoding_name)
-            tokens = enc.encode(str_prompt)
-            if len(tokens) <= limit:
-                return str_prompt
-            truncated = tokens[: limit - 3]
-            clipped_text = enc.decode(truncated)
-            return clipped_text + "..."
-        except Exception:
-            return self._fallback_clip_prompt(str_prompt, limit)
-
-    def _fallback_clip_prompt(self, str_prompt: str, limit: int) -> str:
-        char_limit = limit * 4 if limit * 4 <= 10 else limit * 4 - 10
-        return str_prompt[:char_limit] + "..."
-
-    async def throttle(
-        self,
-        prompt: Any,
-        throttle_notif_callback: Callable[[str], Any] | None = None,
-    ):
-        now = time.time()
-        str_prompt = self._prompt_to_str(prompt)
-        tokens = self.count_token(str_prompt)
-        # Clean up old entries
-        while self.request_times and now - self.request_times[0] > 60:
-            self.request_times.popleft()
-        while self.token_times and now - self.token_times[0][0] > 60:
-            self.token_times.popleft()
-        # Check per-request token limit
-        if tokens > self.max_tokens_per_request:
-            raise ValueError(
-                (
-                    "Request exceeds max_tokens_per_request "
-                    f"({tokens} > {self.max_tokens_per_request})."
-                )
-            )
-        if tokens > self.max_tokens_per_minute:
-            raise ValueError(
-                (
-                    "Request exceeds max_tokens_per_minute "
-                    f"({tokens} > {self.max_tokens_per_minute})."
-                )
-            )
-        # Wait if over per-minute request or token limit
-        while (
-            len(self.request_times) >= self.max_requests_per_minute
-            or sum(t for _, t in self.token_times) + tokens > self.max_tokens_per_minute
-        ):
-            if throttle_notif_callback is not None:
-                if len(self.request_times) >= self.max_requests_per_minute:
-                    rpm = len(self.request_times)
-                    throttle_notif_callback(
-                        f"Max request per minute exceeded: {rpm} of {self.max_requests_per_minute}"
-                    )
-                else:
-                    tpm = sum(t for _, t in self.token_times) + tokens
-                    throttle_notif_callback(
-                        f"Max token per minute exceeded: {tpm} of {self.max_tokens_per_minute}"
-                    )
-            await asyncio.sleep(self.throttle_sleep)
-            now = time.time()
-            while self.request_times and now - self.request_times[0] > 60:
-                self.request_times.popleft()
-            while self.token_times and now - self.token_times[0][0] > 60:
-                self.token_times.popleft()
-        # Record this request
-        self.request_times.append(now)
-        self.token_times.append((now, tokens))
-
-    def _prompt_to_str(self, prompt: Any) -> str:
-        try:
-            return json.dumps(prompt)
-        except Exception:
-            return f"{prompt}"
-
-
-llm_rate_limitter = LLMRateLimitter()

zrb/task/llm/agent.py

@@ -1,204 +0,0 @@
-import inspect
-from collections.abc import Callable
-from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any
-
-from zrb.config.llm_rate_limitter import LLMRateLimitter
-from zrb.context.any_context import AnyContext
-from zrb.task.llm.history_processor import create_summarize_history_processor
-from zrb.task.llm.tool_wrapper import wrap_func, wrap_tool
-
-if TYPE_CHECKING:
-    from pydantic_ai import Agent, Tool
-    from pydantic_ai._agent_graph import HistoryProcessor
-    from pydantic_ai.models import Model
-    from pydantic_ai.output import OutputDataT, OutputSpec
-    from pydantic_ai.settings import ModelSettings
-    from pydantic_ai.toolsets import AbstractToolset
-
-    ToolOrCallable = Tool | Callable
-
-
-def create_agent_instance(
-    ctx: AnyContext,
-    model: "str | Model",
-    rate_limitter: LLMRateLimitter | None = None,
-    output_type: "OutputSpec[OutputDataT]" = str,
-    system_prompt: str = "",
-    model_settings: "ModelSettings | None" = None,
-    tools: list["ToolOrCallable"] = [],
-    toolsets: list["AbstractToolset[None]"] = [],
-    retries: int = 3,
-    yolo_mode: bool | list[str] | None = None,
-    summarization_model: "Model | str | None" = None,
-    summarization_model_settings: "ModelSettings | None" = None,
-    summarization_system_prompt: str | None = None,
-    summarization_retries: int = 2,
-    summarization_token_threshold: int | None = None,
-    history_processors: list["HistoryProcessor"] | None = None,
-    auto_summarize: bool = True,
-) -> "Agent[None, Any]":
-    """Creates a new Agent instance with configured tools and servers."""
-    from pydantic_ai import Agent, RunContext, Tool
-    from pydantic_ai.tools import GenerateToolJsonSchema
-    from pydantic_ai.toolsets import ToolsetTool, WrapperToolset
-
-    @dataclass
-    class ConfirmationWrapperToolset(WrapperToolset):
-        ctx: AnyContext
-        yolo_mode: bool | list[str]
-
-        async def call_tool(
-            self, name: str, tool_args: dict, ctx: RunContext, tool: ToolsetTool[None]
-        ) -> Any:
-            # The `tool` object is passed in. Use it for inspection.
-            # Define a temporary function that performs the actual tool call.
-            async def execute_delegated_tool_call(**params):
-                # Pass all arguments down the chain.
-                return await self.wrapped.call_tool(name, tool_args, ctx, tool)
-
-            # For the confirmation UI, make our temporary function look like the real one.
-            try:
-                execute_delegated_tool_call.__name__ = name
-                execute_delegated_tool_call.__doc__ = tool.function.__doc__
-                execute_delegated_tool_call.__signature__ = inspect.signature(
-                    tool.function
-                )
-            except (AttributeError, TypeError):
-                pass  # Ignore if we can't inspect the original function
-            # Use the existing wrap_func to get the confirmation logic
-            wrapped_executor = wrap_func(
-                execute_delegated_tool_call, self.ctx, self.yolo_mode
-            )
-            # Call the wrapped executor. This will trigger the confirmation prompt.
-            return await wrapped_executor(**tool_args)
-
-    if yolo_mode is None:
-        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, 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, yolo_mode))
-    # Wrap toolsets
-    wrapped_toolsets = [
-        ConfirmationWrapperToolset(wrapped=toolset, ctx=ctx, yolo_mode=yolo_mode)
-        for toolset in toolsets
-    ]
-    # Create History processor with summarizer
-    history_processors = [] if history_processors is None else history_processors
-    if auto_summarize:
-        history_processors += [
-            create_summarize_history_processor(
-                ctx=ctx,
-                system_prompt=system_prompt,
-                rate_limitter=rate_limitter,
-                summarization_model=summarization_model,
-                summarization_model_settings=summarization_model_settings,
-                summarization_system_prompt=summarization_system_prompt,
-                summarization_token_threshold=summarization_token_threshold,
-                summarization_retries=summarization_retries,
-            )
-        ]
-    # Return Agent
-    return Agent[None, Any](
-        model=model,
-        output_type=output_type,
-        instructions=system_prompt,
-        tools=tool_list,
-        toolsets=wrapped_toolsets,
-        model_settings=model_settings,
-        retries=retries,
-        history_processors=history_processors,
-    )
-
-
-def get_agent(
-    ctx: AnyContext,
-    model: "str | Model",
-    rate_limitter: LLMRateLimitter | None = None,
-    output_type: "OutputSpec[OutputDataT]" = str,
-    system_prompt: str = "",
-    model_settings: "ModelSettings | None" = None,
-    tools_attr: (
-        "list[ToolOrCallable] | Callable[[AnyContext], list[ToolOrCallable]]"
-    ) = [],
-    additional_tools: "list[ToolOrCallable]" = [],
-    toolsets_attr: "list[AbstractToolset[None] | str] | Callable[[AnyContext], list[AbstractToolset[None] | str]]" = [],  # noqa
-    additional_toolsets: "list[AbstractToolset[None] | str]" = [],
-    retries: int = 3,
-    yolo_mode: bool | list[str] | None = None,
-    summarization_model: "Model | str | None" = None,
-    summarization_model_settings: "ModelSettings | None" = None,
-    summarization_system_prompt: str | None = None,
-    summarization_retries: int = 2,
-    summarization_token_threshold: int | None = None,
-    history_processors: list["HistoryProcessor"] | None = None,
-) -> "Agent":
-    """Retrieves the configured Agent instance or creates one if necessary."""
-    # Get tools for agent
-    tools = list(tools_attr(ctx) if callable(tools_attr) else tools_attr)
-    tools.extend(additional_tools)
-    # Get Toolsets for agent
-    toolset_or_str_list = list(
-        toolsets_attr(ctx) if callable(toolsets_attr) else toolsets_attr
-    )
-    toolset_or_str_list.extend(additional_toolsets)
-    toolsets = _render_toolset_or_str_list(ctx, toolset_or_str_list)
-    # If no agent provided, create one using the configuration
-    return create_agent_instance(
-        ctx=ctx,
-        model=model,
-        rate_limitter=rate_limitter,
-        output_type=output_type,
-        system_prompt=system_prompt,
-        tools=tools,
-        toolsets=toolsets,
-        model_settings=model_settings,
-        retries=retries,
-        yolo_mode=yolo_mode,
-        summarization_model=summarization_model,
-        summarization_model_settings=summarization_model_settings,
-        summarization_system_prompt=summarization_system_prompt,
-        summarization_retries=summarization_retries,
-        summarization_token_threshold=summarization_token_threshold,
-        history_processors=history_processors,
-    )
-
-
-def _render_toolset_or_str_list(
-    ctx: AnyContext, toolset_or_str_list: list["AbstractToolset[None] | str"]
-) -> list["AbstractToolset[None]"]:
-    from pydantic_ai.mcp import load_mcp_servers
-
-    toolsets = []
-    for toolset_or_str in toolset_or_str_list:
-        if isinstance(toolset_or_str, str):
-            try:
-                servers = load_mcp_servers(toolset_or_str)
-                for server in servers:
-                    toolsets.append(server)
-            except Exception as e:
-                ctx.log_error(f"Invalid MCP Config {toolset_or_str}: {e}")
-            continue
-        toolsets.append(toolset_or_str)
-    return toolsets

zrb/task/llm/agent_runner.py

@@ -1,152 +0,0 @@
-import json
-from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
-
-from zrb.config.llm_rate_limitter import LLMRateLimitter, llm_rate_limitter
-from zrb.context.any_context import AnyContext
-from zrb.task.llm.error import extract_api_error_details
-from zrb.task.llm.print_node import print_node
-from zrb.task.llm.typing import ListOfDict
-from zrb.util.cli.style import stylize_faint
-
-if TYPE_CHECKING:
-    from pydantic_ai import Agent, Tool
-    from pydantic_ai.agent import AgentRun
-    from pydantic_ai.messages import UserContent
-
-    ToolOrCallable = Tool | Callable
-
-
-async def run_agent_iteration(
-    ctx: AnyContext,
-    agent: "Agent[None, Any]",
-    user_prompt: str,
-    attachments: "list[UserContent] | None" = None,
-    history_list: ListOfDict | None = None,
-    rate_limitter: LLMRateLimitter | None = None,
-    max_retry: int = 2,
-    log_indent_level: int = 0,
-) -> "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
-                ),
-                log_indent_level=log_indent_level,
-            )
-        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: LLMRateLimitter,
-    log_indent_level: int,
-) -> "AgentRun":
-    from openai import APIError
-    from pydantic_ai import UsageLimits
-    from pydantic_ai.messages import ModelMessagesTypeAdapter
-
-    agent_payload = _estimate_request_payload(
-        agent, user_prompt, attachments, history_list
-    )
-    callback = _create_print_throttle_notif(ctx)
-    if rate_limitter:
-        await rate_limitter.throttle(agent_payload, callback)
-    else:
-        await llm_rate_limitter.throttle(agent_payload, callback)
-    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),
-            usage_limits=UsageLimits(request_limit=None),  # We don't want limit
-        ) as agent_run:
-            async for node in agent_run:
-                # Each node represents a step in the agent's execution
-                try:
-                    await print_node(
-                        _get_plain_printer(ctx), agent_run, node, log_indent_level
-                    )
-                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 _create_print_throttle_notif(ctx: AnyContext) -> Callable[[str], None]:
-    def _print_throttle_notif(reason: str):
-        ctx.print(stylize_faint(f"  ⌛>> Request Throttled: {reason}"), plain=True)
-
-    return _print_throttle_notif
-
-
-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