Skip to content

Commit 288e732

Browse files
authored
docs: add MLflow AI Gateway section to MLflow integration page (#454)
* docs: add MLflow AI Gateway section to MLflow integration page Adds a subsection showing how to use OpenAIChatGenerator with MLflow AI Gateway's OpenAI-compatible endpoint via api_base_url. Includes setup instructions, standalone and pipeline examples. * docs: add Create Endpoint screenshot to MLflow AI Gateway section
1 parent 794da71 commit 288e732

2 files changed

Lines changed: 63 additions & 0 deletions

File tree

214 KB
Loading

integrations/mlflow.md

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ toc: true
2222
- [Overview](#overview)
2323
- [Installation](#installation)
2424
- [Usage](#usage)
25+
- [Trace a RAG Pipeline](#trace-a-rag-pipeline)
26+
- [Use MLflow AI Gateway as an LLM Backend](#use-mlflow-ai-gateway-as-an-llm-backend)
2527
- [License](#license)
2628

2729
## Overview
@@ -130,6 +132,67 @@ Auto-tracing for Haystack can be disabled by calling:
130132
mlflow.haystack.autolog(disable=True)
131133
```
132134

135+
### Use MLflow AI Gateway as an LLM Backend
136+
137+
[MLflow AI Gateway](https://mlflow.org/docs/latest/genai/governance/ai-gateway/) (MLflow ≥ 3.0) is a database-backed LLM proxy that routes requests to multiple providers — OpenAI, Anthropic, Gemini, Mistral, Bedrock, Ollama, and more — through a single OpenAI-compatible endpoint. Provider API keys are stored encrypted on the server, and features like fallback/retry, traffic splitting, and budget tracking are configured in the MLflow UI with no code changes needed.
138+
139+
Since the gateway exposes an OpenAI-compatible API, you can use Haystack's built-in `OpenAIChatGenerator` with a custom `api_base_url`:
140+
141+
**1. Install MLflow and start the server:**
142+
143+
```bash
144+
pip install mlflow[genai]
145+
mlflow server --host 127.0.0.1 --port 5000
146+
```
147+
148+
**2. Create a gateway endpoint** in the MLflow UI at `http://localhost:5000`. Navigate to **AI Gateway → Create Endpoint**, select a provider and model, and enter your provider API key. See the [MLflow AI Gateway documentation](https://mlflow.org/docs/latest/genai/governance/ai-gateway/endpoints/) for details.
149+
150+
![MLflow AI Gateway — Create Endpoint](../images/mlflow-gateway-create-endpoint.png)
151+
152+
**3. Use the endpoint in a Haystack pipeline:**
153+
154+
```python
155+
from haystack import Pipeline
156+
from haystack.components.builders import ChatPromptBuilder
157+
from haystack.components.generators.chat import OpenAIChatGenerator
158+
from haystack.dataclasses import ChatMessage
159+
from haystack.utils import Secret
160+
161+
pipe = Pipeline()
162+
pipe.add_component("prompt_builder", ChatPromptBuilder())
163+
pipe.add_component(
164+
"llm",
165+
OpenAIChatGenerator(
166+
model="my-chat-endpoint", # your MLflow Gateway endpoint name
167+
api_key=Secret.from_token("unused"), # provider keys live on the server
168+
api_base_url="http://localhost:5000/gateway/openai/v1",
169+
),
170+
)
171+
pipe.connect("prompt_builder", "llm")
172+
173+
messages = [ChatMessage.from_user("What is MLflow AI Gateway?")]
174+
result = pipe.run({"prompt_builder": {"template": messages}})
175+
print(result["llm"]["replies"][0].text)
176+
```
177+
178+
You can also use it standalone:
179+
180+
```python
181+
from haystack.components.generators.chat import OpenAIChatGenerator
182+
from haystack.dataclasses import ChatMessage
183+
from haystack.utils import Secret
184+
185+
generator = OpenAIChatGenerator(
186+
model="my-chat-endpoint",
187+
api_key=Secret.from_token("unused"),
188+
api_base_url="http://localhost:5000/gateway/openai/v1",
189+
)
190+
191+
messages = [ChatMessage.from_user("What is MLflow AI Gateway?")]
192+
result = generator.run(messages)
193+
print(result["replies"][0].text)
194+
```
195+
133196
## License
134197

135198
MLflow is distributed under the terms of the [Apache-2.0](https://spdx.org/licenses/Apache-2.0.html) license.

0 commit comments

Comments
 (0)