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/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..2423212 100644 --- a/app/agent/prompts.py +++ b/app/agent/prompts.py @@ -1,138 +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 -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. +# 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. + +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. - -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. +# 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. -Sempre que você tiver **qualquer dúvida** sobre o que buscar, peça mais detalhes ao usuário. +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. ---- +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. -# 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"). +Whenever you have **any doubt** about what to search for, ask the user for more detail. --- -# 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 (`--`). +# 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 (`--`). + +## 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. -## 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. +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 `period_end` como filtro padrão. Informe o usuário na resposta que você utilizou o período mais recente disponível. +- **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 -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: -- 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. +# 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). + +## `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. + +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. + +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. -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. +## Structured Fields +Fill them **only** based on the tool results obtained in this conversation: +- **`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. -## 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**: Inclui todos os elementos requeridos na resposta final?""" +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?""" diff --git a/app/agent/schemas.py b/app/agent/schemas.py new file mode 100644 index 0000000..1226148 --- /dev/null +++ b/app/agent/schemas.py @@ -0,0 +1,98 @@ +from enum import Enum + +from pydantic import BaseModel, Field + + +class TemporalGranularity(str, Enum): + """Granularity of the data's temporal coverage.""" + + DAY = "day" + MONTH = "month" + YEAR = "year" + + +class TemporalCoverage(BaseModel): + """The interval the SQL query actually filtered on in the answer.""" + + period_start: str = Field( + description=( + "Start of the interval filtered by the SQL query (e.g. '2010' for " + "`ano = 2010` or `ano BETWEEN 2010 AND 2012`). Format it to match `granularity`: " + "YYYY (year), YYYY-MM (month) or YYYY-MM-DD (day) — e.g. '2010', '2010-01', '2010-01-01'. " + "May be narrower than the table's full coverage." + ) + ) + period_end: str = Field( + description=( + "End of the interval filtered by the SQL query (e.g. '2010' for " + "`ano = 2010`; '2012' for `ano BETWEEN 2010 AND 2012`). Format it to match `granularity`: " + "YYYY (year), YYYY-MM (month) or YYYY-MM-DD (day) — e.g. '2012', '2012-01', '2012-01-01'. " + "May be narrower than the table's full coverage." + ) + ) + granularity: TemporalGranularity = Field( + description=( + "Granularity of `period_start`/`period_end`, matching their format: " + "YYYY (year), YYYY-MM (month), YYYY-MM-DD (day)." + ) + ) + + +class DataSource(BaseModel): + """A Base dos Dados table the answer draws on or points the user to.""" + + dataset_id: str = Field( + description=( + "Dataset UUID (the `dataset_id` field from `get_table_details` or the `id` " + "field from `get_dataset_details`), not the BigQuery id (e.g. 'br_bd_diretorios')." + ) + ) + table_id: str = Field( + description=( + "Table UUID (the `id` field from `get_table_details`, or a table's `id` from " + "the tables list of `get_dataset_details`). Not the dataset UUID, not the BigQuery id." + ) + ) + name: str = Field(description="Human-readable name of the table.") + + +class StructuredResponse(BaseModel): + """The agent's structured response for the user interface.""" + + response: str = Field( + description=( + "The prose answer in Markdown, written in the user's language: a direct answer to the question " + "with the data obtained, plus analysis and context. Do NOT repeat the source, period, SQL, " + "or suggestions here — each of those has its own dedicated field." + ) + ) + data_sources: list[DataSource] | None = Field( + default=None, + description=( + "The tables the answer draws on — those you queried, or specific tables you recommend " + "on clarification turns. Leave empty (None) when no table is relevant." + ), + ) + temporal_coverage: TemporalCoverage | None = Field( + default=None, + description=( + "The interval your SQL query actually filtered. Leave empty (None) when no query " + "was run (e.g. a clarification turn) or the answer has no temporal dimension." + ), + ) + sql_queries: list[str] | None = Field( + default=None, + description=( + "The SQL queries whose results the answer is based on, each with " + "inline comments, so the user can reproduce the result. Include every " + "query that contributed to the answer (e.g. one query per metric when the " + "answer combines several), but EXCLUDE exploratory or failed-then-corrected queries. " + "Leave empty (None) when no query was executed." + ), + ) + follow_up_questions: list[str] | None = Field( + default=None, + description=( + "3 suggested follow-up questions (in the user's language) to explore the data further." + ), + ) diff --git a/app/agent/tools/api.py b/app/agent/tools/api.py index 841eafc..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: @@ -173,6 +178,7 @@ async def get_dataset_details(dataset_id: str) -> str: dataset_tables.append( TableOverview( id=table_id, + dataset_id=dataset_id, gcp_id=table_gcp_id, name=table_name, description=table_description, @@ -225,8 +231,6 @@ async def get_table_details(table_id: str) -> str: - period_start / period_end: First and last period covered by the table. Format varies (`2024`, `'2026-04-12'`, etc.) — use the value verbatim, matched to the appropriate temporal column (`ano`, `data`, etc.). - - Next step: Use `execute_bigquery_sql()` to execute queries. """ response = await _client.post( url=GRAPHQL_URL, @@ -296,8 +300,11 @@ async def get_table_details(table_id: str) -> str: ) ) + dataset_id = table["dataset"]["id"].split("DatasetNode:")[-1] + result = Table( id=table_id, + dataset_id=dataset_id, gcp_id=table_gcp_id, name=table_name, description=table_description, 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. 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..d78a820 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 @@ -84,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. @@ -128,7 +100,27 @@ def _process_chunk(chunk: dict[str, Any]) -> StreamEvent | None: - None for ignored chunks """ if "model" in chunk: - ai_messages: list[AIMessage] = chunk["model"]["messages"] + update: dict[str, Any] = chunk["model"] + + # When `response_format` is set (see app.main:91), the model node sets `structured_response` + # (a StructuredResponse) on the turn it produces the final answer. This is the final answer; + # the accompanying structured-output tool call / ToolMessage in `update["messages"]` is + # internal and must not be emitted as a tool_call. + structured: StructuredResponse | None = update.get("structured_response") + + if structured is not None: + response_text = sanitize_markdown_links(structured.response) + structured_response = structured.model_dump() + structured_response["response"] = response_text + return StreamEvent( + type="final_answer", + data=EventData( + content=response_text, + structured_response=structured_response, + ), + ) + + ai_messages: list[AIMessage] = update["messages"] # If no messages are returned, the model returned an empty response # with no tool calls. This also counts as a final (but empty) answer. @@ -145,7 +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 @@ -183,7 +175,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 +184,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 +217,7 @@ async def run_agent( events = [] artifacts = [] assistant_message = "" + structured_response: dict[str, Any] | None = None status: MessageStatus | None = None try: @@ -245,6 +240,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 +276,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, ) 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/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,