Skip to content

fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path#28010

Closed
Iana-Kasimova-TR wants to merge 3 commits into
BerriAI:litellm_internal_stagingfrom
Iana-Kasimova-TR:fix-vertex-ai-gemma-maas-26083
Closed

fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path#28010
Iana-Kasimova-TR wants to merge 3 commits into
BerriAI:litellm_internal_stagingfrom
Iana-Kasimova-TR:fix-vertex-ai-gemma-maas-26083

Conversation

@Iana-Kasimova-TR

Copy link
Copy Markdown

Relevant issues

Fixes #26083

Summary

vertex_ai/google/gemma-4-26b-a4b-it-maas previously fell through to VertexAIModelRoute.NON_GEMINI (legacy chat-bison/text-bison path) and
didn't work. Implements owtaylor's plan from the issue thread:

  • Add google/gemma- to PartnerModelPrefixes so the model is detected as a Vertex partner model.
  • Add it to should_use_openai_handler so it flows through the existing OpenAI-compatible handler (same as gpt-oss/llama3/qwen MaaS).
  • No OpenAIGPTConfig subclass added — base OpenAIGPTConfig works (owtaylor: "works at least basically without that").
  • No "exclude from gemma detection" code needed: get_vertex_ai_model_route checks "gemma/" in model (with a slash), and google/gemma-...
    does not match. Regression-tested.
  • Model metadata added to model_prices_and_context_window.json with supported_regions: ["global"] so the global endpoint URL is used by
    default.

Note: PR #28004 is tagged as fixing #26083 but actually addresses a different self-hosted-URL bug in the model_garden path; this PR is the
real fix for #26083.

Pre-Submission checklist

  • Tests in tests/test_litellm/ (4 detection tests + 5 URL/end-to-end tests)
  • make test-unit passes on touched files (104 vertex_ai tests green)
  • uv run black . applied
  • PR scope isolated to a single issue

Test plan

  • Smoke test against a real Vertex AI project: litellm.completion(model="vertex_ai/google/gemma-4-26b-a4b-it-maas", messages=[{"role": "user", "content": "hello"}], vertex_project=PROJECT)
  • Verify streaming works
  • Verify function-calling and structured-output requests pass through
  • Verify multimodal (image / PDF) input works

Pricing

Pricing entry: $0.15 / 1M input tokens, $0.60 / 1M output tokens. Please verify against
https://cloud.google.com/vertex-ai/generative-ai/pricing before merge.

…AI path

Fixes BerriAI#26083

vertex_ai/google/gemma-4-26b-a4b-it-maas previously fell through to the
NON_GEMINI route. Per owtaylor's plan on BerriAI#26083: add the google/gemma-
prefix to PartnerModelPrefixes so is_vertex_partner_model picks it up
and should_use_openai_handler routes it to the OpenAI-compatible
/endpoints/openapi/chat/completions URL. No gemma-detection exclusion
needed (the "gemma/" check uses a slash, which google/gemma-... doesn't
match). No OpenAIGPTConfig subclass needed — works with the base handler.
@Iana-Kasimova-TR

Copy link
Copy Markdown
Author

@greptileai

@CLAassistant

Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.


Iana seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
You have signed the CLA already but the status is still pending? Let us recheck it.

@codecov

codecov Bot commented May 15, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@greptile-apps

greptile-apps Bot commented May 15, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR fixes routing for vertex_ai/google/gemma-*-maas models, which previously fell through to the legacy NON_GEMINI path. It adds "google/gemma-" to PartnerModelPrefixes and should_use_openai_handler, directing these models through the existing OpenAI-compatible partner-model handler.

  • Adds GEMMA_MAAS_PREFIX = "google/gemma-" to PartnerModelPrefixes and wires it into both detection methods in main.py; the startswith vs. in asymmetry between is_vertex_partner_model and should_use_openai_handler is pre-existing and harmless here.
  • Adds a correctly-categorised model entry (litellm_provider: vertex_ai-openai_models, supported_regions: ["global"]) for vertex_ai/google/gemma-4-26b-a4b-it-maas in both JSON price files.
  • Adds 4 unit detection tests and 5 mock-only URL/end-to-end tests covering the global endpoint path, the region-override behaviour, and a negative regression guard against google/gemini-* false-positive detection.

Confidence Score: 5/5

Safe to merge; the change is a targeted routing addition that follows established partner-model patterns and is well covered by mock tests.

The routing logic is a small, additive change that mirrors what every other OpenAI-compatible MaaS partner already does. The negative regression guard confirms google/gemini-* models are unaffected. The JSON metadata now uses the correct litellm_provider category and standard cloud.google.com source URL, addressing all previously raised concerns.

No files require special attention.

Important Files Changed

Filename Overview
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py Adds GEMMA_MAAS_PREFIX = "google/gemma-" to PartnerModelPrefixes and threads it through both is_vertex_partner_model (startswith) and should_use_openai_handler (substring); logic is consistent with existing partner model patterns and correctly excludes google/gemini-* models.
model_prices_and_context_window.json Adds vertex_ai/google/gemma-4-26b-a4b-it-maas with litellm_provider="vertex_ai-openai_models", supported_regions=["global"], and correct pricing; addresses previously flagged litellm_provider categorisation issue.
litellm/model_prices_and_context_window_backup.json Mirror of the root JSON entry; also uses litellm_provider="vertex_ai-openai_models" and is consistent with the primary file.
tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py Adds 4 new detection/routing tests for Gemma MaaS; remaining diff is black formatting only with no logic changes.
tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py New mock-only test file covering get_vertex_region, create_vertex_url, and an end-to-end acompletion flow; all network calls are properly mocked with no real HTTP requests.
tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/init.py Empty init.py to make the new gemma test directory a Python package.

Reviews (4): Last reviewed commit: "fix(vertex_ai): address greptile feedbac..." | Re-trigger Greptile

Comment thread model_prices_and_context_window.json Outdated
Comment thread model_prices_and_context_window.json Outdated
@greptile-apps

greptile-apps Bot commented May 15, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

Routes vertex_ai/google/gemma-*-maas models through the existing OpenAI-compatible partner-models handler by adding GEMMA_MAAS_PREFIX = "google/gemma-" to PartnerModelPrefixes and wiring it into both is_vertex_partner_model and should_use_openai_handler. A supported_regions: ["global"] entry in model_prices_and_context_window.json ensures the global endpoint URL is selected automatically.

  • Routing fix (vertex_ai_partner_models/main.py): three-line change — one new enum value, one new startswith branch in is_vertex_partner_model, one new entry in the OPENAI_LIKE_VERTEX_PROVIDERS list for should_use_openai_handler. Consistent with the existing pattern for all other partner-model prefixes.
  • Model metadata (model_prices_and_context_window.json): adds pricing, context window, capability flags, and supported_regions: [\"global\"]; litellm_provider is set to \"vertex_ai-language-models\" (Gemini category) rather than \"vertex_ai-openai_models\" as used by the other OpenAI-compatible MaaS entries.
  • Tests: nine new mock-only tests covering partner-model detection, OpenAI-handler selection, route assertion, negative-case (gemini-1.5-pro not misdetected), global URL construction, region resolution, and end-to-end acompletion URL verification.

Confidence Score: 4/5

The core routing change is minimal and correct; the main risk is limited to a metadata misclassification in the JSON entry that does not affect call routing.

The routing logic in main.py is a clean, pattern-consistent change well-covered by the new tests. The only concern is in model_prices_and_context_window.json where litellm_provider uses the Gemini-model category instead of the OpenAI-compatible MaaS partner model category used by all analogous models, causing the key to be stored with its full vertex_ai/ prefix in vertex_language_models rather than stripped and placed in vertex_openai_models.

model_prices_and_context_window.json — litellm_provider value and source URL should be verified before merge.

Important Files Changed

Filename Overview
litellm/llms/vertex_ai/vertex_ai_partner_models/main.py Adds GEMMA_MAAS_PREFIX to PartnerModelPrefixes and wires it into both is_vertex_partner_model and should_use_openai_handler; logic is clean and consistent with all other partner-model prefixes.
model_prices_and_context_window.json Adds pricing/metadata entry for vertex_ai/google/gemma-4-26b-a4b-it-maas; litellm_provider is set to "vertex_ai-language-models" (Gemini category) instead of "vertex_ai-openai_models" as used by other OpenAI-compatible MaaS partner models, and the source URL uses a non-standard docs.cloud.google.com subdomain.
tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/gemma/test_vertex_ai_gemma_global_endpoint.py New mock-only test file covering global URL construction, region resolution, and end-to-end acompletion routing for Gemma MaaS; no real network calls.
tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py Adds four new detection/routing guard tests for Gemma MaaS; remaining diff is black reformatting of existing tests with no logic changes.

Reviews (2): Last reviewed commit: "fix(vertex_ai): route google/gemma-*-maa..." | Re-trigger Greptile

Comment thread model_prices_and_context_window.json Outdated
Comment thread model_prices_and_context_window.json Outdated
@oss-pr-review-agent-shin

Copy link
Copy Markdown
Contributor

🤖 litellm-agent: This PR is currently BLOCKED from merge.

Score: 4/5

Why blocked:

  • 1 unresolved reviewer concern (greptile) (unresolved_concern, -1 pts)

Details: Score docked for: 1 unresolved reviewer concern (greptile). 1 check also red on neighboring PRs (misc / Run tests) — infra-wide noise, no penalty.

Fix the issues above and push an update — the bot will re-review automatically.

Note: This bot is still in beta and might not always work as expected. Please share any feedback via Slack.

@Iana-Kasimova-TR

Copy link
Copy Markdown
Author

@greptileai

1 similar comment
@icep87

icep87 commented May 22, 2026

Copy link
Copy Markdown
Contributor

@greptileai

@icep87

icep87 commented May 22, 2026

Copy link
Copy Markdown
Contributor

Any plans on getting this merge soon, as currently this blocks usage of Gemma from Vertex model garden

@oss-pr-review-agent-shin

Copy link
Copy Markdown
Contributor

🤖 litellm-agent: This PR is currently BLOCKED from merge.

Score: 0/5

Why blocked:

  • karpathy needs_human — Three capability flags — supports_vision, supports_function_calling, supports_tool_choice — are set to true, but the PR's own test plan leaves 'Verify function-calling and structured-output requests pass through' and 'Verify multimodal (image/PDF) input works' unchecked. If Gemma 4 MaaS on Vertex does not expose these capabilities on the global /endpoints/openapi path, callers using get_model_info to gate feature use will silently attempt unsupported calls (karpathy, -2 pts)
  • all Phase B agent checks non-approving (phase_b_none_approved, -5 pts)

Details: Score docked for: karpathy needs_human — Three capability flags — supports_vision, supports_function_calling, supports_tool_choice — are set to true, but the PR's own test plan leaves 'Verify function-calling and structured-output requests pass through' and 'Verify multimodal (image/PDF) input works' unchecked. If Gemma 4 MaaS on Vertex does not expose these capabilities on the global /endpoints/openapi path, callers using get_model_info to gate feature use will silently attempt unsupported calls; all Phase B agent checks non-approving (karpathy + security + coverage gap).

Fix the issues above and push an update — the bot will re-review automatically.

Note: This bot is still in beta and might not always work as expected. Please share any feedback via Slack.

@icep87

icep87 commented May 25, 2026

Copy link
Copy Markdown
Contributor

@Iana-Kasimova-TR Will you be able to fix the issue so it's now blocked?

icep87 added a commit to icep87/litellm that referenced this pull request May 26, 2026
…or Gemma MaaS

Addresses oss-pr-review-agent-shin feedback on PR BerriAI#28010:

   P1 (patch target): Added a comment explaining why patching
   litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler is correct —
   get_async_httpx_client() (defined in http_handler.py) instantiates
   AsyncHTTPHandler within that module's scope, so the definition-site patch
   intercepts it. Without the mock the test raises AuthenticationError,
   confirming it never silently passes.

   P2 (partner-provider regression guard): Added
   test_gemma_routes_through_openai_handler() which calls
   VertexAIPartnerModels.should_use_openai_handler() directly, so if Gemma's
   routing to VertexPartnerProvider.llama ever changes the URL-shape tests
   below it become a real regression guard rather than an unanchored unit test.

   Also added:
   - test_gemma_maas_supports_function_calling / supports_vision — capability
     flag checks via patch.dict(litellm.model_cost)
   - test_vertex_ai_gemma_function_calling_passthrough — tools + tool_choice
     forwarded in the request body
   - test_vertex_ai_gemma_vision_passthrough — image_url part survives
     transformation to the global endpoint
   Added:
   - test_gemma_maas_supports_function_calling — verifies the utility returns True
     when the model_cost entry carries supports_function_calling=true
   - test_gemma_maas_supports_vision — same for supports_vision
   - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice
     appear in the JSON body POSTed to /endpoints/openapi/chat/completions
   - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts
     survive transformation and reach the global endpoint URL
Sameerlite pushed a commit that referenced this pull request May 28, 2026
…AI path - clone of #28010 (#28846)

* fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path

Fixes #26083

vertex_ai/google/gemma-4-26b-a4b-it-maas previously fell through to the
NON_GEMINI route. Per owtaylor's plan on #26083: add the google/gemma-
prefix to PartnerModelPrefixes so is_vertex_partner_model picks it up
and should_use_openai_handler routes it to the OpenAI-compatible
/endpoints/openapi/chat/completions URL. No gemma-detection exclusion
needed (the "gemma/" check uses a slash, which google/gemma-... doesn't
match). No OpenAIGPTConfig subclass needed — works with the base handler.

* fix(vertex_ai): mark gemma-4-26b-a4b-it-maas as vision-capable (empirically verified)

* fix(vertex_ai): address greptile feedback — provider category, canonical URL, sync backup

* test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS

   Addresses oss-pr-review-agent-shin feedback on PR #28010:
   supports_function_calling, supports_tool_choice, and supports_vision were
   marked true but had no tests proving the payloads actually reached the
   OpenAI-compatible endpoint.

   Added:
   - test_gemma_maas_supports_function_calling — verifies the utility returns True
     when the model_cost entry carries supports_function_calling=true
   - test_gemma_maas_supports_vision — same for supports_vision
   - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice
     appear in the JSON body POSTed to /endpoints/openapi/chat/completions
   - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts
     survive transformation and reach the global endpoint URL

* fix: Delete uv.lock

* test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS

Addresses oss-pr-review-agent-shin feedback on PR #28010:

   P1 (patch target): Added a comment explaining why patching
   litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler is correct —
   get_async_httpx_client() (defined in http_handler.py) instantiates
   AsyncHTTPHandler within that module's scope, so the definition-site patch
   intercepts it. Without the mock the test raises AuthenticationError,
   confirming it never silently passes.

   P2 (partner-provider regression guard): Added
   test_gemma_routes_through_openai_handler() which calls
   VertexAIPartnerModels.should_use_openai_handler() directly, so if Gemma's
   routing to VertexPartnerProvider.llama ever changes the URL-shape tests
   below it become a real regression guard rather than an unanchored unit test.

   Also added:
   - test_gemma_maas_supports_function_calling / supports_vision — capability
     flag checks via patch.dict(litellm.model_cost)
   - test_vertex_ai_gemma_function_calling_passthrough — tools + tool_choice
     forwarded in the request body
   - test_vertex_ai_gemma_vision_passthrough — image_url part survives
     transformation to the global endpoint
   Added:
   - test_gemma_maas_supports_function_calling — verifies the utility returns True
     when the model_cost entry carries supports_function_calling=true
   - test_gemma_maas_supports_vision — same for supports_vision
   - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice
     appear in the JSON body POSTed to /endpoints/openapi/chat/completions
   - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts
     survive transformation and reach the global endpoint URL

* fix: proper patch for unit tests

---------

Co-authored-by: Iana <[email protected]>
@mateo-berri mateo-berri mentioned this pull request Jun 2, 2026
mateo-berri added a commit that referenced this pull request Jun 2, 2026
* Cato Networks guardrail, based on Aim (#26597)

* Aim was acquired by Cato Networks, creating Cato Networks guardrail based on Aim

* Add more tests

* Move test so they are reached by codecov coverage

* base URL trailing slashes

* Support Lemonade runtime context metadata (#28135)

* Support Lemonade runtime context metadata

* Add provider hook for runtime model metadata

* Address provider model info review feedback

Keep the runtime model info hook duck-typed instead of extending the base model-info class, and avoid importing ModelInfoBase from Ollama common utilities to reduce CodeQL cyclic-import noise.

Co-authored-by: openhands <[email protected]>

* Fix CI after staging rebase

Relax the Ollama runtime metadata return annotation to match the provider-hook dict response and update the Google Interactions OpenAPI status expectation for the current live spec.

Co-authored-by: openhands <[email protected]>

* Normalize Lemonade runtime model metadata

* Avoid leaking Ollama metadata auth

* Avoid leaking Lemonade metadata auth

---------

Co-authored-by: Graham Neubig <[email protected]>
Co-authored-by: openhands <[email protected]>

* fix(cato): address guardrail review feedback

Use proxy-authenticated user identity, forward moderation hook return values,
and ensure streaming sender tasks are cancelled and awaited on exit.

Co-authored-by: Cursor <[email protected]>

* fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path - clone of #28010 (#28846)

* fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path

Fixes #26083

vertex_ai/google/gemma-4-26b-a4b-it-maas previously fell through to the
NON_GEMINI route. Per owtaylor's plan on #26083: add the google/gemma-
prefix to PartnerModelPrefixes so is_vertex_partner_model picks it up
and should_use_openai_handler routes it to the OpenAI-compatible
/endpoints/openapi/chat/completions URL. No gemma-detection exclusion
needed (the "gemma/" check uses a slash, which google/gemma-... doesn't
match). No OpenAIGPTConfig subclass needed — works with the base handler.

* fix(vertex_ai): mark gemma-4-26b-a4b-it-maas as vision-capable (empirically verified)

* fix(vertex_ai): address greptile feedback — provider category, canonical URL, sync backup

* test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS

   Addresses oss-pr-review-agent-shin feedback on PR #28010:
   supports_function_calling, supports_tool_choice, and supports_vision were
   marked true but had no tests proving the payloads actually reached the
   OpenAI-compatible endpoint.

   Added:
   - test_gemma_maas_supports_function_calling — verifies the utility returns True
     when the model_cost entry carries supports_function_calling=true
   - test_gemma_maas_supports_vision — same for supports_vision
   - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice
     appear in the JSON body POSTed to /endpoints/openapi/chat/completions
   - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts
     survive transformation and reach the global endpoint URL

* fix: Delete uv.lock

* test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS

Addresses oss-pr-review-agent-shin feedback on PR #28010:

   P1 (patch target): Added a comment explaining why patching
   litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler is correct —
   get_async_httpx_client() (defined in http_handler.py) instantiates
   AsyncHTTPHandler within that module's scope, so the definition-site patch
   intercepts it. Without the mock the test raises AuthenticationError,
   confirming it never silently passes.

   P2 (partner-provider regression guard): Added
   test_gemma_routes_through_openai_handler() which calls
   VertexAIPartnerModels.should_use_openai_handler() directly, so if Gemma's
   routing to VertexPartnerProvider.llama ever changes the URL-shape tests
   below it become a real regression guard rather than an unanchored unit test.

   Also added:
   - test_gemma_maas_supports_function_calling / supports_vision — capability
     flag checks via patch.dict(litellm.model_cost)
   - test_vertex_ai_gemma_function_calling_passthrough — tools + tool_choice
     forwarded in the request body
   - test_vertex_ai_gemma_vision_passthrough — image_url part survives
     transformation to the global endpoint
   Added:
   - test_gemma_maas_supports_function_calling — verifies the utility returns True
     when the model_cost entry carries supports_function_calling=true
   - test_gemma_maas_supports_vision — same for supports_vision
   - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice
     appear in the JSON body POSTed to /endpoints/openapi/chat/completions
   - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts
     survive transformation and reach the global endpoint URL

* fix: proper patch for unit tests

---------

Co-authored-by: Iana <[email protected]>

* fix(cato): guardrail all completion choices on output

When n > 1, only choices[0] was analyzed and redacted. Iterate every
Choices entry so block and anonymize actions apply to all completions.

Co-authored-by: Cursor <[email protected]>

* Fix review

* fix(cato_networks): harden output anonymize handling and restructure nested UI routes

Guard against empty redacted_output and empty all_redacted_messages from Cato.
Restructure nested admin UI HTML exports to index.html so extensionless routes work.

Co-authored-by: Cursor <[email protected]>

* Fix mypy

* fix(cato): guard missing policy_drill_down and all_redacted_messages keys

* fix(cato): avoid KeyError bypassing block action on missing analysis_result

* fix(cato): preserve non-text message fields during anonymize

Rebuild redacted messages from the original messages, overwriting only
content, so tool_calls, tool_call_id, name and multimodal fields survive
the anonymize action.

* fix(cato): preserve trailing messages when fewer redacted messages returned

Avoid silently truncating the conversation in _anonymize_request when Cato
returns fewer redacted messages than were sent, and isolate the no-api-key
config test from a pre-existing CATO_API_KEY environment variable.

* fix(cato,model-info): preserve stream block signal on sender teardown; forward api_key in dynamic model-info lookup

Suppress ConnectionClosed (alongside CancelledError) when tearing down the
Cato streaming sender task so a backend ConnectionClosed cannot mask the
original StreamingCallbackError (e.g. a guardrail block) raised by the
receive loop.

Thread api_key through get_model_info -> _get_model_info_helper so an
explicit key reaches a provider's dynamic get_model_info for a caller-supplied
api_base. Previously only api_base was forwarded, so authenticated Ollama and
Lemonade servers at a custom base could only be queried unauthenticated.

* fix(cato): surface mid-stream forwarding errors instead of blocking on recv

If the upstream LLM stream errors mid-flight, the sender task dies before
sending the terminal done frame, so the consumer would block on websocket.recv()
until Cato closes the connection. Race recv against the sender task and raise the
stored sender exception promptly as a StreamingCallbackError.

* fix(cato): drop spoofable end_user_id from guardrail user identity

Only the key/JWT-bound user_email is a trusted identity. end_user_id is
resolved from caller-supplied request fields (OpenAI user param, headers,
metadata), so an authenticated caller with no bound user_email could set it
to another user's email and have LiteLLM forward x-cato-user-email for that
victim, poisoning Cato audit and policy attribution. Forward only user_email
and omit the header otherwise.

* fix(cato): harden output anonymize path against missing content key

* fix(cato): fall back to original message when redacted content key is missing

* refactor(model-info): drop unused api_key from cached model-info helper

_cached_get_model_info_helper is only called by the cost-tracking hot path,
which never authenticates, so the api_key parameter was never populated.
Keeping it in the lru_cache key offered no benefit and risked fragmenting
the high-RPS cache and retaining credential strings per entry.

* fix(cato): preserve None content on tool-call-only choices in output hook

* fix(ollama): respect static-model guard in OllamaConfig.get_model_info

Delegate to OllamaModelInfo.get_model_info so statically-priced Ollama
models short-circuit before the /api/show network call instead of
hitting the server unconditionally.

* fix(lemonade,ollama): treat empty api_key as unset to avoid leaking server creds

An empty-string api_key was treated as an explicit key, so it passed the
guard meant to keep server-side credentials off caller-supplied bases and
then fell back through the env/global key chain. A caller could point
api_base at a server they control and send api_key="" to receive the
configured provider key in the Authorization header. Gate the credential
fallback on the api_key being truthy instead of merely not-None.

* fix(cato): inspect and redact Responses-API input, not just messages

The guardrail only read data["messages"], so /v1/responses requests, which
carry their text in data["input"], reached Cato as an empty message list
and bypassed inspection entirely. Send build_inspection_messages(data) so
both shapes are analyzed, and write anonymized results back with
apply_redacted_messages_back when the request used input.

* perf(utils): keep api_key out of get_model_info lru_cache key

* fix(cato): propagate ssl_verify to streaming WebSocket connection

The streaming hook applied ssl_verify only to the HTTP handler; the
websockets.connect() call used default verification, so a custom Cato
instance behind TLS with a self-signed cert worked for non-streaming
calls but failed every streaming request. Resolve the ssl_verify setting
into the connect() ssl argument, mirroring the HTTP handler.

* refactor(utils): rename shadowing local in _get_model_info_helper

* fix(cato): flatten multimodal chat content before inspection

Chat Completions requests whose message content is a multimodal parts
array were posted to Cato as the raw OpenAI parts, so text inside
content: [{"type":"text", ...}] reached the model without Cato ever
inspecting the string. Flatten each message's list content to plain text
while keeping the list 1:1 with the request so the index-based redaction
write-back stays valid; Responses-API input requests still go through
build_inspection_messages.

* test(lemonade): clear get_model_info cache around api_base test

* fix(cato): inspect and redact Responses-API input even when messages present

_inspection_messages returned early once messages was non-empty, so a
/v1/responses caller could place benign text in messages and disallowed
text in input and have only messages reach Cato while the model used
input. Inspect both fields and write anonymize redactions back to input
as well as the index-aligned messages.

* test(log_db_metrics): assert table_name event_metadata contract

log_db_metrics now emits minimal event_metadata via _safe_db_event_metadata
(table_name only, function_name/function_kwargs/function_args dropped as
redundant with call_type and unsafe to stamp on a span). The success-path
test still asserted function_name membership and crashed with TypeError on
the None metadata returned when no table_name is passed. Pass a table_name
and assert the surfaced contract instead.

* fix(cato): inspect and redact completion prompt and Responses-API instructions

The Cato guardrail only inspected chat messages and the Responses-API input field, so blocked text placed in the legacy /v1/completions prompt or the /v1/responses instructions field reached the model without ever being sent to Cato. Both fields are now appended as synthetic inspection messages, and the anonymize path slices Cato's redactions back to the field they came from.

* fix(cato): serialize non-str/bytes websocket chunks before forwarding

* fix(cato): inspect tool descriptions and tool-call arguments

* fix(cato): map redacted output by assistant index; restore get_model_info.cache_info

* fix(cato): block output even when detection_message is null/empty

A block_action returned by Cato on the output hook whose detection_message
was null or empty was let through to the caller: the truthiness guard on
detection_message skipped the HTTPException and the unblocked response was
returned. Raise the HTTPException directly in _handle_block_action_on_output
so the output path blocks unconditionally, mirroring the input path.

* fix(cato): inspect and redact nested tool param and legacy function descriptions

Tool/function parameter descriptions and the legacy functions[] array are
forwarded to the model but were not seen by Cato, so blocked text hidden there
bypassed inspection and anonymization. Recursively walk every description string
in tools[].function and functions[] schemas for both the analyze payload and the
anonymize write-back.

* fix(cato): traverse schema descriptions iteratively to satisfy recursive detector

The nested walk() generator recursed over tool/function JSON schemas with no
depth bound, which the recursive_detector code-quality gate rejects. Replace it
with an explicit-stack DFS that yields the same (container, key) refs in the
same pre-order, so schema description redaction is unchanged.

* fix(cato): inspect and redact response_format JSON schema descriptions

response_format json_schema descriptions are forwarded to the model, so
blocked text hidden in nested schema descriptions could bypass Cato
inspection and redaction. Extend the schema-description walk to cover
response_format alongside tools and legacy functions.

* fix(cato): skip output rewrite when Cato returns no redaction

Return None from call_cato_guardrail_on_output on monitor/no-action so the
post-call hook only mutates the message when there is an actual redaction,
instead of redundantly re-writing the original content.

* refactor(utils): resolve explicit api_key model info without the cache

Move the model-info build into a non-cached _build_model_info helper and drop
api_key from the lru-cached _cached_get_model_info signature. Both cached
helpers now take the same (model, provider, api_base) key and never forward
api_key, while explicit per-caller keys are resolved through the builder
directly instead of reaching into the cache wrapper's __wrapped__.

* fix(cato): inspect and redact non-description schema string values

Tool, function and response_format JSON schemas forward more than just
description text to the model. enum, const, default, examples and title
values are sent verbatim, so blocked content hidden in any of them
bypassed Cato inspection and redaction. Walk those schema string values
alongside descriptions on both the inspection and anonymize paths.

* fix(model-info): surface swallowed dynamic model-info errors

The provider-specific get_model_info dispatch falls back to the static cost
map when a provider's dynamic lookup raises, which is intentional graceful
degradation. Previously the exception was discarded with a bare debug line,
so a real failure (e.g. a provider whose get_model_info signature does not
accept api_key) was invisible. Log the exception at warning level with the
model and provider context so the fallback is diagnosable.

* fix(cato): inspect and redact Responses API output in post-call hook

The post-call success hook only handled ModelResponse, so /v1/responses
(which returns a ResponsesAPIResponse) bypassed the Cato output guardrail.
Extract and inspect/redact every output_text content block and function-call
arguments string, blocking on a block action, so generated text cannot escape
inspection by using the Responses API.

* chore: reset _experimental/out folder

* chore(ui): remove orphaned prebuilt dashboard chunk files

The _experimental/out manifests are byte-identical to the base branch, so the
served dashboard already matches base. 436 unreferenced Next.js chunk files had
accumulated in the directory and are not loaded by any manifest; removing them
restores the committed UI artifacts to the base build and drops the artifact
churn from this PR's diff.

* fix(guardrails,ollama): forward ssl_verify to Cato init and raise_for_status on /api/show

---------

Co-authored-by: Alex Yaroslavsky <[email protected]>
Co-authored-by: Graham Neubig <[email protected]>
Co-authored-by: Graham Neubig <[email protected]>
Co-authored-by: openhands <[email protected]>
Co-authored-by: Cursor <[email protected]>
Co-authored-by: Piotr Placzko <[email protected]>
Co-authored-by: Iana <[email protected]>
Co-authored-by: mateo-berri <[email protected]>
yugborana pushed a commit to yugborana/litellm that referenced this pull request Jun 2, 2026
* Cato Networks guardrail, based on Aim (BerriAI#26597)

* Aim was acquired by Cato Networks, creating Cato Networks guardrail based on Aim

* Add more tests

* Move test so they are reached by codecov coverage

* base URL trailing slashes

* Support Lemonade runtime context metadata (BerriAI#28135)

* Support Lemonade runtime context metadata

* Add provider hook for runtime model metadata

* Address provider model info review feedback

Keep the runtime model info hook duck-typed instead of extending the base model-info class, and avoid importing ModelInfoBase from Ollama common utilities to reduce CodeQL cyclic-import noise.

Co-authored-by: openhands <[email protected]>

* Fix CI after staging rebase

Relax the Ollama runtime metadata return annotation to match the provider-hook dict response and update the Google Interactions OpenAPI status expectation for the current live spec.

Co-authored-by: openhands <[email protected]>

* Normalize Lemonade runtime model metadata

* Avoid leaking Ollama metadata auth

* Avoid leaking Lemonade metadata auth

---------

Co-authored-by: Graham Neubig <[email protected]>
Co-authored-by: openhands <[email protected]>

* fix(cato): address guardrail review feedback

Use proxy-authenticated user identity, forward moderation hook return values,
and ensure streaming sender tasks are cancelled and awaited on exit.

Co-authored-by: Cursor <[email protected]>

* fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path - clone of BerriAI#28010 (BerriAI#28846)

* fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path

Fixes BerriAI#26083

vertex_ai/google/gemma-4-26b-a4b-it-maas previously fell through to the
NON_GEMINI route. Per owtaylor's plan on BerriAI#26083: add the google/gemma-
prefix to PartnerModelPrefixes so is_vertex_partner_model picks it up
and should_use_openai_handler routes it to the OpenAI-compatible
/endpoints/openapi/chat/completions URL. No gemma-detection exclusion
needed (the "gemma/" check uses a slash, which google/gemma-... doesn't
match). No OpenAIGPTConfig subclass needed — works with the base handler.

* fix(vertex_ai): mark gemma-4-26b-a4b-it-maas as vision-capable (empirically verified)

* fix(vertex_ai): address greptile feedback — provider category, canonical URL, sync backup

* test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS

   Addresses oss-pr-review-agent-shin feedback on PR BerriAI#28010:
   supports_function_calling, supports_tool_choice, and supports_vision were
   marked true but had no tests proving the payloads actually reached the
   OpenAI-compatible endpoint.

   Added:
   - test_gemma_maas_supports_function_calling — verifies the utility returns True
     when the model_cost entry carries supports_function_calling=true
   - test_gemma_maas_supports_vision — same for supports_vision
   - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice
     appear in the JSON body POSTed to /endpoints/openapi/chat/completions
   - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts
     survive transformation and reach the global endpoint URL

* fix: Delete uv.lock

* test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS

Addresses oss-pr-review-agent-shin feedback on PR BerriAI#28010:

   P1 (patch target): Added a comment explaining why patching
   litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler is correct —
   get_async_httpx_client() (defined in http_handler.py) instantiates
   AsyncHTTPHandler within that module's scope, so the definition-site patch
   intercepts it. Without the mock the test raises AuthenticationError,
   confirming it never silently passes.

   P2 (partner-provider regression guard): Added
   test_gemma_routes_through_openai_handler() which calls
   VertexAIPartnerModels.should_use_openai_handler() directly, so if Gemma's
   routing to VertexPartnerProvider.llama ever changes the URL-shape tests
   below it become a real regression guard rather than an unanchored unit test.

   Also added:
   - test_gemma_maas_supports_function_calling / supports_vision — capability
     flag checks via patch.dict(litellm.model_cost)
   - test_vertex_ai_gemma_function_calling_passthrough — tools + tool_choice
     forwarded in the request body
   - test_vertex_ai_gemma_vision_passthrough — image_url part survives
     transformation to the global endpoint
   Added:
   - test_gemma_maas_supports_function_calling — verifies the utility returns True
     when the model_cost entry carries supports_function_calling=true
   - test_gemma_maas_supports_vision — same for supports_vision
   - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice
     appear in the JSON body POSTed to /endpoints/openapi/chat/completions
   - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts
     survive transformation and reach the global endpoint URL

* fix: proper patch for unit tests

---------

Co-authored-by: Iana <[email protected]>

* fix(cato): guardrail all completion choices on output

When n > 1, only choices[0] was analyzed and redacted. Iterate every
Choices entry so block and anonymize actions apply to all completions.

Co-authored-by: Cursor <[email protected]>

* Fix review

* fix(cato_networks): harden output anonymize handling and restructure nested UI routes

Guard against empty redacted_output and empty all_redacted_messages from Cato.
Restructure nested admin UI HTML exports to index.html so extensionless routes work.

Co-authored-by: Cursor <[email protected]>

* Fix mypy

* fix(cato): guard missing policy_drill_down and all_redacted_messages keys

* fix(cato): avoid KeyError bypassing block action on missing analysis_result

* fix(cato): preserve non-text message fields during anonymize

Rebuild redacted messages from the original messages, overwriting only
content, so tool_calls, tool_call_id, name and multimodal fields survive
the anonymize action.

* fix(cato): preserve trailing messages when fewer redacted messages returned

Avoid silently truncating the conversation in _anonymize_request when Cato
returns fewer redacted messages than were sent, and isolate the no-api-key
config test from a pre-existing CATO_API_KEY environment variable.

* fix(cato,model-info): preserve stream block signal on sender teardown; forward api_key in dynamic model-info lookup

Suppress ConnectionClosed (alongside CancelledError) when tearing down the
Cato streaming sender task so a backend ConnectionClosed cannot mask the
original StreamingCallbackError (e.g. a guardrail block) raised by the
receive loop.

Thread api_key through get_model_info -> _get_model_info_helper so an
explicit key reaches a provider's dynamic get_model_info for a caller-supplied
api_base. Previously only api_base was forwarded, so authenticated Ollama and
Lemonade servers at a custom base could only be queried unauthenticated.

* fix(cato): surface mid-stream forwarding errors instead of blocking on recv

If the upstream LLM stream errors mid-flight, the sender task dies before
sending the terminal done frame, so the consumer would block on websocket.recv()
until Cato closes the connection. Race recv against the sender task and raise the
stored sender exception promptly as a StreamingCallbackError.

* fix(cato): drop spoofable end_user_id from guardrail user identity

Only the key/JWT-bound user_email is a trusted identity. end_user_id is
resolved from caller-supplied request fields (OpenAI user param, headers,
metadata), so an authenticated caller with no bound user_email could set it
to another user's email and have LiteLLM forward x-cato-user-email for that
victim, poisoning Cato audit and policy attribution. Forward only user_email
and omit the header otherwise.

* fix(cato): harden output anonymize path against missing content key

* fix(cato): fall back to original message when redacted content key is missing

* refactor(model-info): drop unused api_key from cached model-info helper

_cached_get_model_info_helper is only called by the cost-tracking hot path,
which never authenticates, so the api_key parameter was never populated.
Keeping it in the lru_cache key offered no benefit and risked fragmenting
the high-RPS cache and retaining credential strings per entry.

* fix(cato): preserve None content on tool-call-only choices in output hook

* fix(ollama): respect static-model guard in OllamaConfig.get_model_info

Delegate to OllamaModelInfo.get_model_info so statically-priced Ollama
models short-circuit before the /api/show network call instead of
hitting the server unconditionally.

* fix(lemonade,ollama): treat empty api_key as unset to avoid leaking server creds

An empty-string api_key was treated as an explicit key, so it passed the
guard meant to keep server-side credentials off caller-supplied bases and
then fell back through the env/global key chain. A caller could point
api_base at a server they control and send api_key="" to receive the
configured provider key in the Authorization header. Gate the credential
fallback on the api_key being truthy instead of merely not-None.

* fix(cato): inspect and redact Responses-API input, not just messages

The guardrail only read data["messages"], so /v1/responses requests, which
carry their text in data["input"], reached Cato as an empty message list
and bypassed inspection entirely. Send build_inspection_messages(data) so
both shapes are analyzed, and write anonymized results back with
apply_redacted_messages_back when the request used input.

* perf(utils): keep api_key out of get_model_info lru_cache key

* fix(cato): propagate ssl_verify to streaming WebSocket connection

The streaming hook applied ssl_verify only to the HTTP handler; the
websockets.connect() call used default verification, so a custom Cato
instance behind TLS with a self-signed cert worked for non-streaming
calls but failed every streaming request. Resolve the ssl_verify setting
into the connect() ssl argument, mirroring the HTTP handler.

* refactor(utils): rename shadowing local in _get_model_info_helper

* fix(cato): flatten multimodal chat content before inspection

Chat Completions requests whose message content is a multimodal parts
array were posted to Cato as the raw OpenAI parts, so text inside
content: [{"type":"text", ...}] reached the model without Cato ever
inspecting the string. Flatten each message's list content to plain text
while keeping the list 1:1 with the request so the index-based redaction
write-back stays valid; Responses-API input requests still go through
build_inspection_messages.

* test(lemonade): clear get_model_info cache around api_base test

* fix(cato): inspect and redact Responses-API input even when messages present

_inspection_messages returned early once messages was non-empty, so a
/v1/responses caller could place benign text in messages and disallowed
text in input and have only messages reach Cato while the model used
input. Inspect both fields and write anonymize redactions back to input
as well as the index-aligned messages.

* test(log_db_metrics): assert table_name event_metadata contract

log_db_metrics now emits minimal event_metadata via _safe_db_event_metadata
(table_name only, function_name/function_kwargs/function_args dropped as
redundant with call_type and unsafe to stamp on a span). The success-path
test still asserted function_name membership and crashed with TypeError on
the None metadata returned when no table_name is passed. Pass a table_name
and assert the surfaced contract instead.

* fix(cato): inspect and redact completion prompt and Responses-API instructions

The Cato guardrail only inspected chat messages and the Responses-API input field, so blocked text placed in the legacy /v1/completions prompt or the /v1/responses instructions field reached the model without ever being sent to Cato. Both fields are now appended as synthetic inspection messages, and the anonymize path slices Cato's redactions back to the field they came from.

* fix(cato): serialize non-str/bytes websocket chunks before forwarding

* fix(cato): inspect tool descriptions and tool-call arguments

* fix(cato): map redacted output by assistant index; restore get_model_info.cache_info

* fix(cato): block output even when detection_message is null/empty

A block_action returned by Cato on the output hook whose detection_message
was null or empty was let through to the caller: the truthiness guard on
detection_message skipped the HTTPException and the unblocked response was
returned. Raise the HTTPException directly in _handle_block_action_on_output
so the output path blocks unconditionally, mirroring the input path.

* fix(cato): inspect and redact nested tool param and legacy function descriptions

Tool/function parameter descriptions and the legacy functions[] array are
forwarded to the model but were not seen by Cato, so blocked text hidden there
bypassed inspection and anonymization. Recursively walk every description string
in tools[].function and functions[] schemas for both the analyze payload and the
anonymize write-back.

* fix(cato): traverse schema descriptions iteratively to satisfy recursive detector

The nested walk() generator recursed over tool/function JSON schemas with no
depth bound, which the recursive_detector code-quality gate rejects. Replace it
with an explicit-stack DFS that yields the same (container, key) refs in the
same pre-order, so schema description redaction is unchanged.

* fix(cato): inspect and redact response_format JSON schema descriptions

response_format json_schema descriptions are forwarded to the model, so
blocked text hidden in nested schema descriptions could bypass Cato
inspection and redaction. Extend the schema-description walk to cover
response_format alongside tools and legacy functions.

* fix(cato): skip output rewrite when Cato returns no redaction

Return None from call_cato_guardrail_on_output on monitor/no-action so the
post-call hook only mutates the message when there is an actual redaction,
instead of redundantly re-writing the original content.

* refactor(utils): resolve explicit api_key model info without the cache

Move the model-info build into a non-cached _build_model_info helper and drop
api_key from the lru-cached _cached_get_model_info signature. Both cached
helpers now take the same (model, provider, api_base) key and never forward
api_key, while explicit per-caller keys are resolved through the builder
directly instead of reaching into the cache wrapper's __wrapped__.

* fix(cato): inspect and redact non-description schema string values

Tool, function and response_format JSON schemas forward more than just
description text to the model. enum, const, default, examples and title
values are sent verbatim, so blocked content hidden in any of them
bypassed Cato inspection and redaction. Walk those schema string values
alongside descriptions on both the inspection and anonymize paths.

* fix(model-info): surface swallowed dynamic model-info errors

The provider-specific get_model_info dispatch falls back to the static cost
map when a provider's dynamic lookup raises, which is intentional graceful
degradation. Previously the exception was discarded with a bare debug line,
so a real failure (e.g. a provider whose get_model_info signature does not
accept api_key) was invisible. Log the exception at warning level with the
model and provider context so the fallback is diagnosable.

* fix(cato): inspect and redact Responses API output in post-call hook

The post-call success hook only handled ModelResponse, so /v1/responses
(which returns a ResponsesAPIResponse) bypassed the Cato output guardrail.
Extract and inspect/redact every output_text content block and function-call
arguments string, blocking on a block action, so generated text cannot escape
inspection by using the Responses API.

* chore: reset _experimental/out folder

* chore(ui): remove orphaned prebuilt dashboard chunk files

The _experimental/out manifests are byte-identical to the base branch, so the
served dashboard already matches base. 436 unreferenced Next.js chunk files had
accumulated in the directory and are not loaded by any manifest; removing them
restores the committed UI artifacts to the base build and drops the artifact
churn from this PR's diff.

* fix(guardrails,ollama): forward ssl_verify to Cato init and raise_for_status on /api/show

---------

Co-authored-by: Alex Yaroslavsky <[email protected]>
Co-authored-by: Graham Neubig <[email protected]>
Co-authored-by: Graham Neubig <[email protected]>
Co-authored-by: openhands <[email protected]>
Co-authored-by: Cursor <[email protected]>
Co-authored-by: Piotr Placzko <[email protected]>
Co-authored-by: Iana <[email protected]>
Co-authored-by: mateo-berri <[email protected]>
@Sameerlite

Copy link
Copy Markdown
Collaborator

The fix for #26083 was merged. Closing this

@Sameerlite Sameerlite closed this Jun 3, 2026
fzowl pushed a commit to fzowl/litellm that referenced this pull request Jun 24, 2026
* Cato Networks guardrail, based on Aim (BerriAI#26597)

* Aim was acquired by Cato Networks, creating Cato Networks guardrail based on Aim

* Add more tests

* Move test so they are reached by codecov coverage

* base URL trailing slashes

* Support Lemonade runtime context metadata (BerriAI#28135)

* Support Lemonade runtime context metadata

* Add provider hook for runtime model metadata

* Address provider model info review feedback

Keep the runtime model info hook duck-typed instead of extending the base model-info class, and avoid importing ModelInfoBase from Ollama common utilities to reduce CodeQL cyclic-import noise.

Co-authored-by: openhands <[email protected]>

* Fix CI after staging rebase

Relax the Ollama runtime metadata return annotation to match the provider-hook dict response and update the Google Interactions OpenAPI status expectation for the current live spec.

Co-authored-by: openhands <[email protected]>

* Normalize Lemonade runtime model metadata

* Avoid leaking Ollama metadata auth

* Avoid leaking Lemonade metadata auth

---------

Co-authored-by: Graham Neubig <[email protected]>
Co-authored-by: openhands <[email protected]>

* fix(cato): address guardrail review feedback

Use proxy-authenticated user identity, forward moderation hook return values,
and ensure streaming sender tasks are cancelled and awaited on exit.

Co-authored-by: Cursor <[email protected]>

* fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path - clone of BerriAI#28010 (BerriAI#28846)

* fix(vertex_ai): route google/gemma-*-maas through partner-models OpenAI path

Fixes BerriAI#26083

vertex_ai/google/gemma-4-26b-a4b-it-maas previously fell through to the
NON_GEMINI route. Per owtaylor's plan on BerriAI#26083: add the google/gemma-
prefix to PartnerModelPrefixes so is_vertex_partner_model picks it up
and should_use_openai_handler routes it to the OpenAI-compatible
/endpoints/openapi/chat/completions URL. No gemma-detection exclusion
needed (the "gemma/" check uses a slash, which google/gemma-... doesn't
match). No OpenAIGPTConfig subclass needed — works with the base handler.

* fix(vertex_ai): mark gemma-4-26b-a4b-it-maas as vision-capable (empirically verified)

* fix(vertex_ai): address greptile feedback — provider category, canonical URL, sync backup

* test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS

   Addresses oss-pr-review-agent-shin feedback on PR BerriAI#28010:
   supports_function_calling, supports_tool_choice, and supports_vision were
   marked true but had no tests proving the payloads actually reached the
   OpenAI-compatible endpoint.

   Added:
   - test_gemma_maas_supports_function_calling — verifies the utility returns True
     when the model_cost entry carries supports_function_calling=true
   - test_gemma_maas_supports_vision — same for supports_vision
   - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice
     appear in the JSON body POSTed to /endpoints/openapi/chat/completions
   - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts
     survive transformation and reach the global endpoint URL

* fix: Delete uv.lock

* test(vertex_ai): add function-calling and vision pass-through tests for Gemma MaaS

Addresses oss-pr-review-agent-shin feedback on PR BerriAI#28010:

   P1 (patch target): Added a comment explaining why patching
   litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler is correct —
   get_async_httpx_client() (defined in http_handler.py) instantiates
   AsyncHTTPHandler within that module's scope, so the definition-site patch
   intercepts it. Without the mock the test raises AuthenticationError,
   confirming it never silently passes.

   P2 (partner-provider regression guard): Added
   test_gemma_routes_through_openai_handler() which calls
   VertexAIPartnerModels.should_use_openai_handler() directly, so if Gemma's
   routing to VertexPartnerProvider.llama ever changes the URL-shape tests
   below it become a real regression guard rather than an unanchored unit test.

   Also added:
   - test_gemma_maas_supports_function_calling / supports_vision — capability
     flag checks via patch.dict(litellm.model_cost)
   - test_vertex_ai_gemma_function_calling_passthrough — tools + tool_choice
     forwarded in the request body
   - test_vertex_ai_gemma_vision_passthrough — image_url part survives
     transformation to the global endpoint
   Added:
   - test_gemma_maas_supports_function_calling — verifies the utility returns True
     when the model_cost entry carries supports_function_calling=true
   - test_gemma_maas_supports_vision — same for supports_vision
   - test_vertex_ai_gemma_function_calling_passthrough — verifies tools + tool_choice
     appear in the JSON body POSTed to /endpoints/openapi/chat/completions
   - test_vertex_ai_gemma_vision_passthrough — verifies image_url content parts
     survive transformation and reach the global endpoint URL

* fix: proper patch for unit tests

---------

Co-authored-by: Iana <[email protected]>

* fix(cato): guardrail all completion choices on output

When n > 1, only choices[0] was analyzed and redacted. Iterate every
Choices entry so block and anonymize actions apply to all completions.

Co-authored-by: Cursor <[email protected]>

* Fix review

* fix(cato_networks): harden output anonymize handling and restructure nested UI routes

Guard against empty redacted_output and empty all_redacted_messages from Cato.
Restructure nested admin UI HTML exports to index.html so extensionless routes work.

Co-authored-by: Cursor <[email protected]>

* Fix mypy

* fix(cato): guard missing policy_drill_down and all_redacted_messages keys

* fix(cato): avoid KeyError bypassing block action on missing analysis_result

* fix(cato): preserve non-text message fields during anonymize

Rebuild redacted messages from the original messages, overwriting only
content, so tool_calls, tool_call_id, name and multimodal fields survive
the anonymize action.

* fix(cato): preserve trailing messages when fewer redacted messages returned

Avoid silently truncating the conversation in _anonymize_request when Cato
returns fewer redacted messages than were sent, and isolate the no-api-key
config test from a pre-existing CATO_API_KEY environment variable.

* fix(cato,model-info): preserve stream block signal on sender teardown; forward api_key in dynamic model-info lookup

Suppress ConnectionClosed (alongside CancelledError) when tearing down the
Cato streaming sender task so a backend ConnectionClosed cannot mask the
original StreamingCallbackError (e.g. a guardrail block) raised by the
receive loop.

Thread api_key through get_model_info -> _get_model_info_helper so an
explicit key reaches a provider's dynamic get_model_info for a caller-supplied
api_base. Previously only api_base was forwarded, so authenticated Ollama and
Lemonade servers at a custom base could only be queried unauthenticated.

* fix(cato): surface mid-stream forwarding errors instead of blocking on recv

If the upstream LLM stream errors mid-flight, the sender task dies before
sending the terminal done frame, so the consumer would block on websocket.recv()
until Cato closes the connection. Race recv against the sender task and raise the
stored sender exception promptly as a StreamingCallbackError.

* fix(cato): drop spoofable end_user_id from guardrail user identity

Only the key/JWT-bound user_email is a trusted identity. end_user_id is
resolved from caller-supplied request fields (OpenAI user param, headers,
metadata), so an authenticated caller with no bound user_email could set it
to another user's email and have LiteLLM forward x-cato-user-email for that
victim, poisoning Cato audit and policy attribution. Forward only user_email
and omit the header otherwise.

* fix(cato): harden output anonymize path against missing content key

* fix(cato): fall back to original message when redacted content key is missing

* refactor(model-info): drop unused api_key from cached model-info helper

_cached_get_model_info_helper is only called by the cost-tracking hot path,
which never authenticates, so the api_key parameter was never populated.
Keeping it in the lru_cache key offered no benefit and risked fragmenting
the high-RPS cache and retaining credential strings per entry.

* fix(cato): preserve None content on tool-call-only choices in output hook

* fix(ollama): respect static-model guard in OllamaConfig.get_model_info

Delegate to OllamaModelInfo.get_model_info so statically-priced Ollama
models short-circuit before the /api/show network call instead of
hitting the server unconditionally.

* fix(lemonade,ollama): treat empty api_key as unset to avoid leaking server creds

An empty-string api_key was treated as an explicit key, so it passed the
guard meant to keep server-side credentials off caller-supplied bases and
then fell back through the env/global key chain. A caller could point
api_base at a server they control and send api_key="" to receive the
configured provider key in the Authorization header. Gate the credential
fallback on the api_key being truthy instead of merely not-None.

* fix(cato): inspect and redact Responses-API input, not just messages

The guardrail only read data["messages"], so /v1/responses requests, which
carry their text in data["input"], reached Cato as an empty message list
and bypassed inspection entirely. Send build_inspection_messages(data) so
both shapes are analyzed, and write anonymized results back with
apply_redacted_messages_back when the request used input.

* perf(utils): keep api_key out of get_model_info lru_cache key

* fix(cato): propagate ssl_verify to streaming WebSocket connection

The streaming hook applied ssl_verify only to the HTTP handler; the
websockets.connect() call used default verification, so a custom Cato
instance behind TLS with a self-signed cert worked for non-streaming
calls but failed every streaming request. Resolve the ssl_verify setting
into the connect() ssl argument, mirroring the HTTP handler.

* refactor(utils): rename shadowing local in _get_model_info_helper

* fix(cato): flatten multimodal chat content before inspection

Chat Completions requests whose message content is a multimodal parts
array were posted to Cato as the raw OpenAI parts, so text inside
content: [{"type":"text", ...}] reached the model without Cato ever
inspecting the string. Flatten each message's list content to plain text
while keeping the list 1:1 with the request so the index-based redaction
write-back stays valid; Responses-API input requests still go through
build_inspection_messages.

* test(lemonade): clear get_model_info cache around api_base test

* fix(cato): inspect and redact Responses-API input even when messages present

_inspection_messages returned early once messages was non-empty, so a
/v1/responses caller could place benign text in messages and disallowed
text in input and have only messages reach Cato while the model used
input. Inspect both fields and write anonymize redactions back to input
as well as the index-aligned messages.

* test(log_db_metrics): assert table_name event_metadata contract

log_db_metrics now emits minimal event_metadata via _safe_db_event_metadata
(table_name only, function_name/function_kwargs/function_args dropped as
redundant with call_type and unsafe to stamp on a span). The success-path
test still asserted function_name membership and crashed with TypeError on
the None metadata returned when no table_name is passed. Pass a table_name
and assert the surfaced contract instead.

* fix(cato): inspect and redact completion prompt and Responses-API instructions

The Cato guardrail only inspected chat messages and the Responses-API input field, so blocked text placed in the legacy /v1/completions prompt or the /v1/responses instructions field reached the model without ever being sent to Cato. Both fields are now appended as synthetic inspection messages, and the anonymize path slices Cato's redactions back to the field they came from.

* fix(cato): serialize non-str/bytes websocket chunks before forwarding

* fix(cato): inspect tool descriptions and tool-call arguments

* fix(cato): map redacted output by assistant index; restore get_model_info.cache_info

* fix(cato): block output even when detection_message is null/empty

A block_action returned by Cato on the output hook whose detection_message
was null or empty was let through to the caller: the truthiness guard on
detection_message skipped the HTTPException and the unblocked response was
returned. Raise the HTTPException directly in _handle_block_action_on_output
so the output path blocks unconditionally, mirroring the input path.

* fix(cato): inspect and redact nested tool param and legacy function descriptions

Tool/function parameter descriptions and the legacy functions[] array are
forwarded to the model but were not seen by Cato, so blocked text hidden there
bypassed inspection and anonymization. Recursively walk every description string
in tools[].function and functions[] schemas for both the analyze payload and the
anonymize write-back.

* fix(cato): traverse schema descriptions iteratively to satisfy recursive detector

The nested walk() generator recursed over tool/function JSON schemas with no
depth bound, which the recursive_detector code-quality gate rejects. Replace it
with an explicit-stack DFS that yields the same (container, key) refs in the
same pre-order, so schema description redaction is unchanged.

* fix(cato): inspect and redact response_format JSON schema descriptions

response_format json_schema descriptions are forwarded to the model, so
blocked text hidden in nested schema descriptions could bypass Cato
inspection and redaction. Extend the schema-description walk to cover
response_format alongside tools and legacy functions.

* fix(cato): skip output rewrite when Cato returns no redaction

Return None from call_cato_guardrail_on_output on monitor/no-action so the
post-call hook only mutates the message when there is an actual redaction,
instead of redundantly re-writing the original content.

* refactor(utils): resolve explicit api_key model info without the cache

Move the model-info build into a non-cached _build_model_info helper and drop
api_key from the lru-cached _cached_get_model_info signature. Both cached
helpers now take the same (model, provider, api_base) key and never forward
api_key, while explicit per-caller keys are resolved through the builder
directly instead of reaching into the cache wrapper's __wrapped__.

* fix(cato): inspect and redact non-description schema string values

Tool, function and response_format JSON schemas forward more than just
description text to the model. enum, const, default, examples and title
values are sent verbatim, so blocked content hidden in any of them
bypassed Cato inspection and redaction. Walk those schema string values
alongside descriptions on both the inspection and anonymize paths.

* fix(model-info): surface swallowed dynamic model-info errors

The provider-specific get_model_info dispatch falls back to the static cost
map when a provider's dynamic lookup raises, which is intentional graceful
degradation. Previously the exception was discarded with a bare debug line,
so a real failure (e.g. a provider whose get_model_info signature does not
accept api_key) was invisible. Log the exception at warning level with the
model and provider context so the fallback is diagnosable.

* fix(cato): inspect and redact Responses API output in post-call hook

The post-call success hook only handled ModelResponse, so /v1/responses
(which returns a ResponsesAPIResponse) bypassed the Cato output guardrail.
Extract and inspect/redact every output_text content block and function-call
arguments string, blocking on a block action, so generated text cannot escape
inspection by using the Responses API.

* chore: reset _experimental/out folder

* chore(ui): remove orphaned prebuilt dashboard chunk files

The _experimental/out manifests are byte-identical to the base branch, so the
served dashboard already matches base. 436 unreferenced Next.js chunk files had
accumulated in the directory and are not loaded by any manifest; removing them
restores the committed UI artifacts to the base build and drops the artifact
churn from this PR's diff.

* fix(guardrails,ollama): forward ssl_verify to Cato init and raise_for_status on /api/show

---------

Co-authored-by: Alex Yaroslavsky <[email protected]>
Co-authored-by: Graham Neubig <[email protected]>
Co-authored-by: Graham Neubig <[email protected]>
Co-authored-by: openhands <[email protected]>
Co-authored-by: Cursor <[email protected]>
Co-authored-by: Piotr Placzko <[email protected]>
Co-authored-by: Iana <[email protected]>
Co-authored-by: mateo-berri <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature]: Support Gemma4-26B-A4B-IT MaaS on Vertex AI

4 participants