feat: add draft functionality and manager module

Author: hyacinthus

intentkit/core/draft.py

@@ -0,0 +1,201 @@
+"""Service functions for agent draft operations."""
+
+from __future__ import annotations
+
+from epyxid import XID
+from fastapi import status
+from sqlalchemy import desc, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from intentkit.models.agent import Agent, AgentUserInput
+from intentkit.models.draft import AgentDraft, AgentDraftTable
+from intentkit.utils.error import IntentKitAPIError
+
+
+async def update_agent_draft(
+    *,
+    agent_id: str,
+    user_id: str,
+    input: AgentUserInput,
+    db: AsyncSession,
+) -> AgentDraft:
+    """Update the latest draft for the specified agent with partial field updates.
+
+    This function only updates fields that are explicitly provided in the input,
+    leaving other fields unchanged. This is more efficient than override as it
+    reduces context usage and minimizes the risk of accidentally changing fields.
+    """
+    query = (
+        select(AgentDraftTable)
+        .where(AgentDraftTable.agent_id == agent_id, AgentDraftTable.owner == user_id)
+        .order_by(desc(AgentDraftTable.created_at))
+        .limit(1)
+    )
+
+    result = await db.execute(query)
+    latest_draft = result.scalar_one_or_none()
+
+    if not latest_draft:
+        raise IntentKitAPIError(
+            status.HTTP_404_NOT_FOUND,
+            "DraftNotFound",
+            "No drafts found for this agent",
+        )
+
+    # Get only the fields that are explicitly provided (exclude_unset=True)
+    update_data = input.model_dump(exclude_unset=True)
+
+    if latest_draft.deployed_at is not None:
+        # Create new draft version if current one is deployed
+        draft_id = str(XID())
+
+        # Start with existing draft data and merge updates
+        draft_data = AgentUserInput.model_validate(latest_draft).model_dump()
+        draft_data.update(update_data)
+
+        updated_input = AgentUserInput.model_validate(draft_data)
+
+        draft_table = AgentDraftTable(
+            id=draft_id,
+            agent_id=agent_id,
+            owner=user_id,
+            version=updated_input.hash(),
+            last_draft_id=latest_draft.id,
+            project_id=latest_draft.project_id,
+            **draft_data,
+        )
+
+        db.add(draft_table)
+        await db.commit()
+        await db.refresh(draft_table)
+
+        return AgentDraft.model_validate(draft_table)
+
+    # Update existing draft in-place
+    for key, value in update_data.items():
+        setattr(latest_draft, key, value)
+
+    # Update version hash based on updated data
+    updated_input = AgentUserInput.model_validate(latest_draft)
+    latest_draft.version = updated_input.hash()
+
+    await db.commit()
+    await db.refresh(latest_draft)
+
+    return AgentDraft.model_validate(latest_draft)
+
+
+async def override_agent_draft(
+    *,
+    agent_id: str,
+    user_id: str,
+    input: AgentUserInput,
+    db: AsyncSession,
+) -> AgentDraft:
+    """Override the latest draft for the specified agent."""
+    query = (
+        select(AgentDraftTable)
+        .where(AgentDraftTable.agent_id == agent_id, AgentDraftTable.owner == user_id)
+        .order_by(desc(AgentDraftTable.created_at))
+        .limit(1)
+    )
+
+    result = await db.execute(query)
+    latest_draft = result.scalar_one_or_none()
+
+    if not latest_draft:
+        raise IntentKitAPIError(
+            status.HTTP_404_NOT_FOUND,
+            "DraftNotFound",
+            "No drafts found for this agent",
+        )
+
+    if latest_draft.deployed_at is not None:
+        draft_id = str(XID())
+
+        draft_table = AgentDraftTable(
+            id=draft_id,
+            agent_id=agent_id,
+            owner=user_id,
+            version=input.hash(),
+            last_draft_id=latest_draft.id,
+            project_id=latest_draft.project_id,
+            **input.model_dump(),
+        )
+
+        db.add(draft_table)
+        await db.commit()
+        await db.refresh(draft_table)
+
+        return AgentDraft.model_validate(draft_table)
+
+    for key, value in input.model_dump().items():
+        setattr(latest_draft, key, value)
+
+    latest_draft.version = input.hash()
+
+    await db.commit()
+    await db.refresh(latest_draft)
+
+    return AgentDraft.model_validate(latest_draft)
+
+
+async def get_agent_latest_draft(
+    *,
+    agent_id: str,
+    user_id: str,
+    db: AsyncSession,
+) -> AgentDraft:
+    """Return the latest draft for the specified agent."""
+    query = (
+        select(AgentDraftTable)
+        .where(AgentDraftTable.agent_id == agent_id, AgentDraftTable.owner == user_id)
+        .order_by(desc(AgentDraftTable.created_at))
+        .limit(1)
+    )
+
+    result = await db.execute(query)
+    latest_draft = result.scalar_one_or_none()
+
+    if latest_draft:
+        return AgentDraft.model_validate(latest_draft)
+
+    agent = await Agent.get(agent_id)
+
+    if not agent:
+        raise IntentKitAPIError(
+            status.HTTP_404_NOT_FOUND,
+            "AgentNotFound",
+            "No drafts found for this agent",
+        )
+
+    if agent.owner != user_id:
+        raise IntentKitAPIError(
+            status.HTTP_403_FORBIDDEN,
+            "Forbidden",
+            "Not your agent",
+        )
+
+    draft_id = str(XID())
+
+    agent_dict = agent.model_dump()
+    input_dict: dict[str, object] = {}
+    for key in AgentUserInput.model_fields:
+        if key in agent_dict:
+            input_dict[key] = agent_dict[key]
+    input = AgentUserInput.model_validate(input_dict)
+
+    draft_table = AgentDraftTable(
+        id=draft_id,
+        agent_id=agent_id,
+        owner=user_id,
+        version=input.hash(),
+        deployed_at=agent.updated_at,
+        **input.model_dump(),
+    )
+
+    db.add(draft_table)
+    await db.commit()
+    await db.refresh(draft_table)
+
+    return AgentDraft.model_validate(draft_table)

intentkit/core/draft_chat.py

@@ -0,0 +1,118 @@
+"""Utilities for streaming draft agent conversations."""
+
+import logging
+import time
+from datetime import datetime, timedelta, timezone
+from typing import AsyncGenerator
+
+from langgraph.graph.state import CompiledStateGraph
+from sqlalchemy import desc, select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from intentkit.core.engine import build_agent, stream_agent_raw
+from intentkit.models.agent import Agent
+from intentkit.models.agent_data import AgentData
+from intentkit.models.chat import ChatMessage, ChatMessageCreate
+from intentkit.models.db import get_session
+from intentkit.models.draft import AgentDraft, AgentDraftTable
+from intentkit.utils.error import IntentKitAPIError
+
+logger = logging.getLogger(__name__)
+
+
+_draft_executors: dict[str, CompiledStateGraph] = {}
+_draft_updated_at: dict[str, datetime] = {}
+_draft_cached_at: dict[str, datetime] = {}
+
+_CACHE_TTL = timedelta(days=1)
+
+
+async def stream_draft(
+    agent_id: str, message: ChatMessageCreate
+) -> AsyncGenerator[ChatMessage, None]:
+    """Stream chat messages for the latest draft of an agent."""
+
+    draft = await _get_latest_draft(agent_id)
+    agent = _agent_from_draft(draft)
+    executor, cold_start_cost = await _get_draft_executor(agent, draft)
+
+    if not message.agent_id:
+        message.agent_id = agent.id
+    message.cold_start_cost = cold_start_cost
+
+    async for chat_message in stream_agent_raw(message, agent, executor):
+        yield chat_message
+
+
+async def _get_latest_draft(agent_id: str) -> AgentDraft:
+    async with get_session() as session:
+        result = await _execute_latest_draft_query(session, agent_id)
+        draft_row = result.scalar_one_or_none()
+
+    if not draft_row:
+        raise IntentKitAPIError(
+            status_code=404,
+            key="DraftNotFound",
+            message=f"No draft found for agent {agent_id}",
+        )
+
+    return AgentDraft.model_validate(draft_row)
+
+
+async def _execute_latest_draft_query(session: AsyncSession, agent_id: str):
+    statement = (
+        select(AgentDraftTable)
+        .where(AgentDraftTable.agent_id == agent_id)
+        .order_by(desc(AgentDraftTable.updated_at))
+        .limit(1)
+    )
+    return await session.execute(statement)
+
+
+def _agent_from_draft(draft: AgentDraft) -> Agent:
+    data = draft.model_dump()
+    data.pop("id", None)
+    data.pop("agent_id", None)
+    data.pop("last_draft_id", None)
+    data["id"] = draft.agent_id
+    data["owner"] = draft.owner
+    data["deployed_at"] = draft.deployed_at
+    data["created_at"] = draft.created_at
+    data["updated_at"] = draft.updated_at
+    data["version"] = draft.version
+    return Agent.model_validate(data)
+
+
+async def _get_draft_executor(
+    agent: Agent, draft: AgentDraft
+) -> tuple[CompiledStateGraph, float]:
+    now = datetime.now(timezone.utc)
+    _cleanup_cache(now)
+
+    cached_executor = _draft_executors.get(agent.id)
+    cached_updated = _draft_updated_at.get(agent.id)
+    cold_start_cost = 0.0
+
+    if not cached_executor or cached_updated != draft.updated_at:
+        start = time.perf_counter()
+        agent_data = AgentData(id=agent.id)
+        cached_executor = await build_agent(agent, agent_data)
+        cold_start_cost = time.perf_counter() - start
+        _draft_executors[agent.id] = cached_executor
+        _draft_updated_at[agent.id] = draft.updated_at
+        _draft_cached_at[agent.id] = now
+        logger.info("Initialized draft executor for agent %s", agent.id)
+    else:
+        _draft_cached_at[agent.id] = now
+
+    return cached_executor, cold_start_cost
+
+
+def _cleanup_cache(now: datetime) -> None:
+    expired_before = now - _CACHE_TTL
+    for agent_id, cached_time in list(_draft_cached_at.items()):
+        if cached_time < expired_before:
+            _draft_cached_at.pop(agent_id, None)
+            _draft_updated_at.pop(agent_id, None)
+            _draft_executors.pop(agent_id, None)
+            logger.debug("Removed expired draft executor for agent %s", agent_id)

intentkit/core/manager/__init__.py

@@ -0,0 +1,25 @@
+"""Manager module for agent management operations."""
+
+from intentkit.core.manager.engine import stream_manager
+from intentkit.core.manager.service import (
+    agent_draft_json_schema,
+    get_latest_public_info,
+    get_skills_hierarchical_text,
+)
+from intentkit.core.manager.skills import (
+    get_agent_latest_draft_skill,
+    get_agent_latest_public_info_skill,
+    update_agent_draft_skill,
+    update_public_info_skill,
+)
+
+__all__ = [
+    "stream_manager",
+    "agent_draft_json_schema",
+    "get_skills_hierarchical_text",
+    "get_latest_public_info",
+    "get_agent_latest_draft_skill",
+    "get_agent_latest_public_info_skill",
+    "update_agent_draft_skill",
+    "update_public_info_skill",
+]