refactor(rust_bridge): pass dispatch context functions directly

litellm/chat_completions/dispatch.py

@@ -70,17 +70,26 @@ def _public_request(
     )
 
 
+def _context(request: LiteLLMChatCompletionsRequest) -> Context:
+    return Context(
+        Route.CHAT_COMPLETIONS,
+        provider=request.custom_llm_provider,
+        model=request.model,
+        delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
+    )
+
+
 _DISPATCH: Final = PublicDispatch(
     route=Route.CHAT_COMPLETIONS,
     request=lambda args, kwargs: _public_request(_COMPLETION, args, kwargs),
-    context=lambda request: _context(request),
+    context=_context,
     bypass=lambda request: request.kwargs.get("acompletion") is True,
 )
 
 _ADISPATCH: Final = PublicDispatch(
     route=Route.CHAT_COMPLETIONS,
     request=lambda args, kwargs: _public_request(_ACOMPLETION, args, kwargs),
-    context=lambda request: _context(request),
+    context=_context,
 )
 
 
@@ -109,15 +118,6 @@ async def acompletion(*args: object, **kwargs: object) -> ChatResult:  # kwargs-
     )
 
 
-def _context(request: LiteLLMChatCompletionsRequest) -> Context:
-    return Context(
-        Route.CHAT_COMPLETIONS,
-        provider=request.custom_llm_provider,
-        model=request.model,
-        delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
-    )
-
-
 completion.__doc__ = _PYTHON_COMPLETION.__doc__
 completion.__wrapped__ = _PYTHON_COMPLETION  # pyright: ignore[reportFunctionMemberAccess]  # inspect.signature follows __wrapped__ to the legacy signature
 acompletion.__doc__ = _PYTHON_ACOMPLETION.__doc__

litellm/messages/dispatch.py

@@ -69,17 +69,26 @@ def _public_request(
     )
 
 
+def _context(request: LiteLLMMessagesRequest) -> Context:
+    return Context(
+        Route.MESSAGES,
+        provider=request.custom_llm_provider,
+        model=request.model,
+        delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
+    )
+
+
 _DISPATCH: Final = PublicDispatch(
     route=Route.MESSAGES,
     request=lambda args, kwargs: _public_request(_MESSAGES, args, kwargs),
-    context=lambda request: _context(request),
+    context=_context,
     bypass=lambda request: request.kwargs.get("is_async") is True,
 )
 
 _ADISPATCH: Final = PublicDispatch(
     route=Route.MESSAGES,
     request=lambda args, kwargs: _public_request(_AMESSAGES, args, kwargs),
-    context=lambda request: _context(request),
+    context=_context,
 )
 
 
@@ -108,15 +117,6 @@ async def anthropic_messages(*args: object, **kwargs: object) -> MessagesResult:
     )
 
 
-def _context(request: LiteLLMMessagesRequest) -> Context:
-    return Context(
-        Route.MESSAGES,
-        provider=request.custom_llm_provider,
-        model=request.model,
-        delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
-    )
-
-
 anthropic_messages_handler.__doc__ = _PYTHON_MESSAGES.__doc__
 anthropic_messages_handler.__wrapped__ = _PYTHON_MESSAGES  # pyright: ignore[reportFunctionMemberAccess]  # inspect.signature follows __wrapped__ to the legacy signature
 anthropic_messages.__doc__ = _PYTHON_AMESSAGES.__doc__

litellm/ocr/dispatch.py

@@ -49,17 +49,22 @@ _PYTHON_AOCR: Final = cast(  # cast-ok: forward the original call shape through
     Callable[..., Awaitable[OCRResponse]], main.aocr
 )
 
+
+def _context(request: LiteLLMOcrRequest) -> Context:
+    return Context(Route.OCR, provider=request.custom_llm_provider, model=request.model)
+
+
 _DISPATCH: Final = PublicDispatch(
     route=Route.OCR,
     request=lambda args, kwargs: _public_request("ocr", args, kwargs),
-    context=lambda request: _context(request),
+    context=_context,
     bypass=lambda request: request.kwargs.get("aocr") is True,
 )
 
 _ADISPATCH: Final = PublicDispatch(
     route=Route.OCR,
     request=lambda args, kwargs: _public_request("aocr", args, kwargs),
-    context=lambda request: _context(request),
+    context=_context,
 )
 
 
@@ -84,7 +89,3 @@ async def aocr(*args: object, **kwargs: object) -> OCRResponse:  # kwargs-ok: pr
         binding=NATIVE_AOCR,
         native=lambda hook, request, call_args, call_kwargs: hook(request, call_args, call_kwargs),
     )
-
-
-def _context(request: LiteLLMOcrRequest) -> Context:
-    return Context(Route.OCR, provider=request.custom_llm_provider, model=request.model)

litellm/responses/dispatch.py

@@ -62,17 +62,26 @@ def _public_request(
     )
 
 
+def _context(request: LiteLLMResponsesRequest) -> Context:
+    return Context(
+        Route.RESPONSES,
+        provider=request.custom_llm_provider,
+        model=request.model,
+        delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
+    )
+
+
 _DISPATCH: Final = PublicDispatch(
     route=Route.RESPONSES,
     request=lambda args, kwargs: _public_request(_RESPONSES, args, kwargs),
-    context=lambda request: _context(request),
+    context=_context,
     bypass=lambda request: request.kwargs.get("aresponses") is True,
 )
 
 _ADISPATCH: Final = PublicDispatch(
     route=Route.RESPONSES,
     request=lambda args, kwargs: _public_request(_ARESPONSES, args, kwargs),
-    context=lambda request: _context(request),
+    context=_context,
 )
 
 
@@ -101,15 +110,6 @@ async def aresponses(*args: object, **kwargs: object) -> ResponsesResult:  # kwa
     )
 
 
-def _context(request: LiteLLMResponsesRequest) -> Context:
-    return Context(
-        Route.RESPONSES,
-        provider=request.custom_llm_provider,
-        model=request.model,
-        delivery=Delivery.STREAMING if request.stream else Delivery.COMPLETED,
-    )
-
-
 responses.__doc__ = _PYTHON_RESPONSES.__doc__
 responses.__wrapped__ = _PYTHON_RESPONSES  # pyright: ignore[reportFunctionMemberAccess]  # inspect.signature follows __wrapped__ to the legacy signature
 aresponses.__doc__ = _PYTHON_ARESPONSES.__doc__