Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
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
12 changes: 7 additions & 5 deletions app/agent/tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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={
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()
Expand Down
10 changes: 5 additions & 5 deletions tests/app/agent/test_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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."""
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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"
)


Expand Down