mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
fix(vertex-live): resolve a Live setup model before logging reads it
A client that named a bare gateway alias logged the session as "unknown" and billed nothing, because the model was read off the raw setup frame and the extractor only yields a name when the string already contains "/models/". The rewriter qualifies that same model a few lines later for the upstream, so the supported client form, an alias, was the one that went unbilled. Resolving through the rewriter first means the real model reaches the logging object, and from there the cost map. A route with no rewriter, which is every non-Live passthrough, hands the frame over untouched. (cherry picked from commit 573982803df612fd94144e2e06dd647f8530d4e8)
This commit is contained in:
parent
b1262b05f4
commit
9580b89bb1
2 changed files with 116 additions and 1 deletions
|
|
@ -2090,6 +2090,22 @@ def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Calla
|
|||
return json.dumps({**message, "setup": {**setup, "model": rewritten_model}}) # mutable-ok: one-shot json payload
|
||||
|
||||
|
||||
def _resolved_vertex_live_setup(
|
||||
setup_data: Mapping[str, object], setup_model_rewriter: Callable[[str], str] | None
|
||||
) -> Mapping[str, object]:
|
||||
"""
|
||||
Give the model extractor the same fully qualified path the upstream will receive.
|
||||
|
||||
Clients may name a bare gateway alias, which the rewriter turns into a ``projects/...`` path before
|
||||
it reaches Vertex. The extractor only reads a path containing ``/models/``, so running it on the raw
|
||||
frame logs the session as ``unknown`` at no cost, which is precisely the supported client form
|
||||
"""
|
||||
setup_model: Final = setup_data.get("model")
|
||||
if setup_model_rewriter is None or not isinstance(setup_model, str):
|
||||
return setup_data
|
||||
return {**setup_data, "model": setup_model_rewriter(setup_model)}
|
||||
|
||||
|
||||
def _truncated_close_reason(reason: str) -> str:
|
||||
"""
|
||||
Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character
|
||||
|
|
@ -2314,7 +2330,12 @@ async def websocket_passthrough_request(
|
|||
setup_data,
|
||||
)
|
||||
if isinstance(setup_data, dict) and "model" in setup_data:
|
||||
extracted_model = _extract_model_from_vertex_ai_setup(setup_data)
|
||||
# Resolve the alias first: a client may name a bare gateway model,
|
||||
# which carries no "/models/" for the extractor to read, so reading
|
||||
# the raw frame leaves the session logged as "unknown" and unbilled.
|
||||
extracted_model = _extract_model_from_vertex_ai_setup(
|
||||
_resolved_vertex_live_setup(setup_data, setup_model_rewriter)
|
||||
)
|
||||
if extracted_model:
|
||||
kwargs["model"] = extracted_model
|
||||
kwargs["custom_llm_provider"] = "vertex_ai-language-models"
|
||||
|
|
|
|||
|
|
@ -5030,6 +5030,100 @@ async def test_websocket_passthrough_rewrites_gateway_alias_setup_model():
|
|||
assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"setup_model",
|
||||
["gemini-live-2.5-flash", "models/gemini-live-2.5-flash", "publishers/google/models/gemini-live-2.5-flash"],
|
||||
)
|
||||
def test_vertex_live_setup_model_resolves_before_extraction(setup_model):
|
||||
"""A bare gateway alias left the session logged as ``unknown`` at zero cost.
|
||||
|
||||
The model was read off the raw client frame, and the extractor only yields a name when the string
|
||||
already contains ``/models/``. The rewriter qualifies it a few lines later for the upstream, so a
|
||||
client that addressed the gateway the documented way, by alias, logged no model and therefore
|
||||
resolved no cost-map entry. Resolving first is what puts the real name on the logging object.
|
||||
"""
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
_build_vertex_live_setup_model_rewriter,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
_extract_model_from_vertex_ai_setup,
|
||||
_resolved_vertex_live_setup,
|
||||
)
|
||||
|
||||
rewriter = _build_vertex_live_setup_model_rewriter(
|
||||
vertex_project="proj-db", vertex_location="global", llm_router=None
|
||||
)
|
||||
setup_data = {"model": setup_model}
|
||||
|
||||
resolved = _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup_data, rewriter))
|
||||
|
||||
assert resolved == "gemini-live-2.5-flash", "an unresolved setup model logs the session as 'unknown'"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_logs_a_bare_alias_setup_model():
|
||||
"""End to end through the relay: a bare alias must reach the logging object as a real model name.
|
||||
|
||||
This is the call-site half of the fix. The helper tests above pass even if extraction moves back
|
||||
before the rewrite, so this one drives the real websocket relay and asserts on what got logged,
|
||||
which is the name the cost map is looked up by. An unbilled session logs ``unknown``.
|
||||
"""
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
_build_vertex_live_setup_model_rewriter,
|
||||
)
|
||||
|
||||
upstream_ws = RecordingUpstreamWebSocket()
|
||||
setup_frame = json.dumps({"setup": {"model": "gemini-live-2.5-flash"}})
|
||||
websocket = _client_websocket(
|
||||
AsyncMock(
|
||||
side_effect=[
|
||||
{"type": "websocket.receive", "text": setup_frame},
|
||||
{"type": "websocket.disconnect"},
|
||||
]
|
||||
)
|
||||
)
|
||||
built = []
|
||||
real_logging = litellm.litellm_core_utils.litellm_logging.Logging
|
||||
|
||||
def _capture(*args, **kwargs):
|
||||
obj = real_logging(*args, **kwargs)
|
||||
built.append(obj)
|
||||
return obj
|
||||
|
||||
with _patched_websocket_passthrough_environment(upstream_ws):
|
||||
with patch("litellm.litellm_core_utils.litellm_logging.Logging", side_effect=_capture):
|
||||
await websocket_passthrough_request(
|
||||
websocket=websocket,
|
||||
target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent",
|
||||
custom_headers={"Authorization": "Bearer token"},
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
forward_headers=False,
|
||||
endpoint="/vertex_ai/live",
|
||||
accept_websocket=False,
|
||||
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
|
||||
vertex_project="proj-db", vertex_location="global", llm_router=None
|
||||
),
|
||||
)
|
||||
|
||||
assert built, "the relay should have built a logging object"
|
||||
assert built[0].model == "gemini-live-2.5-flash", "a bare alias must not log as 'unknown'"
|
||||
|
||||
|
||||
def test_vertex_live_setup_resolution_is_inert_without_a_rewriter():
|
||||
"""Non-Live passthrough routes pass no rewriter, so the frame must be handed over untouched."""
|
||||
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
|
||||
_extract_model_from_vertex_ai_setup,
|
||||
_resolved_vertex_live_setup,
|
||||
)
|
||||
|
||||
setup_data = {"model": "projects/p/locations/global/publishers/google/models/gemini-live-2.5-flash"}
|
||||
|
||||
assert _resolved_vertex_live_setup(setup_data, None) is setup_data
|
||||
assert _extract_model_from_vertex_ai_setup(_resolved_vertex_live_setup(setup_data, None)) == (
|
||||
"gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("rcvd_close", [None, "abnormal", "no_status"])
|
||||
async def test_websocket_passthrough_does_not_relay_unsendable_upstream_close(rcvd_close):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue