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..29c9117 100644 --- a/tests/app/agent/tools/test_bigquery.py +++ b/tests/app/agent/tools/test_bigquery.py @@ -461,13 +461,59 @@ 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 + @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, column_name: str + ): + """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 = [ + {"nome_coluna": "col1", "chave": "1", "valor": "Value 1"}, + ] + + 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": 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 + + 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 +541,84 @@ 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 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( + 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(