agentexec 0.1.0__py3-none-any.whl

agentexec/core/worker.py

@@ -0,0 +1,304 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+import multiprocessing as mp
+from collections.abc import Callable
+from multiprocessing.synchronize import Event as EventClass
+from typing import Any, ClassVar, Generic, TypeVar, cast
+
+from sqlalchemy import Engine
+from sqlalchemy.orm import Session, sessionmaker
+
+from agentexec.config import CONF
+from agentexec.core.task import Task, TaskHandler
+
+logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "Worker",
+    "WorkerPool",
+    "get_worker_session",
+]
+
+T = TypeVar("T")
+
+
+class classproperty(Generic[T]):
+    """Decorator for class-level properties.
+
+    Generic decorator that preserves the return type of the decorated method.
+    """
+
+    def __init__(self, func: Callable[[Any], T]) -> None:
+        self.func = func
+
+    def __get__(self, obj: Any, owner: type) -> T:
+        return self.func(owner)
+
+
+def _build_session(engine: Engine) -> Session:
+    """Helper to build a new SQLAlchemy session from engine."""
+    SessionLocal = sessionmaker(bind=engine, autocommit=False, autoflush=False)
+    return cast(Session, SessionLocal())
+
+
+def get_worker_session() -> Session:
+    """Get the current worker's database session.
+
+    Returns:
+        Session from the current worker
+
+    Raises:
+        RuntimeError: If called outside of a worker process
+    """
+    return Worker.session
+
+
+class Worker:
+    """Individual worker process with isolated state.
+
+    Each worker maintains its own database session, event loop, and processes
+    tasks from the queue independently. Workers run in separate processes for
+    true parallelism.
+    """
+
+    # Class-level reference to current worker in this process
+    current: ClassVar[Worker | None] = None
+
+    # Instance attributes
+    _worker_id: int
+    _handlers: dict[str, TaskHandler]
+    _shutdown_event: EventClass
+    _session: Session
+    _loop: asyncio.AbstractEventLoop
+
+    def __init__(
+        self,
+        worker_id: int,
+        engine: Engine,
+        handlers: dict[str, TaskHandler],
+        shutdown_event: EventClass,
+    ):
+        """Initialize worker with isolated state.
+
+        Args:
+            worker_id: Unique identifier for this worker
+            engine: SQLAlchemy engine to create session from
+            handlers: Task handler registry (shared read-only)
+            shutdown_event: Multiprocessing event for coordinated shutdown
+        """
+        self._worker_id = worker_id
+        self._handlers = handlers
+        self._shutdown_event = shutdown_event
+
+        # Create worker-owned resources
+        self._session = _build_session(engine)
+        self._loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(self._loop)
+
+    @classmethod
+    def run_in_process(
+        cls,
+        worker_id: int,
+        engine: Engine,
+        handlers: dict[str, TaskHandler],
+        shutdown_event: EventClass,
+    ) -> None:
+        """Entry point for running a worker in a new process.
+
+        Creates a Worker instance and runs it. This is called by multiprocessing
+        to start a worker in a new process.
+
+        Args:
+            worker_id: Unique identifier for this worker
+            engine: SQLAlchemy engine to create session from
+            handlers: Task handler registry
+            shutdown_event: Multiprocessing event for coordinated shutdown
+        """
+        instance = cls(worker_id, engine, handlers, shutdown_event)
+        Worker.current = instance
+        instance.run()
+
+    @classproperty
+    def session(cls) -> Session:
+        """Get the current worker's database session.
+
+        Returns:
+            Session from the current worker
+
+        Raises:
+            RuntimeError: If called outside of a worker process
+        """
+        if cls.current is None:
+            raise RuntimeError("Worker.session called outside of worker context")
+
+        return cls.current._session
+
+    def run(self) -> None:
+        """Main worker loop - polls queue and processes tasks."""
+        from agentexec.core.queue import dequeue
+
+        logger.info(f"Worker {self._worker_id} starting")
+
+        try:
+            while not self._shutdown_event.is_set():
+                # Dequeue task (blocks with timeout to check shutdown event)
+                if (task := dequeue()) is not None:
+                    self._process_task(task)
+        except Exception as e:
+            logger.exception(f"Worker {self._worker_id} fatal error: {e}")
+            raise
+        finally:
+            self._cleanup()
+
+    def _process_task(self, task: Task) -> None:
+        """Process a single task from the queue.
+
+        Args:
+            task: Task to process
+        """
+        logger.info(f"Worker {self._worker_id} processing task: {task.task_name}")
+
+        try:
+            handler = self._handlers[task.task_name]
+            task.started()
+
+            if asyncio.iscoroutinefunction(handler):
+                self._loop.run_until_complete(handler(**task.handler_kwargs))
+            else:
+                handler(**task.handler_kwargs)
+
+            task.completed()
+            logger.info(f"Worker {self._worker_id} completed task: {task.task_name}")
+
+        except KeyError:
+            logger.error(f"No handler registered for task: {task.task_name}")
+            raise RuntimeError(f"No handler for task: {task.task_name}")
+        except Exception as e:
+            task.errored(e)
+            logger.exception(f"Worker {self._worker_id} error executing task {task.task_name}: {e}")
+
+    def _cleanup(self) -> None:
+        """Clean up worker resources on shutdown."""
+        self._session.close()
+        logger.debug(f"Worker {self._worker_id} closed database session")
+
+        self._loop.close()
+        logger.info(f"Worker {self._worker_id} shutting down")
+
+
+class WorkerPool:
+    """Manages a pool of worker processes for background task execution.
+
+    The WorkerPool coordinates multiple worker processes, handles task handler
+    registration, and manages graceful shutdown. Each worker runs in a separate
+    process with its own isolated state (session, event loop).
+
+    This allows multiple pools to coexist if needed, each managing their own
+    set of workers and handlers.
+
+    Example:
+        from sqlalchemy import create_engine
+
+        engine = create_engine("sqlite:///agents.db")
+        pool = WorkerPool(engine=engine)
+
+        @pool.task("research_company")
+        async def research_company(payload: dict, agent_id: str):
+            company_name = payload["company_name"]
+            # Task implementation
+            pass
+
+        pool.start()
+    """
+
+    _engine: Engine
+    _handlers: dict[str, TaskHandler]
+    _processes: list[mp.Process]
+    _shutdown_event: EventClass
+
+    def __init__(self, engine: Engine) -> None:
+        """Initialize the worker pool.
+
+        Args:
+            engine: SQLAlchemy engine for activity tracking.
+                   Each worker will create its own session from this engine.
+
+        Configuration is loaded from agentexec.CONF:
+        - num_workers: Number of worker processes to spawn
+        - queue_name: Name of the Redis list to use as task queue
+        - redis_url: Redis connection URL (via get_redis())
+        """
+        self._engine = engine
+        self._handlers: dict[str, TaskHandler] = {}
+        self._processes: list[mp.Process] = []
+        self._shutdown_event: EventClass = mp.Event()
+
+    def task(self, name: str) -> Callable[[TaskHandler], TaskHandler]:
+        """Decorator to register a task handler.
+
+        Args:
+            name: Task type name that will be used when enqueueing.
+
+        Returns:
+            Decorator function that registers the handler.
+
+        Example:
+            @pool.task("research_company")
+            async def research_company(payload: dict, agent_id: str):
+                company_name = payload["company_name"]
+                # Implementation
+                pass
+        """
+
+        def decorator(func: TaskHandler) -> TaskHandler:
+            self._handlers[name] = func
+            logger.info(f"Registered task handler: {name}")
+            return func
+
+        return decorator
+
+    def start(self) -> None:
+        """Start the worker processes.
+
+        Spawns N worker processes (configured via num_workers) that will poll
+        the Redis queue and execute registered task handlers.
+
+        Each worker runs independently with its own session and event loop.
+        """
+        logger.info(f"Starting {CONF.num_workers} worker processes")
+
+        for worker_id in range(CONF.num_workers):
+            process = mp.Process(
+                target=Worker.run_in_process,
+                args=(worker_id, self._engine, self._handlers, self._shutdown_event),
+                daemon=False,
+            )
+            process.start()
+            self._processes.append(process)
+            logger.info(f"Started worker process {worker_id} (PID: {process.pid})")
+
+    def shutdown(self, timeout: int | None = None) -> None:
+        """Gracefully shutdown all worker processes.
+
+        Args:
+            timeout: Maximum seconds to wait for workers to finish current task.
+                    Defaults to CONF.graceful_shutdown_timeout.
+        """
+        if timeout is None:
+            timeout = CONF.graceful_shutdown_timeout
+
+        logger.info("Initiating graceful shutdown of worker pool")
+        self._shutdown_event.set()
+
+        for process in self._processes:
+            process.join(timeout=timeout)
+            if process.is_alive():
+                logger.warning(f"Worker process {process.pid} did not stop gracefully, terminating")
+                process.terminate()
+                process.join(timeout=5)
+
+        self._processes.clear()
+        logger.info("Worker pool shutdown complete")

agentexec/runners/__init__.py

@@ -0,0 +1,13 @@
+"""Agent runners with activity tracking and lifecycle management."""
+
+from agentexec.runners.base import BaseAgentRunner
+
+__all__ = ["BaseAgentRunner"]
+
+# OpenAI runner is only available if agents package is installed
+try:
+    from agentexec.runners.openai import OpenAIRunner
+
+    __all__.append("OpenAIRunner")
+except ImportError:
+    pass

agentexec/runners/base.py

@@ -0,0 +1,135 @@
+from __future__ import annotations
+from typing import Any
+from abc import ABC
+import logging
+import uuid
+
+from agentexec import activity
+
+logger = logging.getLogger(__name__)
+
+
+class BaseAgentRunner(ABC):
+    """Abstract base class for agent runners with activity tracking.
+
+    This base class provides:
+    - Automatic activity tracking (QUEUED -> RUNNING -> COMPLETE/ERROR)
+    - Status update tool for agent self-reporting
+    - Common lifecycle management
+    - Error handling and recovery patterns
+
+    Subclasses must implement:
+    - _execute_agent(): The actual agent execution logic
+    - max_turns_exceptions: Tuple of exception classes for max turns errors
+    """
+
+    agent_id: uuid.UUID
+    max_turns_recovery: bool
+    recovery_turns: int
+
+    prompts: _RunnerPrompts
+    tools: _RunnerTools
+
+    def __init__(
+        self,
+        agent_id: uuid.UUID,
+        *,
+        max_turns_recovery: bool = True,
+        recovery_turns: int = 5,
+        wrap_up_prompt: str | None = None,
+        report_status_prompt: str | None = None,
+    ) -> None:
+        """Initialize the runner.
+
+        Args:
+            max_turns_recovery: Enable automatic recovery when max turns exceeded.
+            wrap_up_prompt: Prompt to use for recovery run.
+            recovery_turns: Number of turns allowed for recovery.
+            status_prompt: Instruction snippet about using the status tool.
+        """
+        self.agent_id = agent_id
+        self.max_turns_recovery = max_turns_recovery
+        self.recovery_turns = recovery_turns
+
+        # Tools namespace for accessing runner-provided tools
+        self.prompts = _RunnerPrompts(
+            report_status=report_status_prompt,
+            wrap_up=wrap_up_prompt,
+        )
+        self.tools = _RunnerTools(self.agent_id)
+
+
+class _RunnerPrompts:
+    """Namespace for runner-provided prompts.
+
+    Accessed via runner.prompts.*
+    """
+
+    report_status: str = (
+        "Using report_activity tool:\n"
+        "    - Always report your current activity before you start a new step using the report_activity tool. \n"
+        "    - Include a brief message about the task and context you are operating in (10 words or less). \n"
+        "    - Don't use internal data or underlying system info and instead focus on what the user would care about. \n"
+        "    - This informs the user of your progress as they have no visibility into your internal operations. \n"
+        "    - You can call multiple tools in parallel per step so don't waste an entire step on this.\n"
+        "    - Call it at the top of your list of the next round of tool uses; we should be careful to minimize turns used.\n"
+    )
+    wrap_up: str = "Please summarize your findings."
+
+    def __init__(
+        self,
+        *,
+        wrap_up: str | None = None,
+        report_status: str | None = None,
+    ) -> None:
+        if wrap_up is not None:
+            self.wrap_up = wrap_up
+        if report_status is not None:
+            self.report_status = report_status
+
+
+class _RunnerTools:
+    """Namespace for runner-provided tools.
+
+    Accessed via runner.tools.*
+    """
+
+    _agent_id: uuid.UUID
+
+    def __init__(self, agent_id: uuid.UUID) -> None:
+        self._agent_id = agent_id
+
+    @property
+    def report_status(self) -> Any:
+        """Get the status update tool.
+
+        This tool allows agents to report their progress back to the activity tracker.
+        The tool is bound to the current agent_id during runner.run().
+
+        Subclasses should override this to wrap with framework-specific decorators.
+
+        Returns:
+            Plain function for status updates.
+        """
+        agent_id = self._agent_id
+
+        def report_activity(message: str, percentage: int) -> str:
+            """Report progress and status updates.
+
+            Use this tool to report your progress as you work through the task.
+
+            Args:
+                message: A brief description of what you're currently doing
+                percentage: Your estimated completion percentage (0-100)
+
+            Returns:
+                Confirmation message
+            """
+            activity.update(
+                agent_id=agent_id,
+                message=message,
+                completion_percentage=percentage,
+            )
+            return "Status updated"
+
+        return report_activity

agentexec/runners/openai.py

@@ -0,0 +1,237 @@
+import logging
+import uuid
+from typing import Any, Callable
+
+from agents import Agent, MaxTurnsExceeded, Runner, function_tool
+from agents.items import TResponseInputItem
+from agents.result import RunResult, RunResultStreaming
+
+from agentexec.runners.base import BaseAgentRunner, _RunnerTools
+
+
+logger = logging.getLogger(__name__)
+
+
+def _extract_input(e: MaxTurnsExceeded) -> list[TResponseInputItem]:
+    """
+    Extract the full conversation input history from a `MaxTurnsExceeded` exception.
+
+    Args:
+        e: The MaxTurnsExceeded exception instance
+    Returns:
+        List of TResponseInputItem representing the full conversation history
+    """
+    if not e.run_data:
+        logger.error("No run data available in MaxTurnsExceeded exception")
+        raise
+
+    # Reconstruct the full conversation history
+    final_input: list[TResponseInputItem] = (
+        list(e.run_data.input)
+        if isinstance(e.run_data.input, list)
+        else [{"role": "user", "content": e.run_data.input}]
+    )
+
+    # Add all the conversation items that were generated
+    final_input.extend([item.to_input_item() for item in e.run_data.new_items])
+
+    return final_input
+
+
+class _OpenAIRunnerTools(_RunnerTools):
+    """OpenAI-specific tools wrapper that decorates with @function_tool."""
+
+    @property
+    def report_status(self) -> Any:
+        """Get the status update tool wrapped with @function_tool decorator."""
+        return function_tool(super().report_status)
+
+
+class OpenAIRunner(BaseAgentRunner):
+    """Runner for OpenAI Agents SDK with automatic activity tracking.
+
+    This runner wraps the OpenAI Agents SDK and provides:
+    - Automatic agent_id generation
+    - Activity lifecycle management (QUEUED -> RUNNING -> COMPLETE/ERROR)
+    - Max turns recovery with configurable wrap-up prompts
+    - Status update tool with agent_id pre-baked
+
+    Example:
+        runner = agentexec.OpenAIRunner(
+            max_turns_recovery=True,
+            wrap_up_prompt="Please summarize your findings.",
+            status_prompt="Use update_status(message, percentage) to report progress.",
+        )
+
+        agent = Agent(
+            name="Research Agent",
+            instructions=f"Research companies. {runner.status_prompt}",
+            tools=[runner.tools.report_status],
+            model="gpt-4o",
+        )
+
+        result = await runner.run(
+            session=session,
+            agent=agent,
+            input="Research Acme Corp",
+            agent_id=agent_id,  # Optional
+            max_turns=15,
+        )
+    """
+
+    def __init__(
+        self,
+        agent_id: uuid.UUID,
+        *,
+        max_turns_recovery: bool = False,
+        wrap_up_prompt: str | None = None,
+        recovery_turns: int = 5,
+        report_status_prompt: str | None = None,
+    ) -> None:
+        """Initialize the OpenAI runner.
+
+        Args:
+            agent_id: UUID for tracking this agent's activity.
+            max_turns_recovery: Enable automatic recovery when max turns exceeded.
+            wrap_up_prompt: Prompt to use for recovery run.
+            recovery_turns: Number of turns allowed for recovery.
+            report_status_prompt: Instruction snippet about using the status tool.
+        """
+        super().__init__(
+            agent_id,
+            max_turns_recovery=max_turns_recovery,
+            recovery_turns=recovery_turns,
+            wrap_up_prompt=wrap_up_prompt,
+            report_status_prompt=report_status_prompt,
+        )
+        # Override with OpenAI-specific tools
+        self.tools = _OpenAIRunnerTools(self.agent_id)
+
+    async def run(
+        self,
+        agent: Agent[Any],
+        input: str | list[TResponseInputItem],
+        max_turns: int = 10,
+        context: Any | None = None,
+    ) -> RunResult:
+        """Run the agent with automatic activity tracking.
+
+        Args:
+            session: SQLAlchemy session for database operations.
+            agent: Agent instance.
+            input: User input/prompt for the agent.
+            agent_id: Optional agent ID (generated if not provided).
+            max_turns: Maximum number of agent iterations.
+            agent_type: Optional agent type for activity tracking.
+            context: Optional context for the agent run.
+
+        Returns:
+            Result from the agent execution.
+        """
+        # TODO match method signature of Runner.run
+        try:
+            result = await Runner.run(
+                agent,
+                input,
+                max_turns=max_turns,
+                context=context,
+            )
+        except MaxTurnsExceeded as e:
+            if not self.max_turns_recovery:
+                raise
+
+            logger.info("Max turns exceeded, attempting recovery")
+            final_input = _extract_input(e)
+            final_input.append(
+                {
+                    "role": "user",
+                    "content": self.prompts.wrap_up,
+                }
+            )
+            result = await Runner.run(
+                agent,
+                final_input,
+                max_turns=self.recovery_turns,
+                context=context,
+            )
+        except Exception:
+            raise
+
+        return result
+
+    async def run_streamed(
+        self,
+        agent: Agent[Any],
+        input: str | list[TResponseInputItem],
+        max_turns: int = 10,
+        context: Any | None = None,
+        forwarder: Callable | None = None,
+    ) -> RunResultStreaming:
+        """Run the agent in streaming mode with automatic activity tracking.
+
+        The returned streaming result can be used just like the underlying framework's
+        streaming result. Activity tracking happens automatically.
+
+        Args:
+            session: SQLAlchemy session for database operations.
+            agent: Agent instance.
+            input: User input/prompt for the agent.
+            agent_id: Optional agent ID (generated if not provided).
+            max_turns: Maximum number of agent iterations.
+            agent_type: Optional agent type for activity tracking.
+            context: Optional context for the agent run.
+
+        Returns:
+            Streaming result from the agent execution.
+
+        Example:
+            result = await runner.run_streamed(session, agent, "Research XYZ", agent_id="123")
+            async for event in result.stream_events():
+                print(event)
+        """
+        # TODO match method signature of Runner.run_streamed
+        # TODO I want to defer the `await` to the caller side
+        # TODO forwarder is just a placeholder but we need to come up with a solution for that functionality
+        try:
+            result = Runner.run_streamed(
+                agent,
+                input,
+                max_turns=max_turns,
+                context=context,
+            )
+
+            async for event in result.stream_events():
+                if forwarder:
+                    await forwarder(event)
+                # yield event
+
+            return result
+        except MaxTurnsExceeded as e:
+            if not self.max_turns_recovery:
+                raise
+
+            logger.info("Max turns exceeded, attempting recovery")
+            final_input = _extract_input(e)
+            final_input.append(
+                {
+                    "role": "user",
+                    "content": self.prompts.wrap_up,
+                }
+            )
+            result = Runner.run_streamed(
+                agent,
+                final_input,
+                max_turns=self.recovery_turns,
+                context=context,
+            )
+
+            async for event in result.stream_events():
+                if forwarder:
+                    await forwarder(event)
+                # yield event
+
+            return result
+        except Exception:
+            raise
+
+        return result