Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
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
17 changes: 11 additions & 6 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,7 +238,7 @@ 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):
Expand Down Expand Up @@ -349,7 +350,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 +371,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,7 +569,7 @@ 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
Expand Down
73 changes: 73 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,78 @@ 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


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