Metadata-Version: 2.5
Name: udiagent
Version: 0.8.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>=3.33.0
Requires-Dist: python-dotenv>=1.0.0
Requires-Dist: requests>=2.32.5
Provides-Extra: all
Requires-Dist: datasets>=5.0.1; extra == 'all'
Requires-Dist: duckdb>=1.0.0; extra == 'all'
Requires-Dist: fastapi[standard]>=0.143.0; extra == 'all'
Requires-Dist: huggingface-hub>=2.2.0; extra == 'all'
Requires-Dist: ijson>=3.4.0.post0; extra == 'all'
Requires-Dist: langfuse>=4; extra == 'all'
Requires-Dist: openai[bedrock]; 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]; extra == 'bedrock'
Provides-Extra: benchmark
Requires-Dist: datasets>=5.0.1; extra == 'benchmark'
Requires-Dist: huggingface-hub>=2.2.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.143.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).

### Reasoning models and the Responses API

Five constructor parameters shape LLM requests. Their defaults send exactly
what earlier releases sent, so nothing changes unless you set them. The server
reads them from `UDI_OPENAI_API`, `UDI_REASONING_EFFORT`, `UDI_TEMPERATURE`
(`omit` for `None`), `UDI_MAX_COMPLETION_TOKENS` and `UDI_VIS_TOOL_CHOICE`.

```python
agent = UDIAgent(
    gpt_model_name="gpt-6.1-sol",
    openai_api_key="sk-...",
    openai_api="responses",       # gpt-6.1-sol only calls tools through Responses
    reasoning_effort="low",       # passed through verbatim; None sends nothing
    temperature=None,             # OpenAI rejects temperature when reasoning is on
    max_completion_tokens=16384,  # reasoning tokens count against the cap
    vis_tool_choice="required",   # pick a chart template instead of replying in prose
)
```

The `"responses"` path translates each request and reply, so the orchestrator
and visualization code are unaware of which API ran. Keep the default
`"chat_completions"` for self-hosted backends (vLLM, Ollama) and Bedrock. When
a call stops at the output cap, the server logs a warning naming the model; if
that warning recurs, raise `max_completion_tokens`.

`vis_tool_choice` applies only to the call that picks a chart template. With
the default `"auto"` a model may answer without picking one, and the user gets
a "no chart template fits" message; `gpt-6.1-sol` does this for about one chart
request in seven. `"required"` makes it pick a template.

Usage reporting includes `cache_write_tokens`: models that bill prompt-cache
writes (GPT-5.6 and later, at 1.25× the input rate) report them, and the server
returns them as `X-Usage-Cache-Write-Tokens`.

### 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.                                                                                                                          |
| `UDI_OPENAI_API`               | `chat_completions` | Which OpenAI-compatible endpoint to call: `chat_completions` or `responses`. Some OpenAI models (e.g. `gpt-6.1-sol`) only call tools through `responses`; self-hosted backends (vLLM, Ollama) generally want the default.                                                                                                                                                                             |
| `UDI_REASONING_EFFORT`         | —                  | Reasoning effort sent with every request (`none`, `low`, `medium`, …), passed through as-is so new provider values work without a release. Unset sends nothing, leaving the model's default.                                                                                                                                                                                                          |
| `UDI_TEMPERATURE`              | `0`                | Sampling temperature sent with every request, or `omit` to send none. OpenAI rejects a temperature whenever `UDI_REASONING_EFFORT` is anything but `none`, so set `omit` with reasoning models there.                                                                                                                                                                                                 |
| `UDI_MAX_COMPLETION_TOKENS`    | —                  | Output-token cap for every request, replacing the built-in ones (1024 for tool calls, 16384 for structured JSON). Reasoning tokens count against it, so raise it when reasoning is on.                                                                                                                                                                                                                |
| `UDI_VIS_TOOL_CHOICE`          | `auto`             | Whether the chart-template call may answer without picking a template (`auto`, which then replies that no template fits) or must pick one (`required`). Some models (e.g. `gpt-6.1-sol`) reply in prose to many chart requests unless this is `required`.                                                                                                                                             |
| `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_BASE_URL` is accepted too.                                                                                                                                                                                                                                                                                                        |
| `LANGFUSE_ENVIRONMENT`         | —                  | Tags traces with an environment label (e.g. `production`). Does not by itself enable tracing.                                                                                                                                                                                                                                                                                                         |
| `UDI_ISSUE_REPORT_WEBHOOK_URL` | —                  | Notified of every issue report (`POST /v1/yac/issue_report`): a JSON POST with a Slack-style `text` summary plus the report's note, source, contact and LangFuse trace link — never the session. Works as-is with Slack and Teams incoming webhooks; route it through Zapier or Power Automate for email. Needs LangFuse, since reports are refused without it.                                       |
| `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/issue_report`         | POST   | Record a user's issue report in LangFuse (503 without LangFuse) |
| `/v1/yac/benchmark`            | POST   | Benchmark variant; can skip routing for chart-only items        |
| `/v1/yac/examples`             | GET    | Example prompts from `data/example_prompts.json`                |
| `/v1/yac/structured_functions` | GET    | Structured function registry                                    |
| `/v1/yac/benchmark_analysis`   | GET    | Newest benchmark analysis under `out/` (auth required)          |

### 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:8000 \
  --env-file /home/ec2-user/.env \
  udi-agent
```

The container listens on port 8000 as a non-root user — an unprivileged port,
so it binds on runtimes that don't let non-root processes take ports below 1024.
Map it wherever your reverse proxy expects it (`-p 8007:8000`
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 serves the newest
  `out/**/benchmark_analysis.json` under the package directory, which a wheel
  install doesn't have; it 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

The runner posts each benchmark item to a running server's `/v1/yac/benchmark`
endpoint, so **the server's configuration is what gets measured** — model, API
path, reasoning effort, temperature and output cap all come from its
environment (see [Reasoning models and the Responses API](#reasoning-models-and-the-responses-api)).
Restart the server to benchmark another configuration.

### Step 0: Start the API server

Use `fastapi run`, not `dev`: the reloader can restart the server mid-run. The
runner sends a placeholder bearer token, so set `INSECURE_DEV_MODE=1`.

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

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

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

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

```bash
uv run python -m udiagent.benchmark.runner --path ./data/benchmark_dqvis/small.jsonl --workers 5 \
  --description "what this run is" --out ./out/small-gpt-5.4
```

Resume a run; items that failed or were refused (rate limit) are retried:

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

`--no-orchestrator` skips the routing call and goes straight to chart
generation, to measure that on its own. Use it only on visualization-only data
such as the DQVis fixtures: the server refuses items that expect any other
orchestrator choice, and those items are recorded as failures.

Each run writes `benchmark_results.json` and `benchmark_analysis.json`. The
analysis carries a `summary` — failures, tokens per item (prompt, cached, cache
write, completion, reasoning), cache-hit percentage, LLM calls per item, cost
(from `PRICES` in `benchmark/runner.py`; check them against OpenAI's pricing
page), latency per item (mean, p50, p95) and per LLM operation, and how often
the model asked a clarifying question — and its `metadata` records the model
and settings the server reported.

A clarifying question scores as a correct answer when the fields it offers
include one the expected answer used: the request really was ambiguous, and in
a blind review people rated those questions good. `clarify_percent` is the
share of items answered with a question and `clarify_credited_percent` the
share credited this way.

### Comparing models

The `*_vs_expected` scores compare against the fixtures' expected outputs, which
an older generator produced; read them relatively, not as accuracy. To compare
models against each other instead, run every configuration over the same items
— ideally 3+ repeats each — and compare with one as the baseline:

```bash
# Keep UDI_* settings out of .env: inline values win over it, but a setting
# left there leaks into every configuration that doesn't set its own.
run_config() {  # <label> [ENV=VALUE ...]
  local label=$1; shift
  env INSECURE_DEV_MODE=1 "$@" uv run --extra server fastapi run src/udiagent/server/app.py --port 8007 &
  local pid=$!
  until curl -sf -o /dev/null http://127.0.0.1:8007/v1/yac/examples; do sleep 1; done
  for p in data/benchmark_dqvis/small.jsonl data/benchmark.json; do
    uv run python -m udiagent.benchmark.runner -p "$p" --label "$label" --description "$r" \
      --out "out/bench/$label/$r/$(basename "${p%%.*}")"
  done
  kill $pid; wait $pid
}
for r in r0 r1 r2; do  # vary the order between repeats to spread out API-side load
  run_config gpt-5.4 GPT_MODEL_NAME=gpt-5.4
  run_config gpt-6-luna GPT_MODEL_NAME=gpt-6-luna UDI_REASONING_EFFORT=none
  run_config gpt-6.1-sol GPT_MODEL_NAME=gpt-6.1-sol UDI_OPENAI_API=responses \
    UDI_REASONING_EFFORT=low UDI_TEMPERATURE=omit UDI_MAX_COMPLETION_TOKENS=16384 \
    UDI_VIS_TOOL_CHOICE=required
done
uv run python -m udiagent.benchmark.runner -a compare --baseline gpt-5.4 \
  out/bench/*/*/*/benchmark_analysis.json > out/bench/compare.md
```

`compare` prints, per dataset, every run summary as mean ± sd across repeats,
then an **agreement** table: each configuration's outputs scored item by item
against the baseline's, with the baseline's output standing in for "expected"
in the usual rubric (orchestrator choice, filter, chart spec), plus how often
both picked the same visualization template. A clarifying question about a
field the baseline used counts as agreeing. The baseline's own column scores
its repeats against each other — the noise floor. Models are not deterministic
even at temperature 0, so read a candidate's agreement against that floor, not
against 1.0.

## License

MIT
