mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
Merge branch 'litellm_internal_staging' into litellm_oss_staging
Resolves merge conflict in tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py by keeping both the new bedrock tool-result file/document tests and the transform_response body-leak regression test. Also addresses Greptile P2 comment: when BedrockImageProcessor returns a block with neither 'image' nor 'document' keys on the tool-result path (image_url and file content types), log a warning instead of silently dropping the block. Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
Cursor Agent
commit
1b6ab9facf
48 changed files with 2091 additions and 167 deletions
|
|
@ -224,6 +224,16 @@ AIOHTTP_CONNECTOR_LIMIT_PER_HOST = int(
|
|||
)
|
||||
AIOHTTP_KEEPALIVE_TIMEOUT = int(os.getenv("AIOHTTP_KEEPALIVE_TIMEOUT", 120))
|
||||
AIOHTTP_TTL_DNS_CACHE = int(os.getenv("AIOHTTP_TTL_DNS_CACHE", 300))
|
||||
# TCP keep-alive (SO_KEEPALIVE) — opt-in. Required when running behind NAT/LBs
|
||||
# whose idle timeout is shorter than provider response timeouts (e.g. AWS NAT
|
||||
# Gateway: 350s vs OpenAI/Azure: 600s). Without this, the kernel sends nothing
|
||||
# during a long provider call and the NAT reaps the flow before the response
|
||||
# arrives. Enabling SO_KEEPALIVE makes the kernel emit TCP probes that reset
|
||||
# the NAT idle timer.
|
||||
AIOHTTP_SO_KEEPALIVE = os.getenv("AIOHTTP_SO_KEEPALIVE", "False").lower() == "true"
|
||||
AIOHTTP_TCP_KEEPIDLE = int(os.getenv("AIOHTTP_TCP_KEEPIDLE", 60))
|
||||
AIOHTTP_TCP_KEEPINTVL = int(os.getenv("AIOHTTP_TCP_KEEPINTVL", 30))
|
||||
AIOHTTP_TCP_KEEPCNT = int(os.getenv("AIOHTTP_TCP_KEEPCNT", 5))
|
||||
# enable_cleanup_closed is only needed for Python versions with the SSL leak bug
|
||||
# Fixed in Python 3.12.7+ and 3.13.1+ (see https://github.com/python/cpython/pull/118960)
|
||||
# Reference: https://github.com/aio-libs/aiohttp/blob/master/aiohttp/connector.py#L74-L78
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import asyncio
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, List, Optional
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy.pass_through_endpoints.success_handler import (
|
||||
|
|
@ -29,12 +29,14 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
|
|||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
request_body: dict,
|
||||
model: str,
|
||||
hidden_params: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
self.litellm_logging_obj = litellm_logging_obj
|
||||
self.request_body = request_body
|
||||
self.start_time = datetime.now()
|
||||
self.collected_chunks: List[bytes] = []
|
||||
self.model = model
|
||||
self._hidden_params: Dict[str, Any] = hidden_params or {}
|
||||
|
||||
async def _handle_async_streaming_logging(
|
||||
self,
|
||||
|
|
@ -76,11 +78,13 @@ class GoogleGenAIGenerateContentStreamingIterator(
|
|||
litellm_metadata: dict,
|
||||
custom_llm_provider: str,
|
||||
request_body: Optional[dict] = None,
|
||||
hidden_params: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
super().__init__(
|
||||
litellm_logging_obj=logging_obj,
|
||||
request_body=request_body or {},
|
||||
model=model,
|
||||
hidden_params=hidden_params,
|
||||
)
|
||||
self.response = response
|
||||
self.model = model
|
||||
|
|
@ -130,11 +134,13 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(
|
|||
litellm_metadata: dict,
|
||||
custom_llm_provider: str,
|
||||
request_body: Optional[dict] = None,
|
||||
hidden_params: Optional[Dict[str, Any]] = None,
|
||||
):
|
||||
super().__init__(
|
||||
litellm_logging_obj=logging_obj,
|
||||
request_body=request_body or {},
|
||||
model=model,
|
||||
hidden_params=hidden_params,
|
||||
)
|
||||
self.response = response
|
||||
self.model = model
|
||||
|
|
|
|||
|
|
@ -4046,6 +4046,13 @@ def _convert_to_bedrock_tool_call_result(
|
|||
tool_result_content_blocks.append(
|
||||
BedrockToolResultContentBlock(document=_block["document"])
|
||||
)
|
||||
else:
|
||||
verbose_logger.warning(
|
||||
"Bedrock Converse: unrecognized BedrockContentBlock keys "
|
||||
"%s for image_url tool-result block %s; dropping.",
|
||||
list(_block.keys()),
|
||||
content,
|
||||
)
|
||||
elif content["type"] == "file":
|
||||
# Match the user-message path (_process_file_message): accept
|
||||
# either file_data (base64 data URI) or file_id (server-side
|
||||
|
|
@ -4077,6 +4084,13 @@ def _convert_to_bedrock_tool_call_result(
|
|||
tool_result_content_blocks.append(
|
||||
BedrockToolResultContentBlock(image=_file_block["image"])
|
||||
)
|
||||
else:
|
||||
verbose_logger.warning(
|
||||
"Bedrock Converse: unrecognized BedrockContentBlock keys "
|
||||
"%s for file tool-result block %s; dropping.",
|
||||
list(_file_block.keys()),
|
||||
content,
|
||||
)
|
||||
|
||||
message.get("name", "")
|
||||
id = str(message.get("tool_call_id", str(uuid.uuid4())))
|
||||
|
|
|
|||
|
|
@ -95,6 +95,7 @@ class AzureAIVectorStoreConfig(BaseVectorStoreConfig, BaseAzureLLM):
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
"""
|
||||
Transform search request for Azure AI Search API
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ class BaseVectorStoreConfig:
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
pass
|
||||
|
||||
|
|
@ -70,6 +71,7 @@ class BaseVectorStoreConfig:
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
"""
|
||||
Optional async version of transform_search_vector_store_request.
|
||||
|
|
@ -84,6 +86,7 @@ class BaseVectorStoreConfig:
|
|||
api_base=api_base,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
|
|
|
|||
|
|
@ -1942,8 +1942,8 @@ class AmazonConverseConfig(BaseConfig):
|
|||
completion_response = ConverseResponseBlock(**response.json()) # type: ignore
|
||||
except Exception as e:
|
||||
raise BedrockError(
|
||||
message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
|
||||
response.text, str(e)
|
||||
message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
|
||||
str(e)
|
||||
),
|
||||
status_code=422,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,14 +1,16 @@
|
|||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
|
||||
from copy import deepcopy
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.types.integrations.rag.bedrock_knowledgebase import (
|
||||
BedrockKBContent,
|
||||
BedrockKBResponse,
|
||||
BedrockKBRetrievalConfiguration,
|
||||
BedrockKBResponse,
|
||||
BedrockKBRetrievalQuery,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
|
@ -202,6 +204,7 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
if isinstance(query, list):
|
||||
query = " ".join(query)
|
||||
|
|
@ -213,24 +216,46 @@ class BedrockVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
|
|||
}
|
||||
|
||||
retrieval_config: Dict[str, Any] = {}
|
||||
|
||||
if isinstance(extra_body, dict):
|
||||
retrieval_config = deepcopy(
|
||||
extra_body.get("retrievalConfiguration")
|
||||
or extra_body.get("retrieval_configuration")
|
||||
or {}
|
||||
)
|
||||
max_results = vector_store_search_optional_params.get("max_num_results")
|
||||
if max_results is not None:
|
||||
existing_number_of_results = retrieval_config.get(
|
||||
"vectorSearchConfiguration", {}
|
||||
).get("numberOfResults")
|
||||
if (
|
||||
existing_number_of_results is not None
|
||||
and existing_number_of_results != max_results
|
||||
):
|
||||
verbose_logger.debug(
|
||||
"Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.numberOfResults (%s) with max_num_results=%s",
|
||||
existing_number_of_results,
|
||||
max_results,
|
||||
)
|
||||
retrieval_config.setdefault("vectorSearchConfiguration", {})[
|
||||
"numberOfResults"
|
||||
] = max_results
|
||||
filters = vector_store_search_optional_params.get("filters")
|
||||
if filters is not None:
|
||||
existing_filter = retrieval_config.get("vectorSearchConfiguration", {}).get(
|
||||
"filter"
|
||||
)
|
||||
if existing_filter is not None and existing_filter != filters:
|
||||
verbose_logger.debug(
|
||||
"Overriding extra_body retrievalConfiguration.vectorSearchConfiguration.filter with filters from vector_store_search_optional_params"
|
||||
)
|
||||
retrieval_config.setdefault("vectorSearchConfiguration", {})[
|
||||
"filter"
|
||||
] = filters
|
||||
if retrieval_config:
|
||||
# Create a properly typed retrieval configuration
|
||||
typed_retrieval_config: BedrockKBRetrievalConfiguration = {}
|
||||
if "vectorSearchConfiguration" in retrieval_config:
|
||||
typed_retrieval_config["vectorSearchConfiguration"] = retrieval_config[
|
||||
"vectorSearchConfiguration"
|
||||
]
|
||||
request_body["retrievalConfiguration"] = typed_retrieval_config
|
||||
request_body["retrievalConfiguration"] = cast(
|
||||
BedrockKBRetrievalConfiguration, retrieval_config
|
||||
)
|
||||
|
||||
litellm_logging_obj.model_call_details["query"] = query
|
||||
return url, request_body
|
||||
|
|
|
|||
|
|
@ -1,5 +1,7 @@
|
|||
import asyncio
|
||||
import inspect
|
||||
import os
|
||||
import socket
|
||||
import ssl
|
||||
import sys
|
||||
import time
|
||||
|
|
@ -29,6 +31,10 @@ from litellm.constants import (
|
|||
AIOHTTP_CONNECTOR_LIMIT_PER_HOST,
|
||||
AIOHTTP_KEEPALIVE_TIMEOUT,
|
||||
AIOHTTP_NEEDS_CLEANUP_CLOSED,
|
||||
AIOHTTP_SO_KEEPALIVE,
|
||||
AIOHTTP_TCP_KEEPCNT,
|
||||
AIOHTTP_TCP_KEEPIDLE,
|
||||
AIOHTTP_TCP_KEEPINTVL,
|
||||
AIOHTTP_TTL_DNS_CACHE,
|
||||
COMPLETION_HTTP_FALLBACK_SECONDS,
|
||||
DEFAULT_SSL_CIPHERS,
|
||||
|
|
@ -54,6 +60,57 @@ except Exception:
|
|||
version = "0.0.0"
|
||||
|
||||
|
||||
# aiohttp 3.10+ exposes a `socket_factory` kwarg on TCPConnector. Older
|
||||
# versions don't — detect once and skip the keep-alive wiring there.
|
||||
# https://docs.aiohttp.org/en/stable/client_reference.html#aiohttp.TCPConnector
|
||||
_AIOHTTP_SUPPORTS_SOCKET_FACTORY = (
|
||||
"socket_factory" in inspect.signature(TCPConnector.__init__).parameters
|
||||
)
|
||||
|
||||
|
||||
def _build_aiohttp_keepalive_socket_factory() -> (
|
||||
Optional[Callable[[Tuple[Any, ...]], socket.socket]]
|
||||
):
|
||||
"""
|
||||
Build a socket_factory that enables SO_KEEPALIVE on aiohttp TCP sockets.
|
||||
|
||||
Why: by default, aiohttp creates sockets without SO_KEEPALIVE, so the kernel
|
||||
sends nothing during a long idle TCP connection. NAT/LB hops (e.g. AWS NAT
|
||||
Gateway, 350s idle timeout) reap the flow well before slow provider
|
||||
responses (OpenAI/Azure: up to 600s) arrive. Enabling SO_KEEPALIVE makes
|
||||
the kernel emit TCP probes that reset the NAT idle timer.
|
||||
|
||||
Returns None when AIOHTTP_SO_KEEPALIVE is disabled or aiohttp is too old.
|
||||
"""
|
||||
if not AIOHTTP_SO_KEEPALIVE or not _AIOHTTP_SUPPORTS_SOCKET_FACTORY:
|
||||
return None
|
||||
|
||||
def factory(addr_info: Tuple[Any, ...]) -> socket.socket:
|
||||
family, type_, proto = addr_info[0], addr_info[1], addr_info[2]
|
||||
sock = socket.socket(family=family, type=type_, proto=proto)
|
||||
sock.setblocking(False)
|
||||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
|
||||
# Linux: TCP_KEEPIDLE is idle-before-first-probe.
|
||||
# macOS/Darwin: TCP_KEEPALIVE is the equivalent.
|
||||
if hasattr(socket, "TCP_KEEPIDLE"):
|
||||
sock.setsockopt(
|
||||
socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, AIOHTTP_TCP_KEEPIDLE
|
||||
)
|
||||
elif hasattr(socket, "TCP_KEEPALIVE"):
|
||||
sock.setsockopt(
|
||||
socket.IPPROTO_TCP, socket.TCP_KEEPALIVE, AIOHTTP_TCP_KEEPIDLE
|
||||
)
|
||||
if hasattr(socket, "TCP_KEEPINTVL"):
|
||||
sock.setsockopt(
|
||||
socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, AIOHTTP_TCP_KEEPINTVL
|
||||
)
|
||||
if hasattr(socket, "TCP_KEEPCNT"):
|
||||
sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, AIOHTTP_TCP_KEEPCNT)
|
||||
return sock
|
||||
|
||||
return factory
|
||||
|
||||
|
||||
def get_default_headers() -> dict:
|
||||
"""
|
||||
Get default headers for HTTP requests.
|
||||
|
|
@ -935,6 +992,11 @@ class AsyncHTTPHandler:
|
|||
transport_connector_kwargs["limit_per_host"] = (
|
||||
AIOHTTP_CONNECTOR_LIMIT_PER_HOST
|
||||
)
|
||||
# Returns None when SO_KEEPALIVE is disabled or aiohttp is too old to
|
||||
# accept socket_factory — version detection lives inside the builder.
|
||||
socket_factory = _build_aiohttp_keepalive_socket_factory()
|
||||
if socket_factory is not None:
|
||||
transport_connector_kwargs["socket_factory"] = socket_factory
|
||||
|
||||
return LiteLLMAiohttpTransport(
|
||||
client=lambda: ClientSession(
|
||||
|
|
|
|||
|
|
@ -155,6 +155,30 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
def _google_genai_streaming_hidden_params(
|
||||
*,
|
||||
api_base: str,
|
||||
litellm_params: GenericLiteLLMParams,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
response_headers: httpx.Headers,
|
||||
) -> Dict[str, Any]:
|
||||
"""Pre-stream metadata for proxy response headers (mirrors CustomStreamWrapper._hidden_params)."""
|
||||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
|
||||
_model_info: Dict[str, Any] = dict(
|
||||
getattr(litellm_params, "model_info", None) or {}
|
||||
)
|
||||
_raw_id = _model_info.get("id") or logging_obj.get_router_model_id() or ""
|
||||
_model_id = _raw_id if isinstance(_raw_id, str) else str(_raw_id)
|
||||
return {
|
||||
"model_id": _model_id,
|
||||
"api_base": api_base,
|
||||
"cache_key": "",
|
||||
"response_cost": "",
|
||||
"additional_headers": process_response_headers(response_headers),
|
||||
}
|
||||
|
||||
|
||||
class BaseLLMHTTPHandler:
|
||||
async def _make_common_async_call(
|
||||
self,
|
||||
|
|
@ -8585,6 +8609,7 @@ class BaseLLMHTTPHandler:
|
|||
api_base=api_base,
|
||||
litellm_logging_obj=logging_obj,
|
||||
litellm_params=dict(litellm_params),
|
||||
extra_body=extra_body,
|
||||
)
|
||||
else:
|
||||
(
|
||||
|
|
@ -8597,6 +8622,7 @@ class BaseLLMHTTPHandler:
|
|||
api_base=api_base,
|
||||
litellm_logging_obj=logging_obj,
|
||||
litellm_params=dict(litellm_params),
|
||||
extra_body=extra_body,
|
||||
)
|
||||
all_optional_params: Dict[str, Any] = dict(litellm_params)
|
||||
all_optional_params.update(vector_store_search_optional_params or {})
|
||||
|
|
@ -8697,6 +8723,7 @@ class BaseLLMHTTPHandler:
|
|||
api_base=api_base,
|
||||
litellm_logging_obj=logging_obj,
|
||||
litellm_params=dict(litellm_params),
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
all_optional_params: Dict[str, Any] = dict(litellm_params)
|
||||
|
|
@ -10425,6 +10452,12 @@ class BaseLLMHTTPHandler:
|
|||
litellm_metadata=litellm_metadata or {},
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
request_body=data,
|
||||
hidden_params=_google_genai_streaming_hidden_params(
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
response_headers=response.headers,
|
||||
),
|
||||
)
|
||||
else:
|
||||
response = sync_httpx_client.post(
|
||||
|
|
@ -10534,6 +10567,12 @@ class BaseLLMHTTPHandler:
|
|||
litellm_metadata=litellm_metadata or {},
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
request_body=data,
|
||||
hidden_params=_google_genai_streaming_hidden_params(
|
||||
api_base=api_base,
|
||||
litellm_params=litellm_params,
|
||||
logging_obj=logging_obj,
|
||||
response_headers=response.headers,
|
||||
),
|
||||
)
|
||||
else:
|
||||
response = await async_httpx_client.post(
|
||||
|
|
|
|||
|
|
@ -118,6 +118,7 @@ class GeminiVectorStoreConfig(BaseVectorStoreConfig):
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
"""
|
||||
Transform search request to Gemini's generateContent format.
|
||||
|
|
|
|||
|
|
@ -130,6 +130,7 @@ class MilvusVectorStoreConfig(BaseVectorStoreConfig):
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
"""
|
||||
Transform search request for Azure AI Search API
|
||||
|
|
|
|||
|
|
@ -106,6 +106,7 @@ class OpenAIVectorStoreConfig(BaseVectorStoreConfig):
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
url = f"{api_base}/{vector_store_id}/search"
|
||||
typed_request_body = VectorStoreSearchRequest(
|
||||
|
|
|
|||
|
|
@ -80,6 +80,7 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig):
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
url = f"{api_base}/{vector_store_id}/search"
|
||||
_, request_body = super().transform_search_vector_store_request(
|
||||
|
|
@ -89,5 +90,6 @@ class PGVectorStoreConfig(OpenAIVectorStoreConfig):
|
|||
api_base=api_base,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
litellm_params=litellm_params,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
return url, request_body
|
||||
|
|
|
|||
|
|
@ -102,6 +102,7 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig):
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
"""RAGFlow vector stores are management-only, search is not supported."""
|
||||
raise NotImplementedError(
|
||||
|
|
|
|||
|
|
@ -79,6 +79,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
"""Sync version - generates embedding synchronously."""
|
||||
# For S3 Vectors, vector_store_id should be in format: bucket_name:index_name
|
||||
|
|
@ -140,6 +141,7 @@ class S3VectorsVectorStoreConfig(BaseVectorStoreConfig, BaseAWSLLM):
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict]:
|
||||
"""Async version - generates embedding asynchronously."""
|
||||
# For S3 Vectors, vector_store_id should be in format: bucket_name:index_name
|
||||
|
|
|
|||
|
|
@ -2395,8 +2395,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
completion_response = GenerateContentResponseBody(**raw_response.json()) # type: ignore
|
||||
except Exception as e:
|
||||
raise VertexAIError(
|
||||
message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
|
||||
raw_response.text, str(e)
|
||||
message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
|
||||
str(e)
|
||||
),
|
||||
status_code=422,
|
||||
headers=raw_response.headers,
|
||||
|
|
@ -2530,8 +2530,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
|
||||
except Exception as e:
|
||||
raise VertexAIError(
|
||||
message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
|
||||
completion_response, str(e)
|
||||
message="Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format(
|
||||
str(e)
|
||||
),
|
||||
status_code=422,
|
||||
headers=raw_response.headers,
|
||||
|
|
|
|||
|
|
@ -100,6 +100,7 @@ class VertexVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
"""
|
||||
Transform search request for Vertex AI RAG API
|
||||
|
|
|
|||
|
|
@ -107,6 +107,7 @@ class VertexSearchAPIVectorStoreConfig(BaseVectorStoreConfig, VertexBase):
|
|||
api_base: str,
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
litellm_params: dict,
|
||||
extra_body: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[str, Dict[str, Any]]:
|
||||
"""
|
||||
Transform search request for Vertex AI RAG API
|
||||
|
|
|
|||
|
|
@ -50,8 +50,11 @@ from litellm.proxy._experimental.mcp_server.oauth2_token_cache import resolve_mc
|
|||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
MCP_TOOL_PREFIX_SEPARATOR,
|
||||
add_server_prefix_to_name,
|
||||
compute_short_server_prefix,
|
||||
get_server_prefix,
|
||||
is_short_mcp_tool_prefix_enabled,
|
||||
is_tool_name_prefixed,
|
||||
iter_known_server_prefixes,
|
||||
merge_mcp_headers,
|
||||
normalize_server_name,
|
||||
split_server_prefix_from_name,
|
||||
|
|
@ -106,6 +109,12 @@ if not _separator_probe.is_valid:
|
|||
SEP_986_URL,
|
||||
)
|
||||
|
||||
_AZURE_ENTRA_HOSTS = {
|
||||
"login.microsoftonline.com", # Global
|
||||
"login.microsoftonline.us", # US Government
|
||||
"login.chinacloudapi.cn", # China
|
||||
}
|
||||
|
||||
|
||||
def _warn_on_server_name_fields(
|
||||
*,
|
||||
|
|
@ -364,6 +373,7 @@ class MCPServerManager:
|
|||
aws_session_name=server_config.get("aws_session_name", None),
|
||||
instructions=server_config.get("instructions", None),
|
||||
)
|
||||
self._assign_unique_short_prefix(new_server)
|
||||
self.config_mcp_servers[server_id] = new_server
|
||||
|
||||
# Check if this is an OpenAPI-based server
|
||||
|
|
@ -726,6 +736,7 @@ class MCPServerManager:
|
|||
try:
|
||||
if mcp_server.server_id not in self.registry:
|
||||
new_server = await self.build_mcp_server_from_table(mcp_server)
|
||||
self._assign_unique_short_prefix(new_server)
|
||||
self.registry[mcp_server.server_id] = new_server
|
||||
await self._maybe_register_openapi_tools(new_server)
|
||||
verbose_logger.debug(f"Added MCP Server: {new_server.name}")
|
||||
|
|
@ -738,6 +749,12 @@ class MCPServerManager:
|
|||
try:
|
||||
if mcp_server.server_id in self.registry:
|
||||
new_server = await self.build_mcp_server_from_table(mcp_server)
|
||||
# Carry the previously-resolved short prefix across so the
|
||||
# tool names stay stable for clients holding cached lists.
|
||||
existing_prefix = self.registry[mcp_server.server_id].short_prefix
|
||||
if existing_prefix and not new_server.short_prefix:
|
||||
new_server.short_prefix = existing_prefix
|
||||
self._assign_unique_short_prefix(new_server)
|
||||
self.registry[mcp_server.server_id] = new_server
|
||||
await self._maybe_register_openapi_tools(new_server)
|
||||
verbose_logger.debug(f"Updated MCP Server: {new_server.name}")
|
||||
|
|
@ -1236,7 +1253,11 @@ class MCPServerManager:
|
|||
|
||||
## HANDLE OPENAPI TOOLS
|
||||
if server.spec_path:
|
||||
_tools = global_mcp_tool_registry.list_tools(tool_prefix=server.name)
|
||||
# OpenAPI tools were stored in the registry under the prefix
|
||||
# active at registration time — fetch by that same prefix.
|
||||
_tools = global_mcp_tool_registry.list_tools(
|
||||
tool_prefix=get_server_prefix(server)
|
||||
)
|
||||
tools = global_mcp_tool_registry.convert_tools_to_mcp_sdk_tool_type(
|
||||
_tools
|
||||
)
|
||||
|
|
@ -1488,11 +1509,28 @@ class MCPServerManager:
|
|||
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
|
||||
response = await client.get(server_url)
|
||||
response.raise_for_status()
|
||||
verbose_logger.warning(
|
||||
"MCP OAuth discovery unexpectedly succeeded for %s; server did not challenge",
|
||||
server_url,
|
||||
(
|
||||
authorization_servers,
|
||||
resource_scopes,
|
||||
) = await self._attempt_well_known_discovery(server_url)
|
||||
metadata = await self._fetch_authorization_server_metadata(
|
||||
authorization_servers
|
||||
)
|
||||
raise RuntimeError("OAuth discovery must not succeed without a challenge")
|
||||
if (
|
||||
metadata is None
|
||||
and not resource_scopes
|
||||
and authorization_servers
|
||||
and response.status_code == 200
|
||||
):
|
||||
verbose_logger.warning(
|
||||
"MCP OAuth discovery for %s received 200 OK without RFC 9728 challenge and no discoverable authorization metadata.",
|
||||
server_url,
|
||||
)
|
||||
if metadata is None and resource_scopes:
|
||||
return MCPOAuthMetadata(scopes=resource_scopes)
|
||||
if metadata is not None and resource_scopes:
|
||||
metadata.scopes = resource_scopes
|
||||
return metadata
|
||||
except HTTPStatusError as exc:
|
||||
verbose_logger.debug(
|
||||
"MCP OAuth discovery for %s received status error: %s",
|
||||
|
|
@ -1510,8 +1548,8 @@ class MCPServerManager:
|
|||
header_value
|
||||
)
|
||||
|
||||
authorization_servers: List[str] = []
|
||||
resource_scopes: Optional[List[str]] = None
|
||||
authorization_servers = []
|
||||
resource_scopes = None
|
||||
if resource_metadata_url:
|
||||
(
|
||||
authorization_servers,
|
||||
|
|
@ -1674,6 +1712,9 @@ class MCPServerManager:
|
|||
f"{base}/.well-known/oauth-authorization-server/{path}"
|
||||
)
|
||||
candidate_urls.append(f"{base}/.well-known/openid-configuration/{path}")
|
||||
candidate_urls.append(
|
||||
f"{issuer_url.rstrip('/')}/.well-known/openid-configuration"
|
||||
)
|
||||
candidate_urls.append(f"{base}/.well-known/oauth-authorization-server")
|
||||
candidate_urls.append(f"{base}/.well-known/openid-configuration")
|
||||
candidate_urls.append(issuer_url.rstrip("/"))
|
||||
|
|
@ -1713,7 +1754,28 @@ class MCPServerManager:
|
|||
):
|
||||
return metadata
|
||||
|
||||
return None
|
||||
return self._build_azure_authorization_server_metadata(parsed)
|
||||
|
||||
@staticmethod
|
||||
def _build_azure_authorization_server_metadata(
|
||||
parsed_issuer_url: Any,
|
||||
) -> Optional[MCPOAuthMetadata]:
|
||||
path_parts = [
|
||||
part for part in (parsed_issuer_url.path or "").split("/") if part
|
||||
]
|
||||
if (
|
||||
parsed_issuer_url.netloc not in _AZURE_ENTRA_HOSTS
|
||||
or len(path_parts) != 2
|
||||
or path_parts[1] != "v2.0"
|
||||
):
|
||||
return None
|
||||
|
||||
tenant = path_parts[0]
|
||||
base = f"{parsed_issuer_url.scheme}://{parsed_issuer_url.netloc}/{tenant}"
|
||||
return MCPOAuthMetadata(
|
||||
authorization_url=f"{base}/oauth2/v2.0/authorize",
|
||||
token_url=f"{base}/oauth2/v2.0/token",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _decrypt_credential_field(
|
||||
|
|
@ -1810,6 +1872,63 @@ class MCPServerManager:
|
|||
verbose_logger.warning(f"Error listing tools from {server_name}: {str(e)}")
|
||||
return []
|
||||
|
||||
_SHORT_PREFIX_MAX_REHASH_ATTEMPTS = 1024
|
||||
|
||||
def _assign_unique_short_prefix(self, server: MCPServer) -> None:
|
||||
"""Resolve and cache a collision-free short tool prefix on ``server``.
|
||||
|
||||
Called at registration time for every MCP server entering the
|
||||
registry. Mutates ``server.short_prefix`` in place. No-ops when
|
||||
``LITELLM_USE_SHORT_MCP_TOOL_PREFIX`` is disabled, when the server
|
||||
has no ``server_id`` (synthetic temp-server objects), or when a
|
||||
prefix is already cached.
|
||||
|
||||
Collision strategy: take the natural hash; if it's already used by
|
||||
a *different* server in the combined registry, rehash with an
|
||||
incrementing attempt counter until we find an unused slot. The
|
||||
attempt counter is folded into the hash so the resulting prefix is
|
||||
still deterministic for a given (server_id, set-of-other-server-ids)
|
||||
pair within one process.
|
||||
"""
|
||||
if not is_short_mcp_tool_prefix_enabled():
|
||||
return
|
||||
if server.short_prefix:
|
||||
return
|
||||
if not server.server_id:
|
||||
return
|
||||
|
||||
used: Dict[str, str] = {}
|
||||
for other in self.get_registry().values():
|
||||
if other.server_id == server.server_id:
|
||||
continue
|
||||
if other.short_prefix:
|
||||
used[other.short_prefix] = other.server_id
|
||||
|
||||
for attempt in range(self._SHORT_PREFIX_MAX_REHASH_ATTEMPTS):
|
||||
candidate = compute_short_server_prefix(server.server_id, attempt=attempt)
|
||||
if candidate not in used:
|
||||
server.short_prefix = candidate
|
||||
if attempt > 0:
|
||||
verbose_logger.info(
|
||||
"MCP short-prefix collision resolved for server %s: "
|
||||
"natural hash collided with %s, using rehashed prefix "
|
||||
"%s (attempt=%d).",
|
||||
server.server_id,
|
||||
used.get(
|
||||
compute_short_server_prefix(server.server_id, attempt=0),
|
||||
"<unknown>",
|
||||
),
|
||||
candidate,
|
||||
attempt,
|
||||
)
|
||||
return
|
||||
|
||||
raise RuntimeError(
|
||||
f"Unable to assign a unique short MCP tool prefix for server "
|
||||
f"{server.server_id} after {self._SHORT_PREFIX_MAX_REHASH_ATTEMPTS} "
|
||||
"attempts; the 3-character prefix space is too crowded."
|
||||
)
|
||||
|
||||
def _create_prefixed_tools(
|
||||
self, tools: List[MCPTool], server: MCPServer, add_prefix: bool = True
|
||||
) -> List[MCPTool]:
|
||||
|
|
@ -1838,9 +1957,13 @@ class MCPServerManager:
|
|||
tool_copy.name = name_to_use
|
||||
prefixed_tools.append(tool_copy)
|
||||
|
||||
# Update tool to server mapping for resolution (support both forms)
|
||||
# Register every known prefix form (alias, server_name, server_id,
|
||||
# short ID) so call_tool can resolve regardless of which form a
|
||||
# caller / cached client is using.
|
||||
self.tool_name_to_mcp_server_name_mapping[original_name] = prefix
|
||||
self.tool_name_to_mcp_server_name_mapping[prefixed_name] = prefix
|
||||
for known_prefix in iter_known_server_prefixes(server):
|
||||
qualified = add_server_prefix_to_name(original_name, known_prefix)
|
||||
self.tool_name_to_mcp_server_name_mapping[qualified] = prefix
|
||||
|
||||
verbose_logger.info(
|
||||
f"Successfully fetched {len(prefixed_tools)} tools from server {server.name}"
|
||||
|
|
@ -2601,37 +2724,43 @@ class MCPServerManager:
|
|||
Returns:
|
||||
MCPServer if found, None otherwise
|
||||
"""
|
||||
registry_servers = list(self.get_registry().values())
|
||||
|
||||
# Build prefix → server lookup covering every known form a tool name
|
||||
# may take (alias / server_name / server_id / short ID). This is what
|
||||
# makes the short-prefix mode work without breaking historical names.
|
||||
prefix_to_server: Dict[str, MCPServer] = {}
|
||||
for server in registry_servers:
|
||||
for known_prefix in iter_known_server_prefixes(server):
|
||||
normalised = normalize_server_name(known_prefix)
|
||||
prefix_to_server.setdefault(normalised, server)
|
||||
|
||||
# First try with the original tool name
|
||||
if tool_name in self.tool_name_to_mcp_server_name_mapping:
|
||||
server_name = self.tool_name_to_mcp_server_name_mapping[tool_name]
|
||||
for server in self.get_registry().values():
|
||||
if normalize_server_name(server.name) == normalize_server_name(
|
||||
server_name
|
||||
):
|
||||
normalised_lookup = normalize_server_name(server_name)
|
||||
if normalised_lookup in prefix_to_server:
|
||||
return prefix_to_server[normalised_lookup]
|
||||
for server in registry_servers:
|
||||
if normalize_server_name(server.name) == normalised_lookup:
|
||||
return server
|
||||
|
||||
# If not found and tool name is prefixed, try extracting server name from prefix
|
||||
known_prefixes = {
|
||||
normalize_server_name(get_server_prefix(s))
|
||||
for s in self.get_registry().values()
|
||||
if get_server_prefix(s)
|
||||
}
|
||||
if is_tool_name_prefixed(tool_name, known_server_prefixes=known_prefixes):
|
||||
# If not found and tool name is prefixed, extract the prefix and
|
||||
# match against any known form.
|
||||
if is_tool_name_prefixed(
|
||||
tool_name, known_server_prefixes=set(prefix_to_server.keys())
|
||||
):
|
||||
(
|
||||
original_tool_name,
|
||||
server_name_from_prefix,
|
||||
) = split_server_prefix_from_name(tool_name)
|
||||
if original_tool_name in self.tool_name_to_mcp_server_name_mapping:
|
||||
for server in self.get_registry().values():
|
||||
if server.server_name is None:
|
||||
if normalize_server_name(server.name) == normalize_server_name(
|
||||
server_name_from_prefix
|
||||
):
|
||||
return server
|
||||
elif normalize_server_name(
|
||||
server.server_name
|
||||
) == normalize_server_name(server_name_from_prefix):
|
||||
return server
|
||||
normalised_prefix = normalize_server_name(server_name_from_prefix)
|
||||
matched_server = prefix_to_server.get(normalised_prefix)
|
||||
if matched_server is not None and (
|
||||
original_tool_name in self.tool_name_to_mcp_server_name_mapping
|
||||
or tool_name in self.tool_name_to_mcp_server_name_mapping
|
||||
):
|
||||
return matched_server
|
||||
|
||||
return None
|
||||
|
||||
|
|
@ -2666,6 +2795,9 @@ class MCPServerManager:
|
|||
previous_registry = self.registry
|
||||
new_registry: Dict[str, MCPServer] = {}
|
||||
|
||||
# Stage one: build every server. Stage two assigns short prefixes
|
||||
# against the *full* set so dedup is deterministic regardless of
|
||||
# iteration order.
|
||||
for server in db_mcp_servers:
|
||||
existing_server = previous_registry.get(server.server_id)
|
||||
|
||||
|
|
@ -2689,10 +2821,21 @@ class MCPServerManager:
|
|||
f"Building server from DB: {server.server_id} ({server.server_name})"
|
||||
)
|
||||
new_server = await self.build_mcp_server_from_table(server)
|
||||
# Carry the cached short_prefix from the previous registry entry
|
||||
# (if any) so the prefix is stable across reloads.
|
||||
if existing_server is not None and existing_server.short_prefix:
|
||||
new_server.short_prefix = existing_server.short_prefix
|
||||
new_registry[server.server_id] = new_server
|
||||
await self._maybe_register_openapi_tools(new_server)
|
||||
|
||||
# Swap in the new registry first so _assign_unique_short_prefix
|
||||
# sees the complete set when checking for collisions.
|
||||
self.registry = new_registry
|
||||
for new_server in new_registry.values():
|
||||
self._assign_unique_short_prefix(new_server)
|
||||
# Register OpenAPI tools *after* the final short prefix is assigned
|
||||
# so the tools are stored in the global registry under the same
|
||||
# prefix that lookups will use.
|
||||
await self._maybe_register_openapi_tools(new_server)
|
||||
|
||||
verbose_logger.debug(
|
||||
"MCP registry refreshed (%s servers in registry)", len(new_registry)
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ from litellm.proxy._experimental.mcp_server.utils import (
|
|||
LITELLM_MCP_SERVER_VERSION,
|
||||
add_server_prefix_to_name,
|
||||
get_server_prefix,
|
||||
iter_known_server_prefixes,
|
||||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.ip_address_utils import IPAddressUtils
|
||||
|
|
@ -711,13 +712,7 @@ if MCP_AVAILABLE:
|
|||
for server in allowed_mcp_servers:
|
||||
if server:
|
||||
match_list = [
|
||||
s.lower()
|
||||
for s in [
|
||||
server.alias,
|
||||
server.server_name,
|
||||
server.server_id,
|
||||
]
|
||||
if s is not None
|
||||
s.lower() for s in iter_known_server_prefixes(server) if s
|
||||
]
|
||||
|
||||
if server_or_group.lower() in match_list:
|
||||
|
|
@ -2031,11 +2026,13 @@ if MCP_AVAILABLE:
|
|||
# Remove prefix from tool name for logging and processing
|
||||
original_tool_name, server_name = split_server_prefix_from_name(name)
|
||||
|
||||
# If tool name is unprefixed, resolve its server so we can enforce permissions
|
||||
if not server_name:
|
||||
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
|
||||
if mcp_server:
|
||||
server_name = mcp_server.name
|
||||
# Resolve the actual MCP server up-front so the permission check uses
|
||||
# the canonical server.name even when the tool name is prefixed with a
|
||||
# short ID (LITELLM_USE_SHORT_MCP_TOOL_PREFIX) that doesn't match the
|
||||
# server's display name directly.
|
||||
mcp_server = global_mcp_server_manager._get_mcp_server_from_tool_name(name)
|
||||
if mcp_server is not None:
|
||||
server_name = mcp_server.name
|
||||
|
||||
# Only enforce server-level permissions when we can resolve a server
|
||||
if server_name:
|
||||
|
|
|
|||
|
|
@ -2,10 +2,11 @@
|
|||
MCP Server Utilities
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, Mapping, Optional, Tuple
|
||||
from typing import Any, Dict, Iterator, Mapping, Optional, Tuple
|
||||
|
||||
import os
|
||||
import hashlib
|
||||
import importlib
|
||||
import os
|
||||
|
||||
# Constants
|
||||
LITELLM_MCP_SERVER_NAME = "litellm-mcp-server"
|
||||
|
|
@ -14,6 +15,89 @@ LITELLM_MCP_SERVER_DESCRIPTION = "MCP Server for LiteLLM"
|
|||
MCP_TOOL_PREFIX_SEPARATOR = os.environ.get("MCP_TOOL_PREFIX_SEPARATOR", "-")
|
||||
MCP_TOOL_PREFIX_FORMAT = "{server_name}{separator}{tool_name}"
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Short-ID tool prefix (opt-in)
|
||||
# ---------------------------------------------------------------------------
|
||||
# When LITELLM_USE_SHORT_MCP_TOOL_PREFIX is truthy the prefix attached to MCP
|
||||
# tool / prompt / resource / resource-template names switches from the
|
||||
# (potentially long) human-readable server name to a deterministic three
|
||||
# character ID derived from the server's ``server_id``.
|
||||
#
|
||||
# Why three characters?
|
||||
# * The first character is restricted to 52 alphabetic characters
|
||||
# ([A-Za-z]) and the remaining two characters use the full base62
|
||||
# alphabet ([0-9A-Za-z]). That guarantees the prefix never starts
|
||||
# with a digit so it remains a valid identifier for every model API
|
||||
# (some providers historically required a leading alphabetic char).
|
||||
# * 52 * 62 * 62 = 199_888 distinct IDs. The chance of a real local
|
||||
# tool name happening to begin with the exact prefix LiteLLM assigned
|
||||
# to a given MCP server is negligible in practice.
|
||||
# * The IDs are short enough that prefixed tool names stay well under
|
||||
# the 60-character upper bound enforced by some model APIs (Anthropic
|
||||
# etc.) even for long upstream tool names.
|
||||
# * The mapping is deterministic (SHA-256 of ``server_id`` → three
|
||||
# characters drawn from the alphabets above), so the prefix is stable
|
||||
# across processes, workers and restarts without any persistence
|
||||
# layer. Two servers with different ``server_id`` values can in
|
||||
# principle hash to the same three chars; that natural-hash collision
|
||||
# IS a routing-correctness issue (the second registrant would otherwise
|
||||
# have its tools misrouted to the first), so registration goes through
|
||||
# ``MCPServerManager._assign_unique_short_prefix`` which rehashes with
|
||||
# a deterministic attempt counter until it finds an unused prefix and
|
||||
# caches the result on ``MCPServer.short_prefix``. A collision is
|
||||
# logged at INFO when it happens.
|
||||
#
|
||||
# This flag is intentionally opt-in for the first release so customers can
|
||||
# migrate. It will become the default in a future release.
|
||||
SHORT_MCP_TOOL_PREFIX_LENGTH = 3
|
||||
_BASE62_ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||||
# Subset of _BASE62_ALPHABET used for the *first* character only, to
|
||||
# guarantee the prefix never starts with a digit.
|
||||
_BASE52_ALPHA_ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
|
||||
|
||||
|
||||
def is_short_mcp_tool_prefix_enabled() -> bool:
|
||||
"""Return True when the short-ID tool prefix mode is enabled.
|
||||
|
||||
Read at call time (not import time) so tests and runtime config changes
|
||||
take effect without reimporting the module.
|
||||
"""
|
||||
raw = os.environ.get("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "")
|
||||
return raw.strip().lower() in ("1", "true", "yes", "on")
|
||||
|
||||
|
||||
def compute_short_server_prefix(server_id: str, attempt: int = 0) -> str:
|
||||
"""Derive the deterministic three-character prefix for a server.
|
||||
|
||||
Uses SHA-256 of ``f"{server_id}#{attempt}"`` and folds the first eight
|
||||
bytes into a fixed-length string whose first character is drawn from
|
||||
``_BASE52_ALPHA_ALPHABET`` (so the prefix never starts with a digit)
|
||||
and whose remaining characters are drawn from the full base62
|
||||
alphabet. Pass ``attempt > 0`` to rehash to a different prefix when
|
||||
the natural hash collides with a prefix already assigned to another
|
||||
server (see ``MCPServerManager._assign_unique_short_prefix``). An
|
||||
empty ``server_id`` raises ``ValueError`` — short prefixes require a
|
||||
stable identifier to be deterministic.
|
||||
"""
|
||||
if not server_id:
|
||||
raise ValueError("compute_short_server_prefix requires a non-empty server_id")
|
||||
|
||||
seed = server_id if attempt == 0 else f"{server_id}#{attempt}"
|
||||
digest = hashlib.sha256(seed.encode("utf-8")).digest()
|
||||
value = int.from_bytes(digest[:8], "big")
|
||||
|
||||
# Build chars from least-significant to most-significant; we reverse
|
||||
# at the end so the first emitted char comes from the high-order
|
||||
# bits of the digest (which is the position we constrain to be
|
||||
# alphabetic).
|
||||
chars = []
|
||||
for position in range(SHORT_MCP_TOOL_PREFIX_LENGTH):
|
||||
is_first_char = position == SHORT_MCP_TOOL_PREFIX_LENGTH - 1
|
||||
alphabet = _BASE52_ALPHA_ALPHABET if is_first_char else _BASE62_ALPHABET
|
||||
value, idx = divmod(value, len(alphabet))
|
||||
chars.append(alphabet[idx])
|
||||
return "".join(reversed(chars))
|
||||
|
||||
|
||||
def is_mcp_available() -> bool:
|
||||
"""
|
||||
|
|
@ -82,7 +166,25 @@ def add_server_prefix_to_name(name: str, server_name: str) -> str:
|
|||
|
||||
|
||||
def get_server_prefix(server: Any) -> str:
|
||||
"""Return the prefix for a server: alias if present, else server_name, else server_id"""
|
||||
"""Return the prefix for a server.
|
||||
|
||||
When the short-prefix mode is enabled (``LITELLM_USE_SHORT_MCP_TOOL_PREFIX``)
|
||||
a three-character base62 ID is returned. We prefer the cached
|
||||
``server.short_prefix`` value when set — that field is populated at
|
||||
registration time by ``MCPServerManager._assign_unique_short_prefix``
|
||||
and resolves natural-hash collisions deterministically — and only fall
|
||||
back to the natural hash for ad-hoc / temp-server objects without a
|
||||
cached value. In default mode the historical behaviour is preserved:
|
||||
alias if present, else server_name, else server_id.
|
||||
"""
|
||||
if is_short_mcp_tool_prefix_enabled():
|
||||
cached = getattr(server, "short_prefix", None)
|
||||
if cached:
|
||||
return cached
|
||||
server_id = getattr(server, "server_id", None)
|
||||
if server_id:
|
||||
return compute_short_server_prefix(server_id)
|
||||
|
||||
if hasattr(server, "alias") and server.alias:
|
||||
return server.alias
|
||||
if hasattr(server, "server_name") and server.server_name:
|
||||
|
|
@ -92,6 +194,36 @@ def get_server_prefix(server: Any) -> str:
|
|||
return ""
|
||||
|
||||
|
||||
def iter_known_server_prefixes(server: Any) -> Iterator[str]:
|
||||
"""Yield every prefix form that may appear in tool names for ``server``.
|
||||
|
||||
Always includes the *current* prefix returned by ``get_server_prefix``.
|
||||
Additionally yields the historical (alias / server_name / server_id) and
|
||||
short-ID forms so the routing layer can resolve tool names regardless of
|
||||
which prefix mode was active when the client first observed them.
|
||||
"""
|
||||
seen = set()
|
||||
|
||||
def _emit(value: Optional[str]) -> Iterator[str]:
|
||||
if value and value not in seen:
|
||||
seen.add(value)
|
||||
yield value
|
||||
|
||||
yield from _emit(get_server_prefix(server))
|
||||
yield from _emit(getattr(server, "short_prefix", None))
|
||||
|
||||
server_id = getattr(server, "server_id", None)
|
||||
if server_id:
|
||||
try:
|
||||
yield from _emit(compute_short_server_prefix(server_id))
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
yield from _emit(getattr(server, "alias", None))
|
||||
yield from _emit(getattr(server, "server_name", None))
|
||||
yield from _emit(server_id)
|
||||
|
||||
|
||||
def split_server_prefix_from_name(prefixed_name: str) -> Tuple[str, str]:
|
||||
"""Return the unprefixed name plus the server name used as prefix."""
|
||||
if MCP_TOOL_PREFIX_SEPARATOR in prefixed_name:
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -1,22 +0,0 @@
|
|||
1:"$Sreact.fragment"
|
||||
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
|
||||
7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
|
||||
a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
|
||||
b:"$Sreact.suspense"
|
||||
d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
|
||||
f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
|
||||
11:I[168027,[],"default"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
|
||||
0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true}
|
||||
8:{}
|
||||
9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params"
|
||||
e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]
|
||||
12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
|
||||
c:null
|
||||
10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]]
|
||||
|
|
@ -1,22 +0,0 @@
|
|||
1:"$Sreact.fragment"
|
||||
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
6:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
|
||||
7:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
|
||||
a:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
|
||||
b:"$Sreact.suspense"
|
||||
d:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
|
||||
f:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
|
||||
11:I[168027,[],"default"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
|
||||
0:{"P":null,"b":"zxkD4-EPlgfKHDTw8O869","c":["","chat"],"q":"","i":false,"f":[[["",{"children":["chat",{"children":["__PAGE__",{}]}]},"$undefined","$undefined",true],[["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next","crossOrigin":"$undefined","nonce":"$undefined"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true,"nonce":"$undefined"}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]],"forbidden":"$undefined","unauthorized":"$undefined"}]}]}]}]}]]}],{"children":[["$","$1","c",{"children":[null,["$","$L4",null,{"parallelRouterKey":"children","error":"$undefined","errorStyles":"$undefined","errorScripts":"$undefined","template":["$","$L5",null,{}],"templateStyles":"$undefined","templateScripts":"$undefined","notFound":"$undefined","forbidden":"$undefined","unauthorized":"$undefined"}]]}],{"children":[["$","$1","c",{"children":[["$","$L6",null,{"Component":"$7","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@8","$@9"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true,"nonce":"$undefined"}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true,"nonce":"$undefined"}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true,"nonce":"$undefined"}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true,"nonce":"$undefined"}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true,"nonce":"$undefined"}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true,"nonce":"$undefined"}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true,"nonce":"$undefined"}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true,"nonce":"$undefined"}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true,"nonce":"$undefined"}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true,"nonce":"$undefined"}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true,"nonce":"$undefined"}]],["$","$La",null,{"children":["$","$b",null,{"name":"Next.MetadataOutlet","children":"$@c"}]}]]}],{},null,false,false]},null,false,false]},null,false,false],["$","$1","h",{"children":[null,["$","$Ld",null,{"children":"$Le"}],["$","div",null,{"hidden":true,"children":["$","$Lf",null,{"children":["$","$b",null,{"name":"Next.Metadata","children":"$L10"}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],false]],"m":"$undefined","G":["$11",[]],"S":true}
|
||||
8:{}
|
||||
9:"$0:f:0:1:1:children:1:children:0:props:children:0:props:serverProvidedParams:params"
|
||||
e:[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]
|
||||
12:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
|
||||
c:null
|
||||
10:[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L12","4",{}]]
|
||||
|
|
@ -1,6 +0,0 @@
|
|||
1:"$Sreact.fragment"
|
||||
2:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ViewportBoundary"]
|
||||
3:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"MetadataBoundary"]
|
||||
4:"$Sreact.suspense"
|
||||
5:I[27201,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"IconMark"]
|
||||
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","h",{"children":[null,["$","$L2",null,{"children":[["$","meta","0",{"charSet":"utf-8"}],["$","meta","1",{"name":"viewport","content":"width=device-width, initial-scale=1"}]]}],["$","div",null,{"hidden":true,"children":["$","$L3",null,{"children":["$","$4",null,{"name":"Next.Metadata","children":[["$","title","0",{"children":"LiteLLM Dashboard"}],["$","meta","1",{"name":"description","content":"LiteLLM Proxy Admin UI"}],["$","link","2",{"rel":"icon","href":"/favicon.ico?favicon.1d32c690.ico","sizes":"48x48","type":"image/x-icon"}],["$","link","3",{"rel":"icon","href":"./favicon.ico"}],["$","$L5","4",{}]]}]}]}],["$","meta",null,{"name":"next-size-adjust","content":""}]]}],"loading":null,"isPartial":false}
|
||||
|
|
@ -1,8 +0,0 @@
|
|||
1:"$Sreact.fragment"
|
||||
2:I[867271,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
3:I[71195,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js"],"default"]
|
||||
4:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
5:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
|
||||
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[[["$","link","0",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","precedence":"next"}],["$","link","1",{"rel":"stylesheet","href":"/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","precedence":"next"}],["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","async":true}]],["$","html",null,{"lang":"en","children":["$","body",null,{"className":"inter_5972bc34-module__OU16Qa__className","children":["$","$L2",null,{"children":["$","$L3",null,{"children":["$","$L4",null,{"parallelRouterKey":"children","template":["$","$L5",null,{}],"notFound":[[["$","title",null,{"children":"404: This page could not be found."}],["$","div",null,{"style":{"fontFamily":"system-ui,\"Segoe UI\",Roboto,Helvetica,Arial,sans-serif,\"Apple Color Emoji\",\"Segoe UI Emoji\"","height":"100vh","textAlign":"center","display":"flex","flexDirection":"column","alignItems":"center","justifyContent":"center"},"children":["$","div",null,{"children":[["$","style",null,{"dangerouslySetInnerHTML":{"__html":"body{color:#000;background:#fff;margin:0}.next-error-h1{border-right:1px solid rgba(0,0,0,.3)}@media (prefers-color-scheme:dark){body{color:#fff;background:#000}.next-error-h1{border-right:1px solid rgba(255,255,255,.3)}}"}}],["$","h1",null,{"className":"next-error-h1","style":{"display":"inline-block","margin":"0 20px 0 0","padding":"0 23px 0 0","fontSize":24,"fontWeight":500,"verticalAlign":"top","lineHeight":"49px"},"children":404}],["$","div",null,{"style":{"display":"inline-block"},"children":["$","h2",null,{"style":{"fontSize":14,"fontWeight":400,"lineHeight":"49px","margin":0},"children":"This page could not be found."}]}]]}]}]],[]]}]}]}]}]}]]}],"loading":null,"isPartial":false}
|
||||
|
|
@ -1,4 +0,0 @@
|
|||
:HL["/litellm-asset-prefix/_next/static/chunks/4e20891f2fd03463.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/chunks/e0cb6755699177c1.css","style"]
|
||||
:HL["/litellm-asset-prefix/_next/static/media/83afe278b6a6bb3c-s.p.3a6ba036.woff2","font",{"crossOrigin":"","type":"font/woff2"}]
|
||||
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","tree":{"name":"","paramType":null,"paramKey":"","hasRuntimePrefetch":false,"slots":{"children":{"name":"chat","paramType":null,"paramKey":"chat","hasRuntimePrefetch":false,"slots":{"children":{"name":"__PAGE__","paramType":null,"paramKey":"__PAGE__","hasRuntimePrefetch":false,"slots":null,"isRootLayout":false}},"isRootLayout":false}},"isRootLayout":true},"staleTime":300}
|
||||
|
|
@ -1,9 +0,0 @@
|
|||
1:"$Sreact.fragment"
|
||||
2:I[347257,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"ClientPageRoot"]
|
||||
3:I[321443,["/litellm-asset-prefix/_next/static/chunks/9e09de50158b3159.js","/litellm-asset-prefix/_next/static/chunks/7e5fe5584502da06.js","/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js"],"default"]
|
||||
6:I[897367,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"OutletBoundary"]
|
||||
7:"$Sreact.suspense"
|
||||
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[["$","$L2",null,{"Component":"$3","serverProvidedParams":{"searchParams":{},"params":{},"promises":["$@4","$@5"]}}],[["$","script","script-0",{"src":"/litellm-asset-prefix/_next/static/chunks/0a6c418370a8c183.js","async":true}],["$","script","script-1",{"src":"/litellm-asset-prefix/_next/static/chunks/cf9c81fc7166f4d4.js","async":true}],["$","script","script-2",{"src":"/litellm-asset-prefix/_next/static/chunks/542a1a209eb732c6.js","async":true}],["$","script","script-3",{"src":"/litellm-asset-prefix/_next/static/chunks/00ff280cdb7d7ee5.js","async":true}],["$","script","script-4",{"src":"/litellm-asset-prefix/_next/static/chunks/a1abfc2f35c701cc.js","async":true}],["$","script","script-5",{"src":"/litellm-asset-prefix/_next/static/chunks/4980372eaa37b78b.js","async":true}],["$","script","script-6",{"src":"/litellm-asset-prefix/_next/static/chunks/119ecee91f911bc8.js","async":true}],["$","script","script-7",{"src":"/litellm-asset-prefix/_next/static/chunks/ac3cf77acb5bf234.js","async":true}],["$","script","script-8",{"src":"/litellm-asset-prefix/_next/static/chunks/fcdf7322b0aa3e2e.js","async":true}],["$","script","script-9",{"src":"/litellm-asset-prefix/_next/static/chunks/7e417dd24c8becd0.js","async":true}],["$","script","script-10",{"src":"/litellm-asset-prefix/_next/static/chunks/89aa55578de861b7.js","async":true}]],["$","$L6",null,{"children":["$","$7",null,{"name":"Next.MetadataOutlet","children":"$@8"}]}]]}],"loading":null,"isPartial":false}
|
||||
4:{}
|
||||
5:"$0:rsc:props:children:0:props:serverProvidedParams:params"
|
||||
8:null
|
||||
|
|
@ -1,4 +0,0 @@
|
|||
1:"$Sreact.fragment"
|
||||
2:I[339756,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
3:I[837457,["/litellm-asset-prefix/_next/static/chunks/d96012bcfc98706a.js","/litellm-asset-prefix/_next/static/chunks/dbca964212122d58.js"],"default"]
|
||||
0:{"buildId":"zxkD4-EPlgfKHDTw8O869","rsc":["$","$1","c",{"children":[null,["$","$L2",null,{"parallelRouterKey":"children","template":["$","$L3",null,{}]}]]}],"loading":null,"isPartial":false}
|
||||
|
|
@ -619,6 +619,67 @@ class ProxyBaseLLMRequestProcessing:
|
|||
verbose_proxy_logger.error(f"Error setting custom headers: {e}")
|
||||
return {}
|
||||
|
||||
@staticmethod
|
||||
async def build_litellm_proxy_success_headers_from_llm_response(
|
||||
*,
|
||||
response: Any,
|
||||
request_data: dict,
|
||||
request: Request,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
version: Optional[str],
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> Dict[str, str]:
|
||||
"""
|
||||
Build LiteLLM proxy response headers for routes that call the LLM directly
|
||||
(e.g. Google native :generateContent) instead of base_process_llm_request.
|
||||
"""
|
||||
if isinstance(response, dict):
|
||||
hidden_params = response.get("_hidden_params") or {}
|
||||
else:
|
||||
hidden_params = getattr(response, "_hidden_params", None) or {}
|
||||
if not isinstance(hidden_params, dict):
|
||||
hidden_params = {}
|
||||
|
||||
model_id = ProxyBaseLLMRequestProcessing._get_model_id_from_response(
|
||||
hidden_params, request_data
|
||||
)
|
||||
|
||||
cache_key = hidden_params.get("cache_key", None) or ""
|
||||
api_base = hidden_params.get("api_base", None) or ""
|
||||
response_cost = hidden_params.get("response_cost", None) or ""
|
||||
fastest_response_batch_completion = hidden_params.get(
|
||||
"fastest_response_batch_completion", None
|
||||
)
|
||||
additional_headers = hidden_params.get("additional_headers", {}) or {}
|
||||
|
||||
custom_headers = ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_id=logging_obj.litellm_call_id,
|
||||
model_id=model_id,
|
||||
cache_key=cache_key,
|
||||
api_base=api_base,
|
||||
version=version,
|
||||
response_cost=response_cost,
|
||||
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
|
||||
fastest_response_batch_completion=fastest_response_batch_completion,
|
||||
request_data=request_data,
|
||||
hidden_params=hidden_params,
|
||||
litellm_logging_obj=logging_obj,
|
||||
**additional_headers,
|
||||
)
|
||||
|
||||
callback_headers = await proxy_logging_obj.post_call_response_headers_hook(
|
||||
data=request_data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
response=response,
|
||||
request_headers=dict(request.headers),
|
||||
)
|
||||
if callback_headers:
|
||||
custom_headers.update(callback_headers)
|
||||
|
||||
return custom_headers
|
||||
|
||||
async def common_processing_pre_call_logic(
|
||||
self,
|
||||
request: Request,
|
||||
|
|
@ -875,7 +936,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
else:
|
||||
verbose_proxy_logger.debug(
|
||||
"Request received by LiteLLM:\n%s",
|
||||
json.dumps(self.data, indent=4, default=str),
|
||||
_payload_str,
|
||||
)
|
||||
|
||||
async def base_process_llm_request( # noqa: PLR0915
|
||||
|
|
@ -1511,9 +1572,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
_response = assembled_response
|
||||
try:
|
||||
from litellm.proxy.proxy_server import llm_router as _global_llm_router
|
||||
from litellm.proxy.utils import (
|
||||
_check_and_merge_model_level_guardrails,
|
||||
)
|
||||
from litellm.proxy.utils import _check_and_merge_model_level_guardrails
|
||||
|
||||
guardrail_data = _check_and_merge_model_level_guardrails(
|
||||
data=captured_data, llm_router=_global_llm_router
|
||||
|
|
@ -1690,11 +1749,12 @@ class ProxyBaseLLMRequestProcessing:
|
|||
elif isinstance(e, httpx.HTTPStatusError):
|
||||
# Handle httpx.HTTPStatusError - extract actual error from response
|
||||
# This matches the original behavior before the refactor in commit 511d435f6f
|
||||
error_body = await e.response.aread()
|
||||
http_status_error: httpx.HTTPStatusError = e
|
||||
error_body = await http_status_error.response.aread()
|
||||
error_text = error_body.decode("utf-8")
|
||||
|
||||
raise HTTPException(
|
||||
status_code=e.response.status_code,
|
||||
status_code=http_status_error.response.status_code,
|
||||
detail={"error": error_text},
|
||||
)
|
||||
error_msg = f"{str(e)}"
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from typing import Union
|
||||
from typing import Any, Awaitable, Callable, Optional, Union
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import (
|
||||
DB_CONNECTION_ERROR_TYPES,
|
||||
ProxyErrorTypes,
|
||||
|
|
@ -123,3 +124,138 @@ class PrismaDBExceptionHandler:
|
|||
):
|
||||
return None
|
||||
raise e
|
||||
|
||||
|
||||
# Default fallback timeouts when neither the caller nor the prisma_client
|
||||
# expose `_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`.
|
||||
# Match the auth path's existing defaults so behavior is uniform across read paths.
|
||||
_DEFAULT_RECONNECT_TIMEOUT_SECONDS = 2.0
|
||||
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS = 0.1
|
||||
|
||||
|
||||
def _coerce_timeout(value: Any, fallback: float) -> float:
|
||||
"""Return `value` if it is a real int/float, else `fallback`. Guards
|
||||
against tests that mock `prisma_client` and leave the timeout slots as
|
||||
MagicMock instances."""
|
||||
if isinstance(value, (int, float)) and not isinstance(value, bool):
|
||||
return float(value)
|
||||
return fallback
|
||||
|
||||
|
||||
async def call_with_db_reconnect_retry(
|
||||
prisma_client: Any,
|
||||
coro_factory: Callable[[], Awaitable[Any]],
|
||||
*,
|
||||
reason: str,
|
||||
timeout_seconds: Optional[float] = None,
|
||||
lock_timeout_seconds: Optional[float] = None,
|
||||
) -> Any:
|
||||
"""Run a Prisma read coroutine with one transport-reconnect-and-retry.
|
||||
|
||||
The canonical "self-heal a transient DB transport blip" wrapper used by
|
||||
`PrismaClient.get_generic_data` and other read paths. Mirrors the inline
|
||||
pattern in `auth_checks._fetch_key_object_from_db_with_reconnect` so we
|
||||
have a single implementation rather than three drifting copies.
|
||||
|
||||
Behavior:
|
||||
1. Await `coro_factory()`. On success, return its value.
|
||||
2. On exception, if it is NOT a transport error (per
|
||||
`is_database_transport_error`), re-raise — data-layer errors like
|
||||
`UniqueViolationError` mean the DB is reachable, reconnect would be
|
||||
pointless.
|
||||
3. If `prisma_client` does not expose `attempt_db_reconnect`, re-raise.
|
||||
This guards against partial stand-ins / older clients in tests.
|
||||
4. Call `prisma_client.attempt_db_reconnect(reason=...)`. If it returns
|
||||
False (cooldown / lock contention / reconnect failure), re-raise.
|
||||
5. Otherwise await `coro_factory()` a second time and return / propagate
|
||||
its result. At-most-one retry by construction — no infinite loop.
|
||||
|
||||
`coro_factory` MUST be a zero-arg callable that returns a fresh awaitable
|
||||
on each call. Passing an already-awaited coroutine would fail on retry
|
||||
with `RuntimeError: cannot reuse already awaited coroutine`.
|
||||
|
||||
`reason` should follow `<subsystem>_<operation>_<table>_failure` so
|
||||
telemetry distinguishes between fan-out callers (e.g.
|
||||
`_update_config_from_db` issues four concurrent reads).
|
||||
|
||||
Args:
|
||||
prisma_client: The `PrismaClient` (or stand-in) that owns
|
||||
`attempt_db_reconnect` and the `_db_auth_reconnect_*` defaults.
|
||||
coro_factory: Zero-arg callable returning the read awaitable.
|
||||
reason: Telemetry tag forwarded to `attempt_db_reconnect`.
|
||||
timeout_seconds: Optional override for the reconnect cycle timeout.
|
||||
Defaults to `prisma_client._db_auth_reconnect_timeout_seconds`,
|
||||
then to 2.0s.
|
||||
lock_timeout_seconds: Optional override for how long the helper will
|
||||
wait to acquire the reconnect lock. Defaults to
|
||||
`prisma_client._db_auth_reconnect_lock_timeout_seconds`, then to
|
||||
0.1s.
|
||||
|
||||
Returns:
|
||||
Whatever `coro_factory()` returns (on first or second attempt).
|
||||
|
||||
Raises:
|
||||
Whatever `coro_factory()` raises if the failure is not a transport
|
||||
error, or if the reconnect attempt does not succeed, or if the retry
|
||||
also fails.
|
||||
"""
|
||||
try:
|
||||
return await coro_factory()
|
||||
except Exception as first_exc:
|
||||
if not PrismaDBExceptionHandler.is_database_transport_error(first_exc):
|
||||
raise
|
||||
if not hasattr(prisma_client, "attempt_db_reconnect"):
|
||||
raise
|
||||
|
||||
resolved_timeout = _coerce_timeout(
|
||||
(
|
||||
timeout_seconds
|
||||
if timeout_seconds is not None
|
||||
else getattr(prisma_client, "_db_auth_reconnect_timeout_seconds", None)
|
||||
),
|
||||
_DEFAULT_RECONNECT_TIMEOUT_SECONDS,
|
||||
)
|
||||
resolved_lock_timeout = _coerce_timeout(
|
||||
(
|
||||
lock_timeout_seconds
|
||||
if lock_timeout_seconds is not None
|
||||
else getattr(
|
||||
prisma_client, "_db_auth_reconnect_lock_timeout_seconds", None
|
||||
)
|
||||
),
|
||||
_DEFAULT_RECONNECT_LOCK_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.warning(
|
||||
"DB transport error on read; attempting reconnect-and-retry. reason=%s error=%s",
|
||||
reason,
|
||||
first_exc,
|
||||
)
|
||||
|
||||
# Preserve the original transport error in telemetry. If
|
||||
# `attempt_db_reconnect` itself raises (e.g. lock cancellation, timer
|
||||
# error, unexpected internal failure), surfacing that exception
|
||||
# instead of `first_exc` would mask the actual DB transport problem
|
||||
# in `failure_handler` / `db_exceptions` alerts. Chain the reconnect
|
||||
# error as the cause for debuggability without losing the original.
|
||||
try:
|
||||
did_reconnect = await prisma_client.attempt_db_reconnect(
|
||||
reason=reason,
|
||||
timeout_seconds=resolved_timeout,
|
||||
lock_timeout_seconds=resolved_lock_timeout,
|
||||
)
|
||||
except Exception as reconnect_exc:
|
||||
verbose_proxy_logger.warning(
|
||||
"DB reconnect attempt raised; preserving original transport error. "
|
||||
"reason=%s reconnect_error=%s",
|
||||
reason,
|
||||
reconnect_exc,
|
||||
)
|
||||
raise first_exc from reconnect_exc
|
||||
if not did_reconnect:
|
||||
raise
|
||||
|
||||
# At most one retry. If the retry also raises a transport error, we
|
||||
# propagate — repeated reconnect-loops are the watchdog's job, not
|
||||
# this helper's.
|
||||
return await coro_factory()
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ async def google_generate_content(
|
|||
general_settings,
|
||||
llm_router,
|
||||
proxy_config,
|
||||
proxy_logging_obj,
|
||||
version,
|
||||
)
|
||||
|
||||
|
|
@ -73,6 +74,16 @@ async def google_generate_content(
|
|||
if llm_router is None:
|
||||
raise HTTPException(status_code=500, detail="Router not initialized")
|
||||
response = await llm_router.agenerate_content(**data)
|
||||
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=response,
|
||||
request_data=data,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
logging_obj=logging_obj,
|
||||
version=version,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
fastapi_response.headers.update(success_headers)
|
||||
return response
|
||||
|
||||
|
||||
|
|
@ -95,6 +106,7 @@ async def google_stream_generate_content(
|
|||
general_settings,
|
||||
llm_router,
|
||||
proxy_config,
|
||||
proxy_logging_obj,
|
||||
version,
|
||||
)
|
||||
|
||||
|
|
@ -137,9 +149,24 @@ async def google_stream_generate_content(
|
|||
raise HTTPException(status_code=500, detail="Router not initialized")
|
||||
response = await llm_router.agenerate_content_stream(**data)
|
||||
|
||||
success_headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=response,
|
||||
request_data=data,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
logging_obj=logging_obj,
|
||||
version=version,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# Check if response is an async iterator (streaming response)
|
||||
if response is not None and hasattr(response, "__aiter__"):
|
||||
return StreamingResponse(content=response, media_type="text/event-stream")
|
||||
return StreamingResponse(
|
||||
content=response,
|
||||
media_type="text/event-stream",
|
||||
headers=success_headers,
|
||||
)
|
||||
fastapi_response.headers.update(success_headers)
|
||||
return response
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -742,6 +742,10 @@ async def _initialize_shared_aiohttp_session():
|
|||
try:
|
||||
from aiohttp import ClientSession, TCPConnector
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
_build_aiohttp_keepalive_socket_factory,
|
||||
)
|
||||
|
||||
connector_kwargs: Dict[str, Any] = {
|
||||
"keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT,
|
||||
"ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE,
|
||||
|
|
@ -752,6 +756,9 @@ async def _initialize_shared_aiohttp_session():
|
|||
connector_kwargs["limit"] = AIOHTTP_CONNECTOR_LIMIT
|
||||
if AIOHTTP_CONNECTOR_LIMIT_PER_HOST > 0:
|
||||
connector_kwargs["limit_per_host"] = AIOHTTP_CONNECTOR_LIMIT_PER_HOST
|
||||
socket_factory = _build_aiohttp_keepalive_socket_factory()
|
||||
if socket_factory is not None:
|
||||
connector_kwargs["socket_factory"] = socket_factory
|
||||
|
||||
connector = TCPConnector(**connector_kwargs)
|
||||
session = ClientSession(connector=connector)
|
||||
|
|
|
|||
|
|
@ -106,7 +106,10 @@ from litellm.proxy.db.create_views import (
|
|||
should_create_missing_views,
|
||||
)
|
||||
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
|
||||
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
|
||||
from litellm.proxy.db.exception_handler import (
|
||||
PrismaDBExceptionHandler,
|
||||
call_with_db_reconnect_retry,
|
||||
)
|
||||
from litellm.proxy.db.log_db_metrics import log_db_metrics
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
|
||||
|
|
@ -2779,30 +2782,42 @@ class PrismaClient:
|
|||
table_name: Literal["users", "keys", "config", "spend"],
|
||||
):
|
||||
"""
|
||||
Generic implementation of get data
|
||||
Generic implementation of get data.
|
||||
|
||||
Self-heals across a single transient transport blip via
|
||||
`call_with_db_reconnect_retry`: on `httpx.ReadError` /
|
||||
`ClientNotConnectedError` / similar, attempt one DB reconnect and
|
||||
retry once before surfacing the failure. Restores the 1.82.6 behavior
|
||||
that was lost in 1.83.x — see issue #25143.
|
||||
"""
|
||||
start_time = time.time()
|
||||
try:
|
||||
|
||||
async def _do_query():
|
||||
if table_name == "users":
|
||||
response = await self.db.litellm_usertable.find_first(
|
||||
return await self.db.litellm_usertable.find_first(
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
elif table_name == "keys":
|
||||
response = await self.db.litellm_verificationtoken.find_first( # type: ignore
|
||||
return await self.db.litellm_verificationtoken.find_first( # type: ignore
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
elif table_name == "config":
|
||||
response = await self.db.litellm_config.find_first( # type: ignore
|
||||
return await self.db.litellm_config.find_first( # type: ignore
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
elif table_name == "spend":
|
||||
response = await self.db.l.find_first( # type: ignore
|
||||
return await self.db.l.find_first( # type: ignore
|
||||
where={key: value} # type: ignore
|
||||
)
|
||||
return response
|
||||
except Exception as e:
|
||||
import traceback
|
||||
return None
|
||||
|
||||
try:
|
||||
return await call_with_db_reconnect_retry(
|
||||
self,
|
||||
_do_query,
|
||||
reason=f"prisma_get_generic_data_{table_name}_lookup_failure",
|
||||
)
|
||||
except Exception as e:
|
||||
error_msg = f"LiteLLM Prisma Client Exception get_generic_data: {str(e)}"
|
||||
verbose_proxy_logger.error(error_msg)
|
||||
error_msg = error_msg + "\nException Type: {}".format(type(e))
|
||||
|
|
@ -4204,7 +4219,6 @@ class PrismaClient:
|
|||
)
|
||||
self._reap_all_zombies()
|
||||
self._cleanup_engine_watcher()
|
||||
self._engine_confirmed_dead = False
|
||||
|
||||
async def _do_heavy_reconnect() -> None:
|
||||
db_url = os.getenv("DATABASE_URL", "")
|
||||
|
|
@ -4217,6 +4231,12 @@ class PrismaClient:
|
|||
await self._start_engine_watcher()
|
||||
|
||||
await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout)
|
||||
# Only clear the "dead engine" flag after the heavy reconnect
|
||||
# actually completed. If `_do_heavy_reconnect()` raises (timeout,
|
||||
# missing DATABASE_URL, recreate failure), the flag stays True so
|
||||
# the next attempt re-enters the heavy branch instead of silently
|
||||
# demoting to the lightweight path.
|
||||
self._engine_confirmed_dead = False
|
||||
else:
|
||||
verbose_proxy_logger.debug(
|
||||
"Performing Prisma DB reconnect (engine alive or unknown)."
|
||||
|
|
|
|||
|
|
@ -81,6 +81,12 @@ class MCPServer(BaseModel):
|
|||
# Defaults to the token's expires_in minus the expiry buffer, or
|
||||
# MCP_PER_USER_TOKEN_DEFAULT_TTL when expires_in is absent.
|
||||
token_storage_ttl_seconds: Optional[int] = None
|
||||
# Resolved short-ID tool prefix when LITELLM_USE_SHORT_MCP_TOOL_PREFIX is
|
||||
# enabled. Set by ``MCPServerManager._assign_unique_short_prefix`` at
|
||||
# registration time so that natural-hash collisions between two
|
||||
# different ``server_id`` values are bumped deterministically. Left
|
||||
# ``None`` in default-prefix mode.
|
||||
short_prefix: Optional[str] = None
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True)
|
||||
|
||||
@property
|
||||
|
|
|
|||
|
|
@ -354,6 +354,7 @@ async def test_bedrock_kb_request_body_has_transformed_filters(
|
|||
api_base=api_base,
|
||||
litellm_logging_obj=logging_obj,
|
||||
litellm_params=litellm_params_dict,
|
||||
extra_body=None,
|
||||
)
|
||||
)
|
||||
captured_request_body["url"] = url
|
||||
|
|
|
|||
|
|
@ -4364,3 +4364,39 @@ def test_bedrock_tool_message_image_url_png_still_becomes_image():
|
|||
assert "document" not in block
|
||||
assert block["image"]["format"] == "png"
|
||||
assert block["image"]["source"]["bytes"] == png_b64
|
||||
|
||||
|
||||
def test_transform_response_does_not_leak_body_on_parse_failure():
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
|
||||
leaky_body = {"output": {"message": {"content": [{"text": "secret content"}]}}}
|
||||
|
||||
class MockResponse:
|
||||
def json(self):
|
||||
return leaky_body
|
||||
|
||||
@property
|
||||
def text(self):
|
||||
return json.dumps(leaky_body)
|
||||
|
||||
with patch(
|
||||
"litellm.llms.bedrock.chat.converse_transformation.ConverseResponseBlock",
|
||||
side_effect=KeyError("missing required field"),
|
||||
):
|
||||
with pytest.raises(BedrockError) as exc_info:
|
||||
AmazonConverseConfig()._transform_response(
|
||||
model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
response=MockResponse(),
|
||||
model_response=ModelResponse(),
|
||||
stream=False,
|
||||
logging_obj=None,
|
||||
optional_params={},
|
||||
api_key=None,
|
||||
data=None,
|
||||
messages=[],
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
msg = str(exc_info.value)
|
||||
assert "secret content" not in msg
|
||||
assert "Error converting to valid response block" in msg
|
||||
|
|
|
|||
|
|
@ -21,7 +21,112 @@ def test_transform_search_request():
|
|||
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
|
||||
litellm_logging_obj=mock_log,
|
||||
litellm_params={},
|
||||
extra_body=None,
|
||||
)
|
||||
|
||||
assert url.endswith("/kb123/retrieve")
|
||||
assert body["retrievalQuery"].get("text") == "hello"
|
||||
|
||||
|
||||
def test_transform_search_request_uses_only_retrieval_config_from_extra_body():
|
||||
config = BedrockVectorStoreConfig()
|
||||
mock_log = MagicMock()
|
||||
mock_log.model_call_details = {}
|
||||
|
||||
url, body = config.transform_search_vector_store_request(
|
||||
vector_store_id="kb123",
|
||||
query="hello",
|
||||
vector_store_search_optional_params={},
|
||||
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
|
||||
litellm_logging_obj=mock_log,
|
||||
litellm_params={},
|
||||
extra_body={
|
||||
"retrievalConfiguration": {
|
||||
"vectorSearchConfiguration": {
|
||||
"overrideSearchType": "HYBRID",
|
||||
"numberOfResults": 8,
|
||||
}
|
||||
},
|
||||
"unrelatedField": {"should_not": "be_forwarded"},
|
||||
},
|
||||
)
|
||||
|
||||
assert url.endswith("/kb123/retrieve")
|
||||
assert body["retrievalQuery"].get("text") == "hello"
|
||||
assert (
|
||||
body["retrievalConfiguration"]["vectorSearchConfiguration"][
|
||||
"overrideSearchType"
|
||||
]
|
||||
== "HYBRID"
|
||||
)
|
||||
assert "unrelatedField" not in body
|
||||
|
||||
|
||||
def test_transform_search_request_does_not_mutate_extra_body_and_overrides_number_of_results():
|
||||
config = BedrockVectorStoreConfig()
|
||||
mock_log = MagicMock()
|
||||
mock_log.model_call_details = {}
|
||||
extra_body = {
|
||||
"retrievalConfiguration": {
|
||||
"vectorSearchConfiguration": {
|
||||
"overrideSearchType": "HYBRID",
|
||||
"numberOfResults": 8,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
_, body = config.transform_search_vector_store_request(
|
||||
vector_store_id="kb123",
|
||||
query="hello",
|
||||
vector_store_search_optional_params={"max_num_results": 10},
|
||||
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
|
||||
litellm_logging_obj=mock_log,
|
||||
litellm_params={},
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
assert (
|
||||
body["retrievalConfiguration"]["vectorSearchConfiguration"]["numberOfResults"]
|
||||
== 10
|
||||
)
|
||||
assert (
|
||||
extra_body["retrievalConfiguration"]["vectorSearchConfiguration"][
|
||||
"numberOfResults"
|
||||
]
|
||||
== 8
|
||||
)
|
||||
|
||||
|
||||
def test_transform_search_request_overrides_filter_without_mutating_extra_body():
|
||||
config = BedrockVectorStoreConfig()
|
||||
mock_log = MagicMock()
|
||||
mock_log.model_call_details = {}
|
||||
extra_body = {
|
||||
"retrievalConfiguration": {
|
||||
"vectorSearchConfiguration": {
|
||||
"filter": {"equals": {"key": "tenant", "value": "a"}}
|
||||
}
|
||||
}
|
||||
}
|
||||
new_filter = {"equals": {"key": "tenant", "value": "b"}}
|
||||
|
||||
_, body = config.transform_search_vector_store_request(
|
||||
vector_store_id="kb123",
|
||||
query="hello",
|
||||
vector_store_search_optional_params={"filters": new_filter},
|
||||
api_base="https://bedrock-agent-runtime.us-west-2.amazonaws.com/knowledgebases",
|
||||
litellm_logging_obj=mock_log,
|
||||
litellm_params={},
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
assert (
|
||||
body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"]
|
||||
== new_filter
|
||||
)
|
||||
assert (
|
||||
extra_body["retrievalConfiguration"]["vectorSearchConfiguration"]["filter"][
|
||||
"equals"
|
||||
]["value"]
|
||||
== "a"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,161 @@
|
|||
import socket
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
|
||||
def _invoke_connector_factory(http_handler_module):
|
||||
"""
|
||||
Drive the lambda factory installed on the transport so TCPConnector is
|
||||
actually constructed. _create_aiohttp_transport returns a transport whose
|
||||
_client_factory is the lambda that builds (TCPConnector → ClientSession);
|
||||
invoking it directly avoids relying on _get_valid_client_session's internal
|
||||
branching to trigger connector construction.
|
||||
"""
|
||||
transport = http_handler_module.AsyncHTTPHandler._create_aiohttp_transport(
|
||||
shared_session=None
|
||||
)
|
||||
transport._client_factory()
|
||||
return transport
|
||||
|
||||
|
||||
def test_socket_factory_omitted_when_disabled(monkeypatch):
|
||||
from litellm.llms.custom_httpx import http_handler as http_handler_module
|
||||
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", False)
|
||||
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
|
||||
|
||||
connector_mock = MagicMock(name="connector")
|
||||
session_mock = MagicMock(name="session")
|
||||
|
||||
with patch.object(
|
||||
http_handler_module, "TCPConnector", return_value=connector_mock
|
||||
) as mock_tcp_connector:
|
||||
with patch.object(
|
||||
http_handler_module, "ClientSession", return_value=session_mock
|
||||
):
|
||||
_invoke_connector_factory(http_handler_module)
|
||||
|
||||
assert mock_tcp_connector.call_count >= 1
|
||||
assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
|
||||
|
||||
|
||||
def test_socket_factory_attached_when_enabled(monkeypatch):
|
||||
from litellm.llms.custom_httpx import http_handler as http_handler_module
|
||||
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
|
||||
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
|
||||
|
||||
connector_mock = MagicMock(name="connector")
|
||||
session_mock = MagicMock(name="session")
|
||||
|
||||
with patch.object(
|
||||
http_handler_module, "TCPConnector", return_value=connector_mock
|
||||
) as mock_tcp_connector:
|
||||
with patch.object(
|
||||
http_handler_module, "ClientSession", return_value=session_mock
|
||||
):
|
||||
_invoke_connector_factory(http_handler_module)
|
||||
|
||||
assert mock_tcp_connector.call_count >= 1
|
||||
factory = mock_tcp_connector.call_args.kwargs.get("socket_factory")
|
||||
assert callable(factory)
|
||||
|
||||
|
||||
def test_socket_factory_skipped_on_old_aiohttp(monkeypatch):
|
||||
from litellm.llms.custom_httpx import http_handler as http_handler_module
|
||||
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
|
||||
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", False)
|
||||
|
||||
connector_mock = MagicMock(name="connector")
|
||||
session_mock = MagicMock(name="session")
|
||||
|
||||
with patch.object(
|
||||
http_handler_module, "TCPConnector", return_value=connector_mock
|
||||
) as mock_tcp_connector:
|
||||
with patch.object(
|
||||
http_handler_module, "ClientSession", return_value=session_mock
|
||||
):
|
||||
_invoke_connector_factory(http_handler_module)
|
||||
|
||||
assert mock_tcp_connector.call_count >= 1
|
||||
assert "socket_factory" not in mock_tcp_connector.call_args.kwargs
|
||||
|
||||
|
||||
def test_socket_factory_sets_keepalive_options(monkeypatch):
|
||||
from litellm.llms.custom_httpx import http_handler as http_handler_module
|
||||
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
|
||||
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 45)
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPINTVL", 15)
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPCNT", 4)
|
||||
|
||||
factory = http_handler_module._build_aiohttp_keepalive_socket_factory()
|
||||
assert factory is not None
|
||||
|
||||
addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0))
|
||||
|
||||
fake_sock = MagicMock(spec=socket.socket)
|
||||
with patch("socket.socket", return_value=fake_sock) as sock_ctor:
|
||||
returned = factory(addr_info)
|
||||
|
||||
sock_ctor.assert_called_once_with(
|
||||
family=socket.AF_INET, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP
|
||||
)
|
||||
assert returned is fake_sock
|
||||
fake_sock.setblocking.assert_called_once_with(False)
|
||||
|
||||
setsockopt_calls = {
|
||||
(call.args[0], call.args[1]): call.args[2]
|
||||
for call in fake_sock.setsockopt.call_args_list
|
||||
}
|
||||
assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1
|
||||
|
||||
if hasattr(socket, "TCP_KEEPIDLE"):
|
||||
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE)] == 45
|
||||
elif hasattr(socket, "TCP_KEEPALIVE"):
|
||||
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPALIVE)] == 45
|
||||
if hasattr(socket, "TCP_KEEPINTVL"):
|
||||
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL)] == 15
|
||||
if hasattr(socket, "TCP_KEEPCNT"):
|
||||
assert setsockopt_calls[(socket.IPPROTO_TCP, socket.TCP_KEEPCNT)] == 4
|
||||
|
||||
|
||||
def test_socket_factory_uses_tcp_keepalive_when_keepidle_unavailable(monkeypatch):
|
||||
"""
|
||||
Cover the macOS/Darwin branch: when TCP_KEEPIDLE is missing but TCP_KEEPALIVE
|
||||
is present, the factory should fall back to TCP_KEEPALIVE for the idle timer.
|
||||
Linux CI runners always have TCP_KEEPIDLE, so we patch socket itself to
|
||||
simulate the BSD-derived environment.
|
||||
"""
|
||||
from litellm.llms.custom_httpx import http_handler as http_handler_module
|
||||
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_SO_KEEPALIVE", True)
|
||||
monkeypatch.setattr(http_handler_module, "_AIOHTTP_SUPPORTS_SOCKET_FACTORY", True)
|
||||
monkeypatch.setattr(http_handler_module, "AIOHTTP_TCP_KEEPIDLE", 60)
|
||||
|
||||
factory = http_handler_module._build_aiohttp_keepalive_socket_factory()
|
||||
assert factory is not None
|
||||
|
||||
fake_socket_module = MagicMock(spec=[])
|
||||
fake_socket_module.SOL_SOCKET = socket.SOL_SOCKET
|
||||
fake_socket_module.SO_KEEPALIVE = socket.SO_KEEPALIVE
|
||||
fake_socket_module.IPPROTO_TCP = socket.IPPROTO_TCP
|
||||
fake_socket_module.TCP_KEEPALIVE = getattr(socket, "TCP_KEEPALIVE", 0x10)
|
||||
fake_sock = MagicMock(spec=socket.socket)
|
||||
fake_socket_module.socket = MagicMock(return_value=fake_sock)
|
||||
|
||||
addr_info = (socket.AF_INET, socket.SOCK_STREAM, socket.IPPROTO_TCP, "", ("", 0))
|
||||
|
||||
with patch.object(http_handler_module, "socket", fake_socket_module):
|
||||
factory(addr_info)
|
||||
|
||||
setsockopt_calls = {
|
||||
(call.args[0], call.args[1]): call.args[2]
|
||||
for call in fake_sock.setsockopt.call_args_list
|
||||
}
|
||||
assert setsockopt_calls[(socket.SOL_SOCKET, socket.SO_KEEPALIVE)] == 1
|
||||
assert (
|
||||
setsockopt_calls[(socket.IPPROTO_TCP, fake_socket_module.TCP_KEEPALIVE)] == 60
|
||||
)
|
||||
assert (socket.IPPROTO_TCP, getattr(socket, "TCP_KEEPIDLE", -1)) not in setsockopt_calls
|
||||
|
|
@ -2,12 +2,16 @@ import os
|
|||
import sys
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import (
|
||||
BaseLLMHTTPHandler,
|
||||
_google_genai_streaming_hidden_params,
|
||||
)
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
|
||||
|
||||
|
|
@ -320,3 +324,29 @@ async def test_async_anthropic_messages_handler_header_priority():
|
|||
assert captured_headers["X-Forwarded-Only"] == "keep"
|
||||
assert captured_headers["X-Extra-Only"] == "also-keep"
|
||||
assert captured_headers["X-Provider-Only"] == "keep-this-too"
|
||||
|
||||
|
||||
def test_google_genai_streaming_hidden_params_model_info_and_router_fallback():
|
||||
logging_obj = Mock()
|
||||
logging_obj.get_router_model_id = Mock(return_value="router-model-id")
|
||||
|
||||
from_model_info = _google_genai_streaming_hidden_params(
|
||||
api_base="https://generativelanguage.googleapis.com/v1beta",
|
||||
litellm_params=GenericLiteLLMParams(model_info={"id": "info-id"}),
|
||||
logging_obj=logging_obj,
|
||||
response_headers=httpx.Headers({"x-ratelimit-remaining": "10"}),
|
||||
)
|
||||
assert from_model_info["model_id"] == "info-id"
|
||||
assert (
|
||||
from_model_info["api_base"]
|
||||
== "https://generativelanguage.googleapis.com/v1beta"
|
||||
)
|
||||
assert isinstance(from_model_info["additional_headers"], dict)
|
||||
|
||||
from_router = _google_genai_streaming_hidden_params(
|
||||
api_base="https://x",
|
||||
litellm_params=GenericLiteLLMParams(),
|
||||
logging_obj=logging_obj,
|
||||
response_headers=httpx.Headers({}),
|
||||
)
|
||||
assert from_router["model_id"] == "router-model-id"
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ class TestS3VectorsVectorStoreConfig:
|
|||
api_base="https://s3vectors.us-west-2.api.aws",
|
||||
litellm_logging_obj=mock_logging_obj,
|
||||
litellm_params={},
|
||||
extra_body=None,
|
||||
)
|
||||
|
||||
def test_transform_search_response(self):
|
||||
|
|
|
|||
|
|
@ -4261,3 +4261,32 @@ def test_sync_streaming_uses_custom_client():
|
|||
# Verify that gemini_client is in the partial's keywords
|
||||
assert "gemini_client" in partial_make_sync_call.keywords
|
||||
assert partial_make_sync_call.keywords["gemini_client"] is mock_client
|
||||
|
||||
|
||||
def test_transform_response_does_not_leak_body_on_parse_failure():
|
||||
leaky_body = {"candidates": [{"content": {"parts": [{"text": "secret content"}]}}]}
|
||||
raw_response = MagicMock()
|
||||
raw_response.json.return_value = leaky_body
|
||||
raw_response.text = json.dumps(leaky_body)
|
||||
raw_response.headers = {}
|
||||
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.GenerateContentResponseBody",
|
||||
side_effect=KeyError("missing required field"),
|
||||
):
|
||||
with pytest.raises(VertexAIError) as exc_info:
|
||||
VertexGeminiConfig().transform_response(
|
||||
model="gemini-pro",
|
||||
raw_response=raw_response,
|
||||
model_response=ModelResponse(),
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
msg = str(exc_info.value)
|
||||
assert "secret content" not in msg
|
||||
assert "Error converting to valid response block" in msg
|
||||
|
|
|
|||
|
|
@ -728,6 +728,136 @@ class TestMCPServerManager:
|
|||
]
|
||||
assert scopes == ["read", "write"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_descovery_metadata_probes_well_known_when_server_does_not_challenge(
|
||||
self,
|
||||
):
|
||||
manager = MCPServerManager()
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
|
||||
mock_metadata = MCPOAuthMetadata(
|
||||
scopes=None,
|
||||
authorization_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/authorize",
|
||||
token_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/token",
|
||||
registration_url=None,
|
||||
)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
),
|
||||
patch.object(
|
||||
manager,
|
||||
"_attempt_well_known_discovery",
|
||||
AsyncMock(
|
||||
return_value=(
|
||||
["https://login.microsoftonline.com/test-tenant-id/v2.0"],
|
||||
["api://some-scope/.default"],
|
||||
)
|
||||
),
|
||||
) as mock_well_known,
|
||||
patch.object(
|
||||
manager,
|
||||
"_fetch_authorization_server_metadata",
|
||||
AsyncMock(return_value=mock_metadata),
|
||||
) as mock_fetch_auth,
|
||||
):
|
||||
result = await manager._descovery_metadata("http://localhost:8001/mcp")
|
||||
|
||||
mock_well_known.assert_awaited_once_with("http://localhost:8001/mcp")
|
||||
mock_fetch_auth.assert_awaited_once_with(
|
||||
["https://login.microsoftonline.com/test-tenant-id/v2.0"]
|
||||
)
|
||||
assert result is mock_metadata
|
||||
assert result.scopes == ["api://some-scope/.default"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_single_authorization_server_metadata_supports_azure_issuer_path(
|
||||
self,
|
||||
):
|
||||
manager = MCPServerManager()
|
||||
issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0"
|
||||
|
||||
def build_response(url: str):
|
||||
mock_response = MagicMock()
|
||||
if url == f"{issuer}/.well-known/openid-configuration":
|
||||
mock_response.json.return_value = {
|
||||
"authorization_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize",
|
||||
"token_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token",
|
||||
"scopes_supported": ["api://some-scope/.default"],
|
||||
}
|
||||
mock_response.raise_for_status = MagicMock()
|
||||
else:
|
||||
request = httpx.Request("GET", url)
|
||||
response_obj = httpx.Response(status_code=404, request=request)
|
||||
mock_response.raise_for_status = MagicMock(
|
||||
side_effect=httpx.HTTPStatusError(
|
||||
"not found", request=request, response=response_obj
|
||||
)
|
||||
)
|
||||
return mock_response
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.get = AsyncMock(side_effect=build_response)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
):
|
||||
result = await manager._fetch_single_authorization_server_metadata(issuer)
|
||||
|
||||
assert result is not None
|
||||
assert (
|
||||
result.authorization_url
|
||||
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize"
|
||||
)
|
||||
assert (
|
||||
result.token_url
|
||||
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token"
|
||||
)
|
||||
assert result.scopes == ["api://some-scope/.default"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_single_authorization_server_metadata_derives_azure_metadata(
|
||||
self,
|
||||
):
|
||||
manager = MCPServerManager()
|
||||
issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0"
|
||||
|
||||
request = httpx.Request("GET", issuer)
|
||||
response_obj = httpx.Response(status_code=404, request=request)
|
||||
mock_response = MagicMock()
|
||||
mock_response.raise_for_status = MagicMock(
|
||||
side_effect=httpx.HTTPStatusError(
|
||||
"not found", request=request, response=response_obj
|
||||
)
|
||||
)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.get = AsyncMock(return_value=mock_response)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
|
||||
return_value=mock_client,
|
||||
):
|
||||
result = await manager._fetch_single_authorization_server_metadata(issuer)
|
||||
|
||||
assert result is not None
|
||||
assert (
|
||||
result.authorization_url
|
||||
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize"
|
||||
)
|
||||
assert (
|
||||
result.token_url
|
||||
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_descovery_metadata_falls_back_to_origin_when_no_auth_servers(self):
|
||||
manager = MCPServerManager()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,302 @@
|
|||
"""
|
||||
Tests for the short-ID MCP tool prefix (LITELLM_USE_SHORT_MCP_TOOL_PREFIX).
|
||||
|
||||
The short-prefix mode swaps the historical alias/server_name prefix on
|
||||
tool names for a deterministic three-character base62 ID derived from the
|
||||
server's ``server_id``. This keeps tool names well below the 60-char
|
||||
upper bound enforced by some model APIs while remaining stable across
|
||||
processes/restarts and tolerant of mixed-version clients.
|
||||
"""
|
||||
|
||||
from typing import List
|
||||
|
||||
import pytest
|
||||
from mcp.types import Tool as MCPTool
|
||||
|
||||
from litellm.proxy._experimental.mcp_server.mcp_server_manager import MCPServerManager
|
||||
from litellm.proxy._experimental.mcp_server.utils import (
|
||||
SHORT_MCP_TOOL_PREFIX_LENGTH,
|
||||
add_server_prefix_to_name,
|
||||
compute_short_server_prefix,
|
||||
get_server_prefix,
|
||||
is_short_mcp_tool_prefix_enabled,
|
||||
iter_known_server_prefixes,
|
||||
)
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
|
||||
def _make_server(
|
||||
*,
|
||||
server_id: str = "abcdef-1234",
|
||||
server_name: str = "github_onprem",
|
||||
alias: str = "github_onprem",
|
||||
) -> MCPServer:
|
||||
return MCPServer(
|
||||
server_id=server_id,
|
||||
name=alias or server_name,
|
||||
alias=alias,
|
||||
server_name=server_name,
|
||||
transport="http",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _reset_env(monkeypatch):
|
||||
monkeypatch.delenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", raising=False)
|
||||
yield
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestShortPrefixHelpers:
|
||||
def test_short_prefix_is_three_base62_chars(self):
|
||||
prefix = compute_short_server_prefix("any-server-id")
|
||||
assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH
|
||||
assert prefix.isalnum() and prefix.isascii()
|
||||
|
||||
def test_short_prefix_first_char_is_alphabetic(self):
|
||||
"""The first char must be [A-Za-z] so the prefix is a valid identifier
|
||||
on every model API (some providers historically required the first
|
||||
character of a function name to be alphabetic)."""
|
||||
# Sweep many server_ids and rehash attempts to give us coverage of
|
||||
# every position the high-order bits can land on.
|
||||
for i in range(200):
|
||||
for attempt in range(4):
|
||||
prefix = compute_short_server_prefix(f"server-{i}", attempt=attempt)
|
||||
assert prefix[0].isalpha(), (
|
||||
f"prefix {prefix!r} for server-{i} (attempt={attempt}) "
|
||||
f"starts with a non-alphabetic character"
|
||||
)
|
||||
|
||||
def test_short_prefix_is_deterministic(self):
|
||||
assert compute_short_server_prefix("abc") == compute_short_server_prefix("abc")
|
||||
assert compute_short_server_prefix("abc") != compute_short_server_prefix("abd")
|
||||
|
||||
def test_short_prefix_requires_server_id(self):
|
||||
with pytest.raises(ValueError):
|
||||
compute_short_server_prefix("")
|
||||
|
||||
def test_flag_defaults_to_false(self):
|
||||
assert is_short_mcp_tool_prefix_enabled() is False
|
||||
|
||||
@pytest.mark.parametrize("value", ["1", "true", "TRUE", "yes", "On"])
|
||||
def test_flag_truthy_values(self, monkeypatch, value):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value)
|
||||
assert is_short_mcp_tool_prefix_enabled() is True
|
||||
|
||||
@pytest.mark.parametrize("value", ["0", "false", "no", "off", ""])
|
||||
def test_flag_falsey_values(self, monkeypatch, value):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", value)
|
||||
assert is_short_mcp_tool_prefix_enabled() is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_server_prefix behaviour
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGetServerPrefix:
|
||||
def test_default_mode_uses_alias(self):
|
||||
server = _make_server(alias="github_onprem", server_name="github_onprem")
|
||||
assert get_server_prefix(server) == "github_onprem"
|
||||
|
||||
def test_short_mode_uses_short_id(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
server = _make_server(server_id="abcdef-1234")
|
||||
prefix = get_server_prefix(server)
|
||||
assert prefix == compute_short_server_prefix("abcdef-1234")
|
||||
assert len(prefix) == SHORT_MCP_TOOL_PREFIX_LENGTH
|
||||
|
||||
def test_short_mode_falls_back_when_no_server_id(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
|
||||
class _Bare:
|
||||
alias = "fallback_alias"
|
||||
server_name = None
|
||||
server_id = None
|
||||
|
||||
assert get_server_prefix(_Bare()) == "fallback_alias"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# iter_known_server_prefixes — covers reverse-lookup tolerance
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIterKnownServerPrefixes:
|
||||
def test_default_mode_includes_short_id_too(self):
|
||||
server = _make_server()
|
||||
prefixes = list(iter_known_server_prefixes(server))
|
||||
# Contains the live prefix and every known form so that mixed-mode
|
||||
# clients can be resolved.
|
||||
assert "github_onprem" in prefixes
|
||||
assert compute_short_server_prefix(server.server_id) in prefixes
|
||||
|
||||
def test_short_mode_still_yields_long_forms(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
server = _make_server()
|
||||
prefixes = list(iter_known_server_prefixes(server))
|
||||
assert "github_onprem" in prefixes
|
||||
assert compute_short_server_prefix(server.server_id) in prefixes
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Manager-level behaviour: list + reverse-lookup
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _stub_tools() -> List[MCPTool]:
|
||||
return [
|
||||
MCPTool(name="get_repo", description="", inputSchema={"type": "object"}),
|
||||
MCPTool(name="list_issues", description="", inputSchema={"type": "object"}),
|
||||
]
|
||||
|
||||
|
||||
class TestManagerShortPrefix:
|
||||
def test_list_tools_uses_short_prefix_when_flag_on(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
|
||||
out = manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
short = compute_short_server_prefix(server.server_id)
|
||||
assert {t.name for t in out} == {f"{short}-get_repo", f"{short}-list_issues"}
|
||||
|
||||
def test_call_tool_lookup_resolves_short_prefix(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
manager.registry[server.server_id] = server
|
||||
manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
short = compute_short_server_prefix(server.server_id)
|
||||
resolved = manager._get_mcp_server_from_tool_name(f"{short}-get_repo")
|
||||
assert resolved is server
|
||||
|
||||
def test_call_tool_lookup_resolves_long_prefix_in_short_mode(self, monkeypatch):
|
||||
"""Old clients that cached the long-prefix name must still route."""
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
manager.registry[server.server_id] = server
|
||||
manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
resolved = manager._get_mcp_server_from_tool_name("github_onprem-get_repo")
|
||||
assert resolved is server
|
||||
|
||||
def test_default_mode_unchanged(self):
|
||||
manager = MCPServerManager()
|
||||
server = _make_server()
|
||||
|
||||
out = manager._create_prefixed_tools(_stub_tools(), server)
|
||||
|
||||
assert {t.name for t in out} == {
|
||||
"github_onprem-get_repo",
|
||||
"github_onprem-list_issues",
|
||||
}
|
||||
assert (
|
||||
manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is None
|
||||
) # registry empty
|
||||
manager.registry[server.server_id] = server
|
||||
assert (
|
||||
manager._get_mcp_server_from_tool_name("github_onprem-get_repo") is server
|
||||
)
|
||||
|
||||
def test_total_tool_name_length_short_enough(self, monkeypatch):
|
||||
"""The short prefix keeps tool names under the 60-char limit even
|
||||
when the upstream tool name is itself reasonably long."""
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
long_server_name = "a" * 50
|
||||
server = _make_server(
|
||||
server_id="server-id-1",
|
||||
server_name=long_server_name,
|
||||
alias=long_server_name,
|
||||
)
|
||||
prefix = get_server_prefix(server)
|
||||
full = add_server_prefix_to_name("get_repo", prefix)
|
||||
assert len(full) < 60
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Collision-resolution at registration time
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestShortPrefixCollisionResolution:
|
||||
"""``_assign_unique_short_prefix`` must rehash on collision.
|
||||
|
||||
The dedup path is exercised by forcing two distinct ``server_id``
|
||||
values to both hash to the same natural prefix via a monkeypatched
|
||||
``compute_short_server_prefix``.
|
||||
"""
|
||||
|
||||
def test_no_op_when_flag_off(self):
|
||||
manager = MCPServerManager()
|
||||
server = _make_server(server_id="abc")
|
||||
manager._assign_unique_short_prefix(server)
|
||||
assert server.short_prefix is None
|
||||
|
||||
def test_assigns_natural_hash_when_no_collision(self, monkeypatch):
|
||||
from litellm.proxy._experimental.mcp_server import utils as mcp_utils
|
||||
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server(server_id="abc")
|
||||
manager._assign_unique_short_prefix(server)
|
||||
|
||||
assert server.short_prefix == mcp_utils.compute_short_server_prefix("abc")
|
||||
|
||||
def test_rehashes_when_natural_hash_collides(self, monkeypatch):
|
||||
"""Two server_ids that natural-hash to the same prefix get
|
||||
deterministic, distinct short prefixes."""
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
|
||||
# Force every attempt=0 hash to "AAA" and attempt=1 to "AAB".
|
||||
# That way the second server registered must rehash to "AAB".
|
||||
from litellm.proxy._experimental.mcp_server import utils as mcp_utils
|
||||
|
||||
def _fake_hash(server_id: str, attempt: int = 0) -> str:
|
||||
return "AAA" if attempt == 0 else f"AA{chr(ord('A') + attempt)}"
|
||||
|
||||
monkeypatch.setattr(mcp_utils, "compute_short_server_prefix", _fake_hash)
|
||||
# Also patch the symbol that the manager imported at module load.
|
||||
from litellm.proxy._experimental.mcp_server import (
|
||||
mcp_server_manager as mgr_module,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(mgr_module, "compute_short_server_prefix", _fake_hash)
|
||||
|
||||
manager = MCPServerManager()
|
||||
first = _make_server(server_id="server-1", alias="srv1")
|
||||
second = _make_server(server_id="server-2", alias="srv2")
|
||||
|
||||
# Pretend both are already in the registry so dedup sees both.
|
||||
manager.registry[first.server_id] = first
|
||||
manager._assign_unique_short_prefix(first)
|
||||
manager.registry[second.server_id] = second
|
||||
manager._assign_unique_short_prefix(second)
|
||||
|
||||
assert first.short_prefix == "AAA"
|
||||
assert second.short_prefix == "AAB"
|
||||
assert first.short_prefix != second.short_prefix
|
||||
|
||||
def test_cached_prefix_is_reused(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
manager = MCPServerManager()
|
||||
server = _make_server(server_id="abc")
|
||||
server.short_prefix = "ZZZ" # pretend a previous registration set this
|
||||
|
||||
manager._assign_unique_short_prefix(server)
|
||||
|
||||
assert server.short_prefix == "ZZZ"
|
||||
|
||||
def test_get_server_prefix_prefers_cached(self, monkeypatch):
|
||||
monkeypatch.setenv("LITELLM_USE_SHORT_MCP_TOOL_PREFIX", "true")
|
||||
server = _make_server(server_id="abc")
|
||||
server.short_prefix = "Q9q"
|
||||
|
||||
assert get_server_prefix(server) == "Q9q"
|
||||
|
|
@ -0,0 +1,255 @@
|
|||
"""
|
||||
Unit tests for `call_with_db_reconnect_retry` — the canonical "try DB read,
|
||||
on transport error reconnect once and retry once" helper.
|
||||
|
||||
Covers the regression in issue #25143 where read paths (e.g.
|
||||
`PrismaClient.get_generic_data`) lost their reconnect-and-retry-once branch in
|
||||
LiteLLM 1.83.x and started emitting `db_exceptions` alerts on transient
|
||||
`httpx.ReadError` flaps that used to self-heal in 1.82.6.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from prisma.errors import UniqueViolationError
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../.."))
|
||||
|
||||
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
|
||||
|
||||
|
||||
def _make_client(
|
||||
*,
|
||||
attempt_db_reconnect_return: bool = True,
|
||||
has_attempt_db_reconnect: bool = True,
|
||||
):
|
||||
"""Build a minimal stand-in for PrismaClient that exposes only the surface
|
||||
`call_with_db_reconnect_retry` actually pokes at."""
|
||||
client = MagicMock()
|
||||
if has_attempt_db_reconnect:
|
||||
client.attempt_db_reconnect = AsyncMock(
|
||||
return_value=attempt_db_reconnect_return
|
||||
)
|
||||
else:
|
||||
# `hasattr(client, "attempt_db_reconnect")` must return False — MagicMock
|
||||
# auto-creates attributes, so we wipe it out via `spec`.
|
||||
client = MagicMock(spec=[])
|
||||
client._db_auth_reconnect_timeout_seconds = 2.0
|
||||
client._db_auth_reconnect_lock_timeout_seconds = 0.1
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_returns_value_on_first_success():
|
||||
"""Happy path: factory succeeds first call, no reconnect attempted."""
|
||||
client = _make_client()
|
||||
|
||||
async def _factory():
|
||||
return {"id": 1}
|
||||
|
||||
result = await call_with_db_reconnect_retry(client, _factory, reason="happy_path")
|
||||
|
||||
assert result == {"id": 1}
|
||||
client.attempt_db_reconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_retries_after_transport_error():
|
||||
"""Transport error on first call → reconnect → second call succeeds."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
if len(invocations) == 1:
|
||||
raise httpx.ReadError("transport blip")
|
||||
return {"id": 1}
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="prisma_get_generic_data_config_lookup_failure"
|
||||
)
|
||||
|
||||
assert result == {"id": 1}
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_does_not_retry_on_data_layer_error():
|
||||
"""Data-layer errors (e.g. UniqueViolationError) are NOT transport errors —
|
||||
propagate immediately, do not reconnect."""
|
||||
client = _make_client()
|
||||
|
||||
async def _factory():
|
||||
raise UniqueViolationError(
|
||||
data={"user_facing_error": {"meta": {}}},
|
||||
message="Unique constraint failed",
|
||||
)
|
||||
|
||||
with pytest.raises(UniqueViolationError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="data_layer_test")
|
||||
|
||||
client.attempt_db_reconnect.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_propagates_when_reconnect_fails():
|
||||
"""Transport error, but reconnect returns False → propagate the original
|
||||
exception. Do not call factory a second time."""
|
||||
client = _make_client(attempt_db_reconnect_return=False)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
raise httpx.ReadError("transport blip")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="reconnect_fails")
|
||||
|
||||
assert len(invocations) == 1
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_propagates_after_second_transport_error():
|
||||
"""Transport error, reconnect succeeds, retry also raises transport error →
|
||||
propagate. At most one retry by construction (no infinite loop)."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
invocations = []
|
||||
|
||||
async def _factory():
|
||||
invocations.append(None)
|
||||
raise httpx.ReadError("still failing")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="second_transport_error"
|
||||
)
|
||||
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_skips_when_no_attempt_db_reconnect_attr():
|
||||
"""Older PrismaClient stand-ins / partial mocks may not expose
|
||||
`attempt_db_reconnect`. The helper must not crash — just propagate the
|
||||
original exception. Mirrors the `hasattr` guard from
|
||||
`auth_checks._fetch_key_object_from_db_with_reconnect`."""
|
||||
client = _make_client(has_attempt_db_reconnect=False)
|
||||
|
||||
async def _factory():
|
||||
raise httpx.ReadError("transport blip")
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="no_reconnect_attr")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_invokes_factory_twice_not_same_coro():
|
||||
"""Guard against the obvious bug of awaiting the same coroutine twice
|
||||
(`RuntimeError: cannot reuse already awaited coroutine`). The helper must
|
||||
call the factory a fresh time on retry, not cache an awaitable."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
factory_call_count = 0
|
||||
|
||||
async def _factory():
|
||||
nonlocal factory_call_count
|
||||
factory_call_count += 1
|
||||
if factory_call_count == 1:
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="fresh_coro_on_retry"
|
||||
)
|
||||
|
||||
assert result == "ok"
|
||||
assert factory_call_count == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_passes_explicit_timeouts():
|
||||
"""Explicit timeout_seconds / lock_timeout_seconds override the auth
|
||||
defaults read off the prisma_client object."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
|
||||
async def _factory():
|
||||
if not hasattr(_factory, "_called"):
|
||||
_factory._called = True # type: ignore[attr-defined]
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
result = await call_with_db_reconnect_retry(
|
||||
client,
|
||||
_factory,
|
||||
reason="explicit_timeouts",
|
||||
timeout_seconds=5.5,
|
||||
lock_timeout_seconds=0.25,
|
||||
)
|
||||
|
||||
assert result == "ok"
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["timeout_seconds"] == 5.5
|
||||
assert call_kwargs["lock_timeout_seconds"] == 0.25
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_uses_auth_defaults_when_unset():
|
||||
"""When timeouts are not provided, helper reads
|
||||
`_db_auth_reconnect_timeout_seconds` / `_db_auth_reconnect_lock_timeout_seconds`
|
||||
off the prisma_client (matching the auth path's existing convention)."""
|
||||
client = _make_client(attempt_db_reconnect_return=True)
|
||||
client._db_auth_reconnect_timeout_seconds = 3.0
|
||||
client._db_auth_reconnect_lock_timeout_seconds = 0.5
|
||||
|
||||
async def _factory():
|
||||
if not hasattr(_factory, "_called"):
|
||||
_factory._called = True # type: ignore[attr-defined]
|
||||
raise httpx.ReadError("transport blip")
|
||||
return "ok"
|
||||
|
||||
await call_with_db_reconnect_retry(client, _factory, reason="defaults")
|
||||
|
||||
call_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert call_kwargs["timeout_seconds"] == 3.0
|
||||
assert call_kwargs["lock_timeout_seconds"] == 0.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_call_with_db_reconnect_retry_preserves_original_error_when_reconnect_raises():
|
||||
"""If `attempt_db_reconnect` itself raises (lock cancellation, timer
|
||||
error, unexpected internal failure), the helper must surface the
|
||||
*original* transport error to telemetry — not the reconnect exception.
|
||||
Otherwise `failure_handler` / `db_exceptions` alerts log the wrong
|
||||
error string and the actual DB transport problem becomes invisible.
|
||||
|
||||
The reconnect error is chained as the `__cause__` for debuggability."""
|
||||
client = MagicMock()
|
||||
reconnect_exc = RuntimeError("simulated reconnect lock cancellation")
|
||||
client.attempt_db_reconnect = AsyncMock(side_effect=reconnect_exc)
|
||||
client._db_auth_reconnect_timeout_seconds = 2.0
|
||||
client._db_auth_reconnect_lock_timeout_seconds = 0.1
|
||||
|
||||
original_exc = httpx.ReadError("transport blip")
|
||||
|
||||
async def _factory():
|
||||
raise original_exc
|
||||
|
||||
with pytest.raises(httpx.ReadError) as exc_info:
|
||||
await call_with_db_reconnect_retry(
|
||||
client, _factory, reason="reconnect_itself_raises"
|
||||
)
|
||||
|
||||
assert exc_info.value is original_exc
|
||||
assert exc_info.value.__cause__ is reconnect_exc
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
|
|
@ -5,6 +5,7 @@ import sys
|
|||
import time
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
|
|
@ -358,3 +359,125 @@ async def test_lightweight_reconnect_skips_kill_on_successful_disconnect(
|
|||
await client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
mock_kill.assert_not_called()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_generic_data: transport-reconnect-and-retry coverage (issue #25143)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_generic_data_retries_on_transport_error_for_config_table(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""`get_generic_data(table_name="config")` self-heals on a transient
|
||||
`httpx.ReadError`: reconnect once, retry once, return the row.
|
||||
|
||||
Regression for issue #25143 — the 1.83.x line lost the reconnect-and-retry
|
||||
branch that 1.82.6 had on this method. `_update_config_from_db` fans out
|
||||
four concurrent `get_generic_data` calls, so a single transport flap used
|
||||
to surface as four `db_exceptions` alerts and a stale config window.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
|
||||
expected_row = {"param_name": "general_settings", "param_value": {"foo": "bar"}}
|
||||
invocations: list[None] = []
|
||||
|
||||
async def _flaky_find_first(**kwargs):
|
||||
invocations.append(None)
|
||||
if len(invocations) == 1:
|
||||
raise httpx.ReadError("simulated transport blip")
|
||||
return expected_row
|
||||
|
||||
client.db.litellm_config.find_first = AsyncMock(side_effect=_flaky_find_first)
|
||||
client.attempt_db_reconnect = AsyncMock(return_value=True)
|
||||
|
||||
result = await client.get_generic_data(
|
||||
key="param_name",
|
||||
value="general_settings",
|
||||
table_name="config",
|
||||
)
|
||||
|
||||
assert result == expected_row
|
||||
assert len(invocations) == 2
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
reconnect_kwargs = client.attempt_db_reconnect.await_args.kwargs
|
||||
assert reconnect_kwargs["reason"] == "prisma_get_generic_data_config_lookup_failure"
|
||||
|
||||
# The failure_handler telemetry side-effect must NOT fire on the first
|
||||
# transport blip — only if the post-retry call also fails. Drain the
|
||||
# event loop so any spuriously-spawned task would have run by now.
|
||||
await asyncio.sleep(0)
|
||||
mock_proxy_logging.failure_handler.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_generic_data_propagates_when_reconnect_fails(mock_proxy_logging):
|
||||
"""If reconnect itself does not succeed, propagate the original transport
|
||||
error and let the existing failure_handler / db_exceptions telemetry fire."""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
|
||||
client.db.litellm_config.find_first = AsyncMock(
|
||||
side_effect=httpx.ReadError("simulated transport blip")
|
||||
)
|
||||
client.attempt_db_reconnect = AsyncMock(return_value=False)
|
||||
|
||||
with pytest.raises(httpx.ReadError):
|
||||
await client.get_generic_data(
|
||||
key="param_name",
|
||||
value="general_settings",
|
||||
table_name="config",
|
||||
)
|
||||
|
||||
client.attempt_db_reconnect.assert_awaited_once()
|
||||
# Failure telemetry IS expected here — the read genuinely failed.
|
||||
await asyncio.sleep(0)
|
||||
mock_proxy_logging.failure_handler.assert_called_once()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _engine_confirmed_dead flag-reset bug (B2)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_engine_confirmed_dead_persists_across_failed_heavy_reconnect(
|
||||
mock_proxy_logging,
|
||||
):
|
||||
"""Regression test for the flag-reset bug.
|
||||
|
||||
Before the fix, `_run_reconnect_cycle` cleared
|
||||
`self._engine_confirmed_dead = False` *before* awaiting
|
||||
`_do_heavy_reconnect()`. If the heavy reconnect raised (e.g. timeout,
|
||||
missing DATABASE_URL, recreate failure), the flag was left cleared and the
|
||||
next attempt could demote to the lightweight path even though the engine
|
||||
was genuinely dead.
|
||||
|
||||
The fix moves the reset into the success branch — the flag must stay True
|
||||
when heavy reconnect raises.
|
||||
"""
|
||||
client = PrismaClient(
|
||||
database_url="mock://test", proxy_logging_obj=mock_proxy_logging
|
||||
)
|
||||
client._engine_confirmed_dead = True
|
||||
client._engine_pid = 0 # so `_is_engine_alive` is not consulted
|
||||
|
||||
# Make the heavy reconnect path raise.
|
||||
client.db.recreate_prisma_client = AsyncMock(
|
||||
side_effect=RuntimeError("simulated heavy reconnect failure")
|
||||
)
|
||||
client._start_engine_watcher = AsyncMock()
|
||||
client._cleanup_engine_watcher = MagicMock()
|
||||
client._reap_all_zombies = MagicMock()
|
||||
|
||||
with patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}):
|
||||
with pytest.raises(Exception):
|
||||
await client._run_reconnect_cycle(timeout_seconds=5.0)
|
||||
|
||||
# The flag must STILL be True so the next attempt re-enters the heavy
|
||||
# branch instead of silently demoting to the lightweight path.
|
||||
assert client._engine_confirmed_dead is True
|
||||
|
|
|
|||
|
|
@ -218,6 +218,141 @@ class TestProxyBaseLLMRequestProcessing:
|
|||
headers_with_invalid
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_litellm_proxy_success_headers_from_llm_response(self):
|
||||
"""
|
||||
Google native :generateContent uses this helper instead of base_process_llm_request;
|
||||
ensure x-litellm-* headers and callback hooks merge like the main proxy path.
|
||||
"""
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.headers = {}
|
||||
|
||||
class _FakeGenaiResponse:
|
||||
_hidden_params = {
|
||||
"model_id": "deployment-model-id",
|
||||
"cache_key": "ck-test",
|
||||
"api_base": "https://generativelanguage.googleapis.com/v1beta",
|
||||
"response_cost": 0.001,
|
||||
"additional_headers": {"llm_provider-ratelimit-requests": "1000"},
|
||||
}
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "call-id-test"
|
||||
|
||||
mock_user = MagicMock()
|
||||
mock_user.tpm_limit = None
|
||||
mock_user.rpm_limit = None
|
||||
mock_user.max_budget = None
|
||||
mock_user.spend = 0.0
|
||||
mock_user.allowed_model_region = None
|
||||
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(
|
||||
return_value={"x-ratelimit-remaining-requests": "999"}
|
||||
)
|
||||
|
||||
headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=_FakeGenaiResponse(),
|
||||
request_data={"model": "gemini/gemini-1.5-flash"},
|
||||
request=mock_request,
|
||||
user_api_key_dict=mock_user,
|
||||
logging_obj=logging_obj,
|
||||
version="9.9.9",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert headers["x-litellm-call-id"] == "call-id-test"
|
||||
assert headers["x-litellm-model-id"] == "deployment-model-id"
|
||||
assert headers["x-litellm-version"] == "9.9.9"
|
||||
assert headers["llm_provider-ratelimit-requests"] == "1000"
|
||||
assert headers["x-ratelimit-remaining-requests"] == "999"
|
||||
proxy_logging_obj.post_call_response_headers_hook.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_litellm_proxy_success_headers_streaming_style_iterator(self):
|
||||
"""AsyncGoogleGenAIGenerateContentStreamingIterator sets _hidden_params at init; headers must propagate."""
|
||||
|
||||
class _FakeStreamLike:
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
raise StopAsyncIteration
|
||||
|
||||
_hidden_params = {
|
||||
"model_id": "stream-model-id",
|
||||
"api_base": "https://generativelanguage.googleapis.com/v1beta",
|
||||
"cache_key": "",
|
||||
"response_cost": "",
|
||||
"additional_headers": {"llm_provider-x": "y"},
|
||||
}
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.headers = {}
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "cid-stream"
|
||||
mock_user = MagicMock()
|
||||
mock_user.tpm_limit = None
|
||||
mock_user.rpm_limit = None
|
||||
mock_user.max_budget = None
|
||||
mock_user.spend = 0.0
|
||||
mock_user.allowed_model_region = None
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
|
||||
|
||||
headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=_FakeStreamLike(),
|
||||
request_data={"model": "gemini/gemini-2.0-flash"},
|
||||
request=mock_request,
|
||||
user_api_key_dict=mock_user,
|
||||
logging_obj=logging_obj,
|
||||
version="1.0.0",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert headers["x-litellm-model-id"] == "stream-model-id"
|
||||
assert headers["x-litellm-model-api-base"] == (
|
||||
"https://generativelanguage.googleapis.com/v1beta"
|
||||
)
|
||||
assert headers["llm_provider-x"] == "y"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_build_litellm_proxy_success_headers_no_hidden_params_metadata_fallback(
|
||||
self,
|
||||
):
|
||||
"""When response has no _hidden_params, model_id can still come from litellm_metadata."""
|
||||
|
||||
class _BareResponse:
|
||||
pass
|
||||
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.headers = {}
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.litellm_call_id = "cid-meta"
|
||||
mock_user = MagicMock()
|
||||
mock_user.tpm_limit = None
|
||||
mock_user.rpm_limit = None
|
||||
mock_user.max_budget = None
|
||||
mock_user.spend = 0.0
|
||||
mock_user.allowed_model_region = None
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.post_call_response_headers_hook = AsyncMock(return_value={})
|
||||
|
||||
headers = await ProxyBaseLLMRequestProcessing.build_litellm_proxy_success_headers_from_llm_response(
|
||||
response=_BareResponse(),
|
||||
request_data={
|
||||
"model": "gemini/gemini-1.5-flash",
|
||||
"litellm_metadata": {"model_info": {"id": "meta-model-id"}},
|
||||
},
|
||||
request=mock_request,
|
||||
user_api_key_dict=mock_user,
|
||||
logging_obj=logging_obj,
|
||||
version="1.0.0",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert headers["x-litellm-model-id"] == "meta-model-id"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_litellm_data_to_request_with_stream_timeout_header(self):
|
||||
"""
|
||||
|
|
@ -1158,13 +1293,6 @@ class TestCommonRequestProcessingHelpers:
|
|||
assert mock_tracer.trace.call_count == 4
|
||||
|
||||
# Verify that each call was made with the correct operation name
|
||||
expected_calls = [
|
||||
(("streaming.chunk.yield",), {}),
|
||||
(("streaming.chunk.yield",), {}),
|
||||
(("streaming.chunk.yield",), {}),
|
||||
(("streaming.chunk.yield",), {}),
|
||||
]
|
||||
|
||||
actual_calls = mock_tracer.trace.call_args_list
|
||||
assert len(actual_calls) == 4
|
||||
|
||||
|
|
|
|||
|
|
@ -267,6 +267,7 @@ class TestRAGFlowVectorStore(BaseVectorStoreTest):
|
|||
api_base="http://localhost:9380",
|
||||
litellm_logging_obj=logging_obj,
|
||||
litellm_params={},
|
||||
extra_body=None,
|
||||
)
|
||||
|
||||
def test_transform_search_vector_store_response_not_implemented(self):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue