mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
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:
parent
28755b98a0
commit
170eb7fb95
9 changed files with 429 additions and 13 deletions
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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() {
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
))
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -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"})),
|
||||
|
|
|
|||
117
tests/test_litellm/chat_completions/test_dispatch.py
Normal file
117
tests/test_litellm/chat_completions/test_dispatch.py
Normal 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
|
||||
155
tests/test_litellm/messages/test_dispatch.py
Normal file
155
tests/test_litellm/messages/test_dispatch.py
Normal 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
|
||||
|
|
@ -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),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue