Rifqi Hafizuddin commited on
Commit
db6db73
·
1 Parent(s): b43ecc7

[KM-691] Bound traceability payload + reject contradictory IR filter ranges

Browse files

Traceability: _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 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 any string to CAP_STR; leave numbers/None untouched."""
 
 
 
 
 
 
 
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": _truncate(query) if isinstance(query, str) else None,
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