Skip to content

Python: Add MLflow AI Gateway provider samples#5507

Open
PattaraS wants to merge 3 commits intomicrosoft:mainfrom
PattaraS:add-mlflow-gateway-samples
Open

Python: Add MLflow AI Gateway provider samples#5507
PattaraS wants to merge 3 commits intomicrosoft:mainfrom
PattaraS:add-mlflow-gateway-samples

Conversation

@PattaraS
Copy link
Copy Markdown

Motivation

MLflow AI Gateway (MLflow ≥ 3.0) is an open-source, database-backed LLM proxy that provides a unified API across 20+ providers (OpenAI, Anthropic, Gemini, Mistral, Bedrock, Ollama, etc.) with encrypted secrets management, fallback/retry, traffic splitting, and budget tracking.

Agent Framework users currently have no documented path for using MLflow AI Gateway as their LLM backend. This PR adds a sample to fill that gap.

Description

Adds python/samples/02-agents/providers/mlflow_gateway/ with:

  • mlflow_gateway_with_openai_chat_client.py — Example showing how to route Agent Framework requests through MLflow AI Gateway's OpenAI-compatible endpoint using the existing OpenAIChatClient with a custom base_url. Includes both streaming and non-streaming examples with tool calling.
  • README.md — Setup instructions (install MLflow, start the server, create a gateway endpoint), environment variable configuration, and feature overview.

The pattern follows the established ollama_with_openai_chat_client.py example — no new dependencies are required beyond the existing OpenAI integration.

Also updates python/samples/02-agents/providers/README.md to add the new folder to the index table.

Checklist

  • Sample follows the existing provider folder structure (matches ollama/, openai/, etc.)
  • Sample uses environment variables for configuration
  • README includes setup, prerequisites, and configuration sections
  • Updated parent providers/README.md index
  • No new dependencies required
  • No breaking changes

AI Disclosure

This pull request was AI-assisted by Claude. All content was reviewed and validated by a human contributor.

Adds a new providers/mlflow_gateway/ directory with an example showing
how to route Agent Framework requests through MLflow AI Gateway's
OpenAI-compatible endpoint using the existing OpenAIChatClient.

Follows the same pattern as ollama_with_openai_chat_client.py — no new
dependencies required beyond the OpenAI integration.
Copilot AI review requested due to automatic review settings April 27, 2026 06:31
@moonbox3 moonbox3 added documentation Improvements or additions to documentation python labels Apr 27, 2026
@github-actions github-actions Bot changed the title Add MLflow AI Gateway provider samples Python: Add MLflow AI Gateway provider samples Apr 27, 2026
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Adds a new Python provider sample demonstrating how to route Agent Framework OpenAIChatClient traffic through MLflow AI Gateway’s OpenAI-compatible endpoint, plus documentation and index entry updates.

Changes:

  • Added an MLflow AI Gateway provider sample with both streaming and non-streaming agent runs (including tool calling).
  • Added a provider-specific README describing MLflow AI Gateway setup and required environment variables.
  • Updated the providers index README to include the new mlflow_gateway/ folder.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File Description
python/samples/02-agents/providers/mlflow_gateway/mlflow_gateway_with_openai_chat_client.py New sample showing OpenAIChatClient(base_url=...) configured for MLflow AI Gateway with streaming/non-streaming runs and a simple tool.
python/samples/02-agents/providers/mlflow_gateway/README.md New documentation for prerequisites, configuration, and usage of the MLflow AI Gateway sample.
python/samples/02-agents/providers/README.md Adds mlflow_gateway/ entry to the provider samples index table.

Comment on lines +41 to +44
| File | Description |
|------|-------------|
| [`mlflow_gateway_with_openai_chat_client.py`](mlflow_gateway_with_openai_chat_client.py) | Connect an Agent Framework agent to MLflow AI Gateway via the OpenAI-compatible endpoint. Shows both streaming and non-streaming responses with tool calling. |

Copy link

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The examples table is malformed markdown (each row starts with ||), so it won’t render as a 2‑column table. Use standard table syntax with a single leading | per row (matching other provider READMEs).

Copilot uses AI. Check for mistakes.
Comment on lines +61 to +65
_client = OpenAIChatClient(
api_key="unused", # Provider keys are managed by the MLflow server
base_url=os.getenv("MLFLOW_GATEWAY_ENDPOINT"),
model=os.getenv("MLFLOW_GATEWAY_MODEL"),
)
Copy link

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

base_url comes from os.getenv("MLFLOW_GATEWAY_ENDPOINT"); if it’s unset (or empty), OpenAIChatClient will fall back to the default OpenAI base URL and may send prompts to OpenAI unexpectedly. Consider validating MLFLOW_GATEWAY_ENDPOINT (and MLFLOW_GATEWAY_MODEL) up front and failing fast with a clear message before constructing the client.

Copilot uses AI. Check for mistakes.
Comment on lines +83 to +87
_client = OpenAIChatClient(
api_key="unused", # Provider keys are managed by the MLflow server
base_url=os.getenv("MLFLOW_GATEWAY_ENDPOINT"),
model=os.getenv("MLFLOW_GATEWAY_MODEL"),
)
Copy link

Copilot AI Apr 27, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above: if MLFLOW_GATEWAY_ENDPOINT isn’t set, the client may default to OpenAI’s public endpoint. Validate required env vars once (e.g., in main()) and pass the resolved values into both examples to avoid accidental misrouting.

Copilot uses AI. Check for mistakes.
Address Copilot review feedback: without explicit validation, an unset
or empty MLFLOW_GATEWAY_ENDPOINT would cause OpenAIChatClient to
silently fall back to OpenAI's public endpoint and forward prompts
there. Validate both env vars in main() and pass the resolved values
into both example functions, failing fast with a clear error message.
@PattaraS PattaraS marked this pull request as ready for review April 27, 2026 07:17
@PattaraS
Copy link
Copy Markdown
Author

Thanks for the review! Pushed 48ac8c9:

  1. Env var validation (real issue) — Added a _require_env() helper that fails fast with a clear error if MLFLOW_GATEWAY_ENDPOINT or MLFLOW_GATEWAY_MODEL is missing or empty. Without it, an unset endpoint would cause OpenAIChatClient to silently default to OpenAI's public endpoint — exactly the leak risk you flagged. Validation now happens once in main() and the resolved values are passed into both example functions.
  2. Markdown table — I believe this was a false positive; the table uses single | per row (verified with od -c). Happy to restructure if it actually rendered wrong somewhere — got a screenshot?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation python

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants