agentpool 2.1.9__py3-none-any.whl

agentpool/utils/__init__.py

@@ -0,0 +1,44 @@
+"""Utilities package."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    import jinja2
+
+
+def setup_env(env: jinja2.Environment) -> None:
+    """Used as extension point for the jinjarope environment.
+
+    Args:
+        env: The jinjarope environment to extend
+    """
+    from agentpool.agents.agent import Agent
+    from agentpool.functional import (
+        run_agent,
+        run_agent_sync,
+        get_structured,
+        get_structured_multiple,
+        pick_one,
+    )
+    from agentpool.jinja_filters import (
+        pydantic_playground,
+        pydantic_playground_iframe,
+        pydantic_playground_link,
+        pydantic_playground_url,
+    )
+
+    env.globals |= dict(agent=Agent)
+    env.filters |= {
+        "run_agent": run_agent,
+        "run_agent_sync": run_agent_sync,
+        "pick_one": pick_one,
+        "get_structured": get_structured,
+        "get_structured_multiple": get_structured_multiple,
+        "pydantic_playground": pydantic_playground,
+        "pydantic_playground_url": pydantic_playground_url,
+        "pydantic_playground_iframe": pydantic_playground_iframe,
+        "pydantic_playground_link": pydantic_playground_link,
+    }

agentpool/utils/baseregistry.py

@@ -0,0 +1,185 @@
+"""Base class for component registries."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import MutableMapping
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from psygnal.containers import EventedDict
+
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator, Iterator, Sequence
+
+    from psygnal.containers import DictEvents
+
+
+TKey = TypeVar("TKey", str, int)
+
+
+class AgentPoolError(Exception):
+    """Base exception for all AgentPool errors."""
+
+
+class BaseRegistry[TKey, TItem](MutableMapping[TKey, TItem], ABC):
+    """Base class for registries providing item storage and change notifications.
+
+    This registry implements a dictionary-like interface backed by an EventedDict,
+    providing automatic event emission for all mutations (additions, removals,
+    modifications).
+
+    Features:
+    - Dictionary-like access (registry[key] = item)
+    - Event emission for all changes
+    - Item validation
+    - Type safety
+    - Customizable error handling
+
+    Available events (accessed via .events):
+        - adding(key, value): Before an item is added
+        - added(key, value): After an item is added
+        - removing(key, value): Before an item is removed
+        - removed(key, value): After an item is removed
+        - changing(key, value): Before an item is modified
+        - changed(key, value): After an item is modified
+
+    To implement, override:
+    - _validate_item: Custom validation/transformation of items
+    - _error_class: Custom error type for exceptions
+    """
+
+    def __init__(self) -> None:
+        """Initialize an empty registry."""
+        self._items = EventedDict[TKey, TItem]()
+        self._initialized = False
+        self._configs: dict[TKey, Any] = {}
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}({self._items})"
+
+    @property
+    def is_empty(self) -> bool:
+        """Check if registry has any items."""
+        return not bool(self._items)
+
+    def has_item(self, key: TKey) -> bool:
+        """Check if an item is registered."""
+        return key in self._items
+
+    @property
+    def events(self) -> DictEvents:
+        """Access to all dictionary events."""
+        return self._items.events
+
+    def register(self, key: TKey, item: TItem | Any, replace: bool = False) -> None:
+        """Register an item."""
+        if key in self._items and not replace:
+            msg = f"Item already registered: {key}"
+            raise self._error_class(msg)
+
+        validated_item = self._validate_item(item)
+        self._items[key] = validated_item
+
+    def get(self, key: TKey) -> TItem:  # type: ignore
+        """Get an item by key."""
+        return self[key]
+
+    def list_items(self) -> Sequence[TKey]:
+        """List all registered item keys."""
+        return list(self._items.keys())
+
+    def reset(self) -> None:
+        """Reset registry to initial state."""
+        self._items.clear()
+        self._configs.clear()
+        self._initialized = False
+
+    async def startup(self) -> None:
+        """Initialize all registered items."""
+        if self._initialized:
+            return
+
+        try:
+            for item in self._items.values():
+                await self._initialize_item(item)
+            self._initialized = True
+        except Exception as exc:
+            await self.shutdown()
+            msg = f"Registry startup failed: {exc}"
+            raise self._error_class(msg) from exc
+
+    async def shutdown(self) -> None:
+        """Cleanup all registered items."""
+        if not self._initialized:
+            return
+
+        errors: list[tuple[TKey, Exception]] = []
+
+        for key, item in self._items.items():
+            try:
+                await self._cleanup_item(item)
+            except Exception as exc:  # noqa: BLE001
+                errors.append((key, exc))
+
+        self._initialized = False
+
+        if errors:
+            error_msgs = [f"{key}: {exc}" for key, exc in errors]
+            msg = f"Errors during shutdown: {', '.join(error_msgs)}"
+            raise self._error_class(msg)
+
+    @property
+    def _error_class(self) -> type[AgentPoolError]:
+        """Error class to use for this registry."""
+        return AgentPoolError
+
+    @abstractmethod
+    def _validate_item(self, item: Any) -> TItem:
+        """Validate and possibly transform item before registration."""
+
+    async def _initialize_item(self, item: TItem) -> None:
+        """Initialize an item during startup."""
+        if hasattr(item, "startup") and callable(item.startup):  # pyright: ignore
+            await item.startup()  # pyright: ignore  # ty: ignore
+
+    async def _cleanup_item(self, item: TItem) -> None:
+        """Clean up an item during shutdown."""
+        if hasattr(item, "shutdown") and callable(item.shutdown):  # pyright: ignore
+            await item.shutdown()  # pyright: ignore  # ty: ignore
+
+    # Implementing MutableMapping methods
+    def __getitem__(self, key: TKey) -> TItem:
+        try:
+            return self._items[key]
+        except KeyError as exc:
+            msg = f"Item not found: {key}"
+            raise self._error_class(msg) from exc
+
+    def __setitem__(self, key: TKey, value: Any) -> None:
+        """Support dict-style assignment."""
+        self.register(key, value)
+
+    def __contains__(self, key: object) -> bool:
+        """Support 'in' operator without raising exceptions."""
+        return key in self._items
+
+    def __delitem__(self, key: TKey) -> None:
+        if key in self._items:
+            del self._items[key]
+        else:
+            msg = f"Item not found: {key}"
+            raise self._error_class(msg)
+
+    def __iter__(self) -> Iterator[TKey]:
+        return iter(self._items)
+
+    async def __aiter__(self) -> AsyncIterator[tuple[TKey, TItem]]:
+        """Async iterate over items, ensuring they're initialized."""
+        if not self._initialized:
+            await self.startup()
+        for key, item in self._items.items():
+            yield key, item
+
+    def __len__(self) -> int:
+        return len(self._items)

agentpool/utils/count_tokens.py

@@ -0,0 +1,62 @@
+"""Token counting utilities with fallback strategies."""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from importlib.util import find_spec
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+DEFAULT_TOKEN_MODEL = "gpt-3.5-turbo"
+
+
+@lru_cache
+def has_tiktoken() -> bool:
+    """Check if tiktoken is available."""
+    return bool(find_spec("tiktoken"))
+
+
+def count_tokens(text: str, model: str | None = None) -> int:
+    """Count tokens in text with fallback strategy.
+
+    Args:
+        text: Text to count tokens for
+        model: Optional model name for tiktoken (ignored in fallback)
+
+    Returns:
+        Estimated token count
+    """
+    if has_tiktoken():
+        import tiktoken
+
+        encoding = tiktoken.encoding_for_model(model or DEFAULT_TOKEN_MODEL)
+        return len(encoding.encode(text))
+
+    # Fallback: very rough approximation
+    # Strategies could be:
+    # 1. ~4 chars per token (quick but rough)
+    # 2. Word count * 1.3 (better for English)
+    # 3. Split on common token boundaries
+    return len(text.split()) + len(text) // 4
+
+
+def batch_count_tokens(texts: Sequence[str], model: str | None = None) -> list[int]:
+    """Count tokens for multiple texts.
+
+    Args:
+        texts: Sequence of texts to count
+        model: Optional model name for tiktoken
+
+    Returns:
+        List of token counts
+    """
+    if has_tiktoken():
+        import tiktoken
+
+        encoding = tiktoken.encoding_for_model(model or DEFAULT_TOKEN_MODEL)
+        return [len(encoding.encode(text)) for text in texts]
+
+    return [count_tokens(text) for text in texts]

agentpool/utils/dag.py

@@ -0,0 +1,184 @@
+"""Minimal DAG (Directed Acyclic Graph) implementation.
+
+This module provides a lightweight DAG node class for tracking message flows
+and generating mermaid diagrams. It replaces the bigtree dependency with
+only the functionality actually used.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+
+if TYPE_CHECKING:
+    from collections.abc import Iterable
+
+
+@dataclass
+class DAGNode:
+    """A node in a Directed Acyclic Graph.
+
+    Nodes can have multiple parents and multiple children, representing
+    a DAG structure suitable for tracking message flows between agents.
+
+    Example:
+        >>> a = DAGNode("a")
+        >>> b = DAGNode("b")
+        >>> c = DAGNode("c")
+        >>> c.add_parent(a)
+        >>> c.add_parent(b)
+        >>> a.children
+        [DAGNode(name='c')]
+    """
+
+    name: str
+    """Name/identifier of this node."""
+
+    _parents: list[DAGNode] = field(default_factory=list, repr=False)
+    _children: list[DAGNode] = field(default_factory=list, repr=False)
+
+    @property
+    def parents(self) -> list[DAGNode]:
+        """Get parent nodes."""
+        return list(self._parents)
+
+    @property
+    def children(self) -> list[DAGNode]:
+        """Get child nodes."""
+        return list(self._children)
+
+    @property
+    def is_root(self) -> bool:
+        """Check if this node has no parents."""
+        return len(self._parents) == 0
+
+    @property
+    def is_leaf(self) -> bool:
+        """Check if this node has no children."""
+        return len(self._children) == 0
+
+    def add_parent(self, parent: DAGNode) -> None:
+        """Add a parent node, also adding self as child of parent.
+
+        Args:
+            parent: Node to add as parent
+
+        Raises:
+            ValueError: If adding would create a cycle
+        """
+        if parent is self:
+            msg = "Node cannot be its own parent"
+            raise ValueError(msg)
+        if parent in self._parents:
+            return  # Already a parent
+        if self._is_ancestor_of(parent):
+            msg = "Adding this parent would create a cycle"
+            raise ValueError(msg)
+
+        self._parents.append(parent)
+        parent._children.append(self)
+
+    def add_child(self, child: DAGNode) -> None:
+        """Add a child node, also adding self as parent of child.
+
+        Args:
+            child: Node to add as child
+
+        Raises:
+            ValueError: If adding would create a cycle
+        """
+        child.add_parent(self)
+
+    def _is_ancestor_of(self, node: DAGNode) -> bool:
+        """Check if self is an ancestor of the given node."""
+        visited: set[str] = set()
+
+        def _check(current: DAGNode) -> bool:
+            if current.name in visited:
+                return False
+            visited.add(current.name)
+            if current is self:
+                return True
+            return any(_check(p) for p in current._parents)
+
+        return _check(node)
+
+    def __rshift__(self, other: DAGNode) -> DAGNode:
+        """Set child using >> operator: parent >> child."""
+        other.add_parent(self)
+        return other
+
+    def __lshift__(self, other: DAGNode) -> DAGNode:
+        """Set parent using << operator: child << parent."""
+        self.add_parent(other)
+        return self
+
+
+def dag_iterator(root: DAGNode) -> Iterable[tuple[DAGNode, DAGNode]]:
+    """Iterate through all edges in a DAG starting from a node.
+
+    Traverses both upward (to parents) and downward (to children) to
+    discover all edges reachable from the starting node.
+
+    Args:
+        root: Starting node for iteration
+
+    Yields:
+        Tuples of (parent, child) for each edge in the DAG
+    """
+    visited_nodes: set[str] = set()
+    visited_edges: set[tuple[str, str]] = set()
+
+    def _iterate(node: DAGNode) -> Iterable[tuple[DAGNode, DAGNode]]:
+        node_name = node.name
+        if node_name in visited_nodes:
+            return
+        visited_nodes.add(node_name)
+
+        # Yield edges to parents (upward)
+        for parent in node._parents:
+            edge = (parent.name, node_name)
+            if edge not in visited_edges:
+                visited_edges.add(edge)
+                yield parent, node
+
+        # Yield edges to children (downward)
+        for child in node._children:
+            edge = (node_name, child.name)
+            if edge not in visited_edges:
+                visited_edges.add(edge)
+                yield node, child
+
+        # Recursively visit parents
+        for parent in node._parents:
+            yield from _iterate(parent)
+
+        # Recursively visit children
+        for child in node._children:
+            yield from _iterate(child)
+
+    yield from _iterate(root)
+
+
+def dag_to_list(dag: DAGNode) -> list[tuple[str, str]]:
+    """Export DAG edges as list of (parent_name, child_name) tuples.
+
+    Example:
+        >>> a = DAGNode("a")
+        >>> b = DAGNode("b")
+        >>> c = DAGNode("c")
+        >>> c.add_parent(a)
+        >>> c.add_parent(b)
+        >>> d = DAGNode("d")
+        >>> d.add_parent(c)
+        >>> sorted(dag_to_list(a))
+        [('a', 'c'), ('b', 'c'), ('c', 'd')]
+
+    Args:
+        dag: Any node in the DAG (will traverse to find all edges)
+
+    Returns:
+        List of (parent_name, child_name) tuples for all edges
+    """
+    return [(parent.name, child.name) for parent, child in dag_iterator(dag)]

agentpool/utils/importing.py

@@ -0,0 +1,206 @@
+"""Utilities for importing callables and classes from dotted paths."""
+
+from __future__ import annotations
+
+import importlib
+import inspect
+from pathlib import Path
+import pkgutil
+from types import ModuleType
+from typing import TYPE_CHECKING, Any
+
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Generator, Iterator
+
+
+def get_module_source(
+    import_path: str,
+    recursive: bool = False,
+    include_tests: bool = False,
+) -> str:
+    """Get source code from a module or package."""
+    try:
+        module = importlib.import_module(import_path)
+        sources = _get_sources(module, recursive=recursive, include_tests=include_tests)
+        return "\n\n# " + "-" * 40 + "\n\n".join(sources)
+
+    except ImportError as exc:
+        msg = f"Could not import module: {import_path}"
+        raise ValueError(msg) from exc
+
+
+def _get_sources(
+    module: ModuleType,
+    recursive: bool,
+    include_tests: bool,
+) -> Generator[str]:
+    """Generate source code for a module and optionally its submodules."""
+    # Get the module's source code
+    if hasattr(module, "__file__") and module.__file__:
+        path = Path(module.__file__)
+        if _should_include_file(path, include_tests):
+            yield f"# File: {path}\n{inspect.getsource(module)}"
+
+    # If recursive and it's a package, get all submodules
+    if recursive and hasattr(module, "__path__"):
+        for _, name, _ in pkgutil.iter_modules(module.__path__):
+            submodule_path = f"{module.__name__}.{name}"
+            try:
+                submodule = importlib.import_module(submodule_path)
+                yield from _get_sources(submodule, recursive, include_tests)
+            except ImportError:
+                continue
+
+
+def _should_include_file(path: Path, include_tests: bool) -> bool:
+    """Check if a file should be included in the source."""
+    if not include_tests and any(p.startswith("test") for p in path.parts):
+        return False
+    return path.suffix == ".py"
+
+
+def import_callable(path: str) -> Callable[..., Any]:
+    """Import a callable from a dotted path.
+
+    Supports both dot and colon notation:
+    - Dot notation: module.submodule.Class.method
+    - Colon notation: module.submodule:Class.method
+
+    Examples:
+        >>> import_callable("os.path.join")
+        >>> import_callable("builtins.str.upper")
+        >>> import_callable("sqlalchemy.orm:Session.query")
+
+    Args:
+        path: Import path using dots and/or colon
+
+    Returns:
+        Imported callable
+
+    Raises:
+        ValueError: If path cannot be imported or result isn't callable
+    """
+    if not path:
+        msg = "Import path cannot be empty"
+        raise ValueError(msg)
+
+    # Normalize path - replace colon with dot if present
+    normalized_path = path.replace(":", ".")
+    parts = normalized_path.split(".")
+
+    # Try importing progressively smaller module paths
+    for i in range(len(parts), 0, -1):
+        try:
+            # Try current module path
+            module_path = ".".join(parts[:i])
+            module = importlib.import_module(module_path)
+
+            # Walk remaining parts as attributes
+            obj = module
+            for part in parts[i:]:
+                obj = getattr(obj, part)
+
+            # Check if we got a callable
+            if callable(obj):
+                return obj
+
+            msg = f"Found object at {path} but it isn't callable"
+            raise ValueError(msg)
+
+        except ImportError:
+            # Try next shorter path
+            continue
+        except AttributeError:
+            # Attribute not found - try next shorter path
+            continue
+
+    # If we get here, no import combination worked
+    msg = f"Could not import callable from path: {path}"
+    raise ValueError(msg)
+
+
+def import_class(path: str) -> type:
+    """Import a class from a dotted path.
+
+    Args:
+        path: Dot-separated path to the class
+
+    Returns:
+        The imported class
+
+    Raises:
+        ValueError: If path is invalid or doesn't point to a class
+    """
+    try:
+        obj = import_callable(path)
+        if not isinstance(obj, type):
+            msg = f"{path} is not a class"
+            raise TypeError(msg)  # noqa: TRY301
+    except Exception as exc:
+        msg = f"Failed to import class from {path}"
+        raise ValueError(msg) from exc
+    else:
+        return obj
+
+
+def get_pyobject_members(
+    obj: type | ModuleType | Any,
+    *,
+    include_imported: bool = False,
+) -> Iterator[tuple[str, Callable[..., Any]]]:
+    """Get callable members defined in a Python object.
+
+    Works with modules, classes, and instances. Only returns public callable
+    members (functions, methods, etc.) that are defined in the object's module
+    unless include_imported is True.
+
+    Args:
+        obj: Any Python object to inspect (module, class, instance)
+        include_imported: Whether to include imported/inherited callables
+
+    Yields:
+        Tuples of (name, callable) for each public callable
+
+    Example:
+        >>> class MyClass:
+        ...     def method(self): pass
+        ...     def _private(self): pass
+        >>> for name, func in get_pyobject_members(MyClass()):
+        ...     print(name)
+        method
+
+        >>> import my_module
+        >>> for name, func in get_pyobject_members(my_module):
+        ...     print(name)
+        public_function
+    """
+    # Get the module where the object is defined
+    defining_module = obj.__name__ if isinstance(obj, ModuleType) else obj.__module__
+
+    for name, member in inspect.getmembers(obj, inspect.isroutine):
+        if name.startswith("_"):
+            continue
+
+        # Check if callable is defined in the object's module
+        if include_imported or getattr(member, "__module__", None) == defining_module:
+            yield name, member
+
+
+if __name__ == "__main__":
+    # ATTENTION: Dont modify this script.
+    import sys
+
+    if len(sys.argv) != 2:  # noqa: PLR2004
+        print("Usage: python importing.py <dot.path.to.object>", file=sys.stderr)
+        sys.exit(1)
+
+    dot_path = sys.argv[1]
+
+    try:
+        obj = import_callable(dot_path)
+        source = inspect.getsource(obj)
+        print(source)
+    except Exception as e:  # noqa: BLE001
+        print(f"Error: {e}", file=sys.stderr)
+        sys.exit(1)