Metadata-Version: 2.5
Name: udiagent
Version: 0.5.0
Summary: UDI Agent: LLM-powered data visualization orchestration library
License: MIT
License-File: LICENSE
Requires-Python: >=3.12
Requires-Dist: jsonschema>=4.26.0
Requires-Dist: openai>=2.15.0
Requires-Dist: python-dotenv>=1.0.0
Requires-Dist: requests>=2.32.5
Provides-Extra: all
Requires-Dist: datasets>=4.5.0; extra == 'all'
Requires-Dist: duckdb>=1.0.0; extra == 'all'
Requires-Dist: fastapi[standard]>=0.128.0; extra == 'all'
Requires-Dist: huggingface-hub>=0.36.0; extra == 'all'
Requires-Dist: ijson>=3.4.0.post0; extra == 'all'
Requires-Dist: langfuse>=4; extra == 'all'
Requires-Dist: openai[bedrock]>=2.44.0; extra == 'all'
Requires-Dist: pandas>=2.0.0; extra == 'all'
Requires-Dist: pydantic-settings>=2.0; extra == 'all'
Requires-Dist: pymysql>=1.1.0; extra == 'all'
Requires-Dist: python-dotenv>=1.0.0; extra == 'all'
Requires-Dist: python-jose[cryptography]>=3.5.0; extra == 'all'
Requires-Dist: udi-grammar-py; extra == 'all'
Provides-Extra: bedrock
Requires-Dist: openai[bedrock]>=2.44.0; extra == 'bedrock'
Provides-Extra: benchmark
Requires-Dist: datasets>=4.5.0; extra == 'benchmark'
Requires-Dist: huggingface-hub>=0.36.0; extra == 'benchmark'
Requires-Dist: ijson>=3.4.0.post0; extra == 'benchmark'
Provides-Extra: codegen
Requires-Dist: pandas>=2.0.0; extra == 'codegen'
Requires-Dist: udi-grammar-py; extra == 'codegen'
Provides-Extra: duckdb
Requires-Dist: duckdb>=1.0.0; extra == 'duckdb'
Provides-Extra: langfuse
Requires-Dist: langfuse>=4; extra == 'langfuse'
Provides-Extra: server
Requires-Dist: fastapi[standard]>=0.128.0; extra == 'server'
Requires-Dist: pydantic-settings>=2.0; extra == 'server'
Requires-Dist: python-dotenv>=1.0.0; extra == 'server'
Requires-Dist: python-jose[cryptography]>=3.5.0; extra == 'server'
Provides-Extra: starrocks
Requires-Dist: pymysql>=1.1.0; extra == 'starrocks'
Provides-Extra: test
Requires-Dist: cryptography>=42.0.0; extra == 'test'
Requires-Dist: duckdb>=1.0.0; extra == 'test'
Requires-Dist: pytest>=8.0.0; extra == 'test'
Description-Content-Type: text/markdown

# 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

```bash
# 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`:

```bash
uv sync --extra server --extra langfuse --extra test --extra bedrock   # server + dev
```

## Library Usage

```python
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:

```python
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:

```jsonc
{
  "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:

```bash
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:

```python
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`:

```python
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](#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`:

```python
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`).

```bash
# 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

<!-- BEGIN generated: env vars (scripts/gen_env_docs.py) -->

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

<!-- END generated: env vars -->

### 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:

```ini
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](https://github.com/hms-dbmi/udi-yac/issues/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`](src/udiagent/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`](src/udiagent/query/README.md).**
Local dev instance: [`dev/starrocks/README.md`](../../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.

```bash
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](#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`](../../.github/workflows/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:**

```bash
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:**

```bash
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](#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.

```bash
# 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](https://github.com/hms-dbmi/udi-yac/blob/main/packages/agent/.env.template)
or just write the values below by hand.

Minimum production values:

```ini
# 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](#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:

```ini
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](#third-party-openai-compatible-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:

```ini
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](#amazon-bedrock-with-an-iam-role-sigv4). 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:

   ```bash
   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

```bash
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:

```bash
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:

```ini
[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):

```bash
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:

```bash
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:

```bash
# 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`](src/udiagent/server/app.py) if you want to restrict callers.

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

[`.github/workflows/deploy-agent.yml`](../../.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):

```ini
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`](../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:

```bash
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:

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

## Benchmarking

### Step 0: Start the API server

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

### Step 1: Run tiny benchmark (1 example)

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

### Step 2: Run small benchmark (100 examples)

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

Resume a failed run:

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

## License

MIT
