Add passthrough support for Gemini generateContent and streamGenerateContent APIs by TomeHirata · Pull Request #19425 · mlflow/mlflow

TomeHirata · 2025-12-16T09:23:30Z

🥞 Stacked PR

Use this link to review incremental changes.

stack/gateway/model-passthrough/gemini [Files changed]
- stack/gateway/structured-output [Files changed]
  - stack/gateway/missing-params [Files changed]

Related Issues/PRs

n/a

What changes are proposed in this pull request?

Add following Gemini model passthrough endpoints where request body is just propagated to the provider endpoint:

/gateway/gemini/v1beta/models/{endpoint_name}:generateContent
/gateway/gemini/v1beta/models/{endpoint_name}:streamGenerateContent

from google import genai
from google.genai import types
import mlflow

mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("Gateway")
from mlflow.tracking._tracking_service.utils import _get_store

store = _get_store()

secret = store.create_gateway_secret(
    secret_name="gemini_api",
    secret_value={"api_key": "<secret>"},
    provider="gemini",
)
model_def = store.create_gateway_model_definition(
    name="gemini-2.5-flash",
    secret_id=secret.secret_id,
    provider="gemini",
    model_name="gemini-2.5-flash",
)
endpoint = store.create_gateway_endpoint(
    name="gemini-v1", model_definition_ids=[model_def.model_definition_id]
)

client = genai.Client(
    api_key='dummy',
    http_options={
        'base_url': "http://localhost:5000/gateway/gemini",
    })
response = client.models.generate_content(
    model="gemini-v1",
    contents='Hello',
)
client.close()
print(response.candidates[0].content.parts[0].text)

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

No. You can skip the rest of this section.
Yes. Give a description of this change to be included in the release notes for MLflow users.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

How should the PR be classified in the release notes? Choose one:

rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
rn/feature - A new user-facing feature worth mentioning in the release notes
rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?

Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
Bug fixes, doc updates and new features usually go into minor releases.
Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
Bug fixes and doc updates usually go into patch releases.

Yes (this PR will be cherry-picked and included in the next patch release)
No (this PR will be included in the next minor release)

github-actions · 2025-12-16T09:31:25Z

Documentation preview for 880b16b is available at:

https://pr-19425--mlflow-docs-preview.netlify.app/docs/latest/

More info

Ignore this comment if this PR does not change the documentation.
The preview is updated when a new commit is pushed to this PR.
This comment was created by this workflow run.
The documentation was built by this workflow run.

Copilot

Pull request overview

This PR adds passthrough support for Gemini's generateContent and streamGenerateContent APIs as part of a larger effort to provide native API support for multiple LLM providers. The implementation follows the patterns established in the stack for OpenAI and Anthropic passthrough endpoints.

Key Changes

Added two new Gemini passthrough endpoints that accept raw Gemini API format
Implemented LiteLLM provider as a fallback for unsupported providers
Extended base provider class with passthrough methods for multiple providers

Reviewed changes

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

Show a summary per file

File	Description
mlflow/server/gateway_api.py	Added Gemini passthrough endpoints (generateContent and streamGenerateContent) and helper function for extracting endpoint names
mlflow/gateway/providers/gemini.py	Implemented passthrough methods for Gemini generateContent APIs
mlflow/gateway/providers/base.py	Added abstract passthrough method definitions for all provider types
mlflow/gateway/providers/anthropic.py	Implemented Anthropic Messages API passthrough endpoint
mlflow/gateway/providers/openai.py	Implemented OpenAI passthrough endpoints for chat, embeddings, and responses APIs
mlflow/gateway/providers/litellm.py	New provider using LiteLLM library as fallback for unsupported providers
mlflow/gateway/config.py	Added LiteLLMConfig class and LITELLM provider enum value
mlflow/gateway/provider_registry.py	Registered LiteLLM provider in the provider registry
tests/server/test_gateway_api.py	Comprehensive tests for all new passthrough endpoints including streaming
tests/gateway/providers/test_gemini.py	Unit tests for Gemini passthrough methods
tests/gateway/providers/test_anthropic.py	Unit tests for Anthropic passthrough methods
tests/gateway/providers/test_openai.py	Unit tests for OpenAI passthrough methods
tests/gateway/providers/test_litellm.py	Complete test suite for new LiteLLM provider
docs/api_reference/api_inventory.txt	Updated API inventory with LiteLLMConfig references

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2025-12-16T09:41:32Z

tests/gateway/providers/test_anthropic.py

+        }
+        response = await provider.passthrough_anthropic_messages(payload)
+
+        assert payload["model"] == "claude-2.1"


This assertion validates that the payload has been mutated in place by the passthrough_anthropic_messages method (the model field is added). While the test currently passes, asserting on side effects of the function under test rather than its return value makes the test more fragile. Consider testing only the return value and the mock call arguments.

Suggested change

assert payload["model"] == "claude-2.1"

Copilot · 2025-12-16T09:41:32Z

mlflow/gateway/providers/anthropic.py

+        Accepts raw Anthropic request format and returns raw Anthropic response format.
+        Supports streaming if the 'stream' parameter is set to True.
+        """
+        # Add model name from config


Mutating the input payload dictionary in place could cause unexpected side effects for callers. The payload is modified by adding the 'model' field, which could affect code that reuses the same dictionary. Consider creating a copy of the payload before modifying it, or explicitly document this mutation as part of the function's contract.

Suggested change

# Add model name from config

# Add model name from config

payload = dict(payload) # Avoid mutating the input dictionary

Copilot · 2025-12-16T09:41:33Z

mlflow/gateway/providers/base.py

+            detail=f"The passthrough Anthropic messages route is not implemented for {self.NAME}"
+            "models.",


Missing space between string literals in error message. The error message will read "...for {self.NAME}models." instead of "...for {self.NAME} models."

Suggested change

detail=f"The passthrough Anthropic messages route is not implemented for {self.NAME}"

"models.",

detail=f"The passthrough Anthropic messages route is not implemented for {self.NAME} models.",

Copilot · 2025-12-16T09:41:33Z

mlflow/server/gateway_api.py

+        if "api_base" in auth_config:
+            litellm_config["litellm_api_base"] = auth_config["api_base"]
+        provider_config = LiteLLMConfig(**litellm_config)
+        model_config.provider = Provider.LITELLM


Mutating the model_config.provider field could have unexpected side effects. The model_config object likely comes from the database store and modifying it in-place could affect other parts of the code that reference it. Consider creating a new variable or ensuring this mutation doesn't propagate beyond this function's scope.

Copilot · 2025-12-16T09:41:33Z

tests/gateway/providers/test_anthropic.py

+        }
+        response = await provider.passthrough_anthropic_messages(payload)
+
+        assert payload["model"] == "claude-2.1"


This assertion validates that the payload has been mutated in place by the passthrough_anthropic_messages method (the model field is added). While the test currently passes, asserting on side effects of the function under test rather than its return value makes the test more fragile. Consider testing only the return value and the mock call arguments.

B-Step62

LTGM, left a few minor comments

mlflow/gateway/providers/gemini.py

…Content APIs Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

BenWilson2

Looks good!

TomeHirata · 2025-12-18T00:35:05Z

Bypassing unrelated test failures

This was referenced Dec 16, 2025

Add LiteLLM provider to support many other providers #19394

Merged

Add OpenAI passthrough endpoints for chat, embeddings, and responses #19404

Merged

Add passthrough support for Anthropic Messages API #19423

Merged

TomeHirata force-pushed the stack/gateway/model-passthrough/gemini branch from 4b14fae to 069efbe Compare December 16, 2025 09:24

TomeHirata marked this pull request as ready for review December 16, 2025 09:24

TomeHirata requested review from B-Step62, BenWilson2 and Copilot December 16, 2025 09:24

Copilot started reviewing on behalf of TomeHirata December 16, 2025 09:25 View session

github-actions bot added the rn/feature Mention under Features in Changelogs. label Dec 16, 2025

Copilot AI reviewed Dec 16, 2025

View reviewed changes

TomeHirata force-pushed the stack/gateway/model-passthrough/gemini branch 2 times, most recently from 20b0386 to cf3d4f8 Compare December 17, 2025 04:07

This was referenced Dec 17, 2025

Enhance Anthropic and Gemini providers to support structured outputs #19452

Merged

Add support for top_k, frequency_penalty, and presence_penalty in Gemini and Mistral provider #19453

Merged

TomeHirata force-pushed the stack/gateway/model-passthrough/gemini branch 4 times, most recently from d991ed4 to f07f11e Compare December 17, 2025 08:48

B-Step62 approved these changes Dec 17, 2025

View reviewed changes

mlflow/gateway/providers/gemini.py Outdated Show resolved Hide resolved

mlflow/gateway/providers/gemini.py Outdated Show resolved Hide resolved

github-actions bot assigned B-Step62 Dec 17, 2025

TomeHirata force-pushed the stack/gateway/model-passthrough/gemini branch 2 times, most recently from 0a9ed11 to b23b045 Compare December 17, 2025 13:36

TomeHirata added 4 commits December 18, 2025 01:11

Add passthrough support for Gemini generateContent and streamGenerate…

1e7a9ce

…Content APIs Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

comment

447a7b2

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

comment

f96d4bb

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

comment

880b16b

Signed-off-by: Tomu Hirata <tomu.hirata@gmail.com>

TomeHirata force-pushed the stack/gateway/model-passthrough/gemini branch from b23b045 to 880b16b Compare December 17, 2025 16:13

TomeHirata enabled auto-merge December 17, 2025 16:14

BenWilson2 approved these changes Dec 17, 2025

View reviewed changes

github-actions bot assigned BenWilson2 Dec 17, 2025

TomeHirata disabled auto-merge December 18, 2025 00:35

TomeHirata merged commit b72d49a into mlflow:master Dec 18, 2025
46 of 50 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add passthrough support for Gemini generateContent and streamGenerateContent APIs#19425

Add passthrough support for Gemini generateContent and streamGenerateContent APIs#19425
TomeHirata merged 4 commits intomlflow:masterfrom
TomeHirata:stack/gateway/model-passthrough/gemini

TomeHirata commented Dec 16, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Dec 16, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI Dec 16, 2025

Uh oh!

Copilot AI Dec 16, 2025

Uh oh!

Copilot AI Dec 16, 2025

Uh oh!

Copilot AI Dec 16, 2025

Uh oh!

Copilot AI Dec 16, 2025

Uh oh!

B-Step62 left a comment

Uh oh!

Uh oh!

Uh oh!

BenWilson2 left a comment

Uh oh!

TomeHirata commented Dec 18, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

	# Add model name from config
	# Add model name from config
	payload = dict(payload) # Avoid mutating the input dictionary

		detail=f"The passthrough Anthropic messages route is not implemented for {self.NAME}"
		"models.",

Conversation

TomeHirata commented Dec 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🥞 Stacked PR

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

How should the PR be classified in the release notes? Choose one:

Should this PR be included in the next patch release?

Uh oh!

github-actions bot commented Dec 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Key Changes

Reviewed changes

Uh oh!

Copilot AI Dec 16, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 16, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 16, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 16, 2025

Choose a reason for hiding this comment

Uh oh!

Copilot AI Dec 16, 2025

Choose a reason for hiding this comment

Uh oh!

B-Step62 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

BenWilson2 left a comment

Choose a reason for hiding this comment

Uh oh!

TomeHirata commented Dec 18, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

TomeHirata commented Dec 16, 2025 •

edited

Loading

github-actions bot commented Dec 16, 2025 •

edited

Loading