fix(rag): use router for completion in RAG query pipeline (#19550)

The RAG query endpoint was failing with "Object of type Router is not
JSON serializable" when called through the proxy. This was caused by two
issues:

1. The Router object passed via kwargs was leaking into the request
   payload sent to providers like Bedrock, causing JSON serialization
   errors.

2. The RAG query pipeline was calling litellm.acompletion() directly
   instead of using the router, so virtual model names configured in the
   proxy weren't being resolved to actual provider model IDs.

This fix:
- Extracts the router from kwargs and uses router.acompletion() when
  available, falling back to litellm.acompletion() otherwise
- Adds "Router" to the list of non-serializable types in
  filter_exceptions_from_params as a defensive measure

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Marcos Griselli 2026-01-24 01:11:17 -03:00 committed by GitHub
parent ac0ab214fb
commit 6b1ce4e766
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 21 additions and 8 deletions

View file

@ -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):

View file

@ -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):