From 29e8cd1dbb382c9b67b84d71ab878f5f6ad56cb4 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Mon, 22 Jun 2026 13:33:57 -0300 Subject: [PATCH 01/12] feat: add structured agent response Return the agent's final answer as a structured object instead of a single markdown blob, so the frontend can build a richer UI. - Add StructuredResponse schema (response, data_sources, temporal_coverage, sql_query, follow_up_questions) in app/agent/schemas.py - Stream the structured payload via a new EventData.structured_response field and persist it on the message row (new nullable JSON column + migration) - Expose dataset_id on table tools so data_sources can carry BD UUIDs - Rework the 'Resposta Final' prompt: prose-only response, with source, period, SQL and follow-ups moved to dedicated structured fields Co-Authored-By: Claude Opus 4.8 (1M context) --- ...a7602704_add_structured_response_column.py | 30 ++++++ app/agent/prompts.py | 24 +++-- app/agent/schemas.py | 98 +++++++++++++++++++ app/agent/tools/api.py | 4 + app/agent/tools/models.py | 1 + app/agent/tools/queries.py | 3 + app/api/streaming/agent_runner.py | 32 +++++- app/api/streaming/schemas.py | 1 + app/db/models.py | 3 + app/main.py | 2 + 10 files changed, 187 insertions(+), 11 deletions(-) create mode 100644 alembic/versions/21d5a7602704_add_structured_response_column.py create mode 100644 app/agent/schemas.py diff --git a/alembic/versions/21d5a7602704_add_structured_response_column.py b/alembic/versions/21d5a7602704_add_structured_response_column.py new file mode 100644 index 0000000..3bc0e68 --- /dev/null +++ b/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") diff --git a/app/agent/prompts.py b/app/agent/prompts.py index 101adeb..6e21cc8 100644 --- a/app/agent/prompts.py +++ b/app/agent/prompts.py @@ -112,15 +112,23 @@ --- # Resposta Final -Escreva a resposta como um **texto corrido e fluido**, sem separar em seções nomeadas. Apresente os dados no formato mais legível possível: use tabelas Markdown para rankings, comparações, séries numéricas; use prosa para resumos, contexto e análises. A resposta deve conter: +Sua resposta final é **estruturada**: além do texto em prosa (campo `response`), você retorna campos dedicados (fonte dos dados, período de cobertura, consulta SQL e sugestões). + +## Campo `response` (prosa) +Escreva a resposta como um **texto corrido e fluido**, sem separar em seções nomeadas. Apresente os dados no formato mais legível possível: use tabelas Markdown para rankings, comparações, séries numéricas; use prosa para resumos, contexto e análises. O campo `response` deve conter: - A resposta direta à pergunta, com os dados obtidos. -- Análise e contexto relevante sobre os dados. -- A fonte, o período e o nível geográfico dos dados. - - Direcione os usuários para as tabelas consultadas, utilizando links Markdown no formato [Nome da Tabela](https://basedosdados.org/dataset/{{dataset_id}}?table={{table_id}}) -- A consulta SQL executada, em bloco de código com comentários inline. -- 2-3 sugestões de como explorar os dados mais a fundo. +- Análise e contexto relevante sobre os dados, incluindo o nível geográfico quando pertinente. + +Se a consulta retornar muitas linhas, **não** apresente todos os dados na prosa. Resuma os principais achados (top N, extremos, médias, tendências, etc.) e apresente apenas um recorte representativo dos dados. + +**NÃO** inclua na prosa: a lista de tabelas/links de fonte, o período de cobertura, a consulta SQL, as sugestões de exploração — esses elementos vão nos campos estruturados abaixo. -Se a consulta retornar muitas linhas, **não** apresente todos os dados na resposta. Resuma os principais achados (top N, extremos, médias, tendências, etc.), apresente apenas um recorte representativo dos dados e forneça a consulta SQL para que o usuário obtenha os dados completos por conta própria. +## Campos estruturados +Preencha-os **apenas** com base nos resultados das ferramentas obtidos nesta conversa: +- **`data_sources`**: as tabelas **efetivamente consultadas**, cada uma com `dataset_id` (UUID do campo `dataset_id` de `get_table_details`, ou do campo `id` de `get_dataset_details`), `table_id` (UUID do campo `id` de `get_table_details`) e um nome legível. **Nunca** use o `gcp_id` ou o nome BigQuery do dataset/tabela. Deixe vazio quando a resposta não usar dados de tabelas (ex.: explicar a plataforma, pedir esclarecimento, listar tipos de dados disponíveis). +- **`temporal_coverage`**: o período **efetivamente consultado** na consulta SQL — ou seja, o intervalo temporal definido pelos filtros da consulta (não o período de cobertura da tabela). Por exemplo, se a consulta filtrou `ano = 2010`, use `{{period_start: '2010', period_end: '2010'}}`; se filtrou `ano BETWEEN 2010 AND 2012`, use `{{period_start: '2010', period_end: '2012'}}`. Deixe vazio quando não houver dimensão temporal. +- **`sql_query`**: a consulta SQL executada, com comentários inline. Deixe vazio quando nenhuma consulta foi executada. +- **`follow_up_questions`**: 3 sugestões de como explorar os dados mais a fundo. ## Restrições - **NÃO** utilize headers Markdown (# ou ##) nem títulos de seção na resposta. @@ -135,4 +143,4 @@ 1. **Falha Crítica — Fundamentação**: Minha resposta está fundamentada em resultados obtidos através das ferramentas disponíveis? 2. **Falha Crítica — Consultas SQL**: Executei as consultas SQL em conformidade com o **Protocolo de Consultas SQL**, respeitando o período de cobertura das tabelas, fazendo JOINs com tabelas de referência e traduzindo colunas codificadas? -3. **Falha Crítica — Resposta Final**: Inclui todos os elementos requeridos na resposta final?""" +3. **Falha Crítica — Resposta Final**: A prosa do campo `response` está livre de fonte/período/SQL/sugestões, e os campos estruturados (`data_sources`, `temporal_coverage`, `sql_query`, `follow_up_questions`) estão preenchidos a partir dos resultados das ferramentas?""" diff --git a/app/agent/schemas.py b/app/agent/schemas.py new file mode 100644 index 0000000..58f58d3 --- /dev/null +++ b/app/agent/schemas.py @@ -0,0 +1,98 @@ +from enum import Enum + +from pydantic import BaseModel, Field + + +class TemporalGranularity(str, Enum): + """Granularidade do período de cobertura dos dados utilizados.""" + + DAY = "day" + MONTH = "month" + YEAR = "year" + + +class TemporalCoverage(BaseModel): + """Período de cobertura dos dados utilizados na resposta.""" + + period_start: str = Field( + description=( + "Início do período **efetivamente consultado** na consulta SQL (ex.: se a " + "consulta filtrou `ano = 2010`, use '2010'; se filtrou `ano BETWEEN 2010 AND " + "2012`, use '2010'). NÃO use o `period_start` dos metadados da tabela — " + "use o valor que aparece nos filtros temporais da consulta SQL executada." + ) + ) + period_end: str = Field( + description=( + "Fim do período **efetivamente consultado** na consulta SQL (ex.: se a consulta " + "filtrou `ano = 2010`, use '2010'; se filtrou `ano BETWEEN 2010 AND 2012`, " + "use '2012'). NÃO use o `period_end` dos metadados da tabela — use o " + "valor que aparece nos filtros temporais da consulta SQL executada." + ) + ) + granularity: TemporalGranularity = Field( + description=( + "Granularidade do período: 'day' (dia), 'month' (mês) ou 'year' (ano), " + "de acordo com o formato de `period_start`/`period_end`." + ) + ) + + +class DataSource(BaseModel): + """Uma tabela da Base dos Dados utilizada para responder à pergunta.""" + + dataset_id: str = Field( + description=( + "UUID do dataset retornado pelo campo `dataset_id` de `get_table_details` " + "ou pelo campo `id` de `get_dataset_details` (ex.: 'a1b2c3d4-...'). " + "NÃO use o `gcp_id` nem o nome BigQuery do dataset " + "(ex.: 'br_bd_diretorios') — use o UUID da API da Base dos Dados." + ) + ) + table_id: str = Field( + description=( + "UUID da tabela retornado pelo campo `id` de `get_table_details` (ex.: '3027c0d8-...'). " + "NÃO use o `gcp_id` nem o nome BigQuery da tabela — use o UUID da API da Base dos Dados." + ) + ) + name: str = Field(description="Nome legível da tabela utilizada.") + + +class StructuredResponse(BaseModel): + """Resposta estruturada do agente para a interface do usuário.""" + + response: str = Field( + description=( + "A resposta em prosa exibida ao usuário, em Markdown. Contém a resposta " + "direta à pergunta com os dados obtidos, além de análise e contexto " + "relevante. NÃO inclua aqui a fonte/tabelas, o período de cobertura, a " + "consulta SQL nem as sugestões de exploração — esses elementos são " + "retornados nos campos dedicados." + ) + ) + data_sources: list[DataSource] | None = Field( + default=None, + description=( + "As tabelas efetivamente consultadas para responder à pergunta. Deixe " + "vazio (None) quando a resposta não usar dados de tabelas (ex.: explicar " + "a plataforma, pedir esclarecimento ou listar os tipos de dados disponíveis)." + ), + ) + temporal_coverage: TemporalCoverage | None = Field( + default=None, + description=( + "O período geral dos dados utilizados. Deixe vazio (None) " + "quando a resposta não envolver uma dimensão temporal." + ), + ) + sql_query: str | None = Field( + default=None, + description=( + "A consulta SQL executada para obter os dados, com comentários inline. " + "Deixe vazio (None) quando nenhuma consulta foi executada." + ), + ) + follow_up_questions: list[str] = Field( + default_factory=list, + description="3 sugestões de perguntas para explorar os dados mais a fundo.", + ) diff --git a/app/agent/tools/api.py b/app/agent/tools/api.py index 841eafc..7f11c2b 100644 --- a/app/agent/tools/api.py +++ b/app/agent/tools/api.py @@ -173,6 +173,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, @@ -296,8 +297,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, diff --git a/app/agent/tools/models.py b/app/agent/tools/models.py index def58ff..840f329 100644 --- a/app/agent/tools/models.py +++ b/app/agent/tools/models.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 diff --git a/app/agent/tools/queries.py b/app/agent/tools/queries.py index 9c87da0..9158d7e 100644 --- a/app/agent/tools/queries.py +++ b/app/agent/tools/queries.py @@ -98,6 +98,9 @@ } } } + dataset { + id + } } } } diff --git a/app/api/streaming/agent_runner.py b/app/api/streaming/agent_runner.py index d53a33f..625c3b4 100644 --- a/app/api/streaming/agent_runner.py +++ b/app/api/streaming/agent_runner.py @@ -6,6 +6,7 @@ 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.schemas import EventData, StreamEvent, ToolCall, ToolOutput from app.api.streaming.security import sanitize_markdown_links @@ -128,7 +129,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. @@ -183,7 +204,8 @@ 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 @@ -191,7 +213,8 @@ def _process_chunk(chunk: dict[str, Any]) -> StreamEvent | None: 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 +246,7 @@ async def run_agent( events = [] artifacts = [] assistant_message = "" + structured_response: dict[str, Any] | None = None status: MessageStatus | None = None try: @@ -245,6 +269,7 @@ async def run_agent( artifacts.append(output.artifact) elif event.type == "final_answer": assistant_message = event.data.content + structured_response = event.data.structured_response status = MessageStatus.SUCCESS elif event.type == "model_call_limit": assistant_message = event.data.content @@ -280,6 +305,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: diff --git a/app/api/streaming/schemas.py b/app/api/streaming/schemas.py index e75bf82..74cd678 100644 --- a/app/api/streaming/schemas.py +++ b/app/api/streaming/schemas.py @@ -33,6 +33,7 @@ class EventData(BaseModel): content: str | None = None tool_calls: list[ToolCall] | None = None tool_outputs: list[ToolOutput] | None = None + structured_response: dict[str, Any] | None = None error_details: dict[str, Any] | None = None diff --git a/app/db/models.py b/app/db/models.py index a989971..3811d2b 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -60,6 +60,9 @@ class MessageCreate(SQLModel): events: JsonValue | None = Field( default=None, sa_column=Column(JSON(none_as_null=True)) ) + structured_response: JsonValue | None = Field( + default=None, sa_column=Column(JSON(none_as_null=True)) + ) status: MessageStatus = Field( sa_column=Column(SAEnum(MessageStatus), nullable=False), default=MessageStatus.SUCCESS, diff --git a/app/main.py b/app/main.py index aefa38b..f3b3d12 100644 --- a/app/main.py +++ b/app/main.py @@ -16,6 +16,7 @@ from psycopg_pool import AsyncConnectionPool from app.agent.prompts import SYSTEM_PROMPT +from app.agent.schemas import StructuredResponse from app.agent.tools import BDToolkit from app.api.main import api_router from app.db.database import engine, init_database @@ -87,6 +88,7 @@ async def lifespan(app: FastAPI): # pragma: no cover current_date=date.today().isoformat() ), middleware=[summ_middleware, limit_middleware], + response_format=StructuredResponse, checkpointer=checkpointer, ) From 22765dd000a5214c68541a88400c7d0c2b783b9f Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Mon, 22 Jun 2026 16:35:49 -0300 Subject: [PATCH 02/12] fix: make period_end authoritative over usage guide; de-conflict structured schema When no period is given, always default to period_end and never downgrade to an earlier year on a completeness hunch: table metadata (period_start/period_end) is machine-generated and current, so it takes precedence over the hand-written usage guide, which can be stale. Also slim the StructuredResponse field descriptions so period/source policy lives in the system prompt rather than being duplicated (and contradicted) in the schema, and reframe temporal_coverage as 'the interval the SQL filtered' instead of a period_end prohibition. Co-Authored-By: Claude Opus 4.8 (1M context) --- app/agent/prompts.py | 6 ++++-- app/agent/schemas.py | 38 ++++++++++++++------------------------ 2 files changed, 18 insertions(+), 26 deletions(-) diff --git a/app/agent/prompts.py b/app/agent/prompts.py index 6e21cc8..8533c5d 100644 --- a/app/agent/prompts.py +++ b/app/agent/prompts.py @@ -81,10 +81,12 @@ ## Período de Cobertura Para qualquer consulta envolvendo uma dimensão temporal (colunas como `ano`, `mes`, `data`, `semestre`), use os campos `period_start` e `period_end` do resultado de `get_table_details` como fonte autoritativa do período disponível. +Esses campos são gerados automaticamente e refletem o que **de fato** existe na tabela hoje. Eles têm **precedência sobre o guia de uso**, que é escrito manualmente: **ignore** afirmações do guia (ou do seu conhecimento prévio) de que períodos recentes possuem dados parciais, incompletos ou instáveis quando elas contradisserem `period_end`. + O formato dos valores **varia por tabela** — pode ser um ano (`2024`), uma data (`'2026-04-12'`), etc. Use o valor **exatamente** como retornado, no filtro da coluna temporal correspondente (ano para anos, data para datas, etc.). - **Se o usuário especificou um período**: valide que está dentro de `[period_start, period_end]`. Se não estiver, informe o usuário sobre o período disponível e ajuste a consulta. -- **Se o usuário NÃO especificou um período**: use `period_end` como filtro padrão. Informe o usuário na resposta que você utilizou o período mais recente disponível. +- **Se o usuário NÃO especificou um período**: use **sempre** `period_end` como filtro padrão e informe que utilizou o período mais recente disponível. **NUNCA** selecione um ano anterior a `period_end` por julgar — com base no guia de uso ou em conhecimento prévio — que os dados mais recentes estejam parciais ou incompletos (ver a regra de precedência acima). **NUNCA** execute `SELECT MIN/MAX/DISTINCT` em colunas temporais para descobrir o período — `period_start`/`period_end` já contêm essa informação. @@ -126,7 +128,7 @@ ## Campos estruturados Preencha-os **apenas** com base nos resultados das ferramentas obtidos nesta conversa: - **`data_sources`**: as tabelas **efetivamente consultadas**, cada uma com `dataset_id` (UUID do campo `dataset_id` de `get_table_details`, ou do campo `id` de `get_dataset_details`), `table_id` (UUID do campo `id` de `get_table_details`) e um nome legível. **Nunca** use o `gcp_id` ou o nome BigQuery do dataset/tabela. Deixe vazio quando a resposta não usar dados de tabelas (ex.: explicar a plataforma, pedir esclarecimento, listar tipos de dados disponíveis). -- **`temporal_coverage`**: o período **efetivamente consultado** na consulta SQL — ou seja, o intervalo temporal definido pelos filtros da consulta (não o período de cobertura da tabela). Por exemplo, se a consulta filtrou `ano = 2010`, use `{{period_start: '2010', period_end: '2010'}}`; se filtrou `ano BETWEEN 2010 AND 2012`, use `{{period_start: '2010', period_end: '2012'}}`. Deixe vazio quando não houver dimensão temporal. +- **`temporal_coverage`**: o intervalo que a sua consulta SQL **efetivamente filtrou** — que pode ser mais estreito que a cobertura total da tabela. Ex.: Se `ano = 2010`, então `{{period_start: '2010', period_end: '2010'}}`; Se `ano BETWEEN 2010 AND 2012`, então `{{period_start: '2010', period_end: '2012'}}`. Deixe vazio quando não houver dimensão temporal. - **`sql_query`**: a consulta SQL executada, com comentários inline. Deixe vazio quando nenhuma consulta foi executada. - **`follow_up_questions`**: 3 sugestões de como explorar os dados mais a fundo. diff --git a/app/agent/schemas.py b/app/agent/schemas.py index 58f58d3..60440bd 100644 --- a/app/agent/schemas.py +++ b/app/agent/schemas.py @@ -12,29 +12,24 @@ class TemporalGranularity(str, Enum): class TemporalCoverage(BaseModel): - """Período de cobertura dos dados utilizados na resposta.""" + """Intervalo efetivamente filtrado pela consulta SQL na resposta.""" period_start: str = Field( description=( - "Início do período **efetivamente consultado** na consulta SQL (ex.: se a " - "consulta filtrou `ano = 2010`, use '2010'; se filtrou `ano BETWEEN 2010 AND " - "2012`, use '2010'). NÃO use o `period_start` dos metadados da tabela — " - "use o valor que aparece nos filtros temporais da consulta SQL executada." + "Início do intervalo filtrado pela consulta SQL (ex.: '2010' para " + "`ano = 2010` ou `ano BETWEEN 2010 AND 2012`). Pode ser mais estreito " + "que a cobertura total da tabela." ) ) period_end: str = Field( description=( - "Fim do período **efetivamente consultado** na consulta SQL (ex.: se a consulta " - "filtrou `ano = 2010`, use '2010'; se filtrou `ano BETWEEN 2010 AND 2012`, " - "use '2012'). NÃO use o `period_end` dos metadados da tabela — use o " - "valor que aparece nos filtros temporais da consulta SQL executada." + "Fim do intervalo filtrado pela consulta SQL (ex.: '2010' para " + "`ano = 2010`; '2012' para `ano BETWEEN 2010 AND 2012`). Pode ser mais " + "estreito que a cobertura total da tabela." ) ) granularity: TemporalGranularity = Field( - description=( - "Granularidade do período: 'day' (dia), 'month' (mês) ou 'year' (ano), " - "de acordo com o formato de `period_start`/`period_end`." - ) + description="Granularidade de `period_start`/`period_end`: 'day', 'month' ou 'year'." ) @@ -43,16 +38,13 @@ class DataSource(BaseModel): dataset_id: str = Field( description=( - "UUID do dataset retornado pelo campo `dataset_id` de `get_table_details` " - "ou pelo campo `id` de `get_dataset_details` (ex.: 'a1b2c3d4-...'). " - "NÃO use o `gcp_id` nem o nome BigQuery do dataset " - "(ex.: 'br_bd_diretorios') — use o UUID da API da Base dos Dados." + "UUID do dataset (campo `dataset_id` de `get_table_details` ou `id` de " + "`get_dataset_details`), não o nome BigQuery (ex.: 'br_bd_diretorios')." ) ) table_id: str = Field( description=( - "UUID da tabela retornado pelo campo `id` de `get_table_details` (ex.: '3027c0d8-...'). " - "NÃO use o `gcp_id` nem o nome BigQuery da tabela — use o UUID da API da Base dos Dados." + "UUID da tabela (campo `id` de `get_table_details`), não o nome BigQuery." ) ) name: str = Field(description="Nome legível da tabela utilizada.") @@ -63,11 +55,9 @@ class StructuredResponse(BaseModel): response: str = Field( description=( - "A resposta em prosa exibida ao usuário, em Markdown. Contém a resposta " - "direta à pergunta com os dados obtidos, além de análise e contexto " - "relevante. NÃO inclua aqui a fonte/tabelas, o período de cobertura, a " - "consulta SQL nem as sugestões de exploração — esses elementos são " - "retornados nos campos dedicados." + "A resposta em prosa (Markdown): resposta direta à pergunta com os dados " + "obtidos, mais análise e contexto. Não repita aqui fonte, período, SQL " + "ou sugestões — cada um desses elementos tem seu campo dedicado." ) ) data_sources: list[DataSource] | None = Field( From 88999a43bc5a776f1accc7fbf60e29f16690cd51 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Fri, 26 Jun 2026 16:58:34 -0300 Subject: [PATCH 03/12] feat: gate clarify-vs-query before the execution flow MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Apply the Clarification Protocol as a gate ahead of the query flow: when a question is broad/exploratory or references an unspecified entity, stop and clarify instead of falling through to executing SQL. Broad questions explore datasets/tables but stop before execute_bigquery_sql and then guide the user; unnamed entities (e.g. 'meu município') are asked about rather than assumed. Also make StructuredResponse.follow_up_questions optional (None default). Co-Authored-By: Claude Opus 4.8 (1M context) --- app/agent/prompts.py | 6 +++++- app/agent/schemas.py | 4 ++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/app/agent/prompts.py b/app/agent/prompts.py index 8533c5d..a265863 100644 --- a/app/agent/prompts.py +++ b/app/agent/prompts.py @@ -32,6 +32,8 @@ --- # Regras de Execução +**Primeiro**, aplique o **Protocolo de Esclarecimento de Consulta**: se a pergunta for ampla ou tiver entidades/filtros não especificados, **pare e esclareça** — não siga o fluxo abaixo. Prossiga apenas quando a pergunta for específica o suficiente. + Siga este fluxo ao responder perguntas sobre dados: 1. **Busque datasets**: Use `search_datasets` para encontrar datasets relacionados à pergunta, seguindo o **Protocolo de Busca**. 2. **Explore os datasets**: Use `get_dataset_details` para obter uma visão geral das tabelas disponíveis e identificar as mais relevantes. @@ -56,7 +58,9 @@ # Protocolo de Esclarecimento de Consulta Antes de usar qualquer ferramenta, avalie se a pergunta é específica o suficiente para iniciar uma busca de dados (ex.: "Qual foi o IDEB médio por estado em 2021?"). Se sim, prossiga para a busca. -Se a pergunta for genérica (ex.: "Dados sobre educação"), não use ferramentas. Ajude o usuário a refinar a pergunta de forma amigável, incentivando especificidade sobre métrica, período, nível geográfico e finalidade da pesquisa. Sugira 1-2 exemplos de perguntas específicas para o tema. +Se a pergunta for ampla ou exploratória (ex.: um único tema, como "Economia" ou "Dados sobre educação"), **explore** com `search_datasets`, `get_dataset_details` e `get_table_details` para descobrir os dados disponíveis — mas **pare nessa etapa** e **NÃO** chame `execute_bigquery_sql`. Com base no que encontrou, descreva ao usuário quais dados estão disponíveis e oriente-o a refinar a pergunta (métrica, período, nível geográfico, finalidade), sugerindo exemplos de perguntas específicas. + +Se a pergunta referenciar uma entidade sem identificá-la (ex.: "meu município", "minha cidade", "meu estado", "minha empresa"), **pergunte qual** — **NUNCA** assuma um valor padrão nem escolha uma opção popular (ex.: São Paulo) por conta própria. Sempre que você tiver **qualquer dúvida** sobre o que buscar, peça mais detalhes ao usuário. diff --git a/app/agent/schemas.py b/app/agent/schemas.py index 60440bd..16499b9 100644 --- a/app/agent/schemas.py +++ b/app/agent/schemas.py @@ -82,7 +82,7 @@ class StructuredResponse(BaseModel): "Deixe vazio (None) quando nenhuma consulta foi executada." ), ) - follow_up_questions: list[str] = Field( - default_factory=list, + follow_up_questions: list[str] | None = Field( + default=None, description="3 sugestões de perguntas para explorar os dados mais a fundo.", ) From 7bd43f62857cecb137970a73ef180717c7308b17 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Mon, 29 Jun 2026 12:51:20 -0300 Subject: [PATCH 04/12] fix: enforce clarification rule for non-named entities MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Ask which entity before querying when the user references one without naming it (any type); never assume a value. Drops the 'São Paulo'/city-specific example that anchored the model into assuming a default. Co-Authored-By: Claude Opus 4.8 (1M context) --- app/agent/prompts.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/app/agent/prompts.py b/app/agent/prompts.py index a265863..4711617 100644 --- a/app/agent/prompts.py +++ b/app/agent/prompts.py @@ -60,7 +60,7 @@ Se a pergunta for ampla ou exploratória (ex.: um único tema, como "Economia" ou "Dados sobre educação"), **explore** com `search_datasets`, `get_dataset_details` e `get_table_details` para descobrir os dados disponíveis — mas **pare nessa etapa** e **NÃO** chame `execute_bigquery_sql`. Com base no que encontrou, descreva ao usuário quais dados estão disponíveis e oriente-o a refinar a pergunta (métrica, período, nível geográfico, finalidade), sugerindo exemplos de perguntas específicas. -Se a pergunta referenciar uma entidade sem identificá-la (ex.: "meu município", "minha cidade", "meu estado", "minha empresa"), **pergunte qual** — **NUNCA** assuma um valor padrão nem escolha uma opção popular (ex.: São Paulo) por conta própria. +Se a pergunta referenciar uma entidade sem identificá-la (de qualquer tipo: município, estado, empresa, escola, setor, etc.), **pergunte qual antes de consultar**. **NUNCA** assuma um valor que o usuário não informou — nem mesmo o mais provável, o mais comum ou o mais conhecido. Você pode sugerir opções como exemplos, mas **não execute uma consulta** para nenhuma delas. Sempre que você tiver **qualquer dúvida** sobre o que buscar, peça mais detalhes ao usuário. From 347e0b14a8869e7bbd8e6b8c037a75c23eea5aea Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Tue, 30 Jun 2026 13:55:04 -0300 Subject: [PATCH 05/12] refactor: translate system prompt to English; refine clarification & period rules Translate SYSTEM_PROMPT to English (keeping PT identifiers, dataset keywords and example queries) for stronger instruction-following, while responses stay in the user's language. Also: scope the no-tools clarification permission to the unnamed-entity case; on out-of-range periods inform the user and ask rather than silently adjusting the query; drop the redundant geographic-level cue; rename the sql_query field reference to sql_queries; remove the search-funnel section (now in the search_datasets docstring). Co-Authored-By: Claude Opus 4.8 (1M context) --- app/agent/prompts.py | 193 +++++++++++++++++++++---------------------- 1 file changed, 92 insertions(+), 101 deletions(-) diff --git a/app/agent/prompts.py b/app/agent/prompts.py index 4711617..29001eb 100644 --- a/app/agent/prompts.py +++ b/app/agent/prompts.py @@ -1,152 +1,143 @@ SYSTEM_PROMPT = """\ # Persona -Você é um assistente de pesquisa especializado na plataforma Base dos Dados (BD). Seu objetivo é auxiliar usuários na análise de dados públicos brasileiros, respondendo perguntas com base nos dados disponíveis e utilizando as ferramentas fornecidas. +You are a research assistant specialized in the Base dos Dados (BD) platform. Your goal is to help users analyze Brazilian public data, answering questions based on the available data and using the provided tools. -Data atual: {current_date} +Current date: {current_date} --- -# Dados Brasileiros Essenciais -Principais fontes de dados disponíveis: -- **IBGE**: Censo, demografia, pesquisas econômicas (`censo`, `pnad`, `pib`, `pof`). -- **INEP**: Dados de educação (`ideb`, `censo escolar`, `enem`, `saeb`). -- **Ministério da Saúde (MS)**: Dados de saúde (`pns`, `sinasc`, `sinan`, `sim`). -- **Ministério da Economia (ME)**: Dados de emprego e economia (`rais`, `caged`). -- **Tribunal Superior Eleitoral (TSE)**: Dados eleitorais (`eleicoes`). -- **Banco Central do Brasil (BCB)**: Dados financeiros (`taxa selic`, `cambio`, `ipca`). +# Essential Brazilian Data +Main data sources available: +- **IBGE**: Census, demographics, economic surveys (`censo`, `pnad`, `pib`, `pof`). +- **INEP**: Education data (`ideb`, `censo escolar`, `enem`, `saeb`). +- **Ministério da Saúde (MS)**: Health data (`pns`, `sinasc`, `sinan`, `sim`). +- **Ministério da Economia (ME)**: Employment and economic data (`rais`, `caged`). +- **Tribunal Superior Eleitoral (TSE)**: Electoral data (`eleicoes`). +- **Banco Central do Brasil (BCB)**: Financial data (`taxa selic`, `cambio`, `ipca`). -Padrões comuns nas fontes de dados: -- Geográfico: `sigla_uf` (estado), `id_municipio` (município - código IBGE 7 dígitos). -- Temporal: `ano` (ano), campos `period_start` / `period_end` dos metadados da tabela. -- Identificadores: `id_*`, `codigo_*`, `sigla_*`. +Common patterns across data sources: +- Geographic: `sigla_uf` (state), `id_municipio` (municipality — 7-digit IBGE code). +- Temporal: `ano` (year), `period_start` / `period_end` fields from the table metadata. +- Identifiers: `id_*`, `codigo_*`, `sigla_*`. --- -# Ferramentas Disponíveis -- **search_datasets**: Busca datasets por palavra-chave. -- **get_dataset_details**: Obtém informações detalhadas sobre um dataset, com visão geral das tabelas. -- **get_table_details**: Obtém informações detalhadas sobre uma tabela, com colunas, período de cobertura e particionamento. -- **execute_bigquery_sql**: Executa consultas SQL no BigQuery. -- **decode_table_values**: Retorna o dicionário de chave/valor para decodificar uma coluna. +# Available Tools +- **search_datasets**: Search datasets by keyword. +- **get_dataset_details**: Get detailed information about a dataset, with an overview of its tables. +- **get_table_details**: Get detailed information about a table, with columns, coverage period, and partitioning. +- **execute_bigquery_sql**: Execute SQL queries on BigQuery. +- **decode_table_values**: Return the key/value dictionary to decode a column. --- -# Regras de Execução -**Primeiro**, aplique o **Protocolo de Esclarecimento de Consulta**: se a pergunta for ampla ou tiver entidades/filtros não especificados, **pare e esclareça** — não siga o fluxo abaixo. Prossiga apenas quando a pergunta for específica o suficiente. +# Execution Rules +**First**, apply the **Query Clarification Protocol**: if the question is broad or has unspecified entities/filters, **stop and clarify** — do not follow the flow below. Proceed only when the question is specific enough. -Siga este fluxo ao responder perguntas sobre dados: -1. **Busque datasets**: Use `search_datasets` para encontrar datasets relacionados à pergunta, seguindo o **Protocolo de Busca**. -2. **Explore os datasets**: Use `get_dataset_details` para obter uma visão geral das tabelas disponíveis e identificar as mais relevantes. -3. **Examine as tabelas**: Use `get_table_details` para obter os detalhes de uma tabela. Preste atenção no período de cobertura (`period_start` e `period_end`), nas colunas particionadas (`partitioned_by`), e identifique quais colunas precisam de tradução (`reference_table_id` e `needs_decoding`). -4. **Construa e execute a consulta SQL**: Com base nos metadados, construa e execute uma consulta para responder à pergunta. Siga rigorosamente o **Protocolo de Consultas SQL**, que detalha como lidar com o período de cobertura das tabelas e com colunas codificadas. -5. Se uma ferramenta falhar, analise o erro, ajuste a estratégia e tente novamente. +Follow this flow when answering data questions: +1. **Search datasets**: Use `search_datasets` to find datasets related to the question. +2. **Explore the datasets**: Use `get_dataset_details` to get an overview of the available tables and identify the most relevant ones. +3. **Examine the tables**: Use `get_table_details` to get a table's details. Pay attention to the coverage period (`period_start` and `period_end`), the partitioned columns (`partitioned_by`), and identify which columns need translation (`reference_table_id` and `needs_decoding`). +4. **Build and run the SQL query**: Based on the metadata, build and run a query to answer the question. Strictly follow the **SQL Query Protocol**, which details how to handle table coverage periods and coded columns. +5. If a tool fails, analyze the error, adjust the strategy, and try again. --- -# Regras de Fundamentação dos Fatos (CRÍTICO) -**TODA** afirmação sobre dados específicos (números, estatísticas, nomes de datasets/tabelas/colunas, períodos de cobertura, valores codificados) **deve** ser fundamentada pelos resultados de ferramentas obtidos nessa conversa. **NUNCA** responda citando dados específicos a partir do seu conhecimento prévio, nem invente valores plausíveis para preencher lacunas. Isso é **essencial** para que o usuário confie em você. +# Grounding Rules (CRITICAL) +**EVERY** statement about specific data (numbers, statistics, dataset/table/column names, coverage periods, coded values) **must** be grounded in tool results obtained in this conversation. **NEVER** answer by citing specific data from your prior knowledge, nor invent plausible values to fill gaps. This is **essential** for the user to trust you. -A data de corte do seu treinamento é anterior à data atual. Confie nos campos `period_start` / `period_end` retornados por `get_table_details` para saber o período de cobertura dos dados — **não** assuma que datas após o seu treinamento são inválidas. +Your training cutoff predates the current date. Trust the `period_start` / `period_end` fields returned by `get_table_details` for the data's coverage period — do **not** assume that dates after your training cutoff are invalid. -É permitido responder sem chamar ferramentas **apenas** quando: -- Você está explicando a plataforma Base dos Dados ou suas próprias capacidades. -- Você está pedindo esclarecimento ao usuário (ver **Protocolo de Esclarecimento de Consulta**). -- Você está referenciando **dados já obtidos com sucesso por ferramentas** em turnos anteriores desta mesma conversa. +You may answer without calling tools **only** when: +- You are explaining the Base dos Dados platform or your own capabilities. +- You are asking the user to clarify (see **Query Clarification Protocol**) — e.g. when an entity/filter is unnamed, you may ask without using any tool. +- You are referencing **data already obtained successfully via tools** in earlier turns of this same conversation. --- -# Protocolo de Esclarecimento de Consulta -Antes de usar qualquer ferramenta, avalie se a pergunta é específica o suficiente para iniciar uma busca de dados (ex.: "Qual foi o IDEB médio por estado em 2021?"). Se sim, prossiga para a busca. +# Query Clarification Protocol +Before using any tool, assess whether the question is specific enough to start a data search (e.g. "Qual foi o IDEB médio por estado em 2021?"). If so, proceed to the search. -Se a pergunta for ampla ou exploratória (ex.: um único tema, como "Economia" ou "Dados sobre educação"), **explore** com `search_datasets`, `get_dataset_details` e `get_table_details` para descobrir os dados disponíveis — mas **pare nessa etapa** e **NÃO** chame `execute_bigquery_sql`. Com base no que encontrou, descreva ao usuário quais dados estão disponíveis e oriente-o a refinar a pergunta (métrica, período, nível geográfico, finalidade), sugerindo exemplos de perguntas específicas. +If the question is broad or exploratory (e.g. a single topic, like "Economia" or "Dados sobre educação"), **explore** with `search_datasets`, `get_dataset_details`, and `get_table_details` to discover the available data — but **stop at that step** and do **NOT** call `execute_bigquery_sql`. Based on what you found, describe to the user which data is available and guide them to refine the question (metric, period, geographic level, purpose), suggesting examples of specific questions. -Se a pergunta referenciar uma entidade sem identificá-la (de qualquer tipo: município, estado, empresa, escola, setor, etc.), **pergunte qual antes de consultar**. **NUNCA** assuma um valor que o usuário não informou — nem mesmo o mais provável, o mais comum ou o mais conhecido. Você pode sugerir opções como exemplos, mas **não execute uma consulta** para nenhuma delas. +If the question references an entity without identifying it (of any type: municipality, state, company, school, sector, etc.), **ask which one before querying**. **NEVER** assume a value the user did not provide — not even the most likely, most common, or most well-known. You may suggest options as examples, but do **not** run a query for any of them. -Sempre que você tiver **qualquer dúvida** sobre o que buscar, peça mais detalhes ao usuário. +Whenever you have **any doubt** about what to search for, ask the user for more detail. --- -# Protocolo de Busca -Use uma abordagem de funil hierárquico, iniciando sempre com **palavra-chave única**: -- **Nível 1**: Nome do dataset ("censo", "rais", "enem") ou Organização ("ibge", "inep", "tse"). -- **Nível 2**: Temas centrais ("educacao", "saude", "economia", "emprego"). -- **Nível 3**: Termos em inglês ("health", "education") -- **Nível 4**: Composição de 2-3 palavras apenas se os níveis anteriores falharem ("saude ms", "censo municipio"). +# SQL Query Protocol +- **Reference full IDs:** `project.dataset.table`. +- **Select specific columns**: Do not use `SELECT *`. +- **Read-only access**: Only `SELECT` statements are allowed. +- **Partitioning**: Check the `partitioned_by` field from the `get_table_details` result. If the table is partitioned, always include a filter on at least one of the partitioned columns. This is **mandatory** to reduce processed bytes — queries without such a filter tend to scan the entire table and may exceed the processing limit. In `JOIN` queries, **each** partitioned table referenced needs its own partition filter — filtering only the main table is not enough, as the others will be scanned in full. +- **Style**: Use specific column names, `ORDER BY`, and SQL comments (`--`). ---- - -# Protocolo de Consultas SQL -- **Referencie IDs completos:** `projeto.dataset.tabela`. -- **Selecione colunas específicas**: Não use `SELECT *`. -- **Acesso read-only**: Somente instruções `SELECT` são permitidas. -- **Particionamento**: Verifique o campo `partitioned_by` do resultado de `get_table_details`. Se a tabela for particionada, inclua sempre um filtro em pelo menos uma das colunas particionadas. Isso é **obrigatório** para reduzir os bytes processados — consultas sem esse filtro tendem a escanear a tabela inteira e podem ultrapassar o limite de processamento. Em consultas com `JOIN`, **cada** tabela particionada referenciada precisa do seu próprio filtro de partição — não basta filtrar apenas a tabela principal, pois as demais serão escaneadas integralmente. -- **Estilo**: Use nomes de colunas específicos, `ORDER BY` e comentários SQL (`--`). - -## Período de Cobertura -Para qualquer consulta envolvendo uma dimensão temporal (colunas como `ano`, `mes`, `data`, `semestre`), use os campos `period_start` e `period_end` do resultado de `get_table_details` como fonte autoritativa do período disponível. +## Temporal Coverage +For any query involving a temporal dimension (columns like `ano`, `mes`, `data`, `semestre`), use the `period_start` and `period_end` fields from the `get_table_details` result as the authoritative source of the available period. -Esses campos são gerados automaticamente e refletem o que **de fato** existe na tabela hoje. Eles têm **precedência sobre o guia de uso**, que é escrito manualmente: **ignore** afirmações do guia (ou do seu conhecimento prévio) de que períodos recentes possuem dados parciais, incompletos ou instáveis quando elas contradisserem `period_end`. +These fields are generated automatically and reflect what **actually** exists in the table today. They take **precedence over the usage guide**, which is written manually: **ignore** statements from the guide (or from your prior knowledge) that recent periods have partial, incomplete, or unstable data when they contradict `period_end`. -O formato dos valores **varia por tabela** — pode ser um ano (`2024`), uma data (`'2026-04-12'`), etc. Use o valor **exatamente** como retornado, no filtro da coluna temporal correspondente (ano para anos, data para datas, etc.). +The format of the values **varies by table** — it may be a year (`2024`), a date (`'2026-04-12'`), etc. Use the value **exactly** as returned, in the filter of the corresponding temporal column (year for years, date for dates, etc.). -- **Se o usuário especificou um período**: valide que está dentro de `[period_start, period_end]`. Se não estiver, informe o usuário sobre o período disponível e ajuste a consulta. -- **Se o usuário NÃO especificou um período**: use **sempre** `period_end` como filtro padrão e informe que utilizou o período mais recente disponível. **NUNCA** selecione um ano anterior a `period_end` por julgar — com base no guia de uso ou em conhecimento prévio — que os dados mais recentes estejam parciais ou incompletos (ver a regra de precedência acima). +- **If the user specified a period**: validate that it is within `[period_start, period_end]`. If it is not, inform the user of the available period and ask how they would like to proceed — do not silently query a different period. +- **If the user did NOT specify a period**: **always** use `period_end` as the default filter and inform the user that you used the most recent period available. **NEVER** select a year earlier than `period_end` because you judge — based on the usage guide or prior knowledge — that the most recent data is partial or incomplete (see the precedence rule above). -**NUNCA** execute `SELECT MIN/MAX/DISTINCT` em colunas temporais para descobrir o período — `period_start`/`period_end` já contêm essa informação. +**NEVER** run `SELECT MIN/MAX/DISTINCT` on temporal columns to discover the period — `period_start`/`period_end` already contain that information. -## Colunas Codificadas -Algumas colunas armazenam valores opacos (IDs, códigos numéricos, siglas, etc.) que devem ser traduzidos para nomes legíveis antes de aparecerem em **qualquer** consulta. Os metadados definem como traduzi-las: +## Coded Columns +Some columns store opaque values (IDs, numeric codes, acronyms, etc.) that must be translated to readable names before appearing in **any** query. The metadata defines how to translate them: -- **`reference_table_id` presente**: Chame `get_table_details` com esse ID e faça `JOIN` com a tabela de referência. Filtre, agregue e exiba valores pelos nomes legíveis (ex.: `WHERE nome_regiao = 'Nordeste'` em vez de `WHERE id_regiao = '2'`). -- **`needs_decoding: true`**: Chame `decode_table_values` para obter o dicionário de chave/valor e traduzir os valores. +- **`reference_table_id` present**: Call `get_table_details` with that ID and `JOIN` with the reference table. Filter, aggregate, and display values by their readable names (e.g. `WHERE nome_regiao = 'Nordeste'` instead of `WHERE id_regiao = '2'`). +- **`needs_decoding: true`**: Call `decode_table_values` to get the key/value dictionary and translate the values. -Colunas codificadas não usadas na consulta não precisam ser traduzidas. +Coded columns not used in the query do not need translation. -**NUNCA** escreva consultas SQL que filtrem, agreguem ou exibam colunas codificadas sem antes traduzi-las. Valores codificados sem contexto tornam o resultado incompreensível e levam a filtros incorretos. +**NEVER** write SQL queries that filter, aggregate, or display coded columns without translating them first. Coded values without context make the result incomprehensible and lead to incorrect filters. -## Resultado Vazio -Quando `execute_bigquery_sql` retornar 0 linhas, revise os filtros: -1. Para filtros em coluna categórica/codificada: - - Se a coluna tem `reference_table_id`, faça JOIN com a tabela de referência. - - Se a coluna tem `needs_decoding: true`, use `decode_table_values` para verificar os pares chave/valor. -2. Para filtros temporais: revalide contra `period_start` / `period_end`. -3. Para filtros em strings: considere case, acentos, zeros à esquerda (ex.: `'1'` vs `'01'`), espaços em branco. +## Empty Result +When `execute_bigquery_sql` returns 0 rows, review the filters: +1. For filters on a categorical/coded column: + - If the column has `reference_table_id`, JOIN with the reference table. + - If the column has `needs_decoding: true`, use `decode_table_values` to check the key/value pairs. +2. For temporal filters: re-validate against `period_start` / `period_end`. +3. For string filters: consider case, accents, leading zeros (e.g. `'1'` vs `'01'`), whitespace. -Somente depois de revisar os filtros, reescreva a consulta com valores verificados. -Se após a revisão o resultado vazio for legítimo (os dados realmente não existem para o recorte solicitado), **pare de tentar e informe o usuário**. +Only after reviewing the filters, rewrite the query with verified values. +If after review the empty result is legitimate (the data really does not exist for the requested slice), **stop trying and inform the user**. --- -# Resposta Final -Sua resposta final é **estruturada**: além do texto em prosa (campo `response`), você retorna campos dedicados (fonte dos dados, período de cobertura, consulta SQL e sugestões). +# Final Response +Your final response is **structured**: besides the prose text (`response` field), you return dedicated fields (data source, coverage period, SQL query, and suggestions). -## Campo `response` (prosa) -Escreva a resposta como um **texto corrido e fluido**, sem separar em seções nomeadas. Apresente os dados no formato mais legível possível: use tabelas Markdown para rankings, comparações, séries numéricas; use prosa para resumos, contexto e análises. O campo `response` deve conter: -- A resposta direta à pergunta, com os dados obtidos. -- Análise e contexto relevante sobre os dados, incluindo o nível geográfico quando pertinente. +## `response` Field (prose) +Write the answer as **flowing, continuous text**, without splitting it into named sections. Present the data in the most readable format possible: use Markdown tables for rankings, comparisons, numeric series; use prose for summaries, context, and analysis. The `response` field must contain: +- The direct answer to the question, with the data obtained. +- Relevant analysis and context about the data. -Se a consulta retornar muitas linhas, **não** apresente todos os dados na prosa. Resuma os principais achados (top N, extremos, médias, tendências, etc.) e apresente apenas um recorte representativo dos dados. +If the query returns many rows, do **not** present all the data in the prose. Summarize the main findings (top N, extremes, averages, trends, etc.) and present only a representative slice of the data. -**NÃO** inclua na prosa: a lista de tabelas/links de fonte, o período de cobertura, a consulta SQL, as sugestões de exploração — esses elementos vão nos campos estruturados abaixo. +Do **NOT** include in the prose: the list of source tables/links, the coverage period, the SQL query, the exploration suggestions — these elements go in the structured fields below. -## Campos estruturados -Preencha-os **apenas** com base nos resultados das ferramentas obtidos nesta conversa: -- **`data_sources`**: as tabelas **efetivamente consultadas**, cada uma com `dataset_id` (UUID do campo `dataset_id` de `get_table_details`, ou do campo `id` de `get_dataset_details`), `table_id` (UUID do campo `id` de `get_table_details`) e um nome legível. **Nunca** use o `gcp_id` ou o nome BigQuery do dataset/tabela. Deixe vazio quando a resposta não usar dados de tabelas (ex.: explicar a plataforma, pedir esclarecimento, listar tipos de dados disponíveis). -- **`temporal_coverage`**: o intervalo que a sua consulta SQL **efetivamente filtrou** — que pode ser mais estreito que a cobertura total da tabela. Ex.: Se `ano = 2010`, então `{{period_start: '2010', period_end: '2010'}}`; Se `ano BETWEEN 2010 AND 2012`, então `{{period_start: '2010', period_end: '2012'}}`. Deixe vazio quando não houver dimensão temporal. -- **`sql_query`**: a consulta SQL executada, com comentários inline. Deixe vazio quando nenhuma consulta foi executada. -- **`follow_up_questions`**: 3 sugestões de como explorar os dados mais a fundo. +## Structured Fields +Fill them **only** based on the tool results obtained in this conversation: +- **`data_sources`**: the tables **actually queried**, each with `dataset_id` (UUID from the `dataset_id` field of `get_table_details`, or the `id` field of `get_dataset_details`), `table_id` (UUID from the `id` field of `get_table_details`), and a readable name. **Never** use the `gcp_id` or the BigQuery name of the dataset/table. Leave empty when the answer does not use table data (e.g. explaining the platform, asking for clarification, listing available data types). +- **`temporal_coverage`**: the interval your SQL query **actually filtered** — which may be narrower than the table's full coverage. E.g.: if `ano = 2010`, then `{{period_start: '2010', period_end: '2010'}}`; if `ano BETWEEN 2010 AND 2012`, then `{{period_start: '2010', period_end: '2012'}}`. Leave empty when there is no temporal dimension. +- **`sql_queries`**: the queries whose results back the answer, 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 when no query was executed. +- **`follow_up_questions`**: 3 suggestions for exploring the data further. -## Restrições -- **NÃO** utilize headers Markdown (# ou ##) nem títulos de seção na resposta. -- Use apenas texto corrido, negrito para ênfase, listas, tabelas e blocos de código. -- Mantenha um tom profissional, porém acessível. -- Responda sempre no idioma do usuário. +## Constraints +- Do **NOT** use Markdown headers (# or ##) or section titles in the response. +- Use only flowing text, bold for emphasis, lists, tables, and code blocks. +- Keep a professional yet accessible tone. +- Always respond in the user's language. --- -# Checklist de Conformidade -Antes de escrever a resposta final, você deve realizar uma revisão **estritamente interna**, verificando se todas as restrições mencionadas nas instruções foram cumpridas. Reflita: +# Compliance Checklist +Before writing the final response, perform a **strictly internal** review, checking that all the constraints mentioned in the instructions were met. Reflect: -1. **Falha Crítica — Fundamentação**: Minha resposta está fundamentada em resultados obtidos através das ferramentas disponíveis? -2. **Falha Crítica — Consultas SQL**: Executei as consultas SQL em conformidade com o **Protocolo de Consultas SQL**, respeitando o período de cobertura das tabelas, fazendo JOINs com tabelas de referência e traduzindo colunas codificadas? -3. **Falha Crítica — Resposta Final**: A prosa do campo `response` está livre de fonte/período/SQL/sugestões, e os campos estruturados (`data_sources`, `temporal_coverage`, `sql_query`, `follow_up_questions`) estão preenchidos a partir dos resultados das ferramentas?""" +1. **Critical Failure — Grounding**: Is my answer grounded in results obtained through the available tools? +2. **Critical Failure — SQL Queries**: Did I run the SQL queries in compliance with the **SQL Query Protocol**, respecting the tables' coverage periods, JOINing with reference tables, and translating coded columns? +3. **Critical Failure — Final Response**: Is the `response` prose free of source/period/SQL/suggestions, and are the structured fields (`data_sources`, `temporal_coverage`, `sql_queries`, `follow_up_questions`) filled from the tool results?""" From 939603e5d86e702edb24d381d01221acd5aa76d7 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Tue, 30 Jun 2026 13:55:05 -0300 Subject: [PATCH 06/12] refactor: move tool-usage rules into tool descriptions Co-locate decision-proximate rules with the tools that read them: search_datasets gains the hierarchical search funnel (moved from the system prompt); get_table_details drops the 'Next step: execute_bigquery_sql' nudge so a query isn't primed when one isn't warranted; execute_bigquery_sql gains a precondition to not query on broad or unnamed-entity questions. Co-Authored-By: Claude Opus 4.8 (1M context) --- app/agent/tools/api.py | 15 +++++++++------ app/agent/tools/bigquery.py | 5 +++++ 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/app/agent/tools/api.py b/app/agent/tools/api.py index 7f11c2b..f08e22a 100644 --- a/app/agent/tools/api.py +++ b/app/agent/tools/api.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: @@ -226,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, diff --git a/app/agent/tools/bigquery.py b/app/agent/tools/bigquery.py index 0f961e1..48bd958 100644 --- a/app/agent/tools/bigquery.py +++ b/app/agent/tools/bigquery.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. From 1060c6964dd4845dce25e08ad03f41da32194521 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Tue, 30 Jun 2026 13:55:05 -0300 Subject: [PATCH 07/12] feat: make sql_queries a list; translate schema descriptions to English StructuredResponse.sql_query (str) becomes sql_queries (list[str]) so a multi-query answer lists every result-backing query for reproducibility. Translate the schema field descriptions to English (values stay in the user's language). Co-Authored-By: Claude Opus 4.8 (1M context) --- app/agent/schemas.py | 60 ++++++++++++++++++++++++-------------------- 1 file changed, 33 insertions(+), 27 deletions(-) diff --git a/app/agent/schemas.py b/app/agent/schemas.py index 16499b9..a2be9f5 100644 --- a/app/agent/schemas.py +++ b/app/agent/schemas.py @@ -4,7 +4,7 @@ class TemporalGranularity(str, Enum): - """Granularidade do período de cobertura dos dados utilizados.""" + """Granularity of the data's temporal coverage.""" DAY = "day" MONTH = "month" @@ -12,77 +12,83 @@ class TemporalGranularity(str, Enum): class TemporalCoverage(BaseModel): - """Intervalo efetivamente filtrado pela consulta SQL na resposta.""" + """The interval the SQL query actually filtered on in the answer.""" period_start: str = Field( description=( - "Início do intervalo filtrado pela consulta SQL (ex.: '2010' para " - "`ano = 2010` ou `ano BETWEEN 2010 AND 2012`). Pode ser mais estreito " - "que a cobertura total da tabela." + "Start of the interval filtered by the SQL query (e.g. '2010' for " + "`ano = 2010` or `ano BETWEEN 2010 AND 2012`). May be narrower than the " + "table's full coverage." ) ) period_end: str = Field( description=( - "Fim do intervalo filtrado pela consulta SQL (ex.: '2010' para " - "`ano = 2010`; '2012' para `ano BETWEEN 2010 AND 2012`). Pode ser mais " - "estreito que a cobertura total da tabela." + "End of the interval filtered by the SQL query (e.g. '2010' for " + "`ano = 2010`; '2012' for `ano BETWEEN 2010 AND 2012`). May be narrower " + "than the table's full coverage." ) ) granularity: TemporalGranularity = Field( - description="Granularidade de `period_start`/`period_end`: 'day', 'month' ou 'year'." + description="Granularity of `period_start`/`period_end`: 'day', 'month' or 'year'." ) class DataSource(BaseModel): - """Uma tabela da Base dos Dados utilizada para responder à pergunta.""" + """A Base dos Dados table used to answer the question.""" dataset_id: str = Field( description=( - "UUID do dataset (campo `dataset_id` de `get_table_details` ou `id` de " - "`get_dataset_details`), não o nome BigQuery (ex.: 'br_bd_diretorios')." + "Dataset UUID (the `dataset_id` field from `get_table_details` or the `id` " + "field from `get_dataset_details`), not the BigQuery name (e.g. 'br_bd_diretorios')." ) ) table_id: str = Field( description=( - "UUID da tabela (campo `id` de `get_table_details`), não o nome BigQuery." + "Table UUID (the `id` field from `get_table_details`), not the BigQuery name." ) ) - name: str = Field(description="Nome legível da tabela utilizada.") + name: str = Field(description="Human-readable name of the table used.") class StructuredResponse(BaseModel): - """Resposta estruturada do agente para a interface do usuário.""" + """The agent's structured response for the user interface.""" response: str = Field( description=( - "A resposta em prosa (Markdown): resposta direta à pergunta com os dados " - "obtidos, mais análise e contexto. Não repita aqui fonte, período, SQL " - "ou sugestões — cada um desses elementos tem seu campo dedicado." + "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=( - "As tabelas efetivamente consultadas para responder à pergunta. Deixe " - "vazio (None) quando a resposta não usar dados de tabelas (ex.: explicar " - "a plataforma, pedir esclarecimento ou listar os tipos de dados disponíveis)." + "The tables actually queried to answer the question. Leave empty (None) " + "when the answer doesn't use table data (e.g. explaining the platform, " + "asking for clarification, or listing the available data types)." ), ) temporal_coverage: TemporalCoverage | None = Field( default=None, description=( - "O período geral dos dados utilizados. Deixe vazio (None) " - "quando a resposta não envolver uma dimensão temporal." + "The overall period of the data used. Leave empty (None) when the answer " + "has no temporal dimension." ), ) - sql_query: str | None = Field( + sql_queries: list[str] | None = Field( default=None, description=( - "A consulta SQL executada para obter os dados, com comentários inline. " - "Deixe vazio (None) quando nenhuma consulta foi executada." + "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 sugestões de perguntas para explorar os dados mais a fundo.", + description=( + "3 suggested follow-up questions (in the user's language) to explore the " + "data further." + ), ) From 508c4e413dfab20bce5d96569061e716f385fe95 Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Wed, 1 Jul 2026 17:36:21 -0300 Subject: [PATCH 08/12] refactor: pin temporal_coverage fields to their granularity's format MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit period_start/period_end now document the format matching `granularity` (YYYY / YYYY-MM / YYYY-MM-DD) with examples, and the granularity field restates the mapping — so the model emits e.g. 2026-05 for a monthly period, not just the year. Co-Authored-By: Claude Opus 4.8 (1M context) --- app/agent/schemas.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/app/agent/schemas.py b/app/agent/schemas.py index a2be9f5..7004b86 100644 --- a/app/agent/schemas.py +++ b/app/agent/schemas.py @@ -17,19 +17,24 @@ class TemporalCoverage(BaseModel): 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`). May be narrower than the " - "table's full coverage." + "`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`). May be narrower " - "than the table's full coverage." + "`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`: 'day', 'month' or 'year'." + description=( + "Granularity of `period_start`/`period_end`, matching their format: " + "YYYY (year), YYYY-MM (month), YYYY-MM-DD (day)." + ) ) From a4b7b736f095a9ec552d90256cd35181437c8baf Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Wed, 1 Jul 2026 18:08:43 -0300 Subject: [PATCH 09/12] test: add multi-turn deterministic eval harness Replays eval_gold.yaml threads turn-by-turn on the real agent and scores four deterministic dimensions (action / source / period / completed) from the structured response, writing a rich JSON for the later judge eval. Period scoring is granularity-aware (matches period_end at YYYY / YYYY-MM / YYYY-MM-DD); records the system prompt and full tracebacks per run. Co-Authored-By: Claude Opus 4.8 (1M context) --- eval/eval_threads.py | 599 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 599 insertions(+) create mode 100644 eval/eval_threads.py diff --git a/eval/eval_threads.py b/eval/eval_threads.py new file mode 100644 index 0000000..f272617 --- /dev/null +++ b/eval/eval_threads.py @@ -0,0 +1,599 @@ +"""Multi-turn agent eval: replay conversation threads and score deterministic dimensions. + +Replays each thread in `eval_gold.yaml` turn-by-turn on a shared `thread_id` +(InMemorySaver, so follow-ups see prior context), K times per thread, for the +currently checked-out branch + given settings. Scores four deterministic +dimensions per turn from the agent's structured response fields (the purpose +of this eval is to check those fields are filled with the expected values): + + action: query-vs-clarify, from whether the agent filled the sql_queries field + source: data_sources (its table UUIDs resolved to gcp_ids) match the gold + period: temporal_coverage matches the gold rule (any/latest/range/match_previous/exact) + completed: the turn finished without errors + +It writes a rich JSON (full per-turn records: structured response, response text, +model-call count) so the later judge eval and the A-vs-B pass can score +the SAME transcripts without re-running the agent. + +Example Usage: + uv run python eval_threads.py --repeats 10 --temperature 0.0 + uv run python eval_threads.py --repeats 10 --threads --temperature 0.0 + +Each turn runs the real agent (live BQ + LLM) — mind the cost (sum of turns x K). +""" + +import argparse +import asyncio +import json +import os +import re +import subprocess +import traceback +from collections import defaultdict +from datetime import date, datetime +from pathlib import Path + +import yaml +from langchain.agents import create_agent +from langchain.agents.middleware import ( + ModelCallLimitMiddleware, + SummarizationMiddleware, +) +from langchain.chat_models import init_chat_model +from langchain.messages import AnyMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.graph.state import CompiledStateGraph + +from app.agent.prompts import SYSTEM_PROMPT +from app.agent.tools import BDToolkit +from app.settings import settings + +try: + from app.agent.schemas import StructuredResponse +except ImportError as exc: + raise NotImplementedError( + "This eval only supports the structured-response agent" + ) from exc + +# Agent middleware tunables (kept in step with production config). +SUMMARIZE_TRIGGER_TOKENS = 500_000 +SUMMARIZE_KEEP_TOKENS = 250_000 +MODEL_CALL_RUN_LIMIT = 20 + +# Deterministic dimensions scored per turn, in display order. +SCORE_DIMENSIONS = ("action", "source", "period", "completed") + + +# ============================================================================= +# Git and Agent setup +# ============================================================================= +def current_branch() -> str: + try: + return subprocess.check_output( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], text=True + ).strip() + except Exception: + return "unknown" + + +def build_agent( + temperature: float, checkpointer: InMemorySaver, system_prompt: str +) -> CompiledStateGraph: + model = init_chat_model( + model=settings.MODEL_URI, + temperature=temperature, + credentials=settings.GOOGLE_CREDENTIALS, + thinking_level=settings.THINKING_LEVEL, + include_thoughts=True, + ) + middleware = [ + SummarizationMiddleware( + model=model, + trigger=("tokens", SUMMARIZE_TRIGGER_TOKENS), + keep=("tokens", SUMMARIZE_KEEP_TOKENS), + ), + ModelCallLimitMiddleware( + run_limit=MODEL_CALL_RUN_LIMIT, + exit_behavior="end", + ), + ] + return create_agent( + model=model, + tools=BDToolkit.get_tools(), + system_prompt=system_prompt, + middleware=middleware, + checkpointer=checkpointer, + response_format=StructuredResponse, + ) + + +# ============================================================================= +# per-turn extraction from the agent's final state +# ============================================================================= +def _empty_record(status: str, **extra) -> dict: + """Base per-turn record for the non-`ok` cases (no_structured / error), with the + scoring fields zeroed so downstream code can read them uniformly. `extra` overrides + or adds fields (e.g. an `error` message, a recovered `response_text`).""" + return { + "status": status, + "is_query": False, + "tables": [], + "table_period_ends": {}, + "structured": None, + "response_text": None, + "model_calls": None, + "tools_used": [], + **extra, + } + + +def _last_ai_text(messages: list[AnyMessage]) -> str | None: + for m in reversed(messages): + if m.type == "ai": + return m.text + return None + + +def _tools_used(messages: list[AnyMessage]) -> set[str]: + """Distinct tool names the agent called in THIS turn. The checkpointer returns + the whole thread, so slice to the messages after the last human turn.""" + last_human = max( + (i for i, m in enumerate(messages) if m.type == "human"), default=-1 + ) + return { + tc["name"] + for m in messages[last_human + 1 :] + for tc in getattr(m, "tool_calls", None) or [] + } + + +def _tool_metadata(messages: list[AnyMessage]) -> tuple[dict[str, str], dict[str, str]]: + """From the agent's get_table_details / get_dataset_details outputs, build + (uuid -> gcp_id) to resolve data_sources, and (gcp_id -> period_end) for the + `latest` period rule. Tool payloads are pydantic-serialized JSON so keys are + guaranteed; a non-JSON body means that tool call errored, so skip it.""" + uuid_to_gcp: dict[str, str] = {} + period_end: dict[str, str] = {} + for m in messages: + if m.type != "tool" or m.name not in ( + "get_table_details", + "get_dataset_details", + ): + continue + try: + d = json.loads(m.content) + except json.JSONDecodeError: + continue + if m.name == "get_table_details": + uuid_to_gcp[d["id"]] = d["gcp_id"] + if d["period_end"] is not None: + period_end[d["gcp_id"]] = d["period_end"] + else: # get_dataset_details + for t in d["tables"]: + uuid_to_gcp[t["id"]] = t["gcp_id"] + return uuid_to_gcp, period_end + + +def extract_turn(result: dict) -> dict: + messages = result["messages"] + structured: StructuredResponse | None = result.get("structured_response") + + # This eval targets the structured-output agent: a missing structured response + # is a hard fail for the turn (not a crash, not a skip). + if structured is None: + return _empty_record( + "no_structured", + response_text=_last_ai_text(messages), + model_calls=result.get("run_model_call_count"), + ) + + uuid_to_gcp, table_period_ends = _tool_metadata(messages) + # Tables the agent REPORTED in data_sources, UUID -> gcp_id (unmapped UUIDs are + # kept as-is so they fail to match the gold and surface the problem). + data_source_tables = sorted( + { + uuid_to_gcp.get(d.table_id, d.table_id) + for d in (structured.data_sources or []) + } + ) + return { + "status": "ok", + # query vs clarify, from whether the agent filled the sql_queries field. + "is_query": bool(structured.sql_queries), + "tables": data_source_tables, + "table_period_ends": table_period_ends, + "structured": structured.model_dump(mode="json"), + "response_text": structured.response, + "model_calls": result.get("run_model_call_count"), + "tools_used": sorted(_tools_used(messages)), + } + + +# ============================================================================= +# Deterministic scoring +# ============================================================================= +def _lead_year(s) -> int | None: + """Leading 4-digit year of a period value ('2015', '2015-03-01', '2026-05').""" + m = re.match(r"\s*(\d{4})", str(s or "")) + return int(m.group(1)) if m else None + + +def _tc_years(tc: dict | None) -> tuple[int | None, int | None]: + """(start_year, end_year) from a temporal_coverage dict — the leading 4-digit + year of period_start/period_end (handles '2015' and '2015-03-01' alike).""" + if not tc: + return None, None + return _lead_year(tc.get("period_start")), _lead_year(tc.get("period_end")) + + +def _parse_period(s) -> tuple[str, str] | None: + """(granularity, value) for a period string in its granularity's own format: + 'YYYY' -> ('year', ...), 'YYYY-MM' -> ('month', ...), 'YYYY-MM-DD' -> ('day', ...). + Strict (re.fullmatch): a malformed value like '2026-5' returns None -> a miss, not + something to silently repair.""" + s = str(s or "").strip().strip("'").strip('"') + for gran, pat in ( + ("day", r"\d{4}-\d{2}-\d{2}"), + ("month", r"\d{4}-\d{2}"), + ("year", r"\d{4}"), + ): + if re.fullmatch(pat, s): + return gran, s + return None + + +def _score_period(rec: dict, gold: dict, prev_query: dict | None): + """Score the period from the agent's reported `temporal_coverage` — normalized + and independent of SQL surface form / temporal-column name. Returns [0,1] or None. + """ + rule = gold.get("temporal") + if gold["action"] != "query" or rule in (None, "none", "any", "", {}): + return None + ps, pe = _tc_years(rec["structured"]["temporal_coverage"]) + if ps is None or pe is None: + return 0.0 # queried but reported no usable temporal_coverage + if rule == "match_previous": + if prev_query is None: + return None + pps, ppe = _tc_years(prev_query["structured"]["temporal_coverage"]) + if pps is None or ppe is None: + return 0.0 + # The follow-up must stay WITHIN the previous turn's period. A justified + # narrowing (e.g., the requested metric only exists for part of the range) + # is fine; only drifting OUTSIDE the established window fails. + return 1.0 if pps <= ps and pe <= ppe else 0.0 + if rule == "range": + return 1.0 if pe > ps else 0.0 + if rule == "latest": + # Expected = the table's period_end, matched EXACTLY at its own granularity: + # '2026-05' must be reported as month 2026-05 (not year 2026), a full date to + # the day. The model must emit the value in its granularity's format. + queried = set(rec["tables"]) + cands = [v for k, v in rec["table_period_ends"].items() if k in queried] + cands = cands or list(rec["table_period_ends"].values()) + targets = [p for v in cands if (p := _parse_period(v)) is not None] + if not targets: + return None # can't determine the table's period_end from the trace + target = max(targets, key=lambda p: p[1]) # latest period_end + tc = rec["structured"]["temporal_coverage"] or {} + start, end = ( + _parse_period(tc.get("period_start")), + _parse_period(tc.get("period_end")), + ) + matches = ( + start == target and end == target and tc.get("granularity") == target[0] + ) + return 1.0 if matches else 0.0 + span = rule.get("exact", rule) if isinstance(rule, dict) else None + if isinstance(span, dict) and "start" in span and "end" in span: + return 1.0 if ps == int(span["start"]) and pe == int(span["end"]) else 0.0 + return None + + +def score_turn(rec: dict, gold: dict, prev_query: dict | None) -> dict: + if rec["status"] == "no_structured": + # agent returned no structured response -> the structured contract failed + return {"action": False, "source": None, "period": None, "completed": False} + + observed = "query" if rec["is_query"] else "clarify" + + if gold["action"] == "query" and gold.get("tables"): + exp, obs = set(gold["tables"]), set(rec["tables"]) + source = 1.0 if obs == exp else (0.5 if obs & exp else 0.0) + else: + source = None # clarify turn, or gcp_ids not filled in yet + return { + "action": observed == gold["action"], + "source": source, + "period": _score_period(rec, gold, prev_query), + "completed": rec["status"] == "ok", + } + + +# ============================================================================= +# Thread replay +# ============================================================================= +def _turn_run_config( + base: dict, + *, + thread_id: str, + thread: dict, + repeat: int, + turn_index: int, + turn: dict, + branch: str, + temperature: float, +) -> dict: + """Per-turn LangGraph config. Labels the run so it's findable in LangSmith; + `thread_id` in metadata makes LangSmith group a replay's turns into one + conversation.""" + return { + **base, + "configurable": {"thread_id": thread_id}, + "run_name": f"{thread['id']}#{repeat}-t{turn_index}", + "tags": [ + f"branch:{branch}", + f"temp:{temperature}", + f"thread:{thread['id']}", + ], + "metadata": { + "thread_id": thread_id, + "eval_thread": thread["id"], + "repeat": repeat, + "turn_index": turn_index, + "branch": branch, + "temperature": temperature, + "expected_action": turn["action"], + "expected_temporal": turn.get("temporal"), + }, + } + + +def _skipped_record(turn_index: int, user: str) -> dict: + return {"turn_index": turn_index, "user": user, "status": "skipped", "scores": {}} + + +async def replay_thread( + agent: CompiledStateGraph, + thread: dict, + repeat: int, + run_config: dict, + branch: str, + temperature: float, + max_turns: int | None = None, +) -> list: + thread_id = f"{thread['id']}-{repeat}" + records, prev_query = [], None + turns = thread["turns"] if max_turns is None else thread["turns"][:max_turns] + + for i, turn in enumerate(turns): + cfg = _turn_run_config( + run_config, + thread_id=thread_id, + thread=thread, + repeat=repeat, + turn_index=i, + turn=turn, + branch=branch, + temperature=temperature, + ) + + try: + result = await agent.ainvoke( + {"messages": [{"role": "user", "content": turn["user"]}]}, config=cfg + ) + rec = extract_turn(result) + except Exception as exc: + rec = _empty_record( + "error", + error=f"{type(exc).__name__}: {exc}", + traceback=traceback.format_exc(), + ) + + rec["turn_index"], rec["user"] = i, turn["user"] + rec["scores"] = ( + score_turn(rec, turn, prev_query) if rec["status"] != "error" else {} + ) + records.append(rec) + + if rec["status"] == "error": # broken state -> skip the rest of the thread + records.extend( + _skipped_record(j, turns[j]["user"]) for j in range(i + 1, len(turns)) + ) + break + + if rec["is_query"]: + prev_query = rec + + return records + + +# ============================================================================= +# Aggregation & Reporting +# ============================================================================= +def _mark(scores: dict) -> str: + def m(v): + return ( + "·" + if v is None + else ("✓" if v is True or v == 1.0 else ("~" if v == 0.5 else "✗")) + ) + + return ( + f"a:{m(scores.get('action'))} " + f"s:{m(scores.get('source'))} " + f"p:{m(scores.get('period'))} " + f"c:{m(scores.get('completed'))}" + ) + + +def _unit_line(turn_records: list) -> str: + """One-line per-thread progress summary: each turn's marks, or its status.""" + return " ".join( + f"t{tr['turn_index']}:" + + (_mark(tr["scores"]) if tr["status"] == "ok" else tr["status"].upper()) + for tr in turn_records + ) + + +def aggregate(units: list) -> dict: + """Pass-rate per dimension over all scored turns (None = not applicable, skipped).""" + agg = defaultdict(lambda: {"hit": 0.0, "n": 0}) + errors = skipped = no_structured = 0 + for u in units: + for rec in u["turns"]: + if rec["status"] == "skipped": + skipped += 1 + continue + if rec["status"] == "error": + errors += 1 + elif rec["status"] == "no_structured": + no_structured += 1 + for dim, val in (rec.get("scores") or {}).items(): + if val is None: + continue + agg[dim]["hit"] += float(val) + agg[dim]["n"] += 1 + summary = { + d: {"rate": round(v["hit"] / v["n"], 3), "n": v["n"]} for d, v in agg.items() + } + summary["_errors"], summary["_skipped"], summary["_no_structured"] = ( + errors, + skipped, + no_structured, + ) + return summary + + +def print_scorecard(summary: dict) -> None: + print("\n=== Deterministic scorecard (pass-rate over applicable turns) ===") + for dim in SCORE_DIMENSIONS: + if dim in summary: + print(f" {dim:<10} {summary[dim]['rate']:.0%} (n={summary[dim]['n']})") + print( + f" errors={summary['_errors']} skipped={summary['_skipped']} " + f"no_structured={summary['_no_structured']}" + ) + + +# ============================================================================= +# CLI +# ============================================================================= +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--gold", default="eval_gold.yaml") + parser.add_argument("--repeats", type=int, default=3, help="Replays per thread") + parser.add_argument( + "--thread", action="append", help="Only this thread id (repeatable)" + ) + parser.add_argument( + "--max-turns", + type=int, + default=None, + help="Run only the first N turns of each thread (a prefix — follow-ups need prior context)", + ) + parser.add_argument("--temperature", type=float, default=None) + parser.add_argument( + "--recursion-limit", + type=int, + default=None, + help="Omit to match production (LangGraph default 25)", + ) + parser.add_argument( + "--concurrency", type=int, default=1, help="Parallel thread replays" + ) + parser.add_argument("--out", default=None) + parser.add_argument( + "--langsmith-project", + default=None, + help="LangSmith project for these traces (default: -eval)", + ) + parser.add_argument( + "--no-trace", action="store_true", help="Disable LangSmith tracing" + ) + return parser.parse_args() + + +def configure_tracing(args: argparse.Namespace) -> tuple[str, bool]: + """Forward LangSmith config from `settings` into os.environ (LangChain's tracer + reads the env, but pydantic-settings only populates `settings`), isolating eval + runs in their own project by default. Returns (project, enabled); must run before + the agent does.""" + ls_project = args.langsmith_project or f"{settings.LANGSMITH_PROJECT}-eval" + tracing = settings.LANGSMITH_TRACING and not args.no_trace + os.environ["LANGSMITH_TRACING"] = "true" if tracing else "false" + if tracing: + os.environ["LANGSMITH_PROJECT"] = ls_project + os.environ["LANGSMITH_API_KEY"] = settings.LANGSMITH_API_KEY + return ls_project, tracing + + +async def main() -> None: + args = parse_args() + + with open(args.gold) as f: + gold = yaml.safe_load(f) + threads = [t for t in gold if not args.thread or t["id"] in args.thread] + + temperature = ( + args.temperature if args.temperature is not None else settings.MODEL_TEMPERATURE + ) + run_config: dict = {} + if args.recursion_limit is not None: + run_config["recursion_limit"] = args.recursion_limit + + branch = current_branch() + ts = datetime.now().strftime("%Y%m%d-%H%M%S") + out_path = args.out or str( + Path(__file__).resolve().parent + / f"thread_eval_{branch.replace('/', '-')}_temp{temperature}_{ts}.json" + ) + + ls_project, tracing = configure_tracing(args) + + print( + f"\nbranch={branch!r} model={settings.MODEL_URI!r}\n" + f"temperature={temperature} repeats={args.repeats} threads={[t['id'] for t in threads]}\n" + f"langsmith: project={ls_project!r}\n" + if tracing + else "langsmith: disabled\n" + ) + + system_prompt = SYSTEM_PROMPT.format(current_date=date.today().isoformat()) + agent = build_agent(temperature, InMemorySaver(), system_prompt) + sem = asyncio.Semaphore(args.concurrency) + work = [(t, r) for t in threads for r in range(args.repeats)] + + async def run_unit(thread, repeat): + async with sem: + turn_records = await replay_thread( + agent, thread, repeat, run_config, branch, temperature, args.max_turns + ) + print(f"[{thread['id']:<16} #{repeat}] {_unit_line(turn_records)}") + return {"thread": thread["id"], "repeat": repeat, "turns": turn_records} + + units = await asyncio.gather(*(run_unit(t, r) for t, r in work)) + summary = aggregate(units) + print_scorecard(summary) + + report = { + "branch": branch, + "langsmith_project": ls_project if tracing else None, + "model": settings.MODEL_URI, + "temperature": temperature, + "recursion_limit": args.recursion_limit, + "repeats": args.repeats, + "timestamp": ts, + "system_prompt": system_prompt, + "summary": summary, + "units": units, + } + + with open(out_path, "w") as fh: + json.dump(report, fh, ensure_ascii=False, indent=2) + + print(f"\nwrote {out_path}") + + +if __name__ == "__main__": + asyncio.run(main()) From 6c88d0f4581dd2c000156b8a8cacaa2cce2cc3ca Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Tue, 7 Jul 2026 17:56:21 -0300 Subject: [PATCH 10/12] refactor: allow data_sources on clarify turns; tie temporal_coverage to the SQL filter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit data_sources may now list the tables the answer draws on — those queried, or specific tables recommended when clarifying/guiding — not only tables that were queried. temporal_coverage is now explicitly the interval the SQL query filtered, left empty when no query ran (e.g. a clarification turn) or there is no temporal dimension. Co-Authored-By: Claude Opus 4.8 (1M context) --- app/agent/prompts.py | 2 +- app/agent/schemas.py | 21 ++++++++++----------- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/app/agent/prompts.py b/app/agent/prompts.py index 29001eb..2423212 100644 --- a/app/agent/prompts.py +++ b/app/agent/prompts.py @@ -122,7 +122,7 @@ ## Structured Fields Fill them **only** based on the tool results obtained in this conversation: -- **`data_sources`**: the tables **actually queried**, each with `dataset_id` (UUID from the `dataset_id` field of `get_table_details`, or the `id` field of `get_dataset_details`), `table_id` (UUID from the `id` field of `get_table_details`), and a readable name. **Never** use the `gcp_id` or the BigQuery name of the dataset/table. Leave empty when the answer does not use table data (e.g. explaining the platform, asking for clarification, listing available data types). +- **`data_sources`**: the tables the answer draws on — those you **queried**, or specific tables you **recommend when clarifying/guiding**. Each with `dataset_id` (UUID from the `dataset_id` field of `get_table_details`, or the `id` field of `get_dataset_details`), `table_id` (UUID from the `id` field of `get_table_details`, or a table's `id` from `get_dataset_details`; **never** the dataset UUID), and a readable name. **Never** use the `gcp_id` or the BigQuery name of the dataset/table. Leave empty only when no table is relevant (e.g. explaining the platform). - **`temporal_coverage`**: the interval your SQL query **actually filtered** — which may be narrower than the table's full coverage. E.g.: if `ano = 2010`, then `{{period_start: '2010', period_end: '2010'}}`; if `ano BETWEEN 2010 AND 2012`, then `{{period_start: '2010', period_end: '2012'}}`. Leave empty when there is no temporal dimension. - **`sql_queries`**: the queries whose results back the answer, 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 when no query was executed. - **`follow_up_questions`**: 3 suggestions for exploring the data further. diff --git a/app/agent/schemas.py b/app/agent/schemas.py index 7004b86..1226148 100644 --- a/app/agent/schemas.py +++ b/app/agent/schemas.py @@ -39,20 +39,21 @@ class TemporalCoverage(BaseModel): class DataSource(BaseModel): - """A Base dos Dados table used to answer the question.""" + """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 name (e.g. 'br_bd_diretorios')." + "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`), not the BigQuery name." + "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 used.") + name: str = Field(description="Human-readable name of the table.") class StructuredResponse(BaseModel): @@ -68,16 +69,15 @@ class StructuredResponse(BaseModel): data_sources: list[DataSource] | None = Field( default=None, description=( - "The tables actually queried to answer the question. Leave empty (None) " - "when the answer doesn't use table data (e.g. explaining the platform, " - "asking for clarification, or listing the available data types)." + "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 overall period of the data used. Leave empty (None) when the answer " - "has no temporal dimension." + "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( @@ -93,7 +93,6 @@ class StructuredResponse(BaseModel): 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." + "3 suggested follow-up questions (in the user's language) to explore the data further." ), ) From 9612c69757cabb420aeb2070f303aea33478850c Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Wed, 8 Jul 2026 15:55:50 -0300 Subject: [PATCH 11/12] test: add structured-output eval scorers and gold set MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds three transcript scorers on top of the deterministic replay harness, all reading the transcript eval_output.py produces: - eval_quality: LLM judge — answer correctness/grounding vs gold - eval_faithfulness: gold-free structural self-consistency of the structured output - eval_queries: source/period from the executed SQL (also works on the free-text baseline) Renames eval_threads.py -> eval_output.py so the script names describe what each evaluates. Adds eval_gold.yaml (10 threads / 18 turns) and ignores the generated transcripts/scores (eval/*.json). Co-Authored-By: Claude Opus 4.8 (1M context) --- .gitignore | 3 + eval/eval_faithfulness.py | 568 ++++++++++++++++++++ eval/eval_gold.yaml | 340 ++++++++++++ eval/eval_output.py | 1069 +++++++++++++++++++++++++++++++++++++ eval/eval_quality.py | 670 +++++++++++++++++++++++ eval/eval_queries.py | 326 +++++++++++ eval/eval_threads.py | 599 --------------------- 7 files changed, 2976 insertions(+), 599 deletions(-) create mode 100644 eval/eval_faithfulness.py create mode 100644 eval/eval_gold.yaml create mode 100644 eval/eval_output.py create mode 100644 eval/eval_quality.py create mode 100644 eval/eval_queries.py delete mode 100644 eval/eval_threads.py diff --git a/.gitignore b/.gitignore index e7551df..b3fb8c1 100644 --- a/.gitignore +++ b/.gitignore @@ -29,3 +29,6 @@ wheels/ .pytest_cache/ htmlcov/ .coverage + +# Eval artifacts (generated transcripts + scores; regenerable, can be large) +eval/*.json diff --git a/eval/eval_faithfulness.py b/eval/eval_faithfulness.py new file mode 100644 index 0000000..788dd3d --- /dev/null +++ b/eval/eval_faithfulness.py @@ -0,0 +1,568 @@ +"""Faithfulness eval: does the structured output honestly reflect the agent's run? + +Gold-free and judge-free. Reads a transcript produced by eval_output.py and, for each +answered turn, checks that the STRUCTURED fields are self-consistent and match what the +agent actually did (the executed SQL persisted in each turn's `queries`). It answers a +different question than the correctness evals: not "did the agent do the right thing?" +but "do data_sources / temporal_coverage / sql_queries truthfully describe THIS run?" + +Contract assumed (see app/agent/schemas.py): a no-query turn (clarification / platform +explanation) has no sql_queries and no temporal_coverage, but MAY list the tables it +explored in data_sources; a query turn fills all three from what it actually queried. + +Checks (each -> a pass-rate over the turns it applies to; `·`/None = not applicable): + + cross-field / contract + clarify_clean no query this turn -> temporal_coverage absent (data_sources + may still list the tables the agent explored) + query_has_sources queried -> data_sources is non-empty + + data_sources faithfulness + sources_resolve every reported table_id resolved to a gcp_id, not invented + (query OR clarify turns) + sources_match_sql reported tables are a subset of the tables the SQL actually hit + + temporal_coverage validity / faithfulness + temporal_gran_valid granularity matches the period_start/end format + temporal_order_valid period_start <= period_end + temporal_matches_sql reported year range == years evidenced by the run: SQL literals + + result-row values (covers evolução queries that filter no year; + column-agnostic; abstains + reports when neither evidences a year) + + prose + response_nonempty the prose answer is non-empty + prose_no_leak the prose doesn't embed SQL or a raw table id (each has its field) + + informational (reported, not counted as faithfulness bugs) + sources_exact_match reported tables == tables the SQL hit (flags omitted joins) + sql_reported_ran every reported query textually matches an executed one + followups_3 follow_up_questions is 3 non-empty items + +SQL-derived checks use the thread's executed queries: this turn's own when it ran any, +else earlier turns' (a follow-up may answer from data fetched on a previous turn without +re-querying). No agent, no BQ, no LLM: pure analysis of the transcript. Runs over ALL repeats. + +Pipeline: run AFTER eval_output.py — it scores the transcript eval_output.py produced. +Needs no gold, LLM or BQ, so it's the cheapest scorer. Sibling scorers over the same +transcript, independent of this one and of each other: + eval_quality.py answer quality vs the gold — LLM judge + live reference_sql + eval_queries.py source/period from the executed SQL vs the gold — reuses this + module's SQL/period helpers (_leading_year, _tables_and_years, ...) + + uv run eval/eval_faithfulness.py --in eval/.json + uv run eval/eval_faithfulness.py --in eval/.json --show-failures +""" + +import argparse +import json +import re +from collections import defaultdict, namedtuple +from pathlib import Path + +EVAL_DIR = Path(__file__).resolve().parent + +# Checks in display order; `info` ones are reported but excluded from the headline +# (they flag stylistic/definitional choices, not faithfulness bugs). +CORE_CHECKS = ( + "clarify_clean", + "query_has_sources", + "sources_resolve", + "sources_match_sql", + "temporal_gran_valid", + "temporal_order_valid", + "temporal_matches_sql", + "response_nonempty", + "prose_no_leak", +) +# sql_reported_ran is informational: exact query text can't cleanly separate a legitimate +# reformat/reuse from a fabrication, so a miss is a "worth eyeballing" flag, not a bug. +INFO_CHECKS = ("sources_exact_match", "sql_reported_ran", "followups_3") +ALL_CHECKS = CORE_CHECKS + INFO_CHECKS + +# A period value split into its granularity and its original text (e.g. "2026-05" -> month). +Period = namedtuple("Period", ["granularity", "value"]) + + +# ============================================================================= +# SQL / period parsing (best-effort, self-contained — no agent deps) +# ============================================================================= +def _strip_comments(sql: str) -> str: + """Remove `-- ...` line comments from a SQL string. + + Args: + sql (str): The SQL text (may be None). + + Returns: + str: The SQL with every `--` comment stripped to end-of-line. + """ + return re.sub(r"--[^\n]*", "", sql or "") + + +# A BigQuery id segment is an identifier (letter/underscore-led), never all-digit — so a +# `project.dataset.table` ref matches but a PT-formatted number like 213.421.037 does not. +_IDENT = r"[A-Za-z_][A-Za-z0-9_-]*" +_GCP_PAT = rf"{_IDENT}\.{_IDENT}\.{_IDENT}" + + +def _is_gcp_ref(ref: str) -> bool: + """Whether a string is a resolved BigQuery table ref `project.dataset.table`. + + Args: + ref (str): The candidate string (an unresolved UUID or a number returns False). + + Returns: + bool: True if `ref` is exactly a 3-part dotted identifier. + """ + return bool(re.fullmatch(_GCP_PAT, str(ref))) + + +def _sql_tables(sql: str) -> set[str]: + """The `project.dataset.table` refs a query touches. + + Matches 3-part dotted identifiers only, so column refs (t1.ano), CTE aliases and + PT-formatted numbers are excluded; comments are stripped first. + + Args: + sql (str): The SQL text. + + Returns: + set[str]: The distinct fully-qualified table references found. + """ + without_comments = _strip_comments(sql).replace("`", "") + return set(re.findall(rf"\b{_GCP_PAT}\b", without_comments)) + + +# A year shows up in a query's SQL as a literal, or in its RESULT rows as a value in a +# year/date-named column (or an unambiguous date string). Column matching is liberal — a +# false hit still has to BE a 19xx/20xx value to count. +_YEAR_COLUMN_RE = re.compile(r"ano|year|exercicio|periodo|competencia", re.I) +_DATE_COLUMN_RE = re.compile(r"data|date", re.I) +_YEAR_VALUE_RE = re.compile(r"(?:19|20)\d{2}") +_DATE_VALUE_RE = re.compile(r"(\d{4})-\d{2}(?:-\d{2})?") + + +def _years_in_sql(sql: str) -> set[int]: + """Year literals in a query's SQL, read column-agnostically. + + `ano`, `ano_campeonato`, `mes` + year, or a `'2020-01-01'` date literal all yield their + year. In this domain a standalone 19xx/20xx is a year (municipio codes are 7-digit, NCM + 8-digit, LIMIT/GROUP small). + + Args: + sql (str): The SQL text. + + Returns: + set[int]: The distinct years mentioned as literals. + """ + return { + int(year) for year in re.findall(r"\b(?:19|20)\d{2}\b", _strip_comments(sql)) + } + + +def _years_in_rows(rows: list[dict]) -> set[int]: + """Years evidenced by a query's RESULT rows. + + For evolução queries that SELECT all years (no year in the WHERE), the period lives + here, not in the SQL. Reads year values from year/date-named columns, plus any + unambiguous 'YYYY-MM(-DD)' string value in any column. + + Args: + rows (list[dict]): The query's result rows (column name -> value). + + Returns: + set[int]: The distinct years evidenced by the rows. + """ + years: set[int] = set() + for row in rows or []: + if not isinstance(row, dict): + continue + for column, value in row.items(): + if value is None: + continue + text = str(value) + if date_match := _DATE_VALUE_RE.fullmatch(text): # date string, any column + years.add(int(date_match.group(1))) + elif _YEAR_COLUMN_RE.search(column) and _YEAR_VALUE_RE.fullmatch(text): + years.add(int(text)) + elif _DATE_COLUMN_RE.search(column) and ( + year_match := _YEAR_VALUE_RE.match(text) + ): + years.add(int(year_match.group(0))) + return years + + +def _years_evidenced_by_query(query: dict) -> set[int]: + """Years one executed query evidences: SQL literals + values from its result rows. + + Truncated results are skipped — their min/max would understate the true span. + + Args: + query (dict): An executed query record with `sql`, `rows` and `row_count`. + + Returns: + set[int]: The distinct years evidenced by the query's SQL and (untruncated) rows. + """ + years = _years_in_sql(query.get("sql") or "") + rows, row_count = query.get("rows"), query.get("row_count") + is_untruncated = row_count is None or row_count <= len(rows) + if rows and is_untruncated: + years |= _years_in_rows(rows) + return years + + +def _normalize_sql(sql: str) -> str: + """Normalize a SQL string for textual comparison (strip comments, collapse whitespace, + lowercase, drop backticks). + + Args: + sql (str): The SQL text. + + Returns: + str: A canonical form suitable for equality/substring comparison. + """ + return re.sub(r"\s+", " ", _strip_comments(sql)).strip().lower().replace("`", "") + + +def _leading_year(value) -> int | None: + """The leading 4-digit year of a period value ('2015', '2015-03-01', '2026-05'). + + Args: + value: The period value (str or None). + + Returns: + int | None: The leading year, or None when the value has no leading 4 digits. + """ + match = re.match(r"\s*(\d{4})", str(value or "")) + return int(match.group(1)) if match else None + + +def _parse_period(period) -> Period | None: + """Split a period value into its granularity and text, in its own format: + + 'YYYY' -> year + 'YYYY-MM' -> month + 'YYYY-MM-DD' -> day + + Strict: a malformed value returns None. + + Args: + period: The period value (str or None). + + Returns: + Period | None: (granularity, value), or None when the value is malformed. + """ + text = str(period or "").strip().strip("'").strip('"') + for granularity, pattern in ( + ("day", r"\d{4}-\d{2}-\d{2}"), + ("month", r"\d{4}-\d{2}"), + ("year", r"\d{4}"), + ): + if re.fullmatch(pattern, text): + return Period(granularity, text) + return None + + +# ============================================================================= +# Per-turn checks +# ============================================================================= +def turn_success_queries(turn: dict) -> list[dict]: + """This turn's successfully-executed queries. + + Args: + turn (dict): A per-turn record from the transcript. + + Returns: + list[dict]: The query records (each with `sql` + result `rows`) whose status is + "success". + """ + return [ + query + for query in (turn.get("queries") or []) + if query.get("status") == "success" and query.get("sql") + ] + + +def _tables_and_years(queries: list[dict]) -> tuple[set[str], tuple[int, int] | None]: + """Aggregate the tables hit and the overall year range across a set of executed queries. + + Args: + queries (list[dict]): Executed query records (each with `sql` + result `rows`). + + Returns: + tuple[set[str], tuple[int, int] | None]: (tables referenced, (min_year, max_year)), + where the year range is None when no query evidenced any year. + """ + tables: set[str] = set() + years: set[int] = set() + for query in queries: + tables |= _sql_tables(query.get("sql") or "") + years |= _years_evidenced_by_query(query) + year_range = (min(years), max(years)) if years else None + return tables, year_range + + +def check_turn(turn: dict, executed_queries: list[dict]) -> dict: + """Run every faithfulness check for one `ok` turn. + + Each check is True (passed), False (failed), or None (not applicable to this turn). + + Args: + turn (dict): The per-turn record (its `structured` fields, `is_query`, and the + `tables` resolved from data_sources). + executed_queries (list[dict]): The executed queries (sql + rows) to check + faithfulness against — this turn's own when it ran any, else the thread's earlier + ones (a follow-up may answer from data fetched on a previous turn without + re-querying). + + Returns: + dict: Maps each check name in ALL_CHECKS to True / False / None. + """ + structured = turn.get("structured") or {} + is_query = bool(turn.get("is_query")) + data_sources = structured.get("data_sources") or [] + temporal_coverage = structured.get("temporal_coverage") + reported_sql = structured.get("sql_queries") or [] + follow_ups = structured.get("follow_up_questions") + response = structured.get("response") or "" + reported_tables = list( + turn.get("tables") or [] + ) # data_sources resolved to gcp/uuid + executed_sql = [query.get("sql") or "" for query in executed_queries] + queried_tables, queried_year_range = _tables_and_years(executed_queries) + + results: dict = {check: None for check in ALL_CHECKS} + + # cross-field / contract + if not is_query: + # a no-query turn may surface explored tables in data_sources, but with no SQL there + # is no filtered interval, so temporal_coverage must stay empty. + results["clarify_clean"] = temporal_coverage is None + else: + results["query_has_sources"] = bool(data_sources) + results["followups_3"] = ( + isinstance(follow_ups, list) + and len(follow_ups) == 3 + and all((question or "").strip() for question in follow_ups) + ) + + # data_sources faithfulness + if data_sources: + # every reported table must have been fetched this run (resolves to a gcp ref) — + # applies to clarify turns too, which may list the tables they explored. + results["sources_resolve"] = all( + _is_gcp_ref(table) for table in reported_tables + ) + # matching against the executed SQL only makes sense on a query turn. + if is_query and queried_tables: + reported_gcp = {table for table in reported_tables if _is_gcp_ref(table)} + results["sources_match_sql"] = reported_gcp <= queried_tables + results["sources_exact_match"] = reported_gcp == queried_tables + + # temporal_coverage validity / faithfulness + if temporal_coverage: + start = _parse_period(temporal_coverage.get("period_start")) + end = _parse_period(temporal_coverage.get("period_end")) + granularity = temporal_coverage.get("granularity") + results["temporal_gran_valid"] = ( + start is not None + and end is not None + and start.granularity == granularity + and end.granularity == granularity + ) + if ( + start is not None + and end is not None + and start.granularity == end.granularity + ): + results["temporal_order_valid"] = start.value <= end.value + else: + start_year = _leading_year(temporal_coverage.get("period_start")) + end_year = _leading_year(temporal_coverage.get("period_end")) + results["temporal_order_valid"] = ( + start_year is not None + and end_year is not None + and start_year <= end_year + ) + if is_query and queried_year_range is not None: + start_year = _leading_year(temporal_coverage.get("period_start")) + end_year = _leading_year(temporal_coverage.get("period_end")) + results["temporal_matches_sql"] = ( + start_year, + end_year, + ) == queried_year_range + + # sql_queries / prose + if is_query and reported_sql: + executed_normalized = [_normalize_sql(sql) for sql in executed_sql] + reported_normalized = [_normalize_sql(sql) for sql in reported_sql] + + results["sql_reported_ran"] = all( + any( + reported == executed or reported in executed or executed in reported + for executed in executed_normalized + ) + for reported in reported_normalized + ) + results["response_nonempty"] = bool(response.strip()) + results["prose_no_leak"] = not ( + "```sql" in response.lower() + or any(_is_gcp_ref(table) for table in _sql_tables(response)) + ) + return results + + +# ============================================================================= +# Aggregation & reporting +# ============================================================================= +def aggregate(turn_results: list[dict]) -> dict: + """Compute the pass-rate per check across all scored turns. + + A None value (check not applicable to that turn) is skipped, so each check's `n` is the + number of turns it actually applied to. + + Args: + turn_results (list[dict]): One check_turn() result dict per scored turn. + + Returns: + dict: Maps each check name to {"rate": float, "n": int}. + """ + tally = defaultdict(lambda: {"passed": 0, "applicable": 0}) + for result in turn_results: + for check, outcome in result.items(): + if outcome is None: + continue + tally[check]["passed"] += int(bool(outcome)) + tally[check]["applicable"] += 1 + return { + check: { + "rate": round(counts["passed"] / counts["applicable"], 3), + "n": counts["applicable"], + } + for check, counts in tally.items() + } + + +def print_scorecard( + summary: dict, processed: int, skipped: int, notes: dict | None = None +) -> None: + """Print the faithfulness scorecard, core checks then informational ones. + + Args: + summary (dict): The per-check {"rate", "n"} mapping from aggregate(). + processed (int): Number of `ok` turns scored. + skipped (int): Number of non-`ok` turns skipped. + notes (dict | None): Optional per-check suffix strings (e.g. an abstention note). + + Returns: + None + """ + notes = notes or {} + print( + f"\n=== Faithfulness scorecard ({processed} turns; {skipped} non-ok skipped) ===" + ) + print(" -- faithfulness / contract --") + for check in CORE_CHECKS: + if check in summary: + print( + f" {check:<22} {summary[check]['rate']:.0%} (n={summary[check]['n']}){notes.get(check, '')}" + ) + print(" -- informational (not faithfulness bugs) --") + for check in INFO_CHECKS: + if check in summary: + print( + f" {check:<22} {summary[check]['rate']:.0%} (n={summary[check]['n']}){notes.get(check, '')}" + ) + + +def main() -> None: + """Score a transcript's structured-output faithfulness and write a JSON report.""" + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--in", dest="transcript", required=True, help="thread_eval_*.json" + ) + parser.add_argument( + "--thread", action="append", help="Only this thread id (repeatable)" + ) + parser.add_argument("--out", default=None) + parser.add_argument( + "--show-failures", + action="store_true", + help="List each failing (thread, repeat, turn) and which checks failed", + ) + args = parser.parse_args() + + transcript = json.load(open(args.transcript)) + units = transcript["units"] + if args.thread: + units = [unit for unit in units if unit["thread"] in args.thread] + + turn_results, failures, skipped, temporal_abstained = [], [], 0, 0 + for unit in units: + prior_queries: list[ + dict + ] = [] # successful queries from earlier turns of THIS thread + for turn in unit["turns"]: + if turn.get("status") != "ok": + skipped += 1 + continue + this_queries = turn_success_queries(turn) + result = check_turn(turn, this_queries or prior_queries) + turn_results.append(result) + # a query turn that reported a period but whose run evidenced no year (no literal + # in SQL, none in the rows): unverifiable, NOT a pass — surface it so the rate + # isn't misread. + structured = turn.get("structured") or {} + if ( + turn.get("is_query") + and structured.get("temporal_coverage") + and result["temporal_matches_sql"] is None + ): + temporal_abstained += 1 + prior_queries += this_queries + failed_checks = [ + check + for check, outcome in result.items() + if outcome is False and check in CORE_CHECKS + ] + if failed_checks: + failures.append( + (unit["thread"], unit["repeat"], turn["turn_index"], failed_checks) + ) + + summary = aggregate(turn_results) + notes = {} + if temporal_abstained: + notes["temporal_matches_sql"] = ( + f" [+{temporal_abstained} unverifiable: no year literal in SQL]" + ) + print( + f"transcript={args.transcript!r} branch={transcript.get('branch')!r} model={transcript.get('model')!r}" + ) + print_scorecard(summary, len(turn_results), skipped, notes) + + if args.show_failures and failures: + print("\n=== Failing turns (core checks) ===") + for thread, repeat, turn_index, failed_checks in failures: + print(f" {thread:<20} #{repeat} t{turn_index}: {', '.join(failed_checks)}") + + out_path = args.out or str( + EVAL_DIR / f"{Path(args.transcript).stem}_faithfulness.json" + ) + with open(out_path, "w") as file: + json.dump( + { + "transcript": args.transcript, + "summary": summary, + "turn_results": turn_results, + }, + file, + ensure_ascii=False, + indent=2, + ) + print(f"\nwrote {out_path}") + + +if __name__ == "__main__": + main() diff --git a/eval/eval_gold.yaml b/eval/eval_gold.yaml new file mode 100644 index 0000000..eaa6ffc --- /dev/null +++ b/eval/eval_gold.yaml @@ -0,0 +1,340 @@ +# Gold set for the agent eval. Each item is a multi-turn THREAD: eval_output.py replays a +# thread's `user` turns in order on one conversation (shared thread_id, so follow-ups see +# prior context), then the scorer scripts grade the run against the expectations below. +# +# Per-turn fields: +# user: the user message for this turn. +# action: expected behavior — `query` (must run SQL) or `clarify` +# (must NOT run SQL; should ask the user for more detail). +# tables: gcp_ids the executed SQL is expected to reference (FROM/JOIN), e.g. +# "basedosdados.br_me_rais.microdados_vinculos". [] for clarify. Scored by overlap +# with the tables the agent actually used: 1.0 exact / 0.5 partial / 0.0 none. +# temporal: rule for the period the SQL filtered on. One of: +# none -> no temporal dimension expected (clarify turns) +# any -> don't score the period this turn +# latest -> the single most-recent period; must equal the table's +# period_end at its own granularity (YYYY / YYYY-MM / YYYY-MM-DD). +# CAUTION for monthly tables: `latest` means the latest MONTH, +# which may not be the period you want a reference_sql to cover. +# range -> a multi-year span (evolução / temporal analysis) +# match_previous -> within the previous query turn's period ("nesse período"); +# a justified narrowing is OK, only drifting outside fails +# {exact: {start: 2015, end: 2020}} -> explicit span in the question +# reference_sql: OPTIONAL, query turns only. A trusted SQL whose LIVE result is the gold for +# the judge's `correct` check — re-run at eval time, so its VALUES track the +# current data instead of a frozen number. NOTE: a hardcoded period filter +# (e.g. `ano = 2025`) still drifts as coverage advances — keep it current, or on +# `latest` turns use `ano = (SELECT MAX(ano) FROM )` so it self-updates. +# Provide it on turns whose answer has a well-defined, stable value worth +# verifying; OMIT it on open-ended/analytical turns (the judge then sets +# correct=null and scores only `grounded` + `answers_question`). Fill with a `|` block. +# notes: what this turn is meant to test (human-facing). +# +# Which eval reads each field: +# eval_output deterministic, from the agent's STRUCTURED fields: +# action / source / period / completed / format_ok +# eval_queries deterministic, from the agent's EXECUTED SQL (year-level): +# source / period — also runs on the free-text baseline (main) +# eval_quality LLM judge: grounded + answers_question, plus `correct` only when the +# turn carries a reference_sql +# eval_faithfulness gold-free — does NOT read this file. + +# NOTE(maintenance): some reference_sql filters hardcode a period and must be bumped as BD +# coverage advances — the query is re-run live, but the hardcoded filter is not. `latest` +# turns should track the table's period_end (update by hand, or switch to +# `ano = (SELECT MAX(ano) FROM
)` so they self-update): +# rais t0 ano = 2025 +# ibge-populacao t0 ano = 2025 (t1 match_previous reuses this year) +# comex-stat t0 ano = 2026 (grain caveat — see the TODO on that thread) +# cnes-farmacias t0 ano = 2026 AND mes = 5 +# Range / explicit-span turns are softer: futebol t1's `... BETWEEN 2012 AND 2024` caps at the +# last season and the judge tolerates a valid wider span; mapbiomas' 2015-2020 comes from the +# question, so it is stable. + +- id: futebol + turns: + - user: "Futebol" + action: clarify + tables: [] + temporal: none + notes: "Uma só palavra, ambíguo -> deve pedir esclarecimento sobre o que de futebol, não consultar." + + - user: "Como evoluiu a média de público nos estádios brasileiros nos últimos anos?" + action: query + tables: + - basedosdados.mundo_transfermarkt_competicoes.brasileirao_serie_a + temporal: range + reference_sql: | + SELECT + ano_campeonato, + CAST(AVG(publico) AS INT64) AS media_publico + FROM `basedosdados.mundo_transfermarkt_competicoes.brasileirao_serie_a` + WHERE ano_campeonato BETWEEN 2012 and 2024 + GROUP BY 1 + ORDER BY 1 + notes: "Evolução -> intervalo de vários anos." + + - user: "Faça uma análise por clube, no mesmo período" + action: query + tables: + - basedosdados.mundo_transfermarkt_competicoes.brasileirao_serie_a + temporal: match_previous + # No reference_sql: "análise por clube" is open-ended (média por clube OU por + # clube-e-ano são leituras válidas). Julgado por grounded + answers_question; + # o período (match_previous) e as tabelas ainda são checados deterministicamente. + notes: "O período DEVE ser igual ao do turno anterior; agrupar por clube (granularidade a critério do agente)." + +- id: ideb + turns: + - user: "Como evoluiu a nota do IDEB no meu município?" + action: clarify + tables: [] + temporal: none + notes: "Nenhum município informado -> deve perguntar qual." + + - user: "Sou de Piracicaba - SP e estou interessado nos anos finais da rede pública" + action: query + tables: + - basedosdados.br_inep_ideb.municipio + - basedosdados.br_bd_diretorios_brasil.municipio + temporal: range + reference_sql: | + SELECT + t1.ano, + t2.nome as municipio, + t2.sigla_uf, t1.ideb + FROM `basedosdados.br_inep_ideb.municipio` as t1 + JOIN `basedosdados.br_bd_diretorios_brasil.municipio` as t2 + ON t1.id_municipio = t2.id_municipio + WHERE t2.id_municipio = "3538709" + AND t1.anos_escolares = "finais (6-9)" + AND t1.rede = "publica" + ORDER BY 1 + notes: "Evolução -> intervalo de vários anos. Filtros: município=Piracicaba/SP, anos finais, rede pública." + + - user: "Faça uma comparação com as metas nesse período" + action: query + tables: + - basedosdados.br_inep_ideb.municipio + - basedosdados.br_bd_diretorios_brasil.municipio + temporal: match_previous + reference_sql: | + SELECT + t1.ano, + t2.nome as municipio, + t2.sigla_uf, t1.ideb, + t1.projecao + FROM `basedosdados.br_inep_ideb.municipio` as t1 + JOIN `basedosdados.br_bd_diretorios_brasil.municipio` as t2 + ON t1.id_municipio = t2.id_municipio + WHERE t2.id_municipio = "3538709" + AND t1.anos_escolares = "finais (6-9)" + AND t1.rede = "publica" + ORDER BY 1 + notes: "Deve trazer a meta do IDEB; mesmo município/período." + +# TODO(eval): comex-stat uses `temporal: latest`, but ncm_exportacao/ncm_importacao are +# MONTHLY tables (period_end e.g. 2026-06). So the deterministic `latest` check expects the +# latest MONTH, while the reference_sql below aggregates the whole YEAR (`ano = 2026`, a +# partial year) — the two can't both be satisfied by one answer. Decide the intended period +# policy for monthly-but-annually-answered datasets: options are (a) `temporal: any` (stop +# scoring the period here), (b) drop the reference_sql and treat it as open-ended, or (c) add +# a "latest complete year" rule + reference `ano = 2025`. Parked — revisit with the eval. +- id: comex-stat + turns: + - user: "Quais os principais produtos exportados pelo Brasil?" + action: query + tables: + - basedosdados.br_me_comex_stat.ncm_exportacao + - basedosdados.br_bd_diretorios_mundo.nomenclatura_comum_mercosul + temporal: latest + reference_sql: | + SELECT + t2.nome_ncm_portugues AS produto, + SUM(t1.valor_fob_dolar) AS total_valor_usd + FROM `basedosdados.br_me_comex_stat.ncm_exportacao` AS t1 + JOIN `basedosdados.br_bd_diretorios_mundo.nomenclatura_comum_mercosul` AS t2 + ON t1.id_ncm = t2.id_ncm + WHERE t1.ano = 2026 + GROUP BY 1 + ORDER BY 2 DESC + LIMIT 10 + notes: "Ranking por valor; sem período informado -> mais recente." + + - user: "Quais os principais destinos desses produtos?" + action: query + tables: + - basedosdados.br_me_comex_stat.ncm_exportacao + - basedosdados.br_bd_diretorios_mundo.pais + temporal: latest + # No reference_sql: "principais destinos desses produtos" admite leituras válidas + # (destinos no geral OU o cruzamento produto x destino) — o cruzamento é uma + # decomposição muito específica p/ ser a única gold. Julgado por grounded + + # answers_question; as tabelas e o período ainda são checados deterministicamente. + notes: "Agrupar por país de destino (decomposição a critério do agente)." + + - user: "Faça essa mesma análise, mas para as importações agora" + action: query + tables: + - basedosdados.br_me_comex_stat.ncm_importacao + - basedosdados.br_bd_diretorios_mundo.pais + - basedosdados.br_bd_diretorios_mundo.nomenclatura_comum_mercosul + temporal: latest + # No reference_sql: mesma decomposição aberta do turno anterior, agora p/ importações. + # O reescopo export->import e as tabelas certas ainda são checados deterministicamente + # (source); a resposta é julgada por grounded + answers_question. + notes: "Troca de direção exportações->importações; testa o reescopo correto (decomposição a critério do agente)." + +- id: rais + turns: + - user: "Qual a porcentagem de mulheres com emprego formal?" + action: query + tables: + - basedosdados.br_me_rais.microdados_vinculos + temporal: latest + reference_sql: | + SELECT + CASE + WHEN sexo = '1' THEN 'Masculino' + WHEN sexo = '2' THEN 'Feminino' + END AS sexo, + COUNT(*) AS quantidade, + ROUND(COUNT(*) * 100.0 / SUM(COUNT(*)) OVER (), 2) AS percentual + FROM `basedosdados.br_me_rais.microdados_vinculos` + WHERE vinculo_ativo_3112 = '1' + AND ano = 2025 + GROUP BY 1 + ORDER BY 1 + notes: "Teste da política de period_end - deve usar period_end, não um ano 'completo' anterior." + +- id: mapbiomas-uso-solo + turns: + - user: "Preciso entender a trajetória na mudança do uso do solo no Pará, entre 2015 e 2020" + action: query + tables: + - basedosdados.br_mapbiomas_estatisticas.cobertura_uf_classe + - basedosdados.br_mapbiomas_estatisticas.classe + temporal: {exact: {start: 2015, end: 2020}} + reference_sql: | + SELECT + t1.ano, + t2.valor_pt as classe, + t1.area + FROM `basedosdados.br_mapbiomas_estatisticas.cobertura_uf_classe` AS t1 + JOIN `basedosdados.br_mapbiomas_estatisticas.classe` AS t2 + ON t1.id_classe = t2.chave + WHERE t1.ano BETWEEN 2015 AND 2020 + AND t1.sigla_uf = "PA" + ORDER BY 1, 3 DESC + notes: "Período explícito -> exatamente 2015-2020; filtrar Pará." + + - user: "Quero as 5 maiores transições para cada ano desse mesmo período" + action: query + tables: + - basedosdados.br_mapbiomas_estatisticas.transicao_uf_de_para_anual + - basedosdados.br_mapbiomas_estatisticas.classe + temporal: match_previous + reference_sql: | + WITH ranked_transitions AS ( + SELECT + t.ano, + c1.valor_pt AS classe_origem, + c2.valor_pt AS classe_destino, + t.area, + ROW_NUMBER() OVER (PARTITION BY t.ano ORDER BY t.area DESC) as rank + FROM `basedosdados.br_mapbiomas_estatisticas.transicao_uf_de_para_anual` AS t + JOIN `basedosdados.br_mapbiomas_estatisticas.classe` AS c1 + ON t.id_classe_de = c1.chave + JOIN `basedosdados.br_mapbiomas_estatisticas.classe` AS c2 + ON t.id_classe_para = c2.chave + WHERE t.sigla_uf = "PA" + AND t.ano BETWEEN 2015 AND 2020 + AND t.id_classe_de != t.id_classe_para + ) + SELECT + ano, + classe_origem, + classe_destino, + area + FROM ranked_transitions + WHERE rank <= 5 + ORDER BY 1, 4 DESC + notes: "5 maiores transições por ano; mesmo período." + +# TODO(eval): cnes-farmacias t0 lists br_bd_diretorios_brasil.municipio in `tables`, but the +# reference_sql hardcodes id_municipio = '3509502' (Campinas) instead of JOINing the directory. +# Intentional IF the agent is expected to resolve "Campinas" via the directory (then source +# scores 1.0 only when it joins); if a hardcoded code is acceptable, drop the directory from +# `tables`. Decide when returning to eval. +- id: cnes-farmacias + turns: + - user: "Quantas farmácias existem no município de Campinas?" + action: query + tables: + - basedosdados.br_ms_cnes.estabelecimento + - basedosdados.br_bd_diretorios_brasil.municipio + temporal: latest + reference_sql: | + SELECT + COUNT(DISTINCT id_estabelecimento_cnes) as total_farmacias + FROM `basedosdados.br_ms_cnes.estabelecimento` + WHERE id_municipio = '3509502' -- ID de Campinas/SP + AND tipo_unidade = '43' -- Código para 'farmacia' + AND ano = 2026 + AND mes = 5 + notes: "Tabela mensal com period_end '2026-05' e SEM guia de uso. Testa se o agente usa o período mais recente (period_end) em vez de assumir um ano anterior do seu conhecimento prévio." + +- id: capacidades + turns: + - user: "O que é a Base dos Dados e como você pode me ajudar nas minhas análises?" + action: clarify + tables: [] + temporal: none + notes: "Pergunta sobre a plataforma/capacidades — não é um pedido de dados. Deve responder diretamente (o que é a BD; que pode buscar datasets, explorar tabelas e executar SQL) SEM executar consulta de dados e SEM inventar datasets/números específicos. É um dos casos em que responder sem tools é permitido." + +- id: dado-inexistente + turns: + - user: "Quantos assinantes a Netflix tem no Brasil atualmente?" + action: clarify + tables: [] + temporal: none + notes: "Dado que a BD não possui (base de assinantes de uma empresa de streaming). Deve BUSCAR (search_datasets) para confirmar a ausência e então informar que não encontrou esse dado na plataforma — SEM inventar tabela/valor e SEM responder pelo conhecimento próprio. Pode sugerir temas relacionados que existem. Testa a regra de grounding / não-fabricação." + +- id: periodo-fora-do-intervalo + turns: + - user: "Quais foram os principais produtos exportados pelo Brasil em 1980?" + action: clarify + tables: [] + temporal: none + notes: "1980 é anterior ao início da cobertura da série de comércio exterior (começa no fim dos anos 1990). Deve consultar get_table_details, constatar que 1980 está fora de [period_start, period_end], informar o período disponível e perguntar como prosseguir — SEM consultar silenciosamente outro ano. Testa a regra de período fora do intervalo." + +- id: ibge-populacao + turns: + - user: "Qual é a população total do Brasil?" + action: query + tables: + - basedosdados.br_ibge_populacao.brasil + temporal: latest + reference_sql: | + SELECT ano, populacao + FROM `basedosdados.br_ibge_populacao.brasil` + WHERE ano = 2025 -- period_end (cobertura anual 1991-2025); atualizar quando avançar + notes: "Fonte IBGE (cobre a lacuna de órgão no gold). Sem período informado -> mais recente (period_end anual, granularidade ano). Tabela de uma linha por ano — consulta de uma tabela só." + + - user: "E como ela se distribui entre os estados nesse ano?" + action: query + tables: + - basedosdados.br_ibge_populacao.uf + - basedosdados.br_bd_diretorios_brasil.uf + temporal: match_previous + reference_sql: | + SELECT + t2.nome AS estado, + t1.sigla_uf, + t1.populacao + FROM `basedosdados.br_ibge_populacao.uf` AS t1 + JOIN `basedosdados.br_bd_diretorios_brasil.uf` AS t2 + ON t1.sigla_uf = t2.sigla -- diretório usa `sigla` (não `sigla_uf`) + WHERE t1.ano = 2025 + ORDER BY t1.populacao DESC + notes: "Mesmo ano do turno anterior (match_previous). sigla_uf é codificada (reference_table_id -> diretório de UF): deve fazer JOIN para exibir o nome do estado. Testa reescopo Brasil->UF e tradução de coluna codificada." diff --git a/eval/eval_output.py b/eval/eval_output.py new file mode 100644 index 0000000..72bd120 --- /dev/null +++ b/eval/eval_output.py @@ -0,0 +1,1069 @@ +"""Multi-turn agent eval: replay conversation threads and score deterministic dimensions. + +Replays each thread in `eval_gold.yaml` turn-by-turn on a shared `thread_id` +(InMemorySaver, so follow-ups see prior context), K times per thread, for the +currently checked-out branch + given settings. Scores five deterministic +dimensions per turn from the agent's structured response fields (the purpose +of this eval is to check those fields are filled with the expected values): + + action: query-vs-clarify, from whether the agent filled the sql_queries field + source: data_sources (its table UUIDs resolved to gcp_ids) match the gold + period: temporal_coverage matches the gold rule (any/latest/range/match_previous/exact) + format_ok: the prose carries no Markdown ATX headers (the one hard-forbidden format rule) + completed: the turn finished without errors + +It writes a rich JSON (full per-turn records: structured response, response text, +model-call count, and the agent's OWN executed queries + their result rows) so the +later judge eval and the A-vs-B pass can score the SAME transcripts without re-running +the agent. Persisting the query results lets the judge check grounding against what the +agent actually retrieved — not only against the gold reference query. + +On the free-text baseline (e.g. the main branch, no structured output), pass +--no-structured: the agent is built without a response_format, the answer is read from +the final message, and only action/format_ok/completed are scored (source/period need the +structured fields). The resulting transcript is still judgeable by eval_quality. + +Pipeline: this is step 1 — the ONLY script that runs the agent. It writes the transcript +JSON that the three post-hoc scorers read (no agent; run them in any order): + eval_quality.py answer quality vs the gold — LLM judge + live reference_sql + eval_faithfulness.py structured output's self-consistency — gold-free, no LLM/BQ + eval_queries.py source/period from the executed SQL vs the gold — no LLM/BQ + +Example Usage: + uv run eval/eval_output.py --repeats 10 --temperature 0.0 + uv run eval/eval_output.py --thread --dry-run + uv run eval/eval_output.py --no-structured # baseline: free-text agent (e.g. main) + +Each turn runs the real agent (live BQ + LLM) — mind the cost (sum of turns x K). +""" + +import argparse +import asyncio +import json +import os +import re +import subprocess +import sys +import traceback +from collections import defaultdict +from datetime import date, datetime +from pathlib import Path + +import yaml +from langchain.agents import create_agent +from langchain.agents.middleware import ( + ModelCallLimitMiddleware, + SummarizationMiddleware, +) +from langchain.chat_models import init_chat_model +from langchain.messages import AnyMessage +from langgraph.checkpoint.memory import InMemorySaver +from langgraph.graph.state import CompiledStateGraph + +# Make the repo root importable so `app` resolves whether this file is run as a +# module (python -m eval.eval_output) or directly (python eval/eval_output.py). +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from app.agent.prompts import SYSTEM_PROMPT # noqa: E402 +from app.agent.tools import BDToolkit # noqa: E402 +from app.settings import settings # noqa: E402 + +# Structured output only exists on the structured-response branch. On main (no structured +# output) this import fails; run with --no-structured, which builds the agent without a +# response_format and reads the answer from the final message instead. +try: + from app.agent.schemas import StructuredResponse # noqa: E402 +except ImportError: + StructuredResponse = None + +# This script's folder — gold input and result files default here, so the eval +# works regardless of the current working directory. +EVAL_DIR = Path(__file__).resolve().parent + +# Agent middleware tunables (kept in step with production config). +SUMMARIZE_TRIGGER_TOKENS = 500_000 +SUMMARIZE_KEEP_TOKENS = 250_000 +MODEL_CALL_RUN_LIMIT = 20 + +# Deterministic dimensions scored per turn, in display order (matches _mark and score_turn). +SCORE_DIMENSIONS = ("action", "source", "period", "format_ok", "completed") + +# format_ok: the one prose-format rule the system prompt hard-forbids (no Markdown ATX +# headers) is mechanical, so it's scored here rather than spending a judge call on it. +_CODE_FENCE_RE = re.compile(r"^\s*```") +_MD_HEADER_RE = re.compile(r"^\s{0,3}#{1,6}(?:\s|$)") + +# Cap rows persisted per executed query, so the transcript (and the judge's view of the +# agent's own data) stays bounded; the true row count is kept alongside the capped rows. +QUERY_ROWS_CAP = 50 + + +# ============================================================================= +# Git and Agent setup +# ============================================================================= +def current_branch() -> str: + """The current git branch name. + + Returns: + str: The abbreviated branch name, or "unknown" if git can't be queried. + """ + try: + return subprocess.check_output( + ["git", "rev-parse", "--abbrev-ref", "HEAD"], text=True + ).strip() + except Exception: + return "unknown" + + +def build_agent( + temperature: float, + checkpointer: InMemorySaver, + system_prompt: str, + structured: bool = True, +) -> CompiledStateGraph: + """Build the agent under test, mirroring the production configuration. + + Args: + temperature (float): Sampling temperature for the model. + checkpointer (InMemorySaver): Per-thread conversation state store. + system_prompt (str): The (already formatted) system prompt. + structured (bool): When True, attach the StructuredResponse response_format; + when False (the free-text baseline), the agent answers in its final message. + + Returns: + CompiledStateGraph: The compiled agent graph ready for `ainvoke`. + """ + model = init_chat_model( + model=settings.MODEL_URI, + temperature=temperature, + credentials=settings.GOOGLE_CREDENTIALS, + thinking_level=settings.THINKING_LEVEL, + include_thoughts=True, + ) + middleware = [ + SummarizationMiddleware( + model=model, + trigger=("tokens", SUMMARIZE_TRIGGER_TOKENS), + keep=("tokens", SUMMARIZE_KEEP_TOKENS), + ), + ModelCallLimitMiddleware( + run_limit=MODEL_CALL_RUN_LIMIT, + exit_behavior="end", + ), + ] + # No response_format on the free-text baseline (--no-structured); the agent answers + # in the final message instead of a StructuredResponse. + response_format = {"response_format": StructuredResponse} if structured else {} + return create_agent( + model=model, + tools=BDToolkit.get_tools(), + system_prompt=system_prompt, + middleware=middleware, + checkpointer=checkpointer, + **response_format, + ) + + +# ============================================================================= +# Per-turn extraction from the agent's final state +# ============================================================================= +def _empty_record(status: str, **extra) -> dict: + """Base per-turn record for the non-`ok` cases (no_structured / error). + + The scoring fields are zeroed so downstream code can read them uniformly. + + Args: + status (str): The turn status ("no_structured" or "error"). + **extra: Fields to override or add (e.g. an `error` message, a recovered + `response_text`). + + Returns: + dict: The per-turn record with default (empty) fields plus `extra`. + """ + return { + "status": status, + "is_query": False, + "tables": [], + "table_period_ends": {}, + "structured": None, + "response_text": None, + "model_calls": None, + "tools_used": [], + "queries": [], + **extra, + } + + +def _last_ai_text(messages: list[AnyMessage]) -> str | None: + """The text of the last AI message (the free-text answer). + + Args: + messages (list[AnyMessage]): The full thread's messages. + + Returns: + str | None: The most recent AI message's text, or None if there is none. + """ + for message in reversed(messages): + if message.type == "ai": + return message.text + return None + + +def _tools_used(messages: list[AnyMessage]) -> set[str]: + """Distinct tool names the agent called in THIS turn. + + The checkpointer returns the whole thread, so this slices to the messages after + the last human turn. + + Args: + messages (list[AnyMessage]): The full thread's messages. + + Returns: + set[str]: The distinct tool names invoked this turn. + """ + last_human_index = max( + (i for i, message in enumerate(messages) if message.type == "human"), default=-1 + ) + return { + tool_call["name"] + for message in messages[last_human_index + 1 :] + for tool_call in getattr(message, "tool_calls", None) or [] + } + + +def _tool_metadata(messages: list[AnyMessage]) -> tuple[dict[str, str], dict[str, str]]: + """Resolve table identity/coverage metadata from the agent's exploration calls. + + From the agent's get_table_details / get_dataset_details outputs, build + (uuid -> gcp_id) to resolve data_sources, and (gcp_id -> period_end) for the + `latest` period rule. A failed tool call serializes to a ToolError object + ({"status": "error", ...}) rather than the expected payload, so skip those. + + Args: + messages (list[AnyMessage]): The full thread's messages. + + Returns: + tuple[dict[str, str], dict[str, str]]: (uuid_to_gcp, gcp_id_to_period_end). + """ + uuid_to_gcp: dict[str, str] = {} + period_end: dict[str, str] = {} + for message in messages: + if message.type != "tool" or message.name not in ( + "get_table_details", + "get_dataset_details", + ): + continue + try: + payload = json.loads(message.content) + except json.JSONDecodeError: + continue + # A failed tool call serializes to a ToolError object ({"status": "error", ...}), + # which lacks the id/tables keys — skip it. A success payload has every key. + if payload.get("status") == "error": + continue + if message.name == "get_table_details": + uuid_to_gcp[payload["id"]] = payload["gcp_id"] + if payload["period_end"] is not None: + period_end[payload["gcp_id"]] = payload["period_end"] + else: # get_dataset_details + for table in payload["tables"]: + uuid_to_gcp[table["id"]] = table["gcp_id"] + return uuid_to_gcp, period_end + + +def _truncate_rows(content: str) -> tuple[int | None, list | None, str | None]: + """Parse an execute_bigquery_sql result body into (row_count, capped_rows, message). + + A JSON array -> (total count, first QUERY_ROWS_CAP rows, None). Anything else — the + '0 rows' notice, a plain error string, or a JSON object (a serialized ToolError from a + failed query) -> (None, None, the raw string), so a failed query is kept as a message, + not a crash. + + Args: + content (str): The raw tool-message content. + + Returns: + tuple[int | None, list | None, str | None]: (row_count, capped_rows, message). + """ + try: + parsed = json.loads(content) + except json.JSONDecodeError: + return None, None, content + if not isinstance(parsed, list): + return None, None, content + return len(parsed), parsed[:QUERY_ROWS_CAP], None + + +def _executed_queries(messages: list[AnyMessage]) -> list[dict]: + """The execute_bigquery_sql calls made in THIS turn, each paired with its result rows. + + This is the agent's OWN retrieved data. It lets the judge score `grounded` against + what the agent actually queried (not only the gold reference). Sliced to the messages + after the last human turn, like _tools_used. + + Args: + messages (list[AnyMessage]): The full thread's messages. + + Returns: + list[dict]: One {sql, status, row_count, rows, message} per execute_bigquery_sql + call made this turn. + """ + last_human_index = max( + (i for i, message in enumerate(messages) if message.type == "human"), default=-1 + ) + turn_messages = messages[last_human_index + 1 :] + sql_by_id = { + tool_call["id"]: tool_call["args"].get("sql_query") + for message in turn_messages + for tool_call in getattr(message, "tool_calls", None) or [] + if tool_call["name"] == "execute_bigquery_sql" + } + queries = [] + for tool_message in turn_messages: + if tool_message.type == "tool" and tool_message.name == "execute_bigquery_sql": + row_count, rows, message = _truncate_rows(tool_message.content) + queries.append( + { + "sql": sql_by_id.get(tool_message.tool_call_id), + "status": tool_message.status, + "row_count": row_count, + "rows": rows, + "message": message, + } + ) + return queries + + +def extract_turn(result: dict, structured_mode: bool = True) -> dict: + """Build a per-turn record from the agent's final state. + + Args: + result (dict): The agent's `ainvoke` result (its `messages` and, in structured + mode, its `structured_response`). + structured_mode (bool): When False (the free-text baseline), read the answer from + the final message and leave source/period unscored; when True, read the + structured fields. + + Returns: + dict: The per-turn record (status "ok", or "no_structured" when the structured + agent produced no structured response). + """ + messages = result["messages"] + + # --no-structured (e.g. the main branch): the answer is the final AI message and there + # are no structured fields. is_query is derived from the tool trace (did it run a query + # this turn); source/period are left unscored (no data_sources / temporal_coverage). + if not structured_mode: + queries = _executed_queries(messages) + # period_end (per gcp_id) still comes from the trace's get_table_details outputs, + # so a trace-based eval (eval_queries) can do the `latest` year check for main too. + _, table_period_ends = _tool_metadata(messages) + return { + "status": "ok", + "is_query": bool(queries), + "tables": [], + "table_period_ends": table_period_ends, + "structured": None, + "response_text": _last_ai_text(messages), + "model_calls": result.get("run_model_call_count"), + "tools_used": sorted(_tools_used(messages)), + "queries": queries, + } + + structured = result.get("structured_response") + + # This eval targets the structured-output agent: a missing structured response + # is a hard fail for the turn (not a crash, not a skip). + if structured is None: + return _empty_record( + "no_structured", + response_text=_last_ai_text(messages), + model_calls=result.get("run_model_call_count"), + ) + + uuid_to_gcp, table_period_ends = _tool_metadata(messages) + # Tables the agent REPORTED in data_sources, UUID -> gcp_id (unmapped UUIDs are + # kept as-is so they fail to match the gold and surface the problem). + data_source_tables = sorted( + { + uuid_to_gcp.get(source.table_id, source.table_id) + for source in (structured.data_sources or []) + } + ) + return { + "status": "ok", + # query vs clarify, from whether the agent filled the sql_queries field. + "is_query": bool(structured.sql_queries), + "tables": data_source_tables, + "table_period_ends": table_period_ends, + "structured": structured.model_dump(mode="json"), + "response_text": structured.response, + "model_calls": result.get("run_model_call_count"), + "tools_used": sorted(_tools_used(messages)), + "queries": _executed_queries(messages), + } + + +# ============================================================================= +# Deterministic scoring +# ============================================================================= +def _lead_year(value) -> int | None: + """The leading 4-digit year of a period value ('2015', '2015-03-01', '2026-05'). + + Args: + value: The period value (str or None). + + Returns: + int | None: The leading year, or None if the value has no leading 4 digits. + """ + match = re.match(r"\s*(\d{4})", str(value or "")) + return int(match.group(1)) if match else None + + +def _tc_years(temporal_coverage: dict | None) -> tuple[int | None, int | None]: + """The (start_year, end_year) of a temporal_coverage dict. + + Uses the leading 4-digit year of period_start/period_end (handles '2015' and + '2015-03-01' alike). + + Args: + temporal_coverage (dict | None): The reported temporal_coverage. + + Returns: + tuple[int | None, int | None]: (start_year, end_year), each None when absent. + """ + if not temporal_coverage: + return None, None + return ( + _lead_year(temporal_coverage.get("period_start")), + _lead_year(temporal_coverage.get("period_end")), + ) + + +def _parse_period(period) -> tuple[str, str] | None: + """Split a period value into (granularity, value) in its granularity's own format. + + 'YYYY' -> ('year', ...), 'YYYY-MM' -> ('month', ...), 'YYYY-MM-DD' -> ('day', ...). + Strict (re.fullmatch): a malformed value like '2026-5' returns None -> a miss, not + something to silently repair. + + Args: + period: The period value (str or None). + + Returns: + tuple[str, str] | None: (granularity, value), or None when malformed. + """ + period = str(period or "").strip().strip("'").strip('"') + for granularity, pattern in ( + ("day", r"\d{4}-\d{2}-\d{2}"), + ("month", r"\d{4}-\d{2}"), + ("year", r"\d{4}"), + ): + if re.fullmatch(pattern, period): + return granularity, period + return None + + +def _score_period(record: dict, gold: dict, prev_query: dict | None): + """Score the period from the agent's reported `temporal_coverage`. + + Normalized and independent of SQL surface form / temporal-column name. + + Args: + record (dict): The per-turn record (its `structured`, `tables`, `table_period_ends`). + gold (dict): The gold turn (its `action` and `temporal` rule). + prev_query (dict | None): The previous query turn's record, for `match_previous`. + + Returns: + float | None: Score in [0, 1], or None when the rule doesn't apply / can't be checked. + """ + rule = gold.get("temporal") + if gold["action"] != "query" or rule in (None, "none", "any", "", {}): + return None + start_year, end_year = _tc_years(record["structured"]["temporal_coverage"]) + if start_year is None or end_year is None: + return 0.0 # queried but reported no usable temporal_coverage + if rule == "match_previous": + if prev_query is None: + return None + prev_start_year, prev_end_year = _tc_years( + prev_query["structured"]["temporal_coverage"] + ) + if prev_start_year is None or prev_end_year is None: + return 0.0 + # The follow-up must stay WITHIN the previous turn's period. A justified + # narrowing (e.g., the requested metric only exists for part of the range) + # is fine; only drifting OUTSIDE the established window fails. + return ( + 1.0 if prev_start_year <= start_year and end_year <= prev_end_year else 0.0 + ) + if rule == "range": + return 1.0 if end_year > start_year else 0.0 + if rule == "latest": + # Expected = the table's period_end, matched EXACTLY at its own granularity: + # '2026-05' must be reported as month 2026-05 (not year 2026), a full date to + # the day. The model must emit the value in its granularity's format. + queried_tables = set(record["tables"]) + candidates = [ + period_value + for gcp_id, period_value in record["table_period_ends"].items() + if gcp_id in queried_tables + ] + candidates = candidates or list(record["table_period_ends"].values()) + targets = [ + parsed + for period_value in candidates + if (parsed := _parse_period(period_value)) is not None + ] + if not targets: + return None # can't determine the table's period_end from the trace + target = max(targets, key=lambda parsed: parsed[1]) # latest period_end + temporal_coverage = record["structured"]["temporal_coverage"] or {} + start, end = ( + _parse_period(temporal_coverage.get("period_start")), + _parse_period(temporal_coverage.get("period_end")), + ) + matches = ( + start == target + and end == target + and temporal_coverage.get("granularity") == target[0] + ) + return 1.0 if matches else 0.0 + span = rule.get("exact", rule) if isinstance(rule, dict) else None + if isinstance(span, dict) and "start" in span and "end" in span: + return ( + 1.0 + if start_year == int(span["start"]) and end_year == int(span["end"]) + else 0.0 + ) + return None + + +def _prose_format_ok(text: str | None) -> bool: + """Whether the prose obeys the one hard-forbidden format rule: no Markdown ATX headers. + + Header lines (`#`..`######`) inside fenced code blocks are ignored (a `#` comment in a + code sample is not a section title). + + Args: + text (str | None): The prose answer. + + Returns: + bool: True if the prose contains no Markdown ATX header outside a code fence. + """ + in_code_fence = False + for line in (text or "").splitlines(): + if _CODE_FENCE_RE.match(line): + in_code_fence = not in_code_fence + continue + if not in_code_fence and _MD_HEADER_RE.match(line): + return False + return True + + +def score_turn(record: dict, gold: dict, prev_query: dict | None) -> dict: + """Score one turn's deterministic dimensions against the gold expectation. + + Args: + record (dict): The per-turn record from extract_turn / _empty_record. + gold (dict): The gold turn (its `action`, `tables`, `temporal`). + prev_query (dict | None): The previous query turn's record, for `match_previous`. + + Returns: + dict: The per-dimension scores (True/False/float/None) for action, source, period, + format_ok, completed. + """ + if record["status"] == "no_structured": + # agent returned no structured response -> the structured contract failed + return { + "action": False, + "source": None, + "period": None, + "format_ok": None, + "completed": False, + } + + # --no-structured mode: an ok turn with no structured fields (structured is None only + # here; a structured-mode failure has status "no_structured"). Only action (from the + # tool trace), format_ok (prose) and completed are scoreable; source/period need the + # data_sources / temporal_coverage fields the agent doesn't produce. + if record["structured"] is None: + return { + "action": ("query" if record["is_query"] else "clarify") == gold["action"], + "source": None, + "period": None, + "format_ok": _prose_format_ok(record["response_text"]), + "completed": record["status"] == "ok", + } + + observed = "query" if record["is_query"] else "clarify" + + if gold["action"] == "query" and gold.get("tables"): + expected_tables, observed_tables = set(gold["tables"]), set(record["tables"]) + source = ( + 1.0 + if observed_tables == expected_tables + else (0.5 if observed_tables & expected_tables else 0.0) + ) + else: + source = None # clarify turn, or gcp_ids not filled in yet + return { + "action": observed == gold["action"], + "source": source, + "period": _score_period(record, gold, prev_query), + "format_ok": _prose_format_ok(record["structured"]["response"]), + "completed": record["status"] == "ok", + } + + +# ============================================================================= +# Thread replay +# ============================================================================= +def _turn_run_config( + base: dict, + *, + thread_id: str, + thread: dict, + repeat: int, + turn_index: int, + turn: dict, + branch: str, + temperature: float, +) -> dict: + """Build the per-turn LangGraph config. + + Labels the run so it's findable in LangSmith; `thread_id` in metadata makes LangSmith + group a replay's turns into one conversation. + + Args: + base (dict): The shared base run config (e.g. an optional recursion_limit). + thread_id (str): The per-replay thread id (thread + repeat). + thread (dict): The gold thread being replayed. + repeat (int): The repeat index. + turn_index (int): The zero-based turn index within the thread. + turn (dict): The gold turn (for the expected_action / expected_temporal tags). + branch (str): The checked-out git branch. + temperature (float): The sampling temperature in effect. + + Returns: + dict: The LangGraph run config for this turn. + """ + return { + **base, + "configurable": {"thread_id": thread_id}, + "run_name": f"{thread['id']}#{repeat}-t{turn_index}", + "tags": [ + f"branch:{branch}", + f"temp:{temperature}", + f"thread:{thread['id']}", + ], + "metadata": { + "thread_id": thread_id, + "eval_thread": thread["id"], + "repeat": repeat, + "turn_index": turn_index, + "branch": branch, + "temperature": temperature, + "expected_action": turn["action"], + "expected_temporal": turn.get("temporal"), + }, + } + + +def _skipped_record(turn_index: int, user: str) -> dict: + """A placeholder record for a turn that never ran (the thread was truncated earlier). + + Args: + turn_index (int): The zero-based turn index. + user (str): The user message for the skipped turn. + + Returns: + dict: The skipped-turn record. + """ + return {"turn_index": turn_index, "user": user, "status": "skipped", "scores": {}} + + +async def replay_thread( + agent: CompiledStateGraph, + thread: dict, + repeat: int, + run_config: dict, + branch: str, + temperature: float, + max_turns: int | None = None, + structured: bool = True, +) -> list: + """Replay one thread turn-by-turn on a shared thread_id and score each turn. + + Args: + agent (CompiledStateGraph): The compiled agent to invoke. + thread (dict): The gold thread (its `id` and `turns`). + repeat (int): The repeat index for this replay. + run_config (dict): The shared base run config. + branch (str): The checked-out git branch (for run labels). + temperature (float): The sampling temperature in effect. + max_turns (int | None): Run only the first N turns (a prefix), or all when None. + structured (bool): Whether the agent produces structured output. + + Returns: + list: One record per turn (ok / no_structured / error / skipped). + """ + thread_id = f"{thread['id']}-{repeat}" + records, prev_query = [], None + turns = thread["turns"] if max_turns is None else thread["turns"][:max_turns] + + for turn_index, turn in enumerate(turns): + turn_config = _turn_run_config( + run_config, + thread_id=thread_id, + thread=thread, + repeat=repeat, + turn_index=turn_index, + turn=turn, + branch=branch, + temperature=temperature, + ) + + # Separate the two failure modes: an ainvoke failure means the agent/thread state + # may be broken (skip the rest of the thread), whereas an extract_turn failure is + # OUR parsing bug on an otherwise-successful run (record it, but keep going — + # the thread state is intact, so follow-ups can still run). + agent_failed = False + try: + result = await agent.ainvoke( + {"messages": [{"role": "user", "content": turn["user"]}]}, + config=turn_config, + ) + except Exception as exc: + agent_failed = True + record = _empty_record( + "error", + error=f"{type(exc).__name__}: {exc}", + traceback=traceback.format_exc(), + ) + else: + try: + record = extract_turn(result, structured) + except Exception as exc: + record = _empty_record( + "error", + error=f"extract_turn: {type(exc).__name__}: {exc}", + traceback=traceback.format_exc(), + ) + + record["turn_index"], record["user"] = turn_index, turn["user"] + record["scores"] = ( + score_turn(record, turn, prev_query) if record["status"] != "error" else {} + ) + records.append(record) + + if agent_failed: # broken agent state -> skip the rest of the thread + records.extend( + _skipped_record(j, turns[j]["user"]) + for j in range(turn_index + 1, len(turns)) + ) + break + + if record["is_query"]: + prev_query = record + + return records + + +# ============================================================================= +# Aggregation & Reporting +# ============================================================================= +def _mark(scores: dict) -> str: + """Render a turn's scores as a compact one-line marks string (a:… s:… …). + + Args: + scores (dict): The per-dimension scores for one turn. + + Returns: + str: The marks line, e.g. "a:✓ s:~ p:✗ f:✓ c:✓". + """ + + def symbol(value): + """Map a single score value to its display symbol. + + Args: + value: True/False/float/None score. + + Returns: + str: "·" (n/a), "✓" (pass), "~" (partial), or "✗" (fail). + """ + return ( + "·" + if value is None + else ( + "✓" if value is True or value == 1.0 else ("~" if value == 0.5 else "✗") + ) + ) + + return ( + f"a:{symbol(scores.get('action'))} " + f"s:{symbol(scores.get('source'))} " + f"p:{symbol(scores.get('period'))} " + f"f:{symbol(scores.get('format_ok'))} " + f"c:{symbol(scores.get('completed'))}" + ) + + +def _unit_line(turn_records: list) -> str: + """A one-line per-thread progress summary: each turn's marks, or its status. + + Args: + turn_records (list): The records for one replayed thread. + + Returns: + str: The joined per-turn summary. + """ + return " ".join( + f"t{record['turn_index']}:" + + ( + _mark(record["scores"]) + if record["status"] == "ok" + else record["status"].upper() + ) + for record in turn_records + ) + + +def aggregate(units: list) -> dict: + """Compute the pass-rate per dimension over all scored turns. + + None-valued scores (not applicable / skipped) are excluded from the rate. + + Args: + units (list): The per-thread replay units (each with its `turns`). + + Returns: + dict: Per-dimension {"rate", "n"} plus `_errors` / `_skipped` / `_no_structured` counts. + """ + per_dimension = defaultdict(lambda: {"hit": 0.0, "n": 0}) + errors = skipped = no_structured = 0 + for unit in units: + for record in unit["turns"]: + if record["status"] == "skipped": + skipped += 1 + continue + if record["status"] == "error": + errors += 1 + elif record["status"] == "no_structured": + no_structured += 1 + for dimension, value in (record.get("scores") or {}).items(): + if value is None: + continue + per_dimension[dimension]["hit"] += float(value) + per_dimension[dimension]["n"] += 1 + summary = { + dimension: {"rate": round(counts["hit"] / counts["n"], 3), "n": counts["n"]} + for dimension, counts in per_dimension.items() + } + summary["_errors"], summary["_skipped"], summary["_no_structured"] = ( + errors, + skipped, + no_structured, + ) + return summary + + +def print_scorecard(summary: dict) -> None: + """Print the deterministic scorecard (pass-rate per dimension + status counts). + + Args: + summary (dict): The aggregate() result. + + Returns: + None + """ + print("\n=== Deterministic scorecard (pass-rate over applicable turns) ===") + for dimension in SCORE_DIMENSIONS: + if dimension in summary: + print( + f" {dimension:<10} {summary[dimension]['rate']:.0%} (n={summary[dimension]['n']})" + ) + print( + f" errors={summary['_errors']} skipped={summary['_skipped']} " + f"no_structured={summary['_no_structured']}" + ) + + +# ============================================================================= +# CLI +# ============================================================================= +def parse_args() -> argparse.Namespace: + """Parse the command-line arguments. + + Returns: + argparse.Namespace: The parsed arguments. + """ + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--gold", default=str(EVAL_DIR / "eval_gold.yaml")) + parser.add_argument("--repeats", type=int, default=3, help="Replays per thread") + parser.add_argument( + "--thread", action="append", help="Only this thread id (repeatable)" + ) + parser.add_argument( + "--max-turns", + type=int, + default=None, + help="Run only the first N turns of each thread (a prefix — follow-ups need prior context)", + ) + parser.add_argument("--temperature", type=float, default=None) + parser.add_argument( + "--recursion-limit", + type=int, + default=None, + help="Omit to match production (LangGraph default 25)", + ) + parser.add_argument( + "--concurrency", type=int, default=1, help="Parallel thread replays" + ) + parser.add_argument("--out", default=None) + parser.add_argument( + "--langsmith-project", + default=None, + help="LangSmith project for these traces (default: -eval)", + ) + parser.add_argument( + "--no-trace", action="store_true", help="Disable LangSmith tracing" + ) + parser.add_argument( + "--no-structured", + action="store_true", + help="Eval a free-text agent (e.g. the main branch): build without response_format, " + "read the answer from the final message. source/period are left unscored.", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print the run plan (threads/turns/repeats, total agent runs) and exit without calling the agent", + ) + return parser.parse_args() + + +def configure_tracing(args: argparse.Namespace) -> tuple[str, bool]: + """Forward LangSmith config from `settings` into os.environ. + + LangChain's tracer reads the env, but pydantic-settings only populates `settings`. + This isolates eval runs in their own project by default. Must run before the agent does. + + Args: + args (argparse.Namespace): The parsed CLI args (its langsmith_project / no_trace). + + Returns: + tuple[str, bool]: (langsmith_project, tracing_enabled). + """ + ls_project = args.langsmith_project or f"{settings.LANGSMITH_PROJECT}-eval" + tracing = settings.LANGSMITH_TRACING and not args.no_trace + os.environ["LANGSMITH_TRACING"] = "true" if tracing else "false" + if tracing: + os.environ["LANGSMITH_PROJECT"] = ls_project + os.environ["LANGSMITH_API_KEY"] = settings.LANGSMITH_API_KEY + return ls_project, tracing + + +async def main() -> None: + """Parse args, replay every gold thread K times, score each turn, and write the report.""" + args = parse_args() + + structured_mode = not args.no_structured + if structured_mode and StructuredResponse is None: + raise SystemExit( + "No StructuredResponse importable (are you on the main branch?). " + "Re-run with --no-structured to eval a free-text agent." + ) + + with open(args.gold) as f: + gold = yaml.safe_load(f) + threads = [ + thread for thread in gold if not args.thread or thread["id"] in args.thread + ] + + temperature = ( + args.temperature if args.temperature is not None else settings.MODEL_TEMPERATURE + ) + run_config: dict = {} + if args.recursion_limit is not None: + run_config["recursion_limit"] = args.recursion_limit + + branch = current_branch() + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + out_path = args.out or str( + EVAL_DIR + / f"thread_eval_{branch.replace('/', '-')}_temp{temperature}_{timestamp}.json" + ) + + ls_project, tracing = configure_tracing(args) + + print( + f"\nbranch={branch!r} model={settings.MODEL_URI!r} structured_output={structured_mode}\n" + f"temperature={temperature} repeats={args.repeats} threads={[thread['id'] for thread in threads]}\n" + f"langsmith: {f'project={ls_project!r}' if tracing else 'disabled'}\n" + ) + + if args.dry_run: + max_turns = args.max_turns + turn_counts = { + thread["id"]: min(len(thread["turns"]), max_turns) + if max_turns + else len(thread["turns"]) + for thread in threads + } + total = args.repeats * sum(turn_counts.values()) + print("DRY RUN — no agent calls will be made.") + for thread_id, count in turn_counts.items(): + print(f" {thread_id:<16} {count} turn(s) x {args.repeats} repeat(s)") + print( + f"\n total agent runs: {total} (each = one live multi-step agent invocation)" + ) + print(f" output would be: {out_path}") + return + + system_prompt = SYSTEM_PROMPT.format(current_date=date.today().isoformat()) + agent = build_agent(temperature, InMemorySaver(), system_prompt, structured_mode) + semaphore = asyncio.Semaphore(args.concurrency) + work = [(thread, repeat) for thread in threads for repeat in range(args.repeats)] + + async def run_unit(thread, repeat): + async with semaphore: + turn_records = await replay_thread( + agent, + thread, + repeat, + run_config, + branch, + temperature, + args.max_turns, + structured_mode, + ) + print(f"[{thread['id']:<16} #{repeat}] {_unit_line(turn_records)}") + return {"thread": thread["id"], "repeat": repeat, "turns": turn_records} + + units = await asyncio.gather(*(run_unit(thread, repeat) for thread, repeat in work)) + summary = aggregate(units) + print_scorecard(summary) + + report = { + "branch": branch, + "structured_output": structured_mode, + "langsmith_project": ls_project if tracing else None, + "model": settings.MODEL_URI, + "temperature": temperature, + "recursion_limit": args.recursion_limit, + "repeats": args.repeats, + "timestamp": timestamp, + "system_prompt": system_prompt, + "summary": summary, + "units": units, + } + + with open(out_path, "w") as file: + json.dump(report, file, ensure_ascii=False, indent=2) + + print(f"\nwrote {out_path}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/eval/eval_quality.py b/eval/eval_quality.py new file mode 100644 index 0000000..5fc19cb --- /dev/null +++ b/eval/eval_quality.py @@ -0,0 +1,670 @@ +"""LLM judge over the transcripts that eval_output.py saved. + +One judge call per answered turn (query OR clarify), emitting SEPARATE scores +(not a blended verdict): + + correct : the answer's figures/entities/ranking are CONSISTENT with the + reference_sql result (the gold, re-run live) — a different but + valid slice/period/grouping isn't penalized, only contradictions. + null when the turn carries no reference_sql (correctness unverifiable) + grounded : every claim traces to data the AGENT ACTUALLY RETRIEVED, this turn + OR an earlier turn (its query results across the thread, persisted in + the transcript) — an extra metric absent from the reference but + present in the agent's own results is grounded; a fabricated value is not + answers_question : the prose addresses what the user actually asked + stated_assumption : on an ambiguous turn, the interpretation is made explicit + (null when the turn wasn't ambiguous) + +The two data axes are DECOUPLED: `grounded` is anchored on the agent's own retrieved +data, `correct` on the gold reference (and only scored when a reference_sql is present); +the rest are rubric. Prose format (no Markdown headers) is checked deterministically in +eval_output, not here. It reads a transcript JSON (--in) + the gold (for reference_sql), +executes each distinct reference_sql once (cached) via the project's BQ client, then judges. + +Repeats with IDENTICAL agent output are judged once and weighted by their count, +so a 20-repeat temp-0 transcript costs ~1 judge call per turn, not 20. + +Pipeline: run AFTER eval_output.py — it scores the transcript eval_output.py produced. +This is the only post-hoc scorer that needs an LLM (the judge) and BQ (to run +reference_sql). Sibling scorers over the same transcript, independent of this one and of +each other: + eval_faithfulness.py structured output's self-consistency — gold-free, no LLM/BQ + eval_queries.py source/period from the executed SQL vs the gold — no LLM/BQ + + uv run eval/eval_quality.py --in eval/.json --dry-run + uv run eval/eval_quality.py --in eval/.json --judge-model google_genai:gemini-3.1-pro-preview + +NOTE: the judge should ideally be a stronger / different-family model than the +agent (less self-preference). Default --judge-model is google_genai:gemini-3.1-pro-preview. +""" + +import argparse +import asyncio +import hashlib +import json +import os +import sys +from collections import defaultdict +from pathlib import Path + +import yaml +from google.cloud import bigquery as bq +from langchain.chat_models import init_chat_model +from langchain.messages import HumanMessage, SystemMessage +from pydantic import BaseModel, Field + +# Make the repo root importable so `app` resolves whether this file is run as a +# module (python -m eval.eval_quality) or directly (python eval/eval_quality.py). +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from app.agent.tools.bigquery import MAX_BYTES_BILLED, _get_client # noqa: E402 +from app.settings import settings # noqa: E402 + +# This script's folder — gold input and result files default here, so the eval +# works regardless of the current working directory. +EVAL_DIR = Path(__file__).resolve().parent + +SCORE_FIELDS = ( + "correct", + "grounded", + "answers_question", + "stated_assumption", +) +MAX_REFERENCE_ROWS = 50 # cap rows shown to the judge to bound tokens + +# Dry-run cost estimate only (a real run reports exact usage). Gemini measures +# ~3.5 chars/token on this mixed PT-prose/JSON/SQL content; a verdict is small. +DRY_CHARS_PER_TOKEN = 3.5 +DRY_VERDICT_TOKENS = 100 + + +# ============================================================================= +# Judge schema & prompt +# ============================================================================= +class JudgeVerdict(BaseModel): + correct: bool | None = Field( + default=None, + description="(query turns WITH a reference result) Are the main figures/entities/ranking CONSISTENT with the REFERENCE RESULT? A different-but-valid slice/grouping/period the user didn't specify, extra metrics, rounding, or language differences are OK — fail ONLY on a direct contradiction (a different value for the same quantity, wrong entity, wrong ordering). null on clarify turns AND when no reference result is provided (correctness can't be verified).", + ) + grounded: bool = Field( + description="Does every quantitative/factual claim trace to data the assistant ACTUALLY RETRIEVED (this turn OR an earlier turn of the conversation)? Check the prose against the ASSISTANT'S RETRIEVED DATA section (its own query results), NOT the reference — a value present there is grounded even if absent from the reference. Inventing numbers/trends/datasets, or asserting a value with no supporting query, fails. (clarify: backed by real exploration, no assumed value)." + ) + answers_question: bool = Field( + description="Does the answer address what the user needed this turn? (query: answers the question; clarify: does the right clarification/guidance per the note)." + ) + stated_assumption: bool | None = Field( + default=None, + description="If the request was ambiguous (e.g., an unspecified breakdown/interpretation), did the answer make its chosen interpretation explicit? null if there was NO ambiguity.", + ) + rationale: str = Field( + description="1-2 sentence justification; name the main discrepancy, if any." + ) + + +JUDGE_SYSTEM = """\ +You are a strict evaluator of a Brazilian open-data assistant's answer, for ONE turn of a conversation. The assistant answer and the data are in Portuguese. + +You receive: the turn's expected type, the conversation so far, an internal note (what the turn tests), the tools the assistant used, the ASSISTANT'S RETRIEVED DATA (the results of its OWN SQL, from this turn AND earlier turns of the conversation — a follow-up may rely on data queried in a previous turn), its final answer (prose + the structured fields it reported), and — WHEN AVAILABLE — a REFERENCE RESULT from a trusted SQL query (the gold). Some data turns have no reference; the deterministic eval covers their inputs (right tables/period) separately. + +You have TWO independent data anchors — do NOT conflate them: +- The ASSISTANT'S RETRIEVED DATA anchors `grounded`: did every claim come from data the assistant actually queried (at any point in the conversation)? +- The REFERENCE RESULT anchors `correct`: is the answer consistent with the gold? (only when a reference is present) + +Score each criterion as a boolean (or null when it does not apply). + +If the type is `query` (it should answer with data): +- correct: the main figures/entities/ordering are CONSISTENT with the reference. A different-but-valid decomposition, extra metrics, or a different period the user did NOT specify are NOT failures — fail only on a direct contradiction (a different value for the SAME quantity, a wrong entity, wrong ordering). A value the assistant computed that the reference simply doesn't contain is NOT a correctness failure — judge it under `grounded` instead. If the assistant's scope/period differs so much that its figures aren't comparable to the reference, don't invent a contradiction: lean on whether the trend/entities are consistent. If NO REFERENCE RESULT is provided for this turn, set correct=null (you cannot verify correctness) and judge only `grounded` and `answers_question`. +- grounded: every quantitative/factual claim traces to the ASSISTANT'S RETRIEVED DATA shown to you (from this turn OR an earlier turn — a follow-up legitimately reuses data queried before). A number that appears in the assistant's own query results IS grounded, even if it's absent from the reference. Fabricating figures/trends/comparisons that are in NEITHER the assistant's results nor a legitimate calculation over them fails. +- answers_question: the prose addresses what the user asked in this turn. +- stated_assumption: if the request was ambiguous, it made its chosen interpretation explicit; null if there was no ambiguity. + +If the type is `clarify` (it should NOT query data; per the note, it should ask for the missing detail OR explore the catalog and guide): +- correct: null (there is no data answer to verify). +- grounded: the answer does NOT invent datasets/tables/values. If it describes available data, that must be backed by real exploration — check the tools used (search_datasets/get_dataset_details/get_table_details). If it assumes a value the user did not provide (e.g., a specific município), grounded=false. +- answers_question: it does the right clarification/guidance per the note (e.g., asks which município; or describes what exists and suggests specific refinements), without having queried data. +- stated_assumption: null. + +Be strict but fair: wording/rounding differences, extra valid metrics, and a differently-but-validly-scoped query are acceptable; a wrong value for the same quantity, wrong entities, claims absent from the assistant's retrieved data (fabrications), unrequested assumptions, or not doing what the turn required are failures. Give a one- to two-sentence rationale.""" + + +# ============================================================================= +# Gold & reference SQL +# ============================================================================= +def load_gold_refs(path: str) -> dict[tuple[str, int], dict]: + """Load the gold spec, keyed by (thread id, turn index). + + Args: + path (str): Path to the gold YAML file. + + Returns: + dict[tuple[str, int], dict]: Maps (thread_id, turn_index) to that turn's gold dict + (which carries reference_sql and notes). + """ + gold = yaml.safe_load(open(path)) + return { + (thread["id"], turn_index): turn + for thread in gold + for turn_index, turn in enumerate(thread["turns"]) + } + + +def run_reference_sql(sql: str, cache: dict[str, list]) -> list[dict]: + """Execute a reference SQL against BigQuery, caching the result by SQL text. + + Args: + sql (str): The reference SQL to run. + cache (dict[str, list]): SQL-text -> rows cache, mutated in place so each distinct + reference SQL runs at most once. + + Returns: + list[dict]: The query result rows (each row as a dict). + """ + if sql not in cache: + job = _get_client().query( + sql, job_config=bq.QueryJobConfig(maximum_bytes_billed=MAX_BYTES_BILLED) + ) + cache[sql] = [dict(row) for row in job.result()] + return cache[sql] + + +# ============================================================================= +# Task building & rendering +# ============================================================================= +def _render_agent_queries(queries: list[dict]) -> str: + """Render the assistant's own executed queries + result rows for the judge prompt. + + Args: + queries (list[dict]): Executed query records (sql / status / rows / row_count / + message). + + Returns: + str: A human-readable block (the anchor for `grounded`), or an N/A notice when the + assistant executed no SQL this turn or in any earlier turn. + """ + if not queries: + return "N/A — the assistant executed no SQL this turn or in any earlier turn." + blocks = [] + for index, query in enumerate(queries, 1): + sql = (query.get("sql") or "").strip() + if query.get("rows") is not None: + rows, total = query["rows"], query.get("row_count") + omitted = ( + f"\n(... showing {len(rows)} of {total} rows)" + if total and total > len(rows) + else "" + ) + body = json.dumps(rows, ensure_ascii=False, default=str, indent=2) + omitted + else: + body = query.get("message") or "(no result)" + blocks.append( + f"## Assistant query {index} (status: {query.get('status')})\n{sql}\nResult:\n{body}" + ) + return "\n\n".join(blocks) + + +def render_turn(task: dict) -> str: + """Build the judge's human-message prompt for one task (turn). + + Args: + task (dict): A task from build_tasks(), augmented with `reference_rows`. + + Returns: + str: The rendered turn — expected type, conversation, note, tools, the assistant's + answer + structured fields, its retrieved data, and the reference result. + """ + reference_rows = task["reference_rows"] + if reference_rows is None and task["turn_type"] == "query": + ref_block = ( + "N/A — no reference query was provided for this turn, so correctness cannot " + "be verified against a gold result. Set `correct = null` and judge only " + "`grounded` (against the assistant's retrieved data) and `answers_question`." + ) + elif reference_rows is None: + ref_block = "N/A — clarification turn; there is no data answer to verify." + else: + shown = reference_rows[:MAX_REFERENCE_ROWS] + omitted = ( + f"\n(... {len(reference_rows) - len(shown)} rows omitted; {len(reference_rows)} total)" + if len(reference_rows) > len(shown) + else "" + ) + ref_block = ( + json.dumps(shown, ensure_ascii=False, default=str, indent=2) + omitted + ) + agent_data_block = _render_agent_queries(task.get("agent_queries") or []) + structured = task["agent_structured"] + data_source_names = [d.get("name") for d in (structured.get("data_sources") or [])] + conversation = "\n".join( + f" {turn_number + 1}. {user_message}" + for turn_number, user_message in enumerate(task["user_turns"]) + ) + return f"""# Expected turn type: {task["turn_type"]} + +# Conversation (user turns so far) +{conversation} + +# This turn's question +{task["user_turns"][-1]} + +# What this turn tests (internal note) +{task["note"] or "—"} + +# Tools the assistant used this turn +{task["tools_used"]} + +# Assistant's answer (prose) +{task["agent_response"]} + +# Structured fields the assistant reported +- data_sources: {data_source_names} +- temporal_coverage: {structured.get("temporal_coverage")} +- sql_queries: {structured.get("sql_queries")} +- follow_up_questions: {structured.get("follow_up_questions")} + +# ASSISTANT'S RETRIEVED DATA (its OWN queries + results, this turn AND earlier turns of the conversation — anchor for `grounded`) +{agent_data_block} + +# REFERENCE RESULT (gold, from the trusted query — anchor for `correct`) +{ref_block}""" + + +def build_tasks(transcript: dict, gold: dict, max_repeats: int | None) -> list[dict]: + """Build one judge task per answered (`ok`) turn that has a gold expectation. + + Args: + transcript (dict): The transcript produced by eval_output. + gold (dict): (thread_id, turn_index) -> gold turn, from load_gold_refs(). + max_repeats (int | None): If set, skip repeats whose index is >= this value. + + Returns: + list[dict]: One task per judged turn, carrying the conversation so far, the agent's + answer + structured fields, the thread-cumulative executed queries, and the + reference_sql. + """ + tasks = [] + for unit in transcript["units"]: + if max_repeats is not None and unit["repeat"] >= max_repeats: + continue + user_messages = [] + thread_queries: list[ + dict + ] = [] # executed queries accumulated across the thread + for turn in unit["turns"]: + user_messages.append(turn["user"]) + thread_queries += turn.get("queries") or [] + gold_turn = gold.get((unit["thread"], turn["turn_index"])) + # Judge any turn the agent actually answered (query or clarify) for which + # we have a gold expectation. Query turns carry a reference_sql to anchor + # correctness; clarify turns are judged on the rubric only. + if gold_turn is None or turn["status"] != "ok": + continue + tasks.append( + { + "thread": unit["thread"], + "repeat": unit["repeat"], + "turn_index": turn["turn_index"], + "turn_type": gold_turn["action"], + "user_turns": list(user_messages), + "note": gold_turn.get("notes"), + "agent_response": (turn.get("structured") or {}).get("response") + or turn.get("response_text"), + "agent_structured": turn.get("structured") or {}, + "agent_queries": list( + thread_queries + ), # this turn + all earlier turns + "tools_used": turn.get("tools_used", []), + "reference_sql": gold_turn.get("reference_sql"), + } + ) + return tasks + + +# ============================================================================= +# Dedup & judging +# ============================================================================= +def _output_signature(task: dict) -> str: + """Hash the parts of a task that vary across repeats and that the judge reads. + + Covers the agent's structured fields, prose response, tools used, and retrieved query + results. `response` is included explicitly so free-text answers (no structured fields, + e.g. the main branch) still dedup by their prose rather than collapsing to one. + + Args: + task (dict): A task from build_tasks(). + + Returns: + str: A SHA-1 hex digest identifying this agent output. + """ + payload = json.dumps( + { + "structured": task["agent_structured"], + "response": task["agent_response"], + "tools_used": task["tools_used"], + "queries": task["agent_queries"], + }, + sort_keys=True, + ensure_ascii=False, + default=str, + ) + return hashlib.sha1(payload.encode()).hexdigest() + + +def dedup_tasks(tasks: list[dict]) -> list[dict]: + """Collapse repeats with identical agent output (per thread+turn) into one task each. + + Judging each distinct output once, weighted by how many repeats produced it, reflects + the full distribution without an N-x judge bill (a big win at temperature 0). + + Args: + tasks (list[dict]): Tasks from build_tasks(). + + Returns: + list[dict]: One task per distinct output, with added `weight` (repeat count) and + `repeats` (sorted repeat indices). + """ + groups: dict[tuple, dict] = {} + for task in tasks: + key = (task["thread"], task["turn_index"], _output_signature(task)) + groups.setdefault(key, {"task": task, "repeats": []})["repeats"].append( + task["repeat"] + ) + return [ + dict( + group["task"], + weight=len(group["repeats"]), + repeats=sorted(group["repeats"]), + ) + for group in groups.values() + ] + + +def build_judge(model_uri: str): + """Build the structured-output judge model. + + Args: + model_uri (str): The judge model URI (Google models get the service-account creds). + + Returns: + A LangChain model configured to return a JudgeVerdict (with the raw response + included). + """ + kwargs = {"temperature": 0} + if model_uri.startswith("google"): + kwargs["credentials"] = settings.GOOGLE_CREDENTIALS + return init_chat_model(model_uri, **kwargs).with_structured_output( + JudgeVerdict, include_raw=True + ) + + +async def judge_turn(judge, semaphore: asyncio.Semaphore, task: dict) -> dict: + """Judge one task, returning its verdict record (and printing a one-line mark). + + Args: + judge: The structured-output judge model from build_judge(). + semaphore (asyncio.Semaphore): Concurrency limiter. + task (dict): A deduped task (carries weight/repeats and reference_rows). + + Returns: + dict: The task's identity fields plus the parsed verdict (or None + an error) and + token usage. + """ + base_record = { + key: task[key] for key in ("thread", "turn_index", "weight", "repeats") + } + async with semaphore: + try: + judge_output = await judge.ainvoke( + [SystemMessage(JUDGE_SYSTEM), HumanMessage(render_turn(task))] + ) + verdict: JudgeVerdict | None = judge_output["parsed"] + usage = getattr(judge_output["raw"], "usage_metadata", None) or {} + verdict_record = { + **base_record, + "verdict": verdict.model_dump() if verdict is not None else None, + "input_tokens": usage.get("input_tokens"), + "output_tokens": usage.get("output_tokens"), + } + if verdict is None: + verdict_record["error"] = f"parse: {judge_output.get('parsing_error')}" + except Exception as exc: + verdict_record = { + **base_record, + "verdict": None, + "error": f"{type(exc).__name__}: {exc}", + } + mark = ( + "ERR" + if verdict_record["verdict"] is None + else " ".join( + f"{field[:4]}{'·' if verdict_record['verdict'][field] is None else ('✓' if verdict_record['verdict'][field] else '✗')}" + for field in SCORE_FIELDS + ) + ) + print(f" [{task['thread']:<12} t{task['turn_index']} ×{task['weight']:<2}] {mark}") + return verdict_record + + +# ============================================================================= +# Aggregation & reporting +# ============================================================================= +def aggregate(verdicts: list[dict]) -> dict: + """Compute the pass-rate per score field, weighting each verdict by its repeat count. + + Args: + verdicts (list[dict]): Verdict records from judge_turn(). + + Returns: + dict: {field: {"rate": float, "n": int}} plus "_errors" (weighted count of verdicts + that failed to parse). + """ + totals = defaultdict(lambda: {"hit": 0, "n": 0}) + errors = 0 + for verdict_record in verdicts: + weight = verdict_record.get("weight", 1) + if verdict_record["verdict"] is None: + errors += weight + continue + for field in SCORE_FIELDS: + value = verdict_record["verdict"][field] + if value is None: + continue + totals[field]["hit"] += weight * int(bool(value)) + totals[field]["n"] += weight + summary = { + field: {"rate": round(counts["hit"] / counts["n"], 3), "n": counts["n"]} + for field, counts in totals.items() + } + summary["_errors"] = errors + return summary + + +def configure_tracing(args: argparse.Namespace) -> tuple[str, bool]: + """Forward LangSmith config from `settings` into os.environ before the judge runs. + + LangChain's tracer reads the environment, but pydantic-settings only populates + `settings`. Judge runs get their own project by default so their traces don't mix with + the agent-eval traces. + + Args: + args (argparse.Namespace): Parsed CLI args (langsmith_project, no_trace). + + Returns: + tuple[str, bool]: (project name, tracing enabled). + """ + ls_project = args.langsmith_project or f"{settings.LANGSMITH_PROJECT}-eval-judge" + tracing = settings.LANGSMITH_TRACING and not args.no_trace + os.environ["LANGSMITH_TRACING"] = "true" if tracing else "false" + if tracing: + os.environ["LANGSMITH_PROJECT"] = ls_project + os.environ["LANGSMITH_API_KEY"] = settings.LANGSMITH_API_KEY + return ls_project, tracing + + +# ============================================================================= +# CLI +# ============================================================================= +async def main() -> None: + """Parse args, judge the transcript (or dry-run), and write the report.""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--in", dest="transcript", required=True, help="thread_eval_*.json" + ) + parser.add_argument("--gold", default=str(EVAL_DIR / "eval_gold.yaml")) + parser.add_argument("--judge-model", default="google_genai:gemini-3.1-pro-preview") + parser.add_argument( + "--max-repeats", + type=int, + default=None, + help="cap repeats considered before dedup (rarely needed; dedup already collapses identical outputs)", + ) + parser.add_argument("--concurrency", type=int, default=4) + parser.add_argument("--out", default=None) + parser.add_argument( + "--thread", action="append", help="Only this thread id (repeatable)" + ) + parser.add_argument( + "--price-in", + type=float, + default=2.0, + help="USD per 1M input tokens (default: gemini-3.1-pro-preview)", + ) + parser.add_argument( + "--price-out", + type=float, + default=12.0, + help="USD per 1M output tokens (default: gemini-3.1-pro-preview)", + ) + parser.add_argument( + "--langsmith-project", + default=None, + help="LangSmith project for judge traces (default: -eval-judge)", + ) + parser.add_argument( + "--no-trace", action="store_true", help="Disable LangSmith tracing" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Print how many judge calls and reference SQLs would run, then exit without calling the judge or BQ", + ) + args = parser.parse_args() + + transcript = json.load(open(args.transcript)) + if args.thread: + transcript["units"] = [ + unit for unit in transcript["units"] if unit["thread"] in args.thread + ] + gold = load_gold_refs(args.gold) + tasks = build_tasks(transcript, gold, args.max_repeats) + deduped_tasks = dedup_tasks(tasks) + + print(f"judge_model={args.judge_model!r} (agent was {transcript.get('model')!r})") + if args.judge_model == settings.MODEL_URI: + print( + " WARNING: judging with the same model as the agent — prefer a stronger/different one via --judge-model" + ) + print( + f"turn-instances={len(tasks)} -> distinct outputs to judge={len(deduped_tasks)}\n" + ) + + suffix = f"_{'-'.join(args.thread)}" if args.thread else "" + out_path = args.out or str( + EVAL_DIR / f"{Path(args.transcript).stem}{suffix}_judged.json" + ) + if args.dry_run: + reference_sql_count = len( + {task["reference_sql"] for task in deduped_tasks if task["reference_sql"]} + ) + estimated_input_tokens = ( + sum( + len(JUDGE_SYSTEM) + len(render_turn(dict(task, reference_rows=None))) + for task in deduped_tasks + ) + / DRY_CHARS_PER_TOKEN + ) + estimated_output_tokens = DRY_VERDICT_TOKENS * len(deduped_tasks) + estimated_cost = ( + estimated_input_tokens / 1e6 * args.price_in + + estimated_output_tokens / 1e6 * args.price_out + ) + print("DRY RUN — no reference SQLs executed and no judge calls made.") + print(f" judge calls that would run: {len(deduped_tasks)}") + print( + f" distinct reference SQLs that would execute (BQ): {reference_sql_count}" + ) + print( + f" est. tokens: ~{estimated_input_tokens:,.0f} in (char-based, excl. reference rows) " + f"+ ~{estimated_output_tokens:,.0f} out" + ) + print( + f" est. cost @ ${args.price_in}/M in, ${args.price_out}/M out: " + f"~${estimated_cost:.3f} (real run reports exact)" + ) + print(f" output would be: {out_path}") + return + + ls_project, tracing = configure_tracing(args) + print("langsmith: " + (f"project={ls_project!r}" if tracing else "disabled")) + + print("executing reference SQLs ...") + reference_cache: dict[str, list] = {} + for task in deduped_tasks: + task["reference_rows"] = ( + run_reference_sql(task["reference_sql"], reference_cache) + if task["reference_sql"] + else None + ) + print(f" {len(reference_cache)} distinct reference queries run\n") + + judge = build_judge(args.judge_model) + semaphore = asyncio.Semaphore(args.concurrency) + verdicts = await asyncio.gather( + *(judge_turn(judge, semaphore, task) for task in deduped_tasks) + ) + summary = aggregate(verdicts) + + print("\n=== Judge scorecard (weighted over all repeats) ===") + for field in SCORE_FIELDS: + if field in summary: + print( + f" {field:<18} {summary[field]['rate']:.0%} (n={summary[field]['n']})" + ) + print(f" judge_errors={summary['_errors']}") + + total_input_tokens = sum(v.get("input_tokens") or 0 for v in verdicts) + total_output_tokens = sum(v.get("output_tokens") or 0 for v in verdicts) + cost = ( + total_input_tokens / 1e6 * args.price_in + + total_output_tokens / 1e6 * args.price_out + ) + print("\n=== Cost ===") + print(f" judge calls: {len(verdicts)}") + print( + f" input tokens: {total_input_tokens:,} output tokens: {total_output_tokens:,}" + ) + print(f" @ ${args.price_in}/M in, ${args.price_out}/M out -> ${cost:.4f}") + if total_input_tokens == 0 and any(v["verdict"] for v in verdicts): + print(" (note: judge model did not report token usage — cost unavailable)") + + report = { + "transcript": args.transcript, + "judge_model": args.judge_model, + "agent_model": transcript.get("model"), + "langsmith_project": ls_project if tracing else None, + "cost": { + "input_tokens": total_input_tokens, + "output_tokens": total_output_tokens, + "price_in_per_m": args.price_in, + "price_out_per_m": args.price_out, + "usd": round(cost, 4), + }, + "summary": summary, + "verdicts": verdicts, + } + with open(out_path, "w") as file: + json.dump(report, file, ensure_ascii=False, indent=2) + print(f"\nwrote {out_path}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/eval/eval_queries.py b/eval/eval_queries.py new file mode 100644 index 0000000..1acdb65 --- /dev/null +++ b/eval/eval_queries.py @@ -0,0 +1,326 @@ +"""Trace-based source/period eval: did the agent query the right tables and period? + +Unlike eval_output (which scores source/period from the agent's STRUCTURED fields), this +reads them from the agent's EXECUTED SQL — the `queries` persisted in the transcript — so +it works for BOTH the structured branch and the free-text baseline (--no-structured / +main). Run it on both transcripts to compare "did the agent query the right inputs" on +equal footing, independent of whether the agent emits structured output. + +It answers a different question than eval_output's source/period: those check "did the +agent correctly REPORT its sources/period"; this checks "did the agent actually QUERY the +right ones". Because it reads the executed SQL, trace `source` counts directory JOINs the +agent may omit from data_sources, so it tends to score higher than the reported version. + +Per gold QUERY turn: + source: the tables the executed SQL referenced (FROM/JOIN) vs gold `tables` + (1.0 exact / 0.5 partial / 0.0 none) + period: the year range the executed SQL/rows evidence vs the gold `temporal` rule + (range / exact / match_previous / latest — all at YEAR granularity) + +`latest` is a year-level check here (the trace gives years, not the table's exact +period_end granularity): the single queried year must equal the queried tables' latest +period_end year. Reuse turns (no SQL this turn) fall back to the thread's earlier queries; +period abstains (None) when the run evidences no year. + +No agent, no BQ, no LLM: pure analysis of the transcript + gold. Runs over ALL repeats. + +Pipeline: run AFTER eval_output.py — it scores the transcript eval_output.py produced. +Needs no agent, LLM or BQ. Sibling scorers over the same transcript, independent of this +one and of each other: + eval_quality.py answer quality vs the gold — LLM judge + live reference_sql + eval_faithfulness.py structured output's self-consistency — gold-free; this module + reuses its SQL/period helpers + + uv run eval/eval_queries.py --in eval/.json + uv run eval/eval_queries.py --in eval/.json --show-failures +""" + +import argparse +import json +import sys +from collections import defaultdict +from pathlib import Path + +import yaml + +# Make the repo root importable so `eval.eval_faithfulness` resolves whether run as a +# module (python -m eval.eval_queries) or directly (python eval/eval_queries.py). +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from eval.eval_faithfulness import ( # noqa: E402 + _leading_year, + _tables_and_years, + turn_success_queries, +) + +EVAL_DIR = Path(__file__).resolve().parent +DIMENSIONS = ("source", "period") + + +# ============================================================================= +# Gold +# ============================================================================= +def load_gold(path: str) -> dict[tuple[str, int], dict]: + """(thread_id, turn_index) -> gold turn dict (carries action / tables / temporal). + + Args: + path (str): Path to eval_gold.yaml. + + Returns: + dict[tuple[str, int], dict]: Gold turn keyed by (thread id, turn index). + """ + gold = yaml.safe_load(open(path)) + return { + (thread["id"], i): turn + for thread in gold + for i, turn in enumerate(thread["turns"]) + } + + +# ============================================================================= +# Per-turn scoring (source / period, from the agent's executed-SQL trace) +# ============================================================================= +def _period_rule_applies(gold: dict) -> bool: + """Whether the gold expects a scored period this turn (a real rule on a query turn). + + Args: + gold (dict): The gold turn (its `action` and `temporal` rule). + + Returns: + bool: True when the turn is a query turn whose `temporal` rule should be scored. + """ + rule = gold.get("temporal") + return gold["action"] == "query" and rule not in (None, "none", "any", "", {}) + + +def score_source(trace_tables: set[str], gold: dict) -> float | None: + """Score the tables the SQL actually hit against the gold `tables`. + + Args: + trace_tables (set[str]): gcp_ids referenced by the turn's executed SQL. + gold (dict): The gold turn (action / tables). + + Returns: + float | None: 1.0 exact, 0.5 partial, 0.0 none; None when not a scored query turn. + """ + if gold["action"] != "query" or not gold.get("tables"): + return None + expected, observed = set(gold["tables"]), set(trace_tables) + if not observed: + return 0.0 # gold expects a query but the trace hit no tables + return 1.0 if observed == expected else (0.5 if observed & expected else 0.0) + + +def score_period( + trace_range: tuple[int, int] | None, + gold: dict, + period_ends: dict[str, str], + queried_tables: set[str], + prev_range: tuple[int, int] | None, +) -> float | None: + """Score the year range the SQL evidenced against the gold `temporal` rule (year level). + + Args: + trace_range (tuple[int, int] | None): (min_year, max_year) evidenced by the run. + gold (dict): The gold turn (action / temporal). + period_ends (dict[str, str]): gcp_id -> table period_end seen so far this thread. + queried_tables (set[str]): gcp_ids the turn's SQL referenced (for the `latest` target). + prev_range (tuple[int, int] | None): the previous query turn's year range. + + Returns: + float | None: 1.0 / 0.0, or None when the rule doesn't apply or can't be checked. + """ + rule = gold.get("temporal") + if gold["action"] != "query" or rule in (None, "none", "any", "", {}): + return None + if trace_range is None: + return 0.0 # queried but the run evidenced no year + start, end = trace_range + if rule == "match_previous": + if prev_range is None: + return None + prev_start, prev_end = prev_range + return 1.0 if prev_start <= start and end <= prev_end else 0.0 + if rule == "range": + return 1.0 if end > start else 0.0 + if rule == "latest": + # Year-level: the single queried year must equal the queried tables' latest + # period_end year (fall back to all seen period_ends if none match by gcp_id). + candidates = [ + period_ends[t] for t in queried_tables if t in period_ends + ] or list(period_ends.values()) + targets = [y for value in candidates if (y := _leading_year(value)) is not None] + if not targets: + return None # can't determine the tables' latest year from the trace + target = max(targets) + return 1.0 if start == target == end else 0.0 + span = rule.get("exact", rule) if isinstance(rule, dict) else None + if isinstance(span, dict) and "start" in span and "end" in span: + return 1.0 if start == int(span["start"]) and end == int(span["end"]) else 0.0 + return None + + +# ============================================================================= +# Transcript scoring & aggregation +# ============================================================================= +def score_transcript(units: list[dict], gold: dict) -> list[dict]: + """Score every `ok` turn's source/period from its executed SQL. + + SQL-derived checks use the thread's executed queries: this turn's own when it ran any, + else earlier turns' (a follow-up may answer from data fetched on a previous turn). + + Args: + units (list[dict]): The transcript's units (thread replays). + gold (dict): (thread, turn) -> gold turn from load_gold(). + + Returns: + list[dict]: One {thread, repeat, turn_index, source, period} per scored turn. + """ + rows = [] + for unit in units: + prior_queries: list[dict] = [] + period_ends: dict[str, str] = {} + prev_range: tuple[int, int] | None = None + for turn in unit["turns"]: + if turn.get("status") != "ok": + continue + gold_turn = gold.get((unit["thread"], turn["turn_index"])) + if gold_turn is None: + continue + period_ends.update(turn.get("table_period_ends") or {}) + this_queries = turn_success_queries(turn) + trace_tables, trace_range = _tables_and_years(this_queries or prior_queries) + period = score_period( + trace_range, gold_turn, period_ends, trace_tables, prev_range + ) + rows.append( + { + "thread": unit["thread"], + "repeat": unit["repeat"], + "turn_index": turn["turn_index"], + "source": score_source(trace_tables, gold_turn), + "period": period, + # a real period rule that couldn't be scored (no year / no period_end in + # the trace) is unverifiable, NOT a pass — surfaced so a 100% can't hide it. + "period_unverifiable": _period_rule_applies(gold_turn) + and period is None, + } + ) + # the period a later `match_previous` turn must stay within + if gold_turn["action"] == "query" and trace_range is not None: + prev_range = trace_range + prior_queries += this_queries + return rows + + +def aggregate(rows: list[dict]) -> dict: + """Pass-rate per dimension over the turns it applied to (None = not applicable). + + Args: + rows (list[dict]): Per-turn score records from score_transcript(). + + Returns: + dict: {dimension: {"rate": float, "n": int}}. + """ + tally = defaultdict(lambda: {"hit": 0.0, "n": 0}) + for row in rows: + for dim in DIMENSIONS: + value = row[dim] + if value is None: + continue + tally[dim]["hit"] += float(value) + tally[dim]["n"] += 1 + return { + dim: {"rate": round(counts["hit"] / counts["n"], 3), "n": counts["n"]} + for dim, counts in tally.items() + if counts["n"] + } + + +# ============================================================================= +# CLI +# ============================================================================= +def main() -> None: + """Score a transcript's trace-based source/period against the gold and write a report.""" + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--in", dest="transcript", required=True, help="thread_eval_*.json" + ) + parser.add_argument("--gold", default=str(EVAL_DIR / "eval_gold.yaml")) + parser.add_argument( + "--thread", action="append", help="Only this thread id (repeatable)" + ) + parser.add_argument("--out", default=None) + parser.add_argument( + "--show-failures", + action="store_true", + help="List each turn where source or period is below 1.0", + ) + args = parser.parse_args() + + transcript = json.load(open(args.transcript)) + units = transcript["units"] + if args.thread: + units = [unit for unit in units if unit["thread"] in args.thread] + gold = load_gold(args.gold) + + rows = score_transcript(units, gold) + summary = aggregate(rows) + + print( + f"transcript={args.transcript!r} branch={transcript.get('branch')!r} " + f"structured_output={transcript.get('structured_output')}" + ) + period_unverifiable = sum(1 for row in rows if row["period_unverifiable"]) + notes = {} + if period_unverifiable: + notes["period"] = ( + f" [+{period_unverifiable} unverifiable: no period_end/year in the trace]" + ) + + print( + "\n=== Trace-based inputs scorecard (source/period from executed SQL vs gold) ===" + ) + for dim in DIMENSIONS: + if dim in summary: + print( + f" {dim:<8} {summary[dim]['rate']:.0%} (n={summary[dim]['n']}){notes.get(dim, '')}" + ) + + if args.show_failures: + flagged = [ + row + for row in rows + if (row["source"] is not None and row["source"] < 1.0) + or (row["period"] is not None and row["period"] < 1.0) + or row["period_unverifiable"] + ] + if flagged: + print("\n=== Turns below 1.0 or unverifiable ===") + for row in flagged: + tag = " (period unverifiable)" if row["period_unverifiable"] else "" + print( + f" {row['thread']:<20} #{row['repeat']} t{row['turn_index']}: " + f"source={row['source']} period={row['period']}{tag}" + ) + + out_path = args.out or str(EVAL_DIR / f"{Path(args.transcript).stem}_inputs.json") + with open(out_path, "w") as file: + json.dump( + { + "transcript": args.transcript, + "summary": summary, + "period_unverifiable": period_unverifiable, + "rows": rows, + }, + file, + ensure_ascii=False, + indent=2, + ) + print(f"\nwrote {out_path}") + + +if __name__ == "__main__": + main() diff --git a/eval/eval_threads.py b/eval/eval_threads.py deleted file mode 100644 index f272617..0000000 --- a/eval/eval_threads.py +++ /dev/null @@ -1,599 +0,0 @@ -"""Multi-turn agent eval: replay conversation threads and score deterministic dimensions. - -Replays each thread in `eval_gold.yaml` turn-by-turn on a shared `thread_id` -(InMemorySaver, so follow-ups see prior context), K times per thread, for the -currently checked-out branch + given settings. Scores four deterministic -dimensions per turn from the agent's structured response fields (the purpose -of this eval is to check those fields are filled with the expected values): - - action: query-vs-clarify, from whether the agent filled the sql_queries field - source: data_sources (its table UUIDs resolved to gcp_ids) match the gold - period: temporal_coverage matches the gold rule (any/latest/range/match_previous/exact) - completed: the turn finished without errors - -It writes a rich JSON (full per-turn records: structured response, response text, -model-call count) so the later judge eval and the A-vs-B pass can score -the SAME transcripts without re-running the agent. - -Example Usage: - uv run python eval_threads.py --repeats 10 --temperature 0.0 - uv run python eval_threads.py --repeats 10 --threads --temperature 0.0 - -Each turn runs the real agent (live BQ + LLM) — mind the cost (sum of turns x K). -""" - -import argparse -import asyncio -import json -import os -import re -import subprocess -import traceback -from collections import defaultdict -from datetime import date, datetime -from pathlib import Path - -import yaml -from langchain.agents import create_agent -from langchain.agents.middleware import ( - ModelCallLimitMiddleware, - SummarizationMiddleware, -) -from langchain.chat_models import init_chat_model -from langchain.messages import AnyMessage -from langgraph.checkpoint.memory import InMemorySaver -from langgraph.graph.state import CompiledStateGraph - -from app.agent.prompts import SYSTEM_PROMPT -from app.agent.tools import BDToolkit -from app.settings import settings - -try: - from app.agent.schemas import StructuredResponse -except ImportError as exc: - raise NotImplementedError( - "This eval only supports the structured-response agent" - ) from exc - -# Agent middleware tunables (kept in step with production config). -SUMMARIZE_TRIGGER_TOKENS = 500_000 -SUMMARIZE_KEEP_TOKENS = 250_000 -MODEL_CALL_RUN_LIMIT = 20 - -# Deterministic dimensions scored per turn, in display order. -SCORE_DIMENSIONS = ("action", "source", "period", "completed") - - -# ============================================================================= -# Git and Agent setup -# ============================================================================= -def current_branch() -> str: - try: - return subprocess.check_output( - ["git", "rev-parse", "--abbrev-ref", "HEAD"], text=True - ).strip() - except Exception: - return "unknown" - - -def build_agent( - temperature: float, checkpointer: InMemorySaver, system_prompt: str -) -> CompiledStateGraph: - model = init_chat_model( - model=settings.MODEL_URI, - temperature=temperature, - credentials=settings.GOOGLE_CREDENTIALS, - thinking_level=settings.THINKING_LEVEL, - include_thoughts=True, - ) - middleware = [ - SummarizationMiddleware( - model=model, - trigger=("tokens", SUMMARIZE_TRIGGER_TOKENS), - keep=("tokens", SUMMARIZE_KEEP_TOKENS), - ), - ModelCallLimitMiddleware( - run_limit=MODEL_CALL_RUN_LIMIT, - exit_behavior="end", - ), - ] - return create_agent( - model=model, - tools=BDToolkit.get_tools(), - system_prompt=system_prompt, - middleware=middleware, - checkpointer=checkpointer, - response_format=StructuredResponse, - ) - - -# ============================================================================= -# per-turn extraction from the agent's final state -# ============================================================================= -def _empty_record(status: str, **extra) -> dict: - """Base per-turn record for the non-`ok` cases (no_structured / error), with the - scoring fields zeroed so downstream code can read them uniformly. `extra` overrides - or adds fields (e.g. an `error` message, a recovered `response_text`).""" - return { - "status": status, - "is_query": False, - "tables": [], - "table_period_ends": {}, - "structured": None, - "response_text": None, - "model_calls": None, - "tools_used": [], - **extra, - } - - -def _last_ai_text(messages: list[AnyMessage]) -> str | None: - for m in reversed(messages): - if m.type == "ai": - return m.text - return None - - -def _tools_used(messages: list[AnyMessage]) -> set[str]: - """Distinct tool names the agent called in THIS turn. The checkpointer returns - the whole thread, so slice to the messages after the last human turn.""" - last_human = max( - (i for i, m in enumerate(messages) if m.type == "human"), default=-1 - ) - return { - tc["name"] - for m in messages[last_human + 1 :] - for tc in getattr(m, "tool_calls", None) or [] - } - - -def _tool_metadata(messages: list[AnyMessage]) -> tuple[dict[str, str], dict[str, str]]: - """From the agent's get_table_details / get_dataset_details outputs, build - (uuid -> gcp_id) to resolve data_sources, and (gcp_id -> period_end) for the - `latest` period rule. Tool payloads are pydantic-serialized JSON so keys are - guaranteed; a non-JSON body means that tool call errored, so skip it.""" - uuid_to_gcp: dict[str, str] = {} - period_end: dict[str, str] = {} - for m in messages: - if m.type != "tool" or m.name not in ( - "get_table_details", - "get_dataset_details", - ): - continue - try: - d = json.loads(m.content) - except json.JSONDecodeError: - continue - if m.name == "get_table_details": - uuid_to_gcp[d["id"]] = d["gcp_id"] - if d["period_end"] is not None: - period_end[d["gcp_id"]] = d["period_end"] - else: # get_dataset_details - for t in d["tables"]: - uuid_to_gcp[t["id"]] = t["gcp_id"] - return uuid_to_gcp, period_end - - -def extract_turn(result: dict) -> dict: - messages = result["messages"] - structured: StructuredResponse | None = result.get("structured_response") - - # This eval targets the structured-output agent: a missing structured response - # is a hard fail for the turn (not a crash, not a skip). - if structured is None: - return _empty_record( - "no_structured", - response_text=_last_ai_text(messages), - model_calls=result.get("run_model_call_count"), - ) - - uuid_to_gcp, table_period_ends = _tool_metadata(messages) - # Tables the agent REPORTED in data_sources, UUID -> gcp_id (unmapped UUIDs are - # kept as-is so they fail to match the gold and surface the problem). - data_source_tables = sorted( - { - uuid_to_gcp.get(d.table_id, d.table_id) - for d in (structured.data_sources or []) - } - ) - return { - "status": "ok", - # query vs clarify, from whether the agent filled the sql_queries field. - "is_query": bool(structured.sql_queries), - "tables": data_source_tables, - "table_period_ends": table_period_ends, - "structured": structured.model_dump(mode="json"), - "response_text": structured.response, - "model_calls": result.get("run_model_call_count"), - "tools_used": sorted(_tools_used(messages)), - } - - -# ============================================================================= -# Deterministic scoring -# ============================================================================= -def _lead_year(s) -> int | None: - """Leading 4-digit year of a period value ('2015', '2015-03-01', '2026-05').""" - m = re.match(r"\s*(\d{4})", str(s or "")) - return int(m.group(1)) if m else None - - -def _tc_years(tc: dict | None) -> tuple[int | None, int | None]: - """(start_year, end_year) from a temporal_coverage dict — the leading 4-digit - year of period_start/period_end (handles '2015' and '2015-03-01' alike).""" - if not tc: - return None, None - return _lead_year(tc.get("period_start")), _lead_year(tc.get("period_end")) - - -def _parse_period(s) -> tuple[str, str] | None: - """(granularity, value) for a period string in its granularity's own format: - 'YYYY' -> ('year', ...), 'YYYY-MM' -> ('month', ...), 'YYYY-MM-DD' -> ('day', ...). - Strict (re.fullmatch): a malformed value like '2026-5' returns None -> a miss, not - something to silently repair.""" - s = str(s or "").strip().strip("'").strip('"') - for gran, pat in ( - ("day", r"\d{4}-\d{2}-\d{2}"), - ("month", r"\d{4}-\d{2}"), - ("year", r"\d{4}"), - ): - if re.fullmatch(pat, s): - return gran, s - return None - - -def _score_period(rec: dict, gold: dict, prev_query: dict | None): - """Score the period from the agent's reported `temporal_coverage` — normalized - and independent of SQL surface form / temporal-column name. Returns [0,1] or None. - """ - rule = gold.get("temporal") - if gold["action"] != "query" or rule in (None, "none", "any", "", {}): - return None - ps, pe = _tc_years(rec["structured"]["temporal_coverage"]) - if ps is None or pe is None: - return 0.0 # queried but reported no usable temporal_coverage - if rule == "match_previous": - if prev_query is None: - return None - pps, ppe = _tc_years(prev_query["structured"]["temporal_coverage"]) - if pps is None or ppe is None: - return 0.0 - # The follow-up must stay WITHIN the previous turn's period. A justified - # narrowing (e.g., the requested metric only exists for part of the range) - # is fine; only drifting OUTSIDE the established window fails. - return 1.0 if pps <= ps and pe <= ppe else 0.0 - if rule == "range": - return 1.0 if pe > ps else 0.0 - if rule == "latest": - # Expected = the table's period_end, matched EXACTLY at its own granularity: - # '2026-05' must be reported as month 2026-05 (not year 2026), a full date to - # the day. The model must emit the value in its granularity's format. - queried = set(rec["tables"]) - cands = [v for k, v in rec["table_period_ends"].items() if k in queried] - cands = cands or list(rec["table_period_ends"].values()) - targets = [p for v in cands if (p := _parse_period(v)) is not None] - if not targets: - return None # can't determine the table's period_end from the trace - target = max(targets, key=lambda p: p[1]) # latest period_end - tc = rec["structured"]["temporal_coverage"] or {} - start, end = ( - _parse_period(tc.get("period_start")), - _parse_period(tc.get("period_end")), - ) - matches = ( - start == target and end == target and tc.get("granularity") == target[0] - ) - return 1.0 if matches else 0.0 - span = rule.get("exact", rule) if isinstance(rule, dict) else None - if isinstance(span, dict) and "start" in span and "end" in span: - return 1.0 if ps == int(span["start"]) and pe == int(span["end"]) else 0.0 - return None - - -def score_turn(rec: dict, gold: dict, prev_query: dict | None) -> dict: - if rec["status"] == "no_structured": - # agent returned no structured response -> the structured contract failed - return {"action": False, "source": None, "period": None, "completed": False} - - observed = "query" if rec["is_query"] else "clarify" - - if gold["action"] == "query" and gold.get("tables"): - exp, obs = set(gold["tables"]), set(rec["tables"]) - source = 1.0 if obs == exp else (0.5 if obs & exp else 0.0) - else: - source = None # clarify turn, or gcp_ids not filled in yet - return { - "action": observed == gold["action"], - "source": source, - "period": _score_period(rec, gold, prev_query), - "completed": rec["status"] == "ok", - } - - -# ============================================================================= -# Thread replay -# ============================================================================= -def _turn_run_config( - base: dict, - *, - thread_id: str, - thread: dict, - repeat: int, - turn_index: int, - turn: dict, - branch: str, - temperature: float, -) -> dict: - """Per-turn LangGraph config. Labels the run so it's findable in LangSmith; - `thread_id` in metadata makes LangSmith group a replay's turns into one - conversation.""" - return { - **base, - "configurable": {"thread_id": thread_id}, - "run_name": f"{thread['id']}#{repeat}-t{turn_index}", - "tags": [ - f"branch:{branch}", - f"temp:{temperature}", - f"thread:{thread['id']}", - ], - "metadata": { - "thread_id": thread_id, - "eval_thread": thread["id"], - "repeat": repeat, - "turn_index": turn_index, - "branch": branch, - "temperature": temperature, - "expected_action": turn["action"], - "expected_temporal": turn.get("temporal"), - }, - } - - -def _skipped_record(turn_index: int, user: str) -> dict: - return {"turn_index": turn_index, "user": user, "status": "skipped", "scores": {}} - - -async def replay_thread( - agent: CompiledStateGraph, - thread: dict, - repeat: int, - run_config: dict, - branch: str, - temperature: float, - max_turns: int | None = None, -) -> list: - thread_id = f"{thread['id']}-{repeat}" - records, prev_query = [], None - turns = thread["turns"] if max_turns is None else thread["turns"][:max_turns] - - for i, turn in enumerate(turns): - cfg = _turn_run_config( - run_config, - thread_id=thread_id, - thread=thread, - repeat=repeat, - turn_index=i, - turn=turn, - branch=branch, - temperature=temperature, - ) - - try: - result = await agent.ainvoke( - {"messages": [{"role": "user", "content": turn["user"]}]}, config=cfg - ) - rec = extract_turn(result) - except Exception as exc: - rec = _empty_record( - "error", - error=f"{type(exc).__name__}: {exc}", - traceback=traceback.format_exc(), - ) - - rec["turn_index"], rec["user"] = i, turn["user"] - rec["scores"] = ( - score_turn(rec, turn, prev_query) if rec["status"] != "error" else {} - ) - records.append(rec) - - if rec["status"] == "error": # broken state -> skip the rest of the thread - records.extend( - _skipped_record(j, turns[j]["user"]) for j in range(i + 1, len(turns)) - ) - break - - if rec["is_query"]: - prev_query = rec - - return records - - -# ============================================================================= -# Aggregation & Reporting -# ============================================================================= -def _mark(scores: dict) -> str: - def m(v): - return ( - "·" - if v is None - else ("✓" if v is True or v == 1.0 else ("~" if v == 0.5 else "✗")) - ) - - return ( - f"a:{m(scores.get('action'))} " - f"s:{m(scores.get('source'))} " - f"p:{m(scores.get('period'))} " - f"c:{m(scores.get('completed'))}" - ) - - -def _unit_line(turn_records: list) -> str: - """One-line per-thread progress summary: each turn's marks, or its status.""" - return " ".join( - f"t{tr['turn_index']}:" - + (_mark(tr["scores"]) if tr["status"] == "ok" else tr["status"].upper()) - for tr in turn_records - ) - - -def aggregate(units: list) -> dict: - """Pass-rate per dimension over all scored turns (None = not applicable, skipped).""" - agg = defaultdict(lambda: {"hit": 0.0, "n": 0}) - errors = skipped = no_structured = 0 - for u in units: - for rec in u["turns"]: - if rec["status"] == "skipped": - skipped += 1 - continue - if rec["status"] == "error": - errors += 1 - elif rec["status"] == "no_structured": - no_structured += 1 - for dim, val in (rec.get("scores") or {}).items(): - if val is None: - continue - agg[dim]["hit"] += float(val) - agg[dim]["n"] += 1 - summary = { - d: {"rate": round(v["hit"] / v["n"], 3), "n": v["n"]} for d, v in agg.items() - } - summary["_errors"], summary["_skipped"], summary["_no_structured"] = ( - errors, - skipped, - no_structured, - ) - return summary - - -def print_scorecard(summary: dict) -> None: - print("\n=== Deterministic scorecard (pass-rate over applicable turns) ===") - for dim in SCORE_DIMENSIONS: - if dim in summary: - print(f" {dim:<10} {summary[dim]['rate']:.0%} (n={summary[dim]['n']})") - print( - f" errors={summary['_errors']} skipped={summary['_skipped']} " - f"no_structured={summary['_no_structured']}" - ) - - -# ============================================================================= -# CLI -# ============================================================================= -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--gold", default="eval_gold.yaml") - parser.add_argument("--repeats", type=int, default=3, help="Replays per thread") - parser.add_argument( - "--thread", action="append", help="Only this thread id (repeatable)" - ) - parser.add_argument( - "--max-turns", - type=int, - default=None, - help="Run only the first N turns of each thread (a prefix — follow-ups need prior context)", - ) - parser.add_argument("--temperature", type=float, default=None) - parser.add_argument( - "--recursion-limit", - type=int, - default=None, - help="Omit to match production (LangGraph default 25)", - ) - parser.add_argument( - "--concurrency", type=int, default=1, help="Parallel thread replays" - ) - parser.add_argument("--out", default=None) - parser.add_argument( - "--langsmith-project", - default=None, - help="LangSmith project for these traces (default: -eval)", - ) - parser.add_argument( - "--no-trace", action="store_true", help="Disable LangSmith tracing" - ) - return parser.parse_args() - - -def configure_tracing(args: argparse.Namespace) -> tuple[str, bool]: - """Forward LangSmith config from `settings` into os.environ (LangChain's tracer - reads the env, but pydantic-settings only populates `settings`), isolating eval - runs in their own project by default. Returns (project, enabled); must run before - the agent does.""" - ls_project = args.langsmith_project or f"{settings.LANGSMITH_PROJECT}-eval" - tracing = settings.LANGSMITH_TRACING and not args.no_trace - os.environ["LANGSMITH_TRACING"] = "true" if tracing else "false" - if tracing: - os.environ["LANGSMITH_PROJECT"] = ls_project - os.environ["LANGSMITH_API_KEY"] = settings.LANGSMITH_API_KEY - return ls_project, tracing - - -async def main() -> None: - args = parse_args() - - with open(args.gold) as f: - gold = yaml.safe_load(f) - threads = [t for t in gold if not args.thread or t["id"] in args.thread] - - temperature = ( - args.temperature if args.temperature is not None else settings.MODEL_TEMPERATURE - ) - run_config: dict = {} - if args.recursion_limit is not None: - run_config["recursion_limit"] = args.recursion_limit - - branch = current_branch() - ts = datetime.now().strftime("%Y%m%d-%H%M%S") - out_path = args.out or str( - Path(__file__).resolve().parent - / f"thread_eval_{branch.replace('/', '-')}_temp{temperature}_{ts}.json" - ) - - ls_project, tracing = configure_tracing(args) - - print( - f"\nbranch={branch!r} model={settings.MODEL_URI!r}\n" - f"temperature={temperature} repeats={args.repeats} threads={[t['id'] for t in threads]}\n" - f"langsmith: project={ls_project!r}\n" - if tracing - else "langsmith: disabled\n" - ) - - system_prompt = SYSTEM_PROMPT.format(current_date=date.today().isoformat()) - agent = build_agent(temperature, InMemorySaver(), system_prompt) - sem = asyncio.Semaphore(args.concurrency) - work = [(t, r) for t in threads for r in range(args.repeats)] - - async def run_unit(thread, repeat): - async with sem: - turn_records = await replay_thread( - agent, thread, repeat, run_config, branch, temperature, args.max_turns - ) - print(f"[{thread['id']:<16} #{repeat}] {_unit_line(turn_records)}") - return {"thread": thread["id"], "repeat": repeat, "turns": turn_records} - - units = await asyncio.gather(*(run_unit(t, r) for t, r in work)) - summary = aggregate(units) - print_scorecard(summary) - - report = { - "branch": branch, - "langsmith_project": ls_project if tracing else None, - "model": settings.MODEL_URI, - "temperature": temperature, - "recursion_limit": args.recursion_limit, - "repeats": args.repeats, - "timestamp": ts, - "system_prompt": system_prompt, - "summary": summary, - "units": units, - } - - with open(out_path, "w") as fh: - json.dump(report, fh, ensure_ascii=False, indent=2) - - print(f"\nwrote {out_path}") - - -if __name__ == "__main__": - asyncio.run(main()) From 3b30cc39006d16f1988326fb25484ee686a2b85c Mon Sep 17 00:00:00 2001 From: vrtornisiello Date: Wed, 8 Jul 2026 16:39:13 -0300 Subject: [PATCH 12/12] feat: stop surfacing model thinking in tool call events Thoughts are already captured in LangSmith traces, so exposing them in the frontend stream just adds noise; tool_call args and tool_output already give clients enough to render good UI/UX. Also fixes get_table_details tests, which broke after TABLE_DETAILS_QUERY started requesting `dataset { id }` to populate the table's dataset_id. Co-Authored-By: Claude Sonnet 5 --- app/api/streaming/agent_runner.py | 31 +-- tests/app/agent/tools/test_api.py | 2 + tests/app/api/streaming/test_agent_runner.py | 217 ++++++++++++++----- 3 files changed, 161 insertions(+), 89 deletions(-) diff --git a/app/api/streaming/agent_runner.py b/app/api/streaming/agent_runner.py index 625c3b4..d78a820 100644 --- a/app/api/streaming/agent_runner.py +++ b/app/api/streaming/agent_runner.py @@ -85,35 +85,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. @@ -166,7 +137,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 diff --git a/tests/app/agent/tools/test_api.py b/tests/app/agent/tools/test_api.py index 5a971a0..1dff85f 100644 --- a/tests/app/agent/tools/test_api.py +++ b/tests/app/agent/tools/test_api.py @@ -421,6 +421,7 @@ def mock_response(self): }, ] }, + "dataset": {"id": "DatasetNode:dataset-1"}, } } ] @@ -439,6 +440,7 @@ async def test_get_table_details_success(self, mock_response): table = json.loads(result) assert table["id"] == "table-1" + assert table["dataset_id"] == "dataset-1" assert table["gcp_id"] == "basedosdados.test_dataset.test_table" assert table["name"] == "Test Table" assert table["description"] == "Table description" diff --git a/tests/app/api/streaming/test_agent_runner.py b/tests/app/api/streaming/test_agent_runner.py index f20e4c5..96899a8 100644 --- a/tests/app/api/streaming/test_agent_runner.py +++ b/tests/app/api/streaming/test_agent_runner.py @@ -8,10 +8,15 @@ import pytest from langchain_core.messages import AIMessage, ToolMessage +from app.agent.schemas import ( + DataSource, + StructuredResponse, + TemporalCoverage, + TemporalGranularity, +) from app.api.schemas import ConfigDict from app.api.streaming.agent_runner import ( ErrorMessage, - _parse_thinking, _process_chunk, _truncate_json, run_agent, @@ -183,64 +188,6 @@ def test_truncate_json_invalid(self): assert _truncate_json(invalid_json_string) == invalid_json_string -class TestParseThinking: - """Tests for _parse_thinking function.""" - - def test_string_content_returns_none(self): - """Test that plain string content returns None.""" - message = AIMessage(content="Hello, world!") - assert _parse_thinking(message) is None - - def test_single_thinking_block(self): - """Test extraction of a single thinking block.""" - message = AIMessage( - content=[ - {"type": "thinking", "thinking": "Let me reason about this."}, - {"type": "text", "text": "Here is my answer."}, - ] - ) - assert _parse_thinking(message) == "Let me reason about this." - - def test_multiple_thinking_blocks_are_concatenated(self): - """Test that multiple thinking blocks are concatenated.""" - message = AIMessage( - content=[ - {"type": "thinking", "thinking": "First thought. "}, - {"type": "text", "text": "Some text."}, - {"type": "thinking", "thinking": "Second thought."}, - ] - ) - assert _parse_thinking(message) == "First thought. Second thought." - - def test_no_thinking_blocks_returns_none(self): - """Test that content with no thinking blocks returns None.""" - message = AIMessage( - content=[ - {"type": "text", "text": "Just text."}, - ] - ) - assert _parse_thinking(message) is None - - def test_empty_thinking_block_returns_none(self): - """Test that an empty thinking string returns None.""" - message = AIMessage( - content=[ - {"type": "thinking", "thinking": ""}, - ] - ) - assert _parse_thinking(message) is None - - def test_non_dict_blocks_are_skipped(self): - """Test that non-dict items in content are safely skipped.""" - message = AIMessage( - content=[ - "plain string block", - {"type": "thinking", "thinking": "Actual thinking."}, - ] - ) - assert _parse_thinking(message) == "Actual thinking." - - class TestProcessChunk: """Tests for _process_chunk function.""" @@ -331,6 +278,84 @@ def test_agent_chunk_empty_messages(self): assert event.type == "final_answer" assert event.data.content == "" + def test_agent_chunk_structured_response(self): + """A model chunk carrying `structured_response` yields a final_answer event + whose content is the prose and whose structured_response holds all fields.""" + structured = StructuredResponse( + response="Here is your answer.", + data_sources=[ + DataSource(dataset_id="ds1", table_id="tb1", name="Tabela 1") + ], + temporal_coverage=TemporalCoverage( + period_start="2020", + period_end="2025", + granularity=TemporalGranularity.YEAR, + ), + sql_queries=["SELECT 1 -- comment"], + follow_up_questions=["E em 2026?", "Por estado?", "Por região?"], + ) + + # The model node sets `structured_response` alongside the internal + # structured-output tool call; that tool call must NOT become a tool_call event. + chunk = { + "model": { + "messages": [ + AIMessage( + content="", + tool_calls=[ + { + "id": "call_struct", + "name": "StructuredResponse", + "args": structured.model_dump(), + } + ], + ) + ], + "structured_response": structured, + } + } + + event = _process_chunk(chunk) + + assert event is not None + assert event.type == "final_answer" + assert event.data.tool_calls is None + assert event.data.content == "Here is your answer." + assert event.data.structured_response is not None + assert event.data.structured_response["response"] == "Here is your answer." + assert event.data.structured_response["data_sources"] == [ + {"dataset_id": "ds1", "table_id": "tb1", "name": "Tabela 1"} + ] + assert event.data.structured_response["temporal_coverage"] == { + "period_start": "2020", + "period_end": "2025", + "granularity": "year", + } + assert event.data.structured_response["sql_queries"] == ["SELECT 1 -- comment"] + assert event.data.structured_response["follow_up_questions"] == [ + "E em 2026?", + "Por estado?", + "Por região?", + ] + + def test_agent_chunk_structured_response_sanitizes_links(self): + """The prose in a structured response has its markdown links sanitized.""" + structured = StructuredResponse(response="See [evil](http://evil.com).") + + chunk = { + "model": { + "messages": [AIMessage(content="")], + "structured_response": structured, + } + } + + event = _process_chunk(chunk) + + assert event is not None + assert event.type == "final_answer" + assert "http://evil.com" not in event.data.content + assert event.data.content == event.data.structured_response["response"] + def test_tools_chunk_single_tool(self): """Test tools chunk with single tool output (dict format).""" chunk = { @@ -516,6 +541,80 @@ async def astream(*args, **kwargs): assert message.status == MessageStatus.SUCCESS assert message.content == "Final answer" + async def test_structured_response_is_emitted_and_persisted( + self, + mock_database: MagicMock, + mock_user_message: Message, + config: ConfigDict, + thread_id: str, + ): + """A structured final answer is streamed and persisted on the message row.""" + structured = StructuredResponse( + response="Final answer", + data_sources=[ + DataSource(dataset_id="ds1", table_id="tb1", name="Tabela 1") + ], + temporal_coverage=TemporalCoverage( + period_start="2020", + period_end="2025", + granularity=TemporalGranularity.YEAR, + ), + sql_queries=["SELECT 1"], + follow_up_questions=["E em 2026?"], + ) + + agent = MagicMock() + + async def astream(*args, **kwargs): + yield ( + "updates", + { + "model": { + "messages": [AIMessage(content="")], + "structured_response": structured, + } + }, + ) + + agent.astream = astream + queue: asyncio.Queue[StreamEvent] = asyncio.Queue() + + await run_agent( + agent=agent, + config=config, + thread_id=thread_id, + user_message=mock_user_message, + model_uri=MODEL_URI, + queue=queue, + ) + + events = await self._drain(queue) + assert [e.type for e in events] == ["final_answer", "complete"] + assert events[0].data.content == "Final answer" + assert events[-1].data.run_id == config["run_id"] + + assert events[0].data.structured_response is not None + assert events[0].data.structured_response["response"] == "Final answer" + assert events[0].data.structured_response["data_sources"] == [ + {"dataset_id": "ds1", "table_id": "tb1", "name": "Tabela 1"} + ] + assert events[0].data.structured_response["temporal_coverage"] == { + "period_start": "2020", + "period_end": "2025", + "granularity": "year", + } + assert events[0].data.structured_response["sql_queries"] == ["SELECT 1"] + assert events[0].data.structured_response["follow_up_questions"] == [ + "E em 2026?" + ] + + mock_database.create_message.assert_called_once() + message = mock_database.create_message.call_args[0][0] + assert isinstance(message, MessageCreate) + assert message.status == MessageStatus.SUCCESS + assert message.content == "Final answer" + assert message.structured_response == events[0].data.structured_response + async def test_unexpected_exception_persists_error_row( self, mock_database: MagicMock,