diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 9cb0a00d9fc..00695cbfb5b 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -351,9 +351,9 @@ def filter_exceptions_from_params(data: Any, max_depth: int = 20) -> Any: # Skip callable objects (functions, methods, lambdas) but not classes (type objects) if callable(data) and not isinstance(data, type): return None - # Skip known non-serializable object types (Logging, etc.) + # Skip known non-serializable object types (Logging, Router, etc.) obj_type_name = type(data).__name__ - if obj_type_name in ["Logging", "LiteLLMLoggingObj"]: + if obj_type_name in ["Logging", "LiteLLMLoggingObj", "Router"]: return None if isinstance(data, dict): diff --git a/litellm/rag/main.py b/litellm/rag/main.py index 0ccbb435e53..620097f83dd 100644 --- a/litellm/rag/main.py +++ b/litellm/rag/main.py @@ -198,6 +198,10 @@ async def _execute_query_pipeline( """ Execute the RAG query pipeline. """ + # Extract router from kwargs - use it for completion if available + # to properly resolve virtual model names + router: Optional["Router"] = kwargs.pop("router", None) + # 1. Extract query from last user message query_text = RAGQuery.extract_query_from_messages(messages) if not query_text: @@ -233,12 +237,21 @@ async def _execute_query_pipeline( context_message = RAGQuery.build_context_message(context_chunks) modified_messages = messages[:-1] + [context_message] + [messages[-1]] - response = await litellm.acompletion( - model=model, - messages=modified_messages, - stream=stream, - **kwargs, - ) + # Use router if available to properly resolve virtual model names + if router is not None: + response = await router.acompletion( + model=model, + messages=modified_messages, + stream=stream, + **kwargs, + ) + else: + response = await litellm.acompletion( + model=model, + messages=modified_messages, + stream=stream, + **kwargs, + ) # 5. Attach search results to response if not stream and isinstance(response, ModelResponse):