Merge pull request #15456 from BerriAI/litellm_staging_branch_10_11_2025_p1

Litellm staging branch 10 11 2025 p1
This commit is contained in:
Krish Dholakia 2025-10-12 22:01:45 -07:00 committed by GitHub
commit ff20f8402a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
32 changed files with 2642 additions and 674 deletions

2
.gitignore vendored
View file

@ -97,3 +97,5 @@ litellm_config.yaml
.vscode/launch.json
litellm/proxy/to_delete_loadtest_work/*
update_model_cost_map.py
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py
litellm/proxy/_experimental/out/guardrails/index.html

View file

@ -131,7 +131,9 @@ class ProxyExtrasDBManager:
)
@staticmethod
def _resolve_all_migrations(migrations_dir: str, schema_path: str):
def _resolve_all_migrations(
migrations_dir: str, schema_path: str, mark_all_applied: bool = True
):
"""
1. Compare the current database state to schema.prisma and generate a migration for the diff.
2. Run prisma migrate deploy to apply any pending migrations.
@ -210,6 +212,8 @@ class ProxyExtrasDBManager:
logger.warning("Migration diff application timed out.")
# 3. Mark all migrations as applied
if not mark_all_applied:
return
migration_names = ProxyExtrasDBManager._get_migration_names(migrations_dir)
logger.info(f"Resolving {len(migration_names)} migrations")
for migration_name in migration_names:
@ -263,6 +267,13 @@ class ProxyExtrasDBManager:
logger.info(f"prisma migrate deploy stdout: {result.stdout}")
logger.info("prisma migrate deploy completed")
# Run sanity check to ensure DB matches schema
logger.info("Running post-migration sanity check...")
ProxyExtrasDBManager._resolve_all_migrations(
migrations_dir, schema_path, mark_all_applied=False
)
logger.info("✅ Post-migration sanity check completed")
return True
except subprocess.CalledProcessError as e:
logger.info(f"prisma db error: {e.stderr}, e: {e.stdout}")

View file

@ -57,22 +57,22 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
) -> Tuple[Optional[Any], int]:
"""
Handle raw dict response items from Responses API (e.g., GPT-5 Codex format).
Args:
item: Raw dict response item with 'type' field
index: Current choice index
Returns:
Tuple of (Choice object or None, updated index)
"""
from litellm.types.utils import Choices, Message
item_type = item.get("type")
# Ignore reasoning items for now
if item_type == "reasoning":
return None, index
# Handle message items with output_text content
if item_type == "message":
content_list = item.get("content", [])
@ -83,13 +83,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
response_text = content_item.get("text", "")
msg = Message(
role=item.get("role", "assistant"),
content=response_text if response_text else ""
)
choice = Choices(
message=msg, finish_reason="stop", index=index
content=response_text if response_text else "",
)
choice = Choices(message=msg, finish_reason="stop", index=index)
return choice, index + 1
# Unknown or unsupported type
return None, index
@ -294,8 +292,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
if isinstance(item, ResponseReasoningItem):
for content in item.summary:
response_text = getattr(content, "text", "")
for summary_item in item.summary:
response_text = getattr(summary_item, "text", "")
reasoning_content = response_text if response_text else ""
elif isinstance(item, ResponseOutputMessage):
@ -340,7 +338,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge):
index += 1
elif isinstance(item, dict):
# Handle raw dict responses (e.g., from GPT-5 Codex)
choice, index = self._handle_raw_dict_response_item(item=item, index=index)
choice, index = self._handle_raw_dict_response_item(
item=item, index=index
)
if choice is not None:
choices.append(choice)
else:

View file

@ -86,8 +86,15 @@ class MCPClient:
async def connect(self):
"""Initialize the transport and session."""
if self._session:
verbose_logger.debug(
f"MCP client already connected to {self.server_url or 'stdio'}"
)
return # Already connected
verbose_logger.info(
f"MCP client connecting to {self.server_url or 'stdio'} via {self.transport_type}"
)
try:
if self.transport_type == MCPTransport.stdio:
# For stdio transport, use stdio_client with command-line parameters
@ -107,6 +114,9 @@ class MCPClient:
)
self._session = await self._session_ctx.__aenter__()
await self._session.initialize()
verbose_logger.info(
f"MCP client successfully connected via stdio: {self.stdio_config.get('command', '')}"
)
elif self.transport_type == MCPTransport.sse:
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
@ -122,6 +132,9 @@ class MCPClient:
)
self._session = await self._session_ctx.__aenter__()
await self._session.initialize()
verbose_logger.info(
f"MCP client successfully connected via SSE to {self.server_url}"
)
else: # http
headers = self._get_auth_headers()
httpx_client_factory = self._create_httpx_client_factory()
@ -140,6 +153,9 @@ class MCPClient:
)
self._session = await self._session_ctx.__aenter__()
await self._session.initialize()
verbose_logger.info(
f"MCP client successfully connected via HTTP to {self.server_url}"
)
except ValueError as e:
# Re-raise ValueError exceptions (like missing stdio_config)
verbose_logger.warning(f"MCP client connection failed: {str(e)}")
@ -159,7 +175,12 @@ class MCPClient:
async def disconnect(self):
"""Clean up session and connections."""
verbose_logger.info(
f"MCP client disconnecting from {self.server_url or 'stdio'}"
)
if self._task and not self._task.done():
verbose_logger.debug("MCP client cancelling background task")
self._task.cancel()
try:
await self._task
@ -168,16 +189,24 @@ class MCPClient:
if self._session:
try:
verbose_logger.debug("MCP client closing session")
await self._session_ctx.__aexit__(None, None, None) # type: ignore
except Exception:
except Exception as e:
verbose_logger.debug(
f"Error closing MCP session: {type(e).__name__}: {str(e)}"
)
pass
self._session = None
self._session_ctx = None
if self._transport_ctx:
try:
verbose_logger.debug("MCP client closing transport")
await self._transport_ctx.__aexit__(None, None, None)
except Exception:
except Exception as e:
verbose_logger.debug(
f"Error closing MCP transport: {type(e).__name__}: {str(e)}"
)
pass
self._transport_ctx = None
self._transport = None
@ -261,25 +290,55 @@ class MCPClient:
async def list_tools(self) -> List[MCPTool]:
"""List available tools from the server."""
verbose_logger.debug(
f"MCP client listing tools from {self.server_url or 'stdio'}"
)
if not self._session:
verbose_logger.debug("MCP client session not found, attempting to connect")
try:
await self.connect()
except Exception as e:
verbose_logger.warning(f"MCP client connection failed: {str(e)}")
verbose_logger.error(
f"MCP client connection failed during list_tools: {type(e).__name__}: {str(e)}"
)
return []
if self._session is None:
verbose_logger.warning("MCP client session is not initialized")
verbose_logger.error(
"MCP client session is not initialized after connection attempt"
)
return []
try:
result = await self._session.list_tools()
tool_count = len(result.tools)
tool_names = [tool.name for tool in result.tools]
verbose_logger.info(
f"MCP client listed {tool_count} tools from {self.server_url or 'stdio'}: {tool_names}"
)
return result.tools
except asyncio.CancelledError:
verbose_logger.warning("MCP client list_tools was cancelled")
await self.disconnect()
raise
except Exception as e:
verbose_logger.warning(f"MCP client list_tools failed: {str(e)}")
error_type = type(e).__name__
verbose_logger.error(
f"MCP client list_tools failed - "
f"Error Type: {error_type}, "
f"Error: {str(e)}, "
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
# Check if it's a stream/connection error
if "BrokenResourceError" in error_type or "Broken" in error_type:
verbose_logger.error(
"MCP client detected broken connection/stream during list_tools - "
"the MCP server may have crashed, disconnected, or timed out"
)
await self.disconnect()
# Return empty list instead of raising to allow graceful degradation
return []
@ -290,17 +349,28 @@ class MCPClient:
"""
Call an MCP Tool.
"""
verbose_logger.info(
f"MCP client calling tool '{call_tool_request_params.name}' with arguments: {call_tool_request_params.arguments}"
)
if not self._session:
verbose_logger.warning(
"MCP client session not found, attempting to connect"
)
try:
await self.connect()
except Exception as e:
verbose_logger.warning(f"MCP client connection failed: {str(e)}")
verbose_logger.error(
f"MCP client connection failed before tool call: {type(e).__name__}: {str(e)}"
)
return MCPCallToolResult(
content=[TextContent(type="text", text=f"{str(e)}")], isError=True
)
if self._session is None:
verbose_logger.warning("MCP client session is not initialized")
verbose_logger.error(
"MCP client session is not initialized after connection attempt"
)
return MCPCallToolResult(
content=[
TextContent(
@ -310,22 +380,59 @@ class MCPClient:
isError=True,
)
# Check session and transport state before calling tool
verbose_logger.debug(
f"MCP client state before tool call - "
f"session: {'active' if self._session else 'none'}, "
f"transport: {'active' if self._transport else 'none'}, "
f"session_ctx: {'active' if self._session_ctx else 'none'}, "
f"transport_ctx: {'active' if self._transport_ctx else 'none'}"
)
try:
verbose_logger.debug("MCP client sending tool call to session")
tool_result = await self._session.call_tool(
name=call_tool_request_params.name,
arguments=call_tool_request_params.arguments,
)
verbose_logger.info(
f"MCP client tool call '{call_tool_request_params.name}' completed successfully"
)
return tool_result
except asyncio.CancelledError:
verbose_logger.warning("MCP client tool call was cancelled")
await self.disconnect()
raise
except Exception as e:
verbose_logger.warning(f"MCP client call_tool failed: {str(e)}")
import traceback
error_trace = traceback.format_exc()
verbose_logger.debug(f"MCP client tool call traceback:\n{error_trace}")
# Log detailed error information
error_type = type(e).__name__
verbose_logger.error(
f"MCP client call_tool failed - "
f"Error Type: {error_type}, "
f"Error: {str(e)}, "
f"Tool: {call_tool_request_params.name}, "
f"Server: {self.server_url or 'stdio'}, "
f"Transport: {self.transport_type}"
)
# Check if it's a stream/connection error
if "BrokenResourceError" in error_type or "Broken" in error_type:
verbose_logger.error(
"MCP client detected broken connection/stream - "
"the MCP server may have crashed, disconnected, or timed out. "
"Session and transport will be disconnected."
)
await self.disconnect()
# Return a default error result instead of raising
return MCPCallToolResult(
content=[
TextContent(type="text", text=f"{str(e)}")
TextContent(type="text", text=f"{error_type}: {str(e)}")
], # Empty content for error case
isError=True,
)

View file

@ -1058,6 +1058,103 @@ class MCPServerManager:
)
)
async def _call_regular_mcp_tool(
self,
mcp_server: MCPServer,
original_tool_name: str,
arguments: Dict[str, Any],
tasks: List,
mcp_auth_header: Optional[str],
mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]],
oauth2_headers: Optional[Dict[str, str]],
raw_headers: Optional[Dict[str, str]],
proxy_logging_obj: Optional[ProxyLogging],
) -> CallToolResult:
"""
Call a regular MCP tool using the MCP client.
Args:
mcp_server: The MCP server configuration
original_tool_name: The original tool name (without prefix)
arguments: Tool arguments
tasks: List of async tasks to append to (for during hooks)
mcp_auth_header: MCP auth header (deprecated)
mcp_server_auth_headers: Optional dict of server-specific auth headers
oauth2_headers: Optional OAuth2 headers
raw_headers: Optional raw headers from the request
proxy_logging_obj: Optional ProxyLogging object for hook integration
Returns:
CallToolResult from the MCP server
Raises:
BlockedPiiEntityError: If PII is blocked by guardrails
GuardrailRaisedException: If guardrails block the call
HTTPException: If an HTTP error occurs
"""
# Get server-specific auth header if available
server_auth_header: Optional[Union[Dict[str, str], str]] = None
if mcp_server_auth_headers and mcp_server.alias:
server_auth_header = mcp_server_auth_headers.get(mcp_server.alias)
elif mcp_server_auth_headers and mcp_server.server_name:
server_auth_header = mcp_server_auth_headers.get(mcp_server.server_name)
# Fall back to deprecated mcp_auth_header if no server-specific header found
if server_auth_header is None:
server_auth_header = mcp_auth_header
# oauth2 headers
extra_headers: Optional[Dict[str, str]] = None
if mcp_server.auth_type == MCPAuth.oauth2:
extra_headers = oauth2_headers
if mcp_server.extra_headers and raw_headers:
if extra_headers is None:
extra_headers = {}
for header in mcp_server.extra_headers:
if header in raw_headers:
extra_headers[header] = raw_headers[header]
client = self._create_mcp_client(
server=mcp_server,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
)
call_tool_params = MCPCallToolRequestParams(
name=original_tool_name,
arguments=arguments,
)
async def _call_tool_via_client(client, params):
async with client:
return await client.call_tool(params)
tasks.append(
asyncio.create_task(_call_tool_via_client(client, call_tool_params))
)
# IMPORTANT: Must await tasks INSIDE the context manager to keep connection alive
try:
mcp_responses = await asyncio.gather(*tasks)
except (
BlockedPiiEntityError,
GuardrailRaisedException,
HTTPException,
) as e:
# Re-raise guardrail exceptions to properly fail the MCP call
verbose_logger.error(
f"Guardrail blocked MCP tool call during result check: {str(e)}"
)
raise e
# If proxy_logging_obj is None, the tool call result is at index 0
# If proxy_logging_obj is not None, the tool call result is at index 1 (after the during hook task)
result_index = 1 if proxy_logging_obj else 0
result = mcp_responses[result_index]
return cast(CallToolResult, result)
async def call_tool(
self,
name: str,
@ -1146,48 +1243,19 @@ class MCPServerManager:
)
else:
# For regular MCP servers, use the MCP client
# Get server-specific auth header if available
server_auth_header: Optional[Union[Dict[str, str], str]] = None
if mcp_server_auth_headers and mcp_server.alias:
server_auth_header = mcp_server_auth_headers.get(mcp_server.alias)
elif mcp_server_auth_headers and mcp_server.server_name:
server_auth_header = mcp_server_auth_headers.get(mcp_server.server_name)
# Fall back to deprecated mcp_auth_header if no server-specific header found
if server_auth_header is None:
server_auth_header = mcp_auth_header
# oauth2 headers
extra_headers: Optional[Dict[str, str]] = None
if mcp_server.auth_type == MCPAuth.oauth2:
extra_headers = oauth2_headers
if mcp_server.extra_headers and raw_headers:
if extra_headers is None:
extra_headers = {}
for header in mcp_server.extra_headers:
if header in raw_headers:
extra_headers[header] = raw_headers[header]
client = self._create_mcp_client(
server=mcp_server,
mcp_auth_header=server_auth_header,
extra_headers=extra_headers,
)
call_tool_params = MCPCallToolRequestParams(
name=original_tool_name,
return await self._call_regular_mcp_tool(
mcp_server=mcp_server,
original_tool_name=original_tool_name,
arguments=arguments,
tasks=tasks,
mcp_auth_header=mcp_auth_header,
mcp_server_auth_headers=mcp_server_auth_headers,
oauth2_headers=oauth2_headers,
raw_headers=raw_headers,
proxy_logging_obj=proxy_logging_obj,
)
async def _call_tool_via_client(client, params):
async with client:
return await client.call_tool(params)
tasks.append(
asyncio.create_task(_call_tool_via_client(client, call_tool_params))
)
# For OpenAPI tools, await outside the client context
try:
mcp_responses = await asyncio.gather(*tasks)

View file

@ -11,4 +11,4 @@ mcp_servers:
token_url: https://github.com/login/oauth/access_token
client_id: os.environ/GITHUB_OAUTH_CLIENT_ID
client_secret: os.environ/GITHUB_OAUTH_CLIENT_SECRET
scopes: ["public_repo", "user:email"]
scopes: ["public_repo", "user:email"]

View file

@ -776,6 +776,7 @@ class KeyRequestBase(GenerateRequestBase):
tags: Optional[List[str]] = None
enforced_params: Optional[List[str]] = None
allowed_routes: Optional[list] = []
allowed_passthrough_routes: Optional[list] = None
rpm_limit_type: Optional[
Literal["guaranteed_throughput", "best_effort_throughput"]
] = None # raise an error if 'guaranteed_throughput' is set and we're overallocating rpm
@ -1281,6 +1282,7 @@ class NewTeamRequest(TeamBase):
guardrails: Optional[List[str]] = None
prompts: Optional[List[str]] = None
object_permission: Optional[LiteLLM_ObjectPermissionBase] = None
allowed_passthrough_routes: Optional[list] = None
team_member_budget: Optional[float] = (
None # allow user to set a budget for all team members
)
@ -1336,6 +1338,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase):
team_member_rpm_limit: Optional[int] = None
team_member_tpm_limit: Optional[int] = None
team_member_key_duration: Optional[str] = None
allowed_passthrough_routes: Optional[list] = None
class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase):
@ -1443,7 +1446,7 @@ class LiteLLM_ObjectPermissionTable(LiteLLMPydanticObjectBase):
"1234567890": ["tool_name_1", "tool_name_2"]
}
"""
vector_stores: Optional[List[str]] = []
@ -3138,6 +3141,7 @@ LiteLLM_ManagementEndpoint_MetadataFields_Premium = [
"team_member_key_duration",
"prompts",
"logging",
"allowed_passthrough_routes",
]

View file

@ -52,9 +52,12 @@ class RouteChecks:
if len(valid_token.allowed_routes) == 0:
return True
# explicit check for allowed routes
if route in valid_token.allowed_routes:
return True
# explicit check for allowed routes (exact match or prefix match)
for allowed_route in valid_token.allowed_routes:
if RouteChecks._route_matches_allowed_route(
route=route, allowed_route=allowed_route
):
return True
## check if 'allowed_route' is a field name in LiteLLMRoutes
if any(
@ -62,13 +65,13 @@ class RouteChecks:
for allowed_route in valid_token.allowed_routes
):
for allowed_route in valid_token.allowed_routes:
if allowed_route in LiteLLMRoutes._member_names_:
if allowed_route in LiteLLMRoutes._member_names_:
if RouteChecks.check_route_access(
route=route,
allowed_routes=LiteLLMRoutes._member_map_[allowed_route].value,
):
return True
################################################
# For llm_api_routes, also check registered pass-through endpoints
################################################
@ -76,7 +79,10 @@ class RouteChecks:
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
InitPassThroughEndpointHelpers,
)
if InitPassThroughEndpointHelpers.is_registered_pass_through_route(route=route):
if InitPassThroughEndpointHelpers.is_registered_pass_through_route(
route=route
):
return True
# check if wildcard pattern is allowed
@ -111,6 +117,32 @@ class RouteChecks:
return masker._mask_value(user_id)
@staticmethod
def _raise_admin_only_route_exception(
user_obj: Optional[LiteLLM_UserTable],
route: str,
) -> None:
"""
Raise exception for routes that require proxy admin access
Args:
user_obj (Optional[LiteLLM_UserTable]): The user object
route (str): The route being accessed
Raises:
Exception: With user role and masked user_id information
"""
user_role = "unknown"
user_id = "unknown"
if user_obj is not None:
user_role = user_obj.user_role or "unknown"
user_id = user_obj.user_id or "unknown"
masked_user_id = RouteChecks._mask_user_id(user_id)
raise Exception(
f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={masked_user_id}"
)
@staticmethod
def non_proxy_admin_allowed_routes_check(
user_obj: Optional[LiteLLM_UserTable],
@ -195,16 +227,27 @@ class RouteChecks:
pass
elif route.startswith("/v1/mcp/") or route.startswith("/mcp-rest/"):
pass # authN/authZ handled by api itself
else:
user_role = "unknown"
user_id = "unknown"
if user_obj is not None:
user_role = user_obj.user_role or "unknown"
user_id = user_obj.user_id or "unknown"
elif RouteChecks.check_passthrough_route_access(
route=route, user_api_key_dict=valid_token
):
pass
elif valid_token.allowed_routes is not None:
# check if route is in allowed_routes (exact match or prefix match)
route_allowed = False
for allowed_route in valid_token.allowed_routes:
if RouteChecks._route_matches_allowed_route(
route=route, allowed_route=allowed_route
):
route_allowed = True
break
masked_user_id = RouteChecks._mask_user_id(user_id)
raise Exception(
f"Only proxy admin can be used to generate, delete, update info for new keys/users/teams. Route={route}. Your role={user_role}. Your user_id={masked_user_id}"
if not route_allowed:
RouteChecks._raise_admin_only_route_exception(
user_obj=user_obj, route=route
)
else:
RouteChecks._raise_admin_only_route_exception(
user_obj=user_obj, route=route
)
@staticmethod
@ -347,6 +390,32 @@ class RouteChecks:
# If there's no wildcard, the pattern and route should match exactly
return route == pattern
@staticmethod
def _route_matches_allowed_route(route: str, allowed_route: str) -> bool:
"""
Check if route matches the allowed_route pattern.
Supports both exact match and prefix match.
Examples:
- allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6" -> True (exact match)
- allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-6/v1/chat/completions" -> True (prefix match)
- allowed_route="/fake-openai-proxy-6", route="/fake-openai-proxy-600" -> False (not a valid prefix)
Args:
route: The actual route being accessed
allowed_route: The allowed route pattern
Returns:
bool: True if route matches (exact or prefix), False otherwise
"""
# Exact match
if route == allowed_route:
return True
# Prefix match - ensure we add "/" to prevent false matches like /fake-openai-proxy-600
if route.startswith(allowed_route + "/"):
return True
return False
@staticmethod
def check_route_access(route: str, allowed_routes: List[str]) -> bool:
"""
@ -394,6 +463,44 @@ class RouteChecks:
return False
@staticmethod
def check_passthrough_route_access(
route: str, user_api_key_dict: UserAPIKeyAuth
) -> bool:
"""
Check if route is a passthrough route.
Supports both exact match and prefix match.
"""
metadata = user_api_key_dict.metadata
team_metadata = user_api_key_dict.team_metadata or {}
if metadata is None and team_metadata is None:
return False
if (
"allowed_passthrough_routes" not in metadata
and "allowed_passthrough_routes" not in team_metadata
):
return False
if (
metadata.get("allowed_passthrough_routes") is None
and team_metadata.get("allowed_passthrough_routes") is None
):
return False
allowed_passthrough_routes = (
metadata.get("allowed_passthrough_routes")
or team_metadata.get("allowed_passthrough_routes")
or []
)
# Check if route matches any allowed passthrough route (exact or prefix match)
for allowed_route in allowed_passthrough_routes:
if RouteChecks._route_matches_allowed_route(
route=route, allowed_route=allowed_route
):
return True
return False
@staticmethod
def _is_assistants_api_request(request: Request) -> bool:
"""

View file

@ -47,9 +47,9 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
_add_model_to_db,
)
from litellm.proxy.management_helpers.object_permission_utils import (
_set_object_permission,
attach_object_permission_to_dict,
handle_update_object_permission_common,
_set_object_permission,
)
from litellm.proxy.management_helpers.team_member_permission_checks import (
TeamMemberPermissionChecks,
@ -563,6 +563,7 @@ async def _common_key_generation_helper( # noqa: PLR0915
data_json = data.model_dump(exclude_unset=True, exclude_none=True) # type: ignore
data_json = handle_key_type(data, data_json)
# if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users
if "max_budget" in data_json:
data_json["key_max_budget"] = data_json.pop("max_budget", None)
@ -838,6 +839,7 @@ async def generate_key_fn(
- enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests)
- prompts: Optional[List[str]] - List of prompts that the key is allowed to use.
- allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
- allowed_passthrough_routes: Optional[list] - List of allowed pass through endpoints for the key. Store the actual endpoint or store a wildcard pattern for a set of endpoints. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through endpoints the key can access, without specifying the routes. If allowed_routes is specified, allowed_pass_through_endpoints is ignored.
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "mcp_tool_permissions": {"server_id_1": ["tool1", "tool2"]}}. IF null or {} then no object permission.
- key_type: Optional[str] - Type of key that determines default allowed routes. Options: "llm_api" (can call LLM API routes), "management" (can call management routes), "read_only" (can only call info/read routes), "default" (uses default allowed routes). Defaults to "default".
- prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts.
@ -1114,8 +1116,6 @@ def prepare_metadata_fields(
return non_default_values
async def prepare_key_update_data(
data: Union[UpdateKeyRequest, RegenerateKeyRequest],
existing_key_row: LiteLLM_VerificationToken,
@ -1277,6 +1277,7 @@ async def update_key_fn(
- temp_budget_increase: Optional[float] - Temporary budget increase for the key (Enterprise only).
- temp_budget_expiry: Optional[str] - Expiry time for the temporary budget increase (Enterprise only).
- allowed_routes: Optional[list] - List of allowed routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/chat/completions", "/embeddings", "/keys/*"]
- allowed_passthrough_routes: Optional[list] - List of allowed pass through routes for the key. Store the actual route or store a wildcard pattern for a set of routes. Example - ["/my-custom-endpoint"]. Use this instead of allowed_routes, if you just want to specify which pass through routes the key can access, without specifying the routes. If allowed_routes is specified, allowed_passthrough_routes is ignored.
- prompts: Optional[List[str]] - List of allowed prompts for the key. If specified, the key will only be able to use these specific prompts.
- object_permission: Optional[LiteLLM_ObjectPermissionBase] - key-specific object permission. Example - {"vector_stores": ["vector_store_1", "vector_store_2"], "mcp_tool_permissions": {"server_id_1": ["tool1", "tool2"]}}. IF null or {} then no object permission.
- auto_rotate: Optional[bool] - Whether this key should be automatically rotated
@ -2819,6 +2820,7 @@ async def list_keys(
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)
@router.get(
"/key/aliases",
tags=["key management"],

View file

@ -4,7 +4,6 @@
This is an enterprise feature and requires a premium license.
"""
from litellm._uuid import uuid
from typing import Any, Dict, List, Optional, Set, Tuple
from fastapi import (
@ -17,11 +16,12 @@ from fastapi import (
Request,
Response,
)
from typing_extensions import TypedDict
from pydantic import BaseModel
from typing_extensions import TypedDict
import litellm
from litellm._logging import verbose_proxy_logger
from litellm._uuid import uuid
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import (
LiteLLM_TeamTable,
@ -51,32 +51,31 @@ from litellm.types.proxy.management_endpoints.scim_v2 import *
class UserProvisionerHelpers:
"""Helper methods for user provisioning operations."""
@staticmethod
async def handle_existing_user_by_email(
prisma_client,
new_user_request: NewUserRequest
prisma_client, new_user_request: NewUserRequest
) -> Optional[SCIMUser]:
"""
Check if a user with the given email already exists and update them if found.
Args:
prisma_client: Database client
new_user_request: New user request data
Returns:
SCIMUser if user was updated, None if no existing user found
"""
if not new_user_request.user_email:
return None
existing_user = await prisma_client.db.litellm_usertable.find_first(
where={"user_email": new_user_request.user_email}
)
if not existing_user:
return None
# Update the user
updated_user = await prisma_client.db.litellm_usertable.update(
where={"user_id": existing_user.user_id},
@ -88,12 +87,15 @@ class UserProvisionerHelpers:
"metadata": safe_dumps(new_user_request.metadata),
},
)
return await ScimTransformations.transform_litellm_user_to_scim_user(updated_user)
return await ScimTransformations.transform_litellm_user_to_scim_user(
updated_user
)
class ScimUserData(TypedDict):
"""Typed structure for extracted SCIM user data."""
user_email: Optional[str]
user_alias: Optional[str]
sso_user_id: Optional[str]
@ -105,6 +107,7 @@ class ScimUserData(TypedDict):
class GroupMemberExtractionResult(BaseModel):
"""Result of extracting and processing group members."""
existing_member_ids: List[str]
created_users: List[NewUserResponse]
all_member_ids: List[str] # existing + newly created
@ -121,7 +124,7 @@ scim_router = APIRouter(
async def _get_prisma_client_or_raise_exception():
"""Check if database is connected and raise HTTPException if not."""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(status_code=500, detail={"error": "No database connected"})
return prisma_client
@ -130,32 +133,32 @@ async def _get_prisma_client_or_raise_exception():
async def _check_user_exists(user_id: str):
"""Check if user exists and return user, raise 404 if not found."""
prisma_client = await _get_prisma_client_or_raise_exception()
user = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_id}
)
if not user:
raise HTTPException(
status_code=404, detail={"error": f"User not found with ID: {user_id}"}
)
return user
async def _check_team_exists(team_id: str):
"""Check if team exists and return team, raise 404 if not found."""
prisma_client = await _get_prisma_client_or_raise_exception()
team = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id}
)
if not team:
raise HTTPException(
status_code=404, detail={"error": f"Group not found with ID: {team_id}"}
)
return team
@ -184,7 +187,9 @@ def _extract_scim_user_data(user: SCIMUser) -> ScimUserData:
}
def _build_scim_metadata(given_name: Optional[str], family_name: Optional[str], active: Optional[bool] = None) -> Dict[str, Any]:
def _build_scim_metadata(
given_name: Optional[str], family_name: Optional[str], active: Optional[bool] = None
) -> Dict[str, Any]:
"""Build metadata dictionary with SCIM data."""
metadata: Dict[str, Any] = {
"scim_metadata": LiteLLM_UserScimMetadata(
@ -192,17 +197,17 @@ def _build_scim_metadata(given_name: Optional[str], family_name: Optional[str],
familyName=family_name,
).model_dump()
}
if active is not None:
metadata["scim_active"] = active
return metadata
async def _extract_group_member_ids(group: SCIMGroup) -> GroupMemberExtractionResult:
"""
Extract member IDs from SCIMGroup, creating users that don't exist.
Returns:
GroupMemberExtractionResult with existing members, created users, and all member IDs
"""
@ -210,35 +215,34 @@ async def _extract_group_member_ids(group: SCIMGroup) -> GroupMemberExtractionRe
existing_member_ids = []
created_users = []
all_member_ids = []
if group.members:
for member in group.members:
user_id = member.value
# Check if user exists
user = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": user_id}
)
if user:
existing_member_ids.append(user_id)
all_member_ids.append(user_id)
else:
# Create the user if they don't exist using our helper
created_user = await _create_user_if_not_exists(
user_id=user_id,
created_via="scim_group_membership"
user_id=user_id, created_via="scim_group_membership"
)
if created_user:
created_users.append(created_user)
all_member_ids.append(user_id)
# If creation failed, user is skipped (logged in helper)
return GroupMemberExtractionResult(
existing_member_ids=existing_member_ids,
created_users=created_users,
all_member_ids=all_member_ids
all_member_ids=all_member_ids,
)
@ -246,7 +250,7 @@ async def _get_team_members_display(member_ids: List[str]) -> List[SCIMMember]:
"""Get SCIMMember objects with display names for a list of member IDs."""
prisma_client = await _get_prisma_client_or_raise_exception()
members: List[SCIMMember] = []
for member_id in member_ids:
user = await prisma_client.db.litellm_usertable.find_unique(
where={"user_id": member_id}
@ -254,18 +258,20 @@ async def _get_team_members_display(member_ids: List[str]) -> List[SCIMMember]:
if user:
display_name = user.user_email or user.user_id
members.append(SCIMMember(value=user.user_id, display=display_name))
return members
async def _handle_team_membership_changes(user_id: str, existing_teams: List[str], new_teams: List[str]) -> None:
async def _handle_team_membership_changes(
user_id: str, existing_teams: List[str], new_teams: List[str]
) -> None:
"""Handle adding/removing user from teams based on changes."""
existing_teams_set = set(existing_teams)
new_teams_set = set(new_teams)
teams_to_add = new_teams_set - existing_teams_set
teams_to_remove = existing_teams_set - new_teams_set
if teams_to_add or teams_to_remove:
await patch_team_membership(
user_id=user_id,
@ -274,19 +280,21 @@ async def _handle_team_membership_changes(user_id: str, existing_teams: List[str
)
async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_group") -> Optional[NewUserResponse]:
async def _create_user_if_not_exists(
user_id: str, created_via: str = "scim_group"
) -> Optional[NewUserResponse]:
"""
Helper function to create a user if they don't exist.
Args:
user_id: The user ID to create
created_via: Context for where the user was created from
Returns:
LiteLLM_UserTable if user was created, None if creation failed
"""
from litellm.proxy.management_endpoints.internal_user_endpoints import new_user
try:
# Get default role for new internal users
default_role: Optional[
@ -313,7 +321,7 @@ async def _create_user_if_not_exists(user_id: str, created_via: str = "scim_grou
created_user = await new_user(data=new_user_request)
verbose_proxy_logger.info(f"Created user {user_id} via {created_via}")
return created_user
except Exception as e:
verbose_proxy_logger.exception(f"Failed to create user {user_id}: {e}")
return None
@ -324,7 +332,7 @@ async def _get_team_member_user_ids_from_team(team: LiteLLM_TeamTable) -> List[s
Get the IDs of the members from a team.
Use one source of truth for the member IDs: team.members_with_roles
"""
member_user_ids: List[str] = []
for member in team.members_with_roles or []:
@ -337,7 +345,6 @@ async def _get_team_member_user_ids_from_team(team: LiteLLM_TeamTable) -> List[s
return member_user_ids
# Dependency to set the correct SCIM Content-Type
async def set_scim_content_type(response: Response):
"""Sets the Content-Type header to application/scim+json"""
@ -450,7 +457,7 @@ async def get_user(
verbose_proxy_logger.debug("SCIM GET USER request for user_id=%s", user_id)
try:
user = await _check_user_exists(user_id)
# Convert to SCIM format
scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(user)
return scim_user
@ -458,6 +465,7 @@ async def get_user(
except Exception as e:
raise handle_exception_on_proxy(e)
@scim_router.post(
"/Users",
response_model=SCIMUser,
@ -471,11 +479,9 @@ async def create_user(
Create a user according to SCIM v2 protocol
"""
try:
verbose_proxy_logger.debug(
"SCIM CREATE USER request: %s", user.model_dump()
)
verbose_proxy_logger.debug("SCIM CREATE USER request: %s", user.model_dump())
prisma_client = await _get_prisma_client_or_raise_exception()
# Extract data from SCIM user
user_data = _extract_scim_user_data(user)
@ -487,20 +493,24 @@ async def create_user(
if existing_user:
raise HTTPException(
status_code=409,
detail={"error": f"User already exists with username: {user.userName}"},
detail={
"error": f"User already exists with username: {user.userName}"
},
)
# Create user in database
user_id = user.userName or str(uuid.uuid4())
metadata = _build_scim_metadata(user_data["given_name"], user_data["family_name"])
metadata = _build_scim_metadata(
user_data["given_name"], user_data["family_name"]
)
default_role: Optional[
Literal[
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
]
LitellmUserRoles.PROXY_ADMIN,
LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY,
LitellmUserRoles.INTERNAL_USER,
LitellmUserRoles.INTERNAL_USER_VIEW_ONLY,
]
] = LitellmUserRoles.INTERNAL_USER_VIEW_ONLY
if litellm.default_internal_user_params:
default_role = litellm.default_internal_user_params.get("user_role")
@ -517,22 +527,23 @@ async def create_user(
# Check if user with email already exists and update if found
existing_user_scim = await UserProvisionerHelpers.handle_existing_user_by_email(
prisma_client=prisma_client,
new_user_request=new_user_request
prisma_client=prisma_client, new_user_request=new_user_request
)
if existing_user_scim:
return existing_user_scim
created_user = await new_user(
data=new_user_request,
)
scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(
user=created_user
)
return scim_user
except HTTPException as e: # allow exceptions like SCIMUserAlreadyExists to be raised
except (
HTTPException
) as e: # allow exceptions like SCIMUserAlreadyExists to be raised
raise e
except Exception as e:
raise handle_exception_on_proxy(e)
@ -564,18 +575,16 @@ async def update_user(
# Extract data from SCIM user
user_data = _extract_scim_user_data(user)
# Build metadata with SCIM data
# Build metadata with SCIM data
metadata = _build_scim_metadata(
user_data["given_name"],
user_data["family_name"],
user_data["active"]
user_data["given_name"], user_data["family_name"], user_data["active"]
)
# Handle team membership changes
await _handle_team_membership_changes(
user_id=user_id,
existing_teams=existing_user.teams or [],
new_teams=user_data["teams"]
new_teams=user_data["teams"],
)
# Update user with all new data (full replacement)
@ -590,6 +599,7 @@ async def update_user(
# Serialize metadata to JSON string for Prisma to avoid GraphQL parsing issues
if "metadata" in update_data and isinstance(update_data["metadata"], dict):
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
update_data["metadata"] = safe_dumps(update_data["metadata"])
updated_user = await prisma_client.db.litellm_usertable.update(
@ -598,8 +608,10 @@ async def update_user(
)
# Convert back to SCIM format
scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(updated_user)
scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(
updated_user
)
return scim_user
except Exception as e:
@ -617,9 +629,7 @@ async def delete_user(
"""
Delete a user according to SCIM v2 protocol
"""
verbose_proxy_logger.debug(
"SCIM DELETE USER request for user_id=%s", user_id
)
verbose_proxy_logger.debug("SCIM DELETE USER request for user_id=%s", user_id)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
existing_user = await _check_user_exists(user_id)
@ -668,7 +678,9 @@ def _extract_group_values(value: Any) -> List[str]:
return group_values
def _handle_displayname_update(op_type: str, value: Any, update_data: Dict[str, Any]) -> None:
def _handle_displayname_update(
op_type: str, value: Any, update_data: Dict[str, Any]
) -> None:
"""Handle displayname updates."""
if op_type == "remove":
update_data["user_alias"] = None
@ -676,7 +688,9 @@ def _handle_displayname_update(op_type: str, value: Any, update_data: Dict[str,
update_data["user_alias"] = str(value)
def _handle_externalid_update(op_type: str, value: Any, update_data: Dict[str, Any]) -> None:
def _handle_externalid_update(
op_type: str, value: Any, update_data: Dict[str, Any]
) -> None:
"""Handle externalid updates."""
if op_type == "remove":
update_data["sso_user_id"] = None
@ -697,7 +711,9 @@ def _handle_active_update(op_type: str, value: Any, metadata: Dict[str, Any]) ->
metadata["scim_active"] = bool_val
def _handle_name_update(path: str, op_type: str, value: Any, scim_metadata: Dict[str, Any]) -> None:
def _handle_name_update(
path: str, op_type: str, value: Any, scim_metadata: Dict[str, Any]
) -> None:
"""Handle name field updates (givenName, familyName)."""
if path == "name.givenname":
if op_type == "remove":
@ -711,7 +727,9 @@ def _handle_name_update(path: str, op_type: str, value: Any, scim_metadata: Dict
scim_metadata["familyName"] = str(value)
def _handle_group_operations(op_type: str, value: Any, teams_set: Set[str]) -> Optional[Set[str]]:
def _handle_group_operations(
op_type: str, value: Any, teams_set: Set[str]
) -> Optional[Set[str]]:
"""Handle group/team membership operations."""
group_values = _extract_group_values(value)
if op_type == "replace":
@ -724,7 +742,9 @@ def _handle_group_operations(op_type: str, value: Any, teams_set: Set[str]) -> O
return None
def _handle_generic_metadata(path: str, op_type: str, value: Any, metadata: Dict[str, Any]) -> None:
def _handle_generic_metadata(
path: str, op_type: str, value: Any, metadata: Dict[str, Any]
) -> None:
"""Handle generic metadata operations for unknown paths."""
if op_type == "remove":
metadata.pop(path, None)
@ -769,6 +789,7 @@ def _apply_patch_ops(
update_data["metadata"] = metadata
return update_data, final_team_set
async def patch_team_membership(
user_id: str,
teams_ids_to_add_user_to: List[str],
@ -778,29 +799,35 @@ async def patch_team_membership(
Add or remove user from teams
"""
for _team_id in teams_ids_to_add_user_to:
try:
await team_member_add(
data=TeamMemberAddRequest(
team_id=_team_id,
member=Member(user_id=user_id, role="user"),
),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
)
except Exception as e:
verbose_proxy_logger.exception(f"Error adding user to team {_team_id}: {e}")
try:
await team_member_add(
data=TeamMemberAddRequest(
team_id=_team_id,
member=Member(user_id=user_id, role="user"),
),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN
),
)
except Exception as e:
verbose_proxy_logger.exception(f"Error adding user to team {_team_id}: {e}")
for _team_id in teams_ids_to_remove_user_from:
try:
await team_member_delete(
data=TeamMemberDeleteRequest(team_id=_team_id, user_id=user_id),
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN),
user_api_key_dict=UserAPIKeyAuth(
user_role=LitellmUserRoles.PROXY_ADMIN
),
)
except Exception as e:
verbose_proxy_logger.exception(f"Error removing user from team {_team_id}: {e}")
verbose_proxy_logger.exception(
f"Error removing user from team {_team_id}: {e}"
)
return True
@scim_router.patch(
"/Users/{user_id}",
response_model=SCIMUser,
@ -833,7 +860,7 @@ async def patch_user(
await _handle_team_membership_changes(
user_id=user_id,
existing_teams=existing_user.teams or [],
new_teams=list(final_team_set)
new_teams=list(final_team_set),
)
update_data["teams"] = list(final_team_set)
@ -841,6 +868,7 @@ async def patch_user(
# Serialize metadata to JSON string for Prisma to avoid GraphQL parsing issues
if "metadata" in update_data and isinstance(update_data["metadata"], dict):
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
update_data["metadata"] = safe_dumps(update_data["metadata"])
updated_user = await prisma_client.db.litellm_usertable.update(
@ -848,7 +876,9 @@ async def patch_user(
data=update_data,
)
scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(updated_user)
scim_user = await ScimTransformations.transform_litellm_user_to_scim_user(
updated_user
)
return scim_user
@ -947,9 +977,7 @@ async def get_group(
"""
Get a single group by ID according to SCIM v2 protocol
"""
verbose_proxy_logger.debug(
"SCIM GET GROUP request for group_id=%s", group_id
)
verbose_proxy_logger.debug("SCIM GET GROUP request for group_id=%s", group_id)
try:
team = await _check_team_exists(group_id)
@ -981,9 +1009,9 @@ async def create_group(
)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
# Generate ID if not provided
team_id = group.id or str(uuid.uuid4())
team_id = group.id or group.externalId or str(uuid.uuid4())
# Check if team already exists
existing_team = await prisma_client.db.litellm_teamtable.find_unique(
@ -998,7 +1026,10 @@ async def create_group(
# Extract and process group members (creating users that don't exist)
member_result = await _extract_group_member_ids(group)
members_with_roles = [Member(user_id=member_id, role="user") for member_id in member_result.all_member_ids]
members_with_roles = [
Member(user_id=member_id, role="user")
for member_id in member_result.all_member_ids
]
# Create team in database
created_team = await new_team(
@ -1043,13 +1074,17 @@ async def update_group(
# Extract and process group members (creating users that don't exist)
member_result = await _extract_group_member_ids(group)
verbose_proxy_logger.debug(f"SCIM PUT GROUP all_member_ids: {member_result.all_member_ids}")
verbose_proxy_logger.debug(f"SCIM PUT GROUP created_users: {len(member_result.created_users)}")
verbose_proxy_logger.debug(
f"SCIM PUT GROUP all_member_ids: {member_result.all_member_ids}"
)
verbose_proxy_logger.debug(
f"SCIM PUT GROUP created_users: {len(member_result.created_users)}"
)
# Prepare update data
existing_metadata = existing_team.metadata if existing_team.metadata else {}
updated_metadata = {**existing_metadata, "scim_data": group.model_dump()}
update_data = {
"team_alias": group.displayName,
"metadata": safe_dumps(updated_metadata),
@ -1066,7 +1101,7 @@ async def update_group(
verbose_proxy_logger.debug(f"SCIM PUT GROUP current_members: {current_members}")
final_members = set(member_result.all_member_ids)
verbose_proxy_logger.debug(f"SCIM PUT GROUP final_members: {final_members}")
await _handle_group_membership_changes(
group_id=group_id,
current_members=current_members,
@ -1094,9 +1129,7 @@ async def delete_group(
"""
Delete a group according to SCIM v2 protocol
"""
verbose_proxy_logger.debug(
"SCIM DELETE GROUP request for group_id=%s", group_id
)
verbose_proxy_logger.debug("SCIM DELETE GROUP request for group_id=%s", group_id)
try:
prisma_client = await _get_prisma_client_or_raise_exception()
existing_team = await _check_team_exists(group_id)
@ -1124,21 +1157,19 @@ async def delete_group(
async def _process_group_patch_operations(
patch_ops: SCIMPatchOp,
existing_team,
prisma_client
patch_ops: SCIMPatchOp, existing_team, prisma_client
) -> Tuple[Dict[str, Any], Set[str]]:
"""Process patch operations for a group and return update data and final members."""
update_data: Dict[str, Any] = {}
# Create a fresh copy of existing metadata to avoid Prisma issues
existing_metadata = existing_team.metadata or {}
metadata = dict(existing_metadata) if existing_metadata else {}
# Track member changes
current_members = set(existing_team.members or [])
final_members = current_members.copy()
# Process each patch operation
for op in patch_ops.Operations:
path = (op.path or "").lower()
@ -1169,14 +1200,13 @@ async def _process_group_patch_operations(
else:
# Create the user if they don't exist using our helper
created_user = await _create_user_if_not_exists(
user_id=member_id,
created_via="scim_group_patch"
user_id=member_id, created_via="scim_group_patch"
)
if created_user:
valid_members.append(member_id)
# If creation failed, user is skipped (logged in helper)
if op_type == "replace":
final_members = set(valid_members)
elif op_type == "add":
@ -1194,21 +1224,18 @@ async def _process_group_patch_operations(
# Include metadata in update data if it exists
if metadata:
update_data["metadata"] = metadata
return update_data, final_members
async def _apply_group_patch_updates(
group_id: str,
update_data: Dict[str, Any],
final_members: Set[str],
prisma_client
group_id: str, update_data: Dict[str, Any], final_members: Set[str], prisma_client
):
"""Apply patch updates to the group in the database."""
# Serialize metadata if present
if "metadata" in update_data and isinstance(update_data["metadata"], dict):
update_data["metadata"] = safe_dumps(update_data["metadata"])
# Update members list
update_data["members"] = list(final_members)
@ -1217,22 +1244,20 @@ async def _apply_group_patch_updates(
where={"team_id": group_id},
data=update_data,
)
return updated_team
async def _handle_group_membership_changes(
group_id: str,
current_members: Set[str],
final_members: Set[str]
group_id: str, current_members: Set[str], final_members: Set[str]
):
"""Handle adding/removing members from the group."""
members_to_add = final_members - current_members
members_to_remove = current_members - final_members
verbose_proxy_logger.debug(f"members_to_add: {members_to_add}")
verbose_proxy_logger.debug(f"members_to_remove: {members_to_remove}")
# Use existing helper functions for team membership changes
for member_id in members_to_add:
await patch_team_membership(
@ -1276,7 +1301,7 @@ async def patch_group(
update_data, final_members = await _process_group_patch_operations(
patch_ops, existing_team, prisma_client
)
# Track current members for comparison
current_members = set(await _get_team_member_user_ids_from_team(existing_team))
@ -1286,9 +1311,7 @@ async def patch_group(
)
# Handle user-team relationship changes
await _handle_group_membership_changes(
group_id, current_members, final_members
)
await _handle_group_membership_changes(group_id, current_members, final_members)
# Convert to SCIM format and return
scim_group = await ScimTransformations.transform_litellm_team_to_scim_group(

View file

@ -27,6 +27,7 @@ from litellm.proxy._types import (
CommonProxyErrors,
DeleteTeamRequest,
LiteLLM_AuditLogs,
LiteLLM_ManagementEndpoint_MetadataFields,
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
LiteLLM_ModelTable,
LiteLLM_OrganizationTable,
@ -56,9 +57,6 @@ from litellm.proxy._types import (
UpdateTeamRequest,
UserAPIKeyAuth,
)
from litellm.proxy.management_helpers.object_permission_utils import (
_set_object_permission,
)
from litellm.proxy.auth.auth_checks import (
allowed_route_check_inside_route,
can_org_access_model,
@ -76,6 +74,7 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import (
get_daily_activity,
)
from litellm.proxy.management_helpers.object_permission_utils import (
_set_object_permission,
handle_update_object_permission_common,
)
from litellm.proxy.management_helpers.team_member_permission_checks import (
@ -320,7 +319,9 @@ async def new_team( # noqa: PLR0915
- team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
- team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
- prompts: Optional[List[str]] - List of allowed prompts for the team. If specified, the team will only be able to use these specific prompts.
- allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team.
Returns:
- team_id: (str) Unique team id - used for tracking spend across multiple keys for same team id.
@ -478,7 +479,7 @@ async def new_team( # noqa: PLR0915
## Create Team Member Budget Table
data_json = data.json()
## Handle Object Permission - MCP, Vector Stores etc.
data_json = await _set_object_permission(
data_json=data_json,
@ -514,6 +515,14 @@ async def new_team( # noqa: PLR0915
value=getattr(data, field),
)
for field in LiteLLM_ManagementEndpoint_MetadataFields:
if getattr(data, field, None) is not None:
_set_object_metadata_field(
object_data=complete_team_data,
field_name=field,
value=getattr(data, field),
)
# If budget_duration is set, set `budget_reset_at`
if complete_team_data.budget_duration is not None:
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
@ -619,6 +628,53 @@ async def _update_model_table(
return _model_id
async def fetch_and_validate_organization(
organization_id: str,
existing_team_row: Any,
llm_router: Optional[Router],
prisma_client: Any,
) -> Any:
"""
Fetch and validate an organization for team update operations.
Args:
organization_id: The organization ID to fetch
existing_team_row: The existing team row being updated
llm_router: The LLM router instance
prisma_client: The Prisma database client
Returns:
The organization row from the database
Raises:
HTTPException: If llm_router is None, organization not found, or validation fails
"""
if llm_router is None:
raise HTTPException(
status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}
)
organization_row = await prisma_client.db.litellm_organizationtable.find_unique(
where={"organization_id": organization_id},
include={"litellm_budget_table": True, "users": True},
)
if organization_row is None:
raise HTTPException(
status_code=404,
detail={
"error": f"Organization not found, passed organization_id={organization_id}"
},
)
validate_team_org_change(
team=LiteLLM_TeamTable(**existing_team_row.model_dump()),
organization=LiteLLM_OrganizationTable(**organization_row.model_dump()),
llm_router=llm_router,
)
return organization_row
def validate_team_org_change(
team: LiteLLM_TeamTable, organization: LiteLLM_OrganizationTable, llm_router: Router
@ -754,6 +810,7 @@ async def update_team(
- team_member_rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for individual team members.
- team_member_tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for individual team members.
- team_member_key_duration: Optional[str] - The duration for a team member's key. e.g. "1d", "1w", "1mo"
- allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team.
Example - update team TPM Limit
```
@ -810,25 +867,11 @@ async def update_team(
if (
data.organization_id is not None and len(data.organization_id) > 0
): # allow unsetting the organization_id
if llm_router is None:
raise HTTPException(
status_code=500, detail={"error": CommonProxyErrors.no_llm_router.value}
)
organization_row = await prisma_client.db.litellm_organizationtable.find_unique(
where={"organization_id": data.organization_id},
include={"litellm_budget_table": True, "users": True},
)
if organization_row is None:
raise HTTPException(
status_code=404,
detail={
"error": f"Organization not found, passed organization_id={data.organization_id}"
},
)
validate_team_org_change(
team=LiteLLM_TeamTable(**existing_team_row.model_dump()),
organization=LiteLLM_OrganizationTable(**organization_row.model_dump()),
await fetch_and_validate_organization(
organization_id=data.organization_id,
existing_team_row=existing_team_row,
llm_router=llm_router,
prisma_client=prisma_client,
)
elif data.organization_id is not None and len(data.organization_id) == 0:
# unsetting the organization_id
@ -877,6 +920,13 @@ async def update_team(
field_name=field,
)
for field in LiteLLM_ManagementEndpoint_MetadataFields:
if field in updated_kv and updated_kv[field] is not None:
_update_team_metadata_field(
updated_kv=updated_kv,
field_name=field,
)
if "model_aliases" in updated_kv:
updated_kv.pop("model_aliases")
_model_id = await _update_model_table(

View file

@ -5,7 +5,7 @@ import json
import traceback
from base64 import b64encode
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple, Union
from typing import Any, Dict, List, Optional, Tuple, Union, cast
from urllib.parse import urlencode, urlparse
import httpx
@ -42,6 +42,7 @@ from litellm.passthrough import BasePassthroughUtils
from litellm.proxy._types import (
ConfigFieldInfo,
ConfigFieldUpdate,
LiteLLMRoutes,
PassThroughEndpointResponse,
PassThroughGenericEndpoint,
ProxyException,
@ -66,7 +67,7 @@ router = APIRouter()
pass_through_endpoint_logging = PassThroughEndpointLogging()
# Global registry to track registered pass-through routes and prevent memory leaks
_registered_pass_through_routes: Dict[str, Dict[str, str]] = {}
_registered_pass_through_routes: Dict[str, Dict[str, Union[str, Dict[str, Any]]]] = {}
def get_response_body(response: httpx.Response) -> Optional[dict]:
@ -974,24 +975,74 @@ def create_pass_through_route(
] = None, # if pass-through endpoint is a streaming request
subpath: str = "", # captures sub-paths when include_subpath=True
):
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
InitPassThroughEndpointHelpers,
)
path = request.url.path
if not InitPassThroughEndpointHelpers.is_registered_pass_through_route(
route=path
):
raise HTTPException(
status_code=404,
detail=f"Pass-through endpoint {endpoint} not found. This could have been deleted or not yet added to the proxy.",
)
passthrough_params = (
InitPassThroughEndpointHelpers.get_registered_pass_through_route(
route=path
)
)
target_params = {
"target": target,
"custom_headers": custom_headers,
"forward_headers": _forward_headers,
"merge_query_params": _merge_query_params,
"cost_per_request": cost_per_request,
}
if passthrough_params is not None:
target_params.update(passthrough_params.get("passthrough_params", {}))
# Extract and cast parameters with proper types
param_target = target_params.get("target") or target
param_custom_headers = target_params.get("custom_headers", custom_headers)
param_forward_headers = target_params.get(
"forward_headers", _forward_headers
)
param_merge_query_params = target_params.get(
"merge_query_params", _merge_query_params
)
param_cost_per_request = target_params.get(
"cost_per_request", cost_per_request
)
# Construct the full target URL with subpath if needed
full_target = (
HttpPassThroughEndpointHelpers.construct_target_url_with_subpath(
base_target=target, subpath=subpath, include_subpath=include_subpath
base_target=cast(str, param_target),
subpath=subpath,
include_subpath=include_subpath,
)
)
# Ensure custom_headers is a dict
headers_dict = (
param_custom_headers if isinstance(param_custom_headers, dict) else {}
)
return await pass_through_request( # type: ignore
request=request,
target=full_target,
custom_headers=custom_headers or {},
custom_headers=headers_dict,
user_api_key_dict=user_api_key_dict,
forward_headers=_forward_headers,
merge_query_params=_merge_query_params,
forward_headers=cast(Optional[bool], param_forward_headers),
merge_query_params=cast(Optional[bool], param_merge_query_params),
query_params=query_params,
stream=stream,
custom_body=custom_body,
cost_per_request=cost_per_request,
cost_per_request=cast(Optional[float], param_cost_per_request),
custom_llm_provider=custom_llm_provider,
)
@ -1592,6 +1643,14 @@ class InitPassThroughEndpointHelpers:
"endpoint_id": endpoint_id,
"path": path,
"type": "exact",
"passthrough_params": {
"target": target,
"custom_headers": custom_headers,
"forward_headers": forward_headers,
"merge_query_params": merge_query_params,
"dependencies": dependencies,
"cost_per_request": cost_per_request,
},
}
@staticmethod
@ -1645,6 +1704,14 @@ class InitPassThroughEndpointHelpers:
"endpoint_id": endpoint_id,
"path": path,
"type": "subpath",
"passthrough_params": {
"target": target,
"custom_headers": custom_headers,
"forward_headers": forward_headers,
"merge_query_params": merge_query_params,
"dependencies": dependencies,
"cost_per_request": cost_per_request,
},
}
@staticmethod
@ -1661,6 +1728,11 @@ class InitPassThroughEndpointHelpers:
"Removed pass-through route from registry: %s", key
)
@staticmethod
def clear_all_pass_through_routes():
"""Clear all pass-through routes from the registry"""
_registered_pass_through_routes.clear()
@staticmethod
def is_registered_pass_through_route(route: str) -> bool:
"""
@ -1675,6 +1747,12 @@ class InitPassThroughEndpointHelpers:
Returns:
bool: True if route is a registered pass-through endpoint, False otherwise
"""
## CHECK IF MAPPED PASS THROUGH ENDPOINT
for mapped_route in LiteLLMRoutes.mapped_pass_through_routes.value:
if route.startswith(mapped_route):
return True
# Fast path: check if any registered route key contains this path
# Keys are in format: "{endpoint_id}:exact:{path}" or "{endpoint_id}:subpath:{path}"
# Extract unique paths from keys for quick checking
@ -1683,7 +1761,6 @@ class InitPassThroughEndpointHelpers:
if len(parts) == 3:
route_type = parts[1]
registered_path = parts[2]
if route_type == "exact" and route == registered_path:
return True
elif route_type == "subpath":
@ -1694,11 +1771,42 @@ class InitPassThroughEndpointHelpers:
return False
@staticmethod
def get_registered_pass_through_route(route: str) -> Optional[Dict[str, Any]]:
"""Get passthrough params for a given route"""
for key in _registered_pass_through_routes.keys():
parts = key.split(":", 2) # Split into [endpoint_id, type, path]
if len(parts) == 3:
route_type = parts[1]
registered_path = parts[2]
if route_type == "exact" and route == registered_path:
return _registered_pass_through_routes[key]
elif route_type == "subpath":
if route == registered_path or route.startswith(
registered_path + "/"
):
return _registered_pass_through_routes[key]
return None
def _get_combined_pass_through_endpoints(
pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]],
config_pass_through_endpoints: List[Dict],
):
"""Get combined pass-through endpoints from db + config"""
return pass_through_endpoints + config_pass_through_endpoints
async def initialize_pass_through_endpoints(
pass_through_endpoints: Union[List[Dict], List[PassThroughGenericEndpoint]],
):
"""
1. Create a global list of pass-through endpoints (db + config)
2. Clear all existing pass-through endpoints from the FastAPI app routes
3. Add new endpoints to the in-memory registry
Initialize a list of pass-through endpoints by adding them to the FastAPI app routes
Args:
@ -1711,9 +1819,26 @@ async def initialize_pass_through_endpoints(
verbose_proxy_logger.debug("initializing pass through endpoints")
from litellm.proxy._types import CommonProxyErrors, LiteLLMRoutes
from litellm.proxy.proxy_server import app, premium_user
from litellm.proxy.proxy_server import (
app,
config_passthrough_endpoints,
premium_user,
)
for endpoint in pass_through_endpoints:
## get combined pass-through endpoints from db + config
combined_pass_through_endpoints: List[Union[Dict, PassThroughGenericEndpoint]]
if config_passthrough_endpoints is not None:
combined_pass_through_endpoints = _get_combined_pass_through_endpoints( # type: ignore
pass_through_endpoints, config_passthrough_endpoints
)
else:
combined_pass_through_endpoints = pass_through_endpoints # type: ignore
## clear all existing pass-through endpoints from the FastAPI app routes
InitPassThroughEndpointHelpers.clear_all_pass_through_routes()
for endpoint in combined_pass_through_endpoints:
if isinstance(endpoint, PassThroughGenericEndpoint):
endpoint = endpoint.model_dump()
@ -1818,23 +1943,91 @@ async def _get_pass_through_endpoints_from_db(
return returned_endpoints
async def _filter_endpoints_by_team_allowed_routes(
team_id: str,
pass_through_endpoints: List[PassThroughGenericEndpoint],
prisma_client,
) -> List[PassThroughGenericEndpoint]:
"""
Filter pass-through endpoints based on team's allowed_passthrough_routes metadata.
Args:
team_id: The team ID to check permissions for
pass_through_endpoints: List of endpoints to filter
prisma_client: Database client
Returns:
Filtered list of endpoints based on team permissions
Raises:
HTTPException: If team is not found
"""
# retrieve team from db
team = await prisma_client.db.litellm_teamtable.find_unique(
where={"team_id": team_id},
)
if team is None:
raise HTTPException(
status_code=404,
detail={"error": "Team not found"},
)
# retrieve team metadata
team_metadata = team.metadata
if (
team_metadata is not None
and team_metadata.get("allowed_passthrough_routes") is not None
):
## FILTER pass_through_endpoints by allowed_passthrough_routes
pass_through_endpoints = [
endpoint
for endpoint in pass_through_endpoints
if endpoint.path in team_metadata.get("allowed_passthrough_routes")
]
return pass_through_endpoints
@router.get(
"/config/pass_through_endpoint",
dependencies=[Depends(user_api_key_auth)],
response_model=PassThroughEndpointResponse,
)
@router.get(
"/config/pass_through_endpoint/team/{team_id}",
dependencies=[Depends(user_api_key_auth)],
response_model=PassThroughEndpointResponse,
)
async def get_pass_through_endpoints(
endpoint_id: Optional[str] = None,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
team_id: Optional[str] = None,
):
"""
GET configured pass through endpoint.
If no endpoint_id given, return all configured endpoints.
""" ## Get existing pass-through endpoint field value
from litellm.proxy._types import CommonProxyErrors
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
pass_through_endpoints = await _get_pass_through_endpoints_from_db(
endpoint_id=endpoint_id, user_api_key_dict=user_api_key_dict
)
if team_id is not None:
pass_through_endpoints = await _filter_endpoints_by_team_allowed_routes(
team_id=team_id,
pass_through_endpoints=pass_through_endpoints,
prisma_client=prisma_client,
)
return PassThroughEndpointResponse(endpoints=pass_through_endpoints)
@ -1930,6 +2123,7 @@ async def update_pass_through_endpoints(
field_value=pass_through_endpoint_data,
config_type="general_settings",
)
await update_config_general_settings(
data=updated_data, user_api_key_dict=user_api_key_dict
)

View file

@ -259,9 +259,7 @@ from litellm.proxy.management_endpoints.customer_endpoints import (
from litellm.proxy.management_endpoints.internal_user_endpoints import (
router as internal_user_router,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import (
user_update,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
from litellm.proxy.management_endpoints.key_management_endpoints import (
delete_verification_tokens,
duration_in_seconds,
@ -308,9 +306,7 @@ from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMi
from litellm.proxy.openai_files_endpoints.files_endpoints import (
router as openai_files_router,
)
from litellm.proxy.openai_files_endpoints.files_endpoints import (
set_files_config,
)
from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
passthrough_endpoint_router,
)
@ -580,7 +576,7 @@ async def _initialize_shared_aiohttp_session():
ttl_dns_cache=AIOHTTP_TTL_DNS_CACHE,
enable_cleanup_closed=True,
)
session = ClientSession(connector=connector)
verbose_proxy_logger.info(
f"SESSION REUSE: Created shared aiohttp session for connection pooling (ID: {id(session)})"
@ -723,7 +719,7 @@ async def proxy_startup_event(app: FastAPI):
verbose_proxy_logger.info("SESSION REUSE: Closed shared aiohttp session")
except Exception as e:
verbose_proxy_logger.error(f"Error closing shared aiohttp session: {e}")
await proxy_shutdown_event()
@ -995,13 +991,16 @@ experimental = False
llm_router: Optional[Router] = None
llm_model_list: Optional[list] = None
general_settings: dict = {}
config_passthrough_endpoints: Optional[List[Dict[str, Any]]] = None
callback_settings: dict = {}
log_file = "api_log.json"
worker_config = None
master_key: Optional[str] = None
otel_logging = False
prisma_client: Optional[PrismaClient] = None
shared_aiohttp_session: Optional["ClientSession"] = None # Global shared session for connection reuse
shared_aiohttp_session: Optional["ClientSession"] = (
None # Global shared session for connection reuse
)
user_api_key_cache = DualCache(
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
)
@ -1385,31 +1384,33 @@ async def update_cache( # noqa: PLR0915
"""
if tags is None or response_cost is None:
return
try:
for tag_name in tags:
if not tag_name or not isinstance(tag_name, str):
continue
cache_key = f"tag:{tag_name}"
# Fetch the existing tag object from cache
existing_tag_obj = await user_api_key_cache.async_get_cache(key=cache_key)
existing_tag_obj = await user_api_key_cache.async_get_cache(
key=cache_key
)
if existing_tag_obj is None:
# do nothing if tag not in api key cache
continue
verbose_proxy_logger.debug(
f"_update_tag_cache: existing spend for tag={tag_name}: {existing_tag_obj}; response_cost: {response_cost}"
)
if isinstance(existing_tag_obj, dict):
existing_spend = existing_tag_obj.get("spend", 0) or 0
else:
existing_spend = getattr(existing_tag_obj, "spend", 0) or 0
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
# Update the spend column for the given tag
if isinstance(existing_tag_obj, dict):
existing_tag_obj["spend"] = new_spend
@ -1482,6 +1483,7 @@ async def _run_background_health_check():
from litellm.proxy.health_check_utils.shared_health_check_manager import (
SharedHealthCheckManager,
)
shared_health_manager = SharedHealthCheckManager(
redis_cache=redis_usage_cache,
health_check_ttl=DEFAULT_SHARED_HEALTH_CHECK_TTL,
@ -1502,16 +1504,21 @@ async def _run_background_health_check():
# Use shared health check if available, otherwise fall back to direct health check
# Convert health_check_details to bool for perform_shared_health_check (defaults to True if None)
details_bool = health_check_details if health_check_details is not None else True
details_bool = (
health_check_details if health_check_details is not None else True
)
if shared_health_manager is not None:
try:
healthy_endpoints, unhealthy_endpoints = await shared_health_manager.perform_shared_health_check(
model_list=_llm_model_list, details=details_bool
healthy_endpoints, unhealthy_endpoints = (
await shared_health_manager.perform_shared_health_check(
model_list=_llm_model_list, details=details_bool
)
)
except Exception as e:
verbose_proxy_logger.error(
"Error in shared health check, falling back to direct health check: %s", str(e)
"Error in shared health check, falling back to direct health check: %s",
str(e),
)
healthy_endpoints, unhealthy_endpoints = await perform_health_check(
model_list=_llm_model_list, details=health_check_details
@ -1879,7 +1886,7 @@ class ProxyConfig:
"""
Load config values into proxy global state
"""
global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, callback_settings, proxy_batch_polling_interval
global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, user_custom_sso, user_custom_ui_sso_sign_in_handler, use_background_health_checks, use_shared_health_check, health_check_interval, use_queue, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details, callback_settings, proxy_batch_polling_interval, config_passthrough_endpoints
config: dict = await self.get_config(config_file_path=config_file_path)
@ -2234,9 +2241,13 @@ class ProxyConfig:
## pass through endpoints
if general_settings.get("pass_through_endpoints", None) is not None:
config_passthrough_endpoints = general_settings[
"pass_through_endpoints"
]
await initialize_pass_through_endpoints(
pass_through_endpoints=general_settings["pass_through_endpoints"]
)
## ADMIN UI ACCESS ##
ui_access_mode = general_settings.get(
"ui_access_mode", "all"
@ -3055,7 +3066,9 @@ class ProxyConfig:
return current_config
# For dictionary values, update only non-none values
if isinstance(current_config[param_name], dict) and isinstance(db_param_value, dict):
if isinstance(current_config[param_name], dict) and isinstance(
db_param_value, dict
):
_deep_merge_dicts(current_config[param_name], db_param_value)
else:
# Non-dict or mismatched types: DB value replaces config (unchanged behavior)

View file

@ -58,7 +58,7 @@ class LiteLLMCompletionTransformationHandler:
responses_api_request=responses_api_request,
**kwargs,
)
completion_args = {}
completion_args.update(kwargs)
completion_args.update(litellm_completion_request)
@ -83,6 +83,7 @@ class LiteLLMCompletionTransformationHandler:
elif isinstance(litellm_completion_response, litellm.CustomStreamWrapper):
return LiteLLMCompletionStreamingIterator(
model=model,
litellm_custom_stream_wrapper=litellm_completion_response,
request_input=input,
responses_api_request=responses_api_request,
@ -106,7 +107,7 @@ class LiteLLMCompletionTransformationHandler:
previous_response_id=previous_response_id,
litellm_completion_request=litellm_completion_request,
)
acompletion_args = {}
acompletion_args.update(kwargs)
acompletion_args.update(litellm_completion_request)
@ -130,9 +131,12 @@ class LiteLLMCompletionTransformationHandler:
elif isinstance(litellm_completion_response, litellm.CustomStreamWrapper):
return LiteLLMCompletionStreamingIterator(
model=litellm_completion_request.get("model") or "",
litellm_custom_stream_wrapper=litellm_completion_response,
request_input=request_input,
responses_api_request=responses_api_request,
custom_llm_provider=litellm_completion_request.get("custom_llm_provider"),
custom_llm_provider=litellm_completion_request.get(
"custom_llm_provider"
),
litellm_metadata=kwargs.get("litellm_metadata", {}),
)

View file

@ -1,4 +1,6 @@
from typing import List, Optional, Union
import time
import uuid
from typing import List, Optional, Union, cast
import litellm
from litellm.main import stream_chunk_builder
@ -8,11 +10,23 @@ from litellm.responses.litellm_completion_transformation.transformation import (
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import (
PART_UNION_TYPES,
BaseLiteLLMOpenAIResponseObject,
ContentPartAddedEvent,
ContentPartDoneEvent,
ContentPartDonePartOutputText,
ContentPartDonePartReasoningText,
OutputItemAddedEvent,
OutputItemDoneEvent,
OutputTextDeltaEvent,
OutputTextDoneEvent,
ReasoningSummaryTextDeltaEvent,
ResponseCompletedEvent,
ResponseCreatedEvent,
ResponseInProgressEvent,
ResponseInputParam,
ResponsesAPIOptionalRequestParams,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
ResponsesAPIStreamingResponse,
)
@ -32,12 +46,14 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
def __init__(
self,
model: str,
litellm_custom_stream_wrapper: litellm.CustomStreamWrapper,
request_input: Union[str, ResponseInputParam],
responses_api_request: ResponsesAPIOptionalRequestParams,
custom_llm_provider: Optional[str] = None,
litellm_metadata: Optional[dict] = None,
):
self.model: str = model
self.litellm_custom_stream_wrapper: litellm.CustomStreamWrapper = (
litellm_custom_stream_wrapper
)
@ -50,14 +66,274 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
self.collected_chat_completion_chunks: List[ModelResponseStream] = []
self.finished: bool = False
self.litellm_logging_obj = litellm_custom_stream_wrapper.logging_obj
self.sent_response_created_event: bool = False
self.sent_response_in_progress_event: bool = False
self.sent_output_item_added_event: bool = False
self.sent_content_part_added_event: bool = False
self.sent_output_text_done_event: bool = False
self.sent_output_content_part_done_event: bool = False
self.sent_output_item_done_event: bool = False
self.litellm_model_response: Optional[
Union[ModelResponse, TextCompletionResponse]
] = None
self.final_text: str = ""
def _default_response_created_event_data(self) -> dict:
response_created_event_data = {
"id": f"resp_{str(uuid.uuid4())}",
"object": "response",
"created_at": int(time.time()),
"status": "in_progress",
"error": None,
"incomplete_details": None,
"instructions": self.request_input,
"max_output_tokens": None,
"model": self.model,
"output": [],
"parallel_tool_calls": True,
"previous_response_id": None,
"reasoning": {"effort": None, "summary": None},
"store": True,
}
if "temperature" in self.responses_api_request:
response_created_event_data["temperature"] = self.responses_api_request[
"temperature"
]
if "text" in self.responses_api_request:
response_created_event_data["text"] = self.responses_api_request["text"]
if "tool_choice" in self.responses_api_request:
response_created_event_data["tool_choice"] = self.responses_api_request[
"tool_choice"
]
else:
response_created_event_data["tool_choice"] = "auto"
if "tools" in self.responses_api_request:
response_created_event_data["tools"] = self.responses_api_request["tools"]
else:
response_created_event_data["tools"] = []
if "top_p" in self.responses_api_request:
response_created_event_data["top_p"] = self.responses_api_request["top_p"]
else:
response_created_event_data["top_p"] = 1.0
if "truncation" in self.responses_api_request:
response_created_event_data["truncation"] = self.responses_api_request[
"truncation"
]
if "user" in self.responses_api_request:
response_created_event_data["user"] = self.responses_api_request["user"]
if "metadata" in self.responses_api_request:
response_created_event_data["metadata"] = self.responses_api_request[
"metadata"
]
return response_created_event_data
def create_response_created_event(self) -> ResponseCreatedEvent:
"""
data: {"type":"response.created","response":{"id":"resp_67c9fdcecf488190bdd9a0409de3a1ec07b8b0ad4e5eb654","object":"response","created_at":1741290958,"status":"in_progress","error":null,"incomplete_details":null,"instructions":"You are a helpful assistant.","max_output_tokens":null,"model":"gpt-4.1-2025-04-14","output":[],"parallel_tool_calls":true,"previous_response_id":null,"reasoning":{"effort":null,"summary":null},"store":true,"temperature":1.0,"text":{"format":{"type":"text"}},"tool_choice":"auto","tools":[],"top_p":1.0,"truncation":"disabled","usage":null,"user":null,"metadata":{}}}
"""
response_created_event_data = self._default_response_created_event_data()
return ResponseCreatedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_CREATED,
response=ResponsesAPIResponse(**response_created_event_data),
)
def create_response_in_progress_event(self) -> ResponseInProgressEvent:
response_in_progress_event_data = self._default_response_created_event_data()
response_in_progress_event_data["status"] = "in_progress"
return ResponseInProgressEvent(
type=ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS,
response=ResponsesAPIResponse(**response_in_progress_event_data),
)
def create_output_item_added_event(self) -> OutputItemAddedEvent:
return OutputItemAddedEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
output_index=0,
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": f"msg_{str(uuid.uuid4())}",
"type": "message",
"status": "in_progress",
"role": "assistant",
"content": [],
}
),
)
def create_content_part_added_event(self) -> ContentPartAddedEvent:
return ContentPartAddedEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
item_id=f"msg_{str(uuid.uuid4())}",
output_index=0,
content_index=0,
part=BaseLiteLLMOpenAIResponseObject(
**{"type": "output_text", "text": "", "annotations": []}
),
)
def create_litellm_model_response(
self,
) -> Optional[ModelResponse]:
return cast(
Optional[ModelResponse],
stream_chunk_builder(
chunks=self.collected_chat_completion_chunks,
logging_obj=self.litellm_logging_obj,
),
)
def create_output_text_done_event(
self, litellm_complete_object: ModelResponse
) -> OutputTextDoneEvent:
return OutputTextDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
item_id=f"msg_{str(uuid.uuid4())}",
output_index=0,
content_index=0,
text=getattr(litellm_complete_object.choices[0].message, "content", "") # type: ignore
or "",
)
def create_output_content_part_done_event(
self, litellm_complete_object: ModelResponse
) -> ContentPartDoneEvent:
text = getattr(litellm_complete_object.choices[0].message, "content", "") or "" # type: ignore
reasoning_content = getattr(litellm_complete_object.choices[0].message, "reasoning_content", "") or "" # type: ignore
part: Optional[PART_UNION_TYPES] = None
if reasoning_content:
part = ContentPartDonePartReasoningText(
type="reasoning_text",
reasoning=reasoning_content,
)
else:
part = ContentPartDonePartOutputText(
type="output_text",
text=text,
annotations=[],
logprobs=None,
)
return ContentPartDoneEvent(
type=ResponsesAPIStreamEvents.CONTENT_PART_DONE,
item_id=f"msg_{str(uuid.uuid4())}",
output_index=0,
content_index=0,
part=part,
)
def create_output_item_done_event(
self, litellm_complete_object: ModelResponse
) -> OutputItemDoneEvent:
text = self.litellm_model_response.choices[0].message.content or "" # type: ignore
return OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=0,
sequence_number=1,
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": f"msg_{str(uuid.uuid4())}",
"status": "completed",
"type": "message",
"role": "assistant",
"content": [
{
"type": "output_text",
"text": text,
"annotations": [],
}
],
}
),
)
def return_default_done_events(
self, litellm_complete_object: ModelResponse
) -> Optional[BaseLiteLLMOpenAIResponseObject]:
if self.sent_output_text_done_event is False:
self.sent_output_text_done_event = True
return self.create_output_text_done_event(litellm_complete_object)
if self.sent_output_content_part_done_event is False:
self.sent_output_content_part_done_event = True
return self.create_output_content_part_done_event(litellm_complete_object)
if self.sent_output_item_done_event is False:
self.sent_output_item_done_event = True
return self.create_output_item_done_event(litellm_complete_object)
return None
def return_default_initial_events(
self,
) -> Optional[BaseLiteLLMOpenAIResponseObject]:
if self.sent_response_created_event is False:
self.sent_response_created_event = True
return self.create_response_created_event()
elif self.sent_response_in_progress_event is False:
self.sent_response_in_progress_event = True
return self.create_response_in_progress_event()
elif self.sent_output_item_added_event is False:
self.sent_output_item_added_event = True
return self.create_output_item_added_event()
elif self.sent_content_part_added_event is False:
self.sent_content_part_added_event = True
return self.create_content_part_added_event()
return None
def is_stream_finished(self) -> bool:
if (
self.sent_output_text_done_event is True
and self.sent_output_content_part_done_event is True
and self.sent_output_item_done_event is True
):
return True
return False
def common_done_event_logic(
self, sync_mode: bool = True
) -> BaseLiteLLMOpenAIResponseObject:
if not self.litellm_model_response or isinstance(
self.litellm_model_response, TextCompletionResponse
):
self.litellm_model_response = self.create_litellm_model_response()
if self.litellm_model_response:
done_event = self.return_default_done_events(self.litellm_model_response)
if done_event:
return done_event
else:
if sync_mode:
raise StopIteration
else:
raise StopAsyncIteration
self.finished = self.is_stream_finished()
response_completed_event = self._emit_response_completed_event(
self.litellm_model_response
)
if response_completed_event:
return response_completed_event
else:
if sync_mode:
raise StopIteration
else:
raise StopAsyncIteration
async def __anext__(
self,
) -> Union[ResponsesAPIStreamingResponse, ResponseCompletedEvent]:
) -> Union[
ResponsesAPIStreamingResponse,
ResponseCompletedEvent,
BaseLiteLLMOpenAIResponseObject,
]:
try:
while True:
if self.finished is True:
raise StopAsyncIteration
result = self.return_default_initial_events()
if result:
return result
# Get the next chunk from the stream
try:
chunk = await self.litellm_custom_stream_wrapper.__anext__()
@ -70,12 +346,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if response_api_chunk:
return response_api_chunk
except StopAsyncIteration:
self.finished = True
response_completed_event = self._emit_response_completed_event()
if response_completed_event:
return response_completed_event
else:
raise StopAsyncIteration
return self.common_done_event_logic(sync_mode=False)
except Exception as e:
# Handle HTTP errors
@ -87,12 +358,20 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
def __next__(
self,
) -> Union[ResponsesAPIStreamingResponse, ResponseCompletedEvent]:
) -> Union[
ResponsesAPIStreamingResponse,
ResponseCompletedEvent,
BaseLiteLLMOpenAIResponseObject,
]:
try:
while True:
if self.finished is True:
raise StopIteration
# Get the next chunk from the stream
result = self.return_default_initial_events()
if result:
return result
try:
chunk = self.litellm_custom_stream_wrapper.__next__()
self.collected_chat_completion_chunks.append(chunk)
@ -104,13 +383,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
if response_api_chunk:
return response_api_chunk
except StopIteration:
self.finished = True
response_completed_event = self._emit_response_completed_event()
if response_completed_event:
return response_completed_event
else:
raise StopIteration
return self.common_done_event_logic(sync_mode=True)
except Exception as e:
# Handle HTTP errors
self.finished = True
@ -165,26 +438,33 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator):
chat_completion_delta: ChatCompletionDelta = choice.delta
return chat_completion_delta.content or ""
def _emit_response_completed_event(self) -> Optional[ResponseCompletedEvent]:
litellm_model_response: Optional[
Union[ModelResponse, TextCompletionResponse]
] = stream_chunk_builder(chunks=self.collected_chat_completion_chunks, logging_obj=self.litellm_logging_obj)
if litellm_model_response and isinstance(litellm_model_response, ModelResponse):
def _emit_response_completed_event(
self, litellm_model_response: ModelResponse
) -> Optional[ResponseCompletedEvent]:
if litellm_model_response:
# Add cost to usage object if include_cost_in_streaming_usage is True
if litellm.include_cost_in_streaming_usage and self.litellm_logging_obj is not None:
if (
litellm.include_cost_in_streaming_usage
and self.litellm_logging_obj is not None
):
usage = getattr(litellm_model_response, "usage", None)
if usage is not None:
setattr(
usage, "cost", self.litellm_logging_obj._response_cost_calculator(result=litellm_model_response)
usage,
"cost",
self.litellm_logging_obj._response_cost_calculator(
result=litellm_model_response
),
)
# Transform the response
responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response(
request_input=self.request_input,
chat_completion_response=litellm_model_response,
responses_api_request=self.responses_api_request,
)
# Encode the response ID to match non-streaming behavior
encoded_response = ResponsesAPIRequestUtils._update_responses_api_response_id_with_model_id(
responses_api_response=responses_api_response,

View file

@ -1,19 +1,10 @@
from litellm._uuid import uuid
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Optional,
Union,
cast,
)
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
from litellm._logging import verbose_logger
from litellm.responses.streaming_iterator import (
BaseResponsesAPIStreamingIterator,
)
from litellm._uuid import uuid
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
from litellm.types.llms.openai import (
BaseLiteLLMOpenAIResponseObject,
MCPCallArgumentsDeltaEvent,
MCPCallArgumentsDoneEvent,
MCPCallCompletedEvent,
@ -38,22 +29,24 @@ async def create_mcp_list_tools_events(
mcp_tools_with_litellm_proxy: List[ToolParam],
user_api_key_auth: Any,
base_item_id: str,
pre_processed_mcp_tools: List[Any]
pre_processed_mcp_tools: List[Any],
) -> List[ResponsesAPIStreamingResponse]:
"""Create MCP discovery events using pre-processed tools from the parent"""
events: List[ResponsesAPIStreamingResponse] = []
try:
# Extract MCP server names
mcp_servers = []
for tool in mcp_tools_with_litellm_proxy:
if isinstance(tool, dict) and "server_url" in tool:
server_url = tool.get("server_url")
if isinstance(server_url, str) and server_url.startswith("litellm_proxy/mcp/"):
if isinstance(server_url, str) and server_url.startswith(
"litellm_proxy/mcp/"
):
server_name = server_url.split("/")[-1]
mcp_servers.append(server_name)
# Emit list tools in progress event
in_progress_event = MCPListToolsInProgressEvent(
type=ResponsesAPIStreamEvents.MCP_LIST_TOOLS_IN_PROGRESS,
@ -62,21 +55,21 @@ async def create_mcp_list_tools_events(
item_id=base_item_id,
)
events.append(in_progress_event)
# Use the pre-processed MCP tools that were already fetched, filtered, and deduplicated by the parent
filtered_mcp_tools = pre_processed_mcp_tools
# Convert tools to dict format for the event
mcp_tools_dict = []
for tool in filtered_mcp_tools:
if hasattr(tool, 'model_dump') and callable(getattr(tool, 'model_dump')):
if hasattr(tool, "model_dump") and callable(getattr(tool, "model_dump")):
# Type cast to help mypy understand this is safe after hasattr check
mcp_tools_dict.append(cast(Any, tool).model_dump())
elif hasattr(tool, '__dict__'):
elif hasattr(tool, "__dict__"):
mcp_tools_dict.append(tool.__dict__)
else:
mcp_tools_dict.append({"name": getattr(tool, 'name', str(tool))})
mcp_tools_dict.append({"name": getattr(tool, "name", str(tool))})
# Emit list tools completed event
completed_event = MCPListToolsCompletedEvent(
type=ResponsesAPIStreamEvents.MCP_LIST_TOOLS_COMPLETED,
@ -85,7 +78,7 @@ async def create_mcp_list_tools_events(
item_id=base_item_id,
)
events.append(completed_event)
# Add output_item.done event with the actual tools list (matching OpenAI format)
from litellm.types.llms.openai import OutputItemDoneEvent
@ -95,45 +88,50 @@ async def create_mcp_list_tools_events(
first_tool = mcp_tools_with_litellm_proxy[0]
if isinstance(first_tool, dict):
server_label_value = first_tool.get("server_label", "")
server_label = str(server_label_value) if server_label_value is not None else ""
server_label = (
str(server_label_value) if server_label_value is not None else ""
)
# Format tools for OpenAI output_item.done format
formatted_tools = []
for tool in filtered_mcp_tools:
tool_dict = {
"name": getattr(tool, 'name', 'unknown'),
"description": getattr(tool, 'description', ''),
"name": getattr(tool, "name", "unknown"),
"description": getattr(tool, "description", ""),
"annotations": {"read_only": False},
}
# Add input_schema if available
if hasattr(tool, 'inputSchema'):
tool_dict["input_schema"] = getattr(tool, 'inputSchema')
elif hasattr(tool, 'input_schema'):
tool_dict["input_schema"] = getattr(tool, 'input_schema')
if hasattr(tool, "inputSchema"):
tool_dict["input_schema"] = getattr(tool, "inputSchema")
elif hasattr(tool, "input_schema"):
tool_dict["input_schema"] = getattr(tool, "input_schema")
formatted_tools.append(tool_dict)
# Create the output_item.done event with MCP tools list
output_item_done_event = OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=0,
item={
"id": base_item_id,
"type": "mcp_list_tools",
"server_label": server_label,
"tools": formatted_tools
}
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": base_item_id,
"type": "mcp_list_tools",
"server_label": server_label,
"tools": formatted_tools,
}
),
)
events.append(output_item_done_event)
verbose_logger.debug(f"Created {len(events)} MCP discovery events")
except Exception as e:
verbose_logger.error(f"Error creating MCP list tools events: {e}")
import traceback
traceback.print_exc()
# Emit failed event on error
failed_event = MCPListToolsFailedEvent(
type=ResponsesAPIStreamEvents.MCP_LIST_TOOLS_FAILED,
@ -142,37 +140,39 @@ async def create_mcp_list_tools_events(
item_id=base_item_id,
)
events.append(failed_event)
# Still emit output_item.done event even on failure (with empty tools list)
from litellm.types.llms.openai import OutputItemDoneEvent
output_item_done_event = OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=0,
item={
"id": base_item_id,
"type": "mcp_list_tools",
"server_label": "",
"tools": []
}
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": base_item_id,
"type": "mcp_list_tools",
"server_label": "",
"tools": [],
}
),
)
events.append(output_item_done_event)
return events
def create_mcp_call_events(
tool_name: str,
tool_call_id: str,
tool_name: str,
tool_call_id: str,
arguments: str,
result: Optional[str] = None,
base_item_id: Optional[str] = None,
sequence_start: int = 1
sequence_start: int = 1,
) -> List[ResponsesAPIStreamingResponse]:
"""Create MCP call events following OpenAI's specification"""
events: List[ResponsesAPIStreamingResponse] = []
item_id = base_item_id or f"mcp_{uuid.uuid4().hex[:8]}"
# MCP call in progress event
in_progress_event = MCPCallInProgressEvent(
type=ResponsesAPIStreamEvents.MCP_CALL_IN_PROGRESS,
@ -181,7 +181,7 @@ def create_mcp_call_events(
item_id=item_id,
)
events.append(in_progress_event)
# MCP call arguments delta event (streaming the arguments)
arguments_delta_event = MCPCallArgumentsDeltaEvent(
type=ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DELTA,
@ -191,7 +191,7 @@ def create_mcp_call_events(
sequence_number=sequence_start + 1,
)
events.append(arguments_delta_event)
# MCP call arguments done event
arguments_done_event = MCPCallArgumentsDoneEvent(
type=ResponsesAPIStreamEvents.MCP_CALL_ARGUMENTS_DONE,
@ -201,7 +201,7 @@ def create_mcp_call_events(
sequence_number=sequence_start + 2,
)
events.append(arguments_done_event)
# MCP call completed event (or failed if result indicates failure)
if result is not None:
completed_event = MCPCallCompletedEvent(
@ -211,23 +211,25 @@ def create_mcp_call_events(
output_index=0,
)
events.append(completed_event)
# Add output_item.done event with the tool call result
from litellm.types.llms.openai import OutputItemDoneEvent
output_item_done_event = OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=0,
item={
"id": item_id,
"type": "mcp_call",
"approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}",
"arguments": arguments,
"error": None,
"name": tool_name,
"output": result,
"server_label": "litellm"
},
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": item_id,
"type": "mcp_call",
"approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}",
"arguments": arguments,
"error": None,
"name": tool_name,
"output": result,
"server_label": "litellm",
}
),
)
events.append(output_item_done_event)
else:
@ -238,7 +240,7 @@ def create_mcp_call_events(
output_index=0,
)
events.append(failed_event)
return events
@ -250,51 +252,60 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
3. Handles tool execution and follow-up calls for auto-execute tools
4. Emits tool execution events in the stream
"""
def __init__(
self,
base_iterator: Any, # Can be None - will be created internally
mcp_events: List[ResponsesAPIStreamingResponse],
mcp_tools_with_litellm_proxy: Optional[List[Any]] = None,
user_api_key_auth: Any = None,
original_request_params: Optional[Dict[str, Any]] = None
original_request_params: Optional[Dict[str, Any]] = None,
):
# MCP setup
self.mcp_tools_with_litellm_proxy = mcp_tools_with_litellm_proxy or []
self.user_api_key_auth = user_api_key_auth
self.original_request_params = original_request_params or {}
self.should_auto_execute = self._should_auto_execute_tools()
# Streaming state management
self.phase = "mcp_discovery" # mcp_discovery -> initial_response -> tool_execution -> follow_up_response -> finished
self.finished = False
# Event queues and generation flags
self.mcp_discovery_events: List[ResponsesAPIStreamingResponse] = mcp_events # Pre-generated MCP discovery events
self.mcp_discovery_events: List[ResponsesAPIStreamingResponse] = (
mcp_events # Pre-generated MCP discovery events
)
self.tool_execution_events: List[ResponsesAPIStreamingResponse] = []
self.mcp_discovery_generated = True # Events are already generated
self.mcp_events = mcp_events # Store the initial MCP events for backward compatibility
self.mcp_events = (
mcp_events # Store the initial MCP events for backward compatibility
)
# Iterator references
self.base_iterator: Optional[Union[Any, ResponsesAPIResponse]] = base_iterator # Will be created when needed
self.base_iterator: Optional[Union[Any, ResponsesAPIResponse]] = (
base_iterator # Will be created when needed
)
self.follow_up_iterator: Optional[Any] = None
# Response collection for tool execution
self.collected_response: Optional[ResponsesAPIResponse] = None
# Set up model metadata (will be updated when we get the real iterator)
self.model = self.original_request_params.get('model', 'unknown')
self.model = self.original_request_params.get("model", "unknown")
self.litellm_metadata = {}
self.custom_llm_provider = self.original_request_params.get('custom_llm_provider', None)
self.custom_llm_provider = self.original_request_params.get(
"custom_llm_provider", None
)
# Mark as async iterator
self.is_async = True
def _should_auto_execute_tools(self) -> bool:
"""Check if tools should be auto-executed"""
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
return LiteLLM_Proxy_MCP_Handler._should_auto_execute_tools(
self.mcp_tools_with_litellm_proxy
)
@ -306,45 +317,49 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
"""
Phase-based streaming:
1. mcp_discovery - Emit MCP discovery events
2. initial_response - Stream the first LLM response
2. initial_response - Stream the first LLM response
3. tool_execution - Emit tool execution events
4. follow_up_response - Stream the follow-up response
5. finished - End iteration
"""
# Phase 1: MCP Discovery Events
if self.phase == "mcp_discovery":
# Generate MCP discovery events if not already done
# MCP discovery events are already generated and available
# Emit MCP discovery events
if self.mcp_discovery_events:
return self.mcp_discovery_events.pop(0)
# All MCP discovery events emitted, move to next phase
verbose_logger.debug("MCP discovery phase complete, transitioning to initial_response")
verbose_logger.debug(
"MCP discovery phase complete, transitioning to initial_response"
)
self.phase = "initial_response"
await self._create_initial_response_iterator()
# Fall through to process the initial response immediately
# Phase 2: Initial Response Stream
if self.phase == "initial_response":
if self.base_iterator:
# Check if base_iterator is actually iterable
if hasattr(self.base_iterator, '__anext__'):
if hasattr(self.base_iterator, "__anext__"):
try:
chunk = await cast(Any, self.base_iterator).__anext__() # type: ignore[attr-defined]
# If auto-execution is enabled, check for completed responses
if self.should_auto_execute and self._is_response_completed(chunk):
if self.should_auto_execute and self._is_response_completed(
chunk
):
# Collect the response for tool execution
response_obj = getattr(chunk, 'response', None)
response_obj = getattr(chunk, "response", None)
if isinstance(response_obj, ResponsesAPIResponse):
self.collected_response = response_obj
# Move to tool execution phase after emitting this chunk
self.phase = "tool_execution"
await self._generate_tool_execution_events()
return chunk
except StopAsyncIteration:
# Initial response ended, move to next phase
@ -357,24 +372,26 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
else:
# base_iterator is not async iterable (likely a ResponsesAPIResponse)
# Collect it for tool execution if needed
if self.should_auto_execute and isinstance(self.base_iterator, ResponsesAPIResponse):
if self.should_auto_execute and isinstance(
self.base_iterator, ResponsesAPIResponse
):
self.collected_response = self.base_iterator
self.phase = "tool_execution"
await self._generate_tool_execution_events()
else:
self.phase = "finished"
raise StopAsyncIteration
# Phase 3: Tool Execution Events
if self.phase == "tool_execution":
# Emit any queued tool execution events
if self.tool_execution_events:
return self.tool_execution_events.pop(0)
# Move to follow-up response phase
self.phase = "follow_up_response"
await self._create_follow_up_iterator()
# Phase 4: Follow-up Response Stream
if self.phase == "follow_up_response":
if self.follow_up_iterator:
@ -386,20 +403,22 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
else:
self.phase = "finished"
raise StopAsyncIteration
# Phase 5: Finished
if self.phase == "finished":
raise StopAsyncIteration
# Should not reach here
raise StopAsyncIteration
def _is_response_completed(self, chunk: ResponsesAPIStreamingResponse) -> bool:
"""Check if this chunk indicates the response is completed"""
from litellm.types.llms.openai import ResponsesAPIStreamEvents
return getattr(chunk, 'type', None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
return (
getattr(chunk, "type", None) == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
)
async def _create_initial_response_iterator(self) -> None:
"""Create the initial response iterator by making the first LLM call"""
try:
@ -408,38 +427,45 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
# Make the initial response API call - but avoid the MCP wrapper
params = self.original_request_params.copy()
params['stream'] = True # Ensure streaming
params["stream"] = True # Ensure streaming
# Use the pre-fetched all_tools from original_request_params (no re-processing needed)
params_for_llm = {}
for key, value in params.items():
params_for_llm[key] = value # Copy all params as-is since tools are already processed
tools_count = len(params_for_llm.get('tools', []))
params_for_llm[key] = (
value # Copy all params as-is since tools are already processed
)
tools_count = len(params_for_llm.get("tools", []))
verbose_logger.debug(f"Making LLM call with {tools_count} tools")
response = await aresponses(**params_for_llm)
# Set the base iterator
if hasattr(response, '__aiter__') or hasattr(response, '__iter__'):
if hasattr(response, "__aiter__") or hasattr(response, "__iter__"):
self.base_iterator = response
# Copy metadata from the real iterator
self.model = getattr(response, 'model', self.model)
self.litellm_metadata = getattr(response, 'litellm_metadata', {})
self.custom_llm_provider = getattr(response, 'custom_llm_provider', self.custom_llm_provider)
verbose_logger.debug(f"Created base iterator: {type(self.base_iterator)}")
self.model = getattr(response, "model", self.model)
self.litellm_metadata = getattr(response, "litellm_metadata", {})
self.custom_llm_provider = getattr(
response, "custom_llm_provider", self.custom_llm_provider
)
verbose_logger.debug(
f"Created base iterator: {type(self.base_iterator)}"
)
else:
# Non-streaming response - this shouldn't happen but handle it
verbose_logger.warning(f"Got non-streaming response: {type(response)}")
self.base_iterator = None
self.phase = "finished"
except Exception as e:
verbose_logger.error(f"Error creating initial response iterator: {e}")
import traceback
traceback.print_exc()
self.base_iterator = None
self.phase = "finished"
async def _generate_tool_execution_events(self) -> None:
"""Generate tool execution events and execute tools"""
if not self.collected_response:
@ -447,7 +473,7 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
try:
# Extract tool calls from the response
if self.collected_response is not None:
@ -456,9 +482,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
tool_calls = []
if not tool_calls:
return
for tool_call in tool_calls:
tool_name, tool_arguments, tool_call_id = LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call)
tool_name, tool_arguments, tool_call_id = (
LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call)
)
if tool_name and tool_call_id:
# Create MCP call events for this tool execution
call_events = create_mcp_call_events(
@ -467,34 +495,35 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
arguments=tool_arguments or "{}", # JSON string with arguments
result=None, # Will be set after execution
base_item_id=f"mcp_{uuid.uuid4().hex[:8]}",
sequence_start=len(self.tool_execution_events) + 1
sequence_start=len(self.tool_execution_events) + 1,
)
# Add the in_progress and arguments events (not the completed event yet)
self.tool_execution_events.extend(call_events[:-1])
# Execute the tools
tool_results = await LiteLLM_Proxy_MCP_Handler._execute_tool_calls(
tool_calls=tool_calls,
user_api_key_auth=self.user_api_key_auth
tool_calls=tool_calls, user_api_key_auth=self.user_api_key_auth
)
# Create completion events and output_item.done events for tool execution
for tool_result in tool_results:
tool_call_id = tool_result.get("tool_call_id", "unknown")
result_text = tool_result.get("result", "")
# Find matching tool name and arguments
tool_name = "unknown"
tool_arguments = "{}"
for tool_call in tool_calls:
name, args, call_id = LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call)
name, args, call_id = (
LiteLLM_Proxy_MCP_Handler._extract_tool_call_details(tool_call)
)
if call_id == tool_call_id:
tool_name = name or "unknown"
tool_arguments = args or "{}"
break
item_id = f"mcp_{uuid.uuid4().hex[:8]}"
# Create the completion event
completed_event = MCPCallCompletedEvent(
type=ResponsesAPIStreamEvents.MCP_CALL_COMPLETED,
@ -503,79 +532,84 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
output_index=0,
)
self.tool_execution_events.append(completed_event)
# Create output_item.done event with the tool call result
from litellm.types.llms.openai import OutputItemDoneEvent
output_item_done_event = OutputItemDoneEvent(
type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
output_index=0,
item={
"id": item_id,
"type": "mcp_call",
"approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}",
"arguments": tool_arguments,
"error": None,
"name": tool_name,
"output": result_text,
"server_label": "litellm" # or extract from tool config
},
item=BaseLiteLLMOpenAIResponseObject(
**{
"id": item_id,
"type": "mcp_call",
"approval_request_id": f"mcpr_{uuid.uuid4().hex[:8]}",
"arguments": tool_arguments,
"error": None,
"name": tool_name,
"output": result_text,
"server_label": "litellm", # or extract from tool config
}
),
)
self.tool_execution_events.append(output_item_done_event)
# Store tool results for follow-up call
self.tool_results = tool_results
except Exception as e:
verbose_logger.error(f"Error in tool execution: {e}")
import traceback
traceback.print_exc()
self.tool_results = []
async def _create_follow_up_iterator(self) -> None:
"""Create the follow-up response iterator with tool results"""
if not self.collected_response or not hasattr(self, 'tool_results'):
if not self.collected_response or not hasattr(self, "tool_results"):
return
from litellm.responses.main import aresponses
from litellm.responses.mcp.litellm_proxy_mcp_handler import (
LiteLLM_Proxy_MCP_Handler,
)
try:
# Create follow-up input
if self.collected_response is not None:
follow_up_input = LiteLLM_Proxy_MCP_Handler._create_follow_up_input(
response=self.collected_response, # type: ignore[arg-type]
tool_results=self.tool_results,
original_input=self.original_request_params.get('input')
original_input=self.original_request_params.get("input"),
)
# Make follow-up call with streaming
follow_up_params = self.original_request_params.copy()
follow_up_params.update({
'input': follow_up_input,
'previous_response_id': self.collected_response.id, # type: ignore[attr-defined]
'stream': True
})
follow_up_params.update(
{
"input": follow_up_input,
"previous_response_id": self.collected_response.id, # type: ignore[attr-defined]
"stream": True,
}
)
else:
return
# Remove tool_choice to avoid forcing more tool calls
follow_up_params.pop('tool_choice', None)
follow_up_params.pop("tool_choice", None)
follow_up_response = await aresponses(**follow_up_params)
# Set up the follow-up iterator
if hasattr(follow_up_response, '__aiter__'):
if hasattr(follow_up_response, "__aiter__"):
self.follow_up_iterator = follow_up_response
except Exception as e:
verbose_logger.error(f"Error creating follow-up iterator: {e}")
import traceback
traceback.print_exc()
self.follow_up_iterator = None
def __iter__(self):
return self
@ -583,11 +617,11 @@ class MCPEnhancedStreamingIterator(BaseResponsesAPIStreamingIterator):
# First, emit any queued MCP events
if self.mcp_events: # type: ignore[attr-defined]
return self.mcp_events.pop(0) # type: ignore[attr-defined]
# Then delegate to the base iterator
if not self.is_async:
try:
if self.base_iterator and hasattr(self.base_iterator, '__next__'):
if self.base_iterator and hasattr(self.base_iterator, "__next__"):
return next(cast(Any, self.base_iterator)) # type: ignore[arg-type]
else:
raise StopIteration

View file

@ -93,24 +93,35 @@ class BaseResponsesAPIStreamingIterator:
# Store the completed response
if (
openai_responses_api_chunk
and openai_responses_api_chunk.type
and getattr(openai_responses_api_chunk, "type", None)
== ResponsesAPIStreamEvents.RESPONSE_COMPLETED
):
self.completed_response = openai_responses_api_chunk
# Add cost to usage object if include_cost_in_streaming_usage is True
if litellm.include_cost_in_streaming_usage and self.logging_obj is not None:
response_obj: Optional[ResponsesAPIResponse] = getattr(openai_responses_api_chunk, "response", None)
if (
litellm.include_cost_in_streaming_usage
and self.logging_obj is not None
):
response_obj: Optional[ResponsesAPIResponse] = getattr(
openai_responses_api_chunk, "response", None
)
if response_obj:
usage_obj: Optional[ResponseAPIUsage] = getattr(response_obj, "usage", None)
usage_obj: Optional[ResponseAPIUsage] = getattr(
response_obj, "usage", None
)
if usage_obj is not None:
try:
cost: Optional[float] = self.logging_obj._response_cost_calculator(result=response_obj)
cost: Optional[float] = (
self.logging_obj._response_cost_calculator(
result=response_obj
)
)
if cost is not None:
setattr(usage_obj, "cost", cost)
except Exception:
# If cost calculation fails, continue without cost
pass
self._handle_logging_completed_response()
return openai_responses_api_chunk

View file

@ -43,12 +43,7 @@ from openai.types.responses.response import (
# Handle OpenAI SDK version compatibility for Text type
try:
# fmt: off
from openai.types.responses.response_create_params import ( # type: ignore[attr-defined]
Text as ResponseText, # type: ignore[attr-defined]
)
# fmt: on
from openai.types.responses.response_create_params import ( Text as ResponseText ) # type: ignore[attr-defined] # fmt: skip # isort: skip
except (ImportError, AttributeError):
# Fall back to the concrete config type available in all SDK versions
from openai.types.responses.response_text_config_param import (
@ -992,7 +987,9 @@ class ResponsesAPIOptionalRequestParams(TypedDict, total=False):
prompt_cache_key: Optional[str]
stream_options: Optional[dict]
top_logprobs: Optional[int]
partial_images: Optional[int] # Number of partial images to generate (1-3) for streaming image generation
partial_images: Optional[
int
] # Number of partial images to generate (1-3) for streaming image generation
class ResponsesAPIRequestParams(ResponsesAPIOptionalRequestParams, total=False):
@ -1056,7 +1053,9 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject):
parallel_tool_calls: Optional[bool] = None
temperature: Optional[float] = None
tool_choice: Optional[ToolChoice] = None
tools: Optional[Union[List[Tool], List[ResponseFunctionToolCall], List[Dict[str, Any]]]] = None
tools: Optional[
Union[List[Tool], List[ResponseFunctionToolCall], List[Dict[str, Any]]]
] = None
top_p: Optional[float] = None
max_output_tokens: Optional[int] = None
previous_response_id: Optional[str] = None
@ -1180,13 +1179,27 @@ class ReasoningSummaryTextDeltaEvent(BaseLiteLLMOpenAIResponseObject):
class OutputItemAddedEvent(BaseLiteLLMOpenAIResponseObject):
type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED]
output_index: int
item: Optional[dict]
item: Optional[BaseLiteLLMOpenAIResponseObject]
class OutputItemDoneEvent(BaseLiteLLMOpenAIResponseObject):
type: Literal[ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE]
output_index: int
item: dict
sequence_number: int = 1
item: BaseLiteLLMOpenAIResponseObject
class OpenAIChatCompletionLogprobsContentTopLogprobs(TypedDict, total=False):
bytes: List
logprob: Required[float]
token: Required[str]
class OpenAIChatCompletionLogprobsContent(TypedDict, total=False):
bytes: List
logprob: Required[float]
token: Required[str]
top_logprobs: List[OpenAIChatCompletionLogprobsContentTopLogprobs]
class ContentPartAddedEvent(BaseLiteLLMOpenAIResponseObject):
@ -1194,7 +1207,31 @@ class ContentPartAddedEvent(BaseLiteLLMOpenAIResponseObject):
item_id: str
output_index: int
content_index: int
part: dict
part: BaseLiteLLMOpenAIResponseObject
class ContentPartDonePartOutputText(BaseLiteLLMOpenAIResponseObject):
type: Literal["output_text"]
text: str
annotations: List[BaseLiteLLMOpenAIResponseObject]
logprobs: Optional[List[OpenAIChatCompletionLogprobsContent]]
class ContentPartDonePartRefusal(BaseLiteLLMOpenAIResponseObject):
type: Literal["refusal"]
refusal: str
class ContentPartDonePartReasoningText(BaseLiteLLMOpenAIResponseObject):
type: Literal["reasoning_text"]
reasoning: str
PART_UNION_TYPES = Union[
ContentPartDonePartOutputText,
ContentPartDonePartRefusal,
ContentPartDonePartReasoningText,
]
class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject):
@ -1202,7 +1239,7 @@ class ContentPartDoneEvent(BaseLiteLLMOpenAIResponseObject):
item_id: str
output_index: int
content_index: int
part: dict
part: PART_UNION_TYPES
class OutputTextDeltaEvent(BaseLiteLLMOpenAIResponseObject):
@ -1413,6 +1450,7 @@ ResponsesAPIStreamingResponse = Annotated[
ImageGenerationPartialImageEvent,
ErrorEvent,
GenericEvent,
BaseLiteLLMOpenAIResponseObject,
],
Discriminator("type"),
]
@ -1731,19 +1769,6 @@ class OpenAIModerationResponse(BaseLiteLLMOpenAIResponseObject):
_hidden_params: dict = PrivateAttr(default_factory=dict)
class OpenAIChatCompletionLogprobsContentTopLogprobs(TypedDict, total=False):
bytes: List
logprob: Required[float]
token: Required[str]
class OpenAIChatCompletionLogprobsContent(TypedDict, total=False):
bytes: List
logprob: Required[float]
token: Required[str]
top_logprobs: List[OpenAIChatCompletionLogprobsContentTopLogprobs]
class OpenAIChatCompletionLogprobs(TypedDict, total=False):
content: List[OpenAIChatCompletionLogprobsContent]
refusal: List[OpenAIChatCompletionLogprobsContent]

78
log.txt Normal file
View file

@ -0,0 +1,78 @@
============================= test session starts ==============================
platform darwin -- Python 3.11.4, pytest-7.4.1, pluggy-1.2.0
rootdir: /Users/krrishdholakia/Documents/litellm
plugins: snapshot-0.9.0, cov-5.0.0, timeout-2.2.0, postgresql-7.0.1, respx-0.21.1, asyncio-0.21.1, langsmith-0.3.4, anyio-4.8.0, mock-3.11.1, Faker-25.9.2
asyncio: mode=Mode.STRICT
collected 1 item
tests/llm_translation/test_gemini.py . [100%]
=============================== warnings summary ===============================
tests/llm_translation/base_llm_unit_tests.py:481
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:481: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=6, delay=1)
tests/llm_translation/base_llm_unit_tests.py:523
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:523: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=6, delay=1)
tests/llm_translation/base_llm_unit_tests.py:601
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:601: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=6, delay=1)
tests/llm_translation/base_llm_unit_tests.py:641
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:641: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=6, delay=1)
tests/llm_translation/base_llm_unit_tests.py:650
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:650: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=6, delay=1)
tests/llm_translation/base_llm_unit_tests.py:694
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:694: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=6, delay=1)
tests/llm_translation/base_llm_unit_tests.py:745
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:745: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=6, delay=1)
tests/llm_translation/base_llm_unit_tests.py:783
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:783: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=6, delay=1)
tests/llm_translation/base_llm_unit_tests.py:859
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:859: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=4, delay=2)
tests/llm_translation/base_llm_unit_tests.py:955
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:955: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=4, delay=1)
tests/llm_translation/base_llm_unit_tests.py:1073
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:1073: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=3, delay=1)
tests/llm_translation/base_llm_unit_tests.py:1109
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:1109: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=3, delay=1)
tests/llm_translation/base_llm_unit_tests.py:1232
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/base_llm_unit_tests.py:1232: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=3, delay=1)
tests/llm_translation/test_gemini.py:36
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/test_gemini.py:36: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=3, delay=2)
tests/llm_translation/test_gemini.py:510
/Users/krrishdholakia/Documents/litellm/tests/llm_translation/test_gemini.py:510: PytestUnknownMarkWarning: Unknown pytest.mark.flaky - is this a typo? You can register custom marks to avoid this warning - for details, see https://docs.pytest.org/en/stable/how-to/mark.html
@pytest.mark.flaky(retries=3, delay=2)
tests/llm_translation/test_gemini.py::test_gemini_image_generation_async
/Users/krrishdholakia/Library/Python/3.11/lib/python/site-packages/pydantic/main.py:463: UserWarning: Pydantic serializer warnings:
PydanticSerializationUnexpectedValue(Expected 10 fields but got 7: Expected `Message` - serialized value may not be as expected [input_value=Message(content="Here's t...er_specific_fields=None), input_type=Message])
PydanticSerializationUnexpectedValue(Expected `StreamingChoices` - serialized value may not be as expected [input_value=Choices(finish_reason='st...r_specific_fields=None)), input_type=Choices])
return self.__pydantic_serializer__.to_python(
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
======================== 1 passed, 16 warnings in 5.86s ========================

View file

@ -361,6 +361,7 @@ def test_process_anthropic_headers_with_no_matching_headers():
def test_anthropic_tool_use(tool_type, tool_config, message_content):
"""Test Anthropic tool use with computer use and web fetch tools."""
from litellm import completion
litellm._turn_on_debug()
tools = [tool_config]
@ -1518,3 +1519,126 @@ def test_anthropic_streaming():
role_set_count += 1
assert role_set_count == 1
def test_anthropic_via_responses_api():
from litellm.types.llms.openai import ResponsesAPIStreamEvents
response = litellm.responses(
model="anthropic/claude-sonnet-4-5",
input="Who won the World Cup in 2022?",
max_output_tokens=100,
stream=True,
)
assert response is not None
# Expected event sequence
expected_events = [
ResponsesAPIStreamEvents.RESPONSE_CREATED,
ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS,
ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED,
ResponsesAPIStreamEvents.CONTENT_PART_ADDED,
ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, # Can occur multiple times
ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE,
ResponsesAPIStreamEvents.CONTENT_PART_DONE,
ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE,
ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
]
events_seen = []
text_delta_count = 0
for chunk in response:
print(f"chunk: {chunk}")
# Each chunk should have a type attribute
assert hasattr(chunk, "type"), f"Chunk missing 'type' attribute: {chunk}"
event_type = chunk.type
# Track events seen
if event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA:
text_delta_count += 1
if ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA not in events_seen:
events_seen.append(event_type)
else:
events_seen.append(event_type)
# Assert specific structures for each event type
if event_type == ResponsesAPIStreamEvents.RESPONSE_CREATED:
assert chunk.type == ResponsesAPIStreamEvents.RESPONSE_CREATED
assert hasattr(chunk, "response")
assert chunk.response.status == "in_progress"
assert hasattr(chunk.response, "id")
assert hasattr(chunk.response, "model")
elif event_type == ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS:
assert chunk.type == ResponsesAPIStreamEvents.RESPONSE_IN_PROGRESS
assert hasattr(chunk, "response")
assert chunk.response.status == "in_progress"
elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED:
assert chunk.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED
assert hasattr(chunk, "output_index")
assert hasattr(chunk, "item")
assert chunk.item.type == "message"
assert chunk.item.role == "assistant"
elif event_type == ResponsesAPIStreamEvents.CONTENT_PART_ADDED:
assert chunk.type == ResponsesAPIStreamEvents.CONTENT_PART_ADDED
assert hasattr(chunk, "item_id")
assert hasattr(chunk, "output_index")
assert hasattr(chunk, "content_index")
assert hasattr(chunk, "part")
assert chunk.part.type == "output_text"
elif event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA:
assert chunk.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA
assert hasattr(chunk, "item_id")
assert hasattr(chunk, "output_index")
assert hasattr(chunk, "content_index")
assert hasattr(chunk, "delta")
assert isinstance(chunk.delta, str)
elif event_type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE:
assert chunk.type == ResponsesAPIStreamEvents.OUTPUT_TEXT_DONE
assert hasattr(chunk, "item_id")
assert hasattr(chunk, "output_index")
assert hasattr(chunk, "content_index")
assert hasattr(chunk, "text")
elif event_type == ResponsesAPIStreamEvents.CONTENT_PART_DONE:
assert chunk.type == ResponsesAPIStreamEvents.CONTENT_PART_DONE
assert hasattr(chunk, "item_id")
assert hasattr(chunk, "output_index")
assert hasattr(chunk, "content_index")
assert hasattr(chunk, "part")
assert chunk.part.type == "output_text"
elif event_type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE:
assert chunk.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE
assert hasattr(chunk, "output_index")
assert hasattr(chunk, "item")
assert chunk.item.status == "completed"
elif event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED:
assert chunk.type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
assert hasattr(chunk, "response")
assert chunk.response.status == "completed"
assert hasattr(chunk.response, "usage")
assert hasattr(chunk.response, "output")
# Assert we saw all expected events
print(f"Events seen: {events_seen}")
assert (
events_seen == expected_events
), f"Event sequence mismatch. Expected: {expected_events}, Got: {events_seen}"
# Assert we saw at least one text delta
assert (
text_delta_count > 0
), f"Expected at least one response.output_text.delta event, got {text_delta_count}"
print(f"✓ All {len(events_seen)} events matched expected structure")
print(f"✓ Received {text_delta_count} text delta chunks")

View file

@ -728,22 +728,6 @@ def test_openai_safety_identifier_parameter_sync():
assert request_body["safety_identifier"] == "user_code_123456"
def test_gpt_5_reasoning():
litellm._turn_on_debug()
response = litellm.completion(
model="openai/responses/gpt-5-mini",
messages=[
{
"role": "user",
"content": "Think of the capital of France, and then write it.",
}
],
reasoning_effort="low",
)
print("response: ", response)
assert response.choices[0].message.reasoning_content is not None
def test_gpt_5_reasoning_streaming():
litellm._turn_on_debug()
response = litellm.completion(

View file

@ -9,9 +9,12 @@ from openai import AssistantEventHandler
client = openai.OpenAI(base_url="http://0.0.0.0:4000/openai", api_key="sk-1234")
def test_pass_through_file_operations():
# Create a temporary file
with tempfile.NamedTemporaryFile(mode='w+', suffix='.txt', delete=False) as temp_file:
with tempfile.NamedTemporaryFile(
mode="w+", suffix=".txt", delete=False
) as temp_file:
temp_file.write("This is a test file for the OpenAI Assistants API.")
temp_file.flush()
@ -26,6 +29,7 @@ def test_pass_through_file_operations():
delete_file = client.files.delete(file.id)
print("file deleted", delete_file)
def test_openai_assistants_e2e_operations():
assistant = client.beta.assistants.create(
name="Math Tutor",
@ -98,13 +102,13 @@ def test_openai_assistants_e2e_operations_stream():
stream.until_done()
def test_azure_openai_assistants_e2e_operations_stream():
from openai import AzureOpenAI
client = AzureOpenAI(
base_url="http://0.0.0.0:4000/azure-config-passthrough/openai",
base_url="http://0.0.0.0:4000/azure-config-passthrough/openai",
api_key="sk-1234",
api_version="2025-01-01-preview"
api_version="2025-01-01-preview",
)
assistant = client.beta.assistants.create(
name="Math Tutor",
@ -134,4 +138,4 @@ def test_azure_openai_assistants_e2e_operations_stream():
instructions="Please address the user as Jane Doe. The user has a premium account.",
event_handler=EventHandler(),
) as stream:
stream.until_done()
stream.until_done()

View file

@ -1222,6 +1222,104 @@ class TestMCPServerManager:
"Contact proxy admin to allow this tool" in exc_info.value.detail["error"]
)
@pytest.mark.asyncio
async def test_call_tool_without_broken_pipe_error(self):
"""
Test that call_tool properly uses async context manager to avoid broken pipe errors.
This test ensures that tasks are awaited INSIDE the context manager, keeping the connection alive.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from mcp.types import CallToolResult
manager = MCPServerManager()
# Create a test server
server = MCPServer(
server_id="test-server",
name="test-server",
transport=MCPTransport.http,
url="http://test-server.com",
)
# Register the server and map a tool to it
manager.registry = {"test-server": server}
manager.tool_name_to_mcp_server_name_mapping["test_tool"] = "test-server"
# Create mock client that tracks context manager usage
mock_client = MagicMock()
context_entered = False
context_exited = False
call_tool_called_inside_context = False
async def mock_aenter(self):
nonlocal context_entered
context_entered = True
return self
async def mock_aexit(self, exc_type, exc_val, exc_tb):
nonlocal context_exited
context_exited = True
# Verify that call_tool was called before context exit
assert (
call_tool_called_inside_context
), "call_tool must be awaited inside context manager"
return False
async def mock_call_tool(params):
nonlocal call_tool_called_inside_context
# Verify we're inside the context when this is called
assert context_entered, "call_tool called outside context manager"
assert not context_exited, "call_tool called after context exit"
call_tool_called_inside_context = True
# Return a mock CallToolResult
result = MagicMock(spec=CallToolResult)
result.content = [{"type": "text", "text": "Tool executed successfully"}]
result.isError = False
return result
mock_client.__aenter__ = mock_aenter
mock_client.__aexit__ = mock_aexit
mock_client.call_tool = mock_call_tool
# Mock _create_mcp_client to return our mock client
manager._create_mcp_client = MagicMock(return_value=mock_client)
# Mock user auth with no restrictions
user_api_key_auth = MagicMock()
user_api_key_auth.object_permission = None
user_api_key_auth.object_permission_id = None
# Mock proxy logging
proxy_logging_obj = MagicMock()
proxy_logging_obj._create_mcp_request_object_from_kwargs = MagicMock(
return_value={}
)
proxy_logging_obj._convert_mcp_to_llm_format = MagicMock(return_value={})
proxy_logging_obj.pre_call_hook = AsyncMock(return_value={})
proxy_logging_obj.during_call_hook = AsyncMock(return_value=None)
# Call the tool
result = await manager.call_tool(
name="test_tool",
arguments={"param": "value"},
user_api_key_auth=user_api_key_auth,
proxy_logging_obj=proxy_logging_obj,
)
# Verify the result
assert result is not None
assert result.isError is False
assert len(result.content) > 0
# Verify context manager was used properly
assert context_entered, "Context manager __aenter__ was not called"
assert context_exited, "Context manager __aexit__ was not called"
assert (
call_tool_called_inside_context
), "call_tool was not awaited inside context"
if __name__ == "__main__":
pytest.main([__file__])

View file

@ -130,22 +130,24 @@ def test_virtual_key_allowed_routes_with_litellm_routes_member_name_denied():
assert "Only allowed to call routes: ['info_routes']" in str(exc_info.value)
assert "Tried to call route: /chat/completions" in str(exc_info.value)
@pytest.mark.parametrize("route", [
"/anthropic/v1/messages",
"/anthropic/v1/count_tokens",
"/gemini/v1/models",
"/gemini/countTokens",
])
@pytest.mark.parametrize(
"route",
[
"/anthropic/v1/messages",
"/anthropic/v1/count_tokens",
"/gemini/v1/models",
"/gemini/countTokens",
],
)
def test_virtual_key_llm_api_route_includes_passthrough_prefix(route):
"""
Virtual key with llm_api_routes should allow passthrough routes like /anthropic/v1/messages
Relevant issue: https://github.com/BerriAI/litellm/issues/14017
"""
valid_token = UserAPIKeyAuth(
user_id="test_user", allowed_routes=["llm_api_routes"]
)
valid_token = UserAPIKeyAuth(user_id="test_user", allowed_routes=["llm_api_routes"])
result = RouteChecks.is_virtual_key_allowed_to_call_route(
route=route, valid_token=valid_token
@ -232,7 +234,7 @@ def test_virtual_key_allowed_routes_with_no_member_names_only_explicit():
def test_anthropic_count_tokens_route_is_llm_api_route():
"""Test that /v1/messages/count_tokens is recognized as an LLM API route for Anthropic"""
# Test the core anthropic routes
assert RouteChecks.is_llm_api_route("/v1/messages") is True
assert RouteChecks.is_llm_api_route("/v1/messages/count_tokens") is True
@ -240,11 +242,11 @@ def test_anthropic_count_tokens_route_is_llm_api_route():
def test_anthropic_count_tokens_route_accessible_to_internal_users():
"""Test that internal users can access the Anthropic count_tokens route"""
# Test that the route is recognized as an LLM API route (which means it's accessible to internal users)
# This is the core check that was failing in the original issue
assert RouteChecks.is_llm_api_route("/v1/messages/count_tokens") is True
# Also test that the regular messages route still works
assert RouteChecks.is_llm_api_route("/v1/messages") is True
@ -252,7 +254,7 @@ def test_anthropic_count_tokens_route_accessible_to_internal_users():
def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints():
"""
Test that virtual keys with llm_api_routes permission can access registered pass-through endpoints.
This tests the scenario where a pass-through endpoint is registered from the DB
(e.g., /azure-assistant) and a virtual key with llm_api_routes permission should be able to access
both the exact path and subpaths (e.g., /azure-assistant/openai/assistants).
@ -272,7 +274,7 @@ def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints():
"type": "subpath",
},
}
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
mock_registered_routes,
@ -282,21 +284,21 @@ def test_virtual_key_llm_api_routes_allows_registered_pass_through_endpoints():
user_id="test_user",
allowed_routes=["llm_api_routes"],
)
# Test exact match for registered pass-through endpoint
result1 = RouteChecks.is_virtual_key_allowed_to_call_route(
route="/azure-assistant",
valid_token=valid_token,
)
assert result1 is True
# Test subpath for registered pass-through endpoint with subpath type
result2 = RouteChecks.is_virtual_key_allowed_to_call_route(
route="/custom-endpoint/openai/assistants",
valid_token=valid_token,
)
assert result2 is True
# Test exact match for subpath type
result3 = RouteChecks.is_virtual_key_allowed_to_call_route(
route="/custom-endpoint",
@ -319,7 +321,7 @@ def test_virtual_key_without_llm_api_routes_cannot_access_pass_through():
"type": "exact",
},
}
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes",
mock_registered_routes,
@ -329,12 +331,289 @@ def test_virtual_key_without_llm_api_routes_cannot_access_pass_through():
user_id="test_user",
allowed_routes=["info_routes"],
)
# Test that access is denied
with pytest.raises(Exception) as exc_info:
RouteChecks.is_virtual_key_allowed_to_call_route(
route="/azure-assistant",
valid_token=valid_token,
)
assert "Virtual key is not allowed to call this route" in str(exc_info.value)
def test_check_passthrough_route_access_key_metadata_exact_match():
"""Test that key metadata allowed_passthrough_routes allows exact match"""
# Create a UserAPIKeyAuth with allowed_passthrough_routes in metadata
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={"allowed_passthrough_routes": ["/custom-endpoint"]},
)
# Test exact match
result = RouteChecks.check_passthrough_route_access(
route="/custom-endpoint",
user_api_key_dict=valid_token,
)
assert result is True
def test_check_passthrough_route_access_key_metadata_prefix_match():
"""Test that key metadata allowed_passthrough_routes allows prefix match"""
# Create a UserAPIKeyAuth with allowed_passthrough_routes in metadata
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={"allowed_passthrough_routes": ["/custom-endpoint"]},
)
# Test prefix match
result = RouteChecks.check_passthrough_route_access(
route="/custom-endpoint/v1/chat/completions",
user_api_key_dict=valid_token,
)
assert result is True
def test_check_passthrough_route_access_key_metadata_no_match():
"""Test that key metadata allowed_passthrough_routes denies non-matching routes"""
# Create a UserAPIKeyAuth with allowed_passthrough_routes in metadata
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={"allowed_passthrough_routes": ["/custom-endpoint"]},
)
# Test non-matching route
result = RouteChecks.check_passthrough_route_access(
route="/other-endpoint",
user_api_key_dict=valid_token,
)
assert result is False
def test_check_passthrough_route_access_team_metadata_exact_match():
"""Test that team metadata allowed_passthrough_routes allows exact match"""
# Create a UserAPIKeyAuth with allowed_passthrough_routes in team_metadata
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={},
team_metadata={"allowed_passthrough_routes": ["/team-endpoint"]},
)
# Test exact match
result = RouteChecks.check_passthrough_route_access(
route="/team-endpoint",
user_api_key_dict=valid_token,
)
assert result is True
def test_check_passthrough_route_access_team_metadata_prefix_match():
"""Test that team metadata allowed_passthrough_routes allows prefix match"""
# Create a UserAPIKeyAuth with allowed_passthrough_routes in team_metadata
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={},
team_metadata={"allowed_passthrough_routes": ["/team-endpoint"]},
)
# Test prefix match
result = RouteChecks.check_passthrough_route_access(
route="/team-endpoint/v1/messages",
user_api_key_dict=valid_token,
)
assert result is True
def test_check_passthrough_route_access_team_metadata_no_match():
"""Test that team metadata allowed_passthrough_routes denies non-matching routes"""
# Create a UserAPIKeyAuth with allowed_passthrough_routes in team_metadata
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={},
team_metadata={"allowed_passthrough_routes": ["/team-endpoint"]},
)
# Test non-matching route
result = RouteChecks.check_passthrough_route_access(
route="/other-endpoint",
user_api_key_dict=valid_token,
)
assert result is False
def test_check_passthrough_route_access_key_metadata_takes_precedence():
"""Test that key metadata takes precedence over team metadata"""
# Create a UserAPIKeyAuth with different allowed_passthrough_routes in both metadata
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={"allowed_passthrough_routes": ["/key-endpoint"]},
team_metadata={"allowed_passthrough_routes": ["/team-endpoint"]},
)
# Test that key endpoint is allowed
result1 = RouteChecks.check_passthrough_route_access(
route="/key-endpoint",
user_api_key_dict=valid_token,
)
# Test that team endpoint is NOT allowed (key metadata takes precedence)
result2 = RouteChecks.check_passthrough_route_access(
route="/team-endpoint",
user_api_key_dict=valid_token,
)
assert result1 is True
assert result2 is False
def test_check_passthrough_route_access_no_metadata():
"""Test that route is denied when metadata and team_metadata don't have allowed_passthrough_routes"""
# Create a UserAPIKeyAuth without allowed_passthrough_routes
valid_token = UserAPIKeyAuth(
user_id="test_user",
)
# Test that route is denied
result = RouteChecks.check_passthrough_route_access(
route="/any-endpoint",
user_api_key_dict=valid_token,
)
assert result is False
def test_check_passthrough_route_access_no_allowed_passthrough_routes_key():
"""Test that route is denied when allowed_passthrough_routes is not in metadata"""
# Create a UserAPIKeyAuth with metadata but no allowed_passthrough_routes
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={"other_field": "value"},
team_metadata={},
)
# Test that route is denied
result = RouteChecks.check_passthrough_route_access(
route="/any-endpoint",
user_api_key_dict=valid_token,
)
assert result is False
def test_check_passthrough_route_access_allowed_passthrough_routes_is_none():
"""Test that route is denied when allowed_passthrough_routes is None"""
# Create a UserAPIKeyAuth with allowed_passthrough_routes set to None
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={"allowed_passthrough_routes": None},
team_metadata={"allowed_passthrough_routes": None},
)
# Test that route is denied
result = RouteChecks.check_passthrough_route_access(
route="/any-endpoint",
user_api_key_dict=valid_token,
)
assert result is False
def test_check_passthrough_route_access_multiple_routes():
"""Test that multiple allowed_passthrough_routes work correctly"""
# Create a UserAPIKeyAuth with multiple allowed_passthrough_routes
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={
"allowed_passthrough_routes": [
"/endpoint-1",
"/endpoint-2",
"/endpoint-3",
]
},
)
# Test that all allowed routes work
result1 = RouteChecks.check_passthrough_route_access(
route="/endpoint-1/v1/chat",
user_api_key_dict=valid_token,
)
result2 = RouteChecks.check_passthrough_route_access(
route="/endpoint-2",
user_api_key_dict=valid_token,
)
result3 = RouteChecks.check_passthrough_route_access(
route="/endpoint-3/completions",
user_api_key_dict=valid_token,
)
# Test that non-allowed route fails
result4 = RouteChecks.check_passthrough_route_access(
route="/endpoint-4",
user_api_key_dict=valid_token,
)
assert result1 is True
assert result2 is True
assert result3 is True
assert result4 is False
def test_check_passthrough_route_access_prevents_false_prefix_match():
"""Test that prefix matching doesn't allow false matches like /endpoint vs /endpoint-2"""
# Create a UserAPIKeyAuth with allowed_passthrough_routes
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={"allowed_passthrough_routes": ["/endpoint"]},
)
# Test that /endpoint-2 is NOT allowed (not a valid prefix match)
result = RouteChecks.check_passthrough_route_access(
route="/endpoint-2",
user_api_key_dict=valid_token,
)
assert result is False
# Test that /endpoint/something IS allowed (valid prefix match)
result2 = RouteChecks.check_passthrough_route_access(
route="/endpoint/something",
user_api_key_dict=valid_token,
)
assert result2 is True
def test_check_passthrough_route_access_empty_list():
"""Test that empty allowed_passthrough_routes list denies all routes"""
# Create a UserAPIKeyAuth with empty allowed_passthrough_routes
valid_token = UserAPIKeyAuth(
user_id="test_user",
metadata={"allowed_passthrough_routes": []},
)
# Test that route is denied
result = RouteChecks.check_passthrough_route_access(
route="/any-endpoint",
user_api_key_dict=valid_token,
)
assert result is False

View file

@ -670,7 +670,7 @@ async def test_create_pass_through_route_with_cost_per_request():
# Create a proper UserAPIKeyAuth mock
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.api_key = "test-key"
await endpoint_func(
request=mock_request,
user_api_key_dict=mock_user_api_key_dict,
@ -747,12 +747,12 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
"""
Test that pass_through_request (parent method) correctly includes proxy_server_request
in kwargs passed to the success handler.
Critical Test: Ensures that when pass_through_request is called, the kwargs passed to
Critical Test: Ensures that when pass_through_request is called, the kwargs passed to
downstream methods contain the proxy server request details (url, method, body).
"""
print("running test_pass_through_request_contains_proxy_server_request_in_kwargs")
with patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging:
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.HttpPassThroughEndpointHelpers.non_streaming_http_request_handler"
@ -766,38 +766,44 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
with patch(
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.get_response_body"
) as mock_get_response_body:
# Setup mock for pre_call_hook and post_call_failure_hook
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={"test": "data"})
# Setup mock for pre_call_hook and post_call_failure_hook
mock_proxy_logging.pre_call_hook = AsyncMock(
return_value={"test": "data"}
)
mock_proxy_logging.post_call_failure_hook = AsyncMock()
# Setup mock for http response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {}
mock_response.aread = AsyncMock(return_value=b'{"success": true}')
mock_response.aread = AsyncMock(
return_value=b'{"success": true}'
)
mock_response.text = '{"success": true}'
mock_response.raise_for_status = MagicMock()
# Mock the HTTP request handler directly
mock_http_handler.return_value = mock_response
# Mock response body parser
mock_get_response_body.return_value = {"success": True}
# Mock headers for custom headers
mock_processing.get_custom_headers.return_value = {}
# Mock success handler to capture kwargs
mock_success_handler.return_value = None
# Create mock request
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://test-proxy.com/api/endpoint"
mock_request.body = AsyncMock(return_value=b'{"message": "test request"}')
mock_request.body = AsyncMock(
return_value=b'{"message": "test request"}'
)
mock_request.headers = Headers({})
mock_request.query_params = QueryParams({})
# Create mock user API key dict
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.api_key = "test-api-key"
@ -809,7 +815,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
mock_user_api_key_dict.team_alias = "test-team-alias"
mock_user_api_key_dict.end_user_id = "test-end-user-id"
mock_user_api_key_dict.request_route = "/api/endpoint"
# Call pass_through_request (the parent method)
result = await pass_through_request(
request=mock_request,
@ -817,38 +823,38 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
custom_headers={"X-Custom": "header"},
user_api_key_dict=mock_user_api_key_dict,
)
# Verify the success handler was called
mock_success_handler.assert_called_once()
# Extract the kwargs passed to the success handler
call_kwargs = mock_success_handler.call_args[1]
# Verify that litellm_params exists in kwargs
assert "litellm_params" in call_kwargs
litellm_params = call_kwargs["litellm_params"]
# Verify that proxy_server_request exists in litellm_params
assert "proxy_server_request" in litellm_params
proxy_server_request = litellm_params["proxy_server_request"]
# Verify the proxy_server_request contains expected fields
assert "url" in proxy_server_request
assert "method" in proxy_server_request
assert "body" in proxy_server_request
# Verify the values match the original request
assert proxy_server_request["url"] == str(mock_request.url)
assert proxy_server_request["method"] == mock_request.method
# The body should be the value returned by pre_call_hook, not the original request body
assert proxy_server_request["body"] == {"test": "data"}
# Verify other required kwargs are present
assert "call_type" in call_kwargs
assert call_kwargs["call_type"] == "pass_through_endpoint"
assert "litellm_call_id" in call_kwargs
assert "passthrough_logging_payload" in call_kwargs
# Verify metadata contains user information
assert "metadata" in litellm_params
metadata = litellm_params["metadata"]
@ -862,7 +868,7 @@ async def test_pass_through_request_contains_proxy_server_request_in_kwargs():
async def test_create_pass_through_endpoint():
"""
Test creating a new pass-through endpoint
This test verifies that the create_pass_through_endpoints function:
1. Accepts a PassThroughGenericEndpoint object
2. Auto-generates an ID if not provided
@ -881,36 +887,38 @@ async def test_create_pass_through_endpoint():
)
# Mock the database functions
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config:
with patch(
"litellm.proxy.proxy_server.get_config_general_settings"
) as mock_get_config:
with patch(
"litellm.proxy.proxy_server.update_config_general_settings"
) as mock_update_config:
# Mock existing config (empty list)
mock_get_config.return_value = ConfigFieldInfo(
field_name="pass_through_endpoints",
field_value=[]
field_name="pass_through_endpoints", field_value=[]
)
# Create test endpoint data
test_endpoint = PassThroughGenericEndpoint(
path="/test/endpoint",
target="http://example.com/api",
headers={"Authorization": "Bearer test-token"},
include_subpath=True,
cost_per_request=0.50
cost_per_request=0.50,
)
# Mock user API key dict
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
# Call the create function
result = await create_pass_through_endpoints(
data=test_endpoint,
user_api_key_dict=mock_user_api_key_dict
data=test_endpoint, user_api_key_dict=mock_user_api_key_dict
)
# Verify the result
assert isinstance(result, PassThroughEndpointResponse)
assert len(result.endpoints) == 1
created_endpoint = result.endpoints[0]
assert created_endpoint.path == "/test/endpoint"
assert created_endpoint.target == "http://example.com/api"
@ -918,13 +926,13 @@ async def test_create_pass_through_endpoint():
assert created_endpoint.include_subpath is True
assert created_endpoint.cost_per_request == 0.50
assert created_endpoint.id is not None # Should be auto-generated
# Verify database calls
mock_get_config.assert_called_once_with(
field_name="pass_through_endpoints",
user_api_key_dict=mock_user_api_key_dict
user_api_key_dict=mock_user_api_key_dict,
)
mock_update_config.assert_called_once()
update_call_args = mock_update_config.call_args[1]
assert update_call_args["data"].field_name == "pass_through_endpoints"
@ -937,7 +945,7 @@ async def test_create_pass_through_endpoint():
async def test_update_pass_through_endpoint():
"""
Test updating an existing pass-through endpoint
This test verifies that the update_pass_through_endpoints function:
1. Finds the existing endpoint by ID
2. Updates only the provided fields (partial updates)
@ -957,8 +965,12 @@ async def test_update_pass_through_endpoint():
)
# Mock the database functions
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config:
with patch(
"litellm.proxy.proxy_server.get_config_general_settings"
) as mock_get_config:
with patch(
"litellm.proxy.proxy_server.update_config_general_settings"
) as mock_update_config:
# Create existing endpoint data
existing_endpoint_id = "test-endpoint-123"
existing_endpoints = [
@ -968,53 +980,58 @@ async def test_update_pass_through_endpoint():
"target": "http://example.com/api",
"headers": {"Authorization": "Bearer old-token"},
"include_subpath": False,
"cost_per_request": 0.25
"cost_per_request": 0.25,
}
]
# Mock existing config
mock_get_config.return_value = ConfigFieldInfo(
field_name="pass_through_endpoints",
field_value=existing_endpoints
field_name="pass_through_endpoints", field_value=existing_endpoints
)
# Create update data (partial update)
update_data = PassThroughGenericEndpoint(
path="/test/endpoint", # Keep same path
target="http://newapi.com/v2", # Update target
headers={"Authorization": "Bearer new-token", "X-Custom": "header"}, # Update headers
cost_per_request=0.75 # Update cost
headers={
"Authorization": "Bearer new-token",
"X-Custom": "header",
}, # Update headers
cost_per_request=0.75, # Update cost
# include_subpath not provided - should preserve existing value
)
# Mock user API key dict
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
# Call the update function
result = await update_pass_through_endpoints(
endpoint_id=existing_endpoint_id,
data=update_data,
user_api_key_dict=mock_user_api_key_dict
user_api_key_dict=mock_user_api_key_dict,
)
# Verify the result
assert isinstance(result, PassThroughEndpointResponse)
assert len(result.endpoints) == 1
updated_endpoint = result.endpoints[0]
assert updated_endpoint.id == existing_endpoint_id # ID preserved
assert updated_endpoint.path == "/test/endpoint"
assert updated_endpoint.target == "http://newapi.com/v2" # Updated
assert updated_endpoint.headers == {"Authorization": "Bearer new-token", "X-Custom": "header"} # Updated
assert updated_endpoint.headers == {
"Authorization": "Bearer new-token",
"X-Custom": "header",
} # Updated
assert updated_endpoint.include_subpath is False # Preserved existing value
assert updated_endpoint.cost_per_request == 0.75 # Updated
# Verify database calls
mock_get_config.assert_called_once_with(
field_name="pass_through_endpoints",
user_api_key_dict=mock_user_api_key_dict
user_api_key_dict=mock_user_api_key_dict,
)
mock_update_config.assert_called_once()
update_call_args = mock_update_config.call_args[1]
assert update_call_args["data"].field_name == "pass_through_endpoints"
@ -1042,7 +1059,9 @@ async def test_update_pass_through_endpoint_not_found():
)
# Mock the database functions
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
with patch(
"litellm.proxy.proxy_server.get_config_general_settings"
) as mock_get_config:
# Mock existing config with different endpoint
existing_endpoints = [
{
@ -1051,32 +1070,30 @@ async def test_update_pass_through_endpoint_not_found():
"target": "http://different.com/api",
"headers": {},
"include_subpath": False,
"cost_per_request": 0.0
"cost_per_request": 0.0,
}
]
mock_get_config.return_value = ConfigFieldInfo(
field_name="pass_through_endpoints",
field_value=existing_endpoints
field_name="pass_through_endpoints", field_value=existing_endpoints
)
# Create update data
update_data = PassThroughGenericEndpoint(
path="/test/endpoint",
target="http://newapi.com/v2"
path="/test/endpoint", target="http://newapi.com/v2"
)
# Mock user API key dict
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
# Call the update function with non-existent ID
with pytest.raises(HTTPException) as exc_info:
await update_pass_through_endpoints(
endpoint_id="non-existent-endpoint-123",
data=update_data,
user_api_key_dict=mock_user_api_key_dict
user_api_key_dict=mock_user_api_key_dict,
)
# Verify the exception
assert exc_info.value.status_code == 404
assert "not found" in str(exc_info.value.detail).lower()
@ -1086,7 +1103,7 @@ async def test_update_pass_through_endpoint_not_found():
async def test_delete_pass_through_endpoint():
"""
Test deleting an existing pass-through endpoint
This test verifies that the delete_pass_through_endpoints function:
1. Finds the existing endpoint by ID
2. Removes it from the database
@ -1103,8 +1120,12 @@ async def test_delete_pass_through_endpoint():
)
# Mock the database functions
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
with patch("litellm.proxy.proxy_server.update_config_general_settings") as mock_update_config:
with patch(
"litellm.proxy.proxy_server.get_config_general_settings"
) as mock_get_config:
with patch(
"litellm.proxy.proxy_server.update_config_general_settings"
) as mock_update_config:
# Create existing endpoint data
endpoint_to_delete_id = "test-endpoint-123"
other_endpoint_id = "other-endpoint-456"
@ -1115,7 +1136,7 @@ async def test_delete_pass_through_endpoint():
"target": "http://example.com/api",
"headers": {"Authorization": "Bearer test-token"},
"include_subpath": True,
"cost_per_request": 0.50
"cost_per_request": 0.50,
},
{
"id": other_endpoint_id,
@ -1123,29 +1144,28 @@ async def test_delete_pass_through_endpoint():
"target": "http://other.com/api",
"headers": {},
"include_subpath": False,
"cost_per_request": 0.25
}
"cost_per_request": 0.25,
},
]
# Mock existing config
mock_get_config.return_value = ConfigFieldInfo(
field_name="pass_through_endpoints",
field_value=existing_endpoints
field_name="pass_through_endpoints", field_value=existing_endpoints
)
# Mock user API key dict
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
# Call the delete function
result = await delete_pass_through_endpoints(
endpoint_id=endpoint_to_delete_id,
user_api_key_dict=mock_user_api_key_dict
user_api_key_dict=mock_user_api_key_dict,
)
# Verify the result
assert isinstance(result, PassThroughEndpointResponse)
assert len(result.endpoints) == 1
deleted_endpoint = result.endpoints[0]
assert deleted_endpoint.id == endpoint_to_delete_id
assert deleted_endpoint.path == "/test/endpoint"
@ -1153,13 +1173,13 @@ async def test_delete_pass_through_endpoint():
assert deleted_endpoint.headers == {"Authorization": "Bearer test-token"}
assert deleted_endpoint.include_subpath is True
assert deleted_endpoint.cost_per_request == 0.50
# Verify database calls
mock_get_config.assert_called_once_with(
field_name="pass_through_endpoints",
user_api_key_dict=mock_user_api_key_dict
user_api_key_dict=mock_user_api_key_dict,
)
mock_update_config.assert_called_once()
update_call_args = mock_update_config.call_args[1]
assert update_call_args["data"].field_name == "pass_through_endpoints"
@ -1183,7 +1203,9 @@ async def test_delete_pass_through_endpoint_not_found():
)
# Mock the database functions
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
with patch(
"litellm.proxy.proxy_server.get_config_general_settings"
) as mock_get_config:
# Mock existing config with different endpoint
existing_endpoints = [
{
@ -1192,25 +1214,24 @@ async def test_delete_pass_through_endpoint_not_found():
"target": "http://different.com/api",
"headers": {},
"include_subpath": False,
"cost_per_request": 0.0
"cost_per_request": 0.0,
}
]
mock_get_config.return_value = ConfigFieldInfo(
field_name="pass_through_endpoints",
field_value=existing_endpoints
field_name="pass_through_endpoints", field_value=existing_endpoints
)
# Mock user API key dict
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
# Call the delete function with non-existent ID
with pytest.raises(HTTPException) as exc_info:
await delete_pass_through_endpoints(
endpoint_id="non-existent-endpoint-123",
user_api_key_dict=mock_user_api_key_dict
user_api_key_dict=mock_user_api_key_dict,
)
# Verify the exception
assert exc_info.value.status_code == 400
assert "not found" in str(exc_info.value.detail).lower()
@ -1229,34 +1250,33 @@ async def test_delete_pass_through_endpoint_empty_list():
)
# Mock the database functions
with patch("litellm.proxy.proxy_server.get_config_general_settings") as mock_get_config:
with patch(
"litellm.proxy.proxy_server.get_config_general_settings"
) as mock_get_config:
# Mock empty config
mock_get_config.return_value = ConfigFieldInfo(
field_name="pass_through_endpoints",
field_value=None
field_name="pass_through_endpoints", field_value=None
)
# Mock user API key dict
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
# Call the delete function
with pytest.raises(HTTPException) as exc_info:
await delete_pass_through_endpoints(
endpoint_id="any-endpoint-123",
user_api_key_dict=mock_user_api_key_dict
endpoint_id="any-endpoint-123", user_api_key_dict=mock_user_api_key_dict
)
# Verify the exception
assert exc_info.value.status_code == 400
assert "no pass-through endpoints setup" in str(exc_info.value.detail).lower()
@pytest.mark.asyncio
async def test_pass_through_request_query_params_forwarding():
"""
Test that query parameters from the original request are properly forwarded to the target URL.
This test verifies the fix for the bug where query parameters like api-version were being lost
when forwarding requests to Azure OpenAI and other pass-through endpoints.
"""
@ -1275,42 +1295,57 @@ async def test_pass_through_request_query_params_forwarding():
) as mock_get_response_body:
# Setup mock for pre_call_hook
test_body = {"name": "Azure Assistant", "model": "gpt-4o"}
mock_proxy_logging.pre_call_hook = AsyncMock(return_value=test_body)
mock_proxy_logging.pre_call_hook = AsyncMock(
return_value=test_body
)
# Setup mock for http response
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {"content-type": "application/json"}
mock_response.aread = AsyncMock(return_value=b'{"id": "asst_123", "object": "assistant"}')
mock_response.aread = AsyncMock(
return_value=b'{"id": "asst_123", "object": "assistant"}'
)
mock_response.text = '{"id": "asst_123", "object": "assistant"}'
mock_response.raise_for_status = MagicMock()
# Mock the HTTP request handler to capture the call
mock_http_handler.return_value = mock_response
# Mock response body parser
mock_get_response_body.return_value = {"id": "asst_123", "object": "assistant"}
mock_get_response_body.return_value = {
"id": "asst_123",
"object": "assistant",
}
# Mock headers for custom headers
mock_processing.get_custom_headers.return_value = {}
# Mock success handler
mock_success_handler.return_value = None
# Create mock request with query parameters (Azure API version)
mock_request = MagicMock(spec=Request)
mock_request.method = "POST"
mock_request.url = "http://localhost:4000/azure-assistant/openai/assistants"
mock_request.body = AsyncMock(return_value=json.dumps(test_body).encode())
mock_request.headers = Headers({"Content-Type": "application/json"})
mock_request.url = (
"http://localhost:4000/azure-assistant/openai/assistants"
)
mock_request.body = AsyncMock(
return_value=json.dumps(test_body).encode()
)
mock_request.headers = Headers(
{"Content-Type": "application/json"}
)
# Create QueryParams with api-version parameter
mock_request.query_params = QueryParams([("api-version", "2025-01-01-preview")])
mock_request.query_params = QueryParams(
[("api-version", "2025-01-01-preview")]
)
# Create mock user API key dict
mock_user_api_key_dict = MagicMock()
mock_user_api_key_dict.api_key = "sk-1234"
# Call pass_through_request
result = await pass_through_request(
request=mock_request,
@ -1318,20 +1353,25 @@ async def test_pass_through_request_query_params_forwarding():
custom_headers={"Authorization": "Bearer azure_token"},
user_api_key_dict=mock_user_api_key_dict,
)
# Verify the HTTP handler was called
mock_http_handler.assert_called_once()
# Extract the call arguments to verify query parameters were passed
call_kwargs = mock_http_handler.call_args[1]
# The key assertion: query parameters should be preserved and passed to the HTTP handler
assert "requested_query_params" in call_kwargs
assert call_kwargs["requested_query_params"] == {"api-version": "2025-01-01-preview"}
assert call_kwargs["requested_query_params"] == {
"api-version": "2025-01-01-preview"
}
# Verify the target URL is correct
assert str(call_kwargs["url"]) == "https://krris-m2f9a9i7-eastus2.openai.azure.com/openai/assistants"
assert (
str(call_kwargs["url"])
== "https://krris-m2f9a9i7-eastus2.openai.azure.com/openai/assistants"
)
# Verify the request body is preserved
assert call_kwargs["_parsed_body"] == test_body
@ -1356,15 +1396,16 @@ async def test_pass_through_with_httpbin_redirect():
mock_request.method = "GET"
mock_request.headers = Headers({})
mock_request.query_params = QueryParams("")
# Mock the body method to return empty bytes for GET request
async def mock_body():
return b""
mock_request.body = mock_body
# Mock user API key dict
mock_user_api_key_dict = MagicMock()
try:
# Test with httpbin.org redirect endpoint
# This will redirect to httpbin.org/get
@ -1372,19 +1413,283 @@ async def test_pass_through_with_httpbin_redirect():
request=mock_request,
target="https://httpbin.org/redirect/1",
custom_headers={},
user_api_key_dict=mock_user_api_key_dict
user_api_key_dict=mock_user_api_key_dict,
)
# Should get the final response (200) from /get endpoint, not the redirect (302)
assert response.status_code == 200
# The response should be from the /get endpoint
response_content = response.body.decode('utf-8')
response_content = response.body.decode("utf-8")
# httpbin.org/get returns JSON with info about the request
assert '"url": "https://httpbin.org/get"' in response_content
print("GOT A Response from HTTPBIN=", response_content)
except Exception as e:
# If httpbin.org is not accessible, skip the test
import pytest
pytest.skip(f"Could not reach httpbin.org for integration test: {e}")
@pytest.mark.asyncio
async def test_filter_endpoints_by_team_allowed_routes_with_filter():
"""
Test that _filter_endpoints_by_team_allowed_routes correctly filters endpoints
when team has allowed_passthrough_routes in metadata
"""
from litellm.proxy._types import PassThroughGenericEndpoint
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
_filter_endpoints_by_team_allowed_routes,
)
# Create test endpoints
endpoints = [
PassThroughGenericEndpoint(
id="endpoint-1", path="/api/allowed1", target="http://example.com/api1"
),
PassThroughGenericEndpoint(
id="endpoint-2", path="/api/allowed2", target="http://example.com/api2"
),
PassThroughGenericEndpoint(
id="endpoint-3", path="/api/notallowed", target="http://example.com/api3"
),
]
# Mock prisma client
mock_prisma_client = MagicMock()
mock_team = MagicMock()
mock_team.metadata = {
"allowed_passthrough_routes": ["/api/allowed1", "/api/allowed2"]
}
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=mock_team
)
# Call the function
result = await _filter_endpoints_by_team_allowed_routes(
team_id="test-team-123",
pass_through_endpoints=endpoints,
prisma_client=mock_prisma_client,
)
# Should only return allowed endpoints
assert len(result) == 2
assert result[0].path == "/api/allowed1"
assert result[1].path == "/api/allowed2"
# Verify database call
mock_prisma_client.db.litellm_teamtable.find_unique.assert_called_once_with(
where={"team_id": "test-team-123"}
)
@pytest.mark.asyncio
async def test_filter_endpoints_by_team_allowed_routes_team_not_found():
"""
Test that _filter_endpoints_by_team_allowed_routes raises HTTPException
when team is not found
"""
from fastapi import HTTPException
from litellm.proxy._types import PassThroughGenericEndpoint
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
_filter_endpoints_by_team_allowed_routes,
)
# Create test endpoints
endpoints = [
PassThroughGenericEndpoint(
id="endpoint-1", path="/api/test", target="http://example.com/api"
),
]
# Mock prisma client to return None (team not found)
mock_prisma_client = MagicMock()
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=None)
# Call the function and expect HTTPException
with pytest.raises(HTTPException) as exc_info:
await _filter_endpoints_by_team_allowed_routes(
team_id="non-existent-team",
pass_through_endpoints=endpoints,
prisma_client=mock_prisma_client,
)
# Verify the exception
assert exc_info.value.status_code == 404
assert "Team not found" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_filter_endpoints_by_team_allowed_routes_no_metadata():
"""
Test that _filter_endpoints_by_team_allowed_routes returns all endpoints
when team has no metadata
"""
from litellm.proxy._types import PassThroughGenericEndpoint
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
_filter_endpoints_by_team_allowed_routes,
)
# Create test endpoints
endpoints = [
PassThroughGenericEndpoint(
id="endpoint-1", path="/api/test1", target="http://example.com/api1"
),
PassThroughGenericEndpoint(
id="endpoint-2", path="/api/test2", target="http://example.com/api2"
),
]
# Mock prisma client with team that has None metadata
mock_prisma_client = MagicMock()
mock_team = MagicMock()
mock_team.metadata = None
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=mock_team
)
# Call the function
result = await _filter_endpoints_by_team_allowed_routes(
team_id="test-team-123",
pass_through_endpoints=endpoints,
prisma_client=mock_prisma_client,
)
# Should return all endpoints when no metadata
assert len(result) == 2
assert result[0].path == "/api/test1"
assert result[1].path == "/api/test2"
@pytest.mark.asyncio
async def test_filter_endpoints_by_team_allowed_routes_no_allowed_routes_key():
"""
Test that _filter_endpoints_by_team_allowed_routes returns all endpoints
when team metadata doesn't have allowed_passthrough_routes key
"""
from litellm.proxy._types import PassThroughGenericEndpoint
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
_filter_endpoints_by_team_allowed_routes,
)
# Create test endpoints
endpoints = [
PassThroughGenericEndpoint(
id="endpoint-1", path="/api/test1", target="http://example.com/api1"
),
PassThroughGenericEndpoint(
id="endpoint-2", path="/api/test2", target="http://example.com/api2"
),
]
# Mock prisma client with team that has metadata but no allowed_passthrough_routes
mock_prisma_client = MagicMock()
mock_team = MagicMock()
mock_team.metadata = {"some_other_key": "some_value"}
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=mock_team
)
# Call the function
result = await _filter_endpoints_by_team_allowed_routes(
team_id="test-team-123",
pass_through_endpoints=endpoints,
prisma_client=mock_prisma_client,
)
# Should return all endpoints when allowed_passthrough_routes key doesn't exist
assert len(result) == 2
assert result[0].path == "/api/test1"
assert result[1].path == "/api/test2"
@pytest.mark.asyncio
async def test_filter_endpoints_by_team_allowed_routes_empty_allowed_list():
"""
Test that _filter_endpoints_by_team_allowed_routes returns empty list
when team has empty allowed_passthrough_routes list
"""
from litellm.proxy._types import PassThroughGenericEndpoint
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
_filter_endpoints_by_team_allowed_routes,
)
# Create test endpoints
endpoints = [
PassThroughGenericEndpoint(
id="endpoint-1", path="/api/test1", target="http://example.com/api1"
),
PassThroughGenericEndpoint(
id="endpoint-2", path="/api/test2", target="http://example.com/api2"
),
]
# Mock prisma client with team that has empty allowed_passthrough_routes
mock_prisma_client = MagicMock()
mock_team = MagicMock()
mock_team.metadata = {"allowed_passthrough_routes": []}
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=mock_team
)
# Call the function
result = await _filter_endpoints_by_team_allowed_routes(
team_id="test-team-123",
pass_through_endpoints=endpoints,
prisma_client=mock_prisma_client,
)
# Should return empty list when allowed_passthrough_routes is empty
assert len(result) == 0
@pytest.mark.asyncio
async def test_filter_endpoints_by_team_allowed_routes_partial_match():
"""
Test that _filter_endpoints_by_team_allowed_routes correctly filters
when only some endpoints match allowed routes
"""
from litellm.proxy._types import PassThroughGenericEndpoint
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
_filter_endpoints_by_team_allowed_routes,
)
# Create test endpoints
endpoints = [
PassThroughGenericEndpoint(
id="endpoint-1", path="/api/openai", target="http://example.com/openai"
),
PassThroughGenericEndpoint(
id="endpoint-2",
path="/api/anthropic",
target="http://example.com/anthropic",
),
PassThroughGenericEndpoint(
id="endpoint-3", path="/api/azure", target="http://example.com/azure"
),
PassThroughGenericEndpoint(
id="endpoint-4", path="/api/cohere", target="http://example.com/cohere"
),
]
# Mock prisma client with team that allows only 2 routes
mock_prisma_client = MagicMock()
mock_team = MagicMock()
mock_team.metadata = {"allowed_passthrough_routes": ["/api/openai", "/api/azure"]}
mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(
return_value=mock_team
)
# Call the function
result = await _filter_endpoints_by_team_allowed_routes(
team_id="test-team-123",
pass_through_endpoints=endpoints,
prisma_client=mock_prisma_client,
)
# Should return only the 2 allowed endpoints
assert len(result) == 2
assert result[0].path == "/api/openai"
assert result[1].path == "/api/azure"

View file

@ -4,14 +4,20 @@ Test reasoning content preservation in Responses API transformation
from unittest.mock import AsyncMock
from litellm.types.utils import ModelResponseStream, StreamingChoices, Delta
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
LiteLLMCompletionStreamingIterator,
)
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
from litellm.types.utils import ModelResponse, Choices, Message
from litellm.types.utils import (
Choices,
Delta,
Message,
ModelResponse,
ModelResponseStream,
StreamingChoices,
)
class TestReasoningContentStreaming:
@ -41,6 +47,7 @@ class TestReasoningContentStreaming:
mock_stream = AsyncMock()
iterator = LiteLLMCompletionStreamingIterator(
model="test-model",
litellm_custom_stream_wrapper=mock_stream,
request_input="Test input",
responses_api_request={},
@ -78,6 +85,7 @@ class TestReasoningContentStreaming:
mock_stream = AsyncMock()
iterator = LiteLLMCompletionStreamingIterator(
model="test-model",
litellm_custom_stream_wrapper=mock_stream,
request_input="Test input",
responses_api_request={},
@ -114,6 +122,7 @@ class TestReasoningContentStreaming:
mock_stream = AsyncMock()
iterator = LiteLLMCompletionStreamingIterator(
model="test-model",
litellm_custom_stream_wrapper=mock_stream,
request_input="Test input",
responses_api_request={},
@ -272,6 +281,7 @@ def test_streaming_chunk_id_raw():
)
iterator = LiteLLMCompletionStreamingIterator(
model="test-model",
litellm_custom_stream_wrapper=AsyncMock(),
request_input="Test input",
responses_api_request={},

View file

@ -0,0 +1,69 @@
import React, { useEffect, useState } from "react";
import { Select } from "antd";
import { getPassThroughEndpointsCall } from "../networking";
interface PassThroughRoutesSelectorProps {
onChange: (selectedRoutes: string[]) => void;
value?: string[];
className?: string;
accessToken: string;
placeholder?: string;
disabled?: boolean;
teamId?: string | null;
}
const PassThroughRoutesSelector: React.FC<PassThroughRoutesSelectorProps> = ({
onChange,
value,
className,
accessToken,
placeholder = "Select pass through routes",
disabled = false,
teamId,
}) => {
const [passThroughRoutes, setPassThroughRoutes] = useState<string[]>([]);
const [loading, setLoading] = useState(false);
useEffect(() => {
const fetchPassThroughRoutes = async () => {
if (!accessToken) return;
setLoading(true);
try {
const response = await getPassThroughEndpointsCall(accessToken, teamId);
if (response.endpoints) {
const routes = response.endpoints.map((route: { path: string }) => route.path);
setPassThroughRoutes(routes);
}
} catch (error) {
console.error("Error fetching pass through routes:", error);
} finally {
setLoading(false);
}
};
fetchPassThroughRoutes();
}, [accessToken, teamId]);
return (
<Select
mode="tags"
placeholder={placeholder}
onChange={onChange}
value={value}
loading={loading}
className={className}
options={passThroughRoutes.map((route) => ({
label: route,
value: route,
}))}
optionFilterProp="label"
showSearch
style={{ width: "100%" }}
disabled={disabled}
/>
);
};
export default PassThroughRoutesSelector;

View file

@ -4442,10 +4442,14 @@ export const getGeneralSettingsCall = async (accessToken: string) => {
}
};
export const getPassThroughEndpointsCall = async (accessToken: string) => {
export const getPassThroughEndpointsCall = async (accessToken: string, teamId?: string | null) => {
try {
let url = proxyBaseUrl ? `${proxyBaseUrl}/config/pass_through_endpoint` : `/config/pass_through_endpoint`;
if (teamId) {
url += `/team/${teamId}`;
}
//NotificationsManager.info("Requesting model data");
const response = await fetch(url, {
method: "GET",

View file

@ -19,6 +19,7 @@ import {
getPromptsList,
} from "../networking";
import VectorStoreSelector from "../vector_store_management/VectorStoreSelector";
import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector";
import { Team } from "../key_team_helpers/key_list";
import TeamDropdown from "../common_components/team_dropdown";
import { InfoCircleOutlined } from "@ant-design/icons";
@ -912,6 +913,41 @@ const CreateKey: React.FC<CreateKeyProps> = ({
options={promptsList.map((name) => ({ value: name, label: name }))}
/>
</Form.Item>
<Form.Item
label={
<span>
Allowed Pass Through Routes{" "}
<Tooltip title="Allow this key to use specific pass through routes">
<a
href="https://docs.litellm.ai/docs/proxy/pass_through"
target="_blank"
rel="noopener noreferrer"
onClick={(e) => e.stopPropagation()} // Prevent accordion from collapsing when clicking link
>
<InfoCircleOutlined style={{ marginLeft: "4px" }} />
</a>
</Tooltip>
</span>
}
name="allowed_passthrough_routes"
className="mt-4"
help={
premiumUser
? "Select existing pass through routes or enter new ones"
: "Premium feature - Upgrade to set pass through routes by key"
}
>
<PassThroughRoutesSelector
onChange={(values: string[]) => form.setFieldValue("allowed_passthrough_routes", values)}
value={form.getFieldValue("allowed_passthrough_routes")}
accessToken={accessToken}
placeholder={
!premiumUser ? "Premium feature - Upgrade to set pass through routes by key" : "Select or enter pass through routes"
}
disabled={!premiumUser}
teamId={selectedCreateKeyTeam ? selectedCreateKeyTeam.team_id : null}
/>
</Form.Item>
<Form.Item
label={
<span>

View file

@ -42,6 +42,7 @@ import { fetchMCPAccessGroups } from "../networking";
import { CheckIcon, CopyIcon } from "lucide-react";
import { copyToClipboard as utilCopyToClipboard } from "../../utils/dataUtils";
import NotificationsManager from "../molecules/notifications_manager";
import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector";
export interface TeamMembership {
user_id: string;
@ -667,6 +668,15 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
/>
</Form.Item>
<Form.Item label="Allowed Pass Through Routes" name="allowed_passthrough_routes">
<PassThroughRoutesSelector
onChange={(values: string[]) => form.setFieldValue("allowed_passthrough_routes", values)}
value={form.getFieldValue("allowed_passthrough_routes")}
accessToken={accessToken || ""}
placeholder="Select pass through routes"
/>
</Form.Item>
<Form.Item label="MCP Servers / Access Groups" name="mcp_servers_and_groups">
<MCPServerSelector
onChange={(val) => form.setFieldValue("mcp_servers_and_groups", val)}

View file

@ -15,6 +15,7 @@ import { mapInternalToDisplayNames } from "../callback_info_helpers";
import GuardrailSelector from "@/components/guardrails/GuardrailSelector";
import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings";
import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem";
import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector";
interface KeyEditViewProps {
keyData: KeyResponse;
@ -276,6 +277,24 @@ export function KeyEditView({
</Tooltip>
</Form.Item>
<Form.Item label="Allowed Pass Through Routes" name="allowed_passthrough_routes">
<Tooltip title={!premiumUser ? "Setting allowed pass through routes by key is a premium feature" : ""} placement="top">
<PassThroughRoutesSelector
onChange={(values: string[]) => form.setFieldValue("allowed_passthrough_routes", values)}
value={form.getFieldValue("allowed_passthrough_routes")}
accessToken={accessToken || ""}
placeholder={
!premiumUser
? "Premium feature - Upgrade to set allowed pass through routes by key"
: Array.isArray(keyData.metadata?.allowed_passthrough_routes) && keyData.metadata.allowed_passthrough_routes.length > 0
? `Current: ${keyData.metadata.allowed_passthrough_routes.join(", ")}`
: "Select or enter allowed pass through routes"
}
disabled={!premiumUser}
/>
</Tooltip>
</Form.Item>
<Form.Item label="Vector Stores" name="vector_stores">
<VectorStoreSelector
onChange={(values: string[]) => form.setFieldValue("vector_stores", values)}

View file

@ -627,6 +627,19 @@ export default function KeyInfoView({
</Text>
</div>
<div>
<Text className="font-medium">Allowed Pass Through Routes</Text>
<Text>
{Array.isArray(currentKeyData.metadata?.allowed_passthrough_routes) && currentKeyData.metadata.allowed_passthrough_routes.length > 0
? currentKeyData.metadata.allowed_passthrough_routes.map((route, index) => (
<span key={index} className="px-2 mr-2 py-1 bg-blue-100 rounded text-xs">
{route}
</span>
))
: "No pass through routes specified"}
</Text>
</div>
<div>
<Text className="font-medium">Models</Text>
<div className="flex flex-wrap gap-2 mt-1">