agentpool 2.1.9__py3-none-any.whl

agentpool_storage/sql_provider/sql_provider.py

@@ -0,0 +1,444 @@
+"""SQLModel-based storage provider."""
+
+from __future__ import annotations
+
+from decimal import Decimal
+from typing import TYPE_CHECKING, Any, Self
+
+from pydantic_ai import RunUsage
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlmodel import SQLModel, desc, select
+
+from agentpool.log import get_logger
+from agentpool.messaging import TokenCost
+from agentpool.utils.now import get_now
+from agentpool.utils.parse_time import parse_time_period
+from agentpool_storage.base import StorageProvider
+from agentpool_storage.models import QueryFilters
+from agentpool_storage.sql_provider.models import CommandHistory, Conversation, Message
+from agentpool_storage.sql_provider.utils import (
+    build_message_query,
+    format_conversation,
+    parse_model_info,
+    to_chat_message,
+)
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+    from datetime import datetime
+    from types import TracebackType
+
+    from sqlalchemy import Connection
+
+    from agentpool.common_types import JsonValue
+    from agentpool.messaging import ChatMessage
+    from agentpool_config.session import SessionQuery
+    from agentpool_config.storage import SQLStorageConfig
+    from agentpool_storage.models import ConversationData, StatsFilters
+
+
+logger = get_logger(__name__)
+
+
+class SQLModelProvider(StorageProvider):
+    """Storage provider using SQLModel.
+
+    Can work with any database supported by SQLAlchemy/SQLModel.
+    Provides efficient SQL-based filtering and storage.
+    """
+
+    can_load_history = True
+
+    def __init__(self, config: SQLStorageConfig) -> None:
+        """Initialize provider with async database engine.
+
+        Args:
+            config: Configuration for provider
+        """
+        super().__init__(config)
+        self.engine = config.get_engine()
+        self.auto_migrate = config.auto_migration
+        self.session: AsyncSession | None = None
+
+    async def _init_database(self, auto_migrate: bool = True) -> None:
+        """Initialize database tables and optionally migrate columns.
+
+        Args:
+            auto_migrate: Whether to automatically add missing columns
+        """
+        from sqlmodel import SQLModel
+
+        from agentpool_storage.sql_provider.utils import auto_migrate_columns
+
+        # Create tables if they don't exist
+        async with self.engine.begin() as conn:
+            await conn.run_sync(SQLModel.metadata.create_all)
+
+        # Optionally add missing columns
+        if auto_migrate:
+            async with self.engine.begin() as conn:
+
+                def sync_migrate(sync_conn: Connection) -> None:
+                    auto_migrate_columns(sync_conn, self.engine.dialect)
+
+                await conn.run_sync(sync_migrate)
+
+    async def __aenter__(self) -> Self:
+        """Initialize async database resources."""
+        await self._init_database(auto_migrate=self.auto_migrate)
+        await super().__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Clean up async database resources properly."""
+        await self.engine.dispose()
+        return await super().__aexit__(exc_type, exc_val, exc_tb)
+
+    def cleanup(self) -> None:
+        """Clean up database resources."""
+        # For sync cleanup, just pass - proper cleanup happens in __aexit__
+
+    async def filter_messages(self, query: SessionQuery) -> list[ChatMessage[str]]:
+        """Filter messages using SQL queries."""
+        async with AsyncSession(self.engine) as session:
+            stmt = build_message_query(query)
+            result = await session.execute(stmt)
+            messages = result.scalars().all()
+            return [to_chat_message(msg) for msg in messages]
+
+    async def log_message(
+        self,
+        *,
+        conversation_id: str,
+        message_id: str,
+        content: str,
+        role: str,
+        name: str | None = None,
+        cost_info: TokenCost | None = None,
+        model: str | None = None,
+        response_time: float | None = None,
+        forwarded_from: list[str] | None = None,
+        provider_name: str | None = None,
+        provider_response_id: str | None = None,
+        messages: str | None = None,
+        finish_reason: str | None = None,
+    ) -> None:
+        """Log message to database."""
+        from agentpool_storage.sql_provider.models import Message
+
+        provider, model_name = parse_model_info(model)
+
+        async with AsyncSession(self.engine) as session:
+            msg = Message(
+                conversation_id=conversation_id,
+                id=message_id,
+                content=content,
+                role=role,
+                name=name,
+                model=model,
+                model_provider=provider,
+                model_name=model_name,
+                response_time=response_time,
+                total_tokens=cost_info.token_usage.total_tokens if cost_info else None,
+                input_tokens=cost_info.token_usage.input_tokens if cost_info else None,
+                output_tokens=cost_info.token_usage.output_tokens if cost_info else None,
+                cost=float(cost_info.total_cost) if cost_info else None,
+                forwarded_from=forwarded_from,
+                provider_name=provider_name,
+                provider_response_id=provider_response_id,
+                messages=messages,
+                finish_reason=finish_reason,
+                timestamp=get_now(),
+            )
+            session.add(msg)
+            await session.commit()
+
+    async def log_conversation(
+        self,
+        *,
+        conversation_id: str,
+        node_name: str,
+        start_time: datetime | None = None,
+    ) -> None:
+        """Log conversation to database."""
+        from agentpool_storage.sql_provider.models import Conversation
+
+        async with AsyncSession(self.engine) as session:
+            now = start_time or get_now()
+            convo = Conversation(id=conversation_id, agent_name=node_name, start_time=now)
+            session.add(convo)
+            await session.commit()
+
+    async def update_conversation_title(
+        self,
+        conversation_id: str,
+        title: str,
+    ) -> None:
+        """Update the title of a conversation."""
+        async with AsyncSession(self.engine) as session:
+            result = await session.execute(
+                select(Conversation).where(Conversation.id == conversation_id)
+            )
+            conversation = result.scalar_one_or_none()
+            if conversation:
+                conversation.title = title
+                session.add(conversation)
+                await session.commit()
+
+    async def get_conversation_title(
+        self,
+        conversation_id: str,
+    ) -> str | None:
+        """Get the title of a conversation."""
+        async with AsyncSession(self.engine) as session:
+            result = await session.execute(
+                select(Conversation.title).where(Conversation.id == conversation_id)
+            )
+            return result.scalar_one_or_none()
+
+    async def log_command(
+        self,
+        *,
+        agent_name: str,
+        session_id: str,
+        command: str,
+        context_type: type | None = None,
+        metadata: dict[str, JsonValue] | None = None,
+    ) -> None:
+        """Log command to database."""
+        async with AsyncSession(self.engine) as session:
+            history = CommandHistory(
+                session_id=session_id,
+                agent_name=agent_name,
+                command=command,
+                context_type=context_type.__name__ if context_type else None,
+                context_metadata=metadata or {},
+            )
+            session.add(history)
+            await session.commit()
+
+    async def get_filtered_conversations(
+        self,
+        agent_name: str | None = None,
+        period: str | None = None,
+        since: datetime | None = None,
+        query: str | None = None,
+        model: str | None = None,
+        limit: int | None = None,
+        *,
+        compact: bool = False,
+        include_tokens: bool = False,
+    ) -> list[ConversationData]:
+        """Get filtered conversations with formatted output."""
+        # Convert period to since if provided
+        if period:
+            since = get_now() - parse_time_period(period)
+
+        # Create filters
+        filters = QueryFilters(
+            agent_name=agent_name,
+            since=since,
+            query=query,
+            model=model,
+            limit=limit,
+        )
+
+        # Use existing get_conversations method
+        conversations = await self.get_conversations(filters)
+        return [
+            format_conversation(conv, msgs, compact=compact, include_tokens=include_tokens)
+            for conv, msgs in conversations
+        ]
+
+    async def get_commands(
+        self,
+        agent_name: str,
+        session_id: str,
+        *,
+        limit: int | None = None,
+        current_session_only: bool = False,
+    ) -> list[str]:
+        """Get command history from database."""
+        async with AsyncSession(self.engine) as session:
+            query = select(CommandHistory)
+            if current_session_only:
+                query = query.where(CommandHistory.session_id == str(session_id))
+            else:
+                query = query.where(CommandHistory.agent_name == agent_name)
+
+            query = query.order_by(desc(CommandHistory.timestamp))
+            if limit:
+                query = query.limit(limit)
+
+            result = await session.execute(query)
+            return [h.command for h in result.scalars()]
+
+    async def get_conversations(
+        self,
+        filters: QueryFilters,
+    ) -> list[tuple[ConversationData, Sequence[ChatMessage[str]]]]:
+        """Get filtered conversations using SQL queries."""
+        async with AsyncSession(self.engine) as session:
+            results: list[tuple[ConversationData, Sequence[ChatMessage[str]]]] = []
+
+            # Base conversation query
+            conv_query = select(Conversation)
+
+            if filters.agent_name:
+                conv_query = conv_query.where(Conversation.agent_name == filters.agent_name)
+
+            # Apply time filters if provided
+            if filters.since:
+                conv_query = conv_query.where(Conversation.start_time >= filters.since)
+
+            if filters.limit:
+                conv_query = conv_query.limit(filters.limit)
+
+            conv_result = await session.execute(conv_query)
+            conversations = conv_result.scalars().all()
+
+            for conv in conversations:
+                # Get messages for this conversation
+                msg_query = select(Message).where(Message.conversation_id == conv.id)
+
+                if filters.query:
+                    msg_query = msg_query.where(Message.content.contains(filters.query))  # type: ignore
+                if filters.model:
+                    msg_query = msg_query.where(Message.model_name == filters.model)
+
+                msg_query = msg_query.order_by(Message.timestamp.asc())  # type: ignore
+                msg_result = await session.execute(msg_query)
+                messages = msg_result.scalars().all()
+
+                if not messages:
+                    continue
+
+                chat_messages = [to_chat_message(msg) for msg in messages]
+                conv_data = format_conversation(conv, messages)
+                results.append((conv_data, chat_messages))
+
+            return results
+
+    async def get_conversation_stats(
+        self,
+        filters: StatsFilters,
+    ) -> dict[str, dict[str, Any]]:
+        """Get statistics using SQL aggregations."""
+        from agentpool_storage.sql_provider.models import Conversation, Message
+
+        async with AsyncSession(self.engine) as session:
+            # Base query for stats
+            query = (
+                select(  # type: ignore[call-overload]
+                    Message.model,
+                    Conversation.agent_name,
+                    Message.timestamp,
+                    Message.total_tokens,
+                    Message.input_tokens,
+                    Message.output_tokens,
+                )
+                .join(Conversation, Message.conversation_id == Conversation.id)
+                .where(Message.timestamp > filters.cutoff)
+            )
+
+            if filters.agent_name:
+                query = query.where(Conversation.agent_name == filters.agent_name)
+
+            # Execute query and get raw data
+            result = await session.execute(query)
+            rows = [
+                (
+                    model,
+                    agent,
+                    timestamp,
+                    TokenCost(
+                        token_usage=RunUsage(
+                            input_tokens=prompt or 0,
+                            output_tokens=completion or 0,
+                        ),
+                        total_cost=Decimal(0),  # We don't store this in DB
+                    )
+                    if total or prompt or completion
+                    else None,
+                )
+                for model, agent, timestamp, total, prompt, completion in result.all()
+            ]
+
+        # Use base class aggregation
+        return self.aggregate_stats(rows, filters.group_by)
+
+    async def reset(
+        self,
+        *,
+        agent_name: str | None = None,
+        hard: bool = False,
+    ) -> tuple[int, int]:
+        """Reset database storage."""
+        from sqlalchemy import text
+
+        from agentpool_storage.sql_provider.queries import (
+            DELETE_AGENT_CONVERSATIONS,
+            DELETE_AGENT_MESSAGES,
+            DELETE_ALL_CONVERSATIONS,
+            DELETE_ALL_MESSAGES,
+        )
+
+        async with AsyncSession(self.engine) as session:
+            if hard:
+                if agent_name:
+                    msg = "Hard reset cannot be used with agent_name"
+                    raise ValueError(msg)
+                # Drop and recreate all tables
+                async with self.engine.begin() as conn:
+                    await conn.run_sync(SQLModel.metadata.drop_all)
+                await session.commit()
+                # Recreate schema
+                await self._init_database()
+                return 0, 0
+
+            # Get counts first
+            conv_count, msg_count = await self.get_conversation_counts(agent_name=agent_name)
+
+            # Delete data
+            if agent_name:
+                await session.execute(text(DELETE_AGENT_MESSAGES), {"agent": agent_name})
+                await session.execute(text(DELETE_AGENT_CONVERSATIONS), {"agent": agent_name})
+            else:
+                await session.execute(text(DELETE_ALL_MESSAGES))
+                await session.execute(text(DELETE_ALL_CONVERSATIONS))
+
+            await session.commit()
+        return conv_count, msg_count
+
+    async def get_conversation_counts(
+        self,
+        *,
+        agent_name: str | None = None,
+    ) -> tuple[int, int]:
+        """Get conversation and message counts."""
+        from agentpool_storage.sql_provider import Conversation, Message
+
+        if not self.session:
+            msg = "Session not initialized. Use provider as async context manager."
+            raise RuntimeError(msg)
+        if agent_name:
+            conv_query = select(Conversation).where(Conversation.agent_name == agent_name)
+            msg_query = (
+                select(Message).join(Conversation).where(Conversation.agent_name == agent_name)
+            )
+        else:
+            conv_query = select(Conversation)
+            msg_query = select(Message)
+
+        conv_result = await self.session.execute(conv_query)
+        msg_result = await self.session.execute(msg_query)
+        conv_count = len(conv_result.scalars().all())
+        msg_count = len(msg_result.scalars().all())
+
+        return conv_count, msg_count

agentpool_storage/sql_provider/utils.py

@@ -0,0 +1,234 @@
+"""Utilities for database storage."""
+
+from __future__ import annotations
+
+from datetime import datetime
+from decimal import Decimal
+from typing import TYPE_CHECKING, Any
+
+from pydantic_ai import RunUsage
+from sqlalchemy import JSON, Column, and_, or_
+from sqlalchemy.sql import expression
+from sqlmodel import select
+
+from agentpool.messaging import ChatMessage, TokenCost
+from agentpool.storage import deserialize_messages
+from agentpool_storage.models import ConversationData, MessageData
+from agentpool_storage.sql_provider.models import Conversation
+
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from sqlmodel.sql.expression import SelectOfScalar
+    from tokonomics.toko_types import TokenUsage
+
+    from agentpool_config.session import SessionQuery
+    from agentpool_storage.sql_provider.models import Message
+
+
+def aggregate_token_usage(
+    messages: Sequence[Message | ChatMessage[str]],
+) -> TokenUsage:
+    """Sum up tokens from a sequence of messages."""
+    from agentpool_storage.sql_provider.models import Message
+
+    total = prompt = completion = 0
+    for msg in messages:
+        if isinstance(msg, Message):
+            total += msg.total_tokens or 0
+            prompt += msg.input_tokens or 0
+            completion += msg.output_tokens or 0
+        elif msg.cost_info:
+            total += msg.cost_info.token_usage.total_tokens
+            prompt += msg.cost_info.token_usage.input_tokens
+            completion += msg.cost_info.token_usage.output_tokens
+    return {"total": total, "prompt": prompt, "completion": completion}
+
+
+def to_chat_message(db_message: Message) -> ChatMessage[str]:
+    """Convert database message to ChatMessage."""
+    cost_info = None
+    if db_message.total_tokens is not None:
+        cost_info = TokenCost(
+            token_usage=RunUsage(
+                input_tokens=db_message.input_tokens or 0,
+                output_tokens=db_message.output_tokens or 0,
+            ),
+            total_cost=Decimal(db_message.cost or 0.0),
+        )
+
+    return ChatMessage[str](
+        message_id=db_message.id,
+        conversation_id=db_message.conversation_id,
+        content=db_message.content,
+        role=db_message.role,  # type: ignore
+        name=db_message.name,
+        model_name=db_message.model,
+        cost_info=cost_info,
+        response_time=db_message.response_time,
+        forwarded_from=db_message.forwarded_from or [],
+        timestamp=db_message.timestamp,
+        provider_name=db_message.provider_name,
+        provider_response_id=db_message.provider_response_id,
+        messages=deserialize_messages(db_message.messages),
+        finish_reason=db_message.finish_reason,  # type: ignore
+    )
+
+
+def get_column_default(column: Any) -> str:
+    """Get SQL DEFAULT clause for column."""
+    if column.default is None:
+        return ""
+    if hasattr(column.default, "arg"):
+        # Simple default value
+        return f" DEFAULT {column.default.arg}"
+    if hasattr(column.default, "sqltext"):
+        # Computed default
+        return f" DEFAULT {column.default.sqltext}"
+    return ""
+
+
+def auto_migrate_columns(sync_conn: Any, dialect: Any) -> None:
+    """Automatically add missing columns to existing tables.
+
+    Args:
+        sync_conn: Synchronous database connection
+        dialect: SQLAlchemy dialect for SQL type compilation
+    """
+    from sqlalchemy import inspect
+    from sqlalchemy.sql import text
+    from sqlmodel import SQLModel
+
+    inspector = inspect(sync_conn)
+
+    # For each table in our models
+    for table_name, table in SQLModel.metadata.tables.items():
+        existing = {col["name"] for col in inspector.get_columns(table_name)}
+
+        # For each column in model that doesn't exist in DB
+        for col in table.columns:
+            if col.name not in existing:
+                # Create ALTER TABLE statement based on column type
+                type_sql = col.type.compile(dialect)
+                nullable = "" if col.nullable else " NOT NULL"
+                default = get_column_default(col)
+                sql = (
+                    f"ALTER TABLE {table_name} ADD COLUMN {col.name} {type_sql}{nullable}{default}"
+                )
+                sync_conn.execute(text(sql))
+
+
+def parse_model_info(model: str | None) -> tuple[str | None, str | None]:
+    """Parse model string into provider and name.
+
+    Args:
+        model: Full model string (e.g., "openai:gpt-5", "anthropic:claude-sonnet-4-0")
+
+    Returns:
+        Tuple of (provider, name)
+    """
+    if not model:
+        return None, None
+
+    # Try splitting by ':' or '/'
+    parts = model.split(":") if ":" in model else model.split("/")
+
+    if len(parts) == 2:  # noqa: PLR2004
+        provider, name = parts
+        return provider.lower(), name
+
+    # No provider specified, try to infer
+    name = parts[0]
+    if name.startswith(("gpt-", "text-", "dall-e")):
+        return "openai", name
+    if name.startswith("claude"):
+        return "anthropic", name
+    if name.startswith(("llama", "mistral")):
+        return "meta", name
+
+    return None, name
+
+
+def build_message_query(query: SessionQuery) -> SelectOfScalar[Any]:
+    """Build SQLModel query from SessionQuery."""
+    from agentpool_storage.sql_provider.models import Message
+
+    stmt = select(Message).order_by(Message.timestamp)  # type: ignore
+
+    conditions: list[Any] = []
+    if query.name:
+        conditions.append(Message.conversation_id == query.name)
+    if query.agents:
+        agent_conditions = [Column("name").in_(query.agents)]
+        if query.include_forwarded:
+            agent_conditions.append(
+                and_(
+                    Column("forwarded_from").isnot(None),
+                    expression.cast(Column("forwarded_from"), JSON).contains(list(query.agents)),  # type: ignore
+                )
+            )
+        conditions.append(or_(*agent_conditions))
+    if query.since and (cutoff := query.get_time_cutoff()):
+        conditions.append(Message.timestamp >= cutoff)
+    if query.until:
+        conditions.append(Message.timestamp <= datetime.fromisoformat(query.until))
+    if query.contains:
+        conditions.append(Message.content.contains(query.contains))  # type: ignore
+    if query.roles:
+        conditions.append(Message.role.in_(query.roles))  # type: ignore
+
+    if conditions:
+        stmt = stmt.where(and_(*conditions))
+    if query.limit:
+        stmt = stmt.limit(query.limit)
+
+    return stmt
+
+
+def format_conversation(
+    conv: Conversation | ConversationData,
+    messages: Sequence[Message | ChatMessage[str]],
+    *,
+    include_tokens: bool = False,
+    compact: bool = False,
+) -> ConversationData:
+    """Format SQL conversation model to ConversationData."""
+    msgs = list(messages)
+    if compact and len(msgs) > 1:
+        msgs = [msgs[0], msgs[-1]]
+
+    # Convert messages to ChatMessage format if needed
+    chat_messages = [msg if isinstance(msg, ChatMessage) else to_chat_message(msg) for msg in msgs]
+
+    # Convert Conversation to ConversationData format
+    if isinstance(conv, Conversation):
+        return ConversationData(
+            id=conv.id,
+            agent=conv.agent_name,
+            title=conv.title,
+            start_time=conv.start_time.isoformat(),
+            messages=[
+                MessageData(
+                    role=msg.role,
+                    content=msg.content,
+                    timestamp=msg.timestamp.isoformat(),
+                    model=msg.model_name,
+                    name=msg.name,
+                    token_usage={
+                        "prompt": msg.usage.input_tokens,
+                        "completion": msg.usage.output_tokens,
+                        "total": msg.usage.total_tokens,
+                    },
+                    cost=float(msg.cost_info.total_cost) if msg.cost_info else None,
+                    response_time=msg.response_time,
+                )
+                for msg in chat_messages
+            ],
+            token_usage=aggregate_token_usage(msgs) if include_tokens else None,
+        )
+
+    # If it's already ConversationData, update token_usage if needed
+    if include_tokens and conv["token_usage"] is None:
+        conv["token_usage"] = aggregate_token_usage(msgs)
+    return conv