RaghavRaahul/census-data-agent

Census Data Agent

Live demo

• —URL: https://huggingface.co/spaces/RaghavRaahul/census-data-agent
• —Access: Open in browser, no authentication required.
• —Suggested prompts:
• —"What is the median household income in 2020?"
• —"Show the top 10 block groups by housing vacancy in 2019."

Verification references

• —Deployed chat interface: https://huggingface.co/spaces/RaghavRaahul/census-data-agent
• —Census source (ACS): https://www.census.gov/programs-surveys/acs
• —ACS 2020 group definition for B19013 (median household income): https://api.census.gov/data/2020/acs/acs5/groups/B19013.html

Note: screenshot image files were not committed to keep Hugging Face Space and GitHub repositories synchronized without binary-file push issues.

Quickstart: Snowflake connection test

1. 1.Create .env from the template:

cp .env.example .env

1. 1.Fill in all SNOWFLAKE_* values in .env.
2. 2.Optional for LLM SQL mode: set OPENAI_API_KEY and OPENAI_MODEL.

1. 1.Create a virtualenv and run the connection check:

python3 -m venv .venv
.venv/bin/pip install -r requirements.txt
.venv/bin/python scripts/test_connection.py

If successful, the script prints current Snowflake database, schema, warehouse, and role.

Metadata extraction and chunk pipeline

1. 1.Extract schema metadata from Snowflake:

.venv/bin/python scripts/extract_metadata.py

1. 1.Build enriched column metadata and table chunks:

# Faster, no external Census API labels
.venv/bin/python scripts/build_chunks.py

# Optional: fetch ACS/decennial labels from Census API
.venv/bin/python scripts/build_chunks.py --fetch-labels

# Optional: tune fetch-label performance
.venv/bin/python scripts/build_chunks.py --fetch-labels --max-workers 16

1. 1.Evaluate generated chunks:

.venv/bin/python scripts/evaluate_chunks.py

Artifacts are written to data/:

• —raw_tables.csv
• —raw_columns.csv
• —columns_enriched.csv
• —table_index.jsonl
• —column_label_cache.json (only when --fetch-labels is used)

Run tests

.venv/bin/pytest -q

Retrieval and query planning

Validate lexical + semantic retrieval:

.venv/bin/python scripts/validate_retrieval.py

Generate a constrained pre-SQL query plan from a natural language question:

.venv/bin/python scripts/plan_query.py --query "median household income in 2020"

Generated artifacts:

• —data/retrieval_validation.json
• —data/last_query_plan.json

Run end-to-end retrieval -> planning -> validation -> SQL compilation:

.venv/bin/python scripts/run_query_pipeline.py --query "top 10 block groups by housing vacancy in 2019"

Optionally execute compiled SQL in Snowflake:

.venv/bin/python scripts/run_query_pipeline.py --query "median household income in 2020" --execute

Disable LLM SQL generation (deterministic compiler only):

.venv/bin/python scripts/run_query_pipeline.py --query "median household income in 2020" --execute --no-llm-sql

Generated artifact:

• —data/last_pipeline_output.json

Streamlit app

Run the interactive chat app locally:

.venv/bin/streamlit run streamlit_app.py

The app includes:

• —hybrid retrieval (lexical + semantic)
• —constrained planning + SQL compilation (LLM-first with deterministic fallback)
• —guardrails for out-of-scope and prompt/secret-leak attempts
• —optional debug payload viewer in the sidebar

Complex query support (current scope)

• —Current deterministic compiler supports:
• —single-table rank and aggregate queries
• —optional safe metadata join for geography enrichment when query asks for fields like state/county/tract
• —Join policy is intentionally conservative:
• —only joins to retrieved metadata geographic table
• —only when year and grain match primary table
• —only on CENSUS_BLOCK_GROUP
• —If this join condition is not confidently satisfied, pipeline falls back to non-join deterministic SQL.

Multi-turn context handling (LLM + bounded memory)

• —The app now resolves follow-up questions using an LLM context resolver before running the SQL pipeline.
• —Context strategy:
• —keep a bounded recent-turn window (last 10 turns)
• —maintain a structured rolling summary for older turns
• —resolve latest user follow-up into a standalone resolved_query
• —Safety behavior:
• —low-confidence follow-up resolution triggers clarification instead of forced execution
• —resolved query still goes through existing guardrails, plan validation, and SQL safety checks

Follow-up tests

• —Added tests/test_context_memory.py covering:
• —follow-up carryover (what about 2020?)
• —explicit override (actually use 2019)
• —low-confidence clarification behavior
• —memory compaction and summary retention

LLM SQL reliability notes

• —The LLM SQL path now adds three reliability controls:
• —SQL normalization (identifier quoting + trailing semicolon cleanup)
• —compact few-shot examples for rank and aggregate query shapes
• —one-retry repair loop when first output fails validation
• —Why trailing semicolon is removed:
• —the validator treats semicolons conservatively to reduce multi-statement/injection risk
• —many models append a harmless trailing ;, so normalization removes it to avoid false fallback
• —this keeps strict single-statement safety while improving pass-through consistency

Architecture: Current vs Target

Current implementation

flowchart TD
  userQuery[UserQuery] --> intentHeuristics[IntentHeuristics]
  userQuery --> hybridRetriever[HybridRetriever]
  hybridRetriever --> preSqlPlan[DeterministicPreSQLPlanJSON]
  intentHeuristics --> preSqlPlan
  preSqlPlan --> llmSql[LLMSQLGeneration]
  llmSql --> validator[SQLSafetyValidation]
  validator --> fallback[DeterministicFallbackCompiler]

Target implementation

flowchart TD
  userQuery[UserQuery] --> hybridRetriever[HybridRetriever]
  hybridRetriever --> llmPlanner[LLMPlannerToStructuredJSON]
  llmPlanner --> validator[SchemaAndPolicyValidator]
  validator --> sqlCompiler[TemplateSQLCompiler]
  sqlCompiler --> snowflakeExec[SnowflakeExecution]
  snowflakeExec --> groundedAnswer[GroundedAnswer]

Notes:

• —Current constrained planner is deterministic/rule-based and emits pre-SQL JSON.
• —Current SQL generation attempts LLM first, with strict validation and deterministic fallback.
• —Target state adds LLM structured planning (query -> strict JSON) before SQL compilation.

Running project notes

Progress log and decisions are maintained in RUNNING_LOG.md.

Additional documentation:

• —TEST_RUNS.md
• —LATENCY_CHECKS.md
• —REFLECTION.md
• —STRENGTHS_WEAKNESSES.md
• —TECHNICAL_DEEP_DIVE.md
• —RESUME_TALKING_POINTS.md