contextagent 0.1.0__py3-none-any.whl

agentz/runner/iteration.py

@@ -0,0 +1,142 @@
+"""Iteration management for iterative pipeline workflows."""
+
+import time
+from typing import Any, Optional, Tuple
+
+from loguru import logger
+
+
+class IterationManager:
+    """Manages iteration lifecycle for iterative pipelines.
+
+    Handles:
+    - Starting/ending iterations with printer groups
+    - Constraint checking (max iterations, max time)
+    - Final group management
+    """
+
+    # Constants for iteration group IDs
+    ITERATION_GROUP_PREFIX = "iter"
+    FINAL_GROUP_ID = "iter-final"
+
+    def __init__(
+        self,
+        context: Any,
+        max_iterations: int = 5,
+        max_time_minutes: float = 10,
+        start_group_callback: Optional[callable] = None,
+        end_group_callback: Optional[callable] = None,
+    ):
+        """Initialize iteration manager.
+
+        Args:
+            context: Pipeline context with state and begin_iteration/mark_iteration_complete methods
+            max_iterations: Maximum number of iterations allowed
+            max_time_minutes: Maximum time in minutes
+            start_group_callback: Optional callback(group_id, title, border_style, iteration)
+            end_group_callback: Optional callback(group_id, is_done, title)
+        """
+        self.context = context
+        self.max_iterations = max_iterations
+        self.max_time_minutes = max_time_minutes
+        self.start_group_callback = start_group_callback
+        self.end_group_callback = end_group_callback
+
+        self.start_time: Optional[float] = None
+        self.current_iteration: int = 0
+
+    def start_timer(self) -> None:
+        """Start the iteration timer for constraint checking."""
+        self.start_time = time.time()
+
+    def begin_iteration(self) -> Tuple[Any, str]:
+        """Begin a new iteration with printer group.
+
+        Returns:
+            Tuple of (iteration_record, group_id)
+        """
+        iteration = self.context.begin_iteration()
+        group_id = f"{self.ITERATION_GROUP_PREFIX}-{iteration.index}"
+
+        self.current_iteration = iteration.index
+
+        if self.start_group_callback:
+            self.start_group_callback(
+                group_id,
+                title=f"Iteration {iteration.index}",
+                border_style="white",
+                iteration=iteration.index,
+            )
+
+        return iteration, group_id
+
+    def end_iteration(self, group_id: str) -> None:
+        """End the current iteration and close printer group.
+
+        Args:
+            group_id: Group ID to close
+        """
+        self.context.mark_iteration_complete()
+
+        if self.end_group_callback:
+            self.end_group_callback(group_id, is_done=True)
+
+    def start_final_group(self) -> str:
+        """Start final group for post-iteration work.
+
+        Returns:
+            Final group ID
+        """
+        if self.start_group_callback:
+            self.start_group_callback(
+                self.FINAL_GROUP_ID,
+                title="Final Report",
+                border_style="white",
+            )
+        return self.FINAL_GROUP_ID
+
+    def end_final_group(self, group_id: str) -> None:
+        """End final group.
+
+        Args:
+            group_id: Group ID to close
+        """
+        if self.end_group_callback:
+            self.end_group_callback(group_id, is_done=True)
+
+    def should_continue_iteration(self) -> bool:
+        """Check if iteration should continue.
+
+        Checks:
+        - State not complete
+        - Within max iterations
+        - Within max time
+
+        Returns:
+            True if should continue, False otherwise
+        """
+        # Check state completion
+        if hasattr(self.context, 'state') and self.context.state and self.context.state.complete:
+            return False
+
+        return self.check_constraints()
+
+    def check_constraints(self) -> bool:
+        """Check if we've exceeded our constraints (max iterations or time).
+
+        Returns:
+            True if within constraints, False if exceeded
+        """
+        if self.current_iteration >= self.max_iterations:
+            logger.info("\n=== Ending Iteration Loop ===")
+            logger.info(f"Reached maximum iterations ({self.max_iterations})")
+            return False
+
+        if self.start_time is not None:
+            elapsed_minutes = (time.time() - self.start_time) / 60
+            if elapsed_minutes >= self.max_time_minutes:
+                logger.info("\n=== Ending Iteration Loop ===")
+                logger.info(f"Reached maximum time ({self.max_time_minutes} minutes)")
+                return False
+
+        return True

agentz/runner/patterns.py

@@ -0,0 +1,215 @@
+"""High-level execution patterns for common pipeline workflows."""
+
+import asyncio
+from typing import Any, Dict, List, Optional
+
+from loguru import logger
+
+from agentz.runner.utils import record_structured_payload, serialize_output
+
+
+async def execute_tool_plan(
+    plan: Any,
+    tool_agents: Dict[str, Any],
+    group_id: str,
+    context: Any,
+    agent_step_fn: Any,
+    update_printer_fn: Optional[Any] = None,
+) -> None:
+    """Execute a routing plan with tool agents.
+
+    Args:
+        plan: AgentSelectionPlan with tasks to execute
+        tool_agents: Dict mapping agent names to agent instances
+        group_id: Group ID for printer updates
+        context: Pipeline context with state
+        agent_step_fn: Function to execute agent steps
+        update_printer_fn: Optional function for printer updates
+    """
+    # Import here to avoid circular dependency
+    from agentz.profiles.manager.routing import AgentSelectionPlan, AgentTask
+    from agentz.profiles.base import ToolAgentOutput
+
+    if not isinstance(plan, AgentSelectionPlan) or not plan.tasks:
+        return
+
+    state = context.state
+    state.current_iteration.tools.clear()
+
+    async def run_single(task: AgentTask) -> ToolAgentOutput:
+        agent = tool_agents.get(task.agent)
+        if agent is None:
+            output = ToolAgentOutput(
+                output=f"No implementation found for agent {task.agent}",
+                sources=[],
+            )
+            if update_printer_fn:
+                update_printer_fn(
+                    key=f"{group_id}:tool:{task.agent}",
+                    message=f"Completed {task.agent}",
+                    is_done=True,
+                    group_id=group_id,
+                )
+            return output
+
+        raw_result = await agent_step_fn(
+            agent=agent,
+            instructions=task.model_dump_json(),
+            span_name=task.agent,
+            span_type="tool",
+            output_model=ToolAgentOutput,
+            printer_key=f"tool:{task.agent}",
+            printer_title=f"Tool: {task.agent}",
+            printer_group_id=group_id,
+        )
+
+        if isinstance(raw_result, ToolAgentOutput):
+            output = raw_result
+        elif hasattr(raw_result, "final_output_as"):
+            output = raw_result.final_output_as(ToolAgentOutput)
+        elif hasattr(raw_result, "final_output"):
+            output = ToolAgentOutput(output=str(raw_result.final_output), sources=[])
+        else:
+            output = ToolAgentOutput(output=str(raw_result), sources=[])
+
+        try:
+            state.record_payload(output)
+        except Exception as exc:
+            logger.debug(f"Failed to record tool payload for {task.agent}: {exc}")
+
+        if update_printer_fn:
+            update_printer_fn(
+                key=f"{group_id}:tool:{task.agent}",
+                message=f"Completed {task.agent}",
+                is_done=True,
+                group_id=group_id,
+            )
+        return output
+
+    coroutines = [run_single(task) for task in plan.tasks]
+    for coro in asyncio.as_completed(coroutines):
+        tool_output = await coro
+        state.current_iteration.tools.append(tool_output)
+
+
+async def execute_tools(
+    route_plan: Any,
+    tool_agents: Dict[str, Any],
+    group_id: str,
+    context: Any,
+    agent_step_fn: Any,
+    update_printer_fn: Optional[Any] = None,
+) -> None:
+    """Execute tool agents based on routing plan.
+
+    Args:
+        route_plan: The routing plan (can be AgentSelectionPlan or other)
+        tool_agents: Dict mapping agent names to agent instances
+        group_id: Group ID for printer updates
+        context: Pipeline context with state
+        agent_step_fn: Function to execute agent steps
+        update_printer_fn: Optional function for printer updates
+    """
+    from agentz.profiles.manager.routing import AgentSelectionPlan
+
+    # Retrieve route_plan from payloads if needed
+    plan = None
+    if isinstance(route_plan, AgentSelectionPlan):
+        plan = route_plan
+    elif route_plan and hasattr(context, 'state'):
+        # Try to find AgentSelectionPlan in payloads
+        for payload in context.state.current_iteration.payloads:
+            if isinstance(payload, AgentSelectionPlan):
+                plan = payload
+                break
+
+    if plan and plan.tasks:
+        await execute_tool_plan(
+            plan, tool_agents, group_id, context, agent_step_fn, update_printer_fn
+        )
+
+
+async def run_manager_tool_loop(
+    manager_agents: Dict[str, Any],
+    tool_agents: Dict[str, Any],
+    workflow: List[str],
+    context: Any,
+    agent_step_fn: Any,
+    run_iterative_loop_fn: Any,
+    update_printer_fn: Optional[Any] = None,
+) -> Any:
+    """Execute standard manager-tool iterative pattern.
+
+    This pattern implements: observe → evaluate → route → execute tools → repeat.
+
+    Args:
+        manager_agents: Dict of manager agents (observe, evaluate, routing, writer)
+        tool_agents: Dict of tool agents
+        workflow: List of manager agent names to execute in order (e.g., ["observe", "evaluate", "routing"])
+        context: Pipeline context with state
+        agent_step_fn: Function to execute agent steps
+        run_iterative_loop_fn: Function to run iterative loop
+        update_printer_fn: Optional function for printer updates
+
+    Returns:
+        Result from final step
+
+    Example:
+        result = await run_manager_tool_loop(
+            manager_agents=self.manager_agents,
+            tool_agents=self.tool_agents,
+            workflow=["observe", "evaluate", "routing"],
+            context=self.context,
+            agent_step_fn=self.agent_step,
+            run_iterative_loop_fn=self.run_iterative_loop,
+            update_printer_fn=self.update_printer,
+        )
+    """
+    async def iteration_step(iteration, group_id: str):
+        """Execute manager workflow + tool execution."""
+        previous_output = context.state.query
+
+        # Execute manager workflow in sequence
+        for agent_name in workflow:
+            agent = manager_agents.get(agent_name)
+            if agent is None:
+                logger.warning(f"Manager agent '{agent_name}' not found, skipping")
+                continue
+
+            output = await agent(previous_output)
+
+            # Record observation for first step
+            if agent_name == workflow[0]:
+                iteration.observation = serialize_output(output)
+
+            record_structured_payload(context.state, output, context_label=agent_name)
+            previous_output = output
+
+        # Execute tools if not complete
+        if not context.state.complete and previous_output:
+            await execute_tools(
+                previous_output,
+                tool_agents,
+                group_id,
+                context,
+                agent_step_fn,
+                update_printer_fn,
+            )
+
+    async def final_step(final_group: str):
+        """Generate final report."""
+        if update_printer_fn:
+            update_printer_fn("research", "Research workflow complete", is_done=True)
+        logger.info("Research workflow completed")
+
+        writer = manager_agents.get("writer")
+        if writer:
+            await writer(context.state.findings_text())
+
+    if update_printer_fn:
+        update_printer_fn("research", "Executing research workflow...")
+
+    return await run_iterative_loop_fn(
+        iteration_body=iteration_step,
+        final_body=final_step
+    )

agentz/runner/tracker.py

@@ -0,0 +1,188 @@
+"""Runtime state tracking for agent execution operations."""
+
+from contextlib import nullcontext, contextmanager
+from contextvars import ContextVar
+from typing import Any, Dict, Optional
+
+from agents.tracing.create import trace
+from agentz.utils import Printer
+from agentz.context.data_store import DataStore
+from agentz.artifacts import RunReporter
+
+# Context variable to store the current runtime tracker
+# This allows tools to access the tracker without explicit parameter passing
+_current_runtime_tracker: ContextVar[Optional['RuntimeTracker']] = ContextVar(
+    'current_runtime_tracker',
+    default=None
+)
+
+
+class RuntimeTracker:
+    """Manages runtime state and tracking for agent execution.
+
+    This class encapsulates the runtime infrastructure needed for agent execution including:
+    - Tracing configuration and context creation
+    - Printer for status updates
+    - Reporter for recording execution events
+    - Iteration tracking
+    - Pipeline-scoped data store for sharing objects between agents
+    """
+
+    def __init__(
+        self,
+        printer: Optional[Printer] = None,
+        enable_tracing: bool = True,
+        trace_sensitive: bool = False,
+        iteration: int = 0,
+        experiment_id: Optional[str] = None,
+        reporter: Optional[RunReporter] = None,
+    ):
+        """Initialize runtime tracker.
+
+        Args:
+            printer: Optional Printer instance for status updates
+            enable_tracing: Whether tracing is enabled
+            trace_sensitive: Whether to include sensitive data in traces
+            iteration: Current iteration number (for iterative workflows)
+            experiment_id: Optional experiment ID for data store tracking
+            reporter: Optional reporter for recording execution events
+        """
+        self.printer = printer
+        self.enable_tracing = enable_tracing
+        self.trace_sensitive = trace_sensitive
+        self.iteration = iteration
+        self.reporter = reporter
+        self.data_store = DataStore(experiment_id=experiment_id)
+
+    def trace_context(self, name: str, metadata: Optional[Dict[str, Any]] = None):
+        """Create a trace context manager.
+
+        Args:
+            name: Name for the trace
+            metadata: Optional metadata to attach to trace
+
+        Returns:
+            Trace context manager if tracing enabled, otherwise nullcontext
+        """
+        if self.enable_tracing:
+            return trace(name, metadata=metadata)
+        return nullcontext()
+
+    def span_context(self, span_factory, **kwargs):
+        """Create a span context manager.
+
+        Args:
+            span_factory: Factory function for creating spans (agent_span or function_span)
+            **kwargs: Arguments to pass to span factory
+
+        Returns:
+            Span context manager if tracing enabled, otherwise nullcontext
+        """
+        if self.enable_tracing:
+            return span_factory(**kwargs)
+        return nullcontext()
+
+    def update_printer(
+        self,
+        key: str,
+        message: str,
+        is_done: bool = False,
+        hide_checkmark: bool = False,
+        title: Optional[str] = None,
+        border_style: Optional[str] = None,
+        group_id: Optional[str] = None,
+    ) -> None:
+        """Update printer status if printer is active.
+
+        Args:
+            key: Status key to update
+            message: Status message
+            is_done: Whether the task is complete
+            hide_checkmark: Whether to hide the checkmark when done
+            title: Optional panel title
+            border_style: Optional border color
+            group_id: Optional group to nest this item in
+        """
+        if self.reporter:
+            self.reporter.record_status_update(
+                item_id=key,
+                content=message,
+                is_done=is_done,
+                title=title,
+                border_style=border_style,
+                group_id=group_id,
+            )
+        if self.printer:
+            self.printer.update_item(
+                key,
+                message,
+                is_done=is_done,
+                hide_checkmark=hide_checkmark,
+                title=title,
+                border_style=border_style,
+                group_id=group_id
+            )
+
+    def log_panel(
+        self,
+        title: str,
+        content: str,
+        *,
+        border_style: Optional[str] = None,
+        iteration: Optional[int] = None,
+        group_id: Optional[str] = None,
+    ) -> None:
+        """Proxy helper for rendering standalone panels via the printer."""
+        if self.reporter:
+            self.reporter.record_panel(
+                title=title,
+                content=content,
+                border_style=border_style,
+                iteration=iteration,
+                group_id=group_id,
+            )
+        if self.printer:
+            self.printer.log_panel(
+                title,
+                content,
+                border_style=border_style,
+                iteration=iteration,
+            )
+
+    @contextmanager
+    def activate(self):
+        """Context manager to set this tracker as the current runtime tracker.
+
+        This allows tools to access the tracker via get_current_tracker().
+
+        Example:
+            with tracker.activate():
+                # Tools can now access this tracker
+                result = await agent.run(...)
+        """
+        token = _current_runtime_tracker.set(self)
+        try:
+            yield self
+        finally:
+            _current_runtime_tracker.reset(token)
+
+
+def get_current_tracker() -> Optional[RuntimeTracker]:
+    """Get the current runtime tracker (if any).
+
+    Returns:
+        The current RuntimeTracker or None if not in a runtime context
+    """
+    return _current_runtime_tracker.get()
+
+
+def get_current_data_store() -> Optional[DataStore]:
+    """Get the data store from the current runtime tracker (if any).
+
+    This is a convenience function for tools that need to access the data store.
+
+    Returns:
+        The current DataStore or None if not in a runtime context
+    """
+    tracker = get_current_tracker()
+    return tracker.data_store if tracker else None

agentz/runner/utils.py

@@ -0,0 +1,45 @@
+"""Utility functions for pipeline execution."""
+
+from typing import Any, Optional
+
+from loguru import logger
+from pydantic import BaseModel
+
+
+def record_structured_payload(
+    state: Any,
+    value: object,
+    context_label: Optional[str] = None
+) -> None:
+    """Record a structured payload to the current iteration state.
+
+    Args:
+        state: The state object (typically from context.state)
+        value: The payload to record (typically a BaseModel instance)
+        context_label: Optional label for debugging purposes
+    """
+    if isinstance(value, BaseModel):
+        try:
+            if state:
+                state.record_payload(value)
+        except Exception as exc:
+            if context_label:
+                logger.debug(f"Failed to record payload for {context_label}: {exc}")
+            else:
+                logger.debug(f"Failed to record payload: {exc}")
+
+
+def serialize_output(output: Any) -> str:
+    """Serialize agent output to string for storage.
+
+    Args:
+        output: The output to serialize (BaseModel, str, or other)
+
+    Returns:
+        String representation of the output
+    """
+    if isinstance(output, BaseModel):
+        return output.model_dump_json(indent=2)
+    elif isinstance(output, str):
+        return output
+    return str(output)