From 83e6626bdc0c148d8fe5bee1ef17484d09e4dfc6 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Fri, 2 Oct 2026 16:46:05 -0300 Subject: [PATCH 1/2] fix: make decode_table_values robust to dictionary quirks - Match column names with TRIM, since some dictionaries store them with stray whitespace (e.g. Censo Escolar's "tipo_situacao_funcionamento "), which made the lookup return no entries. - Reject the dataset's own dictionary table as the target, instead of silently returning an empty list. - Return an explicit error when no dictionary entries match, instead of an empty list the agent cannot act on. Co-Authored-By: Claude Opus 5.5 --- app/agent/tools/bigquery.py | 22 +++++- tests/app/agent/tools/test_bigquery.py | 105 ++++++++++++++++++++++++- 2 files changed, 123 insertions(+), 4 deletions(-) diff --git a/app/agent/tools/bigquery.py b/app/agent/tools/bigquery.py index 013580f..c170ef8 100644 --- a/app/agent/tools/bigquery.py +++ b/app/agent/tools/bigquery.py @@ -143,10 +143,18 @@ def decode_table_values( f"Invalid table reference: '{table_gcp_id}'. Expected format: project.dataset.table" ) + if table_name == "dicionario": + raise ValueError( + f"'{table_gcp_id}' is the dataset's dictionary table. " + "Pass the data table that holds the coded column instead." + ) + dict_table_id = f"`{project_name}.{dataset_name}.dicionario`" + # Some dictionaries store column names with stray whitespace + # (e.g. "col "), so compare and return them trimmed. search_query = f""" - SELECT nome_coluna, chave, valor + SELECT TRIM(nome_coluna) AS nome_coluna, chave, valor FROM {dict_table_id} WHERE id_tabela = @table_name """ @@ -156,9 +164,9 @@ def decode_table_values( ] if column_name is not None: - search_query += "AND nome_coluna = @column_name\n" + search_query += "AND TRIM(nome_coluna) = @column_name\n" query_params.append( - bq.ScalarQueryParameter("column_name", "STRING", column_name), + bq.ScalarQueryParameter("column_name", "STRING", column_name.strip()), ) search_query += "ORDER BY nome_coluna, chave" @@ -179,4 +187,12 @@ def decode_table_values( except NotFound as e: raise ValueError("Dictionary table not found for this dataset.") from e + if not results: + target = ( + f"column '{column_name.strip()}' of " if column_name is not None else "" + ) + raise ValueError( + f"No dictionary entries found for {target}table '{table_name}'." + ) + return json.dumps(results, ensure_ascii=False, indent=2, default=str) diff --git a/tests/app/agent/tools/test_bigquery.py b/tests/app/agent/tools/test_bigquery.py index 3f77e10..dd2565b 100644 --- a/tests/app/agent/tools/test_bigquery.py +++ b/tests/app/agent/tools/test_bigquery.py @@ -461,13 +461,41 @@ def test_decode_specific_column( call_args = mock_bigquery_client.query.call_args[0][0] assert "id_tabela = @table_name" in call_args - assert "nome_coluna = @column_name" in call_args + assert "TRIM(nome_coluna) = @column_name" in call_args job_config = mock_bigquery_client.query.call_args[1]["job_config"] param_names = {p.name for p in job_config.query_parameters} assert "table_name" in param_names assert "column_name" in param_names + def test_decode_specific_column_strips_whitespace( + self, mocker: MockerFixture, mock_context: AgentContext + ): + """Test that column names are matched and returned without stray whitespace.""" + mock_query_job = MagicMock() + mock_query_job.result.return_value = [] + + mock_bigquery_client = MagicMock(spec=bq.Client) + mock_bigquery_client.query.return_value = mock_query_job + + mocker.patch( + "app.agent.tools.bigquery._bq_client", return_value=mock_bigquery_client + ) + + _invoke_tool( + decode_table_values, + {"table_gcp_id": "project.dataset.table", "column_name": " col1 "}, + context=mock_context, + ) + + call_args = mock_bigquery_client.query.call_args[0][0] + assert "SELECT TRIM(nome_coluna) AS nome_coluna" in call_args + assert "TRIM(nome_coluna) = @column_name" in call_args + + job_config = mock_bigquery_client.query.call_args[1]["job_config"] + params = {p.name: p.value for p in job_config.query_parameters} + assert params["column_name"] == "col1" + def test_dictionary_not_found( self, mocker: MockerFixture, mock_context: AgentContext ): @@ -495,6 +523,81 @@ def test_dictionary_not_found( assert output["status"] == "error" assert output["message"] == "Dictionary table not found for this dataset." + def test_dictionary_table_passed_as_target( + self, mocker: MockerFixture, mock_context: AgentContext + ): + """Test error when the dictionary table itself is passed instead of a data table.""" + mock_bigquery_client = MagicMock(spec=bq.Client) + + mocker.patch( + "app.agent.tools.bigquery._bq_client", return_value=mock_bigquery_client + ) + + message = _invoke_tool( + decode_table_values, + {"table_gcp_id": "project.dataset.dicionario", "column_name": "col1"}, + context=mock_context, + ) + + output = json.loads(message.content) + + assert output["status"] == "error" + assert "dictionary table" in output["message"] + mock_bigquery_client.query.assert_not_called() + + def test_no_entries_for_column( + self, mocker: MockerFixture, mock_context: AgentContext + ): + """Test error when the dictionary has no entries for the requested column.""" + mock_query_job = MagicMock() + mock_query_job.result.return_value = [] + + mock_bigquery_client = MagicMock(spec=bq.Client) + mock_bigquery_client.query.return_value = mock_query_job + + mocker.patch( + "app.agent.tools.bigquery._bq_client", return_value=mock_bigquery_client + ) + + message = _invoke_tool( + decode_table_values, + {"table_gcp_id": "project.dataset.table", "column_name": "col1"}, + context=mock_context, + ) + + output = json.loads(message.content) + + assert output["status"] == "error" + assert ( + output["message"] + == "No dictionary entries found for column 'col1' of table 'table'." + ) + + def test_no_entries_for_table( + self, mocker: MockerFixture, mock_context: AgentContext + ): + """Test error when the dictionary has no entries for the table.""" + mock_query_job = MagicMock() + mock_query_job.result.return_value = [] + + mock_bigquery_client = MagicMock(spec=bq.Client) + mock_bigquery_client.query.return_value = mock_query_job + + mocker.patch( + "app.agent.tools.bigquery._bq_client", return_value=mock_bigquery_client + ) + + message = _invoke_tool( + decode_table_values, + {"table_gcp_id": "project.dataset.table"}, + context=mock_context, + ) + + output = json.loads(message.content) + + assert output["status"] == "error" + assert output["message"] == "No dictionary entries found for table 'table'." + def test_invalid_table_reference(self, mock_context: AgentContext): """Test error when table reference format is invalid.""" message = _invoke_tool( From a7fce2ffe5e9ac75a3e117cd774764db4ce0e067 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Mon, 5 Oct 2026 11:13:23 -0300 Subject: [PATCH 2/2] test: tighten decode_table_values whitespace and dictionary-table tests Cover both whitespace cases in the column-name test: a stored name with stray whitespace matched by the model's trimmed name, and a passed name with whitespace. The mock now returns a row so the success path is exercised. Assert the full error message when the dictionary table is passed as the target. Co-Authored-By: Claude Opus 5.5 --- tests/app/agent/tools/test_bigquery.py | 33 +++++++++++++++++++++----- 1 file changed, 27 insertions(+), 6 deletions(-) diff --git a/tests/app/agent/tools/test_bigquery.py b/tests/app/agent/tools/test_bigquery.py index dd2565b..29c9117 100644 --- a/tests/app/agent/tools/test_bigquery.py +++ b/tests/app/agent/tools/test_bigquery.py @@ -468,12 +468,26 @@ def test_decode_specific_column( assert "table_name" in param_names assert "column_name" in param_names + @pytest.mark.parametrize( + "column_name", + [ + pytest.param("col1", id="stored-name-has-whitespace"), + pytest.param("col1 ", id="passed-name-has-whitespace"), + ], + ) def test_decode_specific_column_strips_whitespace( - self, mocker: MockerFixture, mock_context: AgentContext + self, mocker: MockerFixture, mock_context: AgentContext, column_name: str ): - """Test that column names are matched and returned without stray whitespace.""" + """Test that column names are matched and selected trimmed of stray whitespace. + + Some dictionaries store names with stray whitespace (e.g. "col1 "), which the + model passes without it, so the query trims the stored names on both sides. The + mock stands in for BigQuery, so it returns the row already trimmed by that query. + """ mock_query_job = MagicMock() - mock_query_job.result.return_value = [] + mock_query_job.result.return_value = [ + {"nome_coluna": "col1", "chave": "1", "valor": "Value 1"}, + ] mock_bigquery_client = MagicMock(spec=bq.Client) mock_bigquery_client.query.return_value = mock_query_job @@ -482,12 +496,16 @@ def test_decode_specific_column_strips_whitespace( "app.agent.tools.bigquery._bq_client", return_value=mock_bigquery_client ) - _invoke_tool( + message = _invoke_tool( decode_table_values, - {"table_gcp_id": "project.dataset.table", "column_name": " col1 "}, + {"table_gcp_id": "project.dataset.table", "column_name": column_name}, context=mock_context, ) + output = json.loads(message.content) + + assert output == [{"nome_coluna": "col1", "chave": "1", "valor": "Value 1"}] + call_args = mock_bigquery_client.query.call_args[0][0] assert "SELECT TRIM(nome_coluna) AS nome_coluna" in call_args assert "TRIM(nome_coluna) = @column_name" in call_args @@ -542,7 +560,10 @@ def test_dictionary_table_passed_as_target( output = json.loads(message.content) assert output["status"] == "error" - assert "dictionary table" in output["message"] + assert output["message"] == ( + "'project.dataset.dicionario' is the dataset's dictionary table. " + "Pass the data table that holds the coded column instead." + ) mock_bigquery_client.query.assert_not_called() def test_no_entries_for_column(