agentpool 2.1.9__py3-none-any.whl → 2.2.3__py3-none-any.whl

agentpool/utils/subprocess_utils.py

@@ -0,0 +1,155 @@
+"""Utilities for subprocess management with async support."""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import anyio
+
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Awaitable, Callable
+
+    from anyio.abc import ByteReceiveStream, Process
+
+
+@dataclass
+class SubprocessError(RuntimeError):
+    """Error raised when a subprocess exits unexpectedly."""
+
+    returncode: int | None
+    stderr: str
+
+    def __str__(self) -> str:
+        msg = f"Subprocess exited unexpectedly (code {self.returncode})"
+        if self.stderr:
+            msg = f"{msg}:\n{self.stderr}"
+        return msg
+
+
+async def read_stream(
+    stream: ByteReceiveStream | None,
+    *,
+    timeout: float = 0.5,
+    max_bytes: int = 65536,
+) -> str:
+    """Read all available data from an anyio byte stream.
+
+    Args:
+        stream: The anyio ByteReceiveStream to read from
+        timeout: Timeout for each read operation
+        max_bytes: Maximum bytes to read total
+
+    Returns:
+        Decoded string content from the stream
+    """
+    if stream is None:
+        return ""
+
+    chunks: list[bytes] = []
+    total_bytes = 0
+
+    try:
+        while total_bytes < max_bytes:
+            with anyio.move_on_after(timeout) as scope:
+                chunk = await stream.receive(4096)
+                if not chunk:
+                    break
+                chunks.append(chunk)
+                total_bytes += len(chunk)
+            if scope.cancelled_caught:
+                break
+    except anyio.EndOfStream:
+        pass
+    except Exception as e:  # noqa: BLE001
+        return f"(failed to read stream: {e})"
+
+    return b"".join(chunks).decode(errors="replace").strip()
+
+
+@contextlib.asynccontextmanager
+async def monitor_process(
+    process: Process,
+    *,
+    context: str = "operation",
+) -> AsyncIterator[None]:
+    """Context manager that monitors a subprocess for unexpected exit.
+
+    Races the wrapped code against process termination. If the process
+    exits before the code completes, raises SubprocessError with stderr.
+
+    Args:
+        process: The anyio Process to monitor
+        context: Description of what's being done (for error messages)
+
+    Raises:
+        SubprocessError: If process exits during the wrapped operation
+
+    Example:
+        ```python
+        async with monitor_process(process, context="initialization"):
+            await do_initialization()
+            await create_session()
+        ```
+    """
+    process_wait_task = asyncio.create_task(process.wait())
+
+    try:
+        yield
+    except BaseException:
+        # If the wrapped code raises, cancel the wait task and re-raise
+        process_wait_task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await process_wait_task
+        raise
+
+    # Check if process died during operation
+    if process_wait_task.done():
+        stderr_output = await read_stream(process.stderr)
+        raise SubprocessError(
+            returncode=process.returncode,
+            stderr=stderr_output,
+        )
+
+    # Operation completed successfully, cancel the wait task
+    process_wait_task.cancel()
+    with contextlib.suppress(asyncio.CancelledError):
+        await process_wait_task
+
+
+async def run_with_process_monitor[T](
+    process: Process,
+    coro: Callable[[], Awaitable[T]],
+    *,
+    context: str = "operation",
+) -> T:
+    """Run a coroutine while monitoring a subprocess for unexpected exit.
+
+    Races the coroutine against process termination. If the process
+    exits before the coroutine completes, raises SubprocessError with stderr.
+
+    Args:
+        process: The anyio Process to monitor
+        coro: Async callable to execute
+        context: Description of what's being done (for error messages)
+
+    Returns:
+        The result of the coroutine
+
+    Raises:
+        SubprocessError: If process exits during execution
+
+    Example:
+        ```python
+        result = await run_with_process_monitor(
+            process,
+            lambda: initialize_connection(),
+            context="initialization",
+        )
+        ```
+    """
+    async with monitor_process(process, context=context):
+        return await coro()

agentpool/utils/token_breakdown.py

@@ -0,0 +1,461 @@
+"""Token breakdown utilities for analyzing context window usage."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from dataclasses import dataclass
+import json
+from typing import TYPE_CHECKING, Any
+
+from pydantic_ai.messages import (
+    ModelRequest,
+    ModelResponse,
+    SystemPromptPart,
+    ThinkingPart,
+    ToolCallPart,
+)
+from pydantic_ai.models import ModelRequestParameters
+from pydantic_ai.tools import ToolDefinition
+from pydantic_ai.usage import RequestUsage, RunUsage
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from pydantic_ai.messages import ModelMessage, TextPart
+    from pydantic_ai.models import Model
+    from pydantic_ai.settings import ModelSettings
+
+    from agentpool.messaging.messages import TokenCost
+
+
+@dataclass
+class TokenUsage:
+    """Single item's token count."""
+
+    token_count: int
+    label: str
+
+
+@dataclass
+class RunTokenUsage:
+    """Token usage for a single agent run."""
+
+    run_id: str | None
+    token_count: int
+    request_count: int
+
+
+@dataclass
+class TokenBreakdown:
+    """Complete token breakdown of context."""
+
+    total_tokens: int
+
+    system_prompts: list[TokenUsage]
+    tool_definitions: list[TokenUsage]
+    runs: list[RunTokenUsage]
+
+    approximate: bool
+
+    @property
+    def system_prompts_tokens(self) -> int:
+        return sum(t.token_count for t in self.system_prompts)
+
+    @property
+    def tool_definitions_tokens(self) -> int:
+        return sum(t.token_count for t in self.tool_definitions)
+
+    @property
+    def conversation_tokens(self) -> int:
+        return sum(r.token_count for r in self.runs)
+
+
+def _normalize_tool_schema(tool: ToolDefinition | dict[str, Any]) -> dict[str, Any]:
+    """Convert a ToolDefinition or dict to a consistent dict format."""
+    if isinstance(tool, ToolDefinition):
+        return {
+            "name": tool.name,
+            "description": tool.description,
+            "parameters": tool.parameters_json_schema,
+        }
+    return tool
+
+
+def count_tokens(text: str, model_name: str = "gpt-4") -> int:
+    """Count tokens using tiktoken.
+
+    Args:
+        text: The text to count tokens for.
+        model_name: The model name for encoding selection.
+
+    Returns:
+        The number of tokens in the text.
+    """
+    try:
+        import tiktoken
+    except ImportError:
+        # Rough approximation: ~4 chars per token
+        return len(text) // 4
+
+    try:
+        encoding = tiktoken.encoding_for_model(model_name)
+    except KeyError:
+        # Fall back to cl100k_base for unknown models
+        encoding = tiktoken.get_encoding("cl100k_base")
+
+    return len(encoding.encode(text))
+
+
+async def calculate_usage_from_parts(
+    input_parts: Sequence[Any],
+    response_parts: Sequence[TextPart | ThinkingPart | ToolCallPart],
+    text_content: str,
+    model_name: str | None = None,
+    provider: str | None = None,
+) -> tuple[RequestUsage, TokenCost | None]:
+    """Calculate token usage and cost from input/output parts.
+
+    This is used by agents that don't receive usage info from the backend
+    (like ACP and AG-UI agents) to approximate token counts.
+
+    Args:
+        input_parts: Input parts (prompts, pending parts) sent to the agent
+        response_parts: Response parts received (text, thinking, tool calls)
+        text_content: The final text content of the response
+        model_name: Model name for token counting and cost calculation
+        provider: Provider name for cost calculation
+
+    Returns:
+        Tuple of (RequestUsage, TokenCost or None)
+    """
+    from agentpool.messaging.messages import TokenCost
+
+    model_for_count = model_name or "gpt-4"
+
+    # Input tokens from prompts
+    input_text = " ".join(str(p) for p in input_parts)
+    input_tokens = count_tokens(input_text, model_for_count)
+
+    # Output tokens from response content
+    output_text = text_content
+    for part in response_parts:
+        if isinstance(part, ThinkingPart) and part.content:
+            output_text += part.content
+        elif isinstance(part, ToolCallPart) and part.args:
+            args_str = json.dumps(part.args) if not isinstance(part.args, str) else part.args
+            output_text += args_str
+    output_tokens = count_tokens(output_text, model_for_count)
+
+    # Build usage
+    usage = RequestUsage(input_tokens=input_tokens, output_tokens=output_tokens)
+    run_usage = RunUsage(input_tokens=input_tokens, output_tokens=output_tokens)
+
+    # Calculate cost
+    cost_info = await TokenCost.from_usage(
+        usage=run_usage,
+        model=model_name or "unknown",
+        provider=provider,
+    )
+
+    return usage, cost_info
+
+
+def _extract_system_prompts(messages: Sequence[ModelMessage]) -> list[str]:
+    """Extract all system prompt contents from messages."""
+    prompts: list[str] = []
+    for message in messages:
+        if isinstance(message, ModelRequest):
+            for part in message.parts:
+                if isinstance(part, SystemPromptPart):
+                    prompts.append(part.content)  # noqa: PERF401
+    return prompts
+
+
+def _group_messages_by_run(
+    messages: Sequence[ModelMessage],
+) -> dict[str | None, list[ModelMessage]]:
+    """Group messages by their run_id."""
+    groups: dict[str | None, list[ModelMessage]] = defaultdict(list)
+    for message in messages:
+        run_id = message.run_id
+        groups[run_id].append(message)
+    return dict(groups)
+
+
+def _messages_to_text(messages: Sequence[ModelMessage]) -> str:
+    """Convert messages to a text representation for token counting."""
+    text_parts: list[str] = []
+    for message in messages:
+        if isinstance(message, ModelRequest):
+            for request_part in message.parts:
+                if hasattr(request_part, "content") and isinstance(request_part.content, str):
+                    text_parts.append(request_part.content)  # noqa: PERF401
+        elif isinstance(message, ModelResponse):
+            if text := message.text:
+                text_parts.append(text)
+            for part in message.parts:
+                if isinstance(part, ToolCallPart):
+                    # Tool call arguments
+                    args = part.args
+                    if isinstance(args, str):
+                        text_parts.append(args)
+                    elif args:
+                        text_parts.append(json.dumps(args))
+    return "\n".join(text_parts)
+
+
+async def get_token_breakdown(
+    model: Model,
+    messages: Sequence[ModelMessage],
+    tool_schemas: Sequence[ToolDefinition | dict[str, Any]] | None = None,
+    model_settings: ModelSettings | None = None,
+) -> TokenBreakdown:
+    """Get a breakdown of token usage by component.
+
+    Uses model.count_tokens() if available, falls back to tiktoken.
+
+    Args:
+        model: The model to use for token counting.
+        messages: The message history to analyze.
+        tool_schemas: Tool definitions or raw JSON schemas.
+        model_settings: Optional model settings.
+
+    Returns:
+        A TokenBreakdown with detailed token usage by component.
+    """
+    tool_schemas = tool_schemas or []
+    approximate = False
+    model_name = model.model_name or "gpt-4"
+
+    # Try to use model.count_tokens(), fall back to tiktoken
+    async def count_tokens_for_messages(msgs: Sequence[ModelMessage]) -> int:
+        nonlocal approximate
+        try:
+            # Build minimal ModelRequestParameters for counting
+            params = ModelRequestParameters()
+            usage = await model.count_tokens(list(msgs), model_settings, params)
+        except NotImplementedError:
+            approximate = True
+            return count_tokens(_messages_to_text(msgs), model_name)
+        else:
+            return usage.input_tokens
+
+    # Extract and count system prompts
+    system_prompt_contents = _extract_system_prompts(messages)
+    system_prompt_usages: list[TokenUsage] = []
+    for i, content in enumerate(system_prompt_contents):
+        token_count = count_tokens(content, model_name)
+        label = content[:50] + "..." if len(content) > 50 else content  # noqa: PLR2004
+        system_prompt_usages.append(
+            TokenUsage(token_count=token_count, label=f"System prompt {i + 1}: {label}")
+        )
+    # Mark as approximate since we're using tiktoken for individual prompts
+    if system_prompt_usages:
+        approximate = True
+
+    # Count tool definition tokens
+    tool_usages: list[TokenUsage] = []
+    for tool in tool_schemas:
+        schema = _normalize_tool_schema(tool)
+        schema_text = json.dumps(schema)
+        token_count = count_tokens(schema_text, model_name)
+        tool_name = schema.get("name", "unknown")
+        tool_usages.append(TokenUsage(token_count=token_count, label=tool_name))
+    if tool_usages:
+        approximate = True
+
+    # Group messages by run and count tokens per run
+    run_groups = _group_messages_by_run(messages)
+    run_usages: list[RunTokenUsage] = []
+    for run_id, run_messages in run_groups.items():
+        token_count = await count_tokens_for_messages(run_messages)
+        request_count = sum(1 for m in run_messages if isinstance(m, ModelRequest))
+        run_usages.append(
+            RunTokenUsage(
+                run_id=run_id,
+                token_count=token_count,
+                request_count=request_count,
+            )
+        )
+
+    # Calculate total
+    total = (
+        sum(u.token_count for u in system_prompt_usages)
+        + sum(u.token_count for u in tool_usages)
+        + sum(r.token_count for r in run_usages)
+    )
+
+    return TokenBreakdown(
+        total_tokens=total,
+        system_prompts=system_prompt_usages,
+        tool_definitions=tool_usages,
+        runs=run_usages,
+        approximate=approximate,
+    )
+
+
+def format_breakdown(breakdown: TokenBreakdown, detailed: bool = False) -> str:
+    """Format a token breakdown for display."""
+    lines: list[str] = []
+    # Header
+    approx_marker = " (approximate)" if breakdown.approximate else ""
+    lines.append(f"Token Breakdown{approx_marker}")
+    lines.append("=" * 50)
+    # Summary
+    lines.append(f"Total tokens: {breakdown.total_tokens:,}")
+    lines.append("")
+    # Category breakdown
+    lines.append("By category:")
+    lines.append(f"  System prompts: {breakdown.system_prompts_tokens:,} tokens")
+    lines.append(f"  Tool definitions: {breakdown.tool_definitions_tokens:,} tokens")
+    lines.append(f"  Conversation: {breakdown.conversation_tokens:,} tokens")
+    if detailed:
+        lines.append("")
+        lines.append("-" * 50)
+        # System prompts detail
+        if breakdown.system_prompts:
+            lines.append("")
+            lines.append("System Prompts:")
+            for sp in breakdown.system_prompts:
+                lines.append(f"  [{sp.token_count:,} tokens] {sp.label}")  # noqa: PERF401
+        # Tool definitions detail
+        if breakdown.tool_definitions:
+            lines.append("")
+            lines.append("Tool Definitions:")
+            for tool in breakdown.tool_definitions:
+                lines.append(f"  [{tool.token_count:,} tokens] {tool.label}")  # noqa: PERF401
+        # Runs detail
+        if breakdown.runs:
+            lines.append("")
+            lines.append("Conversation by Run:")
+            for run in breakdown.runs:
+                run_label = run.run_id[:8] + "..." if run.run_id else "(no run_id)"
+                lines.append(
+                    f"  [{run.token_count:,} tokens] {run_label} ({run.request_count} requests)"
+                )
+
+    lines.append("")
+    return "\n".join(lines)
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    from pydantic_ai.messages import (
+        ImageUrl,
+        ModelRequest,
+        ModelResponse,
+        SystemPromptPart,
+        TextPart,
+        ToolCallPart,
+        ToolReturnPart,
+        UserPromptPart,
+    )
+    from pydantic_ai.models.test import TestModel
+    from pydantic_ai.tools import ToolDefinition
+
+    async def main() -> None:
+        # Create sample tool definitions
+        tool_definitions = [
+            ToolDefinition(
+                name="get_weather",
+                description="Get the current weather for a city.",
+                parameters_json_schema={
+                    "type": "object",
+                    "properties": {
+                        "city": {"type": "string", "description": "The city name"},
+                    },
+                    "required": ["city"],
+                },
+            ),
+            ToolDefinition(
+                name="search_database",
+                description="Search a database with a complex query. Supports filtering, sorting, and pagination.",  # noqa: E501
+                parameters_json_schema={
+                    "type": "object",
+                    "properties": {
+                        "query": {"type": "string", "description": "Search query"},
+                        "filters": {
+                            "type": "object",
+                            "description": "Filter conditions",
+                            "additionalProperties": True,
+                        },
+                        "sort_by": {"type": "string", "description": "Field to sort by"},
+                        "limit": {"type": "integer", "description": "Max results"},
+                        "offset": {"type": "integer", "description": "Skip N results"},
+                    },
+                    "required": ["query"],
+                },
+            ),
+        ]
+
+        # Create sample message history simulating two runs
+        messages: list[ModelMessage] = [
+            # First run
+            ModelRequest(
+                parts=[
+                    SystemPromptPart(
+                        content="You are a helpful assistant with access to weather and time tools."
+                    ),
+                    UserPromptPart(content="What's the weather in Paris?"),
+                ],
+                run_id="run-001-abc",
+            ),
+            ModelResponse(
+                parts=[
+                    ToolCallPart(
+                        tool_name="get_weather", args={"city": "Paris"}, tool_call_id="call-1"
+                    ),
+                ],
+                run_id="run-001-abc",
+            ),
+            ModelRequest(
+                parts=[
+                    ToolReturnPart(
+                        tool_name="get_weather",
+                        content="Sunny, 22°C in Paris",
+                        tool_call_id="call-1",
+                    ),
+                ],
+                run_id="run-001-abc",
+            ),
+            ModelResponse(
+                parts=[
+                    TextPart(content="The weather in Paris is sunny with a temperature of 22°C."),
+                ],
+                run_id="run-001-abc",
+            ),
+            # Second run (continuing conversation) - includes an image
+            ModelRequest(
+                parts=[
+                    UserPromptPart(
+                        content=[
+                            "What time is it in Tokyo?",
+                            ImageUrl(url="https://example.com/tokyo-clock.jpg"),
+                        ]
+                    ),
+                ],
+                run_id="run-002-def",
+            ),
+            ModelResponse(
+                parts=[
+                    TextPart(content="The current time in Tokyo is 14:30 JST."),
+                ],
+                run_id="run-002-def",
+            ),
+        ]
+
+        # Get the breakdown using TestModel (will use tiktoken fallback)
+        model = TestModel()
+        breakdown = await get_token_breakdown(
+            model=model, messages=messages, tool_schemas=tool_definitions
+        )
+        # Print summary view
+        print("SUMMARY VIEW")
+        print(format_breakdown(breakdown, detailed=False))
+        # Print detailed view
+        print("DETAILED VIEW")
+        print(format_breakdown(breakdown, detailed=True))
+
+    asyncio.run(main())