File size: 5,617 Bytes
f873f92
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f282b15
f873f92
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f282b15
f873f92
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f282b15
f873f92
f282b15
f873f92
 
 
f282b15
f873f92
f282b15
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f873f92
f282b15
 
 
 
 
 
f873f92
f282b15
f873f92
 
 
 
 
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
"""TraceabilityStore — the seam the chat pipeline persists provenance through (KM-691).

`ChatHandler` (and the v2 chat endpoint's greeting/cache branches) flush one
`TraceabilityPayload` per assistant turn through this seam, right before the `done`
SSE event; `GET /api/v1/traceability` reads it back by (analysis_id, message_id).

- `NullTraceabilityStore` logs and stores nothing (tests / disabled persistence).
- `PostgresTraceabilityStore` writes one `message_traceability` row per turn
  (dedorch, `AsyncSessionLocal`), mirroring `PostgresReportInputStore`.

`save` must NEVER raise on the caller's path — a persistence failure must not break
the user's answer. `get` is the endpoint read and returns `None` on a miss (→ 404).
"""

from __future__ import annotations

from typing import Protocol, runtime_checkable

from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert

from src.db.postgres.connection import AsyncSessionLocal
from src.db.postgres.models import MessageTraceabilityRow
from src.middlewares.logging import get_logger

from .schemas import TraceabilityPayload

logger = get_logger("traceability_store")


@runtime_checkable
class TraceabilityStore(Protocol):
    """Persist + read one provenance record per assistant `message_id`.

    `save` must never raise on the caller's path. `get` returns the payload for one
    turn, or `None` if none exists yet (the turn is still running or the id is unknown).
    """

    async def save(self, payload: TraceabilityPayload) -> None: ...

    async def get(
        self, analysis_id: str, message_id: str, user_id: str | None = None
    ) -> TraceabilityPayload | None: ...


class NullTraceabilityStore:
    """No-op store: logs the payload, persists nothing. Reads return `None`."""

    async def save(self, payload: TraceabilityPayload) -> None:
        logger.info(
            "traceability produced (not persisted — NullTraceabilityStore)",
            message_id=payload.message_id,
            intent=payload.intent,
            n_tool_calls=len(payload.tool_calls),
        )

    async def get(
        self, analysis_id: str, message_id: str, user_id: str | None = None
    ) -> TraceabilityPayload | None:
        return None


class PostgresTraceabilityStore:
    """Writes/reads `message_traceability` jsonb rows. Upsert on `message_id`."""

    async def save(self, payload: TraceabilityPayload) -> None:
        try:
            data = payload.model_dump(mode="json", by_alias=True)
            async with AsyncSessionLocal() as session:
                stmt = insert(MessageTraceabilityRow).values(
                    message_id=payload.message_id,
                    analysis_id=payload.analysis_id,
                    user_id=payload.user_id,
                    intent=payload.intent,
                    data=data,
                )
                # Idempotent: a re-flushed turn overwrites its own row.
                stmt = stmt.on_conflict_do_update(
                    index_elements=[MessageTraceabilityRow.message_id],
                    set_={"data": stmt.excluded.data, "intent": stmt.excluded.intent},
                )
                await session.execute(stmt)
                await session.commit()
            logger.info(
                "traceability persisted",
                message_id=payload.message_id,
                analysis_id=payload.analysis_id,
                intent=payload.intent,
            )
        except Exception as exc:  # never break the user's answer
            logger.error(
                "traceability persist failed",
                degraded_seam="traceability_persist",
                message_id=payload.message_id,
                error=repr(exc),
            )

    async def get(
        self, analysis_id: str, message_id: str, user_id: str | None = None
    ) -> TraceabilityPayload | None:
        """One turn's payload, or None on a miss.

        `user_id` scopes the read to the turn's owner when a caller supplies one.

        **It stays unused by the endpoint — lead decision, reaffirmed 2026-07-23.**
        `GET /api/v1/traceability` looks up by `(analysis_id, message_id)` alone, and
        `GET /api/v1/charts` by `message_id` alone (the 2026-07-13 decision). Review
        finding F-3 proposed adding the parameter to both; that was **declined**, so
        do not "finish" this by wiring it into the endpoint — reopening it needs the
        same sign-off path as any other locked decision.

        The consequence is deliberate and accepted: this row is reachable by
        `(analysis_id, message_id)` alone, and the payload itself carries `user_id`.
        **The service-secret gate is the control that protects these endpoints**, not
        this predicate — which is why `dataeyond__service__secret` being set matters
        more than it would otherwise (DEV_PLAN #37).

        The parameter is kept because it costs nothing and a future Go-forwarded
        identity (#43) would use it.
        """
        async with AsyncSessionLocal() as session:
            where = [
                MessageTraceabilityRow.message_id == message_id,
                MessageTraceabilityRow.analysis_id == analysis_id,
            ]
            if user_id is not None:
                where.append(MessageTraceabilityRow.user_id == user_id)
            result = await session.execute(
                select(MessageTraceabilityRow.data).where(*where)
            )
            row = result.scalar_one_or_none()
        if row is None:
            return None
        return TraceabilityPayload.model_validate(row)