agentpool 2.1.9__py3-none-any.whl

agentpool_storage/session_store.py

@@ -0,0 +1,262 @@
+"""SQL session store implementation."""
+
+from __future__ import annotations
+
+from datetime import timedelta
+from typing import TYPE_CHECKING, Any, Self
+
+from sqlalchemy import delete, select
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlmodel import SQLModel
+
+from agentpool.log import get_logger
+from agentpool.sessions.models import SessionData
+from agentpool.utils.now import get_now
+from agentpool_storage.sql_provider.models import Session
+
+
+if TYPE_CHECKING:
+    from types import TracebackType
+
+    from sqlalchemy.ext.asyncio import AsyncEngine
+
+    from agentpool_config.storage import SQLStorageConfig
+
+logger = get_logger(__name__)
+
+
+class SQLSessionStore:
+    """SQL-based session store using SQLModel.
+
+    Persists session data to a SQL database, supporting:
+    - SQLite (default)
+    - PostgreSQL
+    - MySQL
+    """
+
+    def __init__(self, config: SQLStorageConfig) -> None:
+        """Initialize SQL session store.
+
+        Args:
+            config: SQL storage configuration with database URL
+        """
+        self._config = config
+        self._engine: AsyncEngine | None = None
+
+    async def __aenter__(self) -> Self:
+        """Initialize database connection and create tables."""
+        self._engine = self._config.get_engine()
+
+        async with self._engine.begin() as conn:
+            await conn.run_sync(SQLModel.metadata.create_all)
+
+        # Auto-migrate missing columns
+        await self._migrate_columns()
+
+        logger.debug("SQL session store initialized", url=self._config.url)
+        return self
+
+    async def _migrate_columns(self) -> None:
+        """Add missing columns to existing tables."""
+        from agentpool_storage.sql_provider.utils import auto_migrate_columns
+
+        if self._engine is None:
+            return
+
+        async with self._engine.begin() as conn:
+
+            def sync_migrate(sync_conn: Any) -> None:
+                auto_migrate_columns(sync_conn, self._engine.dialect)  # type: ignore[union-attr]
+
+            await conn.run_sync(sync_migrate)
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Close database connection."""
+        if self._engine:
+            await self._engine.dispose()
+            self._engine = None
+
+    def _get_engine(self) -> AsyncEngine:
+        """Get engine, raising if not initialized."""
+        if self._engine is None:
+            msg = "Session store not initialized. Use async context manager."
+            raise RuntimeError(msg)
+        return self._engine
+
+    def _to_db_model(self, data: SessionData) -> Session:
+        """Convert SessionData to database model."""
+        return Session(
+            session_id=data.session_id,
+            agent_name=data.agent_name,
+            conversation_id=data.conversation_id,
+            pool_id=data.pool_id,
+            title=data.title,
+            cwd=data.cwd,
+            created_at=data.created_at,
+            last_active=data.last_active,
+            metadata_json=data.metadata,
+        )
+
+    def _from_db_model(self, row: Session) -> SessionData:
+        """Convert database model to SessionData."""
+        return SessionData(
+            session_id=row.session_id,
+            agent_name=row.agent_name,
+            conversation_id=row.conversation_id,
+            pool_id=row.pool_id,
+            title=row.title,
+            cwd=row.cwd,
+            created_at=row.created_at,
+            last_active=row.last_active,
+            metadata=row.metadata_json or {},
+        )
+
+    async def save(self, data: SessionData) -> None:
+        """Save or update session data.
+
+        Uses delete-then-insert for upsert semantics.
+
+        Args:
+            data: Session data to persist
+        """
+        engine = self._get_engine()
+
+        async with AsyncSession(engine) as session:
+            # Delete existing if present (upsert via delete+insert)
+            stmt = delete(Session).where(Session.session_id == data.session_id)  # type: ignore[arg-type]
+            await session.execute(stmt)
+
+            # Insert new/updated
+            db_session = self._to_db_model(data)
+            session.add(db_session)
+            await session.commit()
+            logger.debug("Saved session", session_id=data.session_id)
+
+    async def load(self, session_id: str) -> SessionData | None:
+        """Load session data by ID.
+
+        Args:
+            session_id: Session identifier
+
+        Returns:
+            Session data if found, None otherwise
+        """
+        engine = self._get_engine()
+
+        async with AsyncSession(engine) as session:
+            stmt = select(Session).where(Session.session_id == session_id)  # type: ignore[arg-type]
+            result = await session.execute(stmt)
+            row = result.scalars().first()
+
+            if row is None:
+                return None
+
+            return self._from_db_model(row)
+
+    async def delete(self, session_id: str) -> bool:
+        """Delete a session.
+
+        Args:
+            session_id: Session identifier
+
+        Returns:
+            True if session was deleted, False if not found
+        """
+        engine = self._get_engine()
+
+        async with AsyncSession(engine) as session:
+            stmt = delete(Session).where(Session.session_id == session_id)  # type: ignore[arg-type]
+            result = await session.execute(stmt)
+            await session.commit()
+
+            deleted: bool = result.rowcount > 0  # type: ignore[attr-defined]
+            if deleted:
+                logger.debug("Deleted session", session_id=session_id)
+            return deleted
+
+    async def list_sessions(
+        self,
+        pool_id: str | None = None,
+        agent_name: str | None = None,
+    ) -> list[str]:
+        """List session IDs, optionally filtered.
+
+        Args:
+            pool_id: Filter by pool/manifest ID
+            agent_name: Filter by agent name
+
+        Returns:
+            List of session IDs
+        """
+        engine = self._get_engine()
+
+        async with AsyncSession(engine) as session:
+            stmt = select(Session.session_id)  # type: ignore[call-overload]
+
+            if pool_id is not None:
+                stmt = stmt.where(Session.pool_id == pool_id)
+            if agent_name is not None:
+                stmt = stmt.where(Session.agent_name == agent_name)
+
+            stmt = stmt.order_by(Session.last_active.desc())  # type: ignore[attr-defined]
+            result = await session.execute(stmt)
+            # When selecting a single column, scalars() gives us the values directly
+            return list(result.scalars().all())
+
+    async def cleanup_expired(self, max_age_hours: int = 24) -> int:
+        """Remove sessions older than max_age.
+
+        Args:
+            max_age_hours: Maximum session age in hours
+
+        Returns:
+            Number of sessions removed
+        """
+        engine = self._get_engine()
+        cutoff = get_now() - timedelta(hours=max_age_hours)
+
+        async with AsyncSession(engine) as session:
+            stmt = delete(Session).where(Session.last_active < cutoff)  # type: ignore[arg-type]
+            result = await session.execute(stmt)
+            await session.commit()
+
+            count = result.rowcount or 0  # type: ignore[attr-defined]
+            if count > 0:
+                logger.info("Cleaned up expired sessions", count=count)
+            return count
+
+    async def get_all(
+        self,
+        pool_id: str | None = None,
+        limit: int | None = None,
+    ) -> list[SessionData]:
+        """Get all sessions, optionally filtered.
+
+        Args:
+            pool_id: Filter by pool/manifest ID
+            limit: Maximum number of sessions to return
+
+        Returns:
+            List of session data objects
+        """
+        engine = self._get_engine()
+
+        async with AsyncSession(engine) as session:
+            stmt = select(Session)
+
+            if pool_id is not None:
+                stmt = stmt.where(Session.pool_id == pool_id)  # type: ignore[arg-type]
+
+            stmt = stmt.order_by(Session.last_active.desc())  # type: ignore[attr-defined]
+
+            if limit is not None:
+                stmt = stmt.limit(limit)
+
+            result = await session.execute(stmt)
+            # scalars() gives us actual Session model instances
+            return [self._from_db_model(row) for row in result.scalars().all()]

agentpool_storage/sql_provider/__init__.py

@@ -0,0 +1,21 @@
+"""SQL storage provider package."""
+
+from __future__ import annotations
+
+from agentpool_storage.sql_provider.sql_provider import SQLModelProvider
+from agentpool_storage.sql_provider.models import (
+    Conversation,
+    Message,
+    CommandHistory,
+    MessageLog,
+    ConversationLog,
+)
+
+__all__ = [
+    "CommandHistory",
+    "Conversation",
+    "ConversationLog",
+    "Message",
+    "MessageLog",
+    "SQLModelProvider",
+]

agentpool_storage/sql_provider/cli.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+"""Database migration management script for agentpool.
+
+This script provides easy commands for managing database migrations using Alembic.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+import subprocess
+from typing import Annotated
+
+from rich.console import Console
+from rich.prompt import Confirm
+import typer
+
+
+app = typer.Typer(
+    name="agentpool-db",
+    help="Database migration management for agentpool",
+    rich_markup_mode="rich",
+)
+console = Console()
+
+
+def run_command(cmd: list[str], cwd: Path | None = None) -> int:
+    """Run a command and return the exit code."""
+    console.print(f"[bold blue]Running:[/bold blue] {' '.join(cmd)}")
+    result = subprocess.run(cmd, cwd=cwd, check=False)
+    if result.returncode != 0 and cmd[0] == "alembic":
+        console.print(
+            f"[bold red]Error:[/bold red] Command {cmd[0]!r} not found. "
+            "Make sure alembic is installed."
+        )
+    return result.returncode
+
+
+def get_project_root() -> Path:
+    """Get the project root directory."""
+    # Start from this file and go up to find the project root
+    # This file is in src/agentpool_storage/sql_provider/cli.py
+    return Path(__file__).parent.parent.parent.parent
+
+
+@app.command()
+def status() -> int:
+    """Show current migration status."""
+    project_root = get_project_root()
+    return run_command(["alembic", "current", "-v"], cwd=project_root)
+
+
+@app.command()
+def upgrade(
+    revision: Annotated[str, typer.Argument(help="Target revision")] = "head",
+) -> int:
+    """Upgrade database to latest migration."""
+    project_root = get_project_root()
+    return run_command(["alembic", "upgrade", revision], cwd=project_root)
+
+
+@app.command()
+def downgrade(
+    revision: Annotated[
+        str,
+        typer.Argument(help="Target revision (e.g., -1 for previous, or specific revision ID)"),
+    ],
+) -> int:
+    """Downgrade database."""
+    project_root = get_project_root()
+    return run_command(["alembic", "downgrade", revision], cwd=project_root)
+
+
+@app.command()
+def create(
+    message: Annotated[str, typer.Argument(help="Migration message")],
+    autogenerate: Annotated[
+        bool,
+        typer.Option(
+            "--autogenerate",
+            help="Auto-generate migration based on model changes",
+        ),
+    ] = False,
+) -> int:
+    """Create a new migration."""
+    project_root = get_project_root()
+    cmd = ["alembic", "revision", "-m", message]
+    if autogenerate:
+        cmd.append("--autogenerate")
+    return run_command(cmd, cwd=project_root)
+
+
+@app.command()
+def history() -> int:
+    """Show migration history."""
+    project_root = get_project_root()
+    return run_command(["alembic", "history", "-v"], cwd=project_root)
+
+
+@app.command()
+def current() -> int:
+    """Show current revision."""
+    project_root = get_project_root()
+    return run_command(["alembic", "current"], cwd=project_root)
+
+
+@app.command()
+def reset(
+    force: Annotated[
+        bool,
+        typer.Option("--force", help="Force reset without confirmation"),
+    ] = False,
+) -> int:
+    """Reset database (WARNING: destructive)."""
+    if not force and not Confirm.ask(
+        "[bold red]This will delete all data in the database. Continue?[/bold red]",
+        default=False,
+    ):
+        console.print("[yellow]Aborted.[/yellow]")
+        return 0
+
+    project_root = get_project_root()
+    console.print("[bold yellow]Resetting database...[/bold yellow]")
+
+    # First, drop to base (empty database)
+    result = run_command(["alembic", "downgrade", "base"], cwd=project_root)
+    if result != 0:
+        return result
+
+    # Then upgrade to head
+    return run_command(["alembic", "upgrade", "head"], cwd=project_root)
+
+
+def main() -> None:
+    """Main entry point."""
+    # Set up environment
+    project_root = get_project_root()
+    env = os.environ.copy()
+    env["PYTHONPATH"] = str(project_root / "src")
+    os.environ.update(env)
+
+    app()
+
+
+if __name__ == "__main__":
+    main()

agentpool_storage/sql_provider/models.py

@@ -0,0 +1,249 @@
+"""Database models for agentpool storage."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from typing import TYPE_CHECKING, Any
+from uuid import uuid4
+
+from pydantic import ConfigDict
+from schemez import Schema
+from sqlalchemy import Column, DateTime, Text
+from sqlalchemy.ext.asyncio import AsyncAttrs
+from sqlalchemy.types import TypeDecorator
+from sqlmodel import JSON, Field, SQLModel
+from sqlmodel.main import SQLModelConfig  # type: ignore[attr-defined]
+
+from agentpool.common_types import JsonValue
+from agentpool.utils.now import get_now
+
+
+if TYPE_CHECKING:
+    from sqlalchemy import Dialect
+
+
+class UTCDateTime(TypeDecorator[datetime]):
+    """Stores DateTime as UTC."""
+
+    impl = DateTime
+    cache_ok = True
+
+    def process_bind_param(
+        self,
+        value: datetime | None,
+        dialect: Dialect,
+    ) -> datetime | None:
+        if value is not None:
+            value = value.replace(tzinfo=UTC) if value.tzinfo is None else value.astimezone(UTC)
+        return value
+
+    def process_result_value(
+        self,
+        value: datetime | None,
+        dialect: Dialect,
+    ) -> datetime | None:
+        if value is not None:
+            value = value.replace(tzinfo=UTC)
+        return value
+
+
+class CommandHistory(AsyncAttrs, SQLModel, table=True):
+    """Database model for command history."""
+
+    id: int = Field(default=None, primary_key=True)
+    """Primary key for command history entry"""
+
+    session_id: str = Field(index=True)
+    """ID of the chat session"""
+
+    agent_name: str = Field(index=True)
+    """Name of the agent that executed the command"""
+
+    command: str
+    """The command that was executed"""
+
+    context_type: str | None = Field(default=None, index=True)
+    """Type of the command context (e.g. 'AgentContext', 'PoolSupervisor', etc.)"""
+
+    context_metadata: dict[str, JsonValue] = Field(default_factory=dict, sa_column=Column(JSON))
+    """Additional context information about command execution"""
+
+    timestamp: datetime = Field(
+        sa_column=Column(UTCDateTime, default=get_now), default_factory=get_now
+    )
+    """When the command was executed"""
+
+    model_config = SQLModelConfig(use_attribute_docstrings=True)  # pyright: ignore[reportCallIssue]
+
+
+class MessageLog(Schema):
+    """Raw message log entry."""
+
+    timestamp: datetime
+    """When the message was sent"""
+
+    role: str
+    """Role of the message sender (user/assistant/system)"""
+
+    content: str
+    """Content of the message"""
+
+    token_usage: dict[str, int] | None = None
+    """Token usage statistics as provided by model"""
+
+    cost: float | None = None
+    """Cost of generating this message in USD"""
+
+    model: str | None = None
+    """Name of the model that generated this message"""
+
+    model_config = ConfigDict(frozen=True)
+
+
+class ConversationLog(Schema):
+    """Collection of messages forming a conversation."""
+
+    id: str
+    """Unique identifier for the conversation"""
+
+    agent_name: str
+    """Name of the agent handling the conversation"""
+
+    start_time: datetime
+    """When the conversation started"""
+
+    messages: list[MessageLog]
+    """List of messages in the conversation"""
+
+    model_config = ConfigDict(frozen=True)
+
+
+class Message(AsyncAttrs, SQLModel, table=True):
+    """Database model for message logs."""
+
+    id: str = Field(default_factory=lambda: str(uuid4()), primary_key=True)
+    """Unique identifier for the message"""
+
+    conversation_id: str = Field(index=True)
+    """ID of the conversation this message belongs to"""
+
+    timestamp: datetime = Field(
+        sa_column=Column(UTCDateTime, default=get_now), default_factory=get_now
+    )
+    """When the message was sent"""
+
+    role: str
+    """Role of the message sender (user/assistant/system)"""
+
+    name: str | None = Field(default=None, index=True)
+    """Display name of the sender"""
+
+    content: str
+    """Content of the message"""
+
+    model: str | None = None
+    """Full model identifier (including provider)"""
+
+    model_name: str | None = Field(default=None, index=True)
+    """Name of the model (e.g., "gpt-5")"""
+
+    model_provider: str | None = Field(default=None, index=True)
+    """Provider of the model (e.g., "openai")"""
+
+    forwarded_from: list[str] | None = Field(default=None, sa_column=Column(JSON))
+    """List of agent names that forwarded this message"""
+
+    total_tokens: int | None = Field(default=None, index=True)
+    """Total number of tokens used"""
+
+    input_tokens: int | None = None
+    """Number of tokens in the prompt"""
+
+    output_tokens: int | None = None
+    """Number of tokens in the completion"""
+
+    cost: float | None = Field(default=None, index=True)
+    """Cost of generating this message in USD"""
+
+    response_time: float | None = None
+    """Time taken to generate the response in seconds"""
+
+    provider_name: str | None = Field(default=None, index=True)
+    """Name of the LLM provider that generated the response"""
+
+    provider_response_id: str | None = None
+    """Request ID as specified by the model provider"""
+
+    messages: str | None = None
+    """Serialized pydantic-ai model messages"""
+
+    finish_reason: str | None = None  # Literal pydantic_ai.FinishReason
+    """Reason the model finished generating the response"""
+
+    checkpoint_data: dict[str, Any] | None = Field(default=None, sa_column=Column(JSON))
+    """A dictionary of checkpoints (name -> metadata)."""
+
+    model_config = SQLModelConfig(use_attribute_docstrings=True)  # pyright: ignore[reportCallIssue]
+
+
+class Session(AsyncAttrs, SQLModel, table=True):
+    """Database model for session persistence."""
+
+    session_id: str = Field(primary_key=True)
+    """Unique session identifier."""
+
+    agent_name: str = Field(index=True)
+    """Name of the currently active agent."""
+
+    conversation_id: str = Field(index=True)
+    """Links to conversation in message storage."""
+
+    pool_id: str | None = Field(default=None, index=True)
+    """Optional pool/manifest identifier for multi-pool setups."""
+
+    title: str | None = Field(default=None, index=True)
+    """AI-generated or user-provided title for the conversation."""
+
+    cwd: str | None = Field(default=None, sa_column=Column(Text))
+    """Working directory for the session."""
+
+    created_at: datetime = Field(
+        sa_column=Column(UTCDateTime, index=True),
+        default_factory=get_now,
+    )
+    """When the session was created."""
+
+    last_active: datetime = Field(
+        sa_column=Column(UTCDateTime, index=True),
+        default_factory=get_now,
+    )
+    """Last activity timestamp."""
+
+    metadata_json: dict[str, Any] = Field(default_factory=dict, sa_column=Column(JSON))
+    """Protocol-specific or custom metadata stored as JSON."""
+
+    model_config = SQLModelConfig(use_attribute_docstrings=True)  # pyright: ignore[reportCallIssue]
+
+
+class Conversation(AsyncAttrs, SQLModel, table=True):
+    """Database model for conversations."""
+
+    id: str = Field(primary_key=True)
+    """Unique identifier for the conversation"""
+
+    agent_name: str = Field(index=True)
+    """Name of the agent handling the conversation"""
+
+    title: str | None = Field(default=None, index=True)
+    """Generated title for the conversation"""
+
+    start_time: datetime = Field(sa_column=Column(UTCDateTime, index=True), default_factory=get_now)
+    """When the conversation started"""
+
+    total_tokens: int = 0
+    """Total number of tokens used in this conversation"""
+
+    total_cost: float = 0.0
+    """Total cost of this conversation in USD"""
+
+    model_config = SQLModelConfig(use_attribute_docstrings=True)  # pyright: ignore[reportCallIssue]

agentpool_storage/sql_provider/queries.py

@@ -0,0 +1,15 @@
+"""Some pre-defined Queries."""
+
+from __future__ import annotations
+
+
+DELETE_AGENT_MESSAGES = """\
+DELETE FROM message
+WHERE conversation_id IN (
+    SELECT id FROM conversation WHERE agent_name = :agent
+)
+"""
+DELETE_AGENT_CONVERSATIONS = "DELETE FROM conversation WHERE agent_name = :agent"
+
+DELETE_ALL_MESSAGES = "DELETE FROM message"
+DELETE_ALL_CONVERSATIONS = "DELETE FROM conversation"