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 Amazon Bedrock IAM-role (SigV4) authentication
pip install udiagent[bedrock]

# 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 --extra bedrock   # 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	Static Bedrock API key; for an IAM role see below
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 IAM role (SigV4)

To authenticate with an EC2 instance profile or ECS task role instead of a static key, install udiagent[bedrock] and pass bedrock=True:

agent = UDIAgent(
    gpt_model_name="openai.gpt-oss-120b-1:0",
    bedrock=True,
    bedrock_region="us-east-1",   # or set AWS_REGION
)

Requests are signed with SigV4 using the default AWS credential chain — instance profile, task role, AWS_PROFILE, AWS_ACCESS_KEY_ID, ~/.aws/credentials, SSO — so no key is held by the process. Either parameter opts in; bedrock_region alone is enough.

• Mutually exclusive with openai_api_key and openai_base_url (the OpenAI SDK refuses provider= alongside either). Passing both raises ValueError; on the server, the equivalent env combination refuses startup. For a static Bedrock API key, use openai_base_url as in the table above.
• A region is required and determines the endpoint. It resolves from bedrock_region, then AWS_REGION, AWS_DEFAULT_REGION, ~/.aws/config — there is no instance-metadata fallback. Missing it raises OpenAIError at construction, so the server fails to start rather than failing the first request.
• To reach a PrivateLink/VPC endpoint, set AWS_BEDROCK_BASE_URL — not OPENAI_BASE_URL.
• AWS_BEARER_TOKEN_BEDROCK is deliberately ignored. bedrock=True always means SigV4; a stray bearer token in the environment cannot silently change the auth mode.
• Per-request X-OpenAI-Key keys are refused with a 403. Serving one would build a client against api.openai.com and send the prompt — including the data schema and domains — to a third party, which defeats the reason for routing inference through Bedrock in the first place. The guard runs before orchestration on every endpoint that accepts the header (/v1/yac/completions, /v1/yac/benchmark), and _get_gpt_client enforces the same rule as a backstop for direct library use. Set VITE_UDI_REQUIRE_API_KEY=false in the frontend so users are not prompted for a key the server will reject.
• LangFuse cost attribution reads zero for Bedrock model ids, which are not in LangFuse's pricing table. Traces, latencies, and token counts still work.
• Unresolvable AWS credentials return a 503, not a 500. The SDK resolves credentials per request, so a detached instance profile or an unreachable IMDS fails at request time rather than at startup. The response stays terse because an end user cannot act on it; the operator-facing checklist (role, AWS_REGION, IMDS hop limit) is logged at ERROR. Check the server log first when a Bedrock deployment starts returning 503.

Deployment prerequisites are listed under Deployment.

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	Default	Description
JWT_SECRET_KEY	—	JWT signing key for self-issued tokens. Required unless INSECURE_DEV_MODE=1 or JWT_JWKS_URL is set.
JWT_ALGORITHM	HS256	JWT algorithm. Set to the identity provider's (e.g. RS256) when using JWT_JWKS_URL.
JWT_JWKS_URL	—	Verify externally issued tokens against an identity provider's JWKS endpoint instead of JWT_SECRET_KEY. Mutually exclusive with it.
JWT_ISSUER	—	Expected iss claim; validated only when set.
JWT_AUDIENCE	—	Expected aud claim. Required with JWT_JWKS_URL.
INSECURE_DEV_MODE	0	Skip JWT verification entirely. Development only — never set this in production.
UDI_CORS_ORIGINS	*	Comma-separated list of browser origins allowed to call this server. Defaults to * (any origin). Note that credentialed requests are always permitted, so with * any site can call this server with a user's cookies — name your origins explicitly in production.
GPT_MODEL_NAME	gpt-5.4	Model for orchestration; with a custom base URL, that backend's model id. Callers may override it per-request only by supplying their own X-OpenAI-Key.
OPENAI_API_KEY	—	OpenAI API key. If unset, callers must supply one per-request via the X-OpenAI-Key header.
OPENAI_BASE_URL	—	Root of any OpenAI-compatible backend (Azure AI Foundry, Bedrock, OpenRouter, Ollama, vLLM). Must support function calling and JSON-schema structured outputs. For Bedrock this is the static-API-key path; see UDI_BEDROCK for an IAM role.
UDI_BEDROCK	0	Authenticate to Amazon Bedrock with SigV4 from the default AWS credential chain (EC2 instance profile, ECS task role, AWS_PROFILE, ~/.aws/credentials), so no key is stored. Mutually exclusive with OPENAI_API_KEY and OPENAI_BASE_URL. Per-request X-OpenAI-Key headers are refused with a 403 in this mode, so no prompt reaches api.openai.com. Requires the udiagent[bedrock] extra.
AWS_REGION	—	AWS region for Bedrock; also determines the endpoint. Required with UDI_BEDROCK, though the AWS SDK will also accept AWS_DEFAULT_REGION or ~/.aws/config — there is no instance metadata fallback. Set AWS_BEDROCK_BASE_URL instead to reach a PrivateLink endpoint.
LANGFUSE_PUBLIC_KEY	—	LangFuse public key. Tracing is off unless all three are set.
LANGFUSE_SECRET_KEY	—	LangFuse secret key. Tracing is off unless all three are set.
LANGFUSE_HOST	—	LangFuse instance URL, e.g. https://cloud.langfuse.com.
LANGFUSE_ENVIRONMENT	—	Tags traces with an environment label (e.g. production). Does not by itself enable tracing.
UDI_QUERY_BACKENDS	—	Path to a JSON file mapping package names to StarRocks/DuckDB connections, served via /v1/yac/query and /v1/yac/metadata. Written by the seed scripts — see dev/duckdb/README.md. A relative path is tried against the working directory first, then against packages/agent, where the seed scripts write it.
UDI_METADATA_TTL_SECONDS	3600	TTL for the introspected-metadata cache.
UDI_LOG_DIR	—	Where the rotating log file goes. Defaults to <package root>/logs; file logging is skipped if unwritable.
UDI_DATA_DIR	—	Repo-level dev data for /v1/yac/examples. Defaults to <package root>/data; 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.

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.

Getting the token from the portal to here

The server only ever sees an Authorization: Bearer header; it does not care how that header was produced. Two topologies produce it, and which one applies depends on where the host portal keeps its token.

The host page holds the token. It passes it to the chat as authToken (see UDIChatConfig) and the browser calls this server directly. UDI_CORS_ORIGINS must name the portal's origin, since those are cross-origin requests.

The host's backend holds the token. Portals that keep the IdP token server-side — never in a cookie or local storage the page can read — have nothing to hand the chat. Instead the portal proxies our endpoints from its own backend and attaches the header there:

browser → https://portal.example/api/yac/*  (portal backend attaches the
                                             Authorization header)
        → https://agent.internal/v1/yac/*

Point the chat's apiBaseUrl at that path (/api/yac) and leave authToken unset — the chat then sends no Authorization header of its own, and the proxy supplies it. The browser never talks to this server, so UDI_CORS_ORIGINS is irrelevant to the request path and the agent can stay on a private network. That is the arrangement for the Radiant portal (see #118).

The proxy route must forward the request body, the X-Conversation-Id and X-OpenAI-Key request headers, and the X-Usage-* response headers.

Either way, a verified token means the user is allowed to use this server. By default it does not scope query results: server-side query backends connect with the service credentials in UDI_QUERY_BACKENDS. A StarRocks backend with jwtPassthrough is the exception — it forwards the caller's own token to the database, which then applies that user's grants (see query/README.md). That needs a real token, so it rules out the proxy topology above unless the proxy attaches a token the database can verify too.

Future: role-based model permissions

Today, permission to choose a model is tied to who pays for it: /v1/yac/completions honors a request's model only when the caller also supplies an X-OpenAI-Key. Otherwise the server's GPT_MODEL_NAME applies and the requested model is ignored (and logged).

That's a proxy for the real question, which is whether this user is allowed to pick a model. Deployments that already authenticate users through an identity provider have a better signal available: verify_jwt returns the decoded token, so a claim on it (realm_access.roles, a group, a scope) could authorize model selection independently of key ownership — letting a trusted role pick a model on the server's key, while everyone else stays on the default. That would replace the if x_openai_key gate in yac_completions; nothing else in the request path would change.

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 + bedrock 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.

(Optional) Amazon Bedrock with the instance's IAM role

To drop the static key entirely and authenticate with the EC2 instance profile or ECS task role:

UDI_BEDROCK=1
AWS_REGION=us-east-1
GPT_MODEL_NAME=openai.gpt-oss-120b-1:0
# Leave OPENAI_API_KEY and OPENAI_BASE_URL unset — the server refuses to start
# if either is set alongside UDI_BEDROCK.

Callers cannot opt out of Bedrock: a request carrying X-OpenAI-Key is rejected with a 403 rather than served from api.openai.com, so no prompt or data schema leaves your AWS account. Pair this with VITE_UDI_REQUIRE_API_KEY=false in the chat frontend.

Behavior and constraints: Bedrock with an IAM role. Three infrastructure preconditions have to hold, and none of them fail loudly in an obvious way:

1. The role needs Bedrock permissions. The instance profile in cloudformation/udi-agent.yaml (EC2Role) currently carries only AmazonSSMManagedInstanceCore. Add bedrock:InvokeModel and bedrock:InvokeModelWithResponseStream. Models reachable only through a cross-region inference profile also need the action on the profile ARN and on the underlying foundation-model ARNs in every region the profile spans.

2. The IMDS hop limit must be ≥ 2 for a bridged container. The launch template sets no MetadataOptions, so the AWS default HttpPutResponseHopLimit: 1 applies — Docker's bridge network adds a hop, the IMDSv2 token response never reaches the container, and credential resolution fails. Either add MetadataOptions with HttpPutResponseHopLimit: 2 (which affects only newly launched instances, so follow with an ASG instance refresh) or fix a running instance in place:

   aws ec2 modify-instance-metadata-options --instance-id <id> \
     --http-put-response-hop-limit 2 --http-tokens required
   

3. AWS_REGION must be in the env file. There is no instance-metadata fallback for the region. Put it in /home/ec2-user/.env rather than passing -e AWS_REGION=... to docker run — Docker applies -e after --env-file regardless of order, so a flag would override the file and remove the operator's ability to invoke cross-region.

If any of the three is wrong, requests return 503 with a generic message and the server log carries the specific cause at ERROR — start there rather than from the client response.

A local run with AWS_PROFILE set exercises the identical code path (same default credential chain), so the setup can be validated before deploying.

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. The chat's full variable list lives in packages/chat/.env.example (generated from packages/chat/src/app/envVars.ts).

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