feat(rust-bridge): extend native dispatch foundation to chat completions, responses, and messages (#42805)

* feat(rust-bridge): declare native chat completions and responses bindings

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(dispatch): cover chat completions and messages dispatch

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(rust-bridge): keep secret manager stub formatting unchanged

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(rust-bridge): match stub parameter names and exports to the native surface

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(rust-bridge): name declining entrypoint parameters and export embeddings in the stub

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(rust-bridge): cover embeddings bindings in the route matrix

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(rust-bridge): keep secret manager stub formatting unchanged

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Yujong Lee <yujong@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-23 21:47:27 +00:00 • committed by GitHub
parent 28755b98a0
commit 170eb7fb95
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 429 additions and 13 deletions

View file

@ -27,7 +27,7 @@ mod _native {
use crate::routes::audio_transcription::{atranscription, transcription};
#[pymodule_export]
use crate::routes::chat_completions::{
achat_completions, chat_completions, chat_completions_decline,
achat_completions, acompletion, chat_completions, chat_completions_decline, completion,
};
#[pymodule_export]
use crate::routes::embeddings::{aembedding, embedding};
@ -36,7 +36,7 @@ mod _native {
#[pymodule_export]
use crate::routes::ocr::{aocr, ocr};
#[pymodule_export]
use crate::routes::responses::ResponsesWebSocketConnection;
use crate::routes::responses::{ResponsesWebSocketConnection, aresponses, responses};
#[pymodule_export]
use crate::routes::token_counter::TokenCounter;
#[cfg(feature = "huggingface")]
@ -94,6 +94,10 @@ mod tests {
"chat_completions_decline",
"chat_completions",
"achat_completions",
"completion",
"acompletion",
"responses",
"aresponses",
"ResponsesWebSocketConnection",
"NativeDiagnosticProcessor",
"TokenCounter",

View file

@ -1,3 +1,6 @@
use pyo3::types::{PyDict, PyTuple};
use crate::errors::RustBridgeDeclined;
use crate::logger::{run_async, run_sync};
use litellm_core::chat_completions::{
Error, chat_completions as run_chat_completions, chat_completions_decline_reason,
@ -123,9 +126,58 @@ pub(crate) fn achat_completions<'py>(
)
}
#[pyfunction]
#[pyo3(signature = (request, args, kwargs))]
pub(crate) fn completion(
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
drop((request, args, kwargs));
Err(RustBridgeDeclined::new_err(
"native chat completions route is not implemented",
))
}
#[pyfunction]
#[pyo3(signature = (request, args, kwargs))]
pub(crate) fn acompletion(
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
drop((request, args, kwargs));
Err(RustBridgeDeclined::new_err(
"native chat completions route is not implemented",
))
}
#[cfg(test)]
mod tests {
use pyo3::{prelude::*, types::PyList};
use pyo3::{
prelude::*,
types::{PyDict, PyList, PyTuple},
};
use crate::errors::RustBridgeDeclined;
#[test]
fn both_entrypoints_decline_before_provider_execution() {
Python::initialize();
Python::attach(|py| {
let request = PyDict::new(py);
let args = PyTuple::empty(py);
let kwargs = PyDict::new(py);
for entrypoint in [super::completion, super::acompletion] {
let error = entrypoint(request.clone().into_any(), args.clone(), kwargs.clone())
.expect_err(
"native chat completions must decline until a route machine exists",
);
assert!(error.is_instance_of::<RustBridgeDeclined>(py));
}
});
}
#[test]
fn chat_completions_decline_keeps_existing_reasons() {

View file

@ -6,22 +6,26 @@ use pyo3::{
use crate::errors::RustBridgeDeclined;
#[pyfunction]
#[pyo3(signature = (request, args, kwargs))]
pub(crate) fn embedding(
_request: Bound<'_, PyAny>,
_args: Bound<'_, PyTuple>,
_kwargs: Bound<'_, PyDict>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
drop((request, args, kwargs));
Err(RustBridgeDeclined::new_err(
"native embeddings route is not implemented",
))
}
#[pyfunction]
#[pyo3(signature = (request, args, kwargs))]
pub(crate) fn aembedding(
_request: Bound<'_, PyAny>,
_args: Bound<'_, PyTuple>,
_kwargs: Bound<'_, PyDict>,
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
drop((request, args, kwargs));
Err(RustBridgeDeclined::new_err(
"native embeddings route is not implemented",
))

View file

@ -1,12 +1,41 @@
use litellm_core::responses::websocket::ResponsesWebSocketConnection as RustResponsesWebSocketConnection;
use pyo3::prelude::*;
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use serde_json::Value;
use crate::{
errors::responses_error_to_pyerr,
errors::{RustBridgeDeclined, responses_error_to_pyerr},
marshal::{marshal_headers, optional_timeout},
};
#[pyfunction]
#[pyo3(signature = (request, args, kwargs))]
pub(crate) fn responses(
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
drop((request, args, kwargs));
Err(RustBridgeDeclined::new_err(
"native responses route is not implemented",
))
}
#[pyfunction]
#[pyo3(signature = (request, args, kwargs))]
pub(crate) fn aresponses(
request: Bound<'_, PyAny>,
args: Bound<'_, PyTuple>,
kwargs: Bound<'_, PyDict>,
) -> PyResult<Py<PyAny>> {
drop((request, args, kwargs));
Err(RustBridgeDeclined::new_err(
"native responses route is not implemented",
))
}
#[pyclass]
pub(crate) struct ResponsesWebSocketConnection {
inner: RustResponsesWebSocketConnection,
@ -63,7 +92,28 @@ mod tests {
use std::{ffi::CString, time::Duration};
use futures_util::{SinkExt, StreamExt};
use pyo3::{prelude::*, types::PyDict};
use pyo3::{
prelude::*,
types::{PyDict, PyTuple},
};
use crate::errors::RustBridgeDeclined;
#[test]
fn both_entrypoints_decline_before_provider_execution() {
Python::initialize();
Python::attach(|py| {
let request = PyDict::new(py);
let args = PyTuple::empty(py);
let kwargs = PyDict::new(py);
for entrypoint in [super::responses, super::aresponses] {
let error = entrypoint(request.clone().into_any(), args.clone(), kwargs.clone())
.expect_err("native responses must decline until a route machine exists");
assert!(error.is_instance_of::<RustBridgeDeclined>(py));
}
});
}
use tokio::net::TcpListener;
use tokio_tungstenite::{accept_async, tungstenite::Message};

View file

@ -6,11 +6,14 @@ import httpx
from pydantic import JsonValue
from litellm.llms.base_llm.ocr.transformation import OCRResponse
from litellm.rust_bridge.chat_completions.entrypoints import LiteLLMChatCompletionsRequest
from litellm.rust_bridge.embeddings.entrypoints import LiteLLMEmbeddingRequest
from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest
from litellm.rust_bridge.ocr.entrypoints import LiteLLMOcrRequest
from litellm.rust_bridge.responses.entrypoints import LiteLLMResponsesRequest
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
from litellm.types.utils import EmbeddingResponse
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import EmbeddingResponse, ModelResponse
class RustBridgeDeclined(Exception): ...
class RustUpstreamError(Exception): ...
@ -73,6 +76,26 @@ def atranscription(
optional_params: Mapping[str, object] | None = None,
timeout_seconds: float | None = None,
) -> Future[dict[str, object]]: ...
def completion(
request: LiteLLMChatCompletionsRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ModelResponse: ...
def acompletion(
request: LiteLLMChatCompletionsRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> Coroutine[object, object, ModelResponse]: ...
def responses(
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> ResponsesAPIResponse: ...
def aresponses(
request: LiteLLMResponsesRequest,
args: tuple[object, ...],
kwargs: Mapping[str, object],
) -> Coroutine[object, object, ResponsesAPIResponse]: ...
def messages(
request: LiteLLMMessagesRequest,
args: tuple[object, ...],
@ -381,16 +404,22 @@ __all__ = [
"TokenCounter",
"Tokenizer",
"achat_completions",
"acompletion",
"aembedding",
"amessages",
"aocr",
"aresponses",
"atranscription",
"chat_completions",
"chat_completions_decline",
"completion",
"embedding",
"gil_stats",
"messages",
"ocr",
"process_state_started",
"reserve_process_for_forking",
"responses",
"transcription",
]

View file

@ -106,10 +106,12 @@ Rules: TypeAlias = tuple[Rule, ...]
RULES: Final[Rules] = (
LoggerRule(Rollout.RUST_OPT_IN),
RouteRule(Route.CHAT_COMPLETIONS, Rollout.PYTHON_ONLY),
RouteRule(Route.EMBEDDINGS, Rollout.PYTHON_ONLY),
RouteRule(Route.OCR, Rollout.RUST_REQUIRED, providers=frozenset({"aws_textract"})),
RouteRule(Route.OCR, Rollout.RUST_OPT_OUT),
RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY),
RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY),
RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY),
RouteRule(Route.TOKENIZER, Rollout.PYTHON_ONLY),
RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})),

View file

@ -0,0 +1,117 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import Final
import pytest
import litellm
from litellm.chat_completions import dispatch
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Route, RouteRule, Rules
from litellm.rust_bridge.chat_completions.entrypoints import (
LiteLLMChatCompletionsRequest,
NativeAcompletion,
NativeCompletion,
)
from litellm.rust_bridge.configuration import Rollout
from litellm.types.utils import ModelResponse
MESSAGES: Final = [{"role": "user", "content": "hi"}]
@pytest.mark.asyncio
async def test_public_completion_calls_keep_the_python_result() -> None:
sync_response: Final = litellm.completion(model="openai/test-model", messages=MESSAGES, mock_response="ok")
async_response: Final = await litellm.acompletion(model="openai/test-model", messages=MESSAGES, mock_response="ok")
assert isinstance(sync_response, ModelResponse)
assert isinstance(async_response, ModelResponse)
assert sync_response.choices[0].message.content == "ok"
assert async_response.choices[0].message.content == "ok"
def test_sync_completion_request_projects_public_arguments() -> None:
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),)
expected: Final = ModelResponse()
def native(
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> ModelResponse:
assert request.model == "test-model"
assert request.messages == MESSAGES
assert request.custom_llm_provider == "openai"
assert request.stream is True
return expected
binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None)
binding.override(native)
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
("test-model", MESSAGES),
{"custom_llm_provider": "openai", "stream": True},
python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"),
binding=binding,
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
rules=rules,
)
assert response is expected
@pytest.mark.asyncio
async def test_async_completion_falls_back_after_native_declines() -> None:
from litellm.rust_bridge.bindings import native_exception_types
native_types: Final = native_exception_types()
if native_types is None:
pytest.skip("native bridge is unavailable")
declined, _ = native_types
expected: Final = ModelResponse()
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_OPT_OUT),)
async def native(
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> ModelResponse:
raise declined("unsupported")
async def python(*args: object, **kwargs: object) -> ModelResponse:
return expected
binding: Final[NativeBinding[NativeAcompletion]] = NativeBinding("acompletion", validate=lambda _: None)
binding.override(native)
response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
("test-model", MESSAGES),
{},
python=python,
binding=binding,
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
rules=rules,
)
assert response is expected
def test_internal_acompletion_marker_bypasses_native() -> None:
rules: Final[Rules] = (RouteRule(Route.CHAT_COMPLETIONS, Rollout.RUST_REQUIRED),)
expected: Final = ModelResponse()
def python(*args: object, **kwargs: object) -> ModelResponse:
return expected
def native(
request: LiteLLMChatCompletionsRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> ModelResponse:
pytest.fail("acompletion's inner completion call must stay on Python")
binding: Final[NativeBinding[NativeCompletion]] = NativeBinding("completion", validate=lambda _: None)
binding.override(native)
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
("test-model", MESSAGES),
{"custom_llm_provider": "openai", "acompletion": True},
python=python,
binding=binding,
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
rules=rules,
)
assert response is expected

View file

@ -0,0 +1,155 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import Final
import pytest
from pydantic import TypeAdapter
import litellm
from litellm.messages import dispatch
from litellm.rust_bridge.bindings import NativeBinding
from litellm.rust_bridge.catalog import Route, RouteRule, Rules
from litellm.rust_bridge.configuration import Rollout
from litellm.rust_bridge.messages.entrypoints import (
LiteLLMMessagesRequest,
NativeAmessages,
NativeMessages,
)
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
MESSAGES: Final = [{"role": "user", "content": "hi"}]
@pytest.mark.asyncio
async def test_public_anthropic_messages_keeps_the_python_result() -> None:
response: Final = await litellm.anthropic_messages(
model="anthropic/claude-sonnet-4-5", messages=MESSAGES, max_tokens=10, mock_response="ok"
)
assert isinstance(response, dict)
content: Final = TypeAdapter(list[dict[str, object]]).validate_python(response.get("content", []))
assert content[0]["text"] == "ok"
def test_sync_messages_request_projects_public_arguments() -> None:
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
expected: Final = AnthropicMessagesResponse(model="claude-test")
def native(
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> AnthropicMessagesResponse:
assert request.model == "claude-test"
assert request.messages == MESSAGES
assert request.max_tokens == 10
assert request.custom_llm_provider == "anthropic"
return expected
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None)
binding.override(native)
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
(),
{
"model": "claude-test",
"messages": MESSAGES,
"max_tokens": 10,
"custom_llm_provider": "anthropic",
},
python=lambda *args, **kwargs: pytest.fail("required native route must handle this call"),
binding=binding,
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
rules=rules,
)
assert response is expected
def test_messages_binding_error_delegates_unchanged_to_python() -> None:
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
expected: Final = AnthropicMessagesResponse(model="claude-test")
def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse:
return expected
def native(
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> AnthropicMessagesResponse:
pytest.fail("a call without max_tokens cannot project a request and must stay on Python")
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None)
binding.override(native)
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
(),
{"model": "claude-test", "messages": MESSAGES, "custom_llm_provider": "anthropic"},
python=python,
binding=binding,
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
rules=rules,
)
assert response is expected
@pytest.mark.asyncio
async def test_async_messages_falls_back_after_native_declines() -> None:
from litellm.rust_bridge.bindings import native_exception_types
native_types: Final = native_exception_types()
if native_types is None:
pytest.skip("native bridge is unavailable")
declined, _ = native_types
expected: Final = AnthropicMessagesResponse(model="claude-test")
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_OUT),)
async def native(
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> AnthropicMessagesResponse:
raise declined("unsupported")
async def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse:
return expected
binding: Final[NativeBinding[NativeAmessages]] = NativeBinding("amessages", validate=lambda _: None)
binding.override(native)
response: Final = await dispatch._ADISPATCH.arun( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
(),
{"model": "claude-test", "messages": MESSAGES, "max_tokens": 10},
python=python,
binding=binding,
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
rules=rules,
)
assert response is expected
def test_internal_is_async_marker_bypasses_native() -> None:
rules: Final[Rules] = (RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED),)
expected: Final = AnthropicMessagesResponse(model="claude-test")
def python(*args: object, **kwargs: object) -> AnthropicMessagesResponse:
return expected
def native(
request: LiteLLMMessagesRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
) -> AnthropicMessagesResponse:
pytest.fail("anthropic_messages' inner handler call must stay on Python")
binding: Final[NativeBinding[NativeMessages]] = NativeBinding("messages", validate=lambda _: None)
binding.override(native)
response: Final = dispatch._DISPATCH.run( # pyright: ignore[reportPrivateUsage] # test an explicit route decision
(),
{
"model": "claude-test",
"messages": MESSAGES,
"max_tokens": 10,
"custom_llm_provider": "anthropic",
"is_async": True,
},
python=python,
binding=binding,
native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
rules=rules,
)
assert response is expected

View file

@ -5,6 +5,7 @@ import pytest
from litellm.rust_bridge import bindings
from litellm.rust_bridge.chat_completions import entrypoints as chat_completions
from litellm.rust_bridge.embeddings import entrypoints as embeddings
from litellm.rust_bridge.messages import entrypoints as messages
from litellm.rust_bridge.ocr import entrypoints as ocr
from litellm.rust_bridge.responses import entrypoints as responses
@ -43,6 +44,8 @@ def test_binding_validates_native_attribute(
ROUTE_BINDINGS: Final = (
("completion", chat_completions.NATIVE_COMPLETION),
("acompletion", chat_completions.NATIVE_ACOMPLETION),
("embedding", embeddings.NATIVE_EMBEDDING),
("aembedding", embeddings.NATIVE_AEMBEDDING),
("messages", messages.NATIVE_MESSAGES),
("amessages", messages.NATIVE_AMESSAGES),
("responses", responses.NATIVE_RESPONSES),