Merge pull request #56 from basedosdados/staging

chore: sync main with staging

Author: vrtornisiello

.gitignore

@@ -29,3 +29,6 @@ wheels/
 .pytest_cache/
 htmlcov/
 .coverage
+
+# Eval artifacts (generated transcripts + scores; regenerable, can be large)
+eval/*.json

alembic/versions/21d5a7602704_add_structured_response_column.py

@@ -0,0 +1,30 @@
+"""Add structured_response column to messages table.
+
+Revision ID: 21d5a7602704
+Revises: 19e2c92563e2
+Create Date: 2026-06-22 13:20:35.583925
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "21d5a7602704"
+down_revision: Union[str, Sequence[str], None] = "19e2c92563e2"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    op.add_column(
+        "message",
+        sa.Column("structured_response", sa.JSON(none_as_null=True), nullable=True),
+    )
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    op.drop_column("message", "structured_response")

app/agent/prompts.py

@@ -0,0 +1,98 @@
+from enum import Enum
+
+from pydantic import BaseModel, Field
+
+
+class TemporalGranularity(str, Enum):
+    """Granularity of the data's temporal coverage."""
+
+    DAY = "day"
+    MONTH = "month"
+    YEAR = "year"
+
+
+class TemporalCoverage(BaseModel):
+    """The interval the SQL query actually filtered on in the answer."""
+
+    period_start: str = Field(
+        description=(
+            "Start of the interval filtered by the SQL query (e.g. '2010' for "
+            "`ano = 2010` or `ano BETWEEN 2010 AND 2012`). Format it to match `granularity`: "
+            "YYYY (year), YYYY-MM (month) or YYYY-MM-DD (day) — e.g. '2010', '2010-01', '2010-01-01'. "
+            "May be narrower than the table's full coverage."
+        )
+    )
+    period_end: str = Field(
+        description=(
+            "End of the interval filtered by the SQL query (e.g. '2010' for "
+            "`ano = 2010`; '2012' for `ano BETWEEN 2010 AND 2012`). Format it to match `granularity`: "
+            "YYYY (year), YYYY-MM (month) or YYYY-MM-DD (day) — e.g. '2012', '2012-01', '2012-01-01'. "
+            "May be narrower than the table's full coverage."
+        )
+    )
+    granularity: TemporalGranularity = Field(
+        description=(
+            "Granularity of `period_start`/`period_end`, matching their format: "
+            "YYYY (year), YYYY-MM (month), YYYY-MM-DD (day)."
+        )
+    )
+
+
+class DataSource(BaseModel):
+    """A Base dos Dados table the answer draws on or points the user to."""
+
+    dataset_id: str = Field(
+        description=(
+            "Dataset UUID (the `dataset_id` field from `get_table_details` or the `id` "
+            "field from `get_dataset_details`), not the BigQuery id (e.g. 'br_bd_diretorios')."
+        )
+    )
+    table_id: str = Field(
+        description=(
+            "Table UUID (the `id` field from `get_table_details`, or a table's `id` from "
+            "the tables list of `get_dataset_details`). Not the dataset UUID, not the BigQuery id."
+        )
+    )
+    name: str = Field(description="Human-readable name of the table.")
+
+
+class StructuredResponse(BaseModel):
+    """The agent's structured response for the user interface."""
+
+    response: str = Field(
+        description=(
+            "The prose answer in Markdown, written in the user's language: a direct answer to the question "
+            "with the data obtained, plus analysis and context. Do NOT repeat the source, period, SQL, "
+            "or suggestions here — each of those has its own dedicated field."
+        )
+    )
+    data_sources: list[DataSource] | None = Field(
+        default=None,
+        description=(
+            "The tables the answer draws on — those you queried, or specific tables you recommend "
+            "on clarification turns. Leave empty (None) when no table is relevant."
+        ),
+    )
+    temporal_coverage: TemporalCoverage | None = Field(
+        default=None,
+        description=(
+            "The interval your SQL query actually filtered. Leave empty (None) when no query "
+            "was run (e.g. a clarification turn) or the answer has no temporal dimension."
+        ),
+    )
+    sql_queries: list[str] | None = Field(
+        default=None,
+        description=(
+            "The SQL queries whose results the answer is based on, each with "
+            "inline comments, so the user can reproduce the result. Include every "
+            "query that contributed to the answer (e.g. one query per metric when the "
+            "answer combines several), but EXCLUDE exploratory or failed-then-corrected queries. "
+            "Leave empty (None) when no query was executed."
+        ),
+    )
+    follow_up_questions: list[str] | None = Field(
+        default=None,
+        description=(
+            "3 suggested follow-up questions (in the user's language) to explore the data further."
+        ),
+    )

app/agent/schemas.py

@@ -46,14 +46,19 @@ async def search_datasets(query: str) -> str:
     CRITICAL: Use individual KEYWORDS only, not full sentences. The search engine uses Elasticsearch.
 
     Args:
-        query (str): 2-3 keywords maximum. Use Portuguese terms, organization acronyms, or dataset acronyms.
-            Good Examples: "censo", "educacao", "ibge", "inep", "rais", "saude".
+        query (str): 2-3 keywords maximum. Use Portuguese terms, organization names, or dataset names.
+            Good Examples: "censo", "rais", "ibge", "inep", "educacao", "saude".
             Avoid: "Brazilian population data by municipality".
 
     Returns:
         str: JSON array of datasets. If empty/irrelevant results, try different keywords.
 
-    Strategy: Start with broad terms like "censo", "ibge", "inep", "rais", then get specific if needed.
+    Strategy: hierarchical funnel — ALWAYS start with a SINGLE keyword and broaden a level only if it returns nothing:
+        1. Dataset name ("censo", "rais", "enem") or organization ("ibge", "inep", "tse").
+        2. Core theme ("educacao", "saude", "economia", "emprego").
+        3. English term ("health", "education").
+        4. A 2-3 word combination only if the levels above fail ("saude ms", "censo municipio").
+
     Next step: Use `get_dataset_details()` with returned dataset IDs.
     """
     response = await _client.get(
@@ -91,7 +96,7 @@ async def get_dataset_details(dataset_id: str) -> str:
 
     Args:
         dataset_id (str): Dataset ID obtained from `search_datasets()`.
-            This is typically a UUID-like string, not the human-readable name.
+            This is a UUID-like string, not the human-readable name.
 
     Returns:
         str: JSON object with complete dataset information, including:
@@ -173,6 +178,7 @@ async def get_dataset_details(dataset_id: str) -> str:
         dataset_tables.append(
             TableOverview(
                 id=table_id,
+                dataset_id=dataset_id,
                 gcp_id=table_gcp_id,
                 name=table_name,
                 description=table_description,
@@ -225,8 +231,6 @@ async def get_table_details(table_id: str) -> str:
             - period_start / period_end: First and last period covered by the table.
                 Format varies (`2024`, `'2026-04-12'`, etc.) — use the value verbatim,
                 matched to the appropriate temporal column (`ano`, `data`, etc.).
-
-    Next step: Use `execute_bigquery_sql()` to execute queries.
     """
     response = await _client.post(
         url=GRAPHQL_URL,
@@ -296,8 +300,11 @@ async def get_table_details(table_id: str) -> str:
             )
         )
 
+    dataset_id = table["dataset"]["id"].split("DatasetNode:")[-1]
+
     result = Table(
         id=table_id,
+        dataset_id=dataset_id,
         gcp_id=table_gcp_id,
         name=table_name,
         description=table_description,

app/agent/tools/api.py

@@ -26,6 +26,11 @@ def _get_client() -> bq.Client:  # pragma: no cover
 def execute_bigquery_sql(sql_query: str, config: RunnableConfig) -> str:
     """Execute a SQL query against BigQuery tables from the Base dos Dados database.
 
+    PRECONDITION — only call this when the question is already specific enough to
+    answer with data. For a broad/exploratory question (a bare topic) or one that
+    references an entity the user did not name, do NOT call this tool: explore the
+    metadata and ask the user to refine the question first.
+
     Use AFTER identifying the right datasets and understanding tables structure.
     It includes a 10GB processing limit for safety.
 

app/agent/tools/bigquery.py

@@ -16,6 +16,7 @@ class TableOverview(BaseModel):
     """Basic table information without column details."""
 
     id: str
+    dataset_id: str
     gcp_id: str | None
     name: str
     description: str | None

app/agent/tools/models.py

@@ -98,6 +98,9 @@
                         }
                     }
                 }
+                dataset {
+                    id
+                }
             }
         }
     }

app/agent/tools/queries.py

@@ -6,7 +6,9 @@
 from langgraph.graph.state import CompiledStateGraph
 from loguru import logger
 
+from app.agent.schemas import StructuredResponse
 from app.api.schemas import ConfigDict
+from app.api.streaming.data_sources import resolve_data_source_names
 from app.api.streaming.schemas import EventData, StreamEvent, ToolCall, ToolOutput
 from app.api.streaming.security import sanitize_markdown_links
 from app.db.database import AsyncDatabase, sessionmaker
@@ -84,35 +86,6 @@ def _truncate_json(
     return json.dumps(data, ensure_ascii=False, indent=2)
 
 
-def _parse_thinking(message: AIMessage) -> str | None:
-    """Parse thinking content from an AI message.
-
-    Some models (e.g., Gemini 3) return `message.content` as a list of typed blocks,
-    which may include `{"type": "thinking", "thinking": "..."}` entries. When
-    `content` is a plain string, no thinking is available.
-
-    Args:
-        message (AIMessage): The AI message from where to parse the thinking.
-
-    Returns:
-        str | None: The concatenated thinking text, or None if no thinking blocks exist.
-    """
-    if isinstance(message.content, str):
-        return None
-
-    blocks = [
-        block
-        for block in message.content
-        if isinstance(block, dict)
-        and block.get("type") == "thinking"
-        and isinstance(block.get("thinking"), str)
-    ]
-
-    thinking = "".join(block["thinking"] for block in blocks)
-
-    return thinking or None
-
-
 def _process_chunk(chunk: dict[str, Any]) -> StreamEvent | None:
     """Process a streaming chunk from a react agent workflow into a standardized StreamEvent.
 
@@ -128,7 +101,27 @@ def _process_chunk(chunk: dict[str, Any]) -> StreamEvent | None:
             - None for ignored chunks
     """
     if "model" in chunk:
-        ai_messages: list[AIMessage] = chunk["model"]["messages"]
+        update: dict[str, Any] = chunk["model"]
+
+        # When `response_format` is set (see app.main:91), the model node sets `structured_response`
+        # (a StructuredResponse) on the turn it produces the final answer. This is the final answer;
+        # the accompanying structured-output tool call / ToolMessage in `update["messages"]` is
+        # internal and must not be emitted as a tool_call.
+        structured: StructuredResponse | None = update.get("structured_response")
+
+        if structured is not None:
+            response_text = sanitize_markdown_links(structured.response)
+            structured_response = structured.model_dump()
+            structured_response["response"] = response_text
+            return StreamEvent(
+                type="final_answer",
+                data=EventData(
+                    content=response_text,
+                    structured_response=structured_response,
+                ),
+            )
+
+        ai_messages: list[AIMessage] = update["messages"]
 
         # If no messages are returned, the model returned an empty response
         # with no tool calls. This also counts as a final (but empty) answer.
@@ -145,7 +138,7 @@ def _process_chunk(chunk: dict[str, Any]) -> StreamEvent | None:
                 )
                 for tool_call in message.tool_calls
             ]
-            content = _parse_thinking(message) or message.text
+            content = message.text
         else:
             event_type = "final_answer"
             tool_calls = None
@@ -183,15 +176,17 @@ def _process_chunk(chunk: dict[str, Any]) -> StreamEvent | None:
         ]
 
         return StreamEvent(
-            type="tool_output", data=EventData(tool_outputs=tool_outputs)
+            type="tool_output",
+            data=EventData(tool_outputs=tool_outputs),
         )
     elif "ModelCallLimitMiddleware.before_model" in chunk:
         # before_model runs on every model iteration; only the limit-exceeded
         # path sets jump_to="end", so check that rather than the key's presence.
         update = chunk["ModelCallLimitMiddleware.before_model"] or {}
         if update.get("jump_to") == "end":
             event_data = EventData(
-                content=ErrorMessage.MODEL_CALL_LIMIT_REACHED, tool_calls=None
+                content=ErrorMessage.MODEL_CALL_LIMIT_REACHED,
+                tool_calls=None,
             )
             return StreamEvent(type="model_call_limit", data=event_data)
     return None
@@ -223,6 +218,7 @@ async def run_agent(
     events = []
     artifacts = []
     assistant_message = ""
+    structured_response: dict[str, Any] | None = None
     status: MessageStatus | None = None
 
     try:
@@ -245,6 +241,9 @@ async def run_agent(
                         artifacts.append(output.artifact)
             elif event.type == "final_answer":
                 assistant_message = event.data.content
+                if event.data.structured_response is not None:
+                    await resolve_data_source_names(event.data.structured_response)
+                structured_response = event.data.structured_response
                 status = MessageStatus.SUCCESS
             elif event.type == "model_call_limit":
                 assistant_message = event.data.content
@@ -280,6 +279,7 @@ async def run_agent(
             content=assistant_message,
             artifacts=artifacts or None,
             events=events or None,
+            structured_response=structured_response,
             status=status or MessageStatus.ERROR,
         )
         try: