diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index ed87049f2d5..4580ad17a19 100644 --- a/.github/workflows/test-unit.yml +++ b/.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 diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 874e522af15..409e7a5dbfb 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/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", diff --git a/litellm-rust/crates/python-bridge/src/routes/embeddings.rs b/litellm-rust/crates/python-bridge/src/routes/embeddings.rs new file mode 100644 index 00000000000..7c55613ced1 --- /dev/null +++ b/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> { + 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> { + 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::(py)); + } + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index 8a78a26423d..dd694fa589f 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/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; diff --git a/litellm/__init__.py b/litellm/__init__.py index c8df4394a06..376c5fa9010 100644 --- a/litellm/__init__.py +++ b/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 * diff --git a/litellm/embeddings/__init__.py b/litellm/embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/embeddings/dispatch.py b/litellm/embeddings/dispatch.py new file mode 100644 index 00000000000..bba68d2c0f1 --- /dev/null +++ b/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 diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 0e684d1f10c..b6dd1900150 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/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, diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 1a2d153871d..4b49f8c7cd9 100644 --- a/litellm/rust_bridge/catalog.py +++ b/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), diff --git a/litellm/rust_bridge/embeddings/__init__.py b/litellm/rust_bridge/embeddings/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/rust_bridge/embeddings/entrypoints.py b/litellm/rust_bridge/embeddings/entrypoints.py new file mode 100644 index 00000000000..da17434df02 --- /dev/null +++ b/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) diff --git a/tests/test_litellm/embeddings/test_dispatch.py b/tests/test_litellm/embeddings/test_dispatch.py new file mode 100644 index 00000000000..1062c320cbb --- /dev/null +++ b/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