waersdrtftghjk/us_census_agent

US Census 2019 Assistant

An asynchronous, conversation-aware RAG assistant that answers only from the connected 2019 US Census county dataset. It combines a Gradio chat UI, Chroma Cloud retrieval, LangChain/OpenAI models, session-scoped memory, groundedness evaluation, LangSmith observability, and model-specific concurrency controls.

This README is organized so that an evaluator can run the demo first and a new engineer can then understand how a request moves through the system.

Run the demo

Prerequisites

• —Python 3.11 or newer.
• —An OpenAI API key.
• —A Chroma Cloud collection populated with the Census county records described in Dataset contract. The sibling RAG ingestion project can populate that collection.

From the home_assignment/agent directory:

python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
pip install -r requirements.txt
cp .env.example .env

Add at least these values to .env:

OPENAI_API_KEY=...
CHROMA_API_KEY=...
CHROMA_COLLECTION=...

The example already contains the intended Chroma tenant/database and runtime defaults. Change them if the collection was created elsewhere. Never commit .env or real credentials.

Start the application:

python main.py

Open the address printed by Gradio, normally <http://127.0.0.1:7860>. The UI displays the conversation and temporary queue notices. Use LangSmith for token, cost, latency, error, and trace metrics.

Architecture at a glance

main.py is the composition root. It loads validated settings, creates one explicit application service graph, injects it into CensusAgent, and starts Gradio. Services are ordinary injected objects rather than hidden singletons.

Browser / Gradio UI
        |
        | message + complete session history + preferred language
        v
   CensusAgent
        |
        +--> language resolver
        |
        +--> prior-answer/meta-request router
        |       |
        |       +--> translate / summarize / repeat / feedback
        |       |       -> use referenced assistant message; no Chroma retrieval
        |       |
        |       +--> "same format for another location"
        |               -> retain format only; continue to retrieval
        |
        +--> session-aware intent and safety guardrail
                |
                +--> greeting / out-of-scope / toxic
                |       -> dynamic response; no Chroma retrieval
                |
                +--> contextual analysis of existing discussion
                |       -> answer from conversation memory
                |
                +--> new Census question
                        -> main model calls `retrieve_census_data` tool
                        -> tool observation: exact catalog location resolution
                        -> county records or state rollup
                        -> grounded final-model answer
                        -> utility-model groundedness judge
                        -> correction when rejected
        |
        v
answer + updated language preference

Why routing happens before retrieval

Translation, summarization, shortening, repetition, language feedback, and other operations on a previous answer are not new data questions. MetaFollowUpRouter receives every prior assistant response plus the complete session transcript, selects the referenced response, and sends it through the memory path. An intervening greeting or out-of-scope refusal therefore cannot overwrite the substantive answer the user named.

A retrieval_template decision is intentionally different: the referenced answer contributes structure and tone only. All facts still come from newly retrieved records for the requested location.

Retrieval policy

New Census-data requests use a bounded tool-calling loop. The main model receives only one external capability, retrieve_census_data, and must call it before returning Census facts. The tool result is added to the model's message history as an observation; the model may then produce its final answer, with a maximum of three tool-call rounds. The runtime rejects unsupported tool names and invalid arguments, so the model cannot access arbitrary services or execute code.

At the first retrieval, CensusRetrievalTool pages through the entire configured Chroma collection and caches the resulting catalog for the life of the process. Chroma Cloud limits a single get response, so CHROMA_CATALOG_PAGE_SIZE is capped at 300 and pagination continues until collection.count() records have been loaded.

The production answer path performs deterministic geography resolution against that catalog rather than accepting an arbitrary nearest-neighbor result:

• —A county request returns all exact county matches, with an optional state qualifier used to disambiguate duplicate county names.
• —A full state name or postal code resolves through the StateCode domain enum.
• —Texas resolves to state TX; it cannot fall through to Texas County.
• —Conversational LA resolves to Los Angeles, while explicit Louisiana, LA state, or equivalent state wording resolves to Louisiana.
• —A county FIPS code is accepted when matching FIPS metadata exists.
• —A broad state request creates a numeric rollup from the county records currently present and records how many counties were included.
• —An unknown or unavailable location produces a catalog-grounded availability response instead of unrelated vector results.

Because the complete catalog is cached, restart the agent after re-ingesting or changing the Chroma collection.

Dataset contract

The runtime expects the configured Chroma collection to use this shape:

• —document: county name, such as Los Angeles County.
• —metadata.STATE: state name or postal code used for state resolution.
• —Remaining metadata: scalar Census measurements passed to the answer model as retrieved evidence.

The supplied ingestion pipeline writes renter/housing-burden fields including TOTAL_RENTER_UNITS, RENTER_50_PLUS_PCT, and PCT_OVER_50_PERCENT_BURDEN. The agent does not require every metric to be present: it tells the model to preserve available values and disclose incomplete records.

The runtime and ingestion projects must point to the same Chroma tenant, database, and collection. To compare those identities and inspect LA retrieval independently of the LLM, run:

python diagnose_retrieval.py --top-k 10

Add --full-embedding only when the full raw query embedding is needed. The command also prints unfiltered top-k documents and distances for diagnostic comparison; this semantic-search output is diagnostic and is not the runtime's exact-location policy.

Module map

Models and responsibilities

Model names are environment settings, so compatible models can be substituted without changing orchestration code.

Configuration reference

For LangSmith, use the current variables in .env and restart the process:

LANGSMITH_TRACING=true
LANGSMITH_API_KEY=lsv2_...
LANGSMITH_PROJECT=us-census-2019-assistant

The application accepts the older LANGCHAIN_TRACING_V2, LANGCHAIN_API_KEY, and LANGCHAIN_PROJECT names for migration, but new deployments should use LANGSMITH_*. When enabled, each browser request creates a census-agent-request root trace; its child spans include every LangChain/OpenAI call and the direct census-catalog-retrieval Chroma step. The process also writes one census-agent-startup trace at launch. LangSmith then calculates trace count, error rate, end-to-end and per-step latency, LLM call count, and model token/cost metrics in the project dashboard. If trace delivery fails, the app emits a credential-free langsmith_trace_delivery_failed warning in its server log.

Conversation state, limits, and queueing

Gradio stores history and preferred_language in gr.State, so they are isolated per browser session and cleared by the UI's Clear button. They are not persisted across a page reload or server restart.

LangSmith is the sole observability surface for token, cost, latency, trace, and error metrics; the application does not persist local usage counters.

ModelConcurrency owns two application-scoped semaphores:

• —MAX_CONCURRENT_GPT4O for main-model work.
• —MAX_CONCURRENT_GPT4O_MINI for utility, guardrail-response, and judge work.

If the relevant semaphore is already occupied, the UI immediately displays the configured friendly queue notice and replaces it with the final response after a slot becomes available.

Semaphores are local to one Python process. A multi-worker or replicated deployment needs a distributed concurrency limiter if model concurrency must be global.

Prompts and failure behavior

All LLM prompts, runtime user messages, response purposes, and language aliases are stored in config/prompts.toml. PromptCatalog validates required entries and exact template placeholders at startup, then the same immutable catalog is injected into every component. Restart the app after changing the catalog.

The answer model is instructed to use retrieved/session evidence only. The judge can reject an unsupported or irrelevant answer and trigger a correction grounded in the same evidence. If a dependency fails unexpectedly, the exception is logged and the user receives the configured generic failure message; secrets and internal exception details are not shown in the chat.

Development checks

Install development dependencies and run the existing offline suite:

pip install -r requirements-dev.txt
pytest -q ../tests/agent

Run static formatting/lint checks and compilation:

ruff check census_agent main.py diagnose_retrieval.py
ruff format --check census_agent main.py diagnose_retrieval.py
python -m compileall -q census_agent main.py diagnose_retrieval.py

The unit suite mocks external model and Chroma boundaries. Use diagnose_retrieval.py and the suggested evaluation conversation for connected integration verification.