mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(passthrough): trust the declared Google provider and Vertex-shaped paths on custom bases
The failure path only named the Google provider from the URL when the method was one the success handler tracks or the host was a Google API host, so a Gemini countTokens failure behind a custom GEMINI_API_BASE (or a Vertex one behind a custom base) still logged an empty model. A route that declares gemini or vertex_ai now wins outright, and a URL whose path carries /projects/<p>/locations/<l>/ counts as Vertex on any host, while a tracked method on a non-Google host no longer passes as Vertex.
This commit is contained in:
parent
571927705a
commit
f80801985a
3 changed files with 32 additions and 9 deletions
|
|
@ -9,7 +9,11 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import VERTEX_BATCH_PREDICTION_JOBS_ROUTE
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.vertex_ai.common_utils import get_vertex_location_from_url, get_vertex_model_id_from_url
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
get_vertex_location_from_url,
|
||||
get_vertex_model_id_from_url,
|
||||
get_vertex_project_id_from_url,
|
||||
)
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator as VertexModelResponseIterator,
|
||||
)
|
||||
|
|
@ -563,9 +567,16 @@ class VertexPassthroughLoggingHandler:
|
|||
return None
|
||||
|
||||
@staticmethod
|
||||
def is_google_ai_host(url: str) -> bool:
|
||||
hostname: Final = urlparse(url).hostname
|
||||
return hostname is not None and hostname.endswith((GEMINI_API_HOST_SUFFIX, VERTEX_API_HOST_SUFFIX))
|
||||
def is_google_ai_url(url: str) -> bool:
|
||||
parsed_url: Final = urlparse(url)
|
||||
if parsed_url.hostname is not None and parsed_url.hostname.endswith(
|
||||
(GEMINI_API_HOST_SUFFIX, VERTEX_API_HOST_SUFFIX)
|
||||
):
|
||||
return True
|
||||
return (
|
||||
get_vertex_project_id_from_url(parsed_url.path) is not None
|
||||
and get_vertex_location_from_url(parsed_url.path) is not None
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def custom_llm_provider_from_url(url: str) -> str:
|
||||
|
|
|
|||
|
|
@ -732,12 +732,15 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | Non
|
|||
metadata.setdefault("standard_logging_guardrail_information", list(entries))
|
||||
|
||||
|
||||
_GOOGLE_PASSTHROUGH_PROVIDERS: Final = frozenset(
|
||||
(litellm.LlmProviders.GEMINI.value, litellm.LlmProviders.VERTEX_AI.value)
|
||||
)
|
||||
|
||||
|
||||
def _passthrough_google_provider(url_route: str, custom_llm_provider: str | None) -> str | None:
|
||||
if pass_through_endpoint_logging.is_gemini_route(url_route, custom_llm_provider):
|
||||
return "gemini"
|
||||
if pass_through_endpoint_logging.is_vertex_route(url_route) or VertexPassthroughLoggingHandler.is_google_ai_host(
|
||||
url_route
|
||||
):
|
||||
if custom_llm_provider in _GOOGLE_PASSTHROUGH_PROVIDERS:
|
||||
return custom_llm_provider
|
||||
if VertexPassthroughLoggingHandler.is_google_ai_url(url_route):
|
||||
return VertexPassthroughLoggingHandler.custom_llm_provider_from_url(url_route)
|
||||
return None
|
||||
|
||||
|
|
|
|||
|
|
@ -5821,6 +5821,10 @@ _VERTEX_SEARCH_TARGET = (
|
|||
"https://discoveryengine.googleapis.com/v1/projects/p/locations/global"
|
||||
"/dataStores/default/servingConfigs/default:search"
|
||||
)
|
||||
_RELAYED_GEMINI_COUNT_TOKENS_TARGET = "http://127.0.0.1:30732/v1beta/models/gemini-3.8-flash:countTokens"
|
||||
_RELAYED_VERTEX_COUNT_TOKENS_TARGET = (
|
||||
"http://127.0.0.1:30732/v1/projects/p/locations/us-east5/publishers/google/models/gemini-3.8-flash:countTokens"
|
||||
)
|
||||
|
||||
|
||||
async def _run_passthrough_failure(
|
||||
|
|
@ -5899,6 +5903,8 @@ def _upstream_400_as_event_stream(target: str) -> httpx.Response:
|
|||
(_VERTEX_FLASH_TARGET, None, True, "vertex_ai"),
|
||||
(_GEMINI_FLASH_TARGET.replace(":generateContent", ":countTokens"), "gemini", False, "gemini"),
|
||||
(_VERTEX_FLASH_TARGET.replace(":generateContent", ":countTokens"), None, False, "vertex_ai"),
|
||||
(_RELAYED_GEMINI_COUNT_TOKENS_TARGET, "gemini", False, "gemini"),
|
||||
(_RELAYED_VERTEX_COUNT_TOKENS_TARGET, None, False, "vertex_ai"),
|
||||
],
|
||||
)
|
||||
async def test_passthrough_upstream_error_logs_model_from_url(
|
||||
|
|
@ -6017,7 +6023,10 @@ async def test_passthrough_mid_stream_drop_logs_model_from_url(
|
|||
"",
|
||||
"vertex_ai",
|
||||
),
|
||||
({"contents": []}, _RELAYED_GEMINI_COUNT_TOKENS_TARGET, "gemini", "gemini-3.8-flash", "gemini"),
|
||||
({"contents": []}, _RELAYED_VERTEX_COUNT_TOKENS_TARGET, None, "gemini-3.8-flash", "vertex_ai"),
|
||||
({"contents": []}, "https://api.example.com/v1/models/gpt-x:countTokens", None, "", None),
|
||||
({"contents": []}, "https://api.example.com/v1/models/gpt-x:predict", None, "", None),
|
||||
({"contents": []}, "https://api.example.com/v1/models/gpt-x:thing", None, "", None),
|
||||
(None, None, None, "", None),
|
||||
],
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue