agentpool 2.1.9__py3-none-any.whl

agentpool/utils/signatures.py

@@ -0,0 +1,305 @@
+"""Signature utils."""
+
+from __future__ import annotations
+
+import inspect
+from typing import TYPE_CHECKING, Any, get_args, get_origin
+
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+
+logger = get_logger(__name__)
+
+
+_CONTEXT_TYPE_NAMES = frozenset({
+    "NodeContext",
+    "AgentContext",
+    "RunContext",
+})
+
+
+def is_context_type(annotation: Any) -> bool:
+    """Check if an annotation is a context type (for auto-injection/hiding from agent).
+
+    This detects context types that should be hidden from the agent's view of tool
+    parameters. Both our own context types (AgentContext/NodeContext) and pydantic-ai's
+    RunContext are detected.
+
+    Uses name-based detection to avoid importing context classes, which would trigger
+    Pydantic schema generation for their fields.
+
+    Handles:
+    - Direct AgentContext/NodeContext/RunContext references
+    - String annotations (forward references from `from __future__ import annotations`)
+    - Generic forms like AgentContext[SomeDeps], RunContext[SomeDeps]
+    - Union types containing context types
+    """
+    if annotation is None or annotation is inspect.Parameter.empty:
+        return False
+
+    # Handle string annotations (forward references from `from __future__ import annotations`)
+    # Check if the string matches or starts with a known context type name
+    if isinstance(annotation, str):
+        # Handle plain names like "AgentContext" and generics like "AgentContext[Deps]"
+        base_name = annotation.split("[")[0].strip()
+        return base_name in _CONTEXT_TYPE_NAMES
+
+    # Check direct class match by name
+    if isinstance(annotation, type) and annotation.__name__ in _CONTEXT_TYPE_NAMES:
+        return True
+
+    # Check generic origin (e.g., AgentContext[SomeDeps], RunContext[SomeDeps])
+    origin = get_origin(annotation)
+    if origin is not None:
+        if isinstance(origin, type) and origin.__name__ in _CONTEXT_TYPE_NAMES:
+            return True
+        # Handle Union types
+        if origin is type(None) or str(origin) in ("typing.Union", "types.UnionType"):
+            args = get_args(annotation)
+            return any(is_context_type(arg) for arg in args)
+
+    return False
+
+
+def get_params_matching_predicate(
+    fn: Callable[..., Any],
+    predicate: Callable[[inspect.Parameter], bool],
+) -> set[str]:
+    """Get names of function parameters matching a predicate.
+
+    Args:
+        fn: Function to inspect
+        predicate: Function that takes a Parameter and returns True if it matches
+
+    Returns:
+        Set of parameter names that match the predicate
+    """
+    sig = inspect.signature(fn)
+    return {name for name, param in sig.parameters.items() if predicate(param)}
+
+
+def filter_schema_params(schema: dict[str, Any], params_to_remove: set[str]) -> dict[str, Any]:
+    """Filter parameters from a JSON schema.
+
+    Creates a copy of the schema with specified parameters removed from
+    'properties' and 'required' fields.
+
+    Args:
+        schema: JSON schema with 'properties' dict
+        params_to_remove: Set of parameter names to remove
+
+    Returns:
+        New schema dict with parameters filtered out
+    """
+    if not params_to_remove:
+        return schema
+
+    result = schema.copy()
+    if "properties" in result:
+        result["properties"] = {
+            k: v for k, v in result["properties"].items() if k not in params_to_remove
+        }
+    if "required" in result:
+        result["required"] = [r for r in result["required"] if r not in params_to_remove]
+    return result
+
+
+def create_modified_signature(
+    fn_or_sig: Callable[..., Any] | inspect.Signature,
+    *,
+    remove: str | list[str] | None = None,
+    inject: dict[str, type] | None = None,
+) -> inspect.Signature:
+    """Create a modified signature by removing specified parameters / injecting new ones.
+
+    Args:
+        fn_or_sig: The function or signature to modify.
+        remove: The parameter(s) to remove.
+        inject: The parameter(s) to inject.
+
+    Returns:
+        The modified signature.
+    """
+    sig = fn_or_sig if isinstance(fn_or_sig, inspect.Signature) else inspect.signature(fn_or_sig)
+    rem_keys = [remove] if isinstance(remove, str) else remove or []
+    new_params = [p for p in sig.parameters.values() if p.name not in rem_keys]
+    if inject:
+        injected_params = []
+        for k, v in inject.items():
+            injected_params.append(
+                inspect.Parameter(k, inspect.Parameter.POSITIONAL_OR_KEYWORD, annotation=v)
+            )
+        new_params = injected_params + new_params
+    return sig.replace(parameters=new_params)
+
+
+def modify_signature(
+    fn: Callable[..., Any],
+    *,
+    remove: str | list[str] | None = None,
+    inject: dict[str, type] | None = None,
+) -> None:
+    new_sig = create_modified_signature(fn, remove=remove, inject=inject)
+    update_signature(fn, new_sig)
+
+
+def update_signature(fn: Callable[..., Any], signature: inspect.Signature) -> None:
+    """Update function signature and annotations.
+
+    Note: Setting __annotations__ destroys __annotate__ in Python 3.14+ (PEP 649).
+    Callers using functools.wraps should restore __annotations__ from the original
+    function after calling this function.
+    """
+    fn.__signature__ = signature  # type: ignore
+    fn.__annotations__ = {
+        name: param.annotation for name, param in signature.parameters.items()
+    } | {"return": signature.return_annotation}
+
+
+def create_bound_callable(  # noqa: PLR0915
+    original_callable: Callable[..., Any],
+    by_name: dict[str, Any] | None = None,
+    by_type: dict[type, Any] | None = None,
+    bind_kwargs: bool = False,
+) -> Callable[..., Awaitable[Any]]:
+    """Create a wrapper that pre-binds parameters by name or type.
+
+    Parameters are bound by their position in the function signature. Only
+    positional and positional-or-keyword parameters can be bound by default.
+    If bind_kwargs=True, keyword-only parameters can also be bound using the
+    same by_name/by_type logic. Binding by name takes priority over binding by type.
+
+    Args:
+        original_callable: The original callable that may need parameter binding
+        by_name: Parameters to bind by exact parameter name
+        by_type: Parameters to bind by parameter type annotation
+        bind_kwargs: Whether to also bind keyword-only parameters
+
+    Returns:
+        New callable with parameters pre-bound and proper introspection
+
+    Raises:
+        ValueError: If the callable's signature cannot be inspected
+    """
+    try:
+        sig = inspect.signature(original_callable)
+    except (ValueError, TypeError) as e:
+        msg = f"Cannot inspect signature of {original_callable}. Ensure callable is inspectable."
+        raise ValueError(msg) from e
+
+    # Build position-to-value mapping for positional binding
+    context_values = {}
+    # Build name-to-value mapping for keyword-only binding
+    kwarg_bindings = {}
+
+    for i, param in enumerate(sig.parameters.values()):
+        # Bind positional and positional-or-keyword parameters
+        if param.kind in (
+            inspect.Parameter.POSITIONAL_ONLY,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        ):
+            # Bind by name first (higher priority)
+            if by_name and param.name in by_name:
+                context_values[i] = by_name[param.name]
+            # Then bind by type if not already bound
+            elif by_type and _find_matching_type(param.annotation, by_type) is not None:
+                context_values[i] = _find_matching_type(param.annotation, by_type)
+        # Bind keyword-only parameters if enabled
+        elif bind_kwargs and param.kind == inspect.Parameter.KEYWORD_ONLY:
+            # Bind by name first (higher priority)
+            if by_name and param.name in by_name:
+                kwarg_bindings[param.name] = by_name[param.name]
+            # Then bind by type if not already bound
+            elif by_type and _find_matching_type(param.annotation, by_type) is not None:
+                kwarg_bindings[param.name] = _find_matching_type(param.annotation, by_type)
+
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        # Filter out kwargs that would conflict with bound parameters
+        param_names = list(sig.parameters.keys())
+        bound_param_names = {param_names[i] for i in context_values}
+        bound_kwarg_names = set(kwarg_bindings.keys())
+        filtered_kwargs = {
+            k: v
+            for k, v in kwargs.items()
+            if k not in bound_param_names and k not in bound_kwarg_names
+        }
+
+        # Add bound keyword-only parameters
+        filtered_kwargs.update(kwarg_bindings)
+
+        # Build new_args with context values at correct positions
+        new_args = []
+        arg_index = 0
+        for param_index in range(len(sig.parameters)):
+            if param_index in context_values:
+                new_args.append(context_values[param_index])
+            elif arg_index < len(args):
+                new_args.append(args[arg_index])
+                arg_index += 1
+
+        # Add any remaining positional args
+        if arg_index < len(args):
+            new_args.extend(args[arg_index:])
+
+        if inspect.iscoroutinefunction(original_callable):
+            return await original_callable(*new_args, **filtered_kwargs)
+        return original_callable(*new_args, **filtered_kwargs)
+
+    # Preserve introspection attributes
+    wrapper.__name__ = getattr(original_callable, "__name__", "wrapper")
+    wrapper.__doc__ = getattr(original_callable, "__doc__", None)
+    wrapper.__module__ = getattr(original_callable, "__module__", None)  # type: ignore[assignment]
+    wrapper.__wrapped__ = original_callable  # type: ignore[attr-defined]
+    wrapper.__agentpool_wrapped__ = original_callable  # type: ignore[attr-defined]
+
+    # Create modified signature without context parameters
+    try:
+        params = list(sig.parameters.values())
+        # Remove parameters at context positions and bound kwargs
+        context_positions = set(context_values.keys())
+        bound_kwarg_names = set(kwarg_bindings.keys())
+        new_params = [
+            param
+            for i, param in enumerate(params)
+            if i not in context_positions and param.name not in bound_kwarg_names
+        ]
+        new_sig = sig.replace(parameters=new_params)
+        wrapper.__signature__ = new_sig  # type: ignore[attr-defined]
+        wrapper.__annotations__ = {
+            name: param.annotation for name, param in new_sig.parameters.items()
+        }
+        if sig.return_annotation != inspect.Signature.empty:
+            wrapper.__annotations__["return"] = sig.return_annotation
+
+    except (ValueError, TypeError):
+        logger.debug("Failed to update wrapper signature", original=original_callable)
+
+    return wrapper
+
+
+def _find_matching_type(param_annotation: Any, by_type: dict[type, Any]) -> Any | None:
+    """Find a matching type binding for the given parameter annotation.
+
+    Supports exact matching and generic origin matching.
+
+    Args:
+        param_annotation: The parameter's type annotation
+        by_type: Dictionary of type bindings
+
+    Returns:
+        The bound value if a match is found, None otherwise
+    """
+    # First try exact match
+    if param_annotation in by_type:
+        return by_type[param_annotation]
+
+    # Then try origin type matching for generics
+    param_origin = get_origin(param_annotation)
+    if param_origin is not None and param_origin in by_type:
+        return by_type[param_origin]
+
+    return None

agentpool/utils/streams.py

@@ -0,0 +1,112 @@
+"""Stream utilities for merging async iterators."""
+
+from __future__ import annotations
+
+import asyncio
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+
+@asynccontextmanager
+async def merge_queue_into_iterator[T, V](  # noqa: PLR0915
+    primary_stream: AsyncIterator[T],
+    secondary_queue: asyncio.Queue[V],
+) -> AsyncIterator[AsyncIterator[T | V]]:
+    """Merge a primary async stream with events from a secondary queue.
+
+    Args:
+        primary_stream: The main async iterator (e.g., provider events)
+        secondary_queue: Queue containing secondary events (e.g., progress events)
+
+    Yields:
+        Async iterator that yields events from both sources in real-time.
+        Secondary queue is fully drained before the iterator completes.
+
+    Example:
+        ```python
+        progress_queue: asyncio.Queue[ProgressEvent] = asyncio.Queue()
+
+        async with merge_queue_into_iterator(provider_stream, progress_queue) as events:
+            async for event in events:
+                print(f"Got event: {event}")
+        ```
+    """
+    # Create a queue for all merged events
+    event_queue: asyncio.Queue[V | T | None] = asyncio.Queue()
+    primary_done = asyncio.Event()
+    primary_exception: BaseException | None = None
+    # Track if we've signaled the end of streams
+    end_signaled = False
+
+    # Task to read from primary stream and put into merged queue
+    async def primary_task() -> None:
+        nonlocal primary_exception, end_signaled
+        try:
+            async for event in primary_stream:
+                await event_queue.put(event)
+        except asyncio.CancelledError:
+            # Signal completion and unblock merged_events before re-raising
+            primary_done.set()
+            if not end_signaled:
+                end_signaled = True
+                await event_queue.put(None)
+            raise
+        except BaseException as e:  # noqa: BLE001
+            primary_exception = e
+        finally:
+            primary_done.set()
+
+    # Task to read from secondary queue and put into merged queue
+    async def secondary_task() -> None:
+        nonlocal end_signaled
+        try:
+            while not primary_done.is_set():
+                try:
+                    secondary_event = await asyncio.wait_for(secondary_queue.get(), timeout=0.1)
+                    await event_queue.put(secondary_event)
+                except TimeoutError:
+                    continue
+            # Drain any remaining events after primary completes
+            while not secondary_queue.empty():
+                try:
+                    secondary_event = secondary_queue.get_nowait()
+                    await event_queue.put(secondary_event)
+                except asyncio.QueueEmpty:
+                    break
+            # Now signal end of all events (only if not already signaled)
+            if not end_signaled:
+                end_signaled = True
+                await event_queue.put(None)
+        except asyncio.CancelledError:
+            # Still need to signal completion on cancel (only if not already signaled)
+            if not end_signaled:
+                end_signaled = True
+                await event_queue.put(None)
+
+    # Start both tasks
+    primary_task_obj = asyncio.create_task(primary_task())
+    secondary_task_obj = asyncio.create_task(secondary_task())
+
+    try:
+        # Create async iterator that drains the merged queue
+        async def merged_events() -> AsyncIterator[V | T]:
+            while True:
+                event = await event_queue.get()
+                if event is None:  # End of all streams
+                    break
+                yield event
+            # Re-raise any exception from primary stream after draining
+            if primary_exception is not None:
+                raise primary_exception
+
+        yield merged_events()
+
+    finally:
+        # Clean up tasks - cancel BOTH tasks
+        primary_task_obj.cancel()
+        secondary_task_obj.cancel()
+        await asyncio.gather(primary_task_obj, secondary_task_obj, return_exceptions=True)

agentpool/utils/tasks.py

@@ -0,0 +1,186 @@
+"""Task management mixin."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+import heapq
+from typing import TYPE_CHECKING, Any
+
+import anyio
+
+from agentpool.log import get_logger
+from agentpool.utils.inspection import get_fn_name
+from agentpool.utils.now import get_now
+
+
+if TYPE_CHECKING:
+    from collections.abc import Coroutine
+    from datetime import datetime, timedelta
+
+
+logger = get_logger(__name__)
+
+
+@dataclass(order=True)
+class PrioritizedTask:
+    """Task with priority and optional delay."""
+
+    priority: int
+    execute_at: datetime
+    coroutine: Coroutine[Any, Any, Any] = field(compare=False)
+    name: str | None = field(default=None, compare=False)
+
+
+class TaskManager:
+    """Mixin for managing async tasks.
+
+    Provides utilities for:
+    - Creating and tracking tasks
+    - Fire-and-forget task execution
+    - Running coroutines in sync context
+    - Cleanup of pending tasks
+    """
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self._pending_tasks: set[asyncio.Task[Any]] = set()
+        self._task_queue: list[PrioritizedTask] = []  # heap queue
+        self._scheduler_task: asyncio.Task[Any] | None = None
+
+    def create_task[T](
+        self,
+        coro: Coroutine[Any, Any, T],
+        *,
+        name: str | None = None,
+        priority: int = 0,
+        delay: timedelta | None = None,
+    ) -> asyncio.Task[T]:
+        """Create and track a new task with optional priority and delay.
+
+        Args:
+            coro: Coroutine to run
+            name: Optional name for the task (defaults to coroutine function name)
+            priority: Priority (lower = higher priority, default 0)
+            delay: Optional delay before execution
+        """
+        task_name = name or get_fn_name(coro)
+        task = asyncio.create_task(coro, name=task_name)
+        logger.debug("Created task", name=task.get_name(), priority=priority, delay=delay)
+
+        def _done_callback(t: asyncio.Task[Any]) -> None:
+            logger.debug("Task completed", name=t.get_name())
+            self._pending_tasks.discard(t)
+            if t.exception():
+                logger.error("Task failed", error=t.exception(), name=t.get_name())
+
+        task.add_done_callback(_done_callback)
+        self._pending_tasks.add(task)
+
+        if delay is not None:
+            execute_at = get_now() + delay
+            prio_task = PrioritizedTask(priority, execute_at, coro, name)
+            heapq.heappush(self._task_queue, prio_task)  # Store the coro instead of task
+            if not self._scheduler_task:  # Start scheduler if not running
+                self._scheduler_task = asyncio.create_task(self._run_scheduler())
+            task.cancel()  # Cancel the original task since we'll run it later
+            return task
+
+        return task
+
+    async def _run_scheduler(self) -> None:
+        """Run scheduled tasks when their time comes."""
+        try:
+            while self._task_queue:
+                # Get next task without removing
+                next_task = self._task_queue[0]
+                now = get_now()
+
+                if now >= next_task.execute_at:
+                    # Remove and execute
+                    heapq.heappop(self._task_queue)
+                    # Create new task from stored coroutine
+                    new_task = asyncio.create_task(
+                        next_task.coroutine,
+                        name=next_task.name,
+                    )
+                    self._pending_tasks.add(new_task)
+                    new_task.add_done_callback(self._pending_tasks.discard)
+                else:
+                    # Wait until next task is due
+                    await anyio.sleep((next_task.execute_at - now).total_seconds())
+
+        except Exception:
+            logger.exception("Task scheduler error")
+        finally:
+            self._scheduler_task = None
+
+    def fire_and_forget(self, coro: Coroutine[Any, Any, Any]) -> None:
+        """Run coroutine without waiting for result."""
+        try:
+            loop = asyncio.get_running_loop()
+            task = loop.create_task(coro)
+            self._pending_tasks.add(task)
+            task.add_done_callback(self._pending_tasks.discard)
+        except RuntimeError:
+            # No running loop - use new loop
+            loop = asyncio.new_event_loop()
+            try:
+                loop.run_until_complete(coro)
+            finally:
+                loop.close()
+
+    def run_task_sync[T](self, coro: Coroutine[Any, Any, T]) -> T:
+        """Run coroutine synchronously."""
+        try:
+            loop = asyncio.get_running_loop()
+            if loop.is_running():
+                # Running loop - use thread pool
+                import concurrent.futures
+
+                msg = "Running coroutine in Executor due to active event loop"
+                logger.debug(msg, name=coro.__name__)
+                with concurrent.futures.ThreadPoolExecutor() as pool:
+                    future = pool.submit(lambda: asyncio.run(coro))
+                    return future.result()
+
+            # Existing but not running loop - use task tracking
+            task = loop.create_task(coro)
+            self._pending_tasks.add(task)
+            task.add_done_callback(self._pending_tasks.discard)
+            return loop.run_until_complete(task)
+        except RuntimeError:
+            # No loop - create new one
+            return asyncio.run(coro)
+
+    def run_background(
+        self,
+        coro: Coroutine[Any, Any, Any],
+        name: str | None = None,
+        priority: int = 0,
+        delay: timedelta | None = None,
+    ) -> None:
+        """Run a coroutine in the background and track it."""
+        try:
+            self.create_task(coro, name=name, priority=priority, delay=delay)
+
+        except RuntimeError:
+            # No running loop - use fire_and_forget
+            self.fire_and_forget(coro)
+
+    def is_busy(self) -> bool:
+        """Check if we have any tasks pending."""
+        return bool(self._pending_tasks)
+
+    async def cleanup_tasks(self) -> None:
+        """Wait for all pending tasks to complete."""
+        if self._pending_tasks:
+            await asyncio.gather(*self._pending_tasks, return_exceptions=True)
+        self._pending_tasks.clear()
+
+    async def complete_tasks(self, cancel: bool = False) -> None:
+        """Wait for all pending tasks to complete."""
+        if cancel:
+            for task in self._pending_tasks:
+                task.cancel()
+        if self._pending_tasks:
+            await asyncio.wait(self._pending_tasks)