af53111928
Address all architecture review findings: P0 fixes: - Add API key authentication for admin endpoints (analytics, replay, openapi) and WebSocket connections via ADMIN_API_KEY env var - Add PostgreSQL-backed PgSessionManager and PgInterruptManager for multi-worker production deployments (in-memory defaults preserved) P1 fixes: - Implement actual tool generation in OpenAPI approve_job endpoint using generate_tool_code() and generate_agent_yaml() - Add missing clarification, interrupt_expired, and tool_result message handlers in frontend ChatPage P2 fixes: - Replace monkey-patching on CompiledStateGraph with typed GraphContext - Replace 9-param dispatch_message with WebSocketContext dataclass - Extract duplicate _envelope() into shared app/api_utils.py - Replace mutable module-level counter with crypto.randomUUID() - Remove hardcoded mock data from ReviewPage, use api.ts wrappers - Remove `as any` type escape from ReplayPage All 516 tests passing, 0 TypeScript errors.
99 lines
3.0 KiB
Python
99 lines
3.0 KiB
Python
"""Database connection pool and PostgresSaver checkpoint setup."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
|
from psycopg.rows import dict_row
|
|
from psycopg_pool import AsyncConnectionPool
|
|
|
|
if TYPE_CHECKING:
|
|
from app.config import Settings
|
|
|
|
_CONVERSATIONS_DDL = """
|
|
CREATE TABLE IF NOT EXISTS conversations (
|
|
thread_id TEXT PRIMARY KEY,
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
last_activity TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
total_tokens INTEGER NOT NULL DEFAULT 0,
|
|
total_cost_usd DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
|
status TEXT NOT NULL DEFAULT 'active'
|
|
);
|
|
"""
|
|
|
|
_INTERRUPTS_DDL = """
|
|
CREATE TABLE IF NOT EXISTS active_interrupts (
|
|
interrupt_id TEXT PRIMARY KEY,
|
|
thread_id TEXT NOT NULL REFERENCES conversations(thread_id),
|
|
action TEXT NOT NULL,
|
|
params JSONB NOT NULL DEFAULT '{}',
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
resolved_at TIMESTAMPTZ,
|
|
resolution TEXT
|
|
);
|
|
"""
|
|
|
|
_ANALYTICS_EVENTS_DDL = """
|
|
CREATE TABLE IF NOT EXISTS analytics_events (
|
|
id BIGSERIAL PRIMARY KEY,
|
|
thread_id TEXT NOT NULL,
|
|
event_type TEXT NOT NULL,
|
|
agent_name TEXT,
|
|
tool_name TEXT,
|
|
tokens_used INTEGER NOT NULL DEFAULT 0,
|
|
cost_usd DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
|
duration_ms INTEGER,
|
|
success BOOLEAN,
|
|
error_message TEXT,
|
|
metadata JSONB NOT NULL DEFAULT '{}',
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
|
);
|
|
"""
|
|
|
|
_SESSIONS_DDL = """
|
|
CREATE TABLE IF NOT EXISTS sessions (
|
|
thread_id TEXT PRIMARY KEY,
|
|
last_activity TIMESTAMPTZ NOT NULL DEFAULT NOW(),
|
|
has_pending_interrupt BOOLEAN NOT NULL DEFAULT FALSE,
|
|
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
|
|
);
|
|
"""
|
|
|
|
_CONVERSATIONS_MIGRATION_DDL = """
|
|
ALTER TABLE conversations
|
|
ADD COLUMN IF NOT EXISTS resolution_type TEXT,
|
|
ADD COLUMN IF NOT EXISTS agents_used TEXT[],
|
|
ADD COLUMN IF NOT EXISTS turn_count INTEGER NOT NULL DEFAULT 0,
|
|
ADD COLUMN IF NOT EXISTS ended_at TIMESTAMPTZ;
|
|
"""
|
|
|
|
|
|
async def create_pool(settings: Settings) -> AsyncConnectionPool:
|
|
"""Create an async connection pool with the required psycopg settings."""
|
|
pool = AsyncConnectionPool(
|
|
conninfo=settings.database_url,
|
|
kwargs={"autocommit": True, "row_factory": dict_row},
|
|
min_size=2,
|
|
max_size=10,
|
|
)
|
|
await pool.open()
|
|
return pool
|
|
|
|
|
|
async def create_checkpointer(pool: AsyncConnectionPool) -> AsyncPostgresSaver:
|
|
"""Create and initialize the LangGraph checkpointer."""
|
|
checkpointer = AsyncPostgresSaver(conn=pool)
|
|
await checkpointer.setup()
|
|
return checkpointer
|
|
|
|
|
|
async def setup_app_tables(pool: AsyncConnectionPool) -> None:
|
|
"""Create application-specific tables and apply migrations."""
|
|
async with pool.connection() as conn:
|
|
await conn.execute(_CONVERSATIONS_DDL)
|
|
await conn.execute(_INTERRUPTS_DDL)
|
|
await conn.execute(_SESSIONS_DDL)
|
|
await conn.execute(_ANALYTICS_EVENTS_DDL)
|
|
await conn.execute(_CONVERSATIONS_MIGRATION_DDL)
|