aury-agent 0.0.4__py3-none-any.whl

aury/agents/usage/tracker.py

@@ -0,0 +1,236 @@
+"""Usage tracking for LLM and other services."""
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field, asdict
+from datetime import datetime
+from decimal import Decimal
+from enum import Enum
+from typing import Any
+
+from ..core.event_bus import EventBus, Events
+
+# Alias for backward compatibility
+Bus = EventBus
+
+
+class UsageType(Enum):
+    """Types of usage to track."""
+    LLM_INPUT = "llm_input"
+    LLM_OUTPUT = "llm_output"
+    LLM_CACHE_READ = "llm_cache_read"
+    LLM_CACHE_WRITE = "llm_cache_write"
+    EMBEDDING = "embedding"
+    IMAGE_GEN = "image_gen"
+    SPEECH = "speech"
+    SEARCH = "search"
+    EXTERNAL_API = "external_api"
+
+
+@dataclass
+class UsageEntry:
+    """A single usage record."""
+    id: str
+    type: UsageType
+    provider: str
+    model: str | None = None
+    units: int = 0
+    unit_type: str = "tokens"  # tokens, requests, characters, etc.
+    cost: Decimal | None = None
+    
+    # Context
+    session_id: str | None = None
+    invocation_id: str | None = None
+    tool_id: str | None = None
+    step: int | None = None
+    
+    # Metadata
+    created_at: datetime = field(default_factory=datetime.now)
+    metadata: dict[str, Any] = field(default_factory=dict)
+    
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "type": self.type.value,
+            "provider": self.provider,
+            "model": self.model,
+            "units": self.units,
+            "unit_type": self.unit_type,
+            "cost": float(self.cost) if self.cost else None,
+            "session_id": self.session_id,
+            "invocation_id": self.invocation_id,
+            "tool_id": self.tool_id,
+            "step": self.step,
+            "created_at": self.created_at.isoformat(),
+            "metadata": self.metadata,
+        }
+
+
+class UsageTracker:
+    """Track usage of LLM and other services.
+    
+    Provides methods for recording usage, computing costs,
+    and generating summaries.
+    """
+    
+    def __init__(self, bus: Bus | None = None):
+        self.bus = bus
+        self._entries: list[UsageEntry] = []
+        self._lock = asyncio.Lock()
+        self._counter = 0
+    
+    def _generate_id(self) -> str:
+        self._counter += 1
+        return f"usage_{self._counter:06d}"
+    
+    async def record(self, entry: UsageEntry) -> None:
+        """Record a usage entry."""
+        async with self._lock:
+            self._entries.append(entry)
+        
+        if self.bus:
+            await self.bus.publish(Events.USAGE_RECORDED, entry.to_dict())
+    
+    async def record_llm(
+        self,
+        provider: str,
+        model: str,
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+        cache_read_tokens: int = 0,
+        cache_write_tokens: int = 0,
+        session_id: str | None = None,
+        invocation_id: str | None = None,
+        step: int | None = None,
+        **metadata: Any,
+    ) -> None:
+        """Convenience method for recording LLM usage."""
+        base = {
+            "provider": provider,
+            "model": model,
+            "unit_type": "tokens",
+            "session_id": session_id,
+            "invocation_id": invocation_id,
+            "step": step,
+            "metadata": metadata,
+        }
+        
+        if input_tokens > 0:
+            await self.record(UsageEntry(
+                id=self._generate_id(),
+                type=UsageType.LLM_INPUT,
+                units=input_tokens,
+                **base,
+            ))
+        
+        if output_tokens > 0:
+            await self.record(UsageEntry(
+                id=self._generate_id(),
+                type=UsageType.LLM_OUTPUT,
+                units=output_tokens,
+                **base,
+            ))
+        
+        if cache_read_tokens > 0:
+            await self.record(UsageEntry(
+                id=self._generate_id(),
+                type=UsageType.LLM_CACHE_READ,
+                units=cache_read_tokens,
+                **base,
+            ))
+        
+        if cache_write_tokens > 0:
+            await self.record(UsageEntry(
+                id=self._generate_id(),
+                type=UsageType.LLM_CACHE_WRITE,
+                units=cache_write_tokens,
+                **base,
+            ))
+    
+    async def record_embedding(
+        self,
+        provider: str,
+        model: str,
+        tokens: int,
+        session_id: str | None = None,
+        **metadata: Any,
+    ) -> None:
+        """Record embedding usage."""
+        await self.record(UsageEntry(
+            id=self._generate_id(),
+            type=UsageType.EMBEDDING,
+            provider=provider,
+            model=model,
+            units=tokens,
+            unit_type="tokens",
+            session_id=session_id,
+            metadata=metadata,
+        ))
+    
+    def summarize(
+        self,
+        session_id: str | None = None,
+        invocation_id: str | None = None,
+    ) -> dict[str, Any]:
+        """Generate usage summary."""
+        entries = self._entries
+        
+        if session_id:
+            entries = [e for e in entries if e.session_id == session_id]
+        if invocation_id:
+            entries = [e for e in entries if e.invocation_id == invocation_id]
+        
+        by_type: dict[str, int] = {}
+        by_provider: dict[str, int] = {}
+        by_model: dict[str, int] = {}
+        total_cost = Decimal(0)
+        
+        for e in entries:
+            by_type[e.type.value] = by_type.get(e.type.value, 0) + e.units
+            by_provider[e.provider] = by_provider.get(e.provider, 0) + e.units
+            if e.model:
+                by_model[e.model] = by_model.get(e.model, 0) + e.units
+            if e.cost:
+                total_cost += e.cost
+        
+        return {
+            "total_cost": float(total_cost),
+            "by_type": by_type,
+            "by_provider": by_provider,
+            "by_model": by_model,
+            "entry_count": len(entries),
+            "total_tokens": sum(by_type.values()),
+        }
+    
+    def get_entries(
+        self,
+        session_id: str | None = None,
+        invocation_id: str | None = None,
+        type_filter: UsageType | None = None,
+    ) -> list[UsageEntry]:
+        """Get filtered usage entries."""
+        entries = self._entries
+        
+        if session_id:
+            entries = [e for e in entries if e.session_id == session_id]
+        if invocation_id:
+            entries = [e for e in entries if e.invocation_id == invocation_id]
+        if type_filter:
+            entries = [e for e in entries if e.type == type_filter]
+        
+        return entries
+    
+    def clear(self, session_id: str | None = None) -> int:
+        """Clear usage entries.
+        
+        Returns:
+            Number of entries cleared
+        """
+        if session_id:
+            original = len(self._entries)
+            self._entries = [e for e in self._entries if e.session_id != session_id]
+            return original - len(self._entries)
+        
+        count = len(self._entries)
+        self._entries.clear()
+        return count

aury/agents/workflow/__init__.py

@@ -0,0 +1,85 @@
+"""Workflow system for DAG-based orchestration."""
+from .types import (
+    NodeType,
+    Position,
+    NodeSpec,
+    EdgeSpec,
+    WorkflowSpec,
+    Workflow,
+)
+from .parser import (
+    WorkflowParser,
+    WorkflowValidationError,
+)
+from .expression import (
+    ExpressionEvaluator,
+    ExpressionError,
+)
+from .state import (
+    WorkflowState,
+    MergeStrategy,
+    CollectListStrategy,
+    CollectDictStrategy,
+    FirstSuccessStrategy,
+    get_merge_strategy,
+)
+from .dag import DAGExecutor
+from .executor import WorkflowExecutor
+from ..core.factory import AgentFactory
+from .adapter import WorkflowAgent
+from .dsl import (
+    DSLNode,
+    DSLSequence,
+    DSLStep,
+    DSLParallel,
+    DSLCondition,
+    DSLWorkflow,
+    workflow,
+    step,
+    parallel,
+    condition,
+    skip,
+)
+
+__all__ = [
+    # Types
+    "NodeType",
+    "Position",
+    "NodeSpec",
+    "EdgeSpec",
+    "WorkflowSpec",
+    "Workflow",
+    # Parser
+    "WorkflowParser",
+    "WorkflowValidationError",
+    # Expression
+    "ExpressionEvaluator",
+    "ExpressionError",
+    # State
+    "WorkflowState",
+    "MergeStrategy",
+    "CollectListStrategy",
+    "CollectDictStrategy",
+    "FirstSuccessStrategy",
+    "get_merge_strategy",
+    # DAG
+    "DAGExecutor",
+    # Executor
+    "WorkflowExecutor",
+    # Factory
+    "AgentFactory",
+    # Agent
+    "WorkflowAgent",
+    # DSL
+    "DSLNode",
+    "DSLSequence",
+    "DSLStep",
+    "DSLParallel",
+    "DSLCondition",
+    "DSLWorkflow",
+    "workflow",
+    "step",
+    "parallel",
+    "condition",
+    "skip",
+]

aury/agents/workflow/adapter.py

@@ -0,0 +1,268 @@
+"""Workflow agent - executes DAG-based workflows.
+
+WorkflowAgent uses the unified BaseAgent constructor:
+    __init__(self, ctx: InvocationContext, config: AgentConfig | None = None)
+
+Workflow definition and AgentFactory are passed via config or set after init.
+
+Middleware hooks are called at:
+- on_agent_start/end: workflow start/end
+- on_subagent_start/end: each node execution (in executor)
+"""
+from __future__ import annotations
+
+from typing import Any, AsyncIterator, ClassVar, Literal, TYPE_CHECKING
+
+from ..core.base import BaseAgent, AgentConfig
+from ..core.context import InvocationContext
+from ..core.logging import workflow_logger as logger
+from ..core.types.block import BlockEvent, BlockKind, BlockOp
+from ..middleware import HookAction
+from .types import Workflow
+from .executor import WorkflowExecutor
+
+if TYPE_CHECKING:
+    from ..core.factory import AgentFactory
+    from ..core.types.session import Session
+    from ..backends import Backends
+    from ..core.event_bus import Bus
+    from ..middleware import MiddlewareChain
+
+
+class WorkflowAgent(BaseAgent):
+    """Workflow agent - executes DAG-based workflows.
+    
+    Two ways to create:
+    
+    1. Simple (recommended):
+        agent = WorkflowAgent.create(workflow=wf, agent_factory=factory)
+    
+    2. Advanced (for custom Session/Backends/Bus):
+        ctx = InvocationContext(session=session, backends=backends, bus=bus)
+        agent = WorkflowAgent(ctx, config)
+        agent.set_workflow(workflow, agent_factory)
+    """
+    
+    agent_type: ClassVar[Literal["react", "workflow"]] = "workflow"
+    
+    @classmethod
+    def create(
+        cls,
+        workflow: Workflow,
+        agent_factory: "WorkflowAgentFactory",
+        config: AgentConfig | None = None,
+        *,
+        backends: "Backends | None" = None,
+        session: "Session | None" = None,
+        bus: "Bus | None" = None,
+        middleware: "MiddlewareChain | None" = None,
+    ) -> "WorkflowAgent":
+        """Create WorkflowAgent with minimal boilerplate.
+        
+        Args:
+            workflow: Workflow definition
+            agent_factory: Factory for creating sub-agents
+            config: Agent configuration (optional)
+            backends: Backends container (auto-created if None)
+            session: Session object (auto-created if None)
+            bus: Event bus (auto-created if None)
+            middleware: Middleware chain (optional)
+            
+        Returns:
+            Configured WorkflowAgent ready to run
+            
+        Example:
+            agent = WorkflowAgent.create(
+                workflow=my_workflow,
+                agent_factory=factory,
+            )
+            async for response in agent.run(inputs):
+                print(response)
+        """
+        from ..core.event_bus import Bus
+        from ..core.types.session import Session, generate_id
+        from ..backends import Backends
+        
+        # Auto-create backends if not provided
+        if backends is None:
+            backends = Backends.create_default()
+        
+        # Auto-create missing components
+        if session is None:
+            session = Session(id=generate_id("sess"))
+        if bus is None:
+            bus = Bus()
+        
+        # Build context
+        ctx = InvocationContext(
+            session=session,
+            invocation_id=generate_id("inv"),
+            agent_id=config.name if config else workflow.spec.name,
+            backends=backends,
+            bus=bus,
+            middleware=middleware,
+        )
+        
+        agent = cls(ctx, config)
+        agent.set_workflow(workflow, agent_factory)
+        return agent
+
+    def __init__(
+        self,
+        ctx: InvocationContext,
+        config: AgentConfig | None = None,
+    ):
+        """Initialize WorkflowAgent.
+        
+        Args:
+            ctx: InvocationContext with storage, bus, etc.
+            config: Agent configuration (may contain workflow in metadata)
+        """
+        super().__init__(ctx, config)
+        
+        # Workflow can be set via config.metadata or set_workflow()
+        self.workflow: Workflow | None = self.config.metadata.get('workflow')
+        self.agent_factory: "WorkflowAgentFactory | None" = self.config.metadata.get('agent_factory')
+        self._executor: WorkflowExecutor | None = None
+    
+    def set_workflow(
+        self,
+        workflow: Workflow,
+        agent_factory: "WorkflowAgentFactory",
+    ) -> "WorkflowAgent":
+        """Set workflow definition and factory.
+        
+        Args:
+            workflow: Workflow definition
+            agent_factory: Factory for creating sub-agents
+            
+        Returns:
+            Self for chaining
+        """
+        self.workflow = workflow
+        self.agent_factory = agent_factory
+        return self
+    
+    async def _execute(self, input: Any) -> None:
+        """Execute workflow with middleware hooks.
+        
+        Args:
+            input: Workflow inputs (dict or single value)
+            
+        Raises:
+            ValueError: If workflow or agent_factory not set
+        """
+        # Validate workflow is set
+        if self.workflow is None:
+            raise ValueError("Workflow not set. Call set_workflow() first.")
+        if self.agent_factory is None:
+            raise ValueError("AgentFactory not set. Call set_workflow() first.")
+        
+        inputs = input if isinstance(input, dict) else {"input": input}
+        
+        # Build middleware context
+        mw_context = {
+            "session_id": self.session.id,
+            "invocation_id": self.ctx.invocation_id,
+            "agent_id": self.name,
+            "agent_type": self.agent_type,
+            "workflow_name": self.workflow.spec.name,
+        }
+        
+        # === Middleware: on_agent_start ===
+        if self.middleware:
+            hook_result = await self.middleware.process_agent_start(
+                self.name, inputs, mw_context
+            )
+            if hook_result.action == HookAction.STOP:
+                logger.info("Workflow stopped by middleware on_agent_start")
+                await self.ctx.emit(BlockEvent(
+                    kind=BlockKind.ERROR,
+                    op=BlockOp.APPLY,
+                    data={"message": hook_result.message or "Stopped by middleware"},
+                ))
+                return
+            elif hook_result.action == HookAction.SKIP:
+                logger.info("Workflow skipped by middleware on_agent_start")
+                return
+        
+        try:
+            # Create executor with context (pass middleware)
+            self._executor = WorkflowExecutor(
+                workflow=self.workflow,
+                agent_factory=self.agent_factory,
+                ctx=self.ctx,
+                middleware=self.middleware,
+            )
+            
+            result = await self._executor.execute(inputs)
+            
+            # === Middleware: on_agent_end ===
+            if self.middleware:
+                await self.middleware.process_agent_end(
+                    self.name, result, mw_context
+                )
+        
+        except Exception as e:
+            # === Middleware: on_error ===
+            if self.middleware:
+                processed_error = await self.middleware.process_error(e, mw_context)
+                if processed_error is None:
+                    logger.info("Error suppressed by middleware")
+                    return
+            raise
+    
+    async def pause(self) -> str:
+        """Pause workflow and return invocation ID."""
+        if self._executor:
+            self._executor.pause()
+        return self.ctx.invocation_id
+    
+    async def _resume_internal(self, invocation_id: str) -> None:
+        """Internal resume logic."""
+        if self.workflow is None or self.agent_factory is None:
+            raise ValueError("Workflow not set. Call set_workflow() first.")
+        
+        # Load saved state
+        state_key = f"workflow_state:{invocation_id}"
+        if not self.ctx.backends or not self.ctx.backends.state:
+            raise ValueError("No state backend available")
+        saved_state = await self.ctx.backends.state.get("workflow", state_key)
+        
+        if not saved_state:
+            raise ValueError(f"No saved state for invocation: {invocation_id}")
+        
+        inputs = saved_state.get("inputs", {})
+        
+        # Create executor and resume
+        self._executor = WorkflowExecutor(
+            workflow=self.workflow,
+            agent_factory=self.agent_factory,
+            ctx=self.ctx,
+        )
+        
+        await self._executor.execute(inputs, resume_state=saved_state)
+    
+    async def resume(self, invocation_id: str) -> AsyncIterator[BlockEvent]:
+        """Resume paused workflow."""
+        import asyncio
+        from ..core.context import _emit_queue_var
+        
+        queue: asyncio.Queue[BlockEvent] = asyncio.Queue()
+        token = _emit_queue_var.set(queue)
+        
+        try:
+            exec_task = asyncio.create_task(self._resume_internal(invocation_id))
+            
+            while not exec_task.done() or not queue.empty():
+                try:
+                    block = await asyncio.wait_for(queue.get(), timeout=0.05)
+                    yield block
+                except asyncio.TimeoutError:
+                    continue
+            
+            await exec_task
+            
+        finally:
+            _emit_queue_var.reset(token)
+

aury/agents/workflow/dag.py

@@ -0,0 +1,116 @@
+"""Generic DAG executor."""
+from __future__ import annotations
+
+from typing import Any, Callable, TypeVar
+
+T = TypeVar("T")
+
+
+class DAGExecutor:
+    """Generic DAG parallel executor."""
+    
+    def __init__(
+        self,
+        tasks: list[T],
+        get_task_id: Callable[[T], str],
+        get_dependencies: Callable[[T], list[str]],
+    ):
+        self.tasks = {get_task_id(task): task for task in tasks}
+        self.get_task_id = get_task_id
+        self.get_dependencies = get_dependencies
+        
+        self.completed: set[str] = set()
+        self.failed: set[str] = set()
+        self.running: set[str] = set()
+        self.skipped: set[str] = set()
+    
+    def get_ready_tasks(self) -> list[T]:
+        """Get tasks ready for execution.
+        
+        A task is ready when all its dependencies are completed or skipped.
+        If a dependency has failed, the task will never be ready.
+        """
+        ready = []
+        processed = self.completed | self.failed | self.running | self.skipped
+        
+        for task_id, task in self.tasks.items():
+            if task_id in processed:
+                continue
+            
+            deps = self.get_dependencies(task)
+            deps_satisfied = all(
+                dep_id in (self.completed | self.skipped)
+                for dep_id in deps
+            )
+            
+            if deps_satisfied:
+                ready.append(task)
+        
+        return ready
+    
+    def is_blocked(self) -> bool:
+        """Check if DAG is blocked due to failed dependencies.
+        
+        Returns True if there are unprocessed tasks that can never run
+        because their dependencies have failed.
+        """
+        if not self.failed:
+            return False
+        
+        processed = self.completed | self.failed | self.running | self.skipped
+        
+        for task_id, task in self.tasks.items():
+            if task_id in processed:
+                continue
+            
+            # Check if any dependency has failed
+            deps = self.get_dependencies(task)
+            for dep_id in deps:
+                if dep_id in self.failed:
+                    return True
+        
+        return False
+    
+    def mark_running(self, task_id: str) -> None:
+        self.running.add(task_id)
+    
+    def mark_completed(self, task_id: str) -> None:
+        self.running.discard(task_id)
+        self.completed.add(task_id)
+    
+    def mark_failed(self, task_id: str) -> None:
+        self.running.discard(task_id)
+        self.failed.add(task_id)
+    
+    def mark_skipped(self, task_id: str) -> None:
+        self.skipped.add(task_id)
+    
+    def is_finished(self) -> bool:
+        total_processed = len(self.completed) + len(self.failed) + len(self.skipped)
+        return total_processed == len(self.tasks)
+    
+    def has_failures(self) -> bool:
+        return len(self.failed) > 0
+    
+    def get_status(self) -> dict[str, Any]:
+        pending = (
+            len(self.tasks)
+            - len(self.completed)
+            - len(self.failed)
+            - len(self.skipped)
+            - len(self.running)
+        )
+        return {
+            "total": len(self.tasks),
+            "completed": len(self.completed),
+            "failed": len(self.failed),
+            "skipped": len(self.skipped),
+            "running": len(self.running),
+            "pending": pending,
+        }
+    
+    def reset(self) -> None:
+        self.completed.clear()
+        self.failed.clear()
+        self.running.clear()
+        self.skipped.clear()