diff --git a/packages/ai/src/ai/common/database/db_instance_base.py b/packages/ai/src/ai/common/database/db_instance_base.py index e0a7ebd0f..e77c8d1e3 100644 --- a/packages/ai/src/ai/common/database/db_instance_base.py +++ b/packages/ai/src/ai/common/database/db_instance_base.py @@ -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 @@ -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} @@ -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 @@ -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' return result def _buildSQLQueryOnce( @@ -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)) diff --git a/packages/ai/tests/ai/common/database/test_db_base.py b/packages/ai/tests/ai/common/database/test_db_base.py index a219d36b8..710503043 100644 --- a/packages/ai/tests/ai/common/database/test_db_base.py +++ b/packages/ai/tests/ai/common/database/test_db_base.py @@ -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 # --------------------------------------------------------------------------- @@ -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} + + # --------------------------------------------------------------------------- # _tableExists (engine-backed) # ---------------------------------------------------------------------------