feat(embeddings): add native dispatch foundation (#42799)

* feat(embeddings): add native dispatch foundation

* ci: cover embeddings dispatch tests

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

* style: format embeddings dispatch

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>

.github/workflows/test-unit.yml

@@ -111,6 +111,7 @@ jobs:
               tests/test_litellm/images
               tests/test_litellm/interactions
               tests/test_litellm/messages
+              tests/test_litellm/embeddings
               tests/test_litellm/ocr
               tests/test_litellm/passthrough
               tests/test_litellm/rag

litellm-rust/crates/python-bridge/src/lib.rs

@@ -30,6 +30,8 @@ mod _native {
         achat_completions, chat_completions, chat_completions_decline,
     };
     #[pymodule_export]
+    use crate::routes::embeddings::{aembedding, embedding};
+    #[pymodule_export]
     use crate::routes::messages::{amessages, messages};
     #[pymodule_export]
     use crate::routes::ocr::{aocr, ocr};
@@ -83,6 +85,8 @@ mod tests {
                 "ProcessReservedForForking",
                 "ocr",
                 "aocr",
+                "embedding",
+                "aembedding",
                 "transcription",
                 "atranscription",
                 "messages",

litellm-rust/crates/python-bridge/src/routes/embeddings.rs

@@ -0,0 +1,54 @@
+use pyo3::{
+    prelude::*,
+    types::{PyDict, PyTuple},
+};
+
+use crate::errors::RustBridgeDeclined;
+
+#[pyfunction]
+pub(crate) fn embedding(
+    _request: Bound<'_, PyAny>,
+    _args: Bound<'_, PyTuple>,
+    _kwargs: Bound<'_, PyDict>,
+) -> PyResult<Py<PyAny>> {
+    Err(RustBridgeDeclined::new_err(
+        "native embeddings route is not implemented",
+    ))
+}
+
+#[pyfunction]
+pub(crate) fn aembedding(
+    _request: Bound<'_, PyAny>,
+    _args: Bound<'_, PyTuple>,
+    _kwargs: Bound<'_, PyDict>,
+) -> PyResult<Py<PyAny>> {
+    Err(RustBridgeDeclined::new_err(
+        "native embeddings route is not implemented",
+    ))
+}
+
+#[cfg(test)]
+mod tests {
+    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::embedding, super::aembedding] {
+                let error = entrypoint(request.clone().into_any(), args.clone(), kwargs.clone())
+                    .expect_err("native embeddings must decline until a route machine exists");
+                assert!(error.is_instance_of::<RustBridgeDeclined>(py));
+            }
+        });
+    }
+}

litellm-rust/crates/python-bridge/src/routes/mod.rs

@@ -1,5 +1,6 @@
 pub(crate) mod audio_transcription;
 pub(crate) mod chat_completions;
+pub(crate) mod embeddings;
 pub(crate) mod messages;
 pub(crate) mod ocr;
 pub(crate) mod responses;

litellm/__init__.py

@@ -1458,6 +1458,7 @@ from .skills.main import (
 from .containers.main import *
 from .ocr.dispatch import *
 from .chat_completions.dispatch import *
+from .embeddings.dispatch import *
 from .rust_bridge import rust
 from .rag.main import *
 from .sandbox.main import *

litellm/embeddings/dispatch.py

@@ -0,0 +1,95 @@
+import inspect
+from collections.abc import Awaitable, Callable, Coroutine, Mapping
+from types import MappingProxyType
+from typing import Final, TypeAlias, cast  # noqa: TID251  # native binding selects a sync result or an async awaitable
+
+from litellm import main
+from litellm.rust_bridge.catalog import Route, RouteContext
+from litellm.rust_bridge.dispatch import PublicDispatch, call_hook
+from litellm.rust_bridge.embeddings.entrypoints import (
+    NATIVE_AEMBEDDING,
+    NATIVE_EMBEDDING,
+    LiteLLMEmbeddingRequest,
+)
+from litellm.rust_bridge.public_call import bind, optional_mapping, optional_str, signature
+from litellm.types.utils import EmbeddingResponse
+
+__all__ = ("aembedding", "embedding")
+
+PythonEmbedding: TypeAlias = Callable[..., EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]]
+PythonAembedding: TypeAlias = Callable[..., Awaitable[EmbeddingResponse]]
+
+_PYTHON_EMBEDDING: Final = cast(  # cast-ok: [LIT006] preserve the legacy public callable contract
+    PythonEmbedding, main.embedding
+)
+_PYTHON_AEMBEDDING: Final = cast(  # cast-ok: [LIT006] preserve the legacy public callable contract
+    PythonAembedding, main.aembedding
+)
+_EMBEDDING_SIGNATURE: Final = signature(_PYTHON_EMBEDDING)
+
+
+def _public_request(
+    legacy: inspect.Signature, args: tuple[object, ...], kwargs: Mapping[str, object]
+) -> LiteLLMEmbeddingRequest | None:
+    fields: Final = bind(legacy, args, kwargs)
+    if fields is None:
+        return None
+    model: Final = fields.get("model")
+    if not isinstance(model, str):
+        return None
+    extra: Final = optional_mapping(fields.get("kwargs")) or MappingProxyType({})
+    return LiteLLMEmbeddingRequest(
+        model=model,
+        input=fields.get("input"),
+        api_key=optional_str(fields.get("api_key")),
+        api_base=optional_str(fields.get("api_base")),
+        custom_llm_provider=optional_str(fields.get("custom_llm_provider")),
+        kwargs=extra,
+    )
+
+
+def _context(request: LiteLLMEmbeddingRequest) -> RouteContext:
+    return RouteContext(Route.EMBEDDINGS, provider=request.custom_llm_provider, model=request.model)
+
+
+_DISPATCH: Final = PublicDispatch(
+    route=Route.EMBEDDINGS,
+    request=lambda args, kwargs: _public_request(_EMBEDDING_SIGNATURE, args, kwargs),
+    context=_context,
+    bypass=lambda request: request.kwargs.get("aembedding") is True,
+)
+
+_ADISPATCH: Final = PublicDispatch(
+    route=Route.EMBEDDINGS,
+    request=lambda args, kwargs: _public_request(_EMBEDDING_SIGNATURE, args, kwargs),
+    context=_context,
+)
+
+
+def embedding(
+    *args: object,
+    **kwargs: object,  # kwargs-ok: preserve the public embedding call shape
+) -> EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]:
+    return _DISPATCH.run(
+        args,
+        kwargs,
+        python=_PYTHON_EMBEDDING,
+        binding=NATIVE_EMBEDDING,
+        native=call_hook,
+    )
+
+
+async def aembedding(*args: object, **kwargs: object) -> EmbeddingResponse:  # kwargs-ok: preserve the public call shape
+    return await _ADISPATCH.arun(
+        args,
+        kwargs,
+        python=_PYTHON_AEMBEDDING,
+        binding=NATIVE_AEMBEDDING,
+        native=call_hook,
+    )
+
+
+embedding.__doc__ = _PYTHON_EMBEDDING.__doc__
+embedding.__wrapped__ = _PYTHON_EMBEDDING  # pyright: ignore[reportFunctionMemberAccess]  # preserve the legacy signature
+aembedding.__doc__ = _PYTHON_AEMBEDDING.__doc__
+aembedding.__wrapped__ = _PYTHON_AEMBEDDING  # pyright: ignore[reportFunctionMemberAccess]  # preserve the legacy signature

litellm/rust_bridge/_native.pyi

@@ -6,9 +6,11 @@ import httpx
 from pydantic import JsonValue
 
 from litellm.llms.base_llm.ocr.transformation import OCRResponse
+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.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
+from litellm.types.utils import EmbeddingResponse
 
 class RustBridgeDeclined(Exception): ...
 class RustUpstreamError(Exception): ...
@@ -41,6 +43,16 @@ def aocr(
     args: tuple[object, ...],
     kwargs: dict[str, object],
 ) -> Coroutine[object, object, OCRResponse]: ...
+def embedding(
+    request: LiteLLMEmbeddingRequest,
+    args: tuple[object, ...],
+    kwargs: Mapping[str, object],
+) -> EmbeddingResponse: ...
+def aembedding(
+    request: LiteLLMEmbeddingRequest,
+    args: tuple[object, ...],
+    kwargs: Mapping[str, object],
+) -> Coroutine[object, object, EmbeddingResponse]: ...
 def transcription(
     model: str,
     audio: object,

litellm/rust_bridge/catalog.py

@@ -18,6 +18,7 @@ from litellm.types.secret_managers.main import KeyManagementSystem
 
 class Route(str, Enum):
     CHAT_COMPLETIONS = "chat_completions"
+    EMBEDDINGS = "embeddings"
     MESSAGES = "messages"
     RESPONSES = "responses"
     TRANSCRIPTION = "transcription"
@@ -105,6 +106,7 @@ Rules: TypeAlias = tuple[Rule, ...]
 
 RULES: Final[Rules] = (
     LoggerRule(Rollout.RUST_OPT_IN),
+    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),

litellm/rust_bridge/embeddings/entrypoints.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Mapping
+from dataclasses import dataclass
+from typing import Final, Protocol, cast  # noqa: TID251  # validates dynamically loaded native callables
+
+from litellm.rust_bridge.bindings import NativeBinding
+from litellm.types.utils import EmbeddingResponse
+
+
+@dataclass(frozen=True, slots=True)
+class LiteLLMEmbeddingRequest:
+    model: str
+    input: object
+    api_key: str | None
+    api_base: str | None
+    custom_llm_provider: str | None
+    kwargs: Mapping[str, object]
+
+
+class NativeEmbedding(Protocol):
+    def __call__(
+        self,
+        request: LiteLLMEmbeddingRequest,
+        args: tuple[object, ...],
+        kwargs: Mapping[str, object],
+    ) -> EmbeddingResponse: ...
+
+
+class NativeAembedding(Protocol):
+    def __call__(
+        self,
+        request: LiteLLMEmbeddingRequest,
+        args: tuple[object, ...],
+        kwargs: Mapping[str, object],
+    ) -> Awaitable[EmbeddingResponse]: ...
+
+
+def _embedding_binding(value: object) -> NativeEmbedding | None:
+    if not callable(value):
+        return None
+    return cast("NativeEmbedding", value)  # cast-ok: callable validated at the native binding boundary
+
+
+def _aembedding_binding(value: object) -> NativeAembedding | None:
+    if not callable(value):
+        return None
+    return cast("NativeAembedding", value)  # cast-ok: callable validated at the native binding boundary
+
+
+NATIVE_EMBEDDING: Final = NativeBinding("embedding", validate=_embedding_binding)
+NATIVE_AEMBEDDING: Final = NativeBinding("aembedding", validate=_aembedding_binding)

tests/test_litellm/embeddings/test_dispatch.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, Mapping
+from typing import Final
+
+import pytest
+from pydantic import TypeAdapter
+
+import litellm
+from litellm.embeddings 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.embeddings.entrypoints import LiteLLMEmbeddingRequest
+from litellm.types.utils import EmbeddingResponse
+
+
+@pytest.mark.asyncio
+async def test_public_embedding_calls_keep_the_python_result() -> None:
+    vector: Final = [0.1, 0.2]
+
+    sync_response: Final = litellm.embedding(model="openai/test-model", input="hello", mock_response=vector)
+    async_response: Final = await litellm.aembedding(model="openai/test-model", input="hello", mock_response=vector)
+
+    assert isinstance(sync_response, EmbeddingResponse)
+    rows: Final = TypeAdapter(list[dict[str, object]])
+    assert rows.validate_python(sync_response.model_dump()["data"])[0]["embedding"] == vector
+    assert rows.validate_python(async_response.model_dump()["data"])[0]["embedding"] == vector
+
+
+def test_sync_embedding_request_projects_public_arguments() -> None:
+    rules: Final[Rules] = (RouteRule(Route.EMBEDDINGS, Rollout.RUST_REQUIRED),)
+    expected: Final = EmbeddingResponse(model="test-model", data=[])
+
+    def native(
+        request: LiteLLMEmbeddingRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
+    ) -> EmbeddingResponse:
+        assert request.model == "test-model"
+        assert request.input == "hello"
+        assert request.custom_llm_provider == "openai"
+        return expected
+
+    binding: Final[
+        NativeBinding[Callable[[LiteLLMEmbeddingRequest, tuple[object, ...], Mapping[str, object]], EmbeddingResponse]]
+    ] = NativeBinding("embedding", validate=lambda _: None)
+    binding.override(native)
+    response: Final = dispatch._DISPATCH.run(  # pyright: ignore[reportPrivateUsage]  # test an explicit route decision
+        ("test-model", "hello"),
+        {"custom_llm_provider": "openai", "dimensions": 8},
+        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_embedding_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 = EmbeddingResponse(model="test-model", data=[])
+    rules: Final[Rules] = (RouteRule(Route.EMBEDDINGS, Rollout.RUST_OPT_OUT),)
+
+    async def native(
+        request: LiteLLMEmbeddingRequest, args: tuple[object, ...], kwargs: Mapping[str, object]
+    ) -> EmbeddingResponse:
+        raise declined("unsupported")
+
+    async def python(*args: object, **kwargs: object) -> EmbeddingResponse:
+        return expected
+
+    binding: Final[
+        NativeBinding[
+            Callable[[LiteLLMEmbeddingRequest, tuple[object, ...], Mapping[str, object]], Awaitable[EmbeddingResponse]]
+        ]
+    ] = NativeBinding("aembedding", validate=lambda _: None)
+    binding.override(native)
+    response: Final = await dispatch._ADISPATCH.arun(  # pyright: ignore[reportPrivateUsage]  # test an explicit route decision
+        ("test-model", "hello"),
+        {},
+        python=python,
+        binding=binding,
+        native=lambda hook, request, args, kwargs: hook(request, args, kwargs),
+        rules=rules,
+    )
+
+    assert response is expected