Commit ·
2d6406d
1
Parent(s): 09f8451
[NOTICKET] Adopt verb-first skill naming
Browse files- src/agents/planner/examples.py +13 -13
- src/agents/planner/inputs.py +3 -3
- src/agents/planner/registry.py +12 -12
- src/agents/planner/service.py +1 -1
- src/agents/planner/validator.py +5 -5
- src/catalog/reader.py +2 -2
- src/config/prompts/planner.md +4 -4
- src/query/executor/db.py +1 -1
- src/tools/analytics/aggregation.py +1 -1
- src/tools/data_access.py +33 -33
- src/tools/registry.py +3 -3
src/agents/planner/examples.py
CHANGED
|
@@ -2,7 +2,7 @@
|
|
| 2 |
|
| 3 |
Two illustrative (question -> TaskList) pairs that teach the OUTPUT SHAPE:
|
| 4 |
stages, dependency edges, ordered tool-call chains, inline QueryIR,
|
| 5 |
-
"${t<id>}" placeholders, and the assumed data-flow convention — `
|
| 6 |
pulls rows, then a composite `analyze_*` tool consumes them via a `data` placeholder
|
| 7 |
referencing the upstream result's column aliases (Pattern A; the tool team may
|
| 8 |
instead pick self-fetch by `source_id`, in which case these examples are reshaped
|
|
@@ -21,7 +21,7 @@ from .schemas import Task, TaskList, ToolCall
|
|
| 21 |
# --------------------------------------------------------------------------- #
|
| 22 |
# Example A — exploratory, no modeling.
|
| 23 |
# "Which product categories drove last quarter's revenue?"
|
| 24 |
-
# Shows:
|
| 25 |
# category in one call (no manual per-category queries).
|
| 26 |
# --------------------------------------------------------------------------- #
|
| 27 |
|
|
@@ -35,7 +35,7 @@ _EXAMPLE_A = TaskList(
|
|
| 35 |
id="t1",
|
| 36 |
stage="data_understanding",
|
| 37 |
objective="Confirm the sales source exposes category, revenue, and order date.",
|
| 38 |
-
tool_calls=[ToolCall(tool="
|
| 39 |
expected_output="source_shape",
|
| 40 |
success_criteria="Produced the orders table schema; the 3 needed columns are present.",
|
| 41 |
depends_on=[],
|
|
@@ -47,7 +47,7 @@ _EXAMPLE_A = TaskList(
|
|
| 47 |
objective="Pull last quarter's order-level category and revenue rows.",
|
| 48 |
tool_calls=[
|
| 49 |
ToolCall(
|
| 50 |
-
tool="
|
| 51 |
args={
|
| 52 |
"ir": {
|
| 53 |
"source_id": "src_sales",
|
|
@@ -111,7 +111,7 @@ _EXAMPLE_B = TaskList(
|
|
| 111 |
id="t1",
|
| 112 |
stage="data_understanding",
|
| 113 |
objective="Confirm the sales source exposes order date, revenue, and region.",
|
| 114 |
-
tool_calls=[ToolCall(tool="
|
| 115 |
expected_output="source_shape",
|
| 116 |
success_criteria="Produced the orders table schema; the needed columns are present.",
|
| 117 |
depends_on=[],
|
|
@@ -123,7 +123,7 @@ _EXAMPLE_B = TaskList(
|
|
| 123 |
objective="Pull this year's order dates, revenue, and region.",
|
| 124 |
tool_calls=[
|
| 125 |
ToolCall(
|
| 126 |
-
tool="
|
| 127 |
args={
|
| 128 |
"ir": {
|
| 129 |
"source_id": "src_sales",
|
|
@@ -187,8 +187,8 @@ _EXAMPLE_B = TaskList(
|
|
| 187 |
# Example C — mixed structured + unstructured.
|
| 188 |
# "Revenue dipped in Q1 — what happened?"
|
| 189 |
# Shows: a structured branch (query -> analyze_trend) runs alongside an
|
| 190 |
-
# INDEPENDENT
|
| 191 |
-
#
|
| 192 |
# placeholder — it is a source, not a consumer) and can run in parallel; the
|
| 193 |
# Assembler folds the document context into the explanation.
|
| 194 |
# --------------------------------------------------------------------------- #
|
|
@@ -203,7 +203,7 @@ _EXAMPLE_C = TaskList(
|
|
| 203 |
id="t1",
|
| 204 |
stage="data_understanding",
|
| 205 |
objective="Confirm the sales source exposes order date and revenue.",
|
| 206 |
-
tool_calls=[ToolCall(tool="
|
| 207 |
expected_output="source_shape",
|
| 208 |
success_criteria="Produced the orders table schema; date and revenue columns present.",
|
| 209 |
depends_on=[],
|
|
@@ -215,7 +215,7 @@ _EXAMPLE_C = TaskList(
|
|
| 215 |
objective="Pull Q1 order dates and revenue.",
|
| 216 |
tool_calls=[
|
| 217 |
ToolCall(
|
| 218 |
-
tool="
|
| 219 |
args={
|
| 220 |
"ir": {
|
| 221 |
"source_id": "src_sales",
|
|
@@ -273,7 +273,7 @@ _EXAMPLE_C = TaskList(
|
|
| 273 |
objective="Retrieve qualitative context on Q1 operational events behind a dip.",
|
| 274 |
tool_calls=[
|
| 275 |
ToolCall(
|
| 276 |
-
tool="
|
| 277 |
args={
|
| 278 |
"query": "operational issues, outages, or notable events in Q1 2026",
|
| 279 |
"top_k": 5,
|
|
@@ -308,7 +308,7 @@ _EXAMPLE_D = TaskList(
|
|
| 308 |
id="t1",
|
| 309 |
stage="data_understanding",
|
| 310 |
objective="Confirm the sales source exposes region and revenue.",
|
| 311 |
-
tool_calls=[ToolCall(tool="
|
| 312 |
expected_output="source_shape",
|
| 313 |
success_criteria="Produced the orders table schema; region and revenue present.",
|
| 314 |
depends_on=[],
|
|
@@ -320,7 +320,7 @@ _EXAMPLE_D = TaskList(
|
|
| 320 |
objective="Pull order-level region and revenue.",
|
| 321 |
tool_calls=[
|
| 322 |
ToolCall(
|
| 323 |
-
tool="
|
| 324 |
args={
|
| 325 |
"ir": {
|
| 326 |
"source_id": "src_sales",
|
|
|
|
| 2 |
|
| 3 |
Two illustrative (question -> TaskList) pairs that teach the OUTPUT SHAPE:
|
| 4 |
stages, dependency edges, ordered tool-call chains, inline QueryIR,
|
| 5 |
+
"${t<id>}" placeholders, and the assumed data-flow convention — `retrieve_data`
|
| 6 |
pulls rows, then a composite `analyze_*` tool consumes them via a `data` placeholder
|
| 7 |
referencing the upstream result's column aliases (Pattern A; the tool team may
|
| 8 |
instead pick self-fetch by `source_id`, in which case these examples are reshaped
|
|
|
|
| 21 |
# --------------------------------------------------------------------------- #
|
| 22 |
# Example A — exploratory, no modeling.
|
| 23 |
# "Which product categories drove last quarter's revenue?"
|
| 24 |
+
# Shows: retrieve_data pulls rows -> analyze_aggregate sums revenue per
|
| 25 |
# category in one call (no manual per-category queries).
|
| 26 |
# --------------------------------------------------------------------------- #
|
| 27 |
|
|
|
|
| 35 |
id="t1",
|
| 36 |
stage="data_understanding",
|
| 37 |
objective="Confirm the sales source exposes category, revenue, and order date.",
|
| 38 |
+
tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
|
| 39 |
expected_output="source_shape",
|
| 40 |
success_criteria="Produced the orders table schema; the 3 needed columns are present.",
|
| 41 |
depends_on=[],
|
|
|
|
| 47 |
objective="Pull last quarter's order-level category and revenue rows.",
|
| 48 |
tool_calls=[
|
| 49 |
ToolCall(
|
| 50 |
+
tool="retrieve_data",
|
| 51 |
args={
|
| 52 |
"ir": {
|
| 53 |
"source_id": "src_sales",
|
|
|
|
| 111 |
id="t1",
|
| 112 |
stage="data_understanding",
|
| 113 |
objective="Confirm the sales source exposes order date, revenue, and region.",
|
| 114 |
+
tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
|
| 115 |
expected_output="source_shape",
|
| 116 |
success_criteria="Produced the orders table schema; the needed columns are present.",
|
| 117 |
depends_on=[],
|
|
|
|
| 123 |
objective="Pull this year's order dates, revenue, and region.",
|
| 124 |
tool_calls=[
|
| 125 |
ToolCall(
|
| 126 |
+
tool="retrieve_data",
|
| 127 |
args={
|
| 128 |
"ir": {
|
| 129 |
"source_id": "src_sales",
|
|
|
|
| 187 |
# Example C — mixed structured + unstructured.
|
| 188 |
# "Revenue dipped in Q1 — what happened?"
|
| 189 |
# Shows: a structured branch (query -> analyze_trend) runs alongside an
|
| 190 |
+
# INDEPENDENT retrieve_knowledge branch that pulls qualitative context. Note
|
| 191 |
+
# retrieve_knowledge takes a natural-language `query` (NOT a `${t<id>}` data
|
| 192 |
# placeholder — it is a source, not a consumer) and can run in parallel; the
|
| 193 |
# Assembler folds the document context into the explanation.
|
| 194 |
# --------------------------------------------------------------------------- #
|
|
|
|
| 203 |
id="t1",
|
| 204 |
stage="data_understanding",
|
| 205 |
objective="Confirm the sales source exposes order date and revenue.",
|
| 206 |
+
tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
|
| 207 |
expected_output="source_shape",
|
| 208 |
success_criteria="Produced the orders table schema; date and revenue columns present.",
|
| 209 |
depends_on=[],
|
|
|
|
| 215 |
objective="Pull Q1 order dates and revenue.",
|
| 216 |
tool_calls=[
|
| 217 |
ToolCall(
|
| 218 |
+
tool="retrieve_data",
|
| 219 |
args={
|
| 220 |
"ir": {
|
| 221 |
"source_id": "src_sales",
|
|
|
|
| 273 |
objective="Retrieve qualitative context on Q1 operational events behind a dip.",
|
| 274 |
tool_calls=[
|
| 275 |
ToolCall(
|
| 276 |
+
tool="retrieve_knowledge",
|
| 277 |
args={
|
| 278 |
"query": "operational issues, outages, or notable events in Q1 2026",
|
| 279 |
"top_k": 5,
|
|
|
|
| 308 |
id="t1",
|
| 309 |
stage="data_understanding",
|
| 310 |
objective="Confirm the sales source exposes region and revenue.",
|
| 311 |
+
tool_calls=[ToolCall(tool="check_data", args={"source_id": "src_sales"})],
|
| 312 |
expected_output="source_shape",
|
| 313 |
success_criteria="Produced the orders table schema; region and revenue present.",
|
| 314 |
depends_on=[],
|
|
|
|
| 320 |
objective="Pull order-level region and revenue.",
|
| 321 |
tool_calls=[
|
| 322 |
ToolCall(
|
| 323 |
+
tool="retrieve_data",
|
| 324 |
args={
|
| 325 |
"ir": {
|
| 326 |
"source_id": "src_sales",
|
src/agents/planner/inputs.py
CHANGED
|
@@ -4,9 +4,9 @@
|
|
| 4 |
for the planner prompt. It carries every table + column id/type/PII flag + row
|
| 5 |
counts + low-cardinality top_values, with `sample_values` nulled on PII columns
|
| 6 |
(INV: no PII sample values into the prompt, see doc §13). It also lists the
|
| 7 |
-
available unstructured sources so the planner can plan `
|
| 8 |
|
| 9 |
-
The planner *validator* still checks inline `
|
| 10 |
full `Catalog` via the existing IRValidator — the summary is a prompt input, not
|
| 11 |
the validation source of truth.
|
| 12 |
|
|
@@ -124,7 +124,7 @@ class CatalogSummary(BaseModel):
|
|
| 124 |
lines.append("")
|
| 125 |
|
| 126 |
if self.unstructured_sources:
|
| 127 |
-
lines.append("Unstructured sources (for
|
| 128 |
for src in self.unstructured_sources:
|
| 129 |
lines.append(f" - {src.name} — id={src.source_id}")
|
| 130 |
|
|
|
|
| 4 |
for the planner prompt. It carries every table + column id/type/PII flag + row
|
| 5 |
counts + low-cardinality top_values, with `sample_values` nulled on PII columns
|
| 6 |
(INV: no PII sample values into the prompt, see doc §13). It also lists the
|
| 7 |
+
available unstructured sources so the planner can plan `retrieve_knowledge`.
|
| 8 |
|
| 9 |
+
The planner *validator* still checks inline `retrieve_data` IRs against the
|
| 10 |
full `Catalog` via the existing IRValidator — the summary is a prompt input, not
|
| 11 |
the validation source of truth.
|
| 12 |
|
|
|
|
| 124 |
lines.append("")
|
| 125 |
|
| 126 |
if self.unstructured_sources:
|
| 127 |
+
lines.append("Unstructured sources (for retrieve_knowledge):")
|
| 128 |
for src in self.unstructured_sources:
|
| 129 |
lines.append(f" - {src.name} — id={src.source_id}")
|
| 130 |
|
src/agents/planner/registry.py
CHANGED
|
@@ -7,8 +7,8 @@ outside it).
|
|
| 7 |
`src/tools/registry.py::analytics_registry()` (KM-628), built on the canonical
|
| 8 |
`ToolSpec` (`src/tools/contracts.py`, KM-465/KM-627) and the prompt-style tool
|
| 9 |
descriptions (KM-625). No longer a stub on our side — it tracks the real registry.
|
| 10 |
-
- **Data access (`
|
| 11 |
-
`
|
| 12 |
but their wrappers + `ToolSpec`s haven't landed yet (KM-465 #4). We keep best-guess
|
| 13 |
spec bodies here so the Planner can plan end-to-end — but the NAMES derive from
|
| 14 |
`src.tools.data_access.DATA_ACCESS_TOOLS` (R11), so a tool rename/addition upstream
|
|
@@ -16,10 +16,10 @@ outside it).
|
|
| 16 |
this slice and swap `default_registry()` for the tool team's full composition.
|
| 17 |
|
| 18 |
**Confirmed conventions (KM-465):** Pattern A — `analyze_*` tools take a `data`
|
| 19 |
-
`"${t<id>}"` placeholder pointing at an upstream `
|
| 20 |
self-fetch); resolved to a DataFrame at execution time. `input_schema` is the
|
| 21 |
lightweight `{required, properties}` dict the planner validator (check #8) reads;
|
| 22 |
-
`
|
| 23 |
catalog by the existing IRValidator.
|
| 24 |
|
| 25 |
See AGENT_ARCHITECTURE_CONTEXT_new.md §9.2 / §9.3.
|
|
@@ -38,7 +38,7 @@ from .contracts import ToolRegistry, ToolSpec
|
|
| 38 |
# --------------------------------------------------------------------------- #
|
| 39 |
_DATA_ACCESS_SPEC_BODIES: tuple[ToolSpec, ...] = (
|
| 40 |
ToolSpec(
|
| 41 |
-
name="
|
| 42 |
category="analytics.query",
|
| 43 |
input_schema={"required": ["ir"], "properties": {"ir": {"type": "object"}}},
|
| 44 |
output_kind="table",
|
|
@@ -52,11 +52,11 @@ _DATA_ACCESS_SPEC_BODIES: tuple[ToolSpec, ...] = (
|
|
| 52 |
"(count/sum/avg/min/max/count_distinct). Do NOT use it for richer statistics "
|
| 53 |
"(median/percentile/mode/stddev/skew → analyze_descriptive), trends "
|
| 54 |
"(analyze_trend), correlation, segmentation, or share-of-total; and do NOT "
|
| 55 |
-
"use it to read documents (use
|
| 56 |
),
|
| 57 |
),
|
| 58 |
ToolSpec(
|
| 59 |
-
name="
|
| 60 |
category="retrieval.documents",
|
| 61 |
input_schema={
|
| 62 |
"required": ["query"],
|
|
@@ -71,11 +71,11 @@ _DATA_ACCESS_SPEC_BODIES: tuple[ToolSpec, ...] = (
|
|
| 71 |
"Dense-retrieve the most relevant chunks from the user's unstructured "
|
| 72 |
"sources (PDF/DOCX/TXT) for a natural-language `query`. Use this to pull "
|
| 73 |
"qualitative context into an analysis. Optionally scope to one `source_id`. "
|
| 74 |
-
"Do NOT use it for numbers in tables — that is
|
| 75 |
),
|
| 76 |
),
|
| 77 |
ToolSpec(
|
| 78 |
-
name="
|
| 79 |
category="catalog.introspection",
|
| 80 |
input_schema={
|
| 81 |
"required": [],
|
|
@@ -88,11 +88,11 @@ _DATA_ACCESS_SPEC_BODIES: tuple[ToolSpec, ...] = (
|
|
| 88 |
"data_understanding to discover what exists. With a `source_id`, returns that "
|
| 89 |
"source's tables and columns (names, types, row counts) — use to confirm a "
|
| 90 |
"source's shape before querying it. Cheap. Do NOT use it to fetch data rows "
|
| 91 |
-
"(use
|
| 92 |
),
|
| 93 |
),
|
| 94 |
ToolSpec(
|
| 95 |
-
name="
|
| 96 |
category="catalog.introspection",
|
| 97 |
input_schema={"required": [], "properties": {}},
|
| 98 |
output_kind="table",
|
|
@@ -100,7 +100,7 @@ _DATA_ACCESS_SPEC_BODIES: tuple[ToolSpec, ...] = (
|
|
| 100 |
"List the user's unstructured sources / documents (id, name, type). Use in "
|
| 101 |
"data_understanding to discover what qualitative material exists before "
|
| 102 |
"retrieving from it. Do NOT use it to read document content (use "
|
| 103 |
-
"
|
| 104 |
),
|
| 105 |
),
|
| 106 |
)
|
|
|
|
| 7 |
`src/tools/registry.py::analytics_registry()` (KM-628), built on the canonical
|
| 8 |
`ToolSpec` (`src/tools/contracts.py`, KM-465/KM-627) and the prompt-style tool
|
| 9 |
descriptions (KM-625). No longer a stub on our side — it tracks the real registry.
|
| 10 |
+
- **Data access (`retrieve_data` / `retrieve_knowledge` / `check_data` /
|
| 11 |
+
`check_knowledge`) — spec BODIES still a local stub.** The tool team owns these too,
|
| 12 |
but their wrappers + `ToolSpec`s haven't landed yet (KM-465 #4). We keep best-guess
|
| 13 |
spec bodies here so the Planner can plan end-to-end — but the NAMES derive from
|
| 14 |
`src.tools.data_access.DATA_ACCESS_TOOLS` (R11), so a tool rename/addition upstream
|
|
|
|
| 16 |
this slice and swap `default_registry()` for the tool team's full composition.
|
| 17 |
|
| 18 |
**Confirmed conventions (KM-465):** Pattern A — `analyze_*` tools take a `data`
|
| 19 |
+
`"${t<id>}"` placeholder pointing at an upstream `retrieve_data` output (no
|
| 20 |
self-fetch); resolved to a DataFrame at execution time. `input_schema` is the
|
| 21 |
lightweight `{required, properties}` dict the planner validator (check #8) reads;
|
| 22 |
+
`retrieve_data.args["ir"]` carries an inline QueryIR validated against the
|
| 23 |
catalog by the existing IRValidator.
|
| 24 |
|
| 25 |
See AGENT_ARCHITECTURE_CONTEXT_new.md §9.2 / §9.3.
|
|
|
|
| 38 |
# --------------------------------------------------------------------------- #
|
| 39 |
_DATA_ACCESS_SPEC_BODIES: tuple[ToolSpec, ...] = (
|
| 40 |
ToolSpec(
|
| 41 |
+
name="retrieve_data",
|
| 42 |
category="analytics.query",
|
| 43 |
input_schema={"required": ["ir"], "properties": {"ir": {"type": "object"}}},
|
| 44 |
output_kind="table",
|
|
|
|
| 52 |
"(count/sum/avg/min/max/count_distinct). Do NOT use it for richer statistics "
|
| 53 |
"(median/percentile/mode/stddev/skew → analyze_descriptive), trends "
|
| 54 |
"(analyze_trend), correlation, segmentation, or share-of-total; and do NOT "
|
| 55 |
+
"use it to read documents (use retrieve_knowledge)."
|
| 56 |
),
|
| 57 |
),
|
| 58 |
ToolSpec(
|
| 59 |
+
name="retrieve_knowledge",
|
| 60 |
category="retrieval.documents",
|
| 61 |
input_schema={
|
| 62 |
"required": ["query"],
|
|
|
|
| 71 |
"Dense-retrieve the most relevant chunks from the user's unstructured "
|
| 72 |
"sources (PDF/DOCX/TXT) for a natural-language `query`. Use this to pull "
|
| 73 |
"qualitative context into an analysis. Optionally scope to one `source_id`. "
|
| 74 |
+
"Do NOT use it for numbers in tables — that is retrieve_data's job."
|
| 75 |
),
|
| 76 |
),
|
| 77 |
ToolSpec(
|
| 78 |
+
name="check_data",
|
| 79 |
category="catalog.introspection",
|
| 80 |
input_schema={
|
| 81 |
"required": [],
|
|
|
|
| 88 |
"data_understanding to discover what exists. With a `source_id`, returns that "
|
| 89 |
"source's tables and columns (names, types, row counts) — use to confirm a "
|
| 90 |
"source's shape before querying it. Cheap. Do NOT use it to fetch data rows "
|
| 91 |
+
"(use retrieve_data) or to inspect documents (use check_knowledge)."
|
| 92 |
),
|
| 93 |
),
|
| 94 |
ToolSpec(
|
| 95 |
+
name="check_knowledge",
|
| 96 |
category="catalog.introspection",
|
| 97 |
input_schema={"required": [], "properties": {}},
|
| 98 |
output_kind="table",
|
|
|
|
| 100 |
"List the user's unstructured sources / documents (id, name, type). Use in "
|
| 101 |
"data_understanding to discover what qualitative material exists before "
|
| 102 |
"retrieving from it. Do NOT use it to read document content (use "
|
| 103 |
+
"retrieve_knowledge) or to inspect structured data (use check_data)."
|
| 104 |
),
|
| 105 |
),
|
| 106 |
)
|
src/agents/planner/service.py
CHANGED
|
@@ -9,7 +9,7 @@ static plan.
|
|
| 9 |
|
| 10 |
The service takes the full `Catalog` (not just a `CatalogSummary`): it derives
|
| 11 |
the PII-safe `CatalogSummary` for the prompt, but validation needs the full
|
| 12 |
-
catalog so the existing `IRValidator` can check inline `
|
| 13 |
|
| 14 |
See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3.
|
| 15 |
"""
|
|
|
|
| 9 |
|
| 10 |
The service takes the full `Catalog` (not just a `CatalogSummary`): it derives
|
| 11 |
the PII-safe `CatalogSummary` for the prompt, but validation needs the full
|
| 12 |
+
catalog so the existing `IRValidator` can check inline `retrieve_data` IRs.
|
| 13 |
|
| 14 |
See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3.
|
| 15 |
"""
|
src/agents/planner/validator.py
CHANGED
|
@@ -95,8 +95,8 @@ class PlannerValidator:
|
|
| 95 |
f"source_id {src!r} (known: {sorted(known_sources)})"
|
| 96 |
)
|
| 97 |
|
| 98 |
-
# Check 8b — inline
|
| 99 |
-
if call.tool == "
|
| 100 |
self._validate_inline_ir(task.id, call.args, catalog)
|
| 101 |
|
| 102 |
# Check 7 — success_criteria is checkable.
|
|
@@ -114,20 +114,20 @@ class PlannerValidator:
|
|
| 114 |
raw_ir = args.get("ir")
|
| 115 |
if not isinstance(raw_ir, dict):
|
| 116 |
raise PlannerValidationError(
|
| 117 |
-
f"task {task_id}:
|
| 118 |
f"object, got {type(raw_ir).__name__}"
|
| 119 |
)
|
| 120 |
try:
|
| 121 |
ir = QueryIR.model_validate(raw_ir)
|
| 122 |
except ValidationError as e:
|
| 123 |
raise PlannerValidationError(
|
| 124 |
-
f"task {task_id}:
|
| 125 |
) from e
|
| 126 |
try:
|
| 127 |
self._ir_validator.validate(ir, catalog)
|
| 128 |
except IRValidationError as e:
|
| 129 |
raise PlannerValidationError(
|
| 130 |
-
f"task {task_id}:
|
| 131 |
) from e
|
| 132 |
|
| 133 |
@staticmethod
|
|
|
|
| 95 |
f"source_id {src!r} (known: {sorted(known_sources)})"
|
| 96 |
)
|
| 97 |
|
| 98 |
+
# Check 8b — inline retrieve_data IR validates against the catalog.
|
| 99 |
+
if call.tool == "retrieve_data":
|
| 100 |
self._validate_inline_ir(task.id, call.args, catalog)
|
| 101 |
|
| 102 |
# Check 7 — success_criteria is checkable.
|
|
|
|
| 114 |
raw_ir = args.get("ir")
|
| 115 |
if not isinstance(raw_ir, dict):
|
| 116 |
raise PlannerValidationError(
|
| 117 |
+
f"task {task_id}: retrieve_data.args.ir must be an inline QueryIR "
|
| 118 |
f"object, got {type(raw_ir).__name__}"
|
| 119 |
)
|
| 120 |
try:
|
| 121 |
ir = QueryIR.model_validate(raw_ir)
|
| 122 |
except ValidationError as e:
|
| 123 |
raise PlannerValidationError(
|
| 124 |
+
f"task {task_id}: retrieve_data.args.ir is not a valid QueryIR: {e}"
|
| 125 |
) from e
|
| 126 |
try:
|
| 127 |
self._ir_validator.validate(ir, catalog)
|
| 128 |
except IRValidationError as e:
|
| 129 |
raise PlannerValidationError(
|
| 130 |
+
f"task {task_id}: retrieve_data IR failed catalog validation: {e}"
|
| 131 |
) from e
|
| 132 |
|
| 133 |
@staticmethod
|
src/catalog/reader.py
CHANGED
|
@@ -45,8 +45,8 @@ class MemoizingCatalogReader(CatalogReader):
|
|
| 45 |
|
| 46 |
One per request. The same per-user catalog is otherwise fetched from the
|
| 47 |
catalog DB 4-5x during a single slow-path run (planner load, then
|
| 48 |
-
|
| 49 |
-
|
| 50 |
to one round-trip
|
| 51 |
per distinct source_hint and pins a single consistent snapshot for the whole
|
| 52 |
request (plan-time and execution-time catalogs can no longer diverge).
|
|
|
|
| 45 |
|
| 46 |
One per request. The same per-user catalog is otherwise fetched from the
|
| 47 |
catalog DB 4-5x during a single slow-path run (planner load, then
|
| 48 |
+
check_data's structured read + check_knowledge's unstructured read, then
|
| 49 |
+
retrieve_data's structured read). Wrapping the base reader collapses those
|
| 50 |
to one round-trip
|
| 51 |
per distinct source_hint and pins a single consistent snapshot for the whole
|
| 52 |
request (plan-time and execution-time catalogs can no longer diverge).
|
src/config/prompts/planner.md
CHANGED
|
@@ -10,7 +10,7 @@ only a `TaskList` object that conforms to the provided schema.
|
|
| 10 |
|
| 11 |
1. **Emit intent, never code.** Never write SQL, pandas, or any code. The only
|
| 12 |
query you express is an inline `QueryIR` (a JSON intent object) inside a
|
| 13 |
-
`
|
| 14 |
2. **The plan is static.** There is no replanning and no execution feedback. Plan
|
| 15 |
the whole analysis up front; assume each task runs once, in dependency order.
|
| 16 |
3. **Use only tools from the "Available tools" list.** Never invent a tool name.
|
|
@@ -31,16 +31,16 @@ only a `TaskList` object that conforms to the provided schema.
|
|
| 31 |
- **Wire data between tasks with placeholders.** When a task needs an upstream
|
| 32 |
task's output as an argument, use the string `"${t<id>}"` (e.g. `"${t2}"`) as
|
| 33 |
the argument value. Set `depends_on` accordingly.
|
| 34 |
-
- **Data access vs analytics tools.** `
|
| 35 |
point: use it to select, filter, and pull rows (and simple built-in
|
| 36 |
count/sum/avg/min/max/count_distinct the IR can express). For anything richer —
|
| 37 |
descriptive statistics (median/percentile/mode/std/skew), time trends, group
|
| 38 |
comparisons, share-of-total, correlation, segmentation, or data-quality
|
| 39 |
-
profiling — run `
|
| 40 |
the matching composite `analyze_*` tool via a `"${t<id>}"` `data` argument
|
| 41 |
(referencing the upstream result's column aliases).
|
| 42 |
- **Mixing structured + unstructured.** If qualitative context helps, add a
|
| 43 |
-
`
|
| 44 |
- **CRISP-DM stages.** Tag each task with the stage it serves:
|
| 45 |
`data_understanding`, `data_preparation`, or `evaluation`. (Never `modeling`.)
|
| 46 |
- **success_criteria is a reporting signal**, not a control trigger. State, in
|
|
|
|
| 10 |
|
| 11 |
1. **Emit intent, never code.** Never write SQL, pandas, or any code. The only
|
| 12 |
query you express is an inline `QueryIR` (a JSON intent object) inside a
|
| 13 |
+
`retrieve_data` tool call's `args.ir`.
|
| 14 |
2. **The plan is static.** There is no replanning and no execution feedback. Plan
|
| 15 |
the whole analysis up front; assume each task runs once, in dependency order.
|
| 16 |
3. **Use only tools from the "Available tools" list.** Never invent a tool name.
|
|
|
|
| 31 |
- **Wire data between tasks with placeholders.** When a task needs an upstream
|
| 32 |
task's output as an argument, use the string `"${t<id>}"` (e.g. `"${t2}"`) as
|
| 33 |
the argument value. Set `depends_on` accordingly.
|
| 34 |
+
- **Data access vs analytics tools.** `retrieve_data` is the data-access entry
|
| 35 |
point: use it to select, filter, and pull rows (and simple built-in
|
| 36 |
count/sum/avg/min/max/count_distinct the IR can express). For anything richer —
|
| 37 |
descriptive statistics (median/percentile/mode/std/skew), time trends, group
|
| 38 |
comparisons, share-of-total, correlation, segmentation, or data-quality
|
| 39 |
+
profiling — run `retrieve_data` to fetch the rows, then pass its output to
|
| 40 |
the matching composite `analyze_*` tool via a `"${t<id>}"` `data` argument
|
| 41 |
(referencing the upstream result's column aliases).
|
| 42 |
- **Mixing structured + unstructured.** If qualitative context helps, add a
|
| 43 |
+
`retrieve_knowledge` task against an unstructured source listed in the catalog.
|
| 44 |
- **CRISP-DM stages.** Tag each task with the stage it serves:
|
| 45 |
`data_understanding`, `data_preparation`, or `evaluation`. (Never `modeling`.)
|
| 46 |
- **success_criteria is a reporting signal**, not a control trigger. State, in
|
src/query/executor/db.py
CHANGED
|
@@ -220,7 +220,7 @@ class DbExecutor(BaseExecutor):
|
|
| 220 |
"""Best-effort: warm pooled engines for the catalog's schema sources.
|
| 221 |
|
| 222 |
Called at slow-path entry so the TCP+TLS+auth handshake overlaps the ~4s
|
| 223 |
-
Planner LLM call — by the time `
|
| 224 |
already established. Warming is an optimization, never a requirement, so
|
| 225 |
this never raises and per-source failures are swallowed.
|
| 226 |
"""
|
|
|
|
| 220 |
"""Best-effort: warm pooled engines for the catalog's schema sources.
|
| 221 |
|
| 222 |
Called at slow-path entry so the TCP+TLS+auth handshake overlaps the ~4s
|
| 223 |
+
Planner LLM call — by the time `retrieve_data` runs, the connection is
|
| 224 |
already established. Warming is an optimization, never a requirement, so
|
| 225 |
this never raises and per-source failures are swallowed.
|
| 226 |
"""
|
src/tools/analytics/aggregation.py
CHANGED
|
@@ -3,7 +3,7 @@
|
|
| 3 |
An analytical "family" tool: in ONE call it groups rows by one or more keys
|
| 4 |
and computes aggregates (sum, mean, count, min, max, median, nunique) per
|
| 5 |
group. This is the deterministic compute twin of the Planner's
|
| 6 |
-
`
|
| 7 |
|
| 8 |
STATUS: compute layer only — the function takes an already-materialized
|
| 9 |
DataFrame. The wrapper layer (fetching data from the catalog via source_id,
|
|
|
|
| 3 |
An analytical "family" tool: in ONE call it groups rows by one or more keys
|
| 4 |
and computes aggregates (sum, mean, count, min, max, median, nunique) per
|
| 5 |
group. This is the deterministic compute twin of the Planner's
|
| 6 |
+
`retrieve_data` step — the wrapper layer later maps a QueryIR onto this.
|
| 7 |
|
| 8 |
STATUS: compute layer only — the function takes an already-materialized
|
| 9 |
DataFrame. The wrapper layer (fetching data from the catalog via source_id,
|
src/tools/data_access.py
CHANGED
|
@@ -8,18 +8,18 @@ the runtime/Coordinator supplies them; INV-7 keeps the agent layer
|
|
| 8 |
tool-agnostic).
|
| 9 |
|
| 10 |
Tools implemented here:
|
| 11 |
-
- `
|
| 12 |
→ list sources (id, name, type, table count); with a
|
| 13 |
`source_id` → that source's tables/columns (one row per
|
| 14 |
column, metadata only — exposes `pii_flag`, never
|
| 15 |
sample values).
|
| 16 |
-
- `
|
| 17 |
type).
|
| 18 |
-
- `
|
| 19 |
execute, skipping the planner) and returns rows as
|
| 20 |
`ToolOutput(kind="table")` — the Pattern A handoff the
|
| 21 |
`analyze_*` tools consume.
|
| 22 |
-
- `
|
| 23 |
`ToolOutput(kind="documents")`.
|
| 24 |
|
| 25 |
Frozen guarantee (§8.4): **never throws.** Any failure returns
|
|
@@ -51,7 +51,7 @@ DispatcherFactory = Callable[[Catalog], ExecutorDispatcher]
|
|
| 51 |
# adding/renaming a data-access tool can't silently drift the router out of sync
|
| 52 |
# from the registry (R11). Must match the names in `DataAccessToolInvoker.invoke`.
|
| 53 |
DATA_ACCESS_TOOLS: frozenset[str] = frozenset(
|
| 54 |
-
{"
|
| 55 |
)
|
| 56 |
|
| 57 |
|
|
@@ -77,28 +77,28 @@ class DataAccessToolInvoker:
|
|
| 77 |
) -> None:
|
| 78 |
self._user_id = user_id
|
| 79 |
self._reader = catalog_reader
|
| 80 |
-
#
|
| 81 |
# validator is stateless; the dispatcher is built per-call from the
|
| 82 |
# request's catalog (executors are picked by source_type).
|
| 83 |
self._validator = ir_validator or IRValidator()
|
| 84 |
self._dispatcher_factory: DispatcherFactory = (
|
| 85 |
dispatcher_factory or ExecutorDispatcher
|
| 86 |
)
|
| 87 |
-
#
|
| 88 |
# for tests (the real one pulls PGVector + Redis). Lazy-imported on first
|
| 89 |
# use so importing this module stays cheap.
|
| 90 |
self._retriever = document_retriever
|
| 91 |
|
| 92 |
async def invoke(self, tool_name: str, args: dict[str, Any]) -> ToolOutput:
|
| 93 |
try:
|
| 94 |
-
if tool_name == "
|
| 95 |
-
return await self.
|
| 96 |
-
if tool_name == "
|
| 97 |
-
return await self.
|
| 98 |
-
if tool_name == "
|
| 99 |
-
return await self.
|
| 100 |
-
if tool_name == "
|
| 101 |
-
return await self.
|
| 102 |
return ToolOutput(
|
| 103 |
tool=tool_name, kind="error", error=f"unknown tool {tool_name!r}"
|
| 104 |
)
|
|
@@ -107,7 +107,7 @@ class DataAccessToolInvoker:
|
|
| 107 |
tool=tool_name, kind="error", error=f"{type(exc).__name__}: {exc}"
|
| 108 |
)
|
| 109 |
|
| 110 |
-
async def
|
| 111 |
"""Inspect the user's structured data sources (DB + tabular).
|
| 112 |
|
| 113 |
No `source_id` → an overview: one row per structured source (id, name,
|
|
@@ -116,7 +116,7 @@ class DataAccessToolInvoker:
|
|
| 116 |
|
| 117 |
Pattern A note: schema is catalog metadata only — never returns row
|
| 118 |
data or PII sample values (only the `pii_flag` boolean per column).
|
| 119 |
-
Unstructured documents are covered by `
|
| 120 |
"""
|
| 121 |
structured = await self._reader.read(self._user_id, "structured")
|
| 122 |
source_id = args.get("source_id")
|
|
@@ -127,7 +127,7 @@ class DataAccessToolInvoker:
|
|
| 127 |
for s in structured.sources
|
| 128 |
]
|
| 129 |
return ToolOutput(
|
| 130 |
-
tool="
|
| 131 |
kind="table",
|
| 132 |
columns=["source_id", "name", "source_type", "table_count"],
|
| 133 |
rows=rows,
|
|
@@ -139,7 +139,7 @@ class DataAccessToolInvoker:
|
|
| 139 |
)
|
| 140 |
if source is None:
|
| 141 |
return ToolOutput(
|
| 142 |
-
tool="
|
| 143 |
kind="error",
|
| 144 |
error=f"structured source {source_id!r} not found",
|
| 145 |
)
|
|
@@ -159,7 +159,7 @@ class DataAccessToolInvoker:
|
|
| 159 |
for c in t.columns
|
| 160 |
]
|
| 161 |
return ToolOutput(
|
| 162 |
-
tool="
|
| 163 |
kind="table",
|
| 164 |
columns=[
|
| 165 |
"table_id",
|
|
@@ -181,24 +181,24 @@ class DataAccessToolInvoker:
|
|
| 181 |
},
|
| 182 |
)
|
| 183 |
|
| 184 |
-
async def
|
| 185 |
"""List the user's unstructured sources (documents).
|
| 186 |
|
| 187 |
Documents have no column schema to drill into, so there is no
|
| 188 |
-
`source_id` mode — reading document content is `
|
| 189 |
job.
|
| 190 |
"""
|
| 191 |
unstructured = await self._reader.read(self._user_id, "unstructured")
|
| 192 |
rows = [[s.source_id, s.name, s.source_type] for s in unstructured.sources]
|
| 193 |
return ToolOutput(
|
| 194 |
-
tool="
|
| 195 |
kind="table",
|
| 196 |
columns=["source_id", "name", "source_type"],
|
| 197 |
rows=rows,
|
| 198 |
meta={"source_count": len(unstructured.sources)},
|
| 199 |
)
|
| 200 |
|
| 201 |
-
async def
|
| 202 |
"""Run one validated, single-table QueryIR and return rows as a table.
|
| 203 |
|
| 204 |
This is the spine of the slow path (Pattern A): the `analyze_*` tools
|
|
@@ -211,14 +211,14 @@ class DataAccessToolInvoker:
|
|
| 211 |
raw = args.get("ir")
|
| 212 |
if raw is None:
|
| 213 |
return ToolOutput(
|
| 214 |
-
tool="
|
| 215 |
)
|
| 216 |
|
| 217 |
try:
|
| 218 |
ir = raw if isinstance(raw, QueryIR) else QueryIR.model_validate(raw)
|
| 219 |
except ValidationError as exc:
|
| 220 |
return ToolOutput(
|
| 221 |
-
tool="
|
| 222 |
)
|
| 223 |
|
| 224 |
catalog = await self._reader.read(self._user_id, "structured")
|
|
@@ -227,7 +227,7 @@ class DataAccessToolInvoker:
|
|
| 227 |
self._validator.validate(ir, catalog)
|
| 228 |
except IRValidationError as exc:
|
| 229 |
return ToolOutput(
|
| 230 |
-
tool="
|
| 231 |
kind="error",
|
| 232 |
error=f"IR validation failed: {exc}",
|
| 233 |
)
|
|
@@ -238,7 +238,7 @@ class DataAccessToolInvoker:
|
|
| 238 |
|
| 239 |
if result.error:
|
| 240 |
return ToolOutput(
|
| 241 |
-
tool="
|
| 242 |
)
|
| 243 |
|
| 244 |
# QueryResult.rows is list[dict]; ToolOutput.rows is list[list] ordered
|
|
@@ -250,7 +250,7 @@ class DataAccessToolInvoker:
|
|
| 250 |
[_json_safe(row.get(c)) for c in result.columns] for row in result.rows
|
| 251 |
]
|
| 252 |
return ToolOutput(
|
| 253 |
-
tool="
|
| 254 |
kind="table",
|
| 255 |
columns=result.columns,
|
| 256 |
rows=rows,
|
|
@@ -266,7 +266,7 @@ class DataAccessToolInvoker:
|
|
| 266 |
},
|
| 267 |
)
|
| 268 |
|
| 269 |
-
async def
|
| 270 |
"""Dense-retrieve relevant chunks from the user's unstructured sources.
|
| 271 |
|
| 272 |
Pulls qualitative context (PDF/DOCX/TXT) for a natural-language `query`
|
|
@@ -274,7 +274,7 @@ class DataAccessToolInvoker:
|
|
| 274 |
`source_id` scopes to one source (best-effort metadata filter — the
|
| 275 |
router itself does not yet scope by source, so this prunes the results).
|
| 276 |
|
| 277 |
-
TODO(retrieval scoping): the Planner few-shot has no `
|
| 278 |
example, so `source_id` is rarely emitted today and this post-filter is
|
| 279 |
adequate. If source-scoped retrieval becomes common, push scoping down
|
| 280 |
into RetrievalRouter.retrieve()/DocumentRetriever (WHERE
|
|
@@ -284,7 +284,7 @@ class DataAccessToolInvoker:
|
|
| 284 |
query = args.get("query")
|
| 285 |
if not isinstance(query, str) or not query.strip():
|
| 286 |
return ToolOutput(
|
| 287 |
-
tool="
|
| 288 |
kind="error",
|
| 289 |
error="missing 'query' argument",
|
| 290 |
)
|
|
@@ -315,7 +315,7 @@ class DataAccessToolInvoker:
|
|
| 315 |
for r in results
|
| 316 |
]
|
| 317 |
return ToolOutput(
|
| 318 |
-
tool="
|
| 319 |
kind="documents",
|
| 320 |
value=documents,
|
| 321 |
meta={
|
|
|
|
| 8 |
tool-agnostic).
|
| 9 |
|
| 10 |
Tools implemented here:
|
| 11 |
+
- `check_data` — structured data sources (DB + tabular). No `source_id`
|
| 12 |
→ list sources (id, name, type, table count); with a
|
| 13 |
`source_id` → that source's tables/columns (one row per
|
| 14 |
column, metadata only — exposes `pii_flag`, never
|
| 15 |
sample values).
|
| 16 |
+
- `check_knowledge` — the user's unstructured sources / documents (id, name,
|
| 17 |
type).
|
| 18 |
+
- `retrieve_data` — runs a pre-built `QueryIR` (validate -> dispatch ->
|
| 19 |
execute, skipping the planner) and returns rows as
|
| 20 |
`ToolOutput(kind="table")` — the Pattern A handoff the
|
| 21 |
`analyze_*` tools consume.
|
| 22 |
+
- `retrieve_knowledge` — dense retrieval over unstructured sources, returns
|
| 23 |
`ToolOutput(kind="documents")`.
|
| 24 |
|
| 25 |
Frozen guarantee (§8.4): **never throws.** Any failure returns
|
|
|
|
| 51 |
# adding/renaming a data-access tool can't silently drift the router out of sync
|
| 52 |
# from the registry (R11). Must match the names in `DataAccessToolInvoker.invoke`.
|
| 53 |
DATA_ACCESS_TOOLS: frozenset[str] = frozenset(
|
| 54 |
+
{"check_data", "check_knowledge", "retrieve_data", "retrieve_knowledge"}
|
| 55 |
)
|
| 56 |
|
| 57 |
|
|
|
|
| 77 |
) -> None:
|
| 78 |
self._user_id = user_id
|
| 79 |
self._reader = catalog_reader
|
| 80 |
+
# retrieve_data deps — injectable so tests need no real LLM/DB. The
|
| 81 |
# validator is stateless; the dispatcher is built per-call from the
|
| 82 |
# request's catalog (executors are picked by source_type).
|
| 83 |
self._validator = ir_validator or IRValidator()
|
| 84 |
self._dispatcher_factory: DispatcherFactory = (
|
| 85 |
dispatcher_factory or ExecutorDispatcher
|
| 86 |
)
|
| 87 |
+
# retrieve_knowledge dep — the module singleton by default, injectable
|
| 88 |
# for tests (the real one pulls PGVector + Redis). Lazy-imported on first
|
| 89 |
# use so importing this module stays cheap.
|
| 90 |
self._retriever = document_retriever
|
| 91 |
|
| 92 |
async def invoke(self, tool_name: str, args: dict[str, Any]) -> ToolOutput:
|
| 93 |
try:
|
| 94 |
+
if tool_name == "check_data":
|
| 95 |
+
return await self._check_data(args)
|
| 96 |
+
if tool_name == "check_knowledge":
|
| 97 |
+
return await self._check_knowledge()
|
| 98 |
+
if tool_name == "retrieve_data":
|
| 99 |
+
return await self._retrieve_data(args)
|
| 100 |
+
if tool_name == "retrieve_knowledge":
|
| 101 |
+
return await self._retrieve_knowledge(args)
|
| 102 |
return ToolOutput(
|
| 103 |
tool=tool_name, kind="error", error=f"unknown tool {tool_name!r}"
|
| 104 |
)
|
|
|
|
| 107 |
tool=tool_name, kind="error", error=f"{type(exc).__name__}: {exc}"
|
| 108 |
)
|
| 109 |
|
| 110 |
+
async def _check_data(self, args: dict[str, Any]) -> ToolOutput:
|
| 111 |
"""Inspect the user's structured data sources (DB + tabular).
|
| 112 |
|
| 113 |
No `source_id` → an overview: one row per structured source (id, name,
|
|
|
|
| 116 |
|
| 117 |
Pattern A note: schema is catalog metadata only — never returns row
|
| 118 |
data or PII sample values (only the `pii_flag` boolean per column).
|
| 119 |
+
Unstructured documents are covered by `check_knowledge`.
|
| 120 |
"""
|
| 121 |
structured = await self._reader.read(self._user_id, "structured")
|
| 122 |
source_id = args.get("source_id")
|
|
|
|
| 127 |
for s in structured.sources
|
| 128 |
]
|
| 129 |
return ToolOutput(
|
| 130 |
+
tool="check_data",
|
| 131 |
kind="table",
|
| 132 |
columns=["source_id", "name", "source_type", "table_count"],
|
| 133 |
rows=rows,
|
|
|
|
| 139 |
)
|
| 140 |
if source is None:
|
| 141 |
return ToolOutput(
|
| 142 |
+
tool="check_data",
|
| 143 |
kind="error",
|
| 144 |
error=f"structured source {source_id!r} not found",
|
| 145 |
)
|
|
|
|
| 159 |
for c in t.columns
|
| 160 |
]
|
| 161 |
return ToolOutput(
|
| 162 |
+
tool="check_data",
|
| 163 |
kind="table",
|
| 164 |
columns=[
|
| 165 |
"table_id",
|
|
|
|
| 181 |
},
|
| 182 |
)
|
| 183 |
|
| 184 |
+
async def _check_knowledge(self) -> ToolOutput:
|
| 185 |
"""List the user's unstructured sources (documents).
|
| 186 |
|
| 187 |
Documents have no column schema to drill into, so there is no
|
| 188 |
+
`source_id` mode — reading document content is `retrieve_knowledge`'s
|
| 189 |
job.
|
| 190 |
"""
|
| 191 |
unstructured = await self._reader.read(self._user_id, "unstructured")
|
| 192 |
rows = [[s.source_id, s.name, s.source_type] for s in unstructured.sources]
|
| 193 |
return ToolOutput(
|
| 194 |
+
tool="check_knowledge",
|
| 195 |
kind="table",
|
| 196 |
columns=["source_id", "name", "source_type"],
|
| 197 |
rows=rows,
|
| 198 |
meta={"source_count": len(unstructured.sources)},
|
| 199 |
)
|
| 200 |
|
| 201 |
+
async def _retrieve_data(self, args: dict[str, Any]) -> ToolOutput:
|
| 202 |
"""Run one validated, single-table QueryIR and return rows as a table.
|
| 203 |
|
| 204 |
This is the spine of the slow path (Pattern A): the `analyze_*` tools
|
|
|
|
| 211 |
raw = args.get("ir")
|
| 212 |
if raw is None:
|
| 213 |
return ToolOutput(
|
| 214 |
+
tool="retrieve_data", kind="error", error="missing 'ir' argument"
|
| 215 |
)
|
| 216 |
|
| 217 |
try:
|
| 218 |
ir = raw if isinstance(raw, QueryIR) else QueryIR.model_validate(raw)
|
| 219 |
except ValidationError as exc:
|
| 220 |
return ToolOutput(
|
| 221 |
+
tool="retrieve_data", kind="error", error=f"invalid IR: {exc}"
|
| 222 |
)
|
| 223 |
|
| 224 |
catalog = await self._reader.read(self._user_id, "structured")
|
|
|
|
| 227 |
self._validator.validate(ir, catalog)
|
| 228 |
except IRValidationError as exc:
|
| 229 |
return ToolOutput(
|
| 230 |
+
tool="retrieve_data",
|
| 231 |
kind="error",
|
| 232 |
error=f"IR validation failed: {exc}",
|
| 233 |
)
|
|
|
|
| 238 |
|
| 239 |
if result.error:
|
| 240 |
return ToolOutput(
|
| 241 |
+
tool="retrieve_data", kind="error", error=result.error
|
| 242 |
)
|
| 243 |
|
| 244 |
# QueryResult.rows is list[dict]; ToolOutput.rows is list[list] ordered
|
|
|
|
| 250 |
[_json_safe(row.get(c)) for c in result.columns] for row in result.rows
|
| 251 |
]
|
| 252 |
return ToolOutput(
|
| 253 |
+
tool="retrieve_data",
|
| 254 |
kind="table",
|
| 255 |
columns=result.columns,
|
| 256 |
rows=rows,
|
|
|
|
| 266 |
},
|
| 267 |
)
|
| 268 |
|
| 269 |
+
async def _retrieve_knowledge(self, args: dict[str, Any]) -> ToolOutput:
|
| 270 |
"""Dense-retrieve relevant chunks from the user's unstructured sources.
|
| 271 |
|
| 272 |
Pulls qualitative context (PDF/DOCX/TXT) for a natural-language `query`
|
|
|
|
| 274 |
`source_id` scopes to one source (best-effort metadata filter — the
|
| 275 |
router itself does not yet scope by source, so this prunes the results).
|
| 276 |
|
| 277 |
+
TODO(retrieval scoping): the Planner few-shot has no `retrieve_knowledge`
|
| 278 |
example, so `source_id` is rarely emitted today and this post-filter is
|
| 279 |
adequate. If source-scoped retrieval becomes common, push scoping down
|
| 280 |
into RetrievalRouter.retrieve()/DocumentRetriever (WHERE
|
|
|
|
| 284 |
query = args.get("query")
|
| 285 |
if not isinstance(query, str) or not query.strip():
|
| 286 |
return ToolOutput(
|
| 287 |
+
tool="retrieve_knowledge",
|
| 288 |
kind="error",
|
| 289 |
error="missing 'query' argument",
|
| 290 |
)
|
|
|
|
| 315 |
for r in results
|
| 316 |
]
|
| 317 |
return ToolOutput(
|
| 318 |
+
tool="retrieve_knowledge",
|
| 319 |
kind="documents",
|
| 320 |
value=documents,
|
| 321 |
meta={
|
src/tools/registry.py
CHANGED
|
@@ -8,7 +8,7 @@ reads to choose a tool (KM-625). This replaces the agent team's local stub in
|
|
| 8 |
Conventions (decided with the agent team, KM-465):
|
| 9 |
- **Pattern A** — `analyze_*` tools do NOT self-fetch by `source_id`. Each takes a
|
| 10 |
`data` argument that is a `"${t<id>}"` placeholder pointing at an upstream
|
| 11 |
-
`
|
| 12 |
Column arguments reference the aliases that upstream query produced.
|
| 13 |
- `input_schema` is the lightweight JSON-schema-ish dict the planner validator
|
| 14 |
consumes: `required` (arg names with no default) + `properties` (allowed args).
|
|
@@ -17,8 +17,8 @@ Conventions (decided with the agent team, KM-465):
|
|
| 17 |
- `output_kind` is the `ToolOutput.kind` each tool returns: stats (labelled-metric
|
| 18 |
dict) | table (rows×cols) | series (ordered periods).
|
| 19 |
|
| 20 |
-
The four data-access tools (
|
| 21 |
-
|
| 22 |
`default_registry()` composes both slices.
|
| 23 |
"""
|
| 24 |
|
|
|
|
| 8 |
Conventions (decided with the agent team, KM-465):
|
| 9 |
- **Pattern A** — `analyze_*` tools do NOT self-fetch by `source_id`. Each takes a
|
| 10 |
`data` argument that is a `"${t<id>}"` placeholder pointing at an upstream
|
| 11 |
+
`retrieve_data` table output, resolved to a DataFrame at execution time.
|
| 12 |
Column arguments reference the aliases that upstream query produced.
|
| 13 |
- `input_schema` is the lightweight JSON-schema-ish dict the planner validator
|
| 14 |
consumes: `required` (arg names with no default) + `properties` (allowed args).
|
|
|
|
| 17 |
- `output_kind` is the `ToolOutput.kind` each tool returns: stats (labelled-metric
|
| 18 |
dict) | table (rows×cols) | series (ordered periods).
|
| 19 |
|
| 20 |
+
The four data-access tools (retrieve_data / retrieve_knowledge / check_data /
|
| 21 |
+
check_knowledge) are registered separately once their wrappers land (KM-465 #4);
|
| 22 |
`default_registry()` composes both slices.
|
| 23 |
"""
|
| 24 |
|