feat(spec): resolve agent instructions from MLflow Prompt Registry - #3229
Open
Praneeth16 wants to merge 4 commits into
Open
feat(spec): resolve agent instructions from MLflow Prompt Registry#3229Praneeth16 wants to merge 4 commits into
Praneeth16 wants to merge 4 commits into
Conversation
Agent `instructions:` (legacy alias `prompt:`) could only be inline text
or a bundle-relative file. Teams that manage prompts in the MLflow 3
Prompt Registry had to copy prompt text into each bundle by hand, which
drifts from the registry-of-record and loses versioning/alias control.
Add a third `instructions:` form that references a registry prompt and
resolves to plain text at bundle-load time:
- structured: `instructions: {source: mlflow, reference:
prompts:/name@production, vars: {product: X}}`
- shorthand: `instructions: mlflow+prompts:/name@production`
Resolution is funneled through one shared helper
(`omnigent/spec/mlflow_prompts.py`) called by BOTH parser copies
(`spec/parser.py` and `inner/loader.py`) so the rules can never drift.
`AgentSpec.instructions` stays a plain `str`; no harness sees anything
new. The OSS and Databricks-managed (Unity Catalog) registries share one
code path, selected by config/env (`registry_uri`, tracking URI, auth)
rather than branching business logic.
Safety properties baked in and covered by tests:
- The registry fetch is gated on `expand_env`. A scaffolding /
validation parse (untrusted or HTTP-uploaded bundle, `expand_env=
False`) leaves the reference as its literal string and never
contacts the registry on the author's behalf.
- Prompts resolve by version or alias only. Pinned versions are
immutable so MLflow caches them forever; aliases honor MLflow's own
cache TTL (never wrapped in an unbounded cache) so a re-point is
picked up.
- Chat-style (message-list) prompts are rejected with a clear error;
instructions must be a single text string.
- Registry failures (down, missing alias/prompt, auth) surface a clear
error and never silently fall back to a hardcoded prompt.
- Auth stays in the environment / MLflow unified auth; secrets never
live in YAML. Only the resolved prompt name and version are logged,
never prompt contents.
- `mlflow` is an optional extra, imported lazily; a registry reference
with mlflow absent raises a clear `pip install 'omnigent[mlflow]'`
error. The `_read_contained_file` path-containment model is
untouched.
Closes omnigent-ai#3226
Signed-off-by: praneeth_paikray-data <praneeth.paikray@databricks.com>
added 3 commits
July 24, 2026 17:55
The new optional `mlflow` extra pulls in mlflow-skinny, whose transitive dependency gitpython resolved to 3.1.50. That version carries three OSV advisories fixed in 3.1.51: GHSA-2f96-g7mh-g2hx GHSA-v396-v7q4-x2qj GHSA-956x-8gvw-wg5v Because the extra widened the dependency graph, the Security Scan job's `uv export --frozen --all-extras | pip-audit` step now sees gitpython and fails, cascading into the Security Gate shards and Maintainer Approval. Add a uv constraint (`gitpython>=3.1.51,<4`) rather than a direct dependency so gitpython is only pulled in by extras that already use it, and relock. Resolution lands on gitpython 3.1.52, which OSV reports with zero known vulnerabilities. Only gitpython moves in the lock; no other package churns. Signed-off-by: praneeth_paikray-data <praneeth.paikray@databricks.com>
…curity bugs
Cross-vendor review of the MLflow Prompt Registry resolver surfaced six
blocking bugs where the resolver could silently resolve the wrong prompt,
leak state across registries, stringify a repr into an agent's
instructions, or accept secrets in YAML. Fix each at the source and add a
regression test per bug.
- URI validation was prefix-only, so `prompts:/greeting/1?tag=production`
passed and MLflow silently resolved version 1, ignoring the tag. Require
exactly `prompts:/<name>@<alias>` or `prompts:/<name>/<version>` and
reject any query string, fragment, or trailing junk — MLflow resolves by
version or alias only and would drop the extra selector.
- Resolution mutated `mlflow.set_tracking_uri` / `set_registry_uri`
process globals and never restored them, so a Databricks-UC load
contaminated a later OSS load and concurrent loads raced. Resolve through
a per-call `MlflowClient(tracking_uri=..., registry_uri=...)` instead and
never touch globals.
- MLflow's prompt cache is process-global and keyed by URI only, so
registry B could be served registry A's cached prompt for the same
`prompts:/name/version` URI. Disable the cache (cache_ttl_seconds=0) by
default; callers opt back in with an explicit cache_ttl when a single
registry is in play.
- `str(prompt.format(**vars))` defeated type safety: with a partial render
`format` returns a `PromptVersion` (or a message list), and `str(...)`
stuffed its repr into the instructions. Assert the render is a `str` and
raise otherwise.
- tracking_uri / registry_uri accepted `user:pass@host` credentials,
putting secrets in agent YAML. Reject any userinfo-bearing URI and point
at the environment / MLflow unified auth.
- `if vars:` skipped formatting for an explicit empty mapping; key on
presence (`vars is not None`) so `vars: {}` renders consistently.
`MlflowClient.load_prompt(cache_ttl_seconds=...)` landed in MLflow 3.8.0
(mlflow#18963; absent in 3.7.x), so raise the `mlflow` extra and dev floor
from >=3 to >=3.8 and relock (specifier metadata only; no resolved
versions move, gitpython>=3.1.51 constraint intact).
Co-authored-by: Isaac
Signed-off-by: praneeth_paikray-data <praneeth.paikray@databricks.com>
…a is absent
mlflow is an optional extra, but both prompt-registry test modules imported
it at module top level, so CI pytest groups that don't install
omnigent[mlflow] failed at collection with ModuleNotFoundError instead of
skipping. Guard the optional import with pytest.importorskip('mlflow') as the
first mlflow touch in each module and defer the mlflow.genai symbol imports
until after the guard, so the modules skip cleanly when mlflow is absent and
still run in full where it is installed.
Co-authored-by: Isaac
Signed-off-by: praneeth_paikray-data <praneeth.paikray@databricks.com>
Contributor
|
@Praneeth16 This PR is a Bug fix, Feature, or UI / frontend change but the Demo section is missing or only contains a placeholder. These change types require a screenshot or screen recording so reviewers can see the new behaviour without checking out the branch. Please update the Demo section with:
Use |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Related issue
Closes #3226
Summary
Agent
instructions:(legacy aliasprompt:) could only be inline text or a bundle-relative file. Teams that keep prompts in the MLflow 3 Prompt Registry had to hand-copy prompt text into each bundle, which drifts from the registry-of-record and loses versioning/alias control.This adds a third
instructions:form that references a registry prompt and resolves to plain text at bundle-load time:instructions: {source: mlflow, reference: prompts:/name@production, vars: {product: X}}instructions: mlflow+prompts:/name@productionResolution runs through one shared helper (
omnigent/spec/mlflow_prompts.py) called by BOTH parser copies (spec/parser.pyandinner/loader.py), so the rules can never drift between them.AgentSpec.instructionsstays a plainstrand no harness sees anything new. The OSS and Databricks-managed (Unity Catalog) registries share one code path, selected by config/env (registry_uri, tracking URI, auth) rather than branching business logic. v1 is load-time resolution only.Safety properties (all covered by tests):
expand_env. A scaffolding / validation parse (untrusted or HTTP-uploaded bundle,expand_env=False) leaves the reference as its literal string and never contacts the registry on the author's behalf.mlflowis an optional extra, imported lazily; a registry reference with mlflow absent raises a clearpip install 'omnigent[mlflow]'error. The_read_contained_filepath-containment model is untouched.Test Plan
New tests (run from the worktree root against a throwaway sqlite-backed MLflow registry):
.venv/bin/python -m pytest tests/spec/test_mlflow_prompts.py tests/inner/test_loader_mlflow_prompts.py— 21 tests, all pass. Covers: shorthand + structured detection, malformed-form rejection, non-prompts:/URI rejection, OSS resolve by version / alias / no-vars, chat-prompt rejection, missing alias/prompt failure, missing-mlflow-package install hint, parser + loader wiring, theexpand_env=Falseno-fetch guard (monkeypatches the resolver to assert it is never called), the shared-helper assertion for the loader path, and Databricks-UC backend selection (mocked SDK: assertsregistry_uri="databricks-uc", tracking URI not set,link_to_model=False).tests/spec+tests/innersuites pass (302 tests + 3 subtests)..venv/bin/pre-commit run --files <all changed files>— all hooks pass (ruff-format, ruff-check, lint hooks, version sync, normalize-uv-lock-registry, whitespace/EOF/merge-conflict/large-file checks).Demo
N/A — non-visual change (backend prompt resolution).
Type of change
Test coverage
Coverage notes
Manual verification: registered a text prompt and a chat-style prompt in a local sqlite MLflow registry, then confirmed a bundle with
instructions: mlflow+prompts:/greeting@productionresolves to the rendered text, the chat prompt is rejected with a clear error, and a validation parse (expand_env=False) returns the literal reference without any registry call.Changelog
Agent
instructions:can now reference an MLflow 3 Prompt Registry prompt (mlflow+prompts:/name@aliasor asource: mlflowmapping), resolved to text at load time.