agentpool 2.1.9__py3-none-any.whl

agentpool/running/executor.py

@@ -0,0 +1,284 @@
+"""Function execution management."""
+
+from __future__ import annotations
+
+import asyncio
+import importlib.util
+import inspect
+from pathlib import Path
+import sys
+from typing import TYPE_CHECKING, Any, get_type_hints
+
+import anyio
+
+from agentpool.log import get_logger
+from agentpool.running import with_nodes
+
+
+if TYPE_CHECKING:
+    import os
+
+    from agentpool import AgentPool
+    from agentpool.running.discovery import NodeFunction
+
+
+logger = get_logger(__name__)
+
+
+class ExecutionError(Exception):
+    """Raised when function execution fails."""
+
+
+def _validate_path(path: str | os.PathLike[str]) -> Path:
+    if str(path) == "__main__":
+        # Get the actual file being run
+
+        path = sys.modules["__main__"].__file__  # type: ignore
+        if not path:
+            msg = "Could not determine main module file"
+            raise ValueError(msg)
+    path_obj = Path(path)
+    if not path_obj.exists():
+        msg = f"Module not found: {path}"
+        raise ValueError(msg)
+    return path_obj
+
+
+def discover_functions(path: str | os.PathLike[str]) -> list[NodeFunction]:
+    """Find all node functions in a module.
+
+    Args:
+        path: Path to Python module file
+
+    Returns:
+        List of discovered node functions
+
+    Raises:
+        ImportError: If module cannot be imported
+        ValueError: If path is invalid
+    """
+    path_obj = _validate_path(path)
+    # Import module
+    spec = importlib.util.spec_from_file_location(path_obj.stem, path_obj)
+    if not spec or not spec.loader:
+        msg = f"Could not load module: {path}"
+        raise ImportError(msg)
+
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+
+    # Find decorated functions
+    return [i for name, i in inspect.getmembers(module) if hasattr(i, "_node_function")]
+
+
+def _sort_functions(functions: list[NodeFunction]) -> list[NodeFunction]:
+    """Sort functions by order and dependencies.
+
+    Args:
+        functions: Functions to sort
+
+    Returns:
+        Sorted list of functions
+
+    Raises:
+        ValueError: If there are circular dependencies
+    """
+    # First by explicit order
+    ordered = sorted(functions, key=lambda f: f.name)
+
+    # Then resolve dependencies
+    result = []
+    seen = set()
+    in_progress = set()
+
+    def add_function(func: NodeFunction) -> None:
+        if func.name in seen:
+            return
+        if func.name in in_progress:
+            msg = f"Circular dependency detected: {func.name}"
+            raise ValueError(msg)
+
+        in_progress.add(func.name)
+        # Add dependencies first
+        for dep in func.depends_on:
+            dep_fn = next((f for f in ordered if f.name == dep), None)
+            if not dep_fn:
+                msg = f"Missing dependency {dep} for {func.name}"
+                raise ValueError(msg)
+            add_function(dep_fn)
+
+        result.append(func)
+        in_progress.remove(func.name)
+        seen.add(func.name)
+
+    for func in ordered:
+        add_function(func)
+
+    return result
+
+
+def _group_parallel(
+    sorted_funcs: list[NodeFunction],
+) -> list[list[NodeFunction]]:
+    """Group functions that can run in parallel."""
+    if not sorted_funcs:
+        return []
+
+    # Group by dependency signature
+    by_deps: dict[tuple[str, ...], list[NodeFunction]] = {}
+
+    for func in sorted_funcs:
+        # Use tuple of sorted deps as key for consistent grouping
+        key = tuple(sorted(func.depends_on))
+        if key not in by_deps:
+            by_deps[key] = []
+        by_deps[key].append(func)
+
+    # Convert to list of groups, maintaining order
+    groups = []
+    seen_funcs: set[str] = set()
+
+    for func in sorted_funcs:
+        key = tuple(sorted(func.depends_on))
+        if func.name not in seen_funcs:
+            group = by_deps[key]
+            groups.append(group)
+            seen_funcs.update(f.name for f in group)
+    names = [[f.name for f in g] for g in groups]
+    logger.debug("Grouped functions into groups", num_funcs=len(sorted_funcs), group_names=names)
+    return groups
+
+
+async def execute_single(
+    func: NodeFunction,
+    pool: AgentPool,
+    available_results: dict[str, Any],
+    inputs: dict[str, Any] | None = None,
+) -> tuple[str, Any]:
+    """Execute a single function.
+
+    Args:
+        func: Function to execute
+        pool: Agent pool for injection
+        available_results: Results from previous functions
+        inputs: Optional input overrides
+
+    Returns:
+        Tuple of (function name, result)
+
+    Raises:
+        ExecutionError: If execution fails
+    """
+    logger.debug("Executing function", name=func.name)
+    try:
+        kwargs = func.default_inputs.copy()
+        if inputs:
+            kwargs.update(inputs)
+
+        # Get type hints for the function
+        hints = get_type_hints(func.func)
+
+        # Add and validate dependency results
+        for dep in func.depends_on:
+            if dep not in available_results:
+                msg = f"Missing result from {dep}"
+                raise ExecutionError(msg)  # noqa: TRY301
+
+            value = available_results[dep]
+            if dep in hints:  # If parameter is type hinted
+                _validate_value_type(value, hints[dep], func.name, dep)
+            kwargs[dep] = value
+
+        # Execute with node injection
+        wrapped = with_nodes(pool)(func.func)
+        result = await wrapped(**kwargs)
+
+        # Validate return type if hinted
+        if "return" in hints:
+            _validate_value_type(result, hints["return"], func.name, "return")
+    except Exception as e:
+        msg = f"Error executing {func.name}: {e}"
+        raise ExecutionError(msg) from e
+    else:
+        return func.name, result
+
+
+def _validate_dependency_types(functions: list[NodeFunction]) -> None:
+    """Validate that dependency types match return types."""
+    # Get return types for all functions
+    return_types = {}
+    for func in functions:
+        hints = get_type_hints(func.func)
+        if "return" in hints:
+            return_types[func.name] = hints["return"]
+
+    # Check each function's dependencies
+    for func in functions:
+        hints = get_type_hints(func.func)
+        for dep in func.depends_on:
+            # Only validate if both dependency return type AND parameter are typed
+            if dep in hints and dep in return_types:
+                expected_type = hints[dep]
+                provided_type = return_types[dep]
+                if expected_type != provided_type:
+                    msg = (
+                        f"Type mismatch in {func.name}: "
+                        f"dependency {dep!r} is typed as {expected_type}, "
+                        f"but {dep} returns {provided_type}"
+                    )
+                    raise TypeError(msg)
+
+
+def _validate_value_type(value: Any, expected_type: type, func_name: str, param_name: str) -> None:
+    """Validate that a value matches its expected type."""
+    if not isinstance(value, expected_type):
+        msg = (
+            f"Type error in {func_name}: parameter {param_name!r} "
+            f"expected {expected_type.__name__}, got {type(value).__name__}"
+        )
+        raise TypeError(msg)
+
+
+async def execute_functions(
+    functions: list[NodeFunction],
+    pool: AgentPool,
+    inputs: dict[str, Any] | None = None,
+    parallel: bool = False,
+) -> dict[str, Any]:
+    """Execute discovered functions in the right order."""
+    msg = "Executing functions"
+    logger.info(msg, num_functions=len(functions), parallel=parallel)
+    results: dict[str, Any] = {}
+    # Sort by order/dependencies
+    sorted_funcs = _sort_functions(functions)
+    _validate_dependency_types(sorted_funcs)
+
+    if parallel:
+        # Group functions that can run in parallel
+        groups = _group_parallel(sorted_funcs)
+        for i, group in enumerate(groups):
+            logger.debug(
+                "Executing parallel group",
+                group=i + 1,
+                num_groups=len(groups),
+                names=[f.name for f in group],
+            )
+
+            # Ensure previous results are available
+            logger.debug("Available results", results=sorted(results))
+            # Run group in parallel
+            tasks = [execute_single(func, pool, results, inputs) for func in group]
+            group_results = await asyncio.gather(*tasks)
+            # Update results after group completes
+            results.update(dict(group_results))
+            logger.debug("Group complete", num=i + 1)
+            # Add small delay between groups to ensure timing separation
+            if i < len(groups) - 1:
+                await anyio.sleep(0.02)  # 20ms between groups
+    else:
+        # Execute sequentially
+        for func in sorted_funcs:
+            name, result = await execute_single(func, pool, results, inputs)
+            results[name] = result
+
+    return results

agentpool/running/injection.py

@@ -0,0 +1,111 @@
+"""Agent injection utilities."""
+
+from __future__ import annotations
+
+import inspect
+import typing
+from typing import TYPE_CHECKING, Any
+
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from agentpool import AgentPool, MessageNode
+
+
+logger = get_logger(__name__)
+
+
+def is_node_type(typ: Any) -> bool:
+    """Check if a type is or inherits from MessageNode."""
+    from agentpool import MessageNode
+
+    if typ is MessageNode:
+        return True
+
+    # For "real" types
+    if isinstance(typ, type):
+        return issubclass(typ, MessageNode)
+
+    # For generic types (Agent[T], etc)
+    origin = getattr(typ, "__origin__", None)
+    if origin is not None and isinstance(origin, type):
+        return issubclass(origin, MessageNode)
+
+    return False
+
+
+class NodeInjectionError(Exception):
+    """Raised when agent injection fails."""
+
+
+def inject_nodes[T, **P](
+    func: Callable[P, T],
+    pool: AgentPool,
+    provided_kwargs: dict[str, Any],
+) -> dict[str, MessageNode[Any, Any]]:
+    """Get nodes to inject based on function signature."""
+    hints = typing.get_type_hints(func)
+    params = inspect.signature(func).parameters
+    logger.debug("Injecting nodes", module=func.__module__, name=func.__qualname__, type_hint=hints)
+
+    nodes: dict[str, MessageNode[Any, Any]] = {}
+    for name, param in params.items():
+        if param.kind not in {
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            inspect.Parameter.KEYWORD_ONLY,
+        }:
+            logger.debug("Skippin: wrong parameter kind", name=name, kind=param.kind)
+            continue
+
+        hint = hints.get(name)
+        if hint is None:
+            logger.debug("Skipping: no type hint", name=name)
+            continue
+
+        # Handle Optional/Union types
+        origin = getattr(hint, "__origin__", None)
+        args = getattr(hint, "__args__", ())
+
+        # Check for MessageNode or any of its subclasses
+        is_node = (
+            is_node_type(hint)  # Direct node type
+            or (  # Optional[Node[T]] or Union containing Node
+                origin is not None and any(is_node_type(arg) for arg in args)
+            )
+        )
+
+        if not is_node:
+            msg = "Skipping. Not a node type."
+            logger.debug(msg, name=name, hint=hint, origin=origin, args=args)
+            continue
+
+        logger.debug("Found node parameter", name=name)
+
+        # Check for duplicate parameters
+        if name in provided_kwargs and provided_kwargs[name] is not None:
+            msg = (
+                f"Cannot inject node {name!r}: Parameter already provided.\n"
+                f"Remove the explicit argument or rename the parameter."
+            )
+            logger.error(msg)
+            raise NodeInjectionError(msg)
+
+        # Get node from pool
+        if name not in pool.nodes:
+            available = ", ".join(sorted(pool.nodes))
+            msg = (
+                f"No node named {name!r} found in pool.\n"
+                f"Available nodes: {available}\n"
+                f"Check your YAML configuration or node name."
+            )
+            logger.error(msg)
+            raise NodeInjectionError(msg)
+
+        nodes[name] = pool.nodes[name]
+        logger.debug("Injecting node", node=nodes[name], name=name)
+
+    logger.debug("Injection complete.", nodes=sorted(nodes))
+    return nodes

agentpool/running/run_nodes.py

@@ -0,0 +1,87 @@
+"""Automatic agent function execution."""
+
+from __future__ import annotations
+
+import asyncio
+from typing import TYPE_CHECKING, Any
+
+from agentpool.running.executor import discover_functions, execute_functions
+
+
+if TYPE_CHECKING:
+    from upathtools import JoinablePathLike
+
+    from agentpool.models.manifest import AgentsManifest
+
+
+async def run_nodes_async(
+    config: JoinablePathLike | AgentsManifest,
+    *,
+    module: str | None = None,
+    functions: list[str] | None = None,
+    inputs: dict[str, Any] | None = None,
+    parallel: bool = False,
+) -> dict[str, Any]:
+    """Execute node functions with dependency handling.
+
+    Args:
+        config: Node configuration (path or manifest)
+        module: Optional module to discover functions from
+        functions: Optional list of function names to run (auto-discovers if None)
+        inputs: Optional input values for function parameters
+        parallel: Whether to run independent functions in parallel
+
+    Returns:
+        Dict mapping function names to their results
+
+    Example:
+        ```python
+        @node_function
+        async def analyze(analyzer: Agent) -> str:
+            return await analyzer.run("...")
+
+        results = await run_nodes_async("agents.yml")
+        print(results["analyze"])
+        ```
+    """
+    from agentpool import AgentPool
+
+    # Find functions to run
+    if module:
+        discovered = discover_functions(module)
+    else:
+        # Use calling module
+        import inspect
+
+        frame = inspect.currentframe()
+        while frame:
+            if frame.f_globals.get("__name__") != __name__:
+                break
+            frame = frame.f_back
+        if not frame:
+            msg = "Could not determine calling module"
+            raise RuntimeError(msg)
+        discovered = discover_functions(frame.f_globals["__file__"])
+
+    if functions:
+        discovered = [f for f in discovered if f.name in functions]
+
+    # Run with pool
+    async with AgentPool(config) as pool:
+        return await execute_functions(discovered, pool, inputs=inputs, parallel=parallel)
+
+
+def run_nodes(config: JoinablePathLike | AgentsManifest, **kwargs: Any) -> dict[str, Any]:
+    """Run node functions synchronously.
+
+    Convenience wrapper around run_nodes_async for sync contexts.
+    See run_nodes_async for full documentation.
+
+    Args:
+        config: Agent configuration path
+        **kwargs: Arguments to pass to run_nodes_async
+
+    Returns:
+        Dict mapping function names to their results
+    """
+    return asyncio.run(run_nodes_async(config, **kwargs))

agentpool/server.py

@@ -0,0 +1,122 @@
+"""Base classes for protocol bridge servers."""
+
+from __future__ import annotations
+
+import abc
+from typing import TYPE_CHECKING, Any, Self
+
+from agentpool.log import get_logger
+
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+    from agentpool import AgentPool
+
+
+logger = get_logger(__name__)
+
+
+class ServerBridge(abc.ABC):
+    """Base class for AgentPool bridge servers.
+
+    Provides common lifecycle management, context manager protocol,
+    and running state handling for servers that bridge AgentPool agents
+    to external protocols (ACP, OpenAI API, MCP, etc.).
+    """
+
+    def __init__(self, pool: AgentPool[Any], **kwargs: Any) -> None:
+        """Initialize the server bridge.
+
+        Args:
+            pool: Agent pool to expose via the protocol
+            **kwargs: Additional configuration options
+        """
+        self._pool = pool
+        self._running = False
+        self._server_config = kwargs
+
+    @property
+    def is_running(self) -> bool:
+        """Whether the server is currently running."""
+        return self._running
+
+    @property
+    def pool(self) -> AgentPool[Any]:
+        """The underlying agent pool."""
+        return self._pool
+
+    async def run(self) -> None:
+        """Run the server.
+
+        Template method that handles common lifecycle management
+        and delegates to subclass-specific _run() implementation.
+
+        Raises:
+            RuntimeError: If server is already running
+        """
+        if self._running:
+            msg = "Server is already running"
+            raise RuntimeError(msg)
+
+        logger.info("Starting server", name=self.__class__.__name__)
+        self._running = True
+
+        try:
+            await self._run()
+        except Exception:
+            logger.exception("Server error")
+            raise
+        finally:
+            self._running = False
+            logger.info("Server stopped", name=self.__class__.__name__)
+
+    @abc.abstractmethod
+    async def _run(self) -> None:
+        """Run the server implementation.
+
+        Subclasses must implement this method to handle their
+        specific protocol setup and serving logic.
+        """
+        ...
+
+    async def shutdown(self) -> None:
+        """Shutdown the server.
+
+        Default implementation just sets running state to False.
+        Subclasses can override for custom cleanup logic.
+        """
+        if self._running:
+            logger.info("Shutting down server", name=self.__class__.__name__)
+            self._running = False
+
+    async def __aenter__(self) -> Self:
+        """Async context manager entry.
+
+        Default implementation does nothing. Subclasses can override
+        to perform initialization before the server starts.
+
+        Returns:
+            Self for fluent interface
+        """
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Async context manager exit.
+
+        Default implementation calls shutdown() if server is running.
+        Subclasses can override for custom cleanup logic.
+        """
+        if self._running:
+            await self.shutdown()
+
+    def __repr__(self) -> str:
+        """String representation of the server."""
+        status = "running" if self._running else "stopped"
+        pool_info = f"pool-{len(self._pool.agents)}-agents"
+        return f"{self.__class__.__name__}({pool_info}, {status})"

agentpool/sessions/__init__.py

@@ -0,0 +1,13 @@
+"""Session management package."""
+
+from agentpool.sessions.models import SessionData
+from agentpool.sessions.store import SessionStore
+from agentpool.sessions.manager import SessionManager
+from agentpool.sessions.session import ClientSession
+
+__all__ = [
+    "ClientSession",
+    "SessionData",
+    "SessionManager",
+    "SessionStore",
+]