File size: 6,156 Bytes
5a60e93 f282b15 5a60e93 f282b15 5a60e93 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | """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
|