From 8dda73c935f5e7cfbd26b2b81354b801b4e27f68 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 22 Aug 2026 16:35:57 -0700 Subject: [PATCH 01/55] test: tighten regression tests added in #37974 Five of the tests could pass without the behavior they guard being correct. The passthrough spend tests derived their expected spend in setup_method from whatever cost map was live rather than the pinned checked-in one. The test_main.py cost fixture cleared only one of the two price caches, leaving billing to read stale prices while the assertions read the pinned map. The gpt-5.6 bridge test parametrized over two suffixes the version check discards, so both cases were identical. The anthropic flush helper swallowed the loop-binding RuntimeError it exists to report. The cache-write test pinned a literal 1.25 rate ratio unrelated to the bug it guards. --- .../test_openai_cache_write_cost.py | 2 +- ...erimental_pass_through_messages_handler.py | 2 +- ...test_openai_passthrough_logging_handler.py | 18 +++++---- tests/test_litellm/test_main.py | 40 +++++++++++++------ 4 files changed, 40 insertions(+), 22 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py index 92cb417f96c..18acfeda07d 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_openai_cache_write_cost.py @@ -39,7 +39,7 @@ def test_openai_cache_write_tokens_billed_at_the_cache_creation_rate(local_model input_rate = rates["input_cost_per_token"] cache_write_rate = rates["cache_creation_input_token_cost"] output_rate = rates["output_cost_per_token"] - assert cache_write_rate == pytest.approx(input_rate * 1.25) + assert cache_write_rate > input_rate prompt_tokens = 12317 cache_write_tokens = 12314 diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index 9e58ded81bd..1e6dacbb7cb 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -1187,7 +1187,7 @@ async def _flush_logging_worker(capture: "_SuccessPayloadCapture") -> None: await asyncio.sleep(0) try: await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10.0) - except (asyncio.TimeoutError, RuntimeError): + except asyncio.TimeoutError: pass deadline = asyncio.get_running_loop().time() + 10.0 while not capture.payloads and asyncio.get_running_loop().time() < deadline: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py index 6f9142c85df..8f7146d8647 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_openai_passthrough_logging_handler.py @@ -1830,8 +1830,10 @@ class TestOpenAIPassthroughResponsesStreamingSpendLog: def setup_method(self): self.start_time = datetime.now() self.end_time = datetime.now() + + def _expected_spend(self) -> float: rates = litellm.model_cost[self.MODEL_MAP_KEY] - self.expected_spend = ( + return ( self.INPUT_TOKENS * rates["input_cost_per_token"] + self.OUTPUT_TOKENS * rates["output_cost_per_token"] ) @@ -1912,7 +1914,7 @@ class TestOpenAIPassthroughResponsesStreamingSpendLog: logging_obj.model_call_details["custom_llm_provider"] = "openai" return logging_obj - def test_streamed_responses_passthrough_spend_log_is_priced(self): + def test_streamed_responses_passthrough_spend_log_is_priced(self, local_model_cost_map): """The spend row books the same tokens, spend and `resp_` id as the buffered call.""" result = OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( litellm_logging_obj=self._logging_obj(), @@ -1940,7 +1942,7 @@ class TestOpenAIPassthroughResponsesStreamingSpendLog: assert spend_log_row["prompt_tokens"] == self.INPUT_TOKENS assert spend_log_row["completion_tokens"] == self.OUTPUT_TOKENS assert spend_log_row["total_tokens"] == self.INPUT_TOKENS + self.OUTPUT_TOKENS - assert spend_log_row["spend"] == self.expected_spend + assert spend_log_row["spend"] == pytest.approx(self._expected_spend()) assert spend_log_row["request_id"] == self.RESPONSE_ID assert spend_log_row["model"] == "gpt-4o-mini" @@ -1965,7 +1967,6 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: def setup_method(self): self.start_time = datetime.now() self.end_time = datetime.now() - self.expected_spend = self.PROMPT_TOKENS * litellm.model_cost[self.MODEL]["input_cost_per_token"] self.response_body = { "object": "list", "data": [{"object": "embedding", "index": 0, "embedding": [0.0, 1.0]}], @@ -1974,6 +1975,9 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: } self.request_body = {"model": self.MODEL, "input": "hello"} + def _expected_spend(self) -> float: + return self.PROMPT_TOKENS * litellm.model_cost[self.MODEL]["input_cost_per_token"] + def _create_mock_httpx_response(self) -> httpx.Response: mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 @@ -1999,7 +2003,7 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: ) return logging_obj - def test_embeddings_passthrough_spend_log_is_priced(self): + def test_embeddings_passthrough_spend_log_is_priced(self, local_model_cost_map): """The dispatched call books prompt tokens and cost onto the spend row.""" dispatched = PassThroughEndpointLogging().normalize_llm_passthrough_logging_payload( httpx_response=self._create_mock_httpx_response(), @@ -2018,7 +2022,7 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: ) assert dispatched["standard_logging_response_object"] is not None - assert dispatched["kwargs"]["response_cost"] == self.expected_spend + assert dispatched["kwargs"]["response_cost"] == pytest.approx(self._expected_spend()) spend_log_row = get_logging_payload( kwargs=dispatched["kwargs"], @@ -2029,7 +2033,7 @@ class TestOpenAIPassthroughEmbeddingsSpendLog: assert spend_log_row["prompt_tokens"] == self.PROMPT_TOKENS assert spend_log_row["total_tokens"] == self.PROMPT_TOKENS - assert spend_log_row["spend"] == self.expected_spend + assert spend_log_row["spend"] == pytest.approx(self._expected_spend()) assert spend_log_row["model"] == self.MODEL assert spend_log_row["custom_llm_provider"] == "openai" assert spend_log_row["request_id"] == self.CALL_ID diff --git a/tests/test_litellm/test_main.py b/tests/test_litellm/test_main.py index 99b1cc826aa..2599b81cd2c 100644 --- a/tests/test_litellm/test_main.py +++ b/tests/test_litellm/test_main.py @@ -850,15 +850,25 @@ def test_responses_api_bridge_check_gpt_5_4_tools_with_default_reasoning_routes_ assert model_info.get("mode") == "responses" -@pytest.mark.parametrize("model_name", ["gpt-5.6-sol", "gpt-5.6-luna", "gpt-5.6-terra"]) +@pytest.mark.parametrize( + "model_name, expected_mode", + [ + pytest.param("gpt-5.6-sol", "responses", id="above-boundary-bridges"), + pytest.param("gpt-5.1", None, id="below-boundary-stays-chat"), + ], +) def test_responses_api_bridge_check_gpt_5_6_tools_with_default_reasoning_routes_to_responses( - monkeypatch, model_name + monkeypatch, model_name, expected_mode ): """ - The whole gpt-5.6 family must bridge on function tools alone. The bridge used to - require an explicit reasoning_effort, so a gpt-5.6 call carrying tools and no effort - was rejected with "Function tools with reasoning_effort are not supported for - gpt-5.6-sol in /v1/chat/completions". + gpt-5.6 must bridge on function tools alone. The bridge used to require an explicit + reasoning_effort, so a gpt-5.6 call carrying tools and no effort was rejected with + "Function tools with reasoning_effort are not supported for gpt-5.6-sol in + /v1/chat/completions". + + Paired with a model below the gpt-5.4 boundary, which must still stay on chat. The + gate parses the version and drops any suffix, so the family members bridge + identically and only the boundary distinguishes behaviour. """ import litellm from litellm.main import responses_api_bridge_check @@ -877,7 +887,7 @@ def test_responses_api_bridge_check_gpt_5_6_tools_with_default_reasoning_routes_ ) assert model == model_name - assert model_info.get("mode") == "responses" + assert model_info.get("mode") == expected_mode def test_responses_api_bridge_check_gpt_5_4_tools_with_reasoning_none_stays_chat(): @@ -2865,15 +2875,19 @@ def local_cost_map(monkeypatch): """The prices these tests assert are the checked-in ones. Setting the environment variable alone does not reload the map, so pin the map itself. - ``get_model_info`` is lru_cached, so pinning ``model_cost`` is not enough on its - own: a cached entry warmed against the network-fetched map keeps its old prices - and ``completion_cost`` bills at those while the assertions read the pinned map. - Clear on the way in and out so entries never leak across tests in either direction.""" + Prices are read through two separate lru_caches, so pinning ``model_cost`` is not + enough on its own: an entry warmed against the network-fetched map keeps its old + prices and billing reads those while the assertions read the pinned map. + ``_invalidate_model_cost_lowercase_map`` clears both caches, where + ``get_model_info.cache_clear`` reaches only one. Invalidate on the way in and out + so entries never leak across tests in either direction.""" + from litellm.utils import _invalidate_model_cost_lowercase_map + monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() yield - litellm.get_model_info.cache_clear() + _invalidate_model_cost_lowercase_map() def test_a_streamed_response_bills_the_usage_the_provider_reported(local_cost_map): From e2595e7acf9fb5f8383278a190dc66517c4f607f Mon Sep 17 00:00:00 2001 From: Mihidum Hettiyahandi <55163074+mihidumh@users.noreply.github.com> Date: Wed, 12 Aug 2026 13:21:42 +1000 Subject: [PATCH 02/55] fix(cost): honour deployment custom pricing for OCR calls MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ocr_cost() resolved pricing only via litellm.get_model_info(), a cost map lookup keyed by model name. Custom pricing for a router deployment is registered under the deployment id, and _register_custom_pricing_for_request strips pricing fields from the shared {provider}/{model} key, so the lookup could never see it. An OCR model absent from the cost map therefore billed $0 regardless of configuration, even though ocr_cost_per_page and ocr_cost_per_credit are declared CustomPricingLiteLLMParams fields. Let ocr_cost() take deployment model_info and prefer it over the map when it carries OCR pricing, with completion_cost() extracting it from litellm_logging_obj.litellm_params["metadata"]["model_info"] — the same extraction the video generation path already performs for the same reason. Behaviour is unchanged when no custom pricing is set: the map lookup still runs, and an unpriced model still returns 0.0. Fixes #36608 --- litellm/cost_calculator.py | 40 ++++++- tests/test_litellm/test_ocr_custom_pricing.py | 113 ++++++++++++++++++ 2 files changed, 149 insertions(+), 4 deletions(-) create mode 100644 tests/test_litellm/test_ocr_custom_pricing.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 6536941a094..fb415d69fdf 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -334,6 +334,8 @@ def cost_per_token( response: Any | None = None, ### REQUEST MODEL ### request_model: str | None = None, # original request model for router detection + ### DEPLOYMENT-SPECIFIC PRICING ### + custom_model_info: ModelInfo | None = None, # deployment model_info, for non-token custom pricing ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -541,6 +543,7 @@ def cost_per_token( model=model, custom_llm_provider=custom_llm_provider, response=response, + model_info=custom_model_info, ) elif ( call_type == "aretrieve_batch" @@ -1585,6 +1588,16 @@ def completion_cost( if litellm_logging_obj is not None: request_model_for_cost = litellm_logging_obj.model + # Deployment-specific model_info, for modalities whose pricing is + # not token-based and so cannot travel via custom_cost_per_token + # (e.g. OCR per-page pricing). Same extraction as the video path. + _custom_model_info: ModelInfo | None = None + if custom_pricing and litellm_logging_obj is not None: + _cm_litellm_params = getattr(litellm_logging_obj, "litellm_params", None) + if _cm_litellm_params is not None: + _cm_metadata = _cm_litellm_params.get("metadata", {}) or {} + _custom_model_info = _cm_metadata.get("model_info", None) + ( prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar, @@ -1610,6 +1623,7 @@ def completion_cost( vertex_location=vertex_location, response=completion_response, request_model=request_model_for_cost, + custom_model_info=_custom_model_info, ) # Get additional costs from provider (e.g., routing fees, infrastructure costs) @@ -1843,12 +1857,16 @@ def ocr_cost( model: str, custom_llm_provider: str | None, response: object | None = None, + model_info: ModelInfo | None = None, ) -> tuple[float, float]: """ Args: model: str - model name custom_llm_provider: Optional[str] - custom LLM provider response: Optional[Any] - response object + model_info: Optional[ModelInfo] - deployment-specific model info, used for + custom pricing. Takes precedence over the model cost map, mirroring + the video generation cost path. Returns: Tuple[float, float]: cost of OCR processing @@ -1866,10 +1884,24 @@ def ocr_cost( if response.usage_info is None: raise ValueError("OCR response usage_info is None") - try: - model_info: ModelInfo | None = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - model_info = None + ######################################################### + # Deployment-specific pricing wins over the cost map. + # + # Custom pricing set on a deployment is registered under the router's + # deployment id, while the shared "{provider}/{model}" key has its pricing + # fields stripped (see _register_custom_pricing_for_request). A cost map + # lookup therefore cannot see it, so an OCR model that is not in the map + # bills $0 no matter how it is priced in config. Prefer the caller-supplied + # model_info when it carries OCR pricing. + ######################################################### + has_custom_ocr_pricing: Final[bool] = model_info is not None and ( + model_info.get("ocr_cost_per_page") is not None or model_info.get("ocr_cost_per_credit") is not None + ) + if not has_custom_ocr_pricing: + try: + model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: + model_info = None credits: Final = getattr(response.usage_info, "credits", None) cost_per_credit = None diff --git a/tests/test_litellm/test_ocr_custom_pricing.py b/tests/test_litellm/test_ocr_custom_pricing.py new file mode 100644 index 00000000000..7cd4fc1dec1 --- /dev/null +++ b/tests/test_litellm/test_ocr_custom_pricing.py @@ -0,0 +1,113 @@ +""" +Regression tests: OCR cost must honour deployment-specific custom pricing. + +Before the fix, `ocr_cost()` resolved pricing exclusively through +`litellm.get_model_info(model=..., custom_llm_provider=...)`, i.e. a cost map +lookup keyed by model name. Custom pricing set on a deployment is registered +under the router's deployment id, and the shared "{provider}/{model}" key has +its pricing fields stripped, so the lookup could never see it. An OCR model +absent from the cost map therefore billed $0 no matter how it was priced in +config, even though `ocr_cost_per_page` / `ocr_cost_per_credit` are declared +fields of `CustomPricingLiteLLMParams`. +""" + +import pytest + +import litellm +from litellm.cost_calculator import completion_cost, ocr_cost +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo + +# A model deliberately absent from the cost map. +UNMAPPED_MODEL = "azure_ai/some-unmapped-ocr-model-for-testing" +CUSTOM_COST_PER_PAGE = 0.004 +CUSTOM_COST_PER_CREDIT = 0.25 + + +def _ocr_response(model: str, pages_processed: int = 1, credits: int | None = None) -> OCRResponse: + # NOTE: model_construct() is used rather than OCRResponse(...) because the + # OCRResponse field `object: str = "ocr"` shadows the builtin `object` used + # in the `tables` / `keyValuePairs` annotations above it, so pydantic tries + # to resolve "ocr" as a forward-referenced type and schema building fails. + # That is an unrelated defect; validation is not what these tests exercise. + usage_info = OCRUsageInfo(pages_processed=pages_processed) + if credits is not None: + usage_info.credits = credits + return OCRResponse.model_construct( + pages=[OCRPage(index=i, markdown=f"page {i}") for i in range(pages_processed)], + model=model, + usage_info=usage_info, + ) + + +def test_unmapped_ocr_model_has_no_map_pricing() -> None: + """Guard the premise: the model really is absent from the cost map.""" + assert UNMAPPED_MODEL not in litellm.model_cost + + +@pytest.mark.parametrize("pages_processed", [1, 3, 10]) +def test_ocr_cost_uses_custom_per_page_pricing(pages_processed: int) -> None: + cost, _ = ocr_cost( + model=UNMAPPED_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_MODEL, pages_processed=pages_processed), + model_info={"ocr_cost_per_page": CUSTOM_COST_PER_PAGE}, + ) + assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * pages_processed) + + +def test_ocr_cost_uses_custom_per_credit_pricing() -> None: + cost, _ = ocr_cost( + model=UNMAPPED_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_MODEL, pages_processed=2, credits=4), + model_info={"ocr_cost_per_credit": CUSTOM_COST_PER_CREDIT}, + ) + assert cost == pytest.approx(CUSTOM_COST_PER_CREDIT * 4) + + +def test_unmapped_ocr_model_without_custom_pricing_still_bills_zero() -> None: + """Unchanged behaviour when nothing is configured — no map entry, no override.""" + cost, _ = ocr_cost( + model=UNMAPPED_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_MODEL, pages_processed=5), + ) + assert cost == 0.0 + + +def test_custom_pricing_does_not_override_a_mapped_model_when_absent() -> None: + """model_info without OCR pricing must fall through to the cost map.""" + mapped_model = "mistral/mistral-ocr-4-0" + cost, _ = ocr_cost( + model=mapped_model, + custom_llm_provider="mistral", + response=_ocr_response(mapped_model, pages_processed=2), + model_info={"id": "some-deployment-id"}, + ) + assert cost == pytest.approx(0.004 * 2) + + +def test_ocr_custom_pricing_end_to_end_through_completion_cost() -> None: + """The whole path: litellm_params.metadata.model_info -> ocr_cost.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + logging_obj = LiteLLMLogging( + model=UNMAPPED_MODEL, + messages=[], + stream=False, + call_type="ocr", + start_time=None, + litellm_call_id="test-ocr-custom-pricing", + function_id="1234", + ) + logging_obj.litellm_params = {"metadata": {"model_info": {"ocr_cost_per_page": CUSTOM_COST_PER_PAGE}}} + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_MODEL, pages_processed=3), + model=UNMAPPED_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * 3) From c77b5bada5435d56ce3f77d3e04572f1b5b08dab Mon Sep 17 00:00:00 2001 From: Mihidum Hettiyahandi <55163074+mihidumh@users.noreply.github.com> Date: Mon, 17 Aug 2026 13:50:49 +1000 Subject: [PATCH 03/55] fix(cost): drop the `or {}` default so the metadata read adds no LIT002 The type-discipline gate failed on this PR: LIT002 (mutable-collection construction) totalled 27149 against a ceiling of 27146, both new hits on the `_cm_litellm_params.get("metadata", {}) or {}` line. The two dict literals were only there to make the read total; a truthiness check on the value does the same job and constructs nothing. Behaviour is unchanged for every input: a missing, None or empty `metadata` leaves `_custom_model_info` as None either way. This copies the extraction in the video-generation path a few lines above, which still carries the `or {}` form. That one is inside the gate's existing budget, so it is left alone rather than reformatted in an unrelated PR. scripts/type_discipline_gate.py --base b4f5e46a now reports "every LIT rule is within its codebase ceiling"; tests/test_litellm/test_ocr_custom_pricing.py is 8 passed. --- litellm/cost_calculator.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index fb415d69fdf..cb268513e00 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1590,13 +1590,16 @@ def completion_cost( # Deployment-specific model_info, for modalities whose pricing is # not token-based and so cannot travel via custom_cost_per_token - # (e.g. OCR per-page pricing). Same extraction as the video path. + # (e.g. OCR per-page pricing). Same extraction as the video path + # above, minus its `or {}` default: truthiness on the value adds + # no mutable-collection construction (LIT002) and reads the same. _custom_model_info: ModelInfo | None = None if custom_pricing and litellm_logging_obj is not None: _cm_litellm_params = getattr(litellm_logging_obj, "litellm_params", None) if _cm_litellm_params is not None: - _cm_metadata = _cm_litellm_params.get("metadata", {}) or {} - _custom_model_info = _cm_metadata.get("model_info", None) + _cm_metadata = _cm_litellm_params.get("metadata") + if _cm_metadata: + _custom_model_info = _cm_metadata.get("model_info", None) ( prompt_tokens_cost_usd_dollar, From 78ec018052f4332b38ad7889cc6878ed7ff8f31b Mon Sep 17 00:00:00 2001 From: Mihidum Hettiyahandi <55163074+mihidumh@users.noreply.github.com> Date: Mon, 24 Aug 2026 09:23:39 +1000 Subject: [PATCH 04/55] fix(cost): read deployment model_info from litellm_metadata too Router OCR calls go through _ageneric_api_call_with_fallbacks, which stores the deployment's model_info under litellm_metadata rather than metadata, so the custom OCR pricing was still unreachable on that path. Check both keys, metadata first, mirroring _get_base_model_from_litellm_call_metadata. Adds an end-to-end test for the litellm_metadata shape. --- litellm/cost_calculator.py | 11 +++++-- tests/test_litellm/test_ocr_custom_pricing.py | 30 +++++++++++++++++++ 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index cb268513e00..182c176121d 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1593,13 +1593,18 @@ def completion_cost( # (e.g. OCR per-page pricing). Same extraction as the video path # above, minus its `or {}` default: truthiness on the value adds # no mutable-collection construction (LIT002) and reads the same. + # Checked under both keys: router calls that go through + # `_ageneric_api_call_with_fallbacks` (OCR included) store the + # deployment's model_info under `litellm_metadata`, not `metadata`. _custom_model_info: ModelInfo | None = None if custom_pricing and litellm_logging_obj is not None: _cm_litellm_params = getattr(litellm_logging_obj, "litellm_params", None) if _cm_litellm_params is not None: - _cm_metadata = _cm_litellm_params.get("metadata") - if _cm_metadata: - _custom_model_info = _cm_metadata.get("model_info", None) + for _cm_metadata_key in ("metadata", "litellm_metadata"): + _cm_metadata = _cm_litellm_params.get(_cm_metadata_key) + if _cm_metadata and _cm_metadata.get("model_info") is not None: + _custom_model_info = _cm_metadata.get("model_info") + break ( prompt_tokens_cost_usd_dollar, diff --git a/tests/test_litellm/test_ocr_custom_pricing.py b/tests/test_litellm/test_ocr_custom_pricing.py index 7cd4fc1dec1..4dffc43c23f 100644 --- a/tests/test_litellm/test_ocr_custom_pricing.py +++ b/tests/test_litellm/test_ocr_custom_pricing.py @@ -111,3 +111,33 @@ def test_ocr_custom_pricing_end_to_end_through_completion_cost() -> None: litellm_logging_obj=logging_obj, ) assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * 3) + + +def test_ocr_custom_pricing_end_to_end_via_litellm_metadata() -> None: + """Router OCR calls go through `_ageneric_api_call_with_fallbacks`, which + stores the deployment's model_info under `litellm_metadata` rather than + `metadata`. The extraction must read that key too.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + logging_obj = LiteLLMLogging( + model=UNMAPPED_MODEL, + messages=[], + stream=False, + call_type="ocr", + start_time=None, + litellm_call_id="test-ocr-custom-pricing-litellm-metadata", + function_id="1234", + ) + logging_obj.litellm_params = { + "litellm_metadata": {"model_info": {"ocr_cost_per_page": CUSTOM_COST_PER_PAGE}}, + } + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_MODEL, pages_processed=3), + model=UNMAPPED_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * 3) From eb9beced715cc071914e2248984fb4f43cf482c0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 02:49:09 -0700 Subject: [PATCH 05/55] fix(anthropic): price recovered tokens when a /v1/messages client disconnects mid-stream --- .../anthropic_passthrough_logging_handler.py | 50 ++++++++--- .../messages/test_streaming_iterator.py | 84 +++++++++++++++++++ ...t_anthropic_passthrough_logging_handler.py | 82 +++++++++++++++++- 3 files changed, 202 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index 30b75a7b482..f1387606f4e 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -216,11 +216,16 @@ class AnthropicPassthroughLoggingHandler: model=model, speed=AnthropicPassthroughLoggingHandler._cost_relevant_speed(request_body), ) - if response is None: - return None - AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens( + if not isinstance(response, ModelResponse): + return response + recovered_usage: Final = AnthropicPassthroughLoggingHandler._recover_interrupted_stream_output_tokens( response=response, all_chunks=all_chunks, model=model ) + if recovered_usage is None: + return response + AnthropicPassthroughLoggingHandler._reprice_recovered_stream( + response=response, usage=recovered_usage, model=model, logging_obj=litellm_logging_obj + ) return response @staticmethod @@ -259,7 +264,9 @@ class AnthropicPassthroughLoggingHandler: ) except Exception as e: # noqa: BLE001 # an uncostable partial stream still bills its tokens, at zero cost verbose_proxy_logger.warning( - "Anthropic passthrough: could not cost the partial usage of a failed stream (model=%s): %s", model, e + "Anthropic passthrough: could not cost the partial usage of an interrupted stream (model=%s): %s", + model, + e, ) return 0.0 @@ -359,7 +366,7 @@ class AnthropicPassthroughLoggingHandler: response: ModelResponse | TextCompletionResponse, all_chunks: Sequence[str | bytes], model: str, - ) -> None: + ) -> Usage | None: """ An Anthropic stream interrupted before its terminal ``message_delta`` (client disconnect) carries only the ``message_start`` ``output_tokens`` @@ -369,24 +376,24 @@ class AnthropicPassthroughLoggingHandler: untouched because their terminal ``message_delta`` short-circuits here. """ if not isinstance(response, ModelResponse): - return + return None if not AnthropicPassthroughLoggingHandler._stream_was_interrupted(all_chunks): - return + return None usage: Final = getattr(response, "usage", None) - if usage is None: - return + if not isinstance(usage, Usage): + return None output_text: Final = get_content_from_model_response(response) if not output_text: - return + return None try: recovered_output_tokens = litellm.token_counter(model=model, text=output_text, count_response_tokens=True) except Exception: verbose_proxy_logger.warning( "Could not re-tokenize interrupted stream output; keeping placeholder completion token count." ) - return + return None if recovered_output_tokens <= (usage.completion_tokens or 0): - return + return None usage.completion_tokens = recovered_output_tokens usage.total_tokens = (usage.prompt_tokens or 0) + recovered_output_tokens # Anthropic costing reads completion_tokens_details.text_tokens, so the @@ -395,6 +402,25 @@ class AnthropicPassthroughLoggingHandler: details: Final = getattr(usage, "completion_tokens_details", None) if details is not None and getattr(details, "text_tokens", None) is not None: details.text_tokens = recovered_output_tokens + return usage + + @staticmethod + def _reprice_recovered_stream( + response: ModelResponse, + usage: Usage, + model: str, + logging_obj: LiteLLMLoggingObj, + ) -> None: + hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor + usage.cost = None + hidden_params.pop("response_cost", None) + recovered_cost: Final = AnthropicPassthroughLoggingHandler._cost_partial_stream_or_zero( + partial_response=response, model=model, logging_obj=logging_obj + ) + if recovered_cost <= 0: + return + usage.cost = recovered_cost + hidden_params["response_cost"] = recovered_cost @staticmethod def _create_anthropic_response_logging_payload( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index be33b2ee3b1..89193adca1f 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -1,6 +1,7 @@ import asyncio import json from datetime import datetime +from unittest.mock import patch import pytest @@ -824,6 +825,89 @@ async def test_async_sse_wrapper_bills_partial_when_detached_drains_disabled(mon assert len(streaming_iterator_module._DETACHED_STREAM_DRAINS) == 0 +class _SuccessRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.success_kwargs: list = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_kwargs.append(kwargs) + + +@pytest.mark.asyncio +async def test_client_disconnect_partial_billing_prices_recovered_tokens(monkeypatch): + """ + Regression (LIT-6872): a client disconnect that lands on partial billing + re-tokenizes the buffered text into completion_tokens, but the logged cost + stayed priced at the message_start placeholder (1 output token). The success + row's response_cost must match its recovered completion_tokens. + """ + import litellm + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_MAX_DETACHED_STREAM_DRAINS", 0) + monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) + model = "claude-sonnet-5" + recorder = _SuccessRecorder() + logging_obj = LiteLLMLoggingObj( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id="disconnect_partial_cost", + function_id="disconnect_partial_cost", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params={"custom_llm_provider": "anthropic"}, + custom_llm_provider="anthropic", + ) + iterator = BaseAnthropicMessagesStreamingIterator( + litellm_logging_obj=logging_obj, request_body={"model": model, "stream": True} + ) + sentence = "The history of computing spans centuries of mechanical and electronic invention. " + + async def _stream(): + yield {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 29, "output_tokens": 1}}} + yield {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + for _ in range(100): + yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": sentence}} + yield {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1500}} + yield {"type": "message_stop"} + + enqueued: list = [] + + def _capture(async_coroutine): + enqueued.append(async_coroutine) + + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", side_effect=_capture + ): + gen = iterator.async_sse_wrapper(_stream()) + for _ in range(4): + await gen.__anext__() + await gen.aclose() + for _ in range(500): + if enqueued: + break + await asyncio.sleep(0.01) + + assert len(enqueued) == 1, "client disconnect never reached partial billing" + await enqueued[0] + + assert len(recorder.success_kwargs) == 1 + logged = recorder.success_kwargs[0]["standard_logging_object"] + assert 1 < logged["completion_tokens"] < 1500 + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, prompt_tokens=29, completion_tokens=logged["completion_tokens"] + ) + assert logged["response_cost"] == pytest.approx(prompt_cost + completion_cost) + + @pytest.mark.asyncio async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached(monkeypatch): """ diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index d721be62efe..dc3e94fcff2 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.pass_through_endpoints.llm_provider_handlers.anthropic_passthrough_logging_handler import ( AnthropicPassthroughLoggingHandler, @@ -1551,6 +1552,7 @@ class TestInterruptedStreamOutputTokenRecovery: return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() _MODEL = "claude-3-5-haiku-20241022" + _PRICED_MODEL = "claude-sonnet-5" _OUTPUT_TEXT = ( "The history of computing spans centuries, beginning with mechanical " "calculators and the abacus, advancing through Charles Babbage's " @@ -1559,7 +1561,7 @@ class TestInterruptedStreamOutputTokenRecovery: "century that gave rise to the modern information age." ) - def _interrupted_chunks(self, *, placeholder_output_tokens: int = 2): + def _interrupted_chunks(self, *, placeholder_output_tokens: int = 2, model: str | None = None): from litellm.proxy.pass_through_endpoints.streaming_handler import ( PassThroughStreamingHandler, ) @@ -1574,7 +1576,7 @@ class TestInterruptedStreamOutputTokenRecovery: "id": "msg_interrupted", "type": "message", "role": "assistant", - "model": self._MODEL, + "model": model or self._MODEL, "content": [], "stop_reason": None, "stop_sequence": None, @@ -1676,6 +1678,82 @@ class TestInterruptedStreamOutputTokenRecovery: # provider count is preserved verbatim. assert usage.completion_tokens == final + @pytest.mark.asyncio + async def test_interrupted_stream_logs_cost_of_recovered_tokens(self): + """ + Regression (LIT-6872): stream_chunk_builder stamps usage.cost and + _hidden_params["response_cost"] from the message_start placeholder before + the interrupted stream is re-tokenized, and the success handler prefers + that hidden cost over the recomputed one. The logged cost must price the + recovered completion tokens, not the placeholder. + """ + import litellm + + class _SuccessRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.success_kwargs: list = [] + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.success_kwargs.append(kwargs) + + recorder = _SuccessRecorder() + logging_obj = LiteLLMLoggingObj( + model=self._PRICED_MODEL, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="pass_through_endpoint", + start_time=datetime.now(), + litellm_call_id="lit-6872", + function_id="lit-6872", + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=self._PRICED_MODEL, + user="", + optional_params={}, + litellm_params={"custom_llm_provider": "anthropic"}, + custom_llm_provider="anthropic", + ) + placeholder = 1 + handled = AnthropicPassthroughLoggingHandler._handle_logging_anthropic_collected_chunks( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=MagicMock(), + url_route="/anthropic/v1/messages", + request_body={"model": self._PRICED_MODEL, "stream": True}, + endpoint_type="messages", + start_time=datetime.now(), + all_chunks=self._interrupted_chunks(placeholder_output_tokens=placeholder, model=self._PRICED_MODEL), + end_time=datetime.now(), + ) + await logging_obj.dispatch_success_handlers( + result=handled["result"], + start_time=datetime.now(), + end_time=datetime.now(), + cache_hit=False, + prefer_async_handlers=True, + **handled["kwargs"], + ) + for _ in range(300): + if recorder.success_kwargs: + break + await asyncio.sleep(0.01) + + assert len(recorder.success_kwargs) == 1 + logged = recorder.success_kwargs[0]["standard_logging_object"] + recovered_tokens = handled["result"].usage.completion_tokens + assert recovered_tokens > placeholder + assert logged["completion_tokens"] == recovered_tokens + prompt_cost, completion_cost = litellm.cost_per_token( + model=self._PRICED_MODEL, prompt_tokens=29, completion_tokens=recovered_tokens + ) + _, placeholder_completion_cost = litellm.cost_per_token( + model=self._PRICED_MODEL, prompt_tokens=29, completion_tokens=placeholder + ) + assert logged["response_cost"] == pytest.approx(prompt_cost + completion_cost) + assert logged["response_cost"] > prompt_cost + placeholder_completion_cost + assert handled["result"].usage.cost == pytest.approx(prompt_cost + completion_cost) + class TestStreamFalseDeduplication: """ From bd5f066c67795ee3113f2fde9e70e3f95f3a3685 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 5 Sep 2026 10:27:01 +0000 Subject: [PATCH 06/55] test: deflake two tests whose shared-state leaks failed once and passed on CI rerun The Redis semantic cache tests wrapped the first import of litellm.caching.redis_semantic_cache in patch.dict("sys.modules", ...), which snapshots and restores all of sys.modules on exit. Every module first imported inside the block, including litellm.proxy.proxy_server, was dropped from sys.modules while staying cached as an attribute on the litellm.proxy package. The next test that patched litellm.proxy.proxy_server. hit the stale attribute while production code re-imported a fresh module, so the patch never reached it. Replace the whole-dict patch with MonkeyPatch.setitem on the two redisvl keys only The LangSmith init test globally patched asyncio.get_running_loop while constructing the logger. Any orphaned AsyncHTTPHandler finalized by the cyclic GC during that window also called loop.create_task on the mock, tripping assert_called_once. Run the test under a real event loop and assert on the real task instead of patching asyncio Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- test-quality-budget.json | 2 +- .../caching/test_redis_semantic_cache.py | 141 ++++-------------- .../integrations/test_langsmith_init.py | 23 ++- 3 files changed, 39 insertions(+), 127 deletions(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index 3c12371f02f..f382f2479a9 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -3,7 +3,7 @@ "limit": 733 }, "TQ002": { - "limit": 737 + "limit": 736 }, "TQ003": { "limit": 62 diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index df990c43530..9884e9d9bc0 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -1,3 +1,5 @@ +from collections.abc import Iterator +from contextlib import contextmanager from importlib import import_module import sys from unittest.mock import AsyncMock, MagicMock, patch @@ -5,18 +7,19 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +@contextmanager +def _fake_redisvl_modules(semantic_cache_mock: MagicMock, custom_vectorizer_mock: MagicMock) -> Iterator[None]: + with pytest.MonkeyPatch.context() as mp: + mp.setitem(sys.modules, "redisvl.extensions.llmcache", MagicMock(SemanticCache=semantic_cache_mock)) + mp.setitem(sys.modules, "redisvl.utils.vectorize", MagicMock(CustomTextVectorizer=custom_vectorizer_mock)) + yield + # Tests for RedisSemanticCache def test_redis_semantic_cache_initialization(monkeypatch): # Mock the redisvl import semantic_cache_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock(CustomTextVectorizer=MagicMock()), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, MagicMock()): from litellm.caching.redis_semantic_cache import RedisSemanticCache # Set environment variables @@ -44,15 +47,7 @@ def test_redis_semantic_cache_get_cache(monkeypatch): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache # Set environment variables @@ -110,15 +105,7 @@ def test_redis_semantic_cache_rejects_unscoped_cache_hit(monkeypatch): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -162,15 +149,7 @@ def test_redis_semantic_cache_set_cache_stores_cache_key_filter(monkeypatch): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -210,15 +189,7 @@ def test_redis_semantic_cache_uses_isolated_index_for_old_schema(monkeypatch): ) custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -252,15 +223,7 @@ def test_redis_semantic_cache_overwrites_stale_isolated_index(monkeypatch): ) custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -292,15 +255,7 @@ def test_redis_semantic_cache_reraises_unexpected_isolated_index_error(monkeypat ) custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -369,15 +324,15 @@ def test_redis_semantic_cache_builds_filter_expression(monkeypatch): def __eq__(self, value): return (self.field_name, value) - with patch.dict("sys.modules", {"redisvl.query.filter": MagicMock(Tag=FakeTag)}): - from litellm.caching.redis_semantic_cache import RedisSemanticCache + monkeypatch.setitem(sys.modules, "redisvl.query.filter", MagicMock(Tag=FakeTag)) + from litellm.caching.redis_semantic_cache import RedisSemanticCache - redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) + redis_semantic_cache = RedisSemanticCache.__new__(RedisSemanticCache) - assert redis_semantic_cache._get_cache_key_filter_expression("test_key") == ( - RedisSemanticCache.CACHE_KEY_FIELD_NAME, - "test_key", - ) + assert redis_semantic_cache._get_cache_key_filter_expression("test_key") == ( + RedisSemanticCache.CACHE_KEY_FIELD_NAME, + "test_key", + ) @pytest.mark.asyncio @@ -386,15 +341,7 @@ async def test_redis_semantic_cache_async_get_cache(monkeypatch): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache # Set environment variables @@ -449,15 +396,7 @@ async def test_redis_semantic_cache_async_get_cache_rejects_unscoped_hit(monkeyp semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -499,15 +438,7 @@ async def test_redis_semantic_cache_async_set_cache_stores_cache_key_filter( semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -1255,15 +1186,7 @@ def test_redis_init_defers_redisvl_construction(monkeypatch): semantic_cache_mock = MagicMock() custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") @@ -1291,15 +1214,7 @@ def test_redis_failed_llmcache_build_is_not_memoized(monkeypatch): ) custom_vectorizer_mock = MagicMock() - with patch.dict( - "sys.modules", - { - "redisvl.extensions.llmcache": MagicMock(SemanticCache=semantic_cache_mock), - "redisvl.utils.vectorize": MagicMock( - CustomTextVectorizer=custom_vectorizer_mock - ), - }, - ): + with _fake_redisvl_modules(semantic_cache_mock, custom_vectorizer_mock): from litellm.caching.redis_semantic_cache import RedisSemanticCache monkeypatch.setenv("REDIS_HOST", "localhost") diff --git a/tests/test_litellm/integrations/test_langsmith_init.py b/tests/test_litellm/integrations/test_langsmith_init.py index 025aa86466c..34efc08ac3e 100644 --- a/tests/test_litellm/integrations/test_langsmith_init.py +++ b/tests/test_litellm/integrations/test_langsmith_init.py @@ -1,3 +1,4 @@ +import asyncio import os from unittest.mock import MagicMock, patch @@ -154,24 +155,20 @@ class TestLangsmithLoggerInit: assert logger._start_periodic_flush_task() is None mock_get_running_loop.assert_called_once() - @patch("asyncio.get_running_loop") - def test_langsmith_init_starts_periodic_flush_with_running_loop( - self, mock_get_running_loop - ): + @pytest.mark.asyncio + async def test_langsmith_init_starts_periodic_flush_with_running_loop(self): """Test that init schedules periodic flush when a running loop exists.""" - mock_loop = MagicMock() - mock_task = MagicMock() - mock_loop.create_task.return_value = mock_task - mock_get_running_loop.return_value = mock_loop - logger = LangsmithLogger( langsmith_api_key="test-key", langsmith_project="test-project" ) - assert logger._flush_task == mock_task - mock_loop.create_task.assert_called_once() - scheduled_coro = mock_loop.create_task.call_args.args[0] - scheduled_coro.close() + flush_task = logger._flush_task + assert isinstance(flush_task, asyncio.Task) + assert not flush_task.done() + assert flush_task.get_coro().__qualname__ == "CustomBatchLogger.periodic_flush" + flush_task.cancel() + with pytest.raises(asyncio.CancelledError): + await flush_task @pytest.mark.asyncio async def test_async_log_success_event_lazily_starts_periodic_flush(self): From 8a7dc64ab4d2082f5dd3def35a13eb71294d650c Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 5 Sep 2026 10:36:15 +0000 Subject: [PATCH 07/55] test: assert LangSmith periodic flush by observing a batch send Replace the coroutine __qualname__ check with a functional check: queue one event, run with a short flush interval, and wait for async_send_batch to be awaited by the task init scheduled Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/integrations/test_langsmith_init.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/integrations/test_langsmith_init.py b/tests/test_litellm/integrations/test_langsmith_init.py index 34efc08ac3e..0bc9e279fbf 100644 --- a/tests/test_litellm/integrations/test_langsmith_init.py +++ b/tests/test_litellm/integrations/test_langsmith_init.py @@ -1,6 +1,6 @@ import asyncio import os -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -159,13 +159,15 @@ class TestLangsmithLoggerInit: async def test_langsmith_init_starts_periodic_flush_with_running_loop(self): """Test that init schedules periodic flush when a running loop exists.""" logger = LangsmithLogger( - langsmith_api_key="test-key", langsmith_project="test-project" + langsmith_api_key="test-key", langsmith_project="test-project", flush_interval=0.01 ) + batch_sent = asyncio.Event() + logger.async_send_batch = AsyncMock(side_effect=batch_sent.set) + logger.log_queue.append({"id": "run-id"}) flush_task = logger._flush_task assert isinstance(flush_task, asyncio.Task) - assert not flush_task.done() - assert flush_task.get_coro().__qualname__ == "CustomBatchLogger.periodic_flush" + await asyncio.wait_for(batch_sent.wait(), timeout=5) flush_task.cancel() with pytest.raises(asyncio.CancelledError): await flush_task From 2dabac186abc7329222985b295545b38f60dfcde Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 5 Sep 2026 03:53:40 -0700 Subject: [PATCH 08/55] fix(anthropic): clear the placeholder cost once and let logging price recovered tokens Drop the second pricing pass on interrupted /v1/messages streams: clearing the stale usage.cost and hidden response_cost is enough for the existing success and failure logging to price the recovered usage. Add an iterator test for the upstream-close path the proxy takes on a client disconnect. --- .../anthropic_passthrough_logging_handler.py | 21 +-- .../messages/test_streaming_iterator.py | 121 +++++++++++++++--- ...t_anthropic_passthrough_logging_handler.py | 1 - 3 files changed, 106 insertions(+), 37 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py index f1387606f4e..7cec3bac207 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/anthropic_passthrough_logging_handler.py @@ -223,9 +223,7 @@ class AnthropicPassthroughLoggingHandler: ) if recovered_usage is None: return response - AnthropicPassthroughLoggingHandler._reprice_recovered_stream( - response=response, usage=recovered_usage, model=model, logging_obj=litellm_logging_obj - ) + AnthropicPassthroughLoggingHandler._clear_placeholder_cost(response=response, usage=recovered_usage) return response @staticmethod @@ -405,22 +403,9 @@ class AnthropicPassthroughLoggingHandler: return usage @staticmethod - def _reprice_recovered_stream( - response: ModelResponse, - usage: Usage, - model: str, - logging_obj: LiteLLMLoggingObj, - ) -> None: - hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor + def _clear_placeholder_cost(response: ModelResponse, usage: Usage) -> None: usage.cost = None - hidden_params.pop("response_cost", None) - recovered_cost: Final = AnthropicPassthroughLoggingHandler._cost_partial_stream_or_zero( - partial_response=response, model=model, logging_obj=logging_obj - ) - if recovered_cost <= 0: - return - usage.cost = recovered_cost - hidden_params["response_cost"] = recovered_cost + response._hidden_params.pop("response_cost", None) # pyright: ignore[reportPrivateUsage] # no public accessor @staticmethod def _create_anthropic_response_logging_payload( diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py index 89193adca1f..8043496f299 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/messages/test_streaming_iterator.py @@ -834,6 +834,50 @@ class _SuccessRecorder(CustomLogger): self.success_kwargs.append(kwargs) +def _make_priced_logging_obj(call_id: str, recorder: _SuccessRecorder, model: str) -> LiteLLMLoggingObj: + logging_obj = LiteLLMLoggingObj( + model=model, + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="anthropic_messages", + start_time=datetime.now(), + litellm_call_id=call_id, + function_id=call_id, + dynamic_async_success_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model=model, + user="", + optional_params={}, + litellm_params={"custom_llm_provider": "anthropic"}, + custom_llm_provider="anthropic", + ) + return logging_obj + + +class _UpstreamClosedOnDetach: + """Upstream that yields its events and then, like a socket read, waits until it is closed.""" + + def __init__(self, events: tuple[dict, ...]): + self._events = iter(events) + self._closed = asyncio.Event() + + def __aiter__(self): + return self + + async def __anext__(self) -> dict: + if self._closed.is_set(): + raise StopAsyncIteration + try: + return next(self._events) + except StopIteration: + await self._closed.wait() + raise StopAsyncIteration + + async def aclose(self) -> None: + self._closed.set() + + @pytest.mark.asyncio async def test_client_disconnect_partial_billing_prices_recovered_tokens(monkeypatch): """ @@ -849,25 +893,9 @@ async def test_client_disconnect_partial_billing_prices_recovered_tokens(monkeyp monkeypatch.setattr(streaming_iterator_module, "ANTHROPIC_MESSAGES_STREAM_RELAY_QUEUE_MAXSIZE", 4) model = "claude-sonnet-5" recorder = _SuccessRecorder() - logging_obj = LiteLLMLoggingObj( - model=model, - messages=[{"role": "user", "content": "hi"}], - stream=True, - call_type="anthropic_messages", - start_time=datetime.now(), - litellm_call_id="disconnect_partial_cost", - function_id="disconnect_partial_cost", - dynamic_async_success_callbacks=[recorder], - ) - logging_obj.update_environment_variables( - model=model, - user="", - optional_params={}, - litellm_params={"custom_llm_provider": "anthropic"}, - custom_llm_provider="anthropic", - ) iterator = BaseAnthropicMessagesStreamingIterator( - litellm_logging_obj=logging_obj, request_body={"model": model, "stream": True} + litellm_logging_obj=_make_priced_logging_obj("disconnect_partial_cost", recorder, model), + request_body={"model": model, "stream": True}, ) sentence = "The history of computing spans centuries of mechanical and electronic invention. " @@ -908,6 +936,63 @@ async def test_client_disconnect_partial_billing_prices_recovered_tokens(monkeyp assert logged["response_cost"] == pytest.approx(prompt_cost + completion_cost) +@pytest.mark.asyncio +async def test_proxy_disconnect_closing_upstream_prices_recovered_tokens(): + """ + Regression (LIT-6872), proxy path: after a client disconnect the proxy's + shielded cleanup closes the upstream stream while the pump is still reading + it, so the pump bills the chunks collected so far without ever seeing + message_delta. That row's response_cost must be priced from its recovered + completion_tokens, not from the message_start placeholder. + """ + import litellm + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + model = "claude-sonnet-5" + recorder = _SuccessRecorder() + iterator = BaseAnthropicMessagesStreamingIterator( + litellm_logging_obj=_make_priced_logging_obj("disconnect_upstream_closed", recorder, model), + request_body={"model": model, "stream": True}, + ) + sentence = "The history of computing spans centuries of mechanical and electronic invention. " + upstream = _UpstreamClosedOnDetach( + ( + {"type": "message_start", "message": {"id": "msg_1", "usage": {"input_tokens": 29, "output_tokens": 1}}}, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + *({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": sentence}} for _ in range(6)), + ) + ) + enqueued: list = [] + + def _capture(async_coroutine): + enqueued.append(async_coroutine) + + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue", side_effect=_capture + ): + gen = iterator.async_sse_wrapper(upstream) + for _ in range(4): + await gen.__anext__() + await gen.aclose() + assert not enqueued, "billing must wait for the upstream read to end, not the client detach" + await upstream.aclose() + for _ in range(500): + if enqueued: + break + await asyncio.sleep(0.01) + + assert len(enqueued) == 1, "closing the upstream never reached partial billing" + await enqueued[0] + + assert len(recorder.success_kwargs) == 1 + logged = recorder.success_kwargs[0]["standard_logging_object"] + assert logged["completion_tokens"] > 1 + prompt_cost, completion_cost = litellm.cost_per_token( + model=model, prompt_tokens=29, completion_tokens=logged["completion_tokens"] + ) + assert logged["response_cost"] == pytest.approx(prompt_cost + completion_cost) + + @pytest.mark.asyncio async def test_async_sse_wrapper_aborts_upstream_when_detached_drain_cap_reached(monkeypatch): """ diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py index dc3e94fcff2..2d7397594aa 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_anthropic_passthrough_logging_handler.py @@ -1752,7 +1752,6 @@ class TestInterruptedStreamOutputTokenRecovery: ) assert logged["response_cost"] == pytest.approx(prompt_cost + completion_cost) assert logged["response_cost"] > prompt_cost + placeholder_completion_cost - assert handled["result"].usage.cost == pytest.approx(prompt_cost + completion_cost) class TestStreamFalseDeduplication: From 04cc8f855fc19add1bcc73a094b49b57250babbe Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 01:03:53 -0700 Subject: [PATCH 09/55] fix(proxy): resolve config include directives for bucket-hosted configs A config loaded from a GCS or S3 bucket skipped include processing entirely, so every model, guardrail, and setting behind an `include` was silently dropped. Both bucket types shared the same branch in `get_config`, which never called `_process_includes`, and that helper only ever read from disk. The merge now lives in one async helper that takes the loader as a dependency, so disk and bucket configs share the same semantics: list values extend, everything else overrides, nested includes are followed, and the `include` key is stripped. Bucket entries resolve as object keys relative to the config object's prefix, with a leading `/` meaning the bucket root, and an include that cannot be read now raises instead of being skipped. --- litellm/proxy/common_utils/config_includes.py | 68 +++++++ .../proxy/common_utils/load_config_utils.py | 66 ++++++- litellm/proxy/proxy_server.py | 49 ++--- .../common_utils/test_load_config_utils.py | 180 +++++++++++++++++- .../proxy/proxy_server/test_proxy_config.py | 44 ++++- 5 files changed, 370 insertions(+), 37 deletions(-) create mode 100644 litellm/proxy/common_utils/config_includes.py diff --git a/litellm/proxy/common_utils/config_includes.py b/litellm/proxy/common_utils/config_includes.py new file mode 100644 index 00000000000..d6939db2e95 --- /dev/null +++ b/litellm/proxy/common_utils/config_includes.py @@ -0,0 +1,68 @@ +from collections.abc import Awaitable, Mapping +from types import MappingProxyType +from typing import Final, Protocol + +INCLUDE_KEY: Final = "include" + + +class ConfigLoader(Protocol): + def __call__(self, include_entry: str, /) -> Awaitable[Mapping[str, object]]: ... + + +def _merged_value(base_value: object, included_value: object) -> object: + if isinstance(included_value, list) and isinstance(base_value, list): + return [*base_value, *included_value] # mutable-ok: a merged config value stays the plain list the proxy loads + return included_value + + +def _merged_entry(base: Mapping[str, object], included: Mapping[str, object], key: str) -> object: + if key not in included: + return base[key] + return _merged_value(base.get(key), included[key]) + + +def _merged(base: Mapping[str, object], included: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: _merged_entry(base, included, key) for key in (*base, *included)}) + + +def _without_include(config: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType({key: value for key, value in config.items() if key != INCLUDE_KEY}) + + +def include_entries(config: Mapping[str, object]) -> tuple[str, ...]: + if INCLUDE_KEY not in config: + return () + + entries: Final = config[INCLUDE_KEY] + if not isinstance(entries, list): + raise ValueError("'include' must be a list of file paths") + + paths: Final = tuple(entry for entry in entries if isinstance(entry, str)) + if len(paths) != len(entries): + raise ValueError("'include' must be a list of file paths") + + return paths + + +async def _resolve(config: Mapping[str, object], pending: tuple[str, ...], load: ConfigLoader) -> Mapping[str, object]: + if not pending: + return _without_include(config) + + included: Final = await load(pending[0]) + return await _resolve( + _merged(config, _without_include(included)), + (*pending[1:], *include_entries(included)), + load, + ) + + +async def resolve_includes(config: Mapping[str, object], load: ConfigLoader) -> dict[str, object]: + """ + Merge every config named by the `include` directive into the config that declares it. + + List values are extended and every other value is overridden, an included config may declare + further includes, and `load` decides where an entry is read from, so the same merge applies to + configs on disk and to configs hosted in a bucket. + """ + merged: Final = await _resolve(config, include_entries(config), load) + return dict(merged) # mutable-ok: the proxy mutates the config it loads diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 62649ad6ca1..927deb68826 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -1,9 +1,19 @@ import os -from typing import Final +import posixpath +from collections.abc import Awaitable, Mapping +from typing import Final, Protocol import yaml +from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger +from litellm.proxy.common_utils.config_includes import resolve_includes + +_BUCKET_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, object]) + + +class BucketObjectFetcher(Protocol): + def __call__(self, object_key: str, /) -> Awaitable[Mapping[str, object] | None]: ... def get_file_contents_from_s3(bucket_name, object_key): @@ -62,6 +72,60 @@ async def get_config_file_contents_from_gcs(bucket_name, object_key): return None +def resolve_include_object_key(config_object_key: str, include_entry: str) -> str: + """ + Resolve one `include` entry to the object key it names, relative to the config object's prefix. + + A leading "/" means the bucket root, mirroring how an absolute path on disk ignores the + directory the including config sits in. + """ + if include_entry.startswith("/"): + return posixpath.normpath(include_entry).lstrip("/") + return posixpath.normpath(posixpath.join(posixpath.dirname(config_object_key), include_entry)) + + +async def resolve_bucket_includes( + *, + config: Mapping[str, object], + object_key: str, + fetch: BucketObjectFetcher, +) -> dict[str, object]: + async def load(include_entry: str) -> Mapping[str, object]: + include_key: Final = resolve_include_object_key(object_key, include_entry) + included: Final = await fetch(include_key) + if included is None: + raise FileNotFoundError(f"Included config could not be read from bucket: {include_key}") + return included + + return await resolve_includes(config=config, load=load) + + +async def get_config_from_bucket( + *, + bucket_type: str | None, + bucket_name: str, + object_key: str, +) -> dict[str, object] | None: + async def fetch(key: str) -> Mapping[str, object] | None: + raw: Final = ( + await get_config_file_contents_from_gcs(bucket_name=bucket_name, object_key=key) + if bucket_type == "gcs" + else get_file_contents_from_s3(bucket_name=bucket_name, object_key=key) + ) + if raw is None: + return None + try: + return _BUCKET_CONFIG_ADAPTER.validate_python(raw) + except ValidationError as e: + raise ValueError(f"Config object in bucket is not a YAML mapping: {key}") from e + + config: Final = await fetch(object_key) + if config is None: + return None + + return await resolve_bucket_includes(config=config, object_key=object_key, fetch=fetch) + + def download_python_file_from_s3( bucket_name: str, object_key: str, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ba5714fe950..a49df166dca 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -341,6 +341,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy +from litellm.proxy.common_utils.config_includes import resolve_includes from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router @@ -359,10 +360,7 @@ from litellm.proxy.common_utils.http_parsing_utils import ( check_file_size_under_limit, get_form_data, ) -from litellm.proxy.common_utils.load_config_utils import ( - get_config_file_contents_from_gcs, - get_file_contents_from_s3, -) +from litellm.proxy.common_utils.load_config_utils import get_config_from_bucket from litellm.proxy.common_utils.model_deprecation import collect_model_deprecations from litellm.proxy.common_utils.model_listing_utils import ( TeamModelNameTranslator, @@ -4538,12 +4536,12 @@ class ProxyConfig: if config is None: raise Exception("Config cannot be None or Empty.") # Process includes - config = self._process_includes(config=config, base_dir=os.path.dirname(os.path.abspath(file_path or ""))) + config = await self._process_includes(config=config, base_dir=os.path.dirname(os.path.abspath(file_path or ""))) # verbose_proxy_logger.debug(f"loaded config={json.dumps(config, indent=4)}") return config - def _process_includes(self, config: dict, base_dir: str) -> dict: + async def _process_includes(self, config: dict, base_dir: str) -> dict: """ Process includes by appending their contents to the main config @@ -4558,29 +4556,14 @@ class ProxyConfig: callbacks: ["prometheus"] ``` """ - if "include" not in config: - return config - if not isinstance(config["include"], list): - raise ValueError("'include' must be a list of file paths") - - # Load and append all included files - for include_file in config["include"]: - file_path = os.path.join(base_dir, include_file) + async def load_included(include_file: str) -> Mapping[str, object]: + file_path: Final = os.path.join(base_dir, include_file) if not os.path.exists(file_path): raise FileNotFoundError(f"Included file not found: {file_path}") + return self._load_yaml_file(file_path) - included_config = self._load_yaml_file(file_path) - # Simply update/extend the main config with included config - for key, value in included_config.items(): - if isinstance(value, list) and key in config: - config[key].extend(value) - else: - config[key] = value - - # Remove the include directive - del config["include"] - return config + return await resolve_includes(config=config, load=load_included) async def save_config(self, new_config: dict, include_env_vars: bool = False): global prisma_client, general_settings, user_config_file_path, store_model_in_db @@ -4936,15 +4919,19 @@ class ProxyConfig: global prisma_client, store_model_in_db # Load existing config - if os.environ.get("LITELLM_CONFIG_BUCKET_NAME") is not None: - bucket_name: Final = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") + bucket_name: Final = os.environ.get("LITELLM_CONFIG_BUCKET_NAME") + if bucket_name is not None: object_key: Final = os.environ.get("LITELLM_CONFIG_BUCKET_OBJECT_KEY") bucket_type: Final = os.environ.get("LITELLM_CONFIG_BUCKET_TYPE") verbose_proxy_logger.debug("bucket_name: %s, object_key: %s", bucket_name, object_key) - if bucket_type == "gcs": - config = await get_config_file_contents_from_gcs(bucket_name=bucket_name, object_key=object_key) - else: - config = get_file_contents_from_s3(bucket_name=bucket_name, object_key=object_key) + if object_key is None: + raise Exception("LITELLM_CONFIG_BUCKET_OBJECT_KEY must be set to load the config from a bucket.") + + config = await get_config_from_bucket( + bucket_type=bucket_type, + bucket_name=bucket_name, + object_key=object_key, + ) if config is None: raise Exception("Unable to load config from given source.") diff --git a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py index 524c260e94a..c5d84460d39 100644 --- a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py @@ -1,9 +1,14 @@ +import re from unittest.mock import MagicMock, mock_open, patch import pytest import yaml -from litellm.proxy.common_utils.load_config_utils import get_file_contents_from_s3 +from litellm.proxy.common_utils.load_config_utils import ( + get_config_from_bucket, + get_file_contents_from_s3, + resolve_bucket_includes, +) class TestGetFileContentsFromS3: @@ -83,3 +88,176 @@ class TestGetFileContentsFromS3: # Verify yaml.safe_load was called with the decoded content mock_yaml_load.assert_called_once_with(yaml_content) + + +class TestBucketConfigIncludes: + """`include:` directives in a bucket-hosted config.yaml (LIT-6982). + + They used to be dropped silently: the proxy booted with the root config applied and everything + the included objects declared missing, with nothing logged. + """ + + @staticmethod + def _bucket(objects): + async def fetch(object_key): + return objects.get(object_key) + + return fetch + + @pytest.mark.asyncio + async def test_include_resolves_against_the_config_objects_prefix(self): + merged = await resolve_bucket_includes( + config={"include": ["model_config.yaml"], "general_settings": {"master_key": "sk-1234"}}, + object_key="configs/prod/config.yaml", + fetch=self._bucket( + {"configs/prod/model_config.yaml": {"model_list": [{"model_name": "gpt-4o-mini"}]}} + ), + ) + + assert merged == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "gpt-4o-mini"}], + } + + @pytest.mark.asyncio + async def test_include_with_a_leading_slash_reads_from_the_bucket_root(self): + merged = await resolve_bucket_includes( + config={"include": ["/shared/models.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({"shared/models.yaml": {"model_list": [{"model_name": "shared"}]}}), + ) + + assert merged == {"model_list": [{"model_name": "shared"}]} + + @pytest.mark.asyncio + async def test_include_walks_out_of_the_prefix_with_dot_dot(self): + merged = await resolve_bucket_includes( + config={"include": ["../shared/models.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({"configs/shared/models.yaml": {"model_list": [{"model_name": "shared"}]}}), + ) + + assert merged == {"model_list": [{"model_name": "shared"}]} + + @pytest.mark.asyncio + async def test_included_configs_may_declare_further_includes(self): + merged = await resolve_bucket_includes( + config={"include": ["models.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/models.yaml": { + "include": ["extra/more_models.yaml"], + "model_list": [{"model_name": "first"}], + }, + "configs/extra/more_models.yaml": {"model_list": [{"model_name": "second"}]}, + } + ), + ) + + assert merged == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + @pytest.mark.asyncio + async def test_list_values_are_extended_and_other_values_are_overridden(self): + merged = await resolve_bucket_includes( + config={ + "include": ["models.yaml"], + "model_list": [{"model_name": "from-root"}], + "litellm_settings": {"drop_params": True}, + }, + object_key="config.yaml", + fetch=self._bucket( + { + "models.yaml": { + "model_list": [{"model_name": "from-include"}], + "litellm_settings": {"num_retries": 3}, + } + } + ), + ) + + assert merged == { + "model_list": [{"model_name": "from-root"}, {"model_name": "from-include"}], + "litellm_settings": {"num_retries": 3}, + } + + @pytest.mark.asyncio + async def test_a_missing_included_object_fails_loudly_with_its_key(self): + with pytest.raises(FileNotFoundError, match=re.escape("configs/prod/model_config.yaml")): + await resolve_bucket_includes( + config={"include": ["model_config.yaml"]}, + object_key="configs/prod/config.yaml", + fetch=self._bucket({}), + ) + + @pytest.mark.asyncio + async def test_a_non_list_include_fails_loudly(self): + with pytest.raises(ValueError, match="'include' must be a list of file paths"): + await resolve_bucket_includes( + config={"include": "model_config.yaml"}, + object_key="config.yaml", + fetch=self._bucket({}), + ) + + @pytest.mark.asyncio + async def test_get_config_from_bucket_merges_includes_over_s3(self, monkeypatch): + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", + lambda bucket_name, object_key: objects.get(object_key), + ) + + config = await get_config_from_bucket( + bucket_type="s3", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "included-model"}], + } + + @pytest.mark.asyncio + async def test_get_config_from_bucket_merges_includes_over_gcs(self, monkeypatch): + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + + async def fake_gcs(bucket_name, object_key): + return objects.get(object_key) + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.get_config_file_contents_from_gcs", fake_gcs + ) + + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config == { + "general_settings": {"master_key": "sk-1234"}, + "model_list": [{"model_name": "included-model"}], + } + + @pytest.mark.asyncio + async def test_get_config_from_bucket_returns_none_when_the_root_object_is_missing(self, monkeypatch): + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", + lambda bucket_name, object_key: None, + ) + + assert ( + await get_config_from_bucket( + bucket_type="s3", bucket_name="litellm-configs", object_key="missing.yaml" + ) + is None + ) diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 2babfe432f3..9b9c4af9e42 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -710,22 +710,32 @@ async def test_ProxyConfig__get_config_from_file_missing_path_raises(): # --------------------------------------------------------------------------- -def test_ProxyConfig__process_includes_merges_files(tmp_path): +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_merges_files(tmp_path): inc = tmp_path / "models.yaml" inc.write_text("model_list:\n - model_name: gpt-4\n") pc = ProxyConfig() cfg = {"include": ["models.yaml"], "model_list": [], "litellm_settings": {}} - result = pc._process_includes(cfg, base_dir=str(tmp_path)) + result = await pc._process_includes(cfg, base_dir=str(tmp_path)) assert result == { "model_list": [{"model_name": "gpt-4"}], "litellm_settings": {}, } -def test_ProxyConfig__process_includes_missing_file_raises(tmp_path): +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_missing_file_raises(tmp_path): pc = ProxyConfig() with pytest.raises(FileNotFoundError): - pc._process_includes({"include": ["nope.yaml"]}, base_dir=str(tmp_path)) + await pc._process_includes({"include": ["nope.yaml"]}, base_dir=str(tmp_path)) + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_follows_nested_includes(tmp_path): + (tmp_path / "models.yaml").write_text("include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n") + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: second\n") + result = await ProxyConfig()._process_includes({"include": ["models.yaml"]}, base_dir=str(tmp_path)) + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} # --------------------------------------------------------------------------- @@ -1042,6 +1052,32 @@ async def test_ProxyConfig_get_config_loads_from_file(tmp_path, monkeypatch): } +@pytest.mark.asyncio +async def test_ProxyConfig_get_config_from_a_bucket_merges_includes(monkeypatch): + """A bucket-hosted config.yaml used to drop its `include:` entries silently (LIT-6982).""" + objects = { + "lit6982/config.yaml": { + "include": ["model_config.yaml"], + "general_settings": {"master_key": "sk-1234"}, + }, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", + lambda bucket_name, object_key: objects.get(object_key), + ) + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_NAME", "litellm-configs") + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_OBJECT_KEY", "lit6982/config.yaml") + monkeypatch.setenv("LITELLM_CONFIG_BUCKET_TYPE", "s3") + + cfg = await ProxyConfig().get_config() + + assert cfg["model_list"] == [{"model_name": "included-model"}] + assert "include" not in cfg + + @pytest.mark.asyncio async def test_ProxyConfig_get_config_missing_file_raises(monkeypatch): monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) From 0a763bf00dcc335c0e9735936cc9d8679154332e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 02:05:18 -0700 Subject: [PATCH 10/55] fix(proxy): read a bucket config's include tree off the event loop Reading a config from a bucket ran a blocking boto3 GET straight from the event loop for every object in the include tree, and on GCS it built a new bucket client per object, each one starting a flush task that never ends. S3 reads now go through a worker thread, and one bucket client serves the whole include tree. --- litellm/integrations/gcs_bucket/gcs_bucket.py | 3 +- litellm/proxy/common_utils/config_includes.py | 33 +++++-- .../proxy/common_utils/load_config_utils.py | 65 ++++++++++--- litellm/proxy/proxy_server.py | 18 ++-- .../gcs_bucket/test_gcs_bucket_base.py | 17 ++++ .../common_utils/test_load_config_utils.py | 91 ++++++++++++++++++- .../proxy/proxy_server/test_proxy_config.py | 62 ++++++++++++- 7 files changed, 251 insertions(+), 38 deletions(-) diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 31ceb338dcd..e338f490496 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -29,8 +29,6 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): def __init__(self, bucket_name: str | None = None) -> None: from litellm.proxy.proxy_server import premium_user - super().__init__(bucket_name=bucket_name) - self.batch_size = int(os.getenv("GCS_BATCH_SIZE", GCS_DEFAULT_BATCH_SIZE)) self.flush_interval = int(os.getenv("GCS_FLUSH_INTERVAL", GCS_DEFAULT_FLUSH_INTERVAL_SECONDS)) self.use_batched_logging = ( @@ -38,6 +36,7 @@ class GCSBucketLogger(GCSBucketBase, AdditionalLoggingUtils): ) self.flush_lock = asyncio.Lock() super().__init__( + bucket_name=bucket_name, flush_lock=self.flush_lock, batch_size=self.batch_size, flush_interval=self.flush_interval, diff --git a/litellm/proxy/common_utils/config_includes.py b/litellm/proxy/common_utils/config_includes.py index d6939db2e95..cf39bd6435e 100644 --- a/litellm/proxy/common_utils/config_includes.py +++ b/litellm/proxy/common_utils/config_includes.py @@ -6,7 +6,7 @@ INCLUDE_KEY: Final = "include" class ConfigLoader(Protocol): - def __call__(self, include_entry: str, /) -> Awaitable[Mapping[str, object]]: ... + def __call__(self, include_entry: str, declared_in: str, /) -> Awaitable[tuple[str, Mapping[str, object]]]: ... def _merged_value(base_value: object, included_value: object) -> object: @@ -44,25 +44,40 @@ def include_entries(config: Mapping[str, object]) -> tuple[str, ...]: return paths -async def _resolve(config: Mapping[str, object], pending: tuple[str, ...], load: ConfigLoader) -> Mapping[str, object]: +def _pending_from(config: Mapping[str, object], location: str) -> tuple[tuple[str, str], ...]: + return tuple((entry, location) for entry in include_entries(config)) + + +async def _resolve( + config: Mapping[str, object], + pending: tuple[tuple[str, str], ...], + loaded: frozenset[str], + load: ConfigLoader, +) -> Mapping[str, object]: if not pending: return _without_include(config) - included: Final = await load(pending[0]) + entry, declared_in = pending[0] + location, included = await load(entry, declared_in) + if location in loaded: + return await _resolve(config, pending[1:], loaded, load) + return await _resolve( _merged(config, _without_include(included)), - (*pending[1:], *include_entries(included)), + (*pending[1:], *_pending_from(included, location)), + loaded | frozenset((location,)), load, ) -async def resolve_includes(config: Mapping[str, object], load: ConfigLoader) -> dict[str, object]: +async def resolve_includes(config: Mapping[str, object], location: str, load: ConfigLoader) -> dict[str, object]: """ Merge every config named by the `include` directive into the config that declares it. - List values are extended and every other value is overridden, an included config may declare - further includes, and `load` decides where an entry is read from, so the same merge applies to - configs on disk and to configs hosted in a bucket. + List values are extended and every other value is overridden, each entry is resolved relative to + the config that declares it, a config already pulled in is not merged a second time, and `load` + decides where an entry is read from, so the same merge applies to configs on disk and to configs + hosted in a bucket. """ - merged: Final = await _resolve(config, include_entries(config), load) + merged: Final = await _resolve(config, _pending_from(config, location), frozenset((location,)), load) return dict(merged) # mutable-ok: the proxy mutates the config it loads diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index 927deb68826..e72272b1aca 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -1,7 +1,8 @@ +import asyncio import os import posixpath from collections.abc import Awaitable, Mapping -from typing import Final, Protocol +from typing import TYPE_CHECKING, Final, Protocol import yaml from pydantic import TypeAdapter, ValidationError @@ -9,6 +10,9 @@ from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.config_includes import resolve_includes +if TYPE_CHECKING: + from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase + _BUCKET_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, object]) @@ -16,6 +20,10 @@ class BucketObjectFetcher(Protocol): def __call__(self, object_key: str, /) -> Awaitable[Mapping[str, object] | None]: ... +class BucketObjectReader(Protocol): + def __call__(self, object_key: str, /) -> Awaitable[object | None]: ... + + def get_file_contents_from_s3(bucket_name, object_key): try: # v0 rely on boto3 for authentication - allowing boto3 to handle IAM credentials etc @@ -51,14 +59,22 @@ def get_file_contents_from_s3(bucket_name, object_key): return None -async def get_config_file_contents_from_gcs(bucket_name, object_key): +def gcs_config_bucket(bucket_name: str) -> "GCSBucketBase | None": try: from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger - gcs_bucket: Final = GCSBucketLogger( - bucket_name=bucket_name, - ) - file_contents = await gcs_bucket.download_gcs_object(object_key) + return GCSBucketLogger(bucket_name=bucket_name) + except Exception as e: + verbose_proxy_logger.error("Error creating the GCS client for bucket %s: %s", bucket_name, e) + return None + + +async def get_config_file_contents_from_gcs(bucket_name, object_key, gcs_bucket=None): + try: + bucket: Final = gcs_config_bucket(bucket_name) if gcs_bucket is None else gcs_bucket + if bucket is None: + return None + file_contents = await bucket.download_gcs_object(object_key) if file_contents is None: raise Exception(f"File contents are None for {object_key}") # file_contentis is a bytes object, so we need to convert it to yaml @@ -90,14 +106,35 @@ async def resolve_bucket_includes( object_key: str, fetch: BucketObjectFetcher, ) -> dict[str, object]: - async def load(include_entry: str) -> Mapping[str, object]: - include_key: Final = resolve_include_object_key(object_key, include_entry) + async def load(include_entry: str, declared_in: str) -> tuple[str, Mapping[str, object]]: + include_key: Final = resolve_include_object_key(declared_in, include_entry) included: Final = await fetch(include_key) if included is None: raise FileNotFoundError(f"Included config could not be read from bucket: {include_key}") - return included + return include_key, included - return await resolve_includes(config=config, load=load) + return await resolve_includes(config=config, location=object_key, load=load) + + +def bucket_object_reader(bucket_type: str | None, bucket_name: str) -> BucketObjectReader: + """ + Build one reader for a whole config, so an `include` tree costs one bucket client rather than one per object. + """ + if bucket_type != "gcs": + + async def read_from_s3(object_key: str) -> object | None: + return await asyncio.to_thread(get_file_contents_from_s3, bucket_name, object_key) + + return read_from_s3 + + gcs_bucket: Final = gcs_config_bucket(bucket_name) + + async def read_from_gcs(object_key: str) -> object | None: + if gcs_bucket is None: + return None + return await get_config_file_contents_from_gcs(bucket_name, object_key, gcs_bucket) + + return read_from_gcs async def get_config_from_bucket( @@ -106,12 +143,10 @@ async def get_config_from_bucket( bucket_name: str, object_key: str, ) -> dict[str, object] | None: + read: Final = bucket_object_reader(bucket_type, bucket_name) + async def fetch(key: str) -> Mapping[str, object] | None: - raw: Final = ( - await get_config_file_contents_from_gcs(bucket_name=bucket_name, object_key=key) - if bucket_type == "gcs" - else get_file_contents_from_s3(bucket_name=bucket_name, object_key=key) - ) + raw: Final = await read(key) if raw is None: return None try: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a49df166dca..7ac5e697b61 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4536,12 +4536,12 @@ class ProxyConfig: if config is None: raise Exception("Config cannot be None or Empty.") # Process includes - config = await self._process_includes(config=config, base_dir=os.path.dirname(os.path.abspath(file_path or ""))) + config = await self._process_includes(config=config, config_file_path=os.path.abspath(file_path or "")) # verbose_proxy_logger.debug(f"loaded config={json.dumps(config, indent=4)}") return config - async def _process_includes(self, config: dict, base_dir: str) -> dict: + async def _process_includes(self, config: dict, config_file_path: str) -> dict: """ Process includes by appending their contents to the main config @@ -4557,13 +4557,19 @@ class ProxyConfig: ``` """ - async def load_included(include_file: str) -> Mapping[str, object]: - file_path: Final = os.path.join(base_dir, include_file) + included_config_adapter: Final = TypeAdapter(dict[str, object]) + + async def load_included(include_file: str, declared_in: str) -> tuple[str, Mapping[str, object]]: + file_path: Final = os.path.abspath(os.path.join(os.path.dirname(declared_in), include_file)) if not os.path.exists(file_path): raise FileNotFoundError(f"Included file not found: {file_path}") - return self._load_yaml_file(file_path) + try: + included: Final = included_config_adapter.validate_python(self._load_yaml_file(file_path)) + except ValidationError as e: + raise ValueError(f"Included config file is not a YAML mapping: {file_path}") from e + return file_path, included - return await resolve_includes(config=config, load=load_included) + return await resolve_includes(config=config, location=config_file_path, load=load_included) async def save_config(self, new_config: dict, include_env_vars: bool = False): global prisma_client, general_settings, user_config_file_path, store_model_in_db diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py index 8d662311da1..a458752bed0 100644 --- a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py +++ b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py @@ -128,3 +128,20 @@ class TestGCSBucketBase: assert object_name.endswith("-target_uploadType_media") assert ".." not in object_name assert "?" not in object_name + + +class TestGCSBucketLoggerBucketName: + @pytest.mark.asyncio + async def test_the_bucket_name_it_is_constructed_with_survives(self, monkeypatch): + """Reading config.yaml out of a GCS bucket asks for that bucket, not the logging one (LIT-6982).""" + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + + assert GCSBucketLogger(bucket_name="config-bucket").BUCKET_NAME == "config-bucket" + + @pytest.mark.asyncio + async def test_no_bucket_name_still_falls_back_to_the_environment(self, monkeypatch): + monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket") + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + + assert GCSBucketLogger().BUCKET_NAME == "logging-bucket" diff --git a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py index c5d84460d39..b654320569c 100644 --- a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py @@ -1,4 +1,6 @@ +import asyncio import re +import threading from unittest.mock import MagicMock, mock_open, patch import pytest @@ -157,6 +159,60 @@ class TestBucketConfigIncludes: assert merged == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + @pytest.mark.asyncio + async def test_a_nested_include_resolves_against_the_object_that_declares_it(self): + """A nested `include` names a neighbour of the object declaring it, not of the root config.""" + merged = await resolve_bucket_includes( + config={"include": ["shared/models.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/shared/models.yaml": { + "include": ["more_models.yaml"], + "model_list": [{"model_name": "first"}], + }, + "configs/shared/more_models.yaml": {"model_list": [{"model_name": "second"}]}, + "configs/more_models.yaml": {"model_list": [{"model_name": "wrong-prefix"}]}, + } + ), + ) + + assert merged == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + @pytest.mark.asyncio + async def test_an_object_pulled_in_twice_is_merged_once(self): + merged = await resolve_bucket_includes( + config={"include": ["a.yaml", "b.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/a.yaml": {"include": ["shared.yaml"]}, + "configs/b.yaml": {"include": ["./shared.yaml"]}, + "configs/shared.yaml": {"model_list": [{"model_name": "shared"}]}, + } + ), + ) + + assert merged == {"model_list": [{"model_name": "shared"}]} + + @pytest.mark.asyncio + async def test_a_cycle_between_included_objects_terminates(self): + merged = await asyncio.wait_for( + resolve_bucket_includes( + config={"include": ["a.yaml"]}, + object_key="configs/config.yaml", + fetch=self._bucket( + { + "configs/a.yaml": {"include": ["b.yaml"], "model_list": [{"model_name": "from-a"}]}, + "configs/b.yaml": {"include": ["a.yaml"], "model_list": [{"model_name": "from-b"}]}, + } + ), + ), + timeout=10, + ) + + assert merged == {"model_list": [{"model_name": "from-a"}, {"model_name": "from-b"}]} + @pytest.mark.asyncio async def test_list_values_are_extended_and_other_values_are_overridden(self): merged = await resolve_bucket_includes( @@ -222,6 +278,23 @@ class TestBucketConfigIncludes: "model_list": [{"model_name": "included-model"}], } + @pytest.mark.asyncio + async def test_the_blocking_s3_read_runs_off_the_event_loop_thread(self, monkeypatch): + loop_thread = threading.current_thread() + read_threads = [] + + def record_thread(bucket_name, object_key): + read_threads.append(threading.current_thread()) + return {"model_list": [{"model_name": "a-model"}]} + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", record_thread + ) + + await get_config_from_bucket(bucket_type="s3", bucket_name="litellm-configs", object_key="config.yaml") + + assert read_threads and loop_thread not in read_threads + @pytest.mark.asyncio async def test_get_config_from_bucket_merges_includes_over_gcs(self, monkeypatch): objects = { @@ -232,11 +305,20 @@ class TestBucketConfigIncludes: "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, } - async def fake_gcs(bucket_name, object_key): - return objects.get(object_key) + buckets = [] + + class FakeGCSBucket: + def __init__(self): + self.requested = [] + buckets.append(self) + + async def download_gcs_object(self, object_key): + self.requested.append(object_key) + return yaml.dump(objects[object_key]).encode("utf-8") monkeypatch.setattr( - "litellm.proxy.common_utils.load_config_utils.get_config_file_contents_from_gcs", fake_gcs + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), ) config = await get_config_from_bucket( @@ -247,6 +329,9 @@ class TestBucketConfigIncludes: "general_settings": {"master_key": "sk-1234"}, "model_list": [{"model_name": "included-model"}], } + assert [bucket.requested for bucket in buckets] == [ + ["lit6982/config.yaml", "lit6982/model_config.yaml"] + ] @pytest.mark.asyncio async def test_get_config_from_bucket_returns_none_when_the_root_object_is_missing(self, monkeypatch): diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 9b9c4af9e42..3d9138ad098 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -8,6 +8,7 @@ Pins covered: from __future__ import annotations +import asyncio import json import os import re @@ -716,7 +717,7 @@ async def test_ProxyConfig__process_includes_merges_files(tmp_path): inc.write_text("model_list:\n - model_name: gpt-4\n") pc = ProxyConfig() cfg = {"include": ["models.yaml"], "model_list": [], "litellm_settings": {}} - result = await pc._process_includes(cfg, base_dir=str(tmp_path)) + result = await pc._process_includes(cfg, config_file_path=str(tmp_path / "config.yaml")) assert result == { "model_list": [{"model_name": "gpt-4"}], "litellm_settings": {}, @@ -727,17 +728,72 @@ async def test_ProxyConfig__process_includes_merges_files(tmp_path): async def test_ProxyConfig__process_includes_missing_file_raises(tmp_path): pc = ProxyConfig() with pytest.raises(FileNotFoundError): - await pc._process_includes({"include": ["nope.yaml"]}, base_dir=str(tmp_path)) + await pc._process_includes({"include": ["nope.yaml"]}, config_file_path=str(tmp_path / "config.yaml")) @pytest.mark.asyncio async def test_ProxyConfig__process_includes_follows_nested_includes(tmp_path): (tmp_path / "models.yaml").write_text("include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n") (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: second\n") - result = await ProxyConfig()._process_includes({"include": ["models.yaml"]}, base_dir=str(tmp_path)) + result = await ProxyConfig()._process_includes( + {"include": ["models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_resolves_a_nested_include_next_to_its_own_file(tmp_path): + """A nested `include` names a sibling of the file that declares it, not of the root config.""" + (tmp_path / "shared").mkdir() + (tmp_path / "shared" / "models.yaml").write_text( + "include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n" + ) + (tmp_path / "shared" / "more_models.yaml").write_text("model_list:\n - model_name: second\n") + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: wrong-directory\n") + + result = await ProxyConfig()._process_includes( + {"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_merges_a_shared_file_once(tmp_path): + (tmp_path / "shared.yaml").write_text("model_list:\n - model_name: shared\n") + (tmp_path / "a.yaml").write_text("include:\n - shared.yaml\n") + (tmp_path / "b.yaml").write_text("include:\n - ./shared.yaml\n") + + result = await ProxyConfig()._process_includes( + {"include": ["a.yaml", "b.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "shared"}]} + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_names_the_file_when_it_is_not_a_mapping(tmp_path): + (tmp_path / "models.yaml").write_text("- model_name: gpt-4\n") + + with pytest.raises(ValueError, match=re.escape(str(tmp_path / "models.yaml"))): + await ProxyConfig()._process_includes( + {"include": ["models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_terminates_on_a_cycle(tmp_path): + (tmp_path / "a.yaml").write_text("include:\n - b.yaml\nmodel_list:\n - model_name: from-a\n") + (tmp_path / "b.yaml").write_text("include:\n - a.yaml\nmodel_list:\n - model_name: from-b\n") + + result = await asyncio.wait_for( + ProxyConfig()._process_includes({"include": ["a.yaml"]}, config_file_path=str(tmp_path / "config.yaml")), + timeout=10, + ) + + assert result == {"model_list": [{"model_name": "from-a"}, {"model_name": "from-b"}]} + + # --------------------------------------------------------------------------- # ProxyConfig.save_config # --------------------------------------------------------------------------- From 6d993b50932941525f3b20fecd72fd81219541b5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 02:45:23 -0700 Subject: [PATCH 11/55] feat(proxy): serve registered skills as an Agent Skills well-known index Adds an opt-in Agent Skills discovery index at /.well-known/agent-skills/index.json (and the /.well-known/skills/index.json alias) plus GET /v1/skills/{skill_id}/archive, so `npx skills add ` installs skills uploaded through the Skills Gateway into any agent the CLI supports. The archive route repacks the stored upload so SKILL.md sits at the archive root, with fixed entry timestamps so the SHA-256 digest published in the index reproduces. Both routes are unauthenticated, since discovery clients send no credentials, and stay 404 until an admin sets `litellm_settings.public_skills_index: true`. --- litellm/__init__.py | 1 + litellm/proxy/_lazy_openapi_snapshot.json | 41 +++++ litellm/proxy/discovery_endpoints/__init__.py | 3 +- .../agent_skills_archive.py | 130 +++++++++++++ .../agent_skills_endpoints.py | 171 ++++++++++++++++++ litellm/proxy/proxy_server.py | 6 +- .../agent_skills_endpoints.py | 25 +++ .../test_agent_skills_archive.py | 106 +++++++++++ .../test_agent_skills_endpoints.py | 164 +++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 157 ++++++++++++++++ 10 files changed, 802 insertions(+), 2 deletions(-) create mode 100644 litellm/proxy/discovery_endpoints/agent_skills_archive.py create mode 100644 litellm/proxy/discovery_endpoints/agent_skills_endpoints.py create mode 100644 litellm/types/proxy/discovery_endpoints/agent_skills_endpoints.py create mode 100644 tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_archive.py create mode 100644 tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py diff --git a/litellm/__init__.py b/litellm/__init__.py index 42c0ea881fd..33840c66b7c 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -492,6 +492,7 @@ disable_copilot_system_to_assistant: bool = False # If false (default), convert public_mcp_servers: Optional[List[str]] = None public_mcp_hub_strict_whitelist: bool = True public_model_groups: Optional[List[str]] = None +public_skills_index: bool = False public_agent_groups: Optional[List[str]] = None agent_search_embedding_model: Optional[str] = None mcp_tool_search: Optional[Mapping[str, object]] = None diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 91a97ad6544..1e53049a887 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -5065,6 +5065,47 @@ "anthropic_skills" ] } + }, + "/v1/skills/{skill_id}/archive": { + "get": { + "description": "Stored skill upload, repacked so SKILL.md sits at the archive root.", + "operationId": "agent_skills_archive_v1_skills__skill_id__archive_get", + "parameters": [ + { + "in": "path", + "name": "skill_id", + "required": true, + "schema": { + "title": "Skill Id", + "type": "string" + } + } + ], + "responses": { + "200": { + "content": { + "application/json": { + "schema": {} + } + }, + "description": "Successful Response" + }, + "422": { + "content": { + "application/json": { + "schema": { + "$ref": "#/components/schemas/HTTPValidationError" + } + } + }, + "description": "Validation Error" + } + }, + "summary": "Agent Skills Archive", + "tags": [ + "anthropic_skills" + ] + } } } }, diff --git a/litellm/proxy/discovery_endpoints/__init__.py b/litellm/proxy/discovery_endpoints/__init__.py index a6401c2f1b4..52602f30b77 100644 --- a/litellm/proxy/discovery_endpoints/__init__.py +++ b/litellm/proxy/discovery_endpoints/__init__.py @@ -1,3 +1,4 @@ +from .agent_skills_endpoints import router as agent_skills_discovery_router from .ui_discovery_endpoints import router as ui_discovery_endpoints_router -__all__ = ["ui_discovery_endpoints_router"] +__all__ = ["agent_skills_discovery_router", "ui_discovery_endpoints_router"] diff --git a/litellm/proxy/discovery_endpoints/agent_skills_archive.py b/litellm/proxy/discovery_endpoints/agent_skills_archive.py new file mode 100644 index 00000000000..1f2fca3992e --- /dev/null +++ b/litellm/proxy/discovery_endpoints/agent_skills_archive.py @@ -0,0 +1,130 @@ +"""Repack a stored skill upload into the archive shape Agent Skills clients install from. + +Uploads follow the Anthropic Skills API layout, where every file sits under a single +top-level folder. Discovery clients read ``SKILL.md`` from the archive root, so that +folder is stripped and the zip is rebuilt with fixed entry timestamps, which keeps the +SHA-256 digest published in the index reproducible for identical uploads. +""" + +import hashlib +import io +import re +import zipfile +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final + +import yaml + +MAX_ARCHIVE_UNPACKED_BYTES: Final = 50 * 1024 * 1024 +MAX_ARCHIVE_ENTRIES: Final = 1000 +SKILL_MANIFEST_FILENAME: Final = "SKILL.md" + +_ZIP_ENTRY_TIMESTAMP: Final = (1980, 1, 1, 0, 0, 0) +_ZIP_ENTRY_PERMISSIONS: Final = 0o644 << 16 +_FRONTMATTER_PATTERN: Final = re.compile(r"^---\s*\n(.*?)\n---\s*(?:\n|$)", re.DOTALL) +_WINDOWS_DRIVE_PATTERN: Final = re.compile(r"^[A-Za-z]:") +_EMPTY_FRONTMATTER: Final[Mapping[str, object]] = MappingProxyType({}) + + +@dataclass(frozen=True, slots=True) +class SkillArchive: + content: bytes + digest: str + declared_name: str | None + declared_description: str | None + + +def build_skill_archive(stored_content: bytes) -> SkillArchive | None: + """Return the installable archive for an upload, or None when it holds no root SKILL.md.""" + try: + with zipfile.ZipFile(io.BytesIO(stored_content)) as uploaded: + members: Final = _flattened_members(uploaded) + except (zipfile.BadZipFile, OSError, RuntimeError): + return None + + if members is None: + return None + + frontmatter: Final = _manifest_frontmatter(next(data for name, data in members if name == SKILL_MANIFEST_FILENAME)) + content: Final = _repack(members) + return SkillArchive( + content=content, + digest=f"sha256:{hashlib.sha256(content).hexdigest()}", + declared_name=_frontmatter_text(frontmatter, "name"), + declared_description=_frontmatter_text(frontmatter, "description"), + ) + + +def _flattened_members(uploaded: zipfile.ZipFile) -> tuple[tuple[str, bytes], ...] | None: + infos: Final = tuple(info for info in uploaded.infolist() if not info.is_dir()) + if not infos or len(infos) > MAX_ARCHIVE_ENTRIES: + return None + if sum(info.file_size for info in infos) > MAX_ARCHIVE_UNPACKED_BYTES: + return None + + normalized: Final = tuple((info, _normalized_path(info.filename)) for info in infos) + if any(path is None for _, path in normalized): + return None + + prefix: Final = _common_root_prefix(tuple(path for _, path in normalized if path is not None)) + flattened: Final = tuple((info, path[len(prefix) :]) for info, path in normalized if path is not None) + names: Final = frozenset(name for _, name in flattened) + if SKILL_MANIFEST_FILENAME not in names or len(names) != len(flattened): + return None + + return tuple((name, uploaded.read(info)) for info, name in sorted(flattened, key=lambda member: member[1])) + + +def _common_root_prefix(paths: tuple[str, ...]) -> str: + roots: Final = frozenset(path.split("/", 1)[0] for path in paths) + if len(roots) != 1 or not all("/" in path for path in paths): + return "" + return f"{next(iter(roots))}/" + + +def _normalized_path(raw_path: str) -> str | None: + if not raw_path or "\0" in raw_path or "\\" in raw_path: + return None + if raw_path.startswith("/") or _WINDOWS_DRIVE_PATTERN.match(raw_path): + return None + parts: Final = tuple(part for part in raw_path.split("/") if part) + if not parts or any(part in (".", "..") for part in parts): + return None + return "/".join(parts) + + +def _manifest_frontmatter(manifest: bytes) -> Mapping[str, object]: + match: Final = _FRONTMATTER_PATTERN.match(manifest.decode("utf-8", errors="replace")) + if match is None: + return _EMPTY_FRONTMATTER + try: + parsed: Final = yaml.safe_load(match.group(1)) + except yaml.YAMLError: + return _EMPTY_FRONTMATTER + if not isinstance(parsed, dict): + return _EMPTY_FRONTMATTER + return parsed + + +def _frontmatter_text(frontmatter: Mapping[str, object], key: str) -> str | None: + value: Final = frontmatter.get(key) + if not isinstance(value, str): + return None + return value.strip() or None + + +def _zip_entry(name: str) -> zipfile.ZipInfo: + entry: Final = zipfile.ZipInfo(filename=name, date_time=_ZIP_ENTRY_TIMESTAMP) + entry.compress_type = zipfile.ZIP_DEFLATED + entry.external_attr = _ZIP_ENTRY_PERMISSIONS + return entry + + +def _repack(members: tuple[tuple[str, bytes], ...]) -> bytes: + buffer: Final = io.BytesIO() + with zipfile.ZipFile(buffer, mode="w", compression=zipfile.ZIP_DEFLATED) as repacked: + for name, data in members: + repacked.writestr(_zip_entry(name), data) + return buffer.getvalue() diff --git a/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py new file mode 100644 index 00000000000..d3c3b62fbfa --- /dev/null +++ b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py @@ -0,0 +1,171 @@ +"""Serve skills stored on the proxy as an Agent Skills well-known discovery index. + +``npx skills add -a `` reads ``/.well-known/agent-skills/index.json`` +and downloads each entry's archive. Discovery clients send no credentials, so both +routes are unauthenticated and stay off until ``litellm_settings.public_skills_index`` +is enabled, which publishes every stored skill to anyone who can reach the proxy. +""" + +import re +from collections.abc import Sequence +from itertools import groupby +from operator import itemgetter +from types import MappingProxyType +from typing import Final + +from fastapi import APIRouter, Depends, HTTPException, Request, Response + +import litellm +from litellm._logging import verbose_proxy_logger +from litellm.models.skills import LiteLLM_SkillsTable +from litellm.proxy.discovery_endpoints.agent_skills_archive import SkillArchive, build_skill_archive +from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import ( + MAX_SKILL_DESCRIPTION_LENGTH, + MAX_SKILL_NAME_LENGTH, + AgentSkillsIndex, + AgentSkillsIndexEntry, +) + +MAX_INDEXED_SKILLS: Final = 1000 + +_NON_SLUG_PATTERN: Final = re.compile(r"[^a-z0-9]+") +_FALLBACK_SKILL_NAME: Final = "skill" + +router: Final = APIRouter(tags=["public", "skills"]) # mutable-ok: fastapi types tags as list[str | Enum] + + +def ensure_index_enabled() -> None: + if litellm.public_skills_index is not True: + raise HTTPException(status_code=404, detail="Not Found") + + +async def stored_skills() -> Sequence[LiteLLM_SkillsTable]: + from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler + + return await LiteLLMSkillsHandler.list_skills(limit=MAX_INDEXED_SKILLS) + + +async def stored_skill(skill_id: str) -> LiteLLM_SkillsTable | None: + from litellm.llms.litellm_proxy.skills.handler import LiteLLMSkillsHandler + + try: + return await LiteLLMSkillsHandler.get_skill(skill_id) + except ValueError: + return None + + +@router.get( + "/.well-known/agent-skills/index.json", + response_model=AgentSkillsIndex, + dependencies=(Depends(ensure_index_enabled),), +) +@router.get( + "/.well-known/skills/index.json", + response_model=AgentSkillsIndex, + dependencies=(Depends(ensure_index_enabled),), + include_in_schema=False, +) +async def agent_skills_index( + request: Request, + skills: Sequence[LiteLLM_SkillsTable] = Depends(stored_skills), +) -> AgentSkillsIndex: + """Agent Skills v0.2.0 discovery index over every skill stored on this proxy.""" + from litellm.proxy.utils import get_custom_url + + installable: Final = tuple( + (skill, archive) + for skill, archive in ((skill, _archive_for(skill)) for skill in reversed(skills)) + if archive is not None + ) + names: Final = _deduplicated(tuple(_base_name(skill, archive) for skill, archive in installable)) + + return AgentSkillsIndex( + skills=tuple( + AgentSkillsIndexEntry( + name=name, + type="archive", + description=_description(skill, archive, name), + url=get_custom_url( + request_base_url=str(request.base_url), + route=f"v1/skills/{skill.skill_id}/archive", + ), + digest=archive.digest, + ) + for (skill, archive), name in zip(installable, names, strict=True) + ) + ) + + +@router.get( + "/v1/skills/{skill_id}/archive", + dependencies=(Depends(ensure_index_enabled),), +) +async def agent_skills_archive( + skill_id: str, + skill: LiteLLM_SkillsTable | None = Depends(stored_skill), +) -> Response: + """Stored skill upload, repacked so SKILL.md sits at the archive root.""" + archive: Final = _archive_for(skill) if skill is not None else None + if archive is None: + raise HTTPException(status_code=404, detail=f"No installable skill archive for: {skill_id}") + + return Response( + content=archive.content, + media_type="application/zip", + headers=MappingProxyType({"Content-Disposition": f'attachment; filename="{skill_id}.zip"'}), + ) + + +def _archive_for(skill: LiteLLM_SkillsTable) -> SkillArchive | None: + if skill.file_content is None: + return None + + archive: Final = build_skill_archive(skill.file_content) + if archive is None: + verbose_proxy_logger.warning( + "Agent Skills index: skipping skill %s, its upload is not a zip holding SKILL.md at the root of a " + "single top-level folder", + skill.skill_id, + ) + return archive + + +def _base_name(skill: LiteLLM_SkillsTable, archive: SkillArchive) -> str: + candidates: Final = (archive.declared_name, skill.display_title, skill.skill_id) + return next( + (slug for slug in (_slugify(candidate) for candidate in candidates) if slug is not None), + _FALLBACK_SKILL_NAME, + ) + + +def _slugify(raw: str | None) -> str | None: + if raw is None: + return None + return _NON_SLUG_PATTERN.sub("-", raw.lower()).strip("-")[:MAX_SKILL_NAME_LENGTH].rstrip("-") or None + + +def _deduplicated(names: Sequence[str]) -> tuple[str, ...]: + ordinals: Final = MappingProxyType( + { + position: ordinal + for _, duplicates in groupby(sorted(enumerate(names), key=itemgetter(1)), key=itemgetter(1)) + for ordinal, (position, _) in enumerate(duplicates) + } + ) + return tuple(_with_ordinal(name, ordinals[position]) for position, name in enumerate(names)) + + +def _with_ordinal(name: str, ordinal: int) -> str: + if ordinal == 0: + return name + suffix: Final = f"-{ordinal + 1}" + return f"{name[: MAX_SKILL_NAME_LENGTH - len(suffix)].rstrip('-')}{suffix}" + + +def _description(skill: LiteLLM_SkillsTable, archive: SkillArchive, name: str) -> str: + candidates: Final = (archive.declared_description, skill.description, skill.display_title) + chosen: Final = next( + (candidate.strip() for candidate in candidates if candidate is not None and candidate.strip()), + name, + ) + return chosen[:MAX_SKILL_DESCRIPTION_LENGTH] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ba5714fe950..4f0e5860b84 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -431,7 +431,10 @@ from litellm.proxy.db.proxy_worker_heartbeat import ( ProxyWorkerHeartbeat, ) from litellm.proxy.db.spend_counter_reseed import END_USER_COUNTER_PREFIX, SpendCounterReseed -from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router +from litellm.proxy.discovery_endpoints import ( + agent_skills_discovery_router, + ui_discovery_endpoints_router, +) from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config from litellm.proxy.google_endpoints.endpoints import router as google_router @@ -18370,6 +18373,7 @@ app.include_router(user_agent_analytics_router) app.include_router(gateway_request_router) app.include_router(enterprise_router) app.include_router(ui_discovery_endpoints_router) +app.include_router(agent_skills_discovery_router) # Eager: /models/{name}:method overlaps with the OpenAI /models endpoint. app.include_router(google_router) diff --git a/litellm/types/proxy/discovery_endpoints/agent_skills_endpoints.py b/litellm/types/proxy/discovery_endpoints/agent_skills_endpoints.py new file mode 100644 index 00000000000..0d8bb29e172 --- /dev/null +++ b/litellm/types/proxy/discovery_endpoints/agent_skills_endpoints.py @@ -0,0 +1,25 @@ +"""Agent Skills discovery index, version 0.2.0. + +Schema: https://schemas.agentskills.io/discovery/0.2.0/schema.json +""" + +from typing import Final, Literal + +from pydantic import BaseModel, Field + +AGENT_SKILLS_DISCOVERY_SCHEMA_URL: Final = "https://schemas.agentskills.io/discovery/0.2.0/schema.json" +MAX_SKILL_NAME_LENGTH: Final = 64 +MAX_SKILL_DESCRIPTION_LENGTH: Final = 1024 + + +class AgentSkillsIndexEntry(BaseModel): + name: str + type: Literal["archive"] + description: str + url: str + digest: str + + +class AgentSkillsIndex(BaseModel): + discovery_schema: str = Field(default=AGENT_SKILLS_DISCOVERY_SCHEMA_URL, alias="$schema") + skills: tuple[AgentSkillsIndexEntry, ...] diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_archive.py b/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_archive.py new file mode 100644 index 00000000000..dd5ac230aac --- /dev/null +++ b/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_archive.py @@ -0,0 +1,106 @@ +import hashlib +import io +import zipfile + +from litellm.proxy.discovery_endpoints.agent_skills_archive import ( + MAX_ARCHIVE_ENTRIES, + build_skill_archive, +) + +MANIFEST = b"""--- +name: pdf-summarizer +description: Summarize a PDF into an executive brief. +--- + +Read the PDF, then write the brief. +""" + + +def zip_bytes(files: dict[str, bytes]) -> bytes: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive: + for name, content in files.items(): + archive.writestr(name, content) + return buffer.getvalue() + + +def entries_of(content: bytes) -> dict[str, bytes]: + with zipfile.ZipFile(io.BytesIO(content)) as archive: + return {name: archive.read(name) for name in archive.namelist()} + + +def test_single_top_level_folder_is_stripped_so_skill_md_sits_at_the_root(): + archive = build_skill_archive( + zip_bytes( + { + "pdf-summarizer/SKILL.md": MANIFEST, + "pdf-summarizer/reference.md": b"page citations", + "pdf-summarizer/scripts/extract.py": b"print('hi')", + } + ) + ) + + assert archive is not None + assert entries_of(archive.content) == { + "SKILL.md": MANIFEST, + "reference.md": b"page citations", + "scripts/extract.py": b"print('hi')", + } + + +def test_digest_covers_the_repacked_bytes_and_is_stable_across_builds(): + upload = zip_bytes({"pdf-summarizer/SKILL.md": MANIFEST, "pdf-summarizer/reference.md": b"page citations"}) + + first = build_skill_archive(upload) + second = build_skill_archive(upload) + + assert first is not None and second is not None + assert first.digest == f"sha256:{hashlib.sha256(first.content).hexdigest()}" + assert first.content == second.content + + +def test_an_upload_that_is_already_flat_keeps_every_file_where_it_is(): + archive = build_skill_archive(zip_bytes({"SKILL.md": MANIFEST, "reference.md": b"page citations"})) + + assert archive is not None + assert sorted(entries_of(archive.content)) == ["SKILL.md", "reference.md"] + + +def test_manifest_frontmatter_supplies_the_declared_name_and_description(): + archive = build_skill_archive(zip_bytes({"pdf-summarizer/SKILL.md": MANIFEST})) + + assert archive is not None + assert archive.declared_name == "pdf-summarizer" + assert archive.declared_description == "Summarize a PDF into an executive brief." + + +def test_a_manifest_without_frontmatter_declares_nothing(): + archive = build_skill_archive(zip_bytes({"pdf-summarizer/SKILL.md": b"just prose, no frontmatter"})) + + assert archive is not None + assert archive.declared_name is None + assert archive.declared_description is None + + +def test_a_manifest_buried_below_the_stripped_folder_is_not_installable(): + assert build_skill_archive(zip_bytes({"pdf-summarizer/nested/SKILL.md": MANIFEST})) is None + + +def test_an_upload_with_no_manifest_is_not_installable(): + assert build_skill_archive(zip_bytes({"pdf-summarizer/reference.md": b"page citations"})) is None + + +def test_a_non_zip_upload_is_not_installable(): + assert build_skill_archive(MANIFEST) is None + + +def test_a_path_traversal_entry_is_not_installable(): + assert build_skill_archive(zip_bytes({"SKILL.md": MANIFEST, "../escape.md": b"nope"})) is None + + +def test_an_upload_over_the_entry_cap_is_not_installable(): + files = {"pdf-summarizer/SKILL.md": MANIFEST} | { + f"pdf-summarizer/file-{index}.md": b"x" for index in range(MAX_ARCHIVE_ENTRIES) + } + + assert build_skill_archive(zip_bytes(files)) is None diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py new file mode 100644 index 00000000000..4fbb243a2b3 --- /dev/null +++ b/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py @@ -0,0 +1,164 @@ +import hashlib +import io +import zipfile + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +import litellm +from litellm.models.skills import LiteLLM_SkillsTable +from litellm.proxy.discovery_endpoints.agent_skills_endpoints import ( + router, + stored_skill, + stored_skills, +) +from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import ( + AGENT_SKILLS_DISCOVERY_SCHEMA_URL, +) + +WELL_KNOWN_PATHS = ("/.well-known/agent-skills/index.json", "/.well-known/skills/index.json") + +MANIFEST = b"""--- +name: pdf-summarizer +description: Summarize a PDF into an executive brief. +--- + +Read the PDF, then write the brief. +""" + + +def zip_bytes(files: dict[str, bytes]) -> bytes: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive: + for name, content in files.items(): + archive.writestr(name, content) + return buffer.getvalue() + + +def skill( + skill_id: str, + *, + display_title: str | None = "PDF Summarizer", + description: str | None = None, + files: dict[str, bytes] | None = None, +) -> LiteLLM_SkillsTable: + return LiteLLM_SkillsTable( + skill_id=skill_id, + display_title=display_title, + description=description, + file_content=zip_bytes(files if files is not None else {"pdf-summarizer/SKILL.md": MANIFEST}), + ) + + +def client_for(*skills: LiteLLM_SkillsTable) -> TestClient: + app = FastAPI() + app.include_router(router) + + def _skills() -> tuple[LiteLLM_SkillsTable, ...]: + return skills + + def _skill(skill_id: str) -> LiteLLM_SkillsTable | None: + return next((candidate for candidate in skills if candidate.skill_id == skill_id), None) + + app.dependency_overrides[stored_skills] = _skills + app.dependency_overrides[stored_skill] = _skill + return TestClient(app) + + +@pytest.fixture +def index_enabled(monkeypatch): + monkeypatch.setattr(litellm, "public_skills_index", True) + + +def test_discovery_is_absent_until_public_skills_index_is_enabled(monkeypatch): + monkeypatch.setattr(litellm, "public_skills_index", False) + client = client_for(skill("litellm_skill_1")) + + for path in WELL_KNOWN_PATHS: + assert client.get(path).status_code == 404 + assert client.get("/v1/skills/litellm_skill_1/archive").status_code == 404 + + +@pytest.mark.parametrize("path", WELL_KNOWN_PATHS) +def test_index_publishes_each_stored_skill_in_the_v0_2_0_shape(index_enabled, path): + client = client_for(skill("litellm_skill_1")) + + body = client.get(path).json() + + assert body["$schema"] == AGENT_SKILLS_DISCOVERY_SCHEMA_URL + assert len(body["skills"]) == 1 + entry = body["skills"][0] + assert entry["name"] == "pdf-summarizer" + assert entry["type"] == "archive" + assert entry["description"] == "Summarize a PDF into an executive brief." + assert entry["url"].endswith("/v1/skills/litellm_skill_1/archive") + assert entry["digest"].startswith("sha256:") + + +def test_index_digest_matches_the_bytes_the_archive_route_serves(index_enabled): + client = client_for(skill("litellm_skill_1")) + + entry = client.get(WELL_KNOWN_PATHS[0]).json()["skills"][0] + downloaded = client.get(entry["url"]) + + assert downloaded.status_code == 200 + assert downloaded.headers["content-type"] == "application/zip" + assert entry["digest"] == f"sha256:{hashlib.sha256(downloaded.content).hexdigest()}" + + +def test_install_name_falls_back_to_the_manifest_name_without_a_display_title(index_enabled): + client = client_for(skill("litellm_skill_1", display_title=None)) + + assert client.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["name"] == "pdf-summarizer" + + +@pytest.mark.parametrize( + "manifest, stored_description, expected", + [ + (MANIFEST, "registry copy", "Summarize a PDF into an executive brief."), + (b"no frontmatter here", "registry copy", "registry copy"), + (b"no frontmatter here", None, "PDF Summarizer"), + ], +) +def test_description_prefers_the_manifest_then_the_registry_then_the_title( + index_enabled, manifest, stored_description, expected +): + client = client_for( + skill( + "litellm_skill_1", + description=stored_description, + files={"pdf-summarizer/SKILL.md": manifest}, + ) + ) + + assert client.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["description"] == expected + + +def test_skills_sharing_a_title_get_distinct_install_names(index_enabled): + client = client_for( + skill("litellm_skill_2", files={"pdf-summarizer/SKILL.md": b"second"}), + skill("litellm_skill_1", files={"pdf-summarizer/SKILL.md": b"first"}), + ) + + names = [entry["name"] for entry in client.get(WELL_KNOWN_PATHS[0]).json()["skills"]] + + assert names == ["pdf-summarizer", "pdf-summarizer-2"] + + +def test_uploads_without_a_root_manifest_are_left_out_of_the_index(index_enabled): + client = client_for( + skill("litellm_skill_1"), + skill("litellm_skill_2", files={"pdf-summarizer/reference.md": b"no manifest"}), + ) + + body = client.get(WELL_KNOWN_PATHS[0]).json() + + assert [entry["url"].split("/")[-2] for entry in body["skills"]] == ["litellm_skill_1"] + assert client.get("/v1/skills/litellm_skill_2/archive").status_code == 404 + + +def test_archive_route_404s_for_a_skill_that_does_not_exist(index_enabled): + client = client_for(skill("litellm_skill_1")) + + assert client.get("/v1/skills/litellm_skill_missing/archive").status_code == 404 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b1534c19670..03496f9a507 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -21,6 +21,26 @@ export interface paths { patch?: never; trace?: never; }; + "/.well-known/agent-skills/index.json": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Agent Skills Index + * @description Agent Skills v0.2.0 discovery index over every skill stored on this proxy. + */ + get: operations["agent_skills_index__well_known_agent_skills_index_json_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/.well-known/jwks.json": { parameters: { query?: never; @@ -319,6 +339,26 @@ export interface paths { patch?: never; trace?: never; }; + "/.well-known/skills/index.json": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Agent Skills Index + * @description Agent Skills v0.2.0 discovery index over every skill stored on this proxy. + */ + get: operations["agent_skills_index__well_known_skills_index_json_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/a2a/{agent_id}": { parameters: { query?: never; @@ -19962,6 +20002,26 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/skills/{skill_id}/archive": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Agent Skills Archive + * @description Stored skill upload, repacked so SKILL.md sits at the archive root. + */ + get: operations["agent_skills_archive_v1_skills__skill_id__archive_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/threads": { parameters: { query?: never; @@ -23162,6 +23222,32 @@ export interface components { /** Tags */ tags?: string[]; }; + /** AgentSkillsIndex */ + AgentSkillsIndex: { + /** + * $Schema + * @default https://schemas.agentskills.io/discovery/0.2.0/schema.json + */ + $schema: string; + /** Skills */ + skills: components["schemas"]["AgentSkillsIndexEntry"][]; + }; + /** AgentSkillsIndexEntry */ + AgentSkillsIndexEntry: { + /** Description */ + description: string; + /** Digest */ + digest: string; + /** Name */ + name: string; + /** + * Type + * @constant + */ + type: "archive"; + /** Url */ + url: string; + }; /** * AlertType * @description Enum for alert types and management event types @@ -39831,6 +39917,26 @@ export interface operations { }; }; }; + agent_skills_index__well_known_agent_skills_index_json_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AgentSkillsIndex"]; + }; + }; + }; + }; jwks_json__well_known_jwks_json_get: { parameters: { query?: never; @@ -40168,6 +40274,26 @@ export interface operations { }; }; }; + agent_skills_index__well_known_skills_index_json_get: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["AgentSkillsIndex"]; + }; + }; + }; + }; invoke_agent_a2a_a2a__agent_id__post: { parameters: { query?: never; @@ -64677,6 +64803,37 @@ export interface operations { }; }; }; + agent_skills_archive_v1_skills__skill_id__archive_get: { + parameters: { + query?: never; + header?: never; + path: { + skill_id: string; + }; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; create_threads_v1_threads_post: { parameters: { query?: never; From 12204e523079f07aeefdb2257e2b2c938771aae4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 02:47:47 -0700 Subject: [PATCH 12/55] fix(proxy): resolve disk includes next to the config that declares them Keep reading an include left beside the root config, with a warning naming where it was found, so a nested include written against the old rule still boots. Also build one S3 client per config load rather than one per included object, treat an empty included object as an empty config instead of failing the boot, and point the error a dropped bucket include raises at the bucket error logged with it. --- litellm/proxy/common_utils/config_includes.py | 28 +++++++ .../proxy/common_utils/load_config_utils.py | 77 ++++++++++------- litellm/proxy/proxy_server.py | 4 +- .../common_utils/test_load_config_utils.py | 83 ++++++++++++++----- .../proxy/proxy_server/test_proxy_config.py | 21 ++++- 5 files changed, 160 insertions(+), 53 deletions(-) diff --git a/litellm/proxy/common_utils/config_includes.py b/litellm/proxy/common_utils/config_includes.py index cf39bd6435e..23e6c615a09 100644 --- a/litellm/proxy/common_utils/config_includes.py +++ b/litellm/proxy/common_utils/config_includes.py @@ -1,10 +1,38 @@ +import os from collections.abc import Awaitable, Mapping from types import MappingProxyType from typing import Final, Protocol +from litellm._logging import verbose_proxy_logger + INCLUDE_KEY: Final = "include" +def resolve_include_file_path(include_file: str, declared_in: str, root_config_path: str) -> str: + """ + Resolve one `include` entry to the file it names, next to the config that declares it. + + A config written before nested entries resolved this way can name a file sitting next to the root + config instead, so that file is still read, with a warning naming where it was found. + """ + declared_relative: Final = os.path.abspath(os.path.join(os.path.dirname(declared_in), include_file)) + if os.path.exists(declared_relative): + return declared_relative + + root_relative: Final = os.path.abspath(os.path.join(os.path.dirname(root_config_path), include_file)) + if root_relative == declared_relative or not os.path.exists(root_relative): + return declared_relative + + verbose_proxy_logger.warning( + "Config include '%s' declared in %s was not found next to it, so %s was read instead. " + "Move the included file next to the config that declares it.", + include_file, + declared_in, + root_relative, + ) + return root_relative + + class ConfigLoader(Protocol): def __call__(self, include_entry: str, declared_in: str, /) -> Awaitable[tuple[str, Mapping[str, object]]]: ... diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index e72272b1aca..d62286a50f6 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -2,6 +2,7 @@ import asyncio import os import posixpath from collections.abc import Awaitable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol import yaml @@ -24,7 +25,19 @@ class BucketObjectReader(Protocol): def __call__(self, object_key: str, /) -> Awaitable[object | None]: ... -def get_file_contents_from_s3(bucket_name, object_key): +class SyncBucketObjectReader(Protocol): + def __call__(self, object_key: str, /) -> object | None: ... + + +def _parsed_config(file_contents: str) -> object: + parsed: Final = yaml.safe_load(file_contents) + return MappingProxyType({}) if parsed is None else parsed + + +def s3_object_reader(bucket_name: str) -> SyncBucketObjectReader: + """ + Build one reader for a whole config, so an `include` tree costs one S3 client rather than one per object. + """ try: # v0 rely on boto3 for authentication - allowing boto3 to handle IAM credentials etc import boto3 @@ -39,24 +52,28 @@ def get_file_contents_from_s3(bucket_name, object_key): aws_secret_access_key=credentials.secret_key, aws_session_token=credentials.token, # Optional, if using temporary credentials ) - verbose_proxy_logger.debug("Retrieving %s from S3 bucket: %s", object_key, bucket_name) - response: Final = s3_client.get_object(Bucket=bucket_name, Key=object_key) - verbose_proxy_logger.debug("Response: %s", response) - - # Read the file contents and directly parse YAML - file_contents: Final = response["Body"].read().decode("utf-8") - verbose_proxy_logger.debug("File contents retrieved from S3") - - # Parse YAML directly from string - config: Final = yaml.safe_load(file_contents) - return config - except ImportError as e: # this is most likely if a user is not using the litellm docker container verbose_proxy_logger.error("ImportError: %s", e) + return lambda object_key: None except Exception as e: - verbose_proxy_logger.error("Error retrieving file contents: %s", e) - return None + verbose_proxy_logger.error("Error creating the S3 client for bucket %s: %s", bucket_name, e) + return lambda object_key: None + + def read(object_key: str) -> object | None: + try: + verbose_proxy_logger.debug("Retrieving %s from S3 bucket: %s", object_key, bucket_name) + response: Final = s3_client.get_object(Bucket=bucket_name, Key=object_key) + return _parsed_config(response["Body"].read().decode("utf-8")) + except Exception as e: # noqa: BLE001 # any boto3 error must read as a missing object + verbose_proxy_logger.error("Error retrieving %s from S3 bucket %s: %s", object_key, bucket_name, e) + return None + + return read + + +def get_file_contents_from_s3(bucket_name: str, object_key: str) -> object | None: + return s3_object_reader(bucket_name)(object_key) def gcs_config_bucket(bucket_name: str) -> "GCSBucketBase | None": @@ -64,27 +81,27 @@ def gcs_config_bucket(bucket_name: str) -> "GCSBucketBase | None": from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger return GCSBucketLogger(bucket_name=bucket_name) - except Exception as e: + except Exception as e: # noqa: BLE001 # an unbuildable client must read as an unreadable bucket verbose_proxy_logger.error("Error creating the GCS client for bucket %s: %s", bucket_name, e) return None -async def get_config_file_contents_from_gcs(bucket_name, object_key, gcs_bucket=None): +async def get_config_file_contents_from_gcs( + bucket_name: str, + object_key: str, + gcs_bucket: "GCSBucketBase | None" = None, +) -> object | None: try: bucket: Final = gcs_config_bucket(bucket_name) if gcs_bucket is None else gcs_bucket if bucket is None: return None - file_contents = await bucket.download_gcs_object(object_key) + file_contents: Final = await bucket.download_gcs_object(object_key) if file_contents is None: raise Exception(f"File contents are None for {object_key}") - # file_contentis is a bytes object, so we need to convert it to yaml - file_contents = file_contents.decode("utf-8") - # convert to yaml - config: Final = yaml.safe_load(file_contents) - return config + return _parsed_config(file_contents.decode("utf-8")) except Exception as e: - verbose_proxy_logger.error("Error retrieving file contents: %s", e) + verbose_proxy_logger.error("Error retrieving %s from GCS bucket %s: %s", object_key, bucket_name, e) return None @@ -110,20 +127,24 @@ async def resolve_bucket_includes( include_key: Final = resolve_include_object_key(declared_in, include_entry) included: Final = await fetch(include_key) if included is None: - raise FileNotFoundError(f"Included config could not be read from bucket: {include_key}") + raise FileNotFoundError( + f"Included config could not be read from bucket: {include_key}. " + "The underlying bucket error is logged above." + ) return include_key, included return await resolve_includes(config=config, location=object_key, load=load) -def bucket_object_reader(bucket_type: str | None, bucket_name: str) -> BucketObjectReader: +async def bucket_object_reader(bucket_type: str | None, bucket_name: str) -> BucketObjectReader: """ Build one reader for a whole config, so an `include` tree costs one bucket client rather than one per object. """ if bucket_type != "gcs": + read_object: Final = await asyncio.to_thread(s3_object_reader, bucket_name) async def read_from_s3(object_key: str) -> object | None: - return await asyncio.to_thread(get_file_contents_from_s3, bucket_name, object_key) + return await asyncio.to_thread(read_object, object_key) return read_from_s3 @@ -143,7 +164,7 @@ async def get_config_from_bucket( bucket_name: str, object_key: str, ) -> dict[str, object] | None: - read: Final = bucket_object_reader(bucket_type, bucket_name) + read: Final = await bucket_object_reader(bucket_type, bucket_name) async def fetch(key: str) -> Mapping[str, object] | None: raw: Final = await read(key) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7ac5e697b61..7fada09febf 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -341,7 +341,7 @@ from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import ( AuthCacheInvalidationSubscriber, ) from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy -from litellm.proxy.common_utils.config_includes import resolve_includes +from litellm.proxy.common_utils.config_includes import resolve_include_file_path, resolve_includes from litellm.proxy.common_utils.config_sync_pubsub import ConfigSyncSubscriber from litellm.proxy.common_utils.debug_utils import init_verbose_loggers from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router @@ -4560,7 +4560,7 @@ class ProxyConfig: included_config_adapter: Final = TypeAdapter(dict[str, object]) async def load_included(include_file: str, declared_in: str) -> tuple[str, Mapping[str, object]]: - file_path: Final = os.path.abspath(os.path.join(os.path.dirname(declared_in), include_file)) + file_path: Final = resolve_include_file_path(include_file, declared_in, config_file_path) if not os.path.exists(file_path): raise FileNotFoundError(f"Included file not found: {file_path}") try: diff --git a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py index b654320569c..75aa5a3f0f8 100644 --- a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py @@ -93,11 +93,6 @@ class TestGetFileContentsFromS3: class TestBucketConfigIncludes: - """`include:` directives in a bucket-hosted config.yaml (LIT-6982). - - They used to be dropped silently: the proxy booted with the root config applied and everything - the included objects declared missing, with nothing logged. - """ @staticmethod def _bucket(objects): @@ -161,7 +156,6 @@ class TestBucketConfigIncludes: @pytest.mark.asyncio async def test_a_nested_include_resolves_against_the_object_that_declares_it(self): - """A nested `include` names a neighbour of the object declaring it, not of the root config.""" merged = await resolve_bucket_includes( config={"include": ["shared/models.yaml"]}, object_key="configs/config.yaml", @@ -265,8 +259,8 @@ class TestBucketConfigIncludes: "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, } monkeypatch.setattr( - "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", - lambda bucket_name, object_key: objects.get(object_key), + "litellm.proxy.common_utils.load_config_utils.s3_object_reader", + lambda bucket_name: objects.get, ) config = await get_config_from_bucket( @@ -279,21 +273,72 @@ class TestBucketConfigIncludes: } @pytest.mark.asyncio - async def test_the_blocking_s3_read_runs_off_the_event_loop_thread(self, monkeypatch): + async def test_the_blocking_s3_work_runs_off_the_event_loop_thread(self, monkeypatch): loop_thread = threading.current_thread() - read_threads = [] + threads = [] - def record_thread(bucket_name, object_key): - read_threads.append(threading.current_thread()) - return {"model_list": [{"model_name": "a-model"}]} + def build_reader(bucket_name): + threads.append(threading.current_thread()) - monkeypatch.setattr( - "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", record_thread - ) + def read(object_key): + threads.append(threading.current_thread()) + return {"model_list": [{"model_name": "a-model"}]} + + return read + + monkeypatch.setattr("litellm.proxy.common_utils.load_config_utils.s3_object_reader", build_reader) await get_config_from_bucket(bucket_type="s3", bucket_name="litellm-configs", object_key="config.yaml") - assert read_threads and loop_thread not in read_threads + assert len(threads) == 2 and loop_thread not in threads + + @pytest.mark.asyncio + async def test_one_s3_client_serves_the_whole_include_tree(self, monkeypatch): + objects = { + "lit6982/config.yaml": {"include": ["model_config.yaml"]}, + "lit6982/model_config.yaml": {"model_list": [{"model_name": "included-model"}]}, + } + readers = [] + + def build_reader(bucket_name): + requested = [] + readers.append(requested) + + def read(object_key): + requested.append(object_key) + return objects.get(object_key) + + return read + + monkeypatch.setattr("litellm.proxy.common_utils.load_config_utils.s3_object_reader", build_reader) + + await get_config_from_bucket( + bucket_type="s3", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert readers == [["lit6982/config.yaml", "lit6982/model_config.yaml"]] + + @pytest.mark.asyncio + async def test_an_empty_included_object_merges_as_an_empty_config(self, monkeypatch): + objects = { + "lit6982/config.yaml": "include:\n - empty.yaml\nmodel_list:\n - model_name: only-model\n", + "lit6982/empty.yaml": "", + } + + class FakeGCSBucket: + async def download_gcs_object(self, object_key): + return objects[object_key].encode("utf-8") + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), + ) + + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config == {"model_list": [{"model_name": "only-model"}]} @pytest.mark.asyncio async def test_get_config_from_bucket_merges_includes_over_gcs(self, monkeypatch): @@ -336,8 +381,8 @@ class TestBucketConfigIncludes: @pytest.mark.asyncio async def test_get_config_from_bucket_returns_none_when_the_root_object_is_missing(self, monkeypatch): monkeypatch.setattr( - "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", - lambda bucket_name, object_key: None, + "litellm.proxy.common_utils.load_config_utils.s3_object_reader", + lambda bucket_name: (lambda object_key: None), ) assert ( diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 3d9138ad098..e90f101a555 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -743,7 +743,6 @@ async def test_ProxyConfig__process_includes_follows_nested_includes(tmp_path): @pytest.mark.asyncio async def test_ProxyConfig__process_includes_resolves_a_nested_include_next_to_its_own_file(tmp_path): - """A nested `include` names a sibling of the file that declares it, not of the root config.""" (tmp_path / "shared").mkdir() (tmp_path / "shared" / "models.yaml").write_text( "include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n" @@ -758,6 +757,21 @@ async def test_ProxyConfig__process_includes_resolves_a_nested_include_next_to_i assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_still_reads_a_nested_include_left_beside_the_root_config(tmp_path): + (tmp_path / "shared").mkdir() + (tmp_path / "shared" / "models.yaml").write_text( + "include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n" + ) + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: second\n") + + result = await ProxyConfig()._process_includes( + {"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} + + @pytest.mark.asyncio async def test_ProxyConfig__process_includes_merges_a_shared_file_once(tmp_path): (tmp_path / "shared.yaml").write_text("model_list:\n - model_name: shared\n") @@ -1110,7 +1124,6 @@ async def test_ProxyConfig_get_config_loads_from_file(tmp_path, monkeypatch): @pytest.mark.asyncio async def test_ProxyConfig_get_config_from_a_bucket_merges_includes(monkeypatch): - """A bucket-hosted config.yaml used to drop its `include:` entries silently (LIT-6982).""" objects = { "lit6982/config.yaml": { "include": ["model_config.yaml"], @@ -1121,8 +1134,8 @@ async def test_ProxyConfig_get_config_from_a_bucket_merges_includes(monkeypatch) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", False) monkeypatch.setattr( - "litellm.proxy.common_utils.load_config_utils.get_file_contents_from_s3", - lambda bucket_name, object_key: objects.get(object_key), + "litellm.proxy.common_utils.load_config_utils.s3_object_reader", + lambda bucket_name: objects.get, ) monkeypatch.setenv("LITELLM_CONFIG_BUCKET_NAME", "litellm-configs") monkeypatch.setenv("LITELLM_CONFIG_BUCKET_OBJECT_KEY", "lit6982/config.yaml") From 9440dbac74a72865525dbfdc6ebf4a74f53e146e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 03:20:37 -0700 Subject: [PATCH 13/55] perf: cache repacked Agent Skills archives per skill version The well-known index needs every stored upload repacked to publish its digest, and both routes are unauthenticated, so each request was rebuilding every archive on the event loop. With 21 stored skills the index took ~2s and /health/liveliness on the same worker went from 1ms to 1.7s under two concurrent index requests. Repacking now runs off the event loop and each result is cached per skill version, so a worker builds an archive once until the skill changes. The archive route also declares application/zip in OpenAPI rather than JSON. --- litellm/proxy/_lazy_openapi_snapshot.json | 6 ++- .../agent_skills_endpoints.py | 51 +++++++++++++++---- .../test_agent_skills_endpoints.py | 47 +++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 4 files changed, 92 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 1e53049a887..94fa0b67caf 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -5084,8 +5084,10 @@ "responses": { "200": { "content": { - "application/json": { - "schema": {} + "application/zip": { + "schema": { + "type": "string" + } } }, "description": "Successful Response" diff --git a/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py index d3c3b62fbfa..3084cbfd84f 100644 --- a/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py +++ b/litellm/proxy/discovery_endpoints/agent_skills_endpoints.py @@ -6,6 +6,7 @@ routes are unauthenticated and stay off until ``litellm_settings.public_skills_i is enabled, which publishes every stored skill to anyone who can reach the proxy. """ +import asyncio import re from collections.abc import Sequence from itertools import groupby @@ -17,6 +18,7 @@ from fastapi import APIRouter, Depends, HTTPException, Request, Response import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.in_memory_cache import InMemoryCache from litellm.models.skills import LiteLLM_SkillsTable from litellm.proxy.discovery_endpoints.agent_skills_archive import SkillArchive, build_skill_archive from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import ( @@ -27,6 +29,15 @@ from litellm.types.proxy.discovery_endpoints.agent_skills_endpoints import ( ) MAX_INDEXED_SKILLS: Final = 1000 +MAX_CACHED_ARCHIVES: Final = 128 +MAX_CACHED_ARCHIVE_BYTES: Final = 512 * 1024 +ARCHIVE_CACHE_TTL_SECONDS: Final = 3600 + +_ARCHIVE_CACHE: Final = InMemoryCache( + max_size_in_memory=MAX_CACHED_ARCHIVES, + default_ttl=ARCHIVE_CACHE_TTL_SECONDS, + max_size_per_item=MAX_CACHED_ARCHIVE_BYTES // 1024, +) _NON_SLUG_PATTERN: Final = re.compile(r"[^a-z0-9]+") _FALLBACK_SKILL_NAME: Final = "skill" @@ -34,6 +45,12 @@ _FALLBACK_SKILL_NAME: Final = "skill" router: Final = APIRouter(tags=["public", "skills"]) # mutable-ok: fastapi types tags as list[str | Enum] +class ZipArchiveResponse(Response): + """Response whose OpenAPI entry declares an application/zip download rather than JSON.""" + + media_type = "application/zip" + + def ensure_index_enabled() -> None: if litellm.public_skills_index is not True: raise HTTPException(status_code=404, detail="Not Found") @@ -72,11 +89,7 @@ async def agent_skills_index( """Agent Skills v0.2.0 discovery index over every skill stored on this proxy.""" from litellm.proxy.utils import get_custom_url - installable: Final = tuple( - (skill, archive) - for skill, archive in ((skill, _archive_for(skill)) for skill in reversed(skills)) - if archive is not None - ) + installable: Final = await _installable(skills) names: Final = _deduplicated(tuple(_base_name(skill, archive) for skill, archive in installable)) return AgentSkillsIndex( @@ -99,34 +112,50 @@ async def agent_skills_index( @router.get( "/v1/skills/{skill_id}/archive", dependencies=(Depends(ensure_index_enabled),), + response_class=ZipArchiveResponse, ) async def agent_skills_archive( skill_id: str, skill: LiteLLM_SkillsTable | None = Depends(stored_skill), -) -> Response: +) -> ZipArchiveResponse: """Stored skill upload, repacked so SKILL.md sits at the archive root.""" - archive: Final = _archive_for(skill) if skill is not None else None + archive: Final = await _archive_for(skill) if skill is not None else None if archive is None: raise HTTPException(status_code=404, detail=f"No installable skill archive for: {skill_id}") - return Response( + return ZipArchiveResponse( content=archive.content, - media_type="application/zip", headers=MappingProxyType({"Content-Disposition": f'attachment; filename="{skill_id}.zip"'}), ) -def _archive_for(skill: LiteLLM_SkillsTable) -> SkillArchive | None: +async def _installable( + skills: Sequence[LiteLLM_SkillsTable], +) -> tuple[tuple[LiteLLM_SkillsTable, SkillArchive], ...]: + built: Final = tuple([(skill, await _archive_for(skill)) for skill in reversed(skills)]) + return tuple((skill, archive) for skill, archive in built if archive is not None) + + +async def _archive_for(skill: LiteLLM_SkillsTable) -> SkillArchive | None: if skill.file_content is None: return None - archive: Final = build_skill_archive(skill.file_content) + cache_key: Final = None if skill.updated_at is None else f"{skill.skill_id}:{skill.updated_at.isoformat()}" + cached: Final = None if cache_key is None else _ARCHIVE_CACHE.get_cache(cache_key) + if isinstance(cached, SkillArchive): + return cached + + archive: Final = await asyncio.to_thread(build_skill_archive, skill.file_content) if archive is None: verbose_proxy_logger.warning( "Agent Skills index: skipping skill %s, its upload is not a zip holding SKILL.md at the root of a " "single top-level folder", skill.skill_id, ) + return None + + if cache_key is not None and len(archive.content) <= MAX_CACHED_ARCHIVE_BYTES: + _ARCHIVE_CACHE.set_cache(cache_key, archive) return archive diff --git a/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py b/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py index 4fbb243a2b3..ea889e75ae8 100644 --- a/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py +++ b/tests/test_litellm/proxy/discovery_endpoints/test_agent_skills_endpoints.py @@ -1,6 +1,7 @@ import hashlib import io import zipfile +from datetime import datetime, timezone import pytest from fastapi import FastAPI @@ -42,12 +43,14 @@ def skill( display_title: str | None = "PDF Summarizer", description: str | None = None, files: dict[str, bytes] | None = None, + updated_at: datetime | None = None, ) -> LiteLLM_SkillsTable: return LiteLLM_SkillsTable( skill_id=skill_id, display_title=display_title, description=description, file_content=zip_bytes(files if files is not None else {"pdf-summarizer/SKILL.md": MANIFEST}), + updated_at=updated_at, ) @@ -162,3 +165,47 @@ def test_archive_route_404s_for_a_skill_that_does_not_exist(index_enabled): client = client_for(skill("litellm_skill_1")) assert client.get("/v1/skills/litellm_skill_missing/archive").status_code == 404 + + +def test_a_stored_skill_is_repacked_once_per_version(index_enabled): + stamp = datetime(2026, 9, 6, 9, 0, tzinfo=timezone.utc) + first = client_for(skill("litellm_skill_cached", files={"s/SKILL.md": MANIFEST}, updated_at=stamp)) + published = first.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["digest"] + + unchanged_row = client_for( + skill("litellm_skill_cached", files={"s/SKILL.md": MANIFEST, "s/extra.md": b"rewritten"}, updated_at=stamp) + ) + + assert unchanged_row.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["digest"] == published + assert hashlib.sha256(unchanged_row.get("/v1/skills/litellm_skill_cached/archive").content).hexdigest() == ( + published.removeprefix("sha256:") + ) + + +def test_a_skill_edited_since_the_last_read_is_republished(index_enabled): + stamp = datetime(2026, 9, 6, 9, 0, tzinfo=timezone.utc) + before = client_for(skill("litellm_skill_edited", files={"s/SKILL.md": MANIFEST}, updated_at=stamp)) + published = before.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["digest"] + + after = client_for( + skill( + "litellm_skill_edited", + files={"s/SKILL.md": MANIFEST, "s/extra.md": b"rewritten"}, + updated_at=datetime(2026, 9, 6, 10, 0, tzinfo=timezone.utc), + ) + ) + republished = after.get(WELL_KNOWN_PATHS[0]).json()["skills"][0]["digest"] + + assert republished != published + assert hashlib.sha256(after.get("/v1/skills/litellm_skill_edited/archive").content).hexdigest() == ( + republished.removeprefix("sha256:") + ) + + +def test_openapi_declares_the_archive_route_as_a_zip_download(index_enabled): + schema = client_for(skill("litellm_skill_1")).get("/openapi.json").json() + + content = schema["paths"]["/v1/skills/{skill_id}/archive"]["get"]["responses"]["200"]["content"] + + assert "application/zip" in content + assert "application/json" not in content diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 03496f9a507..137e3e31c08 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -64820,7 +64820,7 @@ export interface operations { [name: string]: unknown; }; content: { - "application/json": unknown; + "application/zip": string; }; }; /** @description Validation Error */ From eca59aa90b9ea6b696040857bbc3b495e717ca84 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 6 Sep 2026 03:35:07 -0700 Subject: [PATCH 14/55] fix(proxy): make an ambiguous config include loud, not silent An include entry that matches both a file next to the config that declares it and one next to the root config now warns naming both, so a config that resolves to a different file than it used to says so instead of quietly serving other models. Also from reviewing that change: - an empty root object in a bucket fails the boot again instead of coming up empty - a YAML syntax error in a bucket object logs its own line naming the object - an include already loaded is skipped before it is read rather than after - reading a config out of GCS builds the plain bucket client, so it needs no enterprise license and starts no flush loop that nothing ever cancels --- litellm/proxy/common_utils/config_includes.py | 59 ++++++++----- .../proxy/common_utils/load_config_utils.py | 46 ++++++---- litellm/proxy/proxy_server.py | 11 +-- .../common_utils/test_load_config_utils.py | 83 +++++++++++++++++++ .../proxy/proxy_server/test_proxy_config.py | 24 ++++++ 5 files changed, 183 insertions(+), 40 deletions(-) diff --git a/litellm/proxy/common_utils/config_includes.py b/litellm/proxy/common_utils/config_includes.py index 23e6c615a09..c1bb5ae952f 100644 --- a/litellm/proxy/common_utils/config_includes.py +++ b/litellm/proxy/common_utils/config_includes.py @@ -13,28 +13,41 @@ def resolve_include_file_path(include_file: str, declared_in: str, root_config_p Resolve one `include` entry to the file it names, next to the config that declares it. A config written before nested entries resolved this way can name a file sitting next to the root - config instead, so that file is still read, with a warning naming where it was found. + config instead, so that file is still read, with a warning naming where it was found. When both + files exist the one next to the declaring config wins and the other is named in a warning. """ declared_relative: Final = os.path.abspath(os.path.join(os.path.dirname(declared_in), include_file)) - if os.path.exists(declared_relative): - return declared_relative - root_relative: Final = os.path.abspath(os.path.join(os.path.dirname(root_config_path), include_file)) if root_relative == declared_relative or not os.path.exists(root_relative): return declared_relative + if not os.path.exists(declared_relative): + verbose_proxy_logger.warning( + "Config include '%s' declared in %s was not found next to it, so %s was read instead. " + "Move the included file next to the config that declares it.", + include_file, + declared_in, + root_relative, + ) + return root_relative + verbose_proxy_logger.warning( - "Config include '%s' declared in %s was not found next to it, so %s was read instead. " - "Move the included file next to the config that declares it.", + "Config include '%s' declared in %s matches two files. %s sits next to that config and was read, " + "so %s was skipped. Rename one of the two to say which one you meant.", include_file, declared_in, + declared_relative, root_relative, ) - return root_relative + return declared_relative -class ConfigLoader(Protocol): - def __call__(self, include_entry: str, declared_in: str, /) -> Awaitable[tuple[str, Mapping[str, object]]]: ... +class IncludeResolver(Protocol): + def __call__(self, include_entry: str, declared_in: str, /) -> str: ... + + +class ConfigReader(Protocol): + def __call__(self, location: str, /) -> Awaitable[Mapping[str, object]]: ... def _merged_value(base_value: object, included_value: object) -> object: @@ -80,32 +93,40 @@ async def _resolve( config: Mapping[str, object], pending: tuple[tuple[str, str], ...], loaded: frozenset[str], - load: ConfigLoader, + resolve: IncludeResolver, + read: ConfigReader, ) -> Mapping[str, object]: if not pending: return _without_include(config) entry, declared_in = pending[0] - location, included = await load(entry, declared_in) + location: Final = resolve(entry, declared_in) if location in loaded: - return await _resolve(config, pending[1:], loaded, load) + return await _resolve(config, pending[1:], loaded, resolve, read) + included: Final = await read(location) return await _resolve( _merged(config, _without_include(included)), (*pending[1:], *_pending_from(included, location)), loaded | frozenset((location,)), - load, + resolve, + read, ) -async def resolve_includes(config: Mapping[str, object], location: str, load: ConfigLoader) -> dict[str, object]: +async def resolve_includes( + config: Mapping[str, object], + location: str, + resolve: IncludeResolver, + read: ConfigReader, +) -> dict[str, object]: """ Merge every config named by the `include` directive into the config that declares it. - List values are extended and every other value is overridden, each entry is resolved relative to - the config that declares it, a config already pulled in is not merged a second time, and `load` - decides where an entry is read from, so the same merge applies to configs on disk and to configs - hosted in a bucket. + List values are extended and every other value is overridden, `resolve` turns each entry into the + location it names relative to the config that declares it, a config already pulled in is neither + read nor merged a second time, and `read` decides where a location is read from, so the same merge + applies to configs on disk and to configs hosted in a bucket. """ - merged: Final = await _resolve(config, _pending_from(config, location), frozenset((location,)), load) + merged: Final = await _resolve(config, _pending_from(config, location), frozenset((location,)), resolve, read) return dict(merged) # mutable-ok: the proxy mutates the config it loads diff --git a/litellm/proxy/common_utils/load_config_utils.py b/litellm/proxy/common_utils/load_config_utils.py index d62286a50f6..4a082eb307b 100644 --- a/litellm/proxy/common_utils/load_config_utils.py +++ b/litellm/proxy/common_utils/load_config_utils.py @@ -29,8 +29,12 @@ class SyncBucketObjectReader(Protocol): def __call__(self, object_key: str, /) -> object | None: ... -def _parsed_config(file_contents: str) -> object: - parsed: Final = yaml.safe_load(file_contents) +def _parsed_config(object_key: str, file_contents: str) -> object | None: + try: + parsed: Final = yaml.safe_load(file_contents) + except yaml.YAMLError as e: + verbose_proxy_logger.error("Config object %s is not valid YAML: %s", object_key, e) + return None return MappingProxyType({}) if parsed is None else parsed @@ -64,11 +68,13 @@ def s3_object_reader(bucket_name: str) -> SyncBucketObjectReader: try: verbose_proxy_logger.debug("Retrieving %s from S3 bucket: %s", object_key, bucket_name) response: Final = s3_client.get_object(Bucket=bucket_name, Key=object_key) - return _parsed_config(response["Body"].read().decode("utf-8")) + file_contents: Final = response["Body"].read().decode("utf-8") except Exception as e: # noqa: BLE001 # any boto3 error must read as a missing object verbose_proxy_logger.error("Error retrieving %s from S3 bucket %s: %s", object_key, bucket_name, e) return None + return _parsed_config(object_key, file_contents) + return read @@ -77,10 +83,16 @@ def get_file_contents_from_s3(bucket_name: str, object_key: str) -> object | Non def gcs_config_bucket(bucket_name: str) -> "GCSBucketBase | None": - try: - from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger + """ + Build a plain GCS client for reading config objects. - return GCSBucketLogger(bucket_name=bucket_name) + Reading a config out of a bucket is not GCS logging, so it neither needs the enterprise license + that gate covers nor the batching task the logger starts and never stops. + """ + try: + from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase + + return GCSBucketBase(bucket_name=bucket_name) except Exception as e: # noqa: BLE001 # an unbuildable client must read as an unreadable bucket verbose_proxy_logger.error("Error creating the GCS client for bucket %s: %s", bucket_name, e) return None @@ -98,12 +110,14 @@ async def get_config_file_contents_from_gcs( file_contents: Final = await bucket.download_gcs_object(object_key) if file_contents is None: raise Exception(f"File contents are None for {object_key}") - return _parsed_config(file_contents.decode("utf-8")) + decoded: Final = file_contents.decode("utf-8") except Exception as e: verbose_proxy_logger.error("Error retrieving %s from GCS bucket %s: %s", object_key, bucket_name, e) return None + return _parsed_config(object_key, decoded) + def resolve_include_object_key(config_object_key: str, include_entry: str) -> str: """ @@ -123,17 +137,19 @@ async def resolve_bucket_includes( object_key: str, fetch: BucketObjectFetcher, ) -> dict[str, object]: - async def load(include_entry: str, declared_in: str) -> tuple[str, Mapping[str, object]]: - include_key: Final = resolve_include_object_key(declared_in, include_entry) + async def read(include_key: str) -> Mapping[str, object]: included: Final = await fetch(include_key) if included is None: raise FileNotFoundError( f"Included config could not be read from bucket: {include_key}. " "The underlying bucket error is logged above." ) - return include_key, included + return included - return await resolve_includes(config=config, location=object_key, load=load) + def resolve(include_entry: str, declared_in: str) -> str: + return resolve_include_object_key(declared_in, include_entry) + + return await resolve_includes(config=config, location=object_key, resolve=resolve, read=read) async def bucket_object_reader(bucket_type: str | None, bucket_name: str) -> BucketObjectReader: @@ -176,7 +192,7 @@ async def get_config_from_bucket( raise ValueError(f"Config object in bucket is not a YAML mapping: {key}") from e config: Final = await fetch(object_key) - if config is None: + if not config: return None return await resolve_bucket_includes(config=config, object_key=object_key, fetch=fetch) @@ -256,11 +272,9 @@ async def download_python_file_from_gcs( bool: True if successful, False otherwise """ try: - from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger + from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase - gcs_bucket: Final = GCSBucketLogger( - bucket_name=bucket_name, - ) + gcs_bucket: Final = GCSBucketBase(bucket_name=bucket_name) file_contents = await gcs_bucket.download_gcs_object(object_key) if file_contents is None: raise Exception(f"File contents are None for {object_key}") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7fada09febf..ae9cd71a956 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -4559,17 +4559,18 @@ class ProxyConfig: included_config_adapter: Final = TypeAdapter(dict[str, object]) - async def load_included(include_file: str, declared_in: str) -> tuple[str, Mapping[str, object]]: - file_path: Final = resolve_include_file_path(include_file, declared_in, config_file_path) + def resolve(include_file: str, declared_in: str) -> str: + return resolve_include_file_path(include_file, declared_in, config_file_path) + + async def read_included(file_path: str) -> Mapping[str, object]: if not os.path.exists(file_path): raise FileNotFoundError(f"Included file not found: {file_path}") try: - included: Final = included_config_adapter.validate_python(self._load_yaml_file(file_path)) + return included_config_adapter.validate_python(self._load_yaml_file(file_path)) except ValidationError as e: raise ValueError(f"Included config file is not a YAML mapping: {file_path}") from e - return file_path, included - return await resolve_includes(config=config, location=config_file_path, load=load_included) + return await resolve_includes(config=config, location=config_file_path, resolve=resolve, read=read_included) async def save_config(self, new_config: dict, include_env_vars: bool = False): global prisma_client, general_settings, user_config_file_path, store_model_in_db diff --git a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py index 75aa5a3f0f8..1042dbe9653 100644 --- a/tests/test_litellm/proxy/common_utils/test_load_config_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_load_config_utils.py @@ -1,4 +1,5 @@ import asyncio +import logging import re import threading from unittest.mock import MagicMock, mock_open, patch @@ -7,6 +8,7 @@ import pytest import yaml from litellm.proxy.common_utils.load_config_utils import ( + gcs_config_bucket, get_config_from_bucket, get_file_contents_from_s3, resolve_bucket_includes, @@ -391,3 +393,84 @@ class TestBucketConfigIncludes: ) is None ) + + @pytest.mark.asyncio + async def test_an_object_pulled_in_twice_is_read_once(self): + objects = { + "configs/a.yaml": {"include": ["shared.yaml"]}, + "configs/b.yaml": {"include": ["./shared.yaml"]}, + "configs/shared.yaml": {"model_list": [{"model_name": "shared"}]}, + } + requested = [] + + async def fetch(object_key): + requested.append(object_key) + return objects.get(object_key) + + await resolve_bucket_includes( + config={"include": ["a.yaml", "b.yaml"]}, + object_key="configs/config.yaml", + fetch=fetch, + ) + + assert requested == ["configs/a.yaml", "configs/b.yaml", "configs/shared.yaml"] + + @pytest.mark.asyncio + async def test_an_empty_root_object_does_not_boot_an_empty_proxy(self, monkeypatch): + class FakeGCSBucket: + async def download_gcs_object(self, object_key): + return b"" + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), + ) + + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config is None + + @pytest.mark.asyncio + async def test_an_object_that_is_not_valid_yaml_is_reported_as_a_yaml_error(self, monkeypatch, caplog): + class FakeGCSBucket: + async def download_gcs_object(self, object_key): + return b"model_list: [\n" + + monkeypatch.setattr( + "litellm.proxy.common_utils.load_config_utils.gcs_config_bucket", + lambda bucket_name: FakeGCSBucket(), + ) + + with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"): + config = await get_config_from_bucket( + bucket_type="gcs", bucket_name="litellm-configs", object_key="lit6982/config.yaml" + ) + + assert config is None + assert [ + record + for record in caplog.records + if "not valid YAML" in record.getMessage() and "lit6982/config.yaml" in record.getMessage() + ] + + +class TestGCSConfigBucketClient: + @pytest.mark.asyncio + async def test_reading_a_config_from_gcs_does_not_need_an_enterprise_license(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", False) + + bucket = gcs_config_bucket("litellm-configs") + + assert bucket is not None + assert bucket.BUCKET_NAME == "litellm-configs" + + @pytest.mark.asyncio + async def test_reading_a_config_from_gcs_starts_no_background_task(self, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + running_before = asyncio.all_tasks() + + gcs_config_bucket("litellm-configs") + + assert asyncio.all_tasks() - running_before == set() diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index e90f101a555..fe64bcc2c10 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -10,6 +10,7 @@ from __future__ import annotations import asyncio import json +import logging import os import re from types import SimpleNamespace @@ -772,6 +773,29 @@ async def test_ProxyConfig__process_includes_still_reads_a_nested_include_left_b assert result == {"model_list": [{"model_name": "first"}, {"model_name": "second"}]} +@pytest.mark.asyncio +async def test_ProxyConfig__process_includes_names_both_files_when_a_nested_include_matches_two(tmp_path, caplog): + (tmp_path / "shared").mkdir() + (tmp_path / "shared" / "models.yaml").write_text( + "include:\n - more_models.yaml\nmodel_list:\n - model_name: first\n" + ) + (tmp_path / "shared" / "more_models.yaml").write_text("model_list:\n - model_name: next-to-the-declaring-file\n") + (tmp_path / "more_models.yaml").write_text("model_list:\n - model_name: next-to-the-root-config\n") + + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + result = await ProxyConfig()._process_includes( + {"include": ["shared/models.yaml"]}, config_file_path=str(tmp_path / "config.yaml") + ) + + assert result == {"model_list": [{"model_name": "first"}, {"model_name": "next-to-the-declaring-file"}]} + assert [ + record + for record in caplog.records + if str(tmp_path / "shared" / "more_models.yaml") in record.getMessage() + and str(tmp_path / "more_models.yaml") in record.getMessage() + ] + + @pytest.mark.asyncio async def test_ProxyConfig__process_includes_merges_a_shared_file_once(tmp_path): (tmp_path / "shared.yaml").write_text("model_list:\n - model_name: shared\n") From ae2f1ae512c1ce5b5f029d8c13d8034e2a2d1525 Mon Sep 17 00:00:00 2001 From: mateo Date: Mon, 7 Sep 2026 09:19:02 +0000 Subject: [PATCH 15/55] chore: drop the budget ratchet from the PR branch, the scheduled automation owns it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- test-quality-budget.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test-quality-budget.json b/test-quality-budget.json index f382f2479a9..3c12371f02f 100644 --- a/test-quality-budget.json +++ b/test-quality-budget.json @@ -3,7 +3,7 @@ "limit": 733 }, "TQ002": { - "limit": 736 + "limit": 737 }, "TQ003": { "limit": 62 From 61e664088c9ad57ec99e4433c380ec00a7b8aa95 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 8 Sep 2026 09:51:09 +0000 Subject: [PATCH 16/55] test: make the explicit stagger offset assertion independent of the wall clock Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../common_utils/test_scheduled_job_stagger.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py index ca4d62737b6..10c61a5dd87 100644 --- a/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py +++ b/tests/test_litellm/proxy/common_utils/test_scheduled_job_stagger.py @@ -64,6 +64,10 @@ def _stagger(scheduler: AsyncIOScheduler, identity: str = "pod-a:1", **overrides return apply_scheduled_job_stagger(scheduler=scheduler, settings=_settings(**overrides), identity=identity) +def _trigger_of(scheduler: AsyncIOScheduler, job_id: str): + return next(job.trigger for job in scheduler.get_jobs() if job.id == job_id) + + def _fire_times(trigger, start: datetime, steps: int) -> tuple[datetime, ...]: """The fire times APScheduler would produce, each computed from the one before it""" return tuple( @@ -133,22 +137,24 @@ def test_default_cron_is_staggered_and_keeps_its_offset_on_every_later_fire(): applied = _stagger(scheduler) assert applied[PTU_ROLLUP_JOB_ID] > 0 - trigger = next(job.trigger for job in scheduler.get_jobs() if job.id == PTU_ROLLUP_JOB_ID) - fires = _fire_times(trigger, datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc), 3) + start = datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc) + fires = _fire_times(_trigger_of(scheduler, PTU_ROLLUP_JOB_ID), start, 3) expected = timedelta(minutes=15) + timedelta(seconds=applied[PTU_ROLLUP_JOB_ID]) assert [fire - fire.replace(hour=0, minute=0, second=0, microsecond=0) for fire in fires] == [expected] * 3 -async def test_explicit_offset_overrides_the_derived_one_and_zero_pins_a_job(): +def test_explicit_offset_overrides_the_derived_one_and_zero_pins_a_job(): scheduler = _with_jobs(_scheduler()) applied = _stagger(scheduler, offsets={"periodic_reload_job": 0, PTU_ROLLUP_JOB_ID: 7}) - unstaggered = _next_run_times(_with_jobs(_scheduler())) - staggered = _next_run_times(scheduler) assert applied["periodic_reload_job"] == 0 assert applied[PTU_ROLLUP_JOB_ID] == 7 - assert staggered[PTU_ROLLUP_JOB_ID] - unstaggered[PTU_ROLLUP_JOB_ID] == timedelta(seconds=7) + + start = datetime(2026, 1, 1, 12, 0, tzinfo=timezone.utc) + staggered = _trigger_of(scheduler, PTU_ROLLUP_JOB_ID) + unstaggered = _trigger_of(_with_jobs(_scheduler()), PTU_ROLLUP_JOB_ID) + assert _fire_times(staggered, start, 1)[0] - _fire_times(unstaggered, start, 1)[0] == timedelta(seconds=7) async def test_disabling_the_stagger_leaves_every_schedule_untouched(): From 39239ab974374b52d37dafb944852c6ce62369f8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 18:36:05 -0700 Subject: [PATCH 17/55] refactor(cost): share the deployment model_info lookup between video and OCR cost paths Extract one immutable helper for reading the deployment's model_info off the logging object, stop rebinding the model_info parameter inside ocr_cost, drop the explanatory comment blocks, and move the OCR custom pricing regression tests into tests/test_litellm/test_cost_calculator.py --- litellm/cost_calculator.py | 94 +++++------- tests/test_litellm/test_cost_calculator.py | 130 ++++++++++++++++ tests/test_litellm/test_ocr_custom_pricing.py | 143 ------------------ 3 files changed, 166 insertions(+), 201 deletions(-) delete mode 100644 tests/test_litellm/test_ocr_custom_pricing.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 496f060d5ca..9aed82f37e3 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -341,7 +341,7 @@ def cost_per_token( ### REQUEST MODEL ### request_model: str | None = None, # original request model for router detection ### DEPLOYMENT-SPECIFIC PRICING ### - custom_model_info: ModelInfo | None = None, # deployment model_info, for non-token custom pricing + custom_model_info: ModelInfo | None = None, ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -1429,20 +1429,7 @@ def completion_cost( ) elif call_type in _VIDEO_CALL_TYPES: ### VIDEO GENERATION COST CALCULATION ### - # Extract custom model_info for deployment-specific pricing - _video_model_info: ModelInfo | None = None - if custom_pricing and litellm_logging_obj is not None: - _litellm_params = getattr(litellm_logging_obj, "litellm_params", None) - if _litellm_params is not None: - _video_model_info = next( - ( - model_info - for _metadata_key in ("metadata", "litellm_metadata") - if (model_info := (_litellm_params.get(_metadata_key) or {}).get("model_info")) - is not None - ), - None, - ) + _video_model_info: ModelInfo | None = _deployment_model_info(litellm_logging_obj, custom_pricing) usage_obj = getattr(completion_response, "usage", None) duration_seconds: float | None = None @@ -1638,24 +1625,6 @@ def completion_cost( if litellm_logging_obj is not None: request_model_for_cost = litellm_logging_obj.model - # Deployment-specific model_info, for modalities whose pricing is - # not token-based and so cannot travel via custom_cost_per_token - # (e.g. OCR per-page pricing). Same extraction as the video path - # above, minus its `or {}` default: truthiness on the value adds - # no mutable-collection construction (LIT002) and reads the same. - # Checked under both keys: router calls that go through - # `_ageneric_api_call_with_fallbacks` (OCR included) store the - # deployment's model_info under `litellm_metadata`, not `metadata`. - _custom_model_info: ModelInfo | None = None - if custom_pricing and litellm_logging_obj is not None: - _cm_litellm_params = getattr(litellm_logging_obj, "litellm_params", None) - if _cm_litellm_params is not None: - for _cm_metadata_key in ("metadata", "litellm_metadata"): - _cm_metadata = _cm_litellm_params.get(_cm_metadata_key) - if _cm_metadata and _cm_metadata.get("model_info") is not None: - _custom_model_info = _cm_metadata.get("model_info") - break - ( prompt_tokens_cost_usd_dollar, completion_tokens_cost_usd_dollar, @@ -1681,7 +1650,7 @@ def completion_cost( vertex_location=vertex_location, response=completion_response, request_model=request_model_for_cost, - custom_model_info=_custom_model_info, + custom_model_info=_deployment_model_info(litellm_logging_obj, custom_pricing), ) # Get additional costs from provider (e.g., routing fees, infrastructure costs) @@ -1911,6 +1880,32 @@ def response_cost_calculator( raise e +def _deployment_model_info( + litellm_logging_obj: LitellmLoggingObject | None, + custom_pricing: bool | None, +) -> ModelInfo | None: + if not custom_pricing or litellm_logging_obj is None: + return None + litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) + if litellm_params is None: + return None + return next( + ( + model_info + for metadata_key in ("metadata", "litellm_metadata") + if (metadata := litellm_params.get(metadata_key)) and (model_info := metadata.get("model_info")) is not None + ), + None, + ) + + +def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelInfo | None: + try: + return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) + except Exception: + return None + + def ocr_cost( model: str, custom_llm_provider: str | None, @@ -1922,9 +1917,8 @@ def ocr_cost( model: str - model name custom_llm_provider: Optional[str] - custom LLM provider response: Optional[Any] - response object - model_info: Optional[ModelInfo] - deployment-specific model info, used for - custom pricing. Takes precedence over the model cost map, mirroring - the video generation cost path. + model_info: Optional[ModelInfo] - deployment-specific model info; its OCR pricing + takes precedence over the model cost map Returns: Tuple[float, float]: cost of OCR processing @@ -1942,34 +1936,18 @@ def ocr_cost( if response.usage_info is None: raise ValueError("OCR response usage_info is None") - ######################################################### - # Deployment-specific pricing wins over the cost map. - # - # Custom pricing set on a deployment is registered under the router's - # deployment id, while the shared "{provider}/{model}" key has its pricing - # fields stripped (see _register_custom_pricing_for_request). A cost map - # lookup therefore cannot see it, so an OCR model that is not in the map - # bills $0 no matter how it is priced in config. Prefer the caller-supplied - # model_info when it carries OCR pricing. - ######################################################### - has_custom_ocr_pricing: Final[bool] = model_info is not None and ( + has_custom_ocr_pricing: Final = model_info is not None and ( model_info.get("ocr_cost_per_page") is not None or model_info.get("ocr_cost_per_credit") is not None ) - if not has_custom_ocr_pricing: - try: - model_info = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) - except Exception: - model_info = None + pricing: Final = model_info if has_custom_ocr_pricing else _cost_map_model_info(model, custom_llm_provider) credits: Final = getattr(response.usage_info, "credits", None) - cost_per_credit = None - if model_info is not None: - cost_per_credit = model_info.get("ocr_cost_per_credit") + cost_per_credit: Final = pricing.get("ocr_cost_per_credit") if pricing is not None else None if credits is not None and cost_per_credit is not None: return cost_per_credit * credits, 0.0 - ocr_cost_per_page: Final = model_info.get("ocr_cost_per_page") if model_info is not None else None - annotation_cost_per_page: Final = model_info.get("annotation_cost_per_page") if model_info is not None else None + ocr_cost_per_page: Final = pricing.get("ocr_cost_per_page") if pricing is not None else None + annotation_cost_per_page: Final = pricing.get("annotation_cost_per_page") if pricing is not None else None annotation_rate: Final = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page pages_processed: Final = response.usage_info.pages_processed diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index f8fa2231597..824108ab0ec 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4645,3 +4645,133 @@ def test_collect_and_combine_realtime_usage_stores_partitioned_text_tokens() -> assert combined.completion_tokens_details.reasoning_tokens == 95 assert combined.completion_tokens_details.text_tokens == 38 assert combined.completion_tokens_details.audio_tokens == 0 + + +UNMAPPED_OCR_MODEL: Final = "azure_ai/some-unmapped-ocr-model-for-testing" +MAPPED_OCR_MODEL: Final = "mistral/mistral-ocr-4-0" + + +def _ocr_response(model: str, pages_processed: int, credits: float | None = None): + from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo + + return OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(pages_processed)], + model=model, + usage_info=OCRUsageInfo(pages_processed=pages_processed, credits=credits), + ) + + +def _ocr_logging_obj(litellm_params: dict): + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model=UNMAPPED_OCR_MODEL, + messages=[], + stream=False, + call_type="ocr", + start_time=None, + litellm_call_id="test-ocr-custom-pricing", + function_id="1234", + ) + logging_obj.litellm_params = litellm_params + return logging_obj + + +@pytest.mark.parametrize("pages_processed", [1, 3, 10]) +def test_ocr_cost_uses_deployment_per_page_pricing_for_unmapped_model(pages_processed: int): + from litellm.cost_calculator import ocr_cost + + assert UNMAPPED_OCR_MODEL not in litellm.model_cost + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=pages_processed), + model_info={"ocr_cost_per_page": 0.004}, + ) + assert cost == pytest.approx(0.004 * pages_processed) + + +def test_ocr_cost_uses_deployment_per_credit_pricing_for_unmapped_model(): + from litellm.cost_calculator import ocr_cost + + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=2, credits=4), + model_info={"ocr_cost_per_credit": 0.25}, + ) + assert cost == pytest.approx(0.25 * 4) + + +def test_ocr_cost_unmapped_model_without_deployment_pricing_bills_zero(): + from litellm.cost_calculator import ocr_cost + + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=5), + model_info={"id": "some-deployment-id"}, + ) + assert cost == 0.0 + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_ocr_cost_deployment_pricing_overrides_cost_map_for_mapped_model(): + from litellm.cost_calculator import ocr_cost + + map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"] + assert map_price is not None + override_price: Final = map_price * 10 + + cost, _ = ocr_cost( + model=MAPPED_OCR_MODEL, + custom_llm_provider="mistral", + response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2), + model_info={"ocr_cost_per_page": override_price}, + ) + assert cost == pytest.approx(override_price * 2) + + +@pytest.mark.usefixtures("_local_model_cost_map") +def test_ocr_cost_falls_through_to_cost_map_when_deployment_has_no_ocr_pricing(): + from litellm.cost_calculator import ocr_cost + + map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"] + assert map_price is not None + + cost, _ = ocr_cost( + model=MAPPED_OCR_MODEL, + custom_llm_provider="mistral", + response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2), + model_info={"id": "some-deployment-id"}, + ) + assert cost == pytest.approx(map_price * 2) + + +@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) +def test_completion_cost_ocr_reads_deployment_pricing_from_logging_metadata(metadata_key: str): + logging_obj = _ocr_logging_obj({metadata_key: {"model_info": {"ocr_cost_per_page": 0.004}}}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(0.004 * 3) + + +def test_completion_cost_ocr_ignores_deployment_pricing_without_custom_pricing_flag(): + logging_obj = _ocr_logging_obj({"metadata": {"model_info": {"ocr_cost_per_page": 0.004}}}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=False, + litellm_logging_obj=logging_obj, + ) + assert cost == 0.0 diff --git a/tests/test_litellm/test_ocr_custom_pricing.py b/tests/test_litellm/test_ocr_custom_pricing.py deleted file mode 100644 index 4dffc43c23f..00000000000 --- a/tests/test_litellm/test_ocr_custom_pricing.py +++ /dev/null @@ -1,143 +0,0 @@ -""" -Regression tests: OCR cost must honour deployment-specific custom pricing. - -Before the fix, `ocr_cost()` resolved pricing exclusively through -`litellm.get_model_info(model=..., custom_llm_provider=...)`, i.e. a cost map -lookup keyed by model name. Custom pricing set on a deployment is registered -under the router's deployment id, and the shared "{provider}/{model}" key has -its pricing fields stripped, so the lookup could never see it. An OCR model -absent from the cost map therefore billed $0 no matter how it was priced in -config, even though `ocr_cost_per_page` / `ocr_cost_per_credit` are declared -fields of `CustomPricingLiteLLMParams`. -""" - -import pytest - -import litellm -from litellm.cost_calculator import completion_cost, ocr_cost -from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo - -# A model deliberately absent from the cost map. -UNMAPPED_MODEL = "azure_ai/some-unmapped-ocr-model-for-testing" -CUSTOM_COST_PER_PAGE = 0.004 -CUSTOM_COST_PER_CREDIT = 0.25 - - -def _ocr_response(model: str, pages_processed: int = 1, credits: int | None = None) -> OCRResponse: - # NOTE: model_construct() is used rather than OCRResponse(...) because the - # OCRResponse field `object: str = "ocr"` shadows the builtin `object` used - # in the `tables` / `keyValuePairs` annotations above it, so pydantic tries - # to resolve "ocr" as a forward-referenced type and schema building fails. - # That is an unrelated defect; validation is not what these tests exercise. - usage_info = OCRUsageInfo(pages_processed=pages_processed) - if credits is not None: - usage_info.credits = credits - return OCRResponse.model_construct( - pages=[OCRPage(index=i, markdown=f"page {i}") for i in range(pages_processed)], - model=model, - usage_info=usage_info, - ) - - -def test_unmapped_ocr_model_has_no_map_pricing() -> None: - """Guard the premise: the model really is absent from the cost map.""" - assert UNMAPPED_MODEL not in litellm.model_cost - - -@pytest.mark.parametrize("pages_processed", [1, 3, 10]) -def test_ocr_cost_uses_custom_per_page_pricing(pages_processed: int) -> None: - cost, _ = ocr_cost( - model=UNMAPPED_MODEL, - custom_llm_provider="azure_ai", - response=_ocr_response(UNMAPPED_MODEL, pages_processed=pages_processed), - model_info={"ocr_cost_per_page": CUSTOM_COST_PER_PAGE}, - ) - assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * pages_processed) - - -def test_ocr_cost_uses_custom_per_credit_pricing() -> None: - cost, _ = ocr_cost( - model=UNMAPPED_MODEL, - custom_llm_provider="azure_ai", - response=_ocr_response(UNMAPPED_MODEL, pages_processed=2, credits=4), - model_info={"ocr_cost_per_credit": CUSTOM_COST_PER_CREDIT}, - ) - assert cost == pytest.approx(CUSTOM_COST_PER_CREDIT * 4) - - -def test_unmapped_ocr_model_without_custom_pricing_still_bills_zero() -> None: - """Unchanged behaviour when nothing is configured — no map entry, no override.""" - cost, _ = ocr_cost( - model=UNMAPPED_MODEL, - custom_llm_provider="azure_ai", - response=_ocr_response(UNMAPPED_MODEL, pages_processed=5), - ) - assert cost == 0.0 - - -def test_custom_pricing_does_not_override_a_mapped_model_when_absent() -> None: - """model_info without OCR pricing must fall through to the cost map.""" - mapped_model = "mistral/mistral-ocr-4-0" - cost, _ = ocr_cost( - model=mapped_model, - custom_llm_provider="mistral", - response=_ocr_response(mapped_model, pages_processed=2), - model_info={"id": "some-deployment-id"}, - ) - assert cost == pytest.approx(0.004 * 2) - - -def test_ocr_custom_pricing_end_to_end_through_completion_cost() -> None: - """The whole path: litellm_params.metadata.model_info -> ocr_cost.""" - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - - logging_obj = LiteLLMLogging( - model=UNMAPPED_MODEL, - messages=[], - stream=False, - call_type="ocr", - start_time=None, - litellm_call_id="test-ocr-custom-pricing", - function_id="1234", - ) - logging_obj.litellm_params = {"metadata": {"model_info": {"ocr_cost_per_page": CUSTOM_COST_PER_PAGE}}} - - cost = completion_cost( - completion_response=_ocr_response(UNMAPPED_MODEL, pages_processed=3), - model=UNMAPPED_MODEL, - custom_llm_provider="azure_ai", - call_type="ocr", - custom_pricing=True, - litellm_logging_obj=logging_obj, - ) - assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * 3) - - -def test_ocr_custom_pricing_end_to_end_via_litellm_metadata() -> None: - """Router OCR calls go through `_ageneric_api_call_with_fallbacks`, which - stores the deployment's model_info under `litellm_metadata` rather than - `metadata`. The extraction must read that key too.""" - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging - - logging_obj = LiteLLMLogging( - model=UNMAPPED_MODEL, - messages=[], - stream=False, - call_type="ocr", - start_time=None, - litellm_call_id="test-ocr-custom-pricing-litellm-metadata", - function_id="1234", - ) - logging_obj.litellm_params = { - "litellm_metadata": {"model_info": {"ocr_cost_per_page": CUSTOM_COST_PER_PAGE}}, - } - - cost = completion_cost( - completion_response=_ocr_response(UNMAPPED_MODEL, pages_processed=3), - model=UNMAPPED_MODEL, - custom_llm_provider="azure_ai", - call_type="ocr", - custom_pricing=True, - litellm_logging_obj=logging_obj, - ) - assert cost == pytest.approx(CUSTOM_COST_PER_PAGE * 3) From 9a78bf638aa4688b790094241b2d3484b587d2f0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 18:55:13 -0700 Subject: [PATCH 18/55] test(cost): type the OCR pricing test helpers --- tests/test_litellm/test_cost_calculator.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 824108ab0ec..a602e6ee819 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -17,6 +17,8 @@ from litellm.cost_calculator import ( handle_realtime_stream_cost_calculation, response_cost_calculator, ) +from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo from litellm.types.llms.openai import OpenAIRealtimeStreamList from litellm.types.utils import ( CacheCreationTokenDetails, @@ -4651,9 +4653,7 @@ UNMAPPED_OCR_MODEL: Final = "azure_ai/some-unmapped-ocr-model-for-testing" MAPPED_OCR_MODEL: Final = "mistral/mistral-ocr-4-0" -def _ocr_response(model: str, pages_processed: int, credits: float | None = None): - from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo - +def _ocr_response(model: str, pages_processed: int, credits: float | None = None) -> OCRResponse: return OCRResponse( pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(pages_processed)], model=model, @@ -4661,9 +4661,7 @@ def _ocr_response(model: str, pages_processed: int, credits: float | None = None ) -def _ocr_logging_obj(litellm_params: dict): - from litellm.litellm_core_utils.litellm_logging import Logging - +def _ocr_logging_obj(litellm_params: dict[str, dict[str, ModelInfo]]) -> Logging: logging_obj = Logging( model=UNMAPPED_OCR_MODEL, messages=[], From a29103cbbffdaaa0044e084b5a1c264bcf67361b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:43:31 -0700 Subject: [PATCH 19/55] fix(cost): read OCR pricing registered under the router deployment id --- litellm/cost_calculator.py | 18 +++++++++++++++--- tests/test_litellm/test_cost_calculator.py | 19 +++++++++++++++++++ 2 files changed, 34 insertions(+), 3 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 9aed82f37e3..afafc12623a 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1429,7 +1429,9 @@ def completion_cost( ) elif call_type in _VIDEO_CALL_TYPES: ### VIDEO GENERATION COST CALCULATION ### - _video_model_info: ModelInfo | None = _deployment_model_info(litellm_logging_obj, custom_pricing) + _video_model_info: ModelInfo | None = _deployment_model_info( + litellm_logging_obj, custom_pricing, router_model_id + ) usage_obj = getattr(completion_response, "usage", None) duration_seconds: float | None = None @@ -1650,7 +1652,7 @@ def completion_cost( vertex_location=vertex_location, response=completion_response, request_model=request_model_for_cost, - custom_model_info=_deployment_model_info(litellm_logging_obj, custom_pricing), + custom_model_info=_deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id), ) # Get additional costs from provider (e.g., routing fees, infrastructure costs) @@ -1883,8 +1885,18 @@ def response_cost_calculator( def _deployment_model_info( litellm_logging_obj: LitellmLoggingObject | None, custom_pricing: bool | None, + router_model_id: str | None, ) -> ModelInfo | None: - if not custom_pricing or litellm_logging_obj is None: + if not custom_pricing: + return None + registered_deployment_info: Final = ( + _cost_map_model_info(router_model_id, None) + if router_model_id is not None and router_model_id in litellm.model_cost + else None + ) + if registered_deployment_info is not None: + return registered_deployment_info + if litellm_logging_obj is None: return None litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) if litellm_params is None: diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index a602e6ee819..04a547c222c 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4761,6 +4761,25 @@ def test_completion_cost_ocr_reads_deployment_pricing_from_logging_metadata(meta assert cost == pytest.approx(0.004 * 3) +def test_completion_cost_ocr_prefers_pricing_registered_under_router_model_id(monkeypatch: pytest.MonkeyPatch): + deployment_id: Final = "ocr-deployment-priced-through-litellm-params" + monkeypatch.setitem( + litellm.model_cost, deployment_id, {"mode": "ocr", "litellm_provider": "azure_ai", "ocr_cost_per_page": 0.05} + ) + logging_obj = _ocr_logging_obj({"metadata": {"model_info": {"mode": "ocr"}}}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + router_model_id=deployment_id, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(0.05 * 3) + + def test_completion_cost_ocr_ignores_deployment_pricing_without_custom_pricing_flag(): logging_obj = _ocr_logging_obj({"metadata": {"model_info": {"ocr_cost_per_page": 0.004}}}) From 2c7751219d1a0e260d3088dc83cbb0df1934ec5c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 8 Sep 2026 20:04:23 -0700 Subject: [PATCH 20/55] fix(ocr): bill request-level OCR pricing and fall back to the map without credits --- litellm/cost_calculator.py | 5 ++- litellm/ocr/main.py | 2 + tests/test_litellm/ocr/test_main.py | 49 ++++++++++++++++++++++ tests/test_litellm/test_cost_calculator.py | 16 +++++++ 4 files changed, 70 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/ocr/test_main.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index afafc12623a..68f52e4a10e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1948,12 +1948,13 @@ def ocr_cost( if response.usage_info is None: raise ValueError("OCR response usage_info is None") + credits: Final = getattr(response.usage_info, "credits", None) has_custom_ocr_pricing: Final = model_info is not None and ( - model_info.get("ocr_cost_per_page") is not None or model_info.get("ocr_cost_per_credit") is not None + model_info.get("ocr_cost_per_page") is not None + or (credits is not None and model_info.get("ocr_cost_per_credit") is not None) ) pricing: Final = model_info if has_custom_ocr_pricing else _cost_map_model_info(model, custom_llm_provider) - credits: Final = getattr(response.usage_info, "credits", None) cost_per_credit: Final = pricing.get("ocr_cost_per_credit") if pricing is not None else None if credits is not None and cost_per_credit is not None: return cost_per_credit * credits, 0.0 diff --git a/litellm/ocr/main.py b/litellm/ocr/main.py index df3f9d2096b..8b0e6950801 100644 --- a/litellm/ocr/main.py +++ b/litellm/ocr/main.py @@ -32,6 +32,7 @@ from litellm.rust_bridge import ocr as rust_ocr_bridge from litellm.rust_bridge.bindings import native_exception_types from litellm.rust_bridge.configuration import rust_enabled from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import CustomPricingLiteLLMParams from litellm.utils import ProviderConfigManager, client ####### ENVIRONMENT VARIABLES ################### @@ -171,6 +172,7 @@ def _prepare_ocr_request( litellm_params={ "litellm_call_id": litellm_call_id, "api_base": api_base, + **litellm_params.model_dump(include=frozenset(CustomPricingLiteLLMParams.model_fields), exclude_none=True), }, custom_llm_provider=custom_llm_provider, ) diff --git a/tests/test_litellm/ocr/test_main.py b/tests/test_litellm/ocr/test_main.py new file mode 100644 index 00000000000..0007d98dcb0 --- /dev/null +++ b/tests/test_litellm/ocr/test_main.py @@ -0,0 +1,49 @@ +from typing import Final + +from litellm.litellm_core_utils.litellm_logging import Logging, use_custom_pricing_for_model +from litellm.ocr.main import _prepare_ocr_request + +OCR_MODEL: Final = "mistral/mistral-ocr-4-1" +DOCUMENT: Final = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} + + +def _logging_obj() -> Logging: + return Logging( + model=OCR_MODEL, + messages=[], + stream=False, + call_type="ocr", + start_time=None, + litellm_call_id="test-ocr-request-pricing", + function_id="1234", + ) + + +def _prepare(kwargs: dict[str, object]) -> Logging: + logging_obj: Final = _logging_obj() + _prepare_ocr_request( + model=OCR_MODEL, + document=dict(DOCUMENT), + api_key="test-key", + api_base=None, + timeout=None, + custom_llm_provider=None, + extra_headers=None, + kwargs={"litellm_logging_obj": logging_obj, **kwargs}, + ) + return logging_obj + + +def test_prepare_ocr_request_forwards_custom_pricing_to_logging_params() -> None: + logging_obj: Final = _prepare({"ocr_cost_per_page": 0.05, "ocr_cost_per_credit": 0.5}) + + assert logging_obj.litellm_params["ocr_cost_per_page"] == 0.05 + assert logging_obj.litellm_params["ocr_cost_per_credit"] == 0.5 + assert use_custom_pricing_for_model(logging_obj.litellm_params) is True + + +def test_prepare_ocr_request_without_custom_pricing_leaves_logging_params_unpriced() -> None: + logging_obj: Final = _prepare({}) + + assert "ocr_cost_per_page" not in logging_obj.litellm_params + assert use_custom_pricing_for_model(logging_obj.litellm_params) is False diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 04a547c222c..cb9217b2f41 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4746,6 +4746,22 @@ def test_ocr_cost_falls_through_to_cost_map_when_deployment_has_no_ocr_pricing() assert cost == pytest.approx(map_price * 2) +@pytest.mark.usefixtures("_local_model_cost_map") +def test_ocr_cost_ignores_deployment_credit_pricing_when_response_reports_no_credits(): + from litellm.cost_calculator import ocr_cost + + map_price: Final = litellm.get_model_info(MAPPED_OCR_MODEL)["ocr_cost_per_page"] + assert map_price is not None + + cost, _ = ocr_cost( + model=MAPPED_OCR_MODEL, + custom_llm_provider="mistral", + response=_ocr_response(MAPPED_OCR_MODEL, pages_processed=2), + model_info={"ocr_cost_per_credit": 0.5}, + ) + assert cost == pytest.approx(map_price * 2) + + @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) def test_completion_cost_ocr_reads_deployment_pricing_from_logging_metadata(metadata_key: str): logging_obj = _ocr_logging_obj({metadata_key: {"model_info": {"ocr_cost_per_page": 0.004}}}) From fdd6f6021642b912c14513d8af91d1c9d65674bd Mon Sep 17 00:00:00 2001 From: mateo Date: Wed, 9 Sep 2026 10:01:22 +0000 Subject: [PATCH 21/55] test: count a zombie grandchild as killed in the fake prisma cli Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/db/conftest.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/test_litellm/proxy/db/conftest.py b/tests/test_litellm/proxy/db/conftest.py index d3226b0ec50..bcb7794a20d 100644 --- a/tests/test_litellm/proxy/db/conftest.py +++ b/tests/test_litellm/proxy/db/conftest.py @@ -35,6 +35,14 @@ DB_ENV_KEYS = ( _db_env_snapshot_key = pytest.StashKey[dict[str, Optional[str]]]() +def _is_zombie(pid: int) -> bool: + try: + stat: Final = Path(f"/proc/{pid}/stat").read_text() + except OSError: + return False + return stat.rpartition(")")[2].split()[0] == "Z" + + def _db_env_snapshot() -> dict[str, Optional[str]]: return {key: os.environ.get(key) for key in DB_ENV_KEYS} @@ -136,6 +144,8 @@ class FakePrismaCli: os.kill(pid, 0) except ProcessLookupError: return True + if _is_zombie(pid): + return True time.sleep(0.05) return False From 76c9342e9da00f0d8814b4c5a672cceff005f1a8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:11:42 -0700 Subject: [PATCH 22/55] fix(proxy): run migrations through python -m prisma when the prisma console script is not on PATH The proxy probed for the Prisma CLI by spawning the bare `prisma` console script, so a launcher whose PATH lacked the interpreter's bin directory printed "prisma package not found", skipped every migration and served traffic against an empty schema. Every Prisma command now falls back to `python -m prisma` when the console script is not on PATH, and the boot probe checks for the console script or the importable package instead of spawning anything. --- .../litellm_proxy_extras/prisma_toolchain.py | 27 +++- litellm/proxy/prisma_migration.py | 4 +- litellm/proxy/proxy_cli.py | 123 +++++++++--------- .../test_prisma_toolchain.py | 95 ++++++++++++++ .../proxy/test_prisma_migration.py | 33 +++++ tests/test_litellm/proxy/test_proxy_cli.py | 60 +++++++++ 6 files changed, 277 insertions(+), 65 deletions(-) diff --git a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py index b51de9609d3..9cd48fcf11a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py +++ b/litellm-proxy-extras/litellm_proxy_extras/prisma_toolchain.py @@ -37,11 +37,13 @@ raised it above the deploy default keeps that larger budget for deploy unless the deploy override says otherwise. """ +import importlib.util import math import os import shutil import signal import subprocess +import sys from collections.abc import Mapping, Sequence from dataclasses import dataclass from pathlib import Path @@ -64,6 +66,7 @@ DEFAULT_PRISMA_BOOTSTRAP_TIMEOUT = 600.0 DEFAULT_PRISMA_MIGRATE_DEPLOY_TIMEOUT = 600.0 BOOTSTRAP_ARG = "--version" +PRISMA_CONSOLE_SCRIPT = "prisma" @dataclass(frozen=True) @@ -184,6 +187,28 @@ def _kill_process_group(process: "subprocess.Popen[str]") -> None: return +def prisma_cli_available() -> bool: + """Whether some way of running the Prisma CLI exists: the console script on PATH or the importable package.""" + if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None: + return True + return importlib.util.find_spec(PRISMA_CONSOLE_SCRIPT) is not None + + +def resolve_prisma_argv(argv: Sequence[str]) -> tuple[str, ...]: + """Route a bare ``prisma`` command through ``python -m prisma`` when the console script is not on PATH. + + The console script and ``python -m prisma`` are the same entry point, but + only the module form survives an interpreter whose ``bin`` directory is + missing from PATH, which is how the proxy gets started under launchers and + init systems. Any other executable name is left untouched. + """ + if not argv or argv[0] != PRISMA_CONSOLE_SCRIPT: + return tuple(argv) + if shutil.which(PRISMA_CONSOLE_SCRIPT) is not None: + return tuple(argv) + return (sys.executable, "-m", PRISMA_CONSOLE_SCRIPT, *argv[1:]) + + def run_prisma( argv: Sequence[str], *, @@ -200,7 +225,7 @@ def run_prisma( text unless ``stdout``/``stderr`` say otherwise. """ with subprocess.Popen( - argv, + resolve_prisma_argv(argv), env=env, stdout=stdout, stderr=stderr, diff --git a/litellm/proxy/prisma_migration.py b/litellm/proxy/prisma_migration.py index 1b95d24c011..7e3aff75cef 100644 --- a/litellm/proxy/prisma_migration.py +++ b/litellm/proxy/prisma_migration.py @@ -14,6 +14,8 @@ sys.path.insert(0, os.path.abspath("./")) from typing import Final +from litellm_proxy_extras.prisma_toolchain import resolve_prisma_argv + from litellm._logging import verbose_proxy_logger from litellm.proxy.proxy_cli import run_server from litellm.secret_managers.main import str_to_bool @@ -29,7 +31,7 @@ def main() -> int: run_server(run_server_args, standalone_mode=False) verbose_proxy_logger.info("Running 'prisma generate'...") - result: Final = subprocess.run(("prisma", "generate"), capture_output=True, text=True) + result: Final = subprocess.run(resolve_prisma_argv(("prisma", "generate")), capture_output=True, text=True) verbose_proxy_logger.info("'prisma generate' stdout: %s", result.stdout) if result.returncode != 0: diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e245367b1b4..d9045c57b41 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -1260,73 +1260,69 @@ def run_server( flush=True, ) sys.exit(1) - try: - from litellm.secret_managers.main import get_secret + from litellm.secret_managers.main import get_secret - connection_url_params: Final = _build_db_connection_url_params( - connection_limit=db_connection_pool_limit, - pool_timeout=db_connection_timeout, - connect_timeout=db_connect_timeout, - socket_timeout=db_socket_timeout, - disable_prepared_statements=db_disable_prepared_statements, - extra_params=db_extra_connection_params, + connection_url_params: Final = _build_db_connection_url_params( + connection_limit=db_connection_pool_limit, + pool_timeout=db_connection_timeout, + connect_timeout=db_connect_timeout, + socket_timeout=db_socket_timeout, + disable_prepared_statements=db_disable_prepared_statements, + extra_params=db_extra_connection_params, + ) + lifetime_params: Final = idle_lifetime_params(general_settings.get("database_max_idle_connection_lifetime")) + if os.getenv("DATABASE_URL", None) is not None: + database_url = get_secret("DATABASE_URL", default_value=None) + resolved_url: Final[str | None] = str(database_url) if database_url else None + pg_options: Final[str] = _pg_options_with_timeouts( + _url_query_value(resolved_url, "options"), + db_statement_timeout, + db_lock_timeout, ) - lifetime_params: Final = idle_lifetime_params( - general_settings.get("database_max_idle_connection_lifetime") + writer_url: Final = ( + _with_query_value(resolved_url, "options", pg_options) + if resolved_url and pg_options + else resolved_url ) - if os.getenv("DATABASE_URL", None) is not None: - database_url = get_secret("DATABASE_URL", default_value=None) - resolved_url: Final[str | None] = str(database_url) if database_url else None - pg_options: Final[str] = _pg_options_with_timeouts( - _url_query_value(resolved_url, "options"), - db_statement_timeout, - db_lock_timeout, - ) - writer_url: Final = ( - _with_query_value(resolved_url, "options", pg_options) - if resolved_url and pg_options - else resolved_url - ) - modified_url = append_query_params( - writer_url, - connection_url_params, - ) - os.environ["DATABASE_URL"] = translate_libpq_ssl_params( - add_missing_query_params(modified_url, lifetime_params) - ) - if os.getenv("DIRECT_URL", None) is not None: - database_url = os.getenv("DIRECT_URL") - modified_url = append_query_params(database_url, connection_url_params) - os.environ["DIRECT_URL"] = translate_libpq_ssl_params( - add_missing_query_params(modified_url, lifetime_params) - ) - # The reader pool is a real pool against the same configured cap, so it - # gets the allowlisted pool params. Schema-affecting ones, including any - # the operator smuggled in through database_extra_connection_params, stay - # on the writer. Anything pinned on the replica URL wins, unlike the - # writer where the config is applied on top. - read_replica_url: Final[str | None] = os.getenv("DATABASE_URL_READ_REPLICA") - if read_replica_url: - reader_options: Final[str] = _pg_options_with_timeouts( - _url_query_value(read_replica_url, "options"), - db_statement_timeout, - db_lock_timeout, - ) - os.environ["DATABASE_URL_READ_REPLICA"] = translate_libpq_ssl_params( + modified_url = append_query_params( + writer_url, + connection_url_params, + ) + os.environ["DATABASE_URL"] = translate_libpq_ssl_params( + add_missing_query_params(modified_url, lifetime_params) + ) + if os.getenv("DIRECT_URL", None) is not None: + database_url = os.getenv("DIRECT_URL") + modified_url = append_query_params(database_url, connection_url_params) + os.environ["DIRECT_URL"] = translate_libpq_ssl_params( + add_missing_query_params(modified_url, lifetime_params) + ) + # The reader pool is a real pool against the same configured cap, so it + # gets the allowlisted pool params. Schema-affecting ones, including any + # the operator smuggled in through database_extra_connection_params, stay + # on the writer. Anything pinned on the replica URL wins, unlike the + # writer where the config is applied on top. + read_replica_url: Final[str | None] = os.getenv("DATABASE_URL_READ_REPLICA") + if read_replica_url: + reader_options: Final[str] = _pg_options_with_timeouts( + _url_query_value(read_replica_url, "options"), + db_statement_timeout, + db_lock_timeout, + ) + os.environ["DATABASE_URL_READ_REPLICA"] = translate_libpq_ssl_params( + add_missing_query_params( add_missing_query_params( - add_missing_query_params( - _with_query_value(read_replica_url, "options", reader_options) - if reader_options - else read_replica_url, - reader_shareable_params(connection_url_params), - ), - lifetime_params, - ) + _with_query_value(read_replica_url, "options", reader_options) + if reader_options + else read_replica_url, + reader_shareable_params(connection_url_params), + ), + lifetime_params, ) - subprocess.run(["prisma"], capture_output=True) - is_prisma_runnable = True - except FileNotFoundError: - is_prisma_runnable = False + ) + from litellm_proxy_extras.prisma_toolchain import prisma_cli_available + + is_prisma_runnable: Final = prisma_cli_available() if is_prisma_runnable: from litellm.proxy.db.check_migration import check_prisma_schema_diff @@ -1375,7 +1371,8 @@ def run_server( ) else: print( - f"Unable to connect to DB. DATABASE_URL found in environment, but prisma package not found." # noqa: F541 + "Unable to connect to DB. DATABASE_URL found in environment, but the prisma CLI is neither on " + "PATH nor importable as a package." ) if port == 4000 and ProxyInitializationHelpers._is_port_in_use(port): port = random.randint(1024, 49152) diff --git a/tests/proxy_migration_tests/test_prisma_toolchain.py b/tests/proxy_migration_tests/test_prisma_toolchain.py index 0ed33193a9b..4e2274cc582 100644 --- a/tests/proxy_migration_tests/test_prisma_toolchain.py +++ b/tests/proxy_migration_tests/test_prisma_toolchain.py @@ -37,7 +37,10 @@ from litellm_proxy_extras.prisma_toolchain import ( node_binary_path, prisma_bootstrap_timeout, prisma_command_timeout, + prisma_cli_available, prisma_migrate_deploy_timeout, + resolve_prisma_argv, + run_prisma, ) from litellm_proxy_extras.utils import ProxyExtrasDBManager @@ -401,3 +404,95 @@ def test_every_prisma_command_timeout_is_overridable(module: str) -> None: f"{module} still hardcodes a Prisma timeout at lines {literals}; " "route it through prisma_command_timeout() so it can be raised without a release" ) + + +FAKE_PRISMA_MODULE_MAIN = """import json +import sys + +print(json.dumps({"module_argv": sys.argv[1:]})) +""" + + +def _write_fake_prisma_module(tmp_path: Path) -> Path: + package_dir = tmp_path / "fakemodule" / "prisma" + package_dir.mkdir(parents=True) + (package_dir / "__init__.py").write_text("") + (package_dir / "__main__.py").write_text(FAKE_PRISMA_MODULE_MAIN) + return package_dir.parent + + +def _empty_bin(tmp_path: Path) -> Path: + bin_dir = tmp_path / "emptybin" + bin_dir.mkdir() + return bin_dir + + +def test_run_prisma_uses_the_module_when_the_console_script_is_not_on_path( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + empty_bin = _empty_bin(tmp_path) + module_root = _write_fake_prisma_module(tmp_path) + monkeypatch.setenv("PATH", str(empty_bin)) + + result = run_prisma( + ["prisma", "migrate", "deploy"], + timeout=60, + env={"PATH": str(empty_bin), "PYTHONPATH": str(module_root)}, + ) + + assert json.loads(result.stdout) == {"module_argv": ["migrate", "deploy"]} + + +def test_run_prisma_prefers_the_console_script_on_path( + toolchain_env: tuple[Path, Path], tmp_path: Path +) -> None: + _, log_path = toolchain_env + module_root = _write_fake_prisma_module(tmp_path) + + result = run_prisma( + ["prisma", "--version"], + timeout=60, + env={**os.environ, "PYTHONPATH": str(module_root)}, + ) + + assert [call["args"] for call in _fake_prisma_calls(log_path)] == [["--version"]] + assert "module_argv" not in result.stdout + + +def test_resolve_prisma_argv_leaves_an_explicit_cli_path_alone( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("PATH", str(_empty_bin(tmp_path))) + explicit = ("/app/.cache/prisma-python/prisma", "migrate", "deploy") + + assert resolve_prisma_argv(explicit) == explicit + + +def test_prisma_cli_is_unavailable_with_neither_script_nor_package( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("PATH", str(_empty_bin(tmp_path))) + monkeypatch.delitem(sys.modules, "prisma", raising=False) + monkeypatch.setattr(sys, "path", []) + + assert prisma_cli_available() is False + + +def test_prisma_cli_is_available_through_the_package_alone( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("PATH", str(_empty_bin(tmp_path))) + monkeypatch.delitem(sys.modules, "prisma", raising=False) + monkeypatch.setattr(sys, "path", [str(_write_fake_prisma_module(tmp_path))]) + + assert prisma_cli_available() is True + + +def test_prisma_cli_is_available_through_the_console_script_alone( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("PATH", str(_write_fake_prisma(tmp_path))) + monkeypatch.delitem(sys.modules, "prisma", raising=False) + monkeypatch.setattr(sys, "path", []) + + assert prisma_cli_available() is True diff --git a/tests/test_litellm/proxy/test_prisma_migration.py b/tests/test_litellm/proxy/test_prisma_migration.py index 729adcfb9e0..3fc69b34213 100644 --- a/tests/test_litellm/proxy/test_prisma_migration.py +++ b/tests/test_litellm/proxy/test_prisma_migration.py @@ -1,4 +1,6 @@ import os +import sys +from pathlib import Path from unittest.mock import MagicMock, patch import pytest @@ -64,3 +66,34 @@ class TestPrismaMigration: prisma_migration.main() mock_subprocess_run.assert_not_called() + + @patch("litellm.proxy.prisma_migration.subprocess.run") # test-quality-ok: the spawned argv is the behavior under test + @patch("litellm.proxy.prisma_migration.run_server") # test-quality-ok: run_server boots the whole proxy + def test_prisma_generate_runs_through_the_module_when_the_cli_is_not_on_path( + self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, tmp_path: Path + ) -> None: + mock_subprocess_run.return_value = MagicMock(returncode=0, stdout="", stderr="") + empty_bin: Path = tmp_path / "emptybin" + empty_bin.mkdir() + + with patch.dict(os.environ, {"PATH": str(empty_bin)}, clear=True): + assert prisma_migration.main() == 0 + + assert mock_subprocess_run.call_args.args[0] == (sys.executable, "-m", "prisma", "generate") + + @patch("litellm.proxy.prisma_migration.subprocess.run") # test-quality-ok: the spawned argv is the behavior under test + @patch("litellm.proxy.prisma_migration.run_server") # test-quality-ok: run_server boots the whole proxy + def test_prisma_generate_runs_the_console_script_when_it_is_on_path( + self, mock_run_server: MagicMock, mock_subprocess_run: MagicMock, tmp_path: Path + ) -> None: + mock_subprocess_run.return_value = MagicMock(returncode=0, stdout="", stderr="") + bin_dir: Path = tmp_path / "bin" + bin_dir.mkdir() + script: Path = bin_dir / "prisma" + script.write_text("#!/bin/sh\nexit 0\n") + script.chmod(0o755) + + with patch.dict(os.environ, {"PATH": str(bin_dir)}, clear=True): + assert prisma_migration.main() == 0 + + assert mock_subprocess_run.call_args.args[0] == ("prisma", "generate") diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index c76ff189a8a..e25e6a59884 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -1940,6 +1940,66 @@ class TestRunServerDbSetup: use_migrate=False, use_v2_resolver=False ) + @patch("atexit.register") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch("litellm.proxy.db.check_migration.check_prisma_schema_diff") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch("litellm.proxy.db.prisma_client.should_update_prisma_schema") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + def test_migrations_run_when_the_prisma_cli_is_not_on_path( + self, + mock_should_update_schema, + mock_check_schema_diff, + mock_setup_database, + mock_atexit_register, + tmp_path, + capsys, + ): + from litellm.proxy.proxy_cli import run_server + + mock_should_update_schema.return_value = True + empty_bin = tmp_path / "emptybin" + empty_bin.mkdir() + + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL") + } + clean_env["DATABASE_URL"] = "postgresql://test:test@localhost:5432/test" + clean_env["PATH"] = str(empty_bin) + + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + patch( # test-quality-ok: same isolation as the sibling CLI tests above + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + ): + mock_get_args.return_value = { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + run_server.main(["--local", "--skip_server_startup"], standalone_mode=False) + + assert "prisma CLI is neither on PATH" not in capsys.readouterr().out + mock_setup_database.assert_called_once_with( + use_migrate=True, use_v2_resolver=False + ) + @patch("subprocess.run") @patch("atexit.register") @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") From 45fad445ebad1af76b4923f4722abbd080db3f04 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:07:16 -0700 Subject: [PATCH 23/55] fix(cost): bill request-level OCR pricing on direct SDK calls --- litellm/cost_calculator.py | 47 ++++++++++++++++++++-- tests/test_litellm/ocr/test_main.py | 18 ++++++++- tests/test_litellm/test_cost_calculator.py | 34 ++++++++++++++-- 3 files changed, 91 insertions(+), 8 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 68f52e4a10e..6d35a9e89aa 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, cast from httpx import Response from pydantic import BaseModel +from typing_extensions import ReadOnly, TypedDict import litellm import litellm._logging @@ -306,6 +307,15 @@ def _transcription_usage_has_token_details( return (prompt_tokens_val > 0) or (completion_tokens_val > 0) +OCRPricingField = Literal["ocr_cost_per_page", "ocr_cost_per_credit", "annotation_cost_per_page"] + + +class OCRPricing(TypedDict, total=False): + ocr_cost_per_page: ReadOnly[float | None] + ocr_cost_per_credit: ReadOnly[float | None] + annotation_cost_per_page: ReadOnly[float | None] + + def cost_per_token( model: str = "", prompt_tokens: int = 0, @@ -341,7 +351,7 @@ def cost_per_token( ### REQUEST MODEL ### request_model: str | None = None, # original request model for router detection ### DEPLOYMENT-SPECIFIC PRICING ### - custom_model_info: ModelInfo | None = None, + custom_model_info: OCRPricing | None = None, ) -> tuple[float, float]: """ Calculates the cost per token for a given model, prompt tokens, and completion tokens. @@ -1652,7 +1662,7 @@ def completion_cost( vertex_location=vertex_location, response=completion_response, request_model=request_model_for_cost, - custom_model_info=_deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id), + custom_model_info=_ocr_model_info(litellm_logging_obj, custom_pricing, router_model_id), ) # Get additional costs from provider (e.g., routing fees, infrastructure costs) @@ -1911,6 +1921,35 @@ def _deployment_model_info( ) +def _ocr_model_info( + litellm_logging_obj: LitellmLoggingObject | None, + custom_pricing: bool | None, + router_model_id: str | None, +) -> OCRPricing | None: + deployment_info: Final = _deployment_model_info(litellm_logging_obj, custom_pricing, router_model_id) + litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) if custom_pricing else None + if litellm_params is None: + return deployment_info + return OCRPricing( + ocr_cost_per_page=_request_or_deployment_price("ocr_cost_per_page", litellm_params, deployment_info), + ocr_cost_per_credit=_request_or_deployment_price("ocr_cost_per_credit", litellm_params, deployment_info), + annotation_cost_per_page=_request_or_deployment_price( + "annotation_cost_per_page", litellm_params, deployment_info + ), + ) + + +def _request_or_deployment_price( + field: OCRPricingField, + litellm_params: Mapping[str, object], + deployment_info: ModelInfo | None, +) -> float | None: + request_price: Final = litellm_params.get(field) + if isinstance(request_price, int | float): + return request_price + return deployment_info.get(field) if deployment_info is not None else None + + def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelInfo | None: try: return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) @@ -1922,14 +1961,14 @@ def ocr_cost( model: str, custom_llm_provider: str | None, response: object | None = None, - model_info: ModelInfo | None = None, + model_info: OCRPricing | None = None, ) -> tuple[float, float]: """ Args: model: str - model name custom_llm_provider: Optional[str] - custom LLM provider response: Optional[Any] - response object - model_info: Optional[ModelInfo] - deployment-specific model info; its OCR pricing + model_info: Optional[OCRPricing] - deployment-specific model info; its OCR pricing takes precedence over the model cost map Returns: diff --git a/tests/test_litellm/ocr/test_main.py b/tests/test_litellm/ocr/test_main.py index 0007d98dcb0..de4b28dafdd 100644 --- a/tests/test_litellm/ocr/test_main.py +++ b/tests/test_litellm/ocr/test_main.py @@ -1,9 +1,13 @@ from typing import Final +import pytest + +import litellm from litellm.litellm_core_utils.litellm_logging import Logging, use_custom_pricing_for_model +from litellm.llms.base_llm.ocr.transformation import OCRPage, OCRResponse, OCRUsageInfo from litellm.ocr.main import _prepare_ocr_request -OCR_MODEL: Final = "mistral/mistral-ocr-4-1" +OCR_MODEL: Final = "mistral/some-unmapped-ocr-model-for-testing" DOCUMENT: Final = {"type": "document_url", "document_url": "https://example.com/doc.pdf"} @@ -47,3 +51,15 @@ def test_prepare_ocr_request_without_custom_pricing_leaves_logging_params_unpric assert "ocr_cost_per_page" not in logging_obj.litellm_params assert use_custom_pricing_for_model(logging_obj.litellm_params) is False + + +def test_direct_ocr_call_bills_request_level_per_page_pricing() -> None: + assert OCR_MODEL not in litellm.model_cost + logging_obj: Final = _prepare({"ocr_cost_per_page": 0.05}) + response: Final = OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)], + model=OCR_MODEL, + usage_info=OCRUsageInfo(pages_processed=3), + ) + + assert logging_obj._response_cost_calculator(result=response) == pytest.approx(0.05 * 3) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index cb9217b2f41..6db91d7a775 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4661,8 +4661,8 @@ def _ocr_response(model: str, pages_processed: int, credits: float | None = None ) -def _ocr_logging_obj(litellm_params: dict[str, dict[str, ModelInfo]]) -> Logging: - logging_obj = Logging( +def _ocr_logging_obj(litellm_params: dict[str, object]) -> Logging: + logging_obj: Final = Logging( model=UNMAPPED_OCR_MODEL, messages=[], stream=False, @@ -4671,7 +4671,7 @@ def _ocr_logging_obj(litellm_params: dict[str, dict[str, ModelInfo]]) -> Logging litellm_call_id="test-ocr-custom-pricing", function_id="1234", ) - logging_obj.litellm_params = litellm_params + logging_obj.update_environment_variables(litellm_params=litellm_params, optional_params={}) return logging_obj @@ -4796,6 +4796,34 @@ def test_completion_cost_ocr_prefers_pricing_registered_under_router_model_id(mo assert cost == pytest.approx(0.05 * 3) +def test_completion_cost_ocr_bills_request_level_pricing_for_direct_sdk_call(): + logging_obj = _ocr_logging_obj({"ocr_cost_per_page": 0.05}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(0.05 * 3) + + +def test_completion_cost_ocr_request_level_pricing_fills_in_deployment_model_info_without_ocr_pricing(): + logging_obj = _ocr_logging_obj({"ocr_cost_per_page": 0.05, "metadata": {"model_info": {"mode": "ocr"}}}) + + cost = completion_cost( + completion_response=_ocr_response(UNMAPPED_OCR_MODEL, pages_processed=3), + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + call_type="ocr", + custom_pricing=True, + litellm_logging_obj=logging_obj, + ) + assert cost == pytest.approx(0.05 * 3) + + def test_completion_cost_ocr_ignores_deployment_pricing_without_custom_pricing_flag(): logging_obj = _ocr_logging_obj({"metadata": {"model_info": {"ocr_cost_per_page": 0.004}}}) From 763270e875f92b4c7eca2d636d56f7e6089ed63d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:24:55 -0700 Subject: [PATCH 24/55] fix(cost): treat annotation-only deployment pricing as custom OCR pricing --- litellm/cost_calculator.py | 1 + tests/test_litellm/test_cost_calculator.py | 18 ++++++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 422e340abf5..09700ad6b72 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1999,6 +1999,7 @@ def ocr_cost( credits: Final = getattr(response.usage_info, "credits", None) has_custom_ocr_pricing: Final = model_info is not None and ( model_info.get("ocr_cost_per_page") is not None + or model_info.get("annotation_cost_per_page") is not None or (credits is not None and model_info.get("ocr_cost_per_credit") is not None) ) pricing: Final = model_info if has_custom_ocr_pricing else _cost_map_model_info(model, custom_llm_provider) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index a0ae9511aff..a05fadce42d 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4795,6 +4795,24 @@ def test_ocr_cost_uses_deployment_per_page_pricing_for_unmapped_model(pages_proc assert cost == pytest.approx(0.004 * pages_processed) +def test_ocr_cost_uses_deployment_annotation_only_pricing_for_unmapped_model(): + from litellm.cost_calculator import ocr_cost + + assert UNMAPPED_OCR_MODEL not in litellm.model_cost + response: Final = OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)], + model=UNMAPPED_OCR_MODEL, + usage_info=OCRUsageInfo(pages_processed=3, pages_processed_annotation=2), + ) + cost, _ = ocr_cost( + model=UNMAPPED_OCR_MODEL, + custom_llm_provider="azure_ai", + response=response, + model_info={"annotation_cost_per_page": 0.01}, + ) + assert cost == pytest.approx(0.01 * 2) + + def test_ocr_cost_uses_deployment_per_credit_pricing_for_unmapped_model(): from litellm.cost_calculator import ocr_cost From d71f4aeff9587ae189446598de23b35e6d8d8142 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 9 Sep 2026 19:47:00 -0700 Subject: [PATCH 25/55] fix(cost): layer deployment OCR rates over the cost map field by field --- litellm/cost_calculator.py | 45 +++++++++------------- tests/test_litellm/test_cost_calculator.py | 18 +++++++++ 2 files changed, 37 insertions(+), 26 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 09700ad6b72..93d9609df9e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1939,24 +1939,22 @@ def _ocr_model_info( litellm_params: Final = getattr(litellm_logging_obj, "litellm_params", None) if custom_pricing else None if litellm_params is None: return deployment_info - return OCRPricing( - ocr_cost_per_page=_request_or_deployment_price("ocr_cost_per_page", litellm_params, deployment_info), - ocr_cost_per_credit=_request_or_deployment_price("ocr_cost_per_credit", litellm_params, deployment_info), - annotation_cost_per_page=_request_or_deployment_price( - "annotation_cost_per_page", litellm_params, deployment_info - ), + return _layered_ocr_pricing(litellm_params, deployment_info) + + +def _first_ocr_price(field: OCRPricingField, *sources: Mapping[str, object] | None) -> float | None: + return next( + (price for source in sources if source is not None and isinstance(price := source.get(field), int | float)), + None, ) -def _request_or_deployment_price( - field: OCRPricingField, - litellm_params: Mapping[str, object], - deployment_info: ModelInfo | None, -) -> float | None: - request_price: Final = litellm_params.get(field) - if isinstance(request_price, int | float): - return request_price - return deployment_info.get(field) if deployment_info is not None else None +def _layered_ocr_pricing(*sources: Mapping[str, object] | None) -> OCRPricing: + return OCRPricing( + ocr_cost_per_page=_first_ocr_price("ocr_cost_per_page", *sources), + ocr_cost_per_credit=_first_ocr_price("ocr_cost_per_credit", *sources), + annotation_cost_per_page=_first_ocr_price("annotation_cost_per_page", *sources), + ) def _cost_map_model_info(model: str, custom_llm_provider: str | None) -> ModelInfo | None: @@ -1977,8 +1975,8 @@ def ocr_cost( model: str - model name custom_llm_provider: Optional[str] - custom LLM provider response: Optional[Any] - response object - model_info: Optional[OCRPricing] - deployment-specific model info; its OCR pricing - takes precedence over the model cost map + model_info: Optional[OCRPricing] - deployment-specific OCR pricing; each rate it sets + overrides the model cost map's, the rest fall back to the map Returns: Tuple[float, float]: cost of OCR processing @@ -1997,19 +1995,14 @@ def ocr_cost( raise ValueError("OCR response usage_info is None") credits: Final = getattr(response.usage_info, "credits", None) - has_custom_ocr_pricing: Final = model_info is not None and ( - model_info.get("ocr_cost_per_page") is not None - or model_info.get("annotation_cost_per_page") is not None - or (credits is not None and model_info.get("ocr_cost_per_credit") is not None) - ) - pricing: Final = model_info if has_custom_ocr_pricing else _cost_map_model_info(model, custom_llm_provider) + pricing: Final = _layered_ocr_pricing(model_info, _cost_map_model_info(model, custom_llm_provider)) - cost_per_credit: Final = pricing.get("ocr_cost_per_credit") if pricing is not None else None + cost_per_credit: Final = pricing.get("ocr_cost_per_credit") if credits is not None and cost_per_credit is not None: return cost_per_credit * credits, 0.0 - ocr_cost_per_page: Final = pricing.get("ocr_cost_per_page") if pricing is not None else None - annotation_cost_per_page: Final = pricing.get("annotation_cost_per_page") if pricing is not None else None + ocr_cost_per_page: Final = pricing.get("ocr_cost_per_page") + annotation_cost_per_page: Final = pricing.get("annotation_cost_per_page") annotation_rate: Final = annotation_cost_per_page if annotation_cost_per_page is not None else ocr_cost_per_page pages_processed: Final = response.usage_info.pages_processed diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index a05fadce42d..d07ff231429 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -4813,6 +4813,24 @@ def test_ocr_cost_uses_deployment_annotation_only_pricing_for_unmapped_model(): assert cost == pytest.approx(0.01 * 2) +def test_ocr_cost_annotation_only_override_keeps_mapped_per_page_rate(): + from litellm.cost_calculator import ocr_cost + + map_price: Final = litellm.model_cost[MAPPED_OCR_MODEL]["ocr_cost_per_page"] + response: Final = OCRResponse( + pages=[OCRPage(index=index, markdown=f"page {index}") for index in range(3)], + model=MAPPED_OCR_MODEL, + usage_info=OCRUsageInfo(pages_processed=3, pages_processed_annotation=2), + ) + cost, _ = ocr_cost( + model=MAPPED_OCR_MODEL, + custom_llm_provider="mistral", + response=response, + model_info={"annotation_cost_per_page": 0.01}, + ) + assert cost == pytest.approx(map_price * 3 + 0.01 * 2) + + def test_ocr_cost_uses_deployment_per_credit_pricing_for_unmapped_model(): from litellm.cost_calculator import ocr_cost From 251bc07e97dee08c6a61e3bf1b1e5d497e3d0c4e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:10:34 -0700 Subject: [PATCH 26/55] chore(cost): drop the section header comment above custom_model_info --- .git-check.out | 58 ++++++++++++++++++++++++++++++++++++++ litellm/cost_calculator.py | 1 - 2 files changed, 58 insertions(+), 1 deletion(-) create mode 100644 .git-check.out diff --git a/.git-check.out b/.git-check.out new file mode 100644 index 00000000000..73092f18281 --- /dev/null +++ b/.git-check.out @@ -0,0 +1,58 @@ +uv sync --inexact --frozen --extra proxy --group proxy-dev --group e2e-dev +Audited 232 packages in 79ms +uv run --no-sync python scripts/prisma_generate_if_needed.py +Environment variables loaded from .env +Prisma schema loaded from litellm/proxy/schema.prisma +Warning: The binaryTargets option is not officially supported by Prisma Client Python. + +✔ Generated Prisma Client Python (v0.11.0) to ./.venv/lib/python3.12/site-packages/prisma in 287ms + +cd ui/litellm-dashboard && ../../scripts/with_dashboard_node.sh npm install --no-audit --no-fund + +changed 1 package in 550ms +bootstrap: .env left untouched +bootstrap: done +./scripts/pre_commit_lint.sh +check: logging full output to /Users/mateo/Development/litellm/.git/worktrees/wt40516/pre_commit_lint.log +check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging: + .git-check.out + litellm/cost_calculator.py + litellm/ocr/main.py + tests/test_litellm/ocr/test_main.py + tests/test_litellm/test_cost_calculator.py +check: linting Python (make lint) +uv sync --inexact --frozen --group proxy-dev --group e2e-dev +Audited 205 packages in 51ms +uv run --no-sync python scripts/prisma_generate_if_needed.py +Prisma client already generated for litellm/proxy/schema.prisma (prisma 0.11.0); skipping prisma generate +cd litellm && uv run --no-sync ruff check . && cd .. +uv run --no-sync python scripts/ruff_strict_gate.py --base "origin/litellm_internal_staging" +uv run --no-sync python scripts/type_discipline_gate.py --base "origin/litellm_internal_staging" +uv run --no-sync python scripts/test_quality_gate.py --base "origin/litellm_internal_staging" +uv run --no-sync python scripts/type_check_gate.py --base "origin/litellm_internal_staging" +uv run --no-sync basedpyright tests/e2e +cd litellm && uv run --no-sync python ../tests/documentation_tests/test_circular_imports.py && cd .. +No LiteLLM type hints found. +provisioning .venv-typecheck (first run installs packages and generates the Prisma client; re-runs are near-instant no-ops) +2 files already formatted +All checks passed! +uv run --no-sync ruff check --config ruff-tests.toml tests +warning: Invalid `# noqa` directive on tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py:44: expected a comma-separated list of codes (e.g., `# noqa: F401, F841`). +All checks passed! +[from litellm import *] OK! no issues! +OK: every strict rule is within its codebase ceiling (base origin/litellm_internal_staging) +0 errors, 0 warnings, 0 notes +OK: every TQ rule is within its test-suite ceiling (base origin/litellm_internal_staging) +OK: every LIT rule is within its codebase ceiling (base origin/litellm_internal_staging) +base counts fetched from CI artifact basedpyright-counts-64ab15e608f7cb59 +OK: every rule is within its basedpyright limit or no higher than base (138494 errors total) +check: ruff format --check (scoped litellm files) +2 files already formatted +check: summary + ran: Python lint (make lint) + skipped: tests/e2e checks (basedpyright + raw HTTP client ban) (no tests/e2e Python files in scope) + ran: test-tree lint (ruff-tests.toml + test-quality budget) + skipped: dashboard lint (prettier + eslint + lint budgets) (no dashboard files in scope) + skipped: dashboard API-type sync (npm run gen:api) (no litellm/proxy, litellm/types, or generator files in scope) +check: PASS +check: full log: /Users/mateo/Development/litellm/.git/worktrees/wt40516/pre_commit_lint.log diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 93d9609df9e..22bdb016dc1 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -354,7 +354,6 @@ def cost_per_token( response: Any | None = None, ### REQUEST MODEL ### request_model: str | None = None, # original request model for router detection - ### DEPLOYMENT-SPECIFIC PRICING ### custom_model_info: OCRPricing | None = None, ) -> tuple[float, float]: """ From e4095077ec05984dc1f042caea3b661e257af4e1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 18:23:57 -0700 Subject: [PATCH 27/55] chore: remove a stray local check log --- .git-check.out | 58 -------------------------------------------------- 1 file changed, 58 deletions(-) delete mode 100644 .git-check.out diff --git a/.git-check.out b/.git-check.out deleted file mode 100644 index 73092f18281..00000000000 --- a/.git-check.out +++ /dev/null @@ -1,58 +0,0 @@ -uv sync --inexact --frozen --extra proxy --group proxy-dev --group e2e-dev -Audited 232 packages in 79ms -uv run --no-sync python scripts/prisma_generate_if_needed.py -Environment variables loaded from .env -Prisma schema loaded from litellm/proxy/schema.prisma -Warning: The binaryTargets option is not officially supported by Prisma Client Python. - -✔ Generated Prisma Client Python (v0.11.0) to ./.venv/lib/python3.12/site-packages/prisma in 287ms - -cd ui/litellm-dashboard && ../../scripts/with_dashboard_node.sh npm install --no-audit --no-fund - -changed 1 package in 550ms -bootstrap: .env left untouched -bootstrap: done -./scripts/pre_commit_lint.sh -check: logging full output to /Users/mateo/Development/litellm/.git/worktrees/wt40516/pre_commit_lint.log -check: nothing staged; scoping to the working tree's diff against the merge base with origin/litellm_internal_staging: - .git-check.out - litellm/cost_calculator.py - litellm/ocr/main.py - tests/test_litellm/ocr/test_main.py - tests/test_litellm/test_cost_calculator.py -check: linting Python (make lint) -uv sync --inexact --frozen --group proxy-dev --group e2e-dev -Audited 205 packages in 51ms -uv run --no-sync python scripts/prisma_generate_if_needed.py -Prisma client already generated for litellm/proxy/schema.prisma (prisma 0.11.0); skipping prisma generate -cd litellm && uv run --no-sync ruff check . && cd .. -uv run --no-sync python scripts/ruff_strict_gate.py --base "origin/litellm_internal_staging" -uv run --no-sync python scripts/type_discipline_gate.py --base "origin/litellm_internal_staging" -uv run --no-sync python scripts/test_quality_gate.py --base "origin/litellm_internal_staging" -uv run --no-sync python scripts/type_check_gate.py --base "origin/litellm_internal_staging" -uv run --no-sync basedpyright tests/e2e -cd litellm && uv run --no-sync python ../tests/documentation_tests/test_circular_imports.py && cd .. -No LiteLLM type hints found. -provisioning .venv-typecheck (first run installs packages and generates the Prisma client; re-runs are near-instant no-ops) -2 files already formatted -All checks passed! -uv run --no-sync ruff check --config ruff-tests.toml tests -warning: Invalid `# noqa` directive on tests/test_litellm/proxy/common_utils/test_user_api_key_cache.py:44: expected a comma-separated list of codes (e.g., `# noqa: F401, F841`). -All checks passed! -[from litellm import *] OK! no issues! -OK: every strict rule is within its codebase ceiling (base origin/litellm_internal_staging) -0 errors, 0 warnings, 0 notes -OK: every TQ rule is within its test-suite ceiling (base origin/litellm_internal_staging) -OK: every LIT rule is within its codebase ceiling (base origin/litellm_internal_staging) -base counts fetched from CI artifact basedpyright-counts-64ab15e608f7cb59 -OK: every rule is within its basedpyright limit or no higher than base (138494 errors total) -check: ruff format --check (scoped litellm files) -2 files already formatted -check: summary - ran: Python lint (make lint) - skipped: tests/e2e checks (basedpyright + raw HTTP client ban) (no tests/e2e Python files in scope) - ran: test-tree lint (ruff-tests.toml + test-quality budget) - skipped: dashboard lint (prettier + eslint + lint budgets) (no dashboard files in scope) - skipped: dashboard API-type sync (npm run gen:api) (no litellm/proxy, litellm/types, or generator files in scope) -check: PASS -check: full log: /Users/mateo/Development/litellm/.git/worktrees/wt40516/pre_commit_lint.log From a0b55fe68da96649973db4eb48d44c671c0160f6 Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 11 Sep 2026 01:30:00 +0000 Subject: [PATCH 28/55] feat(ui): search Key Activity by key alias, key hash, user id, or email Team Usage and the main Usage page render every key in the selected scope with no way to narrow the list. Add a client-side search box above Key Activity that filters the loaded keys by alias, hash, user id, or user email, and expose user_id on the daily activity key metadata so the id is searchable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../common_daily_activity.py | 1 + .../common_daily_activity.py | 1 + .../test_common_daily_activity.py | 5 + .../components/EntityUsage/EntityUsage.tsx | 3 +- .../_components/components/UsagePageView.tsx | 3 +- .../components/KeyActivityPanel.test.tsx | 71 +++++++++++++++ .../UsagePage/components/KeyActivityPanel.tsx | 58 ++++++++++++ .../UsagePage/keyActivityFilter.test.ts | 91 +++++++++++++++++++ .../components/UsagePage/keyActivityFilter.ts | 20 ++++ .../src/components/UsagePage/types.ts | 2 + .../src/components/activity_metrics.test.tsx | 17 ++++ .../src/components/activity_metrics.tsx | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 13 files changed, 273 insertions(+), 2 deletions(-) create mode 100644 ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx create mode 100644 ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx create mode 100644 ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.test.ts create mode 100644 ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.ts diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index a8aef30107c..44ed0017e42 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -127,6 +127,7 @@ def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str return KeyMetadata( key_alias=meta.get("key_alias"), team_id=meta.get("team_id"), + user_id=meta.get("user_id"), user_email=meta.get("user_email"), ) diff --git a/litellm/types/proxy/management_endpoints/common_daily_activity.py b/litellm/types/proxy/management_endpoints/common_daily_activity.py index 2b39c5dbb9b..090e5c42376 100644 --- a/litellm/types/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/types/proxy/management_endpoints/common_daily_activity.py @@ -43,6 +43,7 @@ class KeyMetadata(BaseModel): key_alias: str | None = None team_id: str | None = None + user_id: str | None = None user_email: str | None = None diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index 6cd900cb041..71896a18f48 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -613,6 +613,7 @@ def test_key_metadata_includes_recovered_user_email(): "dirty-key": { "key_alias": "batch-worker", "team_id": "team-1", + "user_id": "alice", "user_email": "alice@example.com", } }, @@ -620,6 +621,7 @@ def test_key_metadata_includes_recovered_user_email(): ) assert meta.key_alias == "batch-worker" + assert meta.user_id == "alice" assert meta.user_email == "alice@example.com" @@ -848,9 +850,11 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): mock_deleted_key.token = "deleted-key-hash" mock_deleted_key.key_alias = "toto-test-2" mock_deleted_key.team_id = "69cd4b77-b095-4489-8c46-4f2f31d840a2" + mock_deleted_key.user_id = "deleted-key-owner" mock_prisma.db.litellm_deletedverificationtoken = MagicMock() mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[mock_deleted_key]) + mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[]) result = await get_daily_activity_aggregated( prisma_client=mock_prisma, @@ -871,6 +875,7 @@ async def test_aggregated_activity_preserves_metadata_for_deleted_keys(): key_data = chat_endpoint.api_key_breakdown["deleted-key-hash"] assert key_data.metadata.key_alias == "toto-test-2" assert key_data.metadata.team_id == "69cd4b77-b095-4489-8c46-4f2f31d840a2" + assert key_data.metadata.user_id == "deleted-key-owner" assert key_data.metrics.spend == 10.0 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx index 273e478528e..6c15b3c418d 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/EntityUsage/EntityUsage.tsx @@ -42,6 +42,7 @@ import { valueFormatterSpend } from "@/components/UsagePage/utils/value_formatte import EndpointUsage from "../EndpointUsage/EndpointUsage"; import ModelViewToggle, { ModelViewType } from "../ModelViewToggle"; import TopKeyView from "@/components/UsagePage/components/EntityUsage/TopKeyView"; +import KeyActivityPanel from "@/components/UsagePage/components/KeyActivityPanel"; import TopModelView from "./TopModelView"; import TeamUserSpendCard from "./TeamUserSpendCard"; @@ -654,7 +655,7 @@ const EntityUsage: React.FC = ({ { key: "keys", label: "Key Activity", - content: , + content: , }, { key: "endpoints", label: "Endpoint Activity", content: }, ]; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx index a92d1209567..de353948db9 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.tsx @@ -30,6 +30,7 @@ import { ActivityMetrics, processActivityData } from "@/components/activity_metr import CloudZeroExportModal from "@/components/cloudzero_export_modal"; import UserDropdown from "@/components/common_components/UserDropdown"; import EntityUsageExportModal from "@/components/EntityUsageExport"; +import KeyActivityPanel from "@/components/UsagePage/components/KeyActivityPanel"; import { Team } from "@/components/key_team_helpers/key_list"; import { gatewayDailyActivityCall, @@ -886,7 +887,7 @@ const UsagePage: React.FC = ({ teams, organizations }) => { - + diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx new file mode 100644 index 00000000000..693ac20a360 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.test.tsx @@ -0,0 +1,71 @@ +import { fireEvent, render, screen } from "@testing-library/react"; +import { describe, expect, it, vi } from "vitest"; + +import type { ModelActivityData } from "../types"; +import KeyActivityPanel from "./KeyActivityPanel"; + +vi.mock("@/components/activity_metrics", () => ({ + ActivityMetrics: ({ modelMetrics }: { modelMetrics: Record }) => ( +
    + {Object.keys(modelMetrics).map((hash) => ( +
  • {hash}
    • + ))} +
+ ), +})); + +function activity(label: string, user_email: string | null, user_id: string | null): ModelActivityData { + return { + label, + key_metadata: { key_alias: label, team_id: "team-1", user_id, user_email }, + total_requests: 1, + total_successful_requests: 1, + total_failed_requests: 0, + total_cache_read_input_tokens: 0, + total_cache_creation_input_tokens: 0, + total_tokens: 10, + prompt_tokens: 5, + completion_tokens: 5, + total_spend: 0.01, + top_api_keys: [], + top_models: [], + daily_data: [], + }; +} + +const keyMetrics: Record = { + "hash-alice": activity("alice-key", "alice@example.com", "user-alice"), + "hash-bob": activity("bob-key", "bob@example.com", "user-bob"), +}; + +describe("KeyActivityPanel", () => { + it("renders every key and the full count before searching", () => { + render(); + expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alicehash-bob"); + expect(screen.getByText("Showing 2 of 2 keys")).toBeInTheDocument(); + }); + + it("narrows the rendered keys to those matching the user email", () => { + render(); + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "bob@example.com" } }); + expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-bob"); + expect(screen.getByTestId("rendered-keys")).not.toHaveTextContent("hash-alice"); + expect(screen.getByText("Showing 1 of 2 keys")).toBeInTheDocument(); + }); + + it("shows an empty state instead of zeroed metrics when nothing matches", () => { + render(); + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "carol" } }); + expect(screen.queryByTestId("rendered-keys")).not.toBeInTheDocument(); + expect(screen.getByText('No keys match "carol" in this date range')).toBeInTheDocument(); + }); + + it("clears the search and restores every key", () => { + render(); + fireEvent.change(screen.getByLabelText("Search keys"), { target: { value: "user-alice" } }); + expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alice"); + fireEvent.click(screen.getByLabelText("Clear key search")); + expect(screen.getByLabelText("Search keys")).toHaveValue(""); + expect(screen.getByTestId("rendered-keys")).toHaveTextContent("hash-alicehash-bob"); + }); +}); diff --git a/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx new file mode 100644 index 00000000000..8287a04d0c7 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/components/KeyActivityPanel.tsx @@ -0,0 +1,58 @@ +import { Search, X } from "lucide-react"; +import React, { useMemo, useState } from "react"; + +import { ActivityMetrics } from "@/components/activity_metrics"; +import { InputGroup, InputGroupAddon, InputGroupButton, InputGroupInput } from "@/components/ui/input-group"; + +import { filterKeyActivity } from "../keyActivityFilter"; +import type { ModelActivityData } from "../types"; + +interface KeyActivityPanelProps { + keyMetrics: Record; + hidePromptCachingMetrics?: boolean; +} + +const KeyActivityPanel: React.FC = ({ keyMetrics, hidePromptCachingMetrics = false }) => { + const [query, setQuery] = useState(""); + const filtered = useMemo(() => filterKeyActivity(keyMetrics, query), [keyMetrics, query]); + const totalKeys = Object.keys(keyMetrics).length; + const shownKeys = Object.keys(filtered).length; + const isFiltering = query.trim() !== ""; + + return ( +
+
+ + + + + setQuery(e.target.value)} + /> + {isFiltering && ( + + setQuery("")}> + + + + )} + + + Showing {shownKeys.toLocaleString()} of {totalKeys.toLocaleString()} keys + +
+ {isFiltering && totalKeys > 0 && shownKeys === 0 ? ( +

+ No keys match "{query.trim()}" in this date range +

+ ) : ( + + )} +
+ ); +}; + +export default KeyActivityPanel; diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.test.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.test.ts new file mode 100644 index 00000000000..ce181f6b0c6 --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.test.ts @@ -0,0 +1,91 @@ +import { describe, expect, it } from "vitest"; + +import { filterKeyActivity, keyActivityMatches } from "./keyActivityFilter"; +import type { KeyMetadata, ModelActivityData } from "./types"; + +function activity(label: string, key_metadata?: KeyMetadata): ModelActivityData { + return { + label, + key_metadata, + total_requests: 1, + total_successful_requests: 1, + total_failed_requests: 0, + total_cache_read_input_tokens: 0, + total_cache_creation_input_tokens: 0, + total_tokens: 10, + prompt_tokens: 5, + completion_tokens: 5, + total_spend: 0.01, + top_api_keys: [], + top_models: [], + daily_data: [], + }; +} + +const aliceMeta: KeyMetadata = { + key_alias: "alice-batch", + team_id: "team-research", + user_id: "user-alice-1234", + user_email: "alice@example.com", +}; +const bobMeta: KeyMetadata = { + key_alias: null, + team_id: "team-research", + user_id: "user-bob-5678", + user_email: "bob@example.com", +}; +const alice = activity("alice-batch (team: research)", aliceMeta); +const bob = activity("bob@example.com (team: research)", bobMeta); +const orphan = activity("key-hash-deadbeef", { key_alias: null, team_id: null }); + +const keyMetrics: Record = { + "hash-alice": alice, + "hash-bob": bob, + deadbeef: orphan, +}; + +describe("keyActivityMatches", () => { + it("matches every key on an empty or whitespace query", () => { + expect(keyActivityMatches("deadbeef", orphan, "")).toBe(true); + expect(keyActivityMatches("deadbeef", orphan, " ")).toBe(true); + }); + + it("matches key alias case-insensitively", () => { + expect(keyActivityMatches("hash-alice", alice, "ALICE-batch")).toBe(true); + expect(keyActivityMatches("hash-bob", bob, "alice-batch")).toBe(false); + }); + + it("matches user email", () => { + expect(keyActivityMatches("hash-bob", bob, "bob@example")).toBe(true); + expect(keyActivityMatches("hash-alice", alice, "bob@example")).toBe(false); + }); + + it("matches user id", () => { + expect(keyActivityMatches("hash-alice", alice, "user-alice-1234")).toBe(true); + expect(keyActivityMatches("hash-bob", bob, "user-alice-1234")).toBe(false); + }); + + it("matches the key hash when the key has no alias or user metadata", () => { + expect(keyActivityMatches("deadbeef", orphan, "dead")).toBe(true); + expect(keyActivityMatches("deadbeef", orphan, "alice")).toBe(false); + }); + + it("trims surrounding whitespace from the query", () => { + expect(keyActivityMatches("hash-alice", alice, " alice@example.com ")).toBe(true); + }); +}); + +describe("filterKeyActivity", () => { + it("returns the same object when the query is blank", () => { + expect(filterKeyActivity(keyMetrics, "")).toBe(keyMetrics); + }); + + it("keeps only the keys matching the query, preserving their hashes", () => { + expect(Object.keys(filterKeyActivity(keyMetrics, "example.com"))).toEqual(["hash-alice", "hash-bob"]); + expect(filterKeyActivity(keyMetrics, "user-bob")).toEqual({ "hash-bob": bob }); + }); + + it("returns an empty record when nothing matches", () => { + expect(filterKeyActivity(keyMetrics, "nobody")).toEqual({}); + }); +}); diff --git a/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.ts b/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.ts new file mode 100644 index 00000000000..1e9a654fb2e --- /dev/null +++ b/ui/litellm-dashboard/src/components/UsagePage/keyActivityFilter.ts @@ -0,0 +1,20 @@ +import type { ModelActivityData } from "./types"; + +export function keyActivityMatches(apiKey: string, data: ModelActivityData, query: string): boolean { + const needle = query.trim().toLowerCase(); + if (needle === "") return true; + const meta = data.key_metadata; + return [apiKey, data.label, meta?.key_alias, meta?.user_id, meta?.user_email].some( + (field) => field?.toLowerCase().includes(needle) ?? false, + ); +} + +export function filterKeyActivity( + keyMetrics: Record, + query: string, +): Record { + if (query.trim() === "") return keyMetrics; + return Object.fromEntries( + Object.entries(keyMetrics).filter(([apiKey, data]) => keyActivityMatches(apiKey, data, query)), + ); +} diff --git a/ui/litellm-dashboard/src/components/UsagePage/types.ts b/ui/litellm-dashboard/src/components/UsagePage/types.ts index a10e9e68c4d..fd4f1350020 100644 --- a/ui/litellm-dashboard/src/components/UsagePage/types.ts +++ b/ui/litellm-dashboard/src/components/UsagePage/types.ts @@ -46,6 +46,7 @@ export interface KeyMetricWithMetadata { export interface KeyMetadata { key_alias: string | null; team_id: string | null; + user_id?: string | null; user_email?: string | null; tags?: { tag: string; usage: number }[]; } @@ -70,6 +71,7 @@ export interface TopModelData { export interface ModelActivityData { label: string; + key_metadata?: KeyMetadata; total_requests: number; total_successful_requests: number; total_failed_requests: number; diff --git a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx index b0fc8dc7866..914fe1872b6 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.test.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.test.tsx @@ -655,6 +655,23 @@ describe("processActivityData", () => { expect(result["key1"].label).toBe("test-key-1 (team_id: team1)"); }); + it("retains the api key metadata so key activity can be searched by user", () => { + const metadata = { key_alias: "test-key-1", team_id: "team1", user_id: "user-1", user_email: "user1@example.com" }; + const withUser: { results: DailyData[] } = { + results: [ + createMockDailyData("2025-01-01", mockDailyActivity.results[0].metrics, { + ...EMPTY_BREAKDOWN, + api_keys: { key1: createMockKeyMetricWithMetadata(metadata, mockDailyActivity.results[0].metrics) }, + }), + ], + }; + + const result = processActivityData(withUser, "api_keys", MOCK_TEAMS); + + expect(result["key1"].key_metadata).toEqual(metadata); + expect(processActivityData(withUser, "models")["key1"]).toBeUndefined(); + }); + it("should process data for models key with data", () => { const dailyActivityWithModels: { results: DailyData[] } = { results: [ diff --git a/ui/litellm-dashboard/src/components/activity_metrics.tsx b/ui/litellm-dashboard/src/components/activity_metrics.tsx index f4348fb65ae..7c40a91be29 100644 --- a/ui/litellm-dashboard/src/components/activity_metrics.tsx +++ b/ui/litellm-dashboard/src/components/activity_metrics.tsx @@ -461,6 +461,7 @@ export const processActivityData = ( : key === "entities" ? (modelData as any).metadata?.agent_name || (modelData as any).metadata?.team_alias || model : model, + ...(key === "api_keys" ? { key_metadata: (modelData as KeyMetricWithMetadata).metadata } : {}), total_requests: 0, total_successful_requests: 0, total_failed_requests: 0, diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 6826cded6f5..fd6f52d4e35 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28414,6 +28414,8 @@ export interface components { team_id?: string | null; /** User Email */ user_email?: string | null; + /** User Id */ + user_id?: string | null; }; /** * KeyMetricWithMetadata From d36cdcb8f9808448eeedc99f50996feb1ea92536 Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 11 Sep 2026 01:30:21 +0000 Subject: [PATCH 29/55] fix(ui): show loading state instead of stale rows while a table search is pending Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../hooks/common/useResourceList.test.tsx | 22 +++++++++++++++++++ .../hooks/common/useResourceList.ts | 4 ++-- .../users/_components/view_users.test.tsx | 11 ++++++++++ .../users/_components/view_users.tsx | 2 +- .../components/TeamsPage/TeamsTable.test.tsx | 8 +++++++ .../src/components/TeamsPage/TeamsTable.tsx | 5 +++-- .../VirtualKeysPage/VirtualKeysTable.test.tsx | 9 ++++++++ .../VirtualKeysPage/VirtualKeysTable.tsx | 5 +++-- 8 files changed, 59 insertions(+), 7 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx index 3ca67b082ec..b0c478920f3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.test.tsx @@ -152,6 +152,28 @@ describe("useResourceList", () => { await waitFor(() => expect(lastCall().page_size).toBe(25)); }); + it("reports loading while a new search request is still pending", async () => { + let resolveSecond: ((value: ResourceListPage) => void) | undefined; + const fetchPage = vi.fn((query: ResourceListQuery) => { + calls.push(query); + if (calls.length === 1) return Promise.resolve(page([{ id: "a" }], 3)); + return new Promise>((resolve) => { + resolveSecond = resolve; + }); + }); + const { result } = renderList({ fetchPage }); + await waitFor(() => expect(result.current.rows).toEqual([{ id: "a" }])); + expect(result.current.isLoading).toBe(false); + + act(() => result.current.onSearchChange("zzz")); + await waitFor(() => expect(lastCall().q).toBe("zzz")); + expect(result.current.isLoading).toBe(true); + + act(() => resolveSecond?.(page([], 0))); + await waitFor(() => expect(result.current.isLoading).toBe(false)); + expect(result.current.rows).toEqual([]); + }); + it("surfaces a failed page as an error instead of empty rows", async () => { const fetchPage = vi.fn(() => Promise.reject(new Error("boom"))); const { result } = renderList({ fetchPage }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts index 8a6376b2248..983e537bd00 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/common/useResourceList.ts @@ -86,7 +86,7 @@ export function useResourceList(options: UseResourceListOptions): Re enabled, placeholderData: (previous) => previous, }; - const { data, isLoading, isFetching, error, refetch: refetchQuery } = useQuery(queryOptions); + const { data, isLoading, isPlaceholderData, isFetching, error, refetch: refetchQuery } = useQuery(queryOptions); const toFirstPage = useCallback(() => setPagination((previous) => ({ ...previous, pageIndex: 0 })), []); @@ -123,7 +123,7 @@ export function useResourceList(options: UseResourceListOptions): Re return { rows, rowCount: data?.meta.total_count ?? 0, - isLoading, + isLoading: isLoading || isPlaceholderData, isFetching, error, refetch, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx index 6b0423067fd..10a13c1dbd7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.test.tsx @@ -344,5 +344,16 @@ describe("ViewUserDashboard", () => { expect(latest[4]).toBeNull(); expect(latest[2]).toBe(1); }); + + it("replaces the previous rows with the loading state while the search request is pending", async () => { + renderDashboard(); + expect(await screen.findByText("test@example.com")).toBeInTheDocument(); + + userListCall.mockReturnValue(new Promise(() => undefined)); + fireEvent.change(screen.getByPlaceholderText("Search by email or ID…"), { target: { value: "zzznomatch" } }); + + expect(await screen.findByText("Loading users…")).toBeInTheDocument(); + expect(screen.queryByText("test@example.com")).not.toBeInTheDocument(); + }); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx index 1ed23e7f523..20ce22b6444 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users.tsx @@ -295,7 +295,7 @@ const ViewUserDashboard: React.FC = ({ { expect(screen.queryByText("Acme Team")).not.toBeInTheDocument(); }); +it("replaces the previous rows with the loading state while a new search is pending", () => { + mockUseTeamsTable.mockReturnValue(teamsResult([mockTeam], {}, { isPlaceholderData: true, isFetching: true })); + renderTable(); + + expect(screen.getByText("Loading teams...")).toBeInTheDocument(); + expect(screen.queryByText("Acme Team")).not.toBeInTheDocument(); +}); + describe("sort contract – only backend-sortable columns are sortable", () => { it("requests the default created_at descending sort on first render", () => { renderTable(); diff --git a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx index 3fb19e522f3..a1ec61daa32 100644 --- a/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx +++ b/ui/litellm-dashboard/src/components/TeamsPage/TeamsTable.tsx @@ -83,7 +83,8 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet const { data: teamsResponse, - isPending: isLoading, + isPending, + isPlaceholderData, isFetching, refetch, } = useTeamsTable(tablePagination.pageIndex + 1, tablePagination.pageSize, teamListOptions); @@ -161,7 +162,7 @@ export function TeamsTable({ userRole, userID, onSelectTeam, onEditTeam, onDelet onColumnFiltersChange={handleColumnFiltersChange} enableColumnResizing columnResizeMode="onChange" - isLoading={isLoading} + isLoading={isPending || isPlaceholderData} loadingMessage="Loading teams..." noDataMessage="No teams found" fillHeight diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx index 881b0b93ff9..bc9017468aa 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.test.tsx @@ -289,6 +289,15 @@ it("should show a loading state on the initial load and hide the data", () => { expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); }); +it("replaces the previous rows with the loading state while a new search is pending", () => { + mockUseKeys.mockReturnValue(keysResult([mockKey], {}, { isPlaceholderData: true, isFetching: true })); + + renderWithProviders(); + + expect(screen.getByText("Loading keys...")).toBeInTheDocument(); + expect(screen.queryByText("Test Key Alias")).not.toBeInTheDocument(); +}); + it("should show 'No keys found' message when the key list is empty", () => { mockUseKeys.mockReturnValue(keysResult([])); diff --git a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx index f28ae27b6bc..1394b7259c7 100644 --- a/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx +++ b/ui/litellm-dashboard/src/components/VirtualKeysPage/VirtualKeysTable.tsx @@ -128,7 +128,8 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { const { data: keys, - isPending: isLoading, + isPending, + isPlaceholderData, isFetching, refetch, } = useKeys(tablePagination.pageIndex + 1, tablePagination.pageSize, keyListOptions); @@ -280,7 +281,7 @@ export function VirtualKeysTable({ headerActions }: VirtualKeysTableProps) { onColumnFiltersChange={handleColumnFiltersChange} enableColumnResizing columnResizeMode="onChange" - isLoading={isLoading} + isLoading={isPending || isPlaceholderData} loadingMessage="Loading keys..." noDataMessage="No keys found" fillHeight From 0b7e305d7c9f8964dc354d3143d152d33778cf0b Mon Sep 17 00:00:00 2001 From: ryan Date: Fri, 11 Sep 2026 01:33:13 +0000 Subject: [PATCH 30/55] chore(proxy): regenerate lazy OpenAPI snapshot for KeyMetadata.user_id Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/_lazy_openapi_snapshot.json | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index f0af17ab818..5d652bda225 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3234,6 +3234,17 @@ } ], "title": "User Email" + }, + "user_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "User Id" } }, "title": "KeyMetadata", From 54e247998e74162d3d1a86d2beafce62412797d5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 20:14:50 -0700 Subject: [PATCH 31/55] fix(redis): count pool wait timeouts as breaker timeouts redis-py's blocking pool reports a saturated pool as ConnectionError chained from asyncio.TimeoutError. The circuit breaker classified that as a hard connectivity failure and opened at once while Redis was healthy. Follow the explicit cause chain so it counts as a timeout and stays behind the timeout_min_duration gate --- litellm/caching/redis_cache.py | 10 +++- .../test_litellm/caching/test_redis_cache.py | 60 +++++++++++++++++++ 2 files changed, 69 insertions(+), 1 deletion(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 6b93529e456..7b7d3841109 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -321,7 +321,15 @@ def _redis_timeout_error_types() -> tuple[type, ...]: def _is_redis_timeout_failure(exc: BaseException) -> bool: - return isinstance(exc, _redis_timeout_error_types()) + """True when ``exc`` or any exception it was explicitly raised ``from`` is a timeout. + + redis-py's blocking pool reports a pool wait timeout as ``ConnectionError`` chained from + ``asyncio.TimeoutError``, which is a busy pool rather than an unreachable Redis. + """ + if isinstance(exc, _redis_timeout_error_types()): + return True + cause: Final = exc.__cause__ + return cause is not None and _is_redis_timeout_failure(cause) class _BreakerMetrics: diff --git a/tests/test_litellm/caching/test_redis_cache.py b/tests/test_litellm/caching/test_redis_cache.py index bcaa58c9c40..2e1dfb77fd2 100644 --- a/tests/test_litellm/caching/test_redis_cache.py +++ b/tests/test_litellm/caching/test_redis_cache.py @@ -1030,3 +1030,63 @@ def test_sync_guard_counts_a_timeout_as_a_timeout(): _run_under_circuit_breaker_sync(breaker, "op", timing_out_call) assert breaker.is_open() is False + + +@pytest.mark.asyncio +async def test_pool_wait_timeout_is_a_timeout_failure_not_hard_connectivity(): + """A saturated blocking pool must not open the breaker before the timeout minimum duration. + + redis-py's async BlockingConnectionPool gives up waiting for a free connection by raising + ConnectionError("No connection available.") chained from asyncio.TimeoutError. Redis itself + is healthy in that case, so the failure has to be classed as a timeout and stay behind the + duration gate instead of being counted as a hard connectivity failure. + """ + from fakeredis import FakeServer + from fakeredis.aioredis import FakeConnection + from redis.asyncio import BlockingConnectionPool, Redis + from redis.exceptions import ConnectionError as RedisConnectionError + + from litellm.caching.redis_cache import RedisCircuitBreaker, _run_under_circuit_breaker + + pool = BlockingConnectionPool(connection_class=FakeConnection, server=FakeServer(), max_connections=1, timeout=0.01) + client = Redis(connection_pool=pool) + breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60, timeout_min_duration=5.0) + + busy_connection = await pool.get_connection() + try: + for _ in range(breaker.failure_threshold * 2): + with pytest.raises(RedisConnectionError, match="No connection available"): + await _run_under_circuit_breaker(breaker, "op", lambda: client.get("k")) + finally: + await pool.release(busy_connection) + + assert breaker.is_open() is False, "a busy pool is a timeout gated on duration, not a dead Redis" + assert await _run_under_circuit_breaker(breaker, "op", lambda: client.get("k")) is None + await client.aclose() + + +def test_timeout_classification_follows_the_explicit_cause_chain_only(): + from redis.exceptions import ConnectionError as RedisConnectionError + + from litellm.caching.redis_cache import _is_redis_timeout_failure + + def raise_chained_from_timeout() -> None: + try: + raise asyncio.TimeoutError() + except asyncio.TimeoutError as err: + raise RedisConnectionError("No connection available.") from err + + def raise_while_handling_timeout() -> None: + try: + raise asyncio.TimeoutError() + except asyncio.TimeoutError: + raise RedisConnectionError("refused") + + with pytest.raises(RedisConnectionError) as chained: + raise_chained_from_timeout() + with pytest.raises(RedisConnectionError) as contextual: + raise_while_handling_timeout() + + assert _is_redis_timeout_failure(chained.value) is True + assert _is_redis_timeout_failure(contextual.value) is False + assert _is_redis_timeout_failure(RedisConnectionError("refused")) is False From 1fe6984ec017f282774c61ab2f7b611a5a05d77a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 20:35:25 -0700 Subject: [PATCH 32/55] refactor(redis): walk the exception cause chain iteratively The recursion detector flags any unignored recursive function, so the timeout classification now walks the explicit cause chain with a bounded generator instead of calling itself --- litellm/caching/redis_cache.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 697aabe7c8d..4ba548978d9 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -16,7 +16,7 @@ import inspect import json import logging import time -from collections.abc import Awaitable, Callable, Sequence +from collections.abc import Awaitable, Callable, Iterator, Sequence from contextvars import ContextVar from dataclasses import dataclass from datetime import timedelta @@ -328,16 +328,26 @@ def _redis_timeout_error_types() -> tuple[type, ...]: return (RedisTimeoutError, TimeoutError) +_MAX_EXCEPTION_CAUSE_DEPTH: Final = 20 + + +def _explicit_causes(exc: BaseException) -> Iterator[BaseException]: + current = exc # rebind-ok: advances one link per iteration of the bounded walk + for _ in range(_MAX_EXCEPTION_CAUSE_DEPTH): + yield current + if current.__cause__ is None: + return + current = current.__cause__ + + def _is_redis_timeout_failure(exc: BaseException) -> bool: """True when ``exc`` or any exception it was explicitly raised ``from`` is a timeout. redis-py's blocking pool reports a pool wait timeout as ``ConnectionError`` chained from ``asyncio.TimeoutError``, which is a busy pool rather than an unreachable Redis. """ - if isinstance(exc, _redis_timeout_error_types()): - return True - cause: Final = exc.__cause__ - return cause is not None and _is_redis_timeout_failure(cause) + timeout_types: Final = _redis_timeout_error_types() + return any(isinstance(link, timeout_types) for link in _explicit_causes(exc)) class _BreakerMetrics: From cf1f7096815603f4f93cb182f012360c90043034 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 10 Sep 2026 20:59:33 -0700 Subject: [PATCH 33/55] fix(redis): treat asyncio.TimeoutError as a timeout on Python 3.10 --- litellm/caching/redis_cache.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index 4ba548978d9..4264801d051 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -317,15 +317,15 @@ def _is_redis_health_failure(exc: BaseException) -> bool: def _redis_timeout_error_types() -> tuple[type, ...]: """Health failures that are timeouts rather than unambiguous connectivity errors. - ``builtins.TimeoutError`` covers ``asyncio.TimeoutError`` and ``socket.timeout`` - (aliases since py3.11 / py3.10). ``redis.exceptions.TimeoutError`` does not subclass - either, so it is listed explicitly. + ``builtins.TimeoutError`` covers ``socket.timeout`` (an alias since py3.10) and, from + py3.11, ``asyncio.TimeoutError``; on py3.10 ``asyncio.TimeoutError`` is still its own + class, so it is listed explicitly. ``redis.exceptions.TimeoutError`` subclasses neither. """ try: from redis.exceptions import TimeoutError as RedisTimeoutError except ImportError: - return (TimeoutError,) - return (RedisTimeoutError, TimeoutError) + return (TimeoutError, asyncio.TimeoutError) + return (RedisTimeoutError, TimeoutError, asyncio.TimeoutError) _MAX_EXCEPTION_CAUSE_DEPTH: Final = 20 From da4052dcd77459108894434716509129b79b1d69 Mon Sep 17 00:00:00 2001 From: mateo Date: Fri, 11 Sep 2026 10:18:54 +0000 Subject: [PATCH 34/55] test: give the fake pooler a readiness budget that survives a loaded CI worker Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/proxy/db/test_pgbouncer.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/proxy/db/test_pgbouncer.py b/tests/test_litellm/proxy/db/test_pgbouncer.py index 7b5a0bf10f9..83b9be8044e 100644 --- a/tests/test_litellm/proxy/db/test_pgbouncer.py +++ b/tests/test_litellm/proxy/db/test_pgbouncer.py @@ -411,7 +411,7 @@ class TestPgBouncerProcess: port=port, socket_path=unix_socket_path(tmp_path, port), restart_delay_seconds=0.1, - ready_timeout_seconds=0.3, + ready_timeout_seconds=3.0, ) assert pooler.start() is None first_pid: Final = pooler.pid @@ -421,7 +421,10 @@ class TestPgBouncerProcess: with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name): os.kill(first_pid, signal.SIGKILL) assert _wait_until(lambda: _listening(wrong_port)) - assert _wait_until(lambda: any("did not start listening" in record.message for record in caplog.records)) + assert _wait_until( + lambda: any("did not start listening" in record.message for record in caplog.records), + timeout_seconds=10.0, + ) port_file.write_text(str(port)) assert _wait_until(lambda: _listening(port)) assert _wait_until(lambda: not _listening(wrong_port)) From b82b31a44f1b44b0a4e4edb324305ac10ce6f533 Mon Sep 17 00:00:00 2001 From: Young Han Date: Wed, 2 Sep 2026 13:22:47 -0700 Subject: [PATCH 35/55] feat(realtime): add Meta Muse Voice transcription --- .../litellm_core_utils/realtime_streaming.py | 87 ++- litellm/llms/meta/__init__.py | 3 + litellm/llms/meta/realtime/__init__.py | 10 + litellm/llms/meta/realtime/handler.py | 661 ++++++++++++++++++ litellm/llms/meta/realtime/transformation.py | 619 ++++++++++++++++ litellm/llms/openai_like/providers.json | 2 +- ...odel_prices_and_context_window_backup.json | 27 +- litellm/realtime_api/main.py | 34 +- litellm/types/llms/openai.py | 2 + litellm/types/realtime.py | 12 +- model_prices_and_context_window.json | 27 +- .../test_realtime_streaming.py | 191 ++++- .../realtime/test_meta_realtime_handler.py | 449 ++++++++++++ .../test_meta_realtime_transformation.py | 299 ++++++++ tests/test_litellm/realtime_api/test_main.py | 80 ++- 15 files changed, 2446 insertions(+), 57 deletions(-) create mode 100644 litellm/llms/meta/__init__.py create mode 100644 litellm/llms/meta/realtime/__init__.py create mode 100644 litellm/llms/meta/realtime/handler.py create mode 100644 litellm/llms/meta/realtime/transformation.py create mode 100644 tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py create mode 100644 tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 4923bdda305..4be7dd6b4ce 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -19,7 +19,7 @@ from litellm.types.llms.openai import ( OpenAIRealtimeStreamResponseBaseObject, OpenAIRealtimeStreamSessionEvents, ) -from litellm.types.realtime import ALL_DELTA_TYPES +from litellm.types.realtime import ALL_DELTA_TYPES, RealtimeInputAudioTranscriptionUsage from .litellm_logging import Logging as LiteLLMLogging from .realtime_errors import client_close_code, realtime_error_event, websocket_close_reason @@ -116,6 +116,10 @@ class RealtimeEventNormalizer(Protocol): def patch_outgoing_session(self, session: dict) -> dict: ... +class RealtimeUsageProvider(Protocol): + def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None: ... + + DefaultLoggedRealTimeEventTypes: Final = [ "session.created", "response.create", @@ -139,6 +143,8 @@ class RealTimeStreaming: force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, + usage_provider: RealtimeUsageProvider | None = None, + exclude_private_content_from_logs: bool = False, ): self.websocket: _ClientWebSocket = websocket self.backend_ws = backend_ws @@ -200,6 +206,10 @@ class RealTimeStreaming: self._is_transcription_session: bool = force_transcription_model is not None # Optional per-provider GA event normalizer (e.g. XAIRealtimeNormalizer). self._event_normalizer = event_normalizer + self._usage_provider: RealtimeUsageProvider | None = ( + usage_provider if usage_provider is not None else provider_config + ) + self._exclude_private_content_from_logs = exclude_private_content_from_logs # Per-connection caps for pre-setup audio frames (message count + total bytes). _MAX_BUFFERED_MESSAGES: int = 200 @@ -237,7 +247,7 @@ class RealTimeStreaming: def _should_store_message( self, - message_obj: dict | OpenAIRealtimeEvents, + message_obj: dict[str, Any] | OpenAIRealtimeEvents, # mutable-ok: existing realtime event contract ) -> bool: _msg_type: Final = message_obj["type"] if "type" in message_obj else None if self.logged_real_time_event_types == "*": @@ -246,16 +256,54 @@ class RealTimeStreaming: return True return False + def _message_for_logging( + self, + message_obj: dict[str, Any], # mutable-ok: existing realtime event contract + ) -> dict[str, Any]: # mutable-ok: logging stores concrete event dictionaries + if not self._exclude_private_content_from_logs: + return message_obj + logged_message: dict[str, Any] = { # mutable-ok: incrementally builds the sanitized event copy + key: message_obj[key] + for key in ( + "type", + "event_id", + "item_id", + "response_id", + "conversation_id", + "session_id", + "content_index", + "output_index", + "model", + "mode", + "usage", + ) + if key in message_obj + } + session: Final = message_obj.get("session") + if isinstance(session, dict): + logged_session: Final[dict[str, Any]] = { # mutable-ok: sanitized JSON session snapshot + key: session[key] for key in ("id", "model", "mode", "type") if key in session + } + if logged_session: + logged_message["session"] = logged_session + return logged_message + def store_message(self, message: str | bytes | dict | OpenAIRealtimeEvents): """Store message in list""" if isinstance(message, bytes): message = message.decode("utf-8") if isinstance(message, dict): # TypedDict union members do not narrow to plain dict for mypy. - message_obj: dict[str, Any] = cast(dict[str, Any], message) + parsed_message_obj: dict[str, Any] = cast( # cast-ok: TypedDict events are JSON dictionaries + dict[str, Any], message + ) else: - message_obj = cast(dict[str, Any], json.loads(cast(str, message))) - self._collect_tool_calls_from_response_done(cast(dict, message_obj)) + parsed_message_obj = cast( # cast-ok: parsed realtime events are JSON dictionaries + dict[str, Any], json.loads(message) + ) + if not self._exclude_private_content_from_logs: + self._collect_tool_calls_from_response_done(parsed_message_obj) + message_obj: Final = self._message_for_logging(parsed_message_obj) if not self._should_store_message(message_obj): return try: @@ -273,6 +321,8 @@ class RealTimeStreaming: def _collect_user_input_from_client_event(self, message: str | dict) -> None: """Extract user text content from client WebSocket events for spend logging.""" + if self._exclude_private_content_from_logs: + return try: if isinstance(message, str): msg_obj = json.loads(message) @@ -309,6 +359,8 @@ class RealTimeStreaming: def _collect_user_input_from_backend_event(self, event_obj: dict | OpenAIRealtimeEvents) -> None: """Extract user voice transcription from backend events for spend logging.""" + if self._exclude_private_content_from_logs: + return try: event_type: Final = event_obj.get("type", "") if event_type == "conversation.item.input_audio_transcription.completed": @@ -364,9 +416,9 @@ class RealTimeStreaming: pass def _flush_unbilled_transcription_usage(self) -> None: - if self.provider_config is None: + if self._usage_provider is None: return - usage: Final = self.provider_config.unbilled_usage_on_session_close(self.model) + usage: Final = self._usage_provider.unbilled_usage_on_session_close(self.model) if usage is None: return flush_event: Final = ( @@ -403,12 +455,27 @@ class RealTimeStreaming: except (AttributeError, TypeError): pass + def _input_for_logging( + self, + message: str | dict, # mutable-ok: existing realtime input contract + ) -> str | dict: # mutable-ok: logging stores concrete event dictionaries + if not self._exclude_private_content_from_logs: + return message + try: + parsed_message: Final[object] = message if isinstance(message, dict) else json.loads(message) + except (json.JSONDecodeError, TypeError): + return {} # mutable-ok: empty JSON logging payload + if not isinstance(parsed_message, dict): + return {} # mutable-ok: empty JSON logging payload + return self._message_for_logging(parsed_message) + def store_input(self, message: str | dict): """Store input message""" - self.input_message = message if isinstance(message, dict) else {} + logged_message: Final[str | dict] = self._input_for_logging(message) # mutable-ok: logging payload + self.input_message = logged_message if isinstance(logged_message, dict) else {} self._collect_user_input_from_client_event(message) if self.logging_obj: - self.logging_obj.pre_call(input=message, api_key="") + self.logging_obj.pre_call(input=logged_message, api_key="") async def log_messages(self): """Log messages in list""" @@ -1009,6 +1076,8 @@ class RealTimeStreaming: self.store_message(event_str) self._capture_transcription_usage(event) await self._send_event_to_client(event, event_str) + if self._is_transcription_session: + continue blocked = await self.run_realtime_guardrails( cast(str, transcript), item_id=cast(str | None, event.get("item_id")), diff --git a/litellm/llms/meta/__init__.py b/litellm/llms/meta/__init__.py new file mode 100644 index 00000000000..7c7d32788a2 --- /dev/null +++ b/litellm/llms/meta/__init__.py @@ -0,0 +1,3 @@ +from .realtime import MetaRealtime, MuseRealtimeAdapter + +__all__ = ("MetaRealtime", "MuseRealtimeAdapter") diff --git a/litellm/llms/meta/realtime/__init__.py b/litellm/llms/meta/realtime/__init__.py new file mode 100644 index 00000000000..6398765da24 --- /dev/null +++ b/litellm/llms/meta/realtime/__init__.py @@ -0,0 +1,10 @@ +from .handler import MetaRealtime, MuseRealtimeAdapter +from .transformation import MuseEventTransformer, MuseProtocolError, MuseSessionConfig + +__all__ = ( + "MetaRealtime", + "MuseEventTransformer", + "MuseProtocolError", + "MuseRealtimeAdapter", + "MuseSessionConfig", +) diff --git a/litellm/llms/meta/realtime/handler.py b/litellm/llms/meta/realtime/handler.py new file mode 100644 index 00000000000..9eb98e4ba0b --- /dev/null +++ b/litellm/llms/meta/realtime/handler.py @@ -0,0 +1,661 @@ +from __future__ import annotations + +import asyncio +import base64 +import binascii +import contextlib +import json +import time +import uuid +from collections.abc import Awaitable, Callable, Mapping +from typing import Final, Protocol +from urllib.parse import urlparse, urlunparse + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm._logging import verbose_proxy_logger +from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming +from litellm.llms.custom_httpx.http_handler import get_shared_realtime_ssl_context +from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage, RealtimeQueryParams + +from .transformation import ( + MUSE_MODEL, + MuseEventTransformer, + MuseProtocolError, + MuseSessionConfig, + encode_event, + error_event, + parse_session_update, + session_created_event, + session_updated_event, +) + +DEFAULT_MUSE_REALTIME_URL: Final = "wss://api.meta.ai/v1/asr/realtime" +_MAX_AUDIO_BACKLOG_SECONDS: Final = 4 +_MAX_PENDING_PROVIDER_EVENTS: Final = 256 +_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + + +class _ProviderWebSocket(Protocol): + async def send(self, message: str | bytes) -> None: ... + + async def recv(self, decode: bool | None = None) -> str | bytes: ... + + async def close(self, code: int = 1000, reason: str = "") -> None: ... + + +class _ClientWebSocketExceptions(Protocol): + ConnectionClosed: type[Exception] + + +class _ClientWebSocket(Protocol): + exceptions: _ClientWebSocketExceptions + + @property + def scope(self) -> Mapping[str, object]: ... + + async def send_text(self, data: str) -> None: ... + + async def receive_text(self) -> str: ... + + async def close(self, code: int = 1000, reason: str | None = None) -> None: ... + + +class WebSocketConnect(Protocol): + def __call__( + self, + url: str, + *, + open_timeout: float, + max_size: int | None, + ssl: object | None, + ) -> Awaitable[_ProviderWebSocket]: ... + + +class MuseAdapterError(RuntimeError): + def __init__(self, message: str, *, close_code: int) -> None: + super().__init__(message) + self.close_code: Final = close_code + + +class MuseRealtimeAdapter: + def __init__( + self, + *, + model: str, + api_key: str, + api_base: str | None = None, + timeout: float | None = None, + websocket_connect: WebSocketConnect | None = None, + monotonic: Callable[[], float] = time.monotonic, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + terminate_client: Callable[[int], Awaitable[None]] | None = None, + ) -> None: + if model.removeprefix("meta/") != MUSE_MODEL: + raise ValueError("unsupported Meta realtime model") + self._model: Final = model.removeprefix("meta/") + self._access_token: Final = normalize_access_token(api_key) + self._url: Final = build_muse_realtime_url(api_base) + self._timeout: Final = timeout or 10.0 + self._websocket_connect = websocket_connect + self._monotonic: Final = monotonic + self._sleep: Final = sleep + self._terminate_client: Final = terminate_client + self._provider_ws: _ProviderWebSocket | None = None + self._config: MuseSessionConfig | None = None + self._session_id: str = f"sess_{uuid.uuid4().hex}" + self._events: Final[asyncio.Queue[str | BaseException]] = asyncio.Queue(maxsize=_MAX_PENDING_PROVIDER_EVENTS) + self._events.put_nowait(encode_event(session_created_event(self._model, self._session_id))) + self._transformer: Final = MuseEventTransformer() + self._audio_condition: Final = asyncio.Condition() + self._pending_audio: bytearray = bytearray() + self._audio_generation: int = 0 + self._flush_requested: bool = False + self._end_requested: bool = False + self._end_stream_sent: bool = False + self._audio_consumed: bool = False + self._closed: bool = False + self._resources_closed: bool = False + self._sender_task: asyncio.Task[None] | None = None + self._receiver_task: asyncio.Task[None] | None = None + self.close_code: int = 1000 + self.close_reason: str = "Session closed" + + async def send(self, message: str | bytes) -> None: + if self._closed: + raise MuseAdapterError("Meta Muse realtime session is closed", close_code=self.close_code) + if isinstance(message, bytes): + await self._reject("invalid_request_error", "invalid_event", "Client events must be JSON text") + return + try: + event: Final = _parse_client_event(message) + event_type: Final = event.get("type") + if event_type in ("session.update", "transcription_session.update"): + await self._handle_session_update(message) + return + if event_type == "input_audio_buffer.append": + await self._handle_audio_append(event) + return + if event_type == "input_audio_buffer.clear": + await self._clear_audio() + return + if event_type == "input_audio_buffer.commit": + await self._commit_audio() + return + if event_type == "input_audio_buffer.end": + await self._end_audio() + return + await self._emit( + error_event( + "invalid_request_error", + "unsupported_event", + f"Event type {event_type!r} is not supported for Meta Muse transcription", + ) + ) + except MuseProtocolError as exc: + await self._reject("invalid_request_error", "invalid_event", str(exc)) + + async def recv(self, decode: bool | None = None) -> str | bytes: + event: Final = await self._events.get() + if isinstance(event, BaseException): + close_code: Final = _exception_close_code(event) if isinstance(event, Exception) else 1011 + if self._terminate_client is not None: + await self._terminate_client(close_code) + raise event + return event.encode("utf-8") if decode is False else event + + async def close(self, code: int = 1000, reason: str = "") -> None: + if self._resources_closed: + return + self._closed = True + self._resources_closed = True + self.close_code = sanitize_close_code(code) + self.close_reason = safe_close_reason(self.close_code) + async with self._audio_condition: + self._end_requested = True + self._audio_condition.notify_all() + tasks: Final = tuple(task for task in (self._sender_task, self._receiver_task) if task is not None) + for task in tasks: + task.cancel() + if tasks: + await asyncio.gather(*tasks, return_exceptions=True) + provider_ws: Final = self._provider_ws + if provider_ws is not None: + with contextlib.suppress(Exception): + await provider_ws.close(code=self.close_code, reason=self.close_reason) + + def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None: + return self._transformer.take_unbilled_usage() + + async def _handle_session_update(self, message: str) -> None: + config: Final = parse_session_update(message, self._model) + if self._config is not None: + if config != self._config: + await self._reject( + "invalid_request_error", + "session_configuration_locked", + "Meta Muse session configuration cannot change after setup", + ) + return + await self._emit(session_updated_event(config, self._session_id)) + return + await self._connect(config) + + async def _connect(self, config: MuseSessionConfig) -> None: + connector: Final = self._websocket_connect or _default_websocket_connect + last_error: Exception | None = None # rebind-ok: records the latest bounded handshake attempt + for attempt in range(2): + provider_ws: _ProviderWebSocket | None = None + try: + provider_ws = await connector( + self._url, + open_timeout=self._timeout, + max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, + ssl=_ssl_config(self._url), + ) + await provider_ws.send(json.dumps(config.handshake(self._access_token), separators=(",", ":"))) + raw_ack: str | bytes = await asyncio.wait_for( # rebind-ok: one response per handshake attempt + provider_ws.recv(), timeout=self._timeout + ) + session_id: str = _parse_handshake_ack(raw_ack) # rebind-ok: one ID per handshake attempt + self._provider_ws = provider_ws + self._config = config + self._transformer.configure(config) + self._session_id = session_id + self._sender_task = asyncio.create_task(self._send_audio(), name="meta-muse-realtime-send") + self._receiver_task = asyncio.create_task(self._receive_events(), name="meta-muse-realtime-receive") + await self._emit(session_updated_event(config, session_id)) + return + except asyncio.CancelledError: + if provider_ws is not None: + with contextlib.suppress(Exception): + await provider_ws.close() + raise + except Exception as exc: # noqa: BLE001 # connector implementations expose heterogeneous transport errors + last_error = exc + if provider_ws is not None: + with contextlib.suppress(Exception): + await provider_ws.close() + close_code: int = _exception_close_code(exc) # rebind-ok: classified per handshake attempt + retryable_transport_error: bool = not isinstance( # rebind-ok: classified per handshake attempt + exc, (MuseAdapterError, MuseProtocolError) + ) + if attempt == 0 and retryable_transport_error and close_code in (1011, 1013): + continue + self.close_code = close_code + self.close_reason = safe_close_reason(close_code) + await self._emit( + error_event( + "server_error" if close_code != 1008 else "invalid_request_error", + "provider_connection_error", + "Meta Muse realtime handshake failed", + ) + ) + await self._events.put(MuseAdapterError("Meta Muse realtime handshake failed", close_code=close_code)) + await self._mark_terminated(close_code) + return + assert last_error is not None + raise MuseAdapterError("Meta Muse realtime handshake failed", close_code=1011) + + async def _handle_audio_append(self, event: Mapping[str, JsonValue]) -> None: + config: Final = self._require_configured() + if self._end_requested or self._end_stream_sent: + await self._reject("invalid_request_error", "input_ended", "Audio input has already ended") + return + audio_value: Final = event.get("audio") + if not isinstance(audio_value, str): + await self._reject("invalid_request_error", "invalid_audio", "Audio must be a base64 string") + return + try: + audio: Final = base64.b64decode(audio_value, validate=True) + except (binascii.Error, ValueError): + await self._reject("invalid_request_error", "invalid_audio", "Audio must be valid base64") + return + if len(audio) % 2: + await self._reject("invalid_request_error", "invalid_audio", "PCM16 audio must contain complete samples") + return + if not audio: + return + max_backlog_bytes: Final = config.bytes_per_second * _MAX_AUDIO_BACKLOG_SECONDS + if len(audio) > max_backlog_bytes: + await self._reject( + "invalid_request_error", + "audio_backlog_exceeded", + "Audio append exceeds the four-second Muse backlog limit", + ) + return + async with self._audio_condition: + await self._audio_condition.wait_for( + lambda: self._closed or len(self._pending_audio) + len(audio) <= max_backlog_bytes + ) + if self._closed: + raise MuseAdapterError("Meta Muse realtime session is closed", close_code=self.close_code) + self._pending_audio.extend(audio) + self._audio_condition.notify_all() + + async def _clear_audio(self) -> None: + self._require_configured() + async with self._audio_condition: + self._pending_audio.clear() + self._audio_generation += 1 + self._flush_requested = False + self._audio_condition.notify_all() + await self._emit( + { # mutable-ok: OpenAI-compatible JSON event + "type": "input_audio_buffer.cleared", + "event_id": f"event_{uuid.uuid4().hex}", + } + ) + + async def _commit_audio(self) -> None: + config: Final = self._require_configured() + previous_item_id, item_id = self._transformer.commit_item() + async with self._audio_condition: + self._flush_requested = True + if config.mode == "PUSH_TO_TALK": + self._end_requested = True + self._audio_condition.notify_all() + await self._emit( + { # mutable-ok: OpenAI-compatible JSON event + "type": "input_audio_buffer.committed", + "event_id": f"event_{uuid.uuid4().hex}", + "previous_item_id": previous_item_id, + "item_id": item_id, + } + ) + + async def _end_audio(self) -> None: + self._require_configured() + async with self._audio_condition: + self._flush_requested = True + self._end_requested = True + self._audio_condition.notify_all() + + async def _send_audio(self) -> None: + config: Final = self._require_configured() + provider_ws: Final = self._require_provider_ws() + pacing_origin: float | None = None # rebind-ok: initialized when the first packet is ready + sent_duration: float = 0.0 # rebind-ok: absolute pacing clock advances after each packet + try: + while True: + packet, pacing_origin, ended = await self._next_audio_packet( + config, + pacing_origin, + sent_duration, + ) + if ended: + break + if packet is None: + continue + await provider_ws.send(packet) + self._audio_consumed = True + sent_duration += len(packet) / config.bytes_per_second + await self._send_end_stream() + except asyncio.CancelledError: + raise + except Exception as exc: # noqa: BLE001 # WebSocket implementations expose heterogeneous transport errors + await self._fail_provider(exc, phase="audio send") + + async def _next_audio_packet( + self, + config: MuseSessionConfig, + pacing_origin: float | None, + sent_duration: float, + ) -> tuple[bytes | None, float | None, bool]: + async with self._audio_condition: + await self._audio_condition.wait_for( + lambda: ( + self._closed + or len(self._pending_audio) >= config.packet_bytes + or (self._flush_requested and bool(self._pending_audio)) + or (self._end_requested and not self._pending_audio) + ) + ) + if self._closed or (self._end_requested and not self._pending_audio): + return None, pacing_origin, True + packet_size: Final = min(config.packet_bytes, len(self._pending_audio)) + if packet_size < config.packet_bytes and not self._flush_requested: + return None, pacing_origin, False + generation: Final = self._audio_generation + current_time: Final = self._monotonic() + effective_origin: Final = ( + current_time - sent_duration + if pacing_origin is None or current_time > pacing_origin + sent_duration + else pacing_origin + ) + deadline: Final = effective_origin + sent_duration + delay: Final = deadline - self._monotonic() + if delay > 0: + await self._sleep(delay) + async with self._audio_condition: + if generation != self._audio_generation: + return None, effective_origin, False + actual_size: Final = min(packet_size, len(self._pending_audio)) + packet: Final = bytes(self._pending_audio[:actual_size]) + del self._pending_audio[:actual_size] + if not self._pending_audio: + self._flush_requested = False + self._audio_condition.notify_all() + return packet or None, effective_origin, False + + async def _send_end_stream(self) -> None: + if self._end_stream_sent: + return + provider_ws: Final = self._require_provider_ws() + await provider_ws.send('{"type":"endStream"}') + self._end_stream_sent = True + + async def _receive_events(self) -> None: + provider_ws: Final = self._require_provider_ws() + try: + while True: + raw: str | bytes = await provider_ws.recv() # rebind-ok: one provider frame per iteration + if not isinstance(raw, str): + raise MuseProtocolError("provider returned a non-text event") + for event in self._transformer.transform(raw): + await self._emit(event) + except asyncio.CancelledError: + raise + except Exception as exc: # noqa: BLE001 # provider close exceptions vary by WebSocket implementation + close_code: Final = _exception_close_code(exc) + if close_code == 1000 and self._end_stream_sent: + await self._mark_terminated(1000) + await self._events.put(MuseAdapterError("Meta Muse realtime session completed", close_code=1000)) + return + failure: Final = MuseAdapterError( + "Meta Muse realtime closed before input ended", + close_code=1011 if close_code == 1000 else close_code, + ) + await self._fail_provider(failure, phase="receive") + + async def _fail_provider(self, exc: Exception, *, phase: str) -> None: + close_code: Final = _exception_close_code(exc) + self.close_code = close_code + self.close_reason = safe_close_reason(close_code) + await self._emit( + error_event( + "server_error", + "provider_connection_error", + f"Meta Muse realtime {phase} failed", + ) + ) + await self._events.put(MuseAdapterError(f"Meta Muse realtime {phase} failed", close_code=close_code)) + await self._mark_terminated(close_code) + + async def _mark_terminated(self, close_code: int) -> None: + self._closed = True + self.close_code = sanitize_close_code(close_code) + self.close_reason = safe_close_reason(self.close_code) + async with self._audio_condition: + self._audio_condition.notify_all() + + async def _terminate(self, close_code: int) -> None: + await self._mark_terminated(close_code) + if self._terminate_client is not None: + await self._terminate_client(self.close_code) + + async def _reject(self, error_type: str, code: str, message: str) -> None: + self.close_code = 1008 + self.close_reason = safe_close_reason(1008) + await self._emit(error_event(error_type, code, message)) + await self._events.put(MuseAdapterError(message, close_code=1008)) + await self._mark_terminated(1008) + + async def _emit(self, event: Mapping[str, object]) -> None: + await self._events.put(encode_event(event)) + + def _require_configured(self) -> MuseSessionConfig: + if self._config is None: + raise MuseProtocolError("send session.update before audio events") + return self._config + + def _require_provider_ws(self) -> _ProviderWebSocket: + if self._provider_ws is None: + raise MuseProtocolError("Meta Muse provider connection is not ready") + return self._provider_ws + + +class MetaRealtime: + async def async_realtime( + self, + model: str, + websocket: _ClientWebSocket, + logging_obj: LiteLLMLogging, + api_base: str | None = None, + api_key: str | None = None, + client: object | None = None, + timeout: float | None = None, + query_params: RealtimeQueryParams | None = None, + user_api_key_dict: object | None = None, + litellm_metadata: Mapping[str, object] | None = None, + websocket_connect: WebSocketConnect | None = None, + monotonic: Callable[[], float] = time.monotonic, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + **kwargs: object, # kwargs-ok: realtime dispatcher forwards provider-neutral options + ) -> None: + if api_key is None or not api_key.strip(): + await _send_client_error_and_close(websocket, "Meta Model API key is required") + return + try: + adapter: Final = MuseRealtimeAdapter( + model=model, + api_key=api_key, + api_base=api_base, + timeout=timeout, + websocket_connect=websocket_connect, + monotonic=monotonic, + sleep=sleep, + terminate_client=lambda code: _close_client(websocket, code), + ) + except ValueError: + await _send_client_error_and_close(websocket, "Invalid Meta Muse realtime configuration") + return + realtime_streaming: Final = RealTimeStreaming( + websocket, + adapter, # pyright: ignore[reportArgumentType] # raw adapter intentionally matches the websocket surface + logging_obj, + model=model, + user_api_key_dict=user_api_key_dict, + request_data={ # mutable-ok: relay request metadata payload + "litellm_metadata": dict(litellm_metadata or {}) # mutable-ok: relay owns its metadata copy + }, + force_transcription_model=model, + usage_provider=adapter, + exclude_private_content_from_logs=True, + ) + try: + await realtime_streaming.bidirectional_forward() + except MuseAdapterError as exc: + adapter.close_code = exc.close_code + adapter.close_reason = safe_close_reason(exc.close_code) + except Exception: # noqa: BLE001 # relay errors are normalized before closing the accepted client socket + adapter.close_code = 1011 + adapter.close_reason = safe_close_reason(1011) + verbose_proxy_logger.exception("Meta Muse realtime session failed") + finally: + await adapter.close(code=adapter.close_code) + await _close_client(websocket, adapter.close_code) + + +def normalize_access_token(api_key: str) -> str: + stripped: Final = api_key.strip() + if not stripped: + raise ValueError("Meta Model API key is required") + parts: Final = stripped.split(None, 1) + if parts[0].casefold() == "bearer": + if len(parts) != 2 or not parts[1].strip(): + raise ValueError("Meta Model API key must include a token after Bearer") + return f"Bearer {parts[1].strip()}" + return f"Bearer {stripped}" + + +def build_muse_realtime_url(api_base: str | None) -> str: + if api_base is None: + return DEFAULT_MUSE_REALTIME_URL + parsed: Final = urlparse(api_base.strip()) + scheme: Final = "wss" if parsed.scheme == "https" else parsed.scheme + if ( + scheme != "wss" + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or parsed.fragment + ): + raise ValueError("Meta api_base must be an absolute wss:// or https:// URL without credentials or a fragment") + netloc: Final = f"{parsed.hostname}:{parsed.port}" if parsed.port is not None else parsed.hostname + return urlunparse((scheme, netloc, "/v1/asr/realtime", "", "", "")) + + +def sanitize_close_code(code: int | None) -> int: + if code is not None and code in (1000, 1008, 1011, 1013): + return code + return 1011 + + +def safe_close_reason(code: int) -> str: + return { # mutable-ok: immutable-by-convention close-reason lookup + 1000: "Session closed", + 1008: "Invalid realtime transcription request", + 1011: "Realtime transcription service error", + 1013: "Realtime transcription service unavailable", + }.get(code, "Realtime transcription service error") + + +def _parse_client_event(payload: str) -> Mapping[str, JsonValue]: + try: + value: Final = _JSON_ADAPTER.validate_json(payload) + except ValidationError: + raise MuseProtocolError("invalid JSON object") from None + if not isinstance(value, dict): + raise MuseProtocolError("message must be a JSON object") + event_type: Final = value.get("type") + if not isinstance(event_type, str) or not event_type: + raise MuseProtocolError("message type must be a non-empty string") + return value + + +def _parse_handshake_ack(raw: str | bytes) -> str: + if not isinstance(raw, str): + raise MuseProtocolError("provider returned a non-text handshake response") + message: Final = _parse_json_object(raw) + if message.get("type") == "error": + raise MuseAdapterError("Meta Muse realtime handshake was rejected", close_code=1008) + session_id: Final = message.get("sessionId") + if not isinstance(session_id, str) or not session_id.strip(): + raise MuseProtocolError("provider returned an invalid handshake response") + return session_id.strip() + + +def _parse_json_object(payload: str) -> Mapping[str, JsonValue]: + try: + value: Final = _JSON_ADAPTER.validate_json(payload) + except ValidationError: + raise MuseProtocolError("invalid provider JSON object") from None + if not isinstance(value, dict): + raise MuseProtocolError("provider message must be a JSON object") + return value + + +def _exception_close_code(exc: Exception) -> int: + code: Final = getattr(exc, "code", None) + if isinstance(exc, MuseAdapterError): + return sanitize_close_code(exc.close_code) + return sanitize_close_code(code if isinstance(code, int) else None) + + +def _ssl_config(url: str) -> object | None: + if not url.startswith("wss://"): + return None + config: Final = get_shared_realtime_ssl_context() + return True if config is False else config + + +async def _default_websocket_connect( + url: str, + *, + open_timeout: float, + max_size: int | None, + ssl: object | None, +) -> _ProviderWebSocket: + import websockets + + connection: Final = await websockets.connect( + url, + open_timeout=open_timeout, + max_size=max_size, + ssl=ssl, # pyright: ignore[reportArgumentType] # shared SSL helper returns the library-supported union + ) + return connection + + +async def _send_client_error_and_close(websocket: _ClientWebSocket, message: str) -> None: + with contextlib.suppress(Exception): + await websocket.send_text(encode_event(error_event("invalid_request_error", "invalid_configuration", message))) + await _close_client(websocket, 1008) + + +async def _close_client(websocket: _ClientWebSocket, code: int) -> None: + with contextlib.suppress(Exception): + await websocket.close(code=sanitize_close_code(code), reason=safe_close_reason(sanitize_close_code(code))) diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py new file mode 100644 index 00000000000..f37295892f8 --- /dev/null +++ b/litellm/llms/meta/realtime/transformation.py @@ -0,0 +1,619 @@ +from __future__ import annotations + +import json +import math +import uuid +from collections import OrderedDict, deque +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final, Literal, TypeAlias + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage + +MUSE_MODEL: Final = "muse-voice-transcribe-1.0" +SUPPORTED_SAMPLE_RATES: Final = frozenset((16_000, 24_000)) +SUPPORTED_MODES: Final = frozenset(("PUSH_TO_TALK", "ENDPOINTING", "DIARIZATION")) +SUPPORTED_LANGUAGES: Final = ( + "Arabic", + "Bengali", + "Dutch", + "English", + "French", + "German", + "Hebrew", + "Hindi", + "Indonesian", + "Italian", + "Japanese", + "Kannada", + "Korean", + "Malay", + "Mandarin Chinese", + "Marathi", + "Polish", + "Portuguese", + "Spanish", + "Tagalog", + "Tamil", + "Telugu", + "Thai", + "Turkish", + "Vietnamese", +) +_LANGUAGE_NAMES: Final = { # mutable-ok: immutable-by-convention language lookup table + language.casefold(): language for language in SUPPORTED_LANGUAGES +} +_LANGUAGE_CODES: Final = { # mutable-ok: immutable-by-convention language lookup table + "ar": "Arabic", + "bn": "Bengali", + "de": "German", + "en": "English", + "es": "Spanish", + "fil": "Tagalog", + "fr": "French", + "he": "Hebrew", + "hi": "Hindi", + "id": "Indonesian", + "it": "Italian", + "iw": "Hebrew", + "ja": "Japanese", + "kn": "Kannada", + "ko": "Korean", + "ms": "Malay", + "mr": "Marathi", + "nl": "Dutch", + "pl": "Polish", + "pt": "Portuguese", + "ta": "Tamil", + "te": "Telugu", + "th": "Thai", + "tl": "Tagalog", + "tr": "Turkish", + "vi": "Vietnamese", + "zh": "Mandarin Chinese", +} +_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +OpenAIEvent: TypeAlias = Mapping[str, object] + + +class MuseProtocolError(ValueError): + pass + + +@dataclass(frozen=True, slots=True) +class MuseSessionConfig: + model: str + mode: Literal["PUSH_TO_TALK", "ENDPOINTING", "DIARIZATION"] + sample_rate: Literal[16000, 24000] + keywords: tuple[str, ...] + language_bias: tuple[str, ...] + + @property + def audio_encoding(self) -> Literal["PCM_16KHZ", "PCM_24KHZ"]: + return "PCM_16KHZ" if self.sample_rate == 16_000 else "PCM_24KHZ" + + @property + def bytes_per_second(self) -> int: + return self.sample_rate * 2 + + @property + def packet_bytes(self) -> int: + return self.bytes_per_second * 80 // 1000 + + def handshake(self, access_token: str) -> Mapping[str, object]: + base: Final[Mapping[str, object]] = { # mutable-ok: JSON wire payload + "mode": self.mode, + "authorization": {"accessToken": access_token}, # mutable-ok: JSON wire payload + "audioEncoding": self.audio_encoding, + "model": self.model, + "partialMode": "CUMULATIVE", + "emitAudioProgress": True, + } + payload: dict[str, object] = dict(base) # mutable-ok: incrementally builds JSON wire payload + if self.keywords: + payload["keywords"] = list(self.keywords) # mutable-ok: JSON arrays require concrete lists + if self.language_bias: + payload["languageBias"] = list(self.language_bias) # mutable-ok: JSON arrays require concrete lists + return payload + + def openai_session(self, session_id: str) -> Mapping[str, object]: + turn_detection: Final[Mapping[str, object] | None] = ( + None if self.mode == "PUSH_TO_TALK" else {"type": "server_vad"} # mutable-ok: JSON wire payload + ) + transcription: dict[str, object] = { # mutable-ok: incrementally builds JSON wire payload + "model": self.model, + } + if self.language_bias: + transcription["language"] = self.language_bias[0] + transcription["language_bias"] = list( # mutable-ok: JSON arrays require concrete lists + self.language_bias + ) + if self.keywords: + transcription["keywords"] = list(self.keywords) # mutable-ok: JSON arrays require concrete lists + return { # mutable-ok: JSON wire payload + "id": session_id, + "object": "realtime.transcription_session", + "type": "transcription", + "model": self.model, + "audio": { # mutable-ok: JSON wire payload + "input": { # mutable-ok: JSON wire payload + "format": {"type": "audio/pcm", "rate": self.sample_rate}, # mutable-ok: JSON wire payload + "transcription": transcription, + "turn_detection": turn_detection, + } + }, + } + + +@dataclass(slots=True) +class _TurnState: + item_id: str | None = None + started: bool = False + start_emitted: bool = False + latest_partial: str | None = None + emitted_partial: str = "" + final_text: str | None = None + completed_signal: bool = False + completed_emitted: bool = False + stopped: bool = False + stopped_emitted: bool = False + speaker: str | None = None + + +def _json_object(payload: str) -> Mapping[str, JsonValue]: + try: + value: Final = _JSON_ADAPTER.validate_json(payload) + except ValidationError: + raise MuseProtocolError("invalid JSON object") from None + if not isinstance(value, dict): + raise MuseProtocolError("message must be a JSON object") + return value + + +def _mapping(value: JsonValue | None, name: str) -> Mapping[str, JsonValue]: + if value is None: + return {} # mutable-ok: empty JSON object + if not isinstance(value, dict): + raise MuseProtocolError(f"{name} must be an object") + return value + + +def _string(value: JsonValue | None, name: str) -> str | None: + if value is None: + return None + if not isinstance(value, str): + raise MuseProtocolError(f"{name} must be a string") + return value + + +def _normalize_model(model: str) -> str: + return model.removeprefix("meta/").strip() + + +def normalize_language(language: str) -> str: + value: Final = language.strip() + if not value: + raise MuseProtocolError("language must be non-empty") + documented_name: Final = _LANGUAGE_NAMES.get(value.casefold()) + if documented_name is not None: + return documented_name + primary: Final = value.replace("_", "-").split("-", 1)[0].casefold() + mapped_name: Final = _LANGUAGE_CODES.get(primary) + if mapped_name is None: + raise MuseProtocolError("unsupported Muse Voice language") + return mapped_name + + +def _normalize_string_sequence(value: JsonValue | None, name: str) -> tuple[str, ...]: + if value is None: + return () + if not isinstance(value, list): + raise MuseProtocolError(f"{name} must be an array of strings") + normalized: list[str] = [] # mutable-ok: deduplicates validated language hints before freezing + for entry in value: + if not isinstance(entry, str) or not entry.strip(): + raise MuseProtocolError(f"{name} entries must be non-empty strings") + item: str = entry.strip() # rebind-ok: normalized once for each hint + if item not in normalized: + normalized.append(item) + return tuple(normalized) + + +def _normalize_language_sequence(value: JsonValue | None) -> tuple[str, ...]: + return tuple(dict.fromkeys(normalize_language(item) for item in _normalize_string_sequence(value, "language_bias"))) + + +def _parse_sample_rate(session: Mapping[str, JsonValue]) -> Literal[16000, 24000]: + beta_format: Final = session.get("input_audio_format") + audio: Final = _mapping(session.get("audio"), "session.audio") + audio_input: Final = _mapping(audio.get("input"), "session.audio.input") + ga_format: Final = audio_input.get("format") + if beta_format is not None and ga_format is not None: + raise MuseProtocolError("input audio format must use either beta or GA layout") + if beta_format is not None: + if beta_format != "pcm16": + raise MuseProtocolError("Muse Voice requires pcm16 input audio") + return 24_000 + if ga_format is None: + return 24_000 + if isinstance(ga_format, str): + if ga_format != "pcm16": + raise MuseProtocolError("Muse Voice requires audio/pcm input audio") + return 24_000 + format_mapping: Final = _mapping(ga_format, "session.audio.input.format") + if format_mapping.get("type") != "audio/pcm": + raise MuseProtocolError("Muse Voice requires audio/pcm input audio") + channels: Final = format_mapping.get("channels", 1) + if isinstance(channels, bool) or channels != 1: + raise MuseProtocolError("Muse Voice requires mono input audio") + rate: Final = format_mapping.get("rate", 24_000) + if isinstance(rate, bool) or not isinstance(rate, int) or rate not in SUPPORTED_SAMPLE_RATES: + raise MuseProtocolError("Muse Voice supports PCM16 at 16000 Hz or 24000 Hz") + return rate + + +def _parse_mode( + session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue] +) -> Literal["PUSH_TO_TALK", "ENDPOINTING", "DIARIZATION"]: + explicit: Final = session.get("mode") + if explicit is not None: + if not isinstance(explicit, str) or explicit.upper() not in SUPPORTED_MODES: + raise MuseProtocolError("unsupported Muse Voice mode") + normalized_mode: Final = explicit.upper() + if normalized_mode == "PUSH_TO_TALK": + return "PUSH_TO_TALK" + if normalized_mode == "DIARIZATION": + return "DIARIZATION" + return "ENDPOINTING" + turn_detection_present: Final = "turn_detection" in session or "turn_detection" in audio_input + turn_detection: Final = session.get("turn_detection", audio_input.get("turn_detection")) + if turn_detection_present and turn_detection is None: + return "PUSH_TO_TALK" + if turn_detection is None: + return "ENDPOINTING" + turn_detection_mapping: Final = _mapping(turn_detection, "turn_detection") + if turn_detection_mapping.get("type") not in (None, "server_vad"): + raise MuseProtocolError("Muse Voice supports server_vad turn detection or null") + return "ENDPOINTING" + + +def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig: + message: Final = _json_object(payload) + if message.get("type") not in ("session.update", "transcription_session.update"): + raise MuseProtocolError("expected session.update") + session: Final = _mapping(message.get("session"), "session") + if not session: + raise MuseProtocolError("session.update requires a session object") + session_type: Final = session.get("type") + if session_type not in (None, "transcription", "realtime"): + raise MuseProtocolError("Muse Voice supports transcription sessions only") + audio: Final = _mapping(session.get("audio"), "session.audio") + audio_input: Final = _mapping(audio.get("input"), "session.audio.input") + beta_transcription: Final = session.get("input_audio_transcription") + ga_transcription: Final = audio_input.get("transcription") + if beta_transcription is not None and ga_transcription is not None: + raise MuseProtocolError("input transcription must use either beta or GA layout") + transcription: Final = _mapping( + beta_transcription if beta_transcription is not None else ga_transcription, + "input audio transcription", + ) + requested_model: Final = _string(transcription.get("model"), "transcription model") + normalized_model: Final = _normalize_model(expected_model) + if normalized_model != MUSE_MODEL: + raise MuseProtocolError("unsupported Meta realtime model") + if requested_model is not None and _normalize_model(requested_model) != normalized_model: + raise MuseProtocolError("realtime session model cannot be changed") + language_value: Final = _string(transcription.get("language"), "language") + explicit_bias: Final = _normalize_language_sequence(transcription.get("language_bias")) + language_bias: Final = tuple( + dict.fromkeys((normalize_language(language_value), *explicit_bias)) + if language_value is not None + else explicit_bias + ) + keywords: Final = _normalize_string_sequence(transcription.get("keywords"), "keywords") + return MuseSessionConfig( + model=normalized_model, + mode=_parse_mode(session, audio_input), + sample_rate=_parse_sample_rate(session), + keywords=keywords, + language_bias=language_bias, + ) + + +def session_created_event(model: str, session_id: str) -> OpenAIEvent: + normalized_model: Final = _normalize_model(model) + default_config: Final = MuseSessionConfig( + model=normalized_model, + mode="ENDPOINTING", + sample_rate=24_000, + keywords=(), + language_bias=(), + ) + return { # mutable-ok: OpenAI-compatible JSON event + "type": "session.created", + "event_id": f"event_{uuid.uuid4().hex}", + "session": default_config.openai_session(session_id), + } + + +def session_updated_event(config: MuseSessionConfig, session_id: str) -> OpenAIEvent: + return { # mutable-ok: OpenAI-compatible JSON event + "type": "session.updated", + "event_id": f"event_{uuid.uuid4().hex}", + "session": config.openai_session(session_id), + } + + +def error_event(error_type: str, code: str, message: str) -> OpenAIEvent: + return { # mutable-ok: OpenAI-compatible JSON event + "type": "error", + "event_id": f"event_{uuid.uuid4().hex}", + "error": { # mutable-ok: nested OpenAI-compatible error object + "type": error_type, + "code": code, + "message": message, + }, + } + + +class MuseEventTransformer: + def __init__(self, *, completed_turn_limit: int = 128) -> None: + self._turns: OrderedDict[str, _TurnState] = OrderedDict() # mutable-ok: ordered active-turn state + self._active_turn_id: str | None = None + self._mode: Literal["PUSH_TO_TALK", "ENDPOINTING", "DIARIZATION"] = "ENDPOINTING" + self._completed_turn_ids: set[str] = set() # mutable-ok: bounded completed-turn membership + self._completed_turn_order: deque[str] = deque( # mutable-ok: bounded completion eviction order + maxlen=completed_turn_limit + ) + self._completed_turn_limit: Final = completed_turn_limit + self._pending_item_ids: deque[str] = deque() # mutable-ok: FIFO commit correlation state + self._last_committed_item_id: str | None = None + self._last_audio_processed_ms: float = 0.0 + self._unassigned_usage_seconds: float = 0.0 + + def configure(self, config: MuseSessionConfig) -> None: + self._mode = config.mode + + def transform(self, payload: str) -> tuple[OpenAIEvent, ...]: + message: Final = _json_object(payload) + event_type: Final = message.get("type") + if event_type == "error": + return (error_event("server_error", "provider_error", "Meta Muse realtime transcription failed"),) + if event_type == "audioProgress": + self._update_audio_progress(message) + return () + if event_type == "speechStart": + self._speech_start(message) + elif event_type == "transcript": + self._transcript(message) + elif event_type == "speaker": + self._speaker(message) + elif event_type == "speechEnd": + self._speech_end(message) + elif event_type == "speechComplete": + self._speech_complete(message) + else: + return () + return self._drain() + + def commit_item(self) -> tuple[str | None, str]: + previous_item_id: Final = self._last_committed_item_id + provider_turn_id: Final = self._active_turn_id + active_turn: Final = self._turns.get(provider_turn_id) if provider_turn_id is not None else None + item_id: Final = ( + active_turn.item_id or provider_turn_id + if active_turn is not None and provider_turn_id is not None + else f"item_{uuid.uuid4().hex}" + ) + if active_turn is not None: + active_turn.item_id = item_id + else: + self._pending_item_ids.append(item_id) + self._last_committed_item_id = item_id + return previous_item_id, item_id + + def take_unbilled_usage(self) -> RealtimeInputAudioTranscriptionUsage | None: + seconds: Final = self._unassigned_usage_seconds + if seconds <= 0: + return None + self._unassigned_usage_seconds = 0.0 + return {"type": "duration", "seconds": seconds} # mutable-ok: typed usage wire payload + + def _turn(self, turn_id: str) -> _TurnState: + if turn_id in self._completed_turn_ids: + raise _CompletedTurn + turn: Final = self._turns.get(turn_id) + if turn is not None: + return turn + created: Final = _TurnState(item_id=self._pending_item_ids.popleft() if self._pending_item_ids else turn_id) + self._turns[turn_id] = created + return created + + def _speech_start(self, message: Mapping[str, JsonValue]) -> None: + turn_id: Final = self._required_turn_id(message, "speechStart") + try: + turn: Final = self._turn(turn_id) + except _CompletedTurn: + return + turn.started = True + self._active_turn_id = turn_id + + def _transcript(self, message: Mapping[str, JsonValue]) -> None: + transcript: Final = message.get("transcript") + if not isinstance(transcript, str): + raise MuseProtocolError("transcript event has invalid transcript") + if not transcript and message.get("turnId") is None and self._active_turn_id is None: + return + turn_id: Final = self._transcript_turn_id(message) + try: + turn: Final = self._turn(turn_id) + except _CompletedTurn: + return + final: Final = message.get("final") is True + if final: + turn.final_text = transcript + turn.completed_signal = True + if self._mode == "PUSH_TO_TALK": + turn.stopped = True + if self._active_turn_id == turn_id: + self._active_turn_id = None + return + if turn.final_text is None: + turn.latest_partial = transcript + + def _speaker(self, message: Mapping[str, JsonValue]) -> None: + turn_id: Final = ( + self._required_turn_id(message, "speaker") if message.get("turnId") is not None else self._active_turn_id + ) + if turn_id is None: + raise MuseProtocolError("speaker event arrived outside an active turn") + label: Final = message.get("label") + if not isinstance(label, str) or not label.strip(): + raise MuseProtocolError("speaker event has invalid label") + try: + turn: Final = self._turn(turn_id) + except _CompletedTurn: + return + turn.speaker = label.strip() + + def _speech_end(self, message: Mapping[str, JsonValue]) -> None: + turn_id: Final = self._required_turn_id(message, "speechEnd") + try: + turn: Final = self._turn(turn_id) + except _CompletedTurn: + return + turn.stopped = True + if self._active_turn_id == turn_id: + self._active_turn_id = None + + def _speech_complete(self, message: Mapping[str, JsonValue]) -> None: + turn_id: Final = self._required_turn_id(message, "speechComplete") + transcript: Final = message.get("transcript") + if not isinstance(transcript, str): + raise MuseProtocolError("speechComplete event has invalid transcript") + try: + turn: Final = self._turn(turn_id) + except _CompletedTurn: + return + turn.final_text = transcript + turn.completed_signal = True + + def _update_audio_progress(self, message: Mapping[str, JsonValue]) -> None: + processed_ms: Final = message.get("audioProcessedMs") + if ( + isinstance(processed_ms, bool) + or not isinstance(processed_ms, (int, float)) + or not math.isfinite(processed_ms) + or processed_ms < 0 + ): + raise MuseProtocolError("audioProgress event has invalid audioProcessedMs") + if processed_ms <= self._last_audio_processed_ms: + return + self._unassigned_usage_seconds += (float(processed_ms) - self._last_audio_processed_ms) / 1000 + self._last_audio_processed_ms = float(processed_ms) + + def _drain(self) -> tuple[OpenAIEvent, ...]: + events: list[OpenAIEvent] = [] # mutable-ok: ordered events are frozen to a tuple before return + while self._turns: + turn_id: str = next(iter(self._turns)) # rebind-ok: selects the next ordered turn + turn: _TurnState = self._turns[turn_id] # rebind-ok: state for the selected turn + has_content: bool = ( # rebind-ok: evaluated for the selected turn + turn.latest_partial is not None or turn.final_text is not None + ) + item_id: str = turn.item_id or turn_id # rebind-ok: selected for each ordered turn + if (turn.started or has_content) and not turn.start_emitted: + turn.start_emitted = True + events.append(self._speech_event("input_audio_buffer.speech_started", item_id)) + if turn.latest_partial is not None and turn.final_text is None: + delta: str = self._new_suffix( # rebind-ok: computed for the selected turn + turn.emitted_partial, turn.latest_partial + ) + if delta: + turn.emitted_partial = turn.latest_partial + events.append( + { # mutable-ok: OpenAI-compatible JSON event + "type": "conversation.item.input_audio_transcription.delta", + "event_id": f"event_{uuid.uuid4().hex}", + "item_id": item_id, + "content_index": 0, + "delta": delta, + } + ) + if turn.stopped and not turn.stopped_emitted: + turn.stopped_emitted = True + events.append(self._speech_event("input_audio_buffer.speech_stopped", item_id)) + if turn.final_text is not None and turn.stopped_emitted and not turn.completed_emitted: + turn.completed_emitted = True + usage: RealtimeInputAudioTranscriptionUsage | None = ( # rebind-ok: usage assigned per turn + self.take_unbilled_usage() + ) + completed_event: dict[str, object] = { # mutable-ok: incrementally builds OpenAI JSON event + "type": "conversation.item.input_audio_transcription.completed", + "event_id": f"event_{uuid.uuid4().hex}", + "item_id": item_id, + "content_index": 0, + "transcript": turn.final_text, + } + if turn.speaker is not None: + completed_event["speaker"] = turn.speaker + if usage is not None: + completed_event["usage"] = usage + events.append(completed_event) + if not (turn.completed_emitted and (turn.stopped or turn.completed_signal)): + break + del self._turns[turn_id] + self._remember_completed(turn_id) + return tuple(events) + + def _remember_completed(self, turn_id: str) -> None: + if turn_id in self._completed_turn_ids: + return + if len(self._completed_turn_order) >= self._completed_turn_limit: + self._completed_turn_ids.discard(self._completed_turn_order.popleft()) + self._completed_turn_order.append(turn_id) + self._completed_turn_ids.add(turn_id) + + def _transcript_turn_id(self, message: Mapping[str, JsonValue]) -> str: + if message.get("turnId") is not None: + return self._required_turn_id(message, "transcript") + if self._active_turn_id is not None: + return self._active_turn_id + if self._mode != "PUSH_TO_TALK": + raise MuseProtocolError("transcript event is missing turnId outside an active turn") + turn_id: Final = f"item_{uuid.uuid4().hex}" + self._active_turn_id = turn_id + return turn_id + + @staticmethod + def _required_turn_id(message: Mapping[str, JsonValue], event: str) -> str: + value: Final = message.get("turnId") + if isinstance(value, bool) or not isinstance(value, (str, int)): + raise MuseProtocolError(f"{event} event has invalid turnId") + turn_id: Final = str(value).strip() + if not turn_id: + raise MuseProtocolError(f"{event} event has invalid turnId") + return turn_id + + @staticmethod + def _speech_event(event_type: str, turn_id: str) -> OpenAIEvent: + return { # mutable-ok: OpenAI-compatible JSON event + "type": event_type, + "event_id": f"event_{uuid.uuid4().hex}", + "item_id": turn_id, + } + + @staticmethod + def _new_suffix(previous: str, current: str) -> str: + if current.startswith(previous): + return current[len(previous) :] + return "" + + +class _CompletedTurn(Exception): + pass + + +def encode_event(event: Mapping[str, object]) -> str: + return json.dumps(event, separators=(",", ":")) diff --git a/litellm/llms/openai_like/providers.json b/litellm/llms/openai_like/providers.json index a458a209ea9..fe10293c420 100644 --- a/litellm/llms/openai_like/providers.json +++ b/litellm/llms/openai_like/providers.json @@ -173,7 +173,7 @@ "api_key_env": "META_API_KEY", "api_base_env": "META_API_BASE", "base_class": "openai_gpt", - "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages"] + "supported_endpoints": ["/v1/chat/completions", "/v1/responses", "/v1/messages", "/v1/realtime"] }, "cognition": { "base_url": "https://api.cognition.ai/v1", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7fa09951eae..4dc6768e12a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -34717,6 +34717,21 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "meta/muse-voice-transcribe-1.0": { + "litellm_provider": "meta", + "mode": "audio_transcription", + "source": "https://dev.meta.ai/docs/speech-to-text", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, "meta_llama/Llama-3.3-70B-Instruct": { "litellm_provider": "meta_llama", "max_input_tokens": 128000, @@ -59098,9 +59113,9 @@ "litellm_provider": "wandb", "mode": "chat", "supports_reasoning": true, - "input_cost_per_token": 0.00000131, - "output_cost_per_token": 0.00000396, - "cache_read_input_token_cost": 0.000000044, + "input_cost_per_token": 1.31e-06, + "output_cost_per_token": 3.96e-06, + "cache_read_input_token_cost": 4.4e-08, "supports_prompt_caching": true, "source": "https://wandb.ai/site/pricing/tokens/" }, @@ -59108,9 +59123,9 @@ "litellm_provider": "wandb", "mode": "chat", "supports_reasoning": true, - "input_cost_per_token": 0.0000001, - "output_cost_per_token": 0.00000015, - "cache_read_input_token_cost": 0.00000005, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 1.5e-07, + "cache_read_input_token_cost": 5e-08, "supports_prompt_caching": true, "source": "https://wandb.ai/site/pricing/tokens/" }, diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 44c47af57f4..b0803e44f6b 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -333,7 +333,7 @@ async def _resolve_vertex_access_token_bounded( @wrapper_client -async def _arealtime( +async def _arealtime( # noqa: C901 # central dispatcher branches once per supported realtime provider model: str, websocket: "WebSocket", # fastapi websocket api_base: str | None = None, @@ -391,7 +391,37 @@ async def _arealtime( model=model, provider=LlmProviders(_custom_llm_provider), ) - if provider_config is not None: + if _custom_llm_provider == LlmProviders.META.value: + if model != "muse-voice-transcribe-1.0": + raise ValueError(f"Unsupported Meta realtime model: {model}") + if query_params is None or query_params.get("intent") != "transcription": + raise ValueError("Meta Muse Voice realtime requires intent=transcription") + + from litellm.llms.meta.realtime.handler import MetaRealtime + + meta_api_key: Final = get_secret_str("META_API_KEY") + dynamic_key_override: Final = dynamic_api_key if dynamic_api_key != meta_api_key else None + resolved_meta_api_key: Final = ( + api_key + or litellm_params.api_key + or dynamic_key_override + or get_secret_str("MODEL_API_KEY") + or dynamic_api_key + or meta_api_key + ) + await MetaRealtime().async_realtime( + model=model, + websocket=websocket, + logging_obj=litellm_logging_obj, + api_base=dynamic_api_base or litellm_params.api_base or api_base, + api_key=resolved_meta_api_key, + client=client, + timeout=timeout, + query_params=query_params, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=_build_litellm_metadata(kwargs), + ) + elif provider_config is not None: await base_llm_http_handler.async_realtime( model=model, websocket=websocket, diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index b7c4371f32f..02a102c9579 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2202,6 +2202,8 @@ class OpenAIRealtimeInputAudioTranscriptionCompleted(TypedDict): item_id: ReadOnly[str] content_index: ReadOnly[int] transcript: ReadOnly[str] + usage: NotRequired[ReadOnly[Mapping[str, object]]] + speaker: NotRequired[ReadOnly[str]] class OpenAIRealtimeUsageTokenDetails(TypedDict): diff --git a/litellm/types/realtime.py b/litellm/types/realtime.py index 17dc70126f3..30db794c96e 100644 --- a/litellm/types/realtime.py +++ b/litellm/types/realtime.py @@ -169,9 +169,19 @@ class RealtimeInputAudioTranscriptionUsageInputTokenDetails(TypedDict): audio_tokens: ReadOnly[int] -class RealtimeInputAudioTranscriptionUsage(TypedDict): +class RealtimeInputAudioTranscriptionTokenUsage(TypedDict): type: ReadOnly[Literal["tokens"]] input_tokens: ReadOnly[int] output_tokens: ReadOnly[int] total_tokens: ReadOnly[int] input_token_details: ReadOnly[RealtimeInputAudioTranscriptionUsageInputTokenDetails] + + +class RealtimeInputAudioTranscriptionDurationUsage(TypedDict): + type: ReadOnly[Literal["duration"]] + seconds: ReadOnly[float] + + +RealtimeInputAudioTranscriptionUsage = ( + RealtimeInputAudioTranscriptionTokenUsage | RealtimeInputAudioTranscriptionDurationUsage +) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7fa09951eae..4dc6768e12a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -34717,6 +34717,21 @@ "supports_web_search": true, "supports_xhigh_reasoning_effort": true }, + "meta/muse-voice-transcribe-1.0": { + "litellm_provider": "meta", + "mode": "audio_transcription", + "source": "https://dev.meta.ai/docs/speech-to-text", + "supported_endpoints": [ + "/v1/realtime" + ], + "supported_modalities": [ + "audio" + ], + "supported_output_modalities": [ + "text" + ], + "supports_audio_input": true + }, "meta_llama/Llama-3.3-70B-Instruct": { "litellm_provider": "meta_llama", "max_input_tokens": 128000, @@ -59098,9 +59113,9 @@ "litellm_provider": "wandb", "mode": "chat", "supports_reasoning": true, - "input_cost_per_token": 0.00000131, - "output_cost_per_token": 0.00000396, - "cache_read_input_token_cost": 0.000000044, + "input_cost_per_token": 1.31e-06, + "output_cost_per_token": 3.96e-06, + "cache_read_input_token_cost": 4.4e-08, "supports_prompt_caching": true, "source": "https://wandb.ai/site/pricing/tokens/" }, @@ -59108,9 +59123,9 @@ "litellm_provider": "wandb", "mode": "chat", "supports_reasoning": true, - "input_cost_per_token": 0.0000001, - "output_cost_per_token": 0.00000015, - "cache_read_input_token_cost": 0.00000005, + "input_cost_per_token": 1e-07, + "output_cost_per_token": 1.5e-07, + "cache_read_input_token_cost": 5e-08, "supports_prompt_caching": true, "source": "https://wandb.ai/site/pricing/tokens/" }, diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 9c0f6f59463..e330cb103b1 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -10,8 +10,6 @@ from websockets.exceptions import ConnectionClosed from websockets.frames import Close import litellm - - from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.realtime_streaming import ( REALTIME_SESSION_SUCCESS_LOGGED_KEY, @@ -20,10 +18,6 @@ from litellm.litellm_core_utils.realtime_streaming import ( ) from litellm.llms.xai.realtime.transformation import XAIRealtimeNormalizer from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.llms.openai import ( - OpenAIRealtimeStreamResponseBaseObject, - OpenAIRealtimeStreamSessionEvents, -) def _make_transcript_event(text: str, item_id: str = "item_x") -> bytes: @@ -161,6 +155,7 @@ async def test_backend_to_client_send_text_receives_str_not_bytes(): logging_obj = MagicMock() logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) await streaming.backend_to_client_send_messages() @@ -812,7 +807,6 @@ async def test_transcription_captured_in_backend_to_client(): Test that conversation.item.input_audio_transcription.completed events from the backend are captured as user input during the WebSocket session. """ - import litellm client_ws = MagicMock() client_ws.send_text = AsyncMock() @@ -838,6 +832,7 @@ async def test_transcription_captured_in_backend_to_client(): logging_obj.model_call_details = {"messages": "default-message-value"} logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) await streaming.backend_to_client_send_messages() @@ -883,6 +878,7 @@ async def test_transcription_session_captures_usage_and_skips_response_create(): logging_obj.model_call_details = {} logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) await streaming.backend_to_client_send_messages() @@ -1100,7 +1096,6 @@ def test_capture_transcription_usage_deduplicates_when_already_stored(): When the event is already in messages (logged via store_message), it must not be appended a second time by _capture_transcription_usage. """ - import litellm streaming = RealTimeStreaming(MagicMock(), MagicMock(), MagicMock()) # Add the event type to the default logged list so _should_store_message returns True. @@ -1409,7 +1404,6 @@ async def test_realtime_guardrail_blocks_prompt_injection(monkeypatch: pytest.Mo ) - @pytest.mark.asyncio async def test_realtime_guardrail_allows_clean_transcript(monkeypatch: pytest.MonkeyPatch): """ @@ -1466,7 +1460,6 @@ async def test_realtime_guardrail_allows_clean_transcript(monkeypatch: pytest.Mo assert len(response_creates) == 1, f"Clean transcript should trigger response.create, got: {sent_to_backend}" - @pytest.mark.asyncio async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ @@ -1560,7 +1553,6 @@ async def test_realtime_text_input_guardrail_blocks_and_returns_error(monkeypatc assert len(original_items) == 0, f"Blocked item should not be forwarded to backend, got: {original_items}" - @pytest.mark.asyncio async def test_realtime_function_call_output_guardrail_blocks_and_returns_error(monkeypatch: pytest.MonkeyPatch): """ @@ -1649,7 +1641,6 @@ async def test_realtime_function_call_output_guardrail_blocks_and_returns_error( assert "test@example.com" not in sanitized_item["output"] - @pytest.mark.asyncio async def test_realtime_function_call_output_guardrail_allows_clean_output(monkeypatch: pytest.MonkeyPatch): """ @@ -1714,7 +1705,6 @@ async def test_realtime_function_call_output_guardrail_allows_clean_output(monke assert len(forwarded) == 1, f"Clean function_call_output should be forwarded, got: {forwarded}" - @pytest.mark.asyncio async def test_realtime_text_input_guardrail_uses_pre_call_mode(monkeypatch: pytest.MonkeyPatch): """ @@ -1750,7 +1740,6 @@ async def test_realtime_text_input_guardrail_uses_pre_call_mode(monkeypatch: pyt ) - @pytest.mark.asyncio async def test_realtime_session_created_injects_session_update_for_audio_guardrail(monkeypatch: pytest.MonkeyPatch): """ @@ -1807,7 +1796,6 @@ async def test_realtime_session_created_injects_session_update_for_audio_guardra ) - @pytest.mark.asyncio async def test_realtime_session_created_does_not_inject_session_update_for_pre_call_only( monkeypatch: pytest.MonkeyPatch, @@ -1852,7 +1840,6 @@ async def test_realtime_session_created_does_not_inject_session_update_for_pre_c assert len(session_updates) == 0, f"pre_call-only guardrail must not inject session.update, got: {sent_to_backend}" - @pytest.mark.asyncio async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monkeypatch: pytest.MonkeyPatch): """Model Armor-style pre_call + post_call must not gate audio VAD.""" @@ -1868,17 +1855,17 @@ async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monke litellm, "callbacks", [ - ModelArmorStyleGuardrail( - guardrail_name="model_armor_all_pre_call", - event_hook=GuardrailEventHooks.pre_call, - default_on=False, - ), - ModelArmorStyleGuardrail( - guardrail_name="model_armor_all_post_call", - event_hook=GuardrailEventHooks.post_call, - default_on=False, - ), - ], + ModelArmorStyleGuardrail( + guardrail_name="model_armor_all_pre_call", + event_hook=GuardrailEventHooks.pre_call, + default_on=False, + ), + ModelArmorStyleGuardrail( + guardrail_name="model_armor_all_post_call", + event_hook=GuardrailEventHooks.post_call, + default_on=False, + ), + ], ) client_ws = MagicMock() @@ -1902,7 +1889,6 @@ async def test_pre_call_and_post_call_guardrails_do_not_disable_server_vad(monke assert streaming._has_audio_transcription_guardrails() is False - @pytest.mark.asyncio async def test_end_session_after_n_fails_closes_connection(monkeypatch: pytest.MonkeyPatch): """ @@ -1949,7 +1935,6 @@ async def test_end_session_after_n_fails_closes_connection(monkeypatch: pytest.M assert streaming._violation_count == 2 - @pytest.mark.asyncio async def test_on_violation_end_session_closes_on_first_fail(monkeypatch: pytest.MonkeyPatch): """ @@ -1995,7 +1980,6 @@ async def test_on_violation_end_session_closes_on_first_fail(monkeypatch: pytest assert streaming._violation_count == 1 - @pytest.mark.asyncio async def test_provider_path_suppresses_duplicate_session_created_after_synthetic(): client_ws = MagicMock() @@ -2956,7 +2940,9 @@ async def test_log_messages_routes_async_logging_through_bounded_worker(): mock_worker.ensure_initialized_and_enqueue.assert_called_once() enqueued = mock_worker.ensure_initialized_and_enqueue.call_args - assert (enqueued.args or tuple(enqueued.kwargs.values()))[0] is logging_obj.dispatch_success_handlers.return_value + assert (enqueued.args or tuple(enqueued.kwargs.values()))[ + 0 + ] is logging_obj.dispatch_success_handlers.return_value logging_obj.dispatch_success_handlers.assert_called_once_with(streaming.messages, prefer_async_handlers=True) logging_obj.success_handler.assert_not_called() # the bare create_task path must no longer be used for success logging @@ -3041,6 +3027,7 @@ async def test_session_close_flushes_unbilled_transcription_usage(): logging_obj: Final = MagicMock() logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() usage: Final[RealtimeInputAudioTranscriptionUsage] = { "type": "tokens", @@ -3116,6 +3103,7 @@ async def test_session_close_flush_noop_without_unbilled_usage(): logging_obj: Final = MagicMock() logging_obj.async_success_handler = AsyncMock() logging_obj.success_handler = MagicMock() + logging_obj.dispatch_success_handlers = AsyncMock() provider_config: Final = MagicMock() provider_config.unbilled_usage_on_session_close = MagicMock(return_value=None) @@ -3412,3 +3400,144 @@ async def test_refused_session_does_not_stamp_the_reservation_ownership_marker() assert session.logging.logged_failures == (upstream_close,) assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details + + +@pytest.mark.asyncio +async def test_separate_usage_provider_flushes_duration_once_without_client_event(): + from typing import Final + + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None)) + logging_obj: Final = MagicMock() + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + usage_provider: Final = MagicMock() + usage_provider.unbilled_usage_on_session_close.return_value = { + "type": "duration", + "seconds": 0.75, + } + + streaming: Final = RealTimeStreaming( + client_ws, + backend_ws, + logging_obj, + model="muse-voice-transcribe-1.0", + usage_provider=usage_provider, + ) + + await streaming.backend_to_client_send_messages() + + usage_provider.unbilled_usage_on_session_close.assert_called_once_with("muse-voice-transcribe-1.0") + duration_events: Final = tuple( + message + for message in streaming.messages + if isinstance(message, dict) and message.get("usage") == {"type": "duration", "seconds": 0.75} + ) + assert len(duration_events) == 1 + assert client_ws.send_text.await_count == 0 + logging_obj.dispatch_success_handlers.assert_called_once_with(streaming.messages, prefer_async_handlers=True) + + +@pytest.mark.asyncio +async def test_transformed_transcription_completion_never_sends_response_create(): + from typing import Final + + completed_event: Final = { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": "event_1", + "item_id": "turn_1", + "content_index": 0, + "transcript": "private transcript", + "usage": {"type": "duration", "seconds": 0.5}, + } + provider_config: Final = MagicMock() + provider_config.requires_session_configuration.return_value = True + provider_config.transform_realtime_response.return_value = { + "response": completed_event, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_conversation_id": None, + "current_item_chunks": None, + "current_delta_type": None, + "session_configuration_request": None, + } + provider_config.transform_realtime_request.return_value = (json.dumps({"type": "response.create"}),) + provider_config.is_setup_message.return_value = False + provider_config.is_content_message.return_value = False + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + + streaming: Final = RealTimeStreaming( + client_ws, + backend_ws, + MagicMock(), + provider_config=provider_config, + model="muse-voice-transcribe-1.0", + force_transcription_model="muse-voice-transcribe-1.0", + ) + + await streaming._handle_provider_config_message("{}") + + assert json.loads(client_ws.send_text.await_args.args[0]) == completed_event + backend_ws.send.assert_not_awaited() + + +def test_private_logging_excludes_audio_transcript_hints_and_provider_body(monkeypatch: pytest.MonkeyPatch): + from typing import Final + + monkeypatch.setattr(litellm, "logged_real_time_event_types", "*") + logging_obj: Final = MagicMock() + logging_obj.model_call_details = {} + streaming: Final = RealTimeStreaming( + MagicMock(), + MagicMock(), + logging_obj, + model="muse-voice-transcribe-1.0", + exclude_private_content_from_logs=True, + ) + audio: Final = "cHJpdmF0ZS1hdWRpbw==" + transcript: Final = "private transcript" + keyword: Final = "private keyword" + provider_body: Final = "private provider body" + + streaming.store_input( + json.dumps( + { + "type": "session.update", + "session": { + "type": "transcription", + "model": "muse-voice-transcribe-1.0", + "mode": "ENDPOINTING", + "audio": {"input": {"transcription": {"keywords": [keyword]}}}, + }, + } + ) + ) + streaming.store_input(json.dumps({"type": "input_audio_buffer.append", "audio": audio})) + streaming.store_message( + { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": "event_1", + "item_id": "turn_1", + "transcript": transcript, + "provider_body": provider_body, + "usage": {"type": "duration", "seconds": 1.0}, + } + ) + + logged_inputs: Final = tuple(call.kwargs["input"] for call in logging_obj.pre_call.call_args_list) + serialized: Final = json.dumps({"inputs": logged_inputs, "messages": streaming.messages}) + assert audio not in serialized + assert transcript not in serialized + assert keyword not in serialized + assert provider_body not in serialized + assert "muse-voice-transcribe-1.0" in serialized + assert "ENDPOINTING" in serialized + assert "turn_1" in serialized + assert '"seconds": 1.0' in serialized + assert streaming.input_messages == [] diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py new file mode 100644 index 00000000000..3d31b3a5534 --- /dev/null +++ b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py @@ -0,0 +1,449 @@ +import asyncio +import base64 +import json +from collections.abc import Callable +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.llms.meta.realtime.handler import ( + DEFAULT_MUSE_REALTIME_URL, + MetaRealtime, + MuseAdapterError, + MuseRealtimeAdapter, + build_muse_realtime_url, + normalize_access_token, + safe_close_reason, + sanitize_close_code, +) +from litellm.llms.meta.realtime.transformation import MUSE_MODEL + + +class FakeProviderWebSocket: + def __init__(self, session_id: str = "provider-session") -> None: + self.sent: list[str | bytes] = [] + self.close_calls: list[tuple[int, str]] = [] + self._session_id: Final = session_id + self._recv_count = 0 + self._closed = asyncio.Event() + + async def send(self, message: str | bytes) -> None: + self.sent.append(message) + + async def recv(self, decode: bool | None = None) -> str | bytes: + self._recv_count += 1 + if self._recv_count == 1: + return json.dumps({"sessionId": self._session_id}) + await self._closed.wait() + raise MuseAdapterError("closed", close_code=1000) + + async def close(self, code: int = 1000, reason: str = "") -> None: + self.close_calls.append((code, reason)) + self._closed.set() + + +class DelayedAckWebSocket(FakeProviderWebSocket): + def __init__(self) -> None: + super().__init__() + self.ack_release = asyncio.Event() + + async def recv(self, decode: bool | None = None) -> str | bytes: + self._recv_count += 1 + if self._recv_count == 1: + await self.ack_release.wait() + return json.dumps({"sessionId": self._session_id}) + await self._closed.wait() + raise MuseAdapterError("closed", close_code=1000) + + +async def _wait_until(predicate: Callable[[], bool]) -> None: + for _ in range(100): + if predicate(): + return + await asyncio.sleep(0) + raise AssertionError("condition did not become true") + + +def _session_update(*, rate: int = 24_000, mode: str = "ENDPOINTING") -> str: + return json.dumps( + { + "type": "session.update", + "session": { + "type": "transcription", + "mode": mode, + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": rate, "channels": 1}, + "transcription": {"model": MUSE_MODEL}, + } + }, + }, + } + ) + + +async def _configured_adapter( + *, + rate: int = 24_000, + mode: str = "ENDPOINTING", + provider_ws: FakeProviderWebSocket | None = None, + monotonic: Callable[[], float] = lambda: 10.0, + sleep: Callable[[float], object] | None = None, +) -> tuple[MuseRealtimeAdapter, FakeProviderWebSocket, dict[str, object]]: + ws: Final = provider_ws or FakeProviderWebSocket() + connect_call: Final[dict[str, object]] = {} + + async def connect(url: str, **kwargs: object) -> FakeProviderWebSocket: + connect_call.update({"url": url, **kwargs}) + return ws + + async def no_sleep(_: float) -> None: + return None + + adapter: Final = MuseRealtimeAdapter( + model=f"meta/{MUSE_MODEL}", + api_key=" raw-token ", + websocket_connect=connect, + monotonic=monotonic, + sleep=sleep or no_sleep, + ) + created: Final = json.loads(await adapter.recv()) + assert created["type"] == "session.created" + await adapter.send(_session_update(rate=rate, mode=mode)) + updated: Final = json.loads(await adapter.recv()) + assert updated["type"] == "session.updated" + return adapter, ws, connect_call + + +@pytest.mark.parametrize( + ("api_key", "expected"), + [ + ("token", "Bearer token"), + (" Bearer token ", "Bearer token"), + ("bearer token", "Bearer token"), + ], +) +def test_normalize_access_token_emits_exactly_one_bearer_prefix(api_key: str, expected: str): + assert normalize_access_token(api_key) == expected + + +@pytest.mark.parametrize("api_key", ["", " ", "Bearer", " bearer "]) +def test_normalize_access_token_rejects_missing_token(api_key: str): + with pytest.raises(ValueError, match=r"token|key is required"): + normalize_access_token(api_key) + + +def test_build_muse_realtime_url_uses_fixed_secure_path(): + assert build_muse_realtime_url(None) == DEFAULT_MUSE_REALTIME_URL + assert build_muse_realtime_url("https://example.test/custom/path?ignored=yes") == ( + "wss://example.test/v1/asr/realtime" + ) + assert build_muse_realtime_url("wss://example.test:8443/other") == ("wss://example.test:8443/v1/asr/realtime") + + +@pytest.mark.parametrize( + "api_base", + [ + "http://example.test", + "ws://example.test", + "wss://user:pass@example.test", + "wss://example.test/path#fragment", + "not-a-url", + ], +) +def test_build_muse_realtime_url_rejects_insecure_or_ambiguous_overrides(api_base: str): + with pytest.raises(ValueError, match="absolute wss:// or https://"): + build_muse_realtime_url(api_base) + + +@pytest.mark.asyncio +async def test_handshake_contains_bearer_only_in_json_body_and_waits_for_ack(): + provider_ws: Final = DelayedAckWebSocket() + connect_call: Final[dict[str, object]] = {} + + async def connect(url: str, **kwargs: object) -> DelayedAckWebSocket: + connect_call.update({"url": url, **kwargs}) + return provider_ws + + adapter: Final = MuseRealtimeAdapter( + model=MUSE_MODEL, + api_key="Bearer private-token", + websocket_connect=connect, + ) + await adapter.recv() + update_task: Final = asyncio.create_task(adapter.send(_session_update())) + await _wait_until(lambda: len(provider_ws.sent) == 1) + + assert connect_call["url"] == DEFAULT_MUSE_REALTIME_URL + assert "additional_headers" not in connect_call + handshake: Final = json.loads(provider_ws.sent[0]) + assert handshake["authorization"] == {"accessToken": "Bearer private-token"} + assert handshake["audioEncoding"] == "PCM_24KHZ" + assert not update_task.done() + assert not any(isinstance(frame, bytes) for frame in provider_ws.sent) + + provider_ws.ack_release.set() + await update_task + updated: Final = json.loads(await adapter.recv()) + assert updated["type"] == "session.updated" + assert updated["session"]["id"] == "provider-session" + await adapter.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("rate", "packet_bytes"), [(16_000, 2_560), (24_000, 3_840)]) +async def test_audio_is_strictly_decoded_and_packetized_as_raw_pcm(rate: int, packet_bytes: int): + adapter, provider_ws, _ = await _configured_adapter(rate=rate) + pcm: Final = (b"\xff\xfe\x00\x80" * (packet_bytes // 2))[: packet_bytes * 2] + + await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(pcm).decode()})) + await _wait_until(lambda: sum(isinstance(frame, bytes) for frame in provider_ws.sent) == 2) + + binary_frames: Final = tuple(frame for frame in provider_ws.sent if isinstance(frame, bytes)) + assert binary_frames == (pcm[:packet_bytes], pcm[packet_bytes:]) + assert b"\xff\xfe" in pcm + await adapter.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("audio", "expected_message"), + [ + ("not base64!", "valid base64"), + (base64.b64encode(b"\x00").decode(), "complete samples"), + ], +) +async def test_invalid_base64_or_odd_pcm_is_rejected(audio: str, expected_message: str): + adapter, provider_ws, _ = await _configured_adapter() + + await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": audio})) + error: Final = json.loads(await adapter.recv()) + + assert error["type"] == "error" + assert error["error"]["code"] == "invalid_audio" + assert expected_message in error["error"]["message"] + assert adapter.close_code == 1008 + assert not any(isinstance(frame, bytes) for frame in provider_ws.sent) + await adapter.close() + + +@pytest.mark.asyncio +async def test_absolute_pacing_delays_only_audio_ahead_of_wall_time(): + sleeps: Final[list[float]] = [] + + async def record_sleep(delay: float) -> None: + sleeps.append(delay) + + adapter, provider_ws, _ = await _configured_adapter(monotonic=lambda: 10.0, sleep=record_sleep) + pcm: Final = b"\x01\x02" * 3_840 + + await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(pcm).decode()})) + await _wait_until(lambda: sum(isinstance(frame, bytes) for frame in provider_ws.sent) == 2) + + assert sleeps == pytest.approx([0.08]) + await adapter.close() + + +@pytest.mark.asyncio +async def test_append_larger_than_four_seconds_is_rejected_without_dropping_prefix(): + adapter, provider_ws, _ = await _configured_adapter(rate=16_000) + oversized_pcm: Final = b"\x00\x00" * (16_000 * 4 + 1) + + await adapter.send( + json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(oversized_pcm).decode()}) + ) + error: Final = json.loads(await adapter.recv()) + + assert error["error"]["code"] == "audio_backlog_exceeded" + assert adapter.close_code == 1008 + assert not any(isinstance(frame, bytes) for frame in provider_ws.sent) + await adapter.close() + + +@pytest.mark.asyncio +async def test_clear_discards_only_unsent_audio(): + adapter, provider_ws, _ = await _configured_adapter() + old_pcm: Final = b"\x01\x02" * 100 + new_pcm: Final = b"\x03\x04" * 100 + + await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(old_pcm).decode()})) + await adapter.send(_event("input_audio_buffer.clear")) + cleared: Final = json.loads(await adapter.recv()) + await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(new_pcm).decode()})) + await adapter.send(_event("input_audio_buffer.commit")) + await _wait_until(lambda: any(isinstance(frame, bytes) for frame in provider_ws.sent)) + + assert cleared["type"] == "input_audio_buffer.cleared" + assert tuple(frame for frame in provider_ws.sent if isinstance(frame, bytes)) == (new_pcm,) + assert '{"type":"endStream"}' not in provider_ws.sent + await adapter.close() + + +@pytest.mark.asyncio +async def test_endpointing_commit_flushes_partial_packet_without_ending_stream(): + adapter, provider_ws, _ = await _configured_adapter(mode="ENDPOINTING") + pcm: Final = b"\x01\x02" * 100 + + await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(pcm).decode()})) + await adapter.send(_event("input_audio_buffer.commit")) + committed: Final = json.loads(await adapter.recv()) + await _wait_until(lambda: any(isinstance(frame, bytes) for frame in provider_ws.sent)) + + assert committed["type"] == "input_audio_buffer.committed" + assert committed["item_id"].startswith("item_") + assert committed["previous_item_id"] is None + assert tuple(frame for frame in provider_ws.sent if isinstance(frame, bytes)) == (pcm,) + assert '{"type":"endStream"}' not in provider_ws.sent + await adapter.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("mode", "terminal_event"), + [("PUSH_TO_TALK", "input_audio_buffer.commit"), ("ENDPOINTING", "input_audio_buffer.end")], +) +async def test_commit_or_end_sends_end_stream_exactly_once(mode: str, terminal_event: str): + adapter, provider_ws, _ = await _configured_adapter(mode=mode) + pcm: Final = b"\x01\x02" * 100 + + await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(pcm).decode()})) + await adapter.send(_event(terminal_event)) + await adapter.send(_event("input_audio_buffer.end")) + await _wait_until(lambda: '{"type":"endStream"}' in provider_ws.sent) + + assert tuple(frame for frame in provider_ws.sent if isinstance(frame, bytes)) == (pcm,) + assert provider_ws.sent.count('{"type":"endStream"}') == 1 + await adapter.close() + + +@pytest.mark.asyncio +async def test_response_create_is_returned_as_error_and_never_sent_upstream(): + adapter, provider_ws, _ = await _configured_adapter() + + await adapter.send(_event("response.create")) + error: Final = json.loads(await adapter.recv()) + + assert error["type"] == "error" + assert error["error"]["code"] == "unsupported_event" + assert not any(isinstance(frame, str) and "response.create" in frame for frame in provider_ws.sent) + await adapter.close() + + +@pytest.mark.asyncio +async def test_close_codes_and_reasons_are_sanitized_without_secret_leakage(): + adapter, provider_ws, _ = await _configured_adapter() + secret: Final = "Bearer private-token" + + await adapter.close(code=4001, reason=f"provider rejected {secret}") + + assert adapter.close_code == 1011 + assert adapter.close_reason == "Realtime transcription service error" + assert provider_ws.close_calls == [(1011, "Realtime transcription service error")] + assert secret not in json.dumps(provider_ws.close_calls) + assert sanitize_close_code(1013) == 1013 + assert safe_close_reason(1008) == "Invalid realtime transcription request" + + +@pytest.mark.asyncio +async def test_handshake_failure_reports_only_exception_type(): + secret: Final = "private-token" + + async def failing_connect(url: str, **kwargs: object) -> FakeProviderWebSocket: + raise RuntimeError(f"failed with {secret}") + + adapter: Final = MuseRealtimeAdapter( + model=MUSE_MODEL, + api_key=secret, + websocket_connect=failing_connect, + ) + await adapter.recv() + + await adapter.send(_session_update()) + error = json.loads(await adapter.recv()) + + assert error["type"] == "error" + assert error["error"]["message"] == "Meta Muse realtime handshake failed" + assert secret not in json.dumps(error) + with pytest.raises(MuseAdapterError) as exc_info: + await adapter.recv() + assert exc_info.value.close_code == 1011 + assert secret not in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_meta_realtime_missing_credentials_closes_client_with_policy_code(): + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + client_ws.close = AsyncMock() + + await MetaRealtime().async_realtime( + model=MUSE_MODEL, + websocket=client_ws, + logging_obj=MagicMock(), + api_key=None, + ) + + client_ws.close.assert_awaited_once_with( + code=1008, + reason="Invalid realtime transcription request", + ) + sent_error: Final = json.loads(client_ws.send_text.await_args.args[0]) + assert sent_error["type"] == "error" + assert sent_error["error"]["code"] == "invalid_configuration" + + +@pytest.mark.asyncio +async def test_meta_realtime_invalid_constructor_input_sends_error_before_close(): + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + client_ws.close = AsyncMock() + + await MetaRealtime().async_realtime( + model=MUSE_MODEL, + websocket=client_ws, + logging_obj=MagicMock(), + api_key="Bearer", + ) + + sent_error: Final = json.loads(client_ws.send_text.await_args.args[0]) + assert sent_error["error"]["message"] == "Invalid Meta Muse realtime configuration" + client_ws.close.assert_awaited_once_with( + code=1008, + reason="Invalid realtime transcription request", + ) + + +@pytest.mark.asyncio +async def test_meta_realtime_enables_private_logging_usage_and_model_enforcement(monkeypatch: pytest.MonkeyPatch): + captured: Final[dict[str, object]] = {} + + class CapturingStreaming: + def __init__(self, websocket, backend_ws, logging_obj, **kwargs): + captured.update({"websocket": websocket, "backend_ws": backend_ws, "logging_obj": logging_obj, **kwargs}) + + async def bidirectional_forward(self) -> None: + return None + + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + client_ws.close = AsyncMock() + monkeypatch.setattr("litellm.llms.meta.realtime.handler.RealTimeStreaming", CapturingStreaming) + + await MetaRealtime().async_realtime( + model=MUSE_MODEL, + websocket=client_ws, + logging_obj=MagicMock(), + api_key="private-token", + ) + + adapter: Final = captured["backend_ws"] + assert isinstance(adapter, MuseRealtimeAdapter) + assert captured["force_transcription_model"] == MUSE_MODEL + assert captured["usage_provider"] is adapter + assert captured["exclude_private_content_from_logs"] is True + client_ws.close.assert_awaited_once_with(code=1000, reason="Session closed") + + +def _event(event_type: str) -> str: + return json.dumps({"type": event_type}) diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py new file mode 100644 index 00000000000..8cc4af836a5 --- /dev/null +++ b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py @@ -0,0 +1,299 @@ +import json + +import pytest + +from litellm.llms.meta.realtime.transformation import ( + MUSE_MODEL, + MuseEventTransformer, + MuseProtocolError, + encode_event, + normalize_language, + parse_session_update, + session_created_event, + session_updated_event, +) + + +def _event(event_type: str, **fields: object) -> str: + return json.dumps({"type": event_type, **fields}) + + +def test_beta_session_builds_authenticated_24khz_handshake_with_hints(): + config = parse_session_update( + _event( + "session.update", + session={ + "type": "transcription", + "input_audio_format": "pcm16", + "turn_detection": {"type": "server_vad"}, + "input_audio_transcription": { + "model": "meta/muse-voice-transcribe-1.0", + "language": "en-US", + "language_bias": ["Spanish", "english", "French"], + "keywords": [" Muse ", "LiteLLM", "Muse"], + "prompt": "must not become a keyword", + }, + }, + ), + "meta/muse-voice-transcribe-1.0", + ) + + assert config.sample_rate == 24_000 + assert config.packet_bytes == 3_840 + assert config.mode == "ENDPOINTING" + assert config.language_bias == ("English", "Spanish", "French") + assert config.keywords == ("Muse", "LiteLLM") + assert config.handshake("Bearer token") == { + "mode": "ENDPOINTING", + "authorization": {"accessToken": "Bearer token"}, + "audioEncoding": "PCM_24KHZ", + "model": MUSE_MODEL, + "partialMode": "CUMULATIVE", + "emitAudioProgress": True, + "keywords": ["Muse", "LiteLLM"], + "languageBias": ["English", "Spanish", "French"], + } + assert "must not become a keyword" not in json.dumps(config.handshake("Bearer token")) + + +def test_ga_session_accepts_16khz_mono_push_to_talk(): + config = parse_session_update( + _event( + "session.update", + session={ + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": 16000, "channels": 1}, + "turn_detection": None, + "transcription": {"model": MUSE_MODEL, "language": "zh-Hans"}, + } + }, + }, + ), + MUSE_MODEL, + ) + + assert config.sample_rate == 16_000 + assert config.packet_bytes == 2_560 + assert config.mode == "PUSH_TO_TALK" + assert config.language_bias == ("Mandarin Chinese",) + assert config.handshake("Bearer token")["audioEncoding"] == "PCM_16KHZ" + + +@pytest.mark.parametrize( + ("source", "expected"), + [ + ("EN_us", "English"), + ("mandarin chinese", "Mandarin Chinese"), + ("fil-PH", "Tagalog"), + ("iw-IL", "Hebrew"), + ("pt-BR", "Portuguese"), + ], +) +def test_language_normalization_uses_official_muse_names(source: str, expected: str): + assert normalize_language(source) == expected + + +@pytest.mark.parametrize( + ("session", "message"), + [ + ({"input_audio_format": "g711_ulaw"}, "requires pcm16"), + ({"audio": {"input": {"format": {"type": "audio/pcm", "rate": 8000}}}}, "16000 Hz or 24000 Hz"), + ( + {"audio": {"input": {"format": {"type": "audio/pcm", "rate": 24000, "channels": 2}}}}, + "requires mono", + ), + ( + {"input_audio_format": "pcm16", "audio": {"input": {"format": {"type": "audio/pcm"}}}}, + "either beta or GA layout", + ), + ({"input_audio_transcription": {"model": "other-model"}}, "cannot be changed"), + ({"input_audio_transcription": {"keywords": ["valid", ""]}}, "non-empty strings"), + ({"input_audio_transcription": {"language": "xx"}}, "unsupported Muse Voice language"), + ], +) +def test_session_rejects_unsupported_audio_model_and_hints(session: dict[str, object], message: str): + with pytest.raises(MuseProtocolError, match=message): + parse_session_update(_event("session.update", session={"type": "transcription", **session}), MUSE_MODEL) + + +def test_session_events_expose_openai_transcription_shapes(): + config = parse_session_update( + _event( + "session.update", + session={ + "mode": "DIARIZATION", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": 24000}, + "transcription": {"model": MUSE_MODEL, "language": "ja", "keywords": ["Meta"]}, + } + }, + }, + ), + MUSE_MODEL, + ) + + created = session_created_event(MUSE_MODEL, "session-before-handshake") + updated = session_updated_event(config, "provider-session") + + assert created["type"] == "session.created" + assert created["session"]["type"] == "transcription" + assert updated["type"] == "session.updated" + assert updated["session"]["id"] == "provider-session" + assert updated["session"]["audio"]["input"]["transcription"] == { + "model": MUSE_MODEL, + "language": "Japanese", + "keywords": ["Meta"], + "language_bias": ["Japanese"], + } + + +def test_turnless_empty_silence_transcript_is_ignored(): + transformer = MuseEventTransformer() + + assert transformer.transform(_event("transcript", transcript="", final=True)) == () + + +def test_transcript_without_speech_start_synthesizes_start_before_delta(): + transformer = MuseEventTransformer() + + events = transformer.transform(_event("transcript", turnId="turn-1", transcript="hello", final=False)) + + assert [event["type"] for event in events] == [ + "input_audio_buffer.speech_started", + "conversation.item.input_audio_transcription.delta", + ] + + +def test_cumulative_partials_emit_only_extensions_and_final_is_authoritative(): + transformer = MuseEventTransformer() + + started = transformer.transform(_event("speechStart", turnId="turn-1")) + first = transformer.transform(_event("transcript", turnId="turn-1", transcript="hello", final=False)) + extension = transformer.transform(_event("transcript", turnId="turn-1", transcript="hello world", final=False)) + rewrite = transformer.transform(_event("transcript", turnId="turn-1", transcript="hullo world", final=False)) + assert transformer.transform(_event("speechComplete", turnId="turn-1", transcript="hullo world")) == () + completed = transformer.transform(_event("speechEnd", turnId="turn-1")) + + assert [event["type"] for event in started] == ["input_audio_buffer.speech_started"] + assert first[0]["delta"] == "hello" + assert extension[0]["delta"] == " world" + assert rewrite == () + assert completed[0]["type"] == "input_audio_buffer.speech_stopped" + assert completed[1]["type"] == "conversation.item.input_audio_transcription.completed" + assert completed[1]["item_id"] == "turn-1" + assert completed[1]["transcript"] == "hullo world" + + +def test_completed_transcript_waits_for_speech_stopped(): + transformer = MuseEventTransformer() + + transformer.transform(_event("speechStart", turnId="turn-1")) + assert transformer.transform(_event("speechComplete", turnId="turn-1", transcript="done")) == () + + released = transformer.transform(_event("speechEnd", turnId="turn-1")) + assert [event["type"] for event in released] == [ + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + + +def test_overlapping_turns_are_emitted_in_provider_turn_order(): + transformer = MuseEventTransformer() + + transformer.transform(_event("speechStart", turnId="turn-a")) + transformer.transform(_event("speechStart", turnId="turn-b")) + assert transformer.transform(_event("transcript", turnId="turn-b", transcript="second", final=False)) == () + assert transformer.transform(_event("speechComplete", turnId="turn-a", transcript="first")) == () + released = transformer.transform(_event("speechEnd", turnId="turn-a")) + + assert [(event["type"], event["item_id"]) for event in released] == [ + ("input_audio_buffer.speech_stopped", "turn-a"), + ("conversation.item.input_audio_transcription.completed", "turn-a"), + ("input_audio_buffer.speech_started", "turn-b"), + ("conversation.item.input_audio_transcription.delta", "turn-b"), + ] + assert transformer.transform(_event("speechComplete", turnId="turn-b", transcript="second final")) == () + final_b = transformer.transform(_event("speechEnd", turnId="turn-b")) + assert final_b[0]["type"] == "input_audio_buffer.speech_stopped" + assert final_b[1]["item_id"] == "turn-b" + assert final_b[1]["transcript"] == "second final" + + +def test_committed_item_id_is_used_for_next_provider_turn(): + transformer = MuseEventTransformer() + + previous_item_id, item_id = transformer.commit_item() + started = transformer.transform(_event("speechStart", turnId="provider-turn")) + transformer.transform(_event("speechComplete", turnId="provider-turn", transcript="hello")) + completed = transformer.transform(_event("speechEnd", turnId="provider-turn")) + + assert previous_item_id is None + assert started[0]["item_id"] == item_id + assert completed[-1]["item_id"] == item_id + + +def test_commit_after_speech_start_reuses_active_item_id(): + transformer = MuseEventTransformer() + + started = transformer.transform(_event("speechStart", turnId="provider-turn")) + previous_item_id, item_id = transformer.commit_item() + transformer.transform(_event("speechComplete", turnId="provider-turn", transcript="hello")) + completed = transformer.transform(_event("speechEnd", turnId="provider-turn")) + + assert previous_item_id is None + assert item_id == "provider-turn" + assert started[0]["item_id"] == item_id + assert completed[-1]["item_id"] == item_id + + +def test_speaker_and_positive_audio_progress_deltas_attach_to_next_completion(): + transformer = MuseEventTransformer() + + transformer.transform(_event("audioProgress", audioProcessedMs=1000)) + transformer.transform(_event("audioProgress", audioProcessedMs=750)) + transformer.transform(_event("audioProgress", audioProcessedMs=1600)) + transformer.transform(_event("speaker", turnId=42, label=" Speaker 2 ")) + transformer.transform(_event("speechComplete", turnId=42, transcript="hello")) + completed = transformer.transform(_event("speechEnd", turnId=42)) + + assert completed[-1]["speaker"] == "Speaker 2" + assert completed[-1]["usage"] == {"type": "duration", "seconds": 1.6} + assert transformer.take_unbilled_usage() is None + + +def test_trailing_audio_progress_is_returned_once(): + transformer = MuseEventTransformer() + + transformer.transform(_event("audioProgress", audioProcessedMs=250)) + + assert transformer.take_unbilled_usage() == {"type": "duration", "seconds": 0.25} + assert transformer.take_unbilled_usage() is None + + +def test_completed_turn_tombstone_suppresses_late_duplicates(): + transformer = MuseEventTransformer() + + transformer.transform(_event("speechComplete", turnId="turn-1", transcript="done")) + + assert transformer.transform(_event("speechComplete", turnId="turn-1", transcript="duplicate")) == () + assert transformer.transform(_event("speaker", turnId="turn-1", label="late")) == () + + +def test_provider_error_is_sanitized_and_encodable(): + token = "private-token" + provider_body = f"authorization failed for Bearer {token}" + transformed = MuseEventTransformer().transform( + _event("error", code="AUTH", message=provider_body, request={"accessToken": token}) + ) + + encoded = encode_event(transformed[0]) + assert json.loads(encoded)["error"] == { + "type": "server_error", + "code": "provider_error", + "message": "Meta Muse realtime transcription failed", + } + assert token not in encoded + assert provider_body not in encoded diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 0827bbcdc38..aa3f7d45d84 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -4,7 +4,6 @@ from types import TracebackType from typing import Final from unittest.mock import MagicMock, patch - import pytest import litellm @@ -152,6 +151,85 @@ async def test_vertex_credential_resolution_bounds_a_thread_offloaded_refresh(): assert time.monotonic() - start < 5 +@pytest.mark.asyncio +async def test_meta_realtime_rejects_missing_transcription_intent(monkeypatch: pytest.MonkeyPatch): + def mock_get_llm_provider(model, api_base, api_key): + return model.removeprefix("meta/"), "meta", api_key, api_base + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + + with pytest.raises(ValueError, match="requires intent=transcription"): + await realtime_main._arealtime.__wrapped__( + model="meta/muse-voice-transcribe-1.0", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + query_params={"model": "meta/muse-voice-transcribe-1.0"}, + ) + + +@pytest.mark.asyncio +async def test_meta_realtime_rejects_unsupported_model_before_connecting(monkeypatch: pytest.MonkeyPatch): + def mock_get_llm_provider(model, api_base, api_key): + return model.removeprefix("meta/"), "meta", api_key, api_base + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + + with pytest.raises(ValueError, match="Unsupported Meta realtime model: other-model"): + await realtime_main._arealtime.__wrapped__( + model="meta/other-model", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + query_params={"model": "meta/other-model", "intent": "transcription"}, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("explicit_key", "model_key", "meta_key", "expected"), + [ + ("explicit", "model-env", "meta-env", "explicit"), + (None, "model-env", "meta-env", "model-env"), + (None, None, "meta-env", "meta-env"), + ], +) +async def test_meta_realtime_credential_precedence_is_forwarded_to_handler( + monkeypatch: pytest.MonkeyPatch, + explicit_key: str | None, + model_key: str | None, + meta_key: str | None, + expected: str, +): + captured: dict[str, object] = {} + + def mock_get_llm_provider(model, api_base, api_key): + return model.removeprefix("meta/"), "meta", meta_key, api_base + + def mock_get_secret_str(name: str): + return {"MODEL_API_KEY": model_key, "META_API_KEY": meta_key}.get(name) + + async def mock_async_realtime(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main, "get_secret_str", mock_get_secret_str) + monkeypatch.setattr( + "litellm.llms.meta.realtime.handler.MetaRealtime.async_realtime", + mock_async_realtime, + ) + + await realtime_main._arealtime.__wrapped__( + model="meta/muse-voice-transcribe-1.0", + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + api_key=explicit_key, + query_params={"model": "meta/muse-voice-transcribe-1.0", "intent": "transcription"}, + ) + + assert captured["model"] == "muse-voice-transcribe-1.0" + assert captured["api_key"] == expected + assert captured["query_params"] == {"model": "muse-voice-transcribe-1.0", "intent": "transcription"} + + @pytest.mark.asyncio async def test_arealtime_vertex_branch_resolves_credentials_under_a_bound(monkeypatch): """The wiring half of the regression: the vertex branch of _arealtime must From 1acb994998704e0e1df3f873b52a9226582eb05b Mon Sep 17 00:00:00 2001 From: Young Han Date: Wed, 2 Sep 2026 13:40:35 -0700 Subject: [PATCH 36/55] fix(realtime): bound Muse audio before decoding --- litellm/llms/meta/realtime/handler.py | 10 +++++++++- .../meta/realtime/test_meta_realtime_handler.py | 17 ++++++++++------- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/litellm/llms/meta/realtime/handler.py b/litellm/llms/meta/realtime/handler.py index 9eb98e4ba0b..9dbe3ea8b5d 100644 --- a/litellm/llms/meta/realtime/handler.py +++ b/litellm/llms/meta/realtime/handler.py @@ -268,6 +268,15 @@ class MuseRealtimeAdapter: if not isinstance(audio_value, str): await self._reject("invalid_request_error", "invalid_audio", "Audio must be a base64 string") return + max_backlog_bytes: Final = config.bytes_per_second * _MAX_AUDIO_BACKLOG_SECONDS + max_encoded_bytes: Final = 4 * ((max_backlog_bytes + 2) // 3) + if len(audio_value) > max_encoded_bytes: + await self._reject( + "invalid_request_error", + "audio_backlog_exceeded", + "Audio append exceeds the four-second backlog limit", + ) + return try: audio: Final = base64.b64decode(audio_value, validate=True) except (binascii.Error, ValueError): @@ -278,7 +287,6 @@ class MuseRealtimeAdapter: return if not audio: return - max_backlog_bytes: Final = config.bytes_per_second * _MAX_AUDIO_BACKLOG_SECONDS if len(audio) > max_backlog_bytes: await self._reject( "invalid_request_error", diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py index 3d31b3a5534..14ba972fe14 100644 --- a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py +++ b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py @@ -3,7 +3,7 @@ import base64 import json from collections.abc import Callable from typing import Final -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -246,18 +246,21 @@ async def test_absolute_pacing_delays_only_audio_ahead_of_wall_time(): @pytest.mark.asyncio -async def test_append_larger_than_four_seconds_is_rejected_without_dropping_prefix(): +async def test_append_larger_than_four_seconds_is_rejected_without_decoding(): adapter, provider_ws, _ = await _configured_adapter(rate=16_000) - oversized_pcm: Final = b"\x00\x00" * (16_000 * 4 + 1) + max_pcm_bytes: Final = 16_000 * 2 * 4 + oversized_audio: Final = "A" * (4 * ((max_pcm_bytes + 2) // 3) + 1) + + with patch( # test-quality-ok: proves rejection happens before an attacker-controlled allocation + "litellm.llms.meta.realtime.handler.base64.b64decode" + ) as decode: + await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": oversized_audio})) - await adapter.send( - json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(oversized_pcm).decode()}) - ) error: Final = json.loads(await adapter.recv()) - assert error["error"]["code"] == "audio_backlog_exceeded" assert adapter.close_code == 1008 assert not any(isinstance(frame, bytes) for frame in provider_ws.sent) + decode.assert_not_called() await adapter.close() From 17fde7a261c6c3aff6a2eea72d0a883fd7c17c51 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 11 Sep 2026 19:40:39 -0700 Subject: [PATCH 37/55] refactor(realtime): move Meta Muse Voice onto BaseRealtimeConfig Replace the hand-rolled Meta realtime handler with a MetaRealtimeConfig that plugs into the shared realtime handler and RealTimeStreaming relay. Clients keep speaking the OpenAI realtime wire: session.update, input_audio_buffer.append/commit and the OpenAI transcription events. Unsupported transcription settings are logged and dropped, matching the Gemini realtime precedent, and the Meta-specific session.mode, keywords, language_bias, DIARIZATION and speaker extensions are removed. Drop the MODEL_API_KEY env var in favor of the standard META_API_KEY, remove the private-logging flag so spend logs record the transcript the same way other realtime models do, and add per-second pricing for muse-voice-transcribe-1.0. The relay now sends raw bytes from transform_realtime_request straight to the backend after pace_backend_send, and transcription sessions never trigger response.create. --- .../litellm_core_utils/realtime_streaming.py | 91 +- .../llms/base_llm/realtime/transformation.py | 8 +- litellm/llms/meta/__init__.py | 3 - litellm/llms/meta/realtime/__init__.py | 10 - litellm/llms/meta/realtime/handler.py | 669 --------------- litellm/llms/meta/realtime/transformation.py | 774 ++++++++++-------- ...odel_prices_and_context_window_backup.json | 1 + litellm/realtime_api/main.py | 34 +- litellm/types/llms/meta.py | 58 ++ litellm/types/llms/openai.py | 1 - litellm/utils.py | 4 + model_prices_and_context_window.json | 1 + .../test_realtime_streaming.py | 104 +-- .../realtime/test_meta_realtime_handler.py | 452 ---------- .../test_meta_realtime_transformation.py | 418 ++++++++-- tests/test_litellm/realtime_api/test_main.py | 64 +- 16 files changed, 889 insertions(+), 1803 deletions(-) delete mode 100644 litellm/llms/meta/__init__.py delete mode 100644 litellm/llms/meta/realtime/__init__.py delete mode 100644 litellm/llms/meta/realtime/handler.py create mode 100644 litellm/types/llms/meta.py delete mode 100644 tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 4be7dd6b4ce..06d9241b826 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -19,7 +19,7 @@ from litellm.types.llms.openai import ( OpenAIRealtimeStreamResponseBaseObject, OpenAIRealtimeStreamSessionEvents, ) -from litellm.types.realtime import ALL_DELTA_TYPES, RealtimeInputAudioTranscriptionUsage +from litellm.types.realtime import ALL_DELTA_TYPES from .litellm_logging import Logging as LiteLLMLogging from .realtime_errors import client_close_code, realtime_error_event, websocket_close_reason @@ -116,10 +116,6 @@ class RealtimeEventNormalizer(Protocol): def patch_outgoing_session(self, session: dict) -> dict: ... -class RealtimeUsageProvider(Protocol): - def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None: ... - - DefaultLoggedRealTimeEventTypes: Final = [ "session.created", "response.create", @@ -143,8 +139,6 @@ class RealTimeStreaming: force_transcription_model: str | None = None, event_normalizer: RealtimeEventNormalizer | None = None, logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, - usage_provider: RealtimeUsageProvider | None = None, - exclude_private_content_from_logs: bool = False, ): self.websocket: _ClientWebSocket = websocket self.backend_ws = backend_ws @@ -206,10 +200,6 @@ class RealTimeStreaming: self._is_transcription_session: bool = force_transcription_model is not None # Optional per-provider GA event normalizer (e.g. XAIRealtimeNormalizer). self._event_normalizer = event_normalizer - self._usage_provider: RealtimeUsageProvider | None = ( - usage_provider if usage_provider is not None else provider_config - ) - self._exclude_private_content_from_logs = exclude_private_content_from_logs # Per-connection caps for pre-setup audio frames (message count + total bytes). _MAX_BUFFERED_MESSAGES: int = 200 @@ -247,7 +237,7 @@ class RealTimeStreaming: def _should_store_message( self, - message_obj: dict[str, Any] | OpenAIRealtimeEvents, # mutable-ok: existing realtime event contract + message_obj: dict | OpenAIRealtimeEvents, ) -> bool: _msg_type: Final = message_obj["type"] if "type" in message_obj else None if self.logged_real_time_event_types == "*": @@ -256,54 +246,16 @@ class RealTimeStreaming: return True return False - def _message_for_logging( - self, - message_obj: dict[str, Any], # mutable-ok: existing realtime event contract - ) -> dict[str, Any]: # mutable-ok: logging stores concrete event dictionaries - if not self._exclude_private_content_from_logs: - return message_obj - logged_message: dict[str, Any] = { # mutable-ok: incrementally builds the sanitized event copy - key: message_obj[key] - for key in ( - "type", - "event_id", - "item_id", - "response_id", - "conversation_id", - "session_id", - "content_index", - "output_index", - "model", - "mode", - "usage", - ) - if key in message_obj - } - session: Final = message_obj.get("session") - if isinstance(session, dict): - logged_session: Final[dict[str, Any]] = { # mutable-ok: sanitized JSON session snapshot - key: session[key] for key in ("id", "model", "mode", "type") if key in session - } - if logged_session: - logged_message["session"] = logged_session - return logged_message - def store_message(self, message: str | bytes | dict | OpenAIRealtimeEvents): """Store message in list""" if isinstance(message, bytes): message = message.decode("utf-8") if isinstance(message, dict): # TypedDict union members do not narrow to plain dict for mypy. - parsed_message_obj: dict[str, Any] = cast( # cast-ok: TypedDict events are JSON dictionaries - dict[str, Any], message - ) + message_obj: dict[str, Any] = cast(dict[str, Any], message) else: - parsed_message_obj = cast( # cast-ok: parsed realtime events are JSON dictionaries - dict[str, Any], json.loads(message) - ) - if not self._exclude_private_content_from_logs: - self._collect_tool_calls_from_response_done(parsed_message_obj) - message_obj: Final = self._message_for_logging(parsed_message_obj) + message_obj = cast(dict[str, Any], json.loads(cast(str, message))) + self._collect_tool_calls_from_response_done(cast(dict, message_obj)) if not self._should_store_message(message_obj): return try: @@ -321,8 +273,6 @@ class RealTimeStreaming: def _collect_user_input_from_client_event(self, message: str | dict) -> None: """Extract user text content from client WebSocket events for spend logging.""" - if self._exclude_private_content_from_logs: - return try: if isinstance(message, str): msg_obj = json.loads(message) @@ -359,8 +309,6 @@ class RealTimeStreaming: def _collect_user_input_from_backend_event(self, event_obj: dict | OpenAIRealtimeEvents) -> None: """Extract user voice transcription from backend events for spend logging.""" - if self._exclude_private_content_from_logs: - return try: event_type: Final = event_obj.get("type", "") if event_type == "conversation.item.input_audio_transcription.completed": @@ -416,9 +364,9 @@ class RealTimeStreaming: pass def _flush_unbilled_transcription_usage(self) -> None: - if self._usage_provider is None: + if self.provider_config is None: return - usage: Final = self._usage_provider.unbilled_usage_on_session_close(self.model) + usage: Final = self.provider_config.unbilled_usage_on_session_close(self.model) if usage is None: return flush_event: Final = ( @@ -455,27 +403,12 @@ class RealTimeStreaming: except (AttributeError, TypeError): pass - def _input_for_logging( - self, - message: str | dict, # mutable-ok: existing realtime input contract - ) -> str | dict: # mutable-ok: logging stores concrete event dictionaries - if not self._exclude_private_content_from_logs: - return message - try: - parsed_message: Final[object] = message if isinstance(message, dict) else json.loads(message) - except (json.JSONDecodeError, TypeError): - return {} # mutable-ok: empty JSON logging payload - if not isinstance(parsed_message, dict): - return {} # mutable-ok: empty JSON logging payload - return self._message_for_logging(parsed_message) - def store_input(self, message: str | dict): """Store input message""" - logged_message: Final[str | dict] = self._input_for_logging(message) # mutable-ok: logging payload - self.input_message = logged_message if isinstance(logged_message, dict) else {} + self.input_message = message if isinstance(message, dict) else {} self._collect_user_input_from_client_event(message) if self.logging_obj: - self.logging_obj.pre_call(input=logged_message, api_key="") + self.logging_obj.pre_call(input=message, api_key="") async def log_messages(self): """Log messages in list""" @@ -512,6 +445,12 @@ class RealTimeStreaming: ) sent = False for msg in transformed: + if isinstance(msg, bytes): + await self.provider_config.pace_backend_send(msg) + await self.backend_ws.send(msg) + self._content_sent_after_setup = True + sent = True + continue try: msg_obj = _decode_json_object(msg) except (json.JSONDecodeError, TypeError): diff --git a/litellm/llms/base_llm/realtime/transformation.py b/litellm/llms/base_llm/realtime/transformation.py index cfcde7c6e9e..e44cccc1a62 100644 --- a/litellm/llms/base_llm/realtime/transformation.py +++ b/litellm/llms/base_llm/realtime/transformation.py @@ -1,4 +1,5 @@ from abc import ABC, abstractmethod +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any import httpx @@ -54,9 +55,12 @@ class BaseRealtimeConfig(ABC): message: str, model: str, session_configuration_request: str | None = None, - ) -> list[str]: + ) -> Sequence[str | bytes]: pass + async def pace_backend_send(self, message: bytes) -> None: + return None + def is_setup_message(self, msg_obj: dict) -> bool: return False @@ -79,7 +83,7 @@ class BaseRealtimeConfig(ABC): model: str, logging_session_id: str, session_configuration_request: str | None = None, - ) -> dict | OpenAIRealtimeStreamSessionEvents | None: + ) -> Mapping[str, object] | OpenAIRealtimeStreamSessionEvents | None: """ Optional hook for providers that defer session setup until client `session.update`. diff --git a/litellm/llms/meta/__init__.py b/litellm/llms/meta/__init__.py deleted file mode 100644 index 7c7d32788a2..00000000000 --- a/litellm/llms/meta/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from .realtime import MetaRealtime, MuseRealtimeAdapter - -__all__ = ("MetaRealtime", "MuseRealtimeAdapter") diff --git a/litellm/llms/meta/realtime/__init__.py b/litellm/llms/meta/realtime/__init__.py deleted file mode 100644 index 6398765da24..00000000000 --- a/litellm/llms/meta/realtime/__init__.py +++ /dev/null @@ -1,10 +0,0 @@ -from .handler import MetaRealtime, MuseRealtimeAdapter -from .transformation import MuseEventTransformer, MuseProtocolError, MuseSessionConfig - -__all__ = ( - "MetaRealtime", - "MuseEventTransformer", - "MuseProtocolError", - "MuseRealtimeAdapter", - "MuseSessionConfig", -) diff --git a/litellm/llms/meta/realtime/handler.py b/litellm/llms/meta/realtime/handler.py deleted file mode 100644 index 9dbe3ea8b5d..00000000000 --- a/litellm/llms/meta/realtime/handler.py +++ /dev/null @@ -1,669 +0,0 @@ -from __future__ import annotations - -import asyncio -import base64 -import binascii -import contextlib -import json -import time -import uuid -from collections.abc import Awaitable, Callable, Mapping -from typing import Final, Protocol -from urllib.parse import urlparse, urlunparse - -from pydantic import JsonValue, TypeAdapter, ValidationError - -from litellm._logging import verbose_proxy_logger -from litellm.constants import REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming -from litellm.llms.custom_httpx.http_handler import get_shared_realtime_ssl_context -from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage, RealtimeQueryParams - -from .transformation import ( - MUSE_MODEL, - MuseEventTransformer, - MuseProtocolError, - MuseSessionConfig, - encode_event, - error_event, - parse_session_update, - session_created_event, - session_updated_event, -) - -DEFAULT_MUSE_REALTIME_URL: Final = "wss://api.meta.ai/v1/asr/realtime" -_MAX_AUDIO_BACKLOG_SECONDS: Final = 4 -_MAX_PENDING_PROVIDER_EVENTS: Final = 256 -_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) - - -class _ProviderWebSocket(Protocol): - async def send(self, message: str | bytes) -> None: ... - - async def recv(self, decode: bool | None = None) -> str | bytes: ... - - async def close(self, code: int = 1000, reason: str = "") -> None: ... - - -class _ClientWebSocketExceptions(Protocol): - ConnectionClosed: type[Exception] - - -class _ClientWebSocket(Protocol): - exceptions: _ClientWebSocketExceptions - - @property - def scope(self) -> Mapping[str, object]: ... - - async def send_text(self, data: str) -> None: ... - - async def receive_text(self) -> str: ... - - async def close(self, code: int = 1000, reason: str | None = None) -> None: ... - - -class WebSocketConnect(Protocol): - def __call__( - self, - url: str, - *, - open_timeout: float, - max_size: int | None, - ssl: object | None, - ) -> Awaitable[_ProviderWebSocket]: ... - - -class MuseAdapterError(RuntimeError): - def __init__(self, message: str, *, close_code: int) -> None: - super().__init__(message) - self.close_code: Final = close_code - - -class MuseRealtimeAdapter: - def __init__( - self, - *, - model: str, - api_key: str, - api_base: str | None = None, - timeout: float | None = None, - websocket_connect: WebSocketConnect | None = None, - monotonic: Callable[[], float] = time.monotonic, - sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, - terminate_client: Callable[[int], Awaitable[None]] | None = None, - ) -> None: - if model.removeprefix("meta/") != MUSE_MODEL: - raise ValueError("unsupported Meta realtime model") - self._model: Final = model.removeprefix("meta/") - self._access_token: Final = normalize_access_token(api_key) - self._url: Final = build_muse_realtime_url(api_base) - self._timeout: Final = timeout or 10.0 - self._websocket_connect = websocket_connect - self._monotonic: Final = monotonic - self._sleep: Final = sleep - self._terminate_client: Final = terminate_client - self._provider_ws: _ProviderWebSocket | None = None - self._config: MuseSessionConfig | None = None - self._session_id: str = f"sess_{uuid.uuid4().hex}" - self._events: Final[asyncio.Queue[str | BaseException]] = asyncio.Queue(maxsize=_MAX_PENDING_PROVIDER_EVENTS) - self._events.put_nowait(encode_event(session_created_event(self._model, self._session_id))) - self._transformer: Final = MuseEventTransformer() - self._audio_condition: Final = asyncio.Condition() - self._pending_audio: bytearray = bytearray() - self._audio_generation: int = 0 - self._flush_requested: bool = False - self._end_requested: bool = False - self._end_stream_sent: bool = False - self._audio_consumed: bool = False - self._closed: bool = False - self._resources_closed: bool = False - self._sender_task: asyncio.Task[None] | None = None - self._receiver_task: asyncio.Task[None] | None = None - self.close_code: int = 1000 - self.close_reason: str = "Session closed" - - async def send(self, message: str | bytes) -> None: - if self._closed: - raise MuseAdapterError("Meta Muse realtime session is closed", close_code=self.close_code) - if isinstance(message, bytes): - await self._reject("invalid_request_error", "invalid_event", "Client events must be JSON text") - return - try: - event: Final = _parse_client_event(message) - event_type: Final = event.get("type") - if event_type in ("session.update", "transcription_session.update"): - await self._handle_session_update(message) - return - if event_type == "input_audio_buffer.append": - await self._handle_audio_append(event) - return - if event_type == "input_audio_buffer.clear": - await self._clear_audio() - return - if event_type == "input_audio_buffer.commit": - await self._commit_audio() - return - if event_type == "input_audio_buffer.end": - await self._end_audio() - return - await self._emit( - error_event( - "invalid_request_error", - "unsupported_event", - f"Event type {event_type!r} is not supported for Meta Muse transcription", - ) - ) - except MuseProtocolError as exc: - await self._reject("invalid_request_error", "invalid_event", str(exc)) - - async def recv(self, decode: bool | None = None) -> str | bytes: - event: Final = await self._events.get() - if isinstance(event, BaseException): - close_code: Final = _exception_close_code(event) if isinstance(event, Exception) else 1011 - if self._terminate_client is not None: - await self._terminate_client(close_code) - raise event - return event.encode("utf-8") if decode is False else event - - async def close(self, code: int = 1000, reason: str = "") -> None: - if self._resources_closed: - return - self._closed = True - self._resources_closed = True - self.close_code = sanitize_close_code(code) - self.close_reason = safe_close_reason(self.close_code) - async with self._audio_condition: - self._end_requested = True - self._audio_condition.notify_all() - tasks: Final = tuple(task for task in (self._sender_task, self._receiver_task) if task is not None) - for task in tasks: - task.cancel() - if tasks: - await asyncio.gather(*tasks, return_exceptions=True) - provider_ws: Final = self._provider_ws - if provider_ws is not None: - with contextlib.suppress(Exception): - await provider_ws.close(code=self.close_code, reason=self.close_reason) - - def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None: - return self._transformer.take_unbilled_usage() - - async def _handle_session_update(self, message: str) -> None: - config: Final = parse_session_update(message, self._model) - if self._config is not None: - if config != self._config: - await self._reject( - "invalid_request_error", - "session_configuration_locked", - "Meta Muse session configuration cannot change after setup", - ) - return - await self._emit(session_updated_event(config, self._session_id)) - return - await self._connect(config) - - async def _connect(self, config: MuseSessionConfig) -> None: - connector: Final = self._websocket_connect or _default_websocket_connect - last_error: Exception | None = None # rebind-ok: records the latest bounded handshake attempt - for attempt in range(2): - provider_ws: _ProviderWebSocket | None = None - try: - provider_ws = await connector( - self._url, - open_timeout=self._timeout, - max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, - ssl=_ssl_config(self._url), - ) - await provider_ws.send(json.dumps(config.handshake(self._access_token), separators=(",", ":"))) - raw_ack: str | bytes = await asyncio.wait_for( # rebind-ok: one response per handshake attempt - provider_ws.recv(), timeout=self._timeout - ) - session_id: str = _parse_handshake_ack(raw_ack) # rebind-ok: one ID per handshake attempt - self._provider_ws = provider_ws - self._config = config - self._transformer.configure(config) - self._session_id = session_id - self._sender_task = asyncio.create_task(self._send_audio(), name="meta-muse-realtime-send") - self._receiver_task = asyncio.create_task(self._receive_events(), name="meta-muse-realtime-receive") - await self._emit(session_updated_event(config, session_id)) - return - except asyncio.CancelledError: - if provider_ws is not None: - with contextlib.suppress(Exception): - await provider_ws.close() - raise - except Exception as exc: # noqa: BLE001 # connector implementations expose heterogeneous transport errors - last_error = exc - if provider_ws is not None: - with contextlib.suppress(Exception): - await provider_ws.close() - close_code: int = _exception_close_code(exc) # rebind-ok: classified per handshake attempt - retryable_transport_error: bool = not isinstance( # rebind-ok: classified per handshake attempt - exc, (MuseAdapterError, MuseProtocolError) - ) - if attempt == 0 and retryable_transport_error and close_code in (1011, 1013): - continue - self.close_code = close_code - self.close_reason = safe_close_reason(close_code) - await self._emit( - error_event( - "server_error" if close_code != 1008 else "invalid_request_error", - "provider_connection_error", - "Meta Muse realtime handshake failed", - ) - ) - await self._events.put(MuseAdapterError("Meta Muse realtime handshake failed", close_code=close_code)) - await self._mark_terminated(close_code) - return - assert last_error is not None - raise MuseAdapterError("Meta Muse realtime handshake failed", close_code=1011) - - async def _handle_audio_append(self, event: Mapping[str, JsonValue]) -> None: - config: Final = self._require_configured() - if self._end_requested or self._end_stream_sent: - await self._reject("invalid_request_error", "input_ended", "Audio input has already ended") - return - audio_value: Final = event.get("audio") - if not isinstance(audio_value, str): - await self._reject("invalid_request_error", "invalid_audio", "Audio must be a base64 string") - return - max_backlog_bytes: Final = config.bytes_per_second * _MAX_AUDIO_BACKLOG_SECONDS - max_encoded_bytes: Final = 4 * ((max_backlog_bytes + 2) // 3) - if len(audio_value) > max_encoded_bytes: - await self._reject( - "invalid_request_error", - "audio_backlog_exceeded", - "Audio append exceeds the four-second backlog limit", - ) - return - try: - audio: Final = base64.b64decode(audio_value, validate=True) - except (binascii.Error, ValueError): - await self._reject("invalid_request_error", "invalid_audio", "Audio must be valid base64") - return - if len(audio) % 2: - await self._reject("invalid_request_error", "invalid_audio", "PCM16 audio must contain complete samples") - return - if not audio: - return - if len(audio) > max_backlog_bytes: - await self._reject( - "invalid_request_error", - "audio_backlog_exceeded", - "Audio append exceeds the four-second Muse backlog limit", - ) - return - async with self._audio_condition: - await self._audio_condition.wait_for( - lambda: self._closed or len(self._pending_audio) + len(audio) <= max_backlog_bytes - ) - if self._closed: - raise MuseAdapterError("Meta Muse realtime session is closed", close_code=self.close_code) - self._pending_audio.extend(audio) - self._audio_condition.notify_all() - - async def _clear_audio(self) -> None: - self._require_configured() - async with self._audio_condition: - self._pending_audio.clear() - self._audio_generation += 1 - self._flush_requested = False - self._audio_condition.notify_all() - await self._emit( - { # mutable-ok: OpenAI-compatible JSON event - "type": "input_audio_buffer.cleared", - "event_id": f"event_{uuid.uuid4().hex}", - } - ) - - async def _commit_audio(self) -> None: - config: Final = self._require_configured() - previous_item_id, item_id = self._transformer.commit_item() - async with self._audio_condition: - self._flush_requested = True - if config.mode == "PUSH_TO_TALK": - self._end_requested = True - self._audio_condition.notify_all() - await self._emit( - { # mutable-ok: OpenAI-compatible JSON event - "type": "input_audio_buffer.committed", - "event_id": f"event_{uuid.uuid4().hex}", - "previous_item_id": previous_item_id, - "item_id": item_id, - } - ) - - async def _end_audio(self) -> None: - self._require_configured() - async with self._audio_condition: - self._flush_requested = True - self._end_requested = True - self._audio_condition.notify_all() - - async def _send_audio(self) -> None: - config: Final = self._require_configured() - provider_ws: Final = self._require_provider_ws() - pacing_origin: float | None = None # rebind-ok: initialized when the first packet is ready - sent_duration: float = 0.0 # rebind-ok: absolute pacing clock advances after each packet - try: - while True: - packet, pacing_origin, ended = await self._next_audio_packet( - config, - pacing_origin, - sent_duration, - ) - if ended: - break - if packet is None: - continue - await provider_ws.send(packet) - self._audio_consumed = True - sent_duration += len(packet) / config.bytes_per_second - await self._send_end_stream() - except asyncio.CancelledError: - raise - except Exception as exc: # noqa: BLE001 # WebSocket implementations expose heterogeneous transport errors - await self._fail_provider(exc, phase="audio send") - - async def _next_audio_packet( - self, - config: MuseSessionConfig, - pacing_origin: float | None, - sent_duration: float, - ) -> tuple[bytes | None, float | None, bool]: - async with self._audio_condition: - await self._audio_condition.wait_for( - lambda: ( - self._closed - or len(self._pending_audio) >= config.packet_bytes - or (self._flush_requested and bool(self._pending_audio)) - or (self._end_requested and not self._pending_audio) - ) - ) - if self._closed or (self._end_requested and not self._pending_audio): - return None, pacing_origin, True - packet_size: Final = min(config.packet_bytes, len(self._pending_audio)) - if packet_size < config.packet_bytes and not self._flush_requested: - return None, pacing_origin, False - generation: Final = self._audio_generation - current_time: Final = self._monotonic() - effective_origin: Final = ( - current_time - sent_duration - if pacing_origin is None or current_time > pacing_origin + sent_duration - else pacing_origin - ) - deadline: Final = effective_origin + sent_duration - delay: Final = deadline - self._monotonic() - if delay > 0: - await self._sleep(delay) - async with self._audio_condition: - if generation != self._audio_generation: - return None, effective_origin, False - actual_size: Final = min(packet_size, len(self._pending_audio)) - packet: Final = bytes(self._pending_audio[:actual_size]) - del self._pending_audio[:actual_size] - if not self._pending_audio: - self._flush_requested = False - self._audio_condition.notify_all() - return packet or None, effective_origin, False - - async def _send_end_stream(self) -> None: - if self._end_stream_sent: - return - provider_ws: Final = self._require_provider_ws() - await provider_ws.send('{"type":"endStream"}') - self._end_stream_sent = True - - async def _receive_events(self) -> None: - provider_ws: Final = self._require_provider_ws() - try: - while True: - raw: str | bytes = await provider_ws.recv() # rebind-ok: one provider frame per iteration - if not isinstance(raw, str): - raise MuseProtocolError("provider returned a non-text event") - for event in self._transformer.transform(raw): - await self._emit(event) - except asyncio.CancelledError: - raise - except Exception as exc: # noqa: BLE001 # provider close exceptions vary by WebSocket implementation - close_code: Final = _exception_close_code(exc) - if close_code == 1000 and self._end_stream_sent: - await self._mark_terminated(1000) - await self._events.put(MuseAdapterError("Meta Muse realtime session completed", close_code=1000)) - return - failure: Final = MuseAdapterError( - "Meta Muse realtime closed before input ended", - close_code=1011 if close_code == 1000 else close_code, - ) - await self._fail_provider(failure, phase="receive") - - async def _fail_provider(self, exc: Exception, *, phase: str) -> None: - close_code: Final = _exception_close_code(exc) - self.close_code = close_code - self.close_reason = safe_close_reason(close_code) - await self._emit( - error_event( - "server_error", - "provider_connection_error", - f"Meta Muse realtime {phase} failed", - ) - ) - await self._events.put(MuseAdapterError(f"Meta Muse realtime {phase} failed", close_code=close_code)) - await self._mark_terminated(close_code) - - async def _mark_terminated(self, close_code: int) -> None: - self._closed = True - self.close_code = sanitize_close_code(close_code) - self.close_reason = safe_close_reason(self.close_code) - async with self._audio_condition: - self._audio_condition.notify_all() - - async def _terminate(self, close_code: int) -> None: - await self._mark_terminated(close_code) - if self._terminate_client is not None: - await self._terminate_client(self.close_code) - - async def _reject(self, error_type: str, code: str, message: str) -> None: - self.close_code = 1008 - self.close_reason = safe_close_reason(1008) - await self._emit(error_event(error_type, code, message)) - await self._events.put(MuseAdapterError(message, close_code=1008)) - await self._mark_terminated(1008) - - async def _emit(self, event: Mapping[str, object]) -> None: - await self._events.put(encode_event(event)) - - def _require_configured(self) -> MuseSessionConfig: - if self._config is None: - raise MuseProtocolError("send session.update before audio events") - return self._config - - def _require_provider_ws(self) -> _ProviderWebSocket: - if self._provider_ws is None: - raise MuseProtocolError("Meta Muse provider connection is not ready") - return self._provider_ws - - -class MetaRealtime: - async def async_realtime( - self, - model: str, - websocket: _ClientWebSocket, - logging_obj: LiteLLMLogging, - api_base: str | None = None, - api_key: str | None = None, - client: object | None = None, - timeout: float | None = None, - query_params: RealtimeQueryParams | None = None, - user_api_key_dict: object | None = None, - litellm_metadata: Mapping[str, object] | None = None, - websocket_connect: WebSocketConnect | None = None, - monotonic: Callable[[], float] = time.monotonic, - sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, - **kwargs: object, # kwargs-ok: realtime dispatcher forwards provider-neutral options - ) -> None: - if api_key is None or not api_key.strip(): - await _send_client_error_and_close(websocket, "Meta Model API key is required") - return - try: - adapter: Final = MuseRealtimeAdapter( - model=model, - api_key=api_key, - api_base=api_base, - timeout=timeout, - websocket_connect=websocket_connect, - monotonic=monotonic, - sleep=sleep, - terminate_client=lambda code: _close_client(websocket, code), - ) - except ValueError: - await _send_client_error_and_close(websocket, "Invalid Meta Muse realtime configuration") - return - realtime_streaming: Final = RealTimeStreaming( - websocket, - adapter, # pyright: ignore[reportArgumentType] # raw adapter intentionally matches the websocket surface - logging_obj, - model=model, - user_api_key_dict=user_api_key_dict, - request_data={ # mutable-ok: relay request metadata payload - "litellm_metadata": dict(litellm_metadata or {}) # mutable-ok: relay owns its metadata copy - }, - force_transcription_model=model, - usage_provider=adapter, - exclude_private_content_from_logs=True, - ) - try: - await realtime_streaming.bidirectional_forward() - except MuseAdapterError as exc: - adapter.close_code = exc.close_code - adapter.close_reason = safe_close_reason(exc.close_code) - except Exception: # noqa: BLE001 # relay errors are normalized before closing the accepted client socket - adapter.close_code = 1011 - adapter.close_reason = safe_close_reason(1011) - verbose_proxy_logger.exception("Meta Muse realtime session failed") - finally: - await adapter.close(code=adapter.close_code) - await _close_client(websocket, adapter.close_code) - - -def normalize_access_token(api_key: str) -> str: - stripped: Final = api_key.strip() - if not stripped: - raise ValueError("Meta Model API key is required") - parts: Final = stripped.split(None, 1) - if parts[0].casefold() == "bearer": - if len(parts) != 2 or not parts[1].strip(): - raise ValueError("Meta Model API key must include a token after Bearer") - return f"Bearer {parts[1].strip()}" - return f"Bearer {stripped}" - - -def build_muse_realtime_url(api_base: str | None) -> str: - if api_base is None: - return DEFAULT_MUSE_REALTIME_URL - parsed: Final = urlparse(api_base.strip()) - scheme: Final = "wss" if parsed.scheme == "https" else parsed.scheme - if ( - scheme != "wss" - or not parsed.hostname - or parsed.username is not None - or parsed.password is not None - or parsed.fragment - ): - raise ValueError("Meta api_base must be an absolute wss:// or https:// URL without credentials or a fragment") - netloc: Final = f"{parsed.hostname}:{parsed.port}" if parsed.port is not None else parsed.hostname - return urlunparse((scheme, netloc, "/v1/asr/realtime", "", "", "")) - - -def sanitize_close_code(code: int | None) -> int: - if code is not None and code in (1000, 1008, 1011, 1013): - return code - return 1011 - - -def safe_close_reason(code: int) -> str: - return { # mutable-ok: immutable-by-convention close-reason lookup - 1000: "Session closed", - 1008: "Invalid realtime transcription request", - 1011: "Realtime transcription service error", - 1013: "Realtime transcription service unavailable", - }.get(code, "Realtime transcription service error") - - -def _parse_client_event(payload: str) -> Mapping[str, JsonValue]: - try: - value: Final = _JSON_ADAPTER.validate_json(payload) - except ValidationError: - raise MuseProtocolError("invalid JSON object") from None - if not isinstance(value, dict): - raise MuseProtocolError("message must be a JSON object") - event_type: Final = value.get("type") - if not isinstance(event_type, str) or not event_type: - raise MuseProtocolError("message type must be a non-empty string") - return value - - -def _parse_handshake_ack(raw: str | bytes) -> str: - if not isinstance(raw, str): - raise MuseProtocolError("provider returned a non-text handshake response") - message: Final = _parse_json_object(raw) - if message.get("type") == "error": - raise MuseAdapterError("Meta Muse realtime handshake was rejected", close_code=1008) - session_id: Final = message.get("sessionId") - if not isinstance(session_id, str) or not session_id.strip(): - raise MuseProtocolError("provider returned an invalid handshake response") - return session_id.strip() - - -def _parse_json_object(payload: str) -> Mapping[str, JsonValue]: - try: - value: Final = _JSON_ADAPTER.validate_json(payload) - except ValidationError: - raise MuseProtocolError("invalid provider JSON object") from None - if not isinstance(value, dict): - raise MuseProtocolError("provider message must be a JSON object") - return value - - -def _exception_close_code(exc: Exception) -> int: - code: Final = getattr(exc, "code", None) - if isinstance(exc, MuseAdapterError): - return sanitize_close_code(exc.close_code) - return sanitize_close_code(code if isinstance(code, int) else None) - - -def _ssl_config(url: str) -> object | None: - if not url.startswith("wss://"): - return None - config: Final = get_shared_realtime_ssl_context() - return True if config is False else config - - -async def _default_websocket_connect( - url: str, - *, - open_timeout: float, - max_size: int | None, - ssl: object | None, -) -> _ProviderWebSocket: - import websockets - - connection: Final = await websockets.connect( - url, - open_timeout=open_timeout, - max_size=max_size, - ssl=ssl, # pyright: ignore[reportArgumentType] # shared SSL helper returns the library-supported union - ) - return connection - - -async def _send_client_error_and_close(websocket: _ClientWebSocket, message: str) -> None: - with contextlib.suppress(Exception): - await websocket.send_text(encode_event(error_event("invalid_request_error", "invalid_configuration", message))) - await _close_client(websocket, 1008) - - -async def _close_client(websocket: _ClientWebSocket, code: int) -> None: - with contextlib.suppress(Exception): - await websocket.close(code=sanitize_close_code(code), reason=safe_close_reason(sanitize_close_code(code))) diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py index f37295892f8..096f7c8f1fe 100644 --- a/litellm/llms/meta/realtime/transformation.py +++ b/litellm/llms/meta/realtime/transformation.py @@ -1,20 +1,51 @@ -from __future__ import annotations - +import asyncio +import base64 +import binascii import json import math -import uuid -from collections import OrderedDict, deque -from collections.abc import Mapping +import time +from collections import deque +from collections.abc import Awaitable, Callable, Iterator, Mapping from dataclasses import dataclass -from typing import Final, Literal, TypeAlias +from types import MappingProxyType +from typing import Final, Literal, cast +from urllib.parse import urlparse, urlunparse from pydantic import JsonValue, TypeAdapter, ValidationError -from litellm.types.realtime import RealtimeInputAudioTranscriptionUsage +from litellm import verbose_logger +from litellm._uuid import uuid +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.meta import ( + MuseAudioEncoding, + MuseHandshake, + MuseMode, + MuseSampleRate, + MuseSessionCreatedEvent, + MuseTranscriptionSession, + MuseTranscriptionSettings, + MuseTurnDetection, +) +from litellm.types.llms.openai import ( + OpenAIRealtimeEvents, + OpenAIRealtimeInputAudioBufferSpeechEvent, + OpenAIRealtimeInputAudioTranscriptionCompleted, + OpenAIRealtimeInputAudioTranscriptionDelta, +) +from litellm.types.realtime import ( + RealtimeErrorDetail, + RealtimeErrorEvent, + RealtimeInputAudioTranscriptionDurationUsage, + RealtimeInputAudioTranscriptionUsage, + RealtimeResponseTransformInput, + RealtimeResponseTypedDict, +) MUSE_MODEL: Final = "muse-voice-transcribe-1.0" +DEFAULT_MUSE_REALTIME_URL: Final = "wss://api.meta.ai/v1/asr/realtime" SUPPORTED_SAMPLE_RATES: Final = frozenset((16_000, 24_000)) -SUPPORTED_MODES: Final = frozenset(("PUSH_TO_TALK", "ENDPOINTING", "DIARIZATION")) SUPPORTED_LANGUAGES: Final = ( "Arabic", "Bengali", @@ -42,40 +73,46 @@ SUPPORTED_LANGUAGES: Final = ( "Turkish", "Vietnamese", ) -_LANGUAGE_NAMES: Final = { # mutable-ok: immutable-by-convention language lookup table - language.casefold(): language for language in SUPPORTED_LANGUAGES -} -_LANGUAGE_CODES: Final = { # mutable-ok: immutable-by-convention language lookup table - "ar": "Arabic", - "bn": "Bengali", - "de": "German", - "en": "English", - "es": "Spanish", - "fil": "Tagalog", - "fr": "French", - "he": "Hebrew", - "hi": "Hindi", - "id": "Indonesian", - "it": "Italian", - "iw": "Hebrew", - "ja": "Japanese", - "kn": "Kannada", - "ko": "Korean", - "ms": "Malay", - "mr": "Marathi", - "nl": "Dutch", - "pl": "Polish", - "pt": "Portuguese", - "ta": "Tamil", - "te": "Telugu", - "th": "Thai", - "tl": "Tagalog", - "tr": "Turkish", - "vi": "Vietnamese", - "zh": "Mandarin Chinese", -} +_LANGUAGE_NAMES: Final = MappingProxyType({language.casefold(): language for language in SUPPORTED_LANGUAGES}) +_LANGUAGE_CODES: Final = MappingProxyType( + { + "ar": "Arabic", + "bn": "Bengali", + "de": "German", + "en": "English", + "es": "Spanish", + "fil": "Tagalog", + "fr": "French", + "he": "Hebrew", + "hi": "Hindi", + "id": "Indonesian", + "it": "Italian", + "iw": "Hebrew", + "ja": "Japanese", + "kn": "Kannada", + "ko": "Korean", + "ms": "Malay", + "mr": "Marathi", + "nl": "Dutch", + "pl": "Polish", + "pt": "Portuguese", + "ta": "Tamil", + "te": "Telugu", + "th": "Thai", + "tl": "Tagalog", + "tr": "Turkish", + "vi": "Vietnamese", + "zh": "Mandarin Chinese", + } +) +_SUPPORTED_TRANSCRIPTION_KEYS: Final = frozenset(("model", "language")) +_MAX_AUDIO_BACKLOG_SECONDS: Final = 4 +_PACKET_MS: Final = 80 +_END_STREAM: Final = '{"type":"endStream"}' +_PROVIDER_ERROR_MESSAGE: Final = "Meta Muse realtime transcription failed" _JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) -OpenAIEvent: TypeAlias = Mapping[str, object] +_EMPTY_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) +_SERVER_VAD: Final[MuseTurnDetection] = {"type": "server_vad"} class MuseProtocolError(ValueError): @@ -85,13 +122,12 @@ class MuseProtocolError(ValueError): @dataclass(frozen=True, slots=True) class MuseSessionConfig: model: str - mode: Literal["PUSH_TO_TALK", "ENDPOINTING", "DIARIZATION"] - sample_rate: Literal[16000, 24000] - keywords: tuple[str, ...] + mode: MuseMode + sample_rate: MuseSampleRate language_bias: tuple[str, ...] @property - def audio_encoding(self) -> Literal["PCM_16KHZ", "PCM_24KHZ"]: + def audio_encoding(self) -> MuseAudioEncoding: return "PCM_16KHZ" if self.sample_rate == 16_000 else "PCM_24KHZ" @property @@ -100,66 +136,52 @@ class MuseSessionConfig: @property def packet_bytes(self) -> int: - return self.bytes_per_second * 80 // 1000 + return self.bytes_per_second * _PACKET_MS // 1000 - def handshake(self, access_token: str) -> Mapping[str, object]: - base: Final[Mapping[str, object]] = { # mutable-ok: JSON wire payload - "mode": self.mode, - "authorization": {"accessToken": access_token}, # mutable-ok: JSON wire payload + @property + def max_encoded_append_bytes(self) -> int: + return 4 * ((self.bytes_per_second * _MAX_AUDIO_BACKLOG_SECONDS + 2) // 3) + + def handshake(self, access_token: str) -> MuseHandshake: + base: Final[MuseHandshake] = { + "authorization": {"accessToken": access_token}, "audioEncoding": self.audio_encoding, "model": self.model, + "mode": self.mode, "partialMode": "CUMULATIVE", "emitAudioProgress": True, } - payload: dict[str, object] = dict(base) # mutable-ok: incrementally builds JSON wire payload - if self.keywords: - payload["keywords"] = list(self.keywords) # mutable-ok: JSON arrays require concrete lists - if self.language_bias: - payload["languageBias"] = list(self.language_bias) # mutable-ok: JSON arrays require concrete lists - return payload + if not self.language_bias: + return base + biased: Final[MuseHandshake] = {**base, "languageBias": self.language_bias} + return biased - def openai_session(self, session_id: str) -> Mapping[str, object]: - turn_detection: Final[Mapping[str, object] | None] = ( - None if self.mode == "PUSH_TO_TALK" else {"type": "server_vad"} # mutable-ok: JSON wire payload - ) - transcription: dict[str, object] = { # mutable-ok: incrementally builds JSON wire payload - "model": self.model, - } - if self.language_bias: - transcription["language"] = self.language_bias[0] - transcription["language_bias"] = list( # mutable-ok: JSON arrays require concrete lists - self.language_bias - ) - if self.keywords: - transcription["keywords"] = list(self.keywords) # mutable-ok: JSON arrays require concrete lists - return { # mutable-ok: JSON wire payload + def openai_session(self, session_id: str) -> MuseTranscriptionSession: + session: Final[MuseTranscriptionSession] = { "id": session_id, "object": "realtime.transcription_session", "type": "transcription", - "model": self.model, - "audio": { # mutable-ok: JSON wire payload - "input": { # mutable-ok: JSON wire payload - "format": {"type": "audio/pcm", "rate": self.sample_rate}, # mutable-ok: JSON wire payload - "transcription": transcription, - "turn_detection": turn_detection, + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": self.sample_rate}, + "transcription": self._transcription_settings(), + "turn_detection": None if self.mode == "PUSH_TO_TALK" else _SERVER_VAD, } }, } + return session + + def _transcription_settings(self) -> MuseTranscriptionSettings: + base: Final[MuseTranscriptionSettings] = {"model": self.model} + if not self.language_bias: + return base + localized: Final[MuseTranscriptionSettings] = {**base, "language": self.language_bias[0]} + return localized -@dataclass(slots=True) -class _TurnState: - item_id: str | None = None - started: bool = False - start_emitted: bool = False - latest_partial: str | None = None - emitted_partial: str = "" - final_text: str | None = None - completed_signal: bool = False - completed_emitted: bool = False - stopped: bool = False - stopped_emitted: bool = False - speaker: str | None = None +_DEFAULT_SESSION_CONFIG: Final = MuseSessionConfig( + model=MUSE_MODEL, mode="ENDPOINTING", sample_rate=24_000, language_bias=() +) def _json_object(payload: str) -> Mapping[str, JsonValue]: @@ -174,7 +196,7 @@ def _json_object(payload: str) -> Mapping[str, JsonValue]: def _mapping(value: JsonValue | None, name: str) -> Mapping[str, JsonValue]: if value is None: - return {} # mutable-ok: empty JSON object + return _EMPTY_OBJECT if not isinstance(value, dict): raise MuseProtocolError(f"{name} must be an object") return value @@ -192,6 +214,10 @@ def _normalize_model(model: str) -> str: return model.removeprefix("meta/").strip() +def _event_id() -> str: + return f"event_{uuid.uuid4().hex}" + + def normalize_language(language: str) -> str: value: Final = language.strip() if not value: @@ -206,26 +232,36 @@ def normalize_language(language: str) -> str: return mapped_name -def _normalize_string_sequence(value: JsonValue | None, name: str) -> tuple[str, ...]: - if value is None: - return () - if not isinstance(value, list): - raise MuseProtocolError(f"{name} must be an array of strings") - normalized: list[str] = [] # mutable-ok: deduplicates validated language hints before freezing - for entry in value: - if not isinstance(entry, str) or not entry.strip(): - raise MuseProtocolError(f"{name} entries must be non-empty strings") - item: str = entry.strip() # rebind-ok: normalized once for each hint - if item not in normalized: - normalized.append(item) - return tuple(normalized) +def normalize_access_token(api_key: str) -> str: + stripped: Final = api_key.strip() + if not stripped: + raise ValueError("Meta API key is required") + parts: Final = stripped.split(None, 1) + if parts[0].casefold() != "bearer": + return f"Bearer {stripped}" + if len(parts) != 2 or not parts[1].strip(): + raise ValueError("Meta API key must include a token after Bearer") + return f"Bearer {parts[1].strip()}" -def _normalize_language_sequence(value: JsonValue | None) -> tuple[str, ...]: - return tuple(dict.fromkeys(normalize_language(item) for item in _normalize_string_sequence(value, "language_bias"))) +def build_muse_realtime_url(api_base: str | None) -> str: + if api_base is None: + return DEFAULT_MUSE_REALTIME_URL + parsed: Final = urlparse(api_base.strip()) + scheme: Final = "wss" if parsed.scheme == "https" else parsed.scheme + if ( + scheme != "wss" + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or parsed.fragment + ): + raise ValueError("Meta api_base must be an absolute wss:// or https:// URL without credentials or a fragment") + netloc: Final = f"{parsed.hostname}:{parsed.port}" if parsed.port is not None else parsed.hostname + return urlunparse((scheme, netloc, "/v1/asr/realtime", "", "", "")) -def _parse_sample_rate(session: Mapping[str, JsonValue]) -> Literal[16000, 24000]: +def _parse_sample_rate(session: Mapping[str, JsonValue]) -> MuseSampleRate: beta_format: Final = session.get("input_audio_format") audio: Final = _mapping(session.get("audio"), "session.audio") audio_input: Final = _mapping(audio.get("input"), "session.audio.input") @@ -251,22 +287,10 @@ def _parse_sample_rate(session: Mapping[str, JsonValue]) -> Literal[16000, 24000 rate: Final = format_mapping.get("rate", 24_000) if isinstance(rate, bool) or not isinstance(rate, int) or rate not in SUPPORTED_SAMPLE_RATES: raise MuseProtocolError("Muse Voice supports PCM16 at 16000 Hz or 24000 Hz") - return rate + return 16_000 if rate == 16_000 else 24_000 -def _parse_mode( - session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue] -) -> Literal["PUSH_TO_TALK", "ENDPOINTING", "DIARIZATION"]: - explicit: Final = session.get("mode") - if explicit is not None: - if not isinstance(explicit, str) or explicit.upper() not in SUPPORTED_MODES: - raise MuseProtocolError("unsupported Muse Voice mode") - normalized_mode: Final = explicit.upper() - if normalized_mode == "PUSH_TO_TALK": - return "PUSH_TO_TALK" - if normalized_mode == "DIARIZATION": - return "DIARIZATION" - return "ENDPOINTING" +def _parse_mode(session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue]) -> MuseMode: turn_detection_present: Final = "turn_detection" in session or "turn_detection" in audio_input turn_detection: Final = session.get("turn_detection", audio_input.get("turn_detection")) if turn_detection_present and turn_detection is None: @@ -286,8 +310,7 @@ def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig session: Final = _mapping(message.get("session"), "session") if not session: raise MuseProtocolError("session.update requires a session object") - session_type: Final = session.get("type") - if session_type not in (None, "transcription", "realtime"): + if session.get("type") not in (None, "transcription", "realtime"): raise MuseProtocolError("Muse Voice supports transcription sessions only") audio: Final = _mapping(session.get("audio"), "session.audio") audio_input: Final = _mapping(audio.get("input"), "session.audio.input") @@ -299,88 +322,144 @@ def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig beta_transcription if beta_transcription is not None else ga_transcription, "input audio transcription", ) + unsupported: Final = tuple(sorted(key for key in transcription if key not in _SUPPORTED_TRANSCRIPTION_KEYS)) + if unsupported: + verbose_logger.warning("Meta realtime: dropping unsupported transcription settings %s", unsupported) requested_model: Final = _string(transcription.get("model"), "transcription model") normalized_model: Final = _normalize_model(expected_model) if normalized_model != MUSE_MODEL: raise MuseProtocolError("unsupported Meta realtime model") if requested_model is not None and _normalize_model(requested_model) != normalized_model: raise MuseProtocolError("realtime session model cannot be changed") - language_value: Final = _string(transcription.get("language"), "language") - explicit_bias: Final = _normalize_language_sequence(transcription.get("language_bias")) - language_bias: Final = tuple( - dict.fromkeys((normalize_language(language_value), *explicit_bias)) - if language_value is not None - else explicit_bias - ) - keywords: Final = _normalize_string_sequence(transcription.get("keywords"), "keywords") + language: Final = _string(transcription.get("language"), "language") return MuseSessionConfig( model=normalized_model, mode=_parse_mode(session, audio_input), sample_rate=_parse_sample_rate(session), - keywords=keywords, - language_bias=language_bias, + language_bias=() if language is None else (normalize_language(language),), ) -def session_created_event(model: str, session_id: str) -> OpenAIEvent: - normalized_model: Final = _normalize_model(model) - default_config: Final = MuseSessionConfig( - model=normalized_model, - mode="ENDPOINTING", - sample_rate=24_000, - keywords=(), - language_bias=(), - ) - return { # mutable-ok: OpenAI-compatible JSON event +def session_created_event(config: MuseSessionConfig, session_id: str) -> MuseSessionCreatedEvent: + event: Final[MuseSessionCreatedEvent] = { "type": "session.created", - "event_id": f"event_{uuid.uuid4().hex}", - "session": default_config.openai_session(session_id), - } - - -def session_updated_event(config: MuseSessionConfig, session_id: str) -> OpenAIEvent: - return { # mutable-ok: OpenAI-compatible JSON event - "type": "session.updated", - "event_id": f"event_{uuid.uuid4().hex}", + "event_id": _event_id(), "session": config.openai_session(session_id), } + return event -def error_event(error_type: str, code: str, message: str) -> OpenAIEvent: - return { # mutable-ok: OpenAI-compatible JSON event - "type": "error", - "event_id": f"event_{uuid.uuid4().hex}", - "error": { # mutable-ok: nested OpenAI-compatible error object - "type": error_type, - "code": code, - "message": message, - }, +def error_event(message: str) -> OpenAIRealtimeEvents: + detail: Final[RealtimeErrorDetail] = {"type": "server_error", "message": message} + event: Final[RealtimeErrorEvent] = {"type": "error", "error": detail} + return cast(OpenAIRealtimeEvents, event) # cast-ok: the union has no error member; the relay only serializes it + + +def _speech_event( + event_type: Literal["input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped"], item_id: str +) -> OpenAIRealtimeInputAudioBufferSpeechEvent: + event: Final[OpenAIRealtimeInputAudioBufferSpeechEvent] = { + "type": event_type, + "event_id": _event_id(), + "item_id": item_id, } + return event + + +def _delta_event(item_id: str, delta: str) -> OpenAIRealtimeInputAudioTranscriptionDelta: + event: Final[OpenAIRealtimeInputAudioTranscriptionDelta] = { + "type": "conversation.item.input_audio_transcription.delta", + "event_id": _event_id(), + "item_id": item_id, + "content_index": 0, + "delta": delta, + } + return event + + +def _completed_event( + item_id: str, transcript: str, usage: RealtimeInputAudioTranscriptionUsage | None +) -> OpenAIRealtimeInputAudioTranscriptionCompleted: + event: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": _event_id(), + "item_id": item_id, + "content_index": 0, + "transcript": transcript, + } + if usage is None: + return event + billed: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = {**event, "usage": usage} + return billed + + +def _required_turn_id(message: Mapping[str, JsonValue], event: str) -> str: + value: Final = message.get("turnId") + if isinstance(value, bool) or not isinstance(value, (str, int)): + raise MuseProtocolError(f"{event} event has invalid turnId") + turn_id: Final = str(value).strip() + if not turn_id: + raise MuseProtocolError(f"{event} event has invalid turnId") + return turn_id + + +def _new_suffix(previous: str, current: str) -> str: + return current[len(previous) :] if current.startswith(previous) else "" + + +@dataclass(slots=True) +class _TurnState: + item_id: str + started: bool = False + start_emitted: bool = False + latest_partial: str | None = None + emitted_partial: str = "" + final_text: str | None = None + completed_signal: bool = False + completed_emitted: bool = False + stopped: bool = False + stopped_emitted: bool = False + + @property + def settled(self) -> bool: + return self.completed_emitted and (self.stopped or self.completed_signal) + + def drain( + self, take_usage: Callable[[], RealtimeInputAudioTranscriptionUsage | None] + ) -> Iterator[OpenAIRealtimeEvents]: + has_content: Final = self.latest_partial is not None or self.final_text is not None + if (self.started or has_content) and not self.start_emitted: + self.start_emitted = True + yield _speech_event("input_audio_buffer.speech_started", self.item_id) + if self.latest_partial is not None and self.final_text is None: + delta: Final = _new_suffix(self.emitted_partial, self.latest_partial) + if delta: + self.emitted_partial = self.latest_partial + yield _delta_event(self.item_id, delta) + if self.stopped and not self.stopped_emitted: + self.stopped_emitted = True + yield _speech_event("input_audio_buffer.speech_stopped", self.item_id) + if self.final_text is not None and self.stopped_emitted and not self.completed_emitted: + self.completed_emitted = True + yield _completed_event(self.item_id, self.final_text, take_usage()) class MuseEventTransformer: def __init__(self, *, completed_turn_limit: int = 128) -> None: - self._turns: OrderedDict[str, _TurnState] = OrderedDict() # mutable-ok: ordered active-turn state + self._turns: dict[str, _TurnState] = {} # mutable-ok: insertion-ordered live turn state machine + self._completed_turns: deque[str] = deque(maxlen=completed_turn_limit) # mutable-ok: bounded tombstones self._active_turn_id: str | None = None - self._mode: Literal["PUSH_TO_TALK", "ENDPOINTING", "DIARIZATION"] = "ENDPOINTING" - self._completed_turn_ids: set[str] = set() # mutable-ok: bounded completed-turn membership - self._completed_turn_order: deque[str] = deque( # mutable-ok: bounded completion eviction order - maxlen=completed_turn_limit - ) - self._completed_turn_limit: Final = completed_turn_limit - self._pending_item_ids: deque[str] = deque() # mutable-ok: FIFO commit correlation state - self._last_committed_item_id: str | None = None + self._mode: MuseMode = "ENDPOINTING" self._last_audio_processed_ms: float = 0.0 - self._unassigned_usage_seconds: float = 0.0 + self._unbilled_seconds: float = 0.0 def configure(self, config: MuseSessionConfig) -> None: self._mode = config.mode - def transform(self, payload: str) -> tuple[OpenAIEvent, ...]: - message: Final = _json_object(payload) + def transform(self, message: Mapping[str, JsonValue]) -> tuple[OpenAIRealtimeEvents, ...]: event_type: Final = message.get("type") if event_type == "error": - return (error_event("server_error", "provider_error", "Meta Muse realtime transcription failed"),) + return (error_event(_PROVIDER_ERROR_MESSAGE),) if event_type == "audioProgress": self._update_audio_progress(message) return () @@ -388,57 +467,38 @@ class MuseEventTransformer: self._speech_start(message) elif event_type == "transcript": self._transcript(message) - elif event_type == "speaker": - self._speaker(message) elif event_type == "speechEnd": self._speech_end(message) elif event_type == "speechComplete": self._speech_complete(message) else: return () - return self._drain() - - def commit_item(self) -> tuple[str | None, str]: - previous_item_id: Final = self._last_committed_item_id - provider_turn_id: Final = self._active_turn_id - active_turn: Final = self._turns.get(provider_turn_id) if provider_turn_id is not None else None - item_id: Final = ( - active_turn.item_id or provider_turn_id - if active_turn is not None and provider_turn_id is not None - else f"item_{uuid.uuid4().hex}" - ) - if active_turn is not None: - active_turn.item_id = item_id - else: - self._pending_item_ids.append(item_id) - self._last_committed_item_id = item_id - return previous_item_id, item_id + return tuple(self._drained_events()) def take_unbilled_usage(self) -> RealtimeInputAudioTranscriptionUsage | None: - seconds: Final = self._unassigned_usage_seconds + seconds: Final = self._unbilled_seconds if seconds <= 0: return None - self._unassigned_usage_seconds = 0.0 - return {"type": "duration", "seconds": seconds} # mutable-ok: typed usage wire payload + self._unbilled_seconds = 0.0 + usage: Final[RealtimeInputAudioTranscriptionDurationUsage] = {"type": "duration", "seconds": seconds} + return usage - def _turn(self, turn_id: str) -> _TurnState: - if turn_id in self._completed_turn_ids: - raise _CompletedTurn - turn: Final = self._turns.get(turn_id) - if turn is not None: - return turn - created: Final = _TurnState(item_id=self._pending_item_ids.popleft() if self._pending_item_ids else turn_id) + def _turn(self, turn_id: str) -> _TurnState | None: + if turn_id in self._completed_turns: + return None + existing: Final = self._turns.get(turn_id) + if existing is not None: + return existing + created: Final = _TurnState(item_id=turn_id) self._turns[turn_id] = created return created def _speech_start(self, message: Mapping[str, JsonValue]) -> None: - turn_id: Final = self._required_turn_id(message, "speechStart") - try: - turn: Final = self._turn(turn_id) - except _CompletedTurn: + turn: Final = self._turn(_required_turn_id(message, "speechStart")) + if turn is None: return turn.started = True - self._active_turn_id = turn_id + self._active_turn_id = turn.item_id def _transcript(self, message: Mapping[str, JsonValue]) -> None: transcript: Final = message.get("transcript") @@ -446,56 +506,34 @@ class MuseEventTransformer: raise MuseProtocolError("transcript event has invalid transcript") if not transcript and message.get("turnId") is None and self._active_turn_id is None: return - turn_id: Final = self._transcript_turn_id(message) - try: - turn: Final = self._turn(turn_id) - except _CompletedTurn: + turn: Final = self._turn(self._transcript_turn_id(message)) + if turn is None: return - final: Final = message.get("final") is True - if final: - turn.final_text = transcript - turn.completed_signal = True - if self._mode == "PUSH_TO_TALK": - turn.stopped = True - if self._active_turn_id == turn_id: - self._active_turn_id = None + if message.get("final") is not True: + if turn.final_text is None: + turn.latest_partial = transcript return - if turn.final_text is None: - turn.latest_partial = transcript - - def _speaker(self, message: Mapping[str, JsonValue]) -> None: - turn_id: Final = ( - self._required_turn_id(message, "speaker") if message.get("turnId") is not None else self._active_turn_id - ) - if turn_id is None: - raise MuseProtocolError("speaker event arrived outside an active turn") - label: Final = message.get("label") - if not isinstance(label, str) or not label.strip(): - raise MuseProtocolError("speaker event has invalid label") - try: - turn: Final = self._turn(turn_id) - except _CompletedTurn: - return - turn.speaker = label.strip() + turn.final_text = transcript + turn.completed_signal = True + if self._mode == "PUSH_TO_TALK": + turn.stopped = True + if self._active_turn_id == turn.item_id: + self._active_turn_id = None def _speech_end(self, message: Mapping[str, JsonValue]) -> None: - turn_id: Final = self._required_turn_id(message, "speechEnd") - try: - turn: Final = self._turn(turn_id) - except _CompletedTurn: + turn: Final = self._turn(_required_turn_id(message, "speechEnd")) + if turn is None: return turn.stopped = True - if self._active_turn_id == turn_id: + if self._active_turn_id == turn.item_id: self._active_turn_id = None def _speech_complete(self, message: Mapping[str, JsonValue]) -> None: - turn_id: Final = self._required_turn_id(message, "speechComplete") transcript: Final = message.get("transcript") if not isinstance(transcript, str): raise MuseProtocolError("speechComplete event has invalid transcript") - try: - turn: Final = self._turn(turn_id) - except _CompletedTurn: + turn: Final = self._turn(_required_turn_id(message, "speechComplete")) + if turn is None: return turn.final_text = transcript turn.completed_signal = True @@ -511,73 +549,21 @@ class MuseEventTransformer: raise MuseProtocolError("audioProgress event has invalid audioProcessedMs") if processed_ms <= self._last_audio_processed_ms: return - self._unassigned_usage_seconds += (float(processed_ms) - self._last_audio_processed_ms) / 1000 + self._unbilled_seconds += (float(processed_ms) - self._last_audio_processed_ms) / 1000 self._last_audio_processed_ms = float(processed_ms) - def _drain(self) -> tuple[OpenAIEvent, ...]: - events: list[OpenAIEvent] = [] # mutable-ok: ordered events are frozen to a tuple before return + def _drained_events(self) -> Iterator[OpenAIRealtimeEvents]: while self._turns: - turn_id: str = next(iter(self._turns)) # rebind-ok: selects the next ordered turn - turn: _TurnState = self._turns[turn_id] # rebind-ok: state for the selected turn - has_content: bool = ( # rebind-ok: evaluated for the selected turn - turn.latest_partial is not None or turn.final_text is not None - ) - item_id: str = turn.item_id or turn_id # rebind-ok: selected for each ordered turn - if (turn.started or has_content) and not turn.start_emitted: - turn.start_emitted = True - events.append(self._speech_event("input_audio_buffer.speech_started", item_id)) - if turn.latest_partial is not None and turn.final_text is None: - delta: str = self._new_suffix( # rebind-ok: computed for the selected turn - turn.emitted_partial, turn.latest_partial - ) - if delta: - turn.emitted_partial = turn.latest_partial - events.append( - { # mutable-ok: OpenAI-compatible JSON event - "type": "conversation.item.input_audio_transcription.delta", - "event_id": f"event_{uuid.uuid4().hex}", - "item_id": item_id, - "content_index": 0, - "delta": delta, - } - ) - if turn.stopped and not turn.stopped_emitted: - turn.stopped_emitted = True - events.append(self._speech_event("input_audio_buffer.speech_stopped", item_id)) - if turn.final_text is not None and turn.stopped_emitted and not turn.completed_emitted: - turn.completed_emitted = True - usage: RealtimeInputAudioTranscriptionUsage | None = ( # rebind-ok: usage assigned per turn - self.take_unbilled_usage() - ) - completed_event: dict[str, object] = { # mutable-ok: incrementally builds OpenAI JSON event - "type": "conversation.item.input_audio_transcription.completed", - "event_id": f"event_{uuid.uuid4().hex}", - "item_id": item_id, - "content_index": 0, - "transcript": turn.final_text, - } - if turn.speaker is not None: - completed_event["speaker"] = turn.speaker - if usage is not None: - completed_event["usage"] = usage - events.append(completed_event) - if not (turn.completed_emitted and (turn.stopped or turn.completed_signal)): - break + turn_id, turn = next(iter(self._turns.items())) + yield from turn.drain(self.take_unbilled_usage) + if not turn.settled: + return del self._turns[turn_id] - self._remember_completed(turn_id) - return tuple(events) - - def _remember_completed(self, turn_id: str) -> None: - if turn_id in self._completed_turn_ids: - return - if len(self._completed_turn_order) >= self._completed_turn_limit: - self._completed_turn_ids.discard(self._completed_turn_order.popleft()) - self._completed_turn_order.append(turn_id) - self._completed_turn_ids.add(turn_id) + self._completed_turns.append(turn_id) def _transcript_turn_id(self, message: Mapping[str, JsonValue]) -> str: if message.get("turnId") is not None: - return self._required_turn_id(message, "transcript") + return _required_turn_id(message, "transcript") if self._active_turn_id is not None: return self._active_turn_id if self._mode != "PUSH_TO_TALK": @@ -586,34 +572,162 @@ class MuseEventTransformer: self._active_turn_id = turn_id return turn_id - @staticmethod - def _required_turn_id(message: Mapping[str, JsonValue], event: str) -> str: - value: Final = message.get("turnId") - if isinstance(value, bool) or not isinstance(value, (str, int)): - raise MuseProtocolError(f"{event} event has invalid turnId") - turn_id: Final = str(value).strip() - if not turn_id: - raise MuseProtocolError(f"{event} event has invalid turnId") - return turn_id - @staticmethod - def _speech_event(event_type: str, turn_id: str) -> OpenAIEvent: - return { # mutable-ok: OpenAI-compatible JSON event - "type": event_type, - "event_id": f"event_{uuid.uuid4().hex}", - "item_id": turn_id, +class MetaRealtimeConfig(BaseRealtimeConfig): + def __init__( + self, + *, + monotonic: Callable[[], float] = time.monotonic, + sleep: Callable[[float], Awaitable[None]] = asyncio.sleep, + ) -> None: + self._monotonic: Final = monotonic + self._sleep: Final = sleep + self._transformer: Final = MuseEventTransformer() + self._access_token: str | None = None + self._config: MuseSessionConfig | None = None + self._pending_audio: bytes = b"" + self._end_stream_sent: bool = False + self._pacing_origin: float | None = None + self._sent_duration: float = 0.0 + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseRealtimeConfig contract + model: str, + api_key: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseRealtimeConfig contract + token: Final = api_key or get_secret_str("META_API_KEY") + if token is None: + raise ValueError("api_key is required for Meta API calls") + self._access_token = normalize_access_token(token) + return headers + + def get_complete_url(self, api_base: str | None, model: str, api_key: str | None = None) -> str: + if _normalize_model(model) != MUSE_MODEL: + raise ValueError(f"Unsupported Meta realtime model: {model}") + return build_muse_realtime_url(api_base) + + def is_setup_message(self, msg_obj: Mapping[str, object]) -> bool: + return "authorization" in msg_obj + + def transform_session_created_event( + self, + model: str, + logging_session_id: str, + session_configuration_request: str | None = None, + ) -> MuseSessionCreatedEvent: + return session_created_event(_DEFAULT_SESSION_CONFIG, logging_session_id) + + def transform_realtime_request( + self, + message: str, + model: str, + session_configuration_request: str | None = None, + ) -> tuple[str | bytes, ...]: + request: Final = _json_object(message) + event_type: Final = request.get("type") + if event_type in ("session.update", "transcription_session.update"): + return self._configure(message, model) + if event_type == "input_audio_buffer.append": + return self._append_audio(request) + if event_type == "input_audio_buffer.commit": + return self._flush_audio(end_stream=self._require_config().mode == "PUSH_TO_TALK") + if event_type == "input_audio_buffer.end": + return self._flush_audio(end_stream=True) + if event_type == "input_audio_buffer.clear": + self._pending_audio = b"" + return () + verbose_logger.debug("Meta realtime: dropping unsupported client event %s", event_type) + return () + + async def pace_backend_send(self, message: bytes) -> None: + now: Final = self._monotonic() + origin: Final = self._pacing_origin + effective_origin: Final = ( + now - self._sent_duration if origin is None or now > origin + self._sent_duration else origin + ) + delay: Final = effective_origin + self._sent_duration - now + if delay > 0: + await self._sleep(delay) + self._pacing_origin = effective_origin + self._sent_duration += len(message) / self._require_config().bytes_per_second + + def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None: + return self._transformer.take_unbilled_usage() + + def transform_realtime_response( + self, + message: str | bytes, + model: str, + logging_obj: LiteLLMLoggingObj, + realtime_response_transform_input: RealtimeResponseTransformInput, + ) -> RealtimeResponseTypedDict: + payload: Final = message.decode("utf-8") if isinstance(message, bytes) else message + result: Final[RealtimeResponseTypedDict] = { + "response": list(self._backend_events(payload)), # mutable-ok: RealtimeResponseTypedDict.response is a list + "current_output_item_id": realtime_response_transform_input.get("current_output_item_id"), + "current_response_id": realtime_response_transform_input.get("current_response_id"), + "current_delta_chunks": realtime_response_transform_input.get("current_delta_chunks"), + "current_conversation_id": realtime_response_transform_input.get("current_conversation_id"), + "current_item_chunks": realtime_response_transform_input.get("current_item_chunks"), + "current_delta_type": realtime_response_transform_input.get("current_delta_type"), + "session_configuration_request": realtime_response_transform_input.get("session_configuration_request"), } + return result - @staticmethod - def _new_suffix(previous: str, current: str) -> str: - if current.startswith(previous): - return current[len(previous) :] - return "" + def _backend_events(self, payload: str) -> tuple[OpenAIRealtimeEvents, ...]: + frame: Final = _json_object(payload) + session_id: Final = frame.get("sessionId") + if session_id is None: + return self._transformer.transform(frame) + if not isinstance(session_id, str) or not session_id.strip(): + raise MuseProtocolError("provider returned an invalid handshake response") + created: Final = session_created_event(self._require_config(), session_id.strip()) + event: Final = cast(OpenAIRealtimeEvents, created) # cast-ok: ReadOnly Muse session vs writable OpenAI fields + return (event,) + def _configure(self, message: str, model: str) -> tuple[str, ...]: + if self._config is not None: + verbose_logger.debug("Meta realtime: ignoring session.update after the Muse handshake was sent") + return () + access_token: Final = self._access_token + if access_token is None: + raise MuseProtocolError("Meta API key was not validated before the session was configured") + config: Final = parse_session_update(message, model) + self._config = config + self._transformer.configure(config) + return (json.dumps(config.handshake(access_token), separators=(",", ":")),) -class _CompletedTurn(Exception): - pass + def _append_audio(self, request: Mapping[str, JsonValue]) -> tuple[bytes, ...]: + config: Final = self._require_config() + encoded: Final = request.get("audio") + if not isinstance(encoded, str): + raise MuseProtocolError("Audio must be a base64 string") + if len(encoded) > config.max_encoded_append_bytes: + raise MuseProtocolError("Audio append exceeds the four-second backlog limit") + try: + audio: Final = base64.b64decode(encoded, validate=True) + except (binascii.Error, ValueError): + raise MuseProtocolError("Audio must be valid base64") from None + if len(audio) % 2: + raise MuseProtocolError("PCM16 audio must contain complete samples") + buffered: Final = self._pending_audio + audio + packet_end: Final = len(buffered) - len(buffered) % config.packet_bytes + self._pending_audio = buffered[packet_end:] + return tuple( + buffered[start : start + config.packet_bytes] for start in range(0, packet_end, config.packet_bytes) + ) + def _flush_audio(self, *, end_stream: bool) -> tuple[str | bytes, ...]: + remainder: Final = self._pending_audio + self._pending_audio = b"" + frames: Final[tuple[bytes, ...]] = (remainder,) if remainder else () + if not end_stream or self._end_stream_sent: + return frames + self._end_stream_sent = True + return (*frames, _END_STREAM) -def encode_event(event: Mapping[str, object]) -> str: - return json.dumps(event, separators=(",", ":")) + def _require_config(self) -> MuseSessionConfig: + if self._config is None: + raise MuseProtocolError("session.update must configure the Muse session before audio is sent") + return self._config diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4dc6768e12a..2ae66254b68 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -34718,6 +34718,7 @@ "supports_xhigh_reasoning_effort": true }, "meta/muse-voice-transcribe-1.0": { + "input_cost_per_second": 0.00005, "litellm_provider": "meta", "mode": "audio_transcription", "source": "https://dev.meta.ai/docs/speech-to-text", diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index b0803e44f6b..44c47af57f4 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -333,7 +333,7 @@ async def _resolve_vertex_access_token_bounded( @wrapper_client -async def _arealtime( # noqa: C901 # central dispatcher branches once per supported realtime provider +async def _arealtime( model: str, websocket: "WebSocket", # fastapi websocket api_base: str | None = None, @@ -391,37 +391,7 @@ async def _arealtime( # noqa: C901 # central dispatcher branches once per supp model=model, provider=LlmProviders(_custom_llm_provider), ) - if _custom_llm_provider == LlmProviders.META.value: - if model != "muse-voice-transcribe-1.0": - raise ValueError(f"Unsupported Meta realtime model: {model}") - if query_params is None or query_params.get("intent") != "transcription": - raise ValueError("Meta Muse Voice realtime requires intent=transcription") - - from litellm.llms.meta.realtime.handler import MetaRealtime - - meta_api_key: Final = get_secret_str("META_API_KEY") - dynamic_key_override: Final = dynamic_api_key if dynamic_api_key != meta_api_key else None - resolved_meta_api_key: Final = ( - api_key - or litellm_params.api_key - or dynamic_key_override - or get_secret_str("MODEL_API_KEY") - or dynamic_api_key - or meta_api_key - ) - await MetaRealtime().async_realtime( - model=model, - websocket=websocket, - logging_obj=litellm_logging_obj, - api_base=dynamic_api_base or litellm_params.api_base or api_base, - api_key=resolved_meta_api_key, - client=client, - timeout=timeout, - query_params=query_params, - user_api_key_dict=kwargs.get("user_api_key_dict"), - litellm_metadata=_build_litellm_metadata(kwargs), - ) - elif provider_config is not None: + if provider_config is not None: await base_llm_http_handler.async_realtime( model=model, websocket=websocket, diff --git a/litellm/types/llms/meta.py b/litellm/types/llms/meta.py new file mode 100644 index 00000000000..d7487bb09c8 --- /dev/null +++ b/litellm/types/llms/meta.py @@ -0,0 +1,58 @@ +from typing import Literal, TypeAlias + +from typing_extensions import NotRequired, ReadOnly, TypedDict + +MuseMode: TypeAlias = Literal["PUSH_TO_TALK", "ENDPOINTING"] +MuseAudioEncoding: TypeAlias = Literal["PCM_16KHZ", "PCM_24KHZ"] +MuseSampleRate: TypeAlias = Literal[16000, 24000] + + +class MuseAuthorization(TypedDict): + accessToken: ReadOnly[str] + + +class MuseHandshake(TypedDict): + authorization: ReadOnly[MuseAuthorization] + audioEncoding: ReadOnly[MuseAudioEncoding] + model: ReadOnly[str] + mode: ReadOnly[MuseMode] + partialMode: ReadOnly[Literal["CUMULATIVE"]] + emitAudioProgress: ReadOnly[bool] + languageBias: NotRequired[ReadOnly[tuple[str, ...]]] + + +class MuseTranscriptionAudioFormat(TypedDict): + type: ReadOnly[Literal["audio/pcm"]] + rate: ReadOnly[MuseSampleRate] + + +class MuseTranscriptionSettings(TypedDict): + model: ReadOnly[str] + language: NotRequired[ReadOnly[str]] + + +class MuseTurnDetection(TypedDict): + type: ReadOnly[Literal["server_vad"]] + + +class MuseTranscriptionAudioInput(TypedDict): + format: ReadOnly[MuseTranscriptionAudioFormat] + transcription: ReadOnly[MuseTranscriptionSettings] + turn_detection: ReadOnly[MuseTurnDetection | None] + + +class MuseTranscriptionAudio(TypedDict): + input: ReadOnly[MuseTranscriptionAudioInput] + + +class MuseTranscriptionSession(TypedDict): + id: ReadOnly[str] + object: ReadOnly[Literal["realtime.transcription_session"]] + type: ReadOnly[Literal["transcription"]] + audio: ReadOnly[MuseTranscriptionAudio] + + +class MuseSessionCreatedEvent(TypedDict): + type: ReadOnly[Literal["session.created"]] + event_id: ReadOnly[str] + session: ReadOnly[MuseTranscriptionSession] diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 02a102c9579..765ecaffdae 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2203,7 +2203,6 @@ class OpenAIRealtimeInputAudioTranscriptionCompleted(TypedDict): content_index: ReadOnly[int] transcript: ReadOnly[str] usage: NotRequired[ReadOnly[Mapping[str, object]]] - speaker: NotRequired[ReadOnly[str]] class OpenAIRealtimeUsageTokenDetails(TypedDict): diff --git a/litellm/utils.py b/litellm/utils.py index 1a77655a5a4..394ab4b4094 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9285,6 +9285,10 @@ class ProviderConfigManager: from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig return GeminiRealtimeConfig() + if LlmProviders.META == provider: + from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig + + return MetaRealtimeConfig() return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4dc6768e12a..2ae66254b68 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -34718,6 +34718,7 @@ "supports_xhigh_reasoning_effort": true }, "meta/muse-voice-transcribe-1.0": { + "input_cost_per_second": 0.00005, "litellm_provider": "meta", "mode": "audio_transcription", "source": "https://dev.meta.ai/docs/speech-to-text", diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index e330cb103b1..295110c6bce 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3124,7 +3124,6 @@ async def test_session_close_flush_noop_without_unbilled_usage(): ) - _UPSTREAM_REFUSAL: Final = "Publisher model `publishers/google/models/gemini-live-2.5-flash` was not found" @@ -3192,9 +3191,7 @@ def _backend_ws_closing_with(*frames: bytes | Exception) -> MagicMock: def _relay_session(client_ws: MagicMock, backend_ws: MagicMock) -> _RelaySession: logging: Final = _RecordingLogging() worker: Final = _InlineLoggingWorker() - streaming: Final = RealTimeStreaming( - client_ws, backend_ws, logging, model="gpt-realtime", logging_worker=worker - ) + streaming: Final = RealTimeStreaming(client_ws, backend_ws, logging, model="gpt-realtime", logging_worker=worker) return _RelaySession(streaming=streaming, logging=logging, worker=worker) @@ -3402,44 +3399,6 @@ async def test_refused_session_does_not_stamp_the_reservation_ownership_marker() assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in session.logging.model_call_details -@pytest.mark.asyncio -async def test_separate_usage_provider_flushes_duration_once_without_client_event(): - from typing import Final - - client_ws: Final = MagicMock() - client_ws.send_text = AsyncMock() - backend_ws: Final = MagicMock() - backend_ws.recv = AsyncMock(side_effect=ConnectionClosed(None, None)) - logging_obj: Final = MagicMock() - logging_obj.model_call_details = {} - logging_obj.dispatch_success_handlers = AsyncMock() - usage_provider: Final = MagicMock() - usage_provider.unbilled_usage_on_session_close.return_value = { - "type": "duration", - "seconds": 0.75, - } - - streaming: Final = RealTimeStreaming( - client_ws, - backend_ws, - logging_obj, - model="muse-voice-transcribe-1.0", - usage_provider=usage_provider, - ) - - await streaming.backend_to_client_send_messages() - - usage_provider.unbilled_usage_on_session_close.assert_called_once_with("muse-voice-transcribe-1.0") - duration_events: Final = tuple( - message - for message in streaming.messages - if isinstance(message, dict) and message.get("usage") == {"type": "duration", "seconds": 0.75} - ) - assert len(duration_events) == 1 - assert client_ws.send_text.await_count == 0 - logging_obj.dispatch_success_handlers.assert_called_once_with(streaming.messages, prefer_async_handlers=True) - - @pytest.mark.asyncio async def test_transformed_transcription_completion_never_sends_response_create(): from typing import Final @@ -3487,57 +3446,26 @@ async def test_transformed_transcription_completion_never_sends_response_create( backend_ws.send.assert_not_awaited() -def test_private_logging_excludes_audio_transcript_hints_and_provider_body(monkeypatch: pytest.MonkeyPatch): +@pytest.mark.asyncio +async def test_provider_bytes_are_sent_raw_after_pacing(): from typing import Final - monkeypatch.setattr(litellm, "logged_real_time_event_types", "*") - logging_obj: Final = MagicMock() - logging_obj.model_call_details = {} + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + provider_config: Final = MagicMock() + provider_config.requires_session_configuration.return_value = True + provider_config.transform_realtime_request.return_value = (b"\x00\x01", '{"type":"endStream"}') + provider_config.pace_backend_send = AsyncMock() + provider_config.is_setup_message.return_value = False streaming: Final = RealTimeStreaming( MagicMock(), + backend_ws, MagicMock(), - logging_obj, + provider_config=provider_config, model="muse-voice-transcribe-1.0", - exclude_private_content_from_logs=True, - ) - audio: Final = "cHJpdmF0ZS1hdWRpbw==" - transcript: Final = "private transcript" - keyword: Final = "private keyword" - provider_body: Final = "private provider body" - - streaming.store_input( - json.dumps( - { - "type": "session.update", - "session": { - "type": "transcription", - "model": "muse-voice-transcribe-1.0", - "mode": "ENDPOINTING", - "audio": {"input": {"transcription": {"keywords": [keyword]}}}, - }, - } - ) - ) - streaming.store_input(json.dumps({"type": "input_audio_buffer.append", "audio": audio})) - streaming.store_message( - { - "type": "conversation.item.input_audio_transcription.completed", - "event_id": "event_1", - "item_id": "turn_1", - "transcript": transcript, - "provider_body": provider_body, - "usage": {"type": "duration", "seconds": 1.0}, - } ) - logged_inputs: Final = tuple(call.kwargs["input"] for call in logging_obj.pre_call.call_args_list) - serialized: Final = json.dumps({"inputs": logged_inputs, "messages": streaming.messages}) - assert audio not in serialized - assert transcript not in serialized - assert keyword not in serialized - assert provider_body not in serialized - assert "muse-voice-transcribe-1.0" in serialized - assert "ENDPOINTING" in serialized - assert "turn_1" in serialized - assert '"seconds": 1.0' in serialized - assert streaming.input_messages == [] + assert await streaming._send_to_backend(json.dumps({"type": "input_audio_buffer.commit"})) is True + + assert [call.args[0] for call in backend_ws.send.await_args_list] == [b"\x00\x01", '{"type":"endStream"}'] + provider_config.pace_backend_send.assert_awaited_once_with(b"\x00\x01") diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py deleted file mode 100644 index 14ba972fe14..00000000000 --- a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_handler.py +++ /dev/null @@ -1,452 +0,0 @@ -import asyncio -import base64 -import json -from collections.abc import Callable -from typing import Final -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from litellm.llms.meta.realtime.handler import ( - DEFAULT_MUSE_REALTIME_URL, - MetaRealtime, - MuseAdapterError, - MuseRealtimeAdapter, - build_muse_realtime_url, - normalize_access_token, - safe_close_reason, - sanitize_close_code, -) -from litellm.llms.meta.realtime.transformation import MUSE_MODEL - - -class FakeProviderWebSocket: - def __init__(self, session_id: str = "provider-session") -> None: - self.sent: list[str | bytes] = [] - self.close_calls: list[tuple[int, str]] = [] - self._session_id: Final = session_id - self._recv_count = 0 - self._closed = asyncio.Event() - - async def send(self, message: str | bytes) -> None: - self.sent.append(message) - - async def recv(self, decode: bool | None = None) -> str | bytes: - self._recv_count += 1 - if self._recv_count == 1: - return json.dumps({"sessionId": self._session_id}) - await self._closed.wait() - raise MuseAdapterError("closed", close_code=1000) - - async def close(self, code: int = 1000, reason: str = "") -> None: - self.close_calls.append((code, reason)) - self._closed.set() - - -class DelayedAckWebSocket(FakeProviderWebSocket): - def __init__(self) -> None: - super().__init__() - self.ack_release = asyncio.Event() - - async def recv(self, decode: bool | None = None) -> str | bytes: - self._recv_count += 1 - if self._recv_count == 1: - await self.ack_release.wait() - return json.dumps({"sessionId": self._session_id}) - await self._closed.wait() - raise MuseAdapterError("closed", close_code=1000) - - -async def _wait_until(predicate: Callable[[], bool]) -> None: - for _ in range(100): - if predicate(): - return - await asyncio.sleep(0) - raise AssertionError("condition did not become true") - - -def _session_update(*, rate: int = 24_000, mode: str = "ENDPOINTING") -> str: - return json.dumps( - { - "type": "session.update", - "session": { - "type": "transcription", - "mode": mode, - "audio": { - "input": { - "format": {"type": "audio/pcm", "rate": rate, "channels": 1}, - "transcription": {"model": MUSE_MODEL}, - } - }, - }, - } - ) - - -async def _configured_adapter( - *, - rate: int = 24_000, - mode: str = "ENDPOINTING", - provider_ws: FakeProviderWebSocket | None = None, - monotonic: Callable[[], float] = lambda: 10.0, - sleep: Callable[[float], object] | None = None, -) -> tuple[MuseRealtimeAdapter, FakeProviderWebSocket, dict[str, object]]: - ws: Final = provider_ws or FakeProviderWebSocket() - connect_call: Final[dict[str, object]] = {} - - async def connect(url: str, **kwargs: object) -> FakeProviderWebSocket: - connect_call.update({"url": url, **kwargs}) - return ws - - async def no_sleep(_: float) -> None: - return None - - adapter: Final = MuseRealtimeAdapter( - model=f"meta/{MUSE_MODEL}", - api_key=" raw-token ", - websocket_connect=connect, - monotonic=monotonic, - sleep=sleep or no_sleep, - ) - created: Final = json.loads(await adapter.recv()) - assert created["type"] == "session.created" - await adapter.send(_session_update(rate=rate, mode=mode)) - updated: Final = json.loads(await adapter.recv()) - assert updated["type"] == "session.updated" - return adapter, ws, connect_call - - -@pytest.mark.parametrize( - ("api_key", "expected"), - [ - ("token", "Bearer token"), - (" Bearer token ", "Bearer token"), - ("bearer token", "Bearer token"), - ], -) -def test_normalize_access_token_emits_exactly_one_bearer_prefix(api_key: str, expected: str): - assert normalize_access_token(api_key) == expected - - -@pytest.mark.parametrize("api_key", ["", " ", "Bearer", " bearer "]) -def test_normalize_access_token_rejects_missing_token(api_key: str): - with pytest.raises(ValueError, match=r"token|key is required"): - normalize_access_token(api_key) - - -def test_build_muse_realtime_url_uses_fixed_secure_path(): - assert build_muse_realtime_url(None) == DEFAULT_MUSE_REALTIME_URL - assert build_muse_realtime_url("https://example.test/custom/path?ignored=yes") == ( - "wss://example.test/v1/asr/realtime" - ) - assert build_muse_realtime_url("wss://example.test:8443/other") == ("wss://example.test:8443/v1/asr/realtime") - - -@pytest.mark.parametrize( - "api_base", - [ - "http://example.test", - "ws://example.test", - "wss://user:pass@example.test", - "wss://example.test/path#fragment", - "not-a-url", - ], -) -def test_build_muse_realtime_url_rejects_insecure_or_ambiguous_overrides(api_base: str): - with pytest.raises(ValueError, match="absolute wss:// or https://"): - build_muse_realtime_url(api_base) - - -@pytest.mark.asyncio -async def test_handshake_contains_bearer_only_in_json_body_and_waits_for_ack(): - provider_ws: Final = DelayedAckWebSocket() - connect_call: Final[dict[str, object]] = {} - - async def connect(url: str, **kwargs: object) -> DelayedAckWebSocket: - connect_call.update({"url": url, **kwargs}) - return provider_ws - - adapter: Final = MuseRealtimeAdapter( - model=MUSE_MODEL, - api_key="Bearer private-token", - websocket_connect=connect, - ) - await adapter.recv() - update_task: Final = asyncio.create_task(adapter.send(_session_update())) - await _wait_until(lambda: len(provider_ws.sent) == 1) - - assert connect_call["url"] == DEFAULT_MUSE_REALTIME_URL - assert "additional_headers" not in connect_call - handshake: Final = json.loads(provider_ws.sent[0]) - assert handshake["authorization"] == {"accessToken": "Bearer private-token"} - assert handshake["audioEncoding"] == "PCM_24KHZ" - assert not update_task.done() - assert not any(isinstance(frame, bytes) for frame in provider_ws.sent) - - provider_ws.ack_release.set() - await update_task - updated: Final = json.loads(await adapter.recv()) - assert updated["type"] == "session.updated" - assert updated["session"]["id"] == "provider-session" - await adapter.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize(("rate", "packet_bytes"), [(16_000, 2_560), (24_000, 3_840)]) -async def test_audio_is_strictly_decoded_and_packetized_as_raw_pcm(rate: int, packet_bytes: int): - adapter, provider_ws, _ = await _configured_adapter(rate=rate) - pcm: Final = (b"\xff\xfe\x00\x80" * (packet_bytes // 2))[: packet_bytes * 2] - - await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(pcm).decode()})) - await _wait_until(lambda: sum(isinstance(frame, bytes) for frame in provider_ws.sent) == 2) - - binary_frames: Final = tuple(frame for frame in provider_ws.sent if isinstance(frame, bytes)) - assert binary_frames == (pcm[:packet_bytes], pcm[packet_bytes:]) - assert b"\xff\xfe" in pcm - await adapter.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("audio", "expected_message"), - [ - ("not base64!", "valid base64"), - (base64.b64encode(b"\x00").decode(), "complete samples"), - ], -) -async def test_invalid_base64_or_odd_pcm_is_rejected(audio: str, expected_message: str): - adapter, provider_ws, _ = await _configured_adapter() - - await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": audio})) - error: Final = json.loads(await adapter.recv()) - - assert error["type"] == "error" - assert error["error"]["code"] == "invalid_audio" - assert expected_message in error["error"]["message"] - assert adapter.close_code == 1008 - assert not any(isinstance(frame, bytes) for frame in provider_ws.sent) - await adapter.close() - - -@pytest.mark.asyncio -async def test_absolute_pacing_delays_only_audio_ahead_of_wall_time(): - sleeps: Final[list[float]] = [] - - async def record_sleep(delay: float) -> None: - sleeps.append(delay) - - adapter, provider_ws, _ = await _configured_adapter(monotonic=lambda: 10.0, sleep=record_sleep) - pcm: Final = b"\x01\x02" * 3_840 - - await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(pcm).decode()})) - await _wait_until(lambda: sum(isinstance(frame, bytes) for frame in provider_ws.sent) == 2) - - assert sleeps == pytest.approx([0.08]) - await adapter.close() - - -@pytest.mark.asyncio -async def test_append_larger_than_four_seconds_is_rejected_without_decoding(): - adapter, provider_ws, _ = await _configured_adapter(rate=16_000) - max_pcm_bytes: Final = 16_000 * 2 * 4 - oversized_audio: Final = "A" * (4 * ((max_pcm_bytes + 2) // 3) + 1) - - with patch( # test-quality-ok: proves rejection happens before an attacker-controlled allocation - "litellm.llms.meta.realtime.handler.base64.b64decode" - ) as decode: - await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": oversized_audio})) - - error: Final = json.loads(await adapter.recv()) - assert error["error"]["code"] == "audio_backlog_exceeded" - assert adapter.close_code == 1008 - assert not any(isinstance(frame, bytes) for frame in provider_ws.sent) - decode.assert_not_called() - await adapter.close() - - -@pytest.mark.asyncio -async def test_clear_discards_only_unsent_audio(): - adapter, provider_ws, _ = await _configured_adapter() - old_pcm: Final = b"\x01\x02" * 100 - new_pcm: Final = b"\x03\x04" * 100 - - await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(old_pcm).decode()})) - await adapter.send(_event("input_audio_buffer.clear")) - cleared: Final = json.loads(await adapter.recv()) - await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(new_pcm).decode()})) - await adapter.send(_event("input_audio_buffer.commit")) - await _wait_until(lambda: any(isinstance(frame, bytes) for frame in provider_ws.sent)) - - assert cleared["type"] == "input_audio_buffer.cleared" - assert tuple(frame for frame in provider_ws.sent if isinstance(frame, bytes)) == (new_pcm,) - assert '{"type":"endStream"}' not in provider_ws.sent - await adapter.close() - - -@pytest.mark.asyncio -async def test_endpointing_commit_flushes_partial_packet_without_ending_stream(): - adapter, provider_ws, _ = await _configured_adapter(mode="ENDPOINTING") - pcm: Final = b"\x01\x02" * 100 - - await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(pcm).decode()})) - await adapter.send(_event("input_audio_buffer.commit")) - committed: Final = json.loads(await adapter.recv()) - await _wait_until(lambda: any(isinstance(frame, bytes) for frame in provider_ws.sent)) - - assert committed["type"] == "input_audio_buffer.committed" - assert committed["item_id"].startswith("item_") - assert committed["previous_item_id"] is None - assert tuple(frame for frame in provider_ws.sent if isinstance(frame, bytes)) == (pcm,) - assert '{"type":"endStream"}' not in provider_ws.sent - await adapter.close() - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("mode", "terminal_event"), - [("PUSH_TO_TALK", "input_audio_buffer.commit"), ("ENDPOINTING", "input_audio_buffer.end")], -) -async def test_commit_or_end_sends_end_stream_exactly_once(mode: str, terminal_event: str): - adapter, provider_ws, _ = await _configured_adapter(mode=mode) - pcm: Final = b"\x01\x02" * 100 - - await adapter.send(json.dumps({"type": "input_audio_buffer.append", "audio": base64.b64encode(pcm).decode()})) - await adapter.send(_event(terminal_event)) - await adapter.send(_event("input_audio_buffer.end")) - await _wait_until(lambda: '{"type":"endStream"}' in provider_ws.sent) - - assert tuple(frame for frame in provider_ws.sent if isinstance(frame, bytes)) == (pcm,) - assert provider_ws.sent.count('{"type":"endStream"}') == 1 - await adapter.close() - - -@pytest.mark.asyncio -async def test_response_create_is_returned_as_error_and_never_sent_upstream(): - adapter, provider_ws, _ = await _configured_adapter() - - await adapter.send(_event("response.create")) - error: Final = json.loads(await adapter.recv()) - - assert error["type"] == "error" - assert error["error"]["code"] == "unsupported_event" - assert not any(isinstance(frame, str) and "response.create" in frame for frame in provider_ws.sent) - await adapter.close() - - -@pytest.mark.asyncio -async def test_close_codes_and_reasons_are_sanitized_without_secret_leakage(): - adapter, provider_ws, _ = await _configured_adapter() - secret: Final = "Bearer private-token" - - await adapter.close(code=4001, reason=f"provider rejected {secret}") - - assert adapter.close_code == 1011 - assert adapter.close_reason == "Realtime transcription service error" - assert provider_ws.close_calls == [(1011, "Realtime transcription service error")] - assert secret not in json.dumps(provider_ws.close_calls) - assert sanitize_close_code(1013) == 1013 - assert safe_close_reason(1008) == "Invalid realtime transcription request" - - -@pytest.mark.asyncio -async def test_handshake_failure_reports_only_exception_type(): - secret: Final = "private-token" - - async def failing_connect(url: str, **kwargs: object) -> FakeProviderWebSocket: - raise RuntimeError(f"failed with {secret}") - - adapter: Final = MuseRealtimeAdapter( - model=MUSE_MODEL, - api_key=secret, - websocket_connect=failing_connect, - ) - await adapter.recv() - - await adapter.send(_session_update()) - error = json.loads(await adapter.recv()) - - assert error["type"] == "error" - assert error["error"]["message"] == "Meta Muse realtime handshake failed" - assert secret not in json.dumps(error) - with pytest.raises(MuseAdapterError) as exc_info: - await adapter.recv() - assert exc_info.value.close_code == 1011 - assert secret not in str(exc_info.value) - - -@pytest.mark.asyncio -async def test_meta_realtime_missing_credentials_closes_client_with_policy_code(): - client_ws: Final = MagicMock() - client_ws.send_text = AsyncMock() - client_ws.close = AsyncMock() - - await MetaRealtime().async_realtime( - model=MUSE_MODEL, - websocket=client_ws, - logging_obj=MagicMock(), - api_key=None, - ) - - client_ws.close.assert_awaited_once_with( - code=1008, - reason="Invalid realtime transcription request", - ) - sent_error: Final = json.loads(client_ws.send_text.await_args.args[0]) - assert sent_error["type"] == "error" - assert sent_error["error"]["code"] == "invalid_configuration" - - -@pytest.mark.asyncio -async def test_meta_realtime_invalid_constructor_input_sends_error_before_close(): - client_ws: Final = MagicMock() - client_ws.send_text = AsyncMock() - client_ws.close = AsyncMock() - - await MetaRealtime().async_realtime( - model=MUSE_MODEL, - websocket=client_ws, - logging_obj=MagicMock(), - api_key="Bearer", - ) - - sent_error: Final = json.loads(client_ws.send_text.await_args.args[0]) - assert sent_error["error"]["message"] == "Invalid Meta Muse realtime configuration" - client_ws.close.assert_awaited_once_with( - code=1008, - reason="Invalid realtime transcription request", - ) - - -@pytest.mark.asyncio -async def test_meta_realtime_enables_private_logging_usage_and_model_enforcement(monkeypatch: pytest.MonkeyPatch): - captured: Final[dict[str, object]] = {} - - class CapturingStreaming: - def __init__(self, websocket, backend_ws, logging_obj, **kwargs): - captured.update({"websocket": websocket, "backend_ws": backend_ws, "logging_obj": logging_obj, **kwargs}) - - async def bidirectional_forward(self) -> None: - return None - - client_ws: Final = MagicMock() - client_ws.send_text = AsyncMock() - client_ws.close = AsyncMock() - monkeypatch.setattr("litellm.llms.meta.realtime.handler.RealTimeStreaming", CapturingStreaming) - - await MetaRealtime().async_realtime( - model=MUSE_MODEL, - websocket=client_ws, - logging_obj=MagicMock(), - api_key="private-token", - ) - - adapter: Final = captured["backend_ws"] - assert isinstance(adapter, MuseRealtimeAdapter) - assert captured["force_transcription_model"] == MUSE_MODEL - assert captured["usage_provider"] is adapter - assert captured["exclude_private_content_from_logs"] is True - client_ws.close.assert_awaited_once_with(code=1000, reason="Session closed") - - -def _event(event_type: str) -> str: - return json.dumps({"type": event_type}) diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py index 8cc4af836a5..058e075860b 100644 --- a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py +++ b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py @@ -1,24 +1,70 @@ +import base64 import json +from typing import Final +from unittest.mock import MagicMock import pytest from litellm.llms.meta.realtime.transformation import ( + DEFAULT_MUSE_REALTIME_URL, MUSE_MODEL, + MetaRealtimeConfig, MuseEventTransformer, MuseProtocolError, - encode_event, + MuseSessionConfig, + build_muse_realtime_url, + normalize_access_token, normalize_language, parse_session_update, session_created_event, - session_updated_event, ) +from litellm.types.realtime import RealtimeResponseTransformInput + +EMPTY_TRANSFORM_INPUT: Final[RealtimeResponseTransformInput] = { + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_item_chunks": None, + "current_conversation_id": None, + "current_delta_type": None, +} def _event(event_type: str, **fields: object) -> str: return json.dumps({"type": event_type, **fields}) -def test_beta_session_builds_authenticated_24khz_handshake_with_hints(): +def _ga_session_update(rate: int = 24_000, turn_detection: object = "server_vad") -> str: + return _event( + "session.update", + session={ + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": rate}, + "turn_detection": None if turn_detection is None else {"type": turn_detection}, + "transcription": {"model": f"meta/{MUSE_MODEL}"}, + } + }, + }, + ) + + +def _configured(rate: int = 24_000, turn_detection: object = "server_vad", **kwargs: object) -> MetaRealtimeConfig: + config = MetaRealtimeConfig(**kwargs) + config.validate_environment({}, MUSE_MODEL, api_key="secret-token") + config.transform_realtime_request(_ga_session_update(rate, turn_detection), MUSE_MODEL) + return config + + +def _backend_events(config: MetaRealtimeConfig, payload: str) -> list[dict[str, object]]: + response = config.transform_realtime_response(payload, MUSE_MODEL, MagicMock(), EMPTY_TRANSFORM_INPUT)["response"] + assert isinstance(response, list) + return response + + +def test_beta_session_translates_language_and_drops_non_openai_hints(): config = parse_session_update( _event( "session.update", @@ -29,8 +75,6 @@ def test_beta_session_builds_authenticated_24khz_handshake_with_hints(): "input_audio_transcription": { "model": "meta/muse-voice-transcribe-1.0", "language": "en-US", - "language_bias": ["Spanish", "english", "French"], - "keywords": [" Muse ", "LiteLLM", "Muse"], "prompt": "must not become a keyword", }, }, @@ -41,8 +85,7 @@ def test_beta_session_builds_authenticated_24khz_handshake_with_hints(): assert config.sample_rate == 24_000 assert config.packet_bytes == 3_840 assert config.mode == "ENDPOINTING" - assert config.language_bias == ("English", "Spanish", "French") - assert config.keywords == ("Muse", "LiteLLM") + assert config.language_bias == ("English",) assert config.handshake("Bearer token") == { "mode": "ENDPOINTING", "authorization": {"accessToken": "Bearer token"}, @@ -50,8 +93,7 @@ def test_beta_session_builds_authenticated_24khz_handshake_with_hints(): "model": MUSE_MODEL, "partialMode": "CUMULATIVE", "emitAudioProgress": True, - "keywords": ["Muse", "LiteLLM"], - "languageBias": ["English", "Spanish", "French"], + "languageBias": ("English",), } assert "must not become a keyword" not in json.dumps(config.handshake("Bearer token")) @@ -79,6 +121,7 @@ def test_ga_session_accepts_16khz_mono_push_to_talk(): assert config.mode == "PUSH_TO_TALK" assert config.language_bias == ("Mandarin Chinese",) assert config.handshake("Bearer token")["audioEncoding"] == "PCM_16KHZ" + assert "languageBias" not in MuseSessionConfig(MUSE_MODEL, "ENDPOINTING", 24_000, ()).handshake("Bearer token") @pytest.mark.parametrize( @@ -109,8 +152,9 @@ def test_language_normalization_uses_official_muse_names(source: str, expected: "either beta or GA layout", ), ({"input_audio_transcription": {"model": "other-model"}}, "cannot be changed"), - ({"input_audio_transcription": {"keywords": ["valid", ""]}}, "non-empty strings"), ({"input_audio_transcription": {"language": "xx"}}, "unsupported Muse Voice language"), + ({"turn_detection": {"type": "semantic_vad"}}, "server_vad turn detection or null"), + ({"type": "realtime", "audio": {"input": {"turn_detection": {"type": "semantic_vad"}}}}, "server_vad"), ], ) def test_session_rejects_unsupported_audio_model_and_hints(session: dict[str, object], message: str): @@ -118,16 +162,15 @@ def test_session_rejects_unsupported_audio_model_and_hints(session: dict[str, ob parse_session_update(_event("session.update", session={"type": "transcription", **session}), MUSE_MODEL) -def test_session_events_expose_openai_transcription_shapes(): +def test_session_created_event_exposes_openai_transcription_shape(): config = parse_session_update( _event( "session.update", session={ - "mode": "DIARIZATION", "audio": { "input": { "format": {"type": "audio/pcm", "rate": 24000}, - "transcription": {"model": MUSE_MODEL, "language": "ja", "keywords": ["Meta"]}, + "transcription": {"model": MUSE_MODEL, "language": "ja"}, } }, }, @@ -135,31 +178,25 @@ def test_session_events_expose_openai_transcription_shapes(): MUSE_MODEL, ) - created = session_created_event(MUSE_MODEL, "session-before-handshake") - updated = session_updated_event(config, "provider-session") + created = session_created_event(config, "provider-session") assert created["type"] == "session.created" + assert created["session"]["id"] == "provider-session" assert created["session"]["type"] == "transcription" - assert updated["type"] == "session.updated" - assert updated["session"]["id"] == "provider-session" - assert updated["session"]["audio"]["input"]["transcription"] == { - "model": MUSE_MODEL, - "language": "Japanese", - "keywords": ["Meta"], - "language_bias": ["Japanese"], - } + assert created["session"]["audio"]["input"]["turn_detection"] == {"type": "server_vad"} + assert created["session"]["audio"]["input"]["transcription"] == {"model": MUSE_MODEL, "language": "Japanese"} def test_turnless_empty_silence_transcript_is_ignored(): transformer = MuseEventTransformer() - assert transformer.transform(_event("transcript", transcript="", final=True)) == () + assert transformer.transform(json.loads(_event("transcript", transcript="", final=True))) == () def test_transcript_without_speech_start_synthesizes_start_before_delta(): transformer = MuseEventTransformer() - events = transformer.transform(_event("transcript", turnId="turn-1", transcript="hello", final=False)) + events = transformer.transform(json.loads(_event("transcript", turnId="turn-1", transcript="hello", final=False))) assert [event["type"] for event in events] == [ "input_audio_buffer.speech_started", @@ -170,12 +207,15 @@ def test_transcript_without_speech_start_synthesizes_start_before_delta(): def test_cumulative_partials_emit_only_extensions_and_final_is_authoritative(): transformer = MuseEventTransformer() - started = transformer.transform(_event("speechStart", turnId="turn-1")) - first = transformer.transform(_event("transcript", turnId="turn-1", transcript="hello", final=False)) - extension = transformer.transform(_event("transcript", turnId="turn-1", transcript="hello world", final=False)) - rewrite = transformer.transform(_event("transcript", turnId="turn-1", transcript="hullo world", final=False)) - assert transformer.transform(_event("speechComplete", turnId="turn-1", transcript="hullo world")) == () - completed = transformer.transform(_event("speechEnd", turnId="turn-1")) + def send(payload: str) -> tuple[dict[str, object], ...]: + return transformer.transform(json.loads(payload)) + + started = send(_event("speechStart", turnId="turn-1")) + first = send(_event("transcript", turnId="turn-1", transcript="hello", final=False)) + extension = send(_event("transcript", turnId="turn-1", transcript="hello world", final=False)) + rewrite = send(_event("transcript", turnId="turn-1", transcript="hullo world", final=False)) + assert send(_event("speechComplete", turnId="turn-1", transcript="hullo world")) == () + completed = send(_event("speechEnd", turnId="turn-1")) assert [event["type"] for event in started] == ["input_audio_buffer.speech_started"] assert first[0]["delta"] == "hello" @@ -190,10 +230,10 @@ def test_cumulative_partials_emit_only_extensions_and_final_is_authoritative(): def test_completed_transcript_waits_for_speech_stopped(): transformer = MuseEventTransformer() - transformer.transform(_event("speechStart", turnId="turn-1")) - assert transformer.transform(_event("speechComplete", turnId="turn-1", transcript="done")) == () + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + assert transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="done"))) == () - released = transformer.transform(_event("speechEnd", turnId="turn-1")) + released = transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) assert [event["type"] for event in released] == [ "input_audio_buffer.speech_stopped", "conversation.item.input_audio_transcription.completed", @@ -203,11 +243,14 @@ def test_completed_transcript_waits_for_speech_stopped(): def test_overlapping_turns_are_emitted_in_provider_turn_order(): transformer = MuseEventTransformer() - transformer.transform(_event("speechStart", turnId="turn-a")) - transformer.transform(_event("speechStart", turnId="turn-b")) - assert transformer.transform(_event("transcript", turnId="turn-b", transcript="second", final=False)) == () - assert transformer.transform(_event("speechComplete", turnId="turn-a", transcript="first")) == () - released = transformer.transform(_event("speechEnd", turnId="turn-a")) + def send(payload: str) -> tuple[dict[str, object], ...]: + return transformer.transform(json.loads(payload)) + + send(_event("speechStart", turnId="turn-a")) + send(_event("speechStart", turnId="turn-b")) + assert send(_event("transcript", turnId="turn-b", transcript="second", final=False)) == () + assert send(_event("speechComplete", turnId="turn-a", transcript="first")) == () + released = send(_event("speechEnd", turnId="turn-a")) assert [(event["type"], event["item_id"]) for event in released] == [ ("input_audio_buffer.speech_stopped", "turn-a"), @@ -215,51 +258,42 @@ def test_overlapping_turns_are_emitted_in_provider_turn_order(): ("input_audio_buffer.speech_started", "turn-b"), ("conversation.item.input_audio_transcription.delta", "turn-b"), ] - assert transformer.transform(_event("speechComplete", turnId="turn-b", transcript="second final")) == () - final_b = transformer.transform(_event("speechEnd", turnId="turn-b")) + assert send(_event("speechComplete", turnId="turn-b", transcript="second final")) == () + final_b = send(_event("speechEnd", turnId="turn-b")) assert final_b[0]["type"] == "input_audio_buffer.speech_stopped" assert final_b[1]["item_id"] == "turn-b" assert final_b[1]["transcript"] == "second final" -def test_committed_item_id_is_used_for_next_provider_turn(): +def test_push_to_talk_final_transcript_completes_without_speech_end(): + transformer = MuseEventTransformer() + transformer.configure(MuseSessionConfig(MUSE_MODEL, "PUSH_TO_TALK", 24_000, ())) + + events = transformer.transform(json.loads(_event("transcript", transcript="hello there", final=True))) + + assert [event["type"] for event in events] == [ + "input_audio_buffer.speech_started", + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + assert events[2]["transcript"] == "hello there" + assert str(events[0]["item_id"]).startswith("item_") + + +def test_positive_audio_progress_deltas_attach_to_next_completion_and_speaker_is_ignored(): transformer = MuseEventTransformer() - previous_item_id, item_id = transformer.commit_item() - started = transformer.transform(_event("speechStart", turnId="provider-turn")) - transformer.transform(_event("speechComplete", turnId="provider-turn", transcript="hello")) - completed = transformer.transform(_event("speechEnd", turnId="provider-turn")) + def send(payload: str) -> tuple[dict[str, object], ...]: + return transformer.transform(json.loads(payload)) - assert previous_item_id is None - assert started[0]["item_id"] == item_id - assert completed[-1]["item_id"] == item_id + send(_event("audioProgress", audioProcessedMs=1000)) + send(_event("audioProgress", audioProcessedMs=750)) + send(_event("audioProgress", audioProcessedMs=1600)) + assert send(_event("speaker", turnId=42, label=" Speaker 2 ")) == () + send(_event("speechComplete", turnId=42, transcript="hello")) + completed = send(_event("speechEnd", turnId=42)) - -def test_commit_after_speech_start_reuses_active_item_id(): - transformer = MuseEventTransformer() - - started = transformer.transform(_event("speechStart", turnId="provider-turn")) - previous_item_id, item_id = transformer.commit_item() - transformer.transform(_event("speechComplete", turnId="provider-turn", transcript="hello")) - completed = transformer.transform(_event("speechEnd", turnId="provider-turn")) - - assert previous_item_id is None - assert item_id == "provider-turn" - assert started[0]["item_id"] == item_id - assert completed[-1]["item_id"] == item_id - - -def test_speaker_and_positive_audio_progress_deltas_attach_to_next_completion(): - transformer = MuseEventTransformer() - - transformer.transform(_event("audioProgress", audioProcessedMs=1000)) - transformer.transform(_event("audioProgress", audioProcessedMs=750)) - transformer.transform(_event("audioProgress", audioProcessedMs=1600)) - transformer.transform(_event("speaker", turnId=42, label=" Speaker 2 ")) - transformer.transform(_event("speechComplete", turnId=42, transcript="hello")) - completed = transformer.transform(_event("speechEnd", turnId=42)) - - assert completed[-1]["speaker"] == "Speaker 2" + assert "speaker" not in completed[-1] assert completed[-1]["usage"] == {"type": "duration", "seconds": 1.6} assert transformer.take_unbilled_usage() is None @@ -267,7 +301,7 @@ def test_speaker_and_positive_audio_progress_deltas_attach_to_next_completion(): def test_trailing_audio_progress_is_returned_once(): transformer = MuseEventTransformer() - transformer.transform(_event("audioProgress", audioProcessedMs=250)) + transformer.transform(json.loads(_event("audioProgress", audioProcessedMs=250))) assert transformer.take_unbilled_usage() == {"type": "duration", "seconds": 0.25} assert transformer.take_unbilled_usage() is None @@ -276,24 +310,244 @@ def test_trailing_audio_progress_is_returned_once(): def test_completed_turn_tombstone_suppresses_late_duplicates(): transformer = MuseEventTransformer() - transformer.transform(_event("speechComplete", turnId="turn-1", transcript="done")) + transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="done"))) + released = transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) - assert transformer.transform(_event("speechComplete", turnId="turn-1", transcript="duplicate")) == () - assert transformer.transform(_event("speaker", turnId="turn-1", label="late")) == () + assert [event["type"] for event in released] == [ + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + assert transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="duplicate"))) == () + assert transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) == () + assert ( + transformer.transform(json.loads(_event("transcript", turnId="turn-1", transcript="late", final=False))) == () + ) def test_provider_error_is_sanitized_and_encodable(): token = "private-token" provider_body = f"authorization failed for Bearer {token}" transformed = MuseEventTransformer().transform( - _event("error", code="AUTH", message=provider_body, request={"accessToken": token}) + json.loads(_event("error", code="AUTH", message=provider_body, request={"accessToken": token})) ) - encoded = encode_event(transformed[0]) + encoded = json.dumps(transformed[0]) assert json.loads(encoded)["error"] == { "type": "server_error", - "code": "provider_error", "message": "Meta Muse realtime transcription failed", } assert token not in encoded assert provider_body not in encoded + + +@pytest.mark.parametrize( + ("raw", "expected"), + [("token", "Bearer token"), (" Bearer token ", "Bearer token"), ("bearer token", "Bearer token")], +) +def test_access_token_normalization_adds_single_bearer_prefix(raw: str, expected: str): + assert normalize_access_token(raw) == expected + + +@pytest.mark.parametrize("raw", ["", " ", "Bearer", " bearer "]) +def test_access_token_normalization_rejects_empty_tokens(raw: str): + with pytest.raises(ValueError, match=r"token|key is required"): + normalize_access_token(raw) + + +@pytest.mark.parametrize( + ("api_base", "expected"), + [ + (None, DEFAULT_MUSE_REALTIME_URL), + ("https://example.test/custom/path?ignored=yes", "wss://example.test/v1/asr/realtime"), + ("wss://example.test:8443/other", "wss://example.test:8443/v1/asr/realtime"), + ], +) +def test_realtime_url_pins_muse_path(api_base: str | None, expected: str): + assert build_muse_realtime_url(api_base) == expected + assert MetaRealtimeConfig().get_complete_url(api_base, f"meta/{MUSE_MODEL}") == expected + + +@pytest.mark.parametrize( + "api_base", + [ + "http://example.test", + "ws://example.test", + "wss://user:pass@example.test", + "wss://example.test/path#fragment", + "not-a-url", + ], +) +def test_realtime_url_rejects_insecure_or_ambiguous_bases(api_base: str): + with pytest.raises(ValueError, match="absolute wss:// or https://"): + build_muse_realtime_url(api_base) + + +def test_unsupported_model_is_rejected_before_connecting(): + with pytest.raises(ValueError, match="Unsupported Meta realtime model: meta/other-model"): + MetaRealtimeConfig().get_complete_url(None, "meta/other-model") + + +def test_missing_api_key_is_rejected(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("META_API_KEY", raising=False) + + with pytest.raises(ValueError, match="api_key is required for Meta API calls"): + MetaRealtimeConfig().validate_environment({}, MUSE_MODEL) + + +def test_bearer_token_travels_only_in_the_json_handshake(): + config = MetaRealtimeConfig() + headers = {"x-existing": "kept"} + + assert config.validate_environment(headers, MUSE_MODEL, api_key="secret-token") == {"x-existing": "kept"} + (handshake,) = config.transform_realtime_request(_ga_session_update(), MUSE_MODEL) + + assert isinstance(handshake, str) + assert json.loads(handshake)["authorization"] == {"accessToken": "Bearer secret-token"} + assert config.is_setup_message(json.loads(handshake)) is True + assert config.is_setup_message({"type": "input_audio_buffer.append"}) is False + assert config.transform_realtime_request(_ga_session_update(), MUSE_MODEL) == () + + +def test_synthetic_session_created_uses_default_transcription_shape(): + created = MetaRealtimeConfig().transform_session_created_event(f"meta/{MUSE_MODEL}", "trace-1") + + assert created["type"] == "session.created" + assert created["session"]["id"] == "trace-1" + assert created["session"]["audio"]["input"]["format"] == {"type": "audio/pcm", "rate": 24000} + assert created["session"]["audio"]["input"]["transcription"] == {"model": MUSE_MODEL} + + +def test_audio_before_session_update_is_rejected(): + config = MetaRealtimeConfig() + config.validate_environment({}, MUSE_MODEL, api_key="secret-token") + + with pytest.raises(MuseProtocolError, match=r"session\.update must configure"): + config.transform_realtime_request(_event("input_audio_buffer.append", audio="AAAA"), MUSE_MODEL) + + +@pytest.mark.parametrize(("rate", "packet_bytes"), [(16_000, 2_560), (24_000, 3_840)]) +def test_pcm_is_packetized_into_raw_binary_frames(rate: int, packet_bytes: int): + config = _configured(rate=rate) + pcm = b"\xff\xfe\x00\x80" * (packet_bytes // 2) + b"\x01\x02\x03\x04" + + frames = config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(pcm).decode()), MUSE_MODEL + ) + remainder = config.transform_realtime_request(_event("input_audio_buffer.commit"), MUSE_MODEL) + + assert frames == (pcm[:packet_bytes], pcm[packet_bytes : packet_bytes * 2]) + assert remainder == (pcm[packet_bytes * 2 :],) + + +@pytest.mark.parametrize( + ("audio", "message"), + [ + ("not base64!", "valid base64"), + (base64.b64encode(b"\x00").decode(), "complete samples"), + (12, "base64 string"), + ("A" * (4 * ((24_000 * 2 * 4 + 2) // 3) + 4), "four-second backlog"), + ], +) +def test_invalid_audio_appends_are_rejected(audio: object, message: str): + config = _configured() + + with pytest.raises(MuseProtocolError, match=message): + config.transform_realtime_request(_event("input_audio_buffer.append", audio=audio), MUSE_MODEL) + + +@pytest.mark.asyncio +async def test_backend_sends_are_paced_to_real_time(): + sleeps: list[float] = [] + + async def record_sleep(delay: float) -> None: + sleeps.append(delay) + + config = _configured(monotonic=lambda: 10.0, sleep=record_sleep) + packet = b"\x01\x02" * 1_920 + + await config.pace_backend_send(packet) + await config.pace_backend_send(packet) + await config.pace_backend_send(packet) + + assert sleeps == pytest.approx([0.08, 0.16]) + + +def test_endpointing_commit_flushes_without_end_stream_but_end_sends_it_once(): + config = _configured(turn_detection="server_vad") + + assert config.transform_realtime_request(_event("input_audio_buffer.commit"), MUSE_MODEL) == () + assert config.transform_realtime_request(_event("input_audio_buffer.end"), MUSE_MODEL) == ('{"type":"endStream"}',) + assert config.transform_realtime_request(_event("input_audio_buffer.end"), MUSE_MODEL) == () + + +def test_push_to_talk_commit_ends_the_stream_once(): + config = _configured(turn_detection=None) + config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(b"\x01\x02").decode()), MUSE_MODEL + ) + + assert config.transform_realtime_request(_event("input_audio_buffer.commit"), MUSE_MODEL) == ( + b"\x01\x02", + '{"type":"endStream"}', + ) + assert config.transform_realtime_request(_event("input_audio_buffer.end"), MUSE_MODEL) == () + + +def test_clear_drops_buffered_remainder_and_unknown_events_are_ignored(): + config = _configured() + config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(b"\x01\x02").decode()), MUSE_MODEL + ) + + assert config.transform_realtime_request(_event("input_audio_buffer.clear"), MUSE_MODEL) == () + assert config.transform_realtime_request(_event("response.create"), MUSE_MODEL) == () + assert config.transform_realtime_request(_event("input_audio_buffer.commit"), MUSE_MODEL) == () + + +def test_provider_ack_becomes_session_created_with_provider_id(): + config = _configured(rate=16_000, turn_detection=None) + + (created,) = _backend_events(config, json.dumps({"sessionId": " provider-session "})) + + assert created["type"] == "session.created" + assert created["session"]["id"] == "provider-session" + assert created["session"]["audio"]["input"]["format"]["rate"] == 16000 + assert created["session"]["audio"]["input"]["turn_detection"] is None + + +def test_provider_turn_events_and_close_usage_flow_through_config(): + config = _configured() + + assert _backend_events(config, json.dumps({"type": "audioProgress", "audioProcessedMs": 1349})) == [] + assert _backend_events(config, _event("speechStart", turnId="t1"))[0]["type"] == "input_audio_buffer.speech_started" + assert _backend_events(config, _event("speechComplete", turnId="t1", transcript="what is the weather")) == [] + completed = _backend_events(config, _event("speechEnd", turnId="t1")) + + assert [event["type"] for event in completed] == [ + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + assert completed[1]["usage"] == {"type": "duration", "seconds": 1.349} + assert config.unbilled_usage_on_session_close(MUSE_MODEL) is None + + assert _backend_events(config, json.dumps({"type": "audioProgress", "audioProcessedMs": 2349})) == [] + assert config.unbilled_usage_on_session_close(MUSE_MODEL) == {"type": "duration", "seconds": 1.0} + + +def test_provider_error_frame_becomes_openai_error_without_leaking_token(): + config = _configured() + + (error,) = _backend_events(config, _event("error", message="bad token secret-token")) + + assert error == { + "type": "error", + "error": {"type": "server_error", "message": "Meta Muse realtime transcription failed"}, + } + assert "secret-token" not in json.dumps(error) + + +def test_invalid_provider_ack_is_rejected(): + config = _configured() + + with pytest.raises(MuseProtocolError, match="invalid handshake response"): + _backend_events(config, json.dumps({"sessionId": ""})) diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index aa3f7d45d84..d3d41c5b54b 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -152,81 +152,29 @@ async def test_vertex_credential_resolution_bounds_a_thread_offloaded_refresh(): @pytest.mark.asyncio -async def test_meta_realtime_rejects_missing_transcription_intent(monkeypatch: pytest.MonkeyPatch): - def mock_get_llm_provider(model, api_base, api_key): - return model.removeprefix("meta/"), "meta", api_key, api_base +async def test_meta_realtime_dispatches_to_base_handler_with_meta_config(monkeypatch: pytest.MonkeyPatch): + from litellm.llms.meta.realtime.transformation import MetaRealtimeConfig - monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) - - with pytest.raises(ValueError, match="requires intent=transcription"): - await realtime_main._arealtime.__wrapped__( - model="meta/muse-voice-transcribe-1.0", - websocket=MagicMock(), - litellm_logging_obj=FakeLogging(), - query_params={"model": "meta/muse-voice-transcribe-1.0"}, - ) - - -@pytest.mark.asyncio -async def test_meta_realtime_rejects_unsupported_model_before_connecting(monkeypatch: pytest.MonkeyPatch): - def mock_get_llm_provider(model, api_base, api_key): - return model.removeprefix("meta/"), "meta", api_key, api_base - - monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) - - with pytest.raises(ValueError, match="Unsupported Meta realtime model: other-model"): - await realtime_main._arealtime.__wrapped__( - model="meta/other-model", - websocket=MagicMock(), - litellm_logging_obj=FakeLogging(), - query_params={"model": "meta/other-model", "intent": "transcription"}, - ) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("explicit_key", "model_key", "meta_key", "expected"), - [ - ("explicit", "model-env", "meta-env", "explicit"), - (None, "model-env", "meta-env", "model-env"), - (None, None, "meta-env", "meta-env"), - ], -) -async def test_meta_realtime_credential_precedence_is_forwarded_to_handler( - monkeypatch: pytest.MonkeyPatch, - explicit_key: str | None, - model_key: str | None, - meta_key: str | None, - expected: str, -): captured: dict[str, object] = {} def mock_get_llm_provider(model, api_base, api_key): - return model.removeprefix("meta/"), "meta", meta_key, api_base + return model.removeprefix("meta/"), "meta", None, api_base - def mock_get_secret_str(name: str): - return {"MODEL_API_KEY": model_key, "META_API_KEY": meta_key}.get(name) - - async def mock_async_realtime(self, **kwargs): + async def mock_async_realtime(**kwargs): captured.update(kwargs) monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) - monkeypatch.setattr(realtime_main, "get_secret_str", mock_get_secret_str) - monkeypatch.setattr( - "litellm.llms.meta.realtime.handler.MetaRealtime.async_realtime", - mock_async_realtime, - ) + monkeypatch.setattr(realtime_main.base_llm_http_handler, "async_realtime", mock_async_realtime) await realtime_main._arealtime.__wrapped__( model="meta/muse-voice-transcribe-1.0", websocket=MagicMock(), litellm_logging_obj=FakeLogging(), - api_key=explicit_key, query_params={"model": "meta/muse-voice-transcribe-1.0", "intent": "transcription"}, ) + assert isinstance(captured["provider_config"], MetaRealtimeConfig) assert captured["model"] == "muse-voice-transcribe-1.0" - assert captured["api_key"] == expected assert captured["query_params"] == {"model": "muse-voice-transcribe-1.0", "intent": "transcription"} From 6c974df9a08be66b892551e18d519ca17a7063d5 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 12 Sep 2026 13:24:00 +0000 Subject: [PATCH 38/55] fix(registry): sync Azure o-series and Together deprecation dates, gpt-5.4-mini/nano context window Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 41 ++++++++++++------- model_prices_and_context_window.json | 41 ++++++++++++------- 2 files changed, 52 insertions(+), 30 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 7fa09951eae..eab0aef124c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -8730,7 +8730,7 @@ "supports_function_calling": true }, "azure/o1": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 7.5e-06, "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", @@ -8748,7 +8748,7 @@ }, "azure/o1-2024-12-17": { "cache_read_input_token_cost": 7.5e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8825,7 +8825,7 @@ "supports_vision": false }, "azure/o3": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8855,7 +8855,7 @@ "supports_vision": true }, "azure/o3-2025-04-16": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8886,7 +8886,7 @@ }, "azure/o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, - "deprecation_date": "2026-12-26", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8923,7 +8923,7 @@ "supports_web_search": true }, "azure/o3-mini": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8940,7 +8940,7 @@ }, "azure/o3-mini-2025-01-31": { "cache_read_input_token_cost": 5.5e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8954,7 +8954,7 @@ "supports_vision": false }, "azure/o3-pro": { - "deprecation_date": "2026-12-17", + "deprecation_date": "2026-11-19", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -8985,7 +8985,7 @@ "supports_vision": true }, "azure/o3-pro-2025-06-10": { - "deprecation_date": "2026-12-17", + "deprecation_date": "2026-11-19", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -9016,7 +9016,7 @@ "supports_vision": true }, "azure/o4-mini": { - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -9047,7 +9047,7 @@ }, "azure/o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -31650,7 +31650,7 @@ "input_cost_per_token_batches": 3.75e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31702,7 +31702,7 @@ "input_cost_per_token_batches": 3.75e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31752,7 +31752,7 @@ "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31801,7 +31801,7 @@ "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -43488,6 +43488,7 @@ "supports_tool_choice": true }, "together_ai/openai/gpt-oss-20b": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 5e-08, "litellm_provider": "together_ai", "max_input_tokens": 131072, @@ -43780,6 +43781,7 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/google/gemma-4-31B-it": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -43794,6 +43796,7 @@ "supports_vision": true }, "together_ai/intfloat/multilingual-e5-large-instruct": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 2e-08, "litellm_provider": "together_ai", "max_input_tokens": 514, @@ -43906,6 +43909,7 @@ "supports_tool_choice": true }, "together_ai/thinkingmachines/Inkling-Small": { + "deprecation_date": "2026-09-14", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", @@ -60832,6 +60836,7 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/moonshotai/Kimi-K2.6": { + "deprecation_date": "2026-08-19", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 4.5e-06, "cache_read_input_token_cost": 2e-07, @@ -60858,6 +60863,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/zai-org/GLM-5": { + "deprecation_date": "2026-06-22", "input_cost_per_token": 1e-06, "output_cost_per_token": 3.2e-06, "litellm_provider": "together_ai", @@ -60866,6 +60872,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/zai-org/GLM-5.1": { + "deprecation_date": "2026-07-10", "input_cost_per_token": 1.4e-06, "output_cost_per_token": 4.4e-06, "cache_read_input_token_cost": 2.6e-07, @@ -60883,6 +60890,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-Coder-Next-FP8": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 5e-07, "output_cost_per_token": 1.2e-06, "litellm_provider": "together_ai", @@ -60891,6 +60899,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-VL-32B-Instruct": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 5e-07, "output_cost_per_token": 1.5e-06, "litellm_provider": "together_ai", @@ -60899,6 +60908,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-VL-8B-Instruct": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 1.8e-07, "output_cost_per_token": 6.8e-07, "litellm_provider": "together_ai", @@ -60931,6 +60941,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/QwQ-32B": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 1.2e-06, "litellm_provider": "together_ai", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 7fa09951eae..eab0aef124c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -8730,7 +8730,7 @@ "supports_function_calling": true }, "azure/o1": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 7.5e-06, "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", @@ -8748,7 +8748,7 @@ }, "azure/o1-2024-12-17": { "cache_read_input_token_cost": 7.5e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.5e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8825,7 +8825,7 @@ "supports_vision": false }, "azure/o3": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8855,7 +8855,7 @@ "supports_vision": true }, "azure/o3-2025-04-16": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 2e-06, "litellm_provider": "azure", @@ -8886,7 +8886,7 @@ }, "azure/o3-deep-research": { "cache_read_input_token_cost": 2.5e-06, - "deprecation_date": "2026-12-26", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8923,7 +8923,7 @@ "supports_web_search": true }, "azure/o3-mini": { - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -8940,7 +8940,7 @@ }, "azure/o3-mini-2025-01-31": { "cache_read_input_token_cost": 5.5e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -8954,7 +8954,7 @@ "supports_vision": false }, "azure/o3-pro": { - "deprecation_date": "2026-12-17", + "deprecation_date": "2026-11-19", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -8985,7 +8985,7 @@ "supports_vision": true }, "azure/o3-pro-2025-06-10": { - "deprecation_date": "2026-12-17", + "deprecation_date": "2026-11-19", "input_cost_per_token": 2e-05, "input_cost_per_token_batches": 1e-05, "litellm_provider": "azure", @@ -9016,7 +9016,7 @@ "supports_vision": true }, "azure/o4-mini": { - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 2.75e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", @@ -9047,7 +9047,7 @@ }, "azure/o4-mini-2025-04-16": { "cache_read_input_token_cost": 2.75e-07, - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.1e-06, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -31650,7 +31650,7 @@ "input_cost_per_token_batches": 3.75e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31702,7 +31702,7 @@ "input_cost_per_token_batches": 3.75e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31752,7 +31752,7 @@ "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31801,7 +31801,7 @@ "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 272000, + "max_input_tokens": 400000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -43488,6 +43488,7 @@ "supports_tool_choice": true }, "together_ai/openai/gpt-oss-20b": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 5e-08, "litellm_provider": "together_ai", "max_input_tokens": 131072, @@ -43780,6 +43781,7 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/google/gemma-4-31B-it": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 3.9e-07, "litellm_provider": "together_ai", "max_input_tokens": 262144, @@ -43794,6 +43796,7 @@ "supports_vision": true }, "together_ai/intfloat/multilingual-e5-large-instruct": { + "deprecation_date": "2026-09-14", "input_cost_per_token": 2e-08, "litellm_provider": "together_ai", "max_input_tokens": 514, @@ -43906,6 +43909,7 @@ "supports_tool_choice": true }, "together_ai/thinkingmachines/Inkling-Small": { + "deprecation_date": "2026-09-14", "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 5e-07, "litellm_provider": "together_ai", @@ -60832,6 +60836,7 @@ "source": "https://docs.together.ai/docs/serverless-models" }, "together_ai/moonshotai/Kimi-K2.6": { + "deprecation_date": "2026-08-19", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 4.5e-06, "cache_read_input_token_cost": 2e-07, @@ -60858,6 +60863,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/zai-org/GLM-5": { + "deprecation_date": "2026-06-22", "input_cost_per_token": 1e-06, "output_cost_per_token": 3.2e-06, "litellm_provider": "together_ai", @@ -60866,6 +60872,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/zai-org/GLM-5.1": { + "deprecation_date": "2026-07-10", "input_cost_per_token": 1.4e-06, "output_cost_per_token": 4.4e-06, "cache_read_input_token_cost": 2.6e-07, @@ -60883,6 +60890,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-Coder-Next-FP8": { + "deprecation_date": "2026-05-14", "input_cost_per_token": 5e-07, "output_cost_per_token": 1.2e-06, "litellm_provider": "together_ai", @@ -60891,6 +60899,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-VL-32B-Instruct": { + "deprecation_date": "2026-02-25", "input_cost_per_token": 5e-07, "output_cost_per_token": 1.5e-06, "litellm_provider": "together_ai", @@ -60899,6 +60908,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/Qwen3-VL-8B-Instruct": { + "deprecation_date": "2026-04-16", "input_cost_per_token": 1.8e-07, "output_cost_per_token": 6.8e-07, "litellm_provider": "together_ai", @@ -60931,6 +60941,7 @@ "source": "https://api.together.xyz/v1/models" }, "together_ai/Qwen/QwQ-32B": { + "deprecation_date": "2025-11-13", "input_cost_per_token": 1.2e-06, "output_cost_per_token": 1.2e-06, "litellm_provider": "together_ai", From a6418b3ff5bd5dccbfc8485eb4a96fcfc65c677b Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 12 Sep 2026 13:45:15 +0000 Subject: [PATCH 39/55] fix(registry): keep gpt-5.4-mini/nano max_input_tokens at 272000 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 8 ++++---- model_prices_and_context_window.json | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index eab0aef124c..277afd94e1b 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -31650,7 +31650,7 @@ "input_cost_per_token_batches": 3.75e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31702,7 +31702,7 @@ "input_cost_per_token_batches": 3.75e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31752,7 +31752,7 @@ "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31801,7 +31801,7 @@ "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index eab0aef124c..277afd94e1b 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -31650,7 +31650,7 @@ "input_cost_per_token_batches": 3.75e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31702,7 +31702,7 @@ "input_cost_per_token_batches": 3.75e-07, "input_cost_per_token_priority": 1.5e-06, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31752,7 +31752,7 @@ "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", @@ -31801,7 +31801,7 @@ "input_cost_per_token_flex": 1e-07, "input_cost_per_token_batches": 1e-07, "litellm_provider": "openai", - "max_input_tokens": 400000, + "max_input_tokens": 272000, "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", From db8201df8f49a169989530bb68475675f7a79370 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 12 Sep 2026 14:44:52 +0000 Subject: [PATCH 40/55] fix(registry): sync azure/us and azure/eu o-series deprecation dates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 12 ++++++------ model_prices_and_context_window.json | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 277afd94e1b..3606cc0ee4f 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -4497,7 +4497,7 @@ }, "azure/eu/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -4543,7 +4543,7 @@ }, "azure/eu/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -9600,7 +9600,7 @@ }, "azure/us/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -9645,7 +9645,7 @@ "supports_vision": false }, "azure/us/o3-2025-04-16": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 2.2e-06, "litellm_provider": "azure", @@ -9676,7 +9676,7 @@ }, "azure/us/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -9693,7 +9693,7 @@ }, "azure/us/o4-mini-2025-04-16": { "cache_read_input_token_cost": 3.1e-07, - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "litellm_provider": "azure", "max_input_tokens": 200000, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 277afd94e1b..3606cc0ee4f 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -4497,7 +4497,7 @@ }, "azure/eu/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -4543,7 +4543,7 @@ }, "azure/eu/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -9600,7 +9600,7 @@ }, "azure/us/o1-2024-12-17": { "cache_read_input_token_cost": 8.25e-06, - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.65e-05, "litellm_provider": "azure", "max_input_tokens": 200000, @@ -9645,7 +9645,7 @@ "supports_vision": false }, "azure/us/o3-2025-04-16": { - "deprecation_date": "2026-10-21", + "deprecation_date": "2026-11-19", "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 2.2e-06, "litellm_provider": "azure", @@ -9676,7 +9676,7 @@ }, "azure/us/o3-mini-2025-01-31": { "cache_read_input_token_cost": 6.05e-07, - "deprecation_date": "2026-10-01", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "input_cost_per_token_batches": 6.05e-07, "litellm_provider": "azure", @@ -9693,7 +9693,7 @@ }, "azure/us/o4-mini-2025-04-16": { "cache_read_input_token_cost": 3.1e-07, - "deprecation_date": "2026-10-16", + "deprecation_date": "2026-11-19", "input_cost_per_token": 1.21e-06, "litellm_provider": "azure", "max_input_tokens": 200000, From 9aa06cb4a3ad91f84ad29f8947e1b1f0262c4e14 Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 12 Sep 2026 19:20:05 +0000 Subject: [PATCH 41/55] fix(registry): correct computer-use-preview provider/schema flag and OpenRouter deepseek-v3.2 / claude-opus-4.6 metadata Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 21 ++++++++++++------- model_prices_and_context_window.json | 21 ++++++++++++------- 2 files changed, 26 insertions(+), 16 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 3606cc0ee4f..d676a3abe15 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14615,7 +14615,7 @@ }, "computer-use-preview": { "input_cost_per_token": 3e-06, - "litellm_provider": "azure", + "litellm_provider": "openai", "max_input_tokens": 8192, "max_output_tokens": 1024, "max_tokens": 1024, @@ -14635,10 +14635,11 @@ "supports_parallel_function_calling": true, "supports_prompt_caching": false, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "source": "https://platform.openai.com/docs/models/computer-use-preview" }, "dall-e-2": { "deprecation_date": "2026-05-12", @@ -39019,7 +39020,9 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "supports_response_schema": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/anthropic/claude-sonnet-4.5": { "input_cost_per_image": 0.0048, @@ -39167,18 +39170,20 @@ }, "openrouter/deepseek/deepseek-v3.2": { "input_cost_per_token": 2.69e-07, - "input_cost_per_token_cache_hit": 2.8e-08, + "input_cost_per_token_cache_hit": 1.345e-07, "litellm_provider": "openrouter", "max_input_tokens": 163840, - "max_output_tokens": 163840, - "max_tokens": 163840, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 4e-07, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_response_schema": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/deepseek/deepseek-v3.2-exp": { "input_cost_per_token": 2.7e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 3606cc0ee4f..d676a3abe15 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14615,7 +14615,7 @@ }, "computer-use-preview": { "input_cost_per_token": 3e-06, - "litellm_provider": "azure", + "litellm_provider": "openai", "max_input_tokens": 8192, "max_output_tokens": 1024, "max_tokens": 1024, @@ -14635,10 +14635,11 @@ "supports_parallel_function_calling": true, "supports_prompt_caching": false, "supports_reasoning": true, - "supports_response_schema": true, + "supports_response_schema": false, "supports_system_messages": true, "supports_tool_choice": true, - "supports_vision": true + "supports_vision": true, + "source": "https://platform.openai.com/docs/models/computer-use-preview" }, "dall-e-2": { "deprecation_date": "2026-05-12", @@ -39019,7 +39020,9 @@ "supports_max_reasoning_effort": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 4096 + "prompt_cache_min_tokens": 4096, + "supports_response_schema": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/anthropic/claude-sonnet-4.5": { "input_cost_per_image": 0.0048, @@ -39167,18 +39170,20 @@ }, "openrouter/deepseek/deepseek-v3.2": { "input_cost_per_token": 2.69e-07, - "input_cost_per_token_cache_hit": 2.8e-08, + "input_cost_per_token_cache_hit": 1.345e-07, "litellm_provider": "openrouter", "max_input_tokens": 163840, - "max_output_tokens": 163840, - "max_tokens": 163840, + "max_output_tokens": 65536, + "max_tokens": 65536, "mode": "chat", "output_cost_per_token": 4e-07, "supports_assistant_prefill": true, "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, - "supports_tool_choice": true + "supports_tool_choice": true, + "supports_response_schema": true, + "source": "https://openrouter.ai/api/v1/models" }, "openrouter/deepseek/deepseek-v3.2-exp": { "input_cost_per_token": 2.7e-07, From 6b78438c9986286d9ba9e34c6a2e360a9a1f131b Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 12 Sep 2026 12:39:59 -0700 Subject: [PATCH 42/55] fix(realtime): close every Muse turn on its own terminal signal Turns no longer wait behind each other in a FIFO queue, so an empty server_vad turn (speechStart then speechEnd with no transcript) cannot stall every later turn, and a PUSH_TO_TALK speechComplete now closes its turn without waiting for a speechEnd that never arrives. Each turn keeps its own idempotent emit state, so late or duplicate speechEnd, speechComplete and transcript frames are no-ops, and finished turns are remembered in a bounded map instead of a separate tombstone deque. The session.created ack and the sanitized error frame are now typed as members of OpenAIRealtimeEvents, which removes the typing.cast calls that the strict ruff budget flagged. --- litellm/llms/meta/realtime/transformation.py | 157 ++++++++-------- litellm/types/llms/meta.py | 37 ---- litellm/types/llms/openai.py | 49 +++++ .../test_meta_realtime_transformation.py | 169 ++++++++++++++---- 4 files changed, 251 insertions(+), 161 deletions(-) diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py index 096f7c8f1fe..c43e1897fbc 100644 --- a/litellm/llms/meta/realtime/transformation.py +++ b/litellm/llms/meta/realtime/transformation.py @@ -4,11 +4,10 @@ import binascii import json import math import time -from collections import deque from collections.abc import Awaitable, Callable, Iterator, Mapping from dataclasses import dataclass from types import MappingProxyType -from typing import Final, Literal, cast +from typing import Final, Literal from urllib.parse import urlparse, urlunparse from pydantic import JsonValue, TypeAdapter, ValidationError @@ -18,25 +17,19 @@ from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.secret_managers.main import get_secret_str -from litellm.types.llms.meta import ( - MuseAudioEncoding, - MuseHandshake, - MuseMode, - MuseSampleRate, - MuseSessionCreatedEvent, - MuseTranscriptionSession, - MuseTranscriptionSettings, - MuseTurnDetection, -) +from litellm.types.llms.meta import MuseAudioEncoding, MuseHandshake, MuseMode, MuseSampleRate from litellm.types.llms.openai import ( + OpenAIRealtimeErrorEvent, OpenAIRealtimeEvents, OpenAIRealtimeInputAudioBufferSpeechEvent, OpenAIRealtimeInputAudioTranscriptionCompleted, OpenAIRealtimeInputAudioTranscriptionDelta, + OpenAIRealtimeServerVadTurnDetection, + OpenAIRealtimeTranscriptionSession, + OpenAIRealtimeTranscriptionSessionCreated, + OpenAIRealtimeTranscriptionSettings, ) from litellm.types.realtime import ( - RealtimeErrorDetail, - RealtimeErrorEvent, RealtimeInputAudioTranscriptionDurationUsage, RealtimeInputAudioTranscriptionUsage, RealtimeResponseTransformInput, @@ -112,7 +105,7 @@ _END_STREAM: Final = '{"type":"endStream"}' _PROVIDER_ERROR_MESSAGE: Final = "Meta Muse realtime transcription failed" _JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _EMPTY_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) -_SERVER_VAD: Final[MuseTurnDetection] = {"type": "server_vad"} +_SERVER_VAD: Final[OpenAIRealtimeServerVadTurnDetection] = {"type": "server_vad"} class MuseProtocolError(ValueError): @@ -156,8 +149,8 @@ class MuseSessionConfig: biased: Final[MuseHandshake] = {**base, "languageBias": self.language_bias} return biased - def openai_session(self, session_id: str) -> MuseTranscriptionSession: - session: Final[MuseTranscriptionSession] = { + def openai_session(self, session_id: str) -> OpenAIRealtimeTranscriptionSession: + session: Final[OpenAIRealtimeTranscriptionSession] = { "id": session_id, "object": "realtime.transcription_session", "type": "transcription", @@ -171,11 +164,11 @@ class MuseSessionConfig: } return session - def _transcription_settings(self) -> MuseTranscriptionSettings: - base: Final[MuseTranscriptionSettings] = {"model": self.model} + def _transcription_settings(self) -> OpenAIRealtimeTranscriptionSettings: + base: Final[OpenAIRealtimeTranscriptionSettings] = {"model": self.model} if not self.language_bias: return base - localized: Final[MuseTranscriptionSettings] = {**base, "language": self.language_bias[0]} + localized: Final[OpenAIRealtimeTranscriptionSettings] = {**base, "language": self.language_bias[0]} return localized @@ -340,8 +333,8 @@ def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig ) -def session_created_event(config: MuseSessionConfig, session_id: str) -> MuseSessionCreatedEvent: - event: Final[MuseSessionCreatedEvent] = { +def session_created_event(config: MuseSessionConfig, session_id: str) -> OpenAIRealtimeTranscriptionSessionCreated: + event: Final[OpenAIRealtimeTranscriptionSessionCreated] = { "type": "session.created", "event_id": _event_id(), "session": config.openai_session(session_id), @@ -349,10 +342,12 @@ def session_created_event(config: MuseSessionConfig, session_id: str) -> MuseSes return event -def error_event(message: str) -> OpenAIRealtimeEvents: - detail: Final[RealtimeErrorDetail] = {"type": "server_error", "message": message} - event: Final[RealtimeErrorEvent] = {"type": "error", "error": detail} - return cast(OpenAIRealtimeEvents, event) # cast-ok: the union has no error member; the relay only serializes it +def error_event(message: str) -> OpenAIRealtimeErrorEvent: + event: Final[OpenAIRealtimeErrorEvent] = { + "type": "error", + "error": {"type": "server_error", "message": message}, + } + return event def _speech_event( @@ -415,14 +410,13 @@ class _TurnState: latest_partial: str | None = None emitted_partial: str = "" final_text: str | None = None - completed_signal: bool = False completed_emitted: bool = False stopped: bool = False stopped_emitted: bool = False - @property - def settled(self) -> bool: - return self.completed_emitted and (self.stopped or self.completed_signal) + def finish(self, transcript: str) -> None: + self.final_text = transcript + self.stopped = True def drain( self, take_usage: Callable[[], RealtimeInputAudioTranscriptionUsage | None] @@ -445,9 +439,9 @@ class _TurnState: class MuseEventTransformer: - def __init__(self, *, completed_turn_limit: int = 128) -> None: - self._turns: dict[str, _TurnState] = {} # mutable-ok: insertion-ordered live turn state machine - self._completed_turns: deque[str] = deque(maxlen=completed_turn_limit) # mutable-ok: bounded tombstones + def __init__(self, *, turn_limit: int = 128) -> None: + self._turns: dict[str, _TurnState] = {} # mutable-ok: bounded, insertion-ordered per-turn emit state + self._turn_limit: Final = turn_limit self._active_turn_id: str | None = None self._mode: MuseMode = "ENDPOINTING" self._last_audio_processed_ms: float = 0.0 @@ -463,17 +457,10 @@ class MuseEventTransformer: if event_type == "audioProgress": self._update_audio_progress(message) return () - if event_type == "speechStart": - self._speech_start(message) - elif event_type == "transcript": - self._transcript(message) - elif event_type == "speechEnd": - self._speech_end(message) - elif event_type == "speechComplete": - self._speech_complete(message) - else: + turn: Final = self._apply_turn_event(event_type, message) + if turn is None: return () - return tuple(self._drained_events()) + return tuple(turn.drain(self.take_unbilled_usage)) def take_unbilled_usage(self) -> RealtimeInputAudioTranscriptionUsage | None: seconds: Final = self._unbilled_seconds @@ -483,60 +470,69 @@ class MuseEventTransformer: usage: Final[RealtimeInputAudioTranscriptionDurationUsage] = {"type": "duration", "seconds": seconds} return usage - def _turn(self, turn_id: str) -> _TurnState | None: - if turn_id in self._completed_turns: - return None + def _apply_turn_event(self, event_type: JsonValue | None, message: Mapping[str, JsonValue]) -> _TurnState | None: + match event_type: + case "speechStart": + return self._speech_start(message) + case "transcript": + return self._transcript(message) + case "speechEnd": + return self._speech_end(message) + case "speechComplete": + return self._speech_complete(message) + case _: + return None + + def _turn(self, turn_id: str) -> _TurnState: existing: Final = self._turns.get(turn_id) if existing is not None: return existing created: Final = _TurnState(item_id=turn_id) self._turns[turn_id] = created + if len(self._turns) > self._turn_limit: + del self._turns[next(iter(self._turns))] return created - def _speech_start(self, message: Mapping[str, JsonValue]) -> None: + def _speech_start(self, message: Mapping[str, JsonValue]) -> _TurnState: turn: Final = self._turn(_required_turn_id(message, "speechStart")) - if turn is None: - return turn.started = True self._active_turn_id = turn.item_id + return turn - def _transcript(self, message: Mapping[str, JsonValue]) -> None: + def _transcript(self, message: Mapping[str, JsonValue]) -> _TurnState | None: transcript: Final = message.get("transcript") if not isinstance(transcript, str): raise MuseProtocolError("transcript event has invalid transcript") if not transcript and message.get("turnId") is None and self._active_turn_id is None: - return + return None turn: Final = self._turn(self._transcript_turn_id(message)) - if turn is None: - return - if message.get("final") is not True: - if turn.final_text is None: - turn.latest_partial = transcript - return - turn.final_text = transcript - turn.completed_signal = True - if self._mode == "PUSH_TO_TALK": - turn.stopped = True - if self._active_turn_id == turn.item_id: - self._active_turn_id = None + if message.get("final") is True: + self._finish(turn, transcript) + elif turn.final_text is None: + turn.latest_partial = transcript + return turn - def _speech_end(self, message: Mapping[str, JsonValue]) -> None: + def _speech_end(self, message: Mapping[str, JsonValue]) -> _TurnState: turn: Final = self._turn(_required_turn_id(message, "speechEnd")) - if turn is None: - return turn.stopped = True - if self._active_turn_id == turn.item_id: - self._active_turn_id = None + self._release_active(turn) + return turn - def _speech_complete(self, message: Mapping[str, JsonValue]) -> None: + def _speech_complete(self, message: Mapping[str, JsonValue]) -> _TurnState: transcript: Final = message.get("transcript") if not isinstance(transcript, str): raise MuseProtocolError("speechComplete event has invalid transcript") turn: Final = self._turn(_required_turn_id(message, "speechComplete")) - if turn is None: - return - turn.final_text = transcript - turn.completed_signal = True + self._finish(turn, transcript) + return turn + + def _finish(self, turn: _TurnState, transcript: str) -> None: + turn.finish(transcript) + self._release_active(turn) + + def _release_active(self, turn: _TurnState) -> None: + if self._active_turn_id == turn.item_id: + self._active_turn_id = None def _update_audio_progress(self, message: Mapping[str, JsonValue]) -> None: processed_ms: Final = message.get("audioProcessedMs") @@ -552,15 +548,6 @@ class MuseEventTransformer: self._unbilled_seconds += (float(processed_ms) - self._last_audio_processed_ms) / 1000 self._last_audio_processed_ms = float(processed_ms) - def _drained_events(self) -> Iterator[OpenAIRealtimeEvents]: - while self._turns: - turn_id, turn = next(iter(self._turns.items())) - yield from turn.drain(self.take_unbilled_usage) - if not turn.settled: - return - del self._turns[turn_id] - self._completed_turns.append(turn_id) - def _transcript_turn_id(self, message: Mapping[str, JsonValue]) -> str: if message.get("turnId") is not None: return _required_turn_id(message, "transcript") @@ -615,7 +602,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): model: str, logging_session_id: str, session_configuration_request: str | None = None, - ) -> MuseSessionCreatedEvent: + ) -> OpenAIRealtimeTranscriptionSessionCreated: return session_created_event(_DEFAULT_SESSION_CONFIG, logging_session_id) def transform_realtime_request( @@ -682,9 +669,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): return self._transformer.transform(frame) if not isinstance(session_id, str) or not session_id.strip(): raise MuseProtocolError("provider returned an invalid handshake response") - created: Final = session_created_event(self._require_config(), session_id.strip()) - event: Final = cast(OpenAIRealtimeEvents, created) # cast-ok: ReadOnly Muse session vs writable OpenAI fields - return (event,) + return (session_created_event(self._require_config(), session_id.strip()),) def _configure(self, message: str, model: str) -> tuple[str, ...]: if self._config is not None: diff --git a/litellm/types/llms/meta.py b/litellm/types/llms/meta.py index d7487bb09c8..40ecd6f7b67 100644 --- a/litellm/types/llms/meta.py +++ b/litellm/types/llms/meta.py @@ -19,40 +19,3 @@ class MuseHandshake(TypedDict): partialMode: ReadOnly[Literal["CUMULATIVE"]] emitAudioProgress: ReadOnly[bool] languageBias: NotRequired[ReadOnly[tuple[str, ...]]] - - -class MuseTranscriptionAudioFormat(TypedDict): - type: ReadOnly[Literal["audio/pcm"]] - rate: ReadOnly[MuseSampleRate] - - -class MuseTranscriptionSettings(TypedDict): - model: ReadOnly[str] - language: NotRequired[ReadOnly[str]] - - -class MuseTurnDetection(TypedDict): - type: ReadOnly[Literal["server_vad"]] - - -class MuseTranscriptionAudioInput(TypedDict): - format: ReadOnly[MuseTranscriptionAudioFormat] - transcription: ReadOnly[MuseTranscriptionSettings] - turn_detection: ReadOnly[MuseTurnDetection | None] - - -class MuseTranscriptionAudio(TypedDict): - input: ReadOnly[MuseTranscriptionAudioInput] - - -class MuseTranscriptionSession(TypedDict): - id: ReadOnly[str] - object: ReadOnly[Literal["realtime.transcription_session"]] - type: ReadOnly[Literal["transcription"]] - audio: ReadOnly[MuseTranscriptionAudio] - - -class MuseSessionCreatedEvent(TypedDict): - type: ReadOnly[Literal["session.created"]] - event_id: ReadOnly[str] - session: ReadOnly[MuseTranscriptionSession] diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 765ecaffdae..f852e125968 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -2188,6 +2188,53 @@ class OpenAIRealtimeInputAudioBufferSpeechEvent(TypedDict): item_id: ReadOnly[str] +class OpenAIRealtimeErrorDetail(TypedDict): + type: ReadOnly[str] + message: ReadOnly[str] + + +class OpenAIRealtimeErrorEvent(TypedDict): + type: ReadOnly[Literal["error"]] + error: ReadOnly[OpenAIRealtimeErrorDetail] + + +class OpenAIRealtimeTranscriptionAudioFormat(TypedDict): + type: ReadOnly[Literal["audio/pcm"]] + rate: ReadOnly[int] + + +class OpenAIRealtimeTranscriptionSettings(TypedDict): + model: ReadOnly[str] + language: NotRequired[ReadOnly[str]] + + +class OpenAIRealtimeServerVadTurnDetection(TypedDict): + type: ReadOnly[Literal["server_vad"]] + + +class OpenAIRealtimeTranscriptionAudioInput(TypedDict): + format: ReadOnly[OpenAIRealtimeTranscriptionAudioFormat] + transcription: ReadOnly[OpenAIRealtimeTranscriptionSettings] + turn_detection: ReadOnly[OpenAIRealtimeServerVadTurnDetection | None] + + +class OpenAIRealtimeTranscriptionAudio(TypedDict): + input: ReadOnly[OpenAIRealtimeTranscriptionAudioInput] + + +class OpenAIRealtimeTranscriptionSession(TypedDict): + id: ReadOnly[str] + object: ReadOnly[Literal["realtime.transcription_session"]] + type: ReadOnly[Literal["transcription"]] + audio: ReadOnly[OpenAIRealtimeTranscriptionAudio] + + +class OpenAIRealtimeTranscriptionSessionCreated(TypedDict): + type: ReadOnly[Literal["session.created"]] + event_id: ReadOnly[str] + session: ReadOnly[OpenAIRealtimeTranscriptionSession] + + class OpenAIRealtimeInputAudioTranscriptionDelta(TypedDict): type: ReadOnly[Literal["conversation.item.input_audio_transcription.delta"]] event_id: ReadOnly[str] @@ -2259,6 +2306,8 @@ OpenAIRealtimeEvents = ( | OpenAIRealtimeInputAudioBufferSpeechEvent | OpenAIRealtimeInputAudioTranscriptionDelta | OpenAIRealtimeInputAudioTranscriptionCompleted + | OpenAIRealtimeTranscriptionSessionCreated + | OpenAIRealtimeErrorEvent ) OpenAIRealtimeStreamList = list[OpenAIRealtimeEvents] diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py index 058e075860b..8b9eaf12dc2 100644 --- a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py +++ b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py @@ -1,4 +1,5 @@ import base64 +import itertools import json from typing import Final from unittest.mock import MagicMock @@ -18,6 +19,7 @@ from litellm.llms.meta.realtime.transformation import ( parse_session_update, session_created_event, ) +from litellm.types.llms.meta import MuseMode from litellm.types.realtime import RealtimeResponseTransformInput EMPTY_TRANSFORM_INPUT: Final[RealtimeResponseTransformInput] = { @@ -214,8 +216,7 @@ def test_cumulative_partials_emit_only_extensions_and_final_is_authoritative(): first = send(_event("transcript", turnId="turn-1", transcript="hello", final=False)) extension = send(_event("transcript", turnId="turn-1", transcript="hello world", final=False)) rewrite = send(_event("transcript", turnId="turn-1", transcript="hullo world", final=False)) - assert send(_event("speechComplete", turnId="turn-1", transcript="hullo world")) == () - completed = send(_event("speechEnd", turnId="turn-1")) + completed = send(_event("speechComplete", turnId="turn-1", transcript="hullo world")) assert [event["type"] for event in started] == ["input_audio_buffer.speech_started"] assert first[0]["delta"] == "hello" @@ -225,44 +226,125 @@ def test_cumulative_partials_emit_only_extensions_and_final_is_authoritative(): assert completed[1]["type"] == "conversation.item.input_audio_transcription.completed" assert completed[1]["item_id"] == "turn-1" assert completed[1]["transcript"] == "hullo world" + assert send(_event("speechEnd", turnId="turn-1")) == () -def test_completed_transcript_waits_for_speech_stopped(): +def test_speech_end_then_speech_complete_emits_stopped_then_completed(): transformer = MuseEventTransformer() transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) - assert transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="done"))) == () + stopped = transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) + completed = transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="done"))) - released = transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) - assert [event["type"] for event in released] == [ + assert [event["type"] for event in stopped] == ["input_audio_buffer.speech_stopped"] + assert [event["type"] for event in completed] == ["conversation.item.input_audio_transcription.completed"] + assert completed[0]["transcript"] == "done" + + +def _typed(events: tuple[dict[str, object], ...]) -> list[tuple[object, object]]: + return [(event["type"], event["item_id"]) for event in events] + + +def test_overlapping_turns_emit_independently_and_correlate_by_item_id(): + transformer = MuseEventTransformer() + + def send(payload: str) -> list[tuple[object, object]]: + return _typed(transformer.transform(json.loads(payload))) + + assert send(_event("speechStart", turnId="turn-a")) == [("input_audio_buffer.speech_started", "turn-a")] + assert send(_event("speechStart", turnId="turn-b")) == [("input_audio_buffer.speech_started", "turn-b")] + assert send(_event("transcript", turnId="turn-b", transcript="second", final=False)) == [ + ("conversation.item.input_audio_transcription.delta", "turn-b") + ] + assert send(_event("speechComplete", turnId="turn-a", transcript="first")) == [ + ("input_audio_buffer.speech_stopped", "turn-a"), + ("conversation.item.input_audio_transcription.completed", "turn-a"), + ] + assert send(_event("speechEnd", turnId="turn-a")) == [] + assert send(_event("speechEnd", turnId="turn-b")) == [("input_audio_buffer.speech_stopped", "turn-b")] + assert send(_event("speechComplete", turnId="turn-b", transcript="second final")) == [ + ("conversation.item.input_audio_transcription.completed", "turn-b") + ] + + +def test_empty_vad_turn_is_closed_and_does_not_block_the_next_turn(): + transformer = MuseEventTransformer() + + def send(payload: str) -> list[tuple[object, object]]: + return _typed(transformer.transform(json.loads(payload))) + + assert send(_event("speechStart", turnId="noise")) == [("input_audio_buffer.speech_started", "noise")] + assert send(_event("speechEnd", turnId="noise")) == [("input_audio_buffer.speech_stopped", "noise")] + assert send(_event("speechStart", turnId="speech")) == [("input_audio_buffer.speech_started", "speech")] + assert send(_event("transcript", turnId="speech", transcript="hello", final=False)) == [ + ("conversation.item.input_audio_transcription.delta", "speech") + ] + assert send(_event("speechEnd", turnId="speech")) == [("input_audio_buffer.speech_stopped", "speech")] + assert send(_event("speechComplete", turnId="speech", transcript="hello world")) == [ + ("conversation.item.input_audio_transcription.completed", "speech") + ] + + +@pytest.mark.parametrize("transcript", ["", "late words"]) +def test_late_speech_complete_after_an_empty_speech_end_completes_that_item(transcript: str): + transformer = MuseEventTransformer() + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) + transformer.transform(json.loads(_event("speechStart", turnId="turn-2"))) + + (completed,) = transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript=transcript))) + + assert completed["type"] == "conversation.item.input_audio_transcription.completed" + assert completed["item_id"] == "turn-1" + assert completed["transcript"] == transcript + + +def test_push_to_talk_speech_complete_closes_the_turn_without_speech_end(): + transformer = MuseEventTransformer() + transformer.configure(MuseSessionConfig(MUSE_MODEL, "PUSH_TO_TALK", 24_000, ())) + + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + transformer.transform(json.loads(_event("transcript", turnId="turn-1", transcript="hel", final=False))) + events = transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="hello"))) + + assert [event["type"] for event in events] == [ "input_audio_buffer.speech_stopped", "conversation.item.input_audio_transcription.completed", ] + assert events[1]["transcript"] == "hello" -def test_overlapping_turns_are_emitted_in_provider_turn_order(): +_TERMINAL_SIGNALS: Final = { + "speechEnd": _event("speechEnd", turnId="turn-1"), + "speechComplete": _event("speechComplete", turnId="turn-1", transcript="final words"), + "final": _event("transcript", turnId="turn-1", transcript="final words", final=True), +} +_TERMINAL_ORDERINGS: Final = tuple( + ordering for size in (1, 2, 3) for ordering in itertools.permutations(_TERMINAL_SIGNALS, size) +) + + +@pytest.mark.parametrize("mode", ["ENDPOINTING", "PUSH_TO_TALK"]) +@pytest.mark.parametrize("ordering", _TERMINAL_ORDERINGS, ids="-".join) +def test_every_terminal_signal_order_closes_the_turn_exactly_once(mode: MuseMode, ordering: tuple[str, ...]): transformer = MuseEventTransformer() + transformer.configure(MuseSessionConfig(MUSE_MODEL, mode, 24_000, ())) + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + transformer.transform(json.loads(_event("transcript", turnId="turn-1", transcript="fin", final=False))) - def send(payload: str) -> tuple[dict[str, object], ...]: - return transformer.transform(json.loads(payload)) - - send(_event("speechStart", turnId="turn-a")) - send(_event("speechStart", turnId="turn-b")) - assert send(_event("transcript", turnId="turn-b", transcript="second", final=False)) == () - assert send(_event("speechComplete", turnId="turn-a", transcript="first")) == () - released = send(_event("speechEnd", turnId="turn-a")) - - assert [(event["type"], event["item_id"]) for event in released] == [ - ("input_audio_buffer.speech_stopped", "turn-a"), - ("conversation.item.input_audio_transcription.completed", "turn-a"), - ("input_audio_buffer.speech_started", "turn-b"), - ("conversation.item.input_audio_transcription.delta", "turn-b"), + emitted = [ + event["type"] for signal in ordering for event in transformer.transform(json.loads(_TERMINAL_SIGNALS[signal])) ] - assert send(_event("speechComplete", turnId="turn-b", transcript="second final")) == () - final_b = send(_event("speechEnd", turnId="turn-b")) - assert final_b[0]["type"] == "input_audio_buffer.speech_stopped" - assert final_b[1]["item_id"] == "turn-b" - assert final_b[1]["transcript"] == "second final" + replayed = [ + event["type"] for signal in ordering for event in transformer.transform(json.loads(_TERMINAL_SIGNALS[signal])) + ] + + has_text = bool(set(ordering) & {"speechComplete", "final"}) + assert emitted == [ + "input_audio_buffer.speech_stopped", + *(["conversation.item.input_audio_transcription.completed"] if has_text else []), + ] + assert replayed == [] def test_push_to_talk_final_transcript_completes_without_speech_end(): @@ -290,12 +372,12 @@ def test_positive_audio_progress_deltas_attach_to_next_completion_and_speaker_is send(_event("audioProgress", audioProcessedMs=750)) send(_event("audioProgress", audioProcessedMs=1600)) assert send(_event("speaker", turnId=42, label=" Speaker 2 ")) == () - send(_event("speechComplete", turnId=42, transcript="hello")) - completed = send(_event("speechEnd", turnId=42)) + completed = send(_event("speechComplete", turnId=42, transcript="hello")) assert "speaker" not in completed[-1] assert completed[-1]["usage"] == {"type": "duration", "seconds": 1.6} assert transformer.take_unbilled_usage() is None + assert send(_event("speechEnd", turnId=42)) == () def test_trailing_audio_progress_is_returned_once(): @@ -307,23 +389,37 @@ def test_trailing_audio_progress_is_returned_once(): assert transformer.take_unbilled_usage() is None -def test_completed_turn_tombstone_suppresses_late_duplicates(): +def test_finished_turn_ignores_late_duplicates(): transformer = MuseEventTransformer() - transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="done"))) - released = transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) + released = transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="done"))) assert [event["type"] for event in released] == [ + "input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped", "conversation.item.input_audio_transcription.completed", ] assert transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="duplicate"))) == () assert transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) == () + assert transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) == () assert ( transformer.transform(json.loads(_event("transcript", turnId="turn-1", transcript="late", final=False))) == () ) +def test_turn_memory_is_bounded_by_turn_limit(): + transformer = MuseEventTransformer(turn_limit=2) + + transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="one"))) + transformer.transform(json.loads(_event("speechComplete", turnId="turn-2", transcript="two"))) + assert transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) == () + transformer.transform(json.loads(_event("speechComplete", turnId="turn-3", transcript="three"))) + + forgotten = transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) + + assert [event["type"] for event in forgotten] == ["input_audio_buffer.speech_stopped"] + + def test_provider_error_is_sanitized_and_encodable(): token = "private-token" provider_body = f"authorization failed for Bearer {token}" @@ -520,14 +616,11 @@ def test_provider_turn_events_and_close_usage_flow_through_config(): assert _backend_events(config, json.dumps({"type": "audioProgress", "audioProcessedMs": 1349})) == [] assert _backend_events(config, _event("speechStart", turnId="t1"))[0]["type"] == "input_audio_buffer.speech_started" - assert _backend_events(config, _event("speechComplete", turnId="t1", transcript="what is the weather")) == [] - completed = _backend_events(config, _event("speechEnd", turnId="t1")) + assert _backend_events(config, _event("speechEnd", turnId="t1"))[0]["type"] == "input_audio_buffer.speech_stopped" + completed = _backend_events(config, _event("speechComplete", turnId="t1", transcript="what is the weather")) - assert [event["type"] for event in completed] == [ - "input_audio_buffer.speech_stopped", - "conversation.item.input_audio_transcription.completed", - ] - assert completed[1]["usage"] == {"type": "duration", "seconds": 1.349} + assert [event["type"] for event in completed] == ["conversation.item.input_audio_transcription.completed"] + assert completed[0]["usage"] == {"type": "duration", "seconds": 1.349} assert config.unbilled_usage_on_session_close(MUSE_MODEL) is None assert _backend_events(config, json.dumps({"type": "audioProgress", "audioProcessedMs": 2349})) == [] From 414442cd06e528079b25a7cc806040daf466cd0a Mon Sep 17 00:00:00 2001 From: mateo Date: Sat, 12 Sep 2026 19:46:13 +0000 Subject: [PATCH 43/55] fix(registry): mark computer-use-preview as supporting pdf input Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/model_prices_and_context_window_backup.json | 1 + model_prices_and_context_window.json | 1 + 2 files changed, 2 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index d676a3abe15..2d8e8071479 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -14633,6 +14633,7 @@ ], "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": false, "supports_reasoning": true, "supports_response_schema": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index d676a3abe15..2d8e8071479 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -14633,6 +14633,7 @@ ], "supports_function_calling": true, "supports_parallel_function_calling": true, + "supports_pdf_input": true, "supports_prompt_caching": false, "supports_reasoning": true, "supports_response_schema": false, From a2e383a1a582c8d05170915ad63828046dd9f823 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 12 Sep 2026 12:53:03 -0700 Subject: [PATCH 44/55] fix(realtime): ignore a late speechStart for a finished Muse turn A duplicate speechStart for a turn that already stopped used to make that closed turn active again, so the next turnless PUSH_TO_TALK transcript was routed to the finished item and dropped. --- litellm/llms/meta/realtime/transformation.py | 2 ++ .../test_meta_realtime_transformation.py | 18 ++++++++++++++++++ 2 files changed, 20 insertions(+) diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py index c43e1897fbc..549db60be85 100644 --- a/litellm/llms/meta/realtime/transformation.py +++ b/litellm/llms/meta/realtime/transformation.py @@ -495,6 +495,8 @@ class MuseEventTransformer: def _speech_start(self, message: Mapping[str, JsonValue]) -> _TurnState: turn: Final = self._turn(_required_turn_id(message, "speechStart")) + if turn.stopped: + return turn turn.started = True self._active_turn_id = turn.item_id return turn diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py index 8b9eaf12dc2..3a4c0591fb5 100644 --- a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py +++ b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py @@ -407,6 +407,24 @@ def test_finished_turn_ignores_late_duplicates(): ) +def test_late_duplicate_speech_start_does_not_capture_the_next_turnless_transcript(): + transformer = MuseEventTransformer() + transformer.configure(MuseSessionConfig(MUSE_MODEL, "PUSH_TO_TALK", 24_000, ())) + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + transformer.transform(json.loads(_event("speechComplete", turnId="turn-1", transcript="first"))) + + assert transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) == () + events = transformer.transform(json.loads(_event("transcript", transcript="second", final=True))) + + assert [event["type"] for event in events] == [ + "input_audio_buffer.speech_started", + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + assert events[2]["transcript"] == "second" + assert events[2]["item_id"] != "turn-1" + + def test_turn_memory_is_bounded_by_turn_limit(): transformer = MuseEventTransformer(turn_limit=2) From f5f81e973aa2804b1c4119109973c9da48398f54 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:43:55 -0700 Subject: [PATCH 45/55] fix(budgets): preserve explicit reset interval clears --- .../budget_management_endpoints.py | 5 +++++ .../internal_user_endpoints.py | 10 ++++++---- .../tag_management_endpoints.py | 1 + litellm/proxy/management_helpers/utils.py | 11 ++++++++++- .../test_budget_endpoints.py | 10 +++++----- .../test_internal_user_endpoints.py | 16 ++++++++++++++++ .../components/AccessGroupBudgetModal.tsx | 6 +++++- .../_components/components/CreateTagModal.tsx | 6 +++++- .../tag-management/_components/tag_info.tsx | 2 +- ui/litellm-dashboard/src/components/Teams.tsx | 2 +- .../budget_duration_dropdown.tsx | 8 ++------ .../components/organisms/create_key_button.tsx | 2 +- .../src/components/tag_management/types.tsx | 2 +- .../src/components/team/EditMembership.tsx | 2 +- .../src/components/team/TeamInfo.tsx | 4 +++- 15 files changed, 63 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 62a24109dbb..81a607aaa43 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -14,6 +14,7 @@ All /budget management endpoints #### BUDGET TABLE MANAGEMENT #### import math from collections.abc import Mapping +from types import MappingProxyType from typing import Final from fastapi import APIRouter, Depends, HTTPException @@ -176,6 +177,10 @@ async def update_budget( recomputed_reset_at: Final = ( {"budget_reset_at": get_budget_reset_time(budget_duration=budget_obj.budget_duration)} if budget_obj.budget_duration is not None and "budget_reset_at" not in budget_obj.model_fields_set + else MappingProxyType({"budget_reset_at": None}) + if "budget_duration" in budget_obj.model_fields_set + and budget_obj.budget_duration is None + and "budget_reset_at" not in budget_obj.model_fields_set else {} ) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 10c11119006..000b7f874ee 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -1254,8 +1254,8 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda fields_set: Final = data.fields_set() if hasattr(data, "fields_set") else set() for k, v in data_json.items(): - if k == "max_budget": - if "max_budget" in fields_set: + if k in ("max_budget", "budget_duration"): + if k in fields_set: non_default_values[k] = v elif k == "model_max_budget": if k in fields_set: @@ -1283,8 +1283,10 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time validate_budget_duration(non_default_values["budget_duration"]) - non_default_values["budget_reset_at"] = get_budget_reset_time( - budget_duration=non_default_values["budget_duration"] + non_default_values["budget_reset_at"] = ( + get_budget_reset_time(budget_duration=non_default_values["budget_duration"]) + if non_default_values["budget_duration"] is not None + else None ) if "max_budget" not in non_default_values: diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index b74aa1a4e16..ab33d4bd766 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -438,6 +438,7 @@ async def update_tag( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, litellm_proxy_admin_name=litellm_proxy_admin_name, + budget_duration_cleared="budget_duration" in tag.model_fields_set and tag.budget_duration is None, ) # Get model names for model_info diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index e2d7262fb69..f3bd4b0f6dd 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -3,6 +3,7 @@ from collections.abc import Callable, Mapping, MutableMapping, Sequence from datetime import datetime from functools import wraps +from types import MappingProxyType from typing import Any, Final, Protocol from fastapi import HTTPException, Request @@ -180,6 +181,7 @@ async def handle_budget_for_entity( user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, litellm_proxy_admin_name: str, + budget_duration_cleared: bool = False, ) -> str | None: """ Common helper to handle budget creation/updates for entities (organizations, tags, etc). @@ -208,7 +210,14 @@ async def handle_budget_for_entity( # Extract budget fields from data _json_data: Final = data.model_dump(exclude_none=True) if hasattr(data, "model_dump") else data - _budget_data: Final = {k: v for k, v in _json_data.items() if k in budget_params} + _budget_data: Final = MappingProxyType( + { + k: _json_data.get(k) + for k in budget_params + if k in _json_data + or (k == "budget_duration" and existing_budget_id is not None and budget_duration_cleared) + } + ) # Check if budget_id is explicitly provided in the data data_budget_id: Final[str | None] = getattr(data, "budget_id", None) diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index 79d62f772bd..4b6815d7552 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -340,7 +340,8 @@ async def test_update_budget_recomputes_reset_at_when_duration_changes( @pytest.mark.asyncio -async def test_update_budget_preserves_explicit_reset_at(client_and_mocks): +@pytest.mark.parametrize("budget_duration", ["1d", None]) +async def test_update_budget_preserves_explicit_reset_at(client_and_mocks, budget_duration): """An explicit budget_reset_at from the caller always wins over recompute.""" client, _, mock_table = client_and_mocks captured = _capture_update_data(mock_table) @@ -350,7 +351,7 @@ async def test_update_budget_preserves_explicit_reset_at(client_and_mocks): "/budget/update", json={ "budget_id": "budget_explicit_reset", - "budget_duration": "1d", + "budget_duration": budget_duration, "budget_reset_at": explicit.isoformat(), }, ) @@ -377,8 +378,7 @@ async def test_update_budget_without_duration_leaves_reset_at_untouched( @pytest.mark.asyncio -async def test_update_budget_duration_none_does_not_recompute(client_and_mocks): - """Clearing budget_duration (explicit null) must not recompute against a None duration.""" +async def test_update_budget_duration_none_clears_obsolete_reset(client_and_mocks): client, _, mock_table = client_and_mocks captured = _capture_update_data(mock_table) @@ -389,7 +389,7 @@ async def test_update_budget_duration_none_does_not_recompute(client_and_mocks): assert resp.status_code == 200, resp.text assert "budget_duration" in captured and captured["budget_duration"] is None - assert "budget_reset_at" not in captured + assert captured["budget_reset_at"] is None @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 92b1ab1586d..0d8b19345f1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -2097,6 +2097,22 @@ def test_update_internal_user_params_reset_max_budget_with_none(): assert non_default_values["user_id"] == "test_user" +def test_update_internal_user_params_explicit_duration_clear_overrides_role_default(monkeypatch): + import litellm + + monkeypatch.setattr(litellm, "internal_user_budget_duration", "30d") + data = UpdateUserRequest( + user_id="duration-clear-test", + user_role=LitellmUserRoles.INTERNAL_USER, + budget_duration=None, + ) + + updated = _update_internal_user_params(data_json=data.model_dump(exclude_unset=True), data=data) + + assert updated["budget_duration"] is None + assert updated["budget_reset_at"] is None + + def test_update_internal_user_params_ignores_other_nones(): """ Test that other fields are still filtered out if None diff --git a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx index 8df2a6620f9..cf5cf78d08a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/models-and-endpoints/components/AccessGroupBudgetModal.tsx @@ -98,7 +98,11 @@ const AccessGroupBudgetModal: React.FC = ({ )} > {({ id, value, onChange }) => ( - + onChange(next ?? undefined)} + /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx index 6f6bdcb6fe4..adb5a167c35 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/components/CreateTagModal.tsx @@ -143,7 +143,11 @@ const CreateTagModal: React.FC = ({ visible, onCancel, onSu )} > {({ id, value, onChange }) => ( - + onChange(next ?? undefined)} + /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx index e8cb358c0cf..1648a99bb0e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/tag-management/_components/tag_info.tsx @@ -28,7 +28,7 @@ const tagEditShape = { description: z.string().optional(), models: z.array(z.string()).optional(), max_budget: z.union([z.string(), z.number()]).optional(), - budget_duration: z.string().optional(), + budget_duration: z.string().nullish(), }; const tagEditSchema = z.object(tagEditShape); diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index f0d21cc5350..dc531ea5dad 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -807,7 +807,7 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser showNeverResets placeholder={budgetDurationPlaceholder} value={value} - onChange={onChange} + onChange={(next) => onChange(next ?? undefined)} /> )} diff --git a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx index 3c38907c597..40ee857634c 100644 --- a/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx +++ b/ui/litellm-dashboard/src/components/common_components/budget_duration_dropdown.tsx @@ -14,7 +14,7 @@ const DURATION_LABELS: Record = { interface BudgetDurationDropdownProps { id?: string; value?: string | null; - onChange?: (value: string | undefined) => void; + onChange?: (value: string | null) => void; className?: string; style?: React.CSSProperties; placeholder?: string; @@ -31,11 +31,7 @@ const BudgetDurationDropdown: React.FC = ({ showNeverResets = false, }) => { return ( - diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 82580fd4667..b5789101f77 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1021,7 +1021,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp value={control.value as string | null | undefined} showNeverResets placeholder="Not set" - onChange={control.onChange} + onChange={(next) => control.onChange(next ?? undefined)} /> )} diff --git a/ui/litellm-dashboard/src/components/tag_management/types.tsx b/ui/litellm-dashboard/src/components/tag_management/types.tsx index 3cf17545fd6..88dfa28204d 100644 --- a/ui/litellm-dashboard/src/components/tag_management/types.tsx +++ b/ui/litellm-dashboard/src/components/tag_management/types.tsx @@ -41,7 +41,7 @@ export interface TagUpdateRequest { soft_budget?: number; tpm_limit?: number; rpm_limit?: number; - budget_duration?: string; + budget_duration?: string | null; } export interface TagDeleteRequest { diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.tsx index 036c4f1cc0b..909b5d56c97 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.tsx @@ -157,7 +157,7 @@ const MemberModal = ({ onChange(next)} + onChange={(next) => onChange(mode === "add" ? next ?? undefined : next)} /> ); default: diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 8e3a17c2622..ffc83d0165e 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -1467,7 +1467,9 @@ const TeamInfoView: React.FC = ({ showNeverResets placeholder="Inherit team reset period" value={value === null ? NEVER_RESETS_BUDGET_DURATION : value} - onChange={(next) => onChange(next === NEVER_RESETS_BUDGET_DURATION ? null : next)} + onChange={(next) => + onChange(next === NEVER_RESETS_BUDGET_DURATION ? null : next ?? undefined) + } /> )} From ad966d834020368e5ee1aa9974bb3ff1b5492cd1 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:43:55 -0700 Subject: [PATCH 46/55] fix(projects): persist explicit budget cap clears --- .../management_endpoints/project_endpoints.py | 5 ++++- .../test_project_endpoints_prisma.py | 15 +++++++++++++++ .../hooks/projects/useUpdateProject.ts | 2 +- .../EditProjectModal.integration.test.tsx | 14 ++++++++++++++ .../ProjectModals/EditProjectModal.tsx | 2 +- .../_components/ProjectModals/ProjectBaseForm.tsx | 15 +++++++++++++-- .../ProjectModals/projectFormSchema.ts | 2 +- .../ProjectModals/projectFormUtils.test.ts | 6 +++--- .../_components/ProjectModals/projectFormUtils.ts | 12 ++++++++++-- 9 files changed, 62 insertions(+), 11 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py index b2eda76f9ae..f40ced302ce 100644 --- a/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py +++ b/enterprise/litellm_enterprise/proxy/management_endpoints/project_endpoints.py @@ -780,7 +780,10 @@ async def update_project( # Handle budget updates budget_fields = LiteLLM_BudgetTable.model_fields.keys() - budget_updates = {k: v for k, v in update_data.items() if k in budget_fields} + budget_updates = { + **{k: v for k, v in update_data.items() if k in budget_fields}, + **({"max_budget": None} if "max_budget" in data.model_fields_set and data.max_budget is None else {}), + } if budget_updates and existing_project.budget_id: # Update existing budget diff --git a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py index c23b203feba..36878fa698c 100644 --- a/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py +++ b/tests/enterprise/litellm_enterprise/proxy/management_endpoints/test_project_endpoints_prisma.py @@ -1292,6 +1292,21 @@ async def test_update_project_leaves_metadata_untouched_when_no_limit_is_sent(mo assert "metadata" not in _written_project_data(mock_prisma) +@pytest.mark.asyncio +async def test_update_project_clears_only_the_explicit_budget_cap(monkeypatch): + mock_prisma = _project_update_mocks(monkeypatch, {}) + mock_prisma.db.litellm_projecttable.find_unique.return_value.budget_id = "budget-clear-test" + mock_prisma.db.litellm_budgettable.update = mock.AsyncMock() + + await _run_project_update("project-clear-test", max_budget=None) + + mock_prisma.db.litellm_budgettable.update.assert_awaited_once_with( + where={"budget_id": "budget-clear-test"}, + data={"max_budget": None, "updated_by": "1234"}, + ) + assert "max_budget" not in _written_project_data(mock_prisma) + + @pytest.mark.parametrize("entry", ["all-proxy-models", "*", "azure/*"]) def test_enforce_project_model_quota_rejects_entries_that_expand_at_request_time(entry): """A quota keyed on a wildcard entry is never applied by the limiter, so it fails loudly.""" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts index 8e6bad04a28..e6dec85128c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/hooks/projects/useUpdateProject.ts @@ -10,7 +10,7 @@ export interface ProjectUpdateParams { description?: string; team_id?: string; models?: string[]; - max_budget?: number; + max_budget?: number | null; blocked?: boolean; guardrails?: string[]; metadata?: Record; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx index 5b84aa15dd1..2e39d45e660 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.integration.test.tsx @@ -102,6 +102,20 @@ describe("EditProjectModal submit payload", () => { }); }); + it("should send an explicit clear after blanking a saved budget", async () => { + const user = setup(); + renderModal(); + + const budgetInput = screen.getByRole("spinbutton", { name: "Max Budget (USD)" }); + await user.clear(budgetInput); + await user.tab(); + expect(budgetInput).toHaveValue(null); + await save(user); + + await waitFor(() => expect(mutate).toHaveBeenCalled()); + expect(JSON.parse(JSON.stringify(variables().params))).toMatchObject({ max_budget: null }); + }); + it("includes the advanced fields once Advanced Settings has been opened, even after collapsing it again", async () => { const user = setup(); renderModal(); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx index 77f28b05ea5..31582da5b14 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/EditProjectModal.tsx @@ -86,7 +86,7 @@ function EditProjectForm({ project, onClose, onSuccess }: Omit onChange(toOptionalNumber(event.target.value))} + value={Number.isNaN(value) ? "" : value ?? ""} + onInput={(event) => { + if (event.currentTarget.validity.badInput || Number.isNaN(value)) { + onChange( + event.currentTarget.validity.badInput + ? Number.NaN + : toOptionalNumber(event.currentTarget.value) ?? null, + ); + } + }} + onChange={(event) => + onChange(event.target.validity.badInput ? Number.NaN : toOptionalNumber(event.target.value) ?? null) + } /> )} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts index d4c85d89616..6e6318a9609 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormSchema.ts @@ -22,7 +22,7 @@ export const projectFormSchema = z .pipe(z.string({ error: "Please select a team" }).min(1, "Please select a team")), description: z.string().optional(), models: z.array(z.string()), - max_budget: z.number().optional(), + max_budget: z.number().nullish(), isBlocked: z.boolean(), guardrails: z.array(z.string()).optional(), modelLimits: z.array(modelLimitSchema).optional(), diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts index 28d791ddf54..e76555b3ca3 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.test.ts @@ -27,9 +27,9 @@ describe("buildProjectCreateParams", () => { expect(result.description).toBe("A description"); }); - it("should pass through max_budget when provided", () => { - const result = buildProjectCreateParams({ ...baseValues, max_budget: 50.0 }); - expect(result.max_budget).toBe(50.0); + it.each([50.0, 1e308])("should preserve a finite max_budget of %s", (maxBudget) => { + const result = buildProjectCreateParams({ ...baseValues, max_budget: maxBudget }); + expect(JSON.parse(JSON.stringify(result)).max_budget).toBe(maxBudget); }); it("should build model_rpm_limit from modelLimits entries", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts index 71c1f5c79af..60e97939e69 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/projects/_components/ProjectModals/projectFormUtils.ts @@ -16,6 +16,11 @@ const buildModelLimitMap = ( const buildMetadata = (entries: ProjectFormValues["metadata"]): Record | undefined => entries && Object.fromEntries(entries.flatMap((entry) => (entry.key ? [[entry.key, entry.value] as const] : []))); +const roundBudget = (value: number): number => { + const rounded = Math.round(value * 100) / 100; + return Number.isFinite(rounded) ? rounded : value; +}; + const buildProjectApiParams = (values: ProjectFormValues, sendEmpty: boolean) => { const limitEntries = values.modelLimits ?? []; const modelRpmLimit = buildModelLimitMap(limitEntries, (entry) => entry.rpm); @@ -35,7 +40,7 @@ const buildProjectApiParams = (values: ProjectFormValues, sendEmpty: boolean) => project_alias: values.project_alias, description: values.description, models: values.models ?? [], - max_budget: values.max_budget === undefined ? undefined : Math.round(values.max_budget * 100) / 100, + max_budget: values.max_budget == null ? undefined : roundBudget(values.max_budget), blocked: values.isBlocked ?? false, ...guardrailsParam, ...(keep(modelRpmLimit) && { model_rpm_limit: modelRpmLimit }), @@ -53,4 +58,7 @@ export const buildProjectCreateParams = (values: ProjectFormValues) => buildProj * /project/update leaves an omitted key untouched, so a limit the operator cleared has to go out as * an explicitly empty map. Omitting it is what silently kept a removed quota enforced. */ -export const buildProjectUpdateParams = (values: ProjectFormValues) => buildProjectApiParams(values, true); +export const buildProjectUpdateParams = (values: ProjectFormValues, savedMaxBudget?: number | null) => ({ + ...buildProjectApiParams(values, true), + ...(values.max_budget == null && savedMaxBudget != null ? { max_budget: null } : {}), +}); From b24480941bd06cb615f0489767d35f1616a775d0 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 13:43:55 -0700 Subject: [PATCH 47/55] fix(ui): retain cleared user budget values after save --- .../user_info_view.integration.test.tsx | 31 ++++++++++++++++++- .../_components/view_users/user_info_view.tsx | 5 +-- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx index 69254b1ffe4..0f1a44851c7 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx @@ -1,4 +1,4 @@ -import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { fireEvent, renderWithProviders as render, screen, waitFor } from "../../../../../../tests/test-utils"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; import { describe, expect, it, vi, beforeEach } from "vitest"; import UserInfoView from "./user_info_view"; @@ -163,6 +163,35 @@ describe("UserInfoView add-to-team form", () => { expect(await openEditor(user)).toHaveValue(42); }); + + it("should keep Unlimited selected after saving and reopening the user", async () => { + const user = setup(); + render(); + + await openEditor(user); + await user.click(screen.getByRole("checkbox", { name: "Unlimited Budget" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalled()); + expect(mockUserUpdateUserCall.mock.calls[0][1]).toMatchObject({ max_budget: null }); + await openEditor(user); + expect(screen.getByRole("checkbox", { name: "Unlimited Budget" })).toBeChecked(); + }); + + it("should keep a cleared reset period after saving and reopening the user", async () => { + const user = setup(); + render(); + + await openEditor(user); + await user.click(screen.getByRole("combobox", { name: "Reset Budget" })); + await user.click(await screen.findByRole("option", { name: "n/a" })); + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(mockUserUpdateUserCall).toHaveBeenCalled()); + expect(mockUserUpdateUserCall.mock.calls[0][1]).toMatchObject({ budget_duration: null }); + await openEditor(user); + expect(screen.getByRole("combobox", { name: "Reset Budget" })).toHaveTextContent("n/a"); + }); }); it("offers only the teams the user is not already a member of", async () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx index eed39c8e585..e083e549552 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx @@ -332,8 +332,9 @@ export default function UserInfoView({ user_email: formValues.user_email ?? userData.user_email, user_alias: formValues.user_alias ?? userData.user_alias, models: formValues.models ?? userData.models, - max_budget: formValues.max_budget ?? userData.max_budget, - budget_duration: formValues.budget_duration ?? userData.budget_duration, + max_budget: formValues.max_budget === undefined ? userData.max_budget : formValues.max_budget, + budget_duration: + formValues.budget_duration === undefined ? userData.budget_duration : formValues.budget_duration, metadata: formValues.metadata ?? userData.metadata, model_max_budget: formValues.model_max_budget ?? userData.model_max_budget, object_permission: mcpEntitlement From 261807114a50e09b7fe6b48a7d8709dbf25e9c02 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 12 Sep 2026 14:10:58 -0700 Subject: [PATCH 48/55] fix(proxy): accept both deferred stream logging arg shapes on native routes (#40869) * fix(proxy): accept both deferred stream logging arg shapes on native routes _arm_deferred_stream_dispatch armed a one-argument closure on every anthropic_messages/aresponses stream that was not a CustomStreamWrapper or a LiteLLMCompletionStreamingIterator. The bridged /v1/messages path returns a plain SSE generator that shares its inner CustomStreamWrapper logging_obj, so it stores (assembled_response, cache_hit) and _fire_deferred_stream_logging raised TypeError, dropping spend logs and callbacks and ending the stream with an error. The closure now dispatches on the stored args shape: a single coroutine is enqueued, a two-tuple runs success handlers, anything else is logged and dropped Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): assert dropped deferred payload via caplog instead of patching the logger Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/proxy/common_request_processing.py | 66 ++++++++-------- .../test_deferred_guardrail_logging.py | 78 ++++++++++++++++++- 2 files changed, 107 insertions(+), 37 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 289fe086379..e6ed60ba177 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3,7 +3,7 @@ import contextlib import json import logging import math -from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine, Mapping, Sequence +from collections.abc import AsyncGenerator, Awaitable, Callable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType @@ -3210,20 +3210,24 @@ class ProxyBaseLLMRequestProcessing: end-of-stream blocks complete, so the spend log sees guardrail_information. - Three closure shapes, matching who owns logging for the stream: + Two closure shapes, matching who owns logging for the stream: - CustomStreamWrapper (chat completions) stores (assembled_response, cache_hit); the closure also runs non-apply_guardrail post-call hooks via _run_deferred_stream_guardrails. - - Bridged /v1/responses (LiteLLMCompletionStreamingIterator) shares - its inner CustomStreamWrapper's logging_obj, so it stores the same - (assembled_response, cache_hit) shape; the closure only dispatches - success logging, matching the route's pre-existing hook surface. - - Native anthropic_messages/aresponses iterators store a single - ready-made logging coroutine to enqueue. + - Every other anthropic_messages/aresponses stream gets a closure + that dispatches on the stored args shape, because the arming site + cannot tell the producers apart: native iterators store a single + ready-made logging coroutine to enqueue, while bridged streams + (LiteLLMCompletionStreamingIterator, and the plain SSE generator + AnthropicStreamWrapper returns for bridged /v1/messages) share + their inner CustomStreamWrapper's logging_obj and so store + (assembled_response, cache_hit); for those the closure only + dispatches success logging, matching the route's pre-existing + hook surface. - Raw async generators from passthrough routes bypass all three and - would orphan the closure, so they are not armed here. + Raw async generators from passthrough routes bypass both and would + orphan the closure, so they are not armed here. The router wraps iterators that cannot carry _hidden_params in HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the @@ -3257,31 +3261,27 @@ class ProxyBaseLLMRequestProcessing: if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response): return - from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, - ) - - if isinstance(unwrapped, LiteLLMCompletionStreamingIterator): - _captured_bridge_logging_obj: Final = logging_obj - - async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None: - await _as_success_dispatcher(_captured_bridge_logging_obj).dispatch_success_handlers( - assembled_response, - cache_hit=cache_hit, - start_time=None, - end_time=None, - prefer_async_handlers=True, - ) - - logging_obj._on_deferred_stream_complete = _on_deferred_bridged_stream_complete - return - from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - async def _on_deferred_native_stream_complete( - logging_coroutine: Coroutine[object, object, object], - ) -> None: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + _captured_native_logging_obj: Final = logging_obj + + async def _on_deferred_native_stream_complete(*args: object) -> None: + match args: + case (logging_coroutine,) if asyncio.iscoroutine(logging_coroutine): + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + case (assembled_response, cache_hit): + await _as_success_dispatcher(_captured_native_logging_obj).dispatch_success_handlers( + assembled_response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + prefer_async_handlers=True, + ) + case _: + verbose_proxy_logger.error( + "Deferred stream logging dropped: unexpected stored args shape %s", + tuple(type(arg).__name__ for arg in args), + ) logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 11c3d2f8b20..c550a0a41d2 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -15,6 +15,7 @@ Streaming: CSW.__anext__ stores args on logging_obj at stream end. """ import asyncio +import logging from typing import Any, Final from unittest.mock import AsyncMock, MagicMock, patch @@ -1422,7 +1423,7 @@ class TestArmDeferredStreamDispatch: async def test_native_stream_closure_enqueues_single_coroutine(self): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER - logging_obj, _ = self._dispatch_recording_logging_obj() + logging_obj, recorded = self._dispatch_recording_logging_obj() async def _agen(): yield b"x" @@ -1433,20 +1434,89 @@ class TestArmDeferredStreamDispatch: user_api_key_dict=MagicMock(), logging_obj=logging_obj, ) - closure = logging_obj._on_deferred_stream_complete - assert closure is not None + assert logging_obj._on_deferred_stream_complete is not None async def _logging_coroutine(): return None coro = _logging_coroutine() + logging_obj._deferred_stream_complete_args = (coro,) with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" ) as mock_enqueue: - await closure(coro) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) mock_enqueue.assert_called_once_with(async_coroutine=coro) + assert recorded == {} coro.close() + @pytest.mark.asyncio + @pytest.mark.parametrize("route_type", ["anthropic_messages", "aresponses"]) + async def test_raw_generator_stream_storing_csw_arg_shape_dispatches_success(self, route_type): + """Bridged /v1/messages returns AnthropicStreamWrapper's plain SSE + generator, which shares its inner CustomStreamWrapper's logging_obj and + so stores (assembled_response, cache_hit). The closure armed for a raw + generator must accept that shape too, or _fire_deferred_stream_logging + raises TypeError and the request loses its spend log and callbacks.""" + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj, recorded = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type=route_type, + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, True) + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue: + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + mock_enqueue.assert_not_called() + assert recorded["result"] is assembled + assert recorded["cache_hit"] is True + assert recorded["prefer_async_handlers"] is True + + @pytest.mark.asyncio + @pytest.mark.parametrize("stored_args", [(object(),), (object(), object(), object())]) + async def test_raw_generator_stream_with_unknown_arg_shape_logs_and_drops(self, stored_args, caplog): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj, recorded = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type="anthropic_messages", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + logging_obj._deferred_stream_complete_args = stored_args + with ( + patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue, + caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), + ): + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + mock_enqueue.assert_not_called() + assert recorded == {} + dropped = [r for r in caplog.records if r.getMessage().startswith("Deferred stream logging dropped")] + assert len(dropped) == 1 + @pytest.mark.asyncio async def test_csw_closure_routes_through_deferred_stream_guardrails(self, monkeypatch): from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper From 09d63b259c515fbb2f772e5aea5d259eb3c6ef41 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 12 Sep 2026 21:11:02 +0000 Subject: [PATCH 49/55] test(fireworks): drop hardcoded price snapshot tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 248 ------------------------------- 1 file changed, 248 deletions(-) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 19ed31c7b22..965f960f49b 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4245,254 +4245,6 @@ def test_deepseek_flash_completion_cost(): assert cost == pytest.approx(1.50, abs=1e-9) -_FIREWORKS_MODELS = [ - ( - "accounts/fireworks/models/glm-5p2", - 1.4e-06, - 4.4e-06, - 1.4e-07, - 1048576, - 131072, - False, - True, - ), - ( - "accounts/fireworks/models/glm-5p1", - 1.4e-06, - 4.4e-06, - 2.6e-07, - 202800, - 131072, - False, - True, - ), - ( - "accounts/fireworks/routers/glm-5p1-fast", - 2.8e-06, - 8.8e-06, - 5.2e-07, - 202800, - 131072, - False, - True, - ), - ( - "accounts/fireworks/models/qwen3p7-plus", - 4e-07, - 1.6e-06, - 8e-08, - 262144, - 65536, - True, - True, - ), - ( - "accounts/fireworks/models/minimax-m3", - 3e-07, - 1.2e-06, - 6e-08, - 512000, - 512000, - True, - True, - ), - ( - "accounts/fireworks/models/minimax-m2p7", - 3e-07, - 1.2e-06, - 6e-08, - 196608, - 196608, - False, - True, - ), - ( - "accounts/fireworks/models/kimi-k2p7-code", - 9.5e-07, - 4e-06, - 1.9e-07, - 262144, - 32768, - True, - True, - ), - ( - "accounts/fireworks/routers/kimi-k2p7-code-fast", - 1.9e-06, - 8e-06, - 3.8e-07, - 262144, - 32768, - True, - True, - ), - ( - "accounts/fireworks/models/kimi-k2p6", - 9.5e-07, - 4e-06, - 1.6e-07, - 262144, - 32768, - True, - True, - ), - ( - "accounts/fireworks/routers/kimi-k2p6-fast", - 2e-06, - 8e-06, - 3e-07, - 262144, - 32768, - True, - True, - ), - ( - "accounts/fireworks/models/gpt-oss-120b", - 1.5e-07, - 6e-07, - 1.5e-08, - 131072, - 32768, - False, - True, - ), - ( - "accounts/fireworks/models/gpt-oss-20b", - 7e-08, - 3e-07, - 3.5e-08, - 131072, - 32768, - False, - True, - ), - ( - "accounts/fireworks/models/deepseek-v4-pro", - 1.74e-06, - 3.48e-06, - 1.45e-07, - 1048576, - 384000, - False, - True, - ), - ( - "accounts/fireworks/models/deepseek-v4-flash", - 1.4e-07, - 2.8e-07, - 2.8e-08, - 1048576, - 384000, - False, - True, - ), -] - -_FIREWORKS_SHORT_FORMS = [ - "glm-5p2", - "glm-5p1", - "qwen3p7-plus", - "minimax-m3", - "minimax-m2p7", - "kimi-k2p7-code", - "kimi-k2p6", - "gpt-oss-120b", - "gpt-oss-20b", - "deepseek-v4-pro", - "deepseek-v4-flash", -] - -_FIREWORKS_ROUTER_SHORT_FORMS = [ - "glm-5p1-fast", - "kimi-k2p6-fast", - "kimi-k2p7-code-fast", -] - - -def _assert_fireworks_entry( - model_cost, - model_path, - expected_input, - expected_output, - expected_cache, - expected_max_input, - expected_max_output, - expected_vision, - expected_reasoning, -): - info = model_cost.get(f"fireworks_ai/{model_path}") - assert info is not None, f"fireworks_ai/{model_path} missing from model cost map" - assert info["litellm_provider"] == "fireworks_ai" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert info["cache_read_input_token_cost"] == expected_cache - assert info["max_input_tokens"] == expected_max_input - assert info["max_output_tokens"] == expected_max_output - assert info["max_tokens"] == expected_max_output - assert info["supports_function_calling"] is True - assert info["supports_tool_choice"] is True - assert info["supports_reasoning"] is expected_reasoning - assert info["supports_response_schema"] is True - assert info["supports_vision"] is expected_vision - - -def test_fireworks_models_in_cost_map(): - import json - from pathlib import Path - - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - for entry in _FIREWORKS_MODELS: - _assert_fireworks_entry(model_cost, *entry) - - for short in _FIREWORKS_SHORT_FORMS: - long_key = f"fireworks_ai/accounts/fireworks/models/{short}" - short_key = f"fireworks_ai/{short}" - assert model_cost.get(short_key) == model_cost.get( - long_key - ), f"short-form {short_key} does not match long-form {long_key}" - - for short in _FIREWORKS_ROUTER_SHORT_FORMS: - long_key = f"fireworks_ai/accounts/fireworks/routers/{short}" - short_key = f"fireworks_ai/{short}" - assert model_cost.get(short_key) == model_cost.get( - long_key - ), f"short-form {short_key} does not match long-form {long_key}" - - -def test_fireworks_models_in_backup_cost_map(): - import json - from pathlib import Path - - json_path = ( - Path(__file__).parents[2] - / "litellm" - / "model_prices_and_context_window_backup.json" - ) - with open(json_path) as f: - model_cost = json.load(f) - - for entry in _FIREWORKS_MODELS: - _assert_fireworks_entry(model_cost, *entry) - - for short in _FIREWORKS_SHORT_FORMS: - long_key = f"fireworks_ai/accounts/fireworks/models/{short}" - short_key = f"fireworks_ai/{short}" - assert model_cost.get(short_key) == model_cost.get( - long_key - ), f"short-form {short_key} does not match long-form {long_key}" - - for short in _FIREWORKS_ROUTER_SHORT_FORMS: - long_key = f"fireworks_ai/accounts/fireworks/routers/{short}" - short_key = f"fireworks_ai/{short}" - assert model_cost.get(short_key) == model_cost.get( - long_key - ), f"short-form {short_key} does not match long-form {long_key}" - - class TestBedrockBaseModelLabelKeepsTools: """Regression for #29618: a Bedrock deployment whose ``base_model`` is a friendly label must not silently drop ``tools``/``tool_choice`` under ``drop_params``.""" From 311d9bba37b44099461e6e699232b0903474fc62 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 12 Sep 2026 14:19:20 -0700 Subject: [PATCH 50/55] perf(policy_engine): dedup attachments in one pass after sorting (#40883) get_attached_policies_with_reasons rescanned the sorted matches with next() once per distinct policy, which is quadratic and misses the one second budget past a few thousand global attachments. Build a policy to broadest attachment map in one pass instead, keeping the specificity sort and result order. Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/policy_engine/attachment_registry.py | 6 +++++- .../policy_engine/test_attachment_registry.py | 16 ++++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/policy_engine/attachment_registry.py b/litellm/proxy/policy_engine/attachment_registry.py index a8ead86ac36..76b2291774e 100644 --- a/litellm/proxy/policy_engine/attachment_registry.py +++ b/litellm/proxy/policy_engine/attachment_registry.py @@ -6,6 +6,7 @@ This allows the same policy to be attached to multiple scopes. """ from datetime import datetime, timezone +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, TypedDict from litellm._logging import verbose_proxy_logger @@ -141,8 +142,11 @@ class AttachmentRegistry: ), key=_attachment_specificity, ) + broadest_attachment_by_policy: Final = MappingProxyType( + {attachment.policy: attachment for attachment in reversed(matching_attachments)} + ) unique_attachments: Final = tuple( - next(attachment for attachment in matching_attachments if attachment.policy == policy_name) + broadest_attachment_by_policy[policy_name] for policy_name in dict.fromkeys(attachment.policy for attachment in matching_attachments) ) diff --git a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py index 87b1bc56659..fa37a02a37c 100644 --- a/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py +++ b/tests/test_litellm/proxy/policy_engine/test_attachment_registry.py @@ -4,6 +4,7 @@ Unit tests for AttachmentRegistry - tests policy attachment matching. Tests the main entry point: get_attached_policies() """ +import time from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock @@ -222,6 +223,21 @@ class TestGetAttachedPolicies: # Should only appear once assert attached.count("multi-policy") == 1 + def test_many_distinct_policies_resolve_in_linear_time(self): + policy_count = 20_000 + registry = AttachmentRegistry() + registry.load_attachments( + [{"policy": f"policy-{index}", "scope": "*"} for index in range(policy_count)] + ) + context = PolicyMatchContext(team_alias="team", key_alias="key", model="gpt-4") + + started = time.perf_counter() + attached = registry.get_attached_policies(context) + elapsed = time.perf_counter() - started + + assert attached == [f"policy-{index}" for index in range(policy_count)] + assert elapsed < 1.0, f"{policy_count} attachments took {elapsed:.2f}s, dedup is no longer one pass" + def test_no_attachments_returns_empty(self): """Test empty attachments returns empty list.""" registry = AttachmentRegistry() From fbcc6021220e54460bd27792dd1426e189abae61 Mon Sep 17 00:00:00 2001 From: kerry Date: Sat, 12 Sep 2026 21:19:35 +0000 Subject: [PATCH 51/55] test(fireworks): stop pinning prices in the cost-map tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm/test_utils.py | 203 +++++++++++++++++++++++++++++++ 1 file changed, 203 insertions(+) diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 965f960f49b..835e87aff88 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4245,6 +4245,209 @@ def test_deepseek_flash_completion_cost(): assert cost == pytest.approx(1.50, abs=1e-9) +_FIREWORKS_MODELS = [ + ( + "accounts/fireworks/models/glm-5p2", + 1048576, + 131072, + False, + True, + ), + ( + "accounts/fireworks/models/glm-5p1", + 202800, + 131072, + False, + True, + ), + ( + "accounts/fireworks/routers/glm-5p1-fast", + 202800, + 131072, + False, + True, + ), + ( + "accounts/fireworks/models/qwen3p7-plus", + 262144, + 65536, + True, + True, + ), + ( + "accounts/fireworks/models/minimax-m3", + 512000, + 512000, + True, + True, + ), + ( + "accounts/fireworks/models/minimax-m2p7", + 196608, + 196608, + False, + True, + ), + ( + "accounts/fireworks/models/kimi-k2p7-code", + 262144, + 32768, + True, + True, + ), + ( + "accounts/fireworks/routers/kimi-k2p7-code-fast", + 262144, + 32768, + True, + True, + ), + ( + "accounts/fireworks/models/kimi-k2p6", + 262144, + 32768, + True, + True, + ), + ( + "accounts/fireworks/routers/kimi-k2p6-fast", + 262144, + 32768, + True, + True, + ), + ( + "accounts/fireworks/models/gpt-oss-120b", + 131072, + 32768, + False, + True, + ), + ( + "accounts/fireworks/models/gpt-oss-20b", + 131072, + 32768, + False, + True, + ), + ( + "accounts/fireworks/models/deepseek-v4-pro", + 1048576, + 384000, + False, + True, + ), + ( + "accounts/fireworks/models/deepseek-v4-flash", + 1048576, + 384000, + False, + True, + ), +] + +_FIREWORKS_SHORT_FORMS = [ + "glm-5p2", + "glm-5p1", + "qwen3p7-plus", + "minimax-m3", + "minimax-m2p7", + "kimi-k2p7-code", + "kimi-k2p6", + "gpt-oss-120b", + "gpt-oss-20b", + "deepseek-v4-pro", + "deepseek-v4-flash", +] + +_FIREWORKS_ROUTER_SHORT_FORMS = [ + "glm-5p1-fast", + "kimi-k2p6-fast", + "kimi-k2p7-code-fast", +] + + +def _assert_fireworks_entry( + model_cost, + model_path, + expected_max_input, + expected_max_output, + expected_vision, + expected_reasoning, +): + info = model_cost.get(f"fireworks_ai/{model_path}") + assert info is not None, f"fireworks_ai/{model_path} missing from model cost map" + assert info["litellm_provider"] == "fireworks_ai" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] > 0 + assert info["output_cost_per_token"] > 0 + assert "cache_read_input_token_cost" in info + assert info["max_input_tokens"] == expected_max_input + assert info["max_output_tokens"] == expected_max_output + assert info["max_tokens"] == expected_max_output + assert info["supports_function_calling"] is True + assert info["supports_tool_choice"] is True + assert info["supports_reasoning"] is expected_reasoning + assert info["supports_response_schema"] is True + assert info["supports_vision"] is expected_vision + + +def test_fireworks_models_in_cost_map(): + import json + from pathlib import Path + + json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" + with open(json_path) as f: + model_cost = json.load(f) + + for entry in _FIREWORKS_MODELS: + _assert_fireworks_entry(model_cost, *entry) + + for short in _FIREWORKS_SHORT_FORMS: + long_key = f"fireworks_ai/accounts/fireworks/models/{short}" + short_key = f"fireworks_ai/{short}" + assert model_cost.get(short_key) == model_cost.get( + long_key + ), f"short-form {short_key} does not match long-form {long_key}" + + for short in _FIREWORKS_ROUTER_SHORT_FORMS: + long_key = f"fireworks_ai/accounts/fireworks/routers/{short}" + short_key = f"fireworks_ai/{short}" + assert model_cost.get(short_key) == model_cost.get( + long_key + ), f"short-form {short_key} does not match long-form {long_key}" + + +def test_fireworks_models_in_backup_cost_map(): + import json + from pathlib import Path + + json_path = ( + Path(__file__).parents[2] + / "litellm" + / "model_prices_and_context_window_backup.json" + ) + with open(json_path) as f: + model_cost = json.load(f) + + for entry in _FIREWORKS_MODELS: + _assert_fireworks_entry(model_cost, *entry) + + for short in _FIREWORKS_SHORT_FORMS: + long_key = f"fireworks_ai/accounts/fireworks/models/{short}" + short_key = f"fireworks_ai/{short}" + assert model_cost.get(short_key) == model_cost.get( + long_key + ), f"short-form {short_key} does not match long-form {long_key}" + + for short in _FIREWORKS_ROUTER_SHORT_FORMS: + long_key = f"fireworks_ai/accounts/fireworks/routers/{short}" + short_key = f"fireworks_ai/{short}" + assert model_cost.get(short_key) == model_cost.get( + long_key + ), f"short-form {short_key} does not match long-form {long_key}" + + class TestBedrockBaseModelLabelKeepsTools: """Regression for #29618: a Bedrock deployment whose ``base_model`` is a friendly label must not silently drop ``tools``/``tool_choice`` under ``drop_params``.""" From 5b36de4646451da7cec224ab481d676315d4ffcb Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 12 Sep 2026 14:39:40 -0700 Subject: [PATCH 52/55] fix(guardrails): log mask when a guardrail adds request keys (#40882) * fix(guardrails): log mask when a guardrail adds request keys _inputs_were_modified only compared keys present in the pre-hook baseline, so a guardrail that injected a new key such as tools was logged as allow. Compare over the union of both key sets, and narrow the pre_call return value to the same prompt-bearing keys the baseline holds so passthrough stays allow. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(guardrails): snapshot apply_guardrail inputs before the hook mutates them Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 29 +++++---- .../guardrail_hooks/azure/prompt_shield.py | 1 + .../integrations/test_custom_guardrail.py | 59 +++++++++++++++++++ 3 files changed, 79 insertions(+), 10 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 8a976a966a6..39adea30828 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -5,6 +5,7 @@ import os import secrets from collections.abc import Mapping from datetime import datetime +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, get_args from litellm._logging import verbose_logger @@ -1279,6 +1280,7 @@ class CustomGuardrail(CustomLogger): guardrail_response: Final = self._summarize_guardrail_response( response=response, original_inputs=original_inputs, + event_type=event_type, ) verbose_logger.debug("Guardrail response: %s", response) @@ -1298,6 +1300,7 @@ class CustomGuardrail(CustomLogger): self, response: object, original_inputs: Mapping[str, object] | None, + event_type: GuardrailEventHooks | None, ) -> object: """Reduce a hook's return value to what is safe to log as ``guardrail_response``. @@ -1305,15 +1308,21 @@ class CustomGuardrail(CustomLogger): returns the (possibly modified) request payload. Neither is a provider verdict, and logging them verbatim ships the user's prompt to every logging sink (OTEL spans, Datadog, spend logs), so both collapse to ``"allow"`` / ``"mask"`` by comparing - against ``original_inputs``, a copy taken before the hook ran. A string result is the - hook's own rejection message (the proxy turns it into a 400), not user input, so it is - logged as is. + against ``original_inputs``, a copy taken before the hook ran. A pre_call baseline only + holds the prompt-bearing keys, so the returned request is narrowed to those same keys + before the comparison. A string result is the hook's own rejection message (the proxy + turns it into a 400), not user input, so it is logged as is. """ if response is None: return {} if original_inputs is None or not isinstance(response, Mapping): return response - return "mask" if self._inputs_were_modified(original_inputs, response) else "allow" + compared_response: Final[Mapping[str, object]] = ( + MappingProxyType({key: value for key, value in response.items() if key in _PRE_CALL_CONTENT_KEYS}) + if event_type == GuardrailEventHooks.pre_call + else response + ) + return "mask" if self._inputs_were_modified(original_inputs, compared_response) else "allow" @staticmethod def _is_guardrail_intervention(e: Exception) -> bool: @@ -1355,8 +1364,8 @@ class CustomGuardrail(CustomLogger): raise e def _inputs_were_modified(self, original_inputs: Mapping[str, object], response: Mapping[str, object]) -> bool: - """True when any baseline key's value differs in ``response`` (mask), False otherwise (allow).""" - return any(response.get(key) != value for key, value in original_inputs.items()) + """True when any key of either mapping differs between them (mask), False otherwise (allow).""" + return any(original_inputs.get(key) != response.get(key) for key in original_inputs.keys() | response.keys()) def mask_content_in_string( self, @@ -1476,13 +1485,13 @@ def _original_inputs_for( ) -> dict | None: # mutable-ok: matches _process_response(original_inputs=) signature """Baseline the hook's return value is compared against to decide "allow" vs "mask". - ``apply_guardrail`` masks a fresh ``inputs`` dict, so that dict is the baseline. Pre-call - hooks edit the request in place and return it, so the baseline is a deep copy of the - prompt-bearing keys taken before the hook runs. + Hooks may edit their argument in place and return it, so the baseline is always a deep + copy taken before the hook runs: the whole ``inputs`` dict for ``apply_guardrail``, the + prompt-bearing request keys for pre-call hooks. """ if func_name == "apply_guardrail": inputs: Final = kwargs.get("inputs") - return inputs if isinstance(inputs, dict) else None + return copy.deepcopy(inputs) if isinstance(inputs, dict) else None if event_type != GuardrailEventHooks.pre_call: return None return {key: copy.deepcopy(value) for key, value in request_data.items() if key in _PRE_CALL_CONTENT_KEYS} diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index de9618a44a1..a0724b75ec7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py @@ -341,6 +341,7 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai guardrail_response: Final = self._summarize_guardrail_response( response=response, original_inputs=original_inputs, + event_type=event_type, ) self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=guardrail_response, diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index bb4822eae57..2fd5fa76d8e 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -3016,3 +3016,62 @@ class TestPreCallHookResponseIsNotLoggedVerbatim: ) assert self._logged_response(data) == "mask" + + @pytest.mark.asyncio + async def test_pre_call_hook_adding_tools_logs_mask(self): + class ToolInjectingGuardrail(CustomGuardrail): + @log_guardrail_information + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: object, + data: dict[str, object], + call_type: str, + ) -> dict[str, object]: + return {**data, "tools": [{"type": "function", "function": {"name": "guardrail_injected_tool"}}]} + + data = self._request() + await ToolInjectingGuardrail(guardrail_name="g").async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=None, data=data, call_type="acompletion" + ) + + assert self._logged_response(data) == "mask" + + @pytest.mark.asyncio + async def test_apply_guardrail_adding_tools_logs_mask(self): + class ToolInjectingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + return {**inputs, "tools": [{"type": "function", "function": {"name": "guardrail_injected_tool"}}]} + + data = self._request() + await ToolInjectingGuardrail(guardrail_name="g").apply_guardrail( + inputs={"texts": ["SECRET_PROMPT"]}, request_data=data, input_type="request" + ) + + assert self._logged_response(data) == "mask" + + @pytest.mark.asyncio + async def test_apply_guardrail_masking_inputs_in_place_logs_mask(self): + class InPlaceMaskingGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + inputs["texts"] = [""] + return inputs + + data = self._request() + await InPlaceMaskingGuardrail(guardrail_name="g").apply_guardrail( + inputs={"texts": ["SECRET_PROMPT"]}, request_data=data, input_type="request" + ) + + assert self._logged_response(data) == "mask" From 147eb23aaba9a2588ed5fee236d1953870ea74ef Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Sat, 12 Sep 2026 14:40:01 -0700 Subject: [PATCH 53/55] bump: litellm-enterprise 0.1.66 -> 0.1.67, litellm-proxy-extras 0.4.96 -> 0.4.97 --- enterprise/pyproject.toml | 4 ++-- litellm-proxy-extras/pyproject.toml | 4 ++-- pyproject.toml | 4 ++-- uv.lock | 6 +++--- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/enterprise/pyproject.toml b/enterprise/pyproject.toml index 903c5155a12..c049bf68c46 100644 --- a/enterprise/pyproject.toml +++ b/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.66" +version = "0.1.67" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.66" +version = "0.1.67" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-enterprise==", diff --git a/litellm-proxy-extras/pyproject.toml b/litellm-proxy-extras/pyproject.toml index 7d4c78088f1..f94591872a4 100644 --- a/litellm-proxy-extras/pyproject.toml +++ b/litellm-proxy-extras/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-proxy-extras" -version = "0.4.96" +version = "0.4.97" description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package." readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.4.96" +version = "0.4.97" version_files = [ "pyproject.toml:^version", "../pyproject.toml:litellm-proxy-extras==", diff --git a/pyproject.toml b/pyproject.toml index 31c498ccbbb..62ce4b4fd61 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,8 +67,8 @@ proxy = [ "azure-identity>=1.25.2,<2.0", "azure-storage-blob>=12.28.0,<13.0", "mcp>=1.28.1,<2.0", - "litellm-proxy-extras==0.4.96", - "litellm-enterprise==0.1.66", + "litellm-proxy-extras==0.4.97", + "litellm-enterprise==0.1.67", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", diff --git a/uv.lock b/uv.lock index 9c659be658f..eb4cdef76f1 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-08T03:56:24.358378Z" +exclude-newer = "2026-09-09T21:39:49.468411Z" exclude-newer-span = "P3D" [manifest] @@ -4773,12 +4773,12 @@ proxy-dev = [ [[package]] name = "litellm-enterprise" -version = "0.1.66" +version = "0.1.67" source = { editable = "enterprise" } [[package]] name = "litellm-proxy-extras" -version = "0.4.96" +version = "0.4.97" source = { editable = "litellm-proxy-extras" } [[package]] From 4647cd121593b46c36d3502d36ee7be3e1070358 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 12 Sep 2026 15:13:44 -0700 Subject: [PATCH 54/55] fix(realtime): keep a Muse turn active for turnless partials after speechEnd Muse partials carry no turnId and belong to the most recent speechStart, and the docs say the model may keep post processing a turn after speechEnd until speechComplete. Releasing the active turn on speechEnd made any partial arriving in that window raise and get dropped in ENDPOINTING mode. The turn now stays active until its speechComplete or final transcript. --- litellm/llms/meta/realtime/transformation.py | 1 - .../test_meta_realtime_transformation.py | 19 +++++++++++++++++++ 2 files changed, 19 insertions(+), 1 deletion(-) diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py index 549db60be85..1b8943f0cee 100644 --- a/litellm/llms/meta/realtime/transformation.py +++ b/litellm/llms/meta/realtime/transformation.py @@ -517,7 +517,6 @@ class MuseEventTransformer: def _speech_end(self, message: Mapping[str, JsonValue]) -> _TurnState: turn: Final = self._turn(_required_turn_id(message, "speechEnd")) turn.stopped = True - self._release_active(turn) return turn def _speech_complete(self, message: Mapping[str, JsonValue]) -> _TurnState: diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py index 3a4c0591fb5..a5d7e47fb65 100644 --- a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py +++ b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py @@ -241,6 +241,25 @@ def test_speech_end_then_speech_complete_emits_stopped_then_completed(): assert completed[0]["transcript"] == "done" +def test_turnless_partial_between_speech_end_and_speech_complete_stays_on_that_turn(): + transformer = MuseEventTransformer() + transformer.transform(json.loads(_event("speechStart", turnId="turn-1"))) + transformer.transform(json.loads(_event("transcript", transcript="what is", final=False))) + transformer.transform(json.loads(_event("speechEnd", turnId="turn-1"))) + + post_processed = transformer.transform( + json.loads(_event("transcript", transcript="what is the weather", final=False)) + ) + completed = transformer.transform( + json.loads(_event("speechComplete", turnId="turn-1", transcript="What is the weather?")) + ) + + assert _typed(post_processed) == [("conversation.item.input_audio_transcription.delta", "turn-1")] + assert post_processed[0]["delta"] == " the weather" + assert _typed(completed) == [("conversation.item.input_audio_transcription.completed", "turn-1")] + assert completed[0]["transcript"] == "What is the weather?" + + def _typed(events: tuple[dict[str, object], ...]) -> list[tuple[object, object]]: return [(event["type"], event["item_id"]) for event in events] From 0e435e41486e860a5c8a998dcfd54ce3e23e84ab Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 12 Sep 2026 15:32:45 -0700 Subject: [PATCH 55/55] fix(realtime): run transcription guardrails on transcription-only sessions The provider_config path skipped run_realtime_guardrails for transcription sessions to avoid sending response.create, which also dropped every realtime_input_transcription guardrail: no violation error reached the client and on_violation / end_session_after_n_fails never fired. Run the guardrail for every completed transcript and only suppress response.create when the session has no assistant turn. --- .../litellm_core_utils/realtime_streaming.py | 4 +- .../test_realtime_streaming.py | 61 +++++++++++++++++++ 2 files changed, 62 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index 06d9241b826..e3f8786a39a 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1015,13 +1015,11 @@ class RealTimeStreaming: self.store_message(event_str) self._capture_transcription_usage(event) await self._send_event_to_client(event, event_str) - if self._is_transcription_session: - continue blocked = await self.run_realtime_guardrails( cast(str, transcript), item_id=cast(str | None, event.get("item_id")), ) - if not blocked: + if not blocked and not self._is_transcription_session: await self._send_to_backend(json.dumps({"type": "response.create"})) continue ## LOGGING diff --git a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py index 295110c6bce..2a33d84ec78 100644 --- a/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py +++ b/tests/test_litellm/litellm_core_utils/test_realtime_streaming.py @@ -3446,6 +3446,67 @@ async def test_transformed_transcription_completion_never_sends_response_create( backend_ws.send.assert_not_awaited() +@pytest.mark.asyncio +async def test_transcription_session_still_runs_transcription_guardrail(monkeypatch: pytest.MonkeyPatch): + class BlockingGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + raise ValueError("blocked transcript") + + guardrail: Final = BlockingGuardrail( + guardrail_name="transcription-blocker", + event_hook=GuardrailEventHooks.realtime_input_transcription, + default_on=True, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + completed_event: Final = { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": "event_1", + "item_id": "turn_1", + "content_index": 0, + "transcript": "blocked transcript", + "usage": {"type": "duration", "seconds": 0.5}, + } + provider_config: Final = MagicMock() + provider_config.requires_session_configuration.return_value = True + provider_config.transform_realtime_response.return_value = { + "response": completed_event, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_conversation_id": None, + "current_item_chunks": None, + "current_delta_type": None, + "session_configuration_request": None, + } + provider_config.transform_realtime_request.return_value = () + provider_config.is_setup_message.return_value = False + provider_config.is_content_message.return_value = False + client_ws: Final = MagicMock() + client_ws.send_text = AsyncMock() + backend_ws: Final = MagicMock() + backend_ws.send = AsyncMock() + + streaming: Final = RealTimeStreaming( + client_ws, + backend_ws, + MagicMock(), + provider_config=provider_config, + model="muse-voice-transcribe-1.0", + force_transcription_model="muse-voice-transcribe-1.0", + ) + + await streaming._handle_provider_config_message("{}") + + sent_to_client: Final = [json.loads(call.args[0]) for call in client_ws.send_text.await_args_list] + assert completed_event in sent_to_client + error_events: Final = [event for event in sent_to_client if event.get("type") == "error"] + assert len(error_events) == 1 + assert error_events[0]["error"]["type"] == "guardrail_violation" + backend_ws.send.assert_not_awaited() + assert streaming._violation_count == 1 + + @pytest.mark.asyncio async def test_provider_bytes_are_sent_raw_after_pacing(): from typing import Final