add source to stream output

src/agents/chat_handler.py

@@ -22,6 +22,7 @@ inject mocks).
 
 from __future__ import annotations
 
+import json
 from collections.abc import AsyncIterator
 from typing import TYPE_CHECKING, Any
 
@@ -47,6 +48,8 @@ class ChatHandler:
     Returns an `AsyncIterator[dict]` of SSE-style events with shape
     `{"event": <name>, "data": <str>}`. Event types:
       - `intent`  — emitted once after classification (JSON-encoded decision)
+      - `sources` — JSON array of source refs (one per structured table, or
+                    per (document_id, page_label) for unstructured)
       - `chunk`   — text fragment of the streaming answer (one per token)
       - `done`    — end of stream (data is empty string)
       - `error`   — failure; data is a user-facing message
@@ -125,6 +128,7 @@ class ChatHandler:
         rewritten = decision.rewritten_query or message
         query_result = None
         chunks: list[DocumentChunk] | None = None
+        raw_chunks: Any = None
 
         # ---- 2. Route ------------------------------------------------
         if decision.source_hint == "structured":
@@ -162,6 +166,12 @@ class ChatHandler:
                 return
         # else: chat path — no context
 
+        # ---- 2b. Emit sources ---------------------------------------
+        sources = _build_sources(
+            decision.source_hint, user_id, query_result, raw_chunks
+        )
+        yield {"event": "sources", "data": json.dumps(sources)}
+
         # ---- 3. Stream answer ----------------------------------------
         try:
             async for token in self._get_answer_agent().astream(
@@ -179,6 +189,54 @@ class ChatHandler:
         yield {"event": "done", "data": ""}
 
 
+def _build_sources(
+    source_hint: str,
+    user_id: str,
+    query_result: Any,
+    raw_chunks: Any,
+) -> list[dict[str, Any]]:
+    """Build the sources payload for the SSE `sources` event.
+
+    - structured: one entry per executed table (table_name only).
+    - unstructured: deduped by (document_id, page_label), Phase 1 shape.
+    - chat or error: empty list.
+    """
+    if source_hint == "structured":
+        if query_result is None or getattr(query_result, "error", None):
+            return []
+        table_name = getattr(query_result, "table_name", "") or ""
+        if not table_name:
+            return []
+        return [{
+            "document_id": f"{user_id}_{table_name}",
+            "filename": table_name,
+            "page_label": None,
+        }]
+
+    if source_hint == "unstructured" and raw_chunks:
+        seen: set[tuple[Any, Any]] = set()
+        sources: list[dict[str, Any]] = []
+        for item in raw_chunks:
+            if isinstance(item, RetrievalResult):
+                data = item.metadata.get("data", {})
+            elif isinstance(item, dict):
+                data = item
+            else:
+                continue
+            key = (data.get("document_id"), data.get("page_label"))
+            if key in seen or key == (None, None):
+                continue
+            seen.add(key)
+            sources.append({
+                "document_id": data.get("document_id"),
+                "filename": data.get("filename", "Unknown"),
+                "page_label": data.get("page_label", "Unknown"),
+            })
+        return sources
+
+    return []
+
+
 def _normalize_chunks(raw: Any) -> list[DocumentChunk]:
     """Convert whatever the retriever returns into list[DocumentChunk].
 

src/api/v1/chat.py

@@ -97,7 +97,8 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
     """Chat endpoint with streaming response.
 
     SSE event sequence:
-      1. sources  — JSON array of source references (empty until retrieval is wired into ChatHandler)
+      1. sources  — JSON array of source refs from ChatHandler (table for
+                    structured; deduped document_id/page_label for unstructured)
       2. chunk    — text fragments of the answer
       3. done     — signals end of stream
     """
@@ -136,14 +137,20 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
 
         async def stream_response():
             full_response = ""
-            yield {"event": "sources", "data": json.dumps([])}
+            sources: List[Dict[str, Any]] = []
             async for event in handler.handle(request.message, request.user_id, history):
-                if event["event"] == "chunk":
+                if event["event"] == "sources":
+                    try:
+                        sources = json.loads(event["data"]) or []
+                    except (TypeError, ValueError):
+                        sources = []
+                    yield event
+                elif event["event"] == "chunk":
                     full_response += event["data"]
                     yield event
                 elif event["event"] == "done":
                     await cache_response(redis, cache_key, full_response)
-                    await save_messages(db, request.room_id, request.message, full_response, sources=[])
+                    await save_messages(db, request.room_id, request.message, full_response, sources=sources)
                     yield event
                 elif event["event"] == "error":
                     yield event

src/query/executor/base.py

@@ -17,6 +17,8 @@ class QueryResult:
     truncated: bool = False
     elapsed_ms: int = 0
     error: str | None = None
+    table_id: str = ""
+    table_name: str = ""
 
 
 class BaseExecutor(ABC):

src/query/executor/db.py

@@ -59,8 +59,12 @@ class DbExecutor(BaseExecutor):
 
     async def run(self, ir: QueryIR) -> QueryResult:
         started = time.perf_counter()
+        table_name = ""
         try:
             source = self._find_source(ir.source_id)
+            table_name = next(
+                (t.name for t in source.tables if t.table_id == ir.table_id), ""
+            )
             if source.source_type != "schema":
                 raise ValueError(
                     f"DbExecutor cannot run on source_type={source.source_type!r}; "
@@ -102,6 +106,8 @@ class DbExecutor(BaseExecutor):
                 row_count=len(capped),
                 truncated=truncated,
                 elapsed_ms=elapsed_ms,
+                table_id=ir.table_id,
+                table_name=table_name,
             )
 
         except Exception as e:
@@ -117,6 +123,8 @@ class DbExecutor(BaseExecutor):
                 backend="sql",
                 elapsed_ms=elapsed_ms,
                 error=str(e),
+                table_id=ir.table_id,
+                table_name=table_name,
             )
 
     # ------------------------------------------------------------------

src/query/executor/tabular.py

@@ -55,8 +55,10 @@ class TabularExecutor(BaseExecutor):
 
     async def run(self, ir: QueryIR) -> QueryResult:
         started = time.perf_counter()
+        table_name = ""
         try:
             source, table = self._lookup(ir)
+            table_name = table.name
             if source.source_type != "tabular":
                 raise ValueError(
                     f"TabularExecutor cannot run on source_type={source.source_type!r}; "
@@ -91,6 +93,8 @@ class TabularExecutor(BaseExecutor):
                 row_count=len(rows),
                 truncated=truncated,
                 elapsed_ms=elapsed_ms,
+                table_id=ir.table_id,
+                table_name=table_name,
             )
 
         except Exception as e:
@@ -106,6 +110,8 @@ class TabularExecutor(BaseExecutor):
                 backend="tabular",
                 elapsed_ms=elapsed_ms,
                 error=str(e),
+                table_id=ir.table_id,
+                table_name=table_name,
             )
 
     # ------------------------------------------------------------------