diff --git a/.gitignore b/.gitignore index 983318caa8875d0794650be06c4805a1ea40cf5a..407d560572b3ef56c9d058343413c5e209f928eb 100644 --- a/.gitignore +++ b/.gitignore @@ -46,6 +46,13 @@ test_tesseract.py # Windows binaries — installed via apt in Docker instead software/ +scripts/ tests/ .claude/ -migratego/ \ No newline at end of file +migratego/ +docs/specs/tabular_parquet_contract.md +docs/specs/tabular_parquet.md + +# Personal / local working docs (not for the shared repo) +AGENT_ARCHITECTURE_CONTEXT_new.md +PROJECT_SUMMARY.md \ No newline at end of file diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 28117880b99e5fd1383753d8b0c322804ff38b4b..1f3fb5c25b94eb3933e9bd37db2cbbf2d25a590e 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -1,7 +1,7 @@ # Architecture — Data Eyond Agentic Service -**Last updated**: 2026-05-07 -**Status**: Design phase — folder skeleton in place, implementation in progress +**Last updated**: 2026-05-20 +**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. --- @@ -21,7 +21,7 @@ A catalog-driven AI service for data analysis. Users upload documents and regist The architecture has two paths: -- **Unstructured** (PDF, DOCX, TXT) — dense similarity over prose chunks (the right primitive for free-form text). +- **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. - **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. 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 ### Ingestion (when user uploads a file or connects a DB) ``` +Structured sources (DB connect / XLSX / CSV / Parquet upload) — Python: source upload / DB connect ↓ introspect schema (DB: information_schema; tabular: file headers + sample rows) @@ -129,7 +130,7 @@ validate (Pydantic) write to catalog store (Postgres jsonb in `data_catalog`, keyed by user_id) ``` -For unstructured files: chunk + embed → PGVector. +**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. ### Query (per user message) @@ -143,7 +144,7 @@ Load chat history IntentRouter LLM → needs_search? source_hint? ↓ ├── chat → ChatbotAgent → SSE stream - ├── unstructured → DocumentRetriever → answerer + ├── unstructured → DocumentRetriever (raw SQL: pgvector `<=>` cosine or `<+>` manhattan) → answerer └── structured → CatalogReader (load full Cs ∪ Ct for user) ↓ diff --git a/PHASE1_TO_PHASE2_REPORT.md b/PHASE1_TO_PHASE2_REPORT.md index 93041c8e9a0b21579fa863cdc18023055d48a2bb..1e4de97392a477af63ec63783637d462e1c60f31 100644 --- a/PHASE1_TO_PHASE2_REPORT.md +++ b/PHASE1_TO_PHASE2_REPORT.md @@ -52,7 +52,7 @@ Everything else — IR validation, SQL/pandas compilation, execution — is dete | `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` | | `src/knowledge/parquet_service.py` | `src/storage/parquet.py` | Parquet upload/download helper | | `src/pipeline/document_pipeline/document_pipeline.py` (folder) | `src/pipeline/document_pipeline.py` (flat) | Single module | -| `src/rag/retrievers/document.py` | `src/retrieval/document.py` | `DocumentRetriever` migrated; tabular file types filtered out of results | +| `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. | | `src/rag/router.py` | `src/retrieval/router.py` | `RetrievalRouter`, Redis-cached, unstructured-only; dead `db: AsyncSession` + `source_hint` params removed | | `src/rag/base.py` (`RetrievalResult`, `BaseRetriever`) | `src/retrieval/base.py` | Same dataclass + ABC | @@ -177,7 +177,7 @@ POST /api/v1/chat/stream │ ├── source_hint == "unstructured" │ → RetrievalRouter.retrieve() [retrieval/router.py, Redis-cached] - │ → DocumentRetriever (PGVector MMR/cosine/etc.) + │ → DocumentRetriever (raw SQL: pgvector `<=>` cosine or `<+>` manhattan) │ → ChatbotAgent.astream(chunks=...) │ └── source_hint == "structured" @@ -237,7 +237,7 @@ POST /api/v1/chat/stream | Top-level chat orchestration | Inline in `api/v1/chat.py` | `agents/chat_handler.py::ChatHandler` | Extracted to a reusable module | | Final answer formation | `agents/chatbot.py` (Phase 1 LangChain) | `agents/chatbot.py::ChatbotAgent` (was `AnswerAgent` in `answer_agent.py` mid-cycle) | Rewritten + renamed | | 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 | -| 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 | +| 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. | | 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 | | DB execution | `query/executors/db_executor.py` | `query/executor/db.py::DbExecutor` | Folder renamed (`executors` → `executor`); sqlglot guard + RO txn + 30 s timeout kept | | 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 **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. 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. + +--- + +## 5. Addendum — post-report changes (2026-05-20, mentor commit `61c746f`) + +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: + +- **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. +- **PGVector collection renamed:** `document_embeddings` → `documents` (to match the Go service's writes). Touched files: `db/postgres/vector_store.py`, `retrieval/document.py`. +- **`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. +- **Intent router defaults flipped.** Ambiguous topical/knowledge questions now prefer `unstructured` (was `structured`). Indonesian few-shot examples added to the prompt. +- **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`. +- **Direction.** The long-term split is **Python = agent/ML layer, Go = data plane**. More pieces are expected to follow doc ingestion out of Python. diff --git a/PROGRESS.md b/PROGRESS.md index 9d8361343924c7f035873f9cd19f600411de2329..82710e060647468fed6d48ca4268861cd707351e 100644 --- a/PROGRESS.md +++ b/PROGRESS.md @@ -2,8 +2,280 @@ 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. -**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) -**Current open PR**: `pr/1` — active. Cleanup PR committed and pushed. +**Last updated**: 2026-06-10 (tool layer complete + hardening/DRY + Langfuse tracing + gated slow-path wiring) +**Current open PR**: `pr/2` — active. + +--- + +## Principal architecture review (2026-06-10) — findings + fix tracker + +A full external review (read the context docs + the slow path, tool layer, query +spine, catalog plumbing, chat endpoint, config/connection layers) landed. It confirmed +the DB-latency diagnosis and surfaced several gaps **not previously tracked here**. +Verified against code before logging. Severity: **critical** / important / nice-to-have. + +**Runtime / latency (the original problem):** +- DB connection handling is the anomaly, NOT cold start. `DbExecutor._run_sync` + (`db.py:192`) → `engine_scope` does `create_engine → connect (TCP+TLS+SCRAM) → 2×SET + → dispose` on EVERY query. Measured ~6–8s for 60 rows; a 2nd warm-session query was + still ~6.6s → per-call handshake, never amortized. `engine_scope`'s connect-once-dispose + semantics were designed for the ingestion pipeline and wrongly inherited by the query path. +- `describe_source` ~3.5s is **planner-induced waste**: every few-shot (`examples.py`) + opens with a `describe_source` task, so the LLM always plans a tool that re-reads from + the catalog DB the same catalog already rendered into its prompt. Its impl does 2 + sequential full-catalog reads (`data_access.py:127-128`). Total catalog reads/request ~5×. +- Azure LLM clients rebuilt per request: `ChatHandler(enable_tracing=True)` is constructed + per request (`chat.py:172`) → fresh Orchestrator/Chatbot → fresh AzureChatOpenAI → fresh + TLS to Azure each call. Planner/Assembler correctly use module singletons; the other two don't. +- Tokens (~13k/request) are NORMAL for this design — do not optimize for $. +- **Reject the scheduled DB-warmer idea**: targets cold start (~1.8s slice) not the per-call + handshake, keeps serverless user DBs awake 24/7 (their compute bill), and decrypts every + tenant's creds on a cron (attack surface). Strictly dominated by an engine cache + + request-scoped pre-connect. + +**Fix tracker (new):** + +| # | Fix | Severity | Owner | Status | +|---|---|---|---|---| +| 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 | `[ ]` | +| 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]` | +| 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 | `[ ]` | +| 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]` | +| 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]` | +| 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]` | +| 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 | `[~]` | +| 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 | `[ ]` | +| R6 | **Hard time budget** — wrap `coordinator.run()` in `asyncio.wait_for` (60–90s). `Constraints.time_budget_seconds` is rendered but not enforced. | important | agent | `[ ]` | +| 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 | `[ ]` | +| R8 | **Catalog upsert race** — per-user advisory lock around read-merge-upsert (`store.py`); concurrent uploads can drop a source. | important | DB | `[ ]` | +| R9 | **`extra="ignore"`** in `settings.py:15` (currently `allow` → typo'd env vars silently swallowed); require Azure keys in prod. | nice-to-have | B | `[ ]` | +| 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 | `[ ]` | +| 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]` | +| 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 | `[ ]` | +| 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 | `[ ]` | +| 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]` | +| 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]` | + +**Slow-path endpoint wiring (2026-06-10):** the Orchestrator→slow-path is now wired +into the live endpoint behind an **env flag**. `settings.enable_slow_path` (env +`ENABLE_SLOW_PATH`, default **off**) is passed to the shared `ChatHandler` in +`api/v1/chat.py`. Flip `ENABLE_SLOW_PATH=true` to route `structured` intents through +Planner→TaskRunner→Assembler and test end-to-end from `/chat/stream` (status progress +events + answer stream). Stays opt-in because `BusinessContext` is still the stub; +fast/unstructured paths unchanged. Verified live via `ChatHandler.handle`. + +**Architecture verdict:** fundamentally sound (catalog-driven IR + deterministic compiler ++ static plan is the right call). Debt is transitional duplication (two planners/registries/ +contract modules — documented, owned) and `ChatHandler` drifting toward a god object +(extract the slow-path composition root + the SSE `_build_sources`/`_normalize_chunks` +mappers when convenient). + +--- + +## What just shipped (2026-06-09/10 — tool layer, tracing, slow-path wiring) + +Big stretch since the slow-path workers landed. The tool layer (teammate-owned) is +now **complete and real**, the slow path is **wired into `ChatHandler` behind a gate**, +and the whole chat pipeline is **traced**. Fast path still untouched; live behavior +unchanged (flags default off). + +**Tool layer — COMPLETE (teammate, KM-624→630).** `src/tools/` was re-created (the +2026-05-11 note about deleting it is superseded). Now teammate-owned: +- `src/tools/analytics/` — the 8 **composite** `analyze_*` computes (descriptive, + aggregate, comparison, contribution, profile, correlation, segment, trend) + + prompt-style DESCRIPTIONs (KM-624/625). +- `src/tools/contracts.py` — canonical `ToolSpec`/`ToolRegistry`/`ToolOutput` (KM-627). + `agents/planner/contracts.py` now just re-exports them + keeps the `BusinessContext` + stub (lead's). +- `src/tools/registry.py::analytics_registry()` (KM-628); `src/tools/invoker.py` + + `src/tools/data_access.py` — `AnalyticsToolInvoker` (KM-629), `DataAccessToolInvoker` + + `CompositeToolInvoker` (KM-630). All never-throw. **Pattern A confirmed** (`analyze_*` + take a `data` `${t}` placeholder from an upstream `query_structured`). +- **Verified live E2E (2026-06-09):** real `query_structured` against a user's Neon + Postgres → `analyze_trend` → Assembler. `analyze_contribution` surfaced a real tool + bug (Decimal vs float in `decomposition.py`) — degrade-and-continue held; **now fixed + by the tool owner** (`_coerce_decimals` in `invoker._materialize`, KM-630 / commit + 1195870), so the whole `analyze_*` family is covered in one place. **Directive:** agent + side does NOT modify `src/tools/` without confirmation. + +**Planner — realigned to the real tools (KM-626).** `registry.py::default_registry()` +composes the real `analytics_registry()` + a local stub for the 4 data-access tools. +Few-shots grown to **A–D**: A `analyze_contribution`, B `analyze_trend`, C mixed +structured+unstructured (`retrieve_documents`, independent branch), D `analyze_aggregate`. +`parallelizable_with` **removed** from `Task` (schema/validator/examples/prompt) — +TaskRunner derives parallelism from `depends_on` alone. + +**Slow-path wiring — built, GATED OFF (KM-626).** `agents/chat_handler.py` gains a +`structured→slow` branch behind `ChatHandler(enable_slow_path=False)`: when on it builds +a per-request `CompositeToolInvoker` (composition root) + `SlowPathCoordinator`, streams +`chat_answer`, persists the `analysis_record`. Two seams isolate the remaining blockers: +- `agents/planner/business_context.py::get_business_context(user_id)` — async stub + `BusinessContext`; TODO(lead) swap for the real read. +- `agents/slow_path/store.py` — `AnalysisStore` Protocol + `NullAnalysisStore` (logs + only). Real store = `analysis_records` table in the catalog DB (Neon `dataeyond`) — + **table not created yet**. `chat_answer` still emitted as one chunk (not token-streamed). + +**Observability — Langfuse tracing wired (KM-631).** `src/observability/langfuse/ +tracing.py` — `RequestTracer`/`NullTracer`/`TracingToolInvoker` + `_redact`. One trace +per request groups Orchestrator.classify, Planner.plan (each retry = its own generation), +Assembler.assemble, Chatbot.astream + tool spans (latency/metadata only). Gated: +`ChatHandler(enable_tracing=False)`; `api/v1/chat.py` opts in (`=True`). PII policy: +Orchestrator+Planner unmasked (question + PII-safe summary); Assembler+Chatbot masked +(see real rows/chunks); tool spans carry name + arg keys + row count only. Zero added +LLM tokens; verified live to US Cloud. + +**Live evals green (2026-06-09, real Azure 4o):** `RUN_PLANNER_EVAL=1` and +`RUN_SLOW_PATH_EVAL=1` both pass — Planner emits valid catalog-consistent `QueryIR` and +wires Pattern A correctly; self-corrects via retry. + +**Open follow-ups:** real `BusinessContext` (lead); create `analysis_records` table + +real `AnalysisStore`; register data-access `ToolSpec`s upstream (`data_access_registry()`) +or keep the planner stub; 4o → GPT-mini deployment swap; flip `enable_slow_path` on once +`BusinessContext` is real. NOTE: 3 test files pre-existing broken from rename rot +(`test_chat_handler.py`, `test_intent_router.py`, `test_answer_agent.py` import the old +`answer_agent`/`intent_router` module names). + +--- + +## What just shipped (2026-06-10 — TAB: tool-layer hardening + DRY) + +Owner-side companion to the agent block above. After the live E2E surfaced real-data +edge cases, the tool layer got a round of correctness hardening. All in TAB-owned paths +(`src/tools/`, `src/catalog/`); no agent-side or API change. + +**JSON-safety across the `analyze_*` family.** Real DB rows carry scalar types that +don't survive the jsonb / SSE round-trip: +- `[KM-630] coerce DB Decimal → float` (commit 1195870) — `_coerce_decimals` in + `invoker._materialize` converts object-columns holding `decimal.Decimal` (asyncpg + returns NUMERIC as `Decimal`) to `float64` before any compute runs. Fixes the + `float + Decimal` TypeError in `decomposition.analyze_contribution` **and** the whole + family in one seam — only touches columns that actually contain a `Decimal`. +- `[KM-624] non-JSON-safe scalars in mode & top_value` (commit 6981ed3) — normalize + numpy / non-native scalars so descriptive + top-value outputs serialize cleanly. + +**Planner↔Tools registry alignment + Timestamp keys** (commit 4bb7623, `fix(tools)`): +- `registry.py` — `analyze_descriptive.required` corrected `["data"]` → `["data", + "column_ids"]` to match the compute signature (`column_ids` has no default). Prevents + the Planner from emitting a call that's missing a required arg. `analyze_profile` stays + `["data"]` (its `column_ids` defaults to `None`). +- `aggregation._clean` — group-by over a datetime column produced `pd.Timestamp` group + keys that aren't JSON-safe; now normalized to `.isoformat()` alongside the existing + numpy `.item()` branch. + +**DRY: single `SAMPLE_LIMIT` constant** (commit 6d46ba5, `[NOTICKET] refactor(catalog)`): +- One source of truth in `catalog/introspect/base.py` (`SAMPLE_LIMIT = 3`, down from 5 — + token cost: sample values feed the planner prompt). Both introspection paths import it: + `catalog/introspect/tabular.py` and `pipeline/db_pipeline/extractor.py` (which dropped + its own local `= 3`). Dependency direction is pipeline→catalog (no circular import). + Stale test `test_sample_values_capped_at_five` updated to assert the real cap (3). + +**Audit result:** Planner↔Tools arg alignment swept end-to-end — 7/8 `analyze_*` tools +already matched; the 1 mismatch (`analyze_descriptive`) is the fix above. Pattern A holds +across all of them. + +--- + +## What just shipped (2026-06-08 — KM-626: slow-path agent layer) + +The rest of the slow path after the Planner (KM-567) — TaskRunner, Assembler, and +the coordinator. Built and tested against +mocks; **not yet wired into the live `ChatHandler`** (waits on the tool team's real +`ToolInvoker` + a real `BusinessContext`). Fast path untouched. + +**Naming:** "Orchestrator" = the entry dispatcher only (`agents/orchestration.py`). +The slow-path **workers** live in **`agents/slow_path/`** — deliberately NOT named +"orchestrator". + +**Files added** (`src/agents/slow_path/`): +- `schemas.py` — `TaskResult`, `RunState`; `TaskSummary`, `AnalysisRecord`, + `AssembledOutput`, `AssemblerNarrative`. Reuses `ToolOutput`. +- `invoker.py` — `ToolInvoker` Protocol only; the tool team owns the impl (KM-418). +- `errors.py` — `SlowPathError`, `AssemblerError`. +- `task_runner.py` — deterministic, 0 LLM: wave-based execution, `${t}` placeholder + resolution, internal `validate_args`, never-throw invoke, status labeling, + degrade-and-continue → `RunState`. +- `assembler.py` + `prompt.py` + `config/prompts/assembler.md` — single LLM call → + `AssemblerNarrative`; code merges with `RunState` to build the `AnalysisRecord` + (structured fields copied, never re-authored). +- `coordinator.py` — `SlowPathCoordinator`: Planner → TaskRunner → Assembler. + +**Tests added** (`tests/agents/slow_path/`, 12 passing; gitignored): schema round-trips ++ chat_answer-first; runner happy/placeholder/parallel/degrade/arg-miss; assembler +narrative-vs-snapshot + question threading; coordinator end-to-end. `ruff` clean; +tool-agnostic (no `src/tools/*` import). + +**Open follow-ups (not blockers):** wire `SlowPathCoordinator` into the expanded +Orchestrator/`ChatHandler` once the real invoker + `BusinessContext` exist; swap the +test `MockToolInvoker` for the tool team's real one (zero agent change, INV-7); 4o → +GPT-mini deployment swap. + +--- + +## What just shipped (2026-06-08 — tool taxonomy + ownership revision) + +Team decisions after the teammate pushed KM-624 (`src/tools/analytics/`): + +- **Composite tools, not atomic.** v1 uses **composite "family" tools** (`analyze_*`), + not the atomic `compute_*` set the earlier draft assumed. One `analyze_*` call does a + whole analytical job (e.g. `analyze_descriptive` subsumes median/mode/stddev/percentile; + `analyze_trend` subsumes `date_trunc`). Tool-taxonomy decision recorded. +- **Tool team owns ALL tools** — compute, data-access (`query_structured`, + `retrieve_documents`, `list_sources`, `describe_source`), the wrapper/invoker layer + (KM-418), and **all tool tests**. The agent team owns nothing below the registry contract. +- **Planner stub realigned to the real tools.** `registry.py` rewritten from the 9 atomic + entries to **12 composite entries** (4 data-access + 8 `analyze_*`); `examples.py` + rewritten (Example A → `analyze_contribution`, Example B → `analyze_trend`); `planner.md` + bullet updated; planner tests updated. 32 passing + 1 gated, `ruff` clean. +- **Open (tool team's call):** Pattern A (analyze_* take a `${t}` `data` placeholder + from an upstream `query_structured`) vs Pattern B (self-fetch by `source_id`). Stub + assumes A; reshaped to match once decided (agent code unaffected, INV-7). +- **New coupling:** the tool team's `query_structured`/`retrieve_documents` are expected + to call our existing `QueryService`/`RetrievalRouter`; `query_structured` stays + inline-`QueryIR` so `IRValidator` still applies. Interface to coordinate. + +**Next (our scope, all mock-able now):** TaskRunner + Assembler against a `MockToolInvoker`, +then Orchestrator slow-path wiring. Stubs still to retire on integration: `contracts.py` +(BusinessContext from lead; ToolSpec/ToolRegistry/ToolOutput from tool team) and `registry.py` +(real registry from tool team). Infra: swap the 4o stand-in for a GPT-mini deployment. + +--- + +## What just shipped (2026-06-05 — Phase 3: Planner agent) + +First slow-path agent (the Planner). A single LLM +call turns BusinessContext + Catalog + ToolRegistry + question + Constraints into a +validated, **static** `TaskList` (DAG of fully-specified tool-call chains). No +replanning (INV-6); tool-agnostic against a registry contract (INV-7). Fast path +(`agents/orchestration.py`, `agents/chatbot.py`, `query/`) untouched. + +**Files added** (`src/agents/planner/`): +- `contracts.py` — **STUB** Pydantic contracts pending reconciliation: `BusinessContext` + (+KeyTerm/DataTableNote/DataColumnNote, lead's), `ToolSpec`/`ToolRegistry` (tool + team KM-608), `ToolOutput` envelope. +- `schemas.py` — `CrispStage`, `ToolCall`, `Task`, `TaskList`. No replan schemas. +- `inputs.py` — `CatalogSummary` (condensed, PII `sample_values` nulled, `from_catalog` + builder + `render`) and `Constraints` (max_tasks=5, modeling_allowed=False). +- `registry.py` — **STUB** v1 P0 registry: query_structured, retrieve_documents, + list_sources, describe_source, compute_median/stddev/percentile/mode, date_trunc. +- `errors.py` — `PlannerError`, `PlannerValidationError`. +- `prompt.py` + `config/prompts/planner.md` — system prompt (INV-1/6/7 + principles) + + per-call human content (context + catalog + tools + constraints + few-shots + question). +- `examples.py` — two few-shots (A exploratory revenue-by-category; B descriptive + monthly-trend-by-region with date_trunc), built from the real `TaskList` schema. +- `validator.py` — `PlannerValidator` running the 8 checks; reuses the existing + `IRValidator` for inline `query_structured` IRs. +- `service.py` — `PlannerService` + `plan_analysis(...)`: chain (mirrors + `query/planner/service.py`) + validate-and-retry loop (max 3, mirrors `QueryService`). + +**Tests added** (`tests/agents/planner/`, 30 passing + 1 gated): `test_schemas.py`, +`test_inputs.py`, `test_validator.py` (one failure per check + happy paths), +`test_service.py` (`_FakeChain` + retry), `test_golden_questions.py` (live eval gated on +`RUN_PLANNER_EVAL=1`). `ruff check` clean on planner paths. + +**Open follow-ups (not blockers):** reconcile `BusinessContext` with the lead and +`ToolRegistry`/`ToolSpec` + real tools with teammate (KM-608); "GPT mini" currently uses +the configured 4o deployment (swap `azure_deployment` when a mini deployment exists). Next: +Orchestrator slow-path expansion + TaskRunner + Assembler. --- @@ -54,7 +326,7 @@ Persistent tracker mirroring the 42-item ownership table in `REPO_CONTEXT.md` "T | # | Item | Owner | Status | Notes | |---|---|---|---|---| | 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). | -| 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). | +| 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). | | 7 | `BaseIntrospector` ABC (`catalog/introspect/base.py`) | B | `[x]` | Pre-existing; signature locked | ### Ingestion — shared catalog plumbing @@ -102,8 +374,8 @@ Persistent tracker mirroring the 42-item ownership table in `REPO_CONTEXT.md` "T | # | Item | Status | Notes | |---|---|---|---| -| 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 | -| 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` | +| 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.) | +| 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`. | | 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`. | ### Agents + chat @@ -112,7 +384,20 @@ Persistent tracker mirroring the 42-item ownership table in `REPO_CONTEXT.md` "T |---|---|---|---| | 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. | | 33 | Guardrails prompt (`config/prompts/guardrails.md`) | `[x]` | PR7-bundle — appended to `chatbot_system.md` so guardrails take precedence in conflict. | -| — | 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. | +| — | 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). | + +### Tools — slow-path "Tools" component (TAB) + +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. + +| # | Item | Owner | Status | Notes | +|---|---|---|---|---| +| — | 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). | +| — | 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). | +| — | Analytics registry (`tools/registry.py`) | TAB | `[x]` | KM-628 — `analytics_registry()`. `analyze_descriptive.required` = `["data","column_ids"]` (aligned to compute signature, commit 4bb7623). | +| — | Invoker layer (`tools/invoker.py`) | TAB | `[x]` | KM-629 — `AnalyticsToolInvoker` (Pattern A: `analyze_*` take a `data` `${t}` placeholder from upstream `query_structured`; `_materialize` → DataFrame, `_coerce_decimals` covers the whole family) + `CompositeToolInvoker` (routes data-access vs analytics by name). | +| — | 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). | +| — | Tool tests (`tests/unit/tools/`) | TAB | `[x]` | analytics + data-access + invoker tests (gitignored). Incl. regression `test_decimal_columns_coerced_for_analyze_contribution`. | ### API surface @@ -129,7 +414,7 @@ Persistent tracker mirroring the 42-item ownership table in `REPO_CONTEXT.md` "T | # | Item | Owner | Status | Notes | |---|---|---|---|---| | 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. | -| 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). | +| 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). | | 40 | IR validator tests (`tests/query/ir/test_validator.py`) | B | `[x]` | PR1 — 19 tests, all rules covered | | — | PII detector tests (`tests/catalog/test_pii_detector.py`) | B | `[x]` | PR1 — 26 tests (parametrized) | | — | Catalog validator tests (`tests/catalog/test_validator.py`) | B | `[x]` | PR1 — 5 tests | diff --git a/REPO_CONTEXT.md b/REPO_CONTEXT.md index 9de3bf7a345db76610aa9f1e0675d1614491e9fd..0618b285ac804a324f80c9563ca520e951b07507 100644 --- a/REPO_CONTEXT.md +++ b/REPO_CONTEXT.md @@ -156,7 +156,7 @@ makes any LLM calls.) | `db/postgres/connection.py` | two async engines: `engine` (app) and `_pgvector_engine` (PGVector) | | `db/postgres/init_db.py` | startup: creates `vector` extension, all tables, HNSW + GIN indexes | | `db/postgres/models.py` | SQLAlchemy app tables (users, rooms, chat messages, …) | -| `db/postgres/vector_store.py` | shared PGVector instance (collection `document_embeddings`) | +| `db/postgres/vector_store.py` | shared PGVector instance (collection `documents` — written by Go ingestion service) | | `db/redis/connection.py` | async Redis client | | `storage/az_blob/az_blob.py` | Azure Blob async wrapper (uploads + Parquet) | | `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 | QueryService | ✅ | plan → validate → retry-on-fail (max 3) → dispatch → execute → `QueryResult` | | `ChatbotAgent` + prompt + guardrails | ✅ | Renamed from `AnswerAgent` in Cleanup PR. Guardrails appended to `chatbot_system.md` | | `ChatHandler` (top-level chat orchestrator) | ✅ | SSE events: `intent` / `chunk` / `done` / `error` | -| `DocumentRetriever` + `RetrievalRouter` (Redis-cached) | ✅ | Migrated from `src/rag/` (now deleted); MMR/cosine/euclidean/manhattan/inner_product | +| `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`. | | `/api/v1/chat/stream` | ✅ | Rewired to `ChatHandler`; Redis cache + fast intent + history + message persistence remain in chat.py | | `/api/v1/db-clients/{id}/ingest` | ✅ | Calls only `on_db_registered`; Phase 1 dual-write removed | | `/api/v1/document/{upload,process,delete}` | ✅ | `/process` triggers `on_tabular_uploaded` for CSV/XLSX | diff --git a/pyproject.toml b/pyproject.toml index b1ea3270cbb7a89b8b7ec7b72774eaeaf4d9957a..ea38cb5468913c984643d8f2c0105e2957bcb364 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -120,7 +120,9 @@ ignore = [ ] [tool.ruff.lint.per-file-ignores] -"tests/**" = ["S101", "S105", "S106"] +# S608: golden compiler tests assert literal SQL strings (incl. concatenated +# suffixes) — they never execute against a DB, so it's a false positive here. +"tests/**" = ["S101", "S105", "S106", "S608"] [tool.mypy] python_version = "3.12" diff --git a/src/agents/chat_handler.py b/src/agents/chat_handler.py index 3f5fa1db86fb8c50bc3768088150b8ab86f9e1b3..b009e1a1e559f041ea4a8dd027f13f3e9d2d5b43 100644 --- a/src/agents/chat_handler.py +++ b/src/agents/chat_handler.py @@ -22,8 +22,9 @@ inject mocks). from __future__ import annotations +import asyncio import json -from collections.abc import AsyncIterator +from collections.abc import AsyncIterator, Callable from typing import TYPE_CHECKING, Any from langchain_core.messages import BaseMessage @@ -38,6 +39,8 @@ if TYPE_CHECKING: from ..catalog.reader import CatalogReader from ..query.service import QueryService from ..retrieval.router import RetrievalRouter + from .slow_path.coordinator import SlowPathCoordinator + from .slow_path.store import AnalysisStore logger = get_logger("chat_handler") @@ -62,12 +65,29 @@ class ChatHandler: catalog_reader: CatalogReader | None = None, query_service: QueryService | None = None, document_retriever: RetrievalRouter | None = None, + *, + enable_slow_path: bool = False, + slow_path_coordinator_factory: ( + Callable[[str], SlowPathCoordinator] | None + ) = None, + analysis_store: AnalysisStore | None = None, + enable_tracing: bool = False, ) -> None: self._intent_router = intent_router self._answer_agent = answer_agent self._catalog_reader = catalog_reader self._query_service = query_service self._document_retriever = document_retriever + # Langfuse tracing (tokens + latency). OFF by default so tests never hit + # Langfuse; the live endpoint opts in with ChatHandler(enable_tracing=True). + self._enable_tracing = enable_tracing + # Slow analytical path (Planner -> TaskRunner -> Assembler). OFF by default: + # gated until the lead's real BusinessContext lands. When True, `structured` + # intents route here instead of the single-query QueryService path. The + # factory + store are injectable for tests. + self._enable_slow_path = enable_slow_path + self._slow_path_factory = slow_path_coordinator_factory + self._analysis_store = analysis_store # ------------------------------------------------------------------ # Lazy default-dep builders @@ -115,9 +135,13 @@ class ChatHandler: user_id: str, history: list[BaseMessage] | None = None, ) -> AsyncIterator[dict[str, Any]]: + tracer = self._make_tracer(user_id, message) + # ---- 1. Classify intent -------------------------------------- try: - decision = await self._get_intent_router().classify(message, history) + oc = tracer.callbacks() # orchestrator: PII-safe, full capture + ckw = {"callbacks": oc} if oc else {} + decision = await self._get_intent_router().classify(message, history, **ckw) except Exception as e: logger.error("intent classification failed", error=str(e)) yield {"event": "error", "data": f"Could not classify message: {e}"} @@ -133,7 +157,20 @@ class ChatHandler: # ---- 2. Route ------------------------------------------------ if decision.source_hint == "structured": try: - catalog = await self._get_catalog_reader().read(user_id, "structured") + # One memoizing reader per request: the same catalog is otherwise + # re-fetched from the catalog DB 4-5x across the slow-path run. This + # collapses those to one round-trip per source_hint and pins a single + # consistent snapshot for plan + execution. + from ..catalog.reader import MemoizingCatalogReader + + req_reader = MemoizingCatalogReader(self._get_catalog_reader()) + catalog = await req_reader.read(user_id, "structured") + if self._enable_slow_path: + async for event in self._run_slow_path( + user_id, rewritten, catalog, tracer, req_reader + ): + yield event + return query_result = await self._get_query_service().run( user_id, rewritten, catalog ) @@ -174,12 +211,16 @@ class ChatHandler: yield {"event": "sources", "data": json.dumps(sources)} # ---- 3. Stream answer ---------------------------------------- + # masked: the answer call sees real query rows / doc chunks (possible PII). + mc = tracer.callbacks(masked=True) + akw = {"callbacks": mc} if mc else {} try: async for token in self._get_answer_agent().astream( message, history=history, query_result=query_result, chunks=chunks, + **akw, ): yield {"event": "chunk", "data": token} except Exception as e: @@ -187,6 +228,150 @@ class ChatHandler: yield {"event": "error", "data": f"Answer generation failed: {e}"} return + tracer.end() + yield {"event": "done", "data": ""} + + # ------------------------------------------------------------------ + # Slow analytical path (gated, off by default) + # ------------------------------------------------------------------ + + def _make_tracer(self, user_id: str, question: str) -> Any: + """One Langfuse trace per request (or a NullTracer when disabled).""" + if not self._enable_tracing: + from ..observability.langfuse.tracing import NullTracer + + return NullTracer() + from ..observability.langfuse.tracing import RequestTracer + + return RequestTracer.start(user_id=user_id, question=question) + + def _get_slow_path_coordinator( + self, user_id: str, tracer: Any = None, catalog_reader: CatalogReader | None = None + ) -> SlowPathCoordinator: + """Build the per-request slow-path coordinator (composition root). + + The data-access tools need the authenticated `user_id` + `CatalogReader`, + so the `CompositeToolInvoker` is constructed per request. The slow-path + agent code stays tool-agnostic (INV-7) — only here, the composition root, + do we name concrete tool implementations. When tracing is active the invoker + is wrapped so each tool call records a metadata-only span. + """ + if self._slow_path_factory is not None: + return self._slow_path_factory(user_id) + + from ..tools.data_access import DataAccessToolInvoker + from ..tools.invoker import AnalyticsToolInvoker, CompositeToolInvoker + from .planner.registry import default_registry + from .planner.service import PlannerService + from .slow_path.assembler import Assembler + from .slow_path.coordinator import SlowPathCoordinator + from .slow_path.task_runner import TaskRunner + + invoker: Any = CompositeToolInvoker( + DataAccessToolInvoker(user_id, catalog_reader or self._get_catalog_reader()), + AnalyticsToolInvoker(), + ) + if tracer is not None and getattr(tracer, "active", False): + from ..observability.langfuse.tracing import TracingToolInvoker + + invoker = TracingToolInvoker(invoker, tracer) + registry = default_registry() + return SlowPathCoordinator( + PlannerService(), TaskRunner(invoker, registry), Assembler(), registry + ) + + def _get_analysis_store(self) -> AnalysisStore: + if self._analysis_store is None: + from .slow_path.store import NullAnalysisStore + + self._analysis_store = NullAnalysisStore() + return self._analysis_store + + async def _run_slow_path( + self, + user_id: str, + query: str, + catalog: Any, + tracer: Any = None, + catalog_reader: CatalogReader | None = None, + ) -> AsyncIterator[dict[str, Any]]: + """Run the slow path and stream its assembled answer as SSE events. + + Context comes from the `get_business_context` seam (a stub today); the + `analysis_record` is persisted via the `AnalysisStore` seam (a no-op today). + `chat_answer` is emitted as a single `chunk` (the Assembler returns the whole + object — true token streaming is a later step). + """ + from .planner.business_context import get_business_context + from .planner.inputs import Constraints + + if tracer is None: + from ..observability.langfuse.tracing import NullTracer + + tracer = NullTracer() + + coordinator = self._get_slow_path_coordinator(user_id, tracer, catalog_reader) + context = await get_business_context(user_id) + + # DB3: warm the user's DB connection in parallel with planning so the + # handshake overlaps the ~4s Planner call. Default path only — an injected + # coordinator factory (tests / custom) may not use the real DbExecutor. + if self._slow_path_factory is None: + from ..query.executor.db import DbExecutor + + asyncio.create_task(DbExecutor.prewarm(catalog, user_id)) # noqa: RUF006 + + pc = tracer.callbacks() # planner: PII-safe, full capture + ac = tracer.callbacks(masked=True) # assembler: sees real rows -> masked + run_kw: dict[str, Any] = {} + if pc: + run_kw["planner_callbacks"] = pc + if ac: + run_kw["assembler_callbacks"] = ac + + # R4: bridge the coordinator's per-stage progress callback to SSE `status` + # events so the stream isn't silent for ~12s (and proxies don't drop the + # idle connection). Status events only appear if the coordinator calls back. + progress_q: asyncio.Queue[str] = asyncio.Queue() + + async def _progress(stage: str) -> None: + await progress_q.put(stage) + + run_task = asyncio.create_task( + coordinator.run( + context, catalog, query, Constraints(), progress=_progress, **run_kw + ) + ) + getter: asyncio.Task = asyncio.create_task(progress_q.get()) + pending: set[asyncio.Task] = {run_task, getter} + while True: + done, pending = await asyncio.wait( + pending, return_when=asyncio.FIRST_COMPLETED + ) + if getter in done: + yield {"event": "status", "data": getter.result()} + getter = asyncio.create_task(progress_q.get()) + pending = pending | {getter} + if run_task in done: + getter.cancel() + while not progress_q.empty(): + yield {"event": "status", "data": progress_q.get_nowait()} + break + + try: + result = run_task.result() + except Exception as e: + logger.error("slow path failed", user_id=user_id, error=str(e)) + yield {"event": "error", "data": f"Analysis failed: {e}"} + return + + yield {"event": "sources", "data": json.dumps([])} # TODO: derive from record + yield {"event": "chunk", "data": result.chat_answer} + try: + await self._get_analysis_store().save(result.analysis_record) + except Exception as e: # persistence must never break the user's answer + logger.error("analysis_record persist failed", user_id=user_id, error=str(e)) + tracer.end() # output omitted (chat_answer may contain PII on Cloud) yield {"event": "done", "data": ""} diff --git a/src/agents/chatbot.py b/src/agents/chatbot.py index ff1fb129a7a08cd0e25dfbd46ac08a93736db1aa..18fa1dd671e4f9d022c481c2b17bfbc3a7ef1fc4 100644 --- a/src/agents/chatbot.py +++ b/src/agents/chatbot.py @@ -119,6 +119,9 @@ def _build_default_chain() -> Runnable: azure_endpoint=settings.azureai_endpoint_url_4o, api_key=settings.azureai_api_key_4o, temperature=0.3, + # Emit token usage on the final streamed chunk (this agent only streams), so + # the fast-path answer reports tokens to Langfuse like the non-streaming calls. + model_kwargs={"stream_options": {"include_usage": True}}, ) prompt = ChatPromptTemplate.from_messages( [ @@ -153,6 +156,7 @@ class ChatbotAgent: history: list[BaseMessage] | None = None, query_result: QueryResult | None = None, chunks: list[DocumentChunk] | None = None, + callbacks: list | None = None, ) -> AsyncIterator[str]: """Stream tokens of the final answer. @@ -165,5 +169,9 @@ class ChatbotAgent: "history": history or [], "context": _build_context_block(query_result, chunks), } - async for token in chain.astream(payload): - yield token + if callbacks: + async for token in chain.astream(payload, config={"callbacks": callbacks}): + yield token + else: + async for token in chain.astream(payload): + yield token diff --git a/src/agents/orchestration.py b/src/agents/orchestration.py index 61ed7cb40383ee87a6bbc3201f3093fb6d7f8c98..53f702d7c267660d038f5ce88b803f304e3b33e0 100644 --- a/src/agents/orchestration.py +++ b/src/agents/orchestration.py @@ -96,11 +96,16 @@ class OrchestratorAgent: self, message: str, history: list[BaseMessage] | None = None, + callbacks: list | None = None, ) -> IntentRouterDecision: chain = self._ensure_chain() - decision: IntentRouterDecision = await chain.ainvoke( - {"message": message, "history": history or []} - ) + payload = {"message": message, "history": history or []} + if callbacks: + decision: IntentRouterDecision = await chain.ainvoke( + payload, config={"callbacks": callbacks} + ) + else: + decision = await chain.ainvoke(payload) logger.info( "intent classified", source_hint=decision.source_hint, diff --git a/src/agents/planner/__init__.py b/src/agents/planner/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e4bd9c086cd0e6a70044f75059398d40b4495ca6 --- /dev/null +++ b/src/agents/planner/__init__.py @@ -0,0 +1,8 @@ +"""Planner agent — slow-path CRISP-DM analysis planner. + +Single LLM call: BusinessContext + CatalogSummary + ToolRegistry + question + +Constraints -> a validated, static `TaskList` (a DAG of fully-specified +tool-call chains). No replanning (INV-6); tool-agnostic (INV-7). + +See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3. +""" diff --git a/src/agents/planner/business_context.py b/src/agents/planner/business_context.py new file mode 100644 index 0000000000000000000000000000000000000000..d73f504fb8e005a920c366a588535805cfa7cf8b --- /dev/null +++ b/src/agents/planner/business_context.py @@ -0,0 +1,31 @@ +"""BusinessContext reader — the single seam the slow path reads context through. + +`get_business_context(user_id)` is the one place the live flow obtains a +`BusinessContext`. Today it returns a minimal STUB so the slow path runs end to +end; when the lead's real Business Understanding source lands, swap the body here +(read the interview / stored context) and nothing upstream changes. + +See AGENT_ARCHITECTURE_CONTEXT_new.md §7.1. +""" + +from __future__ import annotations + +from .contracts import BusinessContext + + +async def get_business_context(user_id: str) -> BusinessContext: + """Return the user's BusinessContext. + + STUB until the lead's real source lands. `project_id` flows through as + `RunState.business_context_id`. Async so the real implementation (a DB / store + read) fits without changing this signature. + + TODO(lead): replace the body with the real read (Business Understanding store). + """ + return BusinessContext( + project_id=user_id, + industry="unknown", + completeness="partial", + business_description="(not yet captured — BusinessContext source pending)", + scale_and_scope="(unknown)", + ) diff --git a/src/agents/planner/contracts.py b/src/agents/planner/contracts.py new file mode 100644 index 0000000000000000000000000000000000000000..3bcc9aab19dc6503cbf4c34c430d3b1d1a56d660 --- /dev/null +++ b/src/agents/planner/contracts.py @@ -0,0 +1,65 @@ +"""Contracts the planner consumes from other teams. + +`BusinessContext` (+ KeyTerm / DataTableNote / DataColumnNote) is still a LOCAL +STUB owned by the lead (interview / Business Understanding); shape mirrors +AGENT_ARCHITECTURE_CONTEXT_new.md §7.1 and must be reconciled before integration. + +`ToolSpec` / `ToolRegistry` / `ToolOutput` are NO LONGER defined here — the tool +team owns them (KM-465). They now live in `src.tools.contracts` and are +re-exported below so existing importers (`registry.py`, the slow-path +`invoker.py` / `schemas.py`) keep working against one shared definition. +""" + +from __future__ import annotations + +from typing import Literal + +from pydantic import BaseModel, Field + +# Canonical tool contracts now owned by the tool team (KM-465). Re-exported here +# for backwards-compatible imports; the definitions live in src.tools.contracts. +from src.tools.contracts import ToolOutput, ToolRegistry, ToolSpec + +__all__ = [ + "KeyTerm", + "DataTableNote", + "DataColumnNote", + "BusinessContext", + "ToolSpec", + "ToolRegistry", + "ToolOutput", +] + +# --------------------------------------------------------------------------- # +# BusinessContext (lead's contract — §7.1) +# --------------------------------------------------------------------------- # + + +class KeyTerm(BaseModel): + term: str + meaning: str + + +class DataTableNote(BaseModel): + table_name: str + row_represents: str + + +class DataColumnNote(BaseModel): + column_name: str + meaning: str + + +class BusinessContext(BaseModel): + project_id: str + industry: str + completeness: Literal["partial", "complete"] + business_description: str + scale_and_scope: str + key_terms: list[KeyTerm] = Field(default_factory=list) + data_overview: list[DataTableNote] = Field(default_factory=list) + data_column_notes: list[DataColumnNote] = Field(default_factory=list) + whats_normal: str = "" + recent_events: str = "" + things_to_watch_for: str = "" + open_questions: list[str] = Field(default_factory=list) diff --git a/src/agents/planner/errors.py b/src/agents/planner/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..2640f9a6c581947c02ce7a37e890a82c1babe78f --- /dev/null +++ b/src/agents/planner/errors.py @@ -0,0 +1,15 @@ +"""Typed errors for the planner agent.""" + +from __future__ import annotations + + +class PlannerError(Exception): + """Base error for the planner agent.""" + + +class PlannerValidationError(PlannerError): + """A TaskList failed one of the planner validator's checks. + + The message is specific enough that the planner can be re-prompted with it + to self-correct (max 3 attempts, see service.py). + """ diff --git a/src/agents/planner/examples.py b/src/agents/planner/examples.py new file mode 100644 index 0000000000000000000000000000000000000000..ca701e6ecd56f3c44ea801f8dd8833cdb9a4841d --- /dev/null +++ b/src/agents/planner/examples.py @@ -0,0 +1,384 @@ +"""Few-shot examples for the planner prompt. + +Two illustrative (question -> TaskList) pairs that teach the OUTPUT SHAPE: +stages, dependency edges, ordered tool-call chains, inline QueryIR, +"${t}" placeholders, and the assumed data-flow convention — `query_structured` +pulls rows, then a composite `analyze_*` tool consumes them via a `data` placeholder +referencing the upstream result's column aliases (Pattern A; the tool team may +instead pick self-fetch by `source_id`, in which case these examples are reshaped +to match — see registry.py). They reference a hypothetical sales catalog +(`src_sales` / `t_orders`); these ids are part of the illustration and are not +validated against the user's real catalog. v1 is descriptive/diagnostic — no +modeling tasks. + +See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3 (Examples A and B). +""" + +from __future__ import annotations + +from .schemas import Task, TaskList, ToolCall + +# --------------------------------------------------------------------------- # +# Example A — exploratory, no modeling. +# "Which product categories drove last quarter's revenue?" +# Shows: query_structured pulls rows -> analyze_contribution computes each +# category's share of the total in one call (no manual per-category + total +# queries). +# --------------------------------------------------------------------------- # + +_EXAMPLE_A = TaskList( + plan_id="example_a", + goal_restated="Identify which product categories contributed most to last quarter's revenue.", + assumptions=["'last quarter' = 2026-01-01 to 2026-03-31."], + open_questions=[], + tasks=[ + Task( + id="t1", + stage="data_understanding", + objective="Confirm the sales source exposes category, revenue, and order date.", + tool_calls=[ToolCall(tool="describe_source", args={"source_id": "src_sales"})], + expected_output="source_shape", + success_criteria="Produced the orders table schema; the 3 needed columns are present.", + depends_on=[], + estimated_cost="low", + ), + Task( + id="t2", + stage="data_preparation", + objective="Pull last quarter's order-level category and revenue rows.", + tool_calls=[ + ToolCall( + tool="query_structured", + args={ + "ir": { + "source_id": "src_sales", + "table_id": "t_orders", + "select": [ + {"kind": "column", "column_id": "c_category", "alias": "category"}, + {"kind": "column", "column_id": "c_revenue", "alias": "revenue"}, + ], + "filters": [ + { + "column_id": "c_order_date", + "op": "between", + "value": ["2026-01-01", "2026-03-31"], + "value_type": "date", + } + ], + "limit": 10000, + } + }, + ) + ], + expected_output="quarter_rows", + success_criteria="Produced last quarter's order rows with category and revenue.", + depends_on=["t1"], + estimated_cost="medium", + ), + Task( + id="t3", + stage="evaluation", + objective="Rank each category's revenue share of the quarter total.", + tool_calls=[ + ToolCall( + tool="analyze_contribution", + args={ + "data": "${t2}", + "dimension": "category", + "value_column": "revenue", + "agg": "sum", + }, + ) + ], + expected_output="category_contribution", + success_criteria="Produced each category's revenue share, ranked high to low.", + depends_on=["t2"], + estimated_cost="low", + ), + ], +) + +# --------------------------------------------------------------------------- # +# Example B — descriptive / trend. +# "How has monthly revenue trended by region this year, and what's unusual?" +# --------------------------------------------------------------------------- # + +_EXAMPLE_B = TaskList( + plan_id="example_b", + goal_restated="Describe this year's monthly revenue trend and flag unusual months.", + assumptions=["'this year' starts 2026-01-01."], + open_questions=["'Unusual' is interpreted as months far from the typical monthly revenue."], + tasks=[ + Task( + id="t1", + stage="data_understanding", + objective="Confirm the sales source exposes order date, revenue, and region.", + tool_calls=[ToolCall(tool="describe_source", args={"source_id": "src_sales"})], + expected_output="source_shape", + success_criteria="Produced the orders table schema; the needed columns are present.", + depends_on=[], + estimated_cost="low", + ), + Task( + id="t2", + stage="data_preparation", + objective="Pull this year's order dates, revenue, and region.", + tool_calls=[ + ToolCall( + tool="query_structured", + args={ + "ir": { + "source_id": "src_sales", + "table_id": "t_orders", + "select": [ + { + "kind": "column", + "column_id": "c_order_date", + "alias": "order_date", + }, + {"kind": "column", "column_id": "c_revenue", "alias": "revenue"}, + {"kind": "column", "column_id": "c_region", "alias": "region"}, + ], + "filters": [ + { + "column_id": "c_order_date", + "op": ">=", + "value": "2026-01-01", + "value_type": "date", + } + ], + "limit": 10000, + } + }, + ) + ], + expected_output="ytd_rows", + success_criteria="Produced this year's order-level rows with date, revenue, region.", + depends_on=["t1"], + estimated_cost="medium", + ), + Task( + id="t3", + stage="evaluation", + objective="Bucket revenue into months and summarize the trend and movement.", + tool_calls=[ + ToolCall( + tool="analyze_trend", + args={ + "data": "${t2}", + "date_column": "order_date", + "value_column": "revenue", + "freq": "month", + "agg": "sum", + }, + ) + ], + expected_output="monthly_trend", + success_criteria=( + "Produced a per-month revenue series with direction and change rate to " + "flag months above/below the typical level." + ), + depends_on=["t2"], + estimated_cost="low", + ), + ], +) + + +# --------------------------------------------------------------------------- # +# Example C — mixed structured + unstructured. +# "Revenue dipped in Q1 — what happened?" +# Shows: a structured branch (query -> analyze_trend) runs alongside an +# INDEPENDENT retrieve_documents branch that pulls qualitative context. Note +# retrieve_documents takes a natural-language `query` (NOT a `${t}` data +# placeholder — it is a source, not a consumer) and can run in parallel; the +# Assembler folds the document context into the explanation. +# --------------------------------------------------------------------------- # + +_EXAMPLE_C = TaskList( + plan_id="example_c", + goal_restated="Explain Q1's revenue dip using both the numbers and the qualitative record.", + assumptions=["'Q1' = 2026-01-01 to 2026-03-31."], + open_questions=[], + tasks=[ + Task( + id="t1", + stage="data_understanding", + objective="Confirm the sales source exposes order date and revenue.", + tool_calls=[ToolCall(tool="describe_source", args={"source_id": "src_sales"})], + expected_output="source_shape", + success_criteria="Produced the orders table schema; date and revenue columns present.", + depends_on=[], + estimated_cost="low", + ), + Task( + id="t2", + stage="data_preparation", + objective="Pull Q1 order dates and revenue.", + tool_calls=[ + ToolCall( + tool="query_structured", + args={ + "ir": { + "source_id": "src_sales", + "table_id": "t_orders", + "select": [ + { + "kind": "column", + "column_id": "c_order_date", + "alias": "order_date", + }, + {"kind": "column", "column_id": "c_revenue", "alias": "revenue"}, + ], + "filters": [ + { + "column_id": "c_order_date", + "op": "between", + "value": ["2026-01-01", "2026-03-31"], + "value_type": "date", + } + ], + "limit": 10000, + } + }, + ) + ], + expected_output="q1_rows", + success_criteria="Produced Q1 order rows with date and revenue.", + depends_on=["t1"], + estimated_cost="medium", + ), + Task( + id="t3", + stage="evaluation", + objective="Summarize the Q1 monthly revenue trend to locate the dip.", + tool_calls=[ + ToolCall( + tool="analyze_trend", + args={ + "data": "${t2}", + "date_column": "order_date", + "value_column": "revenue", + "freq": "month", + "agg": "sum", + }, + ) + ], + expected_output="q1_trend", + success_criteria="Produced a per-month revenue series showing where revenue fell.", + depends_on=["t2"], + estimated_cost="low", + ), + Task( + id="t4", + stage="data_understanding", + objective="Retrieve qualitative context on Q1 operational events behind a dip.", + tool_calls=[ + ToolCall( + tool="retrieve_documents", + args={ + "query": "operational issues, outages, or notable events in Q1 2026", + "top_k": 5, + }, + ) + ], + expected_output="q1_context_chunks", + success_criteria="Produced relevant document chunks about Q1 operations.", + depends_on=[], + estimated_cost="low", + ), + ], +) + + +# --------------------------------------------------------------------------- # +# Example D — group-by aggregation (analyze_aggregate arg shape). +# "What is the average and total order value per region?" +# Shows the EXACT analyze_aggregate args: `aggregations` is an OBJECT mapping each +# column to a LIST of functions ({"revenue": ["mean", "sum"]}), and `group_by` is a +# SEPARATE array — NOT a nested list of metric specs. Supported funcs: sum, mean, +# count, min, max, median, nunique. +# --------------------------------------------------------------------------- # + +_EXAMPLE_D = TaskList( + plan_id="example_d", + goal_restated="Report the average and total order value for each region.", + assumptions=[], + open_questions=[], + tasks=[ + Task( + id="t1", + stage="data_understanding", + objective="Confirm the sales source exposes region and revenue.", + tool_calls=[ToolCall(tool="describe_source", args={"source_id": "src_sales"})], + expected_output="source_shape", + success_criteria="Produced the orders table schema; region and revenue present.", + depends_on=[], + estimated_cost="low", + ), + Task( + id="t2", + stage="data_preparation", + objective="Pull order-level region and revenue.", + tool_calls=[ + ToolCall( + tool="query_structured", + args={ + "ir": { + "source_id": "src_sales", + "table_id": "t_orders", + "select": [ + {"kind": "column", "column_id": "c_region", "alias": "region"}, + {"kind": "column", "column_id": "c_revenue", "alias": "revenue"}, + ], + "limit": 10000, + } + }, + ) + ], + expected_output="region_rows", + success_criteria="Produced order rows with region and revenue.", + depends_on=["t1"], + estimated_cost="medium", + ), + Task( + id="t3", + stage="evaluation", + objective="Aggregate mean and total revenue per region.", + tool_calls=[ + ToolCall( + tool="analyze_aggregate", + args={ + "data": "${t2}", + "aggregations": {"revenue": ["mean", "sum"]}, + "group_by": ["region"], + }, + ) + ], + expected_output="region_aggregates", + success_criteria="Produced one row per region with mean and total revenue.", + depends_on=["t2"], + estimated_cost="low", + ), + ], +) + + +EXAMPLES: list[tuple[str, TaskList]] = [ + ("Which product categories drove last quarter's revenue?", _EXAMPLE_A), + ("How has monthly revenue trended by region this year, and what's unusual?", _EXAMPLE_B), + ("Revenue dipped in Q1 — what happened?", _EXAMPLE_C), + ("What is the average and total order value per region?", _EXAMPLE_D), +] + + +def render_examples() -> str: + """Render the few-shots as text for the planner prompt.""" + blocks: list[str] = [] + for i, (question, plan) in enumerate(EXAMPLES, start=1): + blocks.append( + f"## Example {i}\n\n" + f"Question:\n{question}\n\n" + f"TaskList:\n{plan.model_dump_json(indent=2)}" + ) + return "\n\n".join(blocks) diff --git a/src/agents/planner/inputs.py b/src/agents/planner/inputs.py new file mode 100644 index 0000000000000000000000000000000000000000..95496c628c353b8145f40b00864b66c6233b4769 --- /dev/null +++ b/src/agents/planner/inputs.py @@ -0,0 +1,139 @@ +"""Planner input models — CatalogSummary and Constraints. + +`CatalogSummary` is a condensed, PII-safe view of the user's `Catalog`, built +for the planner prompt. It carries every table + column id/type/PII flag + row +counts + low-cardinality top_values, with `sample_values` nulled on PII columns +(INV: no PII sample values into the prompt, see doc §13). It also lists the +available unstructured sources so the planner can plan `retrieve_documents`. + +The planner *validator* still checks inline `query_structured` IRs against the +full `Catalog` via the existing IRValidator — the summary is a prompt input, not +the validation source of truth. + +See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3. +""" + +from __future__ import annotations + +from typing import Any + +from pydantic import BaseModel, Field + +from ...catalog.models import Catalog, DataType + + +class ColumnSummary(BaseModel): + column_id: str + name: str + data_type: DataType + pii_flag: bool = False + sample_values: list[Any] | None = None # nulled when pii_flag is True + top_values: list[Any] | None = None + + +class TableSummary(BaseModel): + table_id: str + name: str + row_count: int | None = None + columns: list[ColumnSummary] = Field(default_factory=list) + + +class StructuredSourceSummary(BaseModel): + source_id: str + name: str + source_type: str # "schema" | "tabular" + tables: list[TableSummary] = Field(default_factory=list) + + +class UnstructuredSourceSummary(BaseModel): + source_id: str + name: str + + +class CatalogSummary(BaseModel): + structured_sources: list[StructuredSourceSummary] = Field(default_factory=list) + unstructured_sources: list[UnstructuredSourceSummary] = Field(default_factory=list) + + @classmethod + def from_catalog(cls, catalog: Catalog) -> CatalogSummary: + structured: list[StructuredSourceSummary] = [] + unstructured: list[UnstructuredSourceSummary] = [] + + for source in catalog.sources: + if source.source_type == "unstructured": + unstructured.append( + UnstructuredSourceSummary(source_id=source.source_id, name=source.name) + ) + continue + + tables = [ + TableSummary( + table_id=table.table_id, + name=table.name, + row_count=table.row_count, + columns=[ + ColumnSummary( + column_id=col.column_id, + name=col.name, + data_type=col.data_type, + pii_flag=col.pii_flag, + # PII columns leak nothing into the prompt: both + # sample_values and (low-cardinality) top_values are + # suppressed — top_values are the same class of data. + sample_values=None if col.pii_flag else col.sample_values, + top_values=( + None + if col.pii_flag or col.stats is None + else col.stats.top_values + ), + ) + for col in table.columns + ], + ) + for table in source.tables + ] + structured.append( + StructuredSourceSummary( + source_id=source.source_id, + name=source.name, + source_type=source.source_type, + tables=tables, + ) + ) + + return cls(structured_sources=structured, unstructured_sources=unstructured) + + def render(self) -> str: + """Render the summary as compact text for the planner prompt.""" + if not self.structured_sources and not self.unstructured_sources: + return "(catalog is empty — the user has not registered any data yet)" + + lines: list[str] = [] + for source in self.structured_sources: + lines.append(f"Source: {source.name} ({source.source_type}) — id={source.source_id}") + for table in source.tables: + rc = f" ({table.row_count:,} rows)" if table.row_count is not None else "" + lines.append(f" Table: {table.name}{rc} — id={table.table_id}") + for col in table.columns: + samples = "PII (suppressed)" if col.pii_flag else (col.sample_values or []) + top = f", top={col.top_values}" if col.top_values else "" + lines.append( + f" - {col.name} [{col.data_type}]: " + f"samples={samples}{top} — id={col.column_id}" + ) + lines.append("") + + if self.unstructured_sources: + lines.append("Unstructured sources (for retrieve_documents):") + for src in self.unstructured_sources: + lines.append(f" - {src.name} — id={src.source_id}") + + return "\n".join(lines).rstrip() + + +class Constraints(BaseModel): + max_tasks: int = 5 + modeling_allowed: bool = False # no modeling tools in v1 + token_budget: int | None = None + time_budget_seconds: int | None = None + row_budget: int = 10_000 diff --git a/src/agents/planner/prompt.py b/src/agents/planner/prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..0de555427c3ae591d1bb0adb3b97cbfe960a8168 --- /dev/null +++ b/src/agents/planner/prompt.py @@ -0,0 +1,106 @@ +"""Builds the planner LLM human-message content. + +The system prompt (`config/prompts/planner.md`) carries the role, invariants, +and planning principles. This module assembles the per-call human content: +business context + condensed catalog + available tools + constraints + the +few-shot examples + the question (+ the prior error on retry). + +Few-shot examples are rendered from `examples.py` (which builds them from the +real `TaskList` schema) so they cannot drift from the output contract. +""" + +from __future__ import annotations + +from .contracts import BusinessContext, ToolRegistry +from .examples import render_examples +from .inputs import CatalogSummary, Constraints + + +def render_business_context(context: BusinessContext) -> str: + lines = [ + f"Project: {context.project_id} (industry: {context.industry}, " + f"context completeness: {context.completeness})", + f"Business: {context.business_description}", + f"Scale & scope: {context.scale_and_scope}", + ] + if context.key_terms: + lines.append("Key terms:") + lines.extend(f" - {kt.term}: {kt.meaning}" for kt in context.key_terms) + if context.data_overview: + lines.append("Data overview:") + lines.extend( + f" - {n.table_name}: {n.row_represents}" for n in context.data_overview + ) + if context.data_column_notes: + lines.append("Column notes:") + lines.extend( + f" - {n.column_name}: {n.meaning}" for n in context.data_column_notes + ) + if context.whats_normal: + lines.append(f"What's normal: {context.whats_normal}") + if context.recent_events: + lines.append(f"Recent events: {context.recent_events}") + if context.things_to_watch_for: + lines.append(f"Things to watch for: {context.things_to_watch_for}") + if context.open_questions: + lines.append("Known open questions:") + lines.extend(f" - {q}" for q in context.open_questions) + return "\n".join(lines) + + +def render_registry(tools: ToolRegistry) -> str: + if not tools.tools: + return "(no tools available)" + blocks: list[str] = [] + for spec in tools.tools: + required = spec.input_schema.get("required", []) + blocks.append( + f"- {spec.name} (category: {spec.category}, returns: {spec.output_kind})\n" + f" required args: {required}\n" + f" {spec.description}" + ) + return "\n".join(blocks) + + +def render_constraints(constraints: Constraints) -> str: + lines = [ + f"- max_tasks: {constraints.max_tasks}", + f"- modeling_allowed: {constraints.modeling_allowed} " + "(no modeling tools exist in v1 — do not emit modeling tasks)", + f"- row_budget: {constraints.row_budget}", + ] + if constraints.token_budget is not None: + lines.append(f"- token_budget: {constraints.token_budget}") + if constraints.time_budget_seconds is not None: + lines.append(f"- time_budget_seconds: {constraints.time_budget_seconds}") + return "\n".join(lines) + + +def build_planner_prompt( + context: BusinessContext, + catalog: CatalogSummary, + tools: ToolRegistry, + query: str, + constraints: Constraints, + previous_error: str | None = None, +) -> str: + """Return the human-message content for the planner LLM. + + The system prompt (`config/prompts/planner.md`) is loaded separately by + `PlannerService`. + """ + sections = [ + f"# Business context\n\n{render_business_context(context)}", + f"# Catalog\n\n{catalog.render()}", + f"# Available tools\n\n{render_registry(tools)}", + f"# Constraints\n\n{render_constraints(constraints)}", + f"# Examples\n\n{render_examples()}", + f"# Question\n\n{query}", + ] + if previous_error: + sections.append( + "# Previous attempt failed validation\n\n" + f"{previous_error}\n\n" + "Emit a corrected TaskList. Do not repeat the same mistake." + ) + return "\n\n".join(sections) diff --git a/src/agents/planner/registry.py b/src/agents/planner/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..4c7783ec8af1dca2b1c270ba0931b7ced52f3e65 --- /dev/null +++ b/src/agents/planner/registry.py @@ -0,0 +1,132 @@ +"""v1 tool registry the Planner plans against (INV-7: agent never names a tool +outside it). + +**Composed from two slices (2026-06-08):** + +- **Analytics (`analyze_*`) — REAL, tool-team-owned.** Sourced live from + `src/tools/registry.py::analytics_registry()` (KM-628), built on the canonical + `ToolSpec` (`src/tools/contracts.py`, KM-465/KM-627) and the prompt-style tool + descriptions (KM-625). No longer a stub on our side — it tracks the real registry. +- **Data access (`query_structured` / `retrieve_documents` / `list_sources` / + `describe_source`) — spec BODIES still a local stub.** The tool team owns these too, + but their wrappers + `ToolSpec`s haven't landed yet (KM-465 #4). We keep best-guess + spec bodies here so the Planner can plan end-to-end — but the NAMES derive from + `src.tools.data_access.DATA_ACCESS_TOOLS` (R11), so a tool rename/addition upstream + fails loudly here instead of drifting silently. When the real specs ship, delete + this slice and swap `default_registry()` for the tool team's full composition. + +**Confirmed conventions (KM-465):** Pattern A — `analyze_*` tools take a `data` +`"${t}"` placeholder pointing at an upstream `query_structured` output (no +self-fetch); resolved to a DataFrame at execution time. `input_schema` is the +lightweight `{required, properties}` dict the planner validator (check #8) reads; +`query_structured.args["ir"]` carries an inline QueryIR validated against the +catalog by the existing IRValidator. + +See AGENT_ARCHITECTURE_CONTEXT_new.md §9.2 / §9.3. +""" + +from __future__ import annotations + +from src.tools.data_access import DATA_ACCESS_TOOLS +from src.tools.registry import analytics_registry + +from .contracts import ToolRegistry, ToolSpec + +# --------------------------------------------------------------------------- # +# Data-access slice — spec bodies are a LOCAL STUB pending the tool team's real +# specs (KM-465 #4); the canonical NAME SET is `DATA_ACCESS_TOOLS` (tool-owned). +# --------------------------------------------------------------------------- # +_DATA_ACCESS_SPEC_BODIES: tuple[ToolSpec, ...] = ( + ToolSpec( + name="query_structured", + category="analytics.query", + input_schema={"required": ["ir"], "properties": {"ir": {"type": "object"}}}, + output_kind="table", + description=( + "Run one validated, single-table query against a structured source (DB " + "schema or tabular file) and return rows. The `ir` argument is an inline " + "QueryIR (the JSON intent: source_id, table_id, select, filters, group_by, " + "order_by, limit) — never SQL. This is the data-access entry point: use it " + "to select, filter, and pull the rows the analytics (`analyze_*`) tools " + "then consume. It also does simple built-in aggregation the IR can express " + "(count/sum/avg/min/max/count_distinct). Do NOT use it for richer statistics " + "(median/percentile/mode/stddev/skew → analyze_descriptive), trends " + "(analyze_trend), correlation, segmentation, or share-of-total; and do NOT " + "use it to read documents (use retrieve_documents)." + ), + ), + ToolSpec( + name="retrieve_documents", + category="retrieval.documents", + input_schema={ + "required": ["query"], + "properties": { + "query": {"type": "string"}, + "source_id": {"type": "string"}, + "top_k": {"type": "integer"}, + }, + }, + output_kind="documents", + description=( + "Dense-retrieve the most relevant chunks from the user's unstructured " + "sources (PDF/DOCX/TXT) for a natural-language `query`. Use this to pull " + "qualitative context into an analysis. Optionally scope to one `source_id`. " + "Do NOT use it for numbers in tables — that is query_structured's job." + ), + ), + ToolSpec( + name="list_sources", + category="catalog.introspection", + input_schema={"required": [], "properties": {}}, + output_kind="table", + description=( + "List the user's available data sources (id, name, type, table count). Use " + "early in data_understanding when the plan must discover what exists before " + "querying. Cheap. Do NOT use it to read column details (use describe_source)." + ), + ), + ToolSpec( + name="describe_source", + category="catalog.introspection", + input_schema={ + "required": ["source_id"], + "properties": {"source_id": {"type": "string"}}, + }, + output_kind="table", + description=( + "Return the tables and columns (names, types, row counts) of one source by " + "`source_id`. Use in data_understanding to confirm the shape of a source " + "before querying it. Do NOT use it to fetch data rows (use query_structured)." + ), + ), +) + +_DATA_ACCESS_SPECS: dict[str, ToolSpec] = {s.name: s for s in _DATA_ACCESS_SPEC_BODIES} + + +def _data_access_slice() -> list[ToolSpec]: + """Data-access specs in body order, with names checked against the tool layer. + + `DATA_ACCESS_TOOLS` (src.tools.data_access) is the canonical name set; the + spec bodies above are still our local stub. Any mismatch (a tool added, + renamed, or removed upstream) raises here instead of drifting silently. + """ + if set(_DATA_ACCESS_SPECS) != DATA_ACCESS_TOOLS: + missing = sorted(DATA_ACCESS_TOOLS - _DATA_ACCESS_SPECS.keys()) + stale = sorted(_DATA_ACCESS_SPECS.keys() - DATA_ACCESS_TOOLS) + raise RuntimeError( + "planner data-access specs out of sync with " + f"src.tools.data_access.DATA_ACCESS_TOOLS: missing spec for {missing}, " + f"stale spec for {stale}" + ) + return list(_DATA_ACCESS_SPECS.values()) + + +def default_registry() -> ToolRegistry: + """The v1 registry: stub data-access slice + the real analytics slice. + + The analytics tools come live from `src.tools.registry` (the tool team's real + registry); the data-access spec bodies are still a local stub, name-checked + against `DATA_ACCESS_TOOLS`. A fresh instance per call. + """ + return ToolRegistry(tools=[*_data_access_slice(), *analytics_registry().tools]) diff --git a/src/agents/planner/schemas.py b/src/agents/planner/schemas.py new file mode 100644 index 0000000000000000000000000000000000000000..9fb9e5614aa732fb76ddb162975010a298108521 --- /dev/null +++ b/src/agents/planner/schemas.py @@ -0,0 +1,60 @@ +"""Planner output schemas — the `TaskList` contract. + +The planner emits exactly one `TaskList`: a DAG of typed tasks, each an ordered +chain of fully-specified tool calls. This is the static plan the TaskRunner +executes verbatim. There is no replanning (INV-6), so there are no +`ReplanRequest` / `ReplanResponse` schemas. + +See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3. +""" + +from __future__ import annotations + +import re +from typing import Any, Literal + +from pydantic import BaseModel, Field + +# The "${t}" placeholder convention (Pattern A): a ToolCall arg whose value +# matches this pattern refers to an upstream task's output, resolved by the +# TaskRunner at execution time. Single definition — the planner validator and +# the TaskRunner both import it (R11). +PLACEHOLDER_RE = re.compile(r"\$\{(t[^}]+)\}") + +CrispStage = Literal[ + "data_understanding", + "data_preparation", + "modeling", # no tools in v1; the planner does not emit modeling tasks + "evaluation", +] + + +class ToolCall(BaseModel): + """One call to a registry tool with concrete, fully-specified arguments. + + `tool` must exist in the ToolRegistry. `args` is validated against the + tool's input_schema; it may contain "${t}" placeholders that the + TaskRunner resolves from an upstream task's output at execution time. + """ + + tool: str + args: dict[str, Any] = Field(default_factory=dict) + + +class Task(BaseModel): + id: str # "t1", "t2", ... + stage: CrispStage + objective: str # plain-language intent for this step + tool_calls: list[ToolCall] = Field(..., min_length=1) # ordered chain + expected_output: str # named result this task produces + success_criteria: str # REPORTING signal, not a control trigger + depends_on: list[str] = Field(default_factory=list) # task ids + estimated_cost: Literal["low", "medium", "high"] = "low" + + +class TaskList(BaseModel): + plan_id: str + goal_restated: str + assumptions: list[str] = Field(default_factory=list) + open_questions: list[str] = Field(default_factory=list) + tasks: list[Task] = Field(default_factory=list) diff --git a/src/agents/planner/service.py b/src/agents/planner/service.py new file mode 100644 index 0000000000000000000000000000000000000000..37bfb55044a58f155240cb0391bae05527782ca6 --- /dev/null +++ b/src/agents/planner/service.py @@ -0,0 +1,158 @@ +"""PlannerService — single LLM call: context + catalog + tools + question -> TaskList. + +Mirrors `query/planner/service.py` (chain construction) and `query/service.py` +(validate-and-retry loop). The planner LLM emits a `TaskList` via structured +output; the `PlannerValidator` runs the 8 checks; on failure the planner is +re-prompted with the error context, up to `max_retries` (default 3). No +replanning happens at execution time — this loop only hardens the *initial* +static plan. + +The service takes the full `Catalog` (not just a `CatalogSummary`): it derives +the PII-safe `CatalogSummary` for the prompt, but validation needs the full +catalog so the existing `IRValidator` can check inline `query_structured` IRs. + +See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3. +""" + +from __future__ import annotations + +from pathlib import Path + +from langchain_core.messages import SystemMessage +from langchain_core.prompts import ChatPromptTemplate +from langchain_core.runnables import Runnable +from langchain_openai import AzureChatOpenAI + +from src.middlewares.logging import get_logger + +from ...catalog.models import Catalog +from .contracts import BusinessContext, ToolRegistry +from .errors import PlannerError, PlannerValidationError +from .inputs import CatalogSummary, Constraints +from .prompt import build_planner_prompt +from .schemas import TaskList +from .validator import PlannerValidator + +logger = get_logger("planner_agent") + +_PROMPT_PATH = ( + Path(__file__).resolve().parent.parent.parent / "config" / "prompts" / "planner.md" +) + + +def _load_prompt_text() -> str: + return _PROMPT_PATH.read_text(encoding="utf-8") + + +def _build_default_chain() -> Runnable: + from src.config.settings import settings + + llm = AzureChatOpenAI( + azure_deployment=settings.azureai_deployment_name_4o, + openai_api_version=settings.azureai_api_version_4o, + azure_endpoint=settings.azureai_endpoint_url_4o, + api_key=settings.azureai_api_key_4o, + temperature=0, + ) + prompt = ChatPromptTemplate.from_messages( + [ + SystemMessage(content=_load_prompt_text()), + ("human", "{human_content}"), + ] + ) + return prompt | llm.with_structured_output(TaskList) + + +_default_chain: Runnable | None = None + + +def _get_default_chain() -> Runnable: + global _default_chain + if _default_chain is None: + _default_chain = _build_default_chain() + return _default_chain + + +class PlannerService: + """Wraps the planner LLM call + the validate-and-retry loop. + + Inject `structured_chain` and/or `validator` for tests. + """ + + def __init__( + self, + structured_chain: Runnable | None = None, + validator: PlannerValidator | None = None, + max_retries: int = 3, + ) -> None: + self._chain = structured_chain + self._validator = validator or PlannerValidator() + self._max_retries = max(1, max_retries) + + def _ensure_chain(self) -> Runnable: + if self._chain is None: + self._chain = _get_default_chain() + return self._chain + + async def plan( + self, + context: BusinessContext, + catalog: Catalog, + tools: ToolRegistry, + query: str, + constraints: Constraints, + callbacks: list | None = None, + ) -> TaskList: + summary = CatalogSummary.from_catalog(catalog) + chain = self._ensure_chain() + previous_error: str | None = None + + for attempt in range(1, self._max_retries + 1): + human_content = build_planner_prompt( + context, summary, tools, query, constraints, previous_error + ) + # All retry attempts share `callbacks`, so each shows up under the same + # trace — that is how retry token cost becomes visible. + if callbacks: + task_list: TaskList = await chain.ainvoke( + {"human_content": human_content}, config={"callbacks": callbacks} + ) + else: + task_list = await chain.ainvoke({"human_content": human_content}) + try: + self._validator.validate(task_list, tools, catalog, constraints) + except PlannerValidationError as e: + previous_error = str(e) + logger.warning( + "planner validation failed", + project_id=context.project_id, + plan_id=task_list.plan_id, + attempt=attempt, + error=previous_error, + ) + continue + + logger.info( + "analysis planned", + project_id=context.project_id, + plan_id=task_list.plan_id, + n_tasks=len(task_list.tasks), + retry=attempt > 1, + ) + return task_list + + raise PlannerError( + f"planner failed validation after {self._max_retries} attempts; " + f"last error: {previous_error}" + ) + + +async def plan_analysis( + context: BusinessContext, + catalog: Catalog, + tools: ToolRegistry, + query: str, + constraints: Constraints, +) -> TaskList: + """Convenience entry point using the default chain + validator.""" + return await PlannerService().plan(context, catalog, tools, query, constraints) diff --git a/src/agents/planner/validator.py b/src/agents/planner/validator.py new file mode 100644 index 0000000000000000000000000000000000000000..c4333886a30c37225ef0dfef9dfc0d96563827d3 --- /dev/null +++ b/src/agents/planner/validator.py @@ -0,0 +1,229 @@ +"""PlannerValidator — checks a TaskList before it reaches the TaskRunner. + +Runs the 8 checks from AGENT_ARCHITECTURE_CONTEXT_new.md §7.3. On failure it +raises `PlannerValidationError` with a message specific enough that the planner +can be re-prompted to self-correct (the retry loop lives in service.py). + +Check #1 (Pydantic parse) is enforced at the structured-output boundary — by the +time a `TaskList` reaches here it has already parsed; this validator additionally +rejects structurally-invalid plans (duplicate ids, dangling edges, cycles). +""" + +from __future__ import annotations + +from pydantic import ValidationError + +from ...catalog.models import Catalog +from ...query.ir.models import QueryIR +from ...query.ir.validator import IRValidationError, IRValidator +from .contracts import ToolRegistry +from .errors import PlannerValidationError +from .inputs import Constraints +from .schemas import PLACEHOLDER_RE, TaskList + +# Heuristic: a checkable success_criteria mentions a measurable signal. +_CHECKABLE_TOKENS = ("rate", "count", "match", "produced", "above", "below", "equal") + +# DFS colors for cycle detection. +_WHITE, _GREY, _BLACK = 0, 1, 2 + + +class PlannerValidator: + def __init__(self, ir_validator: IRValidator | None = None) -> None: + self._ir_validator = ir_validator or IRValidator() + + def validate( + self, + task_list: TaskList, + registry: ToolRegistry, + catalog: Catalog, + constraints: Constraints, + ) -> None: + tasks = task_list.tasks + + # Check 6 — plan non-empty and within the task cap. + if not tasks: + raise PlannerValidationError("plan is empty: at least one task is required") + if len(tasks) > constraints.max_tasks: + raise PlannerValidationError( + f"plan has {len(tasks)} tasks, exceeds max_tasks={constraints.max_tasks}" + ) + + ids = [t.id for t in tasks] + if len(set(ids)) != len(ids): + dupes = sorted({i for i in ids if ids.count(i) > 1}) + raise PlannerValidationError(f"duplicate task id(s): {dupes}") + id_set = set(ids) + tasks_by_id = {t.id: t for t in tasks} + + known_tools = registry.names() + known_sources = {s.source_id for s in catalog.sources} + + for task in tasks: + for call in task.tool_calls: + # Check 2 — every tool exists in the registry. + if call.tool not in known_tools: + raise PlannerValidationError( + f"task {task.id}: tool {call.tool!r} not in registry " + f"(known: {sorted(known_tools)})" + ) + spec = registry.get(call.tool) + assert spec is not None # guaranteed by the membership check above + + # Check 8a — args carry the required keys and no unknown keys. + required = set(spec.input_schema.get("required", [])) + allowed = set(spec.input_schema.get("properties", {}).keys()) | required + missing = required - set(call.args.keys()) + if missing: + raise PlannerValidationError( + f"task {task.id}: tool {call.tool!r} missing required arg(s): " + f"{sorted(missing)}" + ) + unknown = set(call.args.keys()) - allowed + if unknown: + raise PlannerValidationError( + f"task {task.id}: tool {call.tool!r} has unknown arg(s): " + f"{sorted(unknown)} (allowed: {sorted(allowed)})" + ) + + # Check 3 — concrete source_id args must exist in the catalog. + src = call.args.get("source_id") + if isinstance(src, str) and not _is_placeholder(src): + if src not in known_sources: + raise PlannerValidationError( + f"task {task.id}: tool {call.tool!r} references unknown " + f"source_id {src!r} (known: {sorted(known_sources)})" + ) + + # Check 8b — inline query_structured IR validates against the catalog. + if call.tool == "query_structured": + self._validate_inline_ir(task.id, call.args, catalog) + + # Check 7 — success_criteria is checkable. + if not _is_checkable(task.success_criteria): + raise PlannerValidationError( + f"task {task.id}: success_criteria is not checkable — include a " + f"measurable signal (one of {list(_CHECKABLE_TOKENS)}); " + f"got {task.success_criteria!r}" + ) + + # Check 4 — DAG: edges resolve, placeholders resolve, no cycles. + self._validate_dag(tasks_by_id, id_set) + + def _validate_inline_ir(self, task_id: str, args: dict, catalog: Catalog) -> None: + raw_ir = args.get("ir") + if not isinstance(raw_ir, dict): + raise PlannerValidationError( + f"task {task_id}: query_structured.args.ir must be an inline QueryIR " + f"object, got {type(raw_ir).__name__}" + ) + try: + ir = QueryIR.model_validate(raw_ir) + except ValidationError as e: + raise PlannerValidationError( + f"task {task_id}: query_structured.args.ir is not a valid QueryIR: {e}" + ) from e + try: + self._ir_validator.validate(ir, catalog) + except IRValidationError as e: + raise PlannerValidationError( + f"task {task_id}: query_structured IR failed catalog validation: {e}" + ) from e + + @staticmethod + def _validate_dag(tasks_by_id: dict, id_set: set[str]) -> None: + for task in tasks_by_id.values(): + for dep in task.depends_on: + if dep not in id_set: + raise PlannerValidationError( + f"task {task.id}: depends_on references unknown task {dep!r}" + ) + if dep == task.id: + raise PlannerValidationError( + f"task {task.id}: depends_on includes itself" + ) + + cycle = _find_cycle(tasks_by_id) + if cycle: + raise PlannerValidationError(f"cycle detected in depends_on: {' -> '.join(cycle)}") + + # On an acyclic graph, a placeholder is safe iff its target is a + # transitive ancestor — i.e. guaranteed to have completed before this + # task runs. Requiring a *direct* depends_on would wrongly reject valid + # plans that depend on the target through an intermediate task. + ancestors = _all_ancestors(tasks_by_id) + for task in tasks_by_id.values(): + for ref in _placeholder_refs(task): + if ref not in id_set: + raise PlannerValidationError( + f"task {task.id}: placeholder '${{{ref}}}' references unknown task" + ) + if ref not in ancestors[task.id]: + raise PlannerValidationError( + f"task {task.id}: placeholder '${{{ref}}}' used but {ref!r} is " + f"not a (transitive) dependency — add it to depends_on" + ) + + +def _is_placeholder(value: str) -> bool: + return bool(PLACEHOLDER_RE.fullmatch(value.strip())) + + +def _placeholder_refs(task) -> set[str]: + refs: set[str] = set() + for call in task.tool_calls: + for value in call.args.values(): + if isinstance(value, str): + refs.update(PLACEHOLDER_RE.findall(value)) + return refs + + +def _is_checkable(text: str) -> bool: + low = text.lower() + return any(tok in low for tok in _CHECKABLE_TOKENS) + + +def _find_cycle(tasks_by_id: dict) -> list[str] | None: + color = {tid: _WHITE for tid in tasks_by_id} + stack: list[str] = [] + + def dfs(node: str) -> list[str] | None: + color[node] = _GREY + stack.append(node) + for dep in tasks_by_id[node].depends_on: + if color.get(dep) == _GREY: + idx = stack.index(dep) + return stack[idx:] + [dep] + if color.get(dep) == _WHITE: + found = dfs(dep) + if found: + return found + stack.pop() + color[node] = _BLACK + return None + + for tid in tasks_by_id: + if color[tid] == _WHITE: + found = dfs(tid) + if found: + return found + return None + + +def _all_ancestors(tasks_by_id: dict) -> dict[str, set[str]]: + """ancestors[id] = all tasks reachable by following depends_on edges.""" + cache: dict[str, set[str]] = {} + + def visit(node: str, seen: set[str]) -> set[str]: + if node in cache: + return cache[node] + acc: set[str] = set() + for dep in tasks_by_id[node].depends_on: + if dep in seen or dep not in tasks_by_id: + continue + acc.add(dep) + acc |= visit(dep, seen | {dep}) + cache[node] = acc + return acc + + return {tid: visit(tid, {tid}) for tid in tasks_by_id} diff --git a/src/agents/slow_path/__init__.py b/src/agents/slow_path/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c894a36d6ad76a632df17ebcc24645b0b01ac0dc --- /dev/null +++ b/src/agents/slow_path/__init__.py @@ -0,0 +1,10 @@ +"""Slow-path workers: TaskRunner (deterministic) + Assembler (1 LLM call) + Coordinator. + +These are driven *by* the Orchestrator (the intent-router/dispatcher in +`agents/orchestration.py`); this package is deliberately NOT named "orchestrator" +to keep the dispatcher and the workers from sharing a name. It executes the +Planner's static `TaskList` and assembles the two outputs (`chat_answer` + +`AnalysisRecord`). See AGENT_ARCHITECTURE_CONTEXT_new.md §7.2 / §7.4 / §7.5 / +§8.2–8.4. Tool-agnostic: depends only on the `ToolInvoker` protocol and the +`ToolOutput` envelope, never on a specific tool (INV-7). +""" diff --git a/src/agents/slow_path/assembler.py b/src/agents/slow_path/assembler.py new file mode 100644 index 0000000000000000000000000000000000000000..88ab7aca02b49782502e79338063a703b3c2edc9 --- /dev/null +++ b/src/agents/slow_path/assembler.py @@ -0,0 +1,140 @@ +"""Assembler — single LLM call at the end of the slow path. + +Reads the `RunState` (all `TaskResult`s) + `BusinessContext` and produces an +`AssembledOutput` { chat_answer, analysis_record }. Owns all language/output: prose, +markdown tables, citations, and merging structured + unstructured results. + +The model authors only the *narrative* (`AssemblerNarrative`); this service copies +the structured pass-through (`results_snapshot`, `tasks_run`) and metadata from the +`RunState` so the record stays a faithful source of truth (§8.3, INV-4). + +Chain construction mirrors `agents/planner/service.py`. + +See AGENT_ARCHITECTURE_CONTEXT_new.md §7.5. +""" + +from __future__ import annotations + +from datetime import UTC, datetime +from pathlib import Path + +from langchain_core.messages import SystemMessage +from langchain_core.prompts import ChatPromptTemplate +from langchain_core.runnables import Runnable +from langchain_openai import AzureChatOpenAI + +from src.middlewares.logging import get_logger + +from ..planner.contracts import BusinessContext +from .errors import AssemblerError +from .prompt import build_assembler_prompt +from .schemas import ( + AnalysisRecord, + AssembledOutput, + AssemblerNarrative, + RunState, + TaskSummary, +) + +logger = get_logger("assembler") + +_PROMPT_PATH = ( + Path(__file__).resolve().parent.parent.parent / "config" / "prompts" / "assembler.md" +) + + +def _load_prompt_text() -> str: + return _PROMPT_PATH.read_text(encoding="utf-8") + + +def _build_default_chain() -> Runnable: + from src.config.settings import settings + + llm = AzureChatOpenAI( + azure_deployment=settings.azureai_deployment_name_4o, + openai_api_version=settings.azureai_api_version_4o, + azure_endpoint=settings.azureai_endpoint_url_4o, + api_key=settings.azureai_api_key_4o, + temperature=0, + ) + prompt = ChatPromptTemplate.from_messages( + [ + SystemMessage(content=_load_prompt_text()), + ("human", "{human_content}"), + ] + ) + return prompt | llm.with_structured_output(AssemblerNarrative) + + +_default_chain: Runnable | None = None + + +def _get_default_chain() -> Runnable: + global _default_chain + if _default_chain is None: + _default_chain = _build_default_chain() + return _default_chain + + +class Assembler: + """Wraps the single Assembler LLM call. Inject `structured_chain` for tests.""" + + def __init__(self, structured_chain: Runnable | None = None) -> None: + self._chain = structured_chain + + def _ensure_chain(self) -> Runnable: + if self._chain is None: + self._chain = _get_default_chain() + return self._chain + + async def assemble( + self, + run_state: RunState, + context: BusinessContext, + question: str | None = None, + callbacks: list | None = None, + ) -> AssembledOutput: + chain = self._ensure_chain() + human_content = build_assembler_prompt(run_state, context, question) + try: + if callbacks: + narrative: AssemblerNarrative = await chain.ainvoke( + {"human_content": human_content}, config={"callbacks": callbacks} + ) + else: + narrative = await chain.ainvoke({"human_content": human_content}) + except Exception as exc: # surface as a typed error for the caller + raise AssemblerError(f"assembler call failed: {exc}") from exc + + record = _build_record(narrative, run_state) + logger.info( + "analysis assembled", + plan_id=run_state.plan_id, + business_context_id=run_state.business_context_id, + n_tasks=len(run_state.results), + ) + return AssembledOutput(chat_answer=narrative.chat_answer, analysis_record=record) + + +def _build_record(narrative: AssemblerNarrative, run_state: RunState) -> AnalysisRecord: + tasks_run = [ + TaskSummary( + task_id=task_id, + objective=result.objective, + status=result.status, + tools_used=[o.tool for o in result.outputs], + ) + for task_id, result in run_state.results.items() + ] + return AnalysisRecord( + goal_restated=narrative.goal_restated, + findings=narrative.findings, + caveats=narrative.caveats, + data_used=narrative.data_used, + open_questions=narrative.open_questions, + tasks_run=tasks_run, + results_snapshot=run_state.results, + plan_id=run_state.plan_id, + business_context_id=run_state.business_context_id, + created_at=datetime.now(UTC), + ) diff --git a/src/agents/slow_path/coordinator.py b/src/agents/slow_path/coordinator.py new file mode 100644 index 0000000000000000000000000000000000000000..4fe68a9b56f61d2d17745b2ae57825cc3ae004ba --- /dev/null +++ b/src/agents/slow_path/coordinator.py @@ -0,0 +1,66 @@ +"""SlowPathCoordinator — wires the slow path: Planner -> TaskRunner -> Assembler. + +A thin coordination object. This is the unit the (future) expanded Orchestrator / +ChatHandler will call on a `structured` analytical query. It is built and tested +here but **not yet wired into the live chat flow** — that step waits on the tool +team's real `ToolInvoker` and a real `BusinessContext` source. + +See AGENT_ARCHITECTURE_CONTEXT_new.md §5.2 / §6.1. +""" + +from __future__ import annotations + +from collections.abc import Awaitable, Callable + +from ...catalog.models import Catalog +from ..planner.contracts import BusinessContext, ToolRegistry +from ..planner.inputs import Constraints +from ..planner.service import PlannerService +from .assembler import Assembler +from .schemas import AssembledOutput +from .task_runner import TaskRunner + + +class SlowPathCoordinator: + def __init__( + self, + planner: PlannerService, + task_runner: TaskRunner, + assembler: Assembler, + registry: ToolRegistry, + ) -> None: + self._planner = planner + self._task_runner = task_runner + self._assembler = assembler + self._registry = registry + + async def run( + self, + context: BusinessContext, + catalog: Catalog, + query: str, + constraints: Constraints, + planner_callbacks: list | None = None, + assembler_callbacks: list | None = None, + progress: Callable[[str], Awaitable[None]] | None = None, + ) -> AssembledOutput: + # `progress` (optional) surfaces per-stage status to the caller so a long + # slow-path run isn't a silent ~12s on the wire. Each stage is a single + # awaitable, so the most granular signal we can emit is at stage boundaries. + if progress: + await progress("Planning the analysis…") + plan_kw = {"callbacks": planner_callbacks} if planner_callbacks else {} + task_list = await self._planner.plan( + context, catalog, self._registry, query, constraints, **plan_kw + ) + if progress: + await progress(f"Running {len(task_list.tasks)} analysis steps…") + run_state = await self._task_runner.run( + task_list, business_context_id=context.project_id + ) + if progress: + await progress("Composing the answer…") + asm_kw = {"callbacks": assembler_callbacks} if assembler_callbacks else {} + return await self._assembler.assemble( + run_state, context, question=query, **asm_kw + ) diff --git a/src/agents/slow_path/errors.py b/src/agents/slow_path/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..527b02e68b2fca8aac98feae16423869b8d44b3f --- /dev/null +++ b/src/agents/slow_path/errors.py @@ -0,0 +1,11 @@ +"""Typed errors for the slow-path layer.""" + +from __future__ import annotations + + +class SlowPathError(Exception): + """Base error for the slow-path layer.""" + + +class AssemblerError(SlowPathError): + """The Assembler LLM call could not produce a valid `AssembledOutput`.""" diff --git a/src/agents/slow_path/invoker.py b/src/agents/slow_path/invoker.py new file mode 100644 index 0000000000000000000000000000000000000000..72c714ddc0c3e4165172ba9a9f2a0ab4f04c30e6 --- /dev/null +++ b/src/agents/slow_path/invoker.py @@ -0,0 +1,27 @@ +"""The tool invocation seam (§8.4) — the one interface the TaskRunner calls. + +The agent layer stays tool-agnostic (INV-7) by invoking every tool through this +protocol, never importing a tool module directly. The **tool team owns the +implementation** (KM-418); this file defines only the contract the TaskRunner +depends on. + +Frozen guarantees the implementation must hold: +1. **Never throws.** A tool failure returns `ToolOutput(kind="error", error=...)`, + not an exception — the TaskRunner's degrade-and-continue (§7.4) relies on this. + (The TaskRunner still wraps calls defensively, as a backstop.) +2. **Returns the `ToolOutput` envelope** (§8.1) — structured data only, never + rendered tables or prose (that is the Assembler's job). +3. **`tool_name` comes from the registry** (§9.2); unknown names return an error + envelope rather than throwing. +""" + +from __future__ import annotations + +from typing import Any, Protocol, runtime_checkable + +from ..planner.contracts import ToolOutput + + +@runtime_checkable +class ToolInvoker(Protocol): + async def invoke(self, tool_name: str, args: dict[str, Any]) -> ToolOutput: ... diff --git a/src/agents/slow_path/prompt.py b/src/agents/slow_path/prompt.py new file mode 100644 index 0000000000000000000000000000000000000000..15e380f31b5630e9ae958e469e92bf6bb2224557 --- /dev/null +++ b/src/agents/slow_path/prompt.py @@ -0,0 +1,66 @@ +"""Builds the Assembler LLM human-message content. + +The system prompt (`config/prompts/assembler.md`) carries the role and rules. This +module assembles the per-call human content: the business context + the executed +`RunState` (task objectives, statuses, and structured tool outputs) + the original +question. Tool outputs are rendered compactly as data — the model turns them into +prose and markdown tables. +""" + +from __future__ import annotations + +from ..planner.contracts import BusinessContext, ToolOutput +from ..planner.prompt import render_business_context +from .schemas import RunState, TaskResult + +_MAX_ROWS = 20 + + +def render_run_state(run_state: RunState) -> str: + lines = [f"Plan: {run_state.plan_id}"] + if run_state.open_questions: + lines.append("Open questions carried from the plan:") + lines.extend(f" - {q}" for q in run_state.open_questions) + lines.append("") + lines.append("Task results (in execution order):") + for task_id, result in run_state.results.items(): + lines.append(_render_task(task_id, result)) + return "\n".join(lines) + + +def _render_task(task_id: str, result: TaskResult) -> str: + lines = [f"- [{result.status}] {task_id}: {result.objective}"] + if result.error: + lines.append(f" note: {result.error}") + for output in result.outputs: + lines.append(f" {_render_output(output)}") + return "\n".join(lines) + + +def _render_output(output: ToolOutput) -> str: + if output.kind == "error": + return f"({output.tool}) error: {output.error}" + if output.kind == "table" and output.columns is not None: + header = ", ".join(output.columns) + rows = output.rows or [] + preview = "; ".join( + " | ".join(str(cell) for cell in row) for row in rows[:_MAX_ROWS] + ) + more = "" if len(rows) <= _MAX_ROWS else f" … (+{len(rows) - _MAX_ROWS} more rows)" + return f"({output.tool}) table [{header}]: {preview}{more}" + meta = f" meta={output.meta}" if output.meta else "" + return f"({output.tool}) {output.kind}: {output.value}{meta}" + + +def build_assembler_prompt( + run_state: RunState, + context: BusinessContext, + question: str | None = None, +) -> str: + sections = [ + f"# Business context\n\n{render_business_context(context)}", + f"# Analysis results\n\n{render_run_state(run_state)}", + ] + if question: + sections.append(f"# Original question\n\n{question}") + return "\n\n".join(sections) diff --git a/src/agents/slow_path/schemas.py b/src/agents/slow_path/schemas.py new file mode 100644 index 0000000000000000000000000000000000000000..42d4aa2d729815a4ff8a1fa01a4baac93e562176 --- /dev/null +++ b/src/agents/slow_path/schemas.py @@ -0,0 +1,99 @@ +"""Slow-path execution + output contracts. + +The seams between the three slow-path stages: +- TaskRunner writes `RunState` (a blackboard of `TaskResult`s) — §8.2. +- Assembler reads `RunState` + `BusinessContext` and produces `AssembledOutput` + (`chat_answer` + `AnalysisRecord`) — §8.3. + +`ToolOutput` (the tool -> agent envelope) is reused from the planner contracts so +there is exactly one definition across the layer. + +Note on authorship (§8.3): the Assembler LLM authors only the *narrative* fields +(`AssemblerNarrative`). The `AnalysisRecord`'s structured pass-through fields +(`results_snapshot`, `tasks_run`) and metadata are copied from `RunState` by code, +never re-authored by the model — that is the source of truth the report generator +renders from (INV-4). + +See AGENT_ARCHITECTURE_CONTEXT_new.md §8.2 / §8.3. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Literal + +from pydantic import BaseModel, Field + +from ..planner.contracts import ToolOutput + +TaskStatus = Literal["success", "partial", "failure"] + + +# --------------------------------------------------------------------------- # +# Execution state (TaskRunner -> Assembler) — §8.2 +# --------------------------------------------------------------------------- # + + +class TaskResult(BaseModel): + task_id: str + status: TaskStatus + objective: str + outputs: list[ToolOutput] = Field(default_factory=list) # one per tool_call + note: str | None = None + error: str | None = None + + +class RunState(BaseModel): + plan_id: str + business_context_id: str + results: dict[str, TaskResult] = Field(default_factory=dict) # task_id -> result + open_questions: list[str] = Field(default_factory=list) + + +# --------------------------------------------------------------------------- # +# Assembled output (Assembler -> Orchestrator / memory) — §8.3 +# --------------------------------------------------------------------------- # + + +class TaskSummary(BaseModel): + task_id: str + objective: str + status: TaskStatus + tools_used: list[str] = Field(default_factory=list) + + +class AnalysisRecord(BaseModel): + # Narrative fields — authored by the Assembler LLM. + goal_restated: str + findings: list[str] = Field(default_factory=list) + caveats: list[str] = Field(default_factory=list) + data_used: list[str] = Field(default_factory=list) + open_questions: list[str] = Field(default_factory=list) + # Structured pass-through — NOT re-authored; copied from RunState. + tasks_run: list[TaskSummary] = Field(default_factory=list) + results_snapshot: dict[str, TaskResult] = Field(default_factory=dict) + # Metadata. + plan_id: str + business_context_id: str + created_at: datetime + + +class AssembledOutput(BaseModel): + chat_answer: str # FIRST field — streams via SSE; markdown prose + tables + analysis_record: AnalysisRecord + + +class AssemblerNarrative(BaseModel): + """The subset of `AnalysisRecord` the Assembler LLM actually authors. + + Kept separate from `AssembledOutput` so the model never emits the structured + pass-through fields (which would invite hallucinated numbers); `Assembler` + code merges this with the real `RunState` to build the final record. + """ + + chat_answer: str + goal_restated: str + findings: list[str] = Field(default_factory=list) + caveats: list[str] = Field(default_factory=list) + data_used: list[str] = Field(default_factory=list) + open_questions: list[str] = Field(default_factory=list) diff --git a/src/agents/slow_path/store.py b/src/agents/slow_path/store.py new file mode 100644 index 0000000000000000000000000000000000000000..43181a8615ca0d73848a2cbdedea2c490f2419bc --- /dev/null +++ b/src/agents/slow_path/store.py @@ -0,0 +1,44 @@ +"""AnalysisStore — the seam the slow path persists its AnalysisRecord through. + +The Assembler produces an `AnalysisRecord` (the faithful, structured record of a +run — §8.3, INV-4). Persisting it is a separate concern from streaming the answer, +so it sits behind this one-method seam. + +`NullAnalysisStore` is the default: it logs that a record was produced but stores +nothing, because the backing table does not exist yet. The plan is to store records +in the **same catalog DB** (Neon `dataeyond`, `settings.postgres_connstring`). + +TODO(persistence): add a Postgres-backed `AnalysisStore` writing an +`analysis_records` table in the catalog DB, keyed on +(business_context_id, plan_id, created_at), then inject it into ChatHandler. +""" + +from __future__ import annotations + +from typing import Protocol, runtime_checkable + +from src.middlewares.logging import get_logger + +from .schemas import AnalysisRecord + +logger = get_logger("analysis_store") + + +@runtime_checkable +class AnalysisStore(Protocol): + """Persist a completed analysis. Implementations must never raise on the + caller's path — a persistence failure must not break the user's answer.""" + + async def save(self, record: AnalysisRecord) -> None: ... + + +class NullAnalysisStore: + """Default no-op store: logs the record, persists nothing (no table yet).""" + + async def save(self, record: AnalysisRecord) -> None: + logger.info( + "analysis_record produced (not persisted — no store configured)", + plan_id=record.plan_id, + business_context_id=record.business_context_id, + n_tasks=len(record.tasks_run), + ) diff --git a/src/agents/slow_path/task_runner.py b/src/agents/slow_path/task_runner.py new file mode 100644 index 0000000000000000000000000000000000000000..aeac6ea9a14210043c25327f2cdedfd38557337a --- /dev/null +++ b/src/agents/slow_path/task_runner.py @@ -0,0 +1,164 @@ +"""TaskRunner — deterministic execution of a static `TaskList`. Zero LLM. + +Executes tasks in dependency order, parallelizing each ready "wave" with +`asyncio.gather`. For each task it resolves `${t}` placeholders from upstream +results, does an internal `validate_args`, invokes each tool via the `ToolInvoker` +seam, and records a `TaskResult`. On failure it **degrades and continues**: the +task is marked failed, its dependents are skipped, independent branches keep +running. There is no replanning and no mid-run LLM (INV-6). + +`success_criteria` is *not* machine-evaluated here (it is free text); task status +is derived from tool execution outcomes and carried to the Assembler to report. + +See AGENT_ARCHITECTURE_CONTEXT_new.md §7.4. +""" + +from __future__ import annotations + +import asyncio +from typing import Any + +from src.middlewares.logging import get_logger + +from ..planner.contracts import ToolOutput, ToolRegistry +from ..planner.schemas import PLACEHOLDER_RE, Task +from ..planner.schemas import TaskList as PlanTaskList +from .invoker import ToolInvoker +from .schemas import RunState, TaskResult, TaskStatus + +logger = get_logger("task_runner") + + +class TaskRunner: + """Runs a `TaskList` against a `ToolInvoker`, producing a `RunState`.""" + + def __init__(self, invoker: ToolInvoker, registry: ToolRegistry) -> None: + self._invoker = invoker + self._registry = registry + + async def run(self, task_list: PlanTaskList, business_context_id: str) -> RunState: + tasks_by_id: dict[str, Task] = {t.id: t for t in task_list.tasks} + results: dict[str, TaskResult] = {} + remaining: set[str] = set(tasks_by_id) + + while remaining: + ready = [ + tid + for tid in remaining + if all(dep in results for dep in tasks_by_id[tid].depends_on) + ] + if not ready: + # A dependency points outside the plan (or a cycle slipped past the + # planner validator): nothing more can run. Fail the rest honestly. + for tid in list(remaining): + results[tid] = TaskResult( + task_id=tid, + status="failure", + objective=tasks_by_id[tid].objective, + error="unresolved dependency; task could not run", + ) + remaining.discard(tid) + break + + # Skip any ready task whose dependency failed (degrade-and-continue). + to_run: list[Task] = [] + for tid in ready: + task = tasks_by_id[tid] + failed = [d for d in task.depends_on if results[d].status == "failure"] + if failed: + results[tid] = TaskResult( + task_id=tid, + status="failure", + objective=task.objective, + error=f"skipped: upstream {failed} did not succeed", + ) + remaining.discard(tid) + else: + to_run.append(task) + + if not to_run: + continue # remaining dependents will be re-evaluated (and skipped) + + wave = await asyncio.gather( + *(self._run_task(task, results) for task in to_run) + ) + for tr in wave: + results[tr.task_id] = tr + remaining.discard(tr.task_id) + + return RunState( + plan_id=task_list.plan_id, + business_context_id=business_context_id, + results=results, + open_questions=list(task_list.open_questions), + ) + + async def _run_task(self, task: Task, results: dict[str, TaskResult]) -> TaskResult: + outputs: list[ToolOutput] = [] + for call in task.tool_calls: + resolved = self._resolve_args(call.args, results) + arg_error = self._validate_args(call.tool, resolved) + if arg_error is not None: + outputs.append(ToolOutput(tool=call.tool, kind="error", error=arg_error)) + continue + outputs.append(await self._safe_invoke(call.tool, resolved)) + + status = _label(outputs) + error: str | None = None + if status == "failure": + errs = [o.error for o in outputs if o.kind == "error" and o.error] + error = errs[0] if errs else "all tool calls failed" + return TaskResult( + task_id=task.id, + status=status, + objective=task.objective, + outputs=outputs, + error=error, + ) + + def _resolve_args( + self, args: dict[str, Any], results: dict[str, TaskResult] + ) -> dict[str, Any]: + return {k: self._resolve_value(v, results) for k, v in args.items()} + + @staticmethod + def _resolve_value(value: Any, results: dict[str, TaskResult]) -> Any: + # A data arg is exactly a "${t}" placeholder (Pattern A); resolve it to + # the referenced task's representative output (its last ToolOutput). + # Materializing that envelope into a DataFrame is the invoker's job. + if isinstance(value, str): + match = PLACEHOLDER_RE.fullmatch(value.strip()) + if match: + upstream = results.get(match.group(1)) + if upstream is None or not upstream.outputs: + return None + return upstream.outputs[-1] + return value + + def _validate_args(self, tool: str, resolved: dict[str, Any]) -> str | None: + spec = self._registry.get(tool) + if spec is None: + return f"tool {tool!r} not in registry" + required = spec.input_schema.get("required", []) + missing = [r for r in required if resolved.get(r) is None] + if missing: + return f"missing required arg(s): {sorted(missing)}" + return None + + async def _safe_invoke(self, tool: str, args: dict[str, Any]) -> ToolOutput: + try: + return await self._invoker.invoke(tool, args) + except Exception as exc: # noqa: BLE001 — backstop; the invoker is never-throw (§8.4) + logger.warning("tool invoker raised", tool=tool, error=str(exc)) + return ToolOutput(tool=tool, kind="error", error=f"invoker raised: {exc}") + + +def _label(outputs: list[ToolOutput]) -> TaskStatus: + if not outputs: + return "failure" + errors = sum(1 for o in outputs if o.kind == "error") + if errors == 0: + return "success" + if errors == len(outputs): + return "failure" + return "partial" diff --git a/src/api/v1/chat.py b/src/api/v1/chat.py index 93074b2922e94b339fcda928d8474926f868f433..c1801f844345627efd78759a912ded55e9c79c6e 100644 --- a/src/api/v1/chat.py +++ b/src/api/v1/chat.py @@ -22,6 +22,17 @@ logger = get_logger("chat_api") router = APIRouter(prefix="/api/v1", tags=["Chat"]) +# One shared ChatHandler for the process. It holds no per-request state (user_id +# is passed into handle()), and lazily builds + caches the Orchestrator/Chatbot +# chains — so reusing it keeps the Azure OpenAI clients (and their httpx/TLS pools) +# warm across requests instead of re-handshaking on the first call of every request. +# enable_slow_path is env-gated (ENABLE_SLOW_PATH): when on, structured intents route +# Orchestrator -> Planner -> TaskRunner -> Assembler so the team can test e2e here. +_chat_handler = ChatHandler( + enable_tracing=True, + enable_slow_path=settings.enable_slow_path, +) + _GREETINGS = frozenset(["hi", "hello", "hey", "halo", "hai", "hei"]) _GOODBYES = frozenset(["bye", "goodbye", "thanks", "thank you", "terima kasih", "sampai jumpa"]) @@ -169,7 +180,7 @@ async def chat_stream(request: ChatRequest, db: AsyncSession = Depends(get_db)): return EventSourceResponse(stream_direct()) history = await load_history(db, request.room_id, limit=10) - handler = ChatHandler() + handler = _chat_handler async def stream_response(): 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)): except Exception as e: logger.error("save_messages failed", room_id=request.room_id, error=str(e)) yield event + elif event["event"] == "status": + # slow-path progress ("Planning…", "Running N steps…"): forward + # so the client shows activity and the SSE connection stays alive. + yield event elif event["event"] == "error": yield event return diff --git a/src/catalog/introspect/base.py b/src/catalog/introspect/base.py index ab120e6a7f5eb48cd91bffae39ef0aaf415a554d..34e1aaa1998785a9da82b6ac29cce930946959bb 100644 --- a/src/catalog/introspect/base.py +++ b/src/catalog/introspect/base.py @@ -9,6 +9,11 @@ from abc import ABC, abstractmethod from ..models import Source +# Max sample values stored per column (down from 5 — token cost: sample values +# are fed to the planner prompt). Single source of truth for every introspection +# path (tabular files + DB), so the cap can never drift between them. +SAMPLE_LIMIT = 3 + class BaseIntrospector(ABC): """Abstract base. Subclasses: DatabaseIntrospector, TabularIntrospector.""" diff --git a/src/catalog/introspect/tabular.py b/src/catalog/introspect/tabular.py index 08b205410e01c0119bae2be35ff50c2bf467ee7c..40c913972a2f949a1b521f81564f32a00ab405f6 100644 --- a/src/catalog/introspect/tabular.py +++ b/src/catalog/introspect/tabular.py @@ -26,7 +26,7 @@ from src.middlewares.logging import get_logger from ..models import Column, ColumnStats, DataType, Source, Table from ..pii_detector import PIIDetector -from .base import BaseIntrospector +from .base import SAMPLE_LIMIT, BaseIntrospector logger = get_logger("tabular_introspector") @@ -198,7 +198,7 @@ class TabularIntrospector(BaseIntrospector): (document_id, sheet_name, col_name) if sheet_name else (document_id, col_name) ) - sample_raw = series.dropna().head(3).tolist() + sample_raw = series.dropna().head(SAMPLE_LIMIT).tolist() sample_values: list[Any] | None = [_normalize(v) for v in sample_raw] or None is_numeric = pd.api.types.is_numeric_dtype(series) diff --git a/src/catalog/reader.py b/src/catalog/reader.py index 4e07dbf602def3a05fc61de1f76525e57ed166a4..2c933a62332b752ddac63b065a9e0beb8e26c181 100644 --- a/src/catalog/reader.py +++ b/src/catalog/reader.py @@ -38,3 +38,30 @@ class CatalogReader: filtered = [s for s in catalog.sources if s.source_type == "unstructured"] return catalog.model_copy(update={"sources": filtered}) + + +class MemoizingCatalogReader(CatalogReader): + """Request-scoped CatalogReader that caches each ``read`` by source_hint. + + One per request. The same per-user catalog is otherwise fetched from the + catalog DB 4-5x during a single slow-path run (planner load, then + describe_source's structured+unstructured reads, then query_structured's + structured read). Wrapping the base reader collapses those to one round-trip + per distinct source_hint and pins a single consistent snapshot for the whole + request (plan-time and execution-time catalogs can no longer diverge). + """ + + def __init__(self, inner: CatalogReader) -> None: + # `read` is fully overridden below and delegates to `inner`, so the parent's + # `_store` is never used — carry it through only so this stays a real + # CatalogReader (any inner with a `read` works, including test fakes). + super().__init__(getattr(inner, "_store", None)) + self._inner = inner + self._cache: dict[SourceHint, Catalog] = {} + + async def read(self, user_id: str, source_hint: SourceHint) -> Catalog: + cached = self._cache.get(source_hint) + if cached is None: + cached = await self._inner.read(user_id, source_hint) + self._cache[source_hint] = cached + return cached diff --git a/src/config/prompts/assembler.md b/src/config/prompts/assembler.md new file mode 100644 index 0000000000000000000000000000000000000000..f685231dc50a30a24e2f9d2b3987135db666386a --- /dev/null +++ b/src/config/prompts/assembler.md @@ -0,0 +1,43 @@ +You are the Assembler for Data Eyond, an AI data scientist. A deterministic +TaskRunner has just executed a static analysis plan; you receive its results (the +`RunState`) plus the project's business context. Your job is to turn those results +into a decision-ready answer. + +You produce two things in one structured object: +1. `chat_answer` — a compact, to-the-point reply for the chat, in **markdown** + (prose + tables where useful). +2. The narrative fields of an analysis record: `goal_restated`, `findings`, + `caveats`, `data_used`, `open_questions`. + +# Hard rules (non-negotiable) + +1. **Ground every claim in the provided results.** Use only the numbers, tables, + and values present in the task results. **Never invent, estimate, or extrapolate + a number** that is not in the results. If the data does not answer part of the + question, say so. +2. **Report what failed.** Some tasks may have `status: partial` or `failure`. Do + not pretend they succeeded. Briefly state what could not be completed and how it + limits the answer; put unresolved items in `open_questions`. +3. **Render, don't recompute.** Build markdown tables from the structured task + outputs as they are. Do not do your own arithmetic beyond trivially restating a + value already computed. +4. **No tool/code talk.** Write for a business reader. Do not mention tool names, + task ids, SQL, or internal mechanics in `chat_answer`. + +# How to write + +- **`chat_answer`**: lead with the answer. Add a short markdown table when it makes + the numbers clearer. Keep it tight — this streams into a chat, not a report. +- **`findings`**: the key takeaways, each a single self-contained sentence with the + supporting figure. +- **`caveats`**: data-quality limits, partial/failed steps, assumptions that affect + confidence. +- **`data_used`**: the sources/tables/columns the answer rests on (plain names). +- **`goal_restated`**: one sentence restating the business question you answered. +- **`open_questions`**: anything ambiguous, missing, or worth a follow-up. Fold in + any open questions carried from the plan. Empty list if genuinely none. + +# Output + +Return exactly one structured object with the fields above. Be honest, specific, +and concise. diff --git a/src/config/prompts/planner.md b/src/config/prompts/planner.md new file mode 100644 index 0000000000000000000000000000000000000000..8946de1bbf716ef42d78c444d824de1364e7d1a3 --- /dev/null +++ b/src/config/prompts/planner.md @@ -0,0 +1,56 @@ +You are the Planner for Data Eyond, an AI data scientist that works within the +CRISP-DM lifecycle. Your single job is to turn a business question into a +**static analysis plan**: one `TaskList` that downstream deterministic code +executes exactly as written. + +You plan. You do not execute, and you do not write prose for the user. You emit +only a `TaskList` object that conforms to the provided schema. + +# Hard rules (non-negotiable) + +1. **Emit intent, never code.** Never write SQL, pandas, or any code. The only + query you express is an inline `QueryIR` (a JSON intent object) inside a + `query_structured` tool call's `args.ir`. +2. **The plan is static.** There is no replanning and no execution feedback. Plan + the whole analysis up front; assume each task runs once, in dependency order. +3. **Use only tools from the "Available tools" list.** Never invent a tool name. + Every `tool_calls[].tool` must be one of the listed tool names. +4. **Reference only data that exists.** Every `source_id`, `table_id`, and + `column_id` you put in an inline `QueryIR` must come from the "Catalog" + section. Copy the stable ids verbatim — downstream validation does a literal + id lookup, so a paraphrased name fails. +5. **No modeling in v1.** There are no modeling tools. Do not emit `modeling` + tasks. The product is descriptive/diagnostic only — no predictions, no charts. + +# How to plan + +- **Smallest plan that answers the question.** Do not exceed `Constraints.max_tasks`. +- **A task is an ordered chain of tool calls** with fully-specified arguments. A + simple question is one task with one tool call; a step that needs a follow-up + computation is a short chain or a dependent task. +- **Wire data between tasks with placeholders.** When a task needs an upstream + task's output as an argument, use the string `"${t}"` (e.g. `"${t2}"`) as + the argument value. Set `depends_on` accordingly. +- **Data access vs analytics tools.** `query_structured` is the data-access entry + point: use it to select, filter, and pull rows (and simple built-in + count/sum/avg/min/max/count_distinct the IR can express). For anything richer — + descriptive statistics (median/percentile/mode/std/skew), time trends, group + comparisons, share-of-total, correlation, segmentation, or data-quality + profiling — run `query_structured` to fetch the rows, then pass its output to + the matching composite `analyze_*` tool via a `"${t}"` `data` argument + (referencing the upstream result's column aliases). +- **Mixing structured + unstructured.** If qualitative context helps, add a + `retrieve_documents` task against an unstructured source listed in the catalog. +- **CRISP-DM stages.** Tag each task with the stage it serves: + `data_understanding`, `data_preparation`, or `evaluation`. (Never `modeling`.) +- **success_criteria is a reporting signal**, not a control trigger. State, in + checkable terms (counts, rates, "produced", "above"/"below"), what a good + result looks like. It never causes a retry. +- **Surface uncertainty, don't guess.** If the question is ambiguous or the + catalog can't fully answer it, record it in `open_questions` and plan the best + defensible analysis anyway. Record interpretation choices in `assumptions`. + +# Output + +Return exactly one `TaskList`. The "Examples" section in the human message shows +the required shape. Match it. diff --git a/src/config/settings.py b/src/config/settings.py index 9ff31c93f8a27f5a53e6fe698d617b0dcbcc4554..38a956cec789de70033442abd4a80e0fe1013acd 100644 --- a/src/config/settings.py +++ b/src/config/settings.py @@ -16,6 +16,14 @@ class Settings(BaseSettings): case_sensitive=False, ) + # Feature flags + # Route `structured` chat intents through the analytical SLOW PATH + # (Planner -> TaskRunner -> Assembler) instead of the single-query QueryService. + # Off by default; the team flips ENABLE_SLOW_PATH=true to test end-to-end from + # the /chat/stream endpoint. BusinessContext is still a stub until the lead's + # real source lands, so this stays opt-in. + enable_slow_path: bool = Field(alias="enable_slow_path", default=False) + # Database postgres_connstring: str diff --git a/src/database_client/engine.py b/src/database_client/engine.py new file mode 100644 index 0000000000000000000000000000000000000000..b9b352915b81be45b60177fd0d935b7a7fb7561d --- /dev/null +++ b/src/database_client/engine.py @@ -0,0 +1,168 @@ +"""UserEngineCache — pooled, reused SQLAlchemy engines for users' external DBs. + +The query path (`DbExecutor`) previously built a fresh engine and tore it down on +EVERY query (`db_pipeline_service.engine_scope`), paying a full TCP+TLS+auth +handshake per call (~6-8s measured, dominating slow-path latency). That helper's +connect-once-then-dispose semantics are correct for the *ingestion* pipeline +(infrequent, one connection per run) but wrong for the query path (frequent, +latency-sensitive, repeated to the same DB). + +This module caches one pooled engine per external DB so connections stay warm +across queries. Scope: **postgres / supabase only** (the measured case and the +`schema` source type). Other db_types fall back to the legacy per-call path in +`DbExecutor`, so nothing regresses. + +Safety / multi-tenancy: +- Key = client_id + a hash of the decrypted credentials, so a credential rotation + produces a new key (the stale engine idle-evicts) — a cached engine never serves + rotated creds. +- Read-only + statement_timeout are pinned at connection establishment via libpq + `options` (read-only-at-birth), so they can't be escaped by a reused pooled + connection and cost zero per-query round-trips. +- The caller still re-fetches the DatabaseClient row every query and re-checks + ownership + `active` status — caching the engine never bypasses authorization. +- Bounded LRU + idle TTL cap memory / file descriptors / connections held on the + user's DB. `invalidate(client_id)` disposes eagerly on client update/delete. +""" + +from __future__ import annotations + +import hashlib +import json +import threading +import time +from collections import OrderedDict + +from sqlalchemy import URL, create_engine, event +from sqlalchemy.engine import Engine + +from src.middlewares.logging import get_logger + +logger = get_logger("user_engine_cache") + +_POSTGRES_LIKE = frozenset({"postgres", "supabase"}) +_STATEMENT_TIMEOUT_MS = 30_000 + +# Pool sizing is deliberately small: this is a per-user external DB, often with a +# low max_connections, and we cache many of them. pool_pre_ping drops dead +# connections; pool_recycle bounds connection age so a serverless user DB can still +# autosuspend between bursts. +_POOL_SIZE = 1 +_MAX_OVERFLOW = 2 +_POOL_RECYCLE_SECONDS = 300 + +# Cache bounds across all users. +_MAX_ENGINES = 50 +_IDLE_TTL_SECONDS = 600 + + +def _creds_fingerprint(credentials: dict) -> str: + blob = json.dumps(credentials, sort_keys=True, default=str) + return hashlib.sha256(blob.encode("utf-8")).hexdigest()[:16] + + +class UserEngineCache: + """Process-wide cache of pooled engines for users' external Postgres DBs. + + Thread-safe: `DbExecutor` runs sync DB work in `asyncio.to_thread` worker + threads, so concurrent requests can hit this from multiple threads. + """ + + def __init__(self) -> None: + # key -> (engine, last_used_monotonic) + self._engines: OrderedDict[str, tuple[Engine, float]] = OrderedDict() + self._lock = threading.Lock() + + def get_engine(self, client_id: str, db_type: str, credentials: dict) -> Engine | None: + """Return a pooled engine for (client_id, creds), or None if unsupported. + + None means "not a postgres-like DB" — the caller should use its legacy + per-call path for those (rare, unmeasured) db_types. + """ + if db_type not in _POSTGRES_LIKE: + return None + + key = f"{client_id}:{_creds_fingerprint(credentials)}" + now = time.monotonic() + with self._lock: + self._evict_idle(now) + entry = self._engines.get(key) + if entry is not None: + self._engines[key] = (entry[0], now) + self._engines.move_to_end(key) + return entry[0] + + engine = self._build_engine(credentials) + self._engines[key] = (engine, now) + self._engines.move_to_end(key) + self._evict_overflow() + logger.info("user engine created", client_id=client_id, cached=len(self._engines)) + return engine + + def invalidate(self, client_id: str) -> None: + """Dispose + drop every cached engine for a client (creds rotated/deleted).""" + with self._lock: + stale = [k for k in self._engines if k.startswith(f"{client_id}:")] + for k in stale: + engine, _ = self._engines.pop(k) + engine.dispose() + if stale: + logger.info("user engine invalidated", client_id=client_id, disposed=len(stale)) + + # ------------------------------------------------------------------ + + @staticmethod + def _build_engine(credentials: dict) -> Engine: + # Mirrors db_pipeline_service.connect()'s postgres URL shape, plus a real pool. + query = {"sslmode": credentials["ssl_mode"]} if credentials.get("ssl_mode") else {} + url = URL.create( + drivername="postgresql+psycopg2", + username=credentials["username"], + password=credentials["password"], + host=credentials["host"], + port=credentials["port"], + database=credentials["database"], + query=query, + ) + engine = create_engine( + url, + pool_size=_POOL_SIZE, + max_overflow=_MAX_OVERFLOW, + pool_recycle=_POOL_RECYCLE_SECONDS, + pool_pre_ping=True, + ) + + # Apply read-only + statement_timeout once per PHYSICAL connection via a + # connect event (not per query, so the pooling latency win stays). These are + # ordinary SET commands, NOT libpq startup `options` — Neon's transaction + # pooler rejects `default_transaction_read_only` as a startup parameter but + # accepts it as a SET. Best-effort: the authoritative read-only guarantee is + # the compiler (SELECT-only) + the sqlglot DML guard; statement_timeout is + # backed by the executor's asyncio.wait_for. So a failure here must not break + # the connection. + @event.listens_for(engine, "connect") + def _init_session(dbapi_conn, _record): # noqa: ANN001 + try: + cur = dbapi_conn.cursor() + cur.execute(f"SET statement_timeout = {_STATEMENT_TIMEOUT_MS}") + cur.execute("SET default_transaction_read_only = on") + cur.close() + except Exception as exc: # noqa: BLE001 — best-effort session hardening + logger.warning("session init SET failed", error=str(exc)) + + return engine + + def _evict_idle(self, now: float) -> None: + stale = [k for k, (_, ts) in self._engines.items() if now - ts > _IDLE_TTL_SECONDS] + for k in stale: + engine, _ = self._engines.pop(k) + engine.dispose() + + def _evict_overflow(self) -> None: + while len(self._engines) > _MAX_ENGINES: + _, (engine, _) = self._engines.popitem(last=False) # LRU = oldest end + engine.dispose() + + +# Process-wide singleton consumed by DbExecutor. +user_engine_cache = UserEngineCache() diff --git a/src/observability/langfuse/tracing.py b/src/observability/langfuse/tracing.py new file mode 100644 index 0000000000000000000000000000000000000000..a907a13babfea3bf916e4b346a95380de3bbd1fd --- /dev/null +++ b/src/observability/langfuse/tracing.py @@ -0,0 +1,169 @@ +"""Langfuse request tracing — tokens + latency for the chat pipeline. + +One Langfuse trace per chat request. LangChain LLM calls attach a `CallbackHandler` +(auto-captures prompt/completion tokens + latency); deterministic tool calls are +recorded as metadata-only spans. + +PII policy for Langfuse **Cloud** (data leaves to Langfuse's servers): + - UNMASKED (full input/output): **Orchestrator + Planner** — their inputs are the + user question and a PII-safe `CatalogSummary` (sample values stripped by design). + - MASKED (tokens + latency only; input/output redacted): **Assembler + Chatbot** — + their inputs carry real query rows / document chunks that may contain PII. + - Tool spans carry only metadata (tool name, output kind, row COUNT, status) — + never the rows themselves. + +Everything here is best-effort and **never raises**: if Langfuse is unreachable or +disabled, the chat pipeline runs unchanged. Tracing is created only when the caller +opts in (ChatHandler(enable_tracing=True)); otherwise a `NullTracer` is used. +""" + +from __future__ import annotations + +import contextlib +import functools +import time +from typing import Any + +from src.config.settings import settings +from src.middlewares.logging import get_logger + +logger = get_logger("tracing") + + +def _redact(*, data: Any) -> Any: + """Langfuse MaskFunction: drop the value entirely (used for PII-bearing calls).""" + return "" + + +@functools.cache +def _client() -> Any: + from langfuse import Langfuse + + return Langfuse( + public_key=settings.LANGFUSE_PUBLIC_KEY, + secret_key=settings.LANGFUSE_SECRET_KEY, + host=settings.LANGFUSE_HOST, + ) + + +class _NullSpan: + def end(self, _out: Any) -> None: ... + + +class NullTracer: + """No-op tracer (tracing disabled). Same surface as RequestTracer.""" + + active = False + + def callbacks(self, *, masked: bool = False) -> list: + return [] + + def tool_span(self, tool: str, args: dict) -> Any: + return _NullSpan() + + def end(self, *, output: Any = None) -> None: ... + + +class _ToolSpan: + """A metadata-only span around one tool call. Never records row data.""" + + def __init__(self, trace: Any, tool: str, args: dict) -> None: + self._t0 = time.perf_counter() + self._span = trace.span( + name=f"tool:{tool}", + metadata={"tool": tool, "arg_keys": sorted(args)}, # keys only, no values + ) + + def end(self, out: Any) -> None: + with contextlib.suppress(Exception): # never let a span break the run + kind = getattr(out, "kind", None) + is_err = kind == "error" + meta: dict[str, Any] = { + "kind": kind, + "elapsed_ms": round((time.perf_counter() - self._t0) * 1000), + } + if kind == "table": + meta["rows"] = len(getattr(out, "rows", None) or []) + err_msg = (getattr(out, "error", None) or "")[:300] if is_err else None + if err_msg: + meta["error"] = err_msg + self._span.end( + metadata=meta, + level="ERROR" if is_err else "DEFAULT", + status_message=err_msg, + ) + + +class RequestTracer: + """One Langfuse trace per chat request; hands out callbacks + tool spans.""" + + active = True + + def __init__(self, trace: Any) -> None: + self._trace = trace + + @classmethod + def start( + cls, + *, + user_id: str, + question: str | None = None, + session_id: str | None = None, + ) -> RequestTracer | NullTracer: + try: + trace = _client().trace( + name="chat_request", + user_id=user_id, + session_id=session_id, + input=question, # the user's question (same exposure as Planner prompt) + ) + return cls(trace) + except Exception as e: # never let tracing break the request + logger.warning("tracing disabled (init failed)", error=str(e)) + return NullTracer() + + def callbacks(self, *, masked: bool = False) -> list: + """A LangChain callback nested under this trace. `masked=True` redacts the + call's input/output (tokens + latency are still captured).""" + try: + from langfuse.callback import CallbackHandler + + return [ + CallbackHandler( + stateful_client=self._trace, + mask=_redact if masked else None, + ) + ] + except Exception as e: + logger.warning("tracing handler unavailable", error=str(e)) + return [] + + def tool_span(self, tool: str, args: dict) -> Any: + try: + return _ToolSpan(self._trace, tool, args) + except Exception: + return _NullSpan() + + def end(self, *, output: Any = None) -> None: + # Note: callers pass output=None on PII-bearing paths so no answer text is sent. + with contextlib.suppress(Exception): + if output is not None: + self._trace.update(output=output) + + +class TracingToolInvoker: + """Wraps a ToolInvoker to record a metadata-only span per tool call. + + Implements the ToolInvoker protocol; created at the composition root (ChatHandler) + so the slow-path agent code stays tool-agnostic and tracing-agnostic. + """ + + def __init__(self, inner: Any, tracer: RequestTracer) -> None: + self._inner = inner + self._tracer = tracer + + async def invoke(self, tool_name: str, args: dict[str, Any]) -> Any: + span = self._tracer.tool_span(tool_name, args) + out = await self._inner.invoke(tool_name, args) + span.end(out) + return out diff --git a/src/pipeline/db_pipeline/extractor.py b/src/pipeline/db_pipeline/extractor.py index 4f8140105c78a95ebb2deeb5fdebb75527c60113..a64011b79758ae9057dab366bb771665c5979895 100644 --- a/src/pipeline/db_pipeline/extractor.py +++ b/src/pipeline/db_pipeline/extractor.py @@ -12,12 +12,14 @@ import pandas as pd from sqlalchemy import Date, DateTime, Float, Integer, Numeric, inspect from sqlalchemy.engine import Engine +from src.catalog.introspect.base import SAMPLE_LIMIT from src.middlewares.logging import get_logger logger = get_logger("db_extractor") TOP_VALUES_THRESHOLD = 0.05 # show top values if distinct_ratio <= 5% -SAMPLE_LIMIT = 3 # sample N rows per column (down from 5 — token cost) +# SAMPLE_LIMIT (sample N rows per column) is the shared cap defined in +# catalog.introspect.base — one source of truth for both introspection paths. # Dialects with a single-statement CTE that survives `pd.read_sql`. On these we # fold the stats and sample queries into one round-trip per column. MySQL <8 and diff --git a/src/query/compiler/sql.py b/src/query/compiler/sql.py index 827327e8f0b6a222f174713bfe9ec592d121c47a..707100cb1731e235f41b97ed655f0c1f2b643b13 100644 --- a/src/query/compiler/sql.py +++ b/src/query/compiler/sql.py @@ -30,11 +30,18 @@ from ..ir.models import ( ) from .base import BaseCompiler +# Hard ceiling on rows returned to the agent layer. Every compiled query is +# bounded by this even when the IR sets no limit, so an unbounded SELECT can never +# stream an entire user table over the wire / into memory. The executor caps to +# `row_cap` and flags truncation. +MAX_RESULT_ROWS = 10_000 + @dataclass class CompiledSql: sql: str params: dict[str, Any] = field(default_factory=dict) + row_cap: int = MAX_RESULT_ROWS # executor caps rows to this; flags truncation class SqlCompilerError(Exception): @@ -69,14 +76,14 @@ class SqlCompiler(BaseCompiler): orderby_clause = self._build_orderby( ir.order_by, table, cols_by_id, select_aliases ) - limit_clause = self._build_limit(ir.limit) + limit_clause, row_cap = self._build_limit(ir.limit) parts: list[str] = [select_clause, from_clause] for clause in (where_clause, groupby_clause, orderby_clause, limit_clause): if clause: parts.append(clause) - return CompiledSql(sql=" ".join(parts), params=params) + return CompiledSql(sql=" ".join(parts), params=params, row_cap=row_cap) # ------------------------------------------------------------------ # Catalog lookup @@ -277,10 +284,18 @@ class SqlCompiler(BaseCompiler): parts.append(f"{ref} {ob.dir.upper()}") return "ORDER BY " + ", ".join(parts) - def _build_limit(self, limit: int | None) -> str: + def _build_limit(self, limit: int | None) -> tuple[str, int]: + """Return (LIMIT clause, row_cap). + + Always bounded. An explicit IR limit is honored exactly (capped at + MAX_RESULT_ROWS). When the IR has no limit we still emit + `LIMIT MAX_RESULT_ROWS + 1` — the extra row lets the executor tell + "exactly the cap" from "more rows existed" and flag truncation. + """ if limit is None: - return "" - return f"LIMIT {int(limit)}" + return f"LIMIT {MAX_RESULT_ROWS + 1}", MAX_RESULT_ROWS + row_cap = min(int(limit), MAX_RESULT_ROWS) + return f"LIMIT {row_cap}", row_cap # ------------------------------------------------------------------ # Helpers diff --git a/src/query/executor/db.py b/src/query/executor/db.py index 44e29fa7f272d756d36c8450a8f215b6a76e6357..0d025df9f23437baaace124fc69e5b2599753b12 100644 --- a/src/query/executor/db.py +++ b/src/query/executor/db.py @@ -29,6 +29,7 @@ from sqlalchemy import text from ...catalog.models import Catalog, Source from ...database_client.database_client_service import database_client_service +from ...database_client.engine import user_engine_cache from ...db.postgres.connection import AsyncSessionLocal from ...middlewares.logging import get_logger from ...pipeline.db_pipeline import db_pipeline_service @@ -40,9 +41,7 @@ from .base import BaseExecutor, QueryResult logger = get_logger("db_executor") _QUERY_TIMEOUT_SECONDS = 30 -_ROW_HARD_CAP = 10_000 # belt-and-suspenders cap regardless of LIMIT _DBCLIENT_PREFIX = "dbclient://" -_POSTGRES_LIKE = frozenset({"postgres", "supabase"}) class DbExecutor(BaseExecutor): @@ -86,12 +85,16 @@ class DbExecutor(BaseExecutor): creds = decrypt_credentials_dict(client.credentials) columns, rows = await asyncio.wait_for( - asyncio.to_thread(self._run_sync, client.db_type, creds, compiled), + asyncio.to_thread( + self._run_sync, client_id, client.db_type, creds, compiled + ), timeout=_QUERY_TIMEOUT_SECONDS, ) - truncated = len(rows) > _ROW_HARD_CAP - capped = rows[:_ROW_HARD_CAP] + # The compiler bounded the SQL to `row_cap` (+1 when the IR was + # unbounded). More than row_cap rows means the result was truncated. + truncated = len(rows) > compiled.row_cap + capped = rows[:compiled.row_cap] elapsed_ms = int((time.perf_counter() - started) * 1000) logger.info( "db query complete", @@ -188,16 +191,57 @@ class DbExecutor(BaseExecutor): ) @staticmethod - def _run_sync(db_type: str, creds: dict, compiled: CompiledSql) -> tuple[list[str], list[dict]]: - with db_pipeline_service.engine_scope(db_type, creds) as engine: + def _run_sync( + client_id: str, db_type: str, creds: dict, compiled: CompiledSql + ) -> tuple[list[str], list[dict]]: + engine = user_engine_cache.get_engine(client_id, db_type, creds) + if engine is not None: + # Pooled, reused engine (postgres-like). Read-only + statement_timeout + # are set once per physical connection (connect event in UserEngineCache), + # so no per-query SET round-trips and no dispose — the connection returns + # to the pool warm for the next query. with engine.connect() as conn: - if db_type in _POSTGRES_LIKE: - # session-level read-only + per-statement timeout (ms) - conn.execute(text("SET default_transaction_read_only = on")) - conn.execute( - text(f"SET statement_timeout = {_QUERY_TIMEOUT_SECONDS * 1000}") - ) result = conn.execute(text(compiled.sql), compiled.params) - columns = list(result.keys()) - rows = [dict(row) for row in result.mappings()] - return columns, rows + return list(result.keys()), [dict(row) for row in result.mappings()] + + # Legacy per-call path for non-postgres db_types (connect once, dispose). + # These never set read-only/timeout before, so behavior is unchanged. + with db_pipeline_service.engine_scope(db_type, creds) as eng: + with eng.connect() as conn: + result = conn.execute(text(compiled.sql), compiled.params) + return list(result.keys()), [dict(row) for row in result.mappings()] + + # ------------------------------------------------------------------ + # Speculative pre-connect (DB3) + # ------------------------------------------------------------------ + + @classmethod + async def prewarm(cls, catalog: Catalog, user_id: str) -> None: + """Best-effort: warm pooled engines for the catalog's schema sources. + + Called at slow-path entry so the TCP+TLS+auth handshake overlaps the ~4s + Planner LLM call — by the time `query_structured` runs, the connection is + already established. Warming is an optimization, never a requirement, so + this never raises and per-source failures are swallowed. + """ + for source in catalog.sources: + if source.source_type != "schema": + continue + try: + client_id = cls._parse_client_id(source.location_ref) + client = await cls._fetch_client(client_id) + if client.user_id != user_id: + continue + creds = decrypt_credentials_dict(client.credentials) + await asyncio.to_thread(cls._warm_sync, client_id, client.db_type, creds) + except Exception as exc: # noqa: BLE001 — best-effort warming + logger.info("prewarm skipped", source_id=source.source_id, error=str(exc)) + + @staticmethod + def _warm_sync(client_id: str, db_type: str, creds: dict) -> None: + engine = user_engine_cache.get_engine(client_id, db_type, creds) + if engine is not None: + # Open + return a pooled physical connection: forces the handshake and + # runs the connect-event session SETs, leaving the pool warm. + with engine.connect(): + pass diff --git a/src/tools/__init__.py b/src/tools/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9afe601f1bada0e2a85b4926bc54ad356e67ec62 --- /dev/null +++ b/src/tools/__init__.py @@ -0,0 +1,7 @@ +"""Analytics & utility tools (KM-608). + +Each tool is a deterministic computation (no LLM, no SQL generation) invoked by +the Planner/TaskRunner. The compute layer (calculation logic) lives in per-family +submodules (e.g. `analytics`); the wrapper layer (ToolSpec + ToolOutput + +registry) is added once the Planner seam is settled. +""" diff --git a/src/tools/analytics/__init__.py b/src/tools/analytics/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8f0103aa1ee852e808cb688d96688d06aac718de --- /dev/null +++ b/src/tools/analytics/__init__.py @@ -0,0 +1 @@ +"""Analytics tool family (KM-608).""" diff --git a/src/tools/analytics/aggregation.py b/src/tools/analytics/aggregation.py new file mode 100644 index 0000000000000000000000000000000000000000..ef00ed45f2f84745ebf1f6ac0dc8c205d7a96ce1 --- /dev/null +++ b/src/tools/analytics/aggregation.py @@ -0,0 +1,122 @@ +"""analyze_aggregate — group-by aggregation (KM-608). + +An analytical "family" tool: in ONE call it groups rows by one or more keys +and computes aggregates (sum, mean, count, min, max, median, nunique) per +group. This is the deterministic compute twin of the Planner's +`query_structured` step — the wrapper layer later maps a QueryIR onto this. + +STATUS: compute layer only — the function takes an already-materialized +DataFrame. The wrapper layer (fetching data from the catalog via source_id, +the ToolOutput envelope, ToolSpec registration) is added once the Planner +seam (KM-418) is settled. Keeping compute separate from data-fetching makes +this function easy to unit-test in isolation and stable when wrapped. +""" + +from __future__ import annotations + +import pandas as pd + +from src.tools.analytics.descriptive import ColumnNotFoundError + +# Aggregation functions the tool understands. Whitelisted on purpose so an +# unknown function fails loudly instead of silently doing the wrong thing. +SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median", "nunique") + + +class UnsupportedAggregationError(ValueError): + """Requested aggregation is not in SUPPORTED_AGGS (maps to UNSUPPORTED_AGG).""" + + +def _clean(value: object) -> object: + """Convert numpy/pandas scalars to plain Python so the output is JSON-clean. + + A `group_by` over a datetime column yields `pandas.Timestamp` group keys, + which (like numpy scalars) are not JSON-serializable — normalise those too. + """ + if isinstance(value, pd.Timestamp): + return value.isoformat() + if hasattr(value, "item"): + return value.item() + return value + + +# Prompt-style description read by the Planner to decide WHEN to pick this tool. +# Final destination is ToolSpec.description once the wrapper layer is built. +DESCRIPTION = """\ +Summary: Group-by aggregation. Splits rows by one or more key columns and \ +computes aggregates per group (sum, mean, count, min, max, median, nunique). \ +Returns one row per group. + +USE WHEN the question groups a metric by a category — the tell-tale sign is \ +"per"/"each"/"by" a dimension. Trigger words: "per/each" (per/tiap), "by" \ +(berdasarkan), "breakdown", "total/average ... per ...". + +DON'T USE WHEN: + - it summarizes a column with no grouping -> analyze_descriptive + - it compares two specific groups (A vs B) -> analyze_comparison + - it splits a single total into shares -> analyze_contribution + - the grouping is over time periods -> analyze_trend + +Example questions: + - "total revenue per region" + - "average order value by customer segment" + - "how many distinct products were sold per store?" + - "count of orders for each status" +""" + + +def analyze_aggregate( + df: pd.DataFrame, + aggregations: dict[str, list[str]], + group_by: list[str] | None = None, +) -> list[dict[str, object]]: + """Group-by aggregation over one or many keys. + + Args: + df: already-materialized data (in the real system the wrapper fetches + this from a source_id). + aggregations: which columns to aggregate and how, e.g. + {"revenue": ["sum", "mean"], "order_id": ["count"]}. + group_by: grouping keys. If None/empty, the whole table is aggregated + into a single row. + + Returns: + list[dict]: one row per group. Each row holds the group keys plus a + column per aggregate, named "_" (e.g. "revenue_sum"). + + Raises: + ColumnNotFoundError: if any group_by or aggregated column is absent. + UnsupportedAggregationError: if a requested function is not supported. + """ + group_by = group_by or [] + + # Validate columns first (fail-fast on caller mistakes). + referenced = list(group_by) + list(aggregations.keys()) + missing = [c for c in referenced if c not in df.columns] + if missing: + raise ColumnNotFoundError(f"columns not found: {missing}") + + # Validate aggregation functions. + for col, funcs in aggregations.items(): + bad = [f for f in funcs if f not in SUPPORTED_AGGS] + if bad: + raise UnsupportedAggregationError( + f"unsupported aggregation(s) {bad} for column '{col}'; " + f"supported: {list(SUPPORTED_AGGS)}" + ) + + # Build named aggregations: {"revenue_sum": ("revenue", "sum"), ...} + named = { + f"{col}_{func}": (col, func) + for col, funcs in aggregations.items() + for func in funcs + } + + # No grouping → aggregate the entire table into a single row. + if not group_by: + row = {name: _clean(df[col].agg(func)) for name, (col, func) in named.items()} + return [row] + + # dropna=False keeps groups whose key is null so completeness is honest. + grouped = df.groupby(group_by, dropna=False).agg(**named).reset_index() + return [{k: _clean(v) for k, v in record.items()} for record in grouped.to_dict("records")] diff --git a/src/tools/analytics/comparison.py b/src/tools/analytics/comparison.py new file mode 100644 index 0000000000000000000000000000000000000000..efbc78c269b823a132a9a5950b1b37b85a7e8c73 --- /dev/null +++ b/src/tools/analytics/comparison.py @@ -0,0 +1,137 @@ +"""analyze_comparison — compare a metric across two groups (KM-608). + +An analytical "family" tool: in ONE call it aggregates a value for two groups +of a dimension (e.g. region "A" vs "B", channel "online" vs "store") and +reports the gap between them — absolute difference, percent difference, and +direction. group_a is treated as the baseline. Answers questions like +"how does revenue in region A compare to region B?". + +STATUS: compute layer only — the function takes an already-materialized +DataFrame. The wrapper layer (fetching data from the catalog via source_id, +the ToolOutput envelope, ToolSpec registration) is added once the Planner +seam (KM-418) is settled. Keeping compute separate from data-fetching makes +this function easy to unit-test in isolation and stable when wrapped. +""" + +from __future__ import annotations + +import pandas as pd + +from src.tools.analytics.descriptive import ColumnNotFoundError + +# How to aggregate the value within each group before comparing. +SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median") + + +class UnsupportedAggregationError(ValueError): + """The requested aggregation is not supported (maps to error_code UNSUPPORTED_AGG).""" + + +class GroupNotFoundError(ValueError): + """A requested group value does not occur in the dimension column (maps to GROUP_NOT_FOUND).""" + + + + +# Prompt-style description read by the Planner to decide WHEN to pick this tool. +# Final destination is ToolSpec.description once the wrapper layer is built. +DESCRIPTION = """\ +Summary: Head-to-head comparison of one aggregated metric between TWO specific \ +groups of a dimension (group_a is the baseline). Reports each group's value, \ +the absolute and percent difference, and which side is higher. + +USE WHEN the question pits two named groups against each other. Trigger words: \ +"vs"/"versus", "compare" (bandingkan), "A or B", "difference between" \ +(selisih/beda antara), "higher/lower than". + +SETTING GROUPS: group_a is the BASELINE (the reference). The "comparison" field \ +reads as "group_b is {higher/lower/equal} than group_a", and diff = value_b - \ +value_a. Put the reference/older/expected side in group_a. E.g. "is this year \ +higher than last year" -> group_a=last year, group_b=this year. + +DON'T USE WHEN: + - it aggregates across many groups at once -> analyze_aggregate + - it splits a single total into shares -> analyze_contribution + - it tracks change over time -> analyze_trend + +Example questions: + - "compare revenue between Jakarta and Surabaya" + - "is the average order value higher for members or non-members?" + - "difference in churn between plan A and plan B" + - "male vs female average spend" +""" + + +def analyze_comparison( + df: pd.DataFrame, + dimension: str, + value_column: str, + group_a: object, + group_b: object, + agg: str = "sum", +) -> dict[str, object]: + """Compare one aggregated metric between two groups of a dimension. + + Args: + df: already-materialized data (in the real system the wrapper fetches + this from a source_id). + dimension: the categorical column whose values define the two groups. + value_column: numeric column to aggregate for each group. + group_a: baseline group value (the "from"). + group_b: comparison group value (the "to"). + agg: how to aggregate within each group — one of SUPPORTED_AGGS. + + Returns: + dict with: + dimension, value_column, agg — echo of the chosen settings + group_a, value_a — baseline group + its aggregate + group_b, value_b — comparison group + its aggregate + diff_abs — value_b - value_a + diff_pct — diff_abs / value_a, or None if value_a == 0 + comparison — "higher" | "lower" | "equal" (b relative to a) + + Raises: + ColumnNotFoundError: if dimension or value_column is absent. + UnsupportedAggregationError: if agg is not supported. + GroupNotFoundError: if group_a or group_b has no rows. + """ + missing = [c for c in (dimension, value_column) if c not in df.columns] + if missing: + raise ColumnNotFoundError(f"columns not found: {missing}") + if agg not in SUPPORTED_AGGS: + raise UnsupportedAggregationError( + f"unsupported aggregation '{agg}'; supported: {list(SUPPORTED_AGGS)}" + ) + + rows_a = df.loc[df[dimension] == group_a, value_column] + rows_b = df.loc[df[dimension] == group_b, value_column] + empty = [g for g, rows in ((group_a, rows_a), (group_b, rows_b)) if rows.empty] + if empty: + raise GroupNotFoundError( + f"no rows for group(s) {empty} in column '{dimension}'" + ) + + value_a = float(rows_a.agg(agg)) + value_b = float(rows_b.agg(agg)) + diff_abs = value_b - value_a + diff_pct = (diff_abs / value_a) if value_a != 0 else None + + if diff_abs > 0: + comparison = "higher" + elif diff_abs < 0: + comparison = "lower" + else: + comparison = "equal" + + return { + "dimension": dimension, + "value_column": value_column, + "agg": agg, + "group_a": group_a, + "value_a": value_a, + "group_b": group_b, + "value_b": value_b, + "diff_abs": diff_abs, + "diff_pct": diff_pct, + "comparison": comparison, + } diff --git a/src/tools/analytics/decomposition.py b/src/tools/analytics/decomposition.py new file mode 100644 index 0000000000000000000000000000000000000000..6c787ca03a4ec8966c15a09811c8cd3b56310ede --- /dev/null +++ b/src/tools/analytics/decomposition.py @@ -0,0 +1,135 @@ +"""analyze_contribution — share-of-total per category (KM-608). + +An analytical "family" tool: in ONE call it breaks a total down into each +category's contribution (absolute value, share of total, and running +cumulative share), sorted largest-first. This is a single snapshot — "who +makes up the total right now" — as opposed to analyze_comparison (two groups) +or analyze_trend (movement over time). Answers questions like "which region +drives most of revenue?" and supports Pareto (80/20) reasoning. + +STATUS: compute layer only — the function takes an already-materialized +DataFrame. The wrapper layer (fetching data from the catalog via source_id, +the ToolOutput envelope, ToolSpec registration) is added once the Planner +seam (KM-418) is settled. Keeping compute separate from data-fetching makes +this function easy to unit-test in isolation and stable when wrapped. +""" + +from __future__ import annotations + +import pandas as pd + +from src.tools.analytics.descriptive import ColumnNotFoundError + +# Share-of-total is most meaningful for additive aggregates (sum/count), but +# mean/min/max are allowed; "total" is then the sum of the group aggregates. +SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median") + + +class UnsupportedAggregationError(ValueError): + """The requested aggregation is not supported (maps to error_code UNSUPPORTED_AGG).""" + + +def _clean(value: object) -> object: + """Convert numpy scalars to plain Python so the output is JSON-clean.""" + if hasattr(value, "item"): + return value.item() + return value + + +# Prompt-style description read by the Planner to decide WHEN to pick this tool. +# Final destination is ToolSpec.description once the wrapper layer is built. +DESCRIPTION = """\ +Summary: Breaks a single total into per-category contributions (share of total) \ +in one snapshot, largest first, with cumulative share. Supports Pareto (80/20) \ +reasoning. Optional top_n folds the long tail into "Others". + +USE WHEN the question is about how a whole splits into parts, or which \ +categories dominate. Trigger words: "contribution" (kontribusi), "share" \ +(porsi/proporsi), "% of total" (persen dari total), "Pareto/80-20", "top \ +contributors", "which ... make up most". + +DON'T USE WHEN: + - it pits two specific groups against each other -> analyze_comparison + - it tracks change over time -> analyze_trend + - it just aggregates per group without share-of-total -> analyze_aggregate + +Example questions: + - "which products contribute most to total sales?" + - "what share of revenue comes from each region?" + - "top 5 customers by contribution to profit" + - "do 20% of items make up 80% of revenue?" +""" + + +def analyze_contribution( + df: pd.DataFrame, + dimension: str, + value_column: str, + agg: str = "sum", + top_n: int | None = None, +) -> dict[str, object]: + """Contribution (share of total) of each category, largest first. + + Args: + df: already-materialized data (in the real system the wrapper fetches + this from a source_id). + dimension: categorical column to break the total down by. + value_column: numeric column to aggregate per category. + agg: how to aggregate within each category — one of SUPPORTED_AGGS. + top_n: if set, keep the top N categories and lump the remainder into a + single "Others" row. + + Returns: + dict with: + dimension, value_column, agg — echo of the chosen settings + total — sum of all category aggregates + items — [{"category", "value", "share", + "cumulative_share"}] largest first + + Raises: + ColumnNotFoundError: if dimension or value_column is absent. + UnsupportedAggregationError: if agg is not supported. + """ + missing = [c for c in (dimension, value_column) if c not in df.columns] + if missing: + raise ColumnNotFoundError(f"columns not found: {missing}") + if agg not in SUPPORTED_AGGS: + raise UnsupportedAggregationError( + f"unsupported aggregation '{agg}'; supported: {list(SUPPORTED_AGGS)}" + ) + + grouped = df.groupby(dimension, dropna=False)[value_column].agg(agg) + grouped = grouped.sort_values(ascending=False) + pairs = list(grouped.items()) + + # Optionally collapse the long tail into a single "Others" bucket. + if top_n is not None and len(pairs) > top_n: + head = pairs[:top_n] + others_value = sum(val for _, val in pairs[top_n:]) + head.append(("Others", others_value)) + pairs = head + + total = sum(val for _, val in pairs) + + items = [] + cumulative = 0.0 + for cat, val in pairs: + share = (val / total) if total else None + if share is not None: + cumulative += share + items.append( + { + "category": _clean(cat), + "value": _clean(val), + "share": share, + "cumulative_share": cumulative if total else None, + } + ) + + return { + "dimension": dimension, + "value_column": value_column, + "agg": agg, + "total": _clean(total), + "items": items, + } diff --git a/src/tools/analytics/descriptive.py b/src/tools/analytics/descriptive.py new file mode 100644 index 0000000000000000000000000000000000000000..4940a0564bdc69721ab96f02f922c33529fda1b3 --- /dev/null +++ b/src/tools/analytics/descriptive.py @@ -0,0 +1,154 @@ +"""analyze_descriptive — single/multi-column EDA (KM-608). + +An analytical "family" tool: in ONE call it computes a column's center, +spread, shape, and completeness (mean, median, mode, std, variance, +quartiles, min/max, skew, null_rate). + +STATUS: compute layer only — the function takes an already-materialized +DataFrame. The wrapper layer (fetching data from the catalog via source_id, +the ToolOutput envelope, ToolSpec registration) is added once the Planner +seam (KM-418) is settled. Keeping compute separate from data-fetching makes +this function easy to unit-test in isolation and stable when wrapped. +""" + +from __future__ import annotations + +import pandas as pd + +# Default metrics used when the caller does not narrow them via `metrics`. +DEFAULT_METRICS = ( + "count", + "mean", + "median", + "mode", + "std", + "var", + "q1", + "q3", + "min", + "max", + "skew", + "null_count", + "null_rate", +) + + +class ColumnNotFoundError(ValueError): + """A requested column is absent from the DataFrame (maps to error_code COLUMN_NOT_FOUND).""" + + +def _clean(value: object) -> object: + """Coerce a scalar to a JSON-clean Python value. + + `mode` can be any dtype: an integer column yields `numpy.int64` (NOT + JSON-serializable), a datetime column yields `pandas.Timestamp`. The other + metrics are already wrapped in `float(...)`; mode is the one that needs this. + """ + if value is None: + return None + if isinstance(value, pd.Timestamp): + return value.isoformat() + if hasattr(value, "item"): + return value.item() + return value + + +def _describe_one(series: pd.Series, metrics: tuple[str, ...]) -> dict[str, object]: + """Compute descriptive metrics for a single column. + + Numeric metrics are computed over non-null values. `null_rate` & `count` + are computed over all rows (nulls included) so they reflect completeness + as-is. Undefined cases (e.g. std of a single value) return None — degrade + gracefully instead of raising. + """ + total = len(series) + non_null = series.dropna() + is_numeric = pd.api.types.is_numeric_dtype(series) + + out: dict[str, object] = {} + for m in metrics: + if m == "count": + out["count"] = int(total) + elif m == "null_count": + out["null_count"] = int(series.isna().sum()) + elif m == "null_rate": + out["null_rate"] = float(series.isna().mean()) if total else 0.0 + elif m == "mode": + modes = non_null.mode() + out["mode"] = _clean(modes.iloc[0]) if not modes.empty else None + elif not is_numeric: + out[m] = None + elif m == "mean": + out["mean"] = float(non_null.mean()) if not non_null.empty else None + elif m == "median": + out["median"] = float(non_null.median()) if not non_null.empty else None + elif m == "std": + out["std"] = float(non_null.std()) if non_null.shape[0] > 1 else None + elif m == "var": + out["var"] = float(non_null.var()) if non_null.shape[0] > 1 else None + elif m == "q1": + out["q1"] = float(non_null.quantile(0.25)) if not non_null.empty else None + elif m == "q3": + out["q3"] = float(non_null.quantile(0.75)) if not non_null.empty else None + elif m == "min": + out["min"] = float(non_null.min()) if not non_null.empty else None + elif m == "max": + out["max"] = float(non_null.max()) if not non_null.empty else None + elif m == "skew": + out["skew"] = float(non_null.skew()) if non_null.shape[0] > 2 else None + return out + + +# Prompt-style description read by the Planner to decide WHEN to pick this tool. +# Final destination is ToolSpec.description once the wrapper layer is built. +DESCRIPTION = """\ +Summary: Descriptive statistics (EDA) for one or several columns in a single \ +call — center (mean, median, mode), spread (std, variance, min, max, Q1/Q3 \ +quartiles), distribution shape (skew), and completeness (null count & rate). + +USE WHEN the user asks for an overview, summary, or single-column statistics \ +of ONE or SEVERAL columns as a whole, with NO grouping and NO comparison \ +between groups. Trigger words: "overview/summary" (ringkasan), "average" \ +(rata-rata), "median", "spread/distribution" (sebaran), "how many nulls" \ +(berapa nilai kosong). + +DON'T USE WHEN: + - the question groups by something ("per"/"each"/"by") -> analyze_aggregate + - it compares two specific groups (A vs B) -> analyze_comparison + - it tracks a metric over time -> analyze_trend + - it checks data type, quality, duplicates, outliers, constants -> analyze_profile + +Example questions: + - "what's the average and median customer age?" + - "summarize the income column" + - "how is product price distributed?" + - "how many nulls are in the email column?" +""" + + +def analyze_descriptive( + df: pd.DataFrame, + column_ids: list[str], + metrics: list[str] | None = None, +) -> dict[str, dict[str, object]]: + """Descriptive EDA for one or many columns. + + Args: + df: already-materialized data (in the real system the wrapper fetches + this from a source_id). + column_ids: columns to analyze. + metrics: subset of metrics; defaults to all of DEFAULT_METRICS. + + Returns: + dict: { column_id: { metric: value, ... }, ... } + + Raises: + ColumnNotFoundError: if any column_id is absent from df. + """ + chosen = tuple(metrics) if metrics else DEFAULT_METRICS + + missing = [c for c in column_ids if c not in df.columns] + if missing: + raise ColumnNotFoundError(f"columns not found: {missing}") + + return {col: _describe_one(df[col], chosen) for col in column_ids} diff --git a/src/tools/analytics/quality.py b/src/tools/analytics/quality.py new file mode 100644 index 0000000000000000000000000000000000000000..a930375a8cf224b304bd4402d51551527b86f738 --- /dev/null +++ b/src/tools/analytics/quality.py @@ -0,0 +1,139 @@ +"""analyze_profile — per-column data-quality profile (KM-608). + +An analytical "family" tool: in ONE call it profiles each column's health — +dtype, inferred type, completeness (null count/rate), cardinality (distinct +count/rate, constant flag), and — for numeric columns — min/max/mean plus an +IQR-based outlier count; for non-numeric columns the most frequent value. +Answers "is this data clean enough to analyze?" and surfaces issues (lots of +nulls, a constant column, outliers) before deeper analysis. + +STATUS: compute layer only — the function takes an already-materialized +DataFrame. The wrapper layer (fetching data from the catalog via source_id, +the ToolOutput envelope, ToolSpec registration) is added once the Planner +seam (KM-418) is settled. Keeping compute separate from data-fetching makes +this function easy to unit-test in isolation and stable when wrapped. +""" + +from __future__ import annotations + +import pandas as pd + +from src.tools.analytics.descriptive import ColumnNotFoundError + + +def _clean(value: object) -> object: + """Convert numpy/pandas scalars to plain Python so the output is JSON-clean. + + `top_value` (most frequent value) can be a `pandas.Timestamp` when profiling + a datetime column — neither `Timestamp` nor numpy scalars are JSON-safe. + """ + if isinstance(value, pd.Timestamp): + return value.isoformat() + if hasattr(value, "item"): + return value.item() + return value + + +def _profile_one(series: pd.Series) -> dict[str, object]: + """Build the quality profile for a single column.""" + total = len(series) + non_null = series.dropna() + nn = len(non_null) + distinct = int(series.nunique(dropna=True)) + + is_bool = pd.api.types.is_bool_dtype(series) + is_datetime = pd.api.types.is_datetime64_any_dtype(series) + # bool is technically numeric in pandas; treat it as its own type. + is_numeric = pd.api.types.is_numeric_dtype(series) and not is_bool + + if is_bool: + inferred = "boolean" + elif is_datetime: + inferred = "datetime" + elif is_numeric: + inferred = "numeric" + else: + inferred = "categorical" + + out: dict[str, object] = { + "dtype": str(series.dtype), + "inferred_type": inferred, + "count": int(total), + "null_count": int(series.isna().sum()), + "null_rate": float(series.isna().mean()) if total else 0.0, + "distinct_count": distinct, + "distinct_rate": (distinct / nn) if nn else 0.0, # over non-null values + "is_constant": distinct <= 1, + } + + if is_numeric and nn > 0: + out["min"] = _clean(non_null.min()) + out["max"] = _clean(non_null.max()) + out["mean"] = _clean(non_null.mean()) + # IQR rule: values outside [Q1 - 1.5*IQR, Q3 + 1.5*IQR] are outliers. + # Needs enough points for stable quartiles. + if nn >= 4: + q1 = non_null.quantile(0.25) + q3 = non_null.quantile(0.75) + iqr = q3 - q1 + lower, upper = q1 - 1.5 * iqr, q3 + 1.5 * iqr + out["outlier_count"] = int(((non_null < lower) | (non_null > upper)).sum()) + else: + out["outlier_count"] = None + elif not is_numeric and nn > 0: + counts = non_null.value_counts() + out["top_value"] = _clean(counts.index[0]) + out["top_freq"] = int(counts.iloc[0]) + + return out + + +# Prompt-style description read by the Planner to decide WHEN to pick this tool. +# Final destination is ToolSpec.description once the wrapper layer is built. +DESCRIPTION = """\ +Summary: Per-column data-quality profile. For each column reports dtype, \ +inferred type, completeness (null count/rate), cardinality (distinct count/rate, \ +constant flag), and — for numeric columns — min/max/mean plus an IQR-based \ +outlier count; for non-numeric columns the most frequent value. + +USE WHEN the question is about the HEALTH of the data, not its statistics: \ +missing values, duplicates, data types, outliers, "is this clean enough to \ +analyze". Trigger words: "quality" (kualitas), "missing/nulls" (data kosong), \ +"data type" (tipe data), "duplicates/unique" (duplikat/unik), "outliers". + +DON'T USE WHEN: + - the user wants statistics like mean/median/std/skew -> analyze_descriptive + - it groups or compares -> analyze_aggregate / analyze_comparison + +Example questions: + - "is this dataset clean enough to analyze?" + - "which columns have a lot of missing values?" + - "what are the data types and unique counts per column?" + - "are there outliers in the amount column?" +""" + + +def analyze_profile( + df: pd.DataFrame, + column_ids: list[str] | None = None, +) -> dict[str, dict[str, object]]: + """Per-column data-quality profile. + + Args: + df: already-materialized data (in the real system the wrapper fetches + this from a source_id). + column_ids: columns to profile. If None, every column is profiled. + + Returns: + dict: { column_id: { profile fields, ... }, ... } + + Raises: + ColumnNotFoundError: if any column_id is absent from df. + """ + cols = list(column_ids) if column_ids is not None else list(df.columns) + + missing = [c for c in cols if c not in df.columns] + if missing: + raise ColumnNotFoundError(f"columns not found: {missing}") + + return {col: _profile_one(df[col]) for col in cols} diff --git a/src/tools/analytics/relationship.py b/src/tools/analytics/relationship.py new file mode 100644 index 0000000000000000000000000000000000000000..adb6e8e44259c5d08751c343b3c9d29175a628d5 --- /dev/null +++ b/src/tools/analytics/relationship.py @@ -0,0 +1,132 @@ +"""analyze_correlation — correlation among numeric columns (KM-608). + +An analytical "family" tool: in ONE call it measures how strongly numeric +columns move together. Returns the full correlation matrix plus a list of +column pairs ranked by strength. Answers questions like "does price relate to +units sold?". + +STATUS: compute layer only — the function takes an already-materialized +DataFrame. The wrapper layer (fetching data from the catalog via source_id, +the ToolOutput envelope, ToolSpec registration) is added once the Planner +seam (KM-418) is settled. Keeping compute separate from data-fetching makes +this function easy to unit-test in isolation and stable when wrapped. +""" + +from __future__ import annotations + +import math + +import pandas as pd + +from src.tools.analytics.descriptive import ColumnNotFoundError + +# Correlation methods supported by pandas .corr(). +SUPPORTED_METHODS = ("pearson", "spearman", "kendall") + + +class InvalidMethodError(ValueError): + """The requested method is not supported (maps to error_code INVALID_METHOD).""" + + +class NonNumericColumnError(ValueError): + """A requested column is not numeric (maps to error_code NON_NUMERIC_COLUMN).""" + + +class NotEnoughColumnsError(ValueError): + """Correlation needs at least two numeric columns (maps to NOT_ENOUGH_COLUMNS).""" + + +def _clean(value: object) -> float | None: + """Cast to plain float; NaN (e.g. a zero-variance column) -> None.""" + if value is None: + return None + f = float(value) # type: ignore[arg-type] + return None if math.isnan(f) else f + + +# Prompt-style description read by the Planner to decide WHEN to pick this tool. +# Final destination is ToolSpec.description once the wrapper layer is built. +DESCRIPTION = """\ +Summary: Pairwise correlation across numeric columns (pearson, spearman, or \ +kendall). Returns a correlation matrix plus the strongest pairs ranked by \ +absolute strength. + +USE WHEN the question is about relationship or association between numeric \ +variables. Trigger words: "correlation" (korelasi), "related/relationship" \ +(hubungan/keterkaitan), "does X affect Y", "move together". + +DON'T USE WHEN: + - it implies causation — correlation is not causality; stay descriptive + - it compares two groups of one metric -> analyze_comparison + - it summarizes a single column -> analyze_descriptive + +Example questions: + - "is there a correlation between price and quantity sold?" + - "which variables are most related to revenue?" + - "do age and spending move together?" + - "show the correlation matrix for the numeric columns" +""" + + +def analyze_correlation( + df: pd.DataFrame, + column_ids: list[str] | None = None, + method: str = "pearson", +) -> dict[str, object]: + """Pairwise correlation across numeric columns. + + Args: + df: already-materialized data (in the real system the wrapper fetches + this from a source_id). + column_ids: numeric columns to correlate. If None, every numeric + column in df is used. + method: "pearson" (linear), "spearman" (rank), or "kendall". + + Returns: + dict with: + method — echo of the chosen method + columns — the numeric columns actually correlated + matrix — { col: { col: corr|None } } full square matrix + pairs — [{"a", "b", "corr"}] unique pairs, strongest |corr| first + + Raises: + InvalidMethodError: if method is unknown. + ColumnNotFoundError: if an explicit column is absent. + NonNumericColumnError: if an explicit column is not numeric. + NotEnoughColumnsError: if fewer than two numeric columns remain. + """ + if method not in SUPPORTED_METHODS: + raise InvalidMethodError( + f"unknown method '{method}'; supported: {list(SUPPORTED_METHODS)}" + ) + + if column_ids is None: + cols = [c for c in df.columns if pd.api.types.is_numeric_dtype(df[c])] + else: + missing = [c for c in column_ids if c not in df.columns] + if missing: + raise ColumnNotFoundError(f"columns not found: {missing}") + non_numeric = [ + c for c in column_ids if not pd.api.types.is_numeric_dtype(df[c]) + ] + if non_numeric: + raise NonNumericColumnError(f"columns are not numeric: {non_numeric}") + cols = list(column_ids) + + if len(cols) < 2: + raise NotEnoughColumnsError( + f"need >= 2 numeric columns, got {len(cols)}: {cols}" + ) + + corr = df[cols].corr(method=method) + matrix = {a: {b: _clean(corr.loc[a, b]) for b in cols} for a in cols} + + pairs = [] + for i in range(len(cols)): + for j in range(i + 1, len(cols)): + val = _clean(corr.iloc[i, j]) + if val is not None: + pairs.append({"a": cols[i], "b": cols[j], "corr": val}) + pairs.sort(key=lambda p: abs(p["corr"]), reverse=True) + + return {"method": method, "columns": cols, "matrix": matrix, "pairs": pairs} diff --git a/src/tools/analytics/segmentation.py b/src/tools/analytics/segmentation.py new file mode 100644 index 0000000000000000000000000000000000000000..cf106743a805642caaf38e945f78ee23f47ee59d --- /dev/null +++ b/src/tools/analytics/segmentation.py @@ -0,0 +1,149 @@ +"""analyze_segment — bucket rows into segments (KM-608). + +An analytical "family" tool: in ONE call it bins a numeric column into +segments and reports how rows distribute across them (count, and optionally an +aggregate of another column per segment). Two binning modes: explicit cut +"edges" (e.g. age 0-18-35-60) or equal-frequency "quantile" buckets (quartiles, +deciles). Answers questions like "split customers into age brackets" or "bucket +orders into value tiers". + +STATUS: compute layer only — the function takes an already-materialized +DataFrame. The wrapper layer (fetching data from the catalog via source_id, +the ToolOutput envelope, ToolSpec registration) is added once the Planner +seam (KM-418) is settled. Keeping compute separate from data-fetching makes +this function easy to unit-test in isolation and stable when wrapped. +""" + +from __future__ import annotations + +import math + +import pandas as pd + +from src.tools.analytics.descriptive import ColumnNotFoundError + +# Binning strategies. +SUPPORTED_METHODS = ("edges", "quantile") + +# How to aggregate the value column within each segment. +SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median") + + +class InvalidMethodError(ValueError): + """The requested binning method is not supported (maps to INVALID_METHOD).""" + + +class NonNumericColumnError(ValueError): + """The column to segment on is not numeric (maps to NON_NUMERIC_COLUMN).""" + + +class UnsupportedAggregationError(ValueError): + """The requested aggregation is not supported (maps to UNSUPPORTED_AGG).""" + + +def _clean(value: object) -> object: + """Convert numpy scalars to plain Python; NaN -> None for JSON-clean output.""" + if value is None: + return None + if hasattr(value, "item"): + value = value.item() + if isinstance(value, float) and math.isnan(value): + return None + return value + + +# Prompt-style description read by the Planner to decide WHEN to pick this tool. +# Final destination is ToolSpec.description once the wrapper layer is built. +DESCRIPTION = """\ +Summary: Bins a NUMERIC column into segments and counts how rows distribute \ +across them (optionally aggregating another column per segment). Two modes: \ +explicit cut edges (e.g. age 0-18-35-60) or equal-frequency quantile buckets \ +(quartiles, deciles). + +USE WHEN the question asks to bucket/bracket a continuous number into ranges. \ +Trigger words: "segment" (segmen), "bucket/bracket" (kelompokkan ke rentang), \ +"age groups/tiers" (kelompok umur/tingkatan), "quartiles/deciles", "bins". + +DON'T USE WHEN: + - the category already exists (no binning needed) -> analyze_contribution + - it aggregates by an existing key -> analyze_aggregate + - it compares two named groups -> analyze_comparison + +Example questions: + - "split customers into age brackets 0-18, 18-35, 35-60" + - "bucket orders into value tiers" + - "divide users into spending quartiles" + - "how many customers fall in each income band?" +""" + + +def analyze_segment( + df: pd.DataFrame, + column: str, + bins: list[float] | int, + method: str = "edges", + labels: list[str] | None = None, + value_column: str | None = None, + agg: str = "sum", +) -> dict[str, object]: + """Segment rows by binning a numeric column. + + Args: + df: already-materialized data (in the real system the wrapper fetches + this from a source_id). + column: numeric column to bin on. + bins: for method "edges", the list of cut boundaries (e.g. + [0, 18, 35, 60]); for method "quantile", the number of equal- + frequency buckets (e.g. 4 for quartiles). + method: "edges" (explicit boundaries) or "quantile" (equal frequency). + labels: optional segment names; for "edges" there must be + len(bins) - 1 of them. + value_column: if given, also aggregate this column per segment. + agg: how to aggregate value_column — one of SUPPORTED_AGGS. + + Returns: + dict with: + column, method — echo of the chosen settings + agg — present only when value_column is given + segments — [{"segment", "count", ("value")}], in bin order + + Raises: + ColumnNotFoundError: if column or value_column is absent. + NonNumericColumnError: if column is not numeric. + InvalidMethodError: if method is unknown. + UnsupportedAggregationError: if agg is not supported. + """ + referenced = [column] + ([value_column] if value_column else []) + missing = [c for c in referenced if c not in df.columns] + if missing: + raise ColumnNotFoundError(f"columns not found: {missing}") + if not pd.api.types.is_numeric_dtype(df[column]): + raise NonNumericColumnError(f"column '{column}' is not numeric") + if method not in SUPPORTED_METHODS: + raise InvalidMethodError( + f"unknown method '{method}'; supported: {list(SUPPORTED_METHODS)}" + ) + if value_column is not None and agg not in SUPPORTED_AGGS: + raise UnsupportedAggregationError( + f"unsupported aggregation '{agg}'; supported: {list(SUPPORTED_AGGS)}" + ) + + if method == "edges": + cats = pd.cut(df[column], bins=bins, labels=labels, include_lowest=True) + else: # quantile + cats = pd.qcut(df[column], q=bins, labels=labels, duplicates="drop") + + grouped = df.groupby(cats, observed=False) + counts = grouped.size() + + segments = [] + for seg in counts.index: + row = {"segment": str(seg), "count": int(counts[seg])} + if value_column is not None: + row["value"] = _clean(grouped[value_column].agg(agg).get(seg)) + segments.append(row) + + out: dict[str, object] = {"column": column, "method": method, "segments": segments} + if value_column is not None: + out["agg"] = agg + return out diff --git a/src/tools/analytics/temporal.py b/src/tools/analytics/temporal.py new file mode 100644 index 0000000000000000000000000000000000000000..e36496bf0671b973ffd85809c387b342a01432ba --- /dev/null +++ b/src/tools/analytics/temporal.py @@ -0,0 +1,183 @@ +"""analyze_trend — time-series trend over a period (KM-608). + +An analytical "family" tool: in ONE call it buckets rows into time periods +(day/week/month/quarter/year), aggregates a value per period, and summarizes +the movement (first vs last, absolute & percent change, direction, linear +slope). Answers questions like "how did revenue trend month over month?". + +STATUS: compute layer only — the function takes an already-materialized +DataFrame. The wrapper layer (fetching data from the catalog via source_id, +the ToolOutput envelope, ToolSpec registration) is added once the Planner +seam (KM-418) is settled. Keeping compute separate from data-fetching makes +this function easy to unit-test in isolation and stable when wrapped. +""" + +from __future__ import annotations + +import numpy as np +import pandas as pd + +from src.tools.analytics.descriptive import ColumnNotFoundError + +# Friendly period name -> pandas resample rule. Using the non-deprecated +# pandas 2.2 codes (ME/QE/YE) avoids FutureWarnings. +FREQ_MAP = { + "day": "D", + "week": "W", + "month": "ME", + "quarter": "QE", + "year": "YE", +} + +# How to aggregate the value within each period. +SUPPORTED_AGGS = ("sum", "mean", "count", "min", "max", "median") + + +class InvalidFrequencyError(ValueError): + """The requested period is not in FREQ_MAP (maps to error_code INVALID_FREQUENCY).""" + + +class UnsupportedAggregationError(ValueError): + """The requested aggregation is not supported (maps to error_code UNSUPPORTED_AGG).""" + + +def _clean(value: object) -> object: + """Convert numpy scalars to plain Python; NaN -> None for JSON-clean output.""" + if value is None: + return None + if isinstance(value, float) and np.isnan(value): + return None + if hasattr(value, "item"): + value = value.item() + return None if isinstance(value, float) and np.isnan(value) else value + return value + + +def _period_label(ts: pd.Timestamp, freq: str) -> str: + """Human-readable period label keyed off the friendly frequency name.""" + if freq == "month": + return str(ts.strftime("%Y-%m")) + if freq == "quarter": + return f"{ts.year}-Q{ts.quarter}" + if freq == "year": + return str(ts.strftime("%Y")) + return str(ts.strftime("%Y-%m-%d")) # day / week + + +# Prompt-style description read by the Planner to decide WHEN to pick this tool. +# Final destination is ToolSpec.description once the wrapper layer is built. +DESCRIPTION = """\ +Summary: Time-series trend of one metric over evenly-spaced periods (day, week, \ +month, quarter, year). Reports per-period points plus direction, absolute and \ +percent change, and a linear slope. + +USE WHEN the question is about movement over time — growth, decline, trend, \ +seasonality. Trigger words: "over time" (dari waktu ke waktu), "trend" (tren), \ +"monthly/yearly" (bulanan/tahunan), "growth" (pertumbuhan), "since/last N months". + +DON'T USE WHEN: + - it groups by a non-time category -> analyze_aggregate + - it compares two specific groups (A vs B) -> analyze_comparison + - it summarizes a column with no time axis -> analyze_descriptive + +Example questions: + - "how did monthly revenue change this year?" + - "show the sales trend over the last 12 months" + - "is the number of signups growing quarter over quarter?" + - "yearly profit from 2019 to 2024" +""" + + +def analyze_trend( + df: pd.DataFrame, + date_column: str, + value_column: str, + freq: str = "month", + agg: str = "sum", +) -> dict[str, object]: + """Time-series trend of one value over evenly-spaced periods. + + Args: + df: already-materialized data (in the real system the wrapper fetches + this from a source_id). + date_column: column holding dates/timestamps. + value_column: numeric column to aggregate per period. + freq: period granularity — one of FREQ_MAP keys (default "month"). + agg: how to aggregate within a period — one of SUPPORTED_AGGS. + + Returns: + dict with: + freq, agg — echo of the chosen settings + points — [{"period": str, "value": number|None}, ...] + first, last — value of the first/last non-empty period + change_abs — last - first + change_pct — (last - first) / first, or None if first == 0 + direction — "up" | "down" | "flat" + slope — linear slope across periods, or None if < 2 points + + Raises: + ColumnNotFoundError: if date_column or value_column is absent. + InvalidFrequencyError: if freq is not a known period. + UnsupportedAggregationError: if agg is not supported. + """ + missing = [c for c in (date_column, value_column) if c not in df.columns] + if missing: + raise ColumnNotFoundError(f"columns not found: {missing}") + if freq not in FREQ_MAP: + raise InvalidFrequencyError( + f"unknown frequency '{freq}'; supported: {list(FREQ_MAP)}" + ) + if agg not in SUPPORTED_AGGS: + raise UnsupportedAggregationError( + f"unsupported aggregation '{agg}'; supported: {list(SUPPORTED_AGGS)}" + ) + + # Build a clean datetime-indexed series, then resample into periods. + s = df[[date_column, value_column]].copy() + s[date_column] = pd.to_datetime(s[date_column]) + s = s.dropna(subset=[date_column]).set_index(date_column).sort_index() + resampled = s[value_column].resample(FREQ_MAP[freq]).agg(agg) + + points = [ + {"period": _period_label(ts, freq), "value": _clean(val)} + for ts, val in resampled.items() + ] + + # Summary stats are computed over non-empty periods only. + non_null = resampled.dropna() + first: float | None + last: float | None + change_abs: float | None + change_pct: float | None + slope: float | None + if non_null.empty: + first = last = change_abs = change_pct = slope = None + direction = "flat" + else: + first = float(non_null.iloc[0]) + last = float(non_null.iloc[-1]) + change_abs = last - first + change_pct = (change_abs / first) if first != 0 else None + if change_abs > 0: + direction = "up" + elif change_abs < 0: + direction = "down" + else: + direction = "flat" + if non_null.shape[0] > 1: + x = np.arange(non_null.shape[0]) + slope = float(np.polyfit(x, non_null.to_numpy(dtype=float), 1)[0]) + else: + slope = None + + return { + "freq": freq, + "agg": agg, + "points": points, + "first": first, + "last": last, + "change_abs": change_abs, + "change_pct": change_pct, + "direction": direction, + "slope": slope, + } diff --git a/src/tools/contracts.py b/src/tools/contracts.py new file mode 100644 index 0000000000000000000000000000000000000000..b73e7311964ca336761422247dacf08fadd0a0c4 --- /dev/null +++ b/src/tools/contracts.py @@ -0,0 +1,76 @@ +"""Canonical tool contracts (KM-465 — owned by the tool team). + +These are the single source of truth for the tool <-> agent interface: + + - `ToolSpec` / `ToolRegistry` — the registry contract (§9.2). The concrete v1 + analytics registry instance is built on top of these. + - `ToolOutput` — the tool -> agent output envelope (§8.1). Tools return this at + TaskRunner time; the agent layer plans and degrades against it. + +Ownership note (2026-06-08): the tool team owns these definitions outright; the +agent team's earlier copies in `src/agents/planner/contracts.py` are now thin +re-exports of this module so there is exactly ONE definition across the codebase. +The shapes here are kept byte-for-byte identical to those original stubs so the +already-landed planner / TaskRunner / Assembler keep working unchanged. + +Data-flow decision (KM-465, agreed with the agent team): the analytics tools use +**Pattern A** — `analyze_*` tools do NOT self-fetch by `source_id`; each takes a +`data` argument that is a `"${t}"` placeholder resolved to a DataFrame at +execution time. `analyze_comparison.output_kind` is `"stats"` (a labelled-metric +dict), aligning with the planner registry. + +See AGENT_ARCHITECTURE_CONTEXT_new.md §8.1 / §9.2. +""" + +from __future__ import annotations + +from typing import Any, Literal + +from pydantic import BaseModel, Field + +# --------------------------------------------------------------------------- # +# Tool registry (§9.2) +# --------------------------------------------------------------------------- # + + +class ToolSpec(BaseModel): + name: str + category: str # analytics.query | .aggregation | .timeseries | ... + # JSON-schema-ish dict: {"required": [...], "properties": {arg: {"type": ...}}}. + # VALIDATION CONTRACT — presence only: TaskRunner._validate_args enforces just + # `required` (each must resolve to a non-None arg). The `properties` types are + # DOCUMENTATION for the planner prompt, NOT checked at runtime — a wrong-typed + # arg passes validation and only surfaces (if at all) inside the compute fn. + # Do not assume type-safety here. + input_schema: dict[str, Any] + output_kind: str # the ToolOutput.kind it returns + description: str # prompt-style: what it does, edge cases, what NOT to use it for + phase: Literal["P0", "P1", "P2"] = "P0" + + +class ToolRegistry(BaseModel): + tools: list[ToolSpec] = Field(default_factory=list) + + def names(self) -> set[str]: + return {t.name for t in self.tools} + + def get(self, name: str) -> ToolSpec | None: + for t in self.tools: + if t.name == name: + return t + return None + + +# --------------------------------------------------------------------------- # +# Tool output envelope (§8.1) +# --------------------------------------------------------------------------- # + + +class ToolOutput(BaseModel): + tool: str + kind: Literal["scalar", "table", "stats", "series", "documents", "error"] + value: Any | None = None + columns: list[str] | None = None + rows: list[list[Any]] | None = None + meta: dict[str, Any] = Field(default_factory=dict) + error: str | None = None diff --git a/src/tools/data_access.py b/src/tools/data_access.py new file mode 100644 index 0000000000000000000000000000000000000000..22e0c21ccd468f276ff27e7ef3346ef6e59a95d6 --- /dev/null +++ b/src/tools/data_access.py @@ -0,0 +1,342 @@ +"""DataAccessToolInvoker — the data-access tool family (KM-465 / KM-630). + +Implements the `ToolInvoker` Protocol (src/agents/slow_path/invoker.py) for the +data-access family. Unlike the stateless `AnalyticsToolInvoker`, these tools +read the user's catalog / sources, so the invoker is constructed per-request +with the authenticated `user_id` and its dependencies (dependency injection — +the runtime/Coordinator supplies them; INV-7 keeps the agent layer +tool-agnostic). + +Tools implemented here: +- `list_sources` — the user's data sources (id, name, type, table count). +- `describe_source` — tables/columns of one source (one row per column, + metadata only — exposes `pii_flag`, never sample values). +- `query_structured` — runs a pre-built `QueryIR` (validate -> dispatch -> + execute, skipping the planner) and returns rows as + `ToolOutput(kind="table")` — the Pattern A handoff the + `analyze_*` tools consume. +- `retrieve_documents` — dense retrieval over unstructured sources, returns + `ToolOutput(kind="documents")`. + +Frozen guarantee (§8.4): **never throws.** Any failure returns +`ToolOutput(kind="error", error=...)`. +""" + +from __future__ import annotations + +from collections.abc import Callable +from decimal import Decimal +from typing import Any, Protocol + +from pydantic import ValidationError + +from src.catalog.models import Catalog +from src.catalog.reader import CatalogReader +from src.query.executor.dispatcher import ExecutorDispatcher +from src.query.ir.models import QueryIR +from src.query.ir.validator import IRValidationError, IRValidator +from src.retrieval.base import RetrievalResult +from src.tools.contracts import ToolOutput + +DispatcherFactory = Callable[[Catalog], ExecutorDispatcher] + +# Canonical set of data-access tool names — the single source of truth for which +# tools this invoker serves. `CompositeToolInvoker` imports it to route by name; +# the planner registry should derive its data-access spec names from it (agent -> +# tool is the correct dependency direction). Defining it once here means +# adding/renaming a data-access tool can't silently drift the router out of sync +# from the registry (R11). Must match the names in `DataAccessToolInvoker.invoke`. +DATA_ACCESS_TOOLS: frozenset[str] = frozenset( + {"list_sources", "describe_source", "query_structured", "retrieve_documents"} +) + + +class Retriever(Protocol): + """Minimal interface this invoker needs from the retrieval layer.""" + + async def retrieve( + self, query: str, user_id: str, k: int = 5 + ) -> list[RetrievalResult]: ... + + +class DataAccessToolInvoker: + """Never-throwing invoker for catalog-introspection tools (implements ToolInvoker).""" + + def __init__( + self, + user_id: str, + catalog_reader: CatalogReader, + *, + ir_validator: IRValidator | None = None, + dispatcher_factory: DispatcherFactory | None = None, + document_retriever: Retriever | None = None, + ) -> None: + self._user_id = user_id + self._reader = catalog_reader + # query_structured deps — injectable so tests need no real LLM/DB. The + # validator is stateless; the dispatcher is built per-call from the + # request's catalog (executors are picked by source_type). + self._validator = ir_validator or IRValidator() + self._dispatcher_factory: DispatcherFactory = ( + dispatcher_factory or ExecutorDispatcher + ) + # retrieve_documents dep — the module singleton by default, injectable + # for tests (the real one pulls PGVector + Redis). Lazy-imported on first + # use so importing this module stays cheap. + self._retriever = document_retriever + + async def invoke(self, tool_name: str, args: dict[str, Any]) -> ToolOutput: + try: + if tool_name == "list_sources": + return await self._list_sources() + if tool_name == "describe_source": + return await self._describe_source(args) + if tool_name == "query_structured": + return await self._query_structured(args) + if tool_name == "retrieve_documents": + return await self._retrieve_documents(args) + return ToolOutput( + tool=tool_name, kind="error", error=f"unknown tool {tool_name!r}" + ) + except Exception as exc: # noqa: BLE001 — never-throw seam (§8.4) + return ToolOutput( + tool=tool_name, kind="error", error=f"{type(exc).__name__}: {exc}" + ) + + async def _list_sources(self) -> ToolOutput: + """List the user's data sources (structured + unstructured).""" + structured = await self._reader.read(self._user_id, "structured") + unstructured = await self._reader.read(self._user_id, "unstructured") + sources = list(structured.sources) + list(unstructured.sources) + + rows = [ + [s.source_id, s.name, s.source_type, len(s.tables)] for s in sources + ] + return ToolOutput( + tool="list_sources", + kind="table", + columns=["source_id", "name", "source_type", "table_count"], + rows=rows, + meta={"source_count": len(sources)}, + ) + + async def _describe_source(self, args: dict[str, Any]) -> ToolOutput: + """Describe one source: one row per column across its tables. + + Pattern A note: this is catalog metadata only — never returns row + data or PII sample values (only the `pii_flag` boolean per column). + """ + source_id = args.get("source_id") + if not source_id: + return ToolOutput( + tool="describe_source", + kind="error", + error="missing 'source_id' argument", + ) + + structured = await self._reader.read(self._user_id, "structured") + unstructured = await self._reader.read(self._user_id, "unstructured") + sources = list(structured.sources) + list(unstructured.sources) + + source = next((s for s in sources if s.source_id == source_id), None) + if source is None: + return ToolOutput( + tool="describe_source", + kind="error", + error=f"source {source_id!r} not found", + ) + + rows = [ + [ + t.table_id, + t.name, + t.row_count, + c.column_id, + c.name, + c.data_type, + c.nullable, + c.pii_flag, + ] + for t in source.tables + for c in t.columns + ] + return ToolOutput( + tool="describe_source", + kind="table", + columns=[ + "table_id", + "table_name", + "table_row_count", + "column_id", + "column_name", + "data_type", + "nullable", + "pii_flag", + ], + rows=rows, + meta={ + "source_id": source.source_id, + "source_name": source.name, + "source_type": source.source_type, + "table_count": len(source.tables), + "column_count": len(rows), + }, + ) + + async def _query_structured(self, args: dict[str, Any]) -> ToolOutput: + """Run one validated, single-table QueryIR and return rows as a table. + + This is the spine of the slow path (Pattern A): the `analyze_*` tools + take this output as their `data` arg. We receive an already-built `ir` + from the Planner (never SQL, never an NL question), so we skip the + planner and run validate -> dispatch -> execute directly (the tail of + QueryService.run). Output is `kind="table"` with `columns` + `rows` + (rows are list[list], converted from the executor's list[dict]). + """ + raw = args.get("ir") + if raw is None: + return ToolOutput( + tool="query_structured", kind="error", error="missing 'ir' argument" + ) + + try: + ir = raw if isinstance(raw, QueryIR) else QueryIR.model_validate(raw) + except ValidationError as exc: + return ToolOutput( + tool="query_structured", kind="error", error=f"invalid IR: {exc}" + ) + + catalog = await self._reader.read(self._user_id, "structured") + + try: + self._validator.validate(ir, catalog) + except IRValidationError as exc: + return ToolOutput( + tool="query_structured", + kind="error", + error=f"IR validation failed: {exc}", + ) + + dispatcher = self._dispatcher_factory(catalog) + executor = dispatcher.pick(ir) + result = await executor.run(ir) + + if result.error: + return ToolOutput( + tool="query_structured", kind="error", error=result.error + ) + + # QueryResult.rows is list[dict]; ToolOutput.rows is list[list] ordered + # by `columns` so downstream materialization is positional. DB NUMERIC + # columns arrive as `Decimal` (asyncpg) — coerce to float here so the + # output is JSON-serializable (SSE / analysis_record persistence) and + # plays nicely with the float math in the analyze_* tools. + rows = [ + [_json_safe(row.get(c)) for c in result.columns] for row in result.rows + ] + return ToolOutput( + tool="query_structured", + kind="table", + columns=result.columns, + rows=rows, + meta={ + "source_id": result.source_id, + "source_name": result.source_name, + "table_id": result.table_id, + "table_name": result.table_name, + "backend": result.backend, + "row_count": result.row_count, + "truncated": result.truncated, + "elapsed_ms": result.elapsed_ms, + }, + ) + + async def _retrieve_documents(self, args: dict[str, Any]) -> ToolOutput: + """Dense-retrieve relevant chunks from the user's unstructured sources. + + Pulls qualitative context (PDF/DOCX/TXT) for a natural-language `query` + via the retrieval router. `top_k` caps the number of chunks; optional + `source_id` scopes to one source (best-effort metadata filter — the + router itself does not yet scope by source, so this prunes the results). + + TODO(retrieval scoping): the Planner few-shot has no `retrieve_documents` + example, so `source_id` is rarely emitted today and this post-filter is + adequate. If source-scoped retrieval becomes common, push scoping down + into RetrievalRouter.retrieve()/DocumentRetriever (WHERE + cmetadata->>'source_id' = :source_id) and drop this post-filter — more + correct than pruning an already-top_k'd unscoped result set. + """ + query = args.get("query") + if not isinstance(query, str) or not query.strip(): + return ToolOutput( + tool="retrieve_documents", + kind="error", + error="missing 'query' argument", + ) + + try: + top_k = int(args.get("top_k", 5)) + except (TypeError, ValueError): + top_k = 5 + source_id = args.get("source_id") + + retriever = self._retriever + if retriever is None: + from src.retrieval.router import retrieval_router + + retriever = retrieval_router + + results = await retriever.retrieve(query, self._user_id, top_k) + if source_id: + results = [r for r in results if _result_source_id(r) == source_id] + + documents = [ + { + "content": r.content, + "score": r.score, + "source_type": r.source_type, + "metadata": r.metadata, + } + for r in results + ] + return ToolOutput( + tool="retrieve_documents", + kind="documents", + value=documents, + meta={ + "count": len(documents), + "query": query, + "top_k": top_k, + "source_id": source_id, + }, + ) + + +def _json_safe(value: Any) -> Any: + """Coerce DB scalar types that JSON can't represent into plain Python. + + DB drivers return NUMERIC/DECIMAL as `decimal.Decimal`, which is neither + JSON-serializable nor mixable with `float` math. Convert those to `float`; + everything else passes through unchanged. + """ + if isinstance(value, Decimal): + return float(value) + return value + + +def _result_source_id(result: RetrievalResult) -> str | None: + """Best-effort extraction of a source_id from a retrieval result's metadata. + + The chunk metadata schema is owned by the Go ingestion service; the key may + live at the top level or nested under "data". Returns None if absent. + """ + meta = result.metadata or {} + top = meta.get("source_id") + if isinstance(top, str): + return top + data = meta.get("data") + if isinstance(data, dict): + nested = data.get("source_id") + if isinstance(nested, str): + return nested + return None diff --git a/src/tools/invoker.py b/src/tools/invoker.py new file mode 100644 index 0000000000000000000000000000000000000000..2e1b20e45855966444f80f0aa74653b33c721c4d --- /dev/null +++ b/src/tools/invoker.py @@ -0,0 +1,161 @@ +"""AnalyticsToolInvoker — the runtime seam implementation (KM-465). + +Implements the `ToolInvoker` Protocol the slow-path TaskRunner calls +(src/agents/slow_path/invoker.py). One method, `invoke(tool_name, args)`, does the +whole job for the `analyze_*` family: + +1. Look the tool up in a name -> (compute fn, output_kind) dispatch map; an unknown + name returns an error envelope (never an exception). +2. Materialize the Pattern A `data` argument — which the TaskRunner has already + resolved to the upstream task's `ToolOutput` (kind="table") — into a DataFrame. +3. Call the pure compute function with the remaining args as keyword arguments + (their names match the compute signatures one-to-one). +4. Wrap the result in a `ToolOutput` with the tool's declared `kind`. + +Frozen guarantee (§8.4): **never throws.** Any failure — unknown tool, bad data, +or an exception from compute (e.g. GroupNotFoundError) — comes back as +`ToolOutput(kind="error", error=...)`, so the TaskRunner's degrade-and-continue +keeps working. +""" + +from __future__ import annotations + +import decimal +from collections.abc import Callable +from typing import Any + +import pandas as pd + +from src.tools.analytics import ( + aggregation, + comparison, + decomposition, + descriptive, + quality, + relationship, + segmentation, + temporal, +) +from src.tools.contracts import ToolOutput +from src.tools.data_access import DATA_ACCESS_TOOLS, DataAccessToolInvoker + +# tool name -> (compute callable, ToolOutput.kind it produces). Kept in lockstep +# with src/tools/registry.py output_kind values. +_DISPATCH: dict[str, tuple[Callable[..., Any], str]] = { + "analyze_descriptive": (descriptive.analyze_descriptive, "stats"), + "analyze_aggregate": (aggregation.analyze_aggregate, "table"), + "analyze_comparison": (comparison.analyze_comparison, "stats"), + "analyze_contribution": (decomposition.analyze_contribution, "table"), + "analyze_profile": (quality.analyze_profile, "stats"), + "analyze_correlation": (relationship.analyze_correlation, "stats"), + "analyze_segment": (segmentation.analyze_segment, "table"), + "analyze_trend": (temporal.analyze_trend, "series"), +} + + +class AnalyticsToolInvoker: + """Never-throwing invoker for the `analyze_*` tools (implements ToolInvoker).""" + + async def invoke(self, tool_name: str, args: dict[str, Any]) -> ToolOutput: + entry = _DISPATCH.get(tool_name) + if entry is None: + return ToolOutput( + tool=tool_name, kind="error", error=f"unknown tool {tool_name!r}" + ) + fn, kind = entry + + df, err = _materialize(args.get("data")) + if err is not None: + return ToolOutput(tool=tool_name, kind="error", error=err) + + kwargs = {k: v for k, v in args.items() if k != "data"} + try: + result = fn(df, **kwargs) + except Exception as exc: # noqa: BLE001 — never-throw seam (§8.4) + return ToolOutput( + tool=tool_name, + kind="error", + error=f"{type(exc).__name__}: {exc}", + ) + + return ToolOutput(tool=tool_name, kind=kind, value=result) + + +class CompositeToolInvoker: + """One `invoke()` for the whole tool surface (KM-465 #4). + + The TaskRunner only ever calls one `ToolInvoker`. This composes the two + families behind a single dispatch: the stateless `AnalyticsToolInvoker` + (`analyze_*`) and the per-request stateful `DataAccessToolInvoker` + (catalog/query/retrieval, which need the authenticated `user_id`). Routing + is by tool name; an unknown name falls through to the analytics invoker, + which returns the standard unknown-tool error envelope. + + Constructed per-request — the Coordinator injects the request's `user_id` + and `CatalogReader` into the data-access invoker (INV-7: the agent layer + stays tool-agnostic). + + Frozen guarantee (§8.4): **never throws** — both sub-invokers return + `ToolOutput(kind="error", ...)` on any failure. + """ + + def __init__( + self, + data_access: DataAccessToolInvoker, + analytics: AnalyticsToolInvoker | None = None, + ) -> None: + self._data_access = data_access + self._analytics = analytics or AnalyticsToolInvoker() + + async def invoke(self, tool_name: str, args: dict[str, Any]) -> ToolOutput: + if tool_name in DATA_ACCESS_TOOLS: + return await self._data_access.invoke(tool_name, args) + return await self._analytics.invoke(tool_name, args) + + +def _materialize(data: Any) -> tuple[pd.DataFrame, None] | tuple[None, str]: + """Turn the resolved `data` argument into a DataFrame. + + Accepts the upstream `ToolOutput` (kind="table"), a raw DataFrame, or a + {"columns", "rows"} dict (a serialized table). Returns (df, None) on success + or (None, error_message) on failure — the caller wraps the message. + + Numeric columns are coerced to float (see `_coerce_decimals`): DB NUMERIC + columns arrive as Python `Decimal`, which mixes badly with the float math in + the `analyze_*` compute functions (e.g. `float + Decimal` -> TypeError). + Normalizing here fixes the whole tool family in one place. + """ + if data is None: + return None, "missing 'data' argument (no upstream table to analyze)" + + if isinstance(data, pd.DataFrame): + return _coerce_decimals(data), None + + if isinstance(data, ToolOutput): + if data.kind == "error": + return None, f"upstream data unavailable: {data.error}" + if data.kind != "table" or data.columns is None: + return None, f"cannot materialize 'data' of kind {data.kind!r}" + return _coerce_decimals(pd.DataFrame(data.rows or [], columns=data.columns)), None + + if isinstance(data, dict) and "columns" in data: + df = pd.DataFrame(data.get("rows") or [], columns=data["columns"]) + return _coerce_decimals(df), None + + return None, f"unsupported 'data' type: {type(data).__name__}" + + +def _coerce_decimals(df: pd.DataFrame) -> pd.DataFrame: + """Convert `decimal.Decimal` object-columns to float64 in place. + + DB drivers (asyncpg) return NUMERIC/DECIMAL values as Python `Decimal`, which + land in object-dtype columns. The `analyze_*` compute functions do float math + on these (share-of-total, cumulative sums), and `float + Decimal` raises + TypeError. We only touch columns that actually contain a `Decimal`, so real + string/categorical columns are left untouched. `None`/missing values become + NaN, which the compute functions already handle. + """ + for col in df.columns: + if df[col].dtype == object and df[col].map(lambda v: isinstance(v, decimal.Decimal)).any(): + df[col] = df[col].astype(float) + return df diff --git a/src/tools/registry.py b/src/tools/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..bb7ce1b32f16b2b487292ca250a2412f51f36b8d --- /dev/null +++ b/src/tools/registry.py @@ -0,0 +1,167 @@ +"""Analytics tool registry + +The real registry of the `analyze_*` family, built on the canonical `ToolSpec` +(src/tools/contracts.py) and the prompt-style `DESCRIPTION` constants the Planner +reads to choose a tool (KM-625). This replaces the agent team's local stub in +`src/agents/planner/registry.py` for the analytics slice. + +Conventions (decided with the agent team, KM-465): +- **Pattern A** — `analyze_*` tools do NOT self-fetch by `source_id`. Each takes a + `data` argument that is a `"${t}"` placeholder pointing at an upstream + `query_structured` table output, resolved to a DataFrame at execution time. + Column arguments reference the aliases that upstream query produced. +- `input_schema` is the lightweight JSON-schema-ish dict the planner validator + consumes: `required` (arg names with no default) + `properties` (allowed args). + `required` mirrors each compute function's no-default parameters; value-typing of + placeholder args is deferred to execution time. +- `output_kind` is the `ToolOutput.kind` each tool returns: stats (labelled-metric + dict) | table (rows×cols) | series (ordered periods). + +The four data-access tools (query_structured / retrieve_documents / list_sources / +describe_source) are registered separately once their wrappers land (KM-465 #4); +`default_registry()` composes both slices. +""" + +from __future__ import annotations + +from src.tools.analytics import ( + aggregation, + comparison, + decomposition, + descriptive, + quality, + relationship, + segmentation, + temporal, +) +from src.tools.contracts import ToolRegistry, ToolSpec + +ANALYTICS_TOOLS: list[ToolSpec] = [ + ToolSpec( + name="analyze_descriptive", + category="analytics.descriptive", + input_schema={ + "required": ["data", "column_ids"], + "properties": { + "data": {"type": "string"}, + "column_ids": {"type": "array"}, + "metrics": {"type": "array"}, + }, + }, + output_kind="stats", + description=descriptive.DESCRIPTION, + ), + ToolSpec( + name="analyze_aggregate", + category="analytics.aggregation", + input_schema={ + "required": ["data", "aggregations"], + "properties": { + "data": {"type": "string"}, + "aggregations": {"type": "object"}, + "group_by": {"type": "array"}, + }, + }, + output_kind="table", + description=aggregation.DESCRIPTION, + ), + ToolSpec( + name="analyze_comparison", + category="analytics.comparison", + input_schema={ + "required": ["data", "dimension", "value_column", "group_a", "group_b"], + "properties": { + "data": {"type": "string"}, + "dimension": {"type": "string"}, + "value_column": {"type": "string"}, + "group_a": {}, + "group_b": {}, + "agg": {"type": "string"}, + }, + }, + output_kind="stats", + description=comparison.DESCRIPTION, + ), + ToolSpec( + name="analyze_contribution", + category="analytics.decomposition", + input_schema={ + "required": ["data", "dimension", "value_column"], + "properties": { + "data": {"type": "string"}, + "dimension": {"type": "string"}, + "value_column": {"type": "string"}, + "agg": {"type": "string"}, + "top_n": {"type": "integer"}, + }, + }, + output_kind="table", + description=decomposition.DESCRIPTION, + ), + ToolSpec( + name="analyze_profile", + category="analytics.quality", + input_schema={ + "required": ["data"], + "properties": { + "data": {"type": "string"}, + "column_ids": {"type": "array"}, + }, + }, + output_kind="stats", + description=quality.DESCRIPTION, + ), + ToolSpec( + name="analyze_correlation", + category="analytics.relationship", + input_schema={ + "required": ["data"], + "properties": { + "data": {"type": "string"}, + "column_ids": {"type": "array"}, + "method": {"type": "string"}, + }, + }, + output_kind="stats", + description=relationship.DESCRIPTION, + ), + ToolSpec( + name="analyze_segment", + category="analytics.segmentation", + input_schema={ + "required": ["data", "column", "bins"], + "properties": { + "data": {"type": "string"}, + "column": {"type": "string"}, + "bins": {}, + "method": {"type": "string"}, + "labels": {"type": "array"}, + "value_column": {"type": "string"}, + "agg": {"type": "string"}, + }, + }, + output_kind="table", + description=segmentation.DESCRIPTION, + ), + ToolSpec( + name="analyze_trend", + category="analytics.timeseries", + input_schema={ + "required": ["data", "date_column", "value_column"], + "properties": { + "data": {"type": "string"}, + "date_column": {"type": "string"}, + "value_column": {"type": "string"}, + "freq": {"type": "string"}, + "agg": {"type": "string"}, + }, + }, + output_kind="series", + description=temporal.DESCRIPTION, + ), +] + + +def analytics_registry() -> ToolRegistry: + """The analytics (`analyze_*`) slice of the tool registry (fresh instance).""" + return ToolRegistry(tools=list(ANALYTICS_TOOLS))