Skip to content
Open
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
27 changes: 20 additions & 7 deletions packages/ai/src/ai/common/database/db_instance_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@

from ai.common.schema import Answer, Question, QuestionType
from ai.common.table import Table
from ai.common.utils import parse_bool
from rocketlib.types import IInvokeLLM

from .db_global_base import DEFAULT_MAX_EXECUTE_ROWS, DatabaseGlobalBase
Expand Down Expand Up @@ -237,13 +238,18 @@ def get_sql(self, args):
limit = args.get('limit', 250)
result = self._buildSQLQuery(question, limit=limit)

is_valid = result.get('isValid', '').lower() == 'true'
is_valid = parse_bool(result.get('isValid'))
sql_query = result.get('query', '')

if is_valid and sql_query and is_sql_safe(sql_query):
return {'sql': sql_query, 'valid': True}
elif is_valid and sql_query:
return {'error': 'Generated query contains unsafe SQL', 'sql': sql_query, 'valid': False}
elif result.get('error'):
# EXPLAIN rejected the query on every attempt -- surface the real
# database error instead of the rejected SQL, so callers can tell
# this apart from a genuinely off-topic question.
return {'error': result['error'], 'valid': False}
else:
return {'answer': sql_query, 'valid': False}

Expand Down Expand Up @@ -349,7 +355,7 @@ def _buildSQLQuery(self, question_text: str, *, limit: int = 250) -> dict:
for attempt in range(self.IGlobal.max_validation_attempts):
result = self._buildSQLQueryOnce(question_text, limit=limit, previous_sql=previous_sql, error=last_error)

is_valid = result.get('isValid', '').lower() == 'true'
is_valid = parse_bool(result.get('isValid'))
sql_query = result.get('query', '')

# If the LLM decided the question isn't a DB query, or the safety
Expand All @@ -370,9 +376,13 @@ def _buildSQLQuery(self, question_text: str, *, limit: int = 250) -> dict:
previous_sql = sql_query
last_error = explain_error

warning(
f'SQL validation failed after {self.IGlobal.max_validation_attempts} attempt(s); returning last result.'
)
warning(f'SQL validation failed after {self.IGlobal.max_validation_attempts} attempt(s); rejecting the query.')
# Every EXPLAIN attempt failed. Force isValid False so callers (get_sql,
# get_data, writeQuestions) don't execute a query the database already
# refused; keep the last query for context and carry the EXPLAIN error
# so callers can tell this apart from a non-SQL question.
result['isValid'] = False
result['error'] = last_error or 'SQL validation failed'
Comment thread
coderabbitai[bot] marked this conversation as resolved.
return result

def _buildSQLQueryOnce(
Expand Down Expand Up @@ -564,15 +574,18 @@ def writeQuestions(self, question: Question) -> None:
try:
# Ask the LLM to translate the natural-language question into SQL.
query_json = self._buildSQLQuery(question_text)
is_valid_query = query_json.get('isValid', '').lower() == 'true'
is_valid_query = parse_bool(query_json.get('isValid'))
sql_query = query_json.get('query')

# Execute the query only when the LLM flagged it as valid SQL and
# the safety check passes; otherwise return the LLM's text response.
# When EXPLAIN rejected the query on every attempt, prefer the real
# database error over the rejected SQL text, which reads as prose
# and would be indistinguishable from a genuine off-topic answer.
if is_valid_query and sql_query and is_sql_safe(sql_query):
result = self._executeSQLQuery(sql_query)
else:
result = sql_query
result = query_json.get('error') or sql_query

if 'text' in lanes:
self.instance.writeText(str(result))
Expand Down
92 changes: 92 additions & 0 deletions packages/ai/tests/ai/common/database/test_db_base.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@

from ai.common.database.db_global_base import DatabaseGlobalBase
from ai.common.database.db_instance_base import DatabaseInstanceBase
from ai.common.utils import parse_bool


# ---------------------------------------------------------------------------
Expand Down Expand Up @@ -218,6 +219,97 @@ def test_sanitize_row_scalar_input():
assert DatabaseInstanceBase._sanitize_row(42) == 42


# ---------------------------------------------------------------------------
# _buildSQLQuery (EXPLAIN exhaustion + isValid parsing) — regression for #1601
# ---------------------------------------------------------------------------


class _TestableInstance(DatabaseInstanceBase):
"""Concrete DatabaseInstanceBase that satisfies the two abstract methods."""

def _db_display_name(self):
"""Human-readable name used in tool descriptions."""
return 'TestDB'

def _db_dialect(self):
"""Machine-readable dialect identifier."""
return 'testdb'


class _FakeGlobal:
"""Stub IGlobal: every EXPLAIN attempt rejects the query."""

def __init__(self, max_attempts=2, explain_error='syntax error near FROM'):
self.max_validation_attempts = max_attempts
self._explain_error = explain_error

def _validateQuery(self, sql):
return False, self._explain_error


def _sql_instance(fake_global):
inst = _TestableInstance.__new__(_TestableInstance)
inst.IGlobal = fake_global
return inst


def test_build_sql_query_marks_invalid_after_explain_exhaustion():
"""Every EXPLAIN attempt failing must flip isValid to False, not leave the LLM's 'true'.

Before the fix, exhaustion returned the last LLM response unchanged, so
get_sql/get_data/writeQuestions executed a statement the database had
already refused on every attempt.
"""
inst = _sql_instance(_FakeGlobal())
inst._buildSQLQueryOnce = lambda question_text, *, limit, previous_sql, error: {
'isValid': 'true',
'query': 'SELECT * FROM users',
}

result = inst._buildSQLQuery('all users')

assert parse_bool(result.get('isValid')) is False
assert result.get('error') # the EXPLAIN error is carried for the caller


def test_build_sql_query_accepts_real_json_bool_isvalid():
"""A real JSON bool from the LLM must not crash isValid parsing.

Before the fix, `result.get('isValid', '').lower()` raised AttributeError
the moment the LLM returned a real bool instead of the string 'true'.
"""
fake_global = _FakeGlobal()
fake_global._validateQuery = lambda sql: (True, None)
inst = _sql_instance(fake_global)
inst._buildSQLQueryOnce = lambda question_text, *, limit, previous_sql, error: {
'isValid': True, # real bool, not the string 'true'
'query': 'SELECT 1',
}

result = inst._buildSQLQuery('anything') # must not raise AttributeError

assert parse_bool(result.get('isValid')) is True


def test_get_sql_surfaces_explain_error_not_rejected_sql():
"""get_sql() must return the EXPLAIN error, not the rejected SQL as prose.

Before this fix, the invalid branch always returned {'answer': sql_query,
'valid': False}, so a caller could not tell an EXPLAIN rejection (rejected
SQL text) apart from a genuinely off-topic question (LLM prose answer).
"""
fake_global = _FakeGlobal(explain_error='column "userz" does not exist')
inst = _sql_instance(fake_global)
inst._buildSQLQueryOnce = lambda question_text, *, limit, previous_sql, error: {
'isValid': 'true',
'query': 'SELECT * FROM userz',
}

result = inst.get_sql({'question': 'all users'})

assert result == {'error': 'column "userz" does not exist', 'valid': False}

Comment on lines +294 to +311

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Add regression coverage for writeQuestions() error propagation.

The new test covers get_sql() only. The separate fallback changed in db_instance_base.py Lines [582-588] could regress to emitting rejected SQL while this test still passes. Add a test asserting that writeQuestions() surfaces the EXPLAIN error through its text/answer lanes.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@packages/ai/tests/ai/common/database/test_db_base.py` around lines 294 - 311,
Add regression coverage for writeQuestions() that configures an EXPLAIN failure
and asserts the returned text/answer fields contain the EXPLAIN error rather
than the rejected SQL. Keep the existing get_sql() test unchanged and exercise
the separate fallback path in db_instance_base.py.


# ---------------------------------------------------------------------------
# _tableExists (engine-backed)
# ---------------------------------------------------------------------------
Expand Down
Loading