"""ChartStore — the seam the chat pipeline persists `render_chart` outputs through (S2 visualization, SPINE_V2_PLAN §4.4). `ChatHandler._run_slow_path` scans a completed `AnalysisRecord.results_snapshot` for `ToolOutput(kind="chart")` entries and flushes one `MessageChartRow` per chart through this seam, right before the `done` SSE event (mirrors the traceability flush). Unlike traceability's one row per turn, a turn with multiple `render_chart` calls writes multiple rows sharing one `message_id`. `GET /api/v1/charts` reads them back by (analysis_id, message_id). - `NullChartStore` logs the envelope and stores nothing (tests / disabled persistence). - `PostgresChartStore` writes one `message_charts` row per chart (dedorch, `AsyncSessionLocal`), mirroring `PostgresTraceabilityStore`. `save` must NEVER raise on the caller's path — a chart-persist failure must not break the user's answer. `list_for_message` is the endpoint read; an empty list is a valid result (chartless turn), not an error. """ from __future__ import annotations import uuid from datetime import datetime from typing import Protocol, runtime_checkable from pydantic import BaseModel from sqlalchemy import select from src.db.postgres.connection import AsyncSessionLocal from src.db.postgres.models import MessageChartRow, MessageTraceabilityRow from src.middlewares.logging import get_logger logger = get_logger("charts_store") class ChartRecord(BaseModel): """One persisted chart, as served by `GET /api/v1/charts`.""" chart_id: str chart_type: str title: str | None = None spec: dict created_at: datetime @runtime_checkable class ChartStore(Protocol): """Persist + read `render_chart` outputs for one assistant `message_id`. `save` must never raise on the caller's path. `list_for_message` returns every chart for one turn (possibly empty — a turn need not have asked for a chart); the lookup is by `message_id` alone (Python-minted UUID4, globally unique — lead decision 2026-07-13). `turn_exists` says whether the turn flushed a traceability row, so the endpoint can tell "no charts" from "unknown id". """ async def save( self, *, message_id: str, analysis_id: str, user_id: str, record_id: str | None, envelope: dict, ) -> None: ... async def list_for_message(self, message_id: str) -> list[ChartRecord]: ... async def turn_exists(self, message_id: str) -> bool: ... class NullChartStore: """No-op store: logs the envelope, persists nothing. Reads return `[]`.""" async def save( self, *, message_id: str, analysis_id: str, user_id: str, record_id: str | None, envelope: dict, ) -> None: logger.info( "chart produced (not persisted — NullChartStore)", message_id=message_id, chart_type=envelope.get("chart_type", "unknown"), ) async def list_for_message(self, message_id: str) -> list[ChartRecord]: return [] async def turn_exists(self, message_id: str) -> bool: return False class PostgresChartStore: """Writes/reads `message_charts` rows. One insert per chart (not an upsert).""" async def save( self, *, message_id: str, analysis_id: str, user_id: str, record_id: str | None, envelope: dict, ) -> None: try: async with AsyncSessionLocal() as session: row = MessageChartRow( id=str(uuid.uuid4()), message_id=message_id, analysis_id=analysis_id, user_id=user_id, record_id=record_id, chart_type=envelope.get("chart_type", "unknown"), title=envelope.get("title"), spec=envelope, ) session.add(row) await session.commit() logger.info( "chart persisted", message_id=message_id, analysis_id=analysis_id, chart_type=envelope.get("chart_type", "unknown"), ) except Exception as exc: # never break the user's answer logger.error( "chart persist failed", degraded_seam="chart_persist", message_id=message_id, error=repr(exc), ) async def list_for_message(self, message_id: str) -> list[ChartRecord]: # message_id-only lookup (lead decision 2026-07-13). NOTE for the Harry # migration: the manual DDL's composite index (analysis_id, message_id) # does not serve this predicate — an additive index on (message_id) is # part of the handoff. async with AsyncSessionLocal() as session: result = await session.execute( select(MessageChartRow) .where(MessageChartRow.message_id == message_id) .order_by(MessageChartRow.created_at) ) rows = result.scalars().all() return [ ChartRecord( chart_id=row.id, chart_type=row.chart_type, title=row.title, spec=row.spec, created_at=row.created_at, ) for row in rows ] async def turn_exists(self, message_id: str) -> bool: """True iff the turn flushed its traceability row (written before `done`). Lets the endpoint tri-state a zero-chart GET: `empty` (completed turn, no charts — the common case) vs `not_found` (unknown/mistyped id, or an error turn, which never writes traceability). PK lookup — cheap. """ async with AsyncSessionLocal() as session: result = await session.execute( select(MessageTraceabilityRow.message_id).where( MessageTraceabilityRow.message_id == message_id ) ) return result.scalar_one_or_none() is not None