fix(speech): run the speech bridge's inner completion as an internal sub call

The bridge's inner completion shares the speech call's logging object, and its own success handler ran in a worker thread. That thread could rewrite the speech call's logging payload while the async speech handler was building it, so a proxy spend row sometimes got the provider's response id instead of the call id, and sync loggers got a chat response for a speech request

This is a hand-port of the bridge hunk of #44551, which runs that call as an internal sub call so only the speech call logs. internal_sub_call is copied from main's litellm/_internal_context.py
This commit is contained in:
mateo-berri 2026-11-10 10:13:02 -08:00
parent d94be668d2
commit c5fbd1c3c3
3 changed files with 78 additions and 3 deletions

View file

@ -39,6 +39,16 @@ _service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", de
_resent_on: Final[ContextVar[Callable[[Exception], bool] | None]] = ContextVar("resent_on", default=None)
@contextmanager
def internal_sub_call() -> Generator[None]:
"""A nested call the parent call logs and bills, so the parent must fold this call's cost into its own."""
token: Final = is_internal_call.set(True)
try:
yield
finally:
is_internal_call.reset(token)
@contextmanager
def post_response_phase() -> Generator[None]:
"""Work the caller no longer waits for (success callbacks, response-cache writes), including tasks it spawns."""

View file

@ -8,6 +8,8 @@ from typing import TYPE_CHECKING, Final
from typing_extensions import TypedDict
from litellm._internal_context import internal_sub_call
if TYPE_CHECKING:
import httpx
@ -136,9 +138,10 @@ class SpeechToCompletionBridgeHandler:
voice=voice,
)
result: Final = completion(
**request_data,
)
with internal_sub_call():
result: Final = completion(
**request_data,
)
requested_response_format: Final = optional_params.get("response_format")
if isinstance(result, ModelResponse):

View file

@ -4467,6 +4467,68 @@ async def test_aspeech_gemini_bridge_keeps_proxy_metadata_for_spend_tracking(
assert speech_event.logged_response_cost == pytest.approx(expected_cost)
@dataclass(frozen=True, slots=True)
class _RecordedLoggedCall:
hook: str
response: object
logged_id: object
litellm_call_id: object
def _record_logged_call(hook: str, kwargs: dict[str, object], response_obj: object) -> _RecordedLoggedCall:
logging_payload: Final = kwargs.get("standard_logging_object")
return _RecordedLoggedCall(
hook=hook,
response=response_obj,
logged_id=logging_payload.get("id") if isinstance(logging_payload, dict) else None,
litellm_call_id=kwargs.get("litellm_call_id"),
)
class _LoggedCallRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()
self.events: list[_RecordedLoggedCall] = [] # mutable-ok: test recorder of success-callback events
def log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self.events.append(_record_logged_call("sync", kwargs, response_obj))
async def async_log_success_event(
self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object
) -> None:
self.events.append(_record_logged_call("async", kwargs, response_obj))
@pytest.mark.asyncio
async def test_aspeech_gemini_bridge_logs_only_the_speech_call_under_its_call_id(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.delenv("GEMINI_API_KEY", raising=False)
monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
recorder: Final = _LoggedCallRecorder()
monkeypatch.setattr(litellm, "callbacks", [recorder])
respx_mock.post(
url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash-preview-tts:generateContent.*"
).mock(return_value=httpx.Response(200, json={**_gemini_tts_generate_content_response(), "responseId": "gemini-id"}))
await litellm.aspeech(
model="gemini/gemini-2.5-flash-preview-tts", input="logged once", voice="Kore", api_key="fake-gemini-key"
)
for _ in range(100):
if any(event.hook == "async" for event in recorder.events):
break
await asyncio.sleep(0.05)
await asyncio.sleep(0.2)
assert [(event.hook, type(event.response)) for event in recorder.events] == [
("async", HttpxBinaryResponseContent)
]
assert recorder.events[0].logged_id == recorder.events[0].litellm_call_id
def _stream_builder_text_chunk(model: str, content: str, finish_reason: str | None = None) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-cost",