intentkit 0.5.0__py3-none-any.whl → 0.5.2__py3-none-any.whl

intentkit/models/db.py

@@ -0,0 +1,120 @@
+from contextlib import asynccontextmanager
+from typing import Annotated, AsyncGenerator, Optional
+from urllib.parse import quote_plus
+
+from intentkit.models.db_mig import safe_migrate
+from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.types import Checkpointer
+from psycopg_pool import AsyncConnectionPool
+from pydantic import Field
+from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, create_async_engine
+
+engine = None
+_langgraph_checkpointer: Optional[Checkpointer] = None
+
+
+async def init_db(
+    host: Optional[str],
+    username: Optional[str],
+    password: Optional[str],
+    dbname: Optional[str],
+    port: Annotated[Optional[str], Field(default="5432", description="Database port")],
+    auto_migrate: Annotated[
+        bool, Field(default=True, description="Whether to run migrations automatically")
+    ],
+) -> None:
+    """Initialize the database and handle schema updates.
+
+    Args:
+        host: Database host
+        username: Database username
+        password: Database password
+        dbname: Database name
+        port: Database port (default: 5432)
+        auto_migrate: Whether to run migrations automatically (default: True)
+    """
+    global engine, _langgraph_checkpointer
+    # Initialize psycopg pool and AsyncPostgresSaver if not already initialized
+    if _langgraph_checkpointer is None:
+        if host:
+            pool = AsyncConnectionPool(
+                conninfo=f"postgresql://{username}:{quote_plus(password)}@{host}:{port}/{dbname}",
+                min_size=3,
+                max_size=20,
+                timeout=60,
+                max_idle=30 * 60,
+            )
+            _langgraph_checkpointer = AsyncPostgresSaver(pool)
+            if auto_migrate:
+                await _langgraph_checkpointer.setup()
+        else:
+            _langgraph_checkpointer = InMemorySaver()
+    # Initialize SQLAlchemy engine with pool settings
+    if engine is None:
+        if host:
+            engine = create_async_engine(
+                f"postgresql+asyncpg://{username}:{quote_plus(password)}@{host}:{port}/{dbname}",
+                pool_size=20,  # Increase pool size
+                max_overflow=30,  # Increase max overflow
+                pool_timeout=60,  # Increase timeout
+                pool_pre_ping=True,  # Enable connection health checks
+                pool_recycle=3600,  # Recycle connections after 1 hour
+            )
+        else:
+            engine = create_async_engine(
+                "sqlite+aiosqlite:///:memory:",
+                connect_args={"check_same_thread": False},
+            )
+        if auto_migrate:
+            await safe_migrate(engine)
+
+
+async def get_db() -> AsyncGenerator[AsyncSession, None]:
+    async with AsyncSession(engine) as session:
+        yield session
+
+
+@asynccontextmanager
+async def get_session() -> AsyncSession:
+    """Get a database session using an async context manager.
+
+    This function is designed to be used with the 'async with' statement,
+    ensuring proper session cleanup.
+
+    Returns:
+        AsyncSession: A SQLAlchemy async session that will be automatically closed
+
+    Example:
+        ```python
+        async with get_session() as session:
+            # use session here
+            session.query(...)
+        # session is automatically closed
+        ```
+    """
+    session = AsyncSession(engine)
+    try:
+        yield session
+    finally:
+        await session.close()
+
+
+def get_engine() -> AsyncEngine:
+    """Get the SQLAlchemy async engine.
+
+    Returns:
+        AsyncEngine: The SQLAlchemy async engine
+    """
+    return engine
+
+
+def get_langgraph_checkpointer() -> Checkpointer:
+    """Get the AsyncPostgresSaver instance for langgraph.
+
+    Returns:
+        AsyncPostgresSaver: The AsyncPostgresSaver instance
+    """
+    if _langgraph_checkpointer is None:
+        raise RuntimeError("Database pool not initialized. Call init_db first.")
+    return _langgraph_checkpointer

intentkit/models/db_mig.py

@@ -0,0 +1,102 @@
+"""Database migration utilities."""
+
+import logging
+from typing import Callable
+
+from intentkit.models.base import Base
+from sqlalchemy import Column, MetaData, inspect, text
+
+logger = logging.getLogger(__name__)
+
+
+async def add_column_if_not_exists(
+    conn, dialect, table_name: str, column: Column
+) -> None:
+    """Add a column to a table if it doesn't exist.
+
+    Args:
+        conn: SQLAlchemy conn
+        table_name: Name of the table
+        column: Column to add
+    """
+
+    # Use run_sync to perform inspection on the connection
+    def _get_columns(connection):
+        inspector = inspect(connection)
+        return [c["name"] for c in inspector.get_columns(table_name)]
+
+    columns = await conn.run_sync(_get_columns)
+
+    if column.name not in columns:
+        # Build column definition
+        column_def = f"{column.name} {column.type.compile(dialect)}"
+
+        # Add DEFAULT if specified
+        if column.default is not None:
+            if hasattr(column.default, "arg"):
+                default_value = column.default.arg
+                if not isinstance(default_value, Callable):
+                    if isinstance(default_value, bool):
+                        default_value = str(default_value).lower()
+                    elif isinstance(default_value, str):
+                        default_value = f"'{default_value}'"
+                    elif isinstance(default_value, (list, dict)):
+                        default_value = "'{}'"
+                    column_def += f" DEFAULT {default_value}"
+
+        # Execute ALTER TABLE
+        await conn.execute(text(f"ALTER TABLE {table_name} ADD COLUMN {column_def}"))
+        logger.info(f"Added column {column.name} to table {table_name}")
+
+
+async def update_table_schema(conn, dialect, model_cls) -> None:
+    """Update table schema by adding missing columns from the model.
+
+    Args:
+        conn: SQLAlchemy conn
+        dialect: SQLAlchemy dialect
+        model_cls: SQLAlchemy model class to check for new columns
+    """
+    if not hasattr(model_cls, "__table__"):
+        return
+
+    table_name = model_cls.__tablename__
+    for name, column in model_cls.__table__.columns.items():
+        if name != "id":  # Skip primary key
+            await add_column_if_not_exists(conn, dialect, table_name, column)
+
+
+async def safe_migrate(engine) -> None:
+    """Safely migrate all SQLAlchemy models by adding new columns.
+
+    Args:
+        engine: SQLAlchemy engine
+    """
+    logger.info("Starting database schema migration")
+    dialect = engine.dialect
+
+    async with engine.begin() as conn:
+        try:
+            # Create tables if they don't exist
+            await conn.run_sync(Base.metadata.create_all)
+
+            # Get existing table metadata
+            metadata = MetaData()
+            await conn.run_sync(metadata.reflect)
+
+            # Update schema for all model classes
+            for mapper in Base.registry.mappers:
+                model_cls = mapper.class_
+                if hasattr(model_cls, "__tablename__"):
+                    table_name = model_cls.__tablename__
+                    if table_name in metadata.tables:
+                        # We need a sync wrapper for the async update_table_schema
+                        async def update_table_wrapper():
+                            await update_table_schema(conn, dialect, model_cls)
+
+                        await update_table_wrapper()
+        except Exception as e:
+            logger.error(f"Error updating database schema: {str(e)}")
+            raise
+
+    logger.info("Database schema updated successfully")

intentkit/models/generator.py

@@ -0,0 +1,347 @@
+"""Agent Generation Log Model.
+
+This module defines the database models for logging agent generation operations,
+including token usage, prompts, AI responses, and generation metadata.
+"""
+
+from datetime import datetime, timezone
+from typing import Annotated, Optional
+
+from epyxid import XID
+from intentkit.models.base import Base
+from pydantic import BaseModel, ConfigDict, Field
+from sqlalchemy import (
+    Boolean,
+    Column,
+    DateTime,
+    Integer,
+    String,
+    Text,
+    func,
+    select,
+)
+from sqlalchemy.dialects.postgresql import JSON, JSONB
+from sqlalchemy.ext.asyncio import AsyncSession
+
+
+class AgentGenerationLogTable(Base):
+    """Agent generation log database table model."""
+
+    __tablename__ = "agent_generation_logs"
+
+    id = Column(
+        String,
+        primary_key=True,
+    )
+    user_id = Column(
+        String,
+        nullable=True,
+    )
+    prompt = Column(
+        Text,
+        nullable=False,
+    )
+    existing_agent_id = Column(
+        String,
+        nullable=True,
+    )
+    is_update = Column(
+        Boolean,
+        default=False,
+        nullable=False,
+    )
+    generated_agent_schema = Column(
+        JSON().with_variant(JSONB(), "postgresql"),
+        nullable=True,
+    )
+    identified_skills = Column(
+        JSON().with_variant(JSONB(), "postgresql"),
+        nullable=True,
+    )
+    # LLM API response data
+    llm_model = Column(
+        String,
+        nullable=True,
+    )
+    total_tokens = Column(
+        Integer,
+        default=0,
+    )
+    input_tokens = Column(
+        Integer,
+        default=0,
+    )
+    cached_input_tokens = Column(
+        Integer,
+        default=0,
+    )
+    output_tokens = Column(
+        Integer,
+        default=0,
+    )
+    input_tokens_details = Column(
+        JSON().with_variant(JSONB(), "postgresql"),
+        nullable=True,
+    )
+    completion_tokens_details = Column(
+        JSON().with_variant(JSONB(), "postgresql"),
+        nullable=True,
+    )
+    # Performance metrics
+    generation_time_ms = Column(
+        Integer,
+        nullable=True,
+    )
+    retry_count = Column(
+        Integer,
+        default=0,
+    )
+    validation_errors = Column(
+        JSON().with_variant(JSONB(), "postgresql"),
+        nullable=True,
+    )
+    # Status and results
+    success = Column(
+        Boolean,
+        default=False,
+        nullable=False,
+    )
+    error_message = Column(
+        Text,
+        nullable=True,
+    )
+    # Timestamps
+    created_at = Column(
+        DateTime(timezone=True),
+        nullable=False,
+        server_default=func.now(),
+    )
+    completed_at = Column(
+        DateTime(timezone=True),
+        nullable=True,
+    )
+
+
+class AgentGenerationLogCreate(BaseModel):
+    """Model for creating agent generation log entries."""
+
+    model_config = ConfigDict(
+        use_enum_values=True,
+        from_attributes=True,
+    )
+
+    id: Annotated[
+        str,
+        Field(
+            default_factory=lambda: str(XID()),
+            description="Unique identifier for the generation log",
+        ),
+    ]
+    user_id: Optional[str] = Field(
+        None,
+        description="User ID who initiated the generation",
+    )
+    prompt: str = Field(
+        ...,
+        description="The original prompt used for generation",
+    )
+    existing_agent_id: Optional[str] = Field(
+        None,
+        description="ID of existing agent if this is an update operation",
+    )
+    is_update: bool = Field(
+        False,
+        description="Whether this is an update to existing agent",
+    )
+
+
+class AgentGenerationLog(BaseModel):
+    """Agent generation log model."""
+
+    model_config = ConfigDict(
+        use_enum_values=True,
+        from_attributes=True,
+    )
+
+    id: str
+    user_id: Optional[str] = None
+    prompt: str
+    existing_agent_id: Optional[str] = None
+    is_update: bool = False
+    generated_agent_schema: Optional[dict] = None
+    identified_skills: Optional[dict] = None
+    llm_model: Optional[str] = None
+    total_tokens: int = 0
+    input_tokens: int = 0
+    cached_input_tokens: int = 0
+    output_tokens: int = 0
+    input_tokens_details: Optional[dict] = None
+    completion_tokens_details: Optional[dict] = None
+    generation_time_ms: Optional[int] = None
+    retry_count: int = 0
+    validation_errors: Optional[dict] = None
+    success: bool = False
+    error_message: Optional[str] = None
+    created_at: datetime
+    completed_at: Optional[datetime] = None
+
+    @classmethod
+    async def create(
+        cls,
+        session: AsyncSession,
+        log_data: AgentGenerationLogCreate,
+    ) -> "AgentGenerationLog":
+        """Create a new agent generation log entry.
+
+        Args:
+            session: Database session
+            log_data: Log data to create
+
+        Returns:
+            Created log instance
+        """
+        # Create database record
+        log_record = AgentGenerationLogTable(
+            id=log_data.id,
+            user_id=log_data.user_id,
+            prompt=log_data.prompt,
+            existing_agent_id=log_data.existing_agent_id,
+            is_update=log_data.is_update,
+        )
+
+        session.add(log_record)
+        await session.commit()
+        await session.refresh(log_record)
+
+        return cls.model_validate(log_record)
+
+    async def update_completion(
+        self,
+        session: AsyncSession,
+        generated_agent_schema: Optional[dict] = None,
+        identified_skills: Optional[dict] = None,
+        llm_model: Optional[str] = None,
+        total_tokens: int = 0,
+        input_tokens: int = 0,
+        cached_input_tokens: int = 0,
+        output_tokens: int = 0,
+        input_tokens_details: Optional[dict] = None,
+        completion_tokens_details: Optional[dict] = None,
+        generation_time_ms: Optional[int] = None,
+        retry_count: int = 0,
+        validation_errors: Optional[dict] = None,
+        success: bool = False,
+        error_message: Optional[str] = None,
+    ) -> None:
+        """Update the log entry with completion data.
+
+        Args:
+            session: Database session
+            generated_agent_schema: The generated agent schema
+            identified_skills: Skills identified during generation
+            llm_model: LLM model used
+            total_tokens: Total tokens used
+            input_tokens: Input tokens used
+            cached_input_tokens: Cached input tokens used (for cost calculation)
+            output_tokens: Output tokens used
+            input_tokens_details: Detailed input token breakdown
+            completion_tokens_details: Detailed completion token breakdown
+            generation_time_ms: Generation time in milliseconds
+            retry_count: Number of retries attempted
+            validation_errors: Any validation errors encountered
+            success: Whether generation was successful
+            error_message: Error message if generation failed
+        """
+        # Get the database record
+        log_record = await session.get(AgentGenerationLogTable, self.id)
+        if not log_record:
+            return
+
+        # Update fields
+        log_record.generated_agent_schema = generated_agent_schema
+        log_record.identified_skills = identified_skills
+        log_record.llm_model = llm_model
+        log_record.total_tokens = total_tokens
+        log_record.input_tokens = input_tokens
+        log_record.cached_input_tokens = cached_input_tokens
+        log_record.output_tokens = output_tokens
+        log_record.input_tokens_details = input_tokens_details
+        log_record.completion_tokens_details = completion_tokens_details
+        log_record.generation_time_ms = generation_time_ms
+        log_record.retry_count = retry_count
+        log_record.validation_errors = validation_errors
+        log_record.success = success
+        log_record.error_message = error_message
+        log_record.completed_at = datetime.now(timezone.utc)
+
+        session.add(log_record)
+        await session.commit()
+        await session.refresh(log_record)
+
+        # Update this instance
+        self.generated_agent_schema = log_record.generated_agent_schema
+        self.identified_skills = log_record.identified_skills
+        self.llm_model = log_record.llm_model
+        self.total_tokens = log_record.total_tokens
+        self.input_tokens = log_record.input_tokens
+        self.cached_input_tokens = log_record.cached_input_tokens
+        self.output_tokens = log_record.output_tokens
+        self.input_tokens_details = log_record.input_tokens_details
+        self.completion_tokens_details = log_record.completion_tokens_details
+        self.generation_time_ms = log_record.generation_time_ms
+        self.retry_count = log_record.retry_count
+        self.validation_errors = log_record.validation_errors
+        self.success = log_record.success
+        self.error_message = log_record.error_message
+        self.completed_at = log_record.completed_at
+
+    @classmethod
+    async def get_by_id(
+        cls,
+        session: AsyncSession,
+        log_id: str,
+    ) -> Optional["AgentGenerationLog"]:
+        """Get an agent generation log by ID.
+
+        Args:
+            session: Database session
+            log_id: Log ID
+
+        Returns:
+            Log instance if found, None otherwise
+        """
+        result = await session.execute(
+            select(AgentGenerationLogTable).where(AgentGenerationLogTable.id == log_id)
+        )
+        log_record = result.scalar_one_or_none()
+
+        if log_record:
+            return cls.model_validate(log_record)
+        return None
+
+    @classmethod
+    async def get_by_user(
+        cls,
+        session: AsyncSession,
+        user_id: str,
+        limit: int = 50,
+    ) -> list["AgentGenerationLog"]:
+        """Get agent generation logs for a user.
+
+        Args:
+            session: Database session
+            user_id: User ID
+            limit: Maximum number of logs to return
+
+        Returns:
+            List of log instances
+        """
+        result = await session.execute(
+            select(AgentGenerationLogTable)
+            .where(AgentGenerationLogTable.user_id == user_id)
+            .order_by(AgentGenerationLogTable.created_at.desc())
+            .limit(limit)
+        )
+        log_records = result.scalars().all()
+
+        return [cls.model_validate(record) for record in log_records]