Python API reference

Generated page

Built from the docstrings in packages/agent/src/udiagent on every docs build. For how the pieces fit together, see How the agent works.

udiagent

UDIAgent — LLM-powered data visualization orchestration library.

Classes:

• BudgetExceededError – Raised when an LLM call is refused for quota/rate-limit reasons.
• Orchestrator – Routes user requests to visualization, filter, explanation, and other tool handlers.
• OrchestratorResult – Result of an orchestration run.
• Skill – A skill loaded from a markdown file.
• UDIAgent – UDIAgent for requesting UDI grammar via OpenAI.
• Usage – Aggregated and per-operation token usage for a single Orchestrator.run().

Functions:

• build_rebuff_toolcall – Build a Rebuff tool_call dict without invoking an LLM.
• export_registry_json – Export the function registry as a JSON-serializable dict.
• generate_vis_spec – Generate a visualization spec using the skills pipeline.
• get_function_signatures – Return the function registry as a list of signature descriptions.
• load_grammar – Load a grammar definition by name.
• load_skills – Load all skill .md files from a directory.
• normalize_tool_calls – Reformat tool_calls into valid OpenAI function-calling format and
• parse_schema_from_dict – Parse a UDI data schema dict into a structured representation.
• render_template – Replace {{key}} placeholders in instructions with values from variables.
• segment_structured_text – Segment text into a mixed list of plain strings and structured element objects.
• simplify_data_domains – Simplify data domains for better LLM consumption.
• simplify_data_schema – Simplify the data schema for better LLM consumption.
• split_tool_calls – Split messages with multiple tool calls into separate messages.
• validate_structured_text – Validate that all function references in the text use registered functions.

udiagent.BudgetExceededError

BudgetExceededError(message: str, usage: Usage | None = None)

Bases: Exception

Raised when an LLM call is refused for quota/rate-limit reasons.

Carries the accumulated Usage at the moment of failure so callers can still log / bill for tokens spent before the refusal.

udiagent.Orchestrator

Orchestrator(agent, skills: dict[str, Skill] | None = None, tools: list[dict] | None = None, grammar: dict | None = None)

Routes user requests to visualization, filter, explanation, and other tool handlers.

Parameters:

• agent – A UDIAgent instance.
• skills (dict[str, Skill] | None) – Skill registry. Loaded from bundled data if None.
• tools (list[dict] | None) – Tool definitions list. Defaults to ORCHESTRATOR_TOOLS.
• grammar (dict | None) – Preloaded grammar dict. Auto-loaded if None.

Functions:

• run – Run the orchestrator on a user request.
• run_visualization – Build a chart for the latest user request, skipping the routing call.

udiagent.Orchestrator.run

run(messages: list[dict], data_schema: str, data_domains: str, openai_api_key: str | None = None, model: str | None = None, budget_check: Callable[[Usage], str | None] | None = None, session_id: str | None = None, filters: list[dict] | None = None, issue_reporting: bool = False) -> OrchestratorResult

Run the orchestrator on a user request.

Parameters:

• messages (list[dict]) – Chat history (OpenAI message format).
• data_schema (str) – JSON string describing dataset entities and fields.
• data_domains (str) – JSON string (or list) describing field domains.
• openai_api_key (str | None) – Optional per-request OpenAI key override.
• budget_check (Callable[[Usage], str | None] | None) – Optional callback invoked with the current Usage. Return a non-empty string to short-circuit the run and produce a Rebuff tool_call carrying that message; return None to proceed. Invoked pre-flight and again after the top-level orchestration completion.
• session_id (str | None) – Optional per-conversation ID used to group this run’s LLM calls into a single LangFuse trace/session. No-op when LangFuse is disabled.
• filters (list[dict] | None) – The active filters as the client holds them, edits and all (see udiagent.filters). None means the caller did not say, which the prompt reports as no active filters.
• issue_reporting (bool) – Offer the OfferIssueReport tool, which lets the client open its report form. Pass it only when the client can send reports and something records them.

Returns:

• OrchestratorResult – An OrchestratorResult with tool_calls, orchestrator_choice, and usage.

udiagent.Orchestrator.run_visualization

run_visualization(messages: list[dict], data_schema: str, data_domains: str, openai_api_key: str | None = None, model: str | None = None, session_id: str | None = None) -> OrchestratorResult

Build a chart for the latest user request, skipping the routing call.

For callers that already know the request wants a visualization, such as visualization-only benchmarks measuring chart generation on its own. The latest non-blank user message stands in for the description the routing call would otherwise write.

udiagent.OrchestratorResult

OrchestratorResult(tool_calls: list[dict] = list(), orchestrator_choice: str = 'render-visualization', usage: Usage = Usage()) -> None

Result of an orchestration run.

udiagent.Skill

Skill(name: str, description: str, instructions: str) -> None

A skill loaded from a markdown file.

udiagent.UDIAgent

UDIAgent(gpt_model_name: str, openai_api_key: str | None = None, *, openai_base_url: str | None = None, bedrock: bool = False, bedrock_region: str | None = None, langfuse_public_key: str | None = None, langfuse_secret_key: str | None = None, langfuse_host: str | None = None, langfuse_environment: str | None = None, openai_api: str = 'chat_completions', reasoning_effort: str | None = None, temperature: float | None = 0.0, max_completion_tokens: int | None = None, vis_tool_choice: str = 'auto')

UDIAgent for requesting UDI grammar via OpenAI.

Any OpenAI-compatible backend works: point openai_base_url at its chat-completions root (Azure AI Foundry, Amazon Bedrock, OpenRouter, vLLM, Ollama, LiteLLM, …) and set gpt_model_name to that backend’s model id. The backend must support function calling and JSON-schema structured outputs — both are load-bearing for spec generation.

openai_base_url applies to the default (server-configured) client only; per-request keys still resolve the SDK default, unless the OPENAI_BASE_URL environment variable is set, which the OpenAI SDK applies to every client it builds.

Amazon Bedrock with an IAM role: pass bedrock=True (optionally bedrock_region) instead of openai_api_key. Requests are signed with SigV4 using the default AWS credential chain — an EC2 instance profile or ECS task role — and the endpoint is derived from the region, so no key is held by the process. Requires pip install udiagent[bedrock]. Mutually exclusive with openai_api_key and openai_base_url; for a static Bedrock API key use those instead. Per-request (bring-your-own) keys are refused in this mode rather than served from api.openai.com — see _get_gpt_client.

LangFuse observability is opt-in: pass any of langfuse_public_key, langfuse_secret_key, or langfuse_host to route requests through langfuse.openai.OpenAI. When none are provided, the plain openai client is used and no traces are emitted (even if the langfuse package is installed).

Request shaping applies to every LLM call (see create_completion). The defaults send exactly what this package always sent, so a backend only sees a difference when a deployment opts in:

• openai_api: "chat_completions" (default) or "responses". Some models (e.g. gpt-6.1-sol) only call tools through Responses; self-hosted backends generally want the default.
• reasoning_effort: sent verbatim when set — not validated, so values a provider adds later work without a release. None omits it.
• temperature: 0.0 by default; None omits it, which OpenAI reasoning models require whenever reasoning effort isn’t "none".
• max_completion_tokens: when set, replaces each call’s own output cap (1024 for tool calls, 16384 for structured JSON). Reasoning tokens count against the cap, so raise it when reasoning is on.
• vis_tool_choice: "auto" (default) lets the chart-template call answer without picking a template, which falls back to a “no template fits” message; "required" makes it pick one. Some models (e.g. gpt-6.1-sol) otherwise reply in prose to many chart requests.

Functions:

• create_completion – Make one LLM call, shaped by this agent’s request settings.
• langfuse_urls – (trace_url, session_url) in the LangFuse UI.
• record_issue_report – Record a user’s issue report as an issue-report trace.
• trace – Group every OpenAI call made within this block under one trace.

udiagent.UDIAgent.create_completion

create_completion(client, **kw)

Make one LLM call, shaped by this agent’s request settings.

Takes chat.completions.create kwargs and returns a chat-completion shaped result whichever API is configured, so callers read choices[0].message and usage the same way on both paths.

udiagent.UDIAgent.langfuse_enabled

langfuse_enabled: bool

Whether traces, and so issue reports, are recorded anywhere.

udiagent.UDIAgent.langfuse_urls

langfuse_urls(trace_id: str, session_id: str | None = None) -> tuple[str | None, str | None]

(trace_url, session_url) in the LangFuse UI.

Either is None when there is nothing to link to: no LangFuse, or no session ID. The first call looks the project ID up over the network.

udiagent.UDIAgent.record_issue_report

record_issue_report(*, description: str, contact: str | None, source: str, session: dict, session_id: str | None = None) -> str | None

Record a user’s issue report as an issue-report trace.

session_id (the conversation ID) files it in the same LangFuse session as the turns being reported. The session export goes up as a JSON attachment rather than inline, so a long conversation neither bloats the trace nor hits the ingestion size limit.

Returns the trace ID, or None when LangFuse is disabled.

udiagent.UDIAgent.trace

trace(*, session_id: str | None = None, name: str = 'orchestrator-run')

Group every OpenAI call made within this block under one trace.

When LangFuse is enabled, opens an enclosing span so the langfuse.openai integration nests all generations (orchestration, vis generation, etc.) of a single turn into one trace instead of emitting a separate trace per call. session_id (the frontend’s per-conversation ID) groups successive turns into one LangFuse session. When LangFuse is disabled this is a no-op.

udiagent.Usage

Usage(prompt_tokens: int = 0, completion_tokens: int = 0, total_tokens: int = 0, cached_prompt_tokens: int = 0, cache_write_tokens: int = 0, reasoning_tokens: int = 0, operations: list[dict] = list()) -> None

Aggregated and per-operation token usage for a single Orchestrator.run().

Functions:

• add – Record usage from a chat-completion shaped response (see UDIAgent.create_completion).

udiagent.Usage.add

add(op: str, resp_usage: str) -> None

Record usage from a chat-completion shaped response (see UDIAgent.create_completion).

udiagent.build_rebuff_toolcall

build_rebuff_toolcall(message: str, suggestions: list[str] | None = None, *, reason: str | None = None) -> dict

Build a Rebuff tool_call dict without invoking an LLM.

Matches the shape produced by Orchestrator._handle_rebuff so it can be returned directly alongside (or in place of) normal tool_calls.

reason is an optional machine-readable discriminator (currently only "budget_exceeded") that lets the frontend distinguish a quota refusal from an ordinary rebuff and prompt the user for their own API key. The key is omitted entirely when reason is None so existing rebuffs keep their original payload shape.

udiagent.export_registry_json

export_registry_json()

Export the function registry as a JSON-serializable dict.

Suitable for sharing with the frontend as a contract.

udiagent.generate_vis_spec

generate_vis_spec(agent, messages, data_schema, grammar, config = None, usage = None, openai_api_key = None, data_domains = None, model = None, request_text = None)

Generate a visualization spec using the skills pipeline.

Parameters:

• agent – UDIAgent instance
• messages – chat history (list of dicts with role/content)
• request_text – the user’s own words, when messages ends in a rewrite of them (the orchestrator’s chart description). Optional choices the request must ask for, such as the time span, are checked against it; without it, against the newest user message.
• data_schema – JSON string describing available datasets
• grammar – dict from load_grammar()
• config – optional dict with keys: n: int (default 1) max_corrections: int (default 2)

Returns: {“spec”: dict|str, “valid”: bool, “errors”: list, “corrections”: int}

udiagent.get_function_signatures

get_function_signatures()

Return the function registry as a list of signature descriptions.

Suitable for including in LLM prompts.

udiagent.load_grammar

load_grammar(grammar_name: str = 'udi', base_path: Optional[str | Path] = None) -> dict

Load a grammar definition by name.

If base_path is None, uses the bundled package data. Returns {“schema_dict”: …, “schema_string”: …, “system_prompt”: …}

udiagent.load_skills

load_skills(skills_dir: Optional[str | Path] = None) -> dict[str, Skill]

Load all skill .md files from a directory.

If skills_dir is None, uses the bundled package data. Returns dict mapping skill name -> Skill instance.

udiagent.normalize_tool_calls

normalize_tool_calls(messages: list[dict], *, compact_charts: bool = True) -> list[dict]

Reformat tool_calls into valid OpenAI function-calling format and inject synthetic tool-response messages where missing.

The frontend echoes back dispatched results as flat dicts (e.g. {“name”: “RenderVisualization”, “arguments”: {…}}). The OpenAI API requires each tool_call entry to have “id”, “type”: “function”, and arguments as a JSON string. It also requires a corresponding “tool” role message for every tool_call_id.

With compact_charts (the default), a chart made from a template is shown as the template call behind it rather than its full spec — see compact_visualization_history. Mutates messages; pass a copy.

udiagent.parse_schema_from_dict

parse_schema_from_dict(raw: dict) -> dict

Parse a UDI data schema dict into a structured representation.

Works with an in-memory dict (e.g. the per-request data_schema) instead of a file path. This is the single source of truth for schema parsing; generate_tools.parse_schema loads a file and delegates here. All declared fields are preserved so structured text matches the schema rendered by the frontend.

Returns::

{
    "base_path": str,
    "entities": {
        "name": {
            "url": str,
            "row_count": int,
            "fields": {"field_name": {"type": str, "cardinality": int}},
        },
    },
    "relationships": [
        {
            "from_entity": str, "to_entity": str,
            "from_field": str, "to_field": str,
            "from_cardinality": str, "to_cardinality": str,
        }
    ],
    "links": [  # from each cube's ``udi:linked_dimensions``
        {
            "cube": str, "dimension": str,  # the cube column
            "entity": str, "field": str,    # the line-level column
            "values": {cube_value: [line_level_value, ...]},
        }
    ],
    "prefer": "line_level" | "data_cube",  # only when declared
}

The url and relationships data is required by the visualization template instantiation/validation path (vis_generate) so that a single agent can resolve templates against whatever schema arrives per request.

udiagent.render_template

render_template(instructions, variables)

Replace {{key}} placeholders in instructions with values from variables.

Supports including arbitrary textual data in skill prompts. Unknown placeholders are left as-is.

udiagent.segment_structured_text

segment_structured_text(text, schema)

Segment text into a mixed list of plain strings and structured element objects.

Each structured element {function_name(args)} becomes an object with:

• “expression”: the raw function reference (e.g. “{entity_count()}”)
• “label”: a human-readable description (e.g. “entity count”)
• “value”: the resolved value from the schema

Plain text between function references is kept as-is (strings). Returns (segments list, has_structured_elements bool).

udiagent.simplify_data_domains

simplify_data_domains(data_domains)

Simplify data domains for better LLM consumption.

Accepts the raw domains list (or JSON string) and produces a compact YAML-like text grouped by entity, showing each field’s type and a short domain summary (category count or numeric range).

udiagent.simplify_data_schema

simplify_data_schema(data_schema)

Simplify the data schema for better LLM consumption.

Converts from JSON to a compact YAML-like text, resolves file paths, and removes empty tables and extra information.

udiagent.split_tool_calls

split_tool_calls(messages: list[dict]) -> list[dict]

Split messages with multiple tool calls into separate messages.

Returns a new list — does not mutate the input.

udiagent.validate_structured_text

validate_structured_text(text)

Validate that all function references in the text use registered functions.

Returns list of error strings (empty = valid).

udiagent.api

The HTTP endpoints’ logic, for any web framework.

A host that serves the chat from its own backend (Flask, Django, …) calls these instead of porting the reference server’s handlers: each takes the parsed JSON request body, returns the JSON response body, and raises EndpointError for a failure the client acts on. udiagent.server.app is a thin FastAPI layer over the same functions, so a host and the reference server cannot drift.

A Flask route, for example::

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

What stays with the host is what only it knows: authentication, who may use the server’s OpenAI key, the budget, and which query backends exist.

Classes:

• EndpointError – A failure to answer with status and body().

Functions:

• completions – Run one conversation turn.
• issue_report – Record a user’s issue report, with their session export, in LangFuse.
• metadata – A package’s schema and column values, introspected from its database.
• query – Run every visible chart’s transformation pipeline as SQL, in one batch.
• vis_instantiate – Re-instantiate a template-generated chart with new bindings.

udiagent.api.EndpointError

EndpointError(status: int, code: str, message: str, errors: list[str] | None = None)

Bases: Exception

A failure to answer with status and body().

code is machine-readable. The chat acts on unknown_template: it drops a chart’s tweak controls.

Functions:

• body – The JSON error body: {code, error}, plus errors when there are several.

udiagent.api.EndpointError.body

body() -> dict

The JSON error body: {code, error}, plus errors when there are several.

udiagent.api.completions

completions(orchestrator: Orchestrator, body: dict | YACCompletionRequest, *, openai_api_key: str | None = None, session_id: str | None = None, budget_check: Callable[[Usage], str | None] | None = None) -> tuple[list[dict], dict[str, str]]

Run one conversation turn.

The logic of POST /v1/yac/completions.

The reference server’s policy is built in: a requested model is honored only with a caller’s own key (whoever pays for the tokens picks the model), and the budget is checked only on the server’s key. A quota refusal is a Rebuff in a 200, which the chat shows like any other.

Parameters:

• orchestrator (Orchestrator) – The orchestrator to run. A host that serves members on its own key and everyone else on theirs can keep one for each.
• body (dict | YACCompletionRequest) – The request body (see YACCompletionRequest).
• openai_api_key (str | None) – The caller’s own OpenAI key (the X-OpenAI-Key header), or None to run on the agent’s.
• session_id (str | None) – Groups the conversation’s turns into one LangFuse session (the X-Conversation-Id header).
• budget_check (Callable[[Usage], str | None] | None) – Called with the turn’s Usage before each LLM call; returns a message to refuse with, or None (see Orchestrator.run). Skipped for a caller’s own key.

Returns:

• list[dict] – (tool_calls, headers): the response body, and the X-Usage-*
• dict[str, str] – headers to send with it.

Raises:

• EndpointError – bad_request (422), no_openai_key (401), openai_key_refused (403, a caller’s key on Bedrock), openai_key_rejected (401), llm_call_failed (502) or llm_backend_misconfigured (503).

udiagent.api.issue_report

issue_report(agent, raw: bytes | str, *, session_id: str | None = None, webhook_url: str | None = None, defer: Callable[..., object] | None = None) -> dict

Record a user’s issue report, with their session export, in LangFuse.

The logic of POST /v1/yac/issue_report. The report is filed under the conversation’s LangFuse session, next to the traces of the turns being reported. Without LangFuse there is nowhere to put it, so it is refused rather than accepted and dropped.

Parameters:

• agent – The UDIAgent whose LangFuse records the report.
• raw (bytes | str) – The request body as received, so its size is checked before it is parsed (see YACIssueReportRequest).
• session_id (str | None) – The conversation’s ID (the X-Conversation-Id header).
• webhook_url (str | None) – Where to post a summary (a Slack or Teams incoming webhook), or None for no post.
• defer (Callable…, [object] | None) – Called as defer(fn, *args) to post the summary after the response is sent (FastAPI: BackgroundTasks.add_task). None posts it before returning, which can take up to 10 s.

Returns:

• dict – {traceId} (see IssueReportResponse).

Raises:

• EndpointError – issue_reports_unavailable (503, no LangFuse), too_large (413) or bad_request (422).

udiagent.api.metadata

metadata(engines: dict[str, QueryEngine], caches: dict, package: str | None = None, *, refresh: bool = False, token: str | None = None, principal: str | None = None, ttl_seconds: float = 3600) -> dict

A package’s schema and column values, introspected from its database.

The logic of GET /v1/yac/metadata: the remote-mode replacement for the browser computing them from CSVs. Cached for ttl_seconds.

Parameters:

• engines (dict[str, QueryEngine]) – The host’s query backends, by package name.
• caches (dict) – Where the metadata caches live between requests: an empty dict the host keeps, bounded to the 32 most recently used.
• package (str | None) – The package to describe.
• refresh (bool) – Re-introspect now instead of serving the cache.
• token (str | None) – The caller’s raw JWT, for a backend that forwards it to the database (jwtPassthrough).
• principal (str | None) – Who the caller is (e.g. the token’s sub). A passthrough backend keeps a cache per principal, because the column values are the user’s own rows.

Returns:

• dict – {package, interactive, dataSchema, dataDomains} (see
• dict – MetadataResponse).

Raises:

• EndpointError – no_backend (404) or database_forbidden (403).

udiagent.api.query

query(engines: dict[str, QueryEngine], body: dict | YACQueryRequest, *, token: str | None = None) -> dict

Run every visible chart’s transformation pipeline as SQL, in one batch.

The logic of POST /v1/yac/query: the server-side counterpart of the toolkit’s in-browser executor. A query that fails carries {error} in its result; the others still succeed.

Parameters:

• engines (dict[str, QueryEngine]) – The host’s query backends, by package name.
• body (dict | YACQueryRequest) – The request body (see YACQueryRequest).
• token (str | None) – The caller’s raw JWT, for a backend that forwards it to the database (jwtPassthrough).

Returns:

• dict – {results}, one per query, keyed by vizId (see QueryResponse).

Raises:

• EndpointError – bad_request (422), no_backend (404) or database_forbidden (403).

udiagent.api.vis_instantiate

vis_instantiate(body: dict | YACVisInstantiateRequest) -> dict

Re-instantiate a template-generated chart with new bindings.

The logic of POST /v1/yac/vis_instantiate. No LLM call, so it needs no agent and no OpenAI key.

Parameters:

• body (dict | YACVisInstantiateRequest) – The request body: tool, toolArgs and dataSchema (see YACVisInstantiateRequest).

Returns:

• dict – The response body, shaped as YACVisInstantiateResponse.

Raises:

• EndpointError – bad_request (422), bad_schema (400), unknown_template (404), missing_bindings or invalid_bindings (422), instantiate_failed (500) or templates_unavailable (503).

udiagent.query

Server-side query layer: compiles UDI grammar transformation pipelines to SQL and executes them against a configured backend (StarRocks, DuckDB).

This is “Compiler B” — the server analogue of the toolkit’s Arquero executor (“Compiler A”, packages/grammar/DataSourcesStore.ts), which defines the reference semantics. Parity between the two is enforced by tests/goldens.

Classes:

• DatabaseAuthError – The database rejected the identity forwarded for this request.
• DuckDBConnector – In-process DuckDB. views maps entity/table names to CSV/Parquet file
• MetadataCache – TTL-cached introspection for one engine/package pair.
• QueryEngine –
• StarRocksConnector – StarRocks over the MySQL wire protocol (pymysql, [starrocks] extra).
• UnsupportedQueryError – A spec uses a construct the server-side compiler does not support

Functions:

• use_db_token – Bind the caller’s raw JWT for the duration of one request.

udiagent.query.DatabaseAuthError

Bases: PermissionError

The database rejected the identity forwarded for this request.

Distinct from UnsupportedQueryError: that one is per-visualization and gets reported inside a 200 response, whereas this fails the whole request with a 403 — a credential problem is not a chart-shaped error.

udiagent.query.DuckDBConnector

DuckDBConnector(database: str = ':memory:', views: dict[str, str] | None = None)

In-process DuckDB. views maps entity/table names to CSV/Parquet file paths registered as views — handy for tests and file-backed packages.

udiagent.query.MetadataCache

MetadataCache(engine, package_name: str, ttl_seconds: float = 3600.0, distinct_cap: int = DEFAULT_DISTINCT_CAP)

TTL-cached introspection for one engine/package pair.

get() serves the cached result until ttl_seconds elapse; refresh() forces re-introspection (e.g. after a data load).

udiagent.query.QueryEngine

QueryEngine(connector, table_map: dict[str, str], row_cap: int = DEFAULT_ROW_CAP, entity_schemas: dict[str, dict] | None = None, package_extras: dict | None = None)

Functions:

• run_batch – Run a batched /v1/yac/query request. Each query dict carries
• run_query – offset pages row-level results past the cap (LIMIT cap OFFSET n

row_cap: max rows returned for row-level (non-aggregated) results. entity_schemas: entity name -> extra schema metadata that can’t be introspected from the database — primaryKey/foreignKeys, plus the table description and per-column fields descriptions — merged into each resource by introspect(). The chat’s cross-entity filtering (getEntityRelationship) depends on foreignKeys. package_extras: package-level descriptor keys introspection can’t produce either, such as udi:prefer — merged into the package descriptor by introspect().

udiagent.query.QueryEngine.run_batch

run_batch(queries: list[dict], selections: dict[str, Any] | None = None) -> dict[str, dict]

Run a batched /v1/yac/query request. Each query dict carries {vizId, source, transformation?, displayDataOnly?}. Errors are per-viz: one failing spec doesn’t sink the batch.

udiagent.query.QueryEngine.run_query

run_query(source: list[dict] | dict, transformation: list[dict] | None = None, selections: dict[str, Any] | None = None, display_data_only: bool | None = None, offset: int = 0) -> dict

offset pages row-level results past the cap (LIMIT cap OFFSET n on the display pass). ponytail: pages are only deterministic when the spec carries an orderby — SQL row order is otherwise unspecified.

udiagent.query.StarRocksConnector

StarRocksConnector(host: str, port: int = 9030, user: str = '', password: str = '', database: str | None = None, *, jwt_passthrough: bool = False, principal_field: str = 'sub', **kwargs: Any)

StarRocks over the MySQL wire protocol (pymysql, [starrocks] extra).

Two identity modes, chosen at construction:

• Service credential (default) — one long-lived connection shared by every caller, authenticated with the configured user/password. Unchanged from the original behaviour.
• JWT passthrough (jwt_passthrough=True) — the caller’s own token, bound for the request by use_db_token, is forwarded as the password. StarRocks authenticates it against the same JWKS the agent used and applies that user’s own grants, so authorization lives in the database rather than in this process. One cached connection per principal.

Both modes share one pool: idle connections per identity, each checked out by one query at a time (pymysql connections are not thread-safe) while the pool’s lock is held only to check them out and back in. Queries run in parallel — FastAPI’s threadpool across requests, introspection across tables — and a slow query or connect never stalls another user. Connections are long-lived and so must survive idle timeouts: every execute pings with reconnect and retries once if the socket died mid-query.

Revocation ceiling: MySQL binds identity at CONNECT time and StarRocks does not re-check exp mid-session, so a cached connection is recycled once the token that opened it expires. A grant revoked inside a live session is therefore not seen until the connection is recycled — at most one token lifetime, or _IDLE_SECONDS of inactivity, whichever comes first.

udiagent.query.UnsupportedQueryError

Bases: ValueError

A spec uses a construct the server-side compiler does not support (e.g. legacy raw Arquero expression strings, or ops after kde).

udiagent.query.use_db_token

use_db_token(token: str | None)

Bind the caller’s raw JWT for the duration of one request.

A contextvar rather than a parameter because connector.execute is reached from three independent places — QueryEngine._execute, the compiler’s binby probe, and introspect’s DESCRIBE/DISTINCT passes — so threading a token through would change five signatures across three modules.

Set this in the request handler, never in a FastAPI dependency: sync dependencies run through anyio.to_thread.run_sync in a copied context, so a .set() there never reaches the handler.