agentpool 2.1.9__py3-none-any.whl

agentpool/agents/slashed_agent.py

@@ -0,0 +1,244 @@
+"""Slash command wrapper for Agent that injects command events into streams."""
+
+from __future__ import annotations
+
+import asyncio
+import re
+from typing import TYPE_CHECKING, Any, cast
+
+import anyio
+from slashed.events import (
+    CommandExecutedEvent,
+    CommandOutputEvent as SlashedCommandOutputEvent,
+)
+
+from agentpool.agents.events import CommandCompleteEvent, CommandOutputEvent
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator, Callable
+
+    from slashed import CommandContext, CommandStore
+    from slashed.events import CommandStoreEvent
+
+    from agentpool.agents.base_agent import BaseAgent
+    from agentpool.agents.events import SlashedAgentStreamEvent
+    from agentpool.common_types import PromptCompatible
+
+
+logger = get_logger(__name__)
+SLASH_PATTERN = re.compile(r"^/([\w-]+)(?:\s+(.*))?$")
+
+
+def _parse_slash_command(command_text: str) -> tuple[str, str] | None:
+    """Parse slash command into name and args.
+
+    Args:
+        command_text: Full command text
+
+    Returns:
+        Tuple of (cmd_name, args) or None if invalid
+    """
+    if match := SLASH_PATTERN.match(command_text.strip()):
+        cmd_name = match.group(1)
+        args = match.group(2) or ""
+        return cmd_name, args.strip()
+    return None
+
+
+class SlashedAgent[TDeps, OutputDataT]:
+    """Wrapper around Agent that handles slash commands in streams.
+
+    Uses the "commands first" strategy from the ACP adapter:
+    1. Execute all slash commands first
+    2. Then process remaining content through wrapped agent
+    3. If only commands, end without LLM processing
+    """
+
+    def __init__(
+        self,
+        agent: BaseAgent[TDeps, OutputDataT],
+        command_store: CommandStore | None = None,
+        *,
+        context_data_factory: Callable[[], Any] | None = None,
+    ) -> None:
+        """Initialize with wrapped agent and command store.
+
+        Args:
+            agent: The agent to wrap
+            command_store: Command store for slash commands (creates default if None)
+            context_data_factory: Optional factory for creating command context data
+        """
+        self.agent = agent
+        self._context_data_factory = context_data_factory
+        self._event_queue: asyncio.Queue[CommandStoreEvent] | None = None
+
+        # Create store with our streaming event handler
+        if command_store is None:
+            from slashed import CommandStore
+
+            from agentpool_commands import get_commands
+
+            cmds = get_commands()
+            self.command_store = CommandStore(event_handler=self._emit_event, commands=cmds)
+        else:
+            self.command_store = command_store
+
+    async def _emit_event(self, event: CommandStoreEvent) -> None:
+        """Bridge store events to async queue during command execution."""
+        if self._event_queue:
+            await self._event_queue.put(event)
+
+    def _is_slash_command(self, text: str) -> bool:
+        """Check if text starts with a slash command.
+
+        Args:
+            text: Text to check
+
+        Returns:
+            True if text is a slash command
+        """
+        return bool(SLASH_PATTERN.match(text.strip()))
+
+    async def _execute_slash_command_streaming(
+        self, command_text: str
+    ) -> AsyncGenerator[CommandOutputEvent | CommandCompleteEvent]:
+        """Execute a single slash command and yield events as they happen.
+
+        Args:
+            command_text: Full command text including slash
+
+        Yields:
+            Command output and completion events
+        """
+        parsed = _parse_slash_command(command_text)
+        if not parsed:
+            logger.warning("Invalid slash command", command=command_text)
+            yield CommandCompleteEvent(command="unknown", success=False)
+            return
+
+        cmd_name, args = parsed
+
+        # Set up event queue for this command execution
+        self._event_queue = asyncio.Queue()
+        context_data = (  # Create command context
+            self._context_data_factory() if self._context_data_factory else self.agent.get_context()
+        )
+
+        cmd_ctx = self.command_store.create_context(data=context_data)
+        command_str = f"{cmd_name} {args}".strip()
+        execute_task = asyncio.create_task(self.command_store.execute_command(command_str, cmd_ctx))
+
+        success = True
+        try:
+            # Yield events from queue as command runs
+            while not execute_task.done():
+                try:
+                    # Wait for events with short timeout to check task completion
+                    event = await asyncio.wait_for(self._event_queue.get(), timeout=0.1)
+                    # Convert store events to our stream events
+                    match event:
+                        case SlashedCommandOutputEvent(output=output):
+                            yield CommandOutputEvent(command=cmd_name, output=output)
+                        case CommandExecutedEvent(success=False, error=error) if error:
+                            output = f"Command error: {error}"
+                            yield CommandOutputEvent(command=cmd_name, output=output)
+                            success = False
+                except TimeoutError:
+                    continue
+
+            # Ensure command task completes and handle any remaining events
+            try:
+                await execute_task
+            except Exception as e:
+                logger.exception("Command execution failed", command=cmd_name)
+                success = False
+                yield CommandOutputEvent(command=cmd_name, output=f"Command error: {e}")
+
+            # Drain any remaining events from queue
+            while not self._event_queue.empty():
+                try:
+                    match self._event_queue.get_nowait():
+                        case SlashedCommandOutputEvent(output=output):
+                            yield CommandOutputEvent(command=cmd_name, output=output)
+                except asyncio.QueueEmpty:
+                    break
+
+            # Always yield completion event
+            yield CommandCompleteEvent(command=cmd_name, success=success)
+
+        finally:
+            # Clean up event queue
+            self._event_queue = None
+
+    async def run_stream(
+        self,
+        *prompts: PromptCompatible,
+        **kwargs: Any,
+    ) -> AsyncGenerator[SlashedAgentStreamEvent[OutputDataT]]:
+        """Run agent with slash command support.
+
+        Separates slash commands from regular prompts, executes commands first,
+        then processes remaining content through the wrapped agent.
+
+        Args:
+            *prompts: Input prompts (may include slash commands)
+            **kwargs: Additional arguments passed to agent.run_stream
+
+        Yields:
+            Stream events from command execution and agent processing
+        """
+        # Separate slash commands from regular content
+        commands: list[str] = []
+        regular_prompts: list[Any] = []
+
+        for prompt in prompts:
+            if isinstance(prompt, str) and self._is_slash_command(prompt):
+                logger.debug("Found slash command", command=prompt)
+                commands.append(prompt.strip())
+            else:
+                regular_prompts.append(prompt)
+
+        # Execute all commands first with streaming
+        if commands:
+            for command in commands:
+                logger.info("Processing slash command", command=command)
+                async for cmd_event in self._execute_slash_command_streaming(command):
+                    yield cmd_event
+
+        # If we have regular content, process it through the agent
+        if regular_prompts:
+            logger.debug("Processing prompts through agent", num_prompts=len(regular_prompts))
+            async for event in self.agent.run_stream(*regular_prompts, **kwargs):
+                # ACPAgent always returns str, cast to match OutputDataT
+                yield cast("SlashedAgentStreamEvent[OutputDataT]", event)
+
+        # If we only had commands and no regular content, we're done
+        # (no additional events needed)
+
+
+if __name__ == "__main__":
+    import asyncio
+
+    from agentpool import Agent
+
+    async def main() -> None:
+        agent = Agent("test-agent", model="test", session=False)
+        slashed = SlashedAgent(agent)  # Uses built-in commands by default
+
+        # Add a simple test command that outputs multiple lines
+        @slashed.command_store.command(name="test-streaming", category="test")
+        async def test_streaming(ctx: CommandContext[Any], *args: Any, **kwargs: Any) -> None:
+            """Test command that outputs multiple lines."""
+            await ctx.print("Starting streaming test...")
+            for i in range(3):
+                await ctx.print(f"Output line {i + 1}")
+                await anyio.sleep(0.1)  # Small delay to simulate work
+            await ctx.print("Streaming test complete!")
+
+        print("Testing SlashedAgent streaming:")
+        async for event in slashed.run_stream("/test-streaming"):
+            print(f"Event: {event}")
+
+    anyio.run(main)

agentpool/agents/sys_prompts.py

@@ -0,0 +1,178 @@
+"""System prompt management."""
+
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, Literal
+
+from agentpool import text_templates
+
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from toprompt import AnyPromptType
+
+    from agentpool.agents.base_agent import BaseAgent
+    from agentpool.prompts.manager import PromptManager
+
+
+ToolInjectionMode = Literal["off", "all"]
+ToolUsageStyle = Literal["suggestive", "strict"]
+
+
+class SystemPrompts:
+    """Manages system prompts for an agent."""
+
+    def __init__(
+        self,
+        prompts: AnyPromptType | list[AnyPromptType] | None = None,
+        template: str | None = None,
+        dynamic: bool = True,
+        prompt_manager: PromptManager | None = None,
+        inject_agent_info: bool = True,
+        inject_tools: ToolInjectionMode = "off",
+        tool_usage_style: ToolUsageStyle = "suggestive",
+    ) -> None:
+        """Initialize prompt manager."""
+        from jinjarope import Environment
+        from toprompt import to_prompt
+
+        match prompts:
+            case list():
+                self.prompts = prompts
+            case None:
+                self.prompts = []
+            case _:
+                self.prompts = [prompts]
+        self.prompt_manager = prompt_manager
+        self.template = template
+        self.dynamic = dynamic
+        self.inject_agent_info = inject_agent_info
+        self.inject_tools = inject_tools
+        self.tool_usage_style = tool_usage_style
+        self._cached = False
+        self._env = Environment(enable_async=True)
+        self._env.filters["to_prompt"] = to_prompt
+
+    def __repr__(self) -> str:
+        return (
+            f"SystemPrompts(prompts={len(self.prompts)}, "
+            f"dynamic={self.dynamic}, inject_agent_info={self.inject_agent_info}, "
+            f"inject_tools={self.inject_tools!r})"
+        )
+
+    def __len__(self) -> int:
+        return len(self.prompts)
+
+    def __getitem__(self, idx: int | slice) -> AnyPromptType | list[AnyPromptType]:
+        return self.prompts[idx]
+
+    async def add_by_reference(self, reference: str) -> None:
+        """Add a system prompt using reference syntax.
+
+        Args:
+            reference: [provider:]identifier[@version][?var1=val1,...]
+
+        Examples:
+            await sys_prompts.add_by_reference("code_review?language=python")
+            await sys_prompts.add_by_reference("langfuse:expert@v2")
+        """
+        if not self.prompt_manager:
+            msg = "No prompt_manager available to resolve prompts"
+            raise RuntimeError(msg)
+
+        try:
+            content = await self.prompt_manager.get(reference)
+            self.prompts.append(content)
+        except Exception as e:
+            msg = f"Failed to add prompt {reference!r}"
+            raise RuntimeError(msg) from e
+
+    async def add(
+        self,
+        identifier: str,
+        *,
+        provider: str | None = None,
+        version: str | None = None,
+        variables: dict[str, Any] | None = None,
+    ) -> None:
+        """Add a system prompt.
+
+        Args:
+            identifier: Prompt identifier/name
+            provider: Provider name (None = builtin)
+            version: Optional version string
+            variables: Optional template variables
+
+        Examples:
+            await sys_prompts.add("code_review", variables={"language": "python"})
+            await sys_prompts.add("expert", provider="langfuse", version="v2")
+        """
+        if not self.prompt_manager:
+            msg = "No prompt_manager available to resolve prompts"
+            raise RuntimeError(msg)
+
+        try:
+            content = await self.prompt_manager.get_from(
+                identifier,
+                provider=provider,
+                version=version,
+                variables=variables,
+            )
+            self.prompts.append(content)
+        except Exception as e:
+            ref = f"{provider + ':' if provider else ''}{identifier}"
+            msg = f"Failed to add prompt {ref!r}"
+            raise RuntimeError(msg) from e
+
+    def clear(self) -> None:
+        """Clear all system prompts."""
+        self.prompts = []
+
+    async def refresh_cache(self) -> None:
+        """Force re-evaluation of prompts."""
+        from toprompt import to_prompt
+
+        evaluated = []
+        for prompt in self.prompts:
+            result = await to_prompt(prompt)
+            evaluated.append(result)
+        self.prompts = evaluated
+        self._cached = True
+
+    @asynccontextmanager
+    async def temporary_prompt(
+        self, prompt: AnyPromptType, exclusive: bool = False
+    ) -> AsyncIterator[None]:
+        """Temporarily override system prompts.
+
+        Args:
+            prompt: Single prompt or sequence of prompts to use temporarily
+            exclusive: Whether to only use given prompt. If False, prompt will be
+                       appended to the agents prompts temporarily.
+        """
+        from toprompt import to_prompt
+
+        original_prompts = self.prompts.copy()
+        new_prompt = await to_prompt(prompt)
+        self.prompts = [new_prompt] if not exclusive else [*self.prompts, new_prompt]
+        try:
+            yield
+        finally:
+            self.prompts = original_prompts
+
+    async def format_system_prompt(self, agent: BaseAgent[Any, Any]) -> str:
+        """Format complete system prompt."""
+        if not self.dynamic and not self._cached:
+            await self.refresh_cache()
+        template = self._env.from_string(self.template or text_templates.get_system_prompt())
+        result = await template.render_async(
+            agent=agent,
+            prompts=self.prompts,
+            dynamic=self.dynamic,
+            inject_agent_info=self.inject_agent_info,
+            inject_tools=self.inject_tools,
+            tool_usage_style=self.tool_usage_style,
+        )
+        return result.strip()

agentpool/agents/tool_wrapping.py

@@ -0,0 +1,184 @@
+"""Tool wrapping utilities for pydantic-ai integration."""
+
+from __future__ import annotations
+
+from dataclasses import replace
+from functools import wraps
+import inspect
+import time
+from typing import TYPE_CHECKING, Any
+
+from pydantic_ai import RunContext
+
+from agentpool.agents.context import AgentContext
+from agentpool.tasks import ChainAbortedError, RunAbortedError, ToolSkippedError
+from agentpool.utils.inspection import execute, get_argument_key
+from agentpool.utils.signatures import create_modified_signature, update_signature
+
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+    from agentpool.hooks import AgentHooks
+    from agentpool.tools.base import Tool
+
+
+def wrap_tool[TReturn](  # noqa: PLR0915
+    tool: Tool[TReturn],
+    agent_ctx: AgentContext,
+    hooks: AgentHooks | None = None,
+) -> Callable[..., Awaitable[TReturn | None]]:
+    """Wrap tool with confirmation handling and hooks.
+
+    Strategy:
+    - Tools with RunContext only: Normal pydantic-ai handling
+    - Tools with AgentContext only: Treat as regular tools, inject AgentContext
+    - Tools with both contexts: Present as RunContext-only to pydantic-ai, inject AgentContext
+    - Tools with no context: Normal pydantic-ai handling
+
+    Args:
+        tool: The tool to wrap.
+        agent_ctx: Agent context for confirmation handling and dependency injection.
+        hooks: Optional AgentHooks for pre/post tool execution hooks.
+    """
+    fn = tool.callable
+    run_ctx_key = get_argument_key(fn, RunContext)
+    agent_ctx_key = get_argument_key(fn, AgentContext)
+
+    # Validate parameter order if RunContext is present
+    if run_ctx_key:
+        param_names = list(inspect.signature(fn).parameters.keys())
+        run_ctx_index = param_names.index(run_ctx_key)
+        if run_ctx_index != 0:
+            msg = f"Tool {tool.name!r}: RunContext param {run_ctx_key!r} must come first."
+            raise ValueError(msg)
+
+    async def _execute_with_hooks(
+        execute_fn: Callable[..., Awaitable[TReturn]],
+        tool_input: dict[str, Any],
+        *args: Any,
+        **kwargs: Any,
+    ) -> TReturn | None:
+        """Execute tool with pre/post hooks."""
+        # Pre-tool hooks
+        if hooks:
+            pre_result = await hooks.run_pre_tool_hooks(
+                agent_name=agent_ctx.node_name,
+                tool_name=tool.name,
+                tool_input=tool_input,
+                conversation_id=None,  # Could be passed through if needed
+            )
+            if pre_result.get("decision") == "deny":
+                reason = pre_result.get("reason", "Blocked by pre-tool hook")
+                msg = f"Tool {tool.name} blocked: {reason}"
+                raise ToolSkippedError(msg)
+
+            # Apply modified input if provided
+            if modified := pre_result.get("modified_input"):
+                kwargs.update(modified)
+
+        # Execute the tool
+        start_time = time.perf_counter()
+        result = await execute_fn(*args, **kwargs)
+        duration_ms = (time.perf_counter() - start_time) * 1000
+
+        # Post-tool hooks
+        if hooks:
+            await hooks.run_post_tool_hooks(
+                agent_name=agent_ctx.node_name,
+                tool_name=tool.name,
+                tool_input=tool_input,
+                tool_output=result,
+                duration_ms=duration_ms,
+                conversation_id=None,
+            )
+
+        return result
+
+    if run_ctx_key or agent_ctx_key:
+        # Tool has RunContext and/or AgentContext
+        async def wrapped(ctx: RunContext, *args: Any, **kwargs: Any) -> TReturn | None:  # pyright: ignore
+            result = await agent_ctx.handle_confirmation(tool, kwargs)
+            if result == "allow":
+                # Populate AgentContext with RunContext data if needed
+                if agent_ctx.data is None:
+                    agent_ctx.data = ctx.deps
+
+                if agent_ctx_key:  # inject AgentContext
+                    # Create per-call copy with tool execution fields (avoids race condition)
+                    call_ctx = replace(
+                        agent_ctx,
+                        tool_name=ctx.tool_name,
+                        tool_call_id=ctx.tool_call_id,
+                        tool_input=kwargs.copy(),
+                    )
+                    kwargs[agent_ctx_key] = call_ctx
+
+                tool_input = kwargs.copy()
+                if run_ctx_key:
+                    # Pass RunContext to original function
+                    return await _execute_with_hooks(
+                        lambda *a, **kw: execute(fn, ctx, *a, **kw),
+                        tool_input,
+                        *args,
+                        **kwargs,
+                    )
+                # Don't pass RunContext to original function since it didn't expect it
+                return await _execute_with_hooks(
+                    lambda *a, **kw: execute(fn, *a, **kw),
+                    tool_input,
+                    *args,
+                    **kwargs,
+                )
+            await _handle_confirmation_result(result, tool.name)
+            return None
+
+    else:
+        # Tool has no context - normal function call
+        async def wrapped(*args: Any, **kwargs: Any) -> TReturn | None:  # type: ignore[misc]
+            result = await agent_ctx.handle_confirmation(tool, kwargs)
+            if result == "allow":
+                tool_input = kwargs.copy()
+                return await _execute_with_hooks(
+                    lambda *a, **kw: execute(fn, *a, **kw),
+                    tool_input,
+                    *args,
+                    **kwargs,
+                )
+            await _handle_confirmation_result(result, tool.name)
+            return None
+
+    # Apply wraps first
+    wraps(fn)(wrapped)  # pyright: ignore
+    # Python 3.14: functools.wraps copies __annotate__ but not __annotations__.
+    # Any subsequent assignment to __annotations__ destroys __annotate__ (PEP 649).
+    # Restore from original to preserve deferred annotation evaluation.
+    # TODO: probably review all wraps() calls in the codebase.
+    wrapped.__annotations__ = fn.__annotations__
+    wrapped.__doc__ = tool.description
+    wrapped.__name__ = tool.name
+    # Modify signature for pydantic-ai: hide AgentContext, add RunContext if needed
+    # Must be done AFTER wraps to prevent overwriting
+    if agent_ctx_key and not run_ctx_key:
+        # Tool has AgentContext only - make it appear to have RunContext to pydantic-ai
+        new_sig = create_modified_signature(fn, remove=agent_ctx_key, inject={"ctx": RunContext})
+        update_signature(wrapped, new_sig)
+    elif agent_ctx_key and run_ctx_key:
+        # Tool has both contexts - hide AgentContext from pydantic-ai
+        new_sig = create_modified_signature(fn, remove=agent_ctx_key)
+        update_signature(wrapped, new_sig)
+    return wrapped
+
+
+async def _handle_confirmation_result(result: str, name: str) -> None:
+    """Handle non-allow confirmation results."""
+    match result:
+        case "skip":
+            msg = f"Tool {name} execution skipped"
+            raise ToolSkippedError(msg)
+        case "abort_run":
+            msg = "Run aborted by user"
+            raise RunAbortedError(msg)
+        case "abort_chain":
+            msg = "Agent chain aborted by user"
+            raise ChainAbortedError(msg)

agentpool/base_provider.py

@@ -0,0 +1,28 @@
+"""Base classes for providers."""
+
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any, Self
+
+from pydantic import BaseModel
+
+
+class BaseProvider[ConfigT: BaseModel]:
+    """Base class for all providers."""
+
+    def __init__(self, config: ConfigT) -> None:
+        """Initialize provider with configuration."""
+        self.config = config
+
+    @classmethod
+    @abstractmethod
+    def from_kwargs(cls, **kwargs: Any) -> Self:
+        """Alternative constructor with explicit parameters."""
+        raise NotImplementedError
+
+    def __repr__(self) -> str:
+        """Return string representation of provider with non-default config values."""
+        non_defaults = self.config.model_dump(exclude_defaults=True)
+        fields_str = ", ".join(f"{k}={v!r}" for k, v in non_defaults.items())
+        return f"{self.__class__.__name__}({fields_str})"