diff --git a/app/agent/tools.py b/app/agent/tools.py index 9717618..e65b31a 100644 --- a/app/agent/tools.py +++ b/app/agent/tools.py @@ -280,7 +280,7 @@ def search_datasets(query: str) -> str: Strategy: Start with broad terms like "censo", "ibge", "inep", "rais", then get specific if needed. Next step: Use `get_dataset_details()` with returned dataset IDs. - """ # noqa: E501 + """ response = _http_client.get( url=SEARCH_URL, params={"contains": "tables", "q": query, "page_size": PAGE_SIZE}, @@ -333,7 +333,7 @@ def get_dataset_details(dataset_id: str) -> str: - usage_guide: Provide key information and best practices for using the dataset. Next step: Use `execute_bigquery_sql()` to execute queries. - """ # noqa: E501 + """ response = _http_client.post( url=GRAPHQL_URL, json={ @@ -485,7 +485,7 @@ def execute_bigquery_sql(sql_query: str, config: RunnableConfig) -> str: Returns: str: Query results as JSON array. Empty results return "[]". - """ # noqa: E501 + """ client = get_bigquery_client() job_config = bq.QueryJobConfig(dry_run=True, use_query_cache=False) @@ -544,14 +544,16 @@ def decode_table_values( Returns: str: JSON array with chave (code) and valor (meaning) mappings. """ - # noqa: E501 + if "`" in table_gcp_id: + table_gcp_id = table_gcp_id.replace("`", "") + try: project_name, dataset_name, table_name = table_gcp_id.split(".") except ValueError: raise ToolError( message=f"Invalid table reference: '{table_gcp_id}'", error_type="INVALID_TABLE_REFERENCE", - instructions="Provide a valid table reference in the format `project.dataset.table`", + instructions="Provide a valid table reference in the format project.dataset.table", ) client = get_bigquery_client() diff --git a/tests/app/agent/test_tools.py b/tests/app/agent/test_tools.py index fc669a1..f441018 100644 --- a/tests/app/agent/test_tools.py +++ b/tests/app/agent/test_tools.py @@ -634,14 +634,14 @@ def test_decode_all_columns(self, mocker: MockerFixture, mock_config: dict): assert len(output.results) == 2 assert output.error_details is None - # Verify parameterized table filter was added to query call_args = mock_bigquery_client.query.call_args[0][0] assert "id_tabela = @table_name" in call_args + assert "nome_coluna = @column_name" not in call_args - # Verify query parameters include table_name 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 "nome_coluna = " not in param_names def test_decode_specific_column(self, mocker: MockerFixture, mock_config: dict): """Test decoding a specific column.""" @@ -670,13 +670,13 @@ def test_decode_specific_column(self, mocker: MockerFixture, mock_config: dict): assert output.status == "success" - # Verify parameterized column filter was added to query 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 - # Verify query parameters include column_name 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_dictionary_not_found(self, mocker: MockerFixture, mock_config: dict): @@ -722,7 +722,7 @@ def test_invalid_table_reference(self, mock_config: dict): assert output.error_details.message == "Invalid table reference: 'table'" assert ( output.error_details.instructions - == "Provide a valid table reference in the format `project.dataset.table`" + == "Provide a valid table reference in the format project.dataset.table" )