agentexec 0.1.0__py3-none-any.whl

agentexec/activity/tracker.py

@@ -0,0 +1,267 @@
+import uuid
+
+from sqlalchemy.orm import Session
+
+from agentexec.activity.models import Activity, ActivityLog, Status
+from agentexec.activity.schemas import (
+    ActivityDetailSchema,
+    ActivityListItemSchema,
+    ActivityListSchema,
+)
+
+
+def generate_agent_id() -> uuid.UUID:
+    """Generate a new UUID for an agent.
+
+    This is the centralized function for generating agent IDs.
+    Users can override this if they need custom ID generation logic.
+
+    Returns:
+        A new UUID4 object
+    """
+    return uuid.uuid4()
+
+
+def normalize_agent_id(agent_id: str | uuid.UUID) -> uuid.UUID:
+    """Normalize agent_id to UUID object.
+
+    Args:
+        agent_id: Either a string UUID or UUID object
+
+    Returns:
+        UUID object
+
+    Raises:
+        ValueError: If string is not a valid UUID
+    """
+    if isinstance(agent_id, str):
+        return uuid.UUID(agent_id)
+    return agent_id
+
+
+def create(
+    task_name: str,
+    message: str = "Agent queued",
+    agent_id: str | uuid.UUID | None = None,
+    session: Session | None = None,
+) -> uuid.UUID:
+    """Create a new agent activity record with initial queued status.
+
+    Args:
+        task_name: Name/type of the task (e.g., "research", "analysis")
+        initial_message: Initial log message (default: "Agent queued")
+        agent_id: Optional custom agent ID (string or UUID). If not provided, one will be auto-generated.
+
+    Returns:
+        The agent_id (as UUID object) of the created record
+    """
+    agent_id = normalize_agent_id(agent_id) if agent_id else generate_agent_id()
+
+    if session is None:
+        from agentexec.core.worker import get_worker_session
+        db = get_worker_session()
+    else:
+        db = session
+
+    activity_record = Activity(
+        agent_id=agent_id,
+        agent_type=task_name,
+    )
+    db.add(activity_record)
+    db.flush()
+
+    log = ActivityLog(
+        activity_id=activity_record.id,
+        message=message,
+        status=Status.QUEUED,
+        completion_percentage=0,
+    )
+    db.add(log)
+    db.commit()
+
+    return agent_id
+
+
+def update(
+    agent_id: str | uuid.UUID,
+    message: str,
+    completion_percentage: int | None = None,
+    status: Status | None = None,
+    session: Session | None = None,
+) -> bool:
+    """Update an agent's activity by adding a new log message.
+
+    This function will set the status to RUNNING unless a different status is explicitly provided.
+
+    Args:
+        agent_id: The agent_id of the agent to update
+        message: Log message to append
+        completion_percentage: Optional completion percentage (0-100)
+        status: Optional status to set (default: RUNNING)
+        session: Optional SQLAlchemy session. If not provided, uses global session factory.
+
+    Returns:
+        True if successful
+
+    Raises:
+        ValueError: If agent_id not found
+    """
+    if session is None:
+        from agentexec.core.worker import get_worker_session
+        db = get_worker_session()
+    else:
+        db = session
+
+    Activity.append_log(
+        session=db,
+        agent_id=normalize_agent_id(agent_id),
+        message=message,
+        status=status if status else Status.RUNNING,
+        completion_percentage=completion_percentage,
+    )
+    return True
+
+
+def complete(
+    agent_id: str | uuid.UUID,
+    message: str = "Agent completed",
+    completion_percentage: int = 100,
+    session: Session | None = None,
+) -> bool:
+    """Mark an agent activity as complete.
+
+    Args:
+        agent_id: The agent_id of the agent to mark as complete
+        message: Log message (default: "Agent completed")
+        completion_percentage: Completion percentage (default: 100)
+        session: Optional SQLAlchemy session. If not provided, uses global session factory.
+
+    Returns:
+        True if successful
+
+    Raises:
+        ValueError: If agent_id not found
+    """
+    if session is None:
+        from agentexec.core.worker import get_worker_session
+        db = get_worker_session()
+    else:
+        db = session
+
+    Activity.append_log(
+        session=db,
+        agent_id=normalize_agent_id(agent_id),
+        message=message,
+        status=Status.COMPLETE,
+        completion_percentage=completion_percentage,
+    )
+    return True
+
+
+def error(
+    agent_id: str | uuid.UUID,
+    message: str = "Agent failed",
+    completion_percentage: int = 100,
+    session: Session | None = None,
+) -> bool:
+    """Mark an agent activity as failed.
+
+    Args:
+        agent_id: The agent_id of the agent to mark as failed
+        message: Log message (default: "Agent failed")
+        completion_percentage: Completion percentage (default: 100)
+        session: Optional SQLAlchemy session. If not provided, uses global session factory.
+
+    Returns:
+        True if successful
+
+    Raises:
+        ValueError: If agent_id not found
+    """
+    if session is None:
+        from agentexec.core.worker import get_worker_session
+        db = get_worker_session()
+    else:
+        db = session
+
+    Activity.append_log(
+        session=db,
+        agent_id=normalize_agent_id(agent_id),
+        message=message,
+        status=Status.ERROR,
+        completion_percentage=completion_percentage,
+    )
+    return True
+
+
+def cancel_pending(
+    session: Session | None = None,
+) -> int:
+    """Mark all queued and running agents as canceled.
+
+    Useful during application shutdown to clean up pending tasks.
+
+    Returns:
+        Number of agents that were canceled
+    """
+    if session is None:
+        from agentexec.core.worker import get_worker_session
+        db = get_worker_session()
+    else:
+        db = session
+
+    pending_agent_ids = Activity.get_pending_ids(db)
+    for agent_id in pending_agent_ids:
+        Activity.append_log(
+            session=db,
+            agent_id=agent_id,
+            message="Canceled due to shutdown",
+            status=Status.CANCELED,
+            completion_percentage=None,
+        )
+
+    return len(pending_agent_ids)
+
+
+def list(
+    session: Session,
+    page: int = 1,
+    page_size: int = 50,
+) -> ActivityListSchema:
+    """List activities with pagination.
+
+    Args:
+        session: SQLAlchemy session to use for the query
+        page: Page number (1-indexed)
+        page_size: Number of items per page
+
+    Returns:
+        ActivityList with list of ActivityListItemSchema items
+    """
+    total = session.query(Activity).count()
+    rows = Activity.get_list(session, page=page, page_size=page_size)
+
+    return ActivityListSchema(
+        items=[ActivityListItemSchema.model_validate(row) for row in rows],
+        total=total,
+        page=page,
+        page_size=page_size,
+    )
+
+
+def detail(
+    session: Session,
+    agent_id: str | uuid.UUID,
+) -> ActivityDetailSchema | None:
+    """Get a single activity by agent_id with all logs.
+
+    Args:
+        session: SQLAlchemy session to use for the query
+        agent_id: The agent_id to look up
+
+    Returns:
+        ActivityDetailSchema with full log history, or None if not found
+    """
+    if item := Activity.get_by_agent_id(session, agent_id):
+        return ActivityDetailSchema.model_validate(item)
+    return None

agentexec/config.py

@@ -0,0 +1,72 @@
+from pydantic import AliasChoices, Field
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+
+class Config(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        case_sensitive=False,
+    )
+
+    debug: bool = Field(
+        default=False,
+        description="Enable debug logging",
+        validation_alias=AliasChoices("AGENTEXEC_DEBUG", "DEBUG"),
+    )
+
+    table_prefix: str = Field(
+        default="agentexec_",
+        description="Prefix for database table names",
+        validation_alias="AGENTEXEC_TABLE_PREFIX",
+    )
+    queue_name: str = Field(
+        default="agentexec_tasks",
+        description="Name of the Redis list to use as task queue",
+        validation_alias="AGENTEXEC_QUEUE_NAME",
+    )
+    num_workers: int = Field(
+        default=4,
+        description="Number of worker processes to spawn",
+        validation_alias="AGENTEXEC_NUM_WORKERS",
+    )
+    graceful_shutdown_timeout: int = Field(
+        default=300,
+        description="Maximum seconds to wait for workers to finish on shutdown",
+        validation_alias="AGENTEXEC_GRACEFUL_SHUTDOWN_TIMEOUT",
+    )
+
+    activity_message_create: str = Field(
+        default="Waiting to start.",
+        description="Default message when creating a new agent activity",
+        validation_alias="AGENTEXEC_ACTIVITY_MESSAGE_CREATE",
+    )
+    activity_message_complete: str = Field(
+        default="Completed successfully.",
+        description="Default message when an agent activity completes successfully",
+        validation_alias="AGENTEXEC_ACTIVITY_MESSAGE_COMPLETE",
+    )
+    activity_message_error: str = Field(
+        default="An error occurred during execution.",
+        description="Default message when an agent activity encounters an error",
+        validation_alias="AGENTEXEC_ACTIVITY_MESSAGE_ERROR",
+    )
+
+    redis_url: str = Field(
+        default="redis://localhost:6379/0",
+        description="Redis connection URL",
+        validation_alias=AliasChoices("AGENTEXEC_REDIS_URL", "REDIS_URL"),
+    )
+    redis_pool_size: int = Field(
+        default=10,
+        description="Redis connection pool size",
+        validation_alias=AliasChoices("AGENTEXEC_REDIS_POOL_SIZE", "REDIS_POOL_SIZE"),
+    )
+    redis_pool_timeout: int = Field(
+        default=5,
+        description="Redis connection pool timeout in seconds",
+        validation_alias=AliasChoices("AGENTEXEC_REDIS_POOL_TIMEOUT", "REDIS_POOL_TIMEOUT"),
+    )
+
+
+CONF = Config()

agentexec/core/models.py

@@ -0,0 +1,23 @@
+"""Database configuration for agent-runner.
+
+This module provides the declarative base class for all SQLAlchemy models.
+"""
+
+from sqlalchemy.orm import DeclarativeBase
+
+
+class Base(DeclarativeBase):
+    """Base class for all SQLAlchemy models in agent-runner.
+
+    Users can reference this base in their Alembic migrations to include
+    agent-runner's tables alongside their own application tables.
+
+    Example:
+        # In alembic/env.py
+        import agentexec as ax
+        from models import Base
+
+        target_metadata = [Base.metadata, ax.Base.metadata]
+    """
+
+    pass

agentexec/core/queue.py

@@ -0,0 +1,109 @@
+import logging
+from enum import Enum
+from typing import Any, cast
+
+from agentexec.config import CONF
+from agentexec.core.redis_client import get_redis
+from agentexec.core.task import Task
+
+logger = logging.getLogger(__name__)
+
+
+class Priority(str, Enum):
+    """Task priority levels.
+
+    HIGH: Push to front of queue (processed first).
+    LOW: Push to back of queue (processed later).
+    """
+
+    HIGH = "high"
+    LOW = "low"
+
+
+def enqueue(
+    task_name: str,
+    payload: dict[str, Any],
+    *,
+    priority: Priority = Priority.LOW,
+    queue_name: str | None = None,
+) -> Task:
+    """Enqueue a task for background execution with automatic activity tracking.
+
+    This function automatically creates an activity record for the task and
+    enqueues it to Redis for worker processing.
+
+    Args:
+        task: Task to enqueue.
+        db: Database session for activity tracking.
+        priority: Task priority (Priority.HIGH or Priority.LOW).
+        queue_name: Redis list name to use. If not provided, uses
+                   agentexec.CONF.queue_name.
+
+    Returns:
+        agent_id: UUID for tracking the task's progress.
+
+    Example:
+        task = Task(task_type="research", payload={"company": "Acme"})
+        agent_id = enqueue(task, db=session, priority=Priority.HIGH)
+        # Track progress: ax.activity.get_activity(db, agent_id)
+    """
+
+    try:
+        redis = get_redis()
+        redis_push = {
+            Priority.HIGH: redis.rpush,
+            Priority.LOW: redis.lpush,
+        }[priority]
+    except KeyError:
+        raise ValueError(f"Invalid priority: {priority}. Must be Priority.HIGH or Priority.LOW")
+
+    task = Task.create(
+        task_name=task_name,
+        payload=payload,
+    )
+    redis_push(
+        queue_name or CONF.queue_name,
+        task.model_dump_json(),
+    )
+    logger.info(f"Enqueued task {task.task_name} with agent_id {task.agent_id}")
+    # TODO track errors
+
+    return task
+
+
+def dequeue(queue_name: str | None = None, timeout: int = 1) -> Task | None:
+    """Dequeue a task from the queue.
+
+    Blocks for up to timeout seconds waiting for a task. Returns None if
+    no task is available within the timeout period.
+
+    Args:
+        queue_name: Redis list name to use. If not provided, uses
+                   agentexec.CONF.queue_name.
+        timeout: Maximum seconds to wait for a task.
+
+    Returns:
+        Task if available, None otherwise.
+
+    Example:
+        task = dequeue(timeout=5)
+        if task is not None:
+            # Process task
+            pass
+    """
+    redis = get_redis()
+
+    result = cast(
+        tuple[str, str] | None,
+        redis.brpop(
+            [queue_name or CONF.queue_name],
+            timeout=timeout,
+        ),
+    )
+
+    if result is None:
+        # No task available within the timeout
+        return None
+
+    _, task_json = result
+    return Task.model_validate_json(task_json)

agentexec/core/redis_client.py

@@ -0,0 +1,40 @@
+from redis import Redis
+from agentexec.config import CONF
+
+_redis_client: Redis | None = None
+
+
+def get_redis() -> Redis:
+    """Get the Redis client instance.
+
+    Creates and caches a Redis client on first call. Subsequent calls
+    return the cached instance.
+
+    The client is configured from agentexec.CONF settings:
+    - redis_url: Connection URL
+    - redis_pool_size: Connection pool size
+    - redis_pool_timeout: Connection timeout
+
+    Returns:
+        Redis client instance.
+    """
+    global _redis_client
+
+    if _redis_client is None:
+        _redis_client = Redis.from_url(
+            CONF.redis_url,
+            max_connections=CONF.redis_pool_size,
+            socket_connect_timeout=CONF.redis_pool_timeout,
+            decode_responses=True,  # Automatically decode bytes to strings
+        )
+
+    return _redis_client
+
+
+def close_redis() -> None:
+    """Close the Redis connection and reset the client."""
+    global _redis_client
+
+    if _redis_client is not None:
+        _redis_client.close()
+        _redis_client = None

agentexec/core/task.py

@@ -0,0 +1,132 @@
+from __future__ import annotations
+from typing import Any, Protocol, TypedDict, Unpack
+from uuid import UUID
+from pydantic import BaseModel, Field
+
+from agentexec import activity
+
+
+class TaskHandlerKwargs(TypedDict):
+    """Type for kwargs passed to task handlers.
+
+    Handlers receive the payload as a dict and agent_id as a separate parameter.
+    """
+
+    agent_id: UUID
+    payload: dict[str, Any]
+
+
+class TaskHandler(Protocol):
+    """Protocol for task handler functions.
+
+    Handlers accept **kwargs matching HandlerKwargs structure.
+    Return value is ignored. Can be sync or async.
+    """
+
+    def __call__(self, **kwargs: Unpack[TaskHandlerKwargs]) -> None: ...
+
+
+class Task(BaseModel):
+    """Represents a background task.
+
+    Tasks are serialized to JSON and enqueued to Redis for workers to process.
+    Each task has a type (matching a registered handler), a payload with
+    parameters, and an agent_id for tracking.
+    """
+
+    task_name: str
+    payload: dict[str, Any] = Field(default_factory=dict)
+    agent_id: UUID
+
+    @classmethod
+    def create(cls, task_name: str, payload: dict[str, Any]) -> Task:
+        """Create a new task with automatic activity tracking.
+
+        This is a convenience method that creates both a Task instance and
+        its corresponding activity record in one step.
+
+        Args:
+            task_name: Name/type of the task (e.g., "research", "analysis")
+            payload: Task parameters to pass to the handler
+
+        Returns:
+            Task instance with agent_id set
+
+        Example:
+            task = Task.create("research_company", {"company": "Acme Corp"})
+            # Activity record is created automatically
+        """
+        agent_id = activity.create(
+            task_name=task_name,
+            message="Waiting to start.",
+        )
+        return cls(
+            task_name=task_name,
+            payload=payload,
+            agent_id=agent_id,
+        )
+
+    def started(self) -> None:
+        """Mark the task as started in the activity log.
+
+        Updates the activity status to RUNNING with a starting message.
+
+        Example:
+            task = Task.create("research", {"company": "Acme"})
+            task.started()
+        """
+        activity.update(
+            agent_id=self.agent_id,
+            message="Task started.",
+            completion_percentage=0,
+        )
+
+    def completed(self) -> None:
+        """Mark the task as completed in the activity log.
+
+        Updates the activity status to COMPLETE with a completion message.
+
+        Example:
+            task = Task.create("research", {"company": "Acme"})
+            task.completed()
+        """
+        activity.update(
+            agent_id=self.agent_id,
+            message="Task completed successfully.",
+            completion_percentage=100,
+            status=activity.Status.COMPLETE,
+        )
+
+    def errored(self, exception: Exception) -> None:
+        """Mark the task as errored in the activity log.
+
+        Updates the activity status to ERROR with the provided error message.
+
+        Args:
+            error_message: Description of the error that occurred
+
+        Example:
+            task = Task.create("research", {"company": "Acme"})
+            task.error("Failed to fetch data from API.")
+        """
+        activity.update(
+            agent_id=self.agent_id,
+            message=f"Task failed with error: {exception}",
+            status=activity.Status.ERROR,
+        )
+
+    @property
+    def handler_kwargs(self) -> TaskHandlerKwargs:
+        """Get kwargs to pass to the task handler.
+
+        Builds a dictionary containing the task's payload and agent_id,
+        matching the TaskHandlerKwargs structure expected by handlers.
+
+        Returns:
+            Dict with 'payload' (task parameters) and 'agent_id' (tracking ID)
+
+        Example:
+            kwargs = task._handler_kwargs
+            # Returns: {"payload": {"company": "Acme"}, "agent_id": UUID(...)}
+        """
+        return {"agent_id": self.agent_id, "payload": self.payload}