From 2e9fa3a0430cdbe97847af2bffe7abe2922160ae Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Mon, 2 Nov 2026 02:52:28 +0000 Subject: [PATCH] fix(router): bind and dispatch every factory endpoint from one table through host.invoke on both cores - one table (litellm/router_endpoint_table.py) names each factory-built endpoint, its host op, call type and dispatch; both cores bind and dispatch from it, and an unknown call type is refused when bound - the endpoints that resolved to None without a call now dispatch like their siblings; sync create/get/delete/cancel_interaction answer synchronously - sync search picks among the router's search tools like asearch - ~80 endpoints bound to litellm functions (and agenerate_content's overwritten binding, avector_store_create) now go through host.invoke - routing decisions key on the endpoint's call type, never the host function's __name__ - arerank no longer sends the builtin input; responses get/cancel/delete/list with a plain id and no model call the handler directly - prompt-management groups work without a litellm_call_id or logging object --- litellm/proxy/auth/fallback_model_access.py | 5 +- litellm/router_endpoint_table.py | 345 ++++++ litellm/router_python_core.py | 1100 +++++------------ litellm/router_rust_core.py | 4 +- litellm/router_rust_endpoints.py | 624 ++++------ litellm/router_rust_prompt_management.py | 12 +- litellm/router_utils/endpoint_eligibility.py | 1 + litellm/router_utils/search_api_router.py | 40 +- .../router_surface_check.py | 21 +- .../messages/test_advisor_orchestration.py | 11 +- .../test_vector_store_endpoints.py | 472 +++---- .../test_router/core/endpoints_helpers.py | 15 +- .../test_router/core/fallbacks_helpers.py | 15 - .../test_endpoints_characterization_batch.py | 29 +- ...test_endpoints_characterization_generic.py | 127 +- ...st_endpoints_characterization_inference.py | 5 +- .../core/test_endpoints_dispatch_table.py | 117 ++ ...test_fallbacks_characterization_refusal.py | 16 +- tests/unit/test_router/test_router.py | 3 +- .../unit/test_router/test_router_endpoints.py | 12 +- 20 files changed, 1351 insertions(+), 1623 deletions(-) create mode 100644 litellm/router_endpoint_table.py create mode 100644 tests/unit/test_router/core/test_endpoints_dispatch_table.py diff --git a/litellm/proxy/auth/fallback_model_access.py b/litellm/proxy/auth/fallback_model_access.py index b550d4d6866..3baed125d06 100644 --- a/litellm/proxy/auth/fallback_model_access.py +++ b/litellm/proxy/auth/fallback_model_access.py @@ -19,6 +19,7 @@ from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.auth_checks import can_key_call_resolved_model, can_token_call_search_tool from litellm.router_api import RegistryReader, RoutingContext +from litellm.router_endpoint_table import EndpointHandler from litellm.search import asearch from litellm.types.llms.base import LiteLLMBaseModel @@ -61,7 +62,9 @@ async def is_search_tool_authorized_for_token(*, search_tool_name: str, valid_to def _is_search_request(request_kwargs: Mapping[str, object]) -> bool: - return request_kwargs.get("original_generic_function") is asearch + handler: Final = request_kwargs.get("original_generic_function") + # The router hands over the endpoint's table handler, which keys on the call type, not litellm.search.asearch. + return handler is asearch or (isinstance(handler, EndpointHandler) and handler.call_type == "asearch") def _token_in_metadata(metadata: object) -> UserAPIKeyAuth | None: diff --git a/litellm/router_endpoint_table.py b/litellm/router_endpoint_table.py new file mode 100644 index 00000000000..6d58a0359ab --- /dev/null +++ b/litellm/router_endpoint_table.py @@ -0,0 +1,345 @@ +""" +The router's factory-built endpoints, in one table both router cores bind from. + +Each row names a public router member, the host op that serves it (a litellm function name, or a +dotted path to one under ``litellm``: ``vector_stores.main.asearch``), the call type the routing +decisions key on (stream-timeout split, video-id pin, route-capability pin, the responses socket, +the /v1/messages refusal gate, context compaction, endpoint eligibility), and how the endpoint +dispatches: sync or async, and which helper routes it. A core binds every row and dispatches on +``dispatch`` alone, so an endpoint cannot be bound without a dispatch or dispatched without a row. + +Neither core is imported here: ``PythonCore`` and ``RustCore`` both read this module. +""" + +import random +from collections.abc import Callable, Iterable, Mapping +from types import MappingProxyType +from typing import Final, Literal, Protocol, get_args + +Dispatch = Literal[ + # async + "assistants", + "moderation", + "search", + "responses", + "anthropic_messages", + "generic", + "generic_client", + "generic_provider_client", + "passthrough", + "response_id", + "container", + "vector_store", + "interaction", + "managed_agent", + # sync + "search_sync", + "generic_sync", + "vector_store_sync", + "vector_store_file_sync", + "interaction_sync", + "managed_agent_sync", +] +DISPATCHES: Final[frozenset[str]] = frozenset(get_args(Dispatch)) +SYNC_DISPATCHES: Final[frozenset[str]] = frozenset( + { + "search_sync", + "generic_sync", + "vector_store_sync", + "vector_store_file_sync", + "interaction_sync", + "managed_agent_sync", + } +) + + +class EndpointSpec: + """One factory-built endpoint: ``name`` on the router, served by the host op ``host_op``.""" + + __slots__ = ("call_type", "dispatch", "host_op", "name") + + def __init__(self, name: str, host_op: str, call_type: str, dispatch: Dispatch) -> None: + self.name: Final = name + self.host_op: Final = host_op + self.call_type: Final = call_type + self.dispatch: Final[Dispatch] = dispatch + + @property + def sync(self) -> bool: + return self.dispatch in SYNC_DISPATCHES + + def __repr__(self) -> str: + return f"EndpointSpec({self.name!r}, {self.host_op!r}, {self.call_type!r}, {self.dispatch!r})" + + +def _rows() -> Iterable[EndpointSpec]: + def row(name: str, host_op: str, dispatch: Dispatch, call_type: str | None = None) -> EndpointSpec: + return EndpointSpec(name, host_op, call_type or name, dispatch) + + for name in ( + "acreate_assistants", + "adelete_assistant", + "aget_assistants", + "acreate_thread", + "aget_thread", + "a_add_message", + "aget_messages", + "arun_thread", + ): + yield row(name, name, "assistants", "assistants") + yield row("amoderation", "amoderation", "moderation", "moderation") + yield row("aanthropic_messages", "anthropic_messages", "anthropic_messages", "anthropic_messages") + yield row("anthropic_messages", "anthropic_messages", "anthropic_messages") + yield row("aresponses", "aresponses", "responses") + yield row("responses", "responses", "generic_sync") + for name in ("aget_responses", "acancel_responses", "acompact_responses", "adelete_responses", "alist_input_items"): + yield row(name, name, "response_id") + for name in ("afile_delete", "afile_content"): + yield row(name, name, "generic_provider_client") + yield row("afile_list", "afile_list", "generic", "alist_files") + for name in ("acreate_realtime_client_secret", "arealtime_calls", "acreate_realtime_transcription_session"): + yield row(name, name, "generic_client") + for name in ( + "_arealtime", + "_aresponses_websocket", + "acreate_fine_tuning_job", + "acancel_fine_tuning_job", + "alist_fine_tuning_jobs", + "aretrieve_fine_tuning_job", + "aimage_edit", + "aadapter_generate_content", + "acreate_skill", + "alist_skills", + "aget_skill", + "adelete_skill", + ): + yield row(name, name, "generic") + yield row("allm_passthrough_route", "allm_passthrough_route", "passthrough") + # vector stores (`avector_store_create` is a method on both cores: it routes to one deployment) + for name in ("search", "retrieve", "list", "update", "delete"): + yield row(f"avector_store_{name}", f"vector_stores.main.a{name}", "vector_store") + for name in ("search", "create", "retrieve", "list", "update", "delete"): + yield row(f"vector_store_{name}", f"vector_stores.main.{name}", "vector_store_sync") + for name, target in ( + ("create", "create"), + ("list", "list"), + ("retrieve", "retrieve"), + ("content", "retrieve_content"), + ("update", "update"), + ("delete", "delete"), + ): + yield row(f"avector_store_file_{name}", f"vector_store_files.main.a{target}", "vector_store") + yield row(f"vector_store_file_{name}", f"vector_store_files.main.{target}", "vector_store_file_sync") + # Google GenAI, OCR, search, decisions, systemone, extraction + for name in ("agenerate_content", "agenerate_content_stream"): + yield row(name, f"google_genai.{name}", "generic") + for name in ("generate_content", "generate_content_stream"): + yield row(name, f"google_genai.{name}", "generic_sync") + for module, name in (("ocr", "ocr"), ("decisions", "decisions"), ("systemone", "systemone")): + yield row(f"a{name}", f"{module}.a{name}", "generic") + yield row(name, f"{module}.{name}", "generic_sync") + yield row("aextraction", "extraction.aextraction", "generic") + yield row("extraction", "extraction.extraction", "generic_sync") + yield row("asearch", "search.asearch", "search") + yield row("search", "search.search", "search_sync") + # videos + for name in ( + "generation", + "list", + "status", + "content", + "remix", + "create_character", + "get_character", + "edit", + "extension", + ): + yield row(f"avideo_{name}", f"videos.avideo_{name}", "generic") + yield row(f"video_{name}", f"videos.video_{name}", "generic_sync") + # containers and their files + for name in ("create_container", "list_containers", "retrieve_container", "delete_container"): + yield row(f"a{name}", f"containers.a{name}", "container") + yield row(name, f"containers.{name}", "generic_sync") + for name in ( + "upload_container_file", + "list_container_files", + "retrieve_container_file", + "delete_container_file", + "retrieve_container_file_content", + ): + yield row(f"a{name}", f"containers.endpoint_factory._generated_endpoints.a{name}", "container") + yield row(name, f"containers.endpoint_factory._generated_endpoints.{name}", "generic_sync") + # Google Interactions and Managed Agents + for name in ("create", "get", "delete", "cancel"): + yield row(f"a{name}_interaction", f"interactions.a{name}", "interaction") + yield row(f"{name}_interaction", f"interactions.{name}", "interaction_sync") + for name, target in ( + ("create_agent", "create"), + ("list_agents", "list"), + ("get_agent", "get"), + ("delete_agent", "delete"), + ("list_agent_versions", "list_versions"), + ): + yield row(f"a{name}", f"interactions.agents.a{target}", "managed_agent") + yield row(name, f"interactions.agents.{target}", "managed_agent_sync") + + +ENDPOINTS: Final[Mapping[str, EndpointSpec]] = MappingProxyType({spec.name: spec for spec in _rows()}) + + +def _dispatch_by_call_type() -> Mapping[str, Dispatch]: + table: Final[dict[str, Dispatch]] = {"avector_store_create": "vector_store"} # mutable-ok: built once here + for spec in ENDPOINTS.values(): + known = table.setdefault(spec.call_type, spec.dispatch) + if known != spec.dispatch: + raise RuntimeError(f"call type {spec.call_type!r} dispatches both as {known!r} and {spec.dispatch!r}") + return MappingProxyType(table) + + +# `factory_function(original_function, call_type=...)`: how each call type dispatches. +DISPATCH_BY_CALL_TYPE: Final[Mapping[str, Dispatch]] = _dispatch_by_call_type() + + +def dispatch_for(call_type: str) -> Dispatch: + """How an endpoint of ``call_type`` dispatches; an unknown call type is refused when bound.""" + dispatch: Final = DISPATCH_BY_CALL_TYPE.get(call_type) + if dispatch is None: + raise ValueError(f"no router endpoint dispatches the call type {call_type!r}") + return dispatch + + +class _Invokes(Protocol): + def invoke(self, op: str) -> Callable[..., object]: ... + + +class EndpointHandler: + """An endpoint's handler: the function ``resolve()`` returns, looked up at each call, tagged with + the endpoint's ``call_type`` so the routing decisions key on the endpoint, never on the + function's ``__name__`` (a wrapped or renamed host function keeps them).""" + + __slots__ = ("call_type", "op", "resolve") + + def __init__(self, call_type: str, resolve: Callable[[], Callable[..., object]], op: str | None = None) -> None: + self.call_type: Final = call_type + self.resolve: Final = resolve + self.op: Final = op + + def __call__(self, *args: object, **kwargs: object) -> object: # kwargs-ok: the endpoint's request + return self.resolve()(*args, **kwargs) + + def __deepcopy__(self, memo: dict[int, object]) -> "EndpointHandler": # mutable-ok: copy.deepcopy's memo + return self + + def __repr__(self) -> str: + return f"EndpointHandler({self.call_type!r}, op={self.op!r})" + + +def host_handler(host: _Invokes, op: str, call_type: str) -> EndpointHandler: + """The host op ``op`` as the handler of an endpoint of ``call_type``, resolved at each call.""" + return EndpointHandler(call_type, lambda: host.invoke(op), op) + + +def as_handler(function: Callable[..., object], call_type: str) -> EndpointHandler: + """``function`` as the handler of an endpoint of ``call_type`` (kept as given when it already is one).""" + if isinstance(function, EndpointHandler): + return function + return EndpointHandler(call_type, lambda: function) + + +def handler_call_type(function: object) -> str | None: + """The call type the routing decisions key on: the endpoint's, else the function's ``__name__`` + (a function a caller hands the router's generic call directly).""" + if isinstance(function, EndpointHandler): + return function.call_type + name: Final[object] = getattr(function, "__name__", None) + return name if isinstance(name, str) else None + + +def handler_identity(function: object) -> tuple[str | None, str | None]: + """``(call type, module)`` for ``endpoint_selection``: an endpoint's call type names it alone.""" + if isinstance(function, EndpointHandler): + return function.call_type, None + module: Final[object] = getattr(function, "__module__", None) + return handler_call_type(function), module if isinstance(module, str) else None + + +def with_provider( + custom_llm_provider: str | None, kwargs: Mapping[str, object] +) -> dict[str, object]: # mutable-ok: the call's kwargs + """The request with the wrapper's ``custom_llm_provider`` when the request names none.""" + if custom_llm_provider and "custom_llm_provider" not in kwargs: + return {**kwargs, "custom_llm_provider": custom_llm_provider} + return dict(kwargs) + + +def managed_agent_request( + custom_llm_provider: str | None, kwargs: Mapping[str, object] +) -> dict[str, object]: # mutable-ok: the call's kwargs + """A Managed Agents request: the given provider, else ``gemini``.""" + request: Final = with_provider(custom_llm_provider, kwargs) + request.setdefault("custom_llm_provider", "gemini") + return request + + +def interaction_request( + custom_llm_provider: str | None, kwargs: Mapping[str, object] +) -> dict[str, object]: # mutable-ok: the call's kwargs + """An Interactions request: the given provider, else ``gemini``; an agent name sent as the model + is no model group, so the model is cleared and the call is not routed.""" + request: Final = managed_agent_request(custom_llm_provider, kwargs) + if request.get("agent") and request.get("model") == request.get("agent"): + request["model"] = None + return request + + +_SEARCH_PROTECTED_PARAMS: Final = frozenset({"search_provider", "api_key", "api_base"}) + + +def search_tool_call( + search_tools: Iterable[Mapping[str, object]], search_tool_name: str, kwargs: Mapping[str, object] +) -> Mapping[str, object]: + """``SearchAPIRouter.async_search_with_fallbacks_helper``'s call: one of the search tools named + ``search_tool_name`` at random, with its provider and credentials, its params under the request's.""" + tools: Final = [tool for tool in search_tools if tool.get("search_tool_name") == search_tool_name] + if not tools: + raise ValueError(f"Search tool '{search_tool_name}' not found in router.search_tools") + selected: Final = random.choice(tools) # load balancing across the providers sharing the name + raw_params: Final[object] = selected.get("litellm_params") or {} + params: Final[Mapping[str, object]] = raw_params if isinstance(raw_params, Mapping) else {} # pyright: ignore[reportUnknownVariableType] # a configured tool's params + search_provider: Final = params.get("search_provider") + if not search_provider: + raise ValueError(f"search_provider not found in litellm_params for search tool '{search_tool_name}'") + search_params: Final = { + key: value + for source in (params, kwargs) + for key, value in source.items() + if key not in _SEARCH_PROTECTED_PARAMS and value is not None + } # comprehension-ok: the tool's params, then the request's + return MappingProxyType( + { + "search_provider": search_provider, + "api_key": params.get("api_key"), + "api_base": params.get("api_base"), + **search_params, + } + ) + + +__all__ = [ + "DISPATCHES", + "DISPATCH_BY_CALL_TYPE", + "ENDPOINTS", + "SYNC_DISPATCHES", + "Dispatch", + "EndpointHandler", + "EndpointSpec", + "as_handler", + "dispatch_for", + "handler_call_type", + "handler_identity", + "host_handler", + "interaction_request", + "managed_agent_request", + "search_tool_call", + "with_provider", +] diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py index f26f72b2a3b..57fb3c2dafd 100644 --- a/litellm/router_python_core.py +++ b/litellm/router_python_core.py @@ -256,6 +256,21 @@ from litellm.router_api import ( from litellm.router_api import RoutingContext as RouterContext from litellm.router_api import deployment_id as _deployment_id from litellm.router_deprecated import DeprecatedRouterMembers +from litellm.router_endpoint_table import ( + ENDPOINTS, + SYNC_DISPATCHES, + Dispatch, + EndpointHandler, + as_handler, + dispatch_for, + handler_call_type, + handler_identity, + host_handler, + interaction_request, + managed_agent_request, + search_tool_call, + with_provider, +) from litellm.router_host import router_host from litellm.router_python_host import PythonHost, python_host from litellm.router_strategy.auto_router.auto_router_model_naming import ( @@ -1332,6 +1347,10 @@ def _base_model_widened_window( _SETTING_ATTRIBUTES: Final = frozenset("_core_" + name for name in SETTING_FIELDS) +# A table-bound endpoint: its arguments and result depend on the endpoint, as before the table. +TableEndpoint: TypeAlias = Callable[..., Any] # each endpoint keeps its litellm signature + + class PythonCore(DeprecatedRouterMembers): _core_model_names: set[str] = set() _core_cache_responses: bool | None = False @@ -1340,6 +1359,137 @@ class PythonCore(DeprecatedRouterMembers): _settings_version: int = 0 _core_applied: Mapping[str, object] = MappingProxyType({}) + # The endpoints of litellm/router_endpoint_table.py, bound per instance by _bind_endpoints; declared + # here so type checkers see them on the core and on Router. + _arealtime: TableEndpoint + _aresponses_websocket: TableEndpoint + a_add_message: TableEndpoint + aadapter_generate_content: TableEndpoint + aanthropic_messages: TableEndpoint + acancel_fine_tuning_job: TableEndpoint + acancel_interaction: TableEndpoint + acancel_responses: TableEndpoint + acompact_responses: TableEndpoint + acreate_agent: TableEndpoint + acreate_assistants: TableEndpoint + acreate_container: TableEndpoint + acreate_fine_tuning_job: TableEndpoint + acreate_interaction: TableEndpoint + acreate_realtime_client_secret: TableEndpoint + acreate_realtime_transcription_session: TableEndpoint + acreate_skill: TableEndpoint + acreate_thread: TableEndpoint + adecisions: TableEndpoint + adelete_agent: TableEndpoint + adelete_assistant: TableEndpoint + adelete_container: TableEndpoint + adelete_container_file: TableEndpoint + adelete_interaction: TableEndpoint + adelete_responses: TableEndpoint + adelete_skill: TableEndpoint + aextraction: TableEndpoint + afile_content: TableEndpoint + afile_delete: TableEndpoint + afile_list: TableEndpoint + agenerate_content: TableEndpoint + agenerate_content_stream: TableEndpoint + aget_agent: TableEndpoint + aget_assistants: TableEndpoint + aget_interaction: TableEndpoint + aget_messages: TableEndpoint + aget_responses: TableEndpoint + aget_skill: TableEndpoint + aget_thread: TableEndpoint + aimage_edit: TableEndpoint + alist_agent_versions: TableEndpoint + alist_agents: TableEndpoint + alist_container_files: TableEndpoint + alist_containers: TableEndpoint + alist_fine_tuning_jobs: TableEndpoint + alist_input_items: TableEndpoint + alist_skills: TableEndpoint + allm_passthrough_route: TableEndpoint + amoderation: TableEndpoint + anthropic_messages: TableEndpoint + aocr: TableEndpoint + arealtime_calls: TableEndpoint + aresponses: TableEndpoint + aretrieve_container: TableEndpoint + aretrieve_container_file: TableEndpoint + aretrieve_container_file_content: TableEndpoint + aretrieve_fine_tuning_job: TableEndpoint + arun_thread: TableEndpoint + asearch: TableEndpoint + asystemone: TableEndpoint + aupload_container_file: TableEndpoint + avector_store_delete: TableEndpoint + avector_store_file_content: TableEndpoint + avector_store_file_create: TableEndpoint + avector_store_file_delete: TableEndpoint + avector_store_file_list: TableEndpoint + avector_store_file_retrieve: TableEndpoint + avector_store_file_update: TableEndpoint + avector_store_list: TableEndpoint + avector_store_retrieve: TableEndpoint + avector_store_search: TableEndpoint + avector_store_update: TableEndpoint + avideo_content: TableEndpoint + avideo_create_character: TableEndpoint + avideo_edit: TableEndpoint + avideo_extension: TableEndpoint + avideo_generation: TableEndpoint + avideo_get_character: TableEndpoint + avideo_list: TableEndpoint + avideo_remix: TableEndpoint + avideo_status: TableEndpoint + cancel_interaction: TableEndpoint + create_agent: TableEndpoint + create_container: TableEndpoint + create_interaction: TableEndpoint + decisions: TableEndpoint + delete_agent: TableEndpoint + delete_container: TableEndpoint + delete_container_file: TableEndpoint + delete_interaction: TableEndpoint + extraction: TableEndpoint + generate_content: TableEndpoint + generate_content_stream: TableEndpoint + get_agent: TableEndpoint + get_interaction: TableEndpoint + list_agent_versions: TableEndpoint + list_agents: TableEndpoint + list_container_files: TableEndpoint + list_containers: TableEndpoint + ocr: TableEndpoint + responses: TableEndpoint + retrieve_container: TableEndpoint + retrieve_container_file: TableEndpoint + retrieve_container_file_content: TableEndpoint + search: TableEndpoint + systemone: TableEndpoint + upload_container_file: TableEndpoint + vector_store_create: TableEndpoint + vector_store_delete: TableEndpoint + vector_store_file_content: TableEndpoint + vector_store_file_create: TableEndpoint + vector_store_file_delete: TableEndpoint + vector_store_file_list: TableEndpoint + vector_store_file_retrieve: TableEndpoint + vector_store_file_update: TableEndpoint + vector_store_list: TableEndpoint + vector_store_retrieve: TableEndpoint + vector_store_search: TableEndpoint + vector_store_update: TableEndpoint + video_content: TableEndpoint + video_create_character: TableEndpoint + video_edit: TableEndpoint + video_extension: TableEndpoint + video_generation: TableEndpoint + video_get_character: TableEndpoint + video_list: TableEndpoint + video_remix: TableEndpoint + video_status: TableEndpoint + # -- entry points: everything outside code needs besides the endpoint methods --------------- # See litellm/router_api.py. The members outside code used before these live in # DeprecatedRouterMembers (litellm/router_deprecated.py). @@ -2806,367 +2956,39 @@ class PythonCore(DeprecatedRouterMembers): case _: return None + def _bind_endpoints(self, *, assistants: bool) -> None: + """Bind the factory-built endpoints of ``router_endpoint_table`` (the assistants ones, or + every other one), each served by its host op through ``host.invoke`` at call time.""" + for spec in ENDPOINTS.values(): + if (spec.dispatch == "assistants") is assistants: + setattr( + self, + spec.name, + self._endpoint_wrapper( + host_handler(self._host, spec.host_op, spec.call_type), spec.dispatch, spec.call_type + ), + ) + def initialize_assistants_endpoint(self): ## INITIALIZE PASS THROUGH ASSISTANTS ENDPOINT ## - self.acreate_assistants = self.factory_function(self._host.invoke("acreate_assistants")) - self.adelete_assistant = self.factory_function(self._host.invoke("adelete_assistant")) - self.aget_assistants = self.factory_function(self._host.invoke("aget_assistants")) - self.acreate_thread = self.factory_function(self._host.invoke("acreate_thread")) - self.aget_thread = self.factory_function(self._host.invoke("aget_thread")) - self.a_add_message = self.factory_function(self._host.invoke("a_add_message")) - self.aget_messages = self.factory_function(self._host.invoke("aget_messages")) - self.arun_thread = self.factory_function(self._host.invoke("arun_thread")) - - def _initialize_core_endpoints(self): - """Helper to initialize core router endpoints.""" - self.amoderation = self.factory_function(self._host.invoke("amoderation"), call_type="moderation") - self.aanthropic_messages = self.factory_function( - self._host.invoke("anthropic_messages"), call_type="anthropic_messages" - ) - self.anthropic_messages = self.factory_function( - self._host.invoke("anthropic_messages"), call_type="anthropic_messages" - ) - self.agenerate_content = self.factory_function( - self._host.invoke("agenerate_content"), call_type="agenerate_content" - ) - self.aadapter_generate_content = self.factory_function( - self._host.invoke("aadapter_generate_content"), call_type="aadapter_generate_content" - ) - self.aresponses = self.factory_function(self._host.invoke("aresponses"), call_type="aresponses") - self.afile_delete = self.factory_function(self._host.invoke("afile_delete"), call_type="afile_delete") - self.afile_content = self.factory_function(self._host.invoke("afile_content"), call_type="afile_content") - self.responses = self.factory_function(self._host.invoke("responses"), call_type="responses") - self.aget_responses = self.factory_function(self._host.invoke("aget_responses"), call_type="aget_responses") - self.acancel_responses = self.factory_function( - self._host.invoke("acancel_responses"), call_type="acancel_responses" - ) - self.acompact_responses = self.factory_function( - self._host.invoke("acompact_responses"), call_type="acompact_responses" - ) - self.adelete_responses = self.factory_function( - self._host.invoke("adelete_responses"), call_type="adelete_responses" - ) - self.alist_input_items = self.factory_function( - self._host.invoke("alist_input_items"), call_type="alist_input_items" - ) - self._arealtime = self.factory_function(self._host.invoke("_arealtime"), call_type="_arealtime") - self.acreate_realtime_client_secret = self.factory_function( - self._host.invoke("acreate_realtime_client_secret"), call_type="acreate_realtime_client_secret" - ) - self.arealtime_calls = self.factory_function(self._host.invoke("arealtime_calls"), call_type="arealtime_calls") - self.acreate_realtime_transcription_session = self.factory_function( - self._host.invoke("acreate_realtime_transcription_session"), - call_type="acreate_realtime_transcription_session", - ) - self._aresponses_websocket = self.factory_function( - self._host.invoke("_aresponses_websocket"), call_type="_aresponses_websocket" - ) - self.acreate_fine_tuning_job = self.factory_function( - self._host.invoke("acreate_fine_tuning_job"), call_type="acreate_fine_tuning_job" - ) - self.acancel_fine_tuning_job = self.factory_function( - self._host.invoke("acancel_fine_tuning_job"), call_type="acancel_fine_tuning_job" - ) - self.alist_fine_tuning_jobs = self.factory_function( - self._host.invoke("alist_fine_tuning_jobs"), call_type="alist_fine_tuning_jobs" - ) - self.aretrieve_fine_tuning_job = self.factory_function( - self._host.invoke("aretrieve_fine_tuning_job"), call_type="aretrieve_fine_tuning_job" - ) - self.afile_list = self.factory_function(self._host.invoke("afile_list"), call_type="alist_files") - self.aimage_edit = self.factory_function(self._host.invoke("aimage_edit"), call_type="aimage_edit") - self.allm_passthrough_route = self.factory_function( - self._host.invoke("allm_passthrough_route"), call_type="allm_passthrough_route" - ) - # Note: acancel_batch is defined as a method on the Router class (not using factory_function) - # to properly handle model-to-provider mapping like acreate_batch and aretrieve_batch - - def _initialize_vector_store_endpoints(self): - """Initialize vector store endpoints.""" - from litellm.vector_stores.main import ( - adelete, - alist, - aretrieve, - asearch, - aupdate, - create, - delete, - list, - retrieve, - search, - update, - ) - - self.avector_store_search = self.factory_function(asearch, call_type="avector_store_search") - self.vector_store_search = self.factory_function(search, call_type="vector_store_search") - self.vector_store_create = self.factory_function(create, call_type="vector_store_create") - self.avector_store_retrieve = self.factory_function(aretrieve, call_type="avector_store_retrieve") - self.vector_store_retrieve = self.factory_function(retrieve, call_type="vector_store_retrieve") - self.avector_store_list = self.factory_function(alist, call_type="avector_store_list") - self.vector_store_list = self.factory_function(list, call_type="vector_store_list") - self.avector_store_update = self.factory_function(aupdate, call_type="avector_store_update") - self.vector_store_update = self.factory_function(update, call_type="vector_store_update") - self.avector_store_delete = self.factory_function(adelete, call_type="avector_store_delete") - self.vector_store_delete = self.factory_function(delete, call_type="vector_store_delete") - - def _initialize_vector_store_file_endpoints(self): - """Initialize vector store file endpoints.""" - from litellm.vector_store_files.main import ( - acreate as avector_store_file_create_fn, - ) - from litellm.vector_store_files.main import ( - adelete as avector_store_file_delete_fn, - ) - from litellm.vector_store_files.main import alist as avector_store_file_list_fn - from litellm.vector_store_files.main import ( - aretrieve as avector_store_file_retrieve_fn, - ) - from litellm.vector_store_files.main import ( - aretrieve_content as avector_store_file_content_fn, - ) - from litellm.vector_store_files.main import ( - aupdate as avector_store_file_update_fn, - ) - from litellm.vector_store_files.main import ( - create as vector_store_file_create_fn, - ) - from litellm.vector_store_files.main import ( - delete as vector_store_file_delete_fn, - ) - from litellm.vector_store_files.main import list as vector_store_file_list_fn - from litellm.vector_store_files.main import ( - retrieve as vector_store_file_retrieve_fn, - ) - from litellm.vector_store_files.main import ( - retrieve_content as vector_store_file_content_fn, - ) - from litellm.vector_store_files.main import ( - update as vector_store_file_update_fn, - ) - - self.avector_store_file_create = self.factory_function( - avector_store_file_create_fn, call_type="avector_store_file_create" - ) - self.vector_store_file_create = self.factory_function( - vector_store_file_create_fn, call_type="vector_store_file_create" - ) - self.avector_store_file_list = self.factory_function( - avector_store_file_list_fn, call_type="avector_store_file_list" - ) - self.vector_store_file_list = self.factory_function( - vector_store_file_list_fn, call_type="vector_store_file_list" - ) - self.avector_store_file_retrieve = self.factory_function( - avector_store_file_retrieve_fn, call_type="avector_store_file_retrieve" - ) - self.vector_store_file_retrieve = self.factory_function( - vector_store_file_retrieve_fn, call_type="vector_store_file_retrieve" - ) - self.avector_store_file_content = self.factory_function( - avector_store_file_content_fn, call_type="avector_store_file_content" - ) - self.vector_store_file_content = self.factory_function( - vector_store_file_content_fn, call_type="vector_store_file_content" - ) - self.avector_store_file_update = self.factory_function( - avector_store_file_update_fn, call_type="avector_store_file_update" - ) - self.vector_store_file_update = self.factory_function( - vector_store_file_update_fn, call_type="vector_store_file_update" - ) - self.avector_store_file_delete = self.factory_function( - avector_store_file_delete_fn, call_type="avector_store_file_delete" - ) - self.vector_store_file_delete = self.factory_function( - vector_store_file_delete_fn, call_type="vector_store_file_delete" - ) - - def _initialize_google_genai_endpoints(self): - """Initialize Google GenAI endpoints.""" - from litellm.google_genai import ( - agenerate_content, - agenerate_content_stream, - generate_content, - generate_content_stream, - ) - - self.agenerate_content = self.factory_function(agenerate_content, call_type="agenerate_content") - self.generate_content = self.factory_function(generate_content, call_type="generate_content") - self.agenerate_content_stream = self.factory_function( - agenerate_content_stream, call_type="agenerate_content_stream" - ) - self.generate_content_stream = self.factory_function( - generate_content_stream, call_type="generate_content_stream" - ) - - def _initialize_ocr_search_endpoints(self): - """Initialize OCR and search endpoints.""" - from litellm.ocr import aocr, ocr - - self.aocr = self.factory_function(aocr, call_type="aocr") - self.ocr = self.factory_function(ocr, call_type="ocr") - - from litellm.search import asearch, search - - self.asearch = self.factory_function(asearch, call_type="asearch") - self.search = self.factory_function(search, call_type="search") - - from litellm.decisions import adecisions, decisions - - self.adecisions = self.factory_function(adecisions, call_type="adecisions") - self.decisions = self.factory_function(decisions, call_type="decisions") - - from litellm.systemone import asystemone, systemone - - self.asystemone = self.factory_function(asystemone, call_type="asystemone") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - self.systemone = self.factory_function(systemone, call_type="systemone") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - - from litellm.extraction import aextraction, extraction - - self.aextraction = self.factory_function(aextraction, call_type="aextraction") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - self.extraction = self.factory_function(extraction, call_type="extraction") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - - def _initialize_video_endpoints(self): - """Initialize video endpoints.""" - from litellm.videos import ( - avideo_content, - avideo_create_character, - avideo_edit, - avideo_extension, - avideo_generation, - avideo_get_character, - avideo_list, - avideo_remix, - avideo_status, - video_content, - video_create_character, - video_edit, - video_extension, - video_generation, - video_get_character, - video_list, - video_remix, - video_status, - ) - - self.avideo_generation = self.factory_function(avideo_generation, call_type="avideo_generation") - self.video_generation = self.factory_function(video_generation, call_type="video_generation") - self.avideo_list = self.factory_function(avideo_list, call_type="avideo_list") - self.video_list = self.factory_function(video_list, call_type="video_list") - self.avideo_status = self.factory_function(avideo_status, call_type="avideo_status") - self.video_status = self.factory_function(video_status, call_type="video_status") - self.avideo_content = self.factory_function(avideo_content, call_type="avideo_content") - self.video_content = self.factory_function(video_content, call_type="video_content") - self.avideo_remix = self.factory_function(avideo_remix, call_type="avideo_remix") - self.video_remix = self.factory_function(video_remix, call_type="video_remix") - self.avideo_create_character = self.factory_function( - avideo_create_character, call_type="avideo_create_character" - ) - self.video_create_character = self.factory_function(video_create_character, call_type="video_create_character") - self.avideo_get_character = self.factory_function(avideo_get_character, call_type="avideo_get_character") - self.video_get_character = self.factory_function(video_get_character, call_type="video_get_character") - self.avideo_edit = self.factory_function(avideo_edit, call_type="avideo_edit") - self.video_edit = self.factory_function(video_edit, call_type="video_edit") - self.avideo_extension = self.factory_function(avideo_extension, call_type="avideo_extension") - self.video_extension = self.factory_function(video_extension, call_type="video_extension") - - def _initialize_container_endpoints(self): - """Initialize container endpoints.""" - from litellm.containers import ( - acreate_container, - adelete_container, - alist_containers, - aretrieve_container, - create_container, - delete_container, - list_containers, - retrieve_container, - ) - from litellm.containers.endpoint_factory import ( - _generated_endpoints as container_file_endpoints, - ) - - self.acreate_container = self.factory_function(acreate_container, call_type="acreate_container") - self.create_container = self.factory_function(create_container, call_type="create_container") - self.alist_containers = self.factory_function(alist_containers, call_type="alist_containers") - self.list_containers = self.factory_function(list_containers, call_type="list_containers") - self.aretrieve_container = self.factory_function(aretrieve_container, call_type="aretrieve_container") - self.retrieve_container = self.factory_function(retrieve_container, call_type="retrieve_container") - self.adelete_container = self.factory_function(adelete_container, call_type="adelete_container") - self.delete_container = self.factory_function(delete_container, call_type="delete_container") - - # Auto-register JSON-generated container file endpoints - for name, func in container_file_endpoints.items(): - setattr(self, name, self.factory_function(func, call_type=name)) - - def _initialize_skills_endpoints(self): - """Initialize Anthropic Skills API endpoints.""" - self.acreate_skill = self.factory_function(self._host.invoke("acreate_skill"), call_type="acreate_skill") - self.alist_skills = self.factory_function(self._host.invoke("alist_skills"), call_type="alist_skills") - self.aget_skill = self.factory_function(self._host.invoke("aget_skill"), call_type="aget_skill") - self.adelete_skill = self.factory_function(self._host.invoke("adelete_skill"), call_type="adelete_skill") - - def _initialize_interactions_endpoints(self): - """Initialize Google Interactions API endpoints.""" - from litellm.interactions import acancel as acancel_interaction - from litellm.interactions import acreate as acreate_interaction - from litellm.interactions import adelete as adelete_interaction - from litellm.interactions import aget as aget_interaction - from litellm.interactions import cancel as cancel_interaction - from litellm.interactions import create as create_interaction - from litellm.interactions import delete as delete_interaction - from litellm.interactions import get as get_interaction - - self.acreate_interaction = self.factory_function(acreate_interaction, call_type="acreate_interaction") - self.create_interaction = self.factory_function(create_interaction, call_type="create_interaction") - self.aget_interaction = self.factory_function(aget_interaction, call_type="aget_interaction") - self.get_interaction = self.factory_function(get_interaction, call_type="get_interaction") - self.adelete_interaction = self.factory_function(adelete_interaction, call_type="adelete_interaction") - self.delete_interaction = self.factory_function(delete_interaction, call_type="delete_interaction") - self.acancel_interaction = self.factory_function(acancel_interaction, call_type="acancel_interaction") - self.cancel_interaction = self.factory_function(cancel_interaction, call_type="cancel_interaction") - - def _initialize_managed_agents_endpoints(self): - """Initialize Google Managed Agents API endpoints (v1beta/agents).""" - from litellm.interactions.agents import acreate as acreate_agent - from litellm.interactions.agents import adelete as adelete_agent - from litellm.interactions.agents import aget as aget_agent - from litellm.interactions.agents import alist as alist_agents - from litellm.interactions.agents import alist_versions as alist_agent_versions - from litellm.interactions.agents import create as create_agent - from litellm.interactions.agents import delete as delete_agent - from litellm.interactions.agents import get as get_agent - from litellm.interactions.agents import list as list_agents - from litellm.interactions.agents import list_versions as list_agent_versions - - self.acreate_agent = self.factory_function(acreate_agent, call_type="acreate_agent") - self.create_agent = self.factory_function(create_agent, call_type="create_agent") - self.alist_agents = self.factory_function(alist_agents, call_type="alist_agents") - self.list_agents = self.factory_function(list_agents, call_type="list_agents") - self.aget_agent = self.factory_function(aget_agent, call_type="aget_agent") - self.get_agent = self.factory_function(get_agent, call_type="get_agent") - self.adelete_agent = self.factory_function(adelete_agent, call_type="adelete_agent") - self.delete_agent = self.factory_function(delete_agent, call_type="delete_agent") - self.alist_agent_versions = self.factory_function(alist_agent_versions, call_type="alist_agent_versions") - self.list_agent_versions = self.factory_function(list_agent_versions, call_type="list_agent_versions") - - def _initialize_specialized_endpoints(self): - """Helper to initialize specialized router endpoints (vector store, OCR, search, video, container, skills, interactions).""" - self._initialize_vector_store_endpoints() - self._initialize_vector_store_file_endpoints() - self._initialize_google_genai_endpoints() - self._initialize_ocr_search_endpoints() - # Override vector store methods with router-aware implementations - self._override_vector_store_methods_for_router() - self._initialize_video_endpoints() - self._initialize_container_endpoints() - self._initialize_skills_endpoints() - self._initialize_interactions_endpoints() - self._initialize_managed_agents_endpoints() + self._bind_endpoints(assistants=True) def initialize_router_endpoints(self): - self._initialize_core_endpoints() - self._initialize_specialized_endpoints() + self._bind_endpoints(assistants=False) + + # The endpoint groups are bound together from one table (litellm/router_endpoint_table.py); + # each group's former initializer binds them all. + _initialize_core_endpoints = initialize_router_endpoints + _initialize_specialized_endpoints = initialize_router_endpoints + _initialize_vector_store_endpoints = initialize_router_endpoints + _initialize_vector_store_file_endpoints = initialize_router_endpoints + _initialize_google_genai_endpoints = initialize_router_endpoints + _initialize_ocr_search_endpoints = initialize_router_endpoints + _initialize_video_endpoints = initialize_router_endpoints + _initialize_container_endpoints = initialize_router_endpoints + _initialize_skills_endpoints = initialize_router_endpoints + _initialize_interactions_endpoints = initialize_router_endpoints + _initialize_managed_agents_endpoints = initialize_router_endpoints def validate_fallbacks(self, fallback_param: Sequence[Mapping[str, object]] | None): """ @@ -5940,14 +5762,18 @@ class PythonCore(DeprecatedRouterMembers): messages: list[AllMessageValues], kwargs: dict[str, Any], ): + # The routing function `acompletion` hands along, kept out of `function_setup` (which takes + # the litellm function's name) and handed back for a routed call. + routing_function: Final = kwargs.pop("original_function", None) litellm_logging_object = kwargs.get("litellm_logging_obj", None) if litellm_logging_object is None: litellm_logging_object, kwargs = function_setup( **{ + **kwargs, "original_function": "acompletion", "rules_obj": Rules(), "start_time": get_utc_datetime(), - **kwargs, + "litellm_call_id": kwargs.get("litellm_call_id") or str(uuid.uuid4()), } ) litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object) @@ -6022,7 +5848,6 @@ class PythonCore(DeprecatedRouterMembers): _model_list: Final = self._core_get_model_list(model_name=model) if _model_list is None or len(_model_list) == 0: # if direct call to model - kwargs.pop("original_function") return await self._host.invoke("acompletion")(**kwargs) target_group_timeouts: Final = { @@ -6039,6 +5864,7 @@ class PythonCore(DeprecatedRouterMembers): return await self.async_function_with_fallbacks( **{key: value for key, value in kwargs.items() if key not in ("timeout", "stream_timeout")}, **target_group_timeouts, + original_function=routing_function, ) def image_generation(self, prompt: str, model: str, **kwargs): @@ -6364,7 +6190,6 @@ class PythonCore(DeprecatedRouterMembers): async def arerank(self, model: str, **kwargs): try: kwargs["model"] = model - kwargs["input"] = input kwargs["original_function"] = self._arerank self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) @@ -6795,7 +6620,7 @@ class PythonCore(DeprecatedRouterMembers): passthrough_on_no_deployment: Final = kwargs.pop("passthrough_on_no_deployment", False) function_name: Final = "_ageneric_api_call_with_fallbacks" - handler_name: Final[str | None] = getattr(original_generic_function, "__name__", None) + handler_name: Final[str | None] = handler_call_type(original_generic_function) deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) @@ -6832,7 +6657,7 @@ class PythonCore(DeprecatedRouterMembers): ) custom_llm_provider: Final = provider_for_generic_call(data) - generic_function_name: Final[str] = getattr(original_generic_function, "__name__", "") + generic_function_name: Final[str] = handler_name or "" response_kwargs: Final = { **without_retrieve_stream_defaults(data, kwargs), @@ -7342,9 +7167,7 @@ class PythonCore(DeprecatedRouterMembers): ) carry_superseded_attempt_usage(retry_kwargs) try: - with endpoint_selection( - getattr(handler, "__name__", None), getattr(handler, "__module__", None), model=model_group - ): + with endpoint_selection(*handler_identity(handler), model=model_group): response = await attempt_function(**retry_kwargs) except Exception as retry_error: # noqa: BLE001 # the fallback chain judges the failed retry wrapped = _stream_fallback_error_for_raised(retry_error, model_group, False, e.llm_provider) @@ -7576,7 +7399,7 @@ class PythonCore(DeprecatedRouterMembers): Returns: The response from the handler function """ - handler_name: Final = original_function.__name__ + handler_name: Final = handler_call_type(original_function) metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call") try: verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model) @@ -7960,13 +7783,13 @@ class PythonCore(DeprecatedRouterMembers): try: # If model is None, use the factory function approach (direct SDK call) if model is None: - from litellm.vector_stores.main import acreate - - # Use the factory function to handle the call - factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create") - return await factory_fn(**kwargs) - - from litellm.vector_stores import acreate as avector_store_create_sdk + factory_fn: Final = self._endpoint_wrapper( + host_handler(self._host, "vector_stores.main.acreate", "avector_store_create"), + "vector_store", + "avector_store_create", + ) + created: Final = factory_fn(**kwargs) + return await created if inspect.isawaitable(created) else created parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) deployment: Final = await self._core_async_get_available_deployment( @@ -7998,7 +7821,7 @@ class PythonCore(DeprecatedRouterMembers): custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await avector_store_create_sdk( + pending_response: Final = self._host.invoke("vector_stores.acreate")( **{ **data, "custom_llm_provider": custom_llm_provider, @@ -8007,6 +7830,7 @@ class PythonCore(DeprecatedRouterMembers): **kwargs, } ) + response: Final = await pending_response if inspect.isawaitable(pending_response) else pending_response self.success_calls[model_name] += 1 verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) @@ -8395,380 +8219,122 @@ class PythonCore(DeprecatedRouterMembers): return await original_function(**kwargs) - def factory_function( - self, - original_function: Callable, - call_type: Literal[ - "assistants", - "moderation", - "anthropic_messages", - "aresponses", - "acancel_responses", - "acompact_responses", - "responses", - "aget_responses", - "adelete_responses", - "afile_delete", - "afile_content", - "_arealtime", - "acreate_realtime_client_secret", - "arealtime_calls", - "acreate_realtime_transcription_session", - "_aresponses_websocket", - "acreate_fine_tuning_job", - "acancel_fine_tuning_job", - "alist_fine_tuning_jobs", - "aretrieve_fine_tuning_job", - "alist_files", - "aimage_edit", - "allm_passthrough_route", - "alist_input_items", - "agenerate_content", - "generate_content", - "agenerate_content_stream", - "generate_content_stream", - "avector_store_search", - "avector_store_create", - "avector_store_retrieve", - "avector_store_list", - "avector_store_update", - "avector_store_delete", - "avector_store_file_create", - "avector_store_file_list", - "avector_store_file_retrieve", - "avector_store_file_content", - "avector_store_file_update", - "avector_store_file_delete", - "vector_store_search", - "vector_store_create", - "vector_store_retrieve", - "vector_store_list", - "vector_store_update", - "vector_store_delete", - "vector_store_file_create", - "vector_store_file_list", - "vector_store_file_retrieve", - "vector_store_file_content", - "vector_store_file_update", - "vector_store_file_delete", - "aocr", - "ocr", - "asearch", - "search", - "adecisions", - "decisions", - "asystemone", - "systemone", - "aextraction", - "extraction", - "aadapter_generate_content", - "avideo_generation", - "video_generation", - "avideo_list", - "video_list", - "avideo_status", - "video_status", - "avideo_content", - "video_content", - "avideo_remix", - "video_remix", - "avideo_create_character", - "video_create_character", - "avideo_get_character", - "video_get_character", - "avideo_edit", - "video_edit", - "avideo_extension", - "video_extension", - "acreate_container", - "create_container", - "alist_containers", - "list_containers", - "aretrieve_container", - "retrieve_container", - "adelete_container", - "delete_container", - "aupload_container_file", - "upload_container_file", - "alist_container_files", - "list_container_files", - "aretrieve_container_file", - "retrieve_container_file", - "adelete_container_file", - "delete_container_file", - "acreate_skill", - "alist_skills", - "aget_skill", - "adelete_skill", - "acreate_interaction", - "create_interaction", - "aget_interaction", - "get_interaction", - "adelete_interaction", - "delete_interaction", - "acancel_interaction", - "cancel_interaction", - "acreate_agent", - "create_agent", - "alist_agents", - "list_agents", - "aget_agent", - "get_agent", - "adelete_agent", - "delete_agent", - "alist_agent_versions", - "list_agent_versions", - ] = "assistants", - ): + def factory_function(self, original_function: Callable[..., object], call_type: str = "assistants"): """ - Creates appropriate wrapper functions for different API call types. + Creates the wrapper the router serves an endpoint of ``call_type`` with, dispatched as + ``litellm/router_endpoint_table.py`` says. An unknown call type is refused here, when the + endpoint is bound, never answered with a wrapper that dispatches nowhere. Returns: - A synchronous function for synchronous call types - An asynchronous function for asynchronous call types """ - # Handle synchronous call types - if call_type in ( - "responses", - "generate_content", - "generate_content_stream", - "ocr", - "search", - "decisions", - "systemone", - "extraction", - "video_generation", - "video_list", - "video_status", - "video_content", - "video_remix", - "create_container", - "list_containers", - "retrieve_container", - "delete_container", - ): + return self._endpoint_wrapper(as_handler(original_function, call_type), dispatch_for(call_type), call_type) + + def _endpoint_wrapper(self, handler: EndpointHandler, dispatch: Dispatch, call_type: str) -> Callable[..., object]: + if dispatch in SYNC_DISPATCHES: def sync_wrapper( custom_llm_provider: str | None = None, client: object | None = None, **kwargs, ): - return self._generic_api_call_with_fallbacks(original_function=original_function, **kwargs) + return self._dispatch_sync_endpoint(handler, dispatch, call_type, custom_llm_provider, client, kwargs) return sync_wrapper - if call_type in ( - "vector_store_search", - "vector_store_create", - "vector_store_retrieve", - "vector_store_list", - "vector_store_update", - "vector_store_delete", - ): - - def vector_store_sync_wrapper( - custom_llm_provider: str | None = None, - client: object | None = None, - **kwargs, - ): - provider_kwargs: Final = ( - MappingProxyType({**kwargs, "custom_llm_provider": custom_llm_provider}) - if custom_llm_provider and "custom_llm_provider" not in kwargs - else MappingProxyType(kwargs) - ) - search_kwargs: Final = ( - MappingProxyType( - { - **provider_kwargs, - "_direct_vector_store_embedding_executor": RouterVectorStoreEmbeddingExecutor( - router=self._core_public(), - metadata=self._vector_store_request_metadata(kwargs), - ), - } - ) - if call_type == "vector_store_search" - else provider_kwargs - ) - if search_kwargs.get("model"): - return self._generic_api_call_with_fallbacks(original_function=original_function, **search_kwargs) - if call_type == "vector_store_search": - return original_function(**MappingProxyType({**search_kwargs, "router": self._core_public()})) - return original_function(**search_kwargs) - - return vector_store_sync_wrapper - - if call_type in ( - "vector_store_file_create", - "vector_store_file_list", - "vector_store_file_retrieve", - "vector_store_file_content", - "vector_store_file_update", - "vector_store_file_delete", - ): - - def vector_store_file_sync_wrapper( - custom_llm_provider: str | None = None, - client: object | None = None, - **kwargs, - ): - return original_function( - custom_llm_provider=custom_llm_provider, - client=client, - **kwargs, - ) - - return vector_store_file_sync_wrapper - - if call_type in ( - "create_agent", - "list_agents", - "get_agent", - "delete_agent", - "list_agent_versions", - ): - - def managed_agents_sync_wrapper( - custom_llm_provider: str | None = None, - client: object | None = None, - **kwargs, - ): - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - if "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = "gemini" - return original_function(**kwargs) - - return managed_agents_sync_wrapper - - # Handle asynchronous call types async def async_wrapper( custom_llm_provider: str | None = None, client: AsyncOpenAI | None = None, **kwargs, ): - if call_type == "assistants": + return await self._dispatch_async_endpoint( + handler, dispatch, call_type, custom_llm_provider, client, kwargs + ) + + return async_wrapper + + def _dispatch_sync_endpoint( + self, + handler: EndpointHandler, + dispatch: Dispatch, + call_type: str, + custom_llm_provider: str | None, + client: object | None, + kwargs: dict[str, Any], # mutable-ok: the wrapper's own kwargs, shaped in place as before + ) -> object: + match dispatch: + case "generic_sync": + return self._generic_api_call_with_fallbacks(original_function=handler, **kwargs) + case "search_sync": + return self._search_with_search_tools(original_function=handler, **kwargs) + case "vector_store_sync": + return self._vector_store_sync_call(handler, call_type, custom_llm_provider, kwargs) + case "vector_store_file_sync": + return handler(custom_llm_provider=custom_llm_provider, client=client, **kwargs) + case "interaction_sync": + interaction_kwargs: Final = interaction_request(custom_llm_provider, kwargs) + if interaction_kwargs.get("model"): + return self._generic_api_call_with_fallbacks(original_function=handler, **interaction_kwargs) + return handler(**interaction_kwargs) + case "managed_agent_sync": + return handler(**managed_agent_request(custom_llm_provider, kwargs)) + case _: + raise ValueError(f"{dispatch!r} is not a sync endpoint dispatch") + + async def _dispatch_async_endpoint( + self, + handler: EndpointHandler, + dispatch: Dispatch, + call_type: str, + custom_llm_provider: str | None, + client: AsyncOpenAI | None, + kwargs: dict[str, Any], # mutable-ok: the wrapper's own kwargs, shaped in place as before + ) -> object: + match dispatch: + case "assistants": return await self._pass_through_assistants_endpoint_factory( - original_function=original_function, + original_function=handler, custom_llm_provider=custom_llm_provider, client=client, **kwargs, ) - elif call_type == "moderation": - return await self._pass_through_moderation_endpoint_factory( - original_function=original_function, **kwargs - ) - elif call_type in ("asearch", "search"): - return await self._asearch_with_fallbacks( - original_function=original_function, - **kwargs, - ) - elif call_type in ( - "avector_store_file_create", - "avector_store_file_list", - "avector_store_file_retrieve", - "avector_store_file_content", - "avector_store_file_update", - "avector_store_file_delete", - ): - return await self._init_vector_store_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - elif call_type == "aresponses": - return await self._aresponses_with_streaming_fallbacks( - original_function=original_function, - **kwargs, - ) - elif call_type in ( - "acreate_realtime_client_secret", - "arealtime_calls", - "acreate_realtime_transcription_session", - ): + case "moderation": + return await self._pass_through_moderation_endpoint_factory(original_function=handler, **kwargs) + case "search": + return await self._asearch_with_fallbacks(original_function=handler, **kwargs) + case "responses": + return await self._aresponses_with_streaming_fallbacks(original_function=handler, **kwargs) + case "anthropic_messages": + return await self._aanthropic_messages_with_streaming_fallbacks(original_function=handler, **kwargs) + case "generic": + return await self._ageneric_api_call_with_fallbacks(original_function=handler, **kwargs) + case "generic_client": + return await self._ageneric_api_call_with_fallbacks(original_function=handler, client=client, **kwargs) + case "generic_provider_client": return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, + original_function=handler, + custom_llm_provider=custom_llm_provider, client=client, **kwargs, ) - elif call_type in ( - "anthropic_messages", - "_arealtime", - "_aresponses_websocket", - "acreate_fine_tuning_job", - "acancel_fine_tuning_job", - "alist_fine_tuning_jobs", - "aretrieve_fine_tuning_job", - "alist_files", - "aimage_edit", - "agenerate_content", - "agenerate_content_stream", - "aocr", - "ocr", - "adecisions", - "asystemone", - "aextraction", - "avideo_generation", - "avideo_list", - "avideo_status", - "avideo_content", - "avideo_remix", - "avideo_create_character", - "avideo_get_character", - "avideo_edit", - "avideo_extension", - "acreate_skill", - "alist_skills", - "aget_skill", - "adelete_skill", - ): - return await self._dispatch_generic_call_type( - call_type=call_type, - original_function=original_function, - **kwargs, - ) - elif call_type in ( - "acreate_container", - "alist_containers", - "aretrieve_container", - "adelete_container", - "aupload_container_file", - "alist_container_files", - "aretrieve_container_file", - "adelete_container_file", - "aretrieve_container_file_content", - ): - return await self._init_containers_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - elif call_type == "allm_passthrough_route": + case "passthrough": if client: kwargs["client"] = client return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, + original_function=handler, passthrough_on_no_deployment=True, **kwargs, ) - elif call_type in ( - "aget_responses", - "acancel_responses", - "acompact_responses", - "adelete_responses", - "alist_input_items", - ): + case "response_id": return await self._init_responses_api_endpoints( - original_function=original_function, - **kwargs, + original_function=handler, custom_llm_provider=custom_llm_provider, **kwargs ) - elif call_type in ( - "avector_store_search", - "avector_store_create", - "avector_store_retrieve", - "avector_store_list", - "avector_store_update", - "avector_store_delete", - ): + case "container": + return await self._init_containers_api_endpoints( + original_function=handler, custom_llm_provider=custom_llm_provider, **kwargs + ) + case "vector_store": vector_store_kwargs: Final = ( { **kwargs, @@ -8781,44 +8347,64 @@ class PythonCore(DeprecatedRouterMembers): else kwargs ) return await self._init_vector_store_api_endpoints( - original_function=original_function, + original_function=handler, custom_llm_provider=custom_llm_provider, call_type=call_type, **vector_store_kwargs, ) - elif call_type in ("afile_delete", "afile_content"): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - client=client, - **kwargs, - ) - elif call_type in ( - "acreate_interaction", - "create_interaction", - "aget_interaction", - "adelete_interaction", - "acancel_interaction", - ): + case "interaction": return await self._init_interactions_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, + original_function=handler, custom_llm_provider=custom_llm_provider, **kwargs ) - elif call_type in ( - "acreate_agent", - "alist_agents", - "aget_agent", - "adelete_agent", - "alist_agent_versions", - ): + case "managed_agent": return await self._init_managed_agents_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, + original_function=handler, custom_llm_provider=custom_llm_provider, **kwargs ) + case _: + raise ValueError(f"{dispatch!r} is not an async endpoint dispatch") - return async_wrapper + def _vector_store_sync_call( + self, + handler: EndpointHandler, + call_type: str, + custom_llm_provider: str | None, + kwargs: Mapping[str, object], + ) -> object: + provider_kwargs: Final = MappingProxyType(with_provider(custom_llm_provider, kwargs)) + search_kwargs: Final = ( + MappingProxyType( + { + **provider_kwargs, + "_direct_vector_store_embedding_executor": RouterVectorStoreEmbeddingExecutor( + router=self._core_public(), + metadata=self._vector_store_request_metadata(kwargs), + ), + } + ) + if call_type == "vector_store_search" + else provider_kwargs + ) + if search_kwargs.get("model"): + return self._generic_api_call_with_fallbacks(original_function=handler, **search_kwargs) + if call_type == "vector_store_search": + return handler(**MappingProxyType({**search_kwargs, "router": self._core_public()})) + return handler(**search_kwargs) + + def _search_with_search_tools( + self, + original_function: Callable, + **kwargs: object, # kwargs-ok: the search request + ) -> object: + """The sync ``search``: one attempt on one of the router's search tools named + ``search_tool_name`` (else ``model``), as ``asearch`` picks among them.""" + search_tool_name: Final = kwargs.get("search_tool_name", kwargs.get("model")) + if not search_tool_name: + raise ValueError("search_tool_name or model parameter is required for search") + request: Final = {key: value for key, value in kwargs.items() if key != "model"} + self._update_kwargs_before_fallbacks( + model=search_tool_name, kwargs=request, metadata_variable_name="litellm_metadata" + ) + return original_function(**search_tool_call(self._core_ctx.settings().search_tools, search_tool_name, request)) @staticmethod def _vector_store_request_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: @@ -8918,18 +8504,24 @@ class PythonCore(DeprecatedRouterMembers): async def _init_responses_api_endpoints( self, original_function: Callable, + custom_llm_provider: str | None = None, **kwargs, ): """ Initialize the Responses API endpoints on the router. GET, DELETE, CANCEL Responses API Requests encode the model_id in the response_id, this function decodes the response_id and sets the model to the model_id. + A plain id with no model names no deployment: the handler is called directly, as + ``litellm.aget_responses`` is without a router (the given provider, else the one the id names). """ from litellm.responses.utils import ResponsesAPIRequestUtils model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(kwargs.get("response_id")) if model_id is not None: kwargs["model"] = model_id + if not kwargs.get("model"): + direct_kwargs: Final = {key: value for key, value in kwargs.items() if key != "model"} + return await original_function(**with_provider(custom_llm_provider, direct_kwargs)) return await self._ageneric_api_call_with_fallbacks( original_function=original_function, **kwargs, @@ -8947,22 +8539,15 @@ class PythonCore(DeprecatedRouterMembers): GET, DELETE, CANCEL Interactions API Requests don't need model-based routing, so we call the original function directly with the custom_llm_provider. """ - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - # Default to gemini for interactions API - if "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = "gemini" - # If the proxy accidentally passed agent name as model, clear it - if kwargs.get("agent") and kwargs.get("model") == kwargs.get("agent"): - kwargs["model"] = None + request: Final = interaction_request(custom_llm_provider, kwargs) # Model-based interactions use deployment routing + fallbacks; agent-only calls # must not enter model-group lookup (agent name is not a LiteLLM deployment). - if kwargs.get("model"): + if request.get("model"): return await self._ageneric_api_call_with_fallbacks( original_function=original_function, - **kwargs, + **request, ) - return await original_function(**kwargs) + return await original_function(**request) async def _init_managed_agents_api_endpoints( self, @@ -8976,11 +8561,7 @@ class PythonCore(DeprecatedRouterMembers): CRUD operations for Gemini managed agents don't need model-based routing, so we call the original function directly with the custom_llm_provider. """ - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - if "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = "gemini" - return await original_function(**kwargs) + return await original_function(**managed_agent_request(custom_llm_provider, kwargs)) async def _pass_through_assistants_endpoint_factory( self, @@ -9371,7 +8952,7 @@ class PythonCore(DeprecatedRouterMembers): """ model_group: Final[str | None] = kwargs.get("model") compaction_surface: Final = surface_for_call( - getattr(kwargs.get("original_generic_function") or kwargs.get("original_function"), "__name__", "") + handler_call_type(kwargs.get("original_generic_function") or kwargs.get("original_function")) or "" ) if compaction_surface is not None: kwargs["_context_compaction_state"] = initialize_compaction_state(kwargs, compaction_surface) @@ -9730,8 +9311,7 @@ class PythonCore(DeprecatedRouterMembers): ) try: with endpoint_selection( - getattr(endpoint_function, "__name__", None), - getattr(endpoint_function, "__module__", None), + *handler_identity(endpoint_function), model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group), ): response = original_function(*args, **kwargs) @@ -10541,7 +10121,7 @@ class PythonCore(DeprecatedRouterMembers): get_safeguard_refusal_stop_details, ) - if getattr(original_generic_function, "__name__", "") != "anthropic_messages": + if handler_call_type(original_generic_function) != "anthropic_messages": return False if get_safeguard_refusal_stop_details(response) is None: return False diff --git a/litellm/router_rust_core.py b/litellm/router_rust_core.py index af0da9757c5..ab33fd2f2d2 100644 --- a/litellm/router_rust_core.py +++ b/litellm/router_rust_core.py @@ -1974,9 +1974,7 @@ class RustCore( return await self._drive(plan_for("aspeech"), request) async def arerank(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload - # `arerank` sets kwargs["input"] = input: the Python builtin, which reaches the deployment call. - request: Final = {**request_from_call("arerank", args, kwargs), "input": input} - return await self._drive(plan_for("arerank"), request) + return await self._drive(plan_for("arerank"), request_from_call("arerank", args, kwargs)) async def amoderation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload return await self._endpoints.amoderation(request_from_call("amoderation", args, kwargs)) diff --git a/litellm/router_rust_endpoints.py b/litellm/router_rust_endpoints.py index b6066d6ba46..d27e85f514f 100644 --- a/litellm/router_rust_endpoints.py +++ b/litellm/router_rust_endpoints.py @@ -26,6 +26,14 @@ from litellm._logging import verbose_router_logger from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name from litellm.litellm_core_utils.routing_common_utils import without_router_only_kwargs from litellm.responses.retrieve_stream_query import without_retrieve_stream_defaults +from litellm.router_endpoint_table import ( + ENDPOINTS, + EndpointSpec, + interaction_request, + managed_agent_request, + search_tool_call, + with_provider, +) from litellm.router_rust_shaping import ( ROUTER_CLIENTS_TARGET, STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, @@ -110,6 +118,8 @@ class Plan: # The named parameter the endpoint takes beside **kwargs, sent between the deployment params and # the request (`{**data, name: value, **kwargs}`). named: tuple[str, ...] = () + # The endpoint's call type, which the routing decisions key on (never the host function's name). + call_type: str | None = None @property def sets_headers(self) -> bool: @@ -122,44 +132,36 @@ class Plan: return get_router_metadata_variable_name(function_name=self.function_name) -def _generic(endpoint: str, host_op: str | None = None, *, sync: bool = False) -> Plan: +def _generic(endpoint: str, host_op: str | None = None, *, sync: bool = False, call_type: str | None = None) -> Plan: + op: Final = host_op or endpoint + routed_as: Final = call_type or op.rsplit(".", 1)[-1] if sync: return Plan( endpoint=endpoint, - host_op=host_op or endpoint, + host_op=op, kind="generic_sync", sync=True, prepare_field="litellm_metadata", function_name="generic_api_call", - splits_stream_timeout=_handler_name(host_op or endpoint) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + splits_stream_timeout=routed_as in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, stamp="none", fail_key="group", single_attempt=True, + call_type=routed_as, ) return Plan( endpoint=endpoint, - host_op=host_op or endpoint, + host_op=op, kind="generic", prepare_field="litellm_metadata", function_name="_ageneric_api_call_with_fallbacks", - splits_stream_timeout=_handler_name(host_op or endpoint) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + splits_stream_timeout=routed_as in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, stamp="deployment", fail_key="group", + call_type=routed_as, ) -def _handler_name(host_op: str) -> str: - return host_op.rsplit(".", 1)[-1] - - -def invoked_handler_name(core: "RustCore", plan: Plan) -> str | None: - """The generic handler's ``__name__`` (`original_generic_function.__name__`): the function the - host serves for the op, which decides the stream-timeout split and the responses socket.""" - function: Final[object] = core._host.invoke(plan.host_op) # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # the core's host - name: Final[object] = getattr(function, "__name__", None) - return name if isinstance(name, str) else None - - PLANS: Final[Mapping[str, Plan]] = MappingProxyType( { "acompletion": Plan( @@ -294,166 +296,21 @@ PLANS: Final[Mapping[str, Plan]] = MappingProxyType( # --- the generic endpoints ------------------------------------------------------------------ -# endpoint -> host op, for the async generic endpoints routed through -# `_ageneric_api_call_with_fallbacks` when they carry a model. -_HOST_GENERIC: Final[Mapping[str, str]] = MappingProxyType( +# The table's dispatches the Rust core routes with a generic plan (`_ageneric_api_call_with_fallbacks`, +# or the sync `_generic_api_call_with_fallbacks`) when the request carries a model. +_GENERIC_DISPATCHES: Final = frozenset( { - "afile_delete": "afile_delete", - "afile_content": "afile_content", - "afile_list": "afile_list", - "aget_responses": "aget_responses", - "acancel_responses": "acancel_responses", - "acompact_responses": "acompact_responses", - "adelete_responses": "adelete_responses", - "alist_input_items": "alist_input_items", - "_arealtime": "_arealtime", - "acreate_realtime_client_secret": "acreate_realtime_client_secret", - "arealtime_calls": "arealtime_calls", - "acreate_realtime_transcription_session": "acreate_realtime_transcription_session", - "_aresponses_websocket": "_aresponses_websocket", - "acreate_fine_tuning_job": "acreate_fine_tuning_job", - "acancel_fine_tuning_job": "acancel_fine_tuning_job", - "alist_fine_tuning_jobs": "alist_fine_tuning_jobs", - "aretrieve_fine_tuning_job": "aretrieve_fine_tuning_job", - "aimage_edit": "aimage_edit", - "allm_passthrough_route": "allm_passthrough_route", - "acreate_skill": "acreate_skill", - "alist_skills": "alist_skills", - "aget_skill": "aget_skill", - "adelete_skill": "adelete_skill", - "agenerate_content": "google_genai.agenerate_content", - "agenerate_content_stream": "google_genai.agenerate_content_stream", - "aocr": "ocr.aocr", - "adecisions": "decisions.adecisions", - "asystemone": "systemone.asystemone", - "aextraction": "extraction.aextraction", - **{f"avideo_{name}": f"videos.avideo_{name}" for name in ("generation", "list", "status", "content", "remix")}, - **{ - name: f"containers.{name}" - for name in ("acreate_container", "alist_containers", "aretrieve_container", "adelete_container") - }, - **{ - name: f"containers.endpoint_factory._generated_endpoints.{name}" - for name in ( - "aupload_container_file", - "alist_container_files", - "aretrieve_container_file", - "adelete_container_file", - "aretrieve_container_file_content", - ) - }, - **{ - f"avector_store_{name}": f"vector_stores.main.a{name}" - for name in ("search", "retrieve", "list", "update", "delete") - }, - **{ - f"avector_store_file_{name}": f"vector_store_files.main.a{target}" - for name, target in ( - ("create", "create"), - ("list", "list"), - ("retrieve", "retrieve"), - ("content", "retrieve_content"), - ("update", "update"), - ("delete", "delete"), - ) - }, - "acreate_interaction": "interactions.acreate", - "aget_interaction": "interactions.aget", - "adelete_interaction": "interactions.adelete", - "acancel_interaction": "interactions.acancel", - } -) - -_SYNC_GENERIC: Final[Mapping[str, str]] = MappingProxyType( - { - "generate_content": "google_genai.generate_content", - "generate_content_stream": "google_genai.generate_content_stream", - "ocr": "ocr.ocr", - # The sync wrapper routes a search tool name as a model group, unlike `asearch`. - "search": "search.search", - "decisions": "decisions.decisions", - "systemone": "systemone.systemone", - "extraction": "extraction.extraction", - **{f"video_{name}": f"videos.video_{name}" for name in ("generation", "list", "status", "content", "remix")}, - **{ - name: f"containers.{name}" - for name in ("create_container", "list_containers", "retrieve_container", "delete_container") - }, - **{ - f"vector_store_{name}": f"vector_stores.main.{name}" - for name in ("search", "create", "retrieve", "list", "update", "delete") - }, - } -) - -_CONTAINER_FILE_ROUTED: Final = frozenset( - { - "aupload_container_file", - "alist_container_files", - "aretrieve_container_file", - "adelete_container_file", - "aretrieve_container_file_content", - } -) -_VIDEO_DIRECT: Final = frozenset({"avideo_create_character", "avideo_get_character", "avideo_edit", "avideo_extension"}) -_SYNC_VIDEO_UNBRANCHED: Final = frozenset( - {"video_create_character", "video_get_character", "video_edit", "video_extension"} -) -_SYNC_CONTAINER_FILES_UNBRANCHED: Final = frozenset( - { - "upload_container_file", - "list_container_files", - "retrieve_container_file", - "delete_container_file", - "retrieve_container_file_content", - } -) -_UNBRANCHED: Final = frozenset( - { - "aadapter_generate_content", - *_SYNC_VIDEO_UNBRANCHED, - *_SYNC_CONTAINER_FILES_UNBRANCHED, - "get_interaction", - "delete_interaction", - "cancel_interaction", - } -) -_ASSISTANTS: Final = frozenset( - { - "acreate_assistants", - "adelete_assistant", - "aget_assistants", - "acreate_thread", - "aget_thread", - "a_add_message", - "aget_messages", - "arun_thread", - } -) -_MANAGED_AGENTS: Final[Mapping[str, str]] = MappingProxyType( - { # comprehension-ok: a flat walk of nested rows - f"{prefix}{name}": f"interactions.agents.{prefix}{target}" - for prefix in ("a", "") - for name, target in ( - ("create_agent", "create"), - ("list_agents", "list"), - ("get_agent", "get"), - ("delete_agent", "delete"), - ("list_agent_versions", "list_versions"), - ) - } -) -_SYNC_VECTOR_STORE_FILES: Final[Mapping[str, str]] = MappingProxyType( - { - f"vector_store_file_{name}": f"vector_store_files.main.{target}" - for name, target in ( - ("create", "create"), - ("list", "list"), - ("retrieve", "retrieve"), - ("content", "retrieve_content"), - ("update", "update"), - ("delete", "delete"), - ) + "generic", + "generic_client", + "generic_provider_client", + "passthrough", + "response_id", + "container", + "vector_store", + "interaction", + "generic_sync", + "vector_store_sync", + "interaction_sync", } ) @@ -468,13 +325,10 @@ def plan_for(endpoint: str) -> Plan: def generic_plan(endpoint: str) -> Plan | None: if endpoint in PLANS: return PLANS[endpoint] - if endpoint in _HOST_GENERIC: - return _generic(endpoint, _HOST_GENERIC[endpoint]) - if endpoint in _VIDEO_DIRECT: - return _generic(endpoint, f"videos.{endpoint}") - if endpoint in _SYNC_GENERIC: - return _generic(endpoint, _SYNC_GENERIC[endpoint], sync=True) - return None + spec: Final = ENDPOINTS.get(endpoint) + if spec is None or spec.dispatch not in _GENERIC_DISPATCHES: + return None + return _generic(endpoint, spec.host_op, sync=spec.sync, call_type=spec.call_type) # --- preparing the request ------------------------------------------------------------------- @@ -636,7 +490,7 @@ def attempt_call( kwargs, cache=cache, function_name="generic_api_call", - splits_stream_timeout=invoked_handler_name(core, plan) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + splits_stream_timeout=plan.splits_stream_timeout, ) endpoint = kwargs.get("endpoint") # rebind-ok: rebound as the Python router rebinds it if endpoint and isinstance(endpoint, str): @@ -675,15 +529,14 @@ def _generic_attempt( from litellm.llms.base_llm.passthrough.transformation import replace_path_segment from litellm.router_rust_websocket import turn_admission - handler: Final = invoked_handler_name(core, plan) - socket: Final = handler == "_aresponses_websocket" + socket: Final = plan.call_type == "_aresponses_websocket" update_kwargs_with_deployment( core._shaping, # pyright: ignore[reportPrivateUsage] # the core's own settings deployment, kwargs, cache=core._cache, # pyright: ignore[reportPrivateUsage] # the core's routing cache function_name="_ageneric_api_call_with_fallbacks", - splits_stream_timeout=handler in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, + splits_stream_timeout=plan.splits_stream_timeout, opens_responses_socket=socket, ) model_name: Final = str(data.get("model")) @@ -715,9 +568,6 @@ def _generic_attempt( return AttemptCall(plan.host_op, final, kwargs, model_name, group) -_SEARCH_PROTECTED_PARAMS: Final = frozenset({"search_provider", "api_key", "api_base"}) - - def _search_attempt( core: "RustCore", plan: Plan, @@ -726,33 +576,8 @@ def _search_attempt( ) -> AttemptCall: """`SearchAPIRouter.async_search_with_fallbacks_helper`: one of the search tools named ``group`` at random, called with its provider, credentials and params under the request's.""" - import random - - tools: Final = [ - tool - for tool in core._endpoints.search_tools # pyright: ignore[reportPrivateUsage] # the core's live search tools - if as_dict(tool).get("search_tool_name") == group - ] - if not tools: - raise ValueError(f"Search tool '{group}' not found in router.search_tools") - selected: Final = random.choice(tools) # load balancing, as the Python router picks - params: Final = as_dict(as_dict(selected).get("litellm_params")) - search_provider: Final = params.get("search_provider") - if not search_provider: - raise ValueError(f"search_provider not found in litellm_params for search tool '{group}'") - search_params: Final = { - key: value - for source in (params, kwargs) - for key, value in source.items() - if key not in _SEARCH_PROTECTED_PARAMS and value is not None - } # comprehension-ok: the tool's params, then the request's - final: Final = { - "search_provider": search_provider, - "api_key": params.get("api_key"), - "api_base": params.get("api_base"), - **search_params, - } - return AttemptCall(plan.host_op, final, kwargs, group, group) + tools: Final = [as_dict(tool) for tool in core._endpoints.search_tools] # pyright: ignore[reportPrivateUsage] # the core's live search tools + return AttemptCall(plan.host_op, dict(search_tool_call(tools, group, kwargs)), kwargs, group, group) def unrouted_fields(plan: Plan) -> Mapping[str, object]: @@ -980,25 +805,6 @@ def start_silent_experiments( # --- the endpoints ---------------------------------------------------------------------------- -_KEEPS_CLIENT: Final = frozenset( - {"acreate_realtime_client_secret", "arealtime_calls", "acreate_realtime_transcription_session"} -) -_KEEPS_PROVIDER_AND_CLIENT: Final = frozenset({"afile_delete", "afile_content"}) -_RESPONSE_ID_ENDPOINTS: Final = frozenset( - {"aget_responses", "acancel_responses", "acompact_responses", "adelete_responses", "alist_input_items"} -) -_CONTAINER_ENDPOINTS: Final = frozenset( - {"acreate_container", "alist_containers", "aretrieve_container", "adelete_container", *_CONTAINER_FILE_ROUTED} -) -_VECTOR_STORE_ENDPOINTS: Final = frozenset( - { - *(f"avector_store_{name}" for name in ("search", "create", "retrieve", "list", "update", "delete")), - *(f"avector_store_file_{name}" for name in ("create", "list", "retrieve", "content", "update", "delete")), - } -) -_INTERACTION_ENDPOINTS: Final = frozenset( - {"acreate_interaction", "aget_interaction", "adelete_interaction", "acancel_interaction"} -) _MISSING: Final = object() @@ -1165,16 +971,6 @@ class RustEndpoints: def member(self, name: str) -> Callable[..., object] | None: """The public endpoint ``name``, bound to this core, or None when the core has none.""" - if name in _UNBRANCHED: - return _answers_none - if name in _ASSISTANTS: - return self._assistants(name) - if name in _MANAGED_AGENTS: - return self._managed_agents(name) - if name in _SYNC_VECTOR_STORE_FILES: - return self._sync_vector_store_files(name) - if name == "create_interaction": - return self._factory("create_interaction", "interactions.create") direct: Final = { "avector_store_create": self.avector_store_create, "acreate_file": self.acreate_file, @@ -1189,66 +985,125 @@ class RustEndpoints: "schedule_acompletion": self.schedule_acompletion, "priority_header_update": self.priority_header_update, "aguardrail": self.aguardrail, - "asearch": self.asearch, }.get(name) if direct is not None: return direct - if name == "responses": - return self._sync_factory(name, PLANS["responses"]) - plan: Final = generic_plan(name) - if plan is None or name in PLANS: - return None - if plan.sync: - return self._sync_factory(name, plan) - return self._factory(name, plan.host_op) + spec: Final = ENDPOINTS.get(name) + return None if spec is None else self._bind(spec) # -- the factory wrappers ------------------------------------------------------------------ - def _factory(self, name: str, host_op: str) -> Callable[..., Awaitable[object]]: + def _bind(self, spec: EndpointSpec) -> Callable[..., object]: + """The wrapper serving ``spec``, dispatched as ``router_endpoint_table`` says.""" + if spec.sync: + + def sync_wrapper( + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: the endpoint's request + ) -> object: + return self._dispatch_sync(spec, custom_llm_provider, client, dict(kwargs)) + + sync_wrapper.__name__ = spec.name + return sync_wrapper + async def async_wrapper( custom_llm_provider: str | None = None, client: object = None, **kwargs: object, # kwargs-ok: the endpoint's request ) -> object: - return await self._wrapped(name, host_op, custom_llm_provider, client, dict(kwargs)) + return await self._dispatch_async(spec, custom_llm_provider, client, dict(kwargs)) - async_wrapper.__name__ = name + async_wrapper.__name__ = spec.name return async_wrapper - async def _wrapped( + async def _dispatch_async( self, - name: str, - host_op: str, + spec: EndpointSpec, custom_llm_provider: str | None, client: object, kwargs: dict[str, object], # mutable-ok: the wrapper's own kwargs ) -> object: - plan: Final = generic_plan(name) or _generic(name, host_op) - if name in _CONTAINER_ENDPOINTS: - return await self._containers(plan, custom_llm_provider, kwargs) - if name in _RESPONSE_ID_ENDPOINTS: - from litellm.responses.utils import ResponsesAPIRequestUtils + match spec.dispatch: + case "assistants": + return await self._assistants(spec.host_op, custom_llm_provider, client, kwargs) + case "moderation": + return await self.amoderation(kwargs) + case "search": + return await self.asearch(**kwargs) + case "responses" | "anthropic_messages": + return await self.generic(plan_for(spec.name), kwargs) + case "managed_agent": + return await self._direct(spec.host_op, managed_agent_request(custom_llm_provider, kwargs)) + case _: + pass + plan: Final = plan_for(spec.name) + match spec.dispatch: + case "generic_client": + return await self._ageneric(plan, {**kwargs, "client": client}) + case "generic_provider_client": + return await self._ageneric( + plan, {**kwargs, "custom_llm_provider": custom_llm_provider, "client": client} + ) + case "passthrough": + return await self._ageneric(plan, {**kwargs, "client": client} if client else kwargs, passthrough=True) + case "response_id": + return await self._response_id(plan, custom_llm_provider, kwargs) + case "container": + return await self._containers(plan, custom_llm_provider, kwargs) + case "vector_store": + return await self._vector_stores(spec.name, plan, custom_llm_provider, kwargs) + case "interaction": + return await self._interactions(plan, custom_llm_provider, kwargs) + case _: + return await self._ageneric(plan, kwargs) - model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(kwargs.get("response_id")) # pyright: ignore[reportArgumentType] # any id the caller sends - if model_id is not None: - kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them - "model" - ] = model_id - return await self._ageneric(plan, kwargs) - if name in _VECTOR_STORE_ENDPOINTS: - return await self._vector_stores(name, plan, custom_llm_provider, kwargs) - if name in _INTERACTION_ENDPOINTS or name == "create_interaction": - return await self._interactions(plan, custom_llm_provider, kwargs) - if name in _KEEPS_CLIENT: - return await self._ageneric(plan, {**kwargs, "client": client}) - if name in _KEEPS_PROVIDER_AND_CLIENT: - return await self._ageneric(plan, {**kwargs, "custom_llm_provider": custom_llm_provider, "client": client}) - if name == "allm_passthrough_route": - if client: - kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them - "client" - ] = client - return await self._ageneric(plan, kwargs, passthrough=True) + def _dispatch_sync( + self, + spec: EndpointSpec, + custom_llm_provider: str | None, + client: object, + kwargs: dict[str, object], # mutable-ok: the wrapper's own kwargs + ) -> object: + match spec.dispatch: + case "search_sync": + return self.search(**kwargs) + case "vector_store_file_sync": + return self._invoke(spec.host_op)(custom_llm_provider=custom_llm_provider, client=client, **kwargs) + case "managed_agent_sync": + return self._invoke(spec.host_op)(**managed_agent_request(custom_llm_provider, kwargs)) + case "vector_store_sync": + return self._sync_vector_stores(spec, custom_llm_provider, kwargs) + case "interaction_sync": + request: Final = interaction_request(custom_llm_provider, kwargs) + if request.get("model"): + return self._core._call_sync(plan_for(spec.name), request, _SINGLE_ATTEMPT) # pyright: ignore[reportPrivateUsage] # the endpoints drive their own core + return self._invoke(spec.host_op)(**request) + case _: + if "model" not in kwargs: + raise TypeError( + "_generic_api_call_with_fallbacks() missing 1 required positional argument: 'model'" + ) + return self._core._call_sync(plan_for(spec.name), kwargs, _SINGLE_ATTEMPT) # pyright: ignore[reportPrivateUsage] # see above + + async def _response_id( + self, + plan: Plan, + custom_llm_provider: str | None, + kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place + ) -> object: + """`_init_responses_api_endpoints`: a LiteLLM response id pins the deployment it encodes; a + plain id with no model calls the handler directly, as litellm does without a router.""" + from litellm.responses.utils import ResponsesAPIRequestUtils + + model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(kwargs.get("response_id")) # pyright: ignore[reportArgumentType] # any id the caller sends + if model_id is not None: + kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + "model" + ] = model_id + if not kwargs.get("model"): + direct: Final = {key: value for key, value in kwargs.items() if key != "model"} + return await self._direct(plan.host_op, with_provider(custom_llm_provider, direct)) return await self._ageneric(plan, kwargs) async def generic(self, plan: Plan, request: Mapping[str, object]) -> object: @@ -1275,40 +1130,26 @@ class RustEndpoints: return await self._core._drive(plan, kwargs, passthrough=passthrough_call if passthrough else None) # pyright: ignore[reportPrivateUsage] # see above - def _sync_factory(self, name: str, plan: Plan) -> Callable[..., object]: - is_vector_store: Final = name.startswith("vector_store_") - - def sync_wrapper( - custom_llm_provider: str | None = None, - client: object = None, - **kwargs: object, # kwargs-ok: the endpoint's request - ) -> object: - if not is_vector_store: - if "model" not in kwargs: - raise TypeError( - "_generic_api_call_with_fallbacks() missing 1 required positional argument: 'model'" - ) - return self._core._call_sync(plan, kwargs, _SINGLE_ATTEMPT) # pyright: ignore[reportPrivateUsage] # see above - provider_kwargs: dict[ # mutable-ok: the request kwargs the Python router mutates in place # rebind-ok: rebound as the Python router rebinds it - str, object - ] = ( - {**kwargs, "custom_llm_provider": custom_llm_provider} - if custom_llm_provider and "custom_llm_provider" not in kwargs - else dict(kwargs) - ) - if name == "vector_store_search": - provider_kwargs = { # rebind-ok: rebound as the Python router rebinds it - **provider_kwargs, - "_direct_vector_store_embedding_executor": self._embedding_executor(kwargs), - } - if provider_kwargs.get("model"): - return self._core._call_sync(plan, provider_kwargs, _SINGLE_ATTEMPT) # pyright: ignore[reportPrivateUsage] # see above - if name == "vector_store_search": - return self._invoke(plan.host_op)(**{**provider_kwargs, "router": self._core}) - return self._invoke(plan.host_op)(**provider_kwargs) - - sync_wrapper.__name__ = name - return sync_wrapper + def _sync_vector_stores( + self, + spec: EndpointSpec, + custom_llm_provider: str | None, + kwargs: Mapping[str, object], + ) -> object: + """The sync vector store wrapper: routed with a model, else the handler directly.""" + provider_kwargs: dict[ # mutable-ok: the request kwargs the Python router mutates in place # rebind-ok: rebound as the Python router rebinds it + str, object + ] = with_provider(custom_llm_provider, kwargs) + if spec.name == "vector_store_search": + provider_kwargs = { # rebind-ok: rebound as the Python router rebinds it + **provider_kwargs, + "_direct_vector_store_embedding_executor": self._embedding_executor(kwargs), + } + if provider_kwargs.get("model"): + return self._core._call_sync(plan_for(spec.name), provider_kwargs, _SINGLE_ATTEMPT) # pyright: ignore[reportPrivateUsage] # the endpoints drive their own core + if spec.name == "vector_store_search": + return self._invoke(spec.host_op)(**{**provider_kwargs, "router": self._core}) + return self._invoke(spec.host_op)(**provider_kwargs) def _embedding_executor(self, kwargs: Mapping[str, object]) -> object: from litellm.llms.base_llm.vector_store.transformation import ( @@ -1388,92 +1229,31 @@ class RustEndpoints: kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place ) -> object: """`_init_interactions_api_endpoints`.""" - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them + request: Final = interaction_request(custom_llm_provider, kwargs) + if request.get("model"): + return await self._ageneric(plan, request) + return await self._direct(plan.host_op, request) + + async def _assistants( + self, + op: str, + custom_llm_provider: str | None, + client: object, + kwargs: Mapping[str, object], + ) -> object: + """`_pass_through_assistants_endpoint_factory`.""" + provider: object = custom_llm_provider # rebind-ok: the assistants config's when none is given + sent: Final = dict(kwargs) + if provider is None: + if self._assistants_config is None: + raise Exception( # noqa: TRY002 # the Python router raises a bare Exception + "'custom_llm_provider' must be set. Either via:\n `Router(assistants_config={'custom_llm_provider': ..})` \nor\n `router.arun_thread(custom_llm_provider=..)`" + ) + provider = self._assistants_config[ # rebind-ok: rebound as the Python router rebinds it "custom_llm_provider" - ] = custom_llm_provider - if "custom_llm_provider" not in kwargs: - kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them - "custom_llm_provider" - ] = "gemini" - if kwargs.get("agent") and kwargs.get("model") == kwargs.get("agent"): - kwargs["model"] = ( # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them - None - ) - if kwargs.get("model"): - return await self._ageneric(plan, kwargs) - return await self._direct(plan.host_op, kwargs) - - def _managed_agents(self, name: str) -> Callable[..., object]: - op: Final = _MANAGED_AGENTS[name] - - def prepared( - custom_llm_provider: str | None, - kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place - ) -> dict[str, object]: # mutable-ok: the call's kwargs - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them - "custom_llm_provider" - ] = custom_llm_provider - if "custom_llm_provider" not in kwargs: - kwargs[ # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them - "custom_llm_provider" - ] = "gemini" - return kwargs - - if name.startswith("a"): - - async def managed_agents_wrapper( - custom_llm_provider: str | None = None, - client: object = None, - **kwargs: object, # kwargs-ok: the request - ) -> object: - return await self._direct(op, prepared(custom_llm_provider, dict(kwargs))) - - return managed_agents_wrapper - - def managed_agents_sync_wrapper( - custom_llm_provider: str | None = None, - client: object = None, - **kwargs: object, # kwargs-ok: the request - ) -> object: - return self._invoke(op)(**prepared(custom_llm_provider, dict(kwargs))) - - return managed_agents_sync_wrapper - - def _sync_vector_store_files(self, name: str) -> Callable[..., object]: - op: Final = _SYNC_VECTOR_STORE_FILES[name] - - def vector_store_file_sync_wrapper( - custom_llm_provider: str | None = None, - client: object = None, - **kwargs: object, # kwargs-ok: the request - ) -> object: - return self._invoke(op)(custom_llm_provider=custom_llm_provider, client=client, **kwargs) - - return vector_store_file_sync_wrapper - - def _assistants(self, name: str) -> Callable[..., Awaitable[object]]: - async def assistants_wrapper( - custom_llm_provider: str | None = None, - client: object = None, - **kwargs: object, # kwargs-ok: the request - ) -> object: - """`_pass_through_assistants_endpoint_factory`.""" - provider: object = custom_llm_provider # rebind-ok: the assistants config's when none is given - sent: Final = dict(kwargs) - if provider is None: - if self._assistants_config is None: - raise Exception( # noqa: TRY002 # the Python router raises a bare Exception - "'custom_llm_provider' must be set. Either via:\n `Router(assistants_config={'custom_llm_provider': ..})` \nor\n `router.arun_thread(custom_llm_provider=..)`" - ) - provider = self._assistants_config[ # rebind-ok: rebound as the Python router rebinds it - "custom_llm_provider" - ] - sent.update(as_dict(self._assistants_config.get("litellm_params"))) - return await self._direct(name, {"custom_llm_provider": provider, "client": client, **sent}) - - return assistants_wrapper + ] + sent.update(as_dict(self._assistants_config.get("litellm_params"))) + return await self._direct(op, {"custom_llm_provider": provider, "client": client, **sent}) # -- completions, priority and prompt management ---------------------------------------- @@ -1786,7 +1566,11 @@ class RustEndpoints: async def avector_store_create(self, model: str | None, **kwargs: object) -> object: # kwargs-ok: the request if model is None: - return await self._factory("avector_store_create", "vector_stores.main.acreate")(**kwargs) + provider: Final = kwargs.pop("custom_llm_provider", None) + plan: Final = _generic("avector_store_create", "vector_stores.main.acreate") + return await self._vector_stores( + "avector_store_create", plan, provider if isinstance(provider, str) else None, dict(kwargs) + ) return await self._core._drive(PLANS["avector_store_create"], {**kwargs, "model": model}, _SINGLE_ATTEMPT) # pyright: ignore[reportPrivateUsage] # see above async def aretrieve_batch(self, model: str | None = None, **kwargs: object) -> object: # kwargs-ok: the request @@ -1963,7 +1747,7 @@ class RustEndpoints: async def asearch( self, - custom_llm_provider: str | None = None, + custom_llm_provider: object = None, client: object = None, **kwargs: object, # kwargs-ok: the request ) -> object: @@ -1976,6 +1760,27 @@ class RustEndpoints: request["model"] = search_tool_name return await self._core._drive(PLANS["asearch"], request) # pyright: ignore[reportPrivateUsage] # the endpoints drive their own core + def search( + self, + custom_llm_provider: object = None, + client: object = None, + **kwargs: object, # kwargs-ok: the request + ) -> object: + """The sync ``search``: one attempt on one of the router's search tools named + ``search_tool_name`` (else ``model``), as ``asearch`` picks among them.""" + search_tool_name: Final = kwargs.get("search_tool_name", kwargs.get("model")) + if not search_tool_name: + raise ValueError("search_tool_name or model parameter is required for search") + request: Final = {key: value for key, value in kwargs.items() if key != "model"} + prepared: Final = prepare_request( + PLANS["asearch"], + request, + str(search_tool_name), + self._core._resolved.model_group_alias, # pyright: ignore[reportPrivateUsage] # the endpoints serve their own core + ) + tools: Final = [as_dict(tool) for tool in self.search_tools] + return self._invoke(ENDPOINTS["search"].host_op)(**search_tool_call(tools, str(search_tool_name), prepared)) + # -- guardrails --------------------------------------------------------------------------- async def aguardrail( @@ -2067,11 +1872,6 @@ class RustEndpoints: _SINGLE_ATTEMPT: Final[Mapping[str, object]] = MappingProxyType({"num_retries": 0, "disable_fallbacks": True}) -async def _answers_none(*_args: object, **_kwargs: object) -> None: # kwargs-ok: any request - """An endpoint the Python router's async wrapper has no branch for.""" - return - - def _as_items(value: object) -> list[object]: # mutable-ok: a fresh list return list(value) if isinstance(value, list) else [] # pyright: ignore[reportUnknownArgumentType] # a config list diff --git a/litellm/router_rust_prompt_management.py b/litellm/router_rust_prompt_management.py index 7bf89477cb1..3fead8698cb 100644 --- a/litellm/router_rust_prompt_management.py +++ b/litellm/router_rust_prompt_management.py @@ -37,6 +37,8 @@ def _logging_and_kwargs( kwargs: dict[str, object], # mutable-ok: the request kwargs the Python router mutates in place ) -> tuple[object, dict[str, object]]: # mutable-ok: same """The request's logging object, else the one `function_setup` builds (with the kwargs it returns).""" + import uuid + from litellm.utils import Rules, function_setup, get_utc_datetime logging_obj: Final = kwargs.get("litellm_logging_obj") @@ -44,10 +46,11 @@ def _logging_and_kwargs( return logging_obj, kwargs return function_setup( # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers **{ # pyright: ignore[reportArgumentType] # the request's keyword payload + **kwargs, "original_function": "acompletion", "rules_obj": Rules(), "start_time": get_utc_datetime(), - **kwargs, + "litellm_call_id": kwargs.get("litellm_call_id") or str(uuid.uuid4()), } ) @@ -61,8 +64,10 @@ async def prompt_management_acompletion( """`_prompt_management_factory`, on the request as `acompletion` prepared it.""" from litellm.utils import get_non_default_completion_params - # `acompletion` hands the routing function along; function_setup consumes it. - logging_obj, kwargs = _logging_and_kwargs({**request, "original_function": core.acompletion}) + # The routing function is no request field: the routed call below drives its own plan. + logging_obj, kwargs = _logging_and_kwargs( + {key: value for key, value in request.items() if key != "original_function"} + ) caller_timeout: Final = kwargs.get("timeout") or kwargs.get("request_timeout") caller_stream_timeout: Final = kwargs.get("stream_timeout") specific_deployment: Final = kwargs.pop("specific_deployment", None) @@ -114,7 +119,6 @@ async def prompt_management_acompletion( final["prompt_label"] = prompt_label target_list: Final = core._hosted._core_get_model_list(model_name=str(rendered_model)) # pyright: ignore[reportPrivateUsage, reportUnknownArgumentType] # the core's host registry if target_list is None or len(target_list) == 0: # a direct call to the rendered model - _ = final.pop("original_function") return await core._host.invoke("acompletion")(**final) # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # the core's host timeouts: Final = { key: value diff --git a/litellm/router_utils/endpoint_eligibility.py b/litellm/router_utils/endpoint_eligibility.py index 4a87e582c97..8e21e10451c 100644 --- a/litellm/router_utils/endpoint_eligibility.py +++ b/litellm/router_utils/endpoint_eligibility.py @@ -23,6 +23,7 @@ CHAT_ENDPOINT_BY_CALL: Final[Mapping[str, str]] = MappingProxyType( "agenerate_content": "generateContent", "agenerate_content_stream": "streamGenerateContent", "acreate_interaction": "/v1beta/interactions", + "create_interaction": "/v1beta/interactions", "litellm.interactions.main.acreate": "/v1beta/interactions", "litellm.interactions.main.create": "/v1beta/interactions", "_aresponses_websocket": "/v1/responses", diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py index aa2b29988dc..0783189da85 100644 --- a/litellm/router_utils/search_api_router.py +++ b/litellm/router_utils/search_api_router.py @@ -5,16 +5,15 @@ Handles search tool selection, load balancing, and fallback logic for search req """ import asyncio -import random import traceback from collections.abc import Callable from functools import partial -from types import MappingProxyType from typing import TYPE_CHECKING, Final from litellm._logging import verbose_router_logger from litellm.router_api import ChangeSet from litellm.router_deprecated import renamed_router_param +from litellm.router_endpoint_table import search_tool_call from litellm.types.router import SearchToolLiteLLMParams, SearchToolTypedDict if TYPE_CHECKING: @@ -205,44 +204,13 @@ class SearchAPIRouter: search_tool_name: Final = model # model field contains the search_tool_name try: - # Find matching search tools - matching_tools: Final = SearchAPIRouter.get_matching_search_tools( - ctx=ctx, - search_tool_name=search_tool_name, - ) - - # Simple random selection for load balancing across multiple providers with same name - # For search tools, we use simple random choice since they don't have TPM/RPM constraints - selected_tool: Final = random.choice(matching_tools) - - # Extract search provider and other params from litellm_params - litellm_params: Final = selected_tool.get("litellm_params", {}) - search_provider: Final = litellm_params.get("search_provider") - if not search_provider: - raise ValueError(f"search_provider not found in litellm_params for search tool '{search_tool_name}'") - - api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials( - tool_litellm_params=litellm_params, - ) - protected_params: Final = frozenset(("search_provider", "api_key", "api_base")) - search_params: Final = MappingProxyType( - { - key: value - for params in (litellm_params, kwargs) - for key, value in params.items() - if key not in protected_params and value is not None - } - ) + call: Final = search_tool_call(ctx.settings().search_tools, search_tool_name, kwargs) + search_provider: Final = call["search_provider"] verbose_router_logger.debug("Selected search tool with provider: %s", search_provider) # Call the original search function with the provider config - response: Final = await original_generic_function( - search_provider=search_provider, - api_key=api_key, - api_base=api_base, - **search_params, - ) + response: Final = await original_generic_function(**call) return response diff --git a/tests/code_coverage_tests/router_surface_check.py b/tests/code_coverage_tests/router_surface_check.py index 52bdf100e07..189ed061b05 100644 --- a/tests/code_coverage_tests/router_surface_check.py +++ b/tests/code_coverage_tests/router_surface_check.py @@ -20,18 +20,21 @@ Member uses are found the same way the boundary survey found them: attribute access on a receiver whose name looks like a router (llm_router, litellm_router_instance, router, ...) or on the Router class, plus getattr/setattr/hasattr with a constant name on such a receiver. Members -include the container file endpoints Router.__init__ registers by name. +include the container file endpoints Router.__init__ registers by name and the +endpoints of litellm/router_endpoint_table.py, which the cores bind by name. Run: python tests/code_coverage_tests/router_surface_check.py [--write] --write regenerates the allowlist from the current tree (for shrinking it). """ import ast +import importlib.util import json import os import re import sys -from typing import Dict, List, Set +from collections.abc import Mapping +from typing import Dict, List, Set, cast ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) ALLOWLIST = os.path.join(os.path.dirname(__file__), "router_surface_allowlist.json") @@ -72,6 +75,8 @@ RECEIVER = re.compile( ) # Router.__init__ registers these with setattr(self, name, ...), so they never appear as self.. CONTAINER_FILE_ENDPOINTS = "litellm/containers/endpoints.json" +# Both cores bind the factory-built endpoints from this table with setattr, so they never appear as self.. +ENDPOINT_TABLE = "litellm/router_endpoint_table.py" DYNAMIC_FUNCS = {"getattr", "setattr", "hasattr"} @@ -86,6 +91,17 @@ def _is_test(rel: str) -> bool: return "/tests/" in rel or name.startswith("test_") or name.endswith("_test.py") +def _endpoint_table_names() -> set[str]: + """The factory-built endpoints both cores bind with setattr from litellm/router_endpoint_table.py + (a stdlib-only module, loaded from its file without importing litellm).""" + spec = importlib.util.spec_from_file_location("_router_endpoint_table", os.path.join(ROOT, ENDPOINT_TABLE)) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + endpoints = cast("Mapping[str, object]", module.ENDPOINTS) # cast-ok: the table maps endpoint names to rows + return set(endpoints) + + def router_members() -> Set[str]: members: Set[str] = set() for rel in ROUTER_CLASS_FILES: @@ -112,6 +128,7 @@ def router_members() -> Set[str]: with open(os.path.join(ROOT, CONTAINER_FILE_ENDPOINTS), encoding="utf-8") as f: for endpoint in json.load(f)["endpoints"]: members.update((endpoint["name"], endpoint["async_name"])) + members.update(_endpoint_table_names()) return {m for m in members if not (m.startswith("__") and m.endswith("__"))} diff --git a/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py b/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py index 810dfffaf8b..582649223db 100644 --- a/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py +++ b/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py @@ -1057,9 +1057,10 @@ def _router_with_advisor_deployment( ): """Build a Router whose only deployment is the advisor model on Foundry. - The recorder replaces ``litellm.anthropic_messages`` before construction - because Router binds it at init time, so the returned Router exercises the - real deployment-resolution path and records what it dispatched. + The router looks ``litellm.anthropic_messages`` up through its host at call + time, so each test patches it with the recorder around the call; the router + then exercises the real deployment-resolution path and records what it + dispatched. """ import litellm from litellm.router import Router @@ -1111,6 +1112,7 @@ async def test_advisor_sub_call_routes_through_proxy_router(): side_effect=mock_call, ), patch.object(proxy_server, "llm_router", router), + patch("litellm.anthropic_messages", new=recorder), ): h = AdvisorOrchestrationHandler() result = await h.handle( @@ -1172,6 +1174,7 @@ async def test_advisor_sub_call_routes_through_router_for_alias_and_wildcard(rou side_effect=mock_call, ), patch.object(proxy_server, "llm_router", router), + patch("litellm.anthropic_messages", new=recorder), ): h = AdvisorOrchestrationHandler() await h.handle( @@ -1223,6 +1226,7 @@ async def test_advisor_sub_call_bypasses_router_for_unconfigured_model(): side_effect=mock_call, ), patch.object(proxy_server, "llm_router", router), + patch("litellm.anthropic_messages", new=recorder), ): h = AdvisorOrchestrationHandler() await h.handle( @@ -1278,6 +1282,7 @@ async def test_advisor_sub_call_client_override_bypasses_router(): side_effect=mock_call, ), patch.object(proxy_server, "llm_router", router), + patch("litellm.anthropic_messages", new=recorder), patch.dict(proxy_server.general_settings, {"allow_client_side_credentials": True}), patch.object(litellm, "user_url_validation", False), ): diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 926566605a5..885f38fe39b 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -88,7 +88,9 @@ def test_router_vector_store_search_injects_executor_and_request_metadata(): ) as fallback: assert wrapped(model="vector-alias", vector_store_id="store", query="query") == "routed" assert fallback.call_args.kwargs["model"] == "vector-alias" - assert fallback.call_args.kwargs["original_function"] is original + handler = fallback.call_args.kwargs["original_function"] + assert handler.call_type == "vector_store_search" + assert handler.resolve() is original create_original = MagicMock(return_value="created") wrapped_create = router.factory_function(create_original, call_type="vector_store_create") @@ -98,7 +100,12 @@ def test_router_vector_store_search_injects_executor_and_request_metadata(): router, "_generic_api_call_with_fallbacks", return_value="created-through-router" ) as fallback: assert wrapped_create(model="vector-alias", name="store") == "created-through-router" - fallback.assert_called_once_with(original_function=create_original, model="vector-alias", name="store") + create_kwargs = fallback.call_args.kwargs + assert create_kwargs["original_function"].resolve() is create_original + assert {k: v for k, v in create_kwargs.items() if k != "original_function"} == { + "model": "vector-alias", + "name": "store", + } @pytest.mark.asyncio @@ -303,9 +310,7 @@ async def test_vector_store_file_list_resolves_credentials_from_model_query_para assert result["api_base"] == "https://api.openai.com/v1" assert result["model"] == "openai/gpt-4o-mini" assert "custom_llm_provider" not in result - llm_router.get_deployment_credentials_with_provider.assert_called_once_with( - model_id="team-openai" - ) + llm_router.get_deployment_credentials_with_provider.assert_called_once_with(model_id="team-openai") @pytest.mark.asyncio @@ -337,9 +342,7 @@ async def test_vector_store_file_list_registry_routed_model_skips_key_model_gran assert result["api_key"] == "sk-team-openai" assert result["model"] == "openai/gpt-4o-mini" - llm_router.get_deployment_credentials_with_provider.assert_called_once_with( - model_id="team-openai" - ) + llm_router.get_deployment_credentials_with_provider.assert_called_once_with(model_id="team-openai") @pytest.mark.asyncio @@ -370,9 +373,7 @@ async def test_vector_store_file_list_resolves_single_openai_team_deployment(): assert result["api_base"] == "https://api.openai.com/v1" assert result["model"] == "openai/gpt-4o-mini" assert "custom_llm_provider" not in result - llm_router.get_deployment_credentials_with_provider.assert_called_once_with( - model_id="team-openai", team_id=None - ) + llm_router.get_deployment_credentials_with_provider.assert_called_once_with(model_id="team-openai", team_id=None) @pytest.mark.asyncio @@ -668,9 +669,7 @@ async def test_update_request_data_with_litellm_managed_vector_store_registry(): } mock_registry = MagicMock() - mock_registry.get_litellm_managed_vector_store_from_registry.return_value = ( - mock_vector_store - ) + mock_registry.get_litellm_managed_vector_store_from_registry.return_value = mock_vector_store # Test with vector store registry with patch.object(litellm, "vector_store_registry", mock_registry): @@ -729,16 +728,13 @@ async def test_managed_vector_store_keeps_embedding_reference_and_explicit_confi assert managed_vector_store["litellm_params"]["litellm_embedding_config"] == explicit_config - class TestCheckVectorStorePermission: """Test suite for check_vector_store_permission function.""" def test_permission_allowed_in_key_metadata(self): """Test that permission is allowed when found in key metadata.""" key_metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["read", "write"]} - ] + "allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["read", "write"]}] } result = check_vector_store_permission( @@ -752,11 +748,7 @@ class TestCheckVectorStorePermission: def test_permission_allowed_in_team_metadata(self): """Test that permission is allowed when found in team metadata.""" - team_metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "team-index", "index_permissions": ["write"]} - ] - } + team_metadata = {"allowed_vector_store_indexes": [{"index_name": "team-index", "index_permissions": ["write"]}]} result = check_vector_store_permission( index_name="team-index", @@ -769,11 +761,7 @@ class TestCheckVectorStorePermission: def test_permission_denied_wrong_permission(self): """Test that permission is denied when index exists but wrong permission.""" - key_metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["read"]} - ] - } + key_metadata = {"allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["read"]}]} result = check_vector_store_permission( index_name="my-index", @@ -787,9 +775,7 @@ class TestCheckVectorStorePermission: def test_permission_denied_index_not_found(self): """Test that permission is denied when index doesn't exist.""" key_metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "other-index", "index_permissions": ["read", "write"]} - ] + "allowed_vector_store_indexes": [{"index_name": "other-index", "index_permissions": ["read", "write"]}] } result = check_vector_store_permission( @@ -827,16 +813,8 @@ class TestCheckVectorStorePermission: def test_key_metadata_takes_precedence(self): """Test that key metadata is checked and returns permission successfully.""" - key_metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["read"]} - ] - } - team_metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["write"]} - ] - } + key_metadata = {"allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["read"]}]} + team_metadata = {"allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["write"]}]} # Should find permission in key_metadata (checked first) result = check_vector_store_permission( @@ -850,16 +828,8 @@ class TestCheckVectorStorePermission: def test_team_metadata_as_fallback(self): """Test that team metadata is checked when key metadata doesn't have permission.""" - key_metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "other-index", "index_permissions": ["read"]} - ] - } - team_metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["write"]} - ] - } + key_metadata = {"allowed_vector_store_indexes": [{"index_name": "other-index", "index_permissions": ["read"]}]} + team_metadata = {"allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["write"]}]} # Should find permission in team_metadata result = check_vector_store_permission( @@ -882,28 +852,12 @@ class TestCheckVectorStorePermission: } # Test each index - assert ( - check_vector_store_permission("index-1", "read", key_metadata, None) is True - ) - assert ( - check_vector_store_permission("index-1", "write", key_metadata, None) - is False - ) - assert ( - check_vector_store_permission("index-2", "write", key_metadata, None) - is True - ) - assert ( - check_vector_store_permission("index-2", "read", key_metadata, None) - is False - ) - assert ( - check_vector_store_permission("index-3", "read", key_metadata, None) is True - ) - assert ( - check_vector_store_permission("index-3", "write", key_metadata, None) - is True - ) + assert check_vector_store_permission("index-1", "read", key_metadata, None) is True + assert check_vector_store_permission("index-1", "write", key_metadata, None) is False + assert check_vector_store_permission("index-2", "write", key_metadata, None) is True + assert check_vector_store_permission("index-2", "read", key_metadata, None) is False + assert check_vector_store_permission("index-3", "read", key_metadata, None) is True + assert check_vector_store_permission("index-3", "write", key_metadata, None) is True def test_invalid_metadata_structure(self): """Test handling of invalid metadata structures.""" @@ -972,9 +926,7 @@ class TestIsAllowedToCallVectorStoreEndpoint: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.user_role = None mock_user_api_key.metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["read"]} - ] + "allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["read"]}] } mock_user_api_key.team_metadata = None @@ -1009,9 +961,7 @@ class TestIsAllowedToCallVectorStoreEndpoint: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.user_role = None mock_user_api_key.metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["write"]} - ] + "allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["write"]}] } mock_user_api_key.team_metadata = None @@ -1046,9 +996,7 @@ class TestIsAllowedToCallVectorStoreEndpoint: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.user_role = None mock_user_api_key.metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["read"]} - ] + "allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["read"]}] } mock_user_api_key.team_metadata = None @@ -1107,9 +1055,7 @@ class TestIsAllowedToCallVectorStoreEndpoint: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.user_role = None mock_user_api_key.metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["read", "write"]} - ] + "allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["read", "write"]}] } mock_user_api_key.team_metadata = None @@ -1143,9 +1089,7 @@ class TestIsAllowedToCallVectorStoreEndpoint: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.user_role = None mock_user_api_key.metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["read", "write"]} - ] + "allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["read", "write"]}] } mock_user_api_key.team_metadata = None @@ -1207,9 +1151,7 @@ class TestIsAllowedToCallVectorStoreEndpoint: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.user_role = None mock_user_api_key.metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["read", "write"]} - ] + "allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["read", "write"]}] } mock_user_api_key.team_metadata = None @@ -1278,9 +1220,7 @@ class TestIsAllowedToCallVectorStoreEndpoint: mock_user_api_key.user_role = None mock_user_api_key.metadata = None mock_user_api_key.team_metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "team-index", "index_permissions": ["read"]} - ] + "allowed_vector_store_indexes": [{"index_name": "team-index", "index_permissions": ["read"]}] } # Mock provider config @@ -1351,11 +1291,7 @@ class TestIsAllowedToCallVectorStoreEndpoint: mock_user_api_key = UserAPIKeyAuth( token="sk-test-mock-token-404", key_name="sk-...plNQ", - metadata={ - "allowed_vector_store_indexes": [ - {"index_name": "dall-e-4", "index_permissions": ["write"]} - ] - }, + metadata={"allowed_vector_store_indexes": [{"index_name": "dall-e-4", "index_permissions": ["write"]}]}, spend=0.015, ) mock_provider_config = MagicMock() @@ -1427,12 +1363,8 @@ class TestIndexCreate: } mock_prisma = MagicMock() - mock_prisma.db.litellm_managedvectorstoreindextable.find_unique = AsyncMock( - return_value=None - ) - mock_prisma.db.litellm_managedvectorstoreindextable.create = AsyncMock( - return_value=mock_row - ) + mock_prisma.db.litellm_managedvectorstoreindextable.find_unique = AsyncMock(return_value=None) + mock_prisma.db.litellm_managedvectorstoreindextable.create = AsyncMock(return_value=mock_row) with patch( "litellm.proxy.proxy_server.prisma_client", @@ -1561,9 +1493,7 @@ class TestIsAllowedToCallVectorStoreFilesEndpoint: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.user_role = None mock_user_api_key.metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["read", "write"]} - ] + "allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["read", "write"]}] } mock_user_api_key.team_metadata = None @@ -1588,9 +1518,7 @@ class TestIsAllowedToCallVectorStoreFilesEndpoint: mock_user_api_key = MagicMock(spec=UserAPIKeyAuth) mock_user_api_key.user_role = None mock_user_api_key.metadata = { - "allowed_vector_store_indexes": [ - {"index_name": "my-index", "index_permissions": ["read"]} - ] + "allowed_vector_store_indexes": [{"index_name": "my-index", "index_permissions": ["read"]}] } mock_user_api_key.team_metadata = None @@ -1651,9 +1579,7 @@ class TestVectorStoreManagementEndpointsExist: # Force-register the lazy vector_store_management routes so the # assertions can find them. - already_registered = any( - getattr(r, "path", None) == "/vector_store/new" for r in app.routes - ) + already_registered = any(getattr(r, "path", None) == "/vector_store/new" for r in app.routes) if not already_registered: for feat in LAZY_FEATURES: if feat.name == "vector_store_management": @@ -1682,8 +1608,7 @@ class TestVectorStoreManagementEndpointsExist: # Verify each expected endpoint exists for method, path in expected_endpoints: assert (method, path) in app_routes, ( - f"Expected endpoint {method} {path} not found in registered routes. " - f"Available routes: {app_routes}" + f"Expected endpoint {method} {path} not found in registered routes. Available routes: {app_routes}" ) @@ -1764,27 +1689,15 @@ async def test_vector_store_synchronization_across_instances(): async def mock_delete(where): """Mock delete for removing vector store from DB""" vector_store_id = where.get("vector_store_id") - mock_db_vector_stores[:] = [ - vs - for vs in mock_db_vector_stores - if vs.get("vector_store_id") != vector_store_id - ] + mock_db_vector_stores[:] = [vs for vs in mock_db_vector_stores if vs.get("vector_store_id") != vector_store_id] return None # Create mock prisma client mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( - side_effect=mock_find_unique - ) - mock_prisma_client.db.litellm_managedvectorstorestable.find_many = AsyncMock( - side_effect=mock_find_many - ) - mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock( - side_effect=mock_create - ) - mock_prisma_client.db.litellm_managedvectorstorestable.delete = AsyncMock( - side_effect=mock_delete - ) + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(side_effect=mock_find_unique) + mock_prisma_client.db.litellm_managedvectorstorestable.find_many = AsyncMock(side_effect=mock_find_many) + mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock(side_effect=mock_create) + mock_prisma_client.db.litellm_managedvectorstorestable.delete = AsyncMock(side_effect=mock_delete) # Test vector store data test_vector_store_id = "test-sync-store-001" @@ -1804,18 +1717,13 @@ async def test_vector_store_synchronization_across_instances(): # Step 1: Create vector store on Instance 1 # (Simulate what happens in new_vector_store endpoint) - await mock_prisma_client.db.litellm_managedvectorstorestable.create( - data=test_vector_store - ) + await mock_prisma_client.db.litellm_managedvectorstorestable.create(data=test_vector_store) instance_1_registry.add_vector_store_to_registry(vector_store=test_vector_store) # Verify it's in Instance 1's memory - assert ( - instance_1_registry.get_litellm_managed_vector_store_from_registry( - test_vector_store_id - ) - is not None - ), "Vector store should be in Instance 1's memory" + assert instance_1_registry.get_litellm_managed_vector_store_from_registry(test_vector_store_id) is not None, ( + "Vector store should be in Instance 1's memory" + ) # Verify it's in the database db_store = await mock_prisma_client.db.litellm_managedvectorstorestable.find_unique( @@ -1825,41 +1733,28 @@ async def test_vector_store_synchronization_across_instances(): # Step 2: Instance 2 should be able to find it via database fallback # (Simulate what happens in pop_vector_stores_to_run_with_db_fallback) - found_store = ( - await instance_2_registry.get_litellm_managed_vector_store_from_registry_or_db( - vector_store_id=test_vector_store_id, prisma_client=mock_prisma_client - ) + found_store = await instance_2_registry.get_litellm_managed_vector_store_from_registry_or_db( + vector_store_id=test_vector_store_id, prisma_client=mock_prisma_client ) assert found_store is not None, "Instance 2 should find vector store from database" assert found_store.get("vector_store_id") == test_vector_store_id # Verify it's now cached in Instance 2's memory - assert ( - instance_2_registry.get_litellm_managed_vector_store_from_registry( - test_vector_store_id - ) - is not None - ), "Vector store should now be cached in Instance 2's memory" + assert instance_2_registry.get_litellm_managed_vector_store_from_registry(test_vector_store_id) is not None, ( + "Vector store should now be cached in Instance 2's memory" + ) # Step 3: Test that Instance 2 can list vector stores from database # (Simulate what happens in list_vector_stores endpoint - using DB as source of truth) - vector_stores_from_db = await VectorStoreRegistry.get_vector_stores_from_db( - prisma_client=mock_prisma_client - ) + vector_stores_from_db = await VectorStoreRegistry.get_vector_stores_from_db(prisma_client=mock_prisma_client) # Verify vector store appears in the database list vector_store_ids = [vs.get("vector_store_id") for vs in vector_stores_from_db] - assert ( - test_vector_store_id in vector_store_ids - ), "Instance 2 should see vector store from database" + assert test_vector_store_id in vector_store_ids, "Instance 2 should see vector store from database" # Verify the list endpoint logic: only show DB stores (filter out stale cache) # This simulates what list_vector_stores does - db_vector_store_ids = { - vs.get("vector_store_id") - for vs in vector_stores_from_db - if vs.get("vector_store_id") - } + db_vector_store_ids = {vs.get("vector_store_id") for vs in vector_stores_from_db if vs.get("vector_store_id")} # Instance 2's in-memory cache should only contain stores that exist in DB # (This is what the list endpoint cleanup does) @@ -1869,59 +1764,39 @@ async def test_vector_store_synchronization_across_instances(): instance_2_registry.delete_vector_store_from_registry(vector_store_id=vs_id) # After cleanup, instance 2 should still have the vector store (it's in DB) - assert ( - instance_2_registry.get_litellm_managed_vector_store_from_registry( - test_vector_store_id - ) - is not None - ), "Instance 2 should still have vector store (it exists in DB)" + assert instance_2_registry.get_litellm_managed_vector_store_from_registry(test_vector_store_id) is not None, ( + "Instance 2 should still have vector store (it exists in DB)" + ) # Step 4: Delete vector store on Instance 1 # (Simulate what happens in delete_vector_store endpoint) - await mock_prisma_client.db.litellm_managedvectorstorestable.delete( - where={"vector_store_id": test_vector_store_id} - ) - instance_1_registry.delete_vector_store_from_registry( - vector_store_id=test_vector_store_id - ) + await mock_prisma_client.db.litellm_managedvectorstorestable.delete(where={"vector_store_id": test_vector_store_id}) + instance_1_registry.delete_vector_store_from_registry(vector_store_id=test_vector_store_id) # Verify it's removed from Instance 1's memory - assert ( - instance_1_registry.get_litellm_managed_vector_store_from_registry( - test_vector_store_id - ) - is None - ), "Vector store should be removed from Instance 1's memory" + assert instance_1_registry.get_litellm_managed_vector_store_from_registry(test_vector_store_id) is None, ( + "Vector store should be removed from Instance 1's memory" + ) # Verify it's removed from database - db_store_after_delete = ( - await mock_prisma_client.db.litellm_managedvectorstorestable.find_unique( - where={"vector_store_id": test_vector_store_id} - ) + db_store_after_delete = await mock_prisma_client.db.litellm_managedvectorstorestable.find_unique( + where={"vector_store_id": test_vector_store_id} ) assert db_store_after_delete is None, "Vector store should be removed from database" # Step 5: Instance 2 should NOT show it in the list (database is source of truth) # The list endpoint logic should clean up stale cache entries - vector_stores_from_db_after_delete = ( - await VectorStoreRegistry.get_vector_stores_from_db( - prisma_client=mock_prisma_client - ) + vector_stores_from_db_after_delete = await VectorStoreRegistry.get_vector_stores_from_db( + prisma_client=mock_prisma_client ) # Verify vector store does NOT appear in the database list - vector_store_ids_after_delete = [ - vs.get("vector_store_id") for vs in vector_stores_from_db_after_delete - ] - assert ( - test_vector_store_id not in vector_store_ids_after_delete - ), "Deleted vector store should not be in database" + vector_store_ids_after_delete = [vs.get("vector_store_id") for vs in vector_stores_from_db_after_delete] + assert test_vector_store_id not in vector_store_ids_after_delete, "Deleted vector store should not be in database" # Simulate list endpoint cleanup logic db_vector_store_ids_after_delete = { - vs.get("vector_store_id") - for vs in vector_stores_from_db_after_delete - if vs.get("vector_store_id") + vs.get("vector_store_id") for vs in vector_stores_from_db_after_delete if vs.get("vector_store_id") } # Remove any in-memory vector stores that no longer exist in database @@ -1931,27 +1806,20 @@ async def test_vector_store_synchronization_across_instances(): instance_2_registry.delete_vector_store_from_registry(vector_store_id=vs_id) # Verify it was removed from Instance 2's cache - assert ( - instance_2_registry.get_litellm_managed_vector_store_from_registry( - test_vector_store_id - ) - is None - ), "Deleted vector store should be removed from Instance 2's cache" + assert instance_2_registry.get_litellm_managed_vector_store_from_registry(test_vector_store_id) is None, ( + "Deleted vector store should be removed from Instance 2's cache" + ) # Step 6: Test that using a deleted vector store fails gracefully # (Simulate what happens in pop_vector_stores_to_run_with_db_fallback) non_default_params = {"vector_store_ids": [test_vector_store_id]} - vector_stores_to_run = ( - await instance_2_registry.pop_vector_stores_to_run_with_db_fallback( - non_default_params=non_default_params, - tools=None, - prisma_client=mock_prisma_client, - ) + vector_stores_to_run = await instance_2_registry.pop_vector_stores_to_run_with_db_fallback( + non_default_params=non_default_params, + tools=None, + prisma_client=mock_prisma_client, ) - assert ( - len(vector_stores_to_run) == 0 - ), "Deleted vector store should not be returned when trying to use it" + assert len(vector_stores_to_run) == 0, "Deleted vector store should not be returned when trying to use it" @pytest.mark.asyncio @@ -2020,15 +1888,9 @@ async def test_vector_store_update_and_list_synchronization(): # Create mock prisma client mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_managedvectorstorestable.find_many = AsyncMock( - side_effect=mock_find_many - ) - mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock( - side_effect=mock_create - ) - mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock( - side_effect=mock_update - ) + mock_prisma_client.db.litellm_managedvectorstorestable.find_many = AsyncMock(side_effect=mock_find_many) + mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock(side_effect=mock_create) + mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock(side_effect=mock_update) # Test vector store data test_vector_store_id = "test-update-store-001" @@ -2050,26 +1912,18 @@ async def test_vector_store_update_and_list_synchronization(): } # Step 1: Create vector store on Instance 1 - await mock_prisma_client.db.litellm_managedvectorstorestable.create( - data=test_vector_store - ) + await mock_prisma_client.db.litellm_managedvectorstorestable.create(data=test_vector_store) instance_1_registry.add_vector_store_to_registry(vector_store=test_vector_store) # Step 2: Instance 2 fetches and caches the vector store - vector_stores_from_db = await VectorStoreRegistry.get_vector_stores_from_db( - prisma_client=mock_prisma_client - ) + vector_stores_from_db = await VectorStoreRegistry.get_vector_stores_from_db(prisma_client=mock_prisma_client) for vs in vector_stores_from_db: if vs.get("vector_store_id") == test_vector_store_id: instance_2_registry.add_vector_store_to_registry(vector_store=vs) # Verify both instances have the original data - instance_1_vs = instance_1_registry.get_litellm_managed_vector_store_from_registry( - test_vector_store_id - ) - instance_2_vs = instance_2_registry.get_litellm_managed_vector_store_from_registry( - test_vector_store_id - ) + instance_1_vs = instance_1_registry.get_litellm_managed_vector_store_from_registry(test_vector_store_id) + instance_2_vs = instance_2_registry.get_litellm_managed_vector_store_from_registry(test_vector_store_id) assert instance_1_vs.get("vector_store_name") == original_name assert instance_2_vs.get("vector_store_name") == original_name @@ -2088,29 +1942,21 @@ async def test_vector_store_update_and_list_synchronization(): ) # Verify Instance 1 has the updated data - instance_1_vs_after_update = ( - instance_1_registry.get_litellm_managed_vector_store_from_registry( - test_vector_store_id - ) + instance_1_vs_after_update = instance_1_registry.get_litellm_managed_vector_store_from_registry( + test_vector_store_id ) assert instance_1_vs_after_update.get("vector_store_name") == updated_name # Verify Instance 2 still has stale data in cache - instance_2_vs_before_list = ( - instance_2_registry.get_litellm_managed_vector_store_from_registry( - test_vector_store_id - ) + instance_2_vs_before_list = instance_2_registry.get_litellm_managed_vector_store_from_registry(test_vector_store_id) + assert instance_2_vs_before_list.get("vector_store_name") == original_name, ( + "Instance 2 should still have stale cached data before list operation" ) - assert ( - instance_2_vs_before_list.get("vector_store_name") == original_name - ), "Instance 2 should still have stale cached data before list operation" # Step 4: Instance 2 calls list endpoint (which should sync with database) # This simulates what list_vector_stores endpoint does - vector_stores_from_db_after_update = ( - await VectorStoreRegistry.get_vector_stores_from_db( - prisma_client=mock_prisma_client - ) + vector_stores_from_db_after_update = await VectorStoreRegistry.get_vector_stores_from_db( + prisma_client=mock_prisma_client ) # Build map from database vector stores (database is source of truth) @@ -2126,22 +1972,18 @@ async def test_vector_store_update_and_list_synchronization(): ) # Step 5: Verify Instance 2 now has the updated data - instance_2_vs_after_list = ( - instance_2_registry.get_litellm_managed_vector_store_from_registry( - test_vector_store_id - ) + instance_2_vs_after_list = instance_2_registry.get_litellm_managed_vector_store_from_registry(test_vector_store_id) + assert instance_2_vs_after_list.get("vector_store_name") == updated_name, ( + "Instance 2 should have updated data after list operation syncs with database" ) - assert ( - instance_2_vs_after_list.get("vector_store_name") == updated_name - ), "Instance 2 should have updated data after list operation syncs with database" # Verify the list returned the correct data combined_vector_stores = list(vector_store_map.values()) assert len(combined_vector_stores) == 1 assert combined_vector_stores[0].get("vector_store_id") == test_vector_store_id - assert ( - combined_vector_stores[0].get("vector_store_name") == updated_name - ), "List should return updated data from database" + assert combined_vector_stores[0].get("vector_store_name") == updated_name, ( + "List should return updated data from database" + ) @pytest.mark.asyncio @@ -2185,9 +2027,7 @@ async def test_new_vector_store_persists_embedding_reference_without_credentials } return mock_created_vector_store - mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock( - side_effect=mock_create - ) + mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock(side_effect=mock_create) mock_registry = MagicMock() mock_registry.add_vector_store_to_registry = MagicMock() @@ -2277,9 +2117,7 @@ async def test_new_vector_store_auto_resolves_from_router(): patch("litellm.proxy.proxy_server.llm_router", mock_router), patch.object(litellm, "vector_store_registry", mock_registry), ): - result = await new_vector_store( - vector_store=vector_store_data, user_api_key_dict=mock_user_api_key - ) + result = await new_vector_store(vector_store=vector_store_data, user_api_key_dict=mock_user_api_key) assert result["status"] == "success" # Resolution against the router happens at request-handling time now, @@ -2292,9 +2130,7 @@ async def test_new_vector_store_auto_resolves_from_router(): assert litellm_params_dict["litellm_embedding_model"] == "config-embedding-model" response_vs = result["vector_store"] - assert "router-resolved-api-key" not in _serialize_litellm_params( - response_vs.get("litellm_params") - ) + assert "router-resolved-api-key" not in _serialize_litellm_params(response_vs.get("litellm_params")) def _stub_user_api_key( @@ -2406,9 +2242,7 @@ async def test_create_vector_store_in_db(): mock_created_vector_store = MagicMock() mock_created_vector_store.model_dump.return_value = created_vector_store_data - mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock( - return_value=mock_created_vector_store - ) + mock_prisma_client.db.litellm_managedvectorstorestable.create = AsyncMock(return_value=mock_created_vector_store) mock_registry = MagicMock() mock_registry.add_vector_store_to_registry = MagicMock() @@ -2441,9 +2275,7 @@ async def test_create_vector_store_in_db(): mock_registry.add_vector_store_to_registry.assert_called_once() # Verify that create was called with correct data structure - create_call_args = ( - mock_prisma_client.db.litellm_managedvectorstorestable.create.call_args - ) + create_call_args = mock_prisma_client.db.litellm_managedvectorstorestable.create.call_args create_data = create_call_args.kwargs.get("data", {}) assert create_data["vector_store_id"] == vector_store_id assert create_data["custom_llm_provider"] == custom_llm_provider @@ -2462,9 +2294,7 @@ async def test_create_vector_store_in_db_raises_when_exists(): # Mock that vector store already exists existing_vector_store = MagicMock() - mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( - return_value=existing_vector_store - ) + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=existing_vector_store) with pytest.raises(HTTPException) as exc_info: await create_vector_store_in_db( @@ -2513,16 +2343,12 @@ class TestRedactSensitiveLitellmParams: "api_key": "sk-real-openai-key-12345", "aws_access_key_id": "AKIAIOSFODNN7EXAMPLE", "aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", - "vertex_credentials": ( - '{"type":"service_account","private_key":"-----BEGIN PRIVATE KEY-----..."}' - ), + "vertex_credentials": ('{"type":"service_account","private_key":"-----BEGIN PRIVATE KEY-----..."}'), "azure_authorization_token": "Bearer eyJhbGciOi...", } out = _redact_sensitive_litellm_params(params) for k in params: - assert ( - out[k] == REDACTED_BY_LITELM_STRING - ), f"{k} should be redacted, got {out[k]!r}" + assert out[k] == REDACTED_BY_LITELM_STRING, f"{k} should be redacted, got {out[k]!r}" def test_preserves_non_sensitive_keys(self): from litellm.proxy.vector_store_endpoints.management_endpoints import ( @@ -2663,9 +2489,7 @@ class TestRedactSensitiveLitellmParams: _redact_sensitive_litellm_params, ) - out = _redact_sensitive_litellm_params( - "this is not json but might contain a secret" - ) + out = _redact_sensitive_litellm_params("this is not json but might contain a secret") assert out == REDACTED_BY_LITELM_STRING @@ -2698,9 +2522,7 @@ class TestUpdateVectorStoreAccessControlAndRedaction: ) mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( - return_value=existing_row - ) + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=existing_row) with ( patch( @@ -2720,9 +2542,7 @@ class TestUpdateVectorStoreAccessControlAndRedaction: vector_store_id="vs_other_team", vector_store_description="hijacked", ), - user_api_key_dict=UserAPIKeyAuth( - user_id="attacker", team_id="team-B" - ), + user_api_key_dict=UserAPIKeyAuth(user_id="attacker", team_id="team-B"), ) assert exc_info.value.status_code == 403 # The attacker must NOT see the existing credential in the @@ -2767,12 +2587,8 @@ class TestUpdateVectorStoreAccessControlAndRedaction: ) mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( - return_value=existing_row - ) - mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock( - return_value=updated_row - ) + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock(return_value=updated_row) with ( patch( @@ -2813,17 +2629,11 @@ class TestUpdateVectorStoreAccessControlAndRedaction: from litellm.types.vector_stores import VectorStoreUpdateRequest existing_row = MagicMock() - existing_row.model_dump = MagicMock( - return_value={"vector_store_id": "vs_owned", "team_id": "team-A"} - ) + existing_row.model_dump = MagicMock(return_value={"vector_store_id": "vs_owned", "team_id": "team-A"}) mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock( - return_value=existing_row - ) - mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock( - return_value=None - ) + mock_prisma_client.db.litellm_managedvectorstorestable.find_unique = AsyncMock(return_value=existing_row) + mock_prisma_client.db.litellm_managedvectorstorestable.update = AsyncMock(return_value=None) with ( patch( # test-quality-ok: stubs the auth gate so the test exercises the not-found branch under test @@ -2835,8 +2645,12 @@ class TestUpdateVectorStoreAccessControlAndRedaction: new_callable=AsyncMock, return_value=True, ), - patch("litellm.proxy.proxy_server.prisma_client", mock_prisma_client), # test-quality-ok: proxy_server module global is the endpoint's only injection point - patch("litellm.vector_store_registry", None), # test-quality-ok: litellm module global is the only injection point for the registry + patch( + "litellm.proxy.proxy_server.prisma_client", mock_prisma_client + ), # test-quality-ok: proxy_server module global is the endpoint's only injection point + patch( + "litellm.vector_store_registry", None + ), # test-quality-ok: litellm module global is the only injection point for the registry ): with pytest.raises(HTTPException) as exc_info: await update_vector_store( @@ -3099,7 +2913,9 @@ def test_vector_store_search_rejects_caller_opensearch_index_name_only_on_opense app.dependency_overrides[user_api_key_auth] = lambda: mock_auth try: with ( - patch("litellm.vector_stores.main.asearch", new=fake_search), # test-quality-ok: the search boundary under test + patch( + "litellm.vector_stores.main.asearch", new=fake_search + ), # test-quality-ok: the search boundary under test patch.object(litellm, "vector_store_registry", registry), # test-quality-ok: seeds the store under test patch( # test-quality-ok: store access is not under test, so the request reaches the search boundary "litellm.proxy.vector_store_endpoints.utils.can_user_access_vector_store", @@ -3107,7 +2923,9 @@ def test_vector_store_search_rejects_caller_opensearch_index_name_only_on_opense ), ): router = litellm.Router(model_list=[]) - with patch("litellm.proxy.proxy_server.llm_router", router): # test-quality-ok: router built over the fake search + with patch( + "litellm.proxy.proxy_server.llm_router", router + ): # test-quality-ok: router built over the fake search response = TestClient(app).post( "/v1/vector_stores/kb/search", json={"query": "q", "opensearch_index_name": "another-teams-index"}, @@ -3162,7 +2980,9 @@ class TestConfigOwnedVectorStores: registry = self._registry() prisma = MagicMock() - prisma.db.litellm_managedvectorstorestable.find_many = AsyncMock(return_value=[self._db_row(self.DB_ID, "db-store")]) + prisma.db.litellm_managedvectorstorestable.find_many = AsyncMock( + return_value=[self._db_row(self.DB_ID, "db-store")] + ) with ( patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: proxy_server global, no seam @@ -3171,7 +2991,10 @@ class TestConfigOwnedVectorStores: first = await list_vector_stores(user_api_key_dict=self._admin()) second = await list_vector_stores(user_api_key_dict=self._admin()) - assert [(vs["vector_store_id"], vs["is_config"]) for vs in first["data"]] == [(self.DB_ID, False), (self.CONFIG_ID, True)] + assert [(vs["vector_store_id"], vs["is_config"]) for vs in first["data"]] == [ + (self.DB_ID, False), + (self.CONFIG_ID, True), + ] assert second["data"] == first["data"] assert [vs["vector_store_id"] for vs in registry.vector_stores] == [self.CONFIG_ID, self.DB_ID] @@ -3193,10 +3016,16 @@ class TestConfigOwnedVectorStores: by_id = {vs["vector_store_id"]: vs for vs in response["data"]} assert set(by_id) == {self.CONFIG_ID, self.DB_ID}, response - assert (by_id[self.CONFIG_ID]["vector_store_name"], by_id[self.CONFIG_ID]["is_config"]) == ("config-store", True) + assert (by_id[self.CONFIG_ID]["vector_store_name"], by_id[self.CONFIG_ID]["is_config"]) == ( + "config-store", + True, + ) assert (by_id[self.DB_ID]["vector_store_name"], by_id[self.DB_ID]["is_config"]) == ("db-store", False) assert [vs["vector_store_id"] for vs in registry.vector_stores] == [self.CONFIG_ID, self.DB_ID] - assert registry.get_litellm_managed_vector_store_from_registry(self.CONFIG_ID)["vector_store_name"] == "config-store" + assert ( + registry.get_litellm_managed_vector_store_from_registry(self.CONFIG_ID)["vector_store_name"] + == "config-store" + ) @pytest.mark.asyncio async def test_info_reports_config_ownership(self): @@ -3204,7 +3033,9 @@ class TestConfigOwnedVectorStores: from litellm.types.vector_stores import VectorStoreInfoRequest with ( - patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: proxy_server global, no seam + patch( + "litellm.proxy.proxy_server.prisma_client", MagicMock() + ), # test-quality-ok: proxy_server global, no seam patch.object(litellm, "vector_store_registry", self._registry()), ): config_info = await get_vector_store_info( @@ -3261,7 +3092,10 @@ class TestConfigOwnedVectorStores: assert exc_info.value.status_code == 400, exc_info.value.detail assert exc_info.value.detail["vector_store_id"] == self.CONFIG_ID prisma.db.litellm_managedvectorstorestable.update.assert_not_called() - assert registry.get_litellm_managed_vector_store_from_registry(self.CONFIG_ID)["vector_store_name"] == "config-store" + assert ( + registry.get_litellm_managed_vector_store_from_registry(self.CONFIG_ID)["vector_store_name"] + == "config-store" + ) @pytest.mark.asyncio async def test_delete_of_config_store_is_rejected_and_store_stays_registered(self): @@ -3308,7 +3142,9 @@ class TestConfigOwnedVectorStores: ) assert response["status"] == "success", response - prisma.db.litellm_managedvectorstorestable.delete.assert_awaited_once_with(where={"vector_store_id": self.DB_ID}) + prisma.db.litellm_managedvectorstorestable.delete.assert_awaited_once_with( + where={"vector_store_id": self.DB_ID} + ) assert registry.get_litellm_managed_vector_store_from_registry(self.DB_ID) is None diff --git a/tests/unit/test_router/core/endpoints_helpers.py b/tests/unit/test_router/core/endpoints_helpers.py index ac629352e95..720ac5c912e 100644 --- a/tests/unit/test_router/core/endpoints_helpers.py +++ b/tests/unit/test_router/core/endpoints_helpers.py @@ -5,12 +5,11 @@ Helpers for the endpoint characterization tests (``test_endpoints_characterizati litellm (assistants, skills, realtime, fine-tuning, batches, ...), so every ``host.invoke`` the core makes is answered and recorded. -Some endpoints do not go through ``host.invoke``: the Python core binds litellm's own functions for -them when it is built (vector stores, Google GenAI, OCR, search, videos, containers, interactions, -managed agents). ``patch_direct_endpoints`` replaces those litellm functions, before the core is -built, with recording stand-ins that answer from the same host script, under the same ``__name__`` -(the core routes on the handler's name). A call to one of them is recorded with the function's name -as its ``op``. +Some endpoints are served by a host op named by its dotted path under ``litellm`` (vector stores, +Google GenAI, OCR, search, videos, containers, interactions, managed agents), which ``host.invoke`` +resolves from litellm's modules at call time. ``patch_direct_endpoints`` replaces those litellm +functions with recording stand-ins that answer from the same host script. A call to one of them is +recorded with the function's name as its ``op``. """ import importlib @@ -180,8 +179,8 @@ RECORDED_NAMES: Final[Mapping[tuple[str, str], str]] = { def patch_direct_endpoints(monkeypatch: pytest.MonkeyPatch, host: ScriptedHost) -> None: """Replace litellm's directly bound endpoint functions with recording stand-ins on ``host``. - The stand-in keeps the function's ``__name__`` (the core routes on it); the call is recorded - under ``RECORDED_NAMES`` when the bare name is ambiguous, else under the function's name.""" + The call is recorded under ``RECORDED_NAMES`` when the bare name is ambiguous, else under the + function's name.""" for module_name, attributes in DIRECT_ENDPOINTS.items(): module: Final = importlib.import_module(module_name) for attribute, function_name in attributes.items(): diff --git a/tests/unit/test_router/core/fallbacks_helpers.py b/tests/unit/test_router/core/fallbacks_helpers.py index c88c31a9371..6278b647820 100644 --- a/tests/unit/test_router/core/fallbacks_helpers.py +++ b/tests/unit/test_router/core/fallbacks_helpers.py @@ -4,7 +4,6 @@ from collections.abc import Callable, Mapping from typing import Final from litellm.router_core import RouterCore -from litellm.router_python_host import Endpoint from tests.unit.test_router.core.scripted_host import Outcome, ScriptedHost, deployment MakeCore = Callable[..., RouterCore] @@ -64,20 +63,6 @@ def fails_once_then_answers(error: BaseException) -> Callable[..., Outcome]: return outcome -class NamedEndpointHost(ScriptedHost): - """A ``ScriptedHost`` whose ``anthropic_messages`` endpoint carries that name, as litellm's - own function does. The core recognises a /v1/messages call (and its safeguard-refusal gate) - by the endpoint's ``__name__``; ``ScriptedHost``'s endpoints are all named ``endpoint``.""" - - __slots__ = () - - def invoke(self, op: str) -> Endpoint: - endpoint: Final = super().invoke(op) - if op == "anthropic_messages": - endpoint.__name__ = op # pyright: ignore[reportFunctionMemberAccess] # a plain function - return endpoint - - def anthropic_refusal(*, with_stop_details: bool = True) -> dict[str, object]: # mutable-ok: a provider response response: Final[dict[str, object]] = { "id": "msg_1", diff --git a/tests/unit/test_router/core/test_endpoints_characterization_batch.py b/tests/unit/test_router/core/test_endpoints_characterization_batch.py index ab3543d7ad9..80fb0b7be2a 100644 --- a/tests/unit/test_router/core/test_endpoints_characterization_batch.py +++ b/tests/unit/test_router/core/test_endpoints_characterization_batch.py @@ -464,22 +464,33 @@ async def test_prompt_management_group_with_non_dict_prompt_variables_raises_val await core.acompletion(model="pm", messages=MESSAGES, litellm_call_id="c1") -async def test_prompt_management_group_without_a_call_id_or_logging_object_raises_key_error( +async def test_prompt_management_group_without_a_call_id_or_logging_object_builds_its_own( make_core: MakeCore, host: EndpointHost ): + """With neither, the core builds the logging object under a fresh call id and goes on to + validate the prompt, as it does when the caller sends a call id.""" core: Final = make_core(model_list=[deployment("pm", "p1", model="langfuse/gpt-4o")]) - with pytest.raises(KeyError, match="litellm_call_id"): + with pytest.raises(ValueError, match=re.escape("Prompt ID is not set or not a string. Got=None")): await core.acompletion(model="pm", messages=MESSAGES) - - -async def test_litellm_agent_group_without_a_logging_object_raises_key_error(make_core: MakeCore, host: EndpointHost): - """function_setup consumes original_function, so the direct-call path cannot pop it.""" - core: Final = make_core(model_list=[deployment("pm", "p1", model="litellm_agent/openai/gpt-4o")]) - with pytest.raises(KeyError, match="original_function"): - await core.acompletion(model="pm", messages=MESSAGES, litellm_call_id="c1") assert host.calls == [] +async def test_litellm_agent_group_without_a_call_id_or_logging_object_calls_the_stripped_model_directly( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("pm", "p1", model="litellm_agent/openai/gpt-4o")]) + await core.acompletion(model="pm", messages=MESSAGES) + call: Final = host.calls[0] + assert (call.op, call.model, call.deployment_id, "original_function" in call.kwargs) == ( + "acompletion", + "openai/gpt-4o", + "p1", + False, + ) + sent_call_id: Final = call.kwargs["litellm_call_id"] + assert isinstance(sent_call_id, str) and sent_call_id + + async def test_litellm_agent_group_calls_the_stripped_model_directly_with_the_deployment_params( make_core: MakeCore, host: EndpointHost ): diff --git a/tests/unit/test_router/core/test_endpoints_characterization_generic.py b/tests/unit/test_router/core/test_endpoints_characterization_generic.py index f0f2db3c102..ea4559841d2 100644 --- a/tests/unit/test_router/core/test_endpoints_characterization_generic.py +++ b/tests/unit/test_router/core/test_endpoints_characterization_generic.py @@ -5,8 +5,8 @@ agents, realtime, fine-tuning, files, image edit, LLM pass-through, vector store guardrails): which host op each reaches, the kwargs the core builds, and how an id in the request pins the deployment. -The Python core binds litellm's own functions for some of these endpoints instead of going through -``host.invoke``; the ``host`` fixture replaces those functions with recording stand-ins (see +Some of these endpoints are served by host ops named by a dotted path under ``litellm``; the +``host`` fixture replaces those litellm functions with recording stand-ins (see ``endpoints_helpers.py``), so every call shows up in ``host.calls``. """ @@ -208,39 +208,58 @@ async def test_async_generic_endpoint_is_retried_with_the_router_num_retries(mak @pytest.mark.parametrize( - "endpoint", + ("endpoint", "op"), [ - "aadapter_generate_content", - "video_create_character", - "video_get_character", - "video_edit", - "video_extension", - "upload_container_file", - "list_container_files", - "retrieve_container_file", - "delete_container_file", - "retrieve_container_file_content", - "get_interaction", - "delete_interaction", - "cancel_interaction", + ("aadapter_generate_content", "aadapter_generate_content"), + ("video_create_character", "video_create_character"), + ("video_get_character", "video_get_character"), + ("video_edit", "video_edit"), + ("video_extension", "video_extension"), + ("upload_container_file", "upload_container_file"), + ("list_container_files", "list_container_files"), + ("retrieve_container_file", "retrieve_container_file"), + ("delete_container_file", "delete_container_file"), + ("retrieve_container_file_content", "retrieve_container_file_content"), + ("get_interaction", "interactions.get"), + ("delete_interaction", "interactions.delete"), + ("cancel_interaction", "interactions.cancel"), ], ) -async def test_endpoints_without_a_dispatch_branch_resolve_to_none_without_calling_anything( - make_core: MakeCore, host: EndpointHost, endpoint: str +async def test_endpoints_that_had_no_dispatch_branch_route_like_their_siblings( + make_core: MakeCore, host: EndpointHost, endpoint: str, op: str ): - """These are wrapped by factory_function's async wrapper, which has no branch for them.""" + """The async one routes as the other async generic endpoints do; the sync ones answer + synchronously, routed as the other sync generic endpoints are.""" core: Final = make_core(model_list=[deployment("g", "d1")]) result: Final = getattr(core, endpoint)(model="g", container_id="c", video_id="v", interaction_id="i") - assert inspect.isawaitable(result) - assert await result is None - assert host.calls == [] + assert inspect.isawaitable(result) is endpoint.startswith("a") + response: Final = await result if inspect.isawaitable(result) else result + assert response is not None + assert [(call.op, call.deployment_id) for call in host.calls] == [(op, "d1")] -async def test_sync_create_interaction_is_awaited_and_fails_with_type_error(make_core: MakeCore, host: EndpointHost): +async def test_sync_create_interaction_answers_synchronously_with_the_routed_response( + make_core: MakeCore, host: EndpointHost +): core: Final = make_core(model_list=[deployment("g", "d1")], num_retries=0) - with pytest.raises(TypeError, match="can't be used in 'await' expression"): - await core.create_interaction(model="g", input="x") - assert [call.op for call in host.calls] == ["interactions.create"] + response: Final[object] = getattr(core, "create_interaction")(model="g", input="x") + assert not inspect.isawaitable(response) + assert response is not None + assert [(call.op, call.deployment_id) for call in host.calls] == [("interactions.create", "d1")] + + +async def test_sync_create_interaction_with_an_agent_calls_the_handler_directly_with_gemini( + make_core: MakeCore, host: EndpointHost +): + core: Final = make_core(model_list=[deployment("ag", "d1")]) + getattr(core, "create_interaction")(agent="ag", model="ag", input="x") + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["model"], call.kwargs["custom_llm_provider"]) == ( + "interactions.create", + None, + None, + "gemini", + ) # --- request params, provider and timeouts ------------------------------------------------------ @@ -288,15 +307,25 @@ async def test_generic_timeout_takes_the_router_timeout(make_core: MakeCore, hos @pytest.mark.parametrize( - ("endpoint", "request_body"), [("aresponses", {"input": "hi"}), ("agenerate_content", {"contents": []})] + ("endpoint", "request_body", "expected"), + [ + # A stream-timeout-aware handler (by its endpoint, whatever the host function is named) + # gets the request timeout and the stream timeout apart. + ("aresponses", {"input": "hi"}, (11.0, 3.0)), + ("agenerate_content", {"contents": []}, (3.0, 3.0)), + ], ) -async def test_streaming_generic_call_sends_the_deployment_stream_timeout_as_timeout( - make_core: MakeCore, host: EndpointHost, endpoint: str, request_body: dict[str, object] +async def test_streaming_generic_call_sends_the_deployment_stream_timeout( + make_core: MakeCore, + host: EndpointHost, + endpoint: str, + request_body: dict[str, object], + expected: tuple[float, float], ): core: Final = make_core(model_list=[deployment("g", "d1", timeout=11, stream_timeout=3)]) await _call(core, endpoint, {"model": "g", "stream": True, **request_body}) sent: Final = host.calls[0].kwargs - assert (sent["timeout"], sent["stream_timeout"]) == (3.0, 3.0) + assert (sent["timeout"], sent["stream_timeout"]) == expected async def test_non_streaming_aresponses_sends_the_deployment_timeout(make_core: MakeCore, host: EndpointHost): @@ -354,11 +383,19 @@ async def test_the_response_id_reaches_the_host_still_encoded(make_core: MakeCor assert host.calls[0].kwargs["response_id"] == ENCODED_RESPONSE_ID -async def test_a_plain_response_id_without_a_model_fails_with_type_error(make_core: MakeCore, host: EndpointHost): +async def test_a_plain_response_id_without_a_model_calls_the_handler_directly(make_core: MakeCore, host: EndpointHost): + """As litellm's own function does without a router: no deployment, the given provider.""" core: Final = make_core(model_list=TWO_DEPLOYMENTS) - with pytest.raises(TypeError, match="missing 1 required positional argument: 'model'"): - await core.aget_responses(response_id="resp_plain") - assert host.calls == [] + await _call(core, "aget_responses", {"response_id": "resp_plain", "custom_llm_provider": "azure"}) + call: Final = host.calls[0] + assert (call.op, call.deployment_id, call.kwargs["response_id"], call.kwargs["custom_llm_provider"]) == ( + "aget_responses", + None, + "resp_plain", + "azure", + ) + assert "model" not in call.kwargs + assert "litellm_metadata" not in call.kwargs # --- container ids ------------------------------------------------------------------------ @@ -730,6 +767,28 @@ async def test_asearch_calls_the_named_search_tool_with_its_litellm_params(make_ assert response._hidden_params["additional_headers"]["x-litellm-model-group"] == "s" # pyright: ignore[reportAttributeAccessIssue] # the scripted response +async def test_sync_search_calls_the_named_search_tool_like_asearch(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], search_tools=SEARCH_TOOLS) + response: Final[object] = getattr(core, "search")(model="s", query="q") + assert not inspect.isawaitable(response) + call: Final = host.calls[0] + assert (call.op, call.kwargs["search_provider"], call.kwargs["api_key"], call.kwargs["query"]) == ( + "search", + "tavily", + "tk", + "q", + ) + assert _bucket(call.kwargs, "litellm_metadata")["model_group"] == "s" + assert "model" not in call.kwargs + + +async def test_sync_search_with_an_unknown_tool_raises_value_error(make_core: MakeCore, host: EndpointHost): + core: Final = make_core(model_list=[deployment("g", "d1")], search_tools=SEARCH_TOOLS) + with pytest.raises(ValueError, match=re.escape("Search tool 'g' not found in router.search_tools")): + getattr(core, "search")(search_tool_name="g", query="q") + assert host.calls == [] + + async def test_asearch_with_an_unknown_tool_raises_value_error(make_core: MakeCore, host: EndpointHost): core: Final = make_core(model_list=[deployment("g", "d1")], search_tools=SEARCH_TOOLS, num_retries=0) with pytest.raises(ValueError, match=re.escape("Search tool 'g' not found in router.search_tools")): diff --git a/tests/unit/test_router/core/test_endpoints_characterization_inference.py b/tests/unit/test_router/core/test_endpoints_characterization_inference.py index 8db0d6bea12..44b90331dc1 100644 --- a/tests/unit/test_router/core/test_endpoints_characterization_inference.py +++ b/tests/unit/test_router/core/test_endpoints_characterization_inference.py @@ -106,12 +106,11 @@ async def test_amoderation_is_not_retried(make_core: MakeCore, host: EndpointHos assert host.sleeps == [] -async def test_arerank_sends_the_builtin_input_function_as_input(make_core: MakeCore, host: EndpointHost): - """arerank sets kwargs["input"] = input, the Python builtin, and it reaches the host.""" +async def test_arerank_sends_only_the_request_fields_and_no_input(make_core: MakeCore, host: EndpointHost): core: Final = make_core(model_list=[deployment("g", "d1")]) await core.arerank(model="g", query="q", documents=["a"]) sent: Final = _sent(host) - assert (sent["input"], sent["query"], sent["documents"]) == (input, "q", ["a"]) + assert (sent["query"], sent["documents"], "input" in sent) == ("q", ["a"], False) @pytest.mark.parametrize(("request_voice", "expected"), [(None, "deployment-voice"), ("alloy", "alloy")]) diff --git a/tests/unit/test_router/core/test_endpoints_dispatch_table.py b/tests/unit/test_router/core/test_endpoints_dispatch_table.py new file mode 100644 index 00000000000..362c75d1c98 --- /dev/null +++ b/tests/unit/test_router/core/test_endpoints_dispatch_table.py @@ -0,0 +1,117 @@ +""" +The factory-built endpoints of ``litellm/router_endpoint_table.py`` on each core: every row is +served, answers sync or async as its dispatch says, and reaches the host through ``host.invoke`` +with the row's host op. The host here names every function it serves ``renamed``, so the routing +decisions that depend on the endpoint (stream-timeout split, video-id pin, the responses socket's +turn admission) are shown to key on the endpoint, never on the host function's ``__name__``. +""" + +import inspect +from collections.abc import Callable, Mapping +from typing import Final + +import pytest + +from litellm.router_core import RouterCore +from litellm.router_endpoint_table import DISPATCHES, ENDPOINTS, EndpointSpec, dispatch_for +from litellm.types.videos.utils import encode_video_id_with_provider +from tests.unit.test_router.core.endpoints_helpers import recording_endpoint +from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment + +MakeCore = Callable[..., RouterCore] +HOST_OPS: Final = frozenset(spec.host_op for spec in ENDPOINTS.values()) +SYNC_HOST_OPS: Final = frozenset(spec.host_op for spec in ENDPOINTS.values() if spec.sync) +SEARCH_TOOLS: Final = [{"search_tool_name": "s", "litellm_params": {"search_provider": "tavily", "api_key": "tk"}}] + + +class RenamingHost(ScriptedHost): + """Serves every host op of the endpoint table with a recording function named ``renamed``, + recorded under the op it was invoked as.""" + + __slots__ = () + + def invoke(self, op: str) -> Callable[..., object]: + if op not in HOST_OPS: + return super().invoke(op) + endpoint: Final = recording_endpoint(self, op, is_async=op not in SYNC_HOST_OPS) + endpoint.__name__ = "renamed" + endpoint.__qualname__ = "renamed" + return endpoint + + +@pytest.fixture +def host() -> RenamingHost: + return RenamingHost() + + +async def _settled(result: object) -> object: + return await result if inspect.isawaitable(result) else result + + +def _request(spec: EndpointSpec) -> Mapping[str, object]: + match spec.dispatch: + case "assistants": + return {"custom_llm_provider": "openai"} + case "search" | "search_sync": + return {"model": "s", "query": "q"} + case "managed_agent" | "managed_agent_sync": + return {"name": "a"} + case "vector_store_file_sync": + return {"vector_store_id": "v"} + case _: + return { + "model": "g", + "input": "x", + "video_id": "v", + "container_id": "cntr_native", + "response_id": "resp_plain", + "interaction_id": "i", + } + + +@pytest.mark.parametrize("name", sorted(ENDPOINTS)) +async def test_every_table_endpoint_reaches_the_host_with_its_op(make_core: MakeCore, host: RenamingHost, name: str): + spec: Final = ENDPOINTS[name] + core: Final = make_core(model_list=[deployment("g", "d1")], search_tools=SEARCH_TOOLS) + result: Final[object] = getattr(core, name)(**_request(spec)) + assert inspect.isawaitable(result) is not spec.sync + assert await _settled(result) is not None + assert [call.op for call in host.calls] == [spec.host_op] + + +def test_every_table_row_has_a_known_dispatch_and_every_call_type_one_dispatch(): + assert {spec.dispatch for spec in ENDPOINTS.values()} <= DISPATCHES + assert all(dispatch_for(spec.call_type) == spec.dispatch for spec in ENDPOINTS.values()) + with pytest.raises(ValueError, match="no router endpoint dispatches the call type 'nope'"): + dispatch_for("nope") + + +def test_every_table_endpoint_is_declared_on_the_python_core_for_type_checkers(): + from litellm.router_python_core import PythonCore, TableEndpoint + + declared: Final = {name for name, kind in PythonCore.__annotations__.items() if kind == TableEndpoint} + assert declared == set(ENDPOINTS) + + +async def test_a_renamed_host_function_keeps_the_stream_timeout_split(make_core: MakeCore, host: RenamingHost): + core: Final = make_core(model_list=[deployment("g", "d1", timeout=11, stream_timeout=3)]) + await getattr(core, "aresponses")(model="g", input="hi", stream=True) + sent: Final = host.calls[0].kwargs + assert (sent["timeout"], sent["stream_timeout"]) == (11.0, 3.0) + + +@pytest.mark.parametrize("endpoint", ["avideo_status", "video_status"]) +async def test_a_renamed_host_function_keeps_the_video_id_pin(make_core: MakeCore, host: RenamingHost, endpoint: str): + video_id: Final = encode_video_id_with_provider("video_up", "openai", model_id="d2", deployment_id="d2") + core: Final = make_core(model_list=[deployment("g", "d1"), deployment("g", "d2")]) + for _ in range(6): + await _settled(getattr(core, endpoint)(model="g", video_id=video_id)) + assert host.deployments_called() == ["d2"] * 6 + + +async def test_a_renamed_host_function_keeps_the_responses_socket_turn_admission( + make_core: MakeCore, host: RenamingHost +): + core: Final = make_core(model_list=[deployment("g", "d1")]) + await getattr(core, "_aresponses_websocket")(model="g") + assert callable(host.calls[0].kwargs["turn_admission"]) diff --git a/tests/unit/test_router/core/test_fallbacks_characterization_refusal.py b/tests/unit/test_router/core/test_fallbacks_characterization_refusal.py index d7c893a36f1..7fe1533ce65 100644 --- a/tests/unit/test_router/core/test_fallbacks_characterization_refusal.py +++ b/tests/unit/test_router/core/test_fallbacks_characterization_refusal.py @@ -3,8 +3,8 @@ Characterization of a router core's Anthropic safeguard-refusal fallbacks on /v1 (non-streaming): a ``stop_reason: "refusal"`` response carrying ``stop_details`` re-enters the fallback chain only when a fallback would resolve for it. -The core recognises a /v1/messages call by its endpoint's name, so this module runs on -``NamedEndpointHost`` (a ``ScriptedHost`` whose ``anthropic_messages`` endpoint is named so). +The core recognises a /v1/messages call by the endpoint it serves, never by the host function's +``__name__``: ``ScriptedHost``'s endpoints are all named ``endpoint``, and the refusal gate holds. """ import pytest @@ -13,16 +13,16 @@ from litellm.router_core import RouterCore from tests.unit.test_router.core.fallbacks_helpers import ( MESSAGES, MakeCore, - NamedEndpointHost, anthropic_message, anthropic_refusal, build, ) +from tests.unit.test_router.core.scripted_host import ScriptedHost @pytest.fixture -def host() -> NamedEndpointHost: - return NamedEndpointHost() +def host() -> ScriptedHost: + return ScriptedHost() async def messages(core: RouterCore, **request: object) -> object: # kwargs-ok: request kwargs @@ -40,7 +40,7 @@ async def messages(core: RouterCore, **request: object) -> object: # kwargs-ok: ids=["content-policy-chain", "general-chain", "default-fallbacks"], ) async def test_a_safeguard_refusal_falls_back_when_a_chain_resolves( - make_core: MakeCore, host: NamedEndpointHost, arguments: dict[str, object], expected: list[str] + make_core: MakeCore, host: ScriptedHost, arguments: dict[str, object], expected: list[str] ) -> None: core = build(make_core, **arguments) host.answer("openai/a", anthropic_refusal()) @@ -71,7 +71,7 @@ async def test_a_safeguard_refusal_falls_back_when_a_chain_resolves( ) async def test_a_refusal_is_returned_to_the_caller_when_no_fallback_would_resolve( make_core: MakeCore, - host: NamedEndpointHost, + host: ScriptedHost, arguments: dict[str, object], request_kwargs: dict[str, object], refusal: dict[str, object], @@ -89,7 +89,7 @@ async def test_a_refusal_is_returned_to_the_caller_when_no_fallback_would_resolv @pytest.mark.asyncio async def test_the_last_refusal_is_returned_when_every_fallback_refuses( - make_core: MakeCore, host: NamedEndpointHost + make_core: MakeCore, host: ScriptedHost ) -> None: core = build(make_core, fallbacks=[{"a": ["b"]}]) host.answer("*", anthropic_refusal()) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 9996ffcb098..84d0a58b0e2 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -26481,7 +26481,8 @@ class TestRouterEmbeddingIntegration: assert wrapped(model="vector-alias", vector_store_id="store", query="query") == "routed" assert fallback.call_args.kwargs["model"] == "vector-alias" - assert fallback.call_args.kwargs["original_function"] is original + handler = fallback.call_args.kwargs["original_function"] + assert (handler.resolve(), handler.call_type) == (original, "vector_store_search") assert isinstance( fallback.call_args.kwargs["_direct_vector_store_embedding_executor"], RouterVectorStoreEmbeddingExecutor, diff --git a/tests/unit/test_router/test_router_endpoints.py b/tests/unit/test_router/test_router_endpoints.py index 0c57fb662ac..b5d5753b40a 100644 --- a/tests/unit/test_router/test_router_endpoints.py +++ b/tests/unit/test_router/test_router_endpoints.py @@ -510,15 +510,15 @@ async def test_init_responses_api_endpoints(): return_value=None ) - # Call without a response_id (no model extraction should happen) + # Call without a response_id or a model: nothing names a deployment, so the handler is called + # directly (as litellm's own function is without a router), never routed + direct_handler = AsyncMock() await router._init_responses_api_endpoints( - original_function=AsyncMock(), thread_id="thread_xyz" + original_function=direct_handler, thread_id="thread_xyz" ) - # Verify _ageneric_api_call_with_fallbacks was called but model wasn't changed - first_call_kwargs = router._ageneric_api_call_with_fallbacks.call_args.kwargs - assert "model" not in first_call_kwargs - assert first_call_kwargs["thread_id"] == "thread_xyz" + direct_handler.assert_awaited_once_with(thread_id="thread_xyz") + router._ageneric_api_call_with_fallbacks.assert_not_called() # Reset the mock router._ageneric_api_call_with_fallbacks.reset_mock()