Library usage

Use udiagent directly from Python: the agent, the orchestrator and the vis pipeline take all their configuration as constructor parameters, with no environment variables. The Python API reference lists every public class and function.

from udiagent import UDIAgent, Orchestrator

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

# Create an orchestrator
orch = Orchestrator(agent)

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

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

One agent, any schema

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

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

Visualization template sets and tags

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

Data-shape tag	Used for	Selected when
`line_item`	Tidy, per-record tables (groupby/rollup)	the schema has a resource that is not a cube (or no resources at all)
`data_cube`	Pre-aggregated “powerset” cubes (marginals)	the schema marks a resource `udi:cube` (or declares `udi:dimensions` + `udi:measures`)

A package can hold both kinds of table and then gets both sets. Each template binds only its own kind: validate_bindings rejects a cube template bound to a line-level table and a line-item template bound to a cube, naming the tables that would fit, so the model’s retry can correct it. A template’s kind is read off its text: cube templates use <M> or <MARGINAL…>, and no other template does.

Some charts can be drawn from either kind of table — an event-free survival curve from the line-level event log or from a survival cube. A mixed package can say which to default to with a package-level "udi:prefer": "line_level" (or "data_cube"). The preferred kind’s templates are then listed first, and the prompt tells the model to answer from them whenever they can, keeping the other kind for what only it has. It is a default, not a filter: both sets stay on offer. Ignored when the package holds only one kind of table.

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

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

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

python scripts/regenerate_vis_tools.py

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

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.

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:

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

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

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

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

Amazon Bedrock with an IAM role (SigV4)

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

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

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

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

Deployment prerequisites are listed in the deployment guide.

With LangFuse observability

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

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

Tracing turns on when any of langfuse_public_key, langfuse_secret_key, or langfuse_host is set. langfuse_environment is purely a tag — it labels traces in the LangFuse UI but does not by itself enable tracing. The server is stricter: it refuses to start unless all three LANGFUSE_* variables are set, or none.

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.

Serving the chat from your own backend

The chat calls a handful of HTTP endpoints. A host that serves them from its own web framework (Flask, Django) needs only plain udiagent, not the [server] extra. udiagent.api holds the logic a host would otherwise copy from the reference server: each function takes the parsed JSON body, returns the JSON response body, and raises EndpointError with the status and {code, error, errors} body the chat reads. The reference FastAPI server calls the same functions, so a host answers exactly as it does. The request and response models are in udiagent.models.

Endpoint	Call
`POST /v1/yac/completions`	`completions(orchestrator, body, openai_api_key=, session_id=, budget_check=)`
`POST /v1/yac/vis_instantiate`	`vis_instantiate(body)`
`GET /v1/yac/metadata`	`metadata(engines, caches, package, refresh=, token=, principal=)`
`POST /v1/yac/query`	`query(engines, body, token=)`
`POST /v1/yac/issue_report`	`issue_report(agent, raw_body, session_id=, webhook_url=, defer=)`
`GET /v1/yac/structured_functions`	`udiagent.export_registry_json()`
`GET /v1/yac/examples`	Your own list of prompts about your data.

Authentication, and which callers may use your OpenAI key, stay with your host. So do the query backends: engines maps a package name to a udiagent.query.QueryEngine (see the query README), and caches is an empty dict you keep between requests. You need them only for server-side data; the chat’s default browser mode reads CSVs itself.

from flask import jsonify, request
from udiagent.api import EndpointError, vis_instantiate

@app.post("/v1/yac/vis_instantiate")
def yac_vis_instantiate():
    try:
        return jsonify(vis_instantiate(request.get_json(silent=True) or {}))
    except EndpointError as e:
        return jsonify(e.body()), e.status

completions runs a conversation turn with the reference server’s policy: a requested model is honored only with the caller’s own key, and a budget_check you pass runs only on your key. A quota refusal comes back as a Rebuff, and the X-Usage-* headers come back beside the body. A cross-origin caller can read those headers only if your CORS layer lists them in Access-Control-Expose-Headers. Who may use your OpenAI key stays your decision: pass the caller’s key, or None to run on the agent’s.

from udiagent.api import completions

@app.post("/v1/yac/completions")
def yac_completions():
    try:
        tool_calls, headers = completions(
            orchestrator,
            request.get_json(silent=True) or {},
            openai_api_key=request.headers.get("X-OpenAI-Key"),
            session_id=request.headers.get("X-Conversation-Id"),
        )
    except EndpointError as e:
        return jsonify(e.body()), e.status
    return jsonify(tool_calls), 200, headers

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