Viz.ChatGrid: add Context Management Module for Large SQL-Backed AI Chat Responses

[Walter462]

Problem

Viz.ChatGrid currently passes raw SQL query results directly into the LLM prompt when generating chat responses.
When result sets are large (multi-state, multi-year, etc.), this causes:

• LLM context window overflow (current limit is 4k)
• Truncation of SQL output
• Biased answers caused by SQL ordering (e.g. only the first state like Alabama appears in a multi-state query)

This makes chat responses incorrect, misleading, or incomplete.
The system currently relies on auto-generated SQL filtering and ordering to control what fits into the LLM prompt, which is fragile and unsupervised.

Root Cause

There is no explicit layer responsible for:

• Token budgeting
• Result sampling / ranking
• Summarization
• Format-specific context shaping

The LLM receives raw database output, not curated context.

Proposed Solution

Add a Context Management Module between the database and ChatResponsePrep.

flowchart TB
  user>"USER"]
  user --"natural lang query"--> chat
subgraph "Front__End"
  chat[["CHAT"]]
  end
subgraph "BackEnd"
  LangChain[\"LangChain"/]
  DB[("DB")]
  LangChain --"srl request"-->DB
  ChatResponsePrep[\"Generate Natural language AI response for chat"/]
  OpenAIChatResponse["OpenAI"]
  chat --"natural lang query"--> LangChain
  ChatResponsePrep --"context enriched OpenAI request"--> OpenAIChatResponse
  OpenAIChatResponse --"sql-return context-powered reply"--> ChatResponsePrep
  ContextManagementModule(("Context
  Management
  Module"))
  DB --"raw data"--> ContextManagementModule
  ContextManagementModule --> ChatResponsePrep
  LangChain --"natural lang query"--> ChatResponsePrep 
  ChatResponsePrep --"human-readable reply"-----> chat
end
style ContextManagementModule fill:#000,stroke:#f55,stroke-width:4px,color:#fff,stroke-dasharray: 5 5

Loading

This module should:

• Enforce token budgets
• Select representative rows when full data doesn’t fit
• Perform summarization or aggregation where appropriate
• Provide structured, format-aware context to the LLM

Goals / Acceptance Criteria

The system should guarantee that:

• Multi-group queries (e.g., multi-state, multi-year) always include all groups
• No single category dominates context due to SQL ordering
• The LLM never receives raw, unbounded SQL dumps
• Context fits within LLM token limits while preserving correctness

Possible Implementations (non-binding)

• MCP-based context governance (schema, budgets, coverage guarantees)
• Result ranking + sampling
• Group-aware row selection
• SQL-native aggregation (GROUP BY, COUNT, SUM, AVG, etc.)
• Pre-LLM aggregation (per state, per year, etc.)
• RAG-style retrieval
• LLM-based summarization before final prompt

[pelesh]

@Walter462, thank you for such detailed analysis and suggestions. If you already have ideas how to implement this, feel free to make a pull request.

[tsybinae]

Looks great. Here are a couple database schemas to check if it is database structure dependent. This database contains an open source power flow case, in schema1 divided into 10 tables, in schema2 into 2 tables.

The attachment is a zip file, unzipping it will produce a exago_db1_two_schemas.dump which is a user agnostic packaged postgre database.

exago_db1_two_schemas.zip