File size: 4,983 Bytes
0721bb4 0e02a0f 0721bb4 0066161 0e02a0f 0721bb4 0e02a0f 0721bb4 0066161 0721bb4 0e02a0f 0721bb4 0066161 0721bb4 0e02a0f 0721bb4 0066161 0721bb4 0e02a0f 0721bb4 0066161 0721bb4 0066161 0721bb4 | 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 | """AnalysisStateStore — read/write the per-analysis session state.
The orchestrator gate + Help skill read `AnalysisState` (the locked contract in
`gate.py`) every turn; the Problem Statement skill writes `problem_validated`. The
row shares its id with the chat `rooms` row — one session = one analysis = one
conversation (`analysis_id == room_id`).
Mirrors `PostgresReportInputStore`: each call opens its own `AsyncSession`.
"""
from __future__ import annotations
from sqlalchemy.dialects.postgresql import insert
from src.agents.gate import AnalysisState
from src.db.postgres.connection import AsyncSessionLocal
from src.db.postgres.models import AnalysisStateRow
from src.middlewares.logging import get_logger
logger = get_logger("analysis_state_store")
def _row_to_state(row: AnalysisStateRow) -> AnalysisState:
"""Map a DB row to the frozen `AnalysisState` contract."""
return AnalysisState(
id=row.id,
analysis_title=row.analysis_title,
objective=row.objective or "",
business_questions=list(row.business_questions or []),
user_id=row.user_id,
report_id=row.report_id,
created_at=row.created_at,
updated_at=row.updated_at,
)
class AnalysisStateStore:
"""Read/write the dedorch `analysis` table, keyed by the shared session id."""
async def get(self, analysis_id: str) -> AnalysisState | None:
async with AsyncSessionLocal() as session:
row = await session.get(AnalysisStateRow, analysis_id)
return _row_to_state(row) if row is not None else None
async def ensure(
self,
analysis_id: str,
user_id: str,
analysis_title: str = "New analysis",
) -> AnalysisState:
"""Get-or-create the state row for a session (idempotent, race-safe).
Sessions born from `/room/create` have no `analysis_states` row; without
one the gate redirect-loops and `problem_statement` / `report_id` writes
silently no-op. Called per turn (analysis_id == room_id) so any session is
gate-ready. `INSERT ... ON CONFLICT DO NOTHING` makes concurrent first
turns safe; the row is then read back. Legacy rows created this way carry
no source bindings — binding scoping fail-opens to the whole catalog.
"""
async with AsyncSessionLocal() as session:
# Bare get-or-create row. objective/business_questions are always supplied
# (empty) so the INSERT satisfies dedorch whether or not they are NOT NULL;
# the real goal is set at onboarding by Go and preserved by ON CONFLICT DO NOTHING.
stmt = (
insert(AnalysisStateRow)
.values(
id=analysis_id,
user_id=user_id,
analysis_title=analysis_title,
objective="",
business_questions=[],
)
.on_conflict_do_nothing(index_elements=[AnalysisStateRow.id])
)
await session.execute(stmt)
await session.commit()
row = await session.get(AnalysisStateRow, analysis_id)
return _row_to_state(row)
async def create(
self,
*,
analysis_id: str,
user_id: str,
analysis_title: str = "New analysis",
objective: str = "",
business_questions: list[str] | None = None,
) -> AnalysisState:
"""Create the state row for a new analysis (id shared with its chat room)."""
async with AsyncSessionLocal() as session:
row = AnalysisStateRow(
id=analysis_id,
user_id=user_id,
analysis_title=analysis_title,
objective=objective,
business_questions=business_questions or [],
)
session.add(row)
await session.commit()
await session.refresh(row)
return _row_to_state(row)
async def update(
self,
analysis_id: str,
*,
report_id: str | None = None,
) -> AnalysisState | None:
"""Patch the given fields (only non-None args are written). Returns the row.
Used by the report flow (`report_id` write-back). The analysis goal
(`objective`/`business_questions`) is set at onboarding by Go, not patched here.
Returns None if the analysis doesn't exist.
"""
async with AsyncSessionLocal() as session:
row = await session.get(AnalysisStateRow, analysis_id)
if row is None:
logger.warning(
"analysis row missing — update skipped",
analysis_id=analysis_id,
)
return None
if report_id is not None:
row.report_id = report_id
await session.commit()
await session.refresh(row)
return _row_to_state(row)
|