diff --git a/alembic/versions/91426eac7604_query_handles_query_ref_pk.py b/alembic/versions/91426eac7604_query_handles_query_ref_pk.py new file mode 100644 index 0000000..a9ddc9a --- /dev/null +++ b/alembic/versions/91426eac7604_query_handles_query_ref_pk.py @@ -0,0 +1,37 @@ +"""Make query_ref the sole primary key of query_handles. + +Revision ID: 91426eac7604 +Revises: 4b3d2fa4a75f +Create Date: 2026-08-18 15:02:34.101000 +""" + +from typing import Sequence, Union + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "91426eac7604" +down_revision: Union[str, Sequence[str], None] = "4b3d2fa4a75f" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + op.drop_constraint("query_handles_pkey", "query_handles", type_="primary") + op.create_primary_key("query_handles_pkey", "query_handles", ["query_ref"]) + op.create_index( + op.f("ix_query_handles_message_id"), + "query_handles", + ["message_id"], + unique=False, + ) + + +def downgrade() -> None: + """Downgrade schema.""" + op.drop_index(op.f("ix_query_handles_message_id"), table_name="query_handles") + op.drop_constraint("query_handles_pkey", "query_handles", type_="primary") + op.create_primary_key( + "query_handles_pkey", "query_handles", ["message_id", "query_ref"] + ) diff --git a/alembic/versions/c4e1a9d2f6b8_add_streaming_message_status.py b/alembic/versions/c4e1a9d2f6b8_add_streaming_message_status.py new file mode 100644 index 0000000..dbf3e06 --- /dev/null +++ b/alembic/versions/c4e1a9d2f6b8_add_streaming_message_status.py @@ -0,0 +1,37 @@ +"""Add STREAMING message status. + +Revision ID: c4e1a9d2f6b8 +Revises: 91426eac7604 +Create Date: 2026-08-18 15:02:34.101000 +""" + +from typing import Sequence, Union + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "c4e1a9d2f6b8" +down_revision: Union[str, Sequence[str], None] = "91426eac7604" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Add STREAMING to the messagestatus enum. + + Note: `ALTER TYPE ... ADD VALUE` cannot run inside a transaction block + on PostgreSQL pre-v12, so we use an autocommit block. + """ + with op.get_context().autocommit_block(): + op.execute("ALTER TYPE messagestatus ADD VALUE IF NOT EXISTS 'STREAMING'") + + +def downgrade() -> None: + """Downgrade is intentionally unsupported. + + Postgres cannot drop enum values, and rebuilding the type would + require remapping any STREAMING rows to another status. + """ + raise NotImplementedError( + "Downgrade not supported: removing enum values would silently rewrite rows." + ) diff --git a/alembic/versions/f6ce7837e023_add_query_handles_table.py b/alembic/versions/f6ce7837e023_add_query_handles_table.py index 4fcc13a..d7e613c 100644 --- a/alembic/versions/f6ce7837e023_add_query_handles_table.py +++ b/alembic/versions/f6ce7837e023_add_query_handles_table.py @@ -1,4 +1,4 @@ -"""Add query handles table +"""Add query handles table. Revision ID: f6ce7837e023 Revises: 21d5a7602704 diff --git a/app/agent/prompts.py b/app/agent/prompts.py index c8e219d..42c5f34 100644 --- a/app/agent/prompts.py +++ b/app/agent/prompts.py @@ -5,7 +5,7 @@ # Capabilities -You can search and explore Base dos Dados's datasets and tables, query and analyze the data, translate coded values, and explain the platform and how you work. You have no other abilities — in particular, you cannot generate charts, plots, maps, visualizations, or files. Never offer, promise, or imply an action beyond these, in your prose answer or in the follow-up prompts. Downloads of query results are attached automatically by the interface — you do not generate them and need not offer them. +You can search and explore Base dos Dados's datasets and tables, query and analyze the data, translate coded values, export a query's results as a downloadable file on request, plot a query's results as a chart — including a choropleth map shading Brazilian states or municipalities by a value — and explain the platform and how you work. You cannot produce files other than those exports, or map any geography other than Brazilian states and municipalities. Never offer, promise, or imply an action beyond these, in your prose answer or in the follow-up prompts. # When to act vs. ask @@ -30,6 +30,20 @@ Answer without tools only to explain the platform or your own capabilities, to ask for clarification, or to reuse data you already retrieved earlier in this conversation. +# Exporting results + +`execute_bigquery_sql` returns the `query_ref` in its result. When the user explicitly asks to download or export a result in a specific format (AVRO, CSV, JSON Lines, Parquet), call `export_query_result` with that result's `query_ref` and the format. If the result is not found, use `list_query_results` to look it up. If the tool reports the result expired, re-run the query and export the new result. + +The interface offers the exported file — you do not generate or attach the file yourself. Do not describe the file's contents or claim anything about the download beyond the tool's confirmation. + +# Charting results + +When the user asks for a chart, plot, or visualization, call `chart_query_result` with the result's `query_ref` and a natural-language description of the chart — the mark (bar, line, point, …), what belongs on each axis, titles, labels and any grouping or color. You describe the chart in words; a data visualization specialist turns it into a Vega-Lite spec, so do not write the spec yourself. + +A value that varies across Brazilian states or municipalities can be drawn as a choropleth map (only these two levels, only Brazil). For a map, the result must carry the geographic key to join on — usually `sigla_uf` for states, `id_municipio` for municipalities — so write the query to return one row per state (or municipality) with that column plus the value, and describe it as a map of Brazil shaded by the value. Keep that key column in the result even though it is a code: the map joins on it and supplies the readable name itself, so this is the one case where you do not decode the geographic column away. + +Use the same `query_ref` referencing rules as exports. If the tool reports the result is too large, aggregate further in SQL, chart the smaller result and state this in your response. + # Brazilian data landscape Main data sources available and starting keywords for the search: diff --git a/app/agent/tools/__init__.py b/app/agent/tools/__init__.py index 1e2d5ca..cc97836 100644 --- a/app/agent/tools/__init__.py +++ b/app/agent/tools/__init__.py @@ -2,6 +2,11 @@ from app.agent.tools.api import get_dataset_details, get_table_details, search_datasets from app.agent.tools.bigquery import decode_table_values, execute_bigquery_sql +from app.agent.tools.dataviz import ( + chart_query_result, + export_query_result, + list_query_results, +) class BDToolkit: @@ -19,6 +24,9 @@ def get_tools() -> list[BaseTool]: - get_table_details: Get comprehensive table information. - execute_bigquery_sql: Execute SQL queries against BigQuery tables. - decode_table_values: Decode coded values using dictionary tables. + - list_query_results: List this conversation's exportable/chartable results. + - export_query_result: Offer a result as a downloadable file (AVRO/CSV/JSONL/PARQUET). + - chart_query_result: Render a chart from a result. """ return [ search_datasets, @@ -26,6 +34,9 @@ def get_tools() -> list[BaseTool]: get_table_details, execute_bigquery_sql, decode_table_values, + list_query_results, + export_query_result, + chart_query_result, ] diff --git a/app/agent/tools/bigquery.py b/app/agent/tools/bigquery.py index 236d773..013580f 100644 --- a/app/agent/tools/bigquery.py +++ b/app/agent/tools/bigquery.py @@ -44,7 +44,7 @@ def execute_bigquery_sql( current request — each slug names a separate download. Returns: - JSON object with `row_count` and `rows`. + JSON object with `row_count`, `rows` and `query_ref` (a reference to the query results). """ client = _bq_client() @@ -93,16 +93,19 @@ def execute_bigquery_sql( "the context small. The full result can still be downloaded from the interface." ) - content = json.dumps(payload, ensure_ascii=False, default=str) - # No rows -> nothing to download if not rows: - return content, None + return json.dumps(payload, ensure_ascii=False, default=str), None # Server-minted handle for the anonymous result table BigQuery already materialized # (~24h TTL), so a later export hands back exactly these rows without re-running. query_ref = f"qr_{uuid.uuid4().hex}" + # Surface the handle to the model so it can pass it to export_query_result / + # chart_query_result. The full destination table stays server-side in the artifact. + payload["query_ref"] = query_ref + content = json.dumps(payload, ensure_ascii=False, default=str) + artifact = { "type": "query_result", "query_ref": query_ref, diff --git a/app/agent/tools/dataviz.py b/app/agent/tools/dataviz.py new file mode 100644 index 0000000..72db524 --- /dev/null +++ b/app/agent/tools/dataviz.py @@ -0,0 +1,170 @@ +import asyncio +import json +from typing import Any + +from langchain.tools import ToolRuntime +from langchain_core.tools import tool + +from app.agent.context import AgentContext +from app.agent.tools.exceptions import handle_tool_errors +from app.charts import fetch_chart_data, generate_chart_spec, inject_chart_data +from app.db.database import AsyncDatabase, sessionmaker +from app.exports import ( + OFFERED_EXPORT_FORMATS, + ResultTableExpired, + ResultTooLarge, + is_result_expired, + materialize_export, + sanitize_export_filename, +) +from app.i18n import MessageKey, translate + + +@tool +@handle_tool_errors +async def list_query_results(runtime: ToolRuntime[AgentContext]) -> str: + """List this conversation's query results, including ones from earlier turns. + + Returns: + JSON array of {query_ref, description, ran_at, expired}, oldest first. + """ + async with sessionmaker() as session: + handles = await AsyncDatabase(session).get_query_handles_by_thread( + runtime.context.thread_id + ) + + results = [ + { + "query_ref": handle.query_ref, + "description": handle.slug, + "ran_at": handle.created_at.isoformat(), + "expired": is_result_expired(handle.created_at), + } + for handle in handles + ] + + return json.dumps(results, ensure_ascii=False, default=str) + + +@tool(response_format="content_and_artifact") +@handle_tool_errors(response_format="content_and_artifact") +async def export_query_result( + query_ref: str, file_format: str, runtime: ToolRuntime[AgentContext] +) -> tuple[str, dict[str, Any]]: + """Offer a query's result as a downloadable file in the requested format. + + Args: + query_ref (str): The handle of the result to export. + file_format (str): One of AVRO, CSV, JSONL, PARQUET. + + Returns: + A confirmation that the download is ready, or an error describing what to fix. + """ + file_format = file_format.upper() + + if file_format not in OFFERED_EXPORT_FORMATS: + raise ValueError( + f"Unsupported format '{file_format}'. " + f"Available: {', '.join(OFFERED_EXPORT_FORMATS)}." + ) + + async with sessionmaker() as session: + handle = await AsyncDatabase(session).get_query_handle_from_thread( + query_ref, runtime.context.thread_id + ) + + if handle is None: + raise ValueError( + f"No query result found for '{query_ref}'. " + "Call list_query_results to see the available results." + ) + + if is_result_expired(handle.created_at): + raise ValueError( + f"The result for '{query_ref}' has expired (results are kept ~24h). " + "Re-run the query, then export the new result." + ) + + # The user asked for the file explicitly, so materialize it now rather than lazily on click. + # This lets the card show the exact file size and the real download filename, and surfaces + # an over-limit result up front instead of a download that 400s on click. + try: + exported = await asyncio.to_thread( + materialize_export, + query_ref=handle.query_ref, + destination_table=handle.destination_table, + file_format=file_format, + filename=sanitize_export_filename( + handle.slug, + translate(MessageKey.DEFAULT_EXPORT_FILENAME, runtime.context.language), + ), + message_id=str(handle.message_id), + ) + except ResultTableExpired as e: + raise ValueError( + f"The result for '{query_ref}' has expired (results are kept ~24h). " + "Re-run the query, then export the new result." + ) from e + except ResultTooLarge as e: + raise ValueError( + f"The result for '{query_ref}' is too large to export as a single file. " + "Aggregate or filter it in SQL first, then export the smaller result." + ) from e + + # A client-facing affordance (not a `query_result` handle, so it is not redacted): + artifact = { + "type": "export", + "query_ref": handle.query_ref, + "format": file_format, + "filename": exported.filename, + "size_bytes": exported.size_bytes, + "message_id": str(handle.message_id), + } + + content = json.dumps( + { + "status": "ready", + "query_ref": query_ref, + "format": file_format, + "filename": exported.filename, + "size_bytes": exported.size_bytes, + }, + ensure_ascii=False, + ) + + return content, artifact + + +@tool(response_format="content_and_artifact") +@handle_tool_errors(response_format="content_and_artifact") +async def chart_query_result( + query_ref: str, instructions: str, runtime: ToolRuntime[AgentContext] +) -> tuple[str, dict[str, Any]]: + """Render a chart from a query's result. + + Args: + query_ref (str): The handle of the result to chart. + instructions (str): A natural-language description of the chart. + + Returns: + A confirmation that the chart was rendered, or an error describing what to fix. + """ + handle, columns, rows = await fetch_chart_data(query_ref, runtime.context.thread_id) + + spec = await generate_chart_spec(columns, rows, instructions) + + # A client-facing artifact (not a `query_result` handle, so it is not redacted): the + # interface renders it with Vega-Embed. The data is bound server-side from the exact + # result rows, so the model cannot substitute the numbers. + artifact = { + "type": "chart", + "query_ref": handle.query_ref, + "spec": inject_chart_data(spec, rows), + } + + content = json.dumps( + {"status": "rendered", "query_ref": query_ref, "row_count": len(rows)}, + ensure_ascii=False, + ) + + return content, artifact diff --git a/app/api/routers/chatbot.py b/app/api/routers/chatbot.py index 6e71a88..956d535 100644 --- a/app/api/routers/chatbot.py +++ b/app/api/routers/chatbot.py @@ -1,5 +1,4 @@ import asyncio -import re import uuid from fastapi import APIRouter, BackgroundTasks, HTTPException, Query, status @@ -26,10 +25,10 @@ ) from app.exports import ( OFFERED_EXPORT_FORMATS, - ExportFormat, ResultTableExpired, ResultTooLarge, materialize_export, + sanitize_export_filename, ) from app.i18n import MessageKey, translate from app.settings import settings @@ -239,28 +238,16 @@ def _cleanup(task: asyncio.Task): # pragma: no cover ) -def _sanitize_filename(slug: str, fallback: str) -> str: - """Sanitize a query's slug into a safe base filename. - - Args: - slug (str): The query's slug. - fallback (str): Base name to use when the slug yields nothing filesystem-safe. - - Returns: - str: A filesystem-safe base filename, without extension. - """ - filename = re.sub(r"[^\w-]+", "_", slug).strip("_") - return filename or fallback - - @router.post("/messages/{message_id}/exports") -async def export_message_results( +async def export_message_result( message_id: str, database: AsyncDB, user_id: UserID, query_ref: str, - file_format: ExportFormat = Query("CSV", alias="format"), + file_format: str = Query("CSV", alias="format"), ): + file_format = file_format.upper() + if file_format not in OFFERED_EXPORT_FORMATS: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -272,7 +259,7 @@ async def export_message_results( message, thread = await _authorize_message(database, message_id, user_id) - query_handle = await database.get_query_handle(message.id, query_ref) + query_handle = await database.get_query_handle_from_message(query_ref, message.id) if query_handle is None: raise HTTPException( @@ -286,7 +273,7 @@ async def export_message_results( query_ref=query_handle.query_ref, destination_table=query_handle.destination_table, file_format=file_format, - filename=_sanitize_filename( + filename=sanitize_export_filename( query_handle.slug, translate(MessageKey.DEFAULT_EXPORT_FILENAME, thread.language), ), diff --git a/app/api/streaming/agent_runner.py b/app/api/streaming/agent_runner.py index f3cf440..6108593 100644 --- a/app/api/streaming/agent_runner.py +++ b/app/api/streaming/agent_runner.py @@ -20,27 +20,24 @@ MessageStatus, QueryHandle, ) -from app.exports import CollectedQueryHandle, collect_query_handles +from app.exports import collect_query_handles from app.i18n import LanguageCode, MessageKey, translate def _truncate_json( json_string: str, max_list_len: int = 10, max_str_len: int = 300 ) -> str: - """Iteratively truncates a serialized JSON object by shortening lists and strings - and adding human-readable placeholders. + """Shorten a serialized JSON object's long lists and strings, with placeholders. - Note: - This function only processes JSON objects (dictionaries). If the serialized JSON - represents any other type, the original JSON string will be returned unchanged. + Non-dict JSON is returned unchanged. Args: json_string (str): The serialized JSON to process. - max_list_len (int, optional): The max number of items to keep in a list. Defaults to 10. - max_str_len (int, optional): The max length for any single string. Defaults to 300. + max_list_len (int, optional): Max items to keep in a list. Defaults to 10. + max_str_len (int, optional): Max length for any single string. Defaults to 300. Returns: - str: The truncated, formatted, and serialized JSON object. + str: The truncated, re-serialized JSON object. """ try: data = json.loads(json_string) @@ -83,19 +80,15 @@ def _truncate_json( def _process_chunk(chunk: dict[str, Any], language: LanguageCode) -> StreamEvent | None: - """Process a streaming chunk from a react agent workflow into a StreamEvent. + """Turn a raw agent stream chunk into a StreamEvent. Args: chunk (dict[str, Any]): A raw update chunk from the agent workflow. - language (LanguageCode): A supported language code, for localizing server-emitted content. + language (LanguageCode): Language for localizing server-emitted content. Returns: - StreamEvent | None: Structured event or None if the chunk is ignored: - - "tool_call" for agent messages with tool calls - - "tool_output" for tool execution results - - "final_answer" for agent messages without tool calls - - "model_call_limit" when the model call limit is reached - - None for ignored chunks + StreamEvent | None: The tool_call, tool_output, final_answer or model_call_limit + event, or None for an ignored chunk. """ if "model" in chunk: update: dict[str, Any] = chunk["model"] @@ -189,29 +182,115 @@ def _process_chunk(chunk: dict[str, Any], language: LanguageCode) -> StreamEvent return None +async def _create_placeholder_message( + *, + run_id: str, + thread_id: str, + user_message: Message, + model_uri: str, +) -> None: + """Create the assistant row up front (id == run_id, STREAMING) so query handles + can be persisted against a real FK during the run. + + Raises on failure so the caller aborts the run: persistence is deterministic, + and :func:`_finalize_message` only ever updates this row. + + Args: + run_id (str): The run id, reused as the message id. + thread_id (str): The thread the message belongs to. + user_message (Message): The user message driving the run. + model_uri (str): Model URI. + """ + message_create = MessageCreate( + id=run_id, + thread_id=thread_id, + user_message_id=user_message.id, + model_uri=model_uri, + role=MessageRole.ASSISTANT, + content="", + status=MessageStatus.STREAMING, + ) + async with sessionmaker() as session: + await AsyncDatabase(session).create_message(message_create) + + async def _persist_query_handles( - database: AsyncDatabase, *, - message_id: str, - collected_handles: list[CollectedQueryHandle], -): - """Persist query handles for a message. + run_id: str, + tool_outputs: list[ToolOutput], +) -> None: + """Persist the query handles carried by the tool outputs, scoped to the run's message. + + Best-effort: a failure is logged and never interrupts the run. No dedup — a repeated + query_ref would be a bug, so a duplicate-key insert surfaces it rather than hiding it. Args: - database (AsyncDatabase): The repository to persist through. - message_id (str): The owning message/run the handles are scoped to. - collected_handles: list[CollectedQueryHandle]: Collected query handles. + run_id (str): The owning message/run the handles are scoped to. + tool_outputs (list[ToolOutput]): The tool outputs to scan for query_result artifacts. """ + collected = collect_query_handles(output.artifact for output in tool_outputs) + + if not collected: + return + handles = [ QueryHandle( query_ref=handle.query_ref, - message_id=message_id, + message_id=run_id, slug=handle.slug, destination_table=handle.destination_table, ) - for handle in collected_handles + for handle in collected ] - await database.create_query_handles(handles) + + try: + async with sessionmaker() as session: + await AsyncDatabase(session).create_query_handles(handles) + except Exception: + logger.exception(f"Failed to persist query handles for run {run_id}:") + + +async def _finalize_message( + *, + run_id: str, + content: str, + events: list[dict[str, Any]] | None, + structured_response: dict[str, Any] | None, + status: MessageStatus, +) -> tuple[str | None, dict[str, Any] | None]: + """Write the terminal state onto the placeholder created at run start. + + The placeholder is guaranteed to exist (its creation gates the run), so this only updates. + + Args: + run_id (str): The run/message id to finalize. + content (str): The assistant message content. + events (list[dict[str, Any]] | None): The streamed events to persist. + structured_response (dict[str, Any] | None): The structured response, if any. + status (MessageStatus): The terminal status to write. + + Returns: + tuple[str | None, dict[str, Any] | None]: (message_id, None) on success, or + (None, error_details) if the row could not be written. + """ + try: + async with sessionmaker() as session: + message = await AsyncDatabase(session).update_message( + run_id, + content=content, + events=events, + structured_response=structured_response, + status=status, + ) + if message is None: + logger.error( + f"Placeholder message for run {run_id} vanished before finalize" + ) + return None, {"reason": "persistence_failed"} + return str(message.id), None + except Exception: + logger.exception(f"Failed to persist assistant message for run {run_id}:") + return None, {"reason": "persistence_failed"} async def run_agent( @@ -223,28 +302,55 @@ async def run_agent( model_uri: str, queue: asyncio.Queue[StreamEvent], ): - """Run the agent to completion and push events onto the queue. + """Run the agent to completion, streaming events onto the queue and owning persistence. - Owns persistence: writes the assistant `messages` row in `finally` and - emits a terminal `complete` event carrying either the persisted run_id - (on success) or `error_details` (if persistence fails). Exactly one - `complete` event is emitted per run. + Creates the assistant row up front (STREAMING), persists query handles as tool outputs + arrive, writes the terminal state in `finally`, and emits exactly one `complete` event. Args: - agent (CompiledStateGraph): Agent compiled state graph. + agent (CompiledStateGraph): The compiled agent graph. config (ConfigDict): Config for agent execution. context (AgentContext): The run context. - thread_id (str): Thread unique identifier. - user_message (Message): User message. + thread_id (str): Thread identifier. + user_message (Message): The user message driving the run. model_uri (str): Model URI. - queue (asyncio.Queue[StreamEvent]): Events queue. + queue (asyncio.Queue[StreamEvent]): Queue the events are pushed onto. """ + run_id = config["run_id"] events = [] assistant_message = "" structured_response: dict[str, Any] | None = None - collected_handles: list[CollectedQueryHandle] = [] status: MessageStatus | None = None + # Create the assistant row up front so query handles persist eagerly against a real FK + # during the run. It stays STREAMING until the finally block writes the terminal state. + try: + await _create_placeholder_message( + run_id=run_id, + thread_id=thread_id, + user_message=user_message, + model_uri=model_uri, + ) + except Exception: + logger.exception(f"Failed to create placeholder message for run {run_id}:") + error_details = {"reason": "persistence_failed"} + await queue.put( + StreamEvent( + type="error", + data=EventData( + content=translate(MessageKey.ERROR_UNEXPECTED, context.language), + error_details=error_details, + ), + ) + ) + await queue.put( + StreamEvent( + type="complete", + data=EventData(run_id=None, error_details=error_details), + ) + ) + return + try: async for mode, chunk in agent.astream( # pragma: no cover input={"messages": [{"role": "user", "content": user_message.content}]}, @@ -261,11 +367,11 @@ async def run_agent( continue if event.type == "tool_output": - # Collect every query handle from the execute_bigquery_sql - # tool artifacts, for lazy on-click downloads. - collect_query_handles( - (output.artifact for output in event.data.tool_outputs), - collected_handles, + # Persist each query handle as its tool output arrives, so a later + # tool in the same run (or a later turn) can resolve it from the DB. + await _persist_query_handles( + run_id=run_id, + tool_outputs=event.data.tool_outputs, ) elif event.type == "final_answer": # Resolve data source names to a localized "{dataset_name} — {table_name}". @@ -304,42 +410,13 @@ async def run_agent( events.append(event.model_dump()) await queue.put(event) finally: - message_create = MessageCreate( - id=config["run_id"], - thread_id=thread_id, - user_message_id=user_message.id, - model_uri=model_uri, - role=MessageRole.ASSISTANT, + message_id, error_details = await _finalize_message( + run_id=run_id, content=assistant_message, events=events or None, structured_response=structured_response, status=status or MessageStatus.ERROR, ) - try: - async with sessionmaker() as session: - database = AsyncDatabase(session) - message = await database.create_message(message_create) - message_id = str(message.id) - error_details = None - # Query handles are a best-effort download convenience, persisted after - # the message (its own commit) so a handle failure never loses the message. - if collected_handles: - try: - await _persist_query_handles( - database, - message_id=message_id, - collected_handles=collected_handles, - ) - except Exception: - logger.exception( - f"Failed to persist query handles for run {config['run_id']}:" - ) - except Exception: - logger.exception( - f"Failed to persist assistant message for run {config['run_id']}:" - ) - message_id = None - error_details = {"reason": "persistence_failed"} await queue.put( StreamEvent( type="complete", diff --git a/app/charts.py b/app/charts.py new file mode 100644 index 0000000..0a3dcb1 --- /dev/null +++ b/app/charts.py @@ -0,0 +1,570 @@ +import asyncio +import json +from collections.abc import Callable +from decimal import Decimal +from functools import cache +from typing import Any + +import vl_convert as vlc +from google.api_core.exceptions import NotFound +from google.cloud import bigquery as bq +from langchain_core.messages import HumanMessage, SystemMessage +from langchain_openai import ChatOpenAI +from pydantic import BaseModel, Field, JsonValue + +from app.db.database import AsyncDatabase, sessionmaker +from app.db.models import QueryHandle +from app.exports import ResultTableExpired, is_result_expired +from app.settings import settings + +VEGA_LITE_SCHEMA = "https://vega.github.io/schema/vega-lite/v6.json" + +# Rows shown to the spec generator. It only needs the shape (columns + a few +# example values); the full result is bound afterwards by inject_chart_data. +CHART_SPEC_SAMPLE_ROWS = 10 + +# How many times the generator may retry after a spec fails validation before giving up. +MAX_CHART_SPEC_ATTEMPTS = 3 + +# Keys always stripped anywhere in a model-generated spec for security: `url`/`datasets` +# are external/inline data (SSRF, provenance); `href` makes a mark a clickable link, +# whose `javascript:` URI would be a click-XSS in the viewer's browser. +_UNTRUSTED_KEYS = frozenset({"datasets", "url", "href"}) + +# Geographic assets a spec may inject by name for choropleth maps. The TopoJSON +# files are served as static assets by the website (see settings.GEO_ASSET_URL_BASE). +_GEO_ASSETS: dict[str, dict[str, str]] = { + "brazil_states": { + "file": "brazil_states.topojson", + "feature": "uf", + }, + "brazil_municipalities": { + "file": "brazil_municipalities.topojson", + "feature": "Munic", + }, +} + +# The result rows, referenced from a spec as {"data": {"name": "query_result"}}. +_RESULT_SOURCE = "query_result" + +# Every data-source name a spec may reference; anything else in a `data` node is dropped. +_ALLOWED_SOURCES = frozenset({_RESULT_SOURCE, *_GEO_ASSETS}) + +# System prompt for the chart spec generation. +_CHART_SPEC_INSTRUCTIONS = """\ +You are a data visualization specialist. Given a small, already-aggregated query result and a description of the chart to build, return one complete Vega-Lite v6 spec. Reference the result's columns by their exact names, and never include a data source, dataset, or URL — the exact rows are injected separately. You may use composite views and transforms when they make the chart clearer. + +# Guidelines + +- Follow the description when it is specific. When it leaves the chart type open, pick the clearest form for the data's shape. +- Add a short title, label both axes with human-readable text rather than the raw column names. +- Add a legend whenever a view shows two or more series. + +## Color + +Do not set colors, scales, or ranges for a categorical or single-series encoding — just map a field to the color channel to distinguish series and leave the palette to the defaults. + +For a numeric value mapped to color (any mark colored by a number), set the color scale's `scheme` only when the value's semantic meaning has a color convention worth matching; otherwise leave it to the default. For example: + +- Temperature → blue for cold, red for hot: use "redblue" or "redyellowblue" (add `"reverse": true` if low and high land on the wrong ends). +- A quantity that diverges around a meaningful midpoint → a diverging scheme such as "redblue" or "blueorange" with the scale's `domainMid` set to that midpoint. + +When you do name a scheme, use a real Vega scheme name — never one from another library. Prefer colorblind-safe schemes. + +## Number formatting + +Abbreviate any large-magnitude axis or label with SI notation (the `~s` / `s` format), then rewrite its prefix letters (k, M, G, T) to follow the naming convention the measured quantity uses in the response's language, via an axis `labelExpr` that reformats `datum.value`. Do this on every axis whose numbers reach the thousands or beyond. + +## Log scales + +A log scale is undefined at zero and below, so a single non-positive value in the encoded field collapses the whole axis — no ticks, every point crushed into a corner. Real-world columns routinely carry zeros, nulls, and occasional negatives. So whenever you set `"scale": {"type": "log"}`, first drop the non-positive rows with a `filter` transform over every log-scaled field. If a field genuinely spans zero and cannot be filtered, use a linear scale instead. + +## Size + +The chart is rendered at the container's width, so never set a width. For any chart with many categories along one axis, set an explicit height so the cells or bands are not overly tall — a height noticeably smaller than the chart's width reads best. Otherwise leave the height to default. + +## Choropleth maps + +To shade a Brazilian map by a value, use a `geoshape` mark whose data is the map geometry, injected by name, joined to your result with a `lookup` transform. Never use a URL. Available geometry (each feature's `id` is the join key, its `properties.name` the label): + +- `brazil_states`: the 27 states; `id` is the UF code as text (e.g. "SP"). Join it to your result's UF column (usually `sigla_uf`). +- `brazil_municipalities`: the municipalities; `id` is the 7-digit IBGE code as text (e.g. "3550308"). Join it to your result's municipality id column (usually `id_municipio`). + +Set the geoshape's data to `{"name": "brazil_states"}` (or `"brazil_municipalities"`) and reference your result as `{"name": "query_result"}`; the exact geometry and rows are injected separately. Join with a `lookup` on the geometry's `id`, matching your result's location column as the `key`, and pull the value column via `fields`. Color by that value. Use a `mercator` projection, add a `properties.name` + value tooltip, and set an explicit `height` (Brazil is roughly square, so a height close to the width works well). Examples: + +{ + "data": {"name": "brazil_states"}, + "transform": [{"lookup": "id", "from": {"data": {"name": "query_result"}, "key": "sigla_uf", "fields": ["valor"]}}], + "projection": {"type": "mercator"}, + "mark": "geoshape", + "encoding": { + "color": {"field": "valor", "type": "quantitative", "title": "…", "scale": {"scheme": "yellowgreenblue"}}, + "tooltip": [{"field": "properties.name", "type": "nominal", "title": "Estado"}, {"field": "valor", "type": "quantitative"}] + }, +} + +{ + "data": {"name": "brazil_municipalities"}, + "transform": [{"lookup": "id", "from": {"data": {"name": "query_result"}, "key": "id_municipio", "fields": ["valor"]}}], + "projection": {"type": "mercator"}, + "mark": "geoshape", + "encoding": { + "color": {"field": "valor", "type": "quantitative", "title": "…", "scale": {"scheme": "yellowgreenblue"}}, + "tooltip": [{"field": "properties.name", "type": "nominal", "title": "Município"}, {"field": "valor", "type": "quantitative"}] + }, +} +""" + + +class ChartHandleNotFound(Exception): + """No query result exists for the given query_ref in this thread.""" + + +class ChartResultTooLarge(Exception): + """The result has more rows than a chart should bind.""" + + +class ChartSpecInvalid(Exception): + """The generator could not produce a spec that passes validation.""" + + +class ChartSpec(BaseModel): + """The tool the charting model calls: a free-form Vega-Lite spec.""" + + # Description kept to a concise label per OpenAI's gpt-5.x guidance ("state each + # instruction once"); the how-to lives once in _CHART_SPEC_INSTRUCTIONS (system prompt). + spec: dict[str, Any] = Field(description="The complete Vega-Lite v6 specification.") + + +# =================================================================== +# CHART DATA FETCHING +# =================================================================== +@cache +def _bq_client() -> bq.Client: # pragma: no cover + return bq.Client( + project=settings.GOOGLE_BILLING_PROJECT, + credentials=settings.GOOGLE_CREDENTIALS, + ) + + +def _json_default(value: Any) -> Any: + """Coerce a BigQuery value that JSON can't render natively into a JSON-native one. + + `Decimal` (NUMERIC/BIGNUMERIC) becomes a number so quantitative encodings and color + scales see a number, not a string. Everything else (date/datetime/…) falls back to + its string form, which Vega-Lite parses for temporal encodings. + + Args: + value (Any): The value json.dumps could not serialize. + + Returns: + Any: A JSON-native replacement. + """ + if isinstance(value, Decimal): + if value == value.to_integral_value(): + return int(value) + return float(value) + return str(value) + + +def _fetch_rows( + destination_table: dict[str, Any], +) -> tuple[list[str], list[dict[str, Any]]]: + """Fetch a result table's rows for charting, capped at settings.CHART_MAX_BYTES. + + Fetches row by row and stops as soon as the JSON that would be bound into the spec + exceeds the budget, so a huge result is rejected without materializing all of it. + + Args: + destination_table (dict[str, Any]): `TableReference.to_api_repr()` of the result table. + + Returns: + tuple[list[str], list[dict[str, Any]]]: Column names and row dicts. + + Raises: + ChartResultTooLarge: The bound data would exceed settings.CHART_MAX_BYTES. + ResultTableExpired: The result table no longer exists (~24h TTL). + """ + table_ref = bq.TableReference.from_api_repr(destination_table) + + try: + table = _bq_client().get_table(table_ref) + + columns = [field.name for field in table.schema] + + rows = [] + size = 0 + for row in _bq_client().list_rows(table_ref): + # Serialize once to both measure the row and coerce BigQuery types + # (date/datetime/Decimal/…) to JSON-native values. Raw objects are + # not valid JSON and would fail serialization when the spec is bound. + serialized = json.dumps( + dict(row), ensure_ascii=False, default=_json_default + ) + size += len(serialized.encode()) + if size > settings.CHART_MAX_BYTES: + raise ChartResultTooLarge( + f"The result is too large to chart (over " + f"{settings.CHART_MAX_BYTES // (1024 * 1024)} MB of data). Aggregate " + "or summarize it in SQL first, then chart the smaller result." + ) + rows.append(json.loads(serialized)) + except NotFound as e: + raise ResultTableExpired(str(e)) from e + + return columns, rows + + +async def fetch_chart_data( + query_ref: str, thread_id: str +) -> tuple[QueryHandle, list[str], list[dict[str, Any]]]: + """Resolve a chartable result: authorize the handle, then read its capped rows. + + Args: + query_ref (str): The handle of the result to chart. + thread_id (str): The thread the handle must belong to (authorization). + + Returns: + tuple[QueryHandle, list[str], list[dict[str, Any]]]: The handle, its columns, and rows. + + Raises: + ChartHandleNotFound: No such result in this thread. + ResultTableExpired: The result expired (by age or missing table). + """ + async with sessionmaker() as session: + handle = await AsyncDatabase(session).get_query_handle_from_thread( + query_ref, thread_id + ) + + if handle is None: + raise ChartHandleNotFound( + f"No query result found for '{query_ref}'. " + "Call list_query_results to see the available results." + ) + + if is_result_expired(handle.created_at): + raise ResultTableExpired( + f"The result for '{query_ref}' has expired (results are kept ~24h). " + "Re-run the query, then chart the new result." + ) + + columns, rows = await asyncio.to_thread(_fetch_rows, handle.destination_table) + + return handle, columns, rows + + +# =================================================================== +# CHART DATA INJECTION +# =================================================================== +def _geo_url_node(name: str) -> dict[str, Any]: + """Build a Vega-Lite TopoJSON data node that points at the geometry by URL. + + The URL is the website's static path for the file; the browser fetches (and caches) the + ~megabytes of geometry directly, so it never inflates the spec, the SSE payload, or the + persisted `events`. + + Args: + name (str): A geo source name from `_GEO_ASSETS` (e.g. "brazil_states"). + + Returns: + dict[str, Any]: The `data` node — a topojson `url` plus its topojson `format`. + """ + asset = _GEO_ASSETS[name] + return { + "url": f"{settings.GEO_ASSET_URL_BASE.rstrip('/')}/{asset['file']}", + "format": {"type": "topojson", "feature": asset["feature"]}, + } + + +def _geo_stub_node(name: str) -> dict[str, Any]: + """Build a minimal single-feature TopoJSON node for offline spec validation. + + Args: + name (str): A geo source name from `_GEO_ASSETS` (e.g. "brazil_states"). + + Returns: + dict[str, Any]: A `data` node with a tiny inline TopoJSON and its topojson `format`. + """ + feature = _GEO_ASSETS[name]["feature"] + stub = { + "type": "Topology", + "objects": { + feature: { + "type": "GeometryCollection", + "geometries": [ + {"type": "Polygon", "id": "0", "properties": {}, "arcs": [[0]]} + ], + } + }, + "arcs": [[[0, 0], [1, 0], [0, 1], [0, 0]]], + } + return {"values": stub, "format": {"type": "topojson", "feature": feature}} + + +def _resolve_named_data( + node: JsonValue, + rows: list[dict[str, Any]], + geo_node: Callable[[str], dict[str, Any]], +) -> JsonValue: + """Replace every `{"name": }` data reference with real data. + + `query_result` resolves to the rows; a geo name resolves via a `geo_node` — a URL + node for a render-ready spec, a tiny inline stub for offline validation. Any other + `{"name": ...}` was already dropped by `_sanitize_chart_spec`. + + Args: + node (JsonValue): A spec, or any node within it, to walk. + rows (list[dict[str, Any]]): The exact result rows to bind for `query_result`. + geo_node (Callable[[str], dict[str, Any]]): Resolver for a geo source name. + + Returns: + JsonValue: The same node with its named data references resolved. + """ + if isinstance(node, dict): + if set(node) == {"name"} and node["name"] in _ALLOWED_SOURCES: + if node["name"] == _RESULT_SOURCE: + return {"values": rows} + return geo_node(node["name"]) + return { + key: _resolve_named_data(value, rows, geo_node) + for key, value in node.items() + } + if isinstance(node, list): + return [_resolve_named_data(item, rows, geo_node) for item in node] + return node + + +def inject_chart_data( + spec: dict[str, Any], rows: list[dict[str, Any]] +) -> dict[str, Any]: + """Inject the real geometry and result rows into an already-sanitized spec. + + Takes an already sanitized spec from `generate_chart_spec` and makes it render-ready: + resolves its named sources (a geo name to a TopoJSON URL, `query_result` to the rows) + and binds the rows as the default data for a plain chart that declared none. + + Args: + spec (dict[str, Any]): A sanitized, data-less spec from `generate_chart_spec`. + rows (list[dict[str, Any]]): The exact result rows to render. + + Returns: + dict[str, Any]: A render-ready Vega-Lite spec with the real data bound in. + """ + chart = _resolve_named_data(spec, rows, _geo_url_node) + chart["$schema"] = VEGA_LITE_SCHEMA + chart.setdefault("data", {"values": rows}) + return chart + + +# =================================================================== +# CHART SPEC GENERATION +# =================================================================== +@cache +def _chart_spec_model(): # pragma: no cover + """A model that returns a parsed `ChartSpec` via forced function-calling. + + `method="function_calling"` (not strict json_schema, which 400s on the open spec) and + `include_raw=True` (a bad call surfaces as `parsed=None`, not a raise) are required. + """ + return ChatOpenAI( + api_key=settings.OPENAI_API_KEY, + model=settings.MODEL_URI, + reasoning={ + "effort": "medium", + "summary": "auto", + }, + ).with_structured_output( + ChartSpec, + method="function_calling", + include_raw=True, + ) + + +def _chart_spec_user_prompt( + columns: list[str], + sample: list[dict[str, Any]], + instructions: str, + previous_spec: dict[str, Any] | None, + errors: list[str], +) -> str: + """Build the task-specific user message. + + On a retry, appends the model's own rejected spec and the validator's problems so it + can fix that spec rather than regenerate from scratch. + + Args: + columns (list[str]): The result's column names. + sample (list[dict[str, Any]]): A few example rows, for shape only. + instructions (str): The natural-language description of the chart to produce. + previous_spec (dict[str, Any] | None): The last rejected spec, echoed back on a retry. + errors (list[str]): The reasons the previous spec was rejected; empty on the first try. + + Returns: + str: The user message for the chart-spec model. + """ + prompt = ( + f"What to chart: {instructions}\n\n" + f"Columns: {json.dumps(columns, ensure_ascii=False)}\n\n" + f"Sample rows: {json.dumps(sample, ensure_ascii=False, default=str)}" + ) + + if errors: + # Show the model its own rejected spec so it can edit that spec to fix the + # problems, rather than regenerating from scratch on each retry. + if previous_spec is not None: + prompt += ( + "\n\nYour previous spec:\n" + f"{json.dumps(previous_spec, ensure_ascii=False, default=str)}" + ) + feedback = "- " + "\n- ".join(errors) + prompt += f"\n\nIt was rejected — fix these problems:\n{feedback}" + + return prompt + + +def _sanitize_chart_spec(node: JsonValue) -> JsonValue: + """Sanitize a model spec: drop untrusted keys, keep only allowlisted named sources. + + Every `_UNTRUSTED_KEYS` key (`datasets`, `url`, `href`) is removed outright, and a `data` + node survives only when it is exactly `{"name": }`, so the model can + never supply its own data or a clickable link. + + Args: + node (JsonValue): A spec, or any node within it, to walk. + + Returns: + JsonValue: The same node with untrusted data removed. + """ + if isinstance(node, dict): + out: dict[str, Any] = {} + for key, value in node.items(): + if key in _UNTRUSTED_KEYS: + continue + elif key == "data": + if ( + isinstance(value, dict) + and set(value) == {"name"} + and value["name"] in _ALLOWED_SOURCES + ): + out[key] = {"name": value["name"]} + # any other data node (inline values, url, unknown name) is dropped + continue + out[key] = _sanitize_chart_spec(value) + return out + if isinstance(node, list): + return [_sanitize_chart_spec(item) for item in node] + return node + + +def _collect(node: JsonValue, key: str) -> set[str]: + """Collect every string value stored under `key`, anywhere in the spec. + + Args: + node (JsonValue): A spec, or any node within it, to walk. + key (str): The key whose string values to collect (e.g. "field"). + + Returns: + set[str]: Every string found under `key`. + """ + found: set[str] = set() + if isinstance(node, dict): + for k, value in node.items(): + if k == key and isinstance(value, str): + found.add(value) + else: + found |= _collect(value, key) + elif isinstance(node, list): + for item in node: + found |= _collect(item, key) + return found + + +def _validate_chart_spec(spec: dict[str, Any], columns: list[str]) -> list[str]: + """Return the reasons a spec would not render a correct chart (empty if valid). + + Args: + spec (dict[str, Any]): The candidate spec (already sanitized). + columns (list[str]): The result's real column names. + + Returns: + list[str]: Human-readable problems to feed back to the generator; empty if valid. + """ + errors: list[str] = [] + + referenced = _collect(spec, "field") + derived = _collect(spec, "as") # fields produced by transforms + # A geoshape references its geometry's own fields — the top-level `id` and its labels + # under `properties.*` — which are not result columns; we allow those and flag the rest. + unknown = referenced - set(columns) - derived - {"*", "id"} + missing = sorted(field for field in unknown if not field.startswith("properties.")) + + if missing: + errors.append( + f"Encoding references column(s) not in the result: {missing}. " + f"Available columns: {sorted(columns)}." + ) + + # Resolve named sources (stub geometry, empty rows) so a geoshape/lookup spec renders + # its geometry; a normal chart just gets empty data. Rendering (not just VL→Vega compiling) + # is the only way to validate expression strings in the spec. + compile_spec = _resolve_named_data(spec, [], _geo_stub_node) + compile_spec = {**compile_spec, "$schema": VEGA_LITE_SCHEMA} + compile_spec.setdefault("data", {"values": []}) + + try: + vlc.vegalite_to_svg(json.dumps(compile_spec)) + except Exception as e: + # vl-convert raises on any spec the renderer rejects at build time + errors.append(f"The chart spec does not compile: {e}") + + return errors + + +async def generate_chart_spec( + columns: list[str], rows: list[dict[str, Any]], instructions: str +) -> dict[str, Any]: + """Generate a sanitized Vega-Lite spec for a result, retrying until it validates. + + Args: + columns (list[str]): The result's column names. + rows (list[dict[str, Any]]): The result rows; only a sample is shown to the model. + instructions (str): A natural-language description of the chart to produce. + + Returns: + dict[str, Any]: A validated, sanitized spec (no data bound yet). + + Raises: + ChartSpecInvalid: No attempt produced a spec that validates. + """ + sample = rows[:CHART_SPEC_SAMPLE_ROWS] + previous_spec: dict[str, Any] | None = None + errors: list[str] = [] + + for _attempt in range(MAX_CHART_SPEC_ATTEMPTS): + response = await _chart_spec_model().ainvoke( + [ + SystemMessage(_CHART_SPEC_INSTRUCTIONS), + HumanMessage( + _chart_spec_user_prompt( + columns, sample, instructions, previous_spec, errors + ) + ), + ] + ) + + parsed: ChartSpec | None = response["parsed"] + + if parsed is None: + errors = ["No chart was produced — return a complete Vega-Lite v6 spec."] + else: + spec = _sanitize_chart_spec(parsed.spec) + assert isinstance(spec, dict) # an object spec sanitizes to an object + previous_spec = spec # feed the spec back if it fails to validate + errors = await asyncio.to_thread(_validate_chart_spec, spec, columns) + if not errors: + return spec + + raise ChartSpecInvalid( + f"Could not produce a valid chart after {MAX_CHART_SPEC_ATTEMPTS} attempts: " + + "; ".join(errors) + ) diff --git a/app/db/database.py b/app/db/database.py index 49f5809..596c5ca 100644 --- a/app/db/database.py +++ b/app/db/database.py @@ -1,5 +1,5 @@ from datetime import datetime, timezone -from typing import TypeVar +from typing import Any, TypeVar from uuid import UUID from loguru import logger @@ -18,6 +18,7 @@ FeedbackSyncStatus, Message, MessageCreate, + MessageStatus, QueryHandle, Thread, ThreadCreate, @@ -175,6 +176,44 @@ async def create_message(self, message_create: MessageCreate) -> Message: return message + async def update_message( + self, + message_id: str | UUID, + *, + content: str, + events: Any, + structured_response: Any, + status: MessageStatus, + ) -> Message | None: + """Write the terminal fields of a streaming message row. + + Args: + message_id (str | UUID): The message (run) to finalize. + content (str): The final assistant text. + events (Any): The serialized stream events, or None. + structured_response (Any): The structured response payload, or None. + status (MessageStatus): The terminal status to write. + + Returns: + Message | None: The updated message, or None if it doesn't exist + (e.g. the placeholder was never created). + """ + message = await self.session.get(Message, message_id) + + if message is None: + self.logger.warning(f"Message {message_id} not found for update") + return None + + message.content = content + message.events = events + message.structured_response = structured_response + message.status = status + + await self.session.commit() + await self.session.refresh(message) + + return message + async def get_message(self, message_id: str | UUID) -> Message | None: """Get a message from the messages table. @@ -205,9 +244,14 @@ async def get_messages( """ # Eager-load the query handles so `MessagePublic.downloads` can # derive the download affordance without a lazy load per message. + # STREAMING rows are in-flight placeholders — exclude them so the + # listing shows only completed messages query = ( select(Message) - .where(Message.thread_id == thread_id) + .where( + Message.thread_id == thread_id, + Message.status != MessageStatus.STREAMING, + ) .options(selectinload(Message.query_handles)) ) query = self._apply_order_by(query, Message, order_by) @@ -230,22 +274,29 @@ async def create_query_handles(self, query_handles: list[QueryHandle]) -> None: self.session.add_all(query_handles) await self.session.commit() - async def get_query_handle( - self, message_id: str | UUID, query_ref: str + async def get_query_handle_from_message( + self, + query_ref: str, + message_id: str | UUID, ) -> QueryHandle | None: - """Get a stored query handle by its (`message_id`, `query_ref`). + """Resolve a query handle by `query_ref`, scoped to a message. Args: - message_id (str | UUID): The message the handle is scoped to. query_ref (str): The short handle minted by `execute_bigquery_sql`. + message_id (str | UUID): The message the handle is scoped to. Returns: QueryHandle | None: The handle if found, None otherwise. """ - query_handle = await self.session.get( - QueryHandle, {"message_id": message_id, "query_ref": query_ref} + result = await self.session.execute( + select(QueryHandle).where( + QueryHandle.query_ref == query_ref, + QueryHandle.message_id == message_id, + ) ) + query_handle = result.scalar_one_or_none() + if query_handle is None: self.logger.warning( f"Query handle {query_ref} for message {message_id} not found" @@ -253,6 +304,46 @@ async def get_query_handle( return query_handle + async def get_query_handle_from_thread( + self, query_ref: str, thread_id: str | UUID + ) -> QueryHandle | None: + """Resolve a query handle by `query_ref`, scoped to a thread. + + Args: + query_ref (str): The globally-unique handle to resolve. + thread_id (str | UUID): The thread the handle is scoped to. + + Returns: + QueryHandle | None: The handle if found in this thread, None otherwise. + """ + result = await self.session.execute( + select(QueryHandle) + .join(Message, QueryHandle.message_id == Message.id) + .where(QueryHandle.query_ref == query_ref, Message.thread_id == thread_id) + ) + + return result.scalar_one_or_none() + + async def get_query_handles_by_thread( + self, thread_id: str | UUID + ) -> list[QueryHandle]: + """List every query handle produced in a thread, oldest first. + + Args: + thread_id (str | UUID): The thread whose handles to list. + + Returns: + list[QueryHandle]: The thread's handles, ordered by creation time. + """ + result = await self.session.execute( + select(QueryHandle) + .join(Message, QueryHandle.message_id == Message.id) + .where(Message.thread_id == thread_id) + .order_by(QueryHandle.created_at) + ) + + return list(result.scalars().all()) + # ==================================== Feedback ==================================== async def upsert_feedback( self, feedback_create: FeedbackCreate diff --git a/app/db/models.py b/app/db/models.py index 234a332..c4fb03b 100644 --- a/app/db/models.py +++ b/app/db/models.py @@ -57,6 +57,10 @@ class MessageRole(str, Enum): class MessageStatus(str, Enum): + # Non-terminal: the assistant row is created up front and stays STREAMING while + # the run is in flight (hidden from the thread listing) until a terminal status + # is written. Every other value is terminal. + STREAMING = "STREAMING" ERROR = "ERROR" SUCCESS = "SUCCESS" INTERRUPTED = "INTERRUPTED" @@ -120,15 +124,9 @@ def downloads(self) -> list[dict[str, Any]]: class QueryHandle(SQLModel, table=True): __tablename__ = "query_handles" - # Field order defines the composite PK column order (message_id, query_ref) — keep - # message_id first to match the migration; reordering these fields changes the PK. - message_id: uuid.UUID = Field(foreign_key="message.id", primary_key=True) - - # The `str` column leaves room to shorten it to a model-reproducible - # token should exports ever key off the answer again. + # Globally-unique qr_ minted by execute_bigquery_sql query_ref: str = Field(primary_key=True) - - # The model-generated slug for the query. + message_id: uuid.UUID = Field(foreign_key="message.id", index=True) slug: str destination_table: dict[str, Any] = Field(sa_column=Column(JSON, nullable=False)) created_at: datetime = Field( diff --git a/app/exports.py b/app/exports.py index f9792ce..bb1196c 100644 --- a/app/exports.py +++ b/app/exports.py @@ -1,5 +1,7 @@ +import re from collections.abc import Iterable from dataclasses import dataclass +from datetime import datetime, timedelta, timezone from functools import cache from typing import Any, Literal, TypedDict @@ -90,9 +92,16 @@ class ExportedFile: } -# Formats offered to the frontend. Every format in EXPORT_FORMATS -# is materializable on demand, but only CSV is offered for now. -OFFERED_EXPORT_FORMATS: list[ExportFormat] = ["CSV"] +# Formats offered to the frontend, each materializable on demand via a BigQuery extract job. +OFFERED_EXPORT_FORMATS: list[ExportFormat] = ["AVRO", "CSV", "JSONL", "PARQUET"] + +# BigQuery's anonymous result tables live ~24h. +RESULT_TABLE_TTL = timedelta(hours=24) + + +def is_result_expired(created_at: datetime) -> bool: + """Whether a query handle is old enough that its result table has likely expired.""" + return datetime.now(timezone.utc) - created_at >= RESULT_TABLE_TTL @cache @@ -111,21 +120,18 @@ def _extract_table_to_gcs( ) -> int: """Extract a materialized BigQuery result table to a single GCS object. - Byte-identical to the referenced table — no query is re-run. - Args: destination_table (dict[str, Any]): `TableReference.to_api_repr()` of the result table. object_key (str): Destination object key within the export bucket. file_format (ExportFormat): The output format. Returns: - int: The size in bytes of the written object. + int: Size in bytes of the written object. Raises: - ResultTableExpired: If the result table no longer exists (~24h TTL). - ResultTooLarge: If the result set is over `MAX_EXPORT_BYTES`, or too big - for a single file despite passing that check. - RuntimeError: If the extract reports success but no object was written. + ResultTableExpired: The result table no longer exists (~24h TTL). + ResultTooLarge: The result set is over `MAX_EXPORT_BYTES` or too big for one file. + RuntimeError: The extract reported success but no object was written. """ bucket = settings.GOOGLE_GCS_BUCKET gcs_uri = f"gs://{bucket}/{object_key}" @@ -178,10 +184,7 @@ def materialize_export( ) -> ExportedFile: """Materialize a query handle as a downloadable GCS object. - Extracts the result table (see `_extract_table_to_gcs`) to a deterministic, - message-scoped key (`query_results/{message_id}/{query_ref}.{ext}`; `query_ref` is - only unique within its run). The determinism makes downloads idempotent: a repeat - reuses the existing object, and a lifecycle-deleted one is re-extracted in place. + The object key is deterministic and message-scoped, so a repeat reuses it. Args: query_ref (str): The handle whose result table is exported. @@ -194,8 +197,8 @@ def materialize_export( ExportedFile: The GCS object (bucket, key, filename, mime type, size) to sign. Raises: - ResultTableExpired: the result table no longer exists (~24h TTL). - ResultTooLarge: the result set is too big for a single file. + ResultTableExpired: The result table no longer exists (~24h TTL). + ResultTooLarge: The result set is too big for a single file. """ bucket = settings.GOOGLE_GCS_BUCKET extension = EXPORT_FORMATS[file_format].extension @@ -220,28 +223,43 @@ def materialize_export( ) +def sanitize_export_filename(slug: str, fallback: str) -> str: + """Sanitize a query's slug into a safe base filename (no extension). + + Args: + slug (str): The query's slug. + fallback (str): Base name to use when the slug yields nothing filesystem-safe. + + Returns: + str: A filesystem-safe base filename, without extension. + """ + filename = re.sub(r"[^\w-]+", "_", slug).strip("_") + return filename or fallback + + def collect_query_handles( artifacts: Iterable[JsonValue | None], - collected_handles: list[CollectedQueryHandle], -) -> None: - """Append the `query_result` handles found in a run's tool-output artifacts. +) -> list[CollectedQueryHandle]: + """Return the `query_result` handles found in a run's tool-output artifacts. - Picks out the handles minted by the `execute_bigquery_sql` tool and appends each - (its `query_ref`, display `slug`, and result table) to `collected_handles` in place. + Picks out the handles minted by the `execute_bigquery_sql` tool (each with its + `query_ref`, display `slug`, and result table). Args: artifacts (Iterable[JsonValue | None]): The run's tool-output artifacts. - collected_handles (list[CollectedQueryHandle]): The run's accumulator, appended to in place. + + Returns: + list[CollectedQueryHandle]: The query handles found, in artifact order. """ - for artifact in artifacts: - if isinstance(artifact, dict) and artifact.get("type") == "query_result": - collected_handles.append( - CollectedQueryHandle( - query_ref=artifact["query_ref"], - slug=artifact["slug"], - destination_table=artifact["destination_table"], - ) - ) + return [ + CollectedQueryHandle( + query_ref=artifact["query_ref"], + slug=artifact["slug"], + destination_table=artifact["destination_table"], + ) + for artifact in artifacts + if isinstance(artifact, dict) and artifact.get("type") == "query_result" + ] def query_result_download(query_ref: str, slug: str) -> QueryResultDownload: diff --git a/app/settings.py b/app/settings.py index 430d7fe..90c8620 100644 --- a/app/settings.py +++ b/app/settings.py @@ -27,8 +27,8 @@ class Settings(BaseSettings): AUTH_DEV_MODE: bool = Field( default=False, description=( - "When enabled, bypasses JWT validation and returns AUTH_DEV_USER_ID for all requests. " - "Only works when ENVIRONMENT is set to 'development'. " + "When enabled, bypasses JWT validation and returns AUTH_DEV_USER_ID for " + "all requests. Only works when ENVIRONMENT is set to 'development'. " "WARNING: Must NEVER be enabled in production." ), ) @@ -92,6 +92,19 @@ def SQLALCHEMY_DB_URL(self) -> str: # pragma: no cover GOOGLE_GCS_BUCKET: NonEmptyStr = Field( description="GCS bucket where exported query results are stored." ) + + @computed_field + @cached_property + def GOOGLE_CREDENTIALS(self) -> Credentials: # pragma: no cover + """Google Cloud credentials.""" + return Credentials.from_service_account_file( + filename=self.GOOGLE_SERVICE_ACCOUNT, + scopes=["https://www.googleapis.com/auth/cloud-platform"], + ) + + # ============================================================ + # == Exports & Charts settings == + # ============================================================ SIGNED_URL_TTL_SECONDS: int = Field( default=10 * 60, description="Lifetime of signed URLs generated for downloading exported query results.", @@ -106,15 +119,23 @@ def SQLALCHEMY_DB_URL(self) -> str: # pragma: no cover "file size. Must stay under BigQuery's 1 GB single-file extract limit." ), ) - - @computed_field - @cached_property - def GOOGLE_CREDENTIALS(self) -> Credentials: # pragma: no cover - """Google Cloud credentials.""" - return Credentials.from_service_account_file( - filename=self.GOOGLE_SERVICE_ACCOUNT, - scopes=["https://www.googleapis.com/auth/cloud-platform"], - ) + CHART_MAX_BYTES: int = Field( + default=1 * 1024 * 1024, + gt=0, + le=10 * 1024 * 1024, + description=( + "Largest chart payload, as the JSON size in bytes of the rows bound inline " + "into a chart spec. Bounds the SSE payload, server memory, and in-browser render." + ), + ) + GEO_ASSET_URL_BASE: str = Field( + default="/chatbot/geo", + description=( + "Client-facing base path for chart geometry. A choropleth spec references its " + "TopoJSON as `{base}/{filename}`, which the browser fetches (and caches) directly. " + "Defaults to where the website serves the files statically (public/chatbot/geo/)." + ), + ) # ============================================================ # == LLM settings == @@ -159,8 +180,9 @@ def GOOGLE_CREDENTIALS(self) -> Credentials: # pragma: no cover LOG_ENQUEUE: bool = Field( default=False, description=( - "Whether the messages to be logged should first pass through a multiprocessing-safe queue before reaching the sink. " - "This is useful while logging to a file through multiple processes and also has the advantage of making logging calls non-blocking." + "Whether the messages to be logged should first pass through a multiprocessing-safe queue " + "before reaching the sink. This is useful while logging to a file through multiple processes " + "and also has the advantage of making logging calls non-blocking." ), ) diff --git a/pyproject.toml b/pyproject.toml index c9c621a..bdf1954 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ "pydantic-settings>=2.12.0", "pyjwt>=2.10.1", "sqlmodel>=0.0.31", + "vl-convert-python>=1.9.0.post1", ] [dependency-groups] diff --git a/tests/app/agent/tools/test_bigquery.py b/tests/app/agent/tools/test_bigquery.py index a71a082..3f77e10 100644 --- a/tests/app/agent/tools/test_bigquery.py +++ b/tests/app/agent/tools/test_bigquery.py @@ -109,6 +109,8 @@ def test_successful_query(self, mocker: MockerFixture, mock_context: AgentContex assert output["rows"] == [{"col1": "value1"}, {"col1": "value2"}] assert output["row_count"] == 2 assert re.fullmatch(r"qr_[0-9a-f]{32}", message.artifact["query_ref"]) + # The handle is surfaced in the content too, so the model can reference it. + assert output["query_ref"] == message.artifact["query_ref"] def test_successful_query_exposes_destination_table_on_artifact( self, mocker: MockerFixture, mock_context: AgentContext diff --git a/tests/app/agent/tools/test_dataviz.py b/tests/app/agent/tools/test_dataviz.py new file mode 100644 index 0000000..aef2ed6 --- /dev/null +++ b/tests/app/agent/tools/test_dataviz.py @@ -0,0 +1,282 @@ +import json +import uuid +from contextlib import asynccontextmanager +from datetime import datetime, timedelta, timezone +from unittest.mock import AsyncMock, MagicMock + +from langchain.tools import ToolRuntime +from langchain_core.messages import ToolMessage + +from app.agent.context import AgentContext +from app.agent.tools import dataviz as dataviz_module +from app.agent.tools.dataviz import ( + chart_query_result, + export_query_result, + list_query_results, +) +from app.charts import ChartResultTooLarge +from app.db.models import QueryHandle +from app.exports import ExportedFile, ResultTableExpired, ResultTooLarge + + +def _exported(size_bytes: int = 2048) -> ExportedFile: + return ExportedFile( + bucket="b", + object_key="query_results/m/qr_1.parquet", + filename="resultado.parquet", + mime_type="application/vnd.apache.parquet", + size_bytes=size_bytes, + ) + + +def _runtime(thread_id: str = "test-thread") -> ToolRuntime[AgentContext]: + return ToolRuntime( + state={}, + context=AgentContext(thread_id=thread_id, user_id="test-user", language="pt"), + config={}, + stream_writer=None, + tool_call_id="test-tool-call", + store=None, + ) + + +async def _ainvoke(tool, args: dict) -> ToolMessage: + """Invoke an async tool the way the ToolNode does, returning the ToolMessage.""" + return await tool.ainvoke( + { + "type": "tool_call", + "id": "1", + "name": tool.name, + "args": {**args, "runtime": _runtime()}, + } + ) + + +def _patch_db(monkeypatch, db: MagicMock) -> None: + @asynccontextmanager + async def mock_sessionmaker(): + yield None # AsyncDatabase is mocked, so the session is never used + + monkeypatch.setattr("app.agent.tools.dataviz.sessionmaker", mock_sessionmaker) + monkeypatch.setattr("app.agent.tools.dataviz.AsyncDatabase", lambda session: db) + + +def _handle( + query_ref: str = "qr_1", + slug: str = "resultado", + age: timedelta = timedelta(0), + message_id: uuid.UUID | None = None, +) -> QueryHandle: + return QueryHandle( + query_ref=query_ref, + message_id=message_id or uuid.uuid4(), + slug=slug, + destination_table={"projectId": "p", "datasetId": "d", "tableId": "t"}, + created_at=datetime.now(timezone.utc) - age, + ) + + +class TestListQueryResults: + async def test_lists_thread_handles_with_expiry_flag(self, monkeypatch): + """Each thread handle is listed oldest-first with an expiry flag from its age.""" + fresh = _handle("qr_fresh", "recent", age=timedelta(hours=1)) + stale = _handle("qr_stale", "old", age=timedelta(hours=48)) + + db = MagicMock() + db.get_query_handles_by_thread = AsyncMock(return_value=[fresh, stale]) + _patch_db(monkeypatch, db) + + message = await _ainvoke(list_query_results, {}) + parsed = json.loads(message.content) + + assert [r["query_ref"] for r in parsed] == ["qr_fresh", "qr_stale"] + assert parsed[0]["description"] == "recent" + assert parsed[0]["expired"] is False + assert parsed[1]["expired"] is True + db.get_query_handles_by_thread.assert_awaited_once_with("test-thread") + + async def test_empty_when_no_results(self, monkeypatch): + db = MagicMock() + db.get_query_handles_by_thread = AsyncMock(return_value=[]) + _patch_db(monkeypatch, db) + + message = await _ainvoke(list_query_results, {}) + + assert json.loads(message.content) == [] + + +class TestExportQueryResult: + async def test_returns_export_affordance(self, monkeypatch): + """A valid, fresh handle yields an `export` artifact for the interface.""" + message_id = uuid.uuid4() + handle = _handle( + "qr_1", "resultado", age=timedelta(hours=1), message_id=message_id + ) + + db = MagicMock() + db.get_query_handle_from_thread = AsyncMock(return_value=handle) + _patch_db(monkeypatch, db) + materialize = MagicMock(return_value=_exported(size_bytes=2048)) + monkeypatch.setattr(dataviz_module, "materialize_export", materialize) + + message = await _ainvoke( + export_query_result, {"query_ref": "qr_1", "file_format": "parquet"} + ) + + # Case-insensitive format, thread-scoped resolution, client-facing artifact. + db.get_query_handle_from_thread.assert_awaited_once_with("qr_1", "test-thread") + # The file is materialized eagerly (so the card can show its size and filename), + # with the slug sanitized into the base filename the endpoint would also produce. + assert materialize.call_args.kwargs["file_format"] == "PARQUET" + assert materialize.call_args.kwargs["message_id"] == str(message_id) + assert materialize.call_args.kwargs["filename"] == "resultado" + assert message.artifact == { + "type": "export", + "query_ref": "qr_1", + "format": "PARQUET", + "message_id": str(message_id), + "filename": "resultado.parquet", + "size_bytes": 2048, + } + content = json.loads(message.content) + assert content["status"] == "ready" + assert content["format"] == "PARQUET" + assert content["filename"] == "resultado.parquet" + assert content["size_bytes"] == 2048 + + async def test_unsupported_format_is_rejected(self, monkeypatch): + db = MagicMock() + db.get_query_handle_from_thread = AsyncMock(return_value=_handle()) + _patch_db(monkeypatch, db) + + message = await _ainvoke( + export_query_result, {"query_ref": "qr_1", "file_format": "XLSX"} + ) + content = json.loads(message.content) + + assert content["status"] == "error" + assert "Unsupported format" in content["message"] + # The handle is never even looked up for an unsupported format. + db.get_query_handle_from_thread.assert_not_awaited() + + async def test_missing_ref_is_rejected(self, monkeypatch): + db = MagicMock() + db.get_query_handle_from_thread = AsyncMock(return_value=None) + _patch_db(monkeypatch, db) + + message = await _ainvoke( + export_query_result, {"query_ref": "qr_missing", "file_format": "CSV"} + ) + content = json.loads(message.content) + + assert content["status"] == "error" + assert "No query result found" in content["message"] + + async def test_expired_ref_is_rejected(self, monkeypatch): + handle = _handle("qr_old", "old", age=timedelta(hours=48)) + db = MagicMock() + db.get_query_handle_from_thread = AsyncMock(return_value=handle) + _patch_db(monkeypatch, db) + + message = await _ainvoke( + export_query_result, {"query_ref": "qr_old", "file_format": "CSV"} + ) + content = json.loads(message.content) + + assert content["status"] == "error" + assert "expired" in content["message"] + + async def test_too_large_result_is_reported(self, monkeypatch): + """An over-limit result fails at tool time, not as a click that 400s later.""" + db = MagicMock() + db.get_query_handle_from_thread = AsyncMock(return_value=_handle()) + _patch_db(monkeypatch, db) + monkeypatch.setattr( + dataviz_module, + "materialize_export", + MagicMock(side_effect=ResultTooLarge("over the export limit")), + ) + + message = await _ainvoke( + export_query_result, {"query_ref": "qr_1", "file_format": "CSV"} + ) + content = json.loads(message.content) + + assert content["status"] == "error" + assert "too large" in content["message"] + + async def test_expired_during_materialize_is_reported(self, monkeypatch): + """The table can vanish between the age check and the extract; report it cleanly.""" + db = MagicMock() + db.get_query_handle_from_thread = AsyncMock(return_value=_handle()) + _patch_db(monkeypatch, db) + monkeypatch.setattr( + dataviz_module, + "materialize_export", + MagicMock(side_effect=ResultTableExpired("gone")), + ) + + message = await _ainvoke( + export_query_result, {"query_ref": "qr_1", "file_format": "CSV"} + ) + content = json.loads(message.content) + + assert content["status"] == "error" + assert "expired" in content["message"] + + +class TestChartQueryResult: + async def test_renders_chart_artifact(self, monkeypatch): + """The tool describes the chart; the generated spec is bound to the exact rows.""" + handle = _handle("qr_1", "vendas") + rows = [{"ano": 2025, "total": 10}] + monkeypatch.setattr( + dataviz_module, + "fetch_chart_data", + AsyncMock(return_value=(handle, ["ano", "total"], rows)), + ) + # The spec generator (validate-and-repair loop) is exercised in test_charts; + # here it stands in for the returned, already-validated spec. + monkeypatch.setattr( + dataviz_module, + "generate_chart_spec", + AsyncMock( + return_value={"mark": "bar", "encoding": {"x": {"field": "ano"}}} + ), + ) + + message = await _ainvoke( + chart_query_result, + {"query_ref": "qr_1", "instructions": "a bar chart of total by year"}, + ) + + assert message.artifact["type"] == "chart" + assert message.artifact["query_ref"] == "qr_1" + assert message.artifact["spec"]["mark"] == "bar" + # The card renders the spec; the slug is not part of the artifact. + assert "slug" not in message.artifact + # The server binds the exact rows; the model never supplied them. + assert message.artifact["spec"]["data"] == {"values": rows} + content = json.loads(message.content) + assert content["status"] == "rendered" + assert content["row_count"] == 1 + # generate_chart_spec receives the columns/rows/instructions, not a spec. + dataviz_module.generate_chart_spec.assert_awaited_once_with( + ["ano", "total"], rows, "a bar chart of total by year" + ) + + async def test_too_large_result_is_reported(self, monkeypatch): + monkeypatch.setattr( + dataviz_module, + "fetch_chart_data", + AsyncMock(side_effect=ChartResultTooLarge("2500 rows, over the limit")), + ) + + message = await _ainvoke( + chart_query_result, + {"query_ref": "qr_big", "instructions": "a bar chart"}, + ) + content = json.loads(message.content) + + assert content["status"] == "error" + assert "over the limit" in content["message"] diff --git a/tests/app/agent/tools/test_toolkit.py b/tests/app/agent/tools/test_toolkit.py index c416306..19c2990 100644 --- a/tests/app/agent/tools/test_toolkit.py +++ b/tests/app/agent/tools/test_toolkit.py @@ -8,7 +8,7 @@ def test_get_tools_returns_all_tools(self): """Test that get_tools returns all expected tools.""" tools = BDToolkit.get_tools() - assert len(tools) == 5 + assert len(tools) == 8 tool_names = [tool.name for tool in tools] @@ -17,3 +17,6 @@ def test_get_tools_returns_all_tools(self): assert "get_table_details" in tool_names assert "execute_bigquery_sql" in tool_names assert "decode_table_values" in tool_names + assert "list_query_results" in tool_names + assert "export_query_result" in tool_names + assert "chart_query_result" in tool_names diff --git a/tests/app/api/routers/test_chatbot.py b/tests/app/api/routers/test_chatbot.py index 3cee458..eb29be0 100644 --- a/tests/app/api/routers/test_chatbot.py +++ b/tests/app/api/routers/test_chatbot.py @@ -12,7 +12,6 @@ from pytest_mock import MockerFixture from app.api.dependencies import get_database, get_feedback_sender -from app.api.routers.chatbot import _sanitize_filename from app.api.streaming.schemas import StreamEvent from app.db.database import AsyncDatabase from app.db.models import ( @@ -400,7 +399,7 @@ async def test_list_messages_derives_downloads_from_query_handles( assert download["type"] == "query_result" assert download["query_ref"] == "qr_test" assert download["slug"] == "slug" - assert download["formats"] == ["CSV"] + assert download["formats"] == ["AVRO", "CSV", "JSONL", "PARQUET"] # The internal handles (and their destination tables) never reach the client. assert "query_handles" not in message_json @@ -836,7 +835,7 @@ def test_unsupported_format_is_rejected( downloadable_message: Message, mocker: MockerFixture, ): - """A not-offered format (valid for BigQuery, but not offered) is rejected, not downgraded.""" + """An unsupported format is rejected outright, not downgraded to a default.""" materialize = mocker.patch( "app.api.routers.chatbot.materialize_export", return_value=self._exported(), @@ -844,11 +843,12 @@ def test_unsupported_format_is_rejected( response = client.post( url=f"/api/v1/chatbot/messages/{downloadable_message.id}/exports" - "?query_ref=qr_test&format=PARQUET", + "?query_ref=qr_test&format=XLSX", headers={"Authorization": f"Bearer {access_token}"}, ) assert response.status_code == status.HTTP_400_BAD_REQUEST + assert "Unsupported format" in response.json()["detail"] # The request is rejected outright — no CSV is silently produced. materialize.assert_not_called() @@ -1023,38 +1023,3 @@ def test_export_unauthorized( ) assert response.status_code == status.HTTP_401_UNAUTHORIZED - - -class TestSanitizeFilename: - """Tests for _sanitize_filename — the download filename guard.""" - - @pytest.mark.parametrize( - ("slug", "expected"), - [ - # A clean slug (what the model is asked to produce) passes through unchanged. - ("vendas_por_ano", "vendas_por_ano"), - # Hyphens and digits are allowed. - ("ideb-2021", "ideb-2021"), - # Spaces and punctuation collapse to a single underscore. - ("Vendas por ano", "Vendas_por_ano"), - ("a b", "a_b"), - ("café & leite!", "café_leite"), - # Leading/trailing separators are stripped, not left dangling. - ("_vendas_", "vendas"), - (" vendas ", "vendas"), - # Path separators and traversal are neutralized (no slashes or dots survive). - ("../../etc/passwd", "etc_passwd"), - ("relatorio/2021", "relatorio_2021"), - # File extensions are neutralized. - ("vendas_por_ano.csv", "vendas_por_ano_csv"), - # Accented word characters are preserved (\\w is unicode). - ("população", "população"), - ], - ) - def test_sanitizes_slug(self, slug: str, expected: str): - assert _sanitize_filename(slug, "resultados") == expected - - @pytest.mark.parametrize("slug", ["", " ", "!!!", "/", "..."]) - def test_falls_back_when_nothing_usable_remains(self, slug: str): - """A slug that sanitizes to empty falls back to the provided fallback, never ''.""" - assert _sanitize_filename(slug, "resultados") == "resultados" diff --git a/tests/app/api/streaming/test_agent_runner.py b/tests/app/api/streaming/test_agent_runner.py index dcf1e84..2bc8c15 100644 --- a/tests/app/api/streaming/test_agent_runner.py +++ b/tests/app/api/streaming/test_agent_runner.py @@ -61,17 +61,19 @@ def mock_database( of opening a real DB connection.""" db = MagicMock() - db.create_message = AsyncMock( - return_value=Message( - id=config["run_id"], - thread_id=mock_user_message.thread_id, - user_message_id=mock_user_message.id, - model_uri=mock_user_message.model_uri, - role=MessageRole.ASSISTANT, - content="Mock assistant message", - status=MessageStatus.SUCCESS, - ) + assistant_message = Message( + id=config["run_id"], + thread_id=mock_user_message.thread_id, + user_message_id=mock_user_message.id, + model_uri=mock_user_message.model_uri, + role=MessageRole.ASSISTANT, + content="Mock assistant message", + status=MessageStatus.SUCCESS, ) + # create_message mints the up-front STREAMING placeholder; update_message + # writes the terminal state at run end (the only finalize path). + db.create_message = AsyncMock(return_value=assistant_message) + db.update_message = AsyncMock(return_value=assistant_message) db.create_query_handles = AsyncMock(return_value=None) @asynccontextmanager @@ -615,9 +617,10 @@ async def astream(*args, **kwargs): queue=queue, ) - message = mock_database.create_message.call_args[0][0] - assert message.status == MessageStatus.ERROR - assert message.content == translate(MessageKey.ERROR_UNEXPECTED, "es") + mock_database.update_message.assert_awaited_once() + kwargs = mock_database.update_message.call_args.kwargs + assert kwargs["status"] == MessageStatus.ERROR + assert kwargs["content"] == translate(MessageKey.ERROR_UNEXPECTED, "es") async def test_emits_events_and_persists_success( self, @@ -655,11 +658,16 @@ async def astream(*args, **kwargs): assert events[0].data.content == "Final answer" assert events[-1].data.run_id == config["run_id"] + # Created up front as a STREAMING placeholder, then finalized as SUCCESS. 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" + placeholder = mock_database.create_message.call_args[0][0] + assert isinstance(placeholder, MessageCreate) + assert placeholder.status == MessageStatus.STREAMING + assert placeholder.content == "" + mock_database.update_message.assert_awaited_once() + kwargs = mock_database.update_message.call_args.kwargs + assert kwargs["status"] == MessageStatus.SUCCESS + assert kwargs["content"] == "Final answer" async def test_handle_persist_failure_still_persists_message( self, @@ -721,7 +729,8 @@ async def astream(*args, **kwargs): # The handle write was attempted and failed, but the message still persisted # and the run completes without a persistence error. - mock_database.create_message.assert_called_once() + mock_database.create_message.assert_called_once() # placeholder + mock_database.update_message.assert_awaited_once() # terminal state mock_database.create_query_handles.assert_awaited_once() assert complete.type == "complete" assert complete.data.run_id == config["run_id"] @@ -800,12 +809,17 @@ async def astream(*args, **kwargs): ] assert events[0].data.structured_response["follow_up_prompts"] == ["E em 2026?"] + # Created up front as a STREAMING placeholder, then finalized as SUCCESS. 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 + placeholder = mock_database.create_message.call_args[0][0] + assert isinstance(placeholder, MessageCreate) + assert placeholder.status == MessageStatus.STREAMING + assert placeholder.content == "" + mock_database.update_message.assert_awaited_once() + kwargs = mock_database.update_message.call_args.kwargs + assert kwargs["status"] == MessageStatus.SUCCESS + assert kwargs["content"] == "Final answer" + assert kwargs["structured_response"] == events[0].data.structured_response async def test_executed_query_derives_download_and_persists_handle( self, @@ -980,11 +994,10 @@ async def astream(*args, **kwargs): assert events[0].data.content == translate(MessageKey.ERROR_UNEXPECTED, "pt") assert events[0].data.error_details == {"reason": "agent_failed"} - mock_database.create_message.assert_called_once() - message = mock_database.create_message.call_args[0][0] - assert isinstance(message, MessageCreate) - assert message.status == MessageStatus.ERROR - assert message.content == translate(MessageKey.ERROR_UNEXPECTED, "pt") + mock_database.update_message.assert_awaited_once() + kwargs = mock_database.update_message.call_args.kwargs + assert kwargs["status"] == MessageStatus.ERROR + assert kwargs["content"] == translate(MessageKey.ERROR_UNEXPECTED, "pt") async def test_model_call_limit_persists_with_dedicated_status( self, @@ -1024,21 +1037,21 @@ async def astream(*args, **kwargs): ) assert events[-1].data.run_id == config["run_id"] - mock_database.create_message.assert_called_once() - message = mock_database.create_message.call_args[0][0] - assert isinstance(message, MessageCreate) - assert message.status == MessageStatus.MODEL_CALL_LIMIT - assert message.content == translate(MessageKey.ERROR_MODEL_CALL_LIMIT, "pt") + mock_database.update_message.assert_awaited_once() + kwargs = mock_database.update_message.call_args.kwargs + assert kwargs["status"] == MessageStatus.MODEL_CALL_LIMIT + assert kwargs["content"] == translate(MessageKey.ERROR_MODEL_CALL_LIMIT, "pt") - async def test_complete_still_emitted_when_db_write_fails( + async def test_aborts_run_when_placeholder_cannot_be_written( self, monkeypatch: pytest.MonkeyPatch, mock_user_message: Message, config: ConfigDict, thread_id: str, ): - """If `database.create_message` raises, the consumer must still - receive a `complete` event - otherwise it hangs on `queue.get()`. + """If the placeholder can't be written the DB is unhealthy, so the run aborts + before the model runs — yet still emits error + complete so the consumer does + not hang on `queue.get()`. Persistence stays deterministic: no run without a row. """ db = MagicMock() db.create_message = AsyncMock(side_effect=RuntimeError("db down")) @@ -1055,15 +1068,7 @@ async def mock_sessionmaker(): "app.api.streaming.agent_runner.AsyncDatabase", lambda session: db ) - agent = MagicMock() - - async def astream(*args, **kwargs): - yield ( - "updates", - {"model": {"messages": [AIMessage(content="Final answer")]}}, - ) - - agent.astream = astream + agent = MagicMock() # its astream must never be reached queue: asyncio.Queue[StreamEvent] = asyncio.Queue() await run_agent( @@ -1079,14 +1084,16 @@ async def astream(*args, **kwargs): ) events = await self._drain(queue) - assert [e.type for e in events] == ["final_answer", "complete"] - assert events[0].data.content == "Final answer" + # The model was never invoked; the run ended with error + complete only. + agent.astream.assert_not_called() + assert [e.type for e in events] == ["error", "complete"] + assert events[0].data.content == translate(MessageKey.ERROR_UNEXPECTED, "pt") complete = events[-1] - assert complete.type == "complete" assert complete.data.run_id is None assert complete.data.error_details == {"reason": "persistence_failed"} - assert "db down" not in str(complete.data.error_details) + # No terminal update is attempted — there is no placeholder to update. + db.update_message.assert_not_called() async def test_consumer_cancel_does_not_cancel_producer( self, @@ -1125,11 +1132,16 @@ async def astream(*args, **kwargs): await asyncio.wait_for(task, timeout=2.0) # Producer persisted the message regardless of consumer presence + # Created up front as a STREAMING placeholder, then finalized as SUCCESS. 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" + placeholder = mock_database.create_message.call_args[0][0] + assert isinstance(placeholder, MessageCreate) + assert placeholder.status == MessageStatus.STREAMING + assert placeholder.content == "" + mock_database.update_message.assert_awaited_once() + kwargs = mock_database.update_message.call_args.kwargs + assert kwargs["status"] == MessageStatus.SUCCESS + assert kwargs["content"] == "Final answer" # The complete event is sitting in the queue waiting events = await self._drain(queue) @@ -1180,11 +1192,10 @@ async def astream(*args, **kwargs): assert task.cancelled() - mock_database.create_message.assert_called_once() - message = mock_database.create_message.call_args[0][0] - assert isinstance(message, MessageCreate) - assert message.status == MessageStatus.INTERRUPTED - assert message.content == translate(MessageKey.ERROR_INTERRUPTED, "pt") + mock_database.update_message.assert_awaited_once() + kwargs = mock_database.update_message.call_args.kwargs + assert kwargs["status"] == MessageStatus.INTERRUPTED + assert kwargs["content"] == translate(MessageKey.ERROR_INTERRUPTED, "pt") async def test_cancellation_after_final_answer_preserves_success( self, @@ -1236,8 +1247,134 @@ async def astream(*args, **kwargs): assert task.cancelled() + # Created up front as a STREAMING placeholder, then finalized as SUCCESS. 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" + placeholder = mock_database.create_message.call_args[0][0] + assert isinstance(placeholder, MessageCreate) + assert placeholder.status == MessageStatus.STREAMING + assert placeholder.content == "" + mock_database.update_message.assert_awaited_once() + kwargs = mock_database.update_message.call_args.kwargs + assert kwargs["status"] == MessageStatus.SUCCESS + assert kwargs["content"] == "Final answer" + + async def test_no_handles_persisted_for_non_query_tool_output( + self, + mock_database: MagicMock, + mock_user_message: Message, + config: ConfigDict, + thread_id: str, + ): + """A tool output with no query_result artifact persists no handles.""" + agent = MagicMock() + + async def astream(*args, **kwargs): + yield ( + "updates", + { + "tools": { + "messages": [ + ToolMessage( + content='{"ok": true}', + tool_call_id="1", + name="search_datasets", + status="success", + ) + ] + } + }, + ) + yield ("updates", {"model": {"messages": [AIMessage(content="done")]}}) + + 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, + context=AgentContext( + thread_id="test-thread", user_id="test-user", language="pt" + ), + queue=queue, + ) + + events = await self._drain(queue) + assert any(e.type == "tool_output" for e in events) + mock_database.create_query_handles.assert_not_awaited() + + async def test_finalize_reports_failure_when_placeholder_missing( + self, + mock_database: MagicMock, + mock_user_message: Message, + config: ConfigDict, + thread_id: str, + ): + """If the placeholder vanished mid-run (update finds nothing), report the failure.""" + mock_database.update_message = AsyncMock(return_value=None) + agent = MagicMock() + + async def astream(*args, **kwargs): + yield ( + "updates", + {"model": {"messages": [AIMessage(content="Final answer")]}}, + ) + + 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, + context=AgentContext( + thread_id="test-thread", user_id="test-user", language="pt" + ), + queue=queue, + ) + + complete = (await self._drain(queue))[-1] + assert complete.type == "complete" + assert complete.data.run_id is None + assert complete.data.error_details == {"reason": "persistence_failed"} + + async def test_finalize_reports_failure_when_update_raises( + self, + mock_database: MagicMock, + mock_user_message: Message, + config: ConfigDict, + thread_id: str, + ): + """If the terminal update raises, the run still completes with a persistence error.""" + mock_database.update_message = AsyncMock(side_effect=RuntimeError("db down")) + agent = MagicMock() + + async def astream(*args, **kwargs): + yield ( + "updates", + {"model": {"messages": [AIMessage(content="Final answer")]}}, + ) + + 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, + context=AgentContext( + thread_id="test-thread", user_id="test-user", language="pt" + ), + queue=queue, + ) + + complete = (await self._drain(queue))[-1] + assert complete.type == "complete" + assert complete.data.run_id is None + assert complete.data.error_details == {"reason": "persistence_failed"} diff --git a/tests/app/db/test_database.py b/tests/app/db/test_database.py index f40e108..10e2d95 100644 --- a/tests/app/db/test_database.py +++ b/tests/app/db/test_database.py @@ -12,6 +12,8 @@ FeedbackSyncStatus, Message, MessageCreate, + MessageRole, + MessageStatus, QueryHandle, Thread, ThreadCreate, @@ -220,6 +222,70 @@ async def test_get_messages_not_found( assert isinstance(messages, list) assert len(messages) == 0 + async def test_get_messages_excludes_streaming_placeholders( + self, database: AsyncDatabase, messages_factory: MessagesFactory + ): + """In-flight STREAMING placeholders are hidden from the thread listing.""" + user_message, _ = await messages_factory() + + placeholder = MessageCreate( + thread_id=user_message.thread_id, + user_message_id=user_message.id, + model_uri="mock-model", + role=MessageRole.ASSISTANT, + content="", + status=MessageStatus.STREAMING, + ) + await database.create_message(placeholder) + + messages = await database.get_messages(user_message.thread_id) + + # Only the two completed messages surface; the placeholder is excluded. + assert len(messages) == 2 + assert all(message.status != MessageStatus.STREAMING for message in messages) + + async def test_update_message_writes_terminal_fields( + self, database: AsyncDatabase, user_message: Message + ): + """update_message overwrites the terminal fields of an existing row.""" + placeholder = MessageCreate( + id=uuid.uuid4(), + thread_id=user_message.thread_id, + user_message_id=user_message.id, + model_uri="mock-model", + role=MessageRole.ASSISTANT, + content="", + status=MessageStatus.STREAMING, + ) + created = await database.create_message(placeholder) + + updated = await database.update_message( + created.id, + content="Final answer", + events=[{"type": "final_answer"}], + structured_response={"response": "Final answer"}, + status=MessageStatus.SUCCESS, + ) + + assert updated is not None + assert updated.id == created.id + assert updated.content == "Final answer" + assert updated.status == MessageStatus.SUCCESS + assert updated.events == [{"type": "final_answer"}] + assert updated.structured_response == {"response": "Final answer"} + + async def test_update_message_not_found_returns_none(self, database: AsyncDatabase): + """Updating a non-existent message returns None (finalize then falls back).""" + result = await database.update_message( + uuid.uuid4(), + content="x", + events=None, + structured_response=None, + status=MessageStatus.SUCCESS, + ) + + assert result is None + class TestAsyncDatabaseQueryHandle: """Tests for QueryHandle operations.""" @@ -247,8 +313,12 @@ async def test_create_query_handles_persists( ] ) - first = await database.get_query_handle(assistant_message.id, "qr_a") - second = await database.get_query_handle(assistant_message.id, "qr_b") + first = await database.get_query_handle_from_message( + "qr_a", assistant_message.id + ) + second = await database.get_query_handle_from_message( + "qr_b", assistant_message.id + ) assert first is not None and first.destination_table == self.DESTINATION assert second is not None and second.message_id == assistant_message.id @@ -257,20 +327,68 @@ async def test_create_query_handles_empty_is_noop(self, database: AsyncDatabase) """An empty list persists nothing and does not error.""" await database.create_query_handles([]) - async def test_get_query_handle_found( + async def test_get_query_handle_from_message_found( self, database: AsyncDatabase, query_handle: QueryHandle ): - """An existing handle is returned by its (message_id, query_ref).""" - found = await database.get_query_handle( - query_handle.message_id, query_handle.query_ref + """An existing handle is returned by its (query_ref, message_id).""" + found = await database.get_query_handle_from_message( + query_handle.query_ref, query_handle.message_id ) assert found is not None assert found.query_ref == query_handle.query_ref - async def test_get_query_handle_not_found(self, database: AsyncDatabase): - """A missing (message_id, query_ref) returns None.""" - assert await database.get_query_handle(uuid.uuid4(), "qr_missing") is None + async def test_get_query_handle_from_message_not_found( + self, database: AsyncDatabase + ): + """A missing (query_ref, message_id) returns None.""" + assert ( + await database.get_query_handle_from_message("qr_missing", uuid.uuid4()) + is None + ) + + async def test_get_query_handle_from_thread_found( + self, + database: AsyncDatabase, + query_handle: QueryHandle, + assistant_message: Message, + ): + """A handle resolves by query_ref alone when scoped to its own thread.""" + found = await database.get_query_handle_from_thread( + query_handle.query_ref, assistant_message.thread_id + ) + + assert found is not None + assert found.query_ref == query_handle.query_ref + + async def test_get_query_handle_from_thread_wrong_thread_returns_none( + self, database: AsyncDatabase, query_handle: QueryHandle + ): + """A handle is invisible to a thread it does not belong to (authorization).""" + found = await database.get_query_handle_from_thread( + query_handle.query_ref, uuid.uuid4() + ) + + assert found is None + + async def test_get_query_handles_by_thread( + self, + database: AsyncDatabase, + query_handle: QueryHandle, + assistant_message: Message, + ): + """Every handle produced in a thread is listed.""" + handles = await database.get_query_handles_by_thread( + assistant_message.thread_id + ) + + assert [handle.query_ref for handle in handles] == [query_handle.query_ref] + + async def test_get_query_handles_by_thread_empty( + self, database: AsyncDatabase, thread: Thread + ): + """A thread with no query results returns an empty list.""" + assert await database.get_query_handles_by_thread(thread.id) == [] class TestAsyncDatabaseFeedback: diff --git a/tests/app/test_charts.py b/tests/app/test_charts.py new file mode 100644 index 0000000..58eacc9 --- /dev/null +++ b/tests/app/test_charts.py @@ -0,0 +1,533 @@ +import json +import uuid +from contextlib import asynccontextmanager +from datetime import date, datetime, timedelta, timezone +from decimal import Decimal +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest +from google.api_core.exceptions import NotFound +from langchain_core.messages import SystemMessage + +from app import charts +from app.charts import ( + _CHART_SPEC_INSTRUCTIONS, + MAX_CHART_SPEC_ATTEMPTS, + VEGA_LITE_SCHEMA, + ChartHandleNotFound, + ChartResultTooLarge, + ChartSpec, + ChartSpecInvalid, + _chart_spec_user_prompt, + _collect, + _fetch_rows, + _geo_stub_node, + _geo_url_node, + _resolve_named_data, + _sanitize_chart_spec, + _validate_chart_spec, + fetch_chart_data, + generate_chart_spec, + inject_chart_data, +) +from app.db.models import QueryHandle +from app.exports import ResultTableExpired +from app.settings import settings + +DESTINATION = {"projectId": "p", "datasetId": "d", "tableId": "t"} + + +def _handle(query_ref="qr_1", slug="resultado", age=timedelta(0)): + return QueryHandle( + query_ref=query_ref, + message_id=uuid.uuid4(), + slug=slug, + destination_table=DESTINATION, + created_at=datetime.now(timezone.utc) - age, + ) + + +def _choropleth_spec() -> dict: + """A states choropleth as the model would emit it (named sources, lookup join).""" + return { + "data": {"name": "brazil_states"}, + "transform": [ + { + "lookup": "id", + "from": { + "data": {"name": "query_result"}, + "key": "sigla_uf", + "fields": ["valor"], + }, + } + ], + "projection": {"type": "mercator"}, + "mark": "geoshape", + "encoding": { + "color": {"field": "valor", "type": "quantitative"}, + "tooltip": [{"field": "properties.name", "type": "nominal"}], + }, + } + + +class TestSanitizeChartSpec: + def test_strips_data_datasets_and_urls_recursively(self): + raw = { + "data": {"url": "https://evil.example/top.json"}, + "datasets": {"foo": [1, 2, 3]}, + "layer": [ + { + "mark": "line", + "encoding": {"x": {"field": "ano"}}, + "data": {"url": "https://evil.example/layer.json"}, + } + ], + } + + clean = _sanitize_chart_spec(raw) + + assert "data" not in clean + assert "datasets" not in clean + assert "data" not in clean["layer"][0] + # No URL survives anywhere (external-fetch / SSRF vector). + assert "url" not in json.dumps(clean) + # The presentational parts are untouched. + assert clean["layer"][0]["encoding"] == {"x": {"field": "ano"}} + + def test_keeps_allowlisted_named_sources(self): + clean = _sanitize_chart_spec(_choropleth_spec()) + + assert clean["data"] == {"name": "brazil_states"} + assert clean["transform"][0]["from"]["data"] == {"name": "query_result"} + + def test_drops_unknown_named_source_inline_values_and_url(self): + assert "data" not in _sanitize_chart_spec({"data": {"name": "secret"}}) + assert "data" not in _sanitize_chart_spec({"data": {"values": [{"x": 1}]}}) + assert "data" not in _sanitize_chart_spec({"data": {"url": "http://x"}}) + + def test_strips_href_click_link_but_keeps_a_field_named_href(self): + raw = { + "mark": {"type": "point", "href": "javascript:alert(1)"}, + "encoding": { + "href": {"field": "link"}, # the href channel is dropped + "x": {"field": "href"}, # a column named "href" is a value, not a key + }, + } + + clean = _sanitize_chart_spec(raw) + + assert "href" not in clean["mark"] + assert "href" not in clean["encoding"] + assert clean["encoding"]["x"] == {"field": "href"} + # No clickable-link href survives anywhere (click-XSS vector). + assert "javascript:" not in json.dumps(clean) + + +class TestFetchRows: + def _client(self, rows, columns=("col1",)): + client = MagicMock() + client.get_table.return_value = SimpleNamespace( + schema=[SimpleNamespace(name=name) for name in columns], + ) + client.list_rows.return_value = iter(rows) + return client + + def test_reads_columns_and_rows(self, mocker): + client = self._client([{"col1": "a"}, {"col1": "b"}]) + mocker.patch("app.charts._bq_client", return_value=client) + + columns, rows = _fetch_rows(DESTINATION) + + assert columns == ["col1"] + assert rows == [{"col1": "a"}, {"col1": "b"}] + + def test_allows_a_dense_result_within_budget(self, mocker): + """A high-cardinality result (e.g. a municipal choropleth ~5.5k rows) is allowed.""" + rows = [{"id": i} for i in range(6000)] + client = self._client(rows, columns=("id",)) + mocker.patch("app.charts._bq_client", return_value=client) + + _, got = _fetch_rows(DESTINATION) + + assert len(got) == 6000 + + def test_raises_when_over_byte_budget(self, mocker, monkeypatch): + """The bound data is rejected once it would exceed the payload budget.""" + monkeypatch.setattr( + charts, + "settings", + charts.settings.model_copy(update={"CHART_MAX_BYTES": 32}), + ) + client = self._client([{"col1": "x" * 100}, {"col1": "y" * 100}]) + mocker.patch("app.charts._bq_client", return_value=client) + + with pytest.raises(ChartResultTooLarge): + _fetch_rows(DESTINATION) + + def test_coerces_non_json_types_to_json_native(self, mocker): + """BigQuery date/datetime/Decimal values become JSON-serializable rows. + + A Decimal (NUMERIC/BIGNUMERIC) must become a float, not a string — a quantitative + axis or colour scale binds these values, and a string breaks the numeric encoding. + """ + rows = [ + {"data": date(2025, 12, 31), "temperatura_media": Decimal("24.79")}, + ] + client = self._client(rows, columns=("data", "temperatura_media")) + mocker.patch("app.charts._bq_client", return_value=client) + + _, got = _fetch_rows(DESTINATION) + + assert got == [{"data": "2025-12-31", "temperatura_media": 24.79}] + assert isinstance(got[0]["temperatura_media"], float) + # The bound rows must be plain JSON values, or spec serialization fails. + json.dumps(got) + + def test_missing_table_maps_to_expired(self, mocker): + client = MagicMock() + client.get_table.side_effect = NotFound("gone") + mocker.patch("app.charts._bq_client", return_value=client) + + with pytest.raises(ResultTableExpired): + _fetch_rows(DESTINATION) + + +class TestGeoUrlNode: + def test_points_at_the_static_geo_url(self): + node = _geo_url_node("brazil_states") + + assert node == { + "url": f"{settings.GEO_ASSET_URL_BASE}/brazil_states.topojson", + "format": {"type": "topojson", "feature": "uf"}, + } + + def test_carries_no_inline_geometry(self): + # The whole point: geometry travels by URL, never inline in the spec. + node = _geo_url_node("brazil_municipalities") + + assert "values" not in node + assert node["url"].endswith("/brazil_municipalities.topojson") + + +class TestGeoStubNode: + def test_is_a_minimal_single_feature_topojson(self): + node = _geo_stub_node("brazil_states") + geometries = node["values"]["objects"]["uf"]["geometries"] + + assert node["format"] == {"type": "topojson", "feature": "uf"} + assert node["values"]["type"] == "Topology" + assert len(geometries) == 1 + + def test_keys_the_stub_by_the_assets_feature(self): + node = _geo_stub_node("brazil_municipalities") + + assert node["format"]["feature"] == "Munic" + assert set(node["values"]["objects"]) == {"Munic"} + + +class TestResolveNamedData: + def test_resolves_query_result_to_rows(self): + rows = [{"sigla_uf": "SP", "valor": 1}] + + assert _resolve_named_data({"name": "query_result"}, rows, _geo_url_node) == { + "values": rows + } + + def test_resolves_a_geo_name_via_the_given_resolver(self): + node = _resolve_named_data({"name": "brazil_states"}, [], _geo_url_node) + + assert node == _geo_url_node("brazil_states") + + def test_resolves_references_nested_anywhere(self): + rows = [{"sigla_uf": "SP", "valor": 1}] + spec = {"transform": [{"from": {"data": {"name": "query_result"}}}]} + + resolved = _resolve_named_data(spec, rows, _geo_url_node) + + assert resolved["transform"][0]["from"]["data"] == {"values": rows} + + def test_leaves_unknown_names_and_plain_nodes_untouched(self): + # An unknown name is not resolved here (sanitize drops it earlier). + assert _resolve_named_data({"name": "secret"}, [], _geo_url_node) == { + "name": "secret" + } + # A node with no named source passes through unchanged. + assert _resolve_named_data({"mark": "bar"}, [], _geo_url_node) == { + "mark": "bar" + } + + +class TestInjectChartData: + def test_binds_rows_for_a_plain_chart(self): + # inject_chart_data trusts an already-sanitized spec (see generate_chart_spec); + # it does not strip — it binds the rows a plain chart declared no data for. + spec = {"mark": "bar", "encoding": {"x": {"field": "ano"}}} + rows = [{"ano": 2025, "total": 10}] + + chart = inject_chart_data(spec, rows) + + assert chart["$schema"] == VEGA_LITE_SCHEMA + assert chart["data"] == {"values": rows} + assert chart["mark"] == "bar" + + def test_resolves_named_sources_in_a_choropleth(self): + rows = [{"sigla_uf": "SP", "valor": 1}] + + chart = inject_chart_data(_choropleth_spec(), rows) + + # Top-level geometry becomes a URL node (fetched client-side); the rows fill the + # lookup source inline. + assert chart["data"] == _geo_url_node("brazil_states") + assert chart["transform"][0]["from"]["data"] == {"values": rows} + + +class TestFetchChartData: + def _patch_db(self, monkeypatch, handle): + db = MagicMock() + db.get_query_handle_from_thread = AsyncMock(return_value=handle) + + @asynccontextmanager + async def mock_sessionmaker(): + yield None + + monkeypatch.setattr(charts, "sessionmaker", mock_sessionmaker) + monkeypatch.setattr(charts, "AsyncDatabase", lambda session: db) + return db + + async def test_returns_handle_columns_and_rows(self, monkeypatch): + handle = _handle(age=timedelta(hours=1)) + self._patch_db(monkeypatch, handle) + monkeypatch.setattr( + charts, "_fetch_rows", lambda dest: (["ano"], [{"ano": 2025}]) + ) + + got_handle, columns, rows = await fetch_chart_data("qr_1", "test-thread") + + assert got_handle is handle + assert columns == ["ano"] + assert rows == [{"ano": 2025}] + + async def test_missing_handle_raises(self, monkeypatch): + self._patch_db(monkeypatch, None) + + with pytest.raises(ChartHandleNotFound): + await fetch_chart_data("qr_missing", "test-thread") + + async def test_expired_handle_raises(self, monkeypatch): + self._patch_db(monkeypatch, _handle(age=timedelta(hours=48))) + + with pytest.raises(ResultTableExpired): + await fetch_chart_data("qr_old", "test-thread") + + +class TestCollect: + def test_collects_string_values_under_key_recursively(self): + node = { + "encoding": {"x": {"field": "ano"}, "y": {"field": "total"}}, + "layer": [{"encoding": {"color": {"field": "uf"}}}], + } + + assert _collect(node, "field") == {"ano", "total", "uf"} + + def test_ignores_non_string_values_under_the_key(self): + assert _collect({"field": 5, "x": {"field": "ano"}}, "field") == {"ano"} + + +class TestValidateChartSpec: + """Uses the real vl-convert compiler (in-process).""" + + def test_valid_spec_has_no_errors(self): + spec = { + "mark": "bar", + "encoding": {"x": {"field": "ano"}, "y": {"field": "total"}}, + } + assert _validate_chart_spec(spec, ["ano", "total"]) == [] + + def test_layered_dual_axis_combo_is_valid(self): + """A layer + resolve combo — the case the old allowlist could not express.""" + spec = { + "layer": [ + { + "mark": "line", + "encoding": {"y": {"field": "media", "type": "quantitative"}}, + }, + { + "mark": "bar", + "encoding": {"y": {"field": "var", "type": "quantitative"}}, + }, + ], + "encoding": {"x": {"field": "ano", "type": "ordinal"}}, + "resolve": {"scale": {"y": "independent"}}, + } + assert _validate_chart_spec(spec, ["ano", "media", "var"]) == [] + + def test_missing_column_flagged_across_layers(self): + """A field not in the result, even nested in a layer — else it renders empty.""" + spec = {"layer": [{"mark": "bar", "encoding": {"x": {"field": "vendas"}}}]} + + errors = _validate_chart_spec(spec, ["ano", "total"]) + + assert any("vendas" in error for error in errors) + + def test_transform_derived_field_is_allowed(self): + """A field created by a transform's `as` is not flagged as missing.""" + spec = { + "transform": [{"calculate": "datum.total * 2", "as": "dobro"}], + "mark": "bar", + "encoding": { + "x": {"field": "ano", "type": "ordinal"}, + "y": {"field": "dobro", "type": "quantitative"}, + }, + } + assert _validate_chart_spec(spec, ["ano", "total"]) == [] + + def test_invalid_mark_fails_to_compile(self): + """A structurally invalid spec is caught by the vl-convert compile step.""" + spec = {"mark": "notamark", "encoding": {"x": {"field": "ano"}}} + + errors = _validate_chart_spec(spec, ["ano"]) + + assert any("compile" in error for error in errors) + + def test_unknown_color_scheme_fails_to_compile(self): + """An invalid scheme (d3's RdBu) is rejected by the compile step, like any bad value.""" + spec = { + "mark": "rect", + "encoding": { + "x": {"field": "ano", "type": "ordinal"}, + "y": {"field": "mes", "type": "ordinal"}, + "color": { + "field": "temp", + "type": "quantitative", + "scale": {"scheme": "RdBu"}, + }, + }, + } + + errors = _validate_chart_spec(spec, ["ano", "mes", "temp"]) + + assert any("compile" in error for error in errors) + + def test_choropleth_spec_is_valid(self): + """A choropleth compiles (real geometry) and its geo properties are allowed.""" + assert _validate_chart_spec(_choropleth_spec(), ["sigla_uf", "valor"]) == [] + + +class TestChartSpecUserPrompt: + def test_carries_task_data_on_the_first_attempt(self): + prompt = _chart_spec_user_prompt( + ["ano", "total"], [{"ano": 2025, "total": 10}], "a bar chart", None, [] + ) + + assert "a bar chart" in prompt + assert "ano" in prompt and "total" in prompt + assert "rejected" not in prompt # no retry feedback on the first attempt + + def test_echoes_the_rejected_spec_and_errors_on_retry(self): + prompt = _chart_spec_user_prompt( + ["ano"], [{"ano": 2025}], "a bar chart", {"mark": "nope"}, ["bad column"] + ) + + assert '"nope"' in prompt # the model's own rejected spec, echoed back + assert "bad column" in prompt # the validator's reason + assert "rejected" in prompt + + +def _structured_reply(spec: dict | None) -> dict: + """A `with_structured_output(include_raw=True)` result: a parsed spec, or a parse miss.""" + parsed = ChartSpec(spec=spec) if spec is not None else None + return {"raw": MagicMock(), "parsed": parsed, "parsing_error": None} + + +class TestGenerateChartSpec: + def _model(self, monkeypatch, *messages): + model = MagicMock() + model.ainvoke = AsyncMock(side_effect=list(messages)) + monkeypatch.setattr(charts, "_chart_spec_model", lambda: model) + return model + + async def test_retries_until_valid(self, monkeypatch): + """A spec that fails validation is regenerated with the errors fed back.""" + bad = {"mark": "bar", "encoding": {"x": {"field": "nope"}}} + good = {"mark": "bar", "encoding": {"x": {"field": "ano"}}} + model = self._model( + monkeypatch, _structured_reply(bad), _structured_reply(good) + ) + monkeypatch.setattr( + charts, + "_validate_chart_spec", + lambda spec, columns: [] + if spec["encoding"] == good["encoding"] + else ["bad column"], + ) + + result = await generate_chart_spec(["ano"], [{"ano": 2025}], "a bar chart") + + assert result["encoding"] == good["encoding"] + assert model.ainvoke.await_count == 2 + + async def test_rejected_spec_is_fed_back_on_retry(self, monkeypatch): + """The retry prompt carries the model's own rejected spec plus the errors.""" + bad = {"mark": "bar", "encoding": {"x": {"field": "nope"}}} + good = {"mark": "bar", "encoding": {"x": {"field": "ano"}}} + model = self._model( + monkeypatch, _structured_reply(bad), _structured_reply(good) + ) + monkeypatch.setattr( + charts, + "_validate_chart_spec", + lambda spec, columns: [] + if spec["encoding"] == good["encoding"] + else ["column 'nope' is not in the result"], + ) + + await generate_chart_spec(["ano"], [{"ano": 2025}], "a bar chart") + + system_message, user_message = model.ainvoke.await_args_list[1].args[0] + # The durable how-to is a system message; the retry data rides the user message. + assert isinstance(system_message, SystemMessage) + assert system_message.content == _CHART_SPEC_INSTRUCTIONS + assert '"nope"' in user_message.content # its own rejected spec, echoed back + assert "not in the result" in user_message.content # the validator's reason + + async def test_model_supplied_data_is_stripped(self, monkeypatch): + """A model spec carrying data/url is sanitized (not rejected) before returning.""" + spec = { + "mark": "bar", + "encoding": {"x": {"field": "ano"}}, + "data": {"url": "https://evil.example/x.json"}, + } + model = self._model(monkeypatch, _structured_reply(spec)) + monkeypatch.setattr(charts, "_validate_chart_spec", lambda spec, columns: []) + + result = await generate_chart_spec(["ano"], [{"ano": 2025}], "a bar chart") + + assert "data" not in result + assert "url" not in json.dumps(result) + assert model.ainvoke.await_count == 1 + + async def test_missing_spec_is_retried(self, monkeypatch): + """A reply the parser could not turn into a spec is treated as a failure and retried.""" + empty = _structured_reply(None) + good = _structured_reply({"mark": "bar", "encoding": {"x": {"field": "ano"}}}) + model = self._model(monkeypatch, empty, good) + monkeypatch.setattr(charts, "_validate_chart_spec", lambda spec, columns: []) + + result = await generate_chart_spec(["ano"], [{"ano": 2025}], "a bar chart") + + assert result["mark"] == "bar" + assert model.ainvoke.await_count == 2 + + async def test_raises_after_max_attempts(self, monkeypatch): + spec = {"mark": "bar", "encoding": {}} + messages = [_structured_reply(spec) for _ in range(MAX_CHART_SPEC_ATTEMPTS)] + model = self._model(monkeypatch, *messages) + monkeypatch.setattr( + charts, "_validate_chart_spec", lambda spec, columns: ["always bad"] + ) + + with pytest.raises(ChartSpecInvalid): + await generate_chart_spec(["ano"], [{"ano": 2025}], "a bar chart") + + assert model.ainvoke.await_count == MAX_CHART_SPEC_ATTEMPTS diff --git a/tests/app/test_exports.py b/tests/app/test_exports.py index 5163f61..e38d815 100644 --- a/tests/app/test_exports.py +++ b/tests/app/test_exports.py @@ -15,6 +15,7 @@ collect_query_handles, materialize_export, query_result_download, + sanitize_export_filename, ) from app.settings import settings @@ -31,18 +32,18 @@ def test_export_formats_match_the_advertised_literal(): def test_query_result_download_shape(): - """An executed query becomes one download; only CSV is offered for now.""" + """An executed query becomes one download offering the user-facing formats.""" assert query_result_download("qr_1", "slug") == { "type": "query_result", "query_ref": "qr_1", "slug": "slug", - "formats": ["CSV"], + "formats": ["AVRO", "CSV", "JSONL", "PARQUET"], } class TestCollectQueryHandles: - def test_collect_query_handles_appends_only_query_result_artifacts(self): - """query_result artifacts are appended as handles; other artifacts are ignored.""" + def test_collect_query_handles_returns_only_query_result_artifacts(self): + """query_result artifacts are returned as handles; other artifacts are ignored.""" artifacts = [ { "type": "query_result", @@ -54,10 +55,7 @@ def test_collect_query_handles_appends_only_query_result_artifacts(self): {"type": "file", "id": "x"}, # some other artifact kind ] - collected: list[CollectedQueryHandle] = [] - collect_query_handles(artifacts, collected) - - assert collected == [ + assert collect_query_handles(artifacts) == [ CollectedQueryHandle( query_ref="qr_1", slug="slug", destination_table=DESTINATION ) @@ -66,7 +64,7 @@ def test_collect_query_handles_appends_only_query_result_artifacts(self): def test_collect_query_handles_raises_on_malformed_query_result(self): """A query_result artifact missing a field is a producer bug — fail loud, don't skip.""" with pytest.raises(KeyError): - collect_query_handles([{"type": "query_result", "query_ref": "qr_1"}], []) + collect_query_handles([{"type": "query_result", "query_ref": "qr_1"}]) class TestMaterializeExport: @@ -216,3 +214,38 @@ def test_object_missing_after_extract_raises_runtime_error( with pytest.raises(RuntimeError, match="no file was written"): self._materialize() + + +class TestSanitizeExportFilename: + """The download filename guard, shared by the export tool and the endpoint.""" + + @pytest.mark.parametrize( + ("slug", "expected"), + [ + # A clean slug (what the model is asked to produce) passes through unchanged. + ("vendas_por_ano", "vendas_por_ano"), + # Hyphens and digits are allowed. + ("ideb-2021", "ideb-2021"), + # Spaces and punctuation collapse to a single underscore. + ("Vendas por ano", "Vendas_por_ano"), + ("a b", "a_b"), + ("café & leite!", "café_leite"), + # Leading/trailing separators are stripped, not left dangling. + ("_vendas_", "vendas"), + (" vendas ", "vendas"), + # Path separators and traversal are neutralized (no slashes or dots survive). + ("../../etc/passwd", "etc_passwd"), + ("relatorio/2021", "relatorio_2021"), + # File extensions are neutralized. + ("vendas_por_ano.csv", "vendas_por_ano_csv"), + # Accented word characters are preserved (\w is unicode). + ("população", "população"), + ], + ) + def test_sanitizes_slug(self, slug: str, expected: str): + assert sanitize_export_filename(slug, "resultados") == expected + + @pytest.mark.parametrize("slug", ["", " ", "!!!", "/", "..."]) + def test_falls_back_when_nothing_usable_remains(self, slug: str): + """A slug that sanitizes to empty falls back to the provided fallback, never ''.""" + assert sanitize_export_filename(slug, "resultados") == "resultados" diff --git a/uv.lock b/uv.lock index 0e96e74..571200f 100644 --- a/uv.lock +++ b/uv.lock @@ -150,6 +150,7 @@ dependencies = [ { name = "pydantic-settings" }, { name = "pyjwt" }, { name = "sqlmodel" }, + { name = "vl-convert-python" }, ] [package.dev-dependencies] @@ -184,6 +185,7 @@ requires-dist = [ { name = "pydantic-settings", specifier = ">=2.12.0" }, { name = "pyjwt", specifier = ">=2.10.1" }, { name = "sqlmodel", specifier = ">=0.0.31" }, + { name = "vl-convert-python", specifier = ">=1.9.0.post1" }, ] [package.metadata.requires-dev] @@ -2140,6 +2142,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6a/2a/dc2228b2888f51192c7dc766106cd475f1b768c10caaf9727659726f7391/virtualenv-20.36.1-py3-none-any.whl", hash = "sha256:575a8d6b124ef88f6f51d56d656132389f961062a9177016a50e4f507bbcc19f", size = 6008258, upload-time = "2026-01-09T18:20:59.425Z" }, ] +[[package]] +name = "vl-convert-python" +version = "1.9.0.post1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/93/89/36722344d1758ec2106f4e8eca980f173cfe8f8d0358c1b77cc5d2e035a4/vl_convert_python-1.9.0.post1.tar.gz", hash = "sha256:a5b06b3128037519001166f5341ec7831e19fbd7f3a5f78f73d557ac2d5859ef", size = 4663469, upload-time = "2026-01-21T00:09:55.61Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9f/59/e5862245972ff467d38b0eb5ad28154685e23ecabb47e14f2b6962da7b56/vl_convert_python-1.9.0.post1-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:43e9515f65bbcd317d1ef328787fd7bf0344c2fde9292eb7a0e64d5d3d29fccb", size = 30512930, upload-time = "2026-01-21T00:09:43.198Z" }, + { url = "https://files.pythonhosted.org/packages/62/e6/e7d0b538c2f0daaf120901dc113bd5d5d1fa51a9532fa5ffd90234e8c69e/vl_convert_python-1.9.0.post1-cp37-abi3-macosx_11_0_arm64.whl", hash = "sha256:b0e7a3245f32addec7e7abeb1badf72b1513ed71ba1dba7aca853901217b3f4e", size = 29738742, upload-time = "2026-01-21T00:09:46.016Z" }, + { url = "https://files.pythonhosted.org/packages/b8/e2/5645a1bc174c53ff8cd305ed76a4a76ba36e155302db20b42b7e78daeef8/vl_convert_python-1.9.0.post1-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e6ecfe4b7e2ea9e8c30fd6d6eaea3ef85475be1ad249407d9796dce4ecdb5b32", size = 33366278, upload-time = "2026-01-21T00:09:48.42Z" }, + { url = "https://files.pythonhosted.org/packages/a0/18/88e02899b72fa8273ffb32bde12b0e5776ee0fd9fb29559a49c48ec4c5fa/vl_convert_python-1.9.0.post1-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3c1558fa0055e88c465bd3d71760cde9fa2c94a95f776a0ef9178252fd820b1f", size = 33520215, upload-time = "2026-01-21T00:09:50.992Z" }, + { url = "https://files.pythonhosted.org/packages/2f/db/6e8616587035bf0745d0f10b1791c7e945180ac5d6b28677d2f2b3ca693c/vl_convert_python-1.9.0.post1-cp37-abi3-win_amd64.whl", hash = "sha256:7e263269ac0d304640ca842b44dfe430ed863accd9edecff42e279bfc48ce940", size = 32051516, upload-time = "2026-01-21T00:09:53.47Z" }, +] + [[package]] name = "watchfiles" version = "1.1.1"