UDIAgent

LLM-powered data visualization orchestration library for the Universal Discovery Interface (UDI).

UDIAgent orchestrates LLM calls to generate data visualization specs from natural language queries. It can be used as a standalone Python library or deployed as a FastAPI microservice.

Installation

# Core library only
pip install udiagent

# With the reference FastAPI server
pip install udiagent[server]

# With LangFuse observability
pip install udiagent[langfuse]

# With benchmarking tools
pip install udiagent[benchmark]

# Everything
pip install udiagent[all]

For local development with uv:

uv sync --extra server --extra langfuse --extra test   # server + dev

Library Usage

from udiagent import UDIAgent, Orchestrator

# Initialize the agent with explicit configuration (no environment variables)
agent = UDIAgent(
    gpt_model_name="gpt-5.4",
    openai_api_key="sk-...",
)

# Create an orchestrator
orch = Orchestrator(agent)

# Run a query
result = orch.run(
    messages=[{"role": "user", "content": "Show me a bar chart of donors by sex"}],
    data_schema='{"resources": [...]}',
    data_domains='[{"entity": "donors", "field": "sex", ...}]',
)

# result.tool_calls — list of tool call dicts (e.g. RenderVisualization, FilterData)
# result.orchestrator_choice — "render-visualization", "both", "explain", etc.

One agent, any schema

data_schema / data_domains are per request — a single long-lived Orchestrator serves queries against arbitrary, unrelated datasets with no per-schema setup, regeneration, or restart. The visualization templates are schema-independent: the tool definitions expose free-form entity / field / dimension string arguments, and the schema needed to validate those bindings and fill in a concrete spec is parsed from the data_schema on each call. Just pass a different schema:

orch = Orchestrator(agent)                       # once
orch.run(messages=[...], data_schema=hubmap_schema, data_domains=hubmap_domains)
orch.run(messages=[...], data_schema=penguins_schema, data_domains=penguins_domains)

Visualization template sets and tags

All templates live in one file (data/skills/template_visualizations.json), generated by a single script (scripts/template_viz_generation.py) with the line-item and data-cube variant of each chart type side by side. Every template carries multi-axis tags — a data-shape tag plus a chart-type tag (e.g. ["data_cube", "barchart"]). The orchestrator selects templates per request by the data-shape tag inferred from the incoming schema:

Data-shape tag	Used for	Selected when
line_item	Tidy, per-record tables (groupby/rollup)	default
data_cube	Pre-aggregated "powerset" cubes (marginals)	the schema marks a resource udi:cube (or declares udi:dimensions + udi:measures)

A data cube is a pre-aggregated table with one measure column and several dimension columns, where a row's empty dimensions mean it is aggregated over them. Cube templates read a value by marginal filtering (the active dimensions non-null, every other dimension null) instead of re-aggregating. The marginal filter is built at runtime from the schema's udi:dimensions, so the cube templates work for any cube — mark a resource like this:

{
  "name": "encounter_counts",
  "udi:cube": true,
  "udi:measures": ["cnt"],
  "udi:dimensions": ["period_start_month", "class_display", "gender", "..."],
  "schema": { "fields": [/* cnt + one field per dimension */] },
}

Regenerate the unified template file and the combined typed tool module (schema-free, deterministic) with:

python scripts/regenerate_vis_tools.py

Generation validates every spec against the UDI grammar and reports any non-conforming templates (non-fatal; pass --strict to template_viz_generation.py to fail hard).

Third-party OpenAI-compatible backends

Point openai_base_url at any OpenAI-compatible chat-completions root and set gpt_model_name to that backend's model id:

agent = UDIAgent(
    gpt_model_name="openai.gpt-oss-120b-1:0",
    openai_api_key="...",
    openai_base_url="https://bedrock-mantle.us-east-1.api.aws/openai/v1",
)

Backend	openai_base_url	Notes
Azure AI Foundry	https://<resource>.openai.azure.com/openai/v1	Key is the Foundry API key
Amazon Bedrock	https://bedrock-mantle.<region>.api.aws/openai/v1	Key is a Bedrock API key (see below for IAM roles)
OpenRouter	https://openrouter.ai/api/v1	Model ids are vendor/model
Ollama / vLLM / LMS	http://localhost:11434/v1	No key needed — a placeholder is supplied

The backend must support function calling and JSON-schema structured outputs (response_format: {type: "json_schema", strict: true}); the orchestrator additionally uses tool_choice: "required". Backends missing these will fall through to the non-LLM fallback paths and produce degraded results, so smoke-test one visualization request after switching.

Scoping: passing openai_base_url to the constructor affects only the default client — per-request X-OpenAI-Key callers still reach api.openai.com. Setting the OPENAI_BASE_URL environment variable instead (as the server does) routes every client there, including per-request keys, because the OpenAI SDK applies that variable itself.

Amazon Bedrock with an instance/task role rather than a static key is not wired up: it needs SigV4, i.e. OpenAI(provider=openai.providers.bedrock(...)) plus the openai[bedrock] extra, which is mutually exclusive with api_key/base_url.

With LangFuse observability

LangFuse tracing is opt-in. Install the extra (pip install udiagent[langfuse]) and pass any of the three credentials to UDIAgent:

agent = UDIAgent(
    gpt_model_name="gpt-5.4",
    openai_api_key="sk-...",
    langfuse_public_key="pk-lf-...",
    langfuse_secret_key="sk-lf-...",
    langfuse_host="https://cloud.langfuse.com",  # or your self-hosted URL
    langfuse_environment="production",            # optional; tags traces (e.g. "staging")
)

Tracing turns on when any of langfuse_public_key, langfuse_secret_key, or langfuse_host is set. langfuse_environment is purely a tag — it labels traces in the LangFuse UI but does not by itself enable tracing.

Installing udiagent[langfuse] alone does not enable tracing — credentials must be supplied explicitly. Library consumers who prefer environment-variable configuration should read the env vars themselves and pass the values to the constructor.

Key Classes

Class	Description
UDIAgent	OpenAI client wrapper
Orchestrator	Routes user requests to visualization, filter, explanation, and clarification handlers
OrchestratorResult	Dataclass with tool_calls and orchestrator_choice

Utility Functions

Function	Description
load_grammar()	Load the UDI Grammar JSON schema (bundled with the package)
load_skills()	Load skill prompt templates (bundled with the package)
render_template()	Substitute {{key}} placeholders in a skill instruction template
generate_vis_spec()	Generate a visualization spec using the skills pipeline
simplify_data_domains()	Simplify data domains JSON into compact LLM-friendly text
parse_schema_from_dict()	Parse a data schema dict into structured format

Server Usage

The udiagent.server subpackage provides a reference FastAPI application that wraps the library as a configurable microservice. It reads configuration from environment variables.

Running the Server for Local Development

Pass --extra server on every uv run (it provides the fastapi CLI). uv run prunes any extra you don't name, so a bare uv run fastapi errors with Failed to spawn: fastapi. Add the backend driver extra too when using server-side data (--extra duckdb and/or --extra starrocks).

# Development
uv run --extra server fastapi dev src/udiagent/server/app.py --port 8007

# Production
uv run --extra server fastapi run src/udiagent/server/app.py --port 8007

Server Environment Variables

Variable	Required	Default	Description
OPENAI_API_KEY	No	—	OpenAI API key. If not set, must be provided per-request via X-OpenAI-Key header.
OPENAI_BASE_URL	No	—	OpenAI-compatible backend root (see Third-party backends)
GPT_MODEL_NAME	No	gpt-5.4	Model for orchestration; with a custom base URL, that backend's model id
JWT_SECRET_KEY	Yes*	—	JWT signing key; the server refuses startup when missing (*unless INSECURE_DEV_MODE=1 or JWT_JWKS_URL is set)
JWT_ALGORITHM	No	HS256	JWT algorithm; set to the identity provider's (e.g. RS256) when using JWT_JWKS_URL
JWT_JWKS_URL	No	—	Verify externally issued tokens against an identity provider's JWKS endpoint, instead of JWT_SECRET_KEY (see External identity providers)
JWT_AUDIENCE	Yes*	—	Expected aud claim (*required with JWT_JWKS_URL)
JWT_ISSUER	No	—	Expected iss claim; validated only when set
INSECURE_DEV_MODE	No	0	Set to 1 to skip JWT verification (development only)
LANGFUSE_SECRET_KEY	No	—	LangFuse observability secret key (opt-in; tracing is disabled when unset)
LANGFUSE_PUBLIC_KEY	No	—	LangFuse observability public key (opt-in; tracing is disabled when unset)
LANGFUSE_HOST	No	—	LangFuse instance URL (e.g. https://cloud.langfuse.com)
LANGFUSE_ENVIRONMENT	No	—	Tags traces with an environment label (e.g. production); does not enable tracing
UDI_QUERY_BACKENDS	No	—	Path to a JSON file configuring server-side query backends (see below)
UDI_METADATA_TTL_SECONDS	No	3600	TTL for the introspected-metadata cache
UDI_LOG_DIR	No	<package root>/logs	Where the rotating log file goes; skipped (stream logs only) if unwritable
UDI_DATA_DIR	No	<package root>/data	Repo-level dev data for /v1/yac/examples; set this when installed from a wheel

External identity providers

When the chat is embedded in a portal that already authenticates its users, the server can verify the portal's own tokens instead of issuing its own. Point JWT_JWKS_URL at the identity provider's JWKS endpoint (Keycloak, Globus, Auth0, Entra — anything publishing JWKS) and set JWT_ALGORITHM to the algorithm it signs with:

JWT_JWKS_URL=https://idp.example/realms/udi/protocol/openid-connect/certs
JWT_ALGORITHM=RS256
JWT_AUDIENCE=udi-yac
JWT_ISSUER=https://idp.example/realms/udi   # optional

JWT_AUDIENCE is required in this mode: without it, any token the provider minted for any of its clients would be accepted here. JWT_ISSUER is checked only when set.

JWT_JWKS_URL and JWT_SECRET_KEY are mutually exclusive — setting both refuses startup, as does a JWKS URL with a symmetric JWT_ALGORITHM or without JWT_AUDIENCE.

Keys are fetched on the first authenticated request and cached for 5 minutes. A token whose kid isn't in the cached set triggers an early refresh, so a provider that rotates without pre-publishing its new key doesn't cause an outage until the cache expires; those early refreshes are throttled to one every 10 seconds. While the provider is unreachable, requests return 503.

On the frontend nothing changes: the host portal passes its token to the chat as authToken (see UDIChatConfig), which forwards it as Authorization: Bearer. Note that udi-yac does not run an identity provider of its own — for a standalone deployment with no portal in front, use JWT_SECRET_KEY.

Server Endpoints

Endpoint	Method	Description
/	GET	API status and info
/v1/yac/completions	POST	Main orchestrator — routes user requests to tools
/v1/yac/query	POST	Server-side data: batched grammar→SQL query execution
/v1/yac/metadata	GET	Server-side data: introspected dataSchema/dataDomains
/v1/yac/benchmark	POST	Benchmark variant with optional orchestrator override
/v1/yac/examples	GET	Example prompts from data/example_prompts.json
/v1/yac/structured_functions	GET	Structured function registry
/v1/yac/benchmark_analysis	GET	Latest benchmark analysis results

Server-side query backends

/v1/yac/query and /v1/yac/metadata let data stay on the server: the agent compiles UDI grammar transformation pipelines to SQL and runs them against a configured database (StarRocks, DuckDB) instead of the browser loading CSVs. Configure backends via UDI_QUERY_BACKENDS (a package→backend JSON map) and point the chat at a package with VITE_UDI_REMOTE_PACKAGE.

Full architecture + integration guide: src/udiagent/query/README.md. Local dev instance: dev/starrocks/README.md.

Docker

Build from the repo root, not from packages/agent: uv.lock is the uv workspace lockfile and lives at the root (see root pyproject.toml [tool.uv.workspace]). The .dockerignore at the root keeps the JS half of the monorepo out of the build context.

docker build -f packages/agent/Dockerfile -t udiagent .   # from the repo root
docker run -p 8007:80 --env-file packages/agent/.env udiagent

The image installs the server + langfuse extras. Add --extra duckdb / --extra starrocks to both uv sync lines if the deployment serves server-side query backends.

Deployment Guide

Step-by-step for standing the agent up as a server on a fresh host. There are two supported ways to get the code onto the host — pick one, then follow the shared steps:

	Path A — container from source	Path B — udiagent[server] from PyPI
Use when	deploying a branch or unreleased code; want the exact pinned deps	deploying a released version; no repo checkout; own process manager
Host needs	Docker + a clone of this repo	Python ≥ 3.12 and pip/uv
Dependency versions	exact, from the workspace uv.lock	resolved at install time from pyproject.toml ranges
Process supervision	--restart unless-stopped	yours (systemd unit below)
Dev-data endpoints	work out of the box	need UDI_DATA_DIR (step 3B)
CI deploy	deploy-agent.yml	—

Steps 1–2 and 4–5 apply to both, step 3 splits, step 6 is Path A only, step 7 is client-side.

1. Get the code on the host

Path A — clone the repo:

git clone https://github.com/hms-dbmi/udi-yac.git
cd udi-yac
git checkout <branch>          # e.g. nickakhmetov/third-party-providers-2

Only the Python workspace matters for the agent image — no pnpm install, no uv sync on the host. The build installs dependencies inside the image.

Path B — install the published distribution:

python -m venv /opt/udiagent/venv          # or: uv venv /opt/udiagent/venv
/opt/udiagent/venv/bin/pip install "udiagent[server]"

The PyPI distribution is named udiagent (udi-yac is the npm chat package). Add extras as needed — "udiagent[server,langfuse]", plus duckdb / starrocks for server-side query backends.

2. Write the production env file

Copy the template and fill it in. Environment variables are the server's entire configuration surface — there is no config file format beyond this.

# Path A: on the host, anywhere the container can read it
cp packages/agent/.env.template /home/ec2-user/.env
# Path B: in the directory the service will run from (its CWD is searched)
cp packages/agent/.env.template /opt/udiagent/.env
chmod 600 /home/ec2-user/.env   # or /opt/udiagent/.env

Path B has no repo checkout, so grab the template from GitHub or just write the values below by hand.

Minimum production values:

# Auth — REQUIRED. The server refuses to start with an empty signing key
# unless INSECURE_DEV_MODE=1. Never set INSECURE_DEV_MODE in production.
INSECURE_DEV_MODE=0
JWT_SECRET_KEY=<paste output of: openssl rand -hex 32>

# LLM backend — either a server-held key, or leave blank to require callers
# to send their own via the X-OpenAI-Key header.
OPENAI_API_KEY=sk-...
GPT_MODEL_NAME=gpt-5.4

Full variable list: Server Environment Variables.

Where the values come from, in precedence order: real environment variables (Docker --env-file, systemd EnvironmentFile, shell exports) win, then <package root>/.env (the repo checkout), then ./.env from the process working directory. On Path B the first and third are the usable ones.

(Optional) point at a third-party OpenAI-compatible backend

To serve from Azure AI Foundry, Bedrock, OpenRouter, or a self-hosted Ollama/vLLM, add the backend root and its model id:

OPENAI_BASE_URL=https://openrouter.ai/api/v1
GPT_MODEL_NAME=openai/gpt-oss-120b
OPENAI_API_KEY=<that backend's key>      # may be blank for Ollama/vLLM

Two deployment consequences worth deciding on deliberately:

• OPENAI_BASE_URL is global. The OpenAI SDK reads that env var for every client it builds, so requests that bring their own X-OpenAI-Key are routed to this backend too — a user's personal OpenAI key would be sent to it. If you host a bring-your-own-key deployment, leave OPENAI_BASE_URL unset.
• Capability floor. The backend must support function calling and JSON-schema structured outputs (strict: true); the orchestrator also uses tool_choice: "required". Backends missing these degrade to fallback paths instead of failing loudly, so run step 4's visualization request after switching backends. Details: Third-party backends.

3A. Path A — build and run the container

docker build -f packages/agent/Dockerfile -t udi-agent .     # from repo root

docker stop udi-agent 2>/dev/null; docker rm udi-agent 2>/dev/null
docker run -d \
  --name udi-agent \
  --restart unless-stopped \
  -p 80:80 \
  --env-file /home/ec2-user/.env \
  udi-agent

The container listens on port 80 as a non-root user. Map it wherever your reverse proxy expects it (-p 8007:80 for a local-only port). Logs go to docker logs udi-agent plus a rotating file at /app/packages/agent/logs/udi_agent.log inside the container.

3B. Path B — run the installed distribution

Serve the app module with uvicorn (installed by the server extra). The fastapi run command used for local development takes a file path into the source tree, which an installed wheel doesn't have; uvicorn takes the import path instead:

cd /opt/udiagent                             # .env and relative paths resolve here
UDI_LOG_DIR=/var/log/udiagent \
  ./venv/bin/uvicorn udiagent.server.app:app --host 0.0.0.0 --port 8007

Two path-related settings matter only on this path, because _PACKAGE_ROOT resolves into site-packages when the code is installed rather than checked out. Both variables — and reading .env from the working directory — landed after 0.2.7, so publish a release containing them before relying on Path B (or use Path A, which builds from source):

• UDI_LOG_DIR — without it the rotating log file is written next to site-packages. If that location isn't writable (a system-wide install run by a non-root service user), file logging is skipped with a one-line notice on stdout and the server still boots.
• UDI_DATA_DIR — only the udiagent package's own data (skills, grammar schema) ships in the wheel; the repo's packages/agent/data/ does not. Without this variable /v1/yac/examples returns 404, which the chat UI tolerates by showing no example-prompt suggestions. Point it at a copy of that directory if you want them (it is ~2 MB, mostly benchmark fixtures). /v1/yac/benchmark_analysis likewise reads ./out/ relative to the working directory and is a benchmarking-only endpoint.

A minimal systemd unit:

[Unit]
Description=UDIAgent server
After=network.target

[Service]
User=udiagent
WorkingDirectory=/opt/udiagent
EnvironmentFile=/opt/udiagent/.env
Environment=UDI_LOG_DIR=/var/log/udiagent
ExecStart=/opt/udiagent/venv/bin/uvicorn udiagent.server.app:app --host 0.0.0.0 --port 8007
Restart=always

[Install]
WantedBy=multi-user.target

EnvironmentFile does not do shell quoting or ${VAR} expansion — keep values plain (JWT_SECRET_KEY=abc123, no surrounding quotes).

To upgrade: pip install -U "udiagent[server]" then restart the unit. Pin the version (udiagent[server]==0.2.7) if you want reproducible redeploys, since this path resolves dependency ranges fresh at install time rather than using the repo's lockfile.

4. Verify

Substitute your port (80 for the container as run above, 8007 for the systemd unit):

curl -fsS http://localhost/                          # {"service":"UDIAgent API","status":"running",...}
curl -fsS http://localhost/v1/yac/structured_functions   # 200 — bundled package data loaded

Then exercise the LLM path end to end, which is the only check that catches a bad key, wrong model id, or a backend that can't do structured outputs:

curl -fsS -X POST http://localhost/v1/yac/completions \
  -H 'Content-Type: application/json' \
  -H "Authorization: Bearer $TOKEN" \
  -d '{
    "messages": [{"role": "user", "content": "bar chart of penguins by species"}],
    "dataSchema": "{\"resources\":[{\"name\":\"penguins\",\"schema\":{\"fields\":[{\"name\":\"species\",\"type\":\"string\"},{\"name\":\"body_mass_g\",\"type\":\"number\"}]}}]}",
    "dataDomains": "[]"
  }'

A healthy response is a JSON body with a non-empty tool_calls. $TOKEN is a JWT signed with JWT_SECRET_KEY (HS256 by default) — mint one with the python-jose that the server extra already installed:

# Path A (inside the container)
TOKEN=$(docker exec udi-agent uv run --frozen --no-sync python -c \
  "import os; from jose import jwt; print(jwt.encode({'sub':'smoke-test'}, os.environ['JWT_SECRET_KEY'], algorithm='HS256'))")

# Path B (in the venv, with JWT_SECRET_KEY exported)
TOKEN=$(/opt/udiagent/venv/bin/python -c \
  "import os; from jose import jwt; print(jwt.encode({'sub':'smoke-test'}, os.environ['JWT_SECRET_KEY'], algorithm='HS256'))")

Note the Authorization header is required on /v1/yac/completions even when INSECURE_DEV_MODE=1 — dev mode skips verifying the token, not sending one. A request without the header returns 422, not 401.

Two different failures both return 401; tell them apart by the body — {"detail":"Invalid or expired token"} is a JWT problem, whereas {"error":"No OpenAI API key..."} means auth passed and the LLM credential is the thing that's missing.

Logs: docker logs udi-agent (Path A) or journalctl -u udiagent (Path B), plus the rotating file described in step 3.

5. Terminate TLS in front of it

The server speaks plain HTTP on both paths. The chat frontend is served over HTTPS (GitHub Pages at /udi-yac/), and browsers block HTTPS pages from calling an http:// API — so an HTTP-only agent is unreachable from the deployed chat, not merely insecure. Put nginx/Caddy/an ALB in front with a certificate and proxy to the port from step 3.

CORS is already allow_origins=["*"] with the X-Usage-* headers exposed, so no per-origin configuration is needed; tighten it in server/app.py if you want to restrict callers.

6. Deploy via GitHub Actions (Path A, the EC2 path)

.github/workflows/deploy-agent.yml runs steps 3A–4 on a self-hosted runner: build, restart the container against /home/ec2-user/.env, health-check, prune old images. Trigger it from the Actions tab (workflow_dispatch).

Prerequisites: the self-hosted runner must be registered to this repo (it was previously registered to hms-dbmi/UDIAgent), and /home/ec2-user/.env must exist on that host with step 2's contents. The workflow deploys whatever branch you dispatch it against.

7. Point the chat at it

In packages/chat/.env.local (or the Pages build environment):

VITE_UDI_API_BASE_URL=https://agent.example.org
VITE_UDI_REQUIRE_API_KEY=false     # true = users supply their own OpenAI key

Set VITE_UDI_REQUIRE_API_KEY=true when the server has no OPENAI_API_KEY and expects per-request X-OpenAI-Key headers.

Architecture

Orchestration Flow

User query
  → Orchestrator.run()
    → GPT with ORCHESTRATOR_TOOLS (5 tools: CreateVisualization, FilterData,
      FreeTextExplain, ClarifyVariable, Rebuff)
    → Dispatch each tool call to its handler
    → Return OrchestratorResult(tool_calls, orchestrator_choice)

Visualization Generation

Executes a two-step markdown skill plan via generate_vis_spec (vis_generate.py):

1. generate — LLM produces a UDI Grammar spec from the request, schema, and few-shot examples
2. validate — JSON schema check with a bounded repair-retry loop

Skills live in src/udiagent/data/skills/*.md (YAML frontmatter + prompt body).

Design Principles

• Stateless — All context travels in message history; no server-side session state
• Skills as Markdown — Prompt templates live in .md files with YAML frontmatter
• Per-request key override — Supports both default and per-request OpenAI API keys

Regenerating Template Visualizations and Tool Definitions

The vis pipeline uses two generated artifacts:

• src/udiagent/data/skills/template_visualizations.json — template visualization specs
• src/udiagent/generated_vis_tools.py — typed OpenAI function-calling tool definitions

To regenerate both in one step:

uv pip install -e ".[codegen]"
uv run python scripts/regenerate_vis_tools.py

By default this uses data/data_domains/hubmap_data_schema.json as the schema. To use a different schema:

uv run python scripts/regenerate_vis_tools.py --schema data/data_domains/SenNet_domains.json

Benchmarking

Step 0: Start the API server

uv run --extra server fastapi dev src/udiagent/server/app.py --port 8007 &

Step 1: Run tiny benchmark (1 example)

uv run python -m udiagent.benchmark.runner --no-orchestrator --path ./data/benchmark_dqvis/tiny.jsonl

Step 2: Run small benchmark (100 examples)

uv run python -m udiagent.benchmark.runner --no-orchestrator --path ./data/benchmark_dqvis/small.jsonl --workers 5

Resume a failed run:

uv run python -m udiagent.benchmark.runner --path ./data/benchmark_dqvis/small.jsonl --workers 5 --resume ./out/<TIMESTAMP>/benchmark_results.json

License

MIT