- API versioning: all REST endpoints prefixed with /api/v1/ - Structured logging: replaced stdlib logging with structlog (console/JSON modes) - Alembic migrations: versioned DB schema with initial migration - Error standardization: global exception handlers for consistent envelope format - Interrupt cleanup: asyncio background task for expired interrupt removal - Integration tests: +30 tests (analytics, replay, openapi, error, session APIs) - Frontend tests: +57 tests (all components, pages, useWebSocket hook) - Backend: 557 tests, 89.75% coverage | Frontend: 80 tests, 16 test files
111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
"""Database connection pool and PostgresSaver checkpoint setup."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
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
|
|
|
|
|
|
def run_alembic_migrations(database_url: str) -> None:
|
|
"""Run Alembic migrations to head."""
|
|
from alembic.config import Config
|
|
|
|
from alembic import command
|
|
|
|
alembic_cfg = Config(str(Path(__file__).parent.parent / "alembic.ini"))
|
|
alembic_cfg.set_main_option("sqlalchemy.url", database_url)
|
|
command.upgrade(alembic_cfg, "head")
|
|
|
|
|
|
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)
|