Scaffold server query pipeline & streaming API by sagrimson · Pull Request #353 · datacommonsorg/tools

sagrimson · 2026-06-03T00:43:43Z

Overview

Adds the streaming query API route and supporting server-side pipeline that powers Data Weaver's natural-language-to-chart workflow. A user's query flows through safety checks, analysis, Gemini tool-calling with the Data Commons MCP server, and streams structured chart specifications back to the client via SSE.

Changes Made

Streaming API route (/api/query) — POST endpoint that accepts a natural-language query and streams SSE events (status, analysis, tool_call, chart_spec, metadata, error, complete)
Server pipeline (src/server/)
- safety.ts — layered prompt safety gate (regex patterns + LLM classification)
- analyze/ — extracts places, topic, titles, and date ranges from the query via Gemini
- query/ — Gemini tool-loop that iteratively calls MCP tools to resolve variables
- observations/ — fetches variable metadata/facets from Data Commons
- mcp.ts — JSON-RPC client for the Data Commons MCP server
- gemini.ts — shared Gemini client factory
- dc-api.ts — Data Commons REST API helpers
- config.ts + config/ — service config and skill prompt loading
- types.ts — shared types for the query pipeline and SSE events
Client streaming hook (use-streaming-query.ts) — SSE parser with abort support and typed event dispatch

gemini-code-assist

Code Review

This pull request introduces the 'Atlas' scope component, integrating tldraw to render interactive text and chart cards on a canvas, alongside supporting elements like buttons, skeletons, and custom icons. On the backend, it implements a streaming query API powered by Gemini and Data Commons MCP tools, complete with query analysis, safety filtering, and metadata fetching. The review feedback focuses on aligning with repository conventions—such as renaming files to snake_case and applying category-first naming to components—and ensuring strict TypeScript compliance. Additionally, the reviewer recommends enhancing robustness by adding unmount cleanup hooks to prevent memory leaks, caching compiled regular expressions for better performance, and wrapping JSON parsing in try-catch blocks to prevent runtime TypeErrors.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

beets

Thanks for the PR! The scaffolding and overall flow for the data discovery query pipeline are well-designed. Since this is an initial commit, most of these comments are fine to be added as TODO's in the code for future refactor.

beets · 2026-06-04T23:26:35Z

+      headers: { 'Content-Type': 'application/json' },
+    });
+  }
+  const { query, atlasContext } = body;


selectedEntityDcids is defined in the request body but not yet destructured here. For follow-up queries (where parsed.places is empty), falling back to selectedEntityDcids will ensure previous place context is preserved.

Also, adding a brief verification check ensuring query is a valid string before invoking safety filters will safeguard against runtime TypeErrors.

beets · 2026-06-04T23:27:47Z

+          const entityDcid =
+            parsedResponse.entityDcid || resolvedPlaceDcid || place;


If a place cannot be resolved to a valid DCID, report a warning status or fail gracefully rather than calling the observations API with the raw name.

beets · 2026-06-04T23:38:27Z

+    ? `\n\nDATE CONSTRAINT: The user wants data limited to ${parsed.dateRange.start && parsed.dateRange.end ? `years ${parsed.dateRange.start} through ${parsed.dateRange.end}` : parsed.dateRange.start ? `years from ${parsed.dateRange.start} onward` : `years up to ${parsed.dateRange.end}`}.`
+    : '';


Date constraints might not be limited to years. Could we pass the dates directly in the prompt without assumptions, perhaps as a structured input.

beets · 2026-06-05T13:24:22Z

+    });
+
+    if (response.functionCalls && response.functionCalls.length > 0) {
+      const fc = response.functionCalls[0];


Gemini models frequently output multiple tool calls concurrently. Currently, this extracts only response.functionCalls[0]. As we iterate, let's loop over all calls in response.functionCalls, execute each, and push their tool responses back in sequence to avoid API conversation history mismatch errors. Could add a TODO for now as a reminder.

beets · 2026-06-05T13:30:24Z

+            const cleaned = responseText
+              .replace(/```json/g, '')
+              .replace(/```/g, '')
+              .trim();


Instead of stripping out the json markdown, consider searching for the json block instead to be resilient against unexpected model commentary:

.match(/\{[\s\S]*\}/)

beets · 2026-06-05T13:55:25Z

Should we consolidate all our agent instructions into the skills folder?

beets · 2026-06-05T13:57:44Z

+    .replace(/```json/g, '')
+    .replace(/```/g, '')
+    .trim();


Similar to earlier comment on handling possible LLM preamble.

const jsonMatch = responseText.match(/\{[\s\S]*\}/); const cleaned = jsonMatch ? jsonMatch[0] : responseText; parsedResponse = JSON.parse(cleaned);

beets · 2026-06-05T13:59:12Z

+
+  const response = await genAI.models.generateContent({
+    model: config.models.parseQuery,
+    contents: `${systemPrompt}\n\nQuery: "${query}"\nReturn ONLY its JSON object, no markdown, no other text or explanation.`,


Looks like the return formatting instruction is already part of the system prompt

beets · 2026-06-05T14:03:48Z

@@ -0,0 +1,59 @@
+import { getServiceConfig } from './config';
+
+export interface ObservationResponse {


Nice :) thanks for adding the type for this.

beets · 2026-06-05T14:08:28Z

+  | { type: 'status'; message: string }
+  | { type: 'parsed_query'; data: ParsedQuery }
+  | { type: 'tool_call'; tool: string; args: Record<string, unknown> }
+  | { type: 'query_result'; result: QueryResult; place: string }
+  | { type: 'complete' }
+  | { type: 'error'; message: string };


Could we lift out the types of these events to string constants?

sagrimson requested a review from a team as a code owner June 3, 2026 00:43

gemini-code-assist Bot reviewed Jun 3, 2026

View reviewed changes

sagrimson force-pushed the scott/api-route branch 3 times, most recently from b80edeb to c108463 Compare June 3, 2026 17:16

sagrimson requested review from beets and nick-nlb June 3, 2026 17:46

sagrimson force-pushed the scott/api-route branch from 656111f to b2e7d1e Compare June 3, 2026 23:40

sagrimson added 3 commits June 3, 2026 16:52

add api routes

91a2881

Feedback

7031a5c

Additional feedback, plus fix typo in server/observations/index.ts

a2aab0e

sagrimson force-pushed the scott/api-route branch from b2e7d1e to a2aab0e Compare June 3, 2026 23:52

sagrimson added 3 commits June 4, 2026 15:33

Rename Analyze to ParseQuery

0b3faaf

rename 'query' to 'data_discovery'

7f4e4b3

include metadata in data discovery step

2f80546

beets reviewed Jun 5, 2026

View reviewed changes

		const entityDcid =
		parsedResponse.entityDcid \|\| resolvedPlaceDcid \|\| place;

		? `\n\nDATE CONSTRAINT: The user wants data limited to ${parsed.dateRange.start && parsed.dateRange.end ? `years ${parsed.dateRange.start} through ${parsed.dateRange.end}` : parsed.dateRange.start ? `years from ${parsed.dateRange.start} onward` : `years up to ${parsed.dateRange.end}`}.`
		: '';

		@@ -0,0 +1,59 @@
		import { getServiceConfig } from './config';

		export interface ObservationResponse {

Conversation

sagrimson commented Jun 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Changes Made

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

beets left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

sagrimson commented Jun 3, 2026 •

edited

Loading