diff --git a/app/agent/tools/bigquery.py b/app/agent/tools/bigquery.py index 0478e3e..9aa81af 100644 --- a/app/agent/tools/bigquery.py +++ b/app/agent/tools/bigquery.py @@ -1,4 +1,5 @@ import inspect +import itertools import json import uuid from functools import cache @@ -15,6 +16,11 @@ MAX_BYTES_BILLED = 10 * 10**9 +# Cap on how many rows are serialized into the agent's context. The full result set +# is still materialized in BigQuery's destination table, so downloads (via query_ref) +# get every row regardless — this only keeps a large result from blowing up context. +MAX_CONTEXT_ROWS = 1000 + @cache def _bq_client() -> bq.Client: # pragma: no cover @@ -83,7 +89,9 @@ def execute_bigquery_sql( labels=labels, ), ) - rows = [dict(row) for row in job.result()] + result = job.result() + total_rows = result.total_rows + rows = [dict(row) for row in itertools.islice(result, MAX_CONTEXT_ROWS)] except GoogleAPICallError as e: reason = e.errors[0].get("reason") if getattr(e, "errors", None) else None if reason == "bytesBilledLimitExceeded": @@ -102,9 +110,18 @@ def execute_bigquery_sql( # (~24h TTL), so a later export hands back exactly these rows without re-running. query_ref = f"qr_{uuid.uuid4().hex}" - content = json.dumps( - {"row_count": len(rows), "rows": rows}, ensure_ascii=False, default=str - ) + payload = {"row_count": total_rows, "rows": rows} + + # Surface truncation only when it actually happened, so the agent knows the rows + # it sees are a subset — and that the full set is still available for download. + if total_rows > len(rows): + payload["truncated"] = True + payload["truncation_note"] = ( + f"Only the first {len(rows)} of {total_rows} rows are shown here to keep " + "the context small. The full result can still be downloaded from the interface." + ) + + content = json.dumps(payload, ensure_ascii=False, default=str) artifact = { "type": "query_result", diff --git a/tests/app/agent/tools/test_bigquery.py b/tests/app/agent/tools/test_bigquery.py index 2bace15..c59e9cf 100644 --- a/tests/app/agent/tools/test_bigquery.py +++ b/tests/app/agent/tools/test_bigquery.py @@ -12,11 +12,21 @@ from app.agent.context import AgentContext from app.agent.tools.bigquery import ( MAX_BYTES_BILLED, + MAX_CONTEXT_ROWS, decode_table_values, execute_bigquery_sql, ) +def _mock_result(rows: list[dict], total_rows: int | None = None) -> MagicMock: + """Stand in for BigQuery's RowIterator: iterable over `rows`, and carrying a + `total_rows` count (the full result size, which may exceed the fetched rows).""" + result = MagicMock() + result.__iter__.return_value = iter(rows) + result.total_rows = len(rows) if total_rows is None else total_rows + return result + + @pytest.fixture def mock_context() -> AgentContext: """The run context the agent injects into a tool @@ -69,7 +79,9 @@ def test_successful_query(self, mocker: MockerFixture, mock_context: AgentContex mock_dry_run_query_job.statement_type = "SELECT" mock_query_job = MagicMock() - mock_query_job.result.return_value = [{"col1": "value1"}, {"col1": "value2"}] + mock_query_job.result.return_value = _mock_result( + [{"col1": "value1"}, {"col1": "value2"}] + ) mock_query_job.destination.to_api_repr.return_value = { "projectId": "p", "datasetId": "d", @@ -106,7 +118,7 @@ def test_successful_query_exposes_destination_table_on_artifact( mock_dry_run_query_job.statement_type = "SELECT" mock_query_job = MagicMock() - mock_query_job.result.return_value = [{"col1": "value1"}] + mock_query_job.result.return_value = _mock_result([{"col1": "value1"}]) mock_query_job.destination.to_api_repr.return_value = { "projectId": "p", "datasetId": "d", @@ -146,7 +158,7 @@ def test_successful_query_empty_result( mock_dry_run_query_job.statement_type = "SELECT" mock_query_job = MagicMock() - mock_query_job.result.return_value = [] + mock_query_job.result.return_value = _mock_result([]) mock_bigquery_client = MagicMock(spec=bq.Client) mock_bigquery_client.query.side_effect = [ @@ -170,6 +182,91 @@ def test_successful_query_empty_result( ) assert message.artifact is None + def test_large_result_is_truncated_for_context( + self, mocker: MockerFixture, mock_context: AgentContext + ): + """A result larger than the cap only serializes a prefix, flags the truncation, + and still mints a download handle over the full (materialized) result.""" + total = MAX_CONTEXT_ROWS + 500 + all_rows = [{"n": i} for i in range(total)] + + mock_dry_run_query_job = MagicMock() + mock_dry_run_query_job.statement_type = "SELECT" + + mock_query_job = MagicMock() + mock_query_job.result.return_value = _mock_result(all_rows, total_rows=total) + mock_query_job.destination.to_api_repr.return_value = { + "projectId": "p", + "datasetId": "d", + "tableId": "t", + } + + mock_bigquery_client = MagicMock(spec=bq.Client) + mock_bigquery_client.query.side_effect = [ + mock_dry_run_query_job, + mock_query_job, + ] + + mocker.patch( + "app.agent.tools.bigquery._bq_client", return_value=mock_bigquery_client + ) + + message = _invoke_tool( + execute_bigquery_sql, + {"sql_query": "SELECT * FROM project.dataset.table", "slug": "resultado"}, + context=mock_context, + ) + + output = json.loads(message.content) + + assert output["row_count"] == total + assert len(output["rows"]) == MAX_CONTEXT_ROWS + assert output["rows"][0] == {"n": 0} + assert output["truncated"] is True + assert str(total) in output["truncation_note"] + # Full result is still downloadable. + assert re.fullmatch(r"qr_[0-9a-f]{32}", message.artifact["query_ref"]) + + def test_result_at_cap_is_not_flagged_truncated( + self, mocker: MockerFixture, mock_context: AgentContext + ): + """A result exactly at the cap returns every row and no truncation flag.""" + all_rows = [{"n": i} for i in range(MAX_CONTEXT_ROWS)] + + mock_dry_run_query_job = MagicMock() + mock_dry_run_query_job.statement_type = "SELECT" + + mock_query_job = MagicMock() + mock_query_job.result.return_value = _mock_result(all_rows) + mock_query_job.destination.to_api_repr.return_value = { + "projectId": "p", + "datasetId": "d", + "tableId": "t", + } + + mock_bigquery_client = MagicMock(spec=bq.Client) + mock_bigquery_client.query.side_effect = [ + mock_dry_run_query_job, + mock_query_job, + ] + + mocker.patch( + "app.agent.tools.bigquery._bq_client", return_value=mock_bigquery_client + ) + + message = _invoke_tool( + execute_bigquery_sql, + {"sql_query": "SELECT * FROM project.dataset.table", "slug": "resultado"}, + context=mock_context, + ) + + output = json.loads(message.content) + + assert output["row_count"] == MAX_CONTEXT_ROWS + assert len(output["rows"]) == MAX_CONTEXT_ROWS + assert "truncated" not in output + assert "truncation_note" not in output + def test_forbidden_statement_type( self, mocker: MockerFixture, mock_context: AgentContext ):