Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions backend/app/services/chat/chat_orchestration_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,10 @@
7. Every view is already scoped to the customer you are answering for -- it returns their rows and
no one else's. Never reference `customer_id` in a `WHERE` clause and never invent an id value;
there is no id for you to supply.
8. Use only relation and function names exactly as `get_database_schema` returns them -- never a
plausible-sounding guess (a topic word like "transactions" is not a table name). If
`execute_read_only_sql` rejects a query, its error names what is actually available -- retry
once with a corrected name from that list before telling the customer you cannot answer.
"""


Expand Down
19 changes: 17 additions & 2 deletions backend/app/services/chat/sql_tool_validator.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,8 +70,18 @@ def validate_query(sql: str) -> ValidationResult:
if name in cte_names:
continue
if name not in CURATED_VIEW_NAMES:
# Naming the actual allow-list here, not just rejecting, is load-bearing: the model
# has been observed guessing a plausible-sounding relation name (e.g. `transactions`
# for the topic word in its own system prompt) instead of the real view name
# (`v_transaction_history`) and then giving up on a bare rejection rather than
# retrying. Putting the correction in the same tool result removes the need for it to
# re-call `get_database_schema` or recall the list from earlier in the turn.
return ValidationResult(
ok=False, reason=f"relation '{table.name}' is not available to this assistant"
ok=False,
reason=(
f"relation '{table.name}' is not available to this assistant. "
f"Available views: {', '.join(sorted(CURATED_VIEW_NAMES))}."
),
)

for func in statement.find_all(exp.Func):
Expand All @@ -86,8 +96,13 @@ def validate_query(sql: str) -> ValidationResult:
continue
name = func.name.lower() if isinstance(func, exp.Anonymous) else func.sql_name().lower()
if name not in _ALLOWED_FUNCTIONS:
# Same self-correcting-message rationale as the relation check above.
return ValidationResult(
ok=False, reason=f"function '{name}' is not available to this assistant"
ok=False,
reason=(
f"function '{name}' is not available to this assistant. "
f"Available functions: {', '.join(sorted(_ALLOWED_FUNCTIONS))}."
),
)

return ValidationResult(
Expand Down
18 changes: 18 additions & 0 deletions backend/tests/unit/test_sql_tool_validator.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,3 +92,21 @@ def test_normalized_sql_drops_comments() -> None:
assert result.ok is True
assert result.normalized_sql is not None
assert "sneaky" not in result.normalized_sql


def test_off_allow_list_relation_reason_names_the_available_views() -> None:
"""Observed live: the model guessed `transactions` instead of `v_transaction_history` and
gave up on a bare rejection. The error must name the real options in the same tool result."""
result = validate_query("SELECT * FROM transactions")
assert result.ok is False
assert result.reason is not None
assert "v_transaction_history" in result.reason
assert "v_holdings" in result.reason


def test_disallowed_function_reason_names_the_available_functions() -> None:
result = validate_query("SELECT pg_sleep(5) FROM v_holdings")
assert result.ok is False
assert result.reason is not None
assert "sum" in result.reason
assert "count" in result.reason
Loading