mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
refactor(types): replace Any with proven types in 12 files (#45943)
* refactor(types): replace Any with proven types in 13 files Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): import NotRequired from typing_extensions for Python 3.10 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): hide input values in provider response validation errors Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(types): freeze tool call accumulator and drop duplicate cache check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(types): accept null token text and float usage counts the old parsers handled Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
a7a3a6b375
commit
e0e3523246
12 changed files with 319 additions and 125 deletions
|
|
@ -2507,7 +2507,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
self.callback_duration_ms += (time.time() - _start) * 1000
|
||||
return payload
|
||||
|
||||
def _transform_usage_objects(self, result):
|
||||
def _transform_usage_objects(self, result: object) -> object:
|
||||
if isinstance(result, ResponsesAPIResponse):
|
||||
result = result.model_copy()
|
||||
transformed_usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(result.usage)
|
||||
|
|
@ -2545,13 +2545,13 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
def _success_handler_helper_fn(
|
||||
self,
|
||||
result=None,
|
||||
start_time=None,
|
||||
end_time=None,
|
||||
cache_hit=None,
|
||||
result: object = None,
|
||||
start_time: datetime.datetime | None = None,
|
||||
end_time: datetime.datetime | None = None,
|
||||
cache_hit: bool | None = None,
|
||||
standard_logging_object: StandardLoggingPayload | None = None,
|
||||
build_logging_payload: bool = True,
|
||||
):
|
||||
) -> tuple[datetime.datetime, datetime.datetime, object]:
|
||||
try:
|
||||
if start_time is None:
|
||||
start_time = self.start_time
|
||||
|
|
|
|||
|
|
@ -1,9 +1,10 @@
|
|||
import base64
|
||||
import time
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass, replace
|
||||
from itertools import groupby
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional, TypeAlias, TypedDict, Union, cast
|
||||
from typing import TYPE_CHECKING, Final, Optional, TypeAlias, TypedDict, Union, cast
|
||||
|
||||
from pydantic import BaseModel
|
||||
from typing_extensions import ReadOnly, Required
|
||||
|
|
@ -154,6 +155,15 @@ class _ToolCallChoice(TypedDict, total=False):
|
|||
_ToolCallKey: TypeAlias = tuple[int, int]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _ToolCallAccumulator:
|
||||
id: str | None = None
|
||||
name: str | None = None
|
||||
type: str | None = None
|
||||
custom_name: str | None = None
|
||||
provider_specific_fields: Mapping[str, object] | None = None
|
||||
|
||||
|
||||
class _ToolCallChunk(TypedDict):
|
||||
choices: ReadOnly[Sequence[_ToolCallChoice]]
|
||||
|
||||
|
|
@ -503,7 +513,7 @@ class ChunkProcessor:
|
|||
self, tool_call_chunks: Sequence["_ToolCallChunk"]
|
||||
) -> list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall]:
|
||||
tool_calls_list: list[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall] = []
|
||||
tool_call_map: Final[dict[_ToolCallKey, dict[str, Any]]] = {}
|
||||
tool_call_map: Final[dict[_ToolCallKey, _ToolCallAccumulator]] = {}
|
||||
|
||||
for chunk in tool_call_chunks:
|
||||
choices = chunk["choices"]
|
||||
|
|
@ -537,49 +547,43 @@ class ChunkProcessor:
|
|||
index = (choice_index, getattr(tool_call, "index", 0))
|
||||
|
||||
if index not in tool_call_map:
|
||||
tool_call_map[index] = {
|
||||
"id": None,
|
||||
"name": None,
|
||||
"type": None,
|
||||
"custom_name": None,
|
||||
"provider_specific_fields": None,
|
||||
}
|
||||
tool_call_map[index] = _ToolCallAccumulator()
|
||||
|
||||
# Extract id, type, and function data (handle both dict and object)
|
||||
if isinstance(tool_call, dict):
|
||||
if fragment_id := tool_call.get("id"):
|
||||
tool_call_map[index]["id"] = fragment_id
|
||||
tool_call_map[index] = replace(tool_call_map[index], id=fragment_id)
|
||||
if fragment_type := tool_call.get("type"):
|
||||
tool_call_map[index]["type"] = fragment_type
|
||||
tool_call_map[index] = replace(tool_call_map[index], type=fragment_type)
|
||||
|
||||
function = tool_call.get("function", {})
|
||||
if isinstance(function, dict):
|
||||
if fragment_name := function.get("name"):
|
||||
tool_call_map[index]["name"] = fragment_name
|
||||
tool_call_map[index] = replace(tool_call_map[index], name=fragment_name)
|
||||
else:
|
||||
# function is an object
|
||||
if function_name := getattr(function, "name", None):
|
||||
tool_call_map[index]["name"] = function_name
|
||||
tool_call_map[index] = replace(tool_call_map[index], name=function_name)
|
||||
|
||||
custom = tool_call.get("custom")
|
||||
if isinstance(custom, dict):
|
||||
if custom_name := custom.get("name"):
|
||||
tool_call_map[index]["custom_name"] = custom_name
|
||||
tool_call_map[index] = replace(tool_call_map[index], custom_name=custom_name)
|
||||
else:
|
||||
# tool_call is an object
|
||||
if hasattr(tool_call, "id") and tool_call.id:
|
||||
tool_call_map[index]["id"] = tool_call.id
|
||||
tool_call_map[index] = replace(tool_call_map[index], id=tool_call.id)
|
||||
if hasattr(tool_call, "type") and tool_call.type:
|
||||
tool_call_map[index]["type"] = tool_call.type
|
||||
tool_call_map[index] = replace(tool_call_map[index], type=tool_call.type)
|
||||
if object_function_name := getattr(getattr(tool_call, "function", None), "name", None):
|
||||
tool_call_map[index]["name"] = object_function_name
|
||||
tool_call_map[index] = replace(tool_call_map[index], name=object_function_name)
|
||||
|
||||
object_custom: ChatCompletionDeltaCustomToolCallPayload | None = getattr(
|
||||
tool_call, "custom", None
|
||||
)
|
||||
if object_custom is not None:
|
||||
if getattr(object_custom, "name", None):
|
||||
tool_call_map[index]["custom_name"] = object_custom.name
|
||||
tool_call_map[index] = replace(tool_call_map[index], custom_name=object_custom.name)
|
||||
|
||||
# Preserve provider_specific_fields from streaming chunks
|
||||
provider_fields: object = None
|
||||
|
|
@ -602,12 +606,14 @@ class ChunkProcessor:
|
|||
|
||||
if provider_fields:
|
||||
# Merge provider_specific_fields if multiple chunks have them
|
||||
merged_provider_fields = tool_call_map[index]["provider_specific_fields"]
|
||||
if merged_provider_fields is None:
|
||||
merged_provider_fields = {}
|
||||
tool_call_map[index]["provider_specific_fields"] = merged_provider_fields
|
||||
if isinstance(provider_fields, dict):
|
||||
merged_provider_fields.update(provider_fields)
|
||||
tool_call_map[index] = replace(
|
||||
tool_call_map[index],
|
||||
provider_specific_fields=(
|
||||
{**(tool_call_map[index].provider_specific_fields or {}), **provider_fields}
|
||||
if isinstance(provider_fields, dict)
|
||||
else dict(tool_call_map[index].provider_specific_fields or {})
|
||||
),
|
||||
)
|
||||
|
||||
joined_fragments: Final = self._join_fragments_by_key_and_field(
|
||||
self._iter_tool_call_fragments(tool_call_chunks)
|
||||
|
|
@ -616,37 +622,37 @@ class ChunkProcessor:
|
|||
# Convert the map to a list of tool calls
|
||||
for index in sorted(tool_call_map.keys()):
|
||||
tool_call_data = tool_call_map[index]
|
||||
if tool_call_data["id"] and tool_call_data["custom_name"]:
|
||||
if tool_call_data.id and tool_call_data.custom_name:
|
||||
tool_calls_list.append(
|
||||
ChatCompletionMessageCustomToolCall(
|
||||
id=tool_call_data["id"],
|
||||
id=tool_call_data.id,
|
||||
custom=ChatCompletionCustomToolCallPayload(
|
||||
name=tool_call_data["custom_name"],
|
||||
name=tool_call_data.custom_name,
|
||||
input=joined_fragments.get((index, "custom_input"), ""),
|
||||
),
|
||||
)
|
||||
)
|
||||
elif tool_call_data["id"] and tool_call_data["name"]:
|
||||
elif tool_call_data.id and tool_call_data.name:
|
||||
combined_arguments = joined_fragments.get((index, "arguments"), "") or "{}"
|
||||
|
||||
# Build function - provider_specific_fields should be on tool_call level, not function level
|
||||
function = Function(
|
||||
arguments=combined_arguments,
|
||||
name=tool_call_data["name"],
|
||||
name=tool_call_data.name,
|
||||
)
|
||||
|
||||
# Prepare params for ChatCompletionMessageToolCall
|
||||
tool_call_params = {
|
||||
"id": tool_call_data["id"],
|
||||
"function": function,
|
||||
"type": tool_call_data["type"] or "function",
|
||||
}
|
||||
|
||||
# Add provider_specific_fields if present (for thought signatures in Gemini 3)
|
||||
if tool_call_data.get("provider_specific_fields"):
|
||||
tool_call_params["provider_specific_fields"] = tool_call_data["provider_specific_fields"]
|
||||
|
||||
tool_call = ChatCompletionMessageToolCall(**tool_call_params)
|
||||
tool_call = ChatCompletionMessageToolCall(
|
||||
id=tool_call_data.id,
|
||||
function=function,
|
||||
type=tool_call_data.type or "function",
|
||||
**(
|
||||
{"provider_specific_fields": dict(tool_call_data.provider_specific_fields)}
|
||||
if tool_call_data.provider_specific_fields
|
||||
else {}
|
||||
),
|
||||
)
|
||||
tool_calls_list.append(tool_call)
|
||||
|
||||
return tool_calls_list
|
||||
|
|
|
|||
|
|
@ -474,7 +474,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
Returns:
|
||||
AgentCoreParsedResponse: Parsed response data
|
||||
"""
|
||||
content_type: Final = raw_response.headers.get("content-type", "").lower()
|
||||
content_type: Final = ", ".join(httpx.Headers(raw_response.headers).get_list("content-type")).lower()
|
||||
verbose_logger.debug("AgentCore response Content-Type: %s", content_type)
|
||||
|
||||
# Parse response based on content type
|
||||
|
|
@ -698,7 +698,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
)
|
||||
|
||||
# Check if response is JSON (agent used sync return) instead of SSE
|
||||
content_type: Final = response.headers.get("content-type", "").lower()
|
||||
content_type: Final = ", ".join(httpx.Headers(response.headers).get_list("content-type")).lower()
|
||||
if "application/json" in content_type:
|
||||
verbose_logger.debug(
|
||||
"AgentCore streaming: received JSON response instead of SSE, converting to single-chunk stream"
|
||||
|
|
@ -920,7 +920,7 @@ class AmazonAgentCoreConfig(BaseConfig, BaseAWSLLM):
|
|||
)
|
||||
|
||||
# Check if response is JSON (agent used sync return) instead of SSE
|
||||
content_type: Final = response.headers.get("content-type", "").lower()
|
||||
content_type: Final = ", ".join(httpx.Headers(response.headers).get_list("content-type")).lower()
|
||||
if "application/json" in content_type:
|
||||
verbose_logger.debug(
|
||||
"AgentCore streaming: received JSON response instead of SSE, converting to single-chunk stream"
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, RootModel, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
|
|
@ -12,6 +15,16 @@ else:
|
|||
GenerateContentContentListUnionDict = Any
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _GeminiCountTokensResponse(TypedDict):
|
||||
totalTokens: NotRequired[ReadOnly[int | float]]
|
||||
tokenizer_used: NotRequired[ReadOnly[str]]
|
||||
|
||||
|
||||
class _GeminiCountTokensResponseRoot(RootModel[_GeminiCountTokensResponse]):
|
||||
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
|
||||
|
||||
|
||||
class GoogleAIStudioTokenCounter:
|
||||
def _clean_contents_for_gemini_api(self, contents: Any) -> object:
|
||||
"""
|
||||
|
|
@ -79,7 +92,7 @@ class GoogleAIStudioTokenCounter:
|
|||
|
||||
async def acount_tokens(
|
||||
self,
|
||||
contents: Any,
|
||||
contents: Sequence[Mapping[str, object]] | None,
|
||||
model: str,
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
|
|
@ -143,8 +156,8 @@ class GoogleAIStudioTokenCounter:
|
|||
response.raise_for_status()
|
||||
|
||||
# Parse response
|
||||
result: Final = response.json()
|
||||
return result
|
||||
result: Final = _GeminiCountTokensResponseRoot.model_validate(response.json()).root
|
||||
return dict(result)
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_msg = f"Google Gen AI Studio API error: {e.response.status_code} - {e.response.text}"
|
||||
|
|
|
|||
|
|
@ -41,11 +41,14 @@ Response format:
|
|||
"""
|
||||
|
||||
import base64
|
||||
from collections.abc import Mapping, Sequence
|
||||
from io import BufferedReader, BytesIO
|
||||
from typing import TYPE_CHECKING, Any, Final, cast
|
||||
|
||||
import httpx
|
||||
from httpx._types import RequestFiles
|
||||
from pydantic import ConfigDict, RootModel, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm.images.utils import ImageEditRequestUtils
|
||||
|
|
@ -71,6 +74,52 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _ImageUrl(TypedDict):
|
||||
url: ReadOnly[NotRequired[str | None]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _Image(TypedDict):
|
||||
image_url: ReadOnly[NotRequired[_ImageUrl]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _Message(TypedDict):
|
||||
images: ReadOnly[NotRequired[Sequence[_Image]]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _Choice(TypedDict):
|
||||
message: ReadOnly[NotRequired[_Message]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _TokensDetails(TypedDict):
|
||||
image_tokens: ReadOnly[NotRequired[int | float]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _Usage(TypedDict):
|
||||
prompt_tokens: ReadOnly[NotRequired[int | float]]
|
||||
total_tokens: ReadOnly[NotRequired[int | float]]
|
||||
completion_tokens_details: ReadOnly[NotRequired[_TokensDetails]]
|
||||
prompt_tokens_details: ReadOnly[NotRequired[_TokensDetails | None]]
|
||||
cost: ReadOnly[NotRequired[int | float | str | None]]
|
||||
cost_details: ReadOnly[NotRequired[Mapping[str, object] | None]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _OpenRouterImageEditResponse(TypedDict):
|
||||
choices: ReadOnly[NotRequired[Sequence[_Choice]]]
|
||||
usage: ReadOnly[NotRequired[_Usage | None]]
|
||||
model: ReadOnly[NotRequired[object]]
|
||||
|
||||
|
||||
class _OpenRouterImageEditResponseRoot(RootModel[_OpenRouterImageEditResponse]):
|
||||
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
|
||||
|
||||
|
||||
class OpenRouterImageEditConfig(BaseImageEditConfig):
|
||||
"""
|
||||
Configuration for OpenRouter image editing via chat completions.
|
||||
|
|
@ -201,7 +250,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
|
|||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> ImageResponse:
|
||||
try:
|
||||
response_json: Final = raw_response.json()
|
||||
response_json: Final[object] = raw_response.json()
|
||||
except Exception as e:
|
||||
raise OpenRouterException(
|
||||
message=f"Error parsing OpenRouter response: {e}",
|
||||
|
|
@ -213,7 +262,8 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
|
|||
model_response.data = []
|
||||
|
||||
try:
|
||||
choices: Final = response_json.get("choices", [])
|
||||
payload: Final = _OpenRouterImageEditResponseRoot.model_validate(response_json).root
|
||||
choices: Final = payload.get("choices", [])
|
||||
|
||||
for choice in choices:
|
||||
message = choice.get("message", {})
|
||||
|
|
@ -252,7 +302,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
|
|||
headers={},
|
||||
)
|
||||
|
||||
self._set_usage_and_cost(model_response, response_json, model)
|
||||
self._set_usage_and_cost(model_response, payload, model)
|
||||
return model_response
|
||||
|
||||
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
|
||||
|
|
@ -303,11 +353,11 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
|
|||
def _set_usage_and_cost(
|
||||
self,
|
||||
model_response: ImageResponse,
|
||||
response_json: dict,
|
||||
payload: _OpenRouterImageEditResponse,
|
||||
model: str,
|
||||
) -> None:
|
||||
"""Extract and set usage and cost information from OpenRouter response."""
|
||||
usage_data: Final = response_json.get("usage", {})
|
||||
usage_data: Final = payload.get("usage", {})
|
||||
if usage_data:
|
||||
prompt_tokens: Final = usage_data.get("prompt_tokens", 0)
|
||||
total_tokens: Final = usage_data.get("total_tokens", 0)
|
||||
|
|
@ -347,7 +397,7 @@ class OpenRouterImageEditConfig(BaseImageEditConfig):
|
|||
)
|
||||
response_cost_details.update(cost_details)
|
||||
|
||||
model_response.hidden_params["model"] = response_json.get("model", model)
|
||||
model_response.hidden_params["model"] = payload.get("model", model)
|
||||
|
||||
def _read_image_bytes(self, image: FileTypes) -> bytes:
|
||||
"""Read raw bytes from various image input types."""
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@ from collections.abc import Mapping
|
|||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, RootModel, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
from litellm.llms.base_llm.vector_store.transformation import BaseVectorStoreConfig
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
|
|
@ -22,6 +24,25 @@ else:
|
|||
LiteLLMLoggingObj = Any
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _RAGFlowDataset(TypedDict):
|
||||
id: ReadOnly[NotRequired[str | None]]
|
||||
name: ReadOnly[NotRequired[str | None]]
|
||||
create_time: ReadOnly[NotRequired[int | float | None]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _RAGFlowCreateResponse(TypedDict):
|
||||
data: ReadOnly[NotRequired[_RAGFlowDataset]]
|
||||
|
||||
|
||||
class _RAGFlowCreateResponseRoot(RootModel[_RAGFlowCreateResponse]):
|
||||
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
|
||||
|
||||
|
||||
_RAGFLOW_EMPTY_DATA: Final[_RAGFlowDataset] = {}
|
||||
|
||||
|
||||
class RAGFlowVectorStoreConfig(BaseVectorStoreConfig):
|
||||
"""Vector store configuration for RAGFlow datasets."""
|
||||
|
||||
|
|
@ -185,7 +206,8 @@ class RAGFlowVectorStoreConfig(BaseVectorStoreConfig):
|
|||
headers=response.headers,
|
||||
)
|
||||
|
||||
data: Final = response_json.get("data", {})
|
||||
create_response: Final = _RAGFlowCreateResponseRoot.model_validate(response_json).root
|
||||
data: Final = create_response.get("data", _RAGFLOW_EMPTY_DATA)
|
||||
|
||||
# Extract dataset ID
|
||||
dataset_id: Final = data.get("id")
|
||||
|
|
|
|||
|
|
@ -4,7 +4,8 @@ from collections.abc import AsyncIterator, Iterator, Mapping
|
|||
from typing import Final
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, TypeAdapter
|
||||
from pydantic import ConfigDict, RootModel, TypeAdapter, with_config
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm import verbose_logger
|
||||
|
|
@ -15,6 +16,38 @@ from litellm.types.utils import StreamingChatCompletionChunk
|
|||
_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object], config=ConfigDict(hide_input_in_errors=True))
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _SageMakerToken(TypedDict, total=False):
|
||||
text: ReadOnly[str | None]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _SageMakerChunkData(TypedDict, total=False):
|
||||
token: ReadOnly[_SageMakerToken | None]
|
||||
index: ReadOnly[int | None]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _SageMakerMessagesChunkData(TypedDict, total=False):
|
||||
choices: ReadOnly[Required[list[Mapping[str, object]]]]
|
||||
|
||||
|
||||
class _SageMakerChunkDataRoot(RootModel[_SageMakerChunkData]):
|
||||
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
|
||||
|
||||
|
||||
class _SageMakerMessagesChunkDataRoot(RootModel[_SageMakerMessagesChunkData]):
|
||||
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
|
||||
|
||||
|
||||
def _parse_sagemaker_chunk_data(data: str) -> _SageMakerChunkData:
|
||||
return _SageMakerChunkDataRoot.model_validate(json.loads(data)).root
|
||||
|
||||
|
||||
def _parse_sagemaker_messages_chunk_data(data: str) -> _SageMakerMessagesChunkData:
|
||||
return _SageMakerMessagesChunkDataRoot.model_validate(json.loads(data)).root
|
||||
|
||||
|
||||
def _load_sagemaker_response_stream_shape():
|
||||
try:
|
||||
from botocore.loaders import Loader
|
||||
|
|
@ -62,12 +95,12 @@ class AWSEventStreamDecoder:
|
|||
self.content_blocks: list = []
|
||||
self.is_messages_api = is_messages_api
|
||||
|
||||
def _chunk_parser_messages_api(self, chunk_data: dict) -> StreamingChatCompletionChunk:
|
||||
def _chunk_parser_messages_api(self, chunk_data: _SageMakerMessagesChunkData) -> StreamingChatCompletionChunk:
|
||||
openai_chunk: Final = StreamingChatCompletionChunk(**{"model": self.model, **chunk_data})
|
||||
|
||||
return openai_chunk
|
||||
|
||||
def _chunk_parser(self, chunk_data: dict) -> GChunk:
|
||||
def _chunk_parser(self, chunk_data: _SageMakerChunkData) -> GChunk:
|
||||
verbose_logger.debug("in sagemaker chunk parser, chunk_data %s", chunk_data)
|
||||
_token: Final = chunk_data.get("token", {}) or {}
|
||||
_index: Final = chunk_data.get("index", None) or 0
|
||||
|
|
@ -113,11 +146,12 @@ class AWSEventStreamDecoder:
|
|||
|
||||
# Try to parse the accumulated JSON
|
||||
try:
|
||||
_data = json.loads(accumulated_json)
|
||||
if self.is_messages_api:
|
||||
yield self._chunk_parser_messages_api(chunk_data=_data)
|
||||
yield self._chunk_parser_messages_api(
|
||||
chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
|
||||
)
|
||||
else:
|
||||
yield self._chunk_parser(chunk_data=_data)
|
||||
yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
|
||||
# Reset accumulated_json after successful parsing
|
||||
accumulated_json = ""
|
||||
except json.JSONDecodeError:
|
||||
|
|
@ -127,11 +161,12 @@ class AWSEventStreamDecoder:
|
|||
# Handle any remaining data after the iterator is exhausted
|
||||
if accumulated_json:
|
||||
try:
|
||||
_data = json.loads(accumulated_json)
|
||||
if self.is_messages_api:
|
||||
yield self._chunk_parser_messages_api(chunk_data=_data)
|
||||
yield self._chunk_parser_messages_api(
|
||||
chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
|
||||
)
|
||||
else:
|
||||
yield self._chunk_parser(chunk_data=_data)
|
||||
yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
|
||||
except json.JSONDecodeError:
|
||||
# Handle or log any unparseable data at the end
|
||||
verbose_logger.error("Warning: Unparseable JSON data remained: %s", accumulated_json)
|
||||
|
|
@ -161,11 +196,12 @@ class AWSEventStreamDecoder:
|
|||
accumulated_json += message
|
||||
|
||||
# Try to parse the accumulated JSON
|
||||
_data = json.loads(accumulated_json)
|
||||
if self.is_messages_api:
|
||||
yield self._chunk_parser_messages_api(chunk_data=_data)
|
||||
yield self._chunk_parser_messages_api(
|
||||
chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
|
||||
)
|
||||
else:
|
||||
yield self._chunk_parser(chunk_data=_data)
|
||||
yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
|
||||
# Reset accumulated_json after successful parsing
|
||||
accumulated_json = ""
|
||||
except json.JSONDecodeError:
|
||||
|
|
@ -181,11 +217,12 @@ class AWSEventStreamDecoder:
|
|||
# Handle any remaining data after the iterator is exhausted
|
||||
if accumulated_json:
|
||||
try:
|
||||
_data = json.loads(accumulated_json)
|
||||
if self.is_messages_api:
|
||||
yield self._chunk_parser_messages_api(chunk_data=_data)
|
||||
yield self._chunk_parser_messages_api(
|
||||
chunk_data=_parse_sagemaker_messages_chunk_data(accumulated_json)
|
||||
)
|
||||
else:
|
||||
yield self._chunk_parser(chunk_data=_data)
|
||||
yield self._chunk_parser(chunk_data=_parse_sagemaker_chunk_data(accumulated_json))
|
||||
except json.JSONDecodeError:
|
||||
# Handle or log any unparseable data at the end
|
||||
verbose_logger.error("Warning: Unparseable JSON data remained: %s", accumulated_json)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,9 @@
|
|||
from collections.abc import Sequence
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
from pydantic import ConfigDict, RootModel, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -13,10 +16,6 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
get_async_httpx_client,
|
||||
)
|
||||
from litellm.llms.openai.openai import AllMessageValues
|
||||
from litellm.types.llms.vertex_ai import (
|
||||
CachedContentListAllResponseBody,
|
||||
VertexAICachedContentResponseObject,
|
||||
)
|
||||
from litellm.utils import is_prompt_caching_valid_prompt
|
||||
|
||||
from ..common_utils import VertexAIError, get_vertex_base_url
|
||||
|
|
@ -32,6 +31,31 @@ local_cache_obj: Final = Cache(type=LiteLLMCacheType.LOCAL) # only used for cal
|
|||
MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _CachedContentItem(TypedDict):
|
||||
displayName: ReadOnly[NotRequired[object]]
|
||||
name: ReadOnly[NotRequired[str | None]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _CachedContentListBody(TypedDict):
|
||||
cachedContents: ReadOnly[NotRequired[Sequence[_CachedContentItem]]]
|
||||
nextPageToken: ReadOnly[NotRequired[str | None]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _CreatedCachedContent(TypedDict):
|
||||
name: ReadOnly[NotRequired[str | None]]
|
||||
|
||||
|
||||
class _CachedContentListBodyRoot(RootModel[_CachedContentListBody]):
|
||||
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
|
||||
|
||||
|
||||
class _CreatedCachedContentRoot(RootModel[_CreatedCachedContent]):
|
||||
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
|
||||
|
||||
|
||||
class ContextCachingEndpoints(VertexBase):
|
||||
"""
|
||||
Covers context caching endpoints for Vertex AI + Google AI Studio
|
||||
|
|
@ -157,25 +181,20 @@ class ContextCachingEndpoints(VertexBase):
|
|||
except Exception as e:
|
||||
raise VertexAIError(status_code=500, message=str(e))
|
||||
|
||||
raw_response = resp.json()
|
||||
raw_response = _CachedContentListBodyRoot.model_validate(resp.json()).root
|
||||
logging_obj.post_call(original_response=raw_response)
|
||||
|
||||
if "cachedContents" not in raw_response:
|
||||
return None
|
||||
|
||||
all_cached_items = CachedContentListAllResponseBody(**raw_response)
|
||||
|
||||
if "cachedContents" not in all_cached_items:
|
||||
return None
|
||||
|
||||
# Check current page for matching cache_key
|
||||
for cached_item in all_cached_items["cachedContents"]:
|
||||
for cached_item in raw_response["cachedContents"]:
|
||||
display_name = cached_item.get("displayName")
|
||||
if display_name is not None and display_name == cache_key:
|
||||
return cached_item.get("name")
|
||||
|
||||
# Check if there are more pages
|
||||
page_token = all_cached_items.get("nextPageToken")
|
||||
page_token = raw_response.get("nextPageToken")
|
||||
if not page_token:
|
||||
# No more pages, cache not found
|
||||
break
|
||||
|
|
@ -249,25 +268,20 @@ class ContextCachingEndpoints(VertexBase):
|
|||
except Exception as e:
|
||||
raise VertexAIError(status_code=500, message=str(e))
|
||||
|
||||
raw_response = resp.json()
|
||||
raw_response = _CachedContentListBodyRoot.model_validate(resp.json()).root
|
||||
logging_obj.post_call(original_response=raw_response)
|
||||
|
||||
if "cachedContents" not in raw_response:
|
||||
return None
|
||||
|
||||
all_cached_items = CachedContentListAllResponseBody(**raw_response)
|
||||
|
||||
if "cachedContents" not in all_cached_items:
|
||||
return None
|
||||
|
||||
# Check current page for matching cache_key
|
||||
for cached_item in all_cached_items["cachedContents"]:
|
||||
for cached_item in raw_response["cachedContents"]:
|
||||
display_name = cached_item.get("displayName")
|
||||
if display_name is not None and display_name == cache_key:
|
||||
return cached_item.get("name")
|
||||
|
||||
# Check if there are more pages
|
||||
page_token = all_cached_items.get("nextPageToken")
|
||||
page_token = raw_response.get("nextPageToken")
|
||||
if not page_token:
|
||||
# No more pages, cache not found
|
||||
break
|
||||
|
|
@ -424,14 +438,11 @@ class ContextCachingEndpoints(VertexBase):
|
|||
except httpx.TimeoutException:
|
||||
raise VertexAIError(status_code=408, message="Timeout error occurred.")
|
||||
|
||||
raw_response_cached: Final = response.json()
|
||||
cached_content_response_obj: Final = VertexAICachedContentResponseObject(
|
||||
name=raw_response_cached.get("name"), model=raw_response_cached.get("model")
|
||||
)
|
||||
raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
|
||||
return (
|
||||
non_cached_messages,
|
||||
optional_params,
|
||||
cached_content_response_obj["name"],
|
||||
raw_response_cached.get("name"),
|
||||
)
|
||||
|
||||
async def async_check_and_create_cache(
|
||||
|
|
@ -580,14 +591,11 @@ class ContextCachingEndpoints(VertexBase):
|
|||
except httpx.TimeoutException:
|
||||
raise VertexAIError(status_code=408, message="Timeout error occurred.")
|
||||
|
||||
raw_response_cached: Final = response.json()
|
||||
cached_content_response_obj: Final = VertexAICachedContentResponseObject(
|
||||
name=raw_response_cached.get("name"), model=raw_response_cached.get("model")
|
||||
)
|
||||
raw_response_cached: Final = _CreatedCachedContentRoot.model_validate(response.json()).root
|
||||
return (
|
||||
non_cached_messages,
|
||||
optional_params,
|
||||
cached_content_response_obj["name"],
|
||||
raw_response_cached.get("name"),
|
||||
)
|
||||
|
||||
def get_cache(self):
|
||||
|
|
|
|||
|
|
@ -2209,10 +2209,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
@staticmethod
|
||||
def _set_stream_metadata_on_response(
|
||||
model_response: Any,
|
||||
grounding_metadata: list[dict],
|
||||
url_context_metadata: list[dict],
|
||||
safety_ratings: list[dict],
|
||||
citation_metadata: list[dict],
|
||||
grounding_metadata: Sequence[Mapping[str, object]],
|
||||
url_context_metadata: Sequence[Mapping[str, object]],
|
||||
safety_ratings: Sequence[Mapping[str, object]],
|
||||
citation_metadata: Sequence[Mapping[str, object]],
|
||||
) -> None:
|
||||
setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata)
|
||||
if grounding_metadata:
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal
|
|||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
from pydantic import ConfigDict, RootModel, with_config
|
||||
from typing_extensions import NotRequired, ReadOnly, TypedDict
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
|
@ -65,6 +67,55 @@ from litellm.types.utils import (
|
|||
|
||||
GUARDRAIL_NAME: Final = "model_armor"
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _MatchResult(TypedDict):
|
||||
matchState: NotRequired[ReadOnly[str]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _DeidentifyData(TypedDict):
|
||||
text: NotRequired[ReadOnly[str]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _DeidentifyResult(TypedDict):
|
||||
matchState: NotRequired[ReadOnly[str]]
|
||||
data: NotRequired[ReadOnly[_DeidentifyData]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _SdpResult(TypedDict):
|
||||
inspectResult: NotRequired[ReadOnly[_MatchResult]]
|
||||
deidentifyResult: NotRequired[ReadOnly[_DeidentifyResult]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _Filter(TypedDict):
|
||||
raiFilterResult: NotRequired[ReadOnly[_MatchResult]]
|
||||
piAndJailbreakFilterResult: NotRequired[ReadOnly[_MatchResult]]
|
||||
maliciousUriFilterResult: NotRequired[ReadOnly[_MatchResult]]
|
||||
csamFilterFilterResult: NotRequired[ReadOnly[_MatchResult]]
|
||||
virusScanFilterResult: NotRequired[ReadOnly[_MatchResult]]
|
||||
sdpFilterResult: NotRequired[ReadOnly[_SdpResult | None]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _SanitizationResult(TypedDict):
|
||||
filterResults: NotRequired[ReadOnly[Mapping[str, _Filter] | Sequence[_Filter] | str | int | float | bool | None]]
|
||||
|
||||
|
||||
@with_config(ConfigDict(extra="allow", strict=True, hide_input_in_errors=True))
|
||||
class _ModelArmorResponse(TypedDict):
|
||||
sanitizationResult: NotRequired[ReadOnly[_SanitizationResult]]
|
||||
sanitizedText: NotRequired[ReadOnly[str | None]]
|
||||
text: NotRequired[ReadOnly[str | None]]
|
||||
|
||||
|
||||
class _ModelArmorResponseRoot(RootModel[_ModelArmorResponse]):
|
||||
model_config = ConfigDict(frozen=True, hide_input_in_errors=True)
|
||||
|
||||
|
||||
# Only these carry the finished output; response.created carries an empty body
|
||||
_RESPONSES_TERMINAL_EVENT_TYPES: Final = frozenset({"response.completed", "response.incomplete", "response.failed"})
|
||||
|
||||
|
|
@ -99,10 +150,11 @@ class ModelArmorAPIError(Exception):
|
|||
|
||||
_SCANNED_CONTENT_KEYS: Final = frozenset({"text", "sanitizedText", "findings", "maliciousUriMatchedItems"})
|
||||
|
||||
|
||||
RedactablePayload = dict | list | str | int | float | bool | None
|
||||
|
||||
|
||||
def _redact_scanned_content(payload: RedactablePayload, depth: int = 0) -> RedactablePayload:
|
||||
def _redact_scanned_content(payload: object, depth: int = 0) -> object:
|
||||
if depth >= DEFAULT_MAX_RECURSE_DEPTH:
|
||||
return "[REDACTED]"
|
||||
if isinstance(payload, dict):
|
||||
|
|
@ -223,12 +275,12 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
return f"Model Armor API error (upstream {status_code})"
|
||||
return f"Model Armor API error (upstream {status_code}): {response_text}"
|
||||
|
||||
def _build_block_error_detail(self, message: str, armor_response: RedactablePayload) -> dict:
|
||||
def _build_block_error_detail(self, message: str, armor_response: Mapping[str, object]) -> dict:
|
||||
if self.sanitize_error_detail:
|
||||
return {"error": message}
|
||||
return {"error": message, "model_armor_response": armor_response}
|
||||
|
||||
def _build_logging_response(self, armor_response: RedactablePayload) -> RedactablePayload:
|
||||
def _build_logging_response(self, armor_response: Mapping[str, object]) -> object:
|
||||
if self.sanitize_error_detail:
|
||||
return _redact_scanned_content(armor_response)
|
||||
return armor_response
|
||||
|
|
@ -286,7 +338,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
request_data: dict | None = None,
|
||||
file_bytes: bytes | None = None,
|
||||
file_type: str | None = None,
|
||||
) -> dict:
|
||||
) -> _ModelArmorResponse:
|
||||
"""
|
||||
Make request to Model Armor API. Supports both text and file prompt sanitization.
|
||||
If file_bytes and file_type are provided, file prompt sanitization is performed.
|
||||
|
|
@ -361,8 +413,8 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
|
||||
json_response: Final = response.json()
|
||||
if hasattr(json_response, "__await__"):
|
||||
return await json_response
|
||||
return json_response
|
||||
return _ModelArmorResponseRoot.model_validate(await json_response).root
|
||||
return _ModelArmorResponseRoot.model_validate(json_response).root
|
||||
|
||||
def sanitize_file_prompt(self, file_bytes: bytes, file_type: str, source: str = "user_prompt") -> dict:
|
||||
"""
|
||||
|
|
@ -378,7 +430,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
else:
|
||||
return {"modelResponseData": {"byteItem": {"byteDataType": file_type, "byteData": base64_data}}}
|
||||
|
||||
def _should_block_content(self, armor_response: Mapping[str, object], allow_sanitization: bool = False) -> bool:
|
||||
def _should_block_content(self, armor_response: _ModelArmorResponse, allow_sanitization: bool = False) -> bool:
|
||||
"""Check if Model Armor response indicates content should be blocked, including both inspectResult and deidentifyResult."""
|
||||
for filt in self._filter_result_items(armor_response):
|
||||
# Check RAI, PI/Jailbreak, Malicious URI, CSAM, Virus scan as before
|
||||
|
|
@ -404,7 +456,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
# Fallback dict code removed; all cases handled above
|
||||
return False
|
||||
|
||||
def _get_sanitized_content(self, armor_response: Mapping[str, Any]) -> str | None:
|
||||
def _get_sanitized_content(self, armor_response: _ModelArmorResponse) -> str | None:
|
||||
"""
|
||||
Get the sanitized content from a Model Armor response, if available.
|
||||
Looks for sanitized text in deidentifyResult, and falls back to root-level fields if not found.
|
||||
|
|
@ -434,7 +486,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
return armor_response.get("sanitizedText") or armor_response.get("text")
|
||||
|
||||
@staticmethod
|
||||
def _filter_result_items(armor_response: Mapping[str, Any]) -> Sequence[Any]:
|
||||
def _filter_result_items(armor_response: _ModelArmorResponse) -> Sequence[_Filter]:
|
||||
"""Every filter result in a scan response.
|
||||
|
||||
filterResults is a dict of named filters on most templates and a list on some, so both
|
||||
|
|
@ -447,7 +499,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
return filter_results
|
||||
return []
|
||||
|
||||
def _has_deidentify_match(self, armor_response: Mapping[str, object]) -> bool:
|
||||
def _has_deidentify_match(self, armor_response: _ModelArmorResponse) -> bool:
|
||||
"""Whether an SDP de-identify filter matched, i.e. Model Armor owes this response a redaction."""
|
||||
for filter_entry in self._filter_result_items(armor_response):
|
||||
sdp = filter_entry.get("sdpFilterResult")
|
||||
|
|
@ -457,7 +509,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
|
||||
def _resolve_streaming_outcome(
|
||||
self,
|
||||
armor_response: Mapping[str, object],
|
||||
armor_response: _ModelArmorResponse,
|
||||
assembled_response: object,
|
||||
content: str,
|
||||
) -> tuple[bool, str | None]:
|
||||
|
|
@ -489,7 +541,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
|
|||
return False, sanitized_content
|
||||
|
||||
@staticmethod
|
||||
def _append_armor_response(existing: object, armor_response: Mapping[str, object]) -> object:
|
||||
def _append_armor_response(existing: object, armor_response: object) -> object:
|
||||
"""Accumulate scan responses so a later text scan does not drop an earlier file scan.
|
||||
|
||||
Returns the single response on its own (backward compatible) and a list once a request
|
||||
|
|
|
|||
|
|
@ -245,7 +245,7 @@ class _TransactionFactory(Protocol):
|
|||
|
||||
|
||||
class _ModelTransactionClient(LiteLLMBaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True)
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True, hide_input_in_errors=True)
|
||||
|
||||
tx: _TransactionFactory
|
||||
|
||||
|
|
@ -1856,8 +1856,8 @@ async def delete_team_models(
|
|||
Returns the model_ids that were deleted.
|
||||
"""
|
||||
deleted_model_ids: Final[list[str]] = []
|
||||
async with prisma_client.db.tx() as tx_ctx:
|
||||
tx: Final[_TxModelTables] = tx_ctx
|
||||
transaction_client: Final = _ModelTransactionClient.model_validate(prisma_client.db)
|
||||
async with transaction_client.tx() as tx:
|
||||
for team_id in team_ids:
|
||||
rows = await _get_team_deployments(team_id, prisma_client, table=tx.litellm_proxymodeltable)
|
||||
model_ids = [row.model_id for row in rows]
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ if TYPE_CHECKING:
|
|||
import jwt
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, Response, status
|
||||
from fastapi.responses import RedirectResponse
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._internal_context import with_service_target
|
||||
|
|
@ -237,6 +237,12 @@ class CliSsoTeamDetail(LiteLLMBaseModel):
|
|||
|
||||
|
||||
_CLI_SSO_TEAM_DETAILS_ADAPTER: Final = TypeAdapter(tuple[CliSsoTeamDetail, ...])
|
||||
_CLI_SSO_FLOW_JSON_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(
|
||||
dict[str, object], config=ConfigDict(strict=True, hide_input_in_errors=True)
|
||||
)
|
||||
_SSO_USERINFO_RESPONSE_ADAPTER: Final[TypeAdapter[dict[str, object] | None]] = TypeAdapter(
|
||||
dict[str, object] | None, config=ConfigDict(strict=True, hide_input_in_errors=True)
|
||||
)
|
||||
_TEAMLESS_CLI_SSO_TEAM_DETAIL: Final = CliSsoTeamDetail(team_models=())
|
||||
|
||||
|
||||
|
|
@ -371,7 +377,7 @@ def _get_cli_sso_flow_or_raise(login_id: str | None, cache: DualCache) -> dict:
|
|||
flow = _read_cli_sso_flow(cache, _get_cli_sso_flow_cache_key(cast(str, login_id)))
|
||||
if isinstance(flow, str):
|
||||
try:
|
||||
flow = _as_object(json.loads(flow))
|
||||
flow = _CLI_SSO_FLOW_JSON_ADAPTER.validate_python(json.loads(flow))
|
||||
except ValueError:
|
||||
flow = None
|
||||
if not isinstance(flow, dict) or "poll_secret_hash" not in flow:
|
||||
|
|
@ -4199,7 +4205,7 @@ class SSOAuthenticationHandler:
|
|||
)
|
||||
if resp.status_code == 200:
|
||||
try:
|
||||
userinfo_raw: Final[dict[str, object] | None] = resp.json()
|
||||
userinfo_raw: Final = _SSO_USERINFO_RESPONSE_ADAPTER.validate_python(resp.json())
|
||||
if not userinfo_raw:
|
||||
# JSON null (None) or empty dict ({}) — no identity claims.
|
||||
# Treat as failure so id_token fallback can be attempted.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue