mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
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:
commit
ff20f8402a
32 changed files with 2642 additions and 674 deletions
2
.gitignore
vendored
2
.gitignore
vendored
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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", {}),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
78
log.txt
Normal 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 ========================
|
||||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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__])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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={},
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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)}
|
||||
|
|
|
|||
|
|
@ -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)}
|
||||
|
|
|
|||
|
|
@ -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">
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue