feat: charts and exports

alembic/versions/91426eac7604_query_handles_query_ref_pk.py

@@ -0,0 +1,37 @@
+"""Make query_ref the sole primary key of query_handles.
+
+Revision ID: 91426eac7604
+Revises: 4b3d2fa4a75f
+Create Date: 2026-08-18 15:02:34.101000
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "91426eac7604"
+down_revision: Union[str, Sequence[str], None] = "4b3d2fa4a75f"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    op.drop_constraint("query_handles_pkey", "query_handles", type_="primary")
+    op.create_primary_key("query_handles_pkey", "query_handles", ["query_ref"])
+    op.create_index(
+        op.f("ix_query_handles_message_id"),
+        "query_handles",
+        ["message_id"],
+        unique=False,
+    )
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    op.drop_index(op.f("ix_query_handles_message_id"), table_name="query_handles")
+    op.drop_constraint("query_handles_pkey", "query_handles", type_="primary")
+    op.create_primary_key(
+        "query_handles_pkey", "query_handles", ["message_id", "query_ref"]
+    )

alembic/versions/c4e1a9d2f6b8_add_streaming_message_status.py

@@ -0,0 +1,37 @@
+"""Add STREAMING message status.
+
+Revision ID: c4e1a9d2f6b8
+Revises: 91426eac7604
+Create Date: 2026-08-18 15:02:34.101000
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "c4e1a9d2f6b8"
+down_revision: Union[str, Sequence[str], None] = "91426eac7604"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Add STREAMING to the messagestatus enum.
+
+    Note: `ALTER TYPE ... ADD VALUE` cannot run inside a transaction block
+    on PostgreSQL pre-v12, so we use an autocommit block.
+    """
+    with op.get_context().autocommit_block():
+        op.execute("ALTER TYPE messagestatus ADD VALUE IF NOT EXISTS 'STREAMING'")
+
+
+def downgrade() -> None:
+    """Downgrade is intentionally unsupported.
+
+    Postgres cannot drop enum values, and rebuilding the type would
+    require remapping any STREAMING rows to another status.
+    """
+    raise NotImplementedError(
+        "Downgrade not supported: removing enum values would silently rewrite rows."
+    )

alembic/versions/f6ce7837e023_add_query_handles_table.py

@@ -1,4 +1,4 @@
-"""Add query handles table
+"""Add query handles table.
 
 Revision ID: f6ce7837e023
 Revises: 21d5a7602704

app/agent/prompts.py

@@ -5,7 +5,7 @@
 
 # Capabilities
 
-You can search and explore Base dos Dados's datasets and tables, query and analyze the data, translate coded values, and explain the platform and how you work. You have no other abilities — in particular, you cannot generate charts, plots, maps, visualizations, or files. Never offer, promise, or imply an action beyond these, in your prose answer or in the follow-up prompts. Downloads of query results are attached automatically by the interface — you do not generate them and need not offer them.
+You can search and explore Base dos Dados's datasets and tables, query and analyze the data, translate coded values, export a query's results as a downloadable file on request, plot a query's results as a chart — including a choropleth map shading Brazilian states or municipalities by a value — and explain the platform and how you work. You cannot produce files other than those exports, or map any geography other than Brazilian states and municipalities. Never offer, promise, or imply an action beyond these, in your prose answer or in the follow-up prompts.
 
 # When to act vs. ask
 
@@ -30,6 +30,20 @@
 
 Answer without tools only to explain the platform or your own capabilities, to ask for clarification, or to reuse data you already retrieved earlier in this conversation.
 
+# Exporting results
+
+`execute_bigquery_sql` returns the `query_ref` in its result. When the user explicitly asks to download or export a result in a specific format (AVRO, CSV, JSON Lines, Parquet), call `export_query_result` with that result's `query_ref` and the format. If the result is not found, use `list_query_results` to look it up. If the tool reports the result expired, re-run the query and export the new result.
+
+The interface offers the exported file — you do not generate or attach the file yourself. Do not describe the file's contents or claim anything about the download beyond the tool's confirmation.
+
+# Charting results
+
+When the user asks for a chart, plot, or visualization, call `chart_query_result` with the result's `query_ref` and a natural-language description of the chart — the mark (bar, line, point, …), what belongs on each axis, titles, labels and any grouping or color. You describe the chart in words; a data visualization specialist turns it into a Vega-Lite spec, so do not write the spec yourself.
+
+A value that varies across Brazilian states or municipalities can be drawn as a choropleth map (only these two levels, only Brazil). For a map, the result must carry the geographic key to join on — usually `sigla_uf` for states, `id_municipio` for municipalities — so write the query to return one row per state (or municipality) with that column plus the value, and describe it as a map of Brazil shaded by the value. Keep that key column in the result even though it is a code: the map joins on it and supplies the readable name itself, so this is the one case where you do not decode the geographic column away.
+
+Use the same `query_ref` referencing rules as exports. If the tool reports the result is too large, aggregate further in SQL, chart the smaller result and state this in your response.
+
 # Brazilian data landscape
 
 Main data sources available and starting keywords for the search:

app/agent/tools/__init__.py

@@ -2,6 +2,11 @@
 
 from app.agent.tools.api import get_dataset_details, get_table_details, search_datasets
 from app.agent.tools.bigquery import decode_table_values, execute_bigquery_sql
+from app.agent.tools.dataviz import (
+    chart_query_result,
+    export_query_result,
+    list_query_results,
+)
 
 
 class BDToolkit:
@@ -19,13 +24,19 @@ def get_tools() -> list[BaseTool]:
                 - get_table_details: Get comprehensive table information.
                 - execute_bigquery_sql: Execute SQL queries against BigQuery tables.
                 - decode_table_values: Decode coded values using dictionary tables.
+                - list_query_results: List this conversation's exportable/chartable results.
+                - export_query_result: Offer a result as a downloadable file (AVRO/CSV/JSONL/PARQUET).
+                - chart_query_result: Render a chart from a result.
         """
         return [
             search_datasets,
             get_dataset_details,
             get_table_details,
             execute_bigquery_sql,
             decode_table_values,
+            list_query_results,
+            export_query_result,
+            chart_query_result,
         ]
 
 

app/agent/tools/bigquery.py

@@ -44,7 +44,7 @@ def execute_bigquery_sql(
             current request — each slug names a separate download.
 
     Returns:
-        JSON object with `row_count` and `rows`.
+        JSON object with `row_count`, `rows` and `query_ref` (a reference to the query results).
     """
     client = _bq_client()
 
@@ -93,16 +93,19 @@ def execute_bigquery_sql(
             "the context small. The full result can still be downloaded from the interface."
         )
 
-    content = json.dumps(payload, ensure_ascii=False, default=str)
-
     # No rows -> nothing to download
     if not rows:
-        return content, None
+        return json.dumps(payload, ensure_ascii=False, default=str), None
 
     # Server-minted handle for the anonymous result table BigQuery already materialized
     # (~24h TTL), so a later export hands back exactly these rows without re-running.
     query_ref = f"qr_{uuid.uuid4().hex}"
 
+    # Surface the handle to the model so it can pass it to export_query_result /
+    # chart_query_result. The full destination table stays server-side in the artifact.
+    payload["query_ref"] = query_ref
+    content = json.dumps(payload, ensure_ascii=False, default=str)
+
     artifact = {
         "type": "query_result",
         "query_ref": query_ref,

app/agent/tools/dataviz.py

@@ -0,0 +1,170 @@
+import asyncio
+import json
+from typing import Any
+
+from langchain.tools import ToolRuntime
+from langchain_core.tools import tool
+
+from app.agent.context import AgentContext
+from app.agent.tools.exceptions import handle_tool_errors
+from app.charts import fetch_chart_data, generate_chart_spec, inject_chart_data
+from app.db.database import AsyncDatabase, sessionmaker
+from app.exports import (
+    OFFERED_EXPORT_FORMATS,
+    ResultTableExpired,
+    ResultTooLarge,
+    is_result_expired,
+    materialize_export,
+    sanitize_export_filename,
+)
+from app.i18n import MessageKey, translate
+
+
+@tool
+@handle_tool_errors
+async def list_query_results(runtime: ToolRuntime[AgentContext]) -> str:
+    """List this conversation's query results, including ones from earlier turns.
+
+    Returns:
+        JSON array of {query_ref, description, ran_at, expired}, oldest first.
+    """
+    async with sessionmaker() as session:
+        handles = await AsyncDatabase(session).get_query_handles_by_thread(
+            runtime.context.thread_id
+        )
+
+    results = [
+        {
+            "query_ref": handle.query_ref,
+            "description": handle.slug,
+            "ran_at": handle.created_at.isoformat(),
+            "expired": is_result_expired(handle.created_at),
+        }
+        for handle in handles
+    ]
+
+    return json.dumps(results, ensure_ascii=False, default=str)
+
+
+@tool(response_format="content_and_artifact")
+@handle_tool_errors(response_format="content_and_artifact")
+async def export_query_result(
+    query_ref: str, file_format: str, runtime: ToolRuntime[AgentContext]
+) -> tuple[str, dict[str, Any]]:
+    """Offer a query's result as a downloadable file in the requested format.
+
+    Args:
+        query_ref (str): The handle of the result to export.
+        file_format (str): One of AVRO, CSV, JSONL, PARQUET.
+
+    Returns:
+        A confirmation that the download is ready, or an error describing what to fix.
+    """
+    file_format = file_format.upper()
+
+    if file_format not in OFFERED_EXPORT_FORMATS:
+        raise ValueError(
+            f"Unsupported format '{file_format}'. "
+            f"Available: {', '.join(OFFERED_EXPORT_FORMATS)}."
+        )
+
+    async with sessionmaker() as session:
+        handle = await AsyncDatabase(session).get_query_handle_from_thread(
+            query_ref, runtime.context.thread_id
+        )
+
+    if handle is None:
+        raise ValueError(
+            f"No query result found for '{query_ref}'. "
+            "Call list_query_results to see the available results."
+        )
+
+    if is_result_expired(handle.created_at):
+        raise ValueError(
+            f"The result for '{query_ref}' has expired (results are kept ~24h). "
+            "Re-run the query, then export the new result."
+        )
+
+    # The user asked for the file explicitly, so materialize it now rather than lazily on click.
+    # This lets the card show the exact file size and the real download filename, and surfaces
+    # an over-limit result up front instead of a download that 400s on click.
+    try:
+        exported = await asyncio.to_thread(
+            materialize_export,
+            query_ref=handle.query_ref,
+            destination_table=handle.destination_table,
+            file_format=file_format,
+            filename=sanitize_export_filename(
+                handle.slug,
+                translate(MessageKey.DEFAULT_EXPORT_FILENAME, runtime.context.language),
+            ),
+            message_id=str(handle.message_id),
+        )
+    except ResultTableExpired as e:
+        raise ValueError(
+            f"The result for '{query_ref}' has expired (results are kept ~24h). "
+            "Re-run the query, then export the new result."
+        ) from e
+    except ResultTooLarge as e:
+        raise ValueError(
+            f"The result for '{query_ref}' is too large to export as a single file. "
+            "Aggregate or filter it in SQL first, then export the smaller result."
+        ) from e
+
+    # A client-facing affordance (not a `query_result` handle, so it is not redacted):
+    artifact = {
+        "type": "export",
+        "query_ref": handle.query_ref,
+        "format": file_format,
+        "filename": exported.filename,
+        "size_bytes": exported.size_bytes,
+        "message_id": str(handle.message_id),
+    }
+
+    content = json.dumps(
+        {
+            "status": "ready",
+            "query_ref": query_ref,
+            "format": file_format,
+            "filename": exported.filename,
+            "size_bytes": exported.size_bytes,
+        },
+        ensure_ascii=False,
+    )
+
+    return content, artifact
+
+
+@tool(response_format="content_and_artifact")
+@handle_tool_errors(response_format="content_and_artifact")
+async def chart_query_result(
+    query_ref: str, instructions: str, runtime: ToolRuntime[AgentContext]
+) -> tuple[str, dict[str, Any]]:
+    """Render a chart from a query's result.
+
+    Args:
+        query_ref (str): The handle of the result to chart.
+        instructions (str): A natural-language description of the chart.
+
+    Returns:
+        A confirmation that the chart was rendered, or an error describing what to fix.
+    """
+    handle, columns, rows = await fetch_chart_data(query_ref, runtime.context.thread_id)
+
+    spec = await generate_chart_spec(columns, rows, instructions)
+
+    # A client-facing artifact (not a `query_result` handle, so it is not redacted): the
+    # interface renders it with Vega-Embed. The data is bound server-side from the exact
+    # result rows, so the model cannot substitute the numbers.
+    artifact = {
+        "type": "chart",
+        "query_ref": handle.query_ref,
+        "spec": inject_chart_data(spec, rows),
+    }
+
+    content = json.dumps(
+        {"status": "rendered", "query_ref": query_ref, "row_count": len(rows)},
+        ensure_ascii=False,
+    )
+
+    return content, artifact