Rifqi Hafizuddin commited on
Commit ·
db6db73
1
Parent(s): b43ecc7
[KM-691] Bound traceability payload + reject contradictory IR filter ranges
Browse filesTraceability: _truncate now summarizes embedded ToolOutputs (Pattern-A
analyze_* inputs) instead of re-embedding the full upstream table, and
executed queries get a dedicated 2000-char cap (was clipped at 300).
Planner: IRValidator rejects >=2 disjoint BETWEEN ranges on one column
(all filters are ANDed -> 0 rows), and planner.md teaches multi-period
comparisons to use one spanning range + group-by or one query per period.
- src/config/prompts/planner.md +7 -0
- src/query/ir/validator.py +34 -0
- src/traceability/scratchpad.py +18 -2
src/config/prompts/planner.md
CHANGED
|
@@ -84,6 +84,13 @@ only a `TaskList` object that conforms to the provided schema.
|
|
| 84 |
- For `in`/`not_in` (value is a list) or `between` (value is `[low, high]`), it
|
| 85 |
is still the ELEMENT type — a list of names is `"string"`, a date range is
|
| 86 |
`"date"`. It is **never** `"list"`.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 87 |
- **Use only the `args` a tool lists** — e.g. `analyze_aggregate` takes only
|
| 88 |
`data`/`aggregations`/`group_by`, so never add `order_by`/`limit` to it.
|
| 89 |
- **Top-N ("top/most/least N by <metric>") is a single `retrieve_data` query**,
|
|
|
|
| 84 |
- For `in`/`not_in` (value is a list) or `between` (value is `[low, high]`), it
|
| 85 |
is still the ELEMENT type — a list of names is `"string"`, a date range is
|
| 86 |
`"date"`. It is **never** `"list"`.
|
| 87 |
+
- **Filters are ANDed — never stack two disjoint ranges on one column.** Every
|
| 88 |
+
`filters[]` entry combines with AND, so two non-overlapping `between` ranges on
|
| 89 |
+
the SAME column (e.g. `order_date` in Q1 2025 AND in Q1 2026) match **zero rows**.
|
| 90 |
+
For a multi-period comparison, either (a) use ONE spanning range
|
| 91 |
+
(`between ["2025-01-01", "2026-03-31"]`) and let the analysis step group by the
|
| 92 |
+
period, or (b) emit a SEPARATE `retrieve_data` task per period. Never AND two
|
| 93 |
+
ranges on one column expecting an OR.
|
| 94 |
- **Use only the `args` a tool lists** — e.g. `analyze_aggregate` takes only
|
| 95 |
`data`/`aggregations`/`group_by`, so never add `order_by`/`limit` to it.
|
| 96 |
- **Top-N ("top/most/least N by <metric>") is a single `retrieve_data` query**,
|
src/query/ir/validator.py
CHANGED
|
@@ -5,6 +5,8 @@ is re-prompted with the error context (max 3 retries) — error messages
|
|
| 5 |
must therefore be specific enough that the LLM can self-correct.
|
| 6 |
"""
|
| 7 |
|
|
|
|
|
|
|
| 8 |
from ...catalog.models import Catalog, Column, Source, Table
|
| 9 |
from .models import QueryIR
|
| 10 |
from .operators import (
|
|
@@ -80,6 +82,8 @@ class IRValidator:
|
|
| 80 |
f"(allowed: {sorted(allowed)})"
|
| 81 |
)
|
| 82 |
|
|
|
|
|
|
|
| 83 |
for i, col_id in enumerate(ir.group_by):
|
| 84 |
self._require_column(columns_by_id, col_id, f"group_by[{i}]")
|
| 85 |
|
|
@@ -100,6 +104,36 @@ class IRValidator:
|
|
| 100 |
f"limit {ir.limit} exceeds hard cap {LIMIT_HARD_CAP}"
|
| 101 |
)
|
| 102 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 103 |
def _validate_joins(
|
| 104 |
self, ir: QueryIR, source: Source, columns_by_id: dict[str, Column]
|
| 105 |
) -> None:
|
|
|
|
| 5 |
must therefore be specific enough that the LLM can self-correct.
|
| 6 |
"""
|
| 7 |
|
| 8 |
+
from typing import Any
|
| 9 |
+
|
| 10 |
from ...catalog.models import Catalog, Column, Source, Table
|
| 11 |
from .models import QueryIR
|
| 12 |
from .operators import (
|
|
|
|
| 82 |
f"(allowed: {sorted(allowed)})"
|
| 83 |
)
|
| 84 |
|
| 85 |
+
self._reject_contradictory_ranges(ir)
|
| 86 |
+
|
| 87 |
for i, col_id in enumerate(ir.group_by):
|
| 88 |
self._require_column(columns_by_id, col_id, f"group_by[{i}]")
|
| 89 |
|
|
|
|
| 104 |
f"limit {ir.limit} exceeds hard cap {LIMIT_HARD_CAP}"
|
| 105 |
)
|
| 106 |
|
| 107 |
+
@staticmethod
|
| 108 |
+
def _reject_contradictory_ranges(ir: QueryIR) -> None:
|
| 109 |
+
"""Reject disjoint BETWEEN ranges on the same column.
|
| 110 |
+
|
| 111 |
+
All filters in an IR are ANDed, so two non-overlapping BETWEEN ranges on one
|
| 112 |
+
column (e.g. Q1 2025 AND Q1 2026 on order_date) match no rows — a common LLM
|
| 113 |
+
slip when it means a multi-period comparison. Rejecting lets the planner's
|
| 114 |
+
re-prompt retry correct it (one spanning range + group-by the period, or one
|
| 115 |
+
query per period). Best-effort: incomparable values are skipped, not raised.
|
| 116 |
+
"""
|
| 117 |
+
ranges: dict[str, list[tuple[Any, Any]]] = {}
|
| 118 |
+
for f in ir.filters:
|
| 119 |
+
if f.op == "between" and isinstance(f.value, list) and len(f.value) == 2:
|
| 120 |
+
ranges.setdefault(f.column_id, []).append((f.value[0], f.value[1]))
|
| 121 |
+
for col_id, rs in ranges.items():
|
| 122 |
+
if len(rs) < 2:
|
| 123 |
+
continue
|
| 124 |
+
try:
|
| 125 |
+
# AND of ranges = intersection [max(lo), min(hi)]; empty when lo > hi.
|
| 126 |
+
empty = max(lo for lo, _ in rs) > min(hi for _, hi in rs)
|
| 127 |
+
except TypeError:
|
| 128 |
+
continue # incomparable value types — skip (fail-open)
|
| 129 |
+
if empty:
|
| 130 |
+
raise IRValidationError(
|
| 131 |
+
f"filters place {len(rs)} non-overlapping BETWEEN ranges on column "
|
| 132 |
+
f"{col_id!r}; because all filters are ANDed this matches no rows. "
|
| 133 |
+
"For a multi-period comparison use a single spanning range and group "
|
| 134 |
+
"by the period, or emit one query per period."
|
| 135 |
+
)
|
| 136 |
+
|
| 137 |
def _validate_joins(
|
| 138 |
self, ir: QueryIR, source: Source, columns_by_id: dict[str, Column]
|
| 139 |
) -> None:
|
src/traceability/scratchpad.py
CHANGED
|
@@ -14,6 +14,8 @@ from __future__ import annotations
|
|
| 14 |
|
| 15 |
from typing import Any
|
| 16 |
|
|
|
|
|
|
|
| 17 |
from src.middlewares.logging import get_logger
|
| 18 |
|
| 19 |
from .schemas import PlanningInfo, PlanStep, ToolCallInfo, TraceabilityPayload
|
|
@@ -23,12 +25,26 @@ logger = get_logger("traceability")
|
|
| 23 |
# Truncation caps (bound the JSONB payload) — see plan §3.
|
| 24 |
CAP_PREVIEW_ROWS = 5
|
| 25 |
CAP_STR = 300
|
|
|
|
|
|
|
|
|
|
| 26 |
|
| 27 |
|
| 28 |
def _truncate(obj: Any) -> Any:
|
| 29 |
-
"""Recursively cap
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 30 |
if isinstance(obj, str):
|
| 31 |
return obj[:CAP_STR]
|
|
|
|
|
|
|
|
|
|
|
|
|
| 32 |
if isinstance(obj, dict):
|
| 33 |
return {k: _truncate(v) for k, v in obj.items()}
|
| 34 |
if isinstance(obj, list):
|
|
@@ -121,7 +137,7 @@ class TraceabilityScratchpad:
|
|
| 121 |
"type": "database",
|
| 122 |
"source_id": meta.get("source_id"),
|
| 123 |
"name": table,
|
| 124 |
-
"query":
|
| 125 |
"detail": {"table": table, "row_count": meta.get("row_count")},
|
| 126 |
})
|
| 127 |
|
|
|
|
| 14 |
|
| 15 |
from typing import Any
|
| 16 |
|
| 17 |
+
from pydantic import BaseModel
|
| 18 |
+
|
| 19 |
from src.middlewares.logging import get_logger
|
| 20 |
|
| 21 |
from .schemas import PlanningInfo, PlanStep, ToolCallInfo, TraceabilityPayload
|
|
|
|
| 25 |
# Truncation caps (bound the JSONB payload) — see plan §3.
|
| 26 |
CAP_PREVIEW_ROWS = 5
|
| 27 |
CAP_STR = 300
|
| 28 |
+
# Executed queries get a higher cap than free strings: the SQL/query is the point of
|
| 29 |
+
# the feature, and 300 chars mangles all but the smallest statements.
|
| 30 |
+
CAP_QUERY = 2000
|
| 31 |
|
| 32 |
|
| 33 |
def _truncate(obj: Any) -> Any:
|
| 34 |
+
"""Recursively cap strings to CAP_STR and summarize embedded tool results.
|
| 35 |
+
|
| 36 |
+
A Pattern-A `analyze_*` input carries its upstream `retrieve_data` result as a
|
| 37 |
+
`ToolOutput` (a BaseModel). Left untouched that would re-embed the FULL upstream
|
| 38 |
+
table (all rows + the untruncated query) into the payload, defeating the caps —
|
| 39 |
+
so a tool-result-shaped model is summarized via `_output_to_dict` (preview ≤ 5
|
| 40 |
+
rows, `row_count` kept), and any other BaseModel is dumped then recursed.
|
| 41 |
+
"""
|
| 42 |
if isinstance(obj, str):
|
| 43 |
return obj[:CAP_STR]
|
| 44 |
+
if isinstance(obj, BaseModel):
|
| 45 |
+
if hasattr(obj, "kind") and hasattr(obj, "rows"): # a ToolOutput-shaped result
|
| 46 |
+
return _output_to_dict(obj)
|
| 47 |
+
return _truncate(obj.model_dump(mode="json"))
|
| 48 |
if isinstance(obj, dict):
|
| 49 |
return {k: _truncate(v) for k, v in obj.items()}
|
| 50 |
if isinstance(obj, list):
|
|
|
|
| 137 |
"type": "database",
|
| 138 |
"source_id": meta.get("source_id"),
|
| 139 |
"name": table,
|
| 140 |
+
"query": query[:CAP_QUERY] if isinstance(query, str) else None,
|
| 141 |
"detail": {"table": table, "row_count": meta.get("row_count")},
|
| 142 |
})
|
| 143 |
|