This view is limited to 50 files because it contains too many changes. See the raw diff here.
Files changed (50) hide show
  1. .gitignore +8 -1
  2. ARCHITECTURE.md +6 -5
  3. PHASE1_TO_PHASE2_REPORT.md +16 -3
  4. PROGRESS.md +292 -7
  5. REPO_CONTEXT.md +2 -2
  6. pyproject.toml +3 -1
  7. src/agents/chat_handler.py +188 -3
  8. src/agents/chatbot.py +10 -2
  9. src/agents/orchestration.py +8 -3
  10. src/agents/planner/__init__.py +8 -0
  11. src/agents/planner/business_context.py +31 -0
  12. src/agents/planner/contracts.py +65 -0
  13. src/agents/planner/errors.py +15 -0
  14. src/agents/planner/examples.py +384 -0
  15. src/agents/planner/inputs.py +139 -0
  16. src/agents/planner/prompt.py +106 -0
  17. src/agents/planner/registry.py +132 -0
  18. src/agents/planner/schemas.py +60 -0
  19. src/agents/planner/service.py +158 -0
  20. src/agents/planner/validator.py +229 -0
  21. src/agents/slow_path/__init__.py +10 -0
  22. src/agents/slow_path/assembler.py +140 -0
  23. src/agents/slow_path/coordinator.py +66 -0
  24. src/agents/slow_path/errors.py +11 -0
  25. src/agents/slow_path/invoker.py +27 -0
  26. src/agents/slow_path/prompt.py +66 -0
  27. src/agents/slow_path/schemas.py +99 -0
  28. src/agents/slow_path/store.py +44 -0
  29. src/agents/slow_path/task_runner.py +164 -0
  30. src/api/v1/chat.py +16 -1
  31. src/catalog/introspect/base.py +5 -0
  32. src/catalog/introspect/tabular.py +2 -2
  33. src/catalog/reader.py +27 -0
  34. src/config/prompts/assembler.md +43 -0
  35. src/config/prompts/planner.md +56 -0
  36. src/config/settings.py +8 -0
  37. src/database_client/engine.py +168 -0
  38. src/observability/langfuse/tracing.py +169 -0
  39. src/pipeline/db_pipeline/extractor.py +3 -1
  40. src/query/compiler/sql.py +20 -5
  41. src/query/executor/db.py +60 -16
  42. src/tools/__init__.py +7 -0
  43. src/tools/analytics/__init__.py +1 -0
  44. src/tools/analytics/aggregation.py +122 -0
  45. src/tools/analytics/comparison.py +137 -0
  46. src/tools/analytics/decomposition.py +135 -0
  47. src/tools/analytics/descriptive.py +154 -0
  48. src/tools/analytics/quality.py +139 -0
  49. src/tools/analytics/relationship.py +132 -0
  50. src/tools/analytics/segmentation.py +149 -0
.gitignore CHANGED
@@ -46,6 +46,13 @@ test_tesseract.py
46
  # Windows binaries — installed via apt in Docker instead
47
  software/
48
 
 
49
  tests/
50
  .claude/
51
- migratego/
 
 
 
 
 
 
 
46
  # Windows binaries — installed via apt in Docker instead
47
  software/
48
 
49
+ scripts/
50
  tests/
51
  .claude/
52
+ migratego/
53
+ docs/specs/tabular_parquet_contract.md
54
+ docs/specs/tabular_parquet.md
55
+
56
+ # Personal / local working docs (not for the shared repo)
57
+ AGENT_ARCHITECTURE_CONTEXT_new.md
58
+ PROJECT_SUMMARY.md
ARCHITECTURE.md CHANGED
@@ -1,7 +1,7 @@
1
  # Architecture — Data Eyond Agentic Service
2
 
3
- **Last updated**: 2026-05-07
4
- **Status**: Design phase folder skeleton in place, implementation in progress
5
 
6
  ---
7
 
@@ -21,7 +21,7 @@ A catalog-driven AI service for data analysis. Users upload documents and regist
21
 
22
  The architecture has two paths:
23
 
24
- - **Unstructured** (PDF, DOCX, TXT) — dense similarity over prose chunks (the right primitive for free-form text).
25
  - **Structured** (databases, XLSX, CSV, Parquet) — a per-user **data catalog** describes what tables/columns exist; an LLM produces a structured **JSON intermediate representation (IR)** of the user's intent; a deterministic compiler turns the IR into SQL or pandas operations.
26
 
27
  The LLM produces *intent*, not query syntax. Deterministic code does the rest.
@@ -120,6 +120,7 @@ Compiler and executors are pure code. No LLM in the hot path of query constructi
120
  ### Ingestion (when user uploads a file or connects a DB)
121
 
122
  ```
 
123
  source upload / DB connect
124
 
125
  introspect schema (DB: information_schema; tabular: file headers + sample rows)
@@ -129,7 +130,7 @@ validate (Pydantic)
129
  write to catalog store (Postgres jsonb in `data_catalog`, keyed by user_id)
130
  ```
131
 
132
- For unstructured files: chunk + embed PGVector.
133
 
134
  ### Query (per user message)
135
 
@@ -143,7 +144,7 @@ Load chat history
143
  IntentRouter LLM → needs_search? source_hint?
144
 
145
  ├── chat → ChatbotAgent → SSE stream
146
- ├── unstructured → DocumentRetriever → answerer
147
  └── structured →
148
  CatalogReader (load full Cs ∪ Ct for user)
149
 
 
1
  # Architecture — Data Eyond Agentic Service
2
 
3
+ **Last updated**: 2026-05-20
4
+ **Status**: Phase 2 catalog path shipped; document ingestion has moved to a separate Go service. The long-term split is **Python = agent/ML layer, Go = data plane**; this document covers the Python side only.
5
 
6
  ---
7
 
 
21
 
22
  The architecture has two paths:
23
 
24
+ - **Unstructured** (PDF, DOCX, TXT) — dense similarity over prose chunks (the right primitive for free-form text). **Ingestion is handled by a separate Go service**; this Python service reads embeddings from PGVector at query time.
25
  - **Structured** (databases, XLSX, CSV, Parquet) — a per-user **data catalog** describes what tables/columns exist; an LLM produces a structured **JSON intermediate representation (IR)** of the user's intent; a deterministic compiler turns the IR into SQL or pandas operations.
26
 
27
  The LLM produces *intent*, not query syntax. Deterministic code does the rest.
 
120
  ### Ingestion (when user uploads a file or connects a DB)
121
 
122
  ```
123
+ Structured sources (DB connect / XLSX / CSV / Parquet upload) — Python:
124
  source upload / DB connect
125
 
126
  introspect schema (DB: information_schema; tabular: file headers + sample rows)
 
130
  write to catalog store (Postgres jsonb in `data_catalog`, keyed by user_id)
131
  ```
132
 
133
+ **Unstructured ingestion (PDF / DOCX / TXT) is handled by a separate Go service**, which writes chunks + embeddings into the `documents` collection in PGVector. The Python service does not own this path — it reads only.
134
 
135
  ### Query (per user message)
136
 
 
144
  IntentRouter LLM → needs_search? source_hint?
145
 
146
  ├── chat → ChatbotAgent → SSE stream
147
+ ├── unstructured → DocumentRetriever (raw SQL: pgvector `<=>` cosine or `<+>` manhattan) → answerer
148
  └── structured →
149
  CatalogReader (load full Cs ∪ Ct for user)
150
 
PHASE1_TO_PHASE2_REPORT.md CHANGED
@@ -52,7 +52,7 @@ Everything else — IR validation, SQL/pandas compilation, execution — is dete
52
  | `src/agents/chatbot.py` (Phase 1) → deleted, then `src/agents/answer_agent.py` (`AnswerAgent`) → renamed | `src/agents/chatbot.py::ChatbotAgent` | Final answer formation; streams via `astream` |
53
  | `src/knowledge/parquet_service.py` | `src/storage/parquet.py` | Parquet upload/download helper |
54
  | `src/pipeline/document_pipeline/document_pipeline.py` (folder) | `src/pipeline/document_pipeline.py` (flat) | Single module |
55
- | `src/rag/retrievers/document.py` | `src/retrieval/document.py` | `DocumentRetriever` migrated; tabular file types filtered out of results |
56
  | `src/rag/router.py` | `src/retrieval/router.py` | `RetrievalRouter`, Redis-cached, unstructured-only; dead `db: AsyncSession` + `source_hint` params removed |
57
  | `src/rag/base.py` (`RetrievalResult`, `BaseRetriever`) | `src/retrieval/base.py` | Same dataclass + ABC |
58
 
@@ -177,7 +177,7 @@ POST /api/v1/chat/stream
177
 
178
  ├── source_hint == "unstructured"
179
  │ → RetrievalRouter.retrieve() [retrieval/router.py, Redis-cached]
180
- │ → DocumentRetriever (PGVector MMR/cosine/etc.)
181
  │ → ChatbotAgent.astream(chunks=...)
182
 
183
  └── source_hint == "structured"
@@ -237,7 +237,7 @@ POST /api/v1/chat/stream
237
  | Top-level chat orchestration | Inline in `api/v1/chat.py` | `agents/chat_handler.py::ChatHandler` | Extracted to a reusable module |
238
  | Final answer formation | `agents/chatbot.py` (Phase 1 LangChain) | `agents/chatbot.py::ChatbotAgent` (was `AnswerAgent` in `answer_agent.py` mid-cycle) | Rewritten + renamed |
239
  | Schema retrieval (DB / tabular) | `rag/retrievers/schema.py` + PGVector chunks | **Removed**. Replaced by catalog (`catalog/store.py` jsonb) loaded verbatim into planner prompt | Whole concept replaced |
240
- | Doc retrieval (PDF / DOCX / TXT) | `rag/retrievers/document.py`, `rag/router.py` | `retrieval/document.py`, `retrieval/router.py` | Moved; Redis cache restored; tabular files filtered |
241
  | Query writing | `query/query_executor.py` + `models/sql_query.py` (LLM writes SQL) | `query/planner/service.py` (LLM writes IR) + `query/compiler/sql.py` (deterministic) | LLM emits intent, not SQL |
242
  | DB execution | `query/executors/db_executor.py` | `query/executor/db.py::DbExecutor` | Folder renamed (`executors` → `executor`); sqlglot guard + RO txn + 30 s timeout kept |
243
  | Tabular execution | `query/executors/tabular.py` | `query/executor/tabular.py::TabularExecutor` | Parquet-only; pandas compiler split out |
@@ -258,3 +258,16 @@ POST /api/v1/chat/stream
258
  **Bottom line.** Every Phase 1 file under `src/rag/`, `src/tools/`, `src/query/executors/`, `src/query/query_executor.py`, `src/query/base.py`, `src/api/v1/knowledge.py`, and `src/config/agents/` is gone. Phase 1 introspection helpers under `src/pipeline/db_pipeline/` and `src/database_client/` are still imported by Phase 2 — they were not rewritten, just wrapped. The three LLM call sites are now explicit and the SQL-writing one no longer exists; the planner emits a Pydantic-validated `QueryIR` instead.
259
 
260
  The one filename gotcha to remember: the **intent router** still lives at `src/agents/orchestration.py` as class `OrchestratorAgent` (Phase 1 name kept for import-site compatibility, Phase 2 body). The matching prompt and tests use the `intent_router` name, but the source module does not.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
52
  | `src/agents/chatbot.py` (Phase 1) → deleted, then `src/agents/answer_agent.py` (`AnswerAgent`) → renamed | `src/agents/chatbot.py::ChatbotAgent` | Final answer formation; streams via `astream` |
53
  | `src/knowledge/parquet_service.py` | `src/storage/parquet.py` | Parquet upload/download helper |
54
  | `src/pipeline/document_pipeline/document_pipeline.py` (folder) | `src/pipeline/document_pipeline.py` (flat) | Single module |
55
+ | `src/rag/retrievers/document.py` | `src/retrieval/document.py` | `DocumentRetriever` migrated; tabular file types filtered out of results. **Post-report update (mentor commit 61c746f, 2026-05-20):** rewritten to raw SQL (pgvector `<=>` cosine, `<+>` manhattan only) to dodge asyncpg type-mapping issues with the Go-ingested schema. MMR / euclidean / inner_product dropped. |
56
  | `src/rag/router.py` | `src/retrieval/router.py` | `RetrievalRouter`, Redis-cached, unstructured-only; dead `db: AsyncSession` + `source_hint` params removed |
57
  | `src/rag/base.py` (`RetrievalResult`, `BaseRetriever`) | `src/retrieval/base.py` | Same dataclass + ABC |
58
 
 
177
 
178
  ├── source_hint == "unstructured"
179
  │ → RetrievalRouter.retrieve() [retrieval/router.py, Redis-cached]
180
+ │ → DocumentRetriever (raw SQL: pgvector `<=>` cosine or `<+>` manhattan)
181
  │ → ChatbotAgent.astream(chunks=...)
182
 
183
  └── source_hint == "structured"
 
237
  | Top-level chat orchestration | Inline in `api/v1/chat.py` | `agents/chat_handler.py::ChatHandler` | Extracted to a reusable module |
238
  | Final answer formation | `agents/chatbot.py` (Phase 1 LangChain) | `agents/chatbot.py::ChatbotAgent` (was `AnswerAgent` in `answer_agent.py` mid-cycle) | Rewritten + renamed |
239
  | Schema retrieval (DB / tabular) | `rag/retrievers/schema.py` + PGVector chunks | **Removed**. Replaced by catalog (`catalog/store.py` jsonb) loaded verbatim into planner prompt | Whole concept replaced |
240
+ | Doc retrieval (PDF / DOCX / TXT) | `rag/retrievers/document.py`, `rag/router.py` | `retrieval/document.py`, `retrieval/router.py` | Moved; Redis cache restored; tabular files filtered. **Post-report update:** rewritten to raw SQL (cosine / manhattan only); collection renamed `document_embeddings` → `documents` to match the Go ingestion service. |
241
  | Query writing | `query/query_executor.py` + `models/sql_query.py` (LLM writes SQL) | `query/planner/service.py` (LLM writes IR) + `query/compiler/sql.py` (deterministic) | LLM emits intent, not SQL |
242
  | DB execution | `query/executors/db_executor.py` | `query/executor/db.py::DbExecutor` | Folder renamed (`executors` → `executor`); sqlglot guard + RO txn + 30 s timeout kept |
243
  | Tabular execution | `query/executors/tabular.py` | `query/executor/tabular.py::TabularExecutor` | Parquet-only; pandas compiler split out |
 
258
  **Bottom line.** Every Phase 1 file under `src/rag/`, `src/tools/`, `src/query/executors/`, `src/query/query_executor.py`, `src/query/base.py`, `src/api/v1/knowledge.py`, and `src/config/agents/` is gone. Phase 1 introspection helpers under `src/pipeline/db_pipeline/` and `src/database_client/` are still imported by Phase 2 — they were not rewritten, just wrapped. The three LLM call sites are now explicit and the SQL-writing one no longer exists; the planner emits a Pydantic-validated `QueryIR` instead.
259
 
260
  The one filename gotcha to remember: the **intent router** still lives at `src/agents/orchestration.py` as class `OrchestratorAgent` (Phase 1 name kept for import-site compatibility, Phase 2 body). The matching prompt and tests use the `intent_router` name, but the source module does not.
261
+
262
+ ---
263
+
264
+ ## 5. Addendum — post-report changes (2026-05-20, mentor commit `61c746f`)
265
+
266
+ This report was originally written as a snapshot at Phase 2 completion. The Phase 2 architecture itself hasn't changed, but a few implementation details have shifted as the Go migration progresses. Captured here so the report stays trustworthy:
267
+
268
+ - **Doc ingestion is now a Go service.** PDF/DOCX/TXT chunking + embedding + writes into PGVector are no longer done by Python. The Python service reads only.
269
+ - **PGVector collection renamed:** `document_embeddings` → `documents` (to match the Go service's writes). Touched files: `db/postgres/vector_store.py`, `retrieval/document.py`.
270
+ - **`DocumentRetriever` rewritten to raw SQL.** Uses pgvector operators directly (`<=>` cosine, `<+>` manhattan). The LangChain ORM path couldn't cope with the schema written by the Go service (asyncpg type-mapping issues — id String vs UUID, jsonb_path_match binding quirks). MMR / euclidean / inner_product were dropped as part of the rewrite.
271
+ - **Intent router defaults flipped.** Ambiguous topical/knowledge questions now prefer `unstructured` (was `structured`). Indonesian few-shot examples added to the prompt.
272
+ - **Cache management endpoints added:** `DELETE /api/v1/chat/cache`, `DELETE /api/v1/chat/cache/room/{id}`, `DELETE /api/v1/retrieval/cache/{user_id}`. Redis chat cache now stores `{response, sources}` (was just `response`) so cached replies repopulate `message_sources`.
273
+ - **Direction.** The long-term split is **Python = agent/ML layer, Go = data plane**. More pieces are expected to follow doc ingestion out of Python.
PROGRESS.md CHANGED
@@ -2,8 +2,280 @@
2
 
3
  Persistent tracker mirroring the 42-item ownership table in `REPO_CONTEXT.md` "Team — division of work". Update as PRs land. Future Claude Code sessions read this to know what's already done.
4
 
5
- **Last updated**: 2026-05-12 ([NOTICKET] Cleanup PR landed: ChatHandler wired to chat.py, Phase 1 dual-write dropped from /ingest, on_catalog_rebuild_requested implemented, dead modules deleted, answer_agent→chatbot renamed, retrieval cache restored via RetrievalRouter, top_values added to ColumnStats, lifespan migration, knowledge_router removed)
6
- **Current open PR**: `pr/1` — active. Cleanup PR committed and pushed.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
7
 
8
  ---
9
 
@@ -54,7 +326,7 @@ Persistent tracker mirroring the 42-item ownership table in `REPO_CONTEXT.md` "T
54
  | # | Item | Owner | Status | Notes |
55
  |---|---|---|---|---|
56
  | 5 | DB introspector (`catalog/introspect/database.py`) | DB | `[x]` | PR1 — reuses Phase 1 `database_client_service`, `db_credential_encryption`, `db_pipeline_service.engine_scope`, `extractor.get_schema/profile_column/get_row_count`. PR2a wired FK extraction (was discarded before). |
57
- | 6 | Tabular introspector (`catalog/introspect/tabular.py`) | TAB | `[~]` | PR1-tab — downloads original blob (CSV/XLSX/Parquet), one Table per sheet (XLSX) or one Table (CSV/Parquet). `source_id = document_id`. `fetch_doc`/`fetch_blob` injectable for unit tests (no Settings). |
58
  | 7 | `BaseIntrospector` ABC (`catalog/introspect/base.py`) | B | `[x]` | Pre-existing; signature locked |
59
 
60
  ### Ingestion — shared catalog plumbing
@@ -102,8 +374,8 @@ Persistent tracker mirroring the 42-item ownership table in `REPO_CONTEXT.md` "T
102
 
103
  | # | Item | Status | Notes |
104
  |---|---|---|---|
105
- | 29 | Pandas compiler (`query/compiler/pandas.py`) | `[~]` | PR3-TAB — `CompiledPandas` dataclass; all 12 filter ops; all 6 aggs; group_by via `pd.concat` of Series; alias-aware order_by; `_like_to_regex` (`%`→`.*`, `_`→`.`); pure module-level helpers |
106
- | 30 | Tabular executor (`query/executor/tabular.py`) | `[~]` | PR3-TAB — `fetch_blob` injectable for tests; blob path: single-table → `{uid}/{did}.parquet`, multi-table → `{uid}/{did}__{table.name}.parquet`; `asyncio.to_thread`; 10k row hard cap; errors → `QueryResult.error` |
107
  | 31 | Parquet upload/download wrapper | `[x]` | Moved `knowledge/parquet_service.py` → `storage/parquet.py`. Updated 4 import sites: `pipeline/document_pipeline.py`, `knowledge/processing_service.py`, `query/executor/tabular.py`, `query/executors/tabular.py`. |
108
 
109
  ### Agents + chat
@@ -112,7 +384,20 @@ Persistent tracker mirroring the 42-item ownership table in `REPO_CONTEXT.md` "T
112
  |---|---|---|---|
113
  | 32 | Chatbot agent + prompt (`agents/chatbot.py`, `config/prompts/chatbot_system.md`) | `[x]` | PR7-bundle — `ChatbotAgent` (was `AnswerAgent`) streams tokens, accepts `QueryResult` or list[`DocumentChunk`] or neither. **Cleanup PR**: renamed `answer_agent.py` → `chatbot.py`, `AnswerAgent` → `ChatbotAgent`; Phase 1 `agents/chatbot.py` deleted. |
114
  | 33 | Guardrails prompt (`config/prompts/guardrails.md`) | `[x]` | PR7-bundle — appended to `chatbot_system.md` so guardrails take precedence in conflict. |
115
- | — | Chat handler / orchestrator (`agents/chat_handler.py`) | `[x]` | PR4-bundle — top-level Phase 2 orchestrator. Routes by `source_hint`: chat → AnswerAgent direct; structured → CatalogReader + QueryService; unstructured → DocumentRetriever placeholder + AnswerAgent. Yields `intent` / `chunk` / `done` / `error` SSE-style events. Phase 1 chat.py NOT touched — cleanup PR rewires the API to call this. |
 
 
 
 
 
 
 
 
 
 
 
 
 
116
 
117
  ### API surface
118
 
@@ -129,7 +414,7 @@ Persistent tracker mirroring the 42-item ownership table in `REPO_CONTEXT.md` "T
129
  | # | Item | Owner | Status | Notes |
130
  |---|---|---|---|---|
131
  | 38 | DB compiler golden tests (`tests/query/compiler/test_sql.py`) | DB | `[x]` | PR3-DB — 36 tests across all whitelisted ops, identifier quoting, agg / count_distinct / count(*), order_by alias resolution, parameter sequencing, error paths. Pure-Python, no LLM, no DB. |
132
- | 39 | Pandas compiler golden tests (`tests/unit/query/compiler/test_pandas_compiler.py`) | TAB | `[~]` | PR3-TAB — 43 tests: all 12 filter ops, all 6 aggs, group_by, order_by, limit, aliases, empty DataFrame, error paths. `test_tabular_executor.py` adds 12 more (blob name resolution + happy path + error paths). |
133
  | 40 | IR validator tests (`tests/query/ir/test_validator.py`) | B | `[x]` | PR1 — 19 tests, all rules covered |
134
  | — | PII detector tests (`tests/catalog/test_pii_detector.py`) | B | `[x]` | PR1 — 26 tests (parametrized) |
135
  | — | Catalog validator tests (`tests/catalog/test_validator.py`) | B | `[x]` | PR1 — 5 tests |
 
2
 
3
  Persistent tracker mirroring the 42-item ownership table in `REPO_CONTEXT.md` "Team — division of work". Update as PRs land. Future Claude Code sessions read this to know what's already done.
4
 
5
+ **Last updated**: 2026-06-10 (tool layer complete + hardening/DRY + Langfuse tracing + gated slow-path wiring)
6
+ **Current open PR**: `pr/2` — active.
7
+
8
+ ---
9
+
10
+ ## Principal architecture review (2026-06-10) — findings + fix tracker
11
+
12
+ A full external review (read the context docs + the slow path, tool layer, query
13
+ spine, catalog plumbing, chat endpoint, config/connection layers) landed. It confirmed
14
+ the DB-latency diagnosis and surfaced several gaps **not previously tracked here**.
15
+ Verified against code before logging. Severity: **critical** / important / nice-to-have.
16
+
17
+ **Runtime / latency (the original problem):**
18
+ - DB connection handling is the anomaly, NOT cold start. `DbExecutor._run_sync`
19
+ (`db.py:192`) → `engine_scope` does `create_engine → connect (TCP+TLS+SCRAM) → 2×SET
20
+ → dispose` on EVERY query. Measured ~6–8s for 60 rows; a 2nd warm-session query was
21
+ still ~6.6s → per-call handshake, never amortized. `engine_scope`'s connect-once-dispose
22
+ semantics were designed for the ingestion pipeline and wrongly inherited by the query path.
23
+ - `describe_source` ~3.5s is **planner-induced waste**: every few-shot (`examples.py`)
24
+ opens with a `describe_source` task, so the LLM always plans a tool that re-reads from
25
+ the catalog DB the same catalog already rendered into its prompt. Its impl does 2
26
+ sequential full-catalog reads (`data_access.py:127-128`). Total catalog reads/request ~5×.
27
+ - Azure LLM clients rebuilt per request: `ChatHandler(enable_tracing=True)` is constructed
28
+ per request (`chat.py:172`) → fresh Orchestrator/Chatbot → fresh AzureChatOpenAI → fresh
29
+ TLS to Azure each call. Planner/Assembler correctly use module singletons; the other two don't.
30
+ - Tokens (~13k/request) are NORMAL for this design — do not optimize for $.
31
+ - **Reject the scheduled DB-warmer idea**: targets cold start (~1.8s slice) not the per-call
32
+ handshake, keeps serverless user DBs awake 24/7 (their compute bill), and decrypts every
33
+ tenant's creds on a cron (attack surface). Strictly dominated by an engine cache +
34
+ request-scoped pre-connect.
35
+
36
+ **Fix tracker (new):**
37
+
38
+ | # | Fix | Severity | Owner | Status |
39
+ |---|---|---|---|---|
40
+ | R1 | **AuthN/AuthZ** on data endpoints — reject body-supplied `user_id`/`room_id`, derive identity from a verified token. `/chat/stream` has none (`chat.py:40,128`); tenant isolation is client honesty. **CORRECTION to the review:** `security/auth.py` is a STUB (all `NotImplementedError`); the real JWT impl lives in `src/users/users.py` (`encode_jwt`/`decode_jwt`, HS, env-keyed) **but is unused** — `/login` (`api/v1/users.py`) returns the user profile as plain JSON and mints NO token. So R1 is cross-team: (1) `/login` must issue a JWT, (2) frontend must send it as `Bearer`, (3) data endpoints validate it. **Gates the engine-cache work (DB2).** | **critical** | DB/B + frontend | `[ ]` |
41
+ | R2 | **Always compile a LIMIT** — `sql.py` now emits a bound for every query: explicit limit honored (clamped to `MAX_RESULT_ROWS=10000`), unbounded queries get `LIMIT cap+1` so an unbounded SELECT can't stream a whole table into memory. `CompiledSql.row_cap` carries the cap; `DbExecutor` caps + flags truncation from it (dropped its own `_ROW_HARD_CAP`). Tests updated (`test_sql.py`, +3 cases); `S608` restored to `tests/**` ruff ignore (was dropped). | **critical** | DB | `[x]` |
42
+ | R3 | **Commit `tests/` + minimal CI** — `tests/` is gitignored; the 200+ tests cited as done exist only on laptops (already caused rename rot). GitHub origin carries tests; HF Space gets the Docker build (already doesn't COPY tests). | **critical (process)** | shared | `[ ]` |
43
+ | DB1 | **In-memory `describe_source`** (request-scoped `MemoizingCatalogReader`, `reader.py`) + **LLM-client hoist** (shared module-level `ChatHandler` in `chat.py`). Measured live: `describe_source` 3.5s→~2.0s (structured read now served from the planner's cached snapshot; only the unstructured read remains a round-trip), catalog reads/request ~5→~2. External `query_structured` handshake unchanged (DB2's job) so total slow path is ~flat until DB2. Tests: `tests/catalog/test_reader.py`. | important | agent | `[x]` |
44
+ | DB2 | **Keyed engine cache** — `src/database_client/engine.py::UserEngineCache` (process singleton): pooled engines keyed by `client_id + creds-hash` (rotation auto-invalidates), bounded LRU (50) + 600s idle TTL, `pool_pre_ping` + `pool_recycle=300`. `DbExecutor._run_sync` reuses the warm connection instead of `create_engine→connect→dispose` per query (postgres/supabase only; other db_types keep the legacy path — no regression). **Live-measured: warm `query_structured` 6.6–9.4s → ~2.5s** (the residual is the per-call catalog-DB client fetch + pre-ping, not the external handshake). **Finding:** Neon's transaction pooler REJECTS `default_transaction_read_only` as a libpq startup `option` — caught live; moved read-only + statement_timeout to a per-connection `connect` event (best-effort; authoritative read-only is the SELECT-only compiler + sqlglot guard, see R10). Per-request ownership/active check kept. Proceeded ahead of R1 per owner decision (marginal security delta over the existing no-auth state; auth tracked separately). Tests: `tests/database_client/test_engine.py`. First query/process still cold → DB3. | important | DB | `[x]` |
45
+ | DB3 | **Speculative pre-connect** — `DbExecutor.prewarm(catalog, user_id)` warms the pooled engine for schema sources (fire-and-forget at slow-path entry) so the cold first-query handshake overlaps the ~4s Planner call. Best-effort, never raises; gated to the default path (skipped when a coordinator factory is injected). Verified live through `ChatHandler.handle`. | nice-to-have | DB | `[x]` |
46
+ | R4 | **Per-stage progress events** — `SlowPathCoordinator.run` gained an optional `progress` callback; `ChatHandler` bridges it to SSE `status` events (`chat.py` forwards them). Live: stream now shows `Planning…`→`Running N steps…`→`Composing…` (max wire gap ~4.6s, was ~13s of silence) → fixes proxy idle-timeout + UX. **Deferred:** token-streaming the Assembler answer needs splitting it into a streamed prose call + a structured-record call — that doubles the Assembler LLM calls (cost/latency), so it's a separate decision; the answer is still emitted as one chunk after the (fast ~2.5s) Assembler. Test: `test_chat_handler_wiring.py`. | important | agent | `[~]` |
47
+ | R5 | **Response cache**: key on `user_id` + catalog version; invalidate on ingest. Today `chat:{room_id}:{message}`, 24h TTL, no user (`chat.py:138`) → cross-room replay + stale answers. | important | B | `[ ]` |
48
+ | R6 | **Hard time budget** — wrap `coordinator.run()` in `asyncio.wait_for` (60–90s). `Constraints.time_budget_seconds` is rendered but not enforced. | important | agent | `[ ]` |
49
+ | R7 | **Root-task-failure short-circuit** before the Assembler (templated/fast-path fallback, NOT replanning) — stops paying ~2k tok to narrate an empty RunState. | important | agent | `[ ]` |
50
+ | R8 | **Catalog upsert race** — per-user advisory lock around read-merge-upsert (`store.py`); concurrent uploads can drop a source. | important | DB | `[ ]` |
51
+ | R9 | **`extra="ignore"`** in `settings.py:15` (currently `allow` → typo'd env vars silently swallowed); require Azure keys in prod. | nice-to-have | B | `[ ]` |
52
+ | R10 | **Read-only enforcement is session-state, not a server role.** `REPO_CONTEXT.md` counts "read-only DB credentials" as a defense layer but nothing requests/verifies a read-only role. Either request read-only creds at registration (verify via `SELECT current_setting(...)`) or drop the claim. | important | DB | `[ ]` |
53
+ | R11 | **De-duplicate** `_PLACEHOLDER_RE` (`task_runner.py:31` vs validator) and `_DATA_ACCESS_TOOLS` (invoker vs planner registry) — import one from the other; comments aren't a sync mechanism. **TAB slice done (90e80f9):** canonical `DATA_ACCESS_TOOLS` now lives once in `tools/data_access.py`; `invoker.py` imports it (was a duplicated frozenset synced by comment). **Agent slice done (2026-06-10):** `PLACEHOLDER_RE` single-sourced in `planner/schemas.py` (part of the ToolCall placeholder convention); validator + task_runner import it. `planner/registry.py` keeps local spec *bodies* (stub pending KM-465 #4) but name-checks them against `DATA_ACCESS_TOOLS` in `_data_access_slice()` — upstream rename/add now raises at `default_registry()` instead of drifting silently. Registry output unchanged (same 12 tools, same order). | nice-to-have | agent/tool | `[x]` |
54
+ | R12 | **Doc/process hygiene** — some code docstrings cite internal design specs that are not committed to the repo (design docs are kept out of version control), so the references dangle for anyone but the author; `CLAUDE.md` lists deleted modules (enricher, `pipeline/orchestrator.py`); `main` is 38 commits behind on a dead architecture. | nice-to-have | agent | `[ ]` |
55
+ | R13 | **Pre-existing test failure** (found during R2, NOT caused by it): `tests/query/planner/test_prompt.py::test_render_catalog_with_sources` fails — `query/planner/prompt.py::render_catalog` now renders stable IDs (`src_test_db`) the test asserts are absent. Old query-planner path; confirmed failing on a clean tree. | nice-to-have | DB | `[ ]` |
56
+ | T1 | **`input_schema` is presence-only, not type-checked** — `ToolSpec.input_schema` comment said "validates ToolCall.args", but `TaskRunner._validate_args` only enforces `required` presence; the `properties` types are documentation, never validated at runtime. Clarified the contract in `tools/contracts.py` so nobody assumes type-safety (a wrong-typed arg passes validation, surfaces only inside the compute fn). Doc-only, no behavior change (90e80f9). | nice-to-have | TAB | `[x]` |
57
+ | T2 | **Dead Python embed path?** — `document_pipeline.process()` → `knowledge_processor` → `vector_store.aadd_documents()` still writes PDF/DOCX/TXT embeddings to `langchain_pg_embedding`, contradicting CLAUDE.md's "Go is sole writer, Python reads only". Verified the Go service (`Orchestrator-Agent-Service/internal/documents`) IS a complete ingestion writer to the same tables for all 5 file types (OCR + chunk + embed) → the Python embed branch is very likely redundant. **Blocked on one operational fact:** does the frontend still upload to `/document/process` (Python) or to Go? Park until confirmed — deleting a live ingestion path would break unstructured RAG. The csv/xlsx parquet branch stays regardless (feeds the catalog/tabular path). | nice-to-have | TAB | `[blocked]` |
58
+
59
+ **Slow-path endpoint wiring (2026-06-10):** the Orchestrator→slow-path is now wired
60
+ into the live endpoint behind an **env flag**. `settings.enable_slow_path` (env
61
+ `ENABLE_SLOW_PATH`, default **off**) is passed to the shared `ChatHandler` in
62
+ `api/v1/chat.py`. Flip `ENABLE_SLOW_PATH=true` to route `structured` intents through
63
+ Planner→TaskRunner→Assembler and test end-to-end from `/chat/stream` (status progress
64
+ events + answer stream). Stays opt-in because `BusinessContext` is still the stub;
65
+ fast/unstructured paths unchanged. Verified live via `ChatHandler.handle`.
66
+
67
+ **Architecture verdict:** fundamentally sound (catalog-driven IR + deterministic compiler
68
+ + static plan is the right call). Debt is transitional duplication (two planners/registries/
69
+ contract modules — documented, owned) and `ChatHandler` drifting toward a god object
70
+ (extract the slow-path composition root + the SSE `_build_sources`/`_normalize_chunks`
71
+ mappers when convenient).
72
+
73
+ ---
74
+
75
+ ## What just shipped (2026-06-09/10 — tool layer, tracing, slow-path wiring)
76
+
77
+ Big stretch since the slow-path workers landed. The tool layer (teammate-owned) is
78
+ now **complete and real**, the slow path is **wired into `ChatHandler` behind a gate**,
79
+ and the whole chat pipeline is **traced**. Fast path still untouched; live behavior
80
+ unchanged (flags default off).
81
+
82
+ **Tool layer — COMPLETE (teammate, KM-624→630).** `src/tools/` was re-created (the
83
+ 2026-05-11 note about deleting it is superseded). Now teammate-owned:
84
+ - `src/tools/analytics/` — the 8 **composite** `analyze_*` computes (descriptive,
85
+ aggregate, comparison, contribution, profile, correlation, segment, trend) +
86
+ prompt-style DESCRIPTIONs (KM-624/625).
87
+ - `src/tools/contracts.py` — canonical `ToolSpec`/`ToolRegistry`/`ToolOutput` (KM-627).
88
+ `agents/planner/contracts.py` now just re-exports them + keeps the `BusinessContext`
89
+ stub (lead's).
90
+ - `src/tools/registry.py::analytics_registry()` (KM-628); `src/tools/invoker.py` +
91
+ `src/tools/data_access.py` — `AnalyticsToolInvoker` (KM-629), `DataAccessToolInvoker`
92
+ + `CompositeToolInvoker` (KM-630). All never-throw. **Pattern A confirmed** (`analyze_*`
93
+ take a `data` `${t<id>}` placeholder from an upstream `query_structured`).
94
+ - **Verified live E2E (2026-06-09):** real `query_structured` against a user's Neon
95
+ Postgres → `analyze_trend` → Assembler. `analyze_contribution` surfaced a real tool
96
+ bug (Decimal vs float in `decomposition.py`) — degrade-and-continue held; **now fixed
97
+ by the tool owner** (`_coerce_decimals` in `invoker._materialize`, KM-630 / commit
98
+ 1195870), so the whole `analyze_*` family is covered in one place. **Directive:** agent
99
+ side does NOT modify `src/tools/` without confirmation.
100
+
101
+ **Planner — realigned to the real tools (KM-626).** `registry.py::default_registry()`
102
+ composes the real `analytics_registry()` + a local stub for the 4 data-access tools.
103
+ Few-shots grown to **A–D**: A `analyze_contribution`, B `analyze_trend`, C mixed
104
+ structured+unstructured (`retrieve_documents`, independent branch), D `analyze_aggregate`.
105
+ `parallelizable_with` **removed** from `Task` (schema/validator/examples/prompt) —
106
+ TaskRunner derives parallelism from `depends_on` alone.
107
+
108
+ **Slow-path wiring — built, GATED OFF (KM-626).** `agents/chat_handler.py` gains a
109
+ `structured→slow` branch behind `ChatHandler(enable_slow_path=False)`: when on it builds
110
+ a per-request `CompositeToolInvoker` (composition root) + `SlowPathCoordinator`, streams
111
+ `chat_answer`, persists the `analysis_record`. Two seams isolate the remaining blockers:
112
+ - `agents/planner/business_context.py::get_business_context(user_id)` — async stub
113
+ `BusinessContext`; TODO(lead) swap for the real read.
114
+ - `agents/slow_path/store.py` — `AnalysisStore` Protocol + `NullAnalysisStore` (logs
115
+ only). Real store = `analysis_records` table in the catalog DB (Neon `dataeyond`) —
116
+ **table not created yet**. `chat_answer` still emitted as one chunk (not token-streamed).
117
+
118
+ **Observability — Langfuse tracing wired (KM-631).** `src/observability/langfuse/
119
+ tracing.py` — `RequestTracer`/`NullTracer`/`TracingToolInvoker` + `_redact`. One trace
120
+ per request groups Orchestrator.classify, Planner.plan (each retry = its own generation),
121
+ Assembler.assemble, Chatbot.astream + tool spans (latency/metadata only). Gated:
122
+ `ChatHandler(enable_tracing=False)`; `api/v1/chat.py` opts in (`=True`). PII policy:
123
+ Orchestrator+Planner unmasked (question + PII-safe summary); Assembler+Chatbot masked
124
+ (see real rows/chunks); tool spans carry name + arg keys + row count only. Zero added
125
+ LLM tokens; verified live to US Cloud.
126
+
127
+ **Live evals green (2026-06-09, real Azure 4o):** `RUN_PLANNER_EVAL=1` and
128
+ `RUN_SLOW_PATH_EVAL=1` both pass — Planner emits valid catalog-consistent `QueryIR` and
129
+ wires Pattern A correctly; self-corrects via retry.
130
+
131
+ **Open follow-ups:** real `BusinessContext` (lead); create `analysis_records` table +
132
+ real `AnalysisStore`; register data-access `ToolSpec`s upstream (`data_access_registry()`)
133
+ or keep the planner stub; 4o → GPT-mini deployment swap; flip `enable_slow_path` on once
134
+ `BusinessContext` is real. NOTE: 3 test files pre-existing broken from rename rot
135
+ (`test_chat_handler.py`, `test_intent_router.py`, `test_answer_agent.py` import the old
136
+ `answer_agent`/`intent_router` module names).
137
+
138
+ ---
139
+
140
+ ## What just shipped (2026-06-10 — TAB: tool-layer hardening + DRY)
141
+
142
+ Owner-side companion to the agent block above. After the live E2E surfaced real-data
143
+ edge cases, the tool layer got a round of correctness hardening. All in TAB-owned paths
144
+ (`src/tools/`, `src/catalog/`); no agent-side or API change.
145
+
146
+ **JSON-safety across the `analyze_*` family.** Real DB rows carry scalar types that
147
+ don't survive the jsonb / SSE round-trip:
148
+ - `[KM-630] coerce DB Decimal → float` (commit 1195870) — `_coerce_decimals` in
149
+ `invoker._materialize` converts object-columns holding `decimal.Decimal` (asyncpg
150
+ returns NUMERIC as `Decimal`) to `float64` before any compute runs. Fixes the
151
+ `float + Decimal` TypeError in `decomposition.analyze_contribution` **and** the whole
152
+ family in one seam — only touches columns that actually contain a `Decimal`.
153
+ - `[KM-624] non-JSON-safe scalars in mode & top_value` (commit 6981ed3) — normalize
154
+ numpy / non-native scalars so descriptive + top-value outputs serialize cleanly.
155
+
156
+ **Planner↔Tools registry alignment + Timestamp keys** (commit 4bb7623, `fix(tools)`):
157
+ - `registry.py` — `analyze_descriptive.required` corrected `["data"]` → `["data",
158
+ "column_ids"]` to match the compute signature (`column_ids` has no default). Prevents
159
+ the Planner from emitting a call that's missing a required arg. `analyze_profile` stays
160
+ `["data"]` (its `column_ids` defaults to `None`).
161
+ - `aggregation._clean` — group-by over a datetime column produced `pd.Timestamp` group
162
+ keys that aren't JSON-safe; now normalized to `.isoformat()` alongside the existing
163
+ numpy `.item()` branch.
164
+
165
+ **DRY: single `SAMPLE_LIMIT` constant** (commit 6d46ba5, `[NOTICKET] refactor(catalog)`):
166
+ - One source of truth in `catalog/introspect/base.py` (`SAMPLE_LIMIT = 3`, down from 5 —
167
+ token cost: sample values feed the planner prompt). Both introspection paths import it:
168
+ `catalog/introspect/tabular.py` and `pipeline/db_pipeline/extractor.py` (which dropped
169
+ its own local `= 3`). Dependency direction is pipeline→catalog (no circular import).
170
+ Stale test `test_sample_values_capped_at_five` updated to assert the real cap (3).
171
+
172
+ **Audit result:** Planner↔Tools arg alignment swept end-to-end — 7/8 `analyze_*` tools
173
+ already matched; the 1 mismatch (`analyze_descriptive`) is the fix above. Pattern A holds
174
+ across all of them.
175
+
176
+ ---
177
+
178
+ ## What just shipped (2026-06-08 — KM-626: slow-path agent layer)
179
+
180
+ The rest of the slow path after the Planner (KM-567) — TaskRunner, Assembler, and
181
+ the coordinator. Built and tested against
182
+ mocks; **not yet wired into the live `ChatHandler`** (waits on the tool team's real
183
+ `ToolInvoker` + a real `BusinessContext`). Fast path untouched.
184
+
185
+ **Naming:** "Orchestrator" = the entry dispatcher only (`agents/orchestration.py`).
186
+ The slow-path **workers** live in **`agents/slow_path/`** — deliberately NOT named
187
+ "orchestrator".
188
+
189
+ **Files added** (`src/agents/slow_path/`):
190
+ - `schemas.py` — `TaskResult`, `RunState`; `TaskSummary`, `AnalysisRecord`,
191
+ `AssembledOutput`, `AssemblerNarrative`. Reuses `ToolOutput`.
192
+ - `invoker.py` — `ToolInvoker` Protocol only; the tool team owns the impl (KM-418).
193
+ - `errors.py` — `SlowPathError`, `AssemblerError`.
194
+ - `task_runner.py` — deterministic, 0 LLM: wave-based execution, `${t<id>}` placeholder
195
+ resolution, internal `validate_args`, never-throw invoke, status labeling,
196
+ degrade-and-continue → `RunState`.
197
+ - `assembler.py` + `prompt.py` + `config/prompts/assembler.md` — single LLM call →
198
+ `AssemblerNarrative`; code merges with `RunState` to build the `AnalysisRecord`
199
+ (structured fields copied, never re-authored).
200
+ - `coordinator.py` — `SlowPathCoordinator`: Planner → TaskRunner → Assembler.
201
+
202
+ **Tests added** (`tests/agents/slow_path/`, 12 passing; gitignored): schema round-trips
203
+ + chat_answer-first; runner happy/placeholder/parallel/degrade/arg-miss; assembler
204
+ narrative-vs-snapshot + question threading; coordinator end-to-end. `ruff` clean;
205
+ tool-agnostic (no `src/tools/*` import).
206
+
207
+ **Open follow-ups (not blockers):** wire `SlowPathCoordinator` into the expanded
208
+ Orchestrator/`ChatHandler` once the real invoker + `BusinessContext` exist; swap the
209
+ test `MockToolInvoker` for the tool team's real one (zero agent change, INV-7); 4o →
210
+ GPT-mini deployment swap.
211
+
212
+ ---
213
+
214
+ ## What just shipped (2026-06-08 — tool taxonomy + ownership revision)
215
+
216
+ Team decisions after the teammate pushed KM-624 (`src/tools/analytics/`):
217
+
218
+ - **Composite tools, not atomic.** v1 uses **composite "family" tools** (`analyze_*`),
219
+ not the atomic `compute_*` set the earlier draft assumed. One `analyze_*` call does a
220
+ whole analytical job (e.g. `analyze_descriptive` subsumes median/mode/stddev/percentile;
221
+ `analyze_trend` subsumes `date_trunc`). Tool-taxonomy decision recorded.
222
+ - **Tool team owns ALL tools** — compute, data-access (`query_structured`,
223
+ `retrieve_documents`, `list_sources`, `describe_source`), the wrapper/invoker layer
224
+ (KM-418), and **all tool tests**. The agent team owns nothing below the registry contract.
225
+ - **Planner stub realigned to the real tools.** `registry.py` rewritten from the 9 atomic
226
+ entries to **12 composite entries** (4 data-access + 8 `analyze_*`); `examples.py`
227
+ rewritten (Example A → `analyze_contribution`, Example B → `analyze_trend`); `planner.md`
228
+ bullet updated; planner tests updated. 32 passing + 1 gated, `ruff` clean.
229
+ - **Open (tool team's call):** Pattern A (analyze_* take a `${t<id>}` `data` placeholder
230
+ from an upstream `query_structured`) vs Pattern B (self-fetch by `source_id`). Stub
231
+ assumes A; reshaped to match once decided (agent code unaffected, INV-7).
232
+ - **New coupling:** the tool team's `query_structured`/`retrieve_documents` are expected
233
+ to call our existing `QueryService`/`RetrievalRouter`; `query_structured` stays
234
+ inline-`QueryIR` so `IRValidator` still applies. Interface to coordinate.
235
+
236
+ **Next (our scope, all mock-able now):** TaskRunner + Assembler against a `MockToolInvoker`,
237
+ then Orchestrator slow-path wiring. Stubs still to retire on integration: `contracts.py`
238
+ (BusinessContext from lead; ToolSpec/ToolRegistry/ToolOutput from tool team) and `registry.py`
239
+ (real registry from tool team). Infra: swap the 4o stand-in for a GPT-mini deployment.
240
+
241
+ ---
242
+
243
+ ## What just shipped (2026-06-05 — Phase 3: Planner agent)
244
+
245
+ First slow-path agent (the Planner). A single LLM
246
+ call turns BusinessContext + Catalog + ToolRegistry + question + Constraints into a
247
+ validated, **static** `TaskList` (DAG of fully-specified tool-call chains). No
248
+ replanning (INV-6); tool-agnostic against a registry contract (INV-7). Fast path
249
+ (`agents/orchestration.py`, `agents/chatbot.py`, `query/`) untouched.
250
+
251
+ **Files added** (`src/agents/planner/`):
252
+ - `contracts.py` — **STUB** Pydantic contracts pending reconciliation: `BusinessContext`
253
+ (+KeyTerm/DataTableNote/DataColumnNote, lead's), `ToolSpec`/`ToolRegistry` (tool
254
+ team KM-608), `ToolOutput` envelope.
255
+ - `schemas.py` — `CrispStage`, `ToolCall`, `Task`, `TaskList`. No replan schemas.
256
+ - `inputs.py` — `CatalogSummary` (condensed, PII `sample_values` nulled, `from_catalog`
257
+ builder + `render`) and `Constraints` (max_tasks=5, modeling_allowed=False).
258
+ - `registry.py` — **STUB** v1 P0 registry: query_structured, retrieve_documents,
259
+ list_sources, describe_source, compute_median/stddev/percentile/mode, date_trunc.
260
+ - `errors.py` — `PlannerError`, `PlannerValidationError`.
261
+ - `prompt.py` + `config/prompts/planner.md` — system prompt (INV-1/6/7 + principles) +
262
+ per-call human content (context + catalog + tools + constraints + few-shots + question).
263
+ - `examples.py` — two few-shots (A exploratory revenue-by-category; B descriptive
264
+ monthly-trend-by-region with date_trunc), built from the real `TaskList` schema.
265
+ - `validator.py` — `PlannerValidator` running the 8 checks; reuses the existing
266
+ `IRValidator` for inline `query_structured` IRs.
267
+ - `service.py` — `PlannerService` + `plan_analysis(...)`: chain (mirrors
268
+ `query/planner/service.py`) + validate-and-retry loop (max 3, mirrors `QueryService`).
269
+
270
+ **Tests added** (`tests/agents/planner/`, 30 passing + 1 gated): `test_schemas.py`,
271
+ `test_inputs.py`, `test_validator.py` (one failure per check + happy paths),
272
+ `test_service.py` (`_FakeChain` + retry), `test_golden_questions.py` (live eval gated on
273
+ `RUN_PLANNER_EVAL=1`). `ruff check` clean on planner paths.
274
+
275
+ **Open follow-ups (not blockers):** reconcile `BusinessContext` with the lead and
276
+ `ToolRegistry`/`ToolSpec` + real tools with teammate (KM-608); "GPT mini" currently uses
277
+ the configured 4o deployment (swap `azure_deployment` when a mini deployment exists). Next:
278
+ Orchestrator slow-path expansion + TaskRunner + Assembler.
279
 
280
  ---
281
 
 
326
  | # | Item | Owner | Status | Notes |
327
  |---|---|---|---|---|
328
  | 5 | DB introspector (`catalog/introspect/database.py`) | DB | `[x]` | PR1 — reuses Phase 1 `database_client_service`, `db_credential_encryption`, `db_pipeline_service.engine_scope`, `extractor.get_schema/profile_column/get_row_count`. PR2a wired FK extraction (was discarded before). |
329
+ | 6 | Tabular introspector (`catalog/introspect/tabular.py`) | TAB | `[x]` | PR1-tab — downloads original blob (CSV/XLSX/Parquet), one Table per sheet (XLSX) or one Table (CSV/Parquet). `source_id = document_id`. `fetch_doc`/`fetch_blob` injectable for unit tests (no Settings). **2026-06-10**: sample cap now imports the shared `SAMPLE_LIMIT` (=3) from `catalog/introspect/base.py` — single source of truth across the tabular + DB introspection paths (commit 6d46ba5). |
330
  | 7 | `BaseIntrospector` ABC (`catalog/introspect/base.py`) | B | `[x]` | Pre-existing; signature locked |
331
 
332
  ### Ingestion — shared catalog plumbing
 
374
 
375
  | # | Item | Status | Notes |
376
  |---|---|---|---|
377
+ | 29 | Pandas compiler (`query/compiler/pandas.py`) | `[x]` | PR3-TAB — `CompiledPandas` dataclass; all 12 filter ops; all 6 aggs; group_by via `pd.concat` of Series; alias-aware order_by; `_like_to_regex` (`%`→`.*`, `_`→`.`); pure module-level helpers. (`polars` for large files still deferred — see Planned dependencies.) |
378
+ | 30 | Tabular executor (`query/executor/tabular.py`) | `[x]` | PR3-TAB — `fetch_blob` injectable for tests; blob path: single-table → `{uid}/{did}.parquet`, multi-table → `{uid}/{did}__{table.name}.parquet`; `asyncio.to_thread`; 10k row hard cap; errors → `QueryResult.error`. Dispatcher routes to it by `source_type`. |
379
  | 31 | Parquet upload/download wrapper | `[x]` | Moved `knowledge/parquet_service.py` → `storage/parquet.py`. Updated 4 import sites: `pipeline/document_pipeline.py`, `knowledge/processing_service.py`, `query/executor/tabular.py`, `query/executors/tabular.py`. |
380
 
381
  ### Agents + chat
 
384
  |---|---|---|---|
385
  | 32 | Chatbot agent + prompt (`agents/chatbot.py`, `config/prompts/chatbot_system.md`) | `[x]` | PR7-bundle — `ChatbotAgent` (was `AnswerAgent`) streams tokens, accepts `QueryResult` or list[`DocumentChunk`] or neither. **Cleanup PR**: renamed `answer_agent.py` → `chatbot.py`, `AnswerAgent` → `ChatbotAgent`; Phase 1 `agents/chatbot.py` deleted. |
386
  | 33 | Guardrails prompt (`config/prompts/guardrails.md`) | `[x]` | PR7-bundle — appended to `chatbot_system.md` so guardrails take precedence in conflict. |
387
+ | — | Chat handler / orchestrator (`agents/chat_handler.py`) | `[x]` | PR4-bundle — top-level Phase 2 orchestrator. Routes by `source_hint`: chat → AnswerAgent direct; structured → CatalogReader + QueryService; unstructured → DocumentRetriever placeholder + AnswerAgent. Yields `intent` / `chunk` / `done` / `error` SSE-style events. Phase 1 chat.py NOT touched — cleanup PR rewires the API to call this. **2026-06-09**: gained the gated `structured→slow` branch (`enable_slow_path=False`) + `enable_tracing` (KM-626/631). |
388
+
389
+ ### Tools — slow-path "Tools" component (TAB)
390
+
391
+ New scope after the original 42-item table; added as the tool layer landed (KM-608/624–631). All TAB-owned (`src/tools/`), all never-throw.
392
+
393
+ | # | Item | Owner | Status | Notes |
394
+ |---|---|---|---|---|
395
+ | — | Analytics compute fns (`tools/analytics/`) | TAB | `[x]` | KM-608/624/625 — 8 **composite** `analyze_*` fns (descriptive, aggregate, comparison, contribution, profile, correlation, segment, trend) + prompt-style DESCRIPTIONs. Pure pandas, no I/O. JSON-safe outputs (numpy/Decimal/Timestamp normalized — KM-624 + commit 4bb7623). |
396
+ | — | Tool contracts (`tools/contracts.py`) | TAB | `[x]` | KM-627 — canonical `ToolSpec` / `ToolRegistry` / `ToolOutput`. `agents/planner/contracts.py` re-exports them (+ keeps the lead's `BusinessContext` stub). |
397
+ | — | Analytics registry (`tools/registry.py`) | TAB | `[x]` | KM-628 — `analytics_registry()`. `analyze_descriptive.required` = `["data","column_ids"]` (aligned to compute signature, commit 4bb7623). |
398
+ | — | Invoker layer (`tools/invoker.py`) | TAB | `[x]` | KM-629 — `AnalyticsToolInvoker` (Pattern A: `analyze_*` take a `data` `${t<id>}` placeholder from upstream `query_structured`; `_materialize` → DataFrame, `_coerce_decimals` covers the whole family) + `CompositeToolInvoker` (routes data-access vs analytics by name). |
399
+ | — | Data-access tools (`tools/data_access.py`) | TAB | `[x]` | KM-630 — `DataAccessToolInvoker`: `list_sources` / `describe_source` / `query_structured` / `retrieve_documents`. Per-request DI (`user_id` + `CatalogReader`). `query_structured` calls `IRValidator` + `ExecutorDispatcher` (planner skipped — IR pre-built by the agent Planner). |
400
+ | — | Tool tests (`tests/unit/tools/`) | TAB | `[x]` | analytics + data-access + invoker tests (gitignored). Incl. regression `test_decimal_columns_coerced_for_analyze_contribution`. |
401
 
402
  ### API surface
403
 
 
414
  | # | Item | Owner | Status | Notes |
415
  |---|---|---|---|---|
416
  | 38 | DB compiler golden tests (`tests/query/compiler/test_sql.py`) | DB | `[x]` | PR3-DB — 36 tests across all whitelisted ops, identifier quoting, agg / count_distinct / count(*), order_by alias resolution, parameter sequencing, error paths. Pure-Python, no LLM, no DB. |
417
+ | 39 | Pandas compiler golden tests (`tests/unit/query/compiler/test_pandas_compiler.py`) | TAB | `[x]` | PR3-TAB — 43 tests: all 12 filter ops, all 6 aggs, group_by, order_by, limit, aliases, empty DataFrame, error paths. `test_tabular_executor.py` adds 12 more (blob name resolution + happy path + error paths). |
418
  | 40 | IR validator tests (`tests/query/ir/test_validator.py`) | B | `[x]` | PR1 — 19 tests, all rules covered |
419
  | — | PII detector tests (`tests/catalog/test_pii_detector.py`) | B | `[x]` | PR1 — 26 tests (parametrized) |
420
  | — | Catalog validator tests (`tests/catalog/test_validator.py`) | B | `[x]` | PR1 — 5 tests |
REPO_CONTEXT.md CHANGED
@@ -156,7 +156,7 @@ makes any LLM calls.)
156
  | `db/postgres/connection.py` | two async engines: `engine` (app) and `_pgvector_engine` (PGVector) |
157
  | `db/postgres/init_db.py` | startup: creates `vector` extension, all tables, HNSW + GIN indexes |
158
  | `db/postgres/models.py` | SQLAlchemy app tables (users, rooms, chat messages, …) |
159
- | `db/postgres/vector_store.py` | shared PGVector instance (collection `document_embeddings`) |
160
  | `db/redis/connection.py` | async Redis client |
161
  | `storage/az_blob/az_blob.py` | Azure Blob async wrapper (uploads + Parquet) |
162
  | `middlewares/{cors,logging,rate_limit}.py` | CORS allow-all (POC), structlog JSON, slowapi |
@@ -318,7 +318,7 @@ Single-table only in v1. `having`, `offset`, boolean filter trees, `distinct`, j
318
  | QueryService | ✅ | plan → validate → retry-on-fail (max 3) → dispatch → execute → `QueryResult` |
319
  | `ChatbotAgent` + prompt + guardrails | ✅ | Renamed from `AnswerAgent` in Cleanup PR. Guardrails appended to `chatbot_system.md` |
320
  | `ChatHandler` (top-level chat orchestrator) | ✅ | SSE events: `intent` / `chunk` / `done` / `error` |
321
- | `DocumentRetriever` + `RetrievalRouter` (Redis-cached) | ✅ | Migrated from `src/rag/` (now deleted); MMR/cosine/euclidean/manhattan/inner_product |
322
  | `/api/v1/chat/stream` | ✅ | Rewired to `ChatHandler`; Redis cache + fast intent + history + message persistence remain in chat.py |
323
  | `/api/v1/db-clients/{id}/ingest` | ✅ | Calls only `on_db_registered`; Phase 1 dual-write removed |
324
  | `/api/v1/document/{upload,process,delete}` | ✅ | `/process` triggers `on_tabular_uploaded` for CSV/XLSX |
 
156
  | `db/postgres/connection.py` | two async engines: `engine` (app) and `_pgvector_engine` (PGVector) |
157
  | `db/postgres/init_db.py` | startup: creates `vector` extension, all tables, HNSW + GIN indexes |
158
  | `db/postgres/models.py` | SQLAlchemy app tables (users, rooms, chat messages, …) |
159
+ | `db/postgres/vector_store.py` | shared PGVector instance (collection `documents` — written by Go ingestion service) |
160
  | `db/redis/connection.py` | async Redis client |
161
  | `storage/az_blob/az_blob.py` | Azure Blob async wrapper (uploads + Parquet) |
162
  | `middlewares/{cors,logging,rate_limit}.py` | CORS allow-all (POC), structlog JSON, slowapi |
 
318
  | QueryService | ✅ | plan → validate → retry-on-fail (max 3) → dispatch → execute → `QueryResult` |
319
  | `ChatbotAgent` + prompt + guardrails | ✅ | Renamed from `AnswerAgent` in Cleanup PR. Guardrails appended to `chatbot_system.md` |
320
  | `ChatHandler` (top-level chat orchestrator) | ✅ | SSE events: `intent` / `chunk` / `done` / `error` |
321
+ | `DocumentRetriever` + `RetrievalRouter` (Redis-cached) | ✅ | Migrated from `src/rag/` (now deleted). Mentor commit `61c746f` rewrote to raw SQL (pgvector `<=>` cosine, `<+>` manhattan) to dodge asyncpg type-mapping issues with Go-ingested schema. Methods reduced to `cosine | manhattan`. Collection: `documents`. |
322
  | `/api/v1/chat/stream` | ✅ | Rewired to `ChatHandler`; Redis cache + fast intent + history + message persistence remain in chat.py |
323
  | `/api/v1/db-clients/{id}/ingest` | ✅ | Calls only `on_db_registered`; Phase 1 dual-write removed |
324
  | `/api/v1/document/{upload,process,delete}` | ✅ | `/process` triggers `on_tabular_uploaded` for CSV/XLSX |
pyproject.toml CHANGED
@@ -120,7 +120,9 @@ ignore = [
120
  ]
121
 
122
  [tool.ruff.lint.per-file-ignores]
123
- "tests/**" = ["S101", "S105", "S106"]
 
 
124
 
125
  [tool.mypy]
126
  python_version = "3.12"
 
120
  ]
121
 
122
  [tool.ruff.lint.per-file-ignores]
123
+ # S608: golden compiler tests assert literal SQL strings (incl. concatenated
124
+ # suffixes) — they never execute against a DB, so it's a false positive here.
125
+ "tests/**" = ["S101", "S105", "S106", "S608"]
126
 
127
  [tool.mypy]
128
  python_version = "3.12"
src/agents/chat_handler.py CHANGED
@@ -22,8 +22,9 @@ inject mocks).
22
 
23
  from __future__ import annotations
24
 
 
25
  import json
26
- from collections.abc import AsyncIterator
27
  from typing import TYPE_CHECKING, Any
28
 
29
  from langchain_core.messages import BaseMessage
@@ -38,6 +39,8 @@ if TYPE_CHECKING:
38
  from ..catalog.reader import CatalogReader
39
  from ..query.service import QueryService
40
  from ..retrieval.router import RetrievalRouter
 
 
41
 
42
  logger = get_logger("chat_handler")
43
 
@@ -62,12 +65,29 @@ class ChatHandler:
62
  catalog_reader: CatalogReader | None = None,
63
  query_service: QueryService | None = None,
64
  document_retriever: RetrievalRouter | None = None,
 
 
 
 
 
 
 
65
  ) -> None:
66
  self._intent_router = intent_router
67
  self._answer_agent = answer_agent
68
  self._catalog_reader = catalog_reader
69
  self._query_service = query_service
70
  self._document_retriever = document_retriever
 
 
 
 
 
 
 
 
 
 
71
 
72
  # ------------------------------------------------------------------
73
  # Lazy default-dep builders
@@ -115,9 +135,13 @@ class ChatHandler:
115
  user_id: str,
116
  history: list[BaseMessage] | None = None,
117
  ) -> AsyncIterator[dict[str, Any]]:
 
 
118
  # ---- 1. Classify intent --------------------------------------
119
  try:
120
- decision = await self._get_intent_router().classify(message, history)
 
 
121
  except Exception as e:
122
  logger.error("intent classification failed", error=str(e))
123
  yield {"event": "error", "data": f"Could not classify message: {e}"}
@@ -133,7 +157,20 @@ class ChatHandler:
133
  # ---- 2. Route ------------------------------------------------
134
  if decision.source_hint == "structured":
135
  try:
136
- catalog = await self._get_catalog_reader().read(user_id, "structured")
 
 
 
 
 
 
 
 
 
 
 
 
 
137
  query_result = await self._get_query_service().run(
138
  user_id, rewritten, catalog
139
  )
@@ -174,12 +211,16 @@ class ChatHandler:
174
  yield {"event": "sources", "data": json.dumps(sources)}
175
 
176
  # ---- 3. Stream answer ----------------------------------------
 
 
 
177
  try:
178
  async for token in self._get_answer_agent().astream(
179
  message,
180
  history=history,
181
  query_result=query_result,
182
  chunks=chunks,
 
183
  ):
184
  yield {"event": "chunk", "data": token}
185
  except Exception as e:
@@ -187,6 +228,150 @@ class ChatHandler:
187
  yield {"event": "error", "data": f"Answer generation failed: {e}"}
188
  return
189
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
190
  yield {"event": "done", "data": ""}
191
 
192
 
 
22
 
23
  from __future__ import annotations
24
 
25
+ import asyncio
26
  import json
27
+ from collections.abc import AsyncIterator, Callable
28
  from typing import TYPE_CHECKING, Any
29
 
30
  from langchain_core.messages import BaseMessage
 
39
  from ..catalog.reader import CatalogReader
40
  from ..query.service import QueryService
41
  from ..retrieval.router import RetrievalRouter
42
+ from .slow_path.coordinator import SlowPathCoordinator
43
+ from .slow_path.store import AnalysisStore
44
 
45
  logger = get_logger("chat_handler")
46
 
 
65
  catalog_reader: CatalogReader | None = None,
66
  query_service: QueryService | None = None,
67
  document_retriever: RetrievalRouter | None = None,
68
+ *,
69
+ enable_slow_path: bool = False,
70
+ slow_path_coordinator_factory: (
71
+ Callable[[str], SlowPathCoordinator] | None
72
+ ) = None,
73
+ analysis_store: AnalysisStore | None = None,
74
+ enable_tracing: bool = False,
75
  ) -> None:
76
  self._intent_router = intent_router
77
  self._answer_agent = answer_agent
78
  self._catalog_reader = catalog_reader
79
  self._query_service = query_service
80
  self._document_retriever = document_retriever
81
+ # Langfuse tracing (tokens + latency). OFF by default so tests never hit
82
+ # Langfuse; the live endpoint opts in with ChatHandler(enable_tracing=True).
83
+ self._enable_tracing = enable_tracing
84
+ # Slow analytical path (Planner -> TaskRunner -> Assembler). OFF by default:
85
+ # gated until the lead's real BusinessContext lands. When True, `structured`
86
+ # intents route here instead of the single-query QueryService path. The
87
+ # factory + store are injectable for tests.
88
+ self._enable_slow_path = enable_slow_path
89
+ self._slow_path_factory = slow_path_coordinator_factory
90
+ self._analysis_store = analysis_store
91
 
92
  # ------------------------------------------------------------------
93
  # Lazy default-dep builders
 
135
  user_id: str,
136
  history: list[BaseMessage] | None = None,
137
  ) -> AsyncIterator[dict[str, Any]]:
138
+ tracer = self._make_tracer(user_id, message)
139
+
140
  # ---- 1. Classify intent --------------------------------------
141
  try:
142
+ oc = tracer.callbacks() # orchestrator: PII-safe, full capture
143
+ ckw = {"callbacks": oc} if oc else {}
144
+ decision = await self._get_intent_router().classify(message, history, **ckw)
145
  except Exception as e:
146
  logger.error("intent classification failed", error=str(e))
147
  yield {"event": "error", "data": f"Could not classify message: {e}"}
 
157
  # ---- 2. Route ------------------------------------------------
158
  if decision.source_hint == "structured":
159
  try:
160
+ # One memoizing reader per request: the same catalog is otherwise
161
+ # re-fetched from the catalog DB 4-5x across the slow-path run. This
162
+ # collapses those to one round-trip per source_hint and pins a single
163
+ # consistent snapshot for plan + execution.
164
+ from ..catalog.reader import MemoizingCatalogReader
165
+
166
+ req_reader = MemoizingCatalogReader(self._get_catalog_reader())
167
+ catalog = await req_reader.read(user_id, "structured")
168
+ if self._enable_slow_path:
169
+ async for event in self._run_slow_path(
170
+ user_id, rewritten, catalog, tracer, req_reader
171
+ ):
172
+ yield event
173
+ return
174
  query_result = await self._get_query_service().run(
175
  user_id, rewritten, catalog
176
  )
 
211
  yield {"event": "sources", "data": json.dumps(sources)}
212
 
213
  # ---- 3. Stream answer ----------------------------------------
214
+ # masked: the answer call sees real query rows / doc chunks (possible PII).
215
+ mc = tracer.callbacks(masked=True)
216
+ akw = {"callbacks": mc} if mc else {}
217
  try:
218
  async for token in self._get_answer_agent().astream(
219
  message,
220
  history=history,
221
  query_result=query_result,
222
  chunks=chunks,
223
+ **akw,
224
  ):
225
  yield {"event": "chunk", "data": token}
226
  except Exception as e:
 
228
  yield {"event": "error", "data": f"Answer generation failed: {e}"}
229
  return
230
 
231
+ tracer.end()
232
+ yield {"event": "done", "data": ""}
233
+
234
+ # ------------------------------------------------------------------
235
+ # Slow analytical path (gated, off by default)
236
+ # ------------------------------------------------------------------
237
+
238
+ def _make_tracer(self, user_id: str, question: str) -> Any:
239
+ """One Langfuse trace per request (or a NullTracer when disabled)."""
240
+ if not self._enable_tracing:
241
+ from ..observability.langfuse.tracing import NullTracer
242
+
243
+ return NullTracer()
244
+ from ..observability.langfuse.tracing import RequestTracer
245
+
246
+ return RequestTracer.start(user_id=user_id, question=question)
247
+
248
+ def _get_slow_path_coordinator(
249
+ self, user_id: str, tracer: Any = None, catalog_reader: CatalogReader | None = None
250
+ ) -> SlowPathCoordinator:
251
+ """Build the per-request slow-path coordinator (composition root).
252
+
253
+ The data-access tools need the authenticated `user_id` + `CatalogReader`,
254
+ so the `CompositeToolInvoker` is constructed per request. The slow-path
255
+ agent code stays tool-agnostic (INV-7) — only here, the composition root,
256
+ do we name concrete tool implementations. When tracing is active the invoker
257
+ is wrapped so each tool call records a metadata-only span.
258
+ """
259
+ if self._slow_path_factory is not None:
260
+ return self._slow_path_factory(user_id)
261
+
262
+ from ..tools.data_access import DataAccessToolInvoker
263
+ from ..tools.invoker import AnalyticsToolInvoker, CompositeToolInvoker
264
+ from .planner.registry import default_registry
265
+ from .planner.service import PlannerService
266
+ from .slow_path.assembler import Assembler
267
+ from .slow_path.coordinator import SlowPathCoordinator
268
+ from .slow_path.task_runner import TaskRunner
269
+
270
+ invoker: Any = CompositeToolInvoker(
271
+ DataAccessToolInvoker(user_id, catalog_reader or self._get_catalog_reader()),
272
+ AnalyticsToolInvoker(),
273
+ )
274
+ if tracer is not None and getattr(tracer, "active", False):
275
+ from ..observability.langfuse.tracing import TracingToolInvoker
276
+
277
+ invoker = TracingToolInvoker(invoker, tracer)
278
+ registry = default_registry()
279
+ return SlowPathCoordinator(
280
+ PlannerService(), TaskRunner(invoker, registry), Assembler(), registry
281
+ )
282
+
283
+ def _get_analysis_store(self) -> AnalysisStore:
284
+ if self._analysis_store is None:
285
+ from .slow_path.store import NullAnalysisStore
286
+
287
+ self._analysis_store = NullAnalysisStore()
288
+ return self._analysis_store
289
+
290
+ async def _run_slow_path(
291
+ self,
292
+ user_id: str,
293
+ query: str,
294
+ catalog: Any,
295
+ tracer: Any = None,
296
+ catalog_reader: CatalogReader | None = None,
297
+ ) -> AsyncIterator[dict[str, Any]]:
298
+ """Run the slow path and stream its assembled answer as SSE events.
299
+
300
+ Context comes from the `get_business_context` seam (a stub today); the
301
+ `analysis_record` is persisted via the `AnalysisStore` seam (a no-op today).
302
+ `chat_answer` is emitted as a single `chunk` (the Assembler returns the whole
303
+ object — true token streaming is a later step).
304
+ """
305
+ from .planner.business_context import get_business_context
306
+ from .planner.inputs import Constraints
307
+
308
+ if tracer is None:
309
+ from ..observability.langfuse.tracing import NullTracer
310
+
311
+ tracer = NullTracer()
312
+
313
+ coordinator = self._get_slow_path_coordinator(user_id, tracer, catalog_reader)
314
+ context = await get_business_context(user_id)
315
+
316
+ # DB3: warm the user's DB connection in parallel with planning so the
317
+ # handshake overlaps the ~4s Planner call. Default path only — an injected
318
+ # coordinator factory (tests / custom) may not use the real DbExecutor.
319
+ if self._slow_path_factory is None:
320
+ from ..query.executor.db import DbExecutor
321
+
322
+ asyncio.create_task(DbExecutor.prewarm(catalog, user_id)) # noqa: RUF006
323
+
324
+ pc = tracer.callbacks() # planner: PII-safe, full capture
325
+ ac = tracer.callbacks(masked=True) # assembler: sees real rows -> masked
326
+ run_kw: dict[str, Any] = {}
327
+ if pc:
328
+ run_kw["planner_callbacks"] = pc
329
+ if ac:
330
+ run_kw["assembler_callbacks"] = ac
331
+
332
+ # R4: bridge the coordinator's per-stage progress callback to SSE `status`
333
+ # events so the stream isn't silent for ~12s (and proxies don't drop the
334
+ # idle connection). Status events only appear if the coordinator calls back.
335
+ progress_q: asyncio.Queue[str] = asyncio.Queue()
336
+
337
+ async def _progress(stage: str) -> None:
338
+ await progress_q.put(stage)
339
+
340
+ run_task = asyncio.create_task(
341
+ coordinator.run(
342
+ context, catalog, query, Constraints(), progress=_progress, **run_kw
343
+ )
344
+ )
345
+ getter: asyncio.Task = asyncio.create_task(progress_q.get())
346
+ pending: set[asyncio.Task] = {run_task, getter}
347
+ while True:
348
+ done, pending = await asyncio.wait(
349
+ pending, return_when=asyncio.FIRST_COMPLETED
350
+ )
351
+ if getter in done:
352
+ yield {"event": "status", "data": getter.result()}
353
+ getter = asyncio.create_task(progress_q.get())
354
+ pending = pending | {getter}
355
+ if run_task in done:
356
+ getter.cancel()
357
+ while not progress_q.empty():
358
+ yield {"event": "status", "data": progress_q.get_nowait()}
359
+ break
360
+
361
+ try:
362
+ result = run_task.result()
363
+ except Exception as e:
364
+ logger.error("slow path failed", user_id=user_id, error=str(e))
365
+ yield {"event": "error", "data": f"Analysis failed: {e}"}
366
+ return
367
+
368
+ yield {"event": "sources", "data": json.dumps([])} # TODO: derive from record
369
+ yield {"event": "chunk", "data": result.chat_answer}
370
+ try:
371
+ await self._get_analysis_store().save(result.analysis_record)
372
+ except Exception as e: # persistence must never break the user's answer
373
+ logger.error("analysis_record persist failed", user_id=user_id, error=str(e))
374
+ tracer.end() # output omitted (chat_answer may contain PII on Cloud)
375
  yield {"event": "done", "data": ""}
376
 
377
 
src/agents/chatbot.py CHANGED
@@ -119,6 +119,9 @@ def _build_default_chain() -> Runnable:
119
  azure_endpoint=settings.azureai_endpoint_url_4o,
120
  api_key=settings.azureai_api_key_4o,
121
  temperature=0.3,
 
 
 
122
  )
123
  prompt = ChatPromptTemplate.from_messages(
124
  [
@@ -153,6 +156,7 @@ class ChatbotAgent:
153
  history: list[BaseMessage] | None = None,
154
  query_result: QueryResult | None = None,
155
  chunks: list[DocumentChunk] | None = None,
 
156
  ) -> AsyncIterator[str]:
157
  """Stream tokens of the final answer.
158
 
@@ -165,5 +169,9 @@ class ChatbotAgent:
165
  "history": history or [],
166
  "context": _build_context_block(query_result, chunks),
167
  }
168
- async for token in chain.astream(payload):
169
- yield token
 
 
 
 
 
119
  azure_endpoint=settings.azureai_endpoint_url_4o,
120
  api_key=settings.azureai_api_key_4o,
121
  temperature=0.3,
122
+ # Emit token usage on the final streamed chunk (this agent only streams), so
123
+ # the fast-path answer reports tokens to Langfuse like the non-streaming calls.
124
+ model_kwargs={"stream_options": {"include_usage": True}},
125
  )
126
  prompt = ChatPromptTemplate.from_messages(
127
  [
 
156
  history: list[BaseMessage] | None = None,
157
  query_result: QueryResult | None = None,
158
  chunks: list[DocumentChunk] | None = None,
159
+ callbacks: list | None = None,
160
  ) -> AsyncIterator[str]:
161
  """Stream tokens of the final answer.
162
 
 
169
  "history": history or [],
170
  "context": _build_context_block(query_result, chunks),
171
  }
172
+ if callbacks:
173
+ async for token in chain.astream(payload, config={"callbacks": callbacks}):
174
+ yield token
175
+ else:
176
+ async for token in chain.astream(payload):
177
+ yield token
src/agents/orchestration.py CHANGED
@@ -96,11 +96,16 @@ class OrchestratorAgent:
96
  self,
97
  message: str,
98
  history: list[BaseMessage] | None = None,
 
99
  ) -> IntentRouterDecision:
100
  chain = self._ensure_chain()
101
- decision: IntentRouterDecision = await chain.ainvoke(
102
- {"message": message, "history": history or []}
103
- )
 
 
 
 
104
  logger.info(
105
  "intent classified",
106
  source_hint=decision.source_hint,
 
96
  self,
97
  message: str,
98
  history: list[BaseMessage] | None = None,
99
+ callbacks: list | None = None,
100
  ) -> IntentRouterDecision:
101
  chain = self._ensure_chain()
102
+ payload = {"message": message, "history": history or []}
103
+ if callbacks:
104
+ decision: IntentRouterDecision = await chain.ainvoke(
105
+ payload, config={"callbacks": callbacks}
106
+ )
107
+ else:
108
+ decision = await chain.ainvoke(payload)
109
  logger.info(
110
  "intent classified",
111
  source_hint=decision.source_hint,
src/agents/planner/__init__.py ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ """Planner agent — slow-path CRISP-DM analysis planner.
2
+
3
+ Single LLM call: BusinessContext + CatalogSummary + ToolRegistry + question +
4
+ Constraints -> a validated, static `TaskList` (a DAG of fully-specified
5
+ tool-call chains). No replanning (INV-6); tool-agnostic (INV-7).
6
+
7
+ See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3.
8
+ """
src/agents/planner/business_context.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """BusinessContext reader — the single seam the slow path reads context through.
2
+
3
+ `get_business_context(user_id)` is the one place the live flow obtains a
4
+ `BusinessContext`. Today it returns a minimal STUB so the slow path runs end to
5
+ end; when the lead's real Business Understanding source lands, swap the body here
6
+ (read the interview / stored context) and nothing upstream changes.
7
+
8
+ See AGENT_ARCHITECTURE_CONTEXT_new.md §7.1.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ from .contracts import BusinessContext
14
+
15
+
16
+ async def get_business_context(user_id: str) -> BusinessContext:
17
+ """Return the user's BusinessContext.
18
+
19
+ STUB until the lead's real source lands. `project_id` flows through as
20
+ `RunState.business_context_id`. Async so the real implementation (a DB / store
21
+ read) fits without changing this signature.
22
+
23
+ TODO(lead): replace the body with the real read (Business Understanding store).
24
+ """
25
+ return BusinessContext(
26
+ project_id=user_id,
27
+ industry="unknown",
28
+ completeness="partial",
29
+ business_description="(not yet captured — BusinessContext source pending)",
30
+ scale_and_scope="(unknown)",
31
+ )
src/agents/planner/contracts.py ADDED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Contracts the planner consumes from other teams.
2
+
3
+ `BusinessContext` (+ KeyTerm / DataTableNote / DataColumnNote) is still a LOCAL
4
+ STUB owned by the lead (interview / Business Understanding); shape mirrors
5
+ AGENT_ARCHITECTURE_CONTEXT_new.md §7.1 and must be reconciled before integration.
6
+
7
+ `ToolSpec` / `ToolRegistry` / `ToolOutput` are NO LONGER defined here — the tool
8
+ team owns them (KM-465). They now live in `src.tools.contracts` and are
9
+ re-exported below so existing importers (`registry.py`, the slow-path
10
+ `invoker.py` / `schemas.py`) keep working against one shared definition.
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ from typing import Literal
16
+
17
+ from pydantic import BaseModel, Field
18
+
19
+ # Canonical tool contracts now owned by the tool team (KM-465). Re-exported here
20
+ # for backwards-compatible imports; the definitions live in src.tools.contracts.
21
+ from src.tools.contracts import ToolOutput, ToolRegistry, ToolSpec
22
+
23
+ __all__ = [
24
+ "KeyTerm",
25
+ "DataTableNote",
26
+ "DataColumnNote",
27
+ "BusinessContext",
28
+ "ToolSpec",
29
+ "ToolRegistry",
30
+ "ToolOutput",
31
+ ]
32
+
33
+ # --------------------------------------------------------------------------- #
34
+ # BusinessContext (lead's contract — §7.1)
35
+ # --------------------------------------------------------------------------- #
36
+
37
+
38
+ class KeyTerm(BaseModel):
39
+ term: str
40
+ meaning: str
41
+
42
+
43
+ class DataTableNote(BaseModel):
44
+ table_name: str
45
+ row_represents: str
46
+
47
+
48
+ class DataColumnNote(BaseModel):
49
+ column_name: str
50
+ meaning: str
51
+
52
+
53
+ class BusinessContext(BaseModel):
54
+ project_id: str
55
+ industry: str
56
+ completeness: Literal["partial", "complete"]
57
+ business_description: str
58
+ scale_and_scope: str
59
+ key_terms: list[KeyTerm] = Field(default_factory=list)
60
+ data_overview: list[DataTableNote] = Field(default_factory=list)
61
+ data_column_notes: list[DataColumnNote] = Field(default_factory=list)
62
+ whats_normal: str = ""
63
+ recent_events: str = ""
64
+ things_to_watch_for: str = ""
65
+ open_questions: list[str] = Field(default_factory=list)
src/agents/planner/errors.py ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Typed errors for the planner agent."""
2
+
3
+ from __future__ import annotations
4
+
5
+
6
+ class PlannerError(Exception):
7
+ """Base error for the planner agent."""
8
+
9
+
10
+ class PlannerValidationError(PlannerError):
11
+ """A TaskList failed one of the planner validator's checks.
12
+
13
+ The message is specific enough that the planner can be re-prompted with it
14
+ to self-correct (max 3 attempts, see service.py).
15
+ """
src/agents/planner/examples.py ADDED
@@ -0,0 +1,384 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Few-shot examples for the planner prompt.
2
+
3
+ Two illustrative (question -> TaskList) pairs that teach the OUTPUT SHAPE:
4
+ stages, dependency edges, ordered tool-call chains, inline QueryIR,
5
+ "${t<id>}" placeholders, and the assumed data-flow convention — `query_structured`
6
+ pulls rows, then a composite `analyze_*` tool consumes them via a `data` placeholder
7
+ referencing the upstream result's column aliases (Pattern A; the tool team may
8
+ instead pick self-fetch by `source_id`, in which case these examples are reshaped
9
+ to match — see registry.py). They reference a hypothetical sales catalog
10
+ (`src_sales` / `t_orders`); these ids are part of the illustration and are not
11
+ validated against the user's real catalog. v1 is descriptive/diagnostic — no
12
+ modeling tasks.
13
+
14
+ See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3 (Examples A and B).
15
+ """
16
+
17
+ from __future__ import annotations
18
+
19
+ from .schemas import Task, TaskList, ToolCall
20
+
21
+ # --------------------------------------------------------------------------- #
22
+ # Example A — exploratory, no modeling.
23
+ # "Which product categories drove last quarter's revenue?"
24
+ # Shows: query_structured pulls rows -> analyze_contribution computes each
25
+ # category's share of the total in one call (no manual per-category + total
26
+ # queries).
27
+ # --------------------------------------------------------------------------- #
28
+
29
+ _EXAMPLE_A = TaskList(
30
+ plan_id="example_a",
31
+ goal_restated="Identify which product categories contributed most to last quarter's revenue.",
32
+ assumptions=["'last quarter' = 2026-01-01 to 2026-03-31."],
33
+ open_questions=[],
34
+ tasks=[
35
+ Task(
36
+ id="t1",
37
+ stage="data_understanding",
38
+ objective="Confirm the sales source exposes category, revenue, and order date.",
39
+ tool_calls=[ToolCall(tool="describe_source", args={"source_id": "src_sales"})],
40
+ expected_output="source_shape",
41
+ success_criteria="Produced the orders table schema; the 3 needed columns are present.",
42
+ depends_on=[],
43
+ estimated_cost="low",
44
+ ),
45
+ Task(
46
+ id="t2",
47
+ stage="data_preparation",
48
+ objective="Pull last quarter's order-level category and revenue rows.",
49
+ tool_calls=[
50
+ ToolCall(
51
+ tool="query_structured",
52
+ args={
53
+ "ir": {
54
+ "source_id": "src_sales",
55
+ "table_id": "t_orders",
56
+ "select": [
57
+ {"kind": "column", "column_id": "c_category", "alias": "category"},
58
+ {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
59
+ ],
60
+ "filters": [
61
+ {
62
+ "column_id": "c_order_date",
63
+ "op": "between",
64
+ "value": ["2026-01-01", "2026-03-31"],
65
+ "value_type": "date",
66
+ }
67
+ ],
68
+ "limit": 10000,
69
+ }
70
+ },
71
+ )
72
+ ],
73
+ expected_output="quarter_rows",
74
+ success_criteria="Produced last quarter's order rows with category and revenue.",
75
+ depends_on=["t1"],
76
+ estimated_cost="medium",
77
+ ),
78
+ Task(
79
+ id="t3",
80
+ stage="evaluation",
81
+ objective="Rank each category's revenue share of the quarter total.",
82
+ tool_calls=[
83
+ ToolCall(
84
+ tool="analyze_contribution",
85
+ args={
86
+ "data": "${t2}",
87
+ "dimension": "category",
88
+ "value_column": "revenue",
89
+ "agg": "sum",
90
+ },
91
+ )
92
+ ],
93
+ expected_output="category_contribution",
94
+ success_criteria="Produced each category's revenue share, ranked high to low.",
95
+ depends_on=["t2"],
96
+ estimated_cost="low",
97
+ ),
98
+ ],
99
+ )
100
+
101
+ # --------------------------------------------------------------------------- #
102
+ # Example B — descriptive / trend.
103
+ # "How has monthly revenue trended by region this year, and what's unusual?"
104
+ # --------------------------------------------------------------------------- #
105
+
106
+ _EXAMPLE_B = TaskList(
107
+ plan_id="example_b",
108
+ goal_restated="Describe this year's monthly revenue trend and flag unusual months.",
109
+ assumptions=["'this year' starts 2026-01-01."],
110
+ open_questions=["'Unusual' is interpreted as months far from the typical monthly revenue."],
111
+ tasks=[
112
+ Task(
113
+ id="t1",
114
+ stage="data_understanding",
115
+ objective="Confirm the sales source exposes order date, revenue, and region.",
116
+ tool_calls=[ToolCall(tool="describe_source", args={"source_id": "src_sales"})],
117
+ expected_output="source_shape",
118
+ success_criteria="Produced the orders table schema; the needed columns are present.",
119
+ depends_on=[],
120
+ estimated_cost="low",
121
+ ),
122
+ Task(
123
+ id="t2",
124
+ stage="data_preparation",
125
+ objective="Pull this year's order dates, revenue, and region.",
126
+ tool_calls=[
127
+ ToolCall(
128
+ tool="query_structured",
129
+ args={
130
+ "ir": {
131
+ "source_id": "src_sales",
132
+ "table_id": "t_orders",
133
+ "select": [
134
+ {
135
+ "kind": "column",
136
+ "column_id": "c_order_date",
137
+ "alias": "order_date",
138
+ },
139
+ {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
140
+ {"kind": "column", "column_id": "c_region", "alias": "region"},
141
+ ],
142
+ "filters": [
143
+ {
144
+ "column_id": "c_order_date",
145
+ "op": ">=",
146
+ "value": "2026-01-01",
147
+ "value_type": "date",
148
+ }
149
+ ],
150
+ "limit": 10000,
151
+ }
152
+ },
153
+ )
154
+ ],
155
+ expected_output="ytd_rows",
156
+ success_criteria="Produced this year's order-level rows with date, revenue, region.",
157
+ depends_on=["t1"],
158
+ estimated_cost="medium",
159
+ ),
160
+ Task(
161
+ id="t3",
162
+ stage="evaluation",
163
+ objective="Bucket revenue into months and summarize the trend and movement.",
164
+ tool_calls=[
165
+ ToolCall(
166
+ tool="analyze_trend",
167
+ args={
168
+ "data": "${t2}",
169
+ "date_column": "order_date",
170
+ "value_column": "revenue",
171
+ "freq": "month",
172
+ "agg": "sum",
173
+ },
174
+ )
175
+ ],
176
+ expected_output="monthly_trend",
177
+ success_criteria=(
178
+ "Produced a per-month revenue series with direction and change rate to "
179
+ "flag months above/below the typical level."
180
+ ),
181
+ depends_on=["t2"],
182
+ estimated_cost="low",
183
+ ),
184
+ ],
185
+ )
186
+
187
+
188
+ # --------------------------------------------------------------------------- #
189
+ # Example C — mixed structured + unstructured.
190
+ # "Revenue dipped in Q1 — what happened?"
191
+ # Shows: a structured branch (query -> analyze_trend) runs alongside an
192
+ # INDEPENDENT retrieve_documents branch that pulls qualitative context. Note
193
+ # retrieve_documents takes a natural-language `query` (NOT a `${t<id>}` data
194
+ # placeholder — it is a source, not a consumer) and can run in parallel; the
195
+ # Assembler folds the document context into the explanation.
196
+ # --------------------------------------------------------------------------- #
197
+
198
+ _EXAMPLE_C = TaskList(
199
+ plan_id="example_c",
200
+ goal_restated="Explain Q1's revenue dip using both the numbers and the qualitative record.",
201
+ assumptions=["'Q1' = 2026-01-01 to 2026-03-31."],
202
+ open_questions=[],
203
+ tasks=[
204
+ Task(
205
+ id="t1",
206
+ stage="data_understanding",
207
+ objective="Confirm the sales source exposes order date and revenue.",
208
+ tool_calls=[ToolCall(tool="describe_source", args={"source_id": "src_sales"})],
209
+ expected_output="source_shape",
210
+ success_criteria="Produced the orders table schema; date and revenue columns present.",
211
+ depends_on=[],
212
+ estimated_cost="low",
213
+ ),
214
+ Task(
215
+ id="t2",
216
+ stage="data_preparation",
217
+ objective="Pull Q1 order dates and revenue.",
218
+ tool_calls=[
219
+ ToolCall(
220
+ tool="query_structured",
221
+ args={
222
+ "ir": {
223
+ "source_id": "src_sales",
224
+ "table_id": "t_orders",
225
+ "select": [
226
+ {
227
+ "kind": "column",
228
+ "column_id": "c_order_date",
229
+ "alias": "order_date",
230
+ },
231
+ {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
232
+ ],
233
+ "filters": [
234
+ {
235
+ "column_id": "c_order_date",
236
+ "op": "between",
237
+ "value": ["2026-01-01", "2026-03-31"],
238
+ "value_type": "date",
239
+ }
240
+ ],
241
+ "limit": 10000,
242
+ }
243
+ },
244
+ )
245
+ ],
246
+ expected_output="q1_rows",
247
+ success_criteria="Produced Q1 order rows with date and revenue.",
248
+ depends_on=["t1"],
249
+ estimated_cost="medium",
250
+ ),
251
+ Task(
252
+ id="t3",
253
+ stage="evaluation",
254
+ objective="Summarize the Q1 monthly revenue trend to locate the dip.",
255
+ tool_calls=[
256
+ ToolCall(
257
+ tool="analyze_trend",
258
+ args={
259
+ "data": "${t2}",
260
+ "date_column": "order_date",
261
+ "value_column": "revenue",
262
+ "freq": "month",
263
+ "agg": "sum",
264
+ },
265
+ )
266
+ ],
267
+ expected_output="q1_trend",
268
+ success_criteria="Produced a per-month revenue series showing where revenue fell.",
269
+ depends_on=["t2"],
270
+ estimated_cost="low",
271
+ ),
272
+ Task(
273
+ id="t4",
274
+ stage="data_understanding",
275
+ objective="Retrieve qualitative context on Q1 operational events behind a dip.",
276
+ tool_calls=[
277
+ ToolCall(
278
+ tool="retrieve_documents",
279
+ args={
280
+ "query": "operational issues, outages, or notable events in Q1 2026",
281
+ "top_k": 5,
282
+ },
283
+ )
284
+ ],
285
+ expected_output="q1_context_chunks",
286
+ success_criteria="Produced relevant document chunks about Q1 operations.",
287
+ depends_on=[],
288
+ estimated_cost="low",
289
+ ),
290
+ ],
291
+ )
292
+
293
+
294
+ # --------------------------------------------------------------------------- #
295
+ # Example D — group-by aggregation (analyze_aggregate arg shape).
296
+ # "What is the average and total order value per region?"
297
+ # Shows the EXACT analyze_aggregate args: `aggregations` is an OBJECT mapping each
298
+ # column to a LIST of functions ({"revenue": ["mean", "sum"]}), and `group_by` is a
299
+ # SEPARATE array — NOT a nested list of metric specs. Supported funcs: sum, mean,
300
+ # count, min, max, median, nunique.
301
+ # --------------------------------------------------------------------------- #
302
+
303
+ _EXAMPLE_D = TaskList(
304
+ plan_id="example_d",
305
+ goal_restated="Report the average and total order value for each region.",
306
+ assumptions=[],
307
+ open_questions=[],
308
+ tasks=[
309
+ Task(
310
+ id="t1",
311
+ stage="data_understanding",
312
+ objective="Confirm the sales source exposes region and revenue.",
313
+ tool_calls=[ToolCall(tool="describe_source", args={"source_id": "src_sales"})],
314
+ expected_output="source_shape",
315
+ success_criteria="Produced the orders table schema; region and revenue present.",
316
+ depends_on=[],
317
+ estimated_cost="low",
318
+ ),
319
+ Task(
320
+ id="t2",
321
+ stage="data_preparation",
322
+ objective="Pull order-level region and revenue.",
323
+ tool_calls=[
324
+ ToolCall(
325
+ tool="query_structured",
326
+ args={
327
+ "ir": {
328
+ "source_id": "src_sales",
329
+ "table_id": "t_orders",
330
+ "select": [
331
+ {"kind": "column", "column_id": "c_region", "alias": "region"},
332
+ {"kind": "column", "column_id": "c_revenue", "alias": "revenue"},
333
+ ],
334
+ "limit": 10000,
335
+ }
336
+ },
337
+ )
338
+ ],
339
+ expected_output="region_rows",
340
+ success_criteria="Produced order rows with region and revenue.",
341
+ depends_on=["t1"],
342
+ estimated_cost="medium",
343
+ ),
344
+ Task(
345
+ id="t3",
346
+ stage="evaluation",
347
+ objective="Aggregate mean and total revenue per region.",
348
+ tool_calls=[
349
+ ToolCall(
350
+ tool="analyze_aggregate",
351
+ args={
352
+ "data": "${t2}",
353
+ "aggregations": {"revenue": ["mean", "sum"]},
354
+ "group_by": ["region"],
355
+ },
356
+ )
357
+ ],
358
+ expected_output="region_aggregates",
359
+ success_criteria="Produced one row per region with mean and total revenue.",
360
+ depends_on=["t2"],
361
+ estimated_cost="low",
362
+ ),
363
+ ],
364
+ )
365
+
366
+
367
+ EXAMPLES: list[tuple[str, TaskList]] = [
368
+ ("Which product categories drove last quarter's revenue?", _EXAMPLE_A),
369
+ ("How has monthly revenue trended by region this year, and what's unusual?", _EXAMPLE_B),
370
+ ("Revenue dipped in Q1 — what happened?", _EXAMPLE_C),
371
+ ("What is the average and total order value per region?", _EXAMPLE_D),
372
+ ]
373
+
374
+
375
+ def render_examples() -> str:
376
+ """Render the few-shots as text for the planner prompt."""
377
+ blocks: list[str] = []
378
+ for i, (question, plan) in enumerate(EXAMPLES, start=1):
379
+ blocks.append(
380
+ f"## Example {i}\n\n"
381
+ f"Question:\n{question}\n\n"
382
+ f"TaskList:\n{plan.model_dump_json(indent=2)}"
383
+ )
384
+ return "\n\n".join(blocks)
src/agents/planner/inputs.py ADDED
@@ -0,0 +1,139 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Planner input models — CatalogSummary and Constraints.
2
+
3
+ `CatalogSummary` is a condensed, PII-safe view of the user's `Catalog`, built
4
+ for the planner prompt. It carries every table + column id/type/PII flag + row
5
+ counts + low-cardinality top_values, with `sample_values` nulled on PII columns
6
+ (INV: no PII sample values into the prompt, see doc §13). It also lists the
7
+ available unstructured sources so the planner can plan `retrieve_documents`.
8
+
9
+ The planner *validator* still checks inline `query_structured` IRs against the
10
+ full `Catalog` via the existing IRValidator — the summary is a prompt input, not
11
+ the validation source of truth.
12
+
13
+ See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ from typing import Any
19
+
20
+ from pydantic import BaseModel, Field
21
+
22
+ from ...catalog.models import Catalog, DataType
23
+
24
+
25
+ class ColumnSummary(BaseModel):
26
+ column_id: str
27
+ name: str
28
+ data_type: DataType
29
+ pii_flag: bool = False
30
+ sample_values: list[Any] | None = None # nulled when pii_flag is True
31
+ top_values: list[Any] | None = None
32
+
33
+
34
+ class TableSummary(BaseModel):
35
+ table_id: str
36
+ name: str
37
+ row_count: int | None = None
38
+ columns: list[ColumnSummary] = Field(default_factory=list)
39
+
40
+
41
+ class StructuredSourceSummary(BaseModel):
42
+ source_id: str
43
+ name: str
44
+ source_type: str # "schema" | "tabular"
45
+ tables: list[TableSummary] = Field(default_factory=list)
46
+
47
+
48
+ class UnstructuredSourceSummary(BaseModel):
49
+ source_id: str
50
+ name: str
51
+
52
+
53
+ class CatalogSummary(BaseModel):
54
+ structured_sources: list[StructuredSourceSummary] = Field(default_factory=list)
55
+ unstructured_sources: list[UnstructuredSourceSummary] = Field(default_factory=list)
56
+
57
+ @classmethod
58
+ def from_catalog(cls, catalog: Catalog) -> CatalogSummary:
59
+ structured: list[StructuredSourceSummary] = []
60
+ unstructured: list[UnstructuredSourceSummary] = []
61
+
62
+ for source in catalog.sources:
63
+ if source.source_type == "unstructured":
64
+ unstructured.append(
65
+ UnstructuredSourceSummary(source_id=source.source_id, name=source.name)
66
+ )
67
+ continue
68
+
69
+ tables = [
70
+ TableSummary(
71
+ table_id=table.table_id,
72
+ name=table.name,
73
+ row_count=table.row_count,
74
+ columns=[
75
+ ColumnSummary(
76
+ column_id=col.column_id,
77
+ name=col.name,
78
+ data_type=col.data_type,
79
+ pii_flag=col.pii_flag,
80
+ # PII columns leak nothing into the prompt: both
81
+ # sample_values and (low-cardinality) top_values are
82
+ # suppressed — top_values are the same class of data.
83
+ sample_values=None if col.pii_flag else col.sample_values,
84
+ top_values=(
85
+ None
86
+ if col.pii_flag or col.stats is None
87
+ else col.stats.top_values
88
+ ),
89
+ )
90
+ for col in table.columns
91
+ ],
92
+ )
93
+ for table in source.tables
94
+ ]
95
+ structured.append(
96
+ StructuredSourceSummary(
97
+ source_id=source.source_id,
98
+ name=source.name,
99
+ source_type=source.source_type,
100
+ tables=tables,
101
+ )
102
+ )
103
+
104
+ return cls(structured_sources=structured, unstructured_sources=unstructured)
105
+
106
+ def render(self) -> str:
107
+ """Render the summary as compact text for the planner prompt."""
108
+ if not self.structured_sources and not self.unstructured_sources:
109
+ return "(catalog is empty — the user has not registered any data yet)"
110
+
111
+ lines: list[str] = []
112
+ for source in self.structured_sources:
113
+ lines.append(f"Source: {source.name} ({source.source_type}) — id={source.source_id}")
114
+ for table in source.tables:
115
+ rc = f" ({table.row_count:,} rows)" if table.row_count is not None else ""
116
+ lines.append(f" Table: {table.name}{rc} — id={table.table_id}")
117
+ for col in table.columns:
118
+ samples = "PII (suppressed)" if col.pii_flag else (col.sample_values or [])
119
+ top = f", top={col.top_values}" if col.top_values else ""
120
+ lines.append(
121
+ f" - {col.name} [{col.data_type}]: "
122
+ f"samples={samples}{top} — id={col.column_id}"
123
+ )
124
+ lines.append("")
125
+
126
+ if self.unstructured_sources:
127
+ lines.append("Unstructured sources (for retrieve_documents):")
128
+ for src in self.unstructured_sources:
129
+ lines.append(f" - {src.name} — id={src.source_id}")
130
+
131
+ return "\n".join(lines).rstrip()
132
+
133
+
134
+ class Constraints(BaseModel):
135
+ max_tasks: int = 5
136
+ modeling_allowed: bool = False # no modeling tools in v1
137
+ token_budget: int | None = None
138
+ time_budget_seconds: int | None = None
139
+ row_budget: int = 10_000
src/agents/planner/prompt.py ADDED
@@ -0,0 +1,106 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Builds the planner LLM human-message content.
2
+
3
+ The system prompt (`config/prompts/planner.md`) carries the role, invariants,
4
+ and planning principles. This module assembles the per-call human content:
5
+ business context + condensed catalog + available tools + constraints + the
6
+ few-shot examples + the question (+ the prior error on retry).
7
+
8
+ Few-shot examples are rendered from `examples.py` (which builds them from the
9
+ real `TaskList` schema) so they cannot drift from the output contract.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ from .contracts import BusinessContext, ToolRegistry
15
+ from .examples import render_examples
16
+ from .inputs import CatalogSummary, Constraints
17
+
18
+
19
+ def render_business_context(context: BusinessContext) -> str:
20
+ lines = [
21
+ f"Project: {context.project_id} (industry: {context.industry}, "
22
+ f"context completeness: {context.completeness})",
23
+ f"Business: {context.business_description}",
24
+ f"Scale & scope: {context.scale_and_scope}",
25
+ ]
26
+ if context.key_terms:
27
+ lines.append("Key terms:")
28
+ lines.extend(f" - {kt.term}: {kt.meaning}" for kt in context.key_terms)
29
+ if context.data_overview:
30
+ lines.append("Data overview:")
31
+ lines.extend(
32
+ f" - {n.table_name}: {n.row_represents}" for n in context.data_overview
33
+ )
34
+ if context.data_column_notes:
35
+ lines.append("Column notes:")
36
+ lines.extend(
37
+ f" - {n.column_name}: {n.meaning}" for n in context.data_column_notes
38
+ )
39
+ if context.whats_normal:
40
+ lines.append(f"What's normal: {context.whats_normal}")
41
+ if context.recent_events:
42
+ lines.append(f"Recent events: {context.recent_events}")
43
+ if context.things_to_watch_for:
44
+ lines.append(f"Things to watch for: {context.things_to_watch_for}")
45
+ if context.open_questions:
46
+ lines.append("Known open questions:")
47
+ lines.extend(f" - {q}" for q in context.open_questions)
48
+ return "\n".join(lines)
49
+
50
+
51
+ def render_registry(tools: ToolRegistry) -> str:
52
+ if not tools.tools:
53
+ return "(no tools available)"
54
+ blocks: list[str] = []
55
+ for spec in tools.tools:
56
+ required = spec.input_schema.get("required", [])
57
+ blocks.append(
58
+ f"- {spec.name} (category: {spec.category}, returns: {spec.output_kind})\n"
59
+ f" required args: {required}\n"
60
+ f" {spec.description}"
61
+ )
62
+ return "\n".join(blocks)
63
+
64
+
65
+ def render_constraints(constraints: Constraints) -> str:
66
+ lines = [
67
+ f"- max_tasks: {constraints.max_tasks}",
68
+ f"- modeling_allowed: {constraints.modeling_allowed} "
69
+ "(no modeling tools exist in v1 — do not emit modeling tasks)",
70
+ f"- row_budget: {constraints.row_budget}",
71
+ ]
72
+ if constraints.token_budget is not None:
73
+ lines.append(f"- token_budget: {constraints.token_budget}")
74
+ if constraints.time_budget_seconds is not None:
75
+ lines.append(f"- time_budget_seconds: {constraints.time_budget_seconds}")
76
+ return "\n".join(lines)
77
+
78
+
79
+ def build_planner_prompt(
80
+ context: BusinessContext,
81
+ catalog: CatalogSummary,
82
+ tools: ToolRegistry,
83
+ query: str,
84
+ constraints: Constraints,
85
+ previous_error: str | None = None,
86
+ ) -> str:
87
+ """Return the human-message content for the planner LLM.
88
+
89
+ The system prompt (`config/prompts/planner.md`) is loaded separately by
90
+ `PlannerService`.
91
+ """
92
+ sections = [
93
+ f"# Business context\n\n{render_business_context(context)}",
94
+ f"# Catalog\n\n{catalog.render()}",
95
+ f"# Available tools\n\n{render_registry(tools)}",
96
+ f"# Constraints\n\n{render_constraints(constraints)}",
97
+ f"# Examples\n\n{render_examples()}",
98
+ f"# Question\n\n{query}",
99
+ ]
100
+ if previous_error:
101
+ sections.append(
102
+ "# Previous attempt failed validation\n\n"
103
+ f"{previous_error}\n\n"
104
+ "Emit a corrected TaskList. Do not repeat the same mistake."
105
+ )
106
+ return "\n\n".join(sections)
src/agents/planner/registry.py ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """v1 tool registry the Planner plans against (INV-7: agent never names a tool
2
+ outside it).
3
+
4
+ **Composed from two slices (2026-06-08):**
5
+
6
+ - **Analytics (`analyze_*`) — REAL, tool-team-owned.** Sourced live from
7
+ `src/tools/registry.py::analytics_registry()` (KM-628), built on the canonical
8
+ `ToolSpec` (`src/tools/contracts.py`, KM-465/KM-627) and the prompt-style tool
9
+ descriptions (KM-625). No longer a stub on our side — it tracks the real registry.
10
+ - **Data access (`query_structured` / `retrieve_documents` / `list_sources` /
11
+ `describe_source`) — spec BODIES still a local stub.** The tool team owns these too,
12
+ but their wrappers + `ToolSpec`s haven't landed yet (KM-465 #4). We keep best-guess
13
+ spec bodies here so the Planner can plan end-to-end — but the NAMES derive from
14
+ `src.tools.data_access.DATA_ACCESS_TOOLS` (R11), so a tool rename/addition upstream
15
+ fails loudly here instead of drifting silently. When the real specs ship, delete
16
+ this slice and swap `default_registry()` for the tool team's full composition.
17
+
18
+ **Confirmed conventions (KM-465):** Pattern A — `analyze_*` tools take a `data`
19
+ `"${t<id>}"` placeholder pointing at an upstream `query_structured` output (no
20
+ self-fetch); resolved to a DataFrame at execution time. `input_schema` is the
21
+ lightweight `{required, properties}` dict the planner validator (check #8) reads;
22
+ `query_structured.args["ir"]` carries an inline QueryIR validated against the
23
+ catalog by the existing IRValidator.
24
+
25
+ See AGENT_ARCHITECTURE_CONTEXT_new.md §9.2 / §9.3.
26
+ """
27
+
28
+ from __future__ import annotations
29
+
30
+ from src.tools.data_access import DATA_ACCESS_TOOLS
31
+ from src.tools.registry import analytics_registry
32
+
33
+ from .contracts import ToolRegistry, ToolSpec
34
+
35
+ # --------------------------------------------------------------------------- #
36
+ # Data-access slice — spec bodies are a LOCAL STUB pending the tool team's real
37
+ # specs (KM-465 #4); the canonical NAME SET is `DATA_ACCESS_TOOLS` (tool-owned).
38
+ # --------------------------------------------------------------------------- #
39
+ _DATA_ACCESS_SPEC_BODIES: tuple[ToolSpec, ...] = (
40
+ ToolSpec(
41
+ name="query_structured",
42
+ category="analytics.query",
43
+ input_schema={"required": ["ir"], "properties": {"ir": {"type": "object"}}},
44
+ output_kind="table",
45
+ description=(
46
+ "Run one validated, single-table query against a structured source (DB "
47
+ "schema or tabular file) and return rows. The `ir` argument is an inline "
48
+ "QueryIR (the JSON intent: source_id, table_id, select, filters, group_by, "
49
+ "order_by, limit) — never SQL. This is the data-access entry point: use it "
50
+ "to select, filter, and pull the rows the analytics (`analyze_*`) tools "
51
+ "then consume. It also does simple built-in aggregation the IR can express "
52
+ "(count/sum/avg/min/max/count_distinct). Do NOT use it for richer statistics "
53
+ "(median/percentile/mode/stddev/skew → analyze_descriptive), trends "
54
+ "(analyze_trend), correlation, segmentation, or share-of-total; and do NOT "
55
+ "use it to read documents (use retrieve_documents)."
56
+ ),
57
+ ),
58
+ ToolSpec(
59
+ name="retrieve_documents",
60
+ category="retrieval.documents",
61
+ input_schema={
62
+ "required": ["query"],
63
+ "properties": {
64
+ "query": {"type": "string"},
65
+ "source_id": {"type": "string"},
66
+ "top_k": {"type": "integer"},
67
+ },
68
+ },
69
+ output_kind="documents",
70
+ description=(
71
+ "Dense-retrieve the most relevant chunks from the user's unstructured "
72
+ "sources (PDF/DOCX/TXT) for a natural-language `query`. Use this to pull "
73
+ "qualitative context into an analysis. Optionally scope to one `source_id`. "
74
+ "Do NOT use it for numbers in tables — that is query_structured's job."
75
+ ),
76
+ ),
77
+ ToolSpec(
78
+ name="list_sources",
79
+ category="catalog.introspection",
80
+ input_schema={"required": [], "properties": {}},
81
+ output_kind="table",
82
+ description=(
83
+ "List the user's available data sources (id, name, type, table count). Use "
84
+ "early in data_understanding when the plan must discover what exists before "
85
+ "querying. Cheap. Do NOT use it to read column details (use describe_source)."
86
+ ),
87
+ ),
88
+ ToolSpec(
89
+ name="describe_source",
90
+ category="catalog.introspection",
91
+ input_schema={
92
+ "required": ["source_id"],
93
+ "properties": {"source_id": {"type": "string"}},
94
+ },
95
+ output_kind="table",
96
+ description=(
97
+ "Return the tables and columns (names, types, row counts) of one source by "
98
+ "`source_id`. Use in data_understanding to confirm the shape of a source "
99
+ "before querying it. Do NOT use it to fetch data rows (use query_structured)."
100
+ ),
101
+ ),
102
+ )
103
+
104
+ _DATA_ACCESS_SPECS: dict[str, ToolSpec] = {s.name: s for s in _DATA_ACCESS_SPEC_BODIES}
105
+
106
+
107
+ def _data_access_slice() -> list[ToolSpec]:
108
+ """Data-access specs in body order, with names checked against the tool layer.
109
+
110
+ `DATA_ACCESS_TOOLS` (src.tools.data_access) is the canonical name set; the
111
+ spec bodies above are still our local stub. Any mismatch (a tool added,
112
+ renamed, or removed upstream) raises here instead of drifting silently.
113
+ """
114
+ if set(_DATA_ACCESS_SPECS) != DATA_ACCESS_TOOLS:
115
+ missing = sorted(DATA_ACCESS_TOOLS - _DATA_ACCESS_SPECS.keys())
116
+ stale = sorted(_DATA_ACCESS_SPECS.keys() - DATA_ACCESS_TOOLS)
117
+ raise RuntimeError(
118
+ "planner data-access specs out of sync with "
119
+ f"src.tools.data_access.DATA_ACCESS_TOOLS: missing spec for {missing}, "
120
+ f"stale spec for {stale}"
121
+ )
122
+ return list(_DATA_ACCESS_SPECS.values())
123
+
124
+
125
+ def default_registry() -> ToolRegistry:
126
+ """The v1 registry: stub data-access slice + the real analytics slice.
127
+
128
+ The analytics tools come live from `src.tools.registry` (the tool team's real
129
+ registry); the data-access spec bodies are still a local stub, name-checked
130
+ against `DATA_ACCESS_TOOLS`. A fresh instance per call.
131
+ """
132
+ return ToolRegistry(tools=[*_data_access_slice(), *analytics_registry().tools])
src/agents/planner/schemas.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Planner output schemas — the `TaskList` contract.
2
+
3
+ The planner emits exactly one `TaskList`: a DAG of typed tasks, each an ordered
4
+ chain of fully-specified tool calls. This is the static plan the TaskRunner
5
+ executes verbatim. There is no replanning (INV-6), so there are no
6
+ `ReplanRequest` / `ReplanResponse` schemas.
7
+
8
+ See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ import re
14
+ from typing import Any, Literal
15
+
16
+ from pydantic import BaseModel, Field
17
+
18
+ # The "${t<id>}" placeholder convention (Pattern A): a ToolCall arg whose value
19
+ # matches this pattern refers to an upstream task's output, resolved by the
20
+ # TaskRunner at execution time. Single definition — the planner validator and
21
+ # the TaskRunner both import it (R11).
22
+ PLACEHOLDER_RE = re.compile(r"\$\{(t[^}]+)\}")
23
+
24
+ CrispStage = Literal[
25
+ "data_understanding",
26
+ "data_preparation",
27
+ "modeling", # no tools in v1; the planner does not emit modeling tasks
28
+ "evaluation",
29
+ ]
30
+
31
+
32
+ class ToolCall(BaseModel):
33
+ """One call to a registry tool with concrete, fully-specified arguments.
34
+
35
+ `tool` must exist in the ToolRegistry. `args` is validated against the
36
+ tool's input_schema; it may contain "${t<id>}" placeholders that the
37
+ TaskRunner resolves from an upstream task's output at execution time.
38
+ """
39
+
40
+ tool: str
41
+ args: dict[str, Any] = Field(default_factory=dict)
42
+
43
+
44
+ class Task(BaseModel):
45
+ id: str # "t1", "t2", ...
46
+ stage: CrispStage
47
+ objective: str # plain-language intent for this step
48
+ tool_calls: list[ToolCall] = Field(..., min_length=1) # ordered chain
49
+ expected_output: str # named result this task produces
50
+ success_criteria: str # REPORTING signal, not a control trigger
51
+ depends_on: list[str] = Field(default_factory=list) # task ids
52
+ estimated_cost: Literal["low", "medium", "high"] = "low"
53
+
54
+
55
+ class TaskList(BaseModel):
56
+ plan_id: str
57
+ goal_restated: str
58
+ assumptions: list[str] = Field(default_factory=list)
59
+ open_questions: list[str] = Field(default_factory=list)
60
+ tasks: list[Task] = Field(default_factory=list)
src/agents/planner/service.py ADDED
@@ -0,0 +1,158 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """PlannerService — single LLM call: context + catalog + tools + question -> TaskList.
2
+
3
+ Mirrors `query/planner/service.py` (chain construction) and `query/service.py`
4
+ (validate-and-retry loop). The planner LLM emits a `TaskList` via structured
5
+ output; the `PlannerValidator` runs the 8 checks; on failure the planner is
6
+ re-prompted with the error context, up to `max_retries` (default 3). No
7
+ replanning happens at execution time — this loop only hardens the *initial*
8
+ static plan.
9
+
10
+ The service takes the full `Catalog` (not just a `CatalogSummary`): it derives
11
+ the PII-safe `CatalogSummary` for the prompt, but validation needs the full
12
+ catalog so the existing `IRValidator` can check inline `query_structured` IRs.
13
+
14
+ See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3.
15
+ """
16
+
17
+ from __future__ import annotations
18
+
19
+ from pathlib import Path
20
+
21
+ from langchain_core.messages import SystemMessage
22
+ from langchain_core.prompts import ChatPromptTemplate
23
+ from langchain_core.runnables import Runnable
24
+ from langchain_openai import AzureChatOpenAI
25
+
26
+ from src.middlewares.logging import get_logger
27
+
28
+ from ...catalog.models import Catalog
29
+ from .contracts import BusinessContext, ToolRegistry
30
+ from .errors import PlannerError, PlannerValidationError
31
+ from .inputs import CatalogSummary, Constraints
32
+ from .prompt import build_planner_prompt
33
+ from .schemas import TaskList
34
+ from .validator import PlannerValidator
35
+
36
+ logger = get_logger("planner_agent")
37
+
38
+ _PROMPT_PATH = (
39
+ Path(__file__).resolve().parent.parent.parent / "config" / "prompts" / "planner.md"
40
+ )
41
+
42
+
43
+ def _load_prompt_text() -> str:
44
+ return _PROMPT_PATH.read_text(encoding="utf-8")
45
+
46
+
47
+ def _build_default_chain() -> Runnable:
48
+ from src.config.settings import settings
49
+
50
+ llm = AzureChatOpenAI(
51
+ azure_deployment=settings.azureai_deployment_name_4o,
52
+ openai_api_version=settings.azureai_api_version_4o,
53
+ azure_endpoint=settings.azureai_endpoint_url_4o,
54
+ api_key=settings.azureai_api_key_4o,
55
+ temperature=0,
56
+ )
57
+ prompt = ChatPromptTemplate.from_messages(
58
+ [
59
+ SystemMessage(content=_load_prompt_text()),
60
+ ("human", "{human_content}"),
61
+ ]
62
+ )
63
+ return prompt | llm.with_structured_output(TaskList)
64
+
65
+
66
+ _default_chain: Runnable | None = None
67
+
68
+
69
+ def _get_default_chain() -> Runnable:
70
+ global _default_chain
71
+ if _default_chain is None:
72
+ _default_chain = _build_default_chain()
73
+ return _default_chain
74
+
75
+
76
+ class PlannerService:
77
+ """Wraps the planner LLM call + the validate-and-retry loop.
78
+
79
+ Inject `structured_chain` and/or `validator` for tests.
80
+ """
81
+
82
+ def __init__(
83
+ self,
84
+ structured_chain: Runnable | None = None,
85
+ validator: PlannerValidator | None = None,
86
+ max_retries: int = 3,
87
+ ) -> None:
88
+ self._chain = structured_chain
89
+ self._validator = validator or PlannerValidator()
90
+ self._max_retries = max(1, max_retries)
91
+
92
+ def _ensure_chain(self) -> Runnable:
93
+ if self._chain is None:
94
+ self._chain = _get_default_chain()
95
+ return self._chain
96
+
97
+ async def plan(
98
+ self,
99
+ context: BusinessContext,
100
+ catalog: Catalog,
101
+ tools: ToolRegistry,
102
+ query: str,
103
+ constraints: Constraints,
104
+ callbacks: list | None = None,
105
+ ) -> TaskList:
106
+ summary = CatalogSummary.from_catalog(catalog)
107
+ chain = self._ensure_chain()
108
+ previous_error: str | None = None
109
+
110
+ for attempt in range(1, self._max_retries + 1):
111
+ human_content = build_planner_prompt(
112
+ context, summary, tools, query, constraints, previous_error
113
+ )
114
+ # All retry attempts share `callbacks`, so each shows up under the same
115
+ # trace — that is how retry token cost becomes visible.
116
+ if callbacks:
117
+ task_list: TaskList = await chain.ainvoke(
118
+ {"human_content": human_content}, config={"callbacks": callbacks}
119
+ )
120
+ else:
121
+ task_list = await chain.ainvoke({"human_content": human_content})
122
+ try:
123
+ self._validator.validate(task_list, tools, catalog, constraints)
124
+ except PlannerValidationError as e:
125
+ previous_error = str(e)
126
+ logger.warning(
127
+ "planner validation failed",
128
+ project_id=context.project_id,
129
+ plan_id=task_list.plan_id,
130
+ attempt=attempt,
131
+ error=previous_error,
132
+ )
133
+ continue
134
+
135
+ logger.info(
136
+ "analysis planned",
137
+ project_id=context.project_id,
138
+ plan_id=task_list.plan_id,
139
+ n_tasks=len(task_list.tasks),
140
+ retry=attempt > 1,
141
+ )
142
+ return task_list
143
+
144
+ raise PlannerError(
145
+ f"planner failed validation after {self._max_retries} attempts; "
146
+ f"last error: {previous_error}"
147
+ )
148
+
149
+
150
+ async def plan_analysis(
151
+ context: BusinessContext,
152
+ catalog: Catalog,
153
+ tools: ToolRegistry,
154
+ query: str,
155
+ constraints: Constraints,
156
+ ) -> TaskList:
157
+ """Convenience entry point using the default chain + validator."""
158
+ return await PlannerService().plan(context, catalog, tools, query, constraints)
src/agents/planner/validator.py ADDED
@@ -0,0 +1,229 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """PlannerValidator — checks a TaskList before it reaches the TaskRunner.
2
+
3
+ Runs the 8 checks from AGENT_ARCHITECTURE_CONTEXT_new.md §7.3. On failure it
4
+ raises `PlannerValidationError` with a message specific enough that the planner
5
+ can be re-prompted to self-correct (the retry loop lives in service.py).
6
+
7
+ Check #1 (Pydantic parse) is enforced at the structured-output boundary — by the
8
+ time a `TaskList` reaches here it has already parsed; this validator additionally
9
+ rejects structurally-invalid plans (duplicate ids, dangling edges, cycles).
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ from pydantic import ValidationError
15
+
16
+ from ...catalog.models import Catalog
17
+ from ...query.ir.models import QueryIR
18
+ from ...query.ir.validator import IRValidationError, IRValidator
19
+ from .contracts import ToolRegistry
20
+ from .errors import PlannerValidationError
21
+ from .inputs import Constraints
22
+ from .schemas import PLACEHOLDER_RE, TaskList
23
+
24
+ # Heuristic: a checkable success_criteria mentions a measurable signal.
25
+ _CHECKABLE_TOKENS = ("rate", "count", "match", "produced", "above", "below", "equal")
26
+
27
+ # DFS colors for cycle detection.
28
+ _WHITE, _GREY, _BLACK = 0, 1, 2
29
+
30
+
31
+ class PlannerValidator:
32
+ def __init__(self, ir_validator: IRValidator | None = None) -> None:
33
+ self._ir_validator = ir_validator or IRValidator()
34
+
35
+ def validate(
36
+ self,
37
+ task_list: TaskList,
38
+ registry: ToolRegistry,
39
+ catalog: Catalog,
40
+ constraints: Constraints,
41
+ ) -> None:
42
+ tasks = task_list.tasks
43
+
44
+ # Check 6 — plan non-empty and within the task cap.
45
+ if not tasks:
46
+ raise PlannerValidationError("plan is empty: at least one task is required")
47
+ if len(tasks) > constraints.max_tasks:
48
+ raise PlannerValidationError(
49
+ f"plan has {len(tasks)} tasks, exceeds max_tasks={constraints.max_tasks}"
50
+ )
51
+
52
+ ids = [t.id for t in tasks]
53
+ if len(set(ids)) != len(ids):
54
+ dupes = sorted({i for i in ids if ids.count(i) > 1})
55
+ raise PlannerValidationError(f"duplicate task id(s): {dupes}")
56
+ id_set = set(ids)
57
+ tasks_by_id = {t.id: t for t in tasks}
58
+
59
+ known_tools = registry.names()
60
+ known_sources = {s.source_id for s in catalog.sources}
61
+
62
+ for task in tasks:
63
+ for call in task.tool_calls:
64
+ # Check 2 — every tool exists in the registry.
65
+ if call.tool not in known_tools:
66
+ raise PlannerValidationError(
67
+ f"task {task.id}: tool {call.tool!r} not in registry "
68
+ f"(known: {sorted(known_tools)})"
69
+ )
70
+ spec = registry.get(call.tool)
71
+ assert spec is not None # guaranteed by the membership check above
72
+
73
+ # Check 8a — args carry the required keys and no unknown keys.
74
+ required = set(spec.input_schema.get("required", []))
75
+ allowed = set(spec.input_schema.get("properties", {}).keys()) | required
76
+ missing = required - set(call.args.keys())
77
+ if missing:
78
+ raise PlannerValidationError(
79
+ f"task {task.id}: tool {call.tool!r} missing required arg(s): "
80
+ f"{sorted(missing)}"
81
+ )
82
+ unknown = set(call.args.keys()) - allowed
83
+ if unknown:
84
+ raise PlannerValidationError(
85
+ f"task {task.id}: tool {call.tool!r} has unknown arg(s): "
86
+ f"{sorted(unknown)} (allowed: {sorted(allowed)})"
87
+ )
88
+
89
+ # Check 3 — concrete source_id args must exist in the catalog.
90
+ src = call.args.get("source_id")
91
+ if isinstance(src, str) and not _is_placeholder(src):
92
+ if src not in known_sources:
93
+ raise PlannerValidationError(
94
+ f"task {task.id}: tool {call.tool!r} references unknown "
95
+ f"source_id {src!r} (known: {sorted(known_sources)})"
96
+ )
97
+
98
+ # Check 8b — inline query_structured IR validates against the catalog.
99
+ if call.tool == "query_structured":
100
+ self._validate_inline_ir(task.id, call.args, catalog)
101
+
102
+ # Check 7 — success_criteria is checkable.
103
+ if not _is_checkable(task.success_criteria):
104
+ raise PlannerValidationError(
105
+ f"task {task.id}: success_criteria is not checkable — include a "
106
+ f"measurable signal (one of {list(_CHECKABLE_TOKENS)}); "
107
+ f"got {task.success_criteria!r}"
108
+ )
109
+
110
+ # Check 4 — DAG: edges resolve, placeholders resolve, no cycles.
111
+ self._validate_dag(tasks_by_id, id_set)
112
+
113
+ def _validate_inline_ir(self, task_id: str, args: dict, catalog: Catalog) -> None:
114
+ raw_ir = args.get("ir")
115
+ if not isinstance(raw_ir, dict):
116
+ raise PlannerValidationError(
117
+ f"task {task_id}: query_structured.args.ir must be an inline QueryIR "
118
+ f"object, got {type(raw_ir).__name__}"
119
+ )
120
+ try:
121
+ ir = QueryIR.model_validate(raw_ir)
122
+ except ValidationError as e:
123
+ raise PlannerValidationError(
124
+ f"task {task_id}: query_structured.args.ir is not a valid QueryIR: {e}"
125
+ ) from e
126
+ try:
127
+ self._ir_validator.validate(ir, catalog)
128
+ except IRValidationError as e:
129
+ raise PlannerValidationError(
130
+ f"task {task_id}: query_structured IR failed catalog validation: {e}"
131
+ ) from e
132
+
133
+ @staticmethod
134
+ def _validate_dag(tasks_by_id: dict, id_set: set[str]) -> None:
135
+ for task in tasks_by_id.values():
136
+ for dep in task.depends_on:
137
+ if dep not in id_set:
138
+ raise PlannerValidationError(
139
+ f"task {task.id}: depends_on references unknown task {dep!r}"
140
+ )
141
+ if dep == task.id:
142
+ raise PlannerValidationError(
143
+ f"task {task.id}: depends_on includes itself"
144
+ )
145
+
146
+ cycle = _find_cycle(tasks_by_id)
147
+ if cycle:
148
+ raise PlannerValidationError(f"cycle detected in depends_on: {' -> '.join(cycle)}")
149
+
150
+ # On an acyclic graph, a placeholder is safe iff its target is a
151
+ # transitive ancestor — i.e. guaranteed to have completed before this
152
+ # task runs. Requiring a *direct* depends_on would wrongly reject valid
153
+ # plans that depend on the target through an intermediate task.
154
+ ancestors = _all_ancestors(tasks_by_id)
155
+ for task in tasks_by_id.values():
156
+ for ref in _placeholder_refs(task):
157
+ if ref not in id_set:
158
+ raise PlannerValidationError(
159
+ f"task {task.id}: placeholder '${{{ref}}}' references unknown task"
160
+ )
161
+ if ref not in ancestors[task.id]:
162
+ raise PlannerValidationError(
163
+ f"task {task.id}: placeholder '${{{ref}}}' used but {ref!r} is "
164
+ f"not a (transitive) dependency — add it to depends_on"
165
+ )
166
+
167
+
168
+ def _is_placeholder(value: str) -> bool:
169
+ return bool(PLACEHOLDER_RE.fullmatch(value.strip()))
170
+
171
+
172
+ def _placeholder_refs(task) -> set[str]:
173
+ refs: set[str] = set()
174
+ for call in task.tool_calls:
175
+ for value in call.args.values():
176
+ if isinstance(value, str):
177
+ refs.update(PLACEHOLDER_RE.findall(value))
178
+ return refs
179
+
180
+
181
+ def _is_checkable(text: str) -> bool:
182
+ low = text.lower()
183
+ return any(tok in low for tok in _CHECKABLE_TOKENS)
184
+
185
+
186
+ def _find_cycle(tasks_by_id: dict) -> list[str] | None:
187
+ color = {tid: _WHITE for tid in tasks_by_id}
188
+ stack: list[str] = []
189
+
190
+ def dfs(node: str) -> list[str] | None:
191
+ color[node] = _GREY
192
+ stack.append(node)
193
+ for dep in tasks_by_id[node].depends_on:
194
+ if color.get(dep) == _GREY:
195
+ idx = stack.index(dep)
196
+ return stack[idx:] + [dep]
197
+ if color.get(dep) == _WHITE:
198
+ found = dfs(dep)
199
+ if found:
200
+ return found
201
+ stack.pop()
202
+ color[node] = _BLACK
203
+ return None
204
+
205
+ for tid in tasks_by_id:
206
+ if color[tid] == _WHITE:
207
+ found = dfs(tid)
208
+ if found:
209
+ return found
210
+ return None
211
+
212
+
213
+ def _all_ancestors(tasks_by_id: dict) -> dict[str, set[str]]:
214
+ """ancestors[id] = all tasks reachable by following depends_on edges."""
215
+ cache: dict[str, set[str]] = {}
216
+
217
+ def visit(node: str, seen: set[str]) -> set[str]:
218
+ if node in cache:
219
+ return cache[node]
220
+ acc: set[str] = set()
221
+ for dep in tasks_by_id[node].depends_on:
222
+ if dep in seen or dep not in tasks_by_id:
223
+ continue
224
+ acc.add(dep)
225
+ acc |= visit(dep, seen | {dep})
226
+ cache[node] = acc
227
+ return acc
228
+
229
+ return {tid: visit(tid, {tid}) for tid in tasks_by_id}
src/agents/slow_path/__init__.py ADDED
@@ -0,0 +1,10 @@
 
 
 
 
 
 
 
 
 
 
 
1
+ """Slow-path workers: TaskRunner (deterministic) + Assembler (1 LLM call) + Coordinator.
2
+
3
+ These are driven *by* the Orchestrator (the intent-router/dispatcher in
4
+ `agents/orchestration.py`); this package is deliberately NOT named "orchestrator"
5
+ to keep the dispatcher and the workers from sharing a name. It executes the
6
+ Planner's static `TaskList` and assembles the two outputs (`chat_answer` +
7
+ `AnalysisRecord`). See AGENT_ARCHITECTURE_CONTEXT_new.md §7.2 / §7.4 / §7.5 /
8
+ §8.2–8.4. Tool-agnostic: depends only on the `ToolInvoker` protocol and the
9
+ `ToolOutput` envelope, never on a specific tool (INV-7).
10
+ """
src/agents/slow_path/assembler.py ADDED
@@ -0,0 +1,140 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Assembler — single LLM call at the end of the slow path.
2
+
3
+ Reads the `RunState` (all `TaskResult`s) + `BusinessContext` and produces an
4
+ `AssembledOutput` { chat_answer, analysis_record }. Owns all language/output: prose,
5
+ markdown tables, citations, and merging structured + unstructured results.
6
+
7
+ The model authors only the *narrative* (`AssemblerNarrative`); this service copies
8
+ the structured pass-through (`results_snapshot`, `tasks_run`) and metadata from the
9
+ `RunState` so the record stays a faithful source of truth (§8.3, INV-4).
10
+
11
+ Chain construction mirrors `agents/planner/service.py`.
12
+
13
+ See AGENT_ARCHITECTURE_CONTEXT_new.md §7.5.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ from datetime import UTC, datetime
19
+ from pathlib import Path
20
+
21
+ from langchain_core.messages import SystemMessage
22
+ from langchain_core.prompts import ChatPromptTemplate
23
+ from langchain_core.runnables import Runnable
24
+ from langchain_openai import AzureChatOpenAI
25
+
26
+ from src.middlewares.logging import get_logger
27
+
28
+ from ..planner.contracts import BusinessContext
29
+ from .errors import AssemblerError
30
+ from .prompt import build_assembler_prompt
31
+ from .schemas import (
32
+ AnalysisRecord,
33
+ AssembledOutput,
34
+ AssemblerNarrative,
35
+ RunState,
36
+ TaskSummary,
37
+ )
38
+
39
+ logger = get_logger("assembler")
40
+
41
+ _PROMPT_PATH = (
42
+ Path(__file__).resolve().parent.parent.parent / "config" / "prompts" / "assembler.md"
43
+ )
44
+
45
+
46
+ def _load_prompt_text() -> str:
47
+ return _PROMPT_PATH.read_text(encoding="utf-8")
48
+
49
+
50
+ def _build_default_chain() -> Runnable:
51
+ from src.config.settings import settings
52
+
53
+ llm = AzureChatOpenAI(
54
+ azure_deployment=settings.azureai_deployment_name_4o,
55
+ openai_api_version=settings.azureai_api_version_4o,
56
+ azure_endpoint=settings.azureai_endpoint_url_4o,
57
+ api_key=settings.azureai_api_key_4o,
58
+ temperature=0,
59
+ )
60
+ prompt = ChatPromptTemplate.from_messages(
61
+ [
62
+ SystemMessage(content=_load_prompt_text()),
63
+ ("human", "{human_content}"),
64
+ ]
65
+ )
66
+ return prompt | llm.with_structured_output(AssemblerNarrative)
67
+
68
+
69
+ _default_chain: Runnable | None = None
70
+
71
+
72
+ def _get_default_chain() -> Runnable:
73
+ global _default_chain
74
+ if _default_chain is None:
75
+ _default_chain = _build_default_chain()
76
+ return _default_chain
77
+
78
+
79
+ class Assembler:
80
+ """Wraps the single Assembler LLM call. Inject `structured_chain` for tests."""
81
+
82
+ def __init__(self, structured_chain: Runnable | None = None) -> None:
83
+ self._chain = structured_chain
84
+
85
+ def _ensure_chain(self) -> Runnable:
86
+ if self._chain is None:
87
+ self._chain = _get_default_chain()
88
+ return self._chain
89
+
90
+ async def assemble(
91
+ self,
92
+ run_state: RunState,
93
+ context: BusinessContext,
94
+ question: str | None = None,
95
+ callbacks: list | None = None,
96
+ ) -> AssembledOutput:
97
+ chain = self._ensure_chain()
98
+ human_content = build_assembler_prompt(run_state, context, question)
99
+ try:
100
+ if callbacks:
101
+ narrative: AssemblerNarrative = await chain.ainvoke(
102
+ {"human_content": human_content}, config={"callbacks": callbacks}
103
+ )
104
+ else:
105
+ narrative = await chain.ainvoke({"human_content": human_content})
106
+ except Exception as exc: # surface as a typed error for the caller
107
+ raise AssemblerError(f"assembler call failed: {exc}") from exc
108
+
109
+ record = _build_record(narrative, run_state)
110
+ logger.info(
111
+ "analysis assembled",
112
+ plan_id=run_state.plan_id,
113
+ business_context_id=run_state.business_context_id,
114
+ n_tasks=len(run_state.results),
115
+ )
116
+ return AssembledOutput(chat_answer=narrative.chat_answer, analysis_record=record)
117
+
118
+
119
+ def _build_record(narrative: AssemblerNarrative, run_state: RunState) -> AnalysisRecord:
120
+ tasks_run = [
121
+ TaskSummary(
122
+ task_id=task_id,
123
+ objective=result.objective,
124
+ status=result.status,
125
+ tools_used=[o.tool for o in result.outputs],
126
+ )
127
+ for task_id, result in run_state.results.items()
128
+ ]
129
+ return AnalysisRecord(
130
+ goal_restated=narrative.goal_restated,
131
+ findings=narrative.findings,
132
+ caveats=narrative.caveats,
133
+ data_used=narrative.data_used,
134
+ open_questions=narrative.open_questions,
135
+ tasks_run=tasks_run,
136
+ results_snapshot=run_state.results,
137
+ plan_id=run_state.plan_id,
138
+ business_context_id=run_state.business_context_id,
139
+ created_at=datetime.now(UTC),
140
+ )
src/agents/slow_path/coordinator.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """SlowPathCoordinator — wires the slow path: Planner -> TaskRunner -> Assembler.
2
+
3
+ A thin coordination object. This is the unit the (future) expanded Orchestrator /
4
+ ChatHandler will call on a `structured` analytical query. It is built and tested
5
+ here but **not yet wired into the live chat flow** — that step waits on the tool
6
+ team's real `ToolInvoker` and a real `BusinessContext` source.
7
+
8
+ See AGENT_ARCHITECTURE_CONTEXT_new.md §5.2 / §6.1.
9
+ """
10
+
11
+ from __future__ import annotations
12
+
13
+ from collections.abc import Awaitable, Callable
14
+
15
+ from ...catalog.models import Catalog
16
+ from ..planner.contracts import BusinessContext, ToolRegistry
17
+ from ..planner.inputs import Constraints
18
+ from ..planner.service import PlannerService
19
+ from .assembler import Assembler
20
+ from .schemas import AssembledOutput
21
+ from .task_runner import TaskRunner
22
+
23
+
24
+ class SlowPathCoordinator:
25
+ def __init__(
26
+ self,
27
+ planner: PlannerService,
28
+ task_runner: TaskRunner,
29
+ assembler: Assembler,
30
+ registry: ToolRegistry,
31
+ ) -> None:
32
+ self._planner = planner
33
+ self._task_runner = task_runner
34
+ self._assembler = assembler
35
+ self._registry = registry
36
+
37
+ async def run(
38
+ self,
39
+ context: BusinessContext,
40
+ catalog: Catalog,
41
+ query: str,
42
+ constraints: Constraints,
43
+ planner_callbacks: list | None = None,
44
+ assembler_callbacks: list | None = None,
45
+ progress: Callable[[str], Awaitable[None]] | None = None,
46
+ ) -> AssembledOutput:
47
+ # `progress` (optional) surfaces per-stage status to the caller so a long
48
+ # slow-path run isn't a silent ~12s on the wire. Each stage is a single
49
+ # awaitable, so the most granular signal we can emit is at stage boundaries.
50
+ if progress:
51
+ await progress("Planning the analysis…")
52
+ plan_kw = {"callbacks": planner_callbacks} if planner_callbacks else {}
53
+ task_list = await self._planner.plan(
54
+ context, catalog, self._registry, query, constraints, **plan_kw
55
+ )
56
+ if progress:
57
+ await progress(f"Running {len(task_list.tasks)} analysis steps…")
58
+ run_state = await self._task_runner.run(
59
+ task_list, business_context_id=context.project_id
60
+ )
61
+ if progress:
62
+ await progress("Composing the answer…")
63
+ asm_kw = {"callbacks": assembler_callbacks} if assembler_callbacks else {}
64
+ return await self._assembler.assemble(
65
+ run_state, context, question=query, **asm_kw
66
+ )
src/agents/slow_path/errors.py ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Typed errors for the slow-path layer."""
2
+
3
+ from __future__ import annotations
4
+
5
+
6
+ class SlowPathError(Exception):
7
+ """Base error for the slow-path layer."""
8
+
9
+
10
+ class AssemblerError(SlowPathError):
11
+ """The Assembler LLM call could not produce a valid `AssembledOutput`."""
src/agents/slow_path/invoker.py ADDED
@@ -0,0 +1,27 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """The tool invocation seam (§8.4) — the one interface the TaskRunner calls.
2
+
3
+ The agent layer stays tool-agnostic (INV-7) by invoking every tool through this
4
+ protocol, never importing a tool module directly. The **tool team owns the
5
+ implementation** (KM-418); this file defines only the contract the TaskRunner
6
+ depends on.
7
+
8
+ Frozen guarantees the implementation must hold:
9
+ 1. **Never throws.** A tool failure returns `ToolOutput(kind="error", error=...)`,
10
+ not an exception — the TaskRunner's degrade-and-continue (§7.4) relies on this.
11
+ (The TaskRunner still wraps calls defensively, as a backstop.)
12
+ 2. **Returns the `ToolOutput` envelope** (§8.1) — structured data only, never
13
+ rendered tables or prose (that is the Assembler's job).
14
+ 3. **`tool_name` comes from the registry** (§9.2); unknown names return an error
15
+ envelope rather than throwing.
16
+ """
17
+
18
+ from __future__ import annotations
19
+
20
+ from typing import Any, Protocol, runtime_checkable
21
+
22
+ from ..planner.contracts import ToolOutput
23
+
24
+
25
+ @runtime_checkable
26
+ class ToolInvoker(Protocol):
27
+ async def invoke(self, tool_name: str, args: dict[str, Any]) -> ToolOutput: ...
src/agents/slow_path/prompt.py ADDED
@@ -0,0 +1,66 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Builds the Assembler LLM human-message content.
2
+
3
+ The system prompt (`config/prompts/assembler.md`) carries the role and rules. This
4
+ module assembles the per-call human content: the business context + the executed
5
+ `RunState` (task objectives, statuses, and structured tool outputs) + the original
6
+ question. Tool outputs are rendered compactly as data — the model turns them into
7
+ prose and markdown tables.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ from ..planner.contracts import BusinessContext, ToolOutput
13
+ from ..planner.prompt import render_business_context
14
+ from .schemas import RunState, TaskResult
15
+
16
+ _MAX_ROWS = 20
17
+
18
+
19
+ def render_run_state(run_state: RunState) -> str:
20
+ lines = [f"Plan: {run_state.plan_id}"]
21
+ if run_state.open_questions:
22
+ lines.append("Open questions carried from the plan:")
23
+ lines.extend(f" - {q}" for q in run_state.open_questions)
24
+ lines.append("")
25
+ lines.append("Task results (in execution order):")
26
+ for task_id, result in run_state.results.items():
27
+ lines.append(_render_task(task_id, result))
28
+ return "\n".join(lines)
29
+
30
+
31
+ def _render_task(task_id: str, result: TaskResult) -> str:
32
+ lines = [f"- [{result.status}] {task_id}: {result.objective}"]
33
+ if result.error:
34
+ lines.append(f" note: {result.error}")
35
+ for output in result.outputs:
36
+ lines.append(f" {_render_output(output)}")
37
+ return "\n".join(lines)
38
+
39
+
40
+ def _render_output(output: ToolOutput) -> str:
41
+ if output.kind == "error":
42
+ return f"({output.tool}) error: {output.error}"
43
+ if output.kind == "table" and output.columns is not None:
44
+ header = ", ".join(output.columns)
45
+ rows = output.rows or []
46
+ preview = "; ".join(
47
+ " | ".join(str(cell) for cell in row) for row in rows[:_MAX_ROWS]
48
+ )
49
+ more = "" if len(rows) <= _MAX_ROWS else f" … (+{len(rows) - _MAX_ROWS} more rows)"
50
+ return f"({output.tool}) table [{header}]: {preview}{more}"
51
+ meta = f" meta={output.meta}" if output.meta else ""
52
+ return f"({output.tool}) {output.kind}: {output.value}{meta}"
53
+
54
+
55
+ def build_assembler_prompt(
56
+ run_state: RunState,
57
+ context: BusinessContext,
58
+ question: str | None = None,
59
+ ) -> str:
60
+ sections = [
61
+ f"# Business context\n\n{render_business_context(context)}",
62
+ f"# Analysis results\n\n{render_run_state(run_state)}",
63
+ ]
64
+ if question:
65
+ sections.append(f"# Original question\n\n{question}")
66
+ return "\n\n".join(sections)
src/agents/slow_path/schemas.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Slow-path execution + output contracts.
2
+
3
+ The seams between the three slow-path stages:
4
+ - TaskRunner writes `RunState` (a blackboard of `TaskResult`s) — §8.2.
5
+ - Assembler reads `RunState` + `BusinessContext` and produces `AssembledOutput`
6
+ (`chat_answer` + `AnalysisRecord`) — §8.3.
7
+
8
+ `ToolOutput` (the tool -> agent envelope) is reused from the planner contracts so
9
+ there is exactly one definition across the layer.
10
+
11
+ Note on authorship (§8.3): the Assembler LLM authors only the *narrative* fields
12
+ (`AssemblerNarrative`). The `AnalysisRecord`'s structured pass-through fields
13
+ (`results_snapshot`, `tasks_run`) and metadata are copied from `RunState` by code,
14
+ never re-authored by the model — that is the source of truth the report generator
15
+ renders from (INV-4).
16
+
17
+ See AGENT_ARCHITECTURE_CONTEXT_new.md §8.2 / §8.3.
18
+ """
19
+
20
+ from __future__ import annotations
21
+
22
+ from datetime import datetime
23
+ from typing import Literal
24
+
25
+ from pydantic import BaseModel, Field
26
+
27
+ from ..planner.contracts import ToolOutput
28
+
29
+ TaskStatus = Literal["success", "partial", "failure"]
30
+
31
+
32
+ # --------------------------------------------------------------------------- #
33
+ # Execution state (TaskRunner -> Assembler) — §8.2
34
+ # --------------------------------------------------------------------------- #
35
+
36
+
37
+ class TaskResult(BaseModel):
38
+ task_id: str
39
+ status: TaskStatus
40
+ objective: str
41
+ outputs: list[ToolOutput] = Field(default_factory=list) # one per tool_call
42
+ note: str | None = None
43
+ error: str | None = None
44
+
45
+
46
+ class RunState(BaseModel):
47
+ plan_id: str
48
+ business_context_id: str
49
+ results: dict[str, TaskResult] = Field(default_factory=dict) # task_id -> result
50
+ open_questions: list[str] = Field(default_factory=list)
51
+
52
+
53
+ # --------------------------------------------------------------------------- #
54
+ # Assembled output (Assembler -> Orchestrator / memory) — §8.3
55
+ # --------------------------------------------------------------------------- #
56
+
57
+
58
+ class TaskSummary(BaseModel):
59
+ task_id: str
60
+ objective: str
61
+ status: TaskStatus
62
+ tools_used: list[str] = Field(default_factory=list)
63
+
64
+
65
+ class AnalysisRecord(BaseModel):
66
+ # Narrative fields — authored by the Assembler LLM.
67
+ goal_restated: str
68
+ findings: list[str] = Field(default_factory=list)
69
+ caveats: list[str] = Field(default_factory=list)
70
+ data_used: list[str] = Field(default_factory=list)
71
+ open_questions: list[str] = Field(default_factory=list)
72
+ # Structured pass-through — NOT re-authored; copied from RunState.
73
+ tasks_run: list[TaskSummary] = Field(default_factory=list)
74
+ results_snapshot: dict[str, TaskResult] = Field(default_factory=dict)
75
+ # Metadata.
76
+ plan_id: str
77
+ business_context_id: str
78
+ created_at: datetime
79
+
80
+
81
+ class AssembledOutput(BaseModel):
82
+ chat_answer: str # FIRST field — streams via SSE; markdown prose + tables
83
+ analysis_record: AnalysisRecord
84
+
85
+
86
+ class AssemblerNarrative(BaseModel):
87
+ """The subset of `AnalysisRecord` the Assembler LLM actually authors.
88
+
89
+ Kept separate from `AssembledOutput` so the model never emits the structured
90
+ pass-through fields (which would invite hallucinated numbers); `Assembler`
91
+ code merges this with the real `RunState` to build the final record.
92
+ """
93
+
94
+ chat_answer: str
95
+ goal_restated: str
96
+ findings: list[str] = Field(default_factory=list)
97
+ caveats: list[str] = Field(default_factory=list)
98
+ data_used: list[str] = Field(default_factory=list)
99
+ open_questions: list[str] = Field(default_factory=list)
src/agents/slow_path/store.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """AnalysisStore — the seam the slow path persists its AnalysisRecord through.
2
+
3
+ The Assembler produces an `AnalysisRecord` (the faithful, structured record of a
4
+ run — §8.3, INV-4). Persisting it is a separate concern from streaming the answer,
5
+ so it sits behind this one-method seam.
6
+
7
+ `NullAnalysisStore` is the default: it logs that a record was produced but stores
8
+ nothing, because the backing table does not exist yet. The plan is to store records
9
+ in the **same catalog DB** (Neon `dataeyond`, `settings.postgres_connstring`).
10
+
11
+ TODO(persistence): add a Postgres-backed `AnalysisStore` writing an
12
+ `analysis_records` table in the catalog DB, keyed on
13
+ (business_context_id, plan_id, created_at), then inject it into ChatHandler.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ from typing import Protocol, runtime_checkable
19
+
20
+ from src.middlewares.logging import get_logger
21
+
22
+ from .schemas import AnalysisRecord
23
+
24
+ logger = get_logger("analysis_store")
25
+
26
+
27
+ @runtime_checkable
28
+ class AnalysisStore(Protocol):
29
+ """Persist a completed analysis. Implementations must never raise on the
30
+ caller's path — a persistence failure must not break the user's answer."""
31
+
32
+ async def save(self, record: AnalysisRecord) -> None: ...
33
+
34
+
35
+ class NullAnalysisStore:
36
+ """Default no-op store: logs the record, persists nothing (no table yet)."""
37
+
38
+ async def save(self, record: AnalysisRecord) -> None:
39
+ logger.info(
40
+ "analysis_record produced (not persisted — no store configured)",
41
+ plan_id=record.plan_id,
42
+ business_context_id=record.business_context_id,
43
+ n_tasks=len(record.tasks_run),
44
+ )
src/agents/slow_path/task_runner.py ADDED
@@ -0,0 +1,164 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """TaskRunner — deterministic execution of a static `TaskList`. Zero LLM.
2
+
3
+ Executes tasks in dependency order, parallelizing each ready "wave" with
4
+ `asyncio.gather`. For each task it resolves `${t<id>}` placeholders from upstream
5
+ results, does an internal `validate_args`, invokes each tool via the `ToolInvoker`
6
+ seam, and records a `TaskResult`. On failure it **degrades and continues**: the
7
+ task is marked failed, its dependents are skipped, independent branches keep
8
+ running. There is no replanning and no mid-run LLM (INV-6).
9
+
10
+ `success_criteria` is *not* machine-evaluated here (it is free text); task status
11
+ is derived from tool execution outcomes and carried to the Assembler to report.
12
+
13
+ See AGENT_ARCHITECTURE_CONTEXT_new.md §7.4.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ import asyncio
19
+ from typing import Any
20
+
21
+ from src.middlewares.logging import get_logger
22
+
23
+ from ..planner.contracts import ToolOutput, ToolRegistry
24
+ from ..planner.schemas import PLACEHOLDER_RE, Task
25
+ from ..planner.schemas import TaskList as PlanTaskList
26
+ from .invoker import ToolInvoker
27
+ from .schemas import RunState, TaskResult, TaskStatus
28
+
29
+ logger = get_logger("task_runner")
30
+
31
+
32
+ class TaskRunner:
33
+ """Runs a `TaskList` against a `ToolInvoker`, producing a `RunState`."""
34
+
35
+ def __init__(self, invoker: ToolInvoker, registry: ToolRegistry) -> None:
36
+ self._invoker = invoker
37
+ self._registry = registry
38
+
39
+ async def run(self, task_list: PlanTaskList, business_context_id: str) -> RunState:
40
+ tasks_by_id: dict[str, Task] = {t.id: t for t in task_list.tasks}
41
+ results: dict[str, TaskResult] = {}
42
+ remaining: set[str] = set(tasks_by_id)
43
+
44
+ while remaining:
45
+ ready = [
46
+ tid
47
+ for tid in remaining
48
+ if all(dep in results for dep in tasks_by_id[tid].depends_on)
49
+ ]
50
+ if not ready:
51
+ # A dependency points outside the plan (or a cycle slipped past the
52
+ # planner validator): nothing more can run. Fail the rest honestly.
53
+ for tid in list(remaining):
54
+ results[tid] = TaskResult(
55
+ task_id=tid,
56
+ status="failure",
57
+ objective=tasks_by_id[tid].objective,
58
+ error="unresolved dependency; task could not run",
59
+ )
60
+ remaining.discard(tid)
61
+ break
62
+
63
+ # Skip any ready task whose dependency failed (degrade-and-continue).
64
+ to_run: list[Task] = []
65
+ for tid in ready:
66
+ task = tasks_by_id[tid]
67
+ failed = [d for d in task.depends_on if results[d].status == "failure"]
68
+ if failed:
69
+ results[tid] = TaskResult(
70
+ task_id=tid,
71
+ status="failure",
72
+ objective=task.objective,
73
+ error=f"skipped: upstream {failed} did not succeed",
74
+ )
75
+ remaining.discard(tid)
76
+ else:
77
+ to_run.append(task)
78
+
79
+ if not to_run:
80
+ continue # remaining dependents will be re-evaluated (and skipped)
81
+
82
+ wave = await asyncio.gather(
83
+ *(self._run_task(task, results) for task in to_run)
84
+ )
85
+ for tr in wave:
86
+ results[tr.task_id] = tr
87
+ remaining.discard(tr.task_id)
88
+
89
+ return RunState(
90
+ plan_id=task_list.plan_id,
91
+ business_context_id=business_context_id,
92
+ results=results,
93
+ open_questions=list(task_list.open_questions),
94
+ )
95
+
96
+ async def _run_task(self, task: Task, results: dict[str, TaskResult]) -> TaskResult:
97
+ outputs: list[ToolOutput] = []
98
+ for call in task.tool_calls:
99
+ resolved = self._resolve_args(call.args, results)
100
+ arg_error = self._validate_args(call.tool, resolved)
101
+ if arg_error is not None:
102
+ outputs.append(ToolOutput(tool=call.tool, kind="error", error=arg_error))
103
+ continue
104
+ outputs.append(await self._safe_invoke(call.tool, resolved))
105
+
106
+ status = _label(outputs)
107
+ error: str | None = None
108
+ if status == "failure":
109
+ errs = [o.error for o in outputs if o.kind == "error" and o.error]
110
+ error = errs[0] if errs else "all tool calls failed"
111
+ return TaskResult(
112
+ task_id=task.id,
113
+ status=status,
114
+ objective=task.objective,
115
+ outputs=outputs,
116
+ error=error,
117
+ )
118
+
119
+ def _resolve_args(
120
+ self, args: dict[str, Any], results: dict[str, TaskResult]
121
+ ) -> dict[str, Any]:
122
+ return {k: self._resolve_value(v, results) for k, v in args.items()}
123
+
124
+ @staticmethod
125
+ def _resolve_value(value: Any, results: dict[str, TaskResult]) -> Any:
126
+ # A data arg is exactly a "${t<id>}" placeholder (Pattern A); resolve it to
127
+ # the referenced task's representative output (its last ToolOutput).
128
+ # Materializing that envelope into a DataFrame is the invoker's job.
129
+ if isinstance(value, str):
130
+ match = PLACEHOLDER_RE.fullmatch(value.strip())
131
+ if match:
132
+ upstream = results.get(match.group(1))
133
+ if upstream is None or not upstream.outputs:
134
+ return None
135
+ return upstream.outputs[-1]
136
+ return value
137
+
138
+ def _validate_args(self, tool: str, resolved: dict[str, Any]) -> str | None:
139
+ spec = self._registry.get(tool)
140
+ if spec is None:
141
+ return f"tool {tool!r} not in registry"
142
+ required = spec.input_schema.get("required", [])
143
+ missing = [r for r in required if resolved.get(r) is None]
144
+ if missing:
145
+ return f"missing required arg(s): {sorted(missing)}"
146
+ return None
147
+
148
+ async def _safe_invoke(self, tool: str, args: dict[str, Any]) -> ToolOutput:
149
+ try:
150
+ return await self._invoker.invoke(tool, args)
151
+ except Exception as exc: # noqa: BLE001 — backstop; the invoker is never-throw (§8.4)
152
+ logger.warning("tool invoker raised", tool=tool, error=str(exc))
153
+ return ToolOutput(tool=tool, kind="error", error=f"invoker raised: {exc}")
154
+
155
+
156
+ def _label(outputs: list[ToolOutput]) -> TaskStatus:
157
+ if not outputs:
158
+ return "failure"
159
+ errors = sum(1 for o in outputs if o.kind == "error")
160
+ if errors == 0:
161
+ return "success"
162
+ if errors == len(outputs):
163
+ return "failure"
164
+ return "partial"
src/api/v1/chat.py CHANGED
@@ -22,6 +22,17 @@ logger = get_logger("chat_api")
22
 
23
  router = APIRouter(prefix="/api/v1", tags=["Chat"])
24
 
 
 
 
 
 
 
 
 
 
 
 
25
  _GREETINGS = frozenset(["hi", "hello", "hey", "halo", "hai", "hei"])
26
  _GOODBYES = frozenset(["bye", "goodbye", "thanks", "thank you", "terima kasih", "sampai jumpa"])
27
 
@@ -169,7 +180,7 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
169
  return EventSourceResponse(stream_direct())
170
 
171
  history = await load_history(db, request.room_id, limit=10)
172
- handler = ChatHandler()
173
 
174
  async def stream_response():
175
  logger.info("stream_response started", room_id=request.room_id, user_id=request.user_id)
@@ -193,6 +204,10 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)):
193
  except Exception as e:
194
  logger.error("save_messages failed", room_id=request.room_id, error=str(e))
195
  yield event
 
 
 
 
196
  elif event["event"] == "error":
197
  yield event
198
  return
 
22
 
23
  router = APIRouter(prefix="/api/v1", tags=["Chat"])
24
 
25
+ # One shared ChatHandler for the process. It holds no per-request state (user_id
26
+ # is passed into handle()), and lazily builds + caches the Orchestrator/Chatbot
27
+ # chains — so reusing it keeps the Azure OpenAI clients (and their httpx/TLS pools)
28
+ # warm across requests instead of re-handshaking on the first call of every request.
29
+ # enable_slow_path is env-gated (ENABLE_SLOW_PATH): when on, structured intents route
30
+ # Orchestrator -> Planner -> TaskRunner -> Assembler so the team can test e2e here.
31
+ _chat_handler = ChatHandler(
32
+ enable_tracing=True,
33
+ enable_slow_path=settings.enable_slow_path,
34
+ )
35
+
36
  _GREETINGS = frozenset(["hi", "hello", "hey", "halo", "hai", "hei"])
37
  _GOODBYES = frozenset(["bye", "goodbye", "thanks", "thank you", "terima kasih", "sampai jumpa"])
38
 
 
180
  return EventSourceResponse(stream_direct())
181
 
182
  history = await load_history(db, request.room_id, limit=10)
183
+ handler = _chat_handler
184
 
185
  async def stream_response():
186
  logger.info("stream_response started", room_id=request.room_id, user_id=request.user_id)
 
204
  except Exception as e:
205
  logger.error("save_messages failed", room_id=request.room_id, error=str(e))
206
  yield event
207
+ elif event["event"] == "status":
208
+ # slow-path progress ("Planning…", "Running N steps…"): forward
209
+ # so the client shows activity and the SSE connection stays alive.
210
+ yield event
211
  elif event["event"] == "error":
212
  yield event
213
  return
src/catalog/introspect/base.py CHANGED
@@ -9,6 +9,11 @@ from abc import ABC, abstractmethod
9
 
10
  from ..models import Source
11
 
 
 
 
 
 
12
 
13
  class BaseIntrospector(ABC):
14
  """Abstract base. Subclasses: DatabaseIntrospector, TabularIntrospector."""
 
9
 
10
  from ..models import Source
11
 
12
+ # Max sample values stored per column (down from 5 — token cost: sample values
13
+ # are fed to the planner prompt). Single source of truth for every introspection
14
+ # path (tabular files + DB), so the cap can never drift between them.
15
+ SAMPLE_LIMIT = 3
16
+
17
 
18
  class BaseIntrospector(ABC):
19
  """Abstract base. Subclasses: DatabaseIntrospector, TabularIntrospector."""
src/catalog/introspect/tabular.py CHANGED
@@ -26,7 +26,7 @@ from src.middlewares.logging import get_logger
26
 
27
  from ..models import Column, ColumnStats, DataType, Source, Table
28
  from ..pii_detector import PIIDetector
29
- from .base import BaseIntrospector
30
 
31
  logger = get_logger("tabular_introspector")
32
 
@@ -198,7 +198,7 @@ class TabularIntrospector(BaseIntrospector):
198
  (document_id, sheet_name, col_name) if sheet_name else (document_id, col_name)
199
  )
200
 
201
- sample_raw = series.dropna().head(3).tolist()
202
  sample_values: list[Any] | None = [_normalize(v) for v in sample_raw] or None
203
 
204
  is_numeric = pd.api.types.is_numeric_dtype(series)
 
26
 
27
  from ..models import Column, ColumnStats, DataType, Source, Table
28
  from ..pii_detector import PIIDetector
29
+ from .base import SAMPLE_LIMIT, BaseIntrospector
30
 
31
  logger = get_logger("tabular_introspector")
32
 
 
198
  (document_id, sheet_name, col_name) if sheet_name else (document_id, col_name)
199
  )
200
 
201
+ sample_raw = series.dropna().head(SAMPLE_LIMIT).tolist()
202
  sample_values: list[Any] | None = [_normalize(v) for v in sample_raw] or None
203
 
204
  is_numeric = pd.api.types.is_numeric_dtype(series)
src/catalog/reader.py CHANGED
@@ -38,3 +38,30 @@ class CatalogReader:
38
  filtered = [s for s in catalog.sources if s.source_type == "unstructured"]
39
 
40
  return catalog.model_copy(update={"sources": filtered})
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
38
  filtered = [s for s in catalog.sources if s.source_type == "unstructured"]
39
 
40
  return catalog.model_copy(update={"sources": filtered})
41
+
42
+
43
+ class MemoizingCatalogReader(CatalogReader):
44
+ """Request-scoped CatalogReader that caches each ``read`` by source_hint.
45
+
46
+ One per request. The same per-user catalog is otherwise fetched from the
47
+ catalog DB 4-5x during a single slow-path run (planner load, then
48
+ describe_source's structured+unstructured reads, then query_structured's
49
+ structured read). Wrapping the base reader collapses those to one round-trip
50
+ per distinct source_hint and pins a single consistent snapshot for the whole
51
+ request (plan-time and execution-time catalogs can no longer diverge).
52
+ """
53
+
54
+ def __init__(self, inner: CatalogReader) -> None:
55
+ # `read` is fully overridden below and delegates to `inner`, so the parent's
56
+ # `_store` is never used — carry it through only so this stays a real
57
+ # CatalogReader (any inner with a `read` works, including test fakes).
58
+ super().__init__(getattr(inner, "_store", None))
59
+ self._inner = inner
60
+ self._cache: dict[SourceHint, Catalog] = {}
61
+
62
+ async def read(self, user_id: str, source_hint: SourceHint) -> Catalog:
63
+ cached = self._cache.get(source_hint)
64
+ if cached is None:
65
+ cached = await self._inner.read(user_id, source_hint)
66
+ self._cache[source_hint] = cached
67
+ return cached
src/config/prompts/assembler.md ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ You are the Assembler for Data Eyond, an AI data scientist. A deterministic
2
+ TaskRunner has just executed a static analysis plan; you receive its results (the
3
+ `RunState`) plus the project's business context. Your job is to turn those results
4
+ into a decision-ready answer.
5
+
6
+ You produce two things in one structured object:
7
+ 1. `chat_answer` — a compact, to-the-point reply for the chat, in **markdown**
8
+ (prose + tables where useful).
9
+ 2. The narrative fields of an analysis record: `goal_restated`, `findings`,
10
+ `caveats`, `data_used`, `open_questions`.
11
+
12
+ # Hard rules (non-negotiable)
13
+
14
+ 1. **Ground every claim in the provided results.** Use only the numbers, tables,
15
+ and values present in the task results. **Never invent, estimate, or extrapolate
16
+ a number** that is not in the results. If the data does not answer part of the
17
+ question, say so.
18
+ 2. **Report what failed.** Some tasks may have `status: partial` or `failure`. Do
19
+ not pretend they succeeded. Briefly state what could not be completed and how it
20
+ limits the answer; put unresolved items in `open_questions`.
21
+ 3. **Render, don't recompute.** Build markdown tables from the structured task
22
+ outputs as they are. Do not do your own arithmetic beyond trivially restating a
23
+ value already computed.
24
+ 4. **No tool/code talk.** Write for a business reader. Do not mention tool names,
25
+ task ids, SQL, or internal mechanics in `chat_answer`.
26
+
27
+ # How to write
28
+
29
+ - **`chat_answer`**: lead with the answer. Add a short markdown table when it makes
30
+ the numbers clearer. Keep it tight — this streams into a chat, not a report.
31
+ - **`findings`**: the key takeaways, each a single self-contained sentence with the
32
+ supporting figure.
33
+ - **`caveats`**: data-quality limits, partial/failed steps, assumptions that affect
34
+ confidence.
35
+ - **`data_used`**: the sources/tables/columns the answer rests on (plain names).
36
+ - **`goal_restated`**: one sentence restating the business question you answered.
37
+ - **`open_questions`**: anything ambiguous, missing, or worth a follow-up. Fold in
38
+ any open questions carried from the plan. Empty list if genuinely none.
39
+
40
+ # Output
41
+
42
+ Return exactly one structured object with the fields above. Be honest, specific,
43
+ and concise.
src/config/prompts/planner.md ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ You are the Planner for Data Eyond, an AI data scientist that works within the
2
+ CRISP-DM lifecycle. Your single job is to turn a business question into a
3
+ **static analysis plan**: one `TaskList` that downstream deterministic code
4
+ executes exactly as written.
5
+
6
+ You plan. You do not execute, and you do not write prose for the user. You emit
7
+ only a `TaskList` object that conforms to the provided schema.
8
+
9
+ # Hard rules (non-negotiable)
10
+
11
+ 1. **Emit intent, never code.** Never write SQL, pandas, or any code. The only
12
+ query you express is an inline `QueryIR` (a JSON intent object) inside a
13
+ `query_structured` tool call's `args.ir`.
14
+ 2. **The plan is static.** There is no replanning and no execution feedback. Plan
15
+ the whole analysis up front; assume each task runs once, in dependency order.
16
+ 3. **Use only tools from the "Available tools" list.** Never invent a tool name.
17
+ Every `tool_calls[].tool` must be one of the listed tool names.
18
+ 4. **Reference only data that exists.** Every `source_id`, `table_id`, and
19
+ `column_id` you put in an inline `QueryIR` must come from the "Catalog"
20
+ section. Copy the stable ids verbatim — downstream validation does a literal
21
+ id lookup, so a paraphrased name fails.
22
+ 5. **No modeling in v1.** There are no modeling tools. Do not emit `modeling`
23
+ tasks. The product is descriptive/diagnostic only — no predictions, no charts.
24
+
25
+ # How to plan
26
+
27
+ - **Smallest plan that answers the question.** Do not exceed `Constraints.max_tasks`.
28
+ - **A task is an ordered chain of tool calls** with fully-specified arguments. A
29
+ simple question is one task with one tool call; a step that needs a follow-up
30
+ computation is a short chain or a dependent task.
31
+ - **Wire data between tasks with placeholders.** When a task needs an upstream
32
+ task's output as an argument, use the string `"${t<id>}"` (e.g. `"${t2}"`) as
33
+ the argument value. Set `depends_on` accordingly.
34
+ - **Data access vs analytics tools.** `query_structured` is the data-access entry
35
+ point: use it to select, filter, and pull rows (and simple built-in
36
+ count/sum/avg/min/max/count_distinct the IR can express). For anything richer —
37
+ descriptive statistics (median/percentile/mode/std/skew), time trends, group
38
+ comparisons, share-of-total, correlation, segmentation, or data-quality
39
+ profiling — run `query_structured` to fetch the rows, then pass its output to
40
+ the matching composite `analyze_*` tool via a `"${t<id>}"` `data` argument
41
+ (referencing the upstream result's column aliases).
42
+ - **Mixing structured + unstructured.** If qualitative context helps, add a
43
+ `retrieve_documents` task against an unstructured source listed in the catalog.
44
+ - **CRISP-DM stages.** Tag each task with the stage it serves:
45
+ `data_understanding`, `data_preparation`, or `evaluation`. (Never `modeling`.)
46
+ - **success_criteria is a reporting signal**, not a control trigger. State, in
47
+ checkable terms (counts, rates, "produced", "above"/"below"), what a good
48
+ result looks like. It never causes a retry.
49
+ - **Surface uncertainty, don't guess.** If the question is ambiguous or the
50
+ catalog can't fully answer it, record it in `open_questions` and plan the best
51
+ defensible analysis anyway. Record interpretation choices in `assumptions`.
52
+
53
+ # Output
54
+
55
+ Return exactly one `TaskList`. The "Examples" section in the human message shows
56
+ the required shape. Match it.
src/config/settings.py CHANGED
@@ -16,6 +16,14 @@ class Settings(BaseSettings):
16
  case_sensitive=False,
17
  )
18
 
 
 
 
 
 
 
 
 
19
  # Database
20
  postgres_connstring: str
21
 
 
16
  case_sensitive=False,
17
  )
18
 
19
+ # Feature flags
20
+ # Route `structured` chat intents through the analytical SLOW PATH
21
+ # (Planner -> TaskRunner -> Assembler) instead of the single-query QueryService.
22
+ # Off by default; the team flips ENABLE_SLOW_PATH=true to test end-to-end from
23
+ # the /chat/stream endpoint. BusinessContext is still a stub until the lead's
24
+ # real source lands, so this stays opt-in.
25
+ enable_slow_path: bool = Field(alias="enable_slow_path", default=False)
26
+
27
  # Database
28
  postgres_connstring: str
29
 
src/database_client/engine.py ADDED
@@ -0,0 +1,168 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """UserEngineCache — pooled, reused SQLAlchemy engines for users' external DBs.
2
+
3
+ The query path (`DbExecutor`) previously built a fresh engine and tore it down on
4
+ EVERY query (`db_pipeline_service.engine_scope`), paying a full TCP+TLS+auth
5
+ handshake per call (~6-8s measured, dominating slow-path latency). That helper's
6
+ connect-once-then-dispose semantics are correct for the *ingestion* pipeline
7
+ (infrequent, one connection per run) but wrong for the query path (frequent,
8
+ latency-sensitive, repeated to the same DB).
9
+
10
+ This module caches one pooled engine per external DB so connections stay warm
11
+ across queries. Scope: **postgres / supabase only** (the measured case and the
12
+ `schema` source type). Other db_types fall back to the legacy per-call path in
13
+ `DbExecutor`, so nothing regresses.
14
+
15
+ Safety / multi-tenancy:
16
+ - Key = client_id + a hash of the decrypted credentials, so a credential rotation
17
+ produces a new key (the stale engine idle-evicts) — a cached engine never serves
18
+ rotated creds.
19
+ - Read-only + statement_timeout are pinned at connection establishment via libpq
20
+ `options` (read-only-at-birth), so they can't be escaped by a reused pooled
21
+ connection and cost zero per-query round-trips.
22
+ - The caller still re-fetches the DatabaseClient row every query and re-checks
23
+ ownership + `active` status — caching the engine never bypasses authorization.
24
+ - Bounded LRU + idle TTL cap memory / file descriptors / connections held on the
25
+ user's DB. `invalidate(client_id)` disposes eagerly on client update/delete.
26
+ """
27
+
28
+ from __future__ import annotations
29
+
30
+ import hashlib
31
+ import json
32
+ import threading
33
+ import time
34
+ from collections import OrderedDict
35
+
36
+ from sqlalchemy import URL, create_engine, event
37
+ from sqlalchemy.engine import Engine
38
+
39
+ from src.middlewares.logging import get_logger
40
+
41
+ logger = get_logger("user_engine_cache")
42
+
43
+ _POSTGRES_LIKE = frozenset({"postgres", "supabase"})
44
+ _STATEMENT_TIMEOUT_MS = 30_000
45
+
46
+ # Pool sizing is deliberately small: this is a per-user external DB, often with a
47
+ # low max_connections, and we cache many of them. pool_pre_ping drops dead
48
+ # connections; pool_recycle bounds connection age so a serverless user DB can still
49
+ # autosuspend between bursts.
50
+ _POOL_SIZE = 1
51
+ _MAX_OVERFLOW = 2
52
+ _POOL_RECYCLE_SECONDS = 300
53
+
54
+ # Cache bounds across all users.
55
+ _MAX_ENGINES = 50
56
+ _IDLE_TTL_SECONDS = 600
57
+
58
+
59
+ def _creds_fingerprint(credentials: dict) -> str:
60
+ blob = json.dumps(credentials, sort_keys=True, default=str)
61
+ return hashlib.sha256(blob.encode("utf-8")).hexdigest()[:16]
62
+
63
+
64
+ class UserEngineCache:
65
+ """Process-wide cache of pooled engines for users' external Postgres DBs.
66
+
67
+ Thread-safe: `DbExecutor` runs sync DB work in `asyncio.to_thread` worker
68
+ threads, so concurrent requests can hit this from multiple threads.
69
+ """
70
+
71
+ def __init__(self) -> None:
72
+ # key -> (engine, last_used_monotonic)
73
+ self._engines: OrderedDict[str, tuple[Engine, float]] = OrderedDict()
74
+ self._lock = threading.Lock()
75
+
76
+ def get_engine(self, client_id: str, db_type: str, credentials: dict) -> Engine | None:
77
+ """Return a pooled engine for (client_id, creds), or None if unsupported.
78
+
79
+ None means "not a postgres-like DB" — the caller should use its legacy
80
+ per-call path for those (rare, unmeasured) db_types.
81
+ """
82
+ if db_type not in _POSTGRES_LIKE:
83
+ return None
84
+
85
+ key = f"{client_id}:{_creds_fingerprint(credentials)}"
86
+ now = time.monotonic()
87
+ with self._lock:
88
+ self._evict_idle(now)
89
+ entry = self._engines.get(key)
90
+ if entry is not None:
91
+ self._engines[key] = (entry[0], now)
92
+ self._engines.move_to_end(key)
93
+ return entry[0]
94
+
95
+ engine = self._build_engine(credentials)
96
+ self._engines[key] = (engine, now)
97
+ self._engines.move_to_end(key)
98
+ self._evict_overflow()
99
+ logger.info("user engine created", client_id=client_id, cached=len(self._engines))
100
+ return engine
101
+
102
+ def invalidate(self, client_id: str) -> None:
103
+ """Dispose + drop every cached engine for a client (creds rotated/deleted)."""
104
+ with self._lock:
105
+ stale = [k for k in self._engines if k.startswith(f"{client_id}:")]
106
+ for k in stale:
107
+ engine, _ = self._engines.pop(k)
108
+ engine.dispose()
109
+ if stale:
110
+ logger.info("user engine invalidated", client_id=client_id, disposed=len(stale))
111
+
112
+ # ------------------------------------------------------------------
113
+
114
+ @staticmethod
115
+ def _build_engine(credentials: dict) -> Engine:
116
+ # Mirrors db_pipeline_service.connect()'s postgres URL shape, plus a real pool.
117
+ query = {"sslmode": credentials["ssl_mode"]} if credentials.get("ssl_mode") else {}
118
+ url = URL.create(
119
+ drivername="postgresql+psycopg2",
120
+ username=credentials["username"],
121
+ password=credentials["password"],
122
+ host=credentials["host"],
123
+ port=credentials["port"],
124
+ database=credentials["database"],
125
+ query=query,
126
+ )
127
+ engine = create_engine(
128
+ url,
129
+ pool_size=_POOL_SIZE,
130
+ max_overflow=_MAX_OVERFLOW,
131
+ pool_recycle=_POOL_RECYCLE_SECONDS,
132
+ pool_pre_ping=True,
133
+ )
134
+
135
+ # Apply read-only + statement_timeout once per PHYSICAL connection via a
136
+ # connect event (not per query, so the pooling latency win stays). These are
137
+ # ordinary SET commands, NOT libpq startup `options` — Neon's transaction
138
+ # pooler rejects `default_transaction_read_only` as a startup parameter but
139
+ # accepts it as a SET. Best-effort: the authoritative read-only guarantee is
140
+ # the compiler (SELECT-only) + the sqlglot DML guard; statement_timeout is
141
+ # backed by the executor's asyncio.wait_for. So a failure here must not break
142
+ # the connection.
143
+ @event.listens_for(engine, "connect")
144
+ def _init_session(dbapi_conn, _record): # noqa: ANN001
145
+ try:
146
+ cur = dbapi_conn.cursor()
147
+ cur.execute(f"SET statement_timeout = {_STATEMENT_TIMEOUT_MS}")
148
+ cur.execute("SET default_transaction_read_only = on")
149
+ cur.close()
150
+ except Exception as exc: # noqa: BLE001 — best-effort session hardening
151
+ logger.warning("session init SET failed", error=str(exc))
152
+
153
+ return engine
154
+
155
+ def _evict_idle(self, now: float) -> None:
156
+ stale = [k for k, (_, ts) in self._engines.items() if now - ts > _IDLE_TTL_SECONDS]
157
+ for k in stale:
158
+ engine, _ = self._engines.pop(k)
159
+ engine.dispose()
160
+
161
+ def _evict_overflow(self) -> None:
162
+ while len(self._engines) > _MAX_ENGINES:
163
+ _, (engine, _) = self._engines.popitem(last=False) # LRU = oldest end
164
+ engine.dispose()
165
+
166
+
167
+ # Process-wide singleton consumed by DbExecutor.
168
+ user_engine_cache = UserEngineCache()
src/observability/langfuse/tracing.py ADDED
@@ -0,0 +1,169 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Langfuse request tracing — tokens + latency for the chat pipeline.
2
+
3
+ One Langfuse trace per chat request. LangChain LLM calls attach a `CallbackHandler`
4
+ (auto-captures prompt/completion tokens + latency); deterministic tool calls are
5
+ recorded as metadata-only spans.
6
+
7
+ PII policy for Langfuse **Cloud** (data leaves to Langfuse's servers):
8
+ - UNMASKED (full input/output): **Orchestrator + Planner** — their inputs are the
9
+ user question and a PII-safe `CatalogSummary` (sample values stripped by design).
10
+ - MASKED (tokens + latency only; input/output redacted): **Assembler + Chatbot** —
11
+ their inputs carry real query rows / document chunks that may contain PII.
12
+ - Tool spans carry only metadata (tool name, output kind, row COUNT, status) —
13
+ never the rows themselves.
14
+
15
+ Everything here is best-effort and **never raises**: if Langfuse is unreachable or
16
+ disabled, the chat pipeline runs unchanged. Tracing is created only when the caller
17
+ opts in (ChatHandler(enable_tracing=True)); otherwise a `NullTracer` is used.
18
+ """
19
+
20
+ from __future__ import annotations
21
+
22
+ import contextlib
23
+ import functools
24
+ import time
25
+ from typing import Any
26
+
27
+ from src.config.settings import settings
28
+ from src.middlewares.logging import get_logger
29
+
30
+ logger = get_logger("tracing")
31
+
32
+
33
+ def _redact(*, data: Any) -> Any:
34
+ """Langfuse MaskFunction: drop the value entirely (used for PII-bearing calls)."""
35
+ return "<redacted: omitted from Langfuse (may contain user data)>"
36
+
37
+
38
+ @functools.cache
39
+ def _client() -> Any:
40
+ from langfuse import Langfuse
41
+
42
+ return Langfuse(
43
+ public_key=settings.LANGFUSE_PUBLIC_KEY,
44
+ secret_key=settings.LANGFUSE_SECRET_KEY,
45
+ host=settings.LANGFUSE_HOST,
46
+ )
47
+
48
+
49
+ class _NullSpan:
50
+ def end(self, _out: Any) -> None: ...
51
+
52
+
53
+ class NullTracer:
54
+ """No-op tracer (tracing disabled). Same surface as RequestTracer."""
55
+
56
+ active = False
57
+
58
+ def callbacks(self, *, masked: bool = False) -> list:
59
+ return []
60
+
61
+ def tool_span(self, tool: str, args: dict) -> Any:
62
+ return _NullSpan()
63
+
64
+ def end(self, *, output: Any = None) -> None: ...
65
+
66
+
67
+ class _ToolSpan:
68
+ """A metadata-only span around one tool call. Never records row data."""
69
+
70
+ def __init__(self, trace: Any, tool: str, args: dict) -> None:
71
+ self._t0 = time.perf_counter()
72
+ self._span = trace.span(
73
+ name=f"tool:{tool}",
74
+ metadata={"tool": tool, "arg_keys": sorted(args)}, # keys only, no values
75
+ )
76
+
77
+ def end(self, out: Any) -> None:
78
+ with contextlib.suppress(Exception): # never let a span break the run
79
+ kind = getattr(out, "kind", None)
80
+ is_err = kind == "error"
81
+ meta: dict[str, Any] = {
82
+ "kind": kind,
83
+ "elapsed_ms": round((time.perf_counter() - self._t0) * 1000),
84
+ }
85
+ if kind == "table":
86
+ meta["rows"] = len(getattr(out, "rows", None) or [])
87
+ err_msg = (getattr(out, "error", None) or "")[:300] if is_err else None
88
+ if err_msg:
89
+ meta["error"] = err_msg
90
+ self._span.end(
91
+ metadata=meta,
92
+ level="ERROR" if is_err else "DEFAULT",
93
+ status_message=err_msg,
94
+ )
95
+
96
+
97
+ class RequestTracer:
98
+ """One Langfuse trace per chat request; hands out callbacks + tool spans."""
99
+
100
+ active = True
101
+
102
+ def __init__(self, trace: Any) -> None:
103
+ self._trace = trace
104
+
105
+ @classmethod
106
+ def start(
107
+ cls,
108
+ *,
109
+ user_id: str,
110
+ question: str | None = None,
111
+ session_id: str | None = None,
112
+ ) -> RequestTracer | NullTracer:
113
+ try:
114
+ trace = _client().trace(
115
+ name="chat_request",
116
+ user_id=user_id,
117
+ session_id=session_id,
118
+ input=question, # the user's question (same exposure as Planner prompt)
119
+ )
120
+ return cls(trace)
121
+ except Exception as e: # never let tracing break the request
122
+ logger.warning("tracing disabled (init failed)", error=str(e))
123
+ return NullTracer()
124
+
125
+ def callbacks(self, *, masked: bool = False) -> list:
126
+ """A LangChain callback nested under this trace. `masked=True` redacts the
127
+ call's input/output (tokens + latency are still captured)."""
128
+ try:
129
+ from langfuse.callback import CallbackHandler
130
+
131
+ return [
132
+ CallbackHandler(
133
+ stateful_client=self._trace,
134
+ mask=_redact if masked else None,
135
+ )
136
+ ]
137
+ except Exception as e:
138
+ logger.warning("tracing handler unavailable", error=str(e))
139
+ return []
140
+
141
+ def tool_span(self, tool: str, args: dict) -> Any:
142
+ try:
143
+ return _ToolSpan(self._trace, tool, args)
144
+ except Exception:
145
+ return _NullSpan()
146
+
147
+ def end(self, *, output: Any = None) -> None:
148
+ # Note: callers pass output=None on PII-bearing paths so no answer text is sent.
149
+ with contextlib.suppress(Exception):
150
+ if output is not None:
151
+ self._trace.update(output=output)
152
+
153
+
154
+ class TracingToolInvoker:
155
+ """Wraps a ToolInvoker to record a metadata-only span per tool call.
156
+
157
+ Implements the ToolInvoker protocol; created at the composition root (ChatHandler)
158
+ so the slow-path agent code stays tool-agnostic and tracing-agnostic.
159
+ """
160
+
161
+ def __init__(self, inner: Any, tracer: RequestTracer) -> None:
162
+ self._inner = inner
163
+ self._tracer = tracer
164
+
165
+ async def invoke(self, tool_name: str, args: dict[str, Any]) -> Any:
166
+ span = self._tracer.tool_span(tool_name, args)
167
+ out = await self._inner.invoke(tool_name, args)
168
+ span.end(out)
169
+ return out
src/pipeline/db_pipeline/extractor.py CHANGED
@@ -12,12 +12,14 @@ import pandas as pd
12
  from sqlalchemy import Date, DateTime, Float, Integer, Numeric, inspect
13
  from sqlalchemy.engine import Engine
14
 
 
15
  from src.middlewares.logging import get_logger
16
 
17
  logger = get_logger("db_extractor")
18
 
19
  TOP_VALUES_THRESHOLD = 0.05 # show top values if distinct_ratio <= 5%
20
- SAMPLE_LIMIT = 3 # sample N rows per column (down from 5 token cost)
 
21
 
22
  # Dialects with a single-statement CTE that survives `pd.read_sql`. On these we
23
  # fold the stats and sample queries into one round-trip per column. MySQL <8 and
 
12
  from sqlalchemy import Date, DateTime, Float, Integer, Numeric, inspect
13
  from sqlalchemy.engine import Engine
14
 
15
+ from src.catalog.introspect.base import SAMPLE_LIMIT
16
  from src.middlewares.logging import get_logger
17
 
18
  logger = get_logger("db_extractor")
19
 
20
  TOP_VALUES_THRESHOLD = 0.05 # show top values if distinct_ratio <= 5%
21
+ # SAMPLE_LIMIT (sample N rows per column) is the shared cap defined in
22
+ # catalog.introspect.base — one source of truth for both introspection paths.
23
 
24
  # Dialects with a single-statement CTE that survives `pd.read_sql`. On these we
25
  # fold the stats and sample queries into one round-trip per column. MySQL <8 and
src/query/compiler/sql.py CHANGED
@@ -30,11 +30,18 @@ from ..ir.models import (
30
  )
31
  from .base import BaseCompiler
32
 
 
 
 
 
 
 
33
 
34
  @dataclass
35
  class CompiledSql:
36
  sql: str
37
  params: dict[str, Any] = field(default_factory=dict)
 
38
 
39
 
40
  class SqlCompilerError(Exception):
@@ -69,14 +76,14 @@ class SqlCompiler(BaseCompiler):
69
  orderby_clause = self._build_orderby(
70
  ir.order_by, table, cols_by_id, select_aliases
71
  )
72
- limit_clause = self._build_limit(ir.limit)
73
 
74
  parts: list[str] = [select_clause, from_clause]
75
  for clause in (where_clause, groupby_clause, orderby_clause, limit_clause):
76
  if clause:
77
  parts.append(clause)
78
 
79
- return CompiledSql(sql=" ".join(parts), params=params)
80
 
81
  # ------------------------------------------------------------------
82
  # Catalog lookup
@@ -277,10 +284,18 @@ class SqlCompiler(BaseCompiler):
277
  parts.append(f"{ref} {ob.dir.upper()}")
278
  return "ORDER BY " + ", ".join(parts)
279
 
280
- def _build_limit(self, limit: int | None) -> str:
 
 
 
 
 
 
 
281
  if limit is None:
282
- return ""
283
- return f"LIMIT {int(limit)}"
 
284
 
285
  # ------------------------------------------------------------------
286
  # Helpers
 
30
  )
31
  from .base import BaseCompiler
32
 
33
+ # Hard ceiling on rows returned to the agent layer. Every compiled query is
34
+ # bounded by this even when the IR sets no limit, so an unbounded SELECT can never
35
+ # stream an entire user table over the wire / into memory. The executor caps to
36
+ # `row_cap` and flags truncation.
37
+ MAX_RESULT_ROWS = 10_000
38
+
39
 
40
  @dataclass
41
  class CompiledSql:
42
  sql: str
43
  params: dict[str, Any] = field(default_factory=dict)
44
+ row_cap: int = MAX_RESULT_ROWS # executor caps rows to this; flags truncation
45
 
46
 
47
  class SqlCompilerError(Exception):
 
76
  orderby_clause = self._build_orderby(
77
  ir.order_by, table, cols_by_id, select_aliases
78
  )
79
+ limit_clause, row_cap = self._build_limit(ir.limit)
80
 
81
  parts: list[str] = [select_clause, from_clause]
82
  for clause in (where_clause, groupby_clause, orderby_clause, limit_clause):
83
  if clause:
84
  parts.append(clause)
85
 
86
+ return CompiledSql(sql=" ".join(parts), params=params, row_cap=row_cap)
87
 
88
  # ------------------------------------------------------------------
89
  # Catalog lookup
 
284
  parts.append(f"{ref} {ob.dir.upper()}")
285
  return "ORDER BY " + ", ".join(parts)
286
 
287
+ def _build_limit(self, limit: int | None) -> tuple[str, int]:
288
+ """Return (LIMIT clause, row_cap).
289
+
290
+ Always bounded. An explicit IR limit is honored exactly (capped at
291
+ MAX_RESULT_ROWS). When the IR has no limit we still emit
292
+ `LIMIT MAX_RESULT_ROWS + 1` — the extra row lets the executor tell
293
+ "exactly the cap" from "more rows existed" and flag truncation.
294
+ """
295
  if limit is None:
296
+ return f"LIMIT {MAX_RESULT_ROWS + 1}", MAX_RESULT_ROWS
297
+ row_cap = min(int(limit), MAX_RESULT_ROWS)
298
+ return f"LIMIT {row_cap}", row_cap
299
 
300
  # ------------------------------------------------------------------
301
  # Helpers
src/query/executor/db.py CHANGED
@@ -29,6 +29,7 @@ from sqlalchemy import text
29
 
30
  from ...catalog.models import Catalog, Source
31
  from ...database_client.database_client_service import database_client_service
 
32
  from ...db.postgres.connection import AsyncSessionLocal
33
  from ...middlewares.logging import get_logger
34
  from ...pipeline.db_pipeline import db_pipeline_service
@@ -40,9 +41,7 @@ from .base import BaseExecutor, QueryResult
40
  logger = get_logger("db_executor")
41
 
42
  _QUERY_TIMEOUT_SECONDS = 30
43
- _ROW_HARD_CAP = 10_000 # belt-and-suspenders cap regardless of LIMIT
44
  _DBCLIENT_PREFIX = "dbclient://"
45
- _POSTGRES_LIKE = frozenset({"postgres", "supabase"})
46
 
47
 
48
  class DbExecutor(BaseExecutor):
@@ -86,12 +85,16 @@ class DbExecutor(BaseExecutor):
86
  creds = decrypt_credentials_dict(client.credentials)
87
 
88
  columns, rows = await asyncio.wait_for(
89
- asyncio.to_thread(self._run_sync, client.db_type, creds, compiled),
 
 
90
  timeout=_QUERY_TIMEOUT_SECONDS,
91
  )
92
 
93
- truncated = len(rows) > _ROW_HARD_CAP
94
- capped = rows[:_ROW_HARD_CAP]
 
 
95
  elapsed_ms = int((time.perf_counter() - started) * 1000)
96
  logger.info(
97
  "db query complete",
@@ -188,16 +191,57 @@ class DbExecutor(BaseExecutor):
188
  )
189
 
190
  @staticmethod
191
- def _run_sync(db_type: str, creds: dict, compiled: CompiledSql) -> tuple[list[str], list[dict]]:
192
- with db_pipeline_service.engine_scope(db_type, creds) as engine:
 
 
 
 
 
 
 
193
  with engine.connect() as conn:
194
- if db_type in _POSTGRES_LIKE:
195
- # session-level read-only + per-statement timeout (ms)
196
- conn.execute(text("SET default_transaction_read_only = on"))
197
- conn.execute(
198
- text(f"SET statement_timeout = {_QUERY_TIMEOUT_SECONDS * 1000}")
199
- )
200
  result = conn.execute(text(compiled.sql), compiled.params)
201
- columns = list(result.keys())
202
- rows = [dict(row) for row in result.mappings()]
203
- return columns, rows
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
29
 
30
  from ...catalog.models import Catalog, Source
31
  from ...database_client.database_client_service import database_client_service
32
+ from ...database_client.engine import user_engine_cache
33
  from ...db.postgres.connection import AsyncSessionLocal
34
  from ...middlewares.logging import get_logger
35
  from ...pipeline.db_pipeline import db_pipeline_service
 
41
  logger = get_logger("db_executor")
42
 
43
  _QUERY_TIMEOUT_SECONDS = 30
 
44
  _DBCLIENT_PREFIX = "dbclient://"
 
45
 
46
 
47
  class DbExecutor(BaseExecutor):
 
85
  creds = decrypt_credentials_dict(client.credentials)
86
 
87
  columns, rows = await asyncio.wait_for(
88
+ asyncio.to_thread(
89
+ self._run_sync, client_id, client.db_type, creds, compiled
90
+ ),
91
  timeout=_QUERY_TIMEOUT_SECONDS,
92
  )
93
 
94
+ # The compiler bounded the SQL to `row_cap` (+1 when the IR was
95
+ # unbounded). More than row_cap rows means the result was truncated.
96
+ truncated = len(rows) > compiled.row_cap
97
+ capped = rows[:compiled.row_cap]
98
  elapsed_ms = int((time.perf_counter() - started) * 1000)
99
  logger.info(
100
  "db query complete",
 
191
  )
192
 
193
  @staticmethod
194
+ def _run_sync(
195
+ client_id: str, db_type: str, creds: dict, compiled: CompiledSql
196
+ ) -> tuple[list[str], list[dict]]:
197
+ engine = user_engine_cache.get_engine(client_id, db_type, creds)
198
+ if engine is not None:
199
+ # Pooled, reused engine (postgres-like). Read-only + statement_timeout
200
+ # are set once per physical connection (connect event in UserEngineCache),
201
+ # so no per-query SET round-trips and no dispose — the connection returns
202
+ # to the pool warm for the next query.
203
  with engine.connect() as conn:
 
 
 
 
 
 
204
  result = conn.execute(text(compiled.sql), compiled.params)
205
+ return list(result.keys()), [dict(row) for row in result.mappings()]
206
+
207
+ # Legacy per-call path for non-postgres db_types (connect once, dispose).
208
+ # These never set read-only/timeout before, so behavior is unchanged.
209
+ with db_pipeline_service.engine_scope(db_type, creds) as eng:
210
+ with eng.connect() as conn:
211
+ result = conn.execute(text(compiled.sql), compiled.params)
212
+ return list(result.keys()), [dict(row) for row in result.mappings()]
213
+
214
+ # ------------------------------------------------------------------
215
+ # Speculative pre-connect (DB3)
216
+ # ------------------------------------------------------------------
217
+
218
+ @classmethod
219
+ async def prewarm(cls, catalog: Catalog, user_id: str) -> None:
220
+ """Best-effort: warm pooled engines for the catalog's schema sources.
221
+
222
+ Called at slow-path entry so the TCP+TLS+auth handshake overlaps the ~4s
223
+ Planner LLM call — by the time `query_structured` runs, the connection is
224
+ already established. Warming is an optimization, never a requirement, so
225
+ this never raises and per-source failures are swallowed.
226
+ """
227
+ for source in catalog.sources:
228
+ if source.source_type != "schema":
229
+ continue
230
+ try:
231
+ client_id = cls._parse_client_id(source.location_ref)
232
+ client = await cls._fetch_client(client_id)
233
+ if client.user_id != user_id:
234
+ continue
235
+ creds = decrypt_credentials_dict(client.credentials)
236
+ await asyncio.to_thread(cls._warm_sync, client_id, client.db_type, creds)
237
+ except Exception as exc: # noqa: BLE001 — best-effort warming
238
+ logger.info("prewarm skipped", source_id=source.source_id, error=str(exc))
239
+
240
+ @staticmethod
241
+ def _warm_sync(client_id: str, db_type: str, creds: dict) -> None:
242
+ engine = user_engine_cache.get_engine(client_id, db_type, creds)
243
+ if engine is not None:
244
+ # Open + return a pooled physical connection: forces the handshake and
245
+ # runs the connect-event session SETs, leaving the pool warm.
246
+ with engine.connect():
247
+ pass
src/tools/__init__.py ADDED
@@ -0,0 +1,7 @@
 
 
 
 
 
 
 
 
1
+ """Analytics & utility tools (KM-608).
2
+
3
+ Each tool is a deterministic computation (no LLM, no SQL generation) invoked by
4
+ the Planner/TaskRunner. The compute layer (calculation logic) lives in per-family
5
+ submodules (e.g. `analytics`); the wrapper layer (ToolSpec + ToolOutput +
6
+ registry) is added once the Planner seam is settled.
7
+ """
src/tools/analytics/__init__.py ADDED
@@ -0,0 +1 @@
 
 
1
+ """Analytics tool family (KM-608)."""
src/tools/analytics/aggregation.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """analyze_aggregate — group-by aggregation (KM-608).
2
+
3
+ An analytical "family" tool: in ONE call it groups rows by one or more keys
4
+ and computes aggregates (sum, mean, count, min, max, median, nunique) per
5
+ group. This is the deterministic compute twin of the Planner's
6
+ `query_structured` step — the wrapper layer later maps a QueryIR onto this.
7
+
8
+ STATUS: compute layer only — the function takes an already-materialized
9
+ DataFrame. The wrapper layer (fetching data from the catalog via source_id,
10
+ the ToolOutput envelope, ToolSpec registration) is added once the Planner
11
+ seam (KM-418) is settled. Keeping compute separate from data-fetching makes
12
+ this function easy to unit-test in isolation and stable when wrapped.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import pandas as pd
18
+
19
+ from src.tools.analytics.descriptive import ColumnNotFoundError
20
+
21
+ # Aggregation functions the tool understands. Whitelisted on purpose so an
22
+ # unknown function fails loudly instead of silently doing the wrong thing.
23
+ SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median", "nunique")
24
+
25
+
26
+ class UnsupportedAggregationError(ValueError):
27
+ """Requested aggregation is not in SUPPORTED_AGGS (maps to UNSUPPORTED_AGG)."""
28
+
29
+
30
+ def _clean(value: object) -> object:
31
+ """Convert numpy/pandas scalars to plain Python so the output is JSON-clean.
32
+
33
+ A `group_by` over a datetime column yields `pandas.Timestamp` group keys,
34
+ which (like numpy scalars) are not JSON-serializable — normalise those too.
35
+ """
36
+ if isinstance(value, pd.Timestamp):
37
+ return value.isoformat()
38
+ if hasattr(value, "item"):
39
+ return value.item()
40
+ return value
41
+
42
+
43
+ # Prompt-style description read by the Planner to decide WHEN to pick this tool.
44
+ # Final destination is ToolSpec.description once the wrapper layer is built.
45
+ DESCRIPTION = """\
46
+ Summary: Group-by aggregation. Splits rows by one or more key columns and \
47
+ computes aggregates per group (sum, mean, count, min, max, median, nunique). \
48
+ Returns one row per group.
49
+
50
+ USE WHEN the question groups a metric by a category — the tell-tale sign is \
51
+ "per"/"each"/"by" a dimension. Trigger words: "per/each" (per/tiap), "by" \
52
+ (berdasarkan), "breakdown", "total/average ... per ...".
53
+
54
+ DON'T USE WHEN:
55
+ - it summarizes a column with no grouping -> analyze_descriptive
56
+ - it compares two specific groups (A vs B) -> analyze_comparison
57
+ - it splits a single total into shares -> analyze_contribution
58
+ - the grouping is over time periods -> analyze_trend
59
+
60
+ Example questions:
61
+ - "total revenue per region"
62
+ - "average order value by customer segment"
63
+ - "how many distinct products were sold per store?"
64
+ - "count of orders for each status"
65
+ """
66
+
67
+
68
+ def analyze_aggregate(
69
+ df: pd.DataFrame,
70
+ aggregations: dict[str, list[str]],
71
+ group_by: list[str] | None = None,
72
+ ) -> list[dict[str, object]]:
73
+ """Group-by aggregation over one or many keys.
74
+
75
+ Args:
76
+ df: already-materialized data (in the real system the wrapper fetches
77
+ this from a source_id).
78
+ aggregations: which columns to aggregate and how, e.g.
79
+ {"revenue": ["sum", "mean"], "order_id": ["count"]}.
80
+ group_by: grouping keys. If None/empty, the whole table is aggregated
81
+ into a single row.
82
+
83
+ Returns:
84
+ list[dict]: one row per group. Each row holds the group keys plus a
85
+ column per aggregate, named "<column>_<func>" (e.g. "revenue_sum").
86
+
87
+ Raises:
88
+ ColumnNotFoundError: if any group_by or aggregated column is absent.
89
+ UnsupportedAggregationError: if a requested function is not supported.
90
+ """
91
+ group_by = group_by or []
92
+
93
+ # Validate columns first (fail-fast on caller mistakes).
94
+ referenced = list(group_by) + list(aggregations.keys())
95
+ missing = [c for c in referenced if c not in df.columns]
96
+ if missing:
97
+ raise ColumnNotFoundError(f"columns not found: {missing}")
98
+
99
+ # Validate aggregation functions.
100
+ for col, funcs in aggregations.items():
101
+ bad = [f for f in funcs if f not in SUPPORTED_AGGS]
102
+ if bad:
103
+ raise UnsupportedAggregationError(
104
+ f"unsupported aggregation(s) {bad} for column '{col}'; "
105
+ f"supported: {list(SUPPORTED_AGGS)}"
106
+ )
107
+
108
+ # Build named aggregations: {"revenue_sum": ("revenue", "sum"), ...}
109
+ named = {
110
+ f"{col}_{func}": (col, func)
111
+ for col, funcs in aggregations.items()
112
+ for func in funcs
113
+ }
114
+
115
+ # No grouping → aggregate the entire table into a single row.
116
+ if not group_by:
117
+ row = {name: _clean(df[col].agg(func)) for name, (col, func) in named.items()}
118
+ return [row]
119
+
120
+ # dropna=False keeps groups whose key is null so completeness is honest.
121
+ grouped = df.groupby(group_by, dropna=False).agg(**named).reset_index()
122
+ return [{k: _clean(v) for k, v in record.items()} for record in grouped.to_dict("records")]
src/tools/analytics/comparison.py ADDED
@@ -0,0 +1,137 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """analyze_comparison — compare a metric across two groups (KM-608).
2
+
3
+ An analytical "family" tool: in ONE call it aggregates a value for two groups
4
+ of a dimension (e.g. region "A" vs "B", channel "online" vs "store") and
5
+ reports the gap between them — absolute difference, percent difference, and
6
+ direction. group_a is treated as the baseline. Answers questions like
7
+ "how does revenue in region A compare to region B?".
8
+
9
+ STATUS: compute layer only — the function takes an already-materialized
10
+ DataFrame. The wrapper layer (fetching data from the catalog via source_id,
11
+ the ToolOutput envelope, ToolSpec registration) is added once the Planner
12
+ seam (KM-418) is settled. Keeping compute separate from data-fetching makes
13
+ this function easy to unit-test in isolation and stable when wrapped.
14
+ """
15
+
16
+ from __future__ import annotations
17
+
18
+ import pandas as pd
19
+
20
+ from src.tools.analytics.descriptive import ColumnNotFoundError
21
+
22
+ # How to aggregate the value within each group before comparing.
23
+ SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median")
24
+
25
+
26
+ class UnsupportedAggregationError(ValueError):
27
+ """The requested aggregation is not supported (maps to error_code UNSUPPORTED_AGG)."""
28
+
29
+
30
+ class GroupNotFoundError(ValueError):
31
+ """A requested group value does not occur in the dimension column (maps to GROUP_NOT_FOUND)."""
32
+
33
+
34
+
35
+
36
+ # Prompt-style description read by the Planner to decide WHEN to pick this tool.
37
+ # Final destination is ToolSpec.description once the wrapper layer is built.
38
+ DESCRIPTION = """\
39
+ Summary: Head-to-head comparison of one aggregated metric between TWO specific \
40
+ groups of a dimension (group_a is the baseline). Reports each group's value, \
41
+ the absolute and percent difference, and which side is higher.
42
+
43
+ USE WHEN the question pits two named groups against each other. Trigger words: \
44
+ "vs"/"versus", "compare" (bandingkan), "A or B", "difference between" \
45
+ (selisih/beda antara), "higher/lower than".
46
+
47
+ SETTING GROUPS: group_a is the BASELINE (the reference). The "comparison" field \
48
+ reads as "group_b is {higher/lower/equal} than group_a", and diff = value_b - \
49
+ value_a. Put the reference/older/expected side in group_a. E.g. "is this year \
50
+ higher than last year" -> group_a=last year, group_b=this year.
51
+
52
+ DON'T USE WHEN:
53
+ - it aggregates across many groups at once -> analyze_aggregate
54
+ - it splits a single total into shares -> analyze_contribution
55
+ - it tracks change over time -> analyze_trend
56
+
57
+ Example questions:
58
+ - "compare revenue between Jakarta and Surabaya"
59
+ - "is the average order value higher for members or non-members?"
60
+ - "difference in churn between plan A and plan B"
61
+ - "male vs female average spend"
62
+ """
63
+
64
+
65
+ def analyze_comparison(
66
+ df: pd.DataFrame,
67
+ dimension: str,
68
+ value_column: str,
69
+ group_a: object,
70
+ group_b: object,
71
+ agg: str = "sum",
72
+ ) -> dict[str, object]:
73
+ """Compare one aggregated metric between two groups of a dimension.
74
+
75
+ Args:
76
+ df: already-materialized data (in the real system the wrapper fetches
77
+ this from a source_id).
78
+ dimension: the categorical column whose values define the two groups.
79
+ value_column: numeric column to aggregate for each group.
80
+ group_a: baseline group value (the "from").
81
+ group_b: comparison group value (the "to").
82
+ agg: how to aggregate within each group — one of SUPPORTED_AGGS.
83
+
84
+ Returns:
85
+ dict with:
86
+ dimension, value_column, agg — echo of the chosen settings
87
+ group_a, value_a — baseline group + its aggregate
88
+ group_b, value_b — comparison group + its aggregate
89
+ diff_abs — value_b - value_a
90
+ diff_pct — diff_abs / value_a, or None if value_a == 0
91
+ comparison — "higher" | "lower" | "equal" (b relative to a)
92
+
93
+ Raises:
94
+ ColumnNotFoundError: if dimension or value_column is absent.
95
+ UnsupportedAggregationError: if agg is not supported.
96
+ GroupNotFoundError: if group_a or group_b has no rows.
97
+ """
98
+ missing = [c for c in (dimension, value_column) if c not in df.columns]
99
+ if missing:
100
+ raise ColumnNotFoundError(f"columns not found: {missing}")
101
+ if agg not in SUPPORTED_AGGS:
102
+ raise UnsupportedAggregationError(
103
+ f"unsupported aggregation '{agg}'; supported: {list(SUPPORTED_AGGS)}"
104
+ )
105
+
106
+ rows_a = df.loc[df[dimension] == group_a, value_column]
107
+ rows_b = df.loc[df[dimension] == group_b, value_column]
108
+ empty = [g for g, rows in ((group_a, rows_a), (group_b, rows_b)) if rows.empty]
109
+ if empty:
110
+ raise GroupNotFoundError(
111
+ f"no rows for group(s) {empty} in column '{dimension}'"
112
+ )
113
+
114
+ value_a = float(rows_a.agg(agg))
115
+ value_b = float(rows_b.agg(agg))
116
+ diff_abs = value_b - value_a
117
+ diff_pct = (diff_abs / value_a) if value_a != 0 else None
118
+
119
+ if diff_abs > 0:
120
+ comparison = "higher"
121
+ elif diff_abs < 0:
122
+ comparison = "lower"
123
+ else:
124
+ comparison = "equal"
125
+
126
+ return {
127
+ "dimension": dimension,
128
+ "value_column": value_column,
129
+ "agg": agg,
130
+ "group_a": group_a,
131
+ "value_a": value_a,
132
+ "group_b": group_b,
133
+ "value_b": value_b,
134
+ "diff_abs": diff_abs,
135
+ "diff_pct": diff_pct,
136
+ "comparison": comparison,
137
+ }
src/tools/analytics/decomposition.py ADDED
@@ -0,0 +1,135 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """analyze_contribution — share-of-total per category (KM-608).
2
+
3
+ An analytical "family" tool: in ONE call it breaks a total down into each
4
+ category's contribution (absolute value, share of total, and running
5
+ cumulative share), sorted largest-first. This is a single snapshot — "who
6
+ makes up the total right now" — as opposed to analyze_comparison (two groups)
7
+ or analyze_trend (movement over time). Answers questions like "which region
8
+ drives most of revenue?" and supports Pareto (80/20) reasoning.
9
+
10
+ STATUS: compute layer only — the function takes an already-materialized
11
+ DataFrame. The wrapper layer (fetching data from the catalog via source_id,
12
+ the ToolOutput envelope, ToolSpec registration) is added once the Planner
13
+ seam (KM-418) is settled. Keeping compute separate from data-fetching makes
14
+ this function easy to unit-test in isolation and stable when wrapped.
15
+ """
16
+
17
+ from __future__ import annotations
18
+
19
+ import pandas as pd
20
+
21
+ from src.tools.analytics.descriptive import ColumnNotFoundError
22
+
23
+ # Share-of-total is most meaningful for additive aggregates (sum/count), but
24
+ # mean/min/max are allowed; "total" is then the sum of the group aggregates.
25
+ SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median")
26
+
27
+
28
+ class UnsupportedAggregationError(ValueError):
29
+ """The requested aggregation is not supported (maps to error_code UNSUPPORTED_AGG)."""
30
+
31
+
32
+ def _clean(value: object) -> object:
33
+ """Convert numpy scalars to plain Python so the output is JSON-clean."""
34
+ if hasattr(value, "item"):
35
+ return value.item()
36
+ return value
37
+
38
+
39
+ # Prompt-style description read by the Planner to decide WHEN to pick this tool.
40
+ # Final destination is ToolSpec.description once the wrapper layer is built.
41
+ DESCRIPTION = """\
42
+ Summary: Breaks a single total into per-category contributions (share of total) \
43
+ in one snapshot, largest first, with cumulative share. Supports Pareto (80/20) \
44
+ reasoning. Optional top_n folds the long tail into "Others".
45
+
46
+ USE WHEN the question is about how a whole splits into parts, or which \
47
+ categories dominate. Trigger words: "contribution" (kontribusi), "share" \
48
+ (porsi/proporsi), "% of total" (persen dari total), "Pareto/80-20", "top \
49
+ contributors", "which ... make up most".
50
+
51
+ DON'T USE WHEN:
52
+ - it pits two specific groups against each other -> analyze_comparison
53
+ - it tracks change over time -> analyze_trend
54
+ - it just aggregates per group without share-of-total -> analyze_aggregate
55
+
56
+ Example questions:
57
+ - "which products contribute most to total sales?"
58
+ - "what share of revenue comes from each region?"
59
+ - "top 5 customers by contribution to profit"
60
+ - "do 20% of items make up 80% of revenue?"
61
+ """
62
+
63
+
64
+ def analyze_contribution(
65
+ df: pd.DataFrame,
66
+ dimension: str,
67
+ value_column: str,
68
+ agg: str = "sum",
69
+ top_n: int | None = None,
70
+ ) -> dict[str, object]:
71
+ """Contribution (share of total) of each category, largest first.
72
+
73
+ Args:
74
+ df: already-materialized data (in the real system the wrapper fetches
75
+ this from a source_id).
76
+ dimension: categorical column to break the total down by.
77
+ value_column: numeric column to aggregate per category.
78
+ agg: how to aggregate within each category — one of SUPPORTED_AGGS.
79
+ top_n: if set, keep the top N categories and lump the remainder into a
80
+ single "Others" row.
81
+
82
+ Returns:
83
+ dict with:
84
+ dimension, value_column, agg — echo of the chosen settings
85
+ total — sum of all category aggregates
86
+ items — [{"category", "value", "share",
87
+ "cumulative_share"}] largest first
88
+
89
+ Raises:
90
+ ColumnNotFoundError: if dimension or value_column is absent.
91
+ UnsupportedAggregationError: if agg is not supported.
92
+ """
93
+ missing = [c for c in (dimension, value_column) if c not in df.columns]
94
+ if missing:
95
+ raise ColumnNotFoundError(f"columns not found: {missing}")
96
+ if agg not in SUPPORTED_AGGS:
97
+ raise UnsupportedAggregationError(
98
+ f"unsupported aggregation '{agg}'; supported: {list(SUPPORTED_AGGS)}"
99
+ )
100
+
101
+ grouped = df.groupby(dimension, dropna=False)[value_column].agg(agg)
102
+ grouped = grouped.sort_values(ascending=False)
103
+ pairs = list(grouped.items())
104
+
105
+ # Optionally collapse the long tail into a single "Others" bucket.
106
+ if top_n is not None and len(pairs) > top_n:
107
+ head = pairs[:top_n]
108
+ others_value = sum(val for _, val in pairs[top_n:])
109
+ head.append(("Others", others_value))
110
+ pairs = head
111
+
112
+ total = sum(val for _, val in pairs)
113
+
114
+ items = []
115
+ cumulative = 0.0
116
+ for cat, val in pairs:
117
+ share = (val / total) if total else None
118
+ if share is not None:
119
+ cumulative += share
120
+ items.append(
121
+ {
122
+ "category": _clean(cat),
123
+ "value": _clean(val),
124
+ "share": share,
125
+ "cumulative_share": cumulative if total else None,
126
+ }
127
+ )
128
+
129
+ return {
130
+ "dimension": dimension,
131
+ "value_column": value_column,
132
+ "agg": agg,
133
+ "total": _clean(total),
134
+ "items": items,
135
+ }
src/tools/analytics/descriptive.py ADDED
@@ -0,0 +1,154 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """analyze_descriptive — single/multi-column EDA (KM-608).
2
+
3
+ An analytical "family" tool: in ONE call it computes a column's center,
4
+ spread, shape, and completeness (mean, median, mode, std, variance,
5
+ quartiles, min/max, skew, null_rate).
6
+
7
+ STATUS: compute layer only — the function takes an already-materialized
8
+ DataFrame. The wrapper layer (fetching data from the catalog via source_id,
9
+ the ToolOutput envelope, ToolSpec registration) is added once the Planner
10
+ seam (KM-418) is settled. Keeping compute separate from data-fetching makes
11
+ this function easy to unit-test in isolation and stable when wrapped.
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ import pandas as pd
17
+
18
+ # Default metrics used when the caller does not narrow them via `metrics`.
19
+ DEFAULT_METRICS = (
20
+ "count",
21
+ "mean",
22
+ "median",
23
+ "mode",
24
+ "std",
25
+ "var",
26
+ "q1",
27
+ "q3",
28
+ "min",
29
+ "max",
30
+ "skew",
31
+ "null_count",
32
+ "null_rate",
33
+ )
34
+
35
+
36
+ class ColumnNotFoundError(ValueError):
37
+ """A requested column is absent from the DataFrame (maps to error_code COLUMN_NOT_FOUND)."""
38
+
39
+
40
+ def _clean(value: object) -> object:
41
+ """Coerce a scalar to a JSON-clean Python value.
42
+
43
+ `mode` can be any dtype: an integer column yields `numpy.int64` (NOT
44
+ JSON-serializable), a datetime column yields `pandas.Timestamp`. The other
45
+ metrics are already wrapped in `float(...)`; mode is the one that needs this.
46
+ """
47
+ if value is None:
48
+ return None
49
+ if isinstance(value, pd.Timestamp):
50
+ return value.isoformat()
51
+ if hasattr(value, "item"):
52
+ return value.item()
53
+ return value
54
+
55
+
56
+ def _describe_one(series: pd.Series, metrics: tuple[str, ...]) -> dict[str, object]:
57
+ """Compute descriptive metrics for a single column.
58
+
59
+ Numeric metrics are computed over non-null values. `null_rate` & `count`
60
+ are computed over all rows (nulls included) so they reflect completeness
61
+ as-is. Undefined cases (e.g. std of a single value) return None — degrade
62
+ gracefully instead of raising.
63
+ """
64
+ total = len(series)
65
+ non_null = series.dropna()
66
+ is_numeric = pd.api.types.is_numeric_dtype(series)
67
+
68
+ out: dict[str, object] = {}
69
+ for m in metrics:
70
+ if m == "count":
71
+ out["count"] = int(total)
72
+ elif m == "null_count":
73
+ out["null_count"] = int(series.isna().sum())
74
+ elif m == "null_rate":
75
+ out["null_rate"] = float(series.isna().mean()) if total else 0.0
76
+ elif m == "mode":
77
+ modes = non_null.mode()
78
+ out["mode"] = _clean(modes.iloc[0]) if not modes.empty else None
79
+ elif not is_numeric:
80
+ out[m] = None
81
+ elif m == "mean":
82
+ out["mean"] = float(non_null.mean()) if not non_null.empty else None
83
+ elif m == "median":
84
+ out["median"] = float(non_null.median()) if not non_null.empty else None
85
+ elif m == "std":
86
+ out["std"] = float(non_null.std()) if non_null.shape[0] > 1 else None
87
+ elif m == "var":
88
+ out["var"] = float(non_null.var()) if non_null.shape[0] > 1 else None
89
+ elif m == "q1":
90
+ out["q1"] = float(non_null.quantile(0.25)) if not non_null.empty else None
91
+ elif m == "q3":
92
+ out["q3"] = float(non_null.quantile(0.75)) if not non_null.empty else None
93
+ elif m == "min":
94
+ out["min"] = float(non_null.min()) if not non_null.empty else None
95
+ elif m == "max":
96
+ out["max"] = float(non_null.max()) if not non_null.empty else None
97
+ elif m == "skew":
98
+ out["skew"] = float(non_null.skew()) if non_null.shape[0] > 2 else None
99
+ return out
100
+
101
+
102
+ # Prompt-style description read by the Planner to decide WHEN to pick this tool.
103
+ # Final destination is ToolSpec.description once the wrapper layer is built.
104
+ DESCRIPTION = """\
105
+ Summary: Descriptive statistics (EDA) for one or several columns in a single \
106
+ call — center (mean, median, mode), spread (std, variance, min, max, Q1/Q3 \
107
+ quartiles), distribution shape (skew), and completeness (null count & rate).
108
+
109
+ USE WHEN the user asks for an overview, summary, or single-column statistics \
110
+ of ONE or SEVERAL columns as a whole, with NO grouping and NO comparison \
111
+ between groups. Trigger words: "overview/summary" (ringkasan), "average" \
112
+ (rata-rata), "median", "spread/distribution" (sebaran), "how many nulls" \
113
+ (berapa nilai kosong).
114
+
115
+ DON'T USE WHEN:
116
+ - the question groups by something ("per"/"each"/"by") -> analyze_aggregate
117
+ - it compares two specific groups (A vs B) -> analyze_comparison
118
+ - it tracks a metric over time -> analyze_trend
119
+ - it checks data type, quality, duplicates, outliers, constants -> analyze_profile
120
+
121
+ Example questions:
122
+ - "what's the average and median customer age?"
123
+ - "summarize the income column"
124
+ - "how is product price distributed?"
125
+ - "how many nulls are in the email column?"
126
+ """
127
+
128
+
129
+ def analyze_descriptive(
130
+ df: pd.DataFrame,
131
+ column_ids: list[str],
132
+ metrics: list[str] | None = None,
133
+ ) -> dict[str, dict[str, object]]:
134
+ """Descriptive EDA for one or many columns.
135
+
136
+ Args:
137
+ df: already-materialized data (in the real system the wrapper fetches
138
+ this from a source_id).
139
+ column_ids: columns to analyze.
140
+ metrics: subset of metrics; defaults to all of DEFAULT_METRICS.
141
+
142
+ Returns:
143
+ dict: { column_id: { metric: value, ... }, ... }
144
+
145
+ Raises:
146
+ ColumnNotFoundError: if any column_id is absent from df.
147
+ """
148
+ chosen = tuple(metrics) if metrics else DEFAULT_METRICS
149
+
150
+ missing = [c for c in column_ids if c not in df.columns]
151
+ if missing:
152
+ raise ColumnNotFoundError(f"columns not found: {missing}")
153
+
154
+ return {col: _describe_one(df[col], chosen) for col in column_ids}
src/tools/analytics/quality.py ADDED
@@ -0,0 +1,139 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """analyze_profile — per-column data-quality profile (KM-608).
2
+
3
+ An analytical "family" tool: in ONE call it profiles each column's health —
4
+ dtype, inferred type, completeness (null count/rate), cardinality (distinct
5
+ count/rate, constant flag), and — for numeric columns — min/max/mean plus an
6
+ IQR-based outlier count; for non-numeric columns the most frequent value.
7
+ Answers "is this data clean enough to analyze?" and surfaces issues (lots of
8
+ nulls, a constant column, outliers) before deeper analysis.
9
+
10
+ STATUS: compute layer only — the function takes an already-materialized
11
+ DataFrame. The wrapper layer (fetching data from the catalog via source_id,
12
+ the ToolOutput envelope, ToolSpec registration) is added once the Planner
13
+ seam (KM-418) is settled. Keeping compute separate from data-fetching makes
14
+ this function easy to unit-test in isolation and stable when wrapped.
15
+ """
16
+
17
+ from __future__ import annotations
18
+
19
+ import pandas as pd
20
+
21
+ from src.tools.analytics.descriptive import ColumnNotFoundError
22
+
23
+
24
+ def _clean(value: object) -> object:
25
+ """Convert numpy/pandas scalars to plain Python so the output is JSON-clean.
26
+
27
+ `top_value` (most frequent value) can be a `pandas.Timestamp` when profiling
28
+ a datetime column — neither `Timestamp` nor numpy scalars are JSON-safe.
29
+ """
30
+ if isinstance(value, pd.Timestamp):
31
+ return value.isoformat()
32
+ if hasattr(value, "item"):
33
+ return value.item()
34
+ return value
35
+
36
+
37
+ def _profile_one(series: pd.Series) -> dict[str, object]:
38
+ """Build the quality profile for a single column."""
39
+ total = len(series)
40
+ non_null = series.dropna()
41
+ nn = len(non_null)
42
+ distinct = int(series.nunique(dropna=True))
43
+
44
+ is_bool = pd.api.types.is_bool_dtype(series)
45
+ is_datetime = pd.api.types.is_datetime64_any_dtype(series)
46
+ # bool is technically numeric in pandas; treat it as its own type.
47
+ is_numeric = pd.api.types.is_numeric_dtype(series) and not is_bool
48
+
49
+ if is_bool:
50
+ inferred = "boolean"
51
+ elif is_datetime:
52
+ inferred = "datetime"
53
+ elif is_numeric:
54
+ inferred = "numeric"
55
+ else:
56
+ inferred = "categorical"
57
+
58
+ out: dict[str, object] = {
59
+ "dtype": str(series.dtype),
60
+ "inferred_type": inferred,
61
+ "count": int(total),
62
+ "null_count": int(series.isna().sum()),
63
+ "null_rate": float(series.isna().mean()) if total else 0.0,
64
+ "distinct_count": distinct,
65
+ "distinct_rate": (distinct / nn) if nn else 0.0, # over non-null values
66
+ "is_constant": distinct <= 1,
67
+ }
68
+
69
+ if is_numeric and nn > 0:
70
+ out["min"] = _clean(non_null.min())
71
+ out["max"] = _clean(non_null.max())
72
+ out["mean"] = _clean(non_null.mean())
73
+ # IQR rule: values outside [Q1 - 1.5*IQR, Q3 + 1.5*IQR] are outliers.
74
+ # Needs enough points for stable quartiles.
75
+ if nn >= 4:
76
+ q1 = non_null.quantile(0.25)
77
+ q3 = non_null.quantile(0.75)
78
+ iqr = q3 - q1
79
+ lower, upper = q1 - 1.5 * iqr, q3 + 1.5 * iqr
80
+ out["outlier_count"] = int(((non_null < lower) | (non_null > upper)).sum())
81
+ else:
82
+ out["outlier_count"] = None
83
+ elif not is_numeric and nn > 0:
84
+ counts = non_null.value_counts()
85
+ out["top_value"] = _clean(counts.index[0])
86
+ out["top_freq"] = int(counts.iloc[0])
87
+
88
+ return out
89
+
90
+
91
+ # Prompt-style description read by the Planner to decide WHEN to pick this tool.
92
+ # Final destination is ToolSpec.description once the wrapper layer is built.
93
+ DESCRIPTION = """\
94
+ Summary: Per-column data-quality profile. For each column reports dtype, \
95
+ inferred type, completeness (null count/rate), cardinality (distinct count/rate, \
96
+ constant flag), and — for numeric columns — min/max/mean plus an IQR-based \
97
+ outlier count; for non-numeric columns the most frequent value.
98
+
99
+ USE WHEN the question is about the HEALTH of the data, not its statistics: \
100
+ missing values, duplicates, data types, outliers, "is this clean enough to \
101
+ analyze". Trigger words: "quality" (kualitas), "missing/nulls" (data kosong), \
102
+ "data type" (tipe data), "duplicates/unique" (duplikat/unik), "outliers".
103
+
104
+ DON'T USE WHEN:
105
+ - the user wants statistics like mean/median/std/skew -> analyze_descriptive
106
+ - it groups or compares -> analyze_aggregate / analyze_comparison
107
+
108
+ Example questions:
109
+ - "is this dataset clean enough to analyze?"
110
+ - "which columns have a lot of missing values?"
111
+ - "what are the data types and unique counts per column?"
112
+ - "are there outliers in the amount column?"
113
+ """
114
+
115
+
116
+ def analyze_profile(
117
+ df: pd.DataFrame,
118
+ column_ids: list[str] | None = None,
119
+ ) -> dict[str, dict[str, object]]:
120
+ """Per-column data-quality profile.
121
+
122
+ Args:
123
+ df: already-materialized data (in the real system the wrapper fetches
124
+ this from a source_id).
125
+ column_ids: columns to profile. If None, every column is profiled.
126
+
127
+ Returns:
128
+ dict: { column_id: { profile fields, ... }, ... }
129
+
130
+ Raises:
131
+ ColumnNotFoundError: if any column_id is absent from df.
132
+ """
133
+ cols = list(column_ids) if column_ids is not None else list(df.columns)
134
+
135
+ missing = [c for c in cols if c not in df.columns]
136
+ if missing:
137
+ raise ColumnNotFoundError(f"columns not found: {missing}")
138
+
139
+ return {col: _profile_one(df[col]) for col in cols}
src/tools/analytics/relationship.py ADDED
@@ -0,0 +1,132 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """analyze_correlation — correlation among numeric columns (KM-608).
2
+
3
+ An analytical "family" tool: in ONE call it measures how strongly numeric
4
+ columns move together. Returns the full correlation matrix plus a list of
5
+ column pairs ranked by strength. Answers questions like "does price relate to
6
+ units sold?".
7
+
8
+ STATUS: compute layer only — the function takes an already-materialized
9
+ DataFrame. The wrapper layer (fetching data from the catalog via source_id,
10
+ the ToolOutput envelope, ToolSpec registration) is added once the Planner
11
+ seam (KM-418) is settled. Keeping compute separate from data-fetching makes
12
+ this function easy to unit-test in isolation and stable when wrapped.
13
+ """
14
+
15
+ from __future__ import annotations
16
+
17
+ import math
18
+
19
+ import pandas as pd
20
+
21
+ from src.tools.analytics.descriptive import ColumnNotFoundError
22
+
23
+ # Correlation methods supported by pandas .corr().
24
+ SUPPORTED_METHODS = ("pearson", "spearman", "kendall")
25
+
26
+
27
+ class InvalidMethodError(ValueError):
28
+ """The requested method is not supported (maps to error_code INVALID_METHOD)."""
29
+
30
+
31
+ class NonNumericColumnError(ValueError):
32
+ """A requested column is not numeric (maps to error_code NON_NUMERIC_COLUMN)."""
33
+
34
+
35
+ class NotEnoughColumnsError(ValueError):
36
+ """Correlation needs at least two numeric columns (maps to NOT_ENOUGH_COLUMNS)."""
37
+
38
+
39
+ def _clean(value: object) -> float | None:
40
+ """Cast to plain float; NaN (e.g. a zero-variance column) -> None."""
41
+ if value is None:
42
+ return None
43
+ f = float(value) # type: ignore[arg-type]
44
+ return None if math.isnan(f) else f
45
+
46
+
47
+ # Prompt-style description read by the Planner to decide WHEN to pick this tool.
48
+ # Final destination is ToolSpec.description once the wrapper layer is built.
49
+ DESCRIPTION = """\
50
+ Summary: Pairwise correlation across numeric columns (pearson, spearman, or \
51
+ kendall). Returns a correlation matrix plus the strongest pairs ranked by \
52
+ absolute strength.
53
+
54
+ USE WHEN the question is about relationship or association between numeric \
55
+ variables. Trigger words: "correlation" (korelasi), "related/relationship" \
56
+ (hubungan/keterkaitan), "does X affect Y", "move together".
57
+
58
+ DON'T USE WHEN:
59
+ - it implies causation — correlation is not causality; stay descriptive
60
+ - it compares two groups of one metric -> analyze_comparison
61
+ - it summarizes a single column -> analyze_descriptive
62
+
63
+ Example questions:
64
+ - "is there a correlation between price and quantity sold?"
65
+ - "which variables are most related to revenue?"
66
+ - "do age and spending move together?"
67
+ - "show the correlation matrix for the numeric columns"
68
+ """
69
+
70
+
71
+ def analyze_correlation(
72
+ df: pd.DataFrame,
73
+ column_ids: list[str] | None = None,
74
+ method: str = "pearson",
75
+ ) -> dict[str, object]:
76
+ """Pairwise correlation across numeric columns.
77
+
78
+ Args:
79
+ df: already-materialized data (in the real system the wrapper fetches
80
+ this from a source_id).
81
+ column_ids: numeric columns to correlate. If None, every numeric
82
+ column in df is used.
83
+ method: "pearson" (linear), "spearman" (rank), or "kendall".
84
+
85
+ Returns:
86
+ dict with:
87
+ method — echo of the chosen method
88
+ columns — the numeric columns actually correlated
89
+ matrix — { col: { col: corr|None } } full square matrix
90
+ pairs — [{"a", "b", "corr"}] unique pairs, strongest |corr| first
91
+
92
+ Raises:
93
+ InvalidMethodError: if method is unknown.
94
+ ColumnNotFoundError: if an explicit column is absent.
95
+ NonNumericColumnError: if an explicit column is not numeric.
96
+ NotEnoughColumnsError: if fewer than two numeric columns remain.
97
+ """
98
+ if method not in SUPPORTED_METHODS:
99
+ raise InvalidMethodError(
100
+ f"unknown method '{method}'; supported: {list(SUPPORTED_METHODS)}"
101
+ )
102
+
103
+ if column_ids is None:
104
+ cols = [c for c in df.columns if pd.api.types.is_numeric_dtype(df[c])]
105
+ else:
106
+ missing = [c for c in column_ids if c not in df.columns]
107
+ if missing:
108
+ raise ColumnNotFoundError(f"columns not found: {missing}")
109
+ non_numeric = [
110
+ c for c in column_ids if not pd.api.types.is_numeric_dtype(df[c])
111
+ ]
112
+ if non_numeric:
113
+ raise NonNumericColumnError(f"columns are not numeric: {non_numeric}")
114
+ cols = list(column_ids)
115
+
116
+ if len(cols) < 2:
117
+ raise NotEnoughColumnsError(
118
+ f"need >= 2 numeric columns, got {len(cols)}: {cols}"
119
+ )
120
+
121
+ corr = df[cols].corr(method=method)
122
+ matrix = {a: {b: _clean(corr.loc[a, b]) for b in cols} for a in cols}
123
+
124
+ pairs = []
125
+ for i in range(len(cols)):
126
+ for j in range(i + 1, len(cols)):
127
+ val = _clean(corr.iloc[i, j])
128
+ if val is not None:
129
+ pairs.append({"a": cols[i], "b": cols[j], "corr": val})
130
+ pairs.sort(key=lambda p: abs(p["corr"]), reverse=True)
131
+
132
+ return {"method": method, "columns": cols, "matrix": matrix, "pairs": pairs}
src/tools/analytics/segmentation.py ADDED
@@ -0,0 +1,149 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """analyze_segment — bucket rows into segments (KM-608).
2
+
3
+ An analytical "family" tool: in ONE call it bins a numeric column into
4
+ segments and reports how rows distribute across them (count, and optionally an
5
+ aggregate of another column per segment). Two binning modes: explicit cut
6
+ "edges" (e.g. age 0-18-35-60) or equal-frequency "quantile" buckets (quartiles,
7
+ deciles). Answers questions like "split customers into age brackets" or "bucket
8
+ orders into value tiers".
9
+
10
+ STATUS: compute layer only — the function takes an already-materialized
11
+ DataFrame. The wrapper layer (fetching data from the catalog via source_id,
12
+ the ToolOutput envelope, ToolSpec registration) is added once the Planner
13
+ seam (KM-418) is settled. Keeping compute separate from data-fetching makes
14
+ this function easy to unit-test in isolation and stable when wrapped.
15
+ """
16
+
17
+ from __future__ import annotations
18
+
19
+ import math
20
+
21
+ import pandas as pd
22
+
23
+ from src.tools.analytics.descriptive import ColumnNotFoundError
24
+
25
+ # Binning strategies.
26
+ SUPPORTED_METHODS = ("edges", "quantile")
27
+
28
+ # How to aggregate the value column within each segment.
29
+ SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median")
30
+
31
+
32
+ class InvalidMethodError(ValueError):
33
+ """The requested binning method is not supported (maps to INVALID_METHOD)."""
34
+
35
+
36
+ class NonNumericColumnError(ValueError):
37
+ """The column to segment on is not numeric (maps to NON_NUMERIC_COLUMN)."""
38
+
39
+
40
+ class UnsupportedAggregationError(ValueError):
41
+ """The requested aggregation is not supported (maps to UNSUPPORTED_AGG)."""
42
+
43
+
44
+ def _clean(value: object) -> object:
45
+ """Convert numpy scalars to plain Python; NaN -> None for JSON-clean output."""
46
+ if value is None:
47
+ return None
48
+ if hasattr(value, "item"):
49
+ value = value.item()
50
+ if isinstance(value, float) and math.isnan(value):
51
+ return None
52
+ return value
53
+
54
+
55
+ # Prompt-style description read by the Planner to decide WHEN to pick this tool.
56
+ # Final destination is ToolSpec.description once the wrapper layer is built.
57
+ DESCRIPTION = """\
58
+ Summary: Bins a NUMERIC column into segments and counts how rows distribute \
59
+ across them (optionally aggregating another column per segment). Two modes: \
60
+ explicit cut edges (e.g. age 0-18-35-60) or equal-frequency quantile buckets \
61
+ (quartiles, deciles).
62
+
63
+ USE WHEN the question asks to bucket/bracket a continuous number into ranges. \
64
+ Trigger words: "segment" (segmen), "bucket/bracket" (kelompokkan ke rentang), \
65
+ "age groups/tiers" (kelompok umur/tingkatan), "quartiles/deciles", "bins".
66
+
67
+ DON'T USE WHEN:
68
+ - the category already exists (no binning needed) -> analyze_contribution
69
+ - it aggregates by an existing key -> analyze_aggregate
70
+ - it compares two named groups -> analyze_comparison
71
+
72
+ Example questions:
73
+ - "split customers into age brackets 0-18, 18-35, 35-60"
74
+ - "bucket orders into value tiers"
75
+ - "divide users into spending quartiles"
76
+ - "how many customers fall in each income band?"
77
+ """
78
+
79
+
80
+ def analyze_segment(
81
+ df: pd.DataFrame,
82
+ column: str,
83
+ bins: list[float] | int,
84
+ method: str = "edges",
85
+ labels: list[str] | None = None,
86
+ value_column: str | None = None,
87
+ agg: str = "sum",
88
+ ) -> dict[str, object]:
89
+ """Segment rows by binning a numeric column.
90
+
91
+ Args:
92
+ df: already-materialized data (in the real system the wrapper fetches
93
+ this from a source_id).
94
+ column: numeric column to bin on.
95
+ bins: for method "edges", the list of cut boundaries (e.g.
96
+ [0, 18, 35, 60]); for method "quantile", the number of equal-
97
+ frequency buckets (e.g. 4 for quartiles).
98
+ method: "edges" (explicit boundaries) or "quantile" (equal frequency).
99
+ labels: optional segment names; for "edges" there must be
100
+ len(bins) - 1 of them.
101
+ value_column: if given, also aggregate this column per segment.
102
+ agg: how to aggregate value_column — one of SUPPORTED_AGGS.
103
+
104
+ Returns:
105
+ dict with:
106
+ column, method — echo of the chosen settings
107
+ agg — present only when value_column is given
108
+ segments — [{"segment", "count", ("value")}], in bin order
109
+
110
+ Raises:
111
+ ColumnNotFoundError: if column or value_column is absent.
112
+ NonNumericColumnError: if column is not numeric.
113
+ InvalidMethodError: if method is unknown.
114
+ UnsupportedAggregationError: if agg is not supported.
115
+ """
116
+ referenced = [column] + ([value_column] if value_column else [])
117
+ missing = [c for c in referenced if c not in df.columns]
118
+ if missing:
119
+ raise ColumnNotFoundError(f"columns not found: {missing}")
120
+ if not pd.api.types.is_numeric_dtype(df[column]):
121
+ raise NonNumericColumnError(f"column '{column}' is not numeric")
122
+ if method not in SUPPORTED_METHODS:
123
+ raise InvalidMethodError(
124
+ f"unknown method '{method}'; supported: {list(SUPPORTED_METHODS)}"
125
+ )
126
+ if value_column is not None and agg not in SUPPORTED_AGGS:
127
+ raise UnsupportedAggregationError(
128
+ f"unsupported aggregation '{agg}'; supported: {list(SUPPORTED_AGGS)}"
129
+ )
130
+
131
+ if method == "edges":
132
+ cats = pd.cut(df[column], bins=bins, labels=labels, include_lowest=True)
133
+ else: # quantile
134
+ cats = pd.qcut(df[column], q=bins, labels=labels, duplicates="drop")
135
+
136
+ grouped = df.groupby(cats, observed=False)
137
+ counts = grouped.size()
138
+
139
+ segments = []
140
+ for seg in counts.index:
141
+ row = {"segment": str(seg), "count": int(counts[seg])}
142
+ if value_column is not None:
143
+ row["value"] = _clean(grouped[value_column].agg(agg).get(seg))
144
+ segments.append(row)
145
+
146
+ out: dict[str, object] = {"column": column, "method": method, "segments": segments}
147
+ if value_column is not None:
148
+ out["agg"] = agg
149
+ return out