mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
Merge remote-tracking branch 'origin/litellm_internal_staging' into litellm_registry_audit_bedrock_sol_anthropic_1hr
This commit is contained in:
commit
9300018414
269 changed files with 16477 additions and 2342 deletions
23
.github/pull_request_template.md
vendored
23
.github/pull_request_template.md
vendored
|
|
@ -1,7 +1,10 @@
|
|||
<!-- The whole description's target audience is humans, not AI agents: write it in plain, simple,
|
||||
everyday engineering language, extremely parsable and readable at a glance. This goes double for
|
||||
the TLDR, User Flow, and Caveats sections -->
|
||||
|
||||
## TLDR
|
||||
|
||||
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max
|
||||
This section must be extremely human parsable, comprehensible, and readable: its target audience is humans, not AI agents -->
|
||||
<!-- Fill in the bullets below and keep each one short and concrete: one line per bullet, roughly 10 words max -->
|
||||
|
||||
Problem this solves:
|
||||
|
||||
|
|
@ -110,8 +113,20 @@ If you're seeing a delay in your PR being merged, ping the LiteLLM Team on [Slac
|
|||
|
||||
## Caveats (if any)
|
||||
|
||||
<!-- Short bullet points, just like the TLDR: one line per bullet, roughly 10 words max
|
||||
<!-- Group caveats under severity subheadings (### Severe, ### High, ### Medium, ### Low), with
|
||||
short bullet points inside each, just like the TLDR: one line per bullet, roughly 10 words max
|
||||
Call out known limitations, follow-up work, or anything a reviewer should watch out for
|
||||
Include only the tiers that have caveats; drop the empty ones
|
||||
- Severe: inherent to what the PR deliberately ships, there even when the code works as intended:
|
||||
it can degrade or take down a running deployment (e.g. a slow or table-locking boot migration),
|
||||
rewrite data by design, break an existing workflow on purpose, or change auth behavior. An
|
||||
operator must plan around it before rollout
|
||||
- High: an unintended hole: a correctness, security, data-loss, or backward-compatibility bug,
|
||||
unsafe to ship as is
|
||||
- Medium: a real gap someone can hit, but with a workaround or a narrow blast radius
|
||||
- Low: anything else worth noting: naming, cleanup, an edge case nobody hits
|
||||
Nest bullets as deep as helps: hierarchy beats one long line when it makes things clearer to a
|
||||
human reader
|
||||
Leave this section empty if there are none -->
|
||||
|
||||
## QA runbook
|
||||
|
|
@ -134,6 +149,6 @@ Example checklists:
|
|||
- [ ] Sanity check: this test makes sense to add and is not hand-wavey (e.g., assert actual expected spend instead of just spend > 0) or potentially flaky
|
||||
-->
|
||||
|
||||
### Final Attestation
|
||||
## Final Attestation
|
||||
|
||||
- [ ] The tests check the right things, including the edge cases, and regressions in the respective real-world customer use-cases are not possible after this PR
|
||||
|
|
|
|||
1
.github/workflows/test-unit.yml
vendored
1
.github/workflows/test-unit.yml
vendored
|
|
@ -164,6 +164,7 @@ jobs:
|
|||
tests/test_litellm/proxy/public_endpoints
|
||||
tests/test_litellm/proxy/prompts
|
||||
tests/test_litellm/proxy/rag_endpoints
|
||||
tests/test_litellm/proxy/rerank_endpoints
|
||||
tests/test_litellm/proxy/realtime_endpoints
|
||||
tests/test_litellm/proxy/ui_crud_endpoints
|
||||
tests/test_litellm/proxy/config_resolvers
|
||||
|
|
|
|||
|
|
@ -37,13 +37,14 @@ If you're resolving a linear ticket, in the "## Linear ticket" section of the PR
|
|||
|
||||
Never use `pytest` commands or the like as "Screenshots / Proof of Fix". We prefer curl'ing a live proxy instance running on localhost:4000 (I like to run it with `python litellm/proxy/proxy_cli.py --config litellm/proxy/dev_config.yaml --detailed_debug --reload --use_v2_migration_resolver 2>&1 | tee litellm.log`; the Admin UI dev server is `npm run dev` in `ui/litellm-dashboard`, served on port 3000) and showing both the command run and the output. Also, it should hit real LLM provider APIs, not mocks, and cost real $$$ because that is the most realistic test. The proof of fix should be exactly what the end user / customer would see / do. The run logs in PR #27703 is a prime example of how to do it (not a huge fan of using a python test script that future me and the team will have no visibility into; I prefer just curl commands or a short list of bash commands (e.g., using `for`)). If it's a UI thing, just tell me which URLs to go to (e.g., http://localhost:4000/ui/?page=logs), where to click, what fields to fill out, etc. along with the other commands to run in an ordered list, and I'll do it myself and post the screenshots after you make the PR
|
||||
|
||||
If you ever make public-facing PR descriptions, comments, issues, commit messages, etc., always follow these guidelines to sound less AI-y:
|
||||
If you ever write any human-facing text (pull requests, issues, commit messages, discussion posts, github comments, release notes, docs, etc.), always follow these guidelines to sound less AI-y:
|
||||
- don't use emojis
|
||||
- don't use "—". Instead, reach for ",", ".", conjunction words, ":", ";", etc. in descending order of preference: vary among them, weighted toward the front of the list, and skip "," where it would cause a comma splice or the sentence is getting long. Overusing any one of them, ";" especially, also feels AI-y. A word cap does not penalize you for adding more sentences: when writing under tight word budgets, prefer a period split or a conjunction over ";", and keep to at most one ";" per message
|
||||
- don't use the pattern "It's not X, it's Y", "You're not X, you're Y", etc.
|
||||
- don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
|
||||
- unless explicitly asked, don't use bulleted or numbered lists unless it would be nonsensical not to. Instead, prefer prose
|
||||
- don't add a trailing "." at the end of paragraphs (just like this file). That means every paragraph, not just the last one (of the markdown file, PR description, GitHub comment, etc.). Rule of thumb: if you're adding new line(s) before the next sentence, don't add a "."
|
||||
- don't use →. Instead, prefer not to use arrows, and if need be, use -> instead
|
||||
- use plain, simple, everyday engineering language: the common phrase engineers actually say over rare compact phrasing, in grammatically complete sentences. When explicitly asked to use bullets or ordered lists and structure legitimately helps the reader, prefer nested bullets (any depth is fine) over dense lines in a flat structure
|
||||
|
||||
Don't hesitate to use values in .env to get needed API keys and other secrets, as long as you never add them to conversation history, commit them, or include them in GitHub issues / PRs
|
||||
|
||||
|
|
|
|||
|
|
@ -1,18 +1,18 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"limit": 19949
|
||||
"limit": 18505
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2566
|
||||
"limit": 2564
|
||||
},
|
||||
"reportAssignmentType": {
|
||||
"limit": 320
|
||||
},
|
||||
"reportAttributeAccessIssue": {
|
||||
"limit": 488
|
||||
"limit": 483
|
||||
},
|
||||
"reportCallIssue": {
|
||||
"limit": 114
|
||||
"limit": 113
|
||||
},
|
||||
"reportConstantRedefinition": {
|
||||
"limit": 40
|
||||
|
|
@ -24,7 +24,7 @@
|
|||
"limit": 19
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 6049
|
||||
"limit": 5976
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 7
|
||||
|
|
@ -45,7 +45,7 @@
|
|||
"limit": 35
|
||||
},
|
||||
"reportInvalidTypeForm": {
|
||||
"limit": 35
|
||||
"limit": 34
|
||||
},
|
||||
"reportInvalidTypeVarUse": {
|
||||
"limit": 2
|
||||
|
|
@ -54,10 +54,10 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportMissingParameterType": {
|
||||
"limit": 5661
|
||||
"limit": 5659
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15555
|
||||
"limit": 15504
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 40
|
||||
|
|
@ -72,7 +72,7 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportOptionalMemberAccess": {
|
||||
"limit": 1061
|
||||
"limit": 1058
|
||||
},
|
||||
"reportOptionalOperand": {
|
||||
"limit": 0
|
||||
|
|
@ -84,7 +84,7 @@
|
|||
"limit": 56
|
||||
},
|
||||
"reportPrivateUsage": {
|
||||
"limit": 1822
|
||||
"limit": 1810
|
||||
},
|
||||
"reportRedeclaration": {
|
||||
"limit": 8
|
||||
|
|
@ -99,31 +99,31 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 44655
|
||||
"limit": 44530
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 109
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 39009
|
||||
"limit": 38828
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 19883
|
||||
"limit": 19847
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 30569
|
||||
"limit": 30386
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 117
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"limit": 699
|
||||
"limit": 697
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"limit": 5
|
||||
},
|
||||
"reportUnnecessaryIsInstance": {
|
||||
"limit": 836
|
||||
"limit": 833
|
||||
},
|
||||
"reportUntypedBaseClass": {
|
||||
"limit": 0
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ from litellm.constants import (
|
|||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.integrations.prometheus import PrometheusLogger
|
||||
from litellm.proxy._types import LiteLLM_ManagedObjectTable
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
|
|
@ -351,7 +353,7 @@ class CheckBatchCost:
|
|||
return isinstance(error, (NotFoundError, openai.NotFoundError)) and output_file_id in str(error)
|
||||
|
||||
async def _finalize_unbilled_terminal_job(
|
||||
self, job: "LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
|
||||
self, job: "prisma_models.LiteLLM_ManagedObjectTable", response: "LiteLLMBatch"
|
||||
) -> None:
|
||||
"""Persist a terminal batch that has nothing billable, converting any raw
|
||||
provider file ids to managed ids, and take it out of the poll page."""
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.59"
|
||||
version = "0.1.60"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.9"
|
||||
|
|
@ -26,7 +26,7 @@ required-version = ">=0.10.9"
|
|||
module-root = ""
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.59"
|
||||
version = "0.1.60"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
"../pyproject.toml:litellm-enterprise==",
|
||||
|
|
|
|||
|
|
@ -199,6 +199,7 @@ standard_logging_payload_excluded_fields: Optional[List[str]] = (
|
|||
None # Fields to exclude from StandardLoggingPayload before callbacks receive it
|
||||
)
|
||||
log_raw_request_response: bool = False
|
||||
log_client_error_tracebacks: bool = False
|
||||
request_correlation_in_logs: bool = False
|
||||
redact_messages_in_exceptions: Optional[bool] = False
|
||||
redact_user_api_key_info: Optional[bool] = False
|
||||
|
|
@ -463,6 +464,11 @@ prometheus_metrics_config: Optional[List] = None
|
|||
prometheus_exclude_metrics: Optional[List[str]] = None
|
||||
prometheus_exclude_labels: Optional[List[str]] = None
|
||||
prometheus_emit_stream_label: bool = False
|
||||
prometheus_deployment_and_latency_caller_identity: Literal[
|
||||
"api_key_alias",
|
||||
"user_email",
|
||||
"both",
|
||||
] = "api_key_alias"
|
||||
# Opt-in: emit `rate_limit_category` and `rate_limit_type` labels on
|
||||
# `litellm_proxy_failed_requests_metric`. Off by default to preserve the
|
||||
# pre-unification label set so existing dashboards / recording rules keyed on
|
||||
|
|
@ -1628,6 +1634,9 @@ if TYPE_CHECKING:
|
|||
AmazonMantleMessagesConfig as AmazonMantleMessagesConfig,
|
||||
)
|
||||
from .llms.together_ai.chat import TogetherAIConfig as TogetherAIConfig
|
||||
from .llms.together_ai.chat.transformation import (
|
||||
TogetherAIChatConfig as TogetherAIChatConfig,
|
||||
)
|
||||
from .llms.nlp_cloud.chat.handler import NLPCloudConfig as NLPCloudConfig
|
||||
from .llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig as VertexGeminiConfig,
|
||||
|
|
|
|||
|
|
@ -177,6 +177,7 @@ LLM_CONFIG_NAMES: Final = (
|
|||
"AmazonAnthropicClaudeMessagesConfig",
|
||||
"AmazonMantleMessagesConfig",
|
||||
"TogetherAIConfig",
|
||||
"TogetherAIChatConfig",
|
||||
"NLPCloudConfig",
|
||||
"VertexGeminiConfig",
|
||||
"GoogleAIStudioGeminiConfig",
|
||||
|
|
@ -741,6 +742,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
|
|||
"AmazonMantleMessagesConfig",
|
||||
),
|
||||
"TogetherAIConfig": (".llms.together_ai.chat", "TogetherAIConfig"),
|
||||
"TogetherAIChatConfig": (
|
||||
".llms.together_ai.chat.transformation",
|
||||
"TogetherAIChatConfig",
|
||||
),
|
||||
"NLPCloudConfig": (".llms.nlp_cloud.chat.handler", "NLPCloudConfig"),
|
||||
"VertexGeminiConfig": (
|
||||
".llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini",
|
||||
|
|
|
|||
|
|
@ -18,6 +18,14 @@ already does when one of its pooled connections errors), leaving every other nod
|
|||
connections untouched. Every other branch (MOVED, ASK, CLUSTERDOWN, slot-not-covered,
|
||||
retry-exhaustion) is unchanged from upstream, since those already carry real evidence the
|
||||
topology changed.
|
||||
|
||||
redis-py 8.x fixed this upstream with gentler machinery than this override's
|
||||
``node.disconnect()`` (which also kills connections other coroutines are mid-operation
|
||||
on, so one timeout cascades into a reconnect storm and, with TLS, a fresh handshake per
|
||||
killed connection): it marks in-use connections for reconnect only after their current
|
||||
operation completes, disconnects only the idle pooled ones, and defers reinitialization
|
||||
to the outer retry loop. When the installed ``ClusterNode`` has that per-connection
|
||||
recovery API, the factory returns the base ``RedisCluster`` unmodified.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
|
|
@ -72,8 +80,16 @@ class _ClusterAttrs(Protocol):
|
|||
_VERIFIED_REDIS_VERSIONS: Final = frozenset({"5.3.1"})
|
||||
|
||||
|
||||
def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]:
|
||||
"""Builds the ``RedisCluster`` subclass with the per-node isolation fix.
|
||||
def get_litellm_async_redis_cluster_class(
|
||||
cluster_node_class: type | None = None,
|
||||
) -> type["_AsyncRedisClusterType"]:
|
||||
"""Returns the base ``RedisCluster`` when the installed redis-py already recovers a
|
||||
node-level connection error per-connection (8.x+), else builds the ``RedisCluster``
|
||||
subclass with the per-node isolation fix for older versions whose upstream branch
|
||||
tears down the whole cluster client.
|
||||
|
||||
``cluster_node_class`` exists for dependency injection in tests; production callers
|
||||
leave it unset and the installed ``ClusterNode`` is used.
|
||||
|
||||
Imported lazily because this module is reachable from a base ``import litellm`` while
|
||||
redis is not a base dependency. Cheap to call repeatedly: the underlying redis
|
||||
|
|
@ -81,7 +97,10 @@ def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]:
|
|||
"""
|
||||
import redis
|
||||
from redis.asyncio.cluster import (
|
||||
RedisCluster as _BaseAsyncRedisCluster, # pyright: ignore[reportUnknownVariableType] # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin
|
||||
ClusterNode as _AsyncClusterNode, # pyright: ignore[reportUnknownVariableType] # redis-py ships no resolvable stub for this class under the repo's current (stale) types-redis pin
|
||||
)
|
||||
from redis.asyncio.cluster import (
|
||||
RedisCluster as _BaseAsyncRedisCluster, # pyright: ignore[reportUnknownVariableType] # same stale-stub gap as the import above
|
||||
)
|
||||
from redis.cluster import get_node_name
|
||||
from redis.commands import READ_COMMANDS
|
||||
|
|
@ -98,6 +117,15 @@ def get_litellm_async_redis_cluster_class() -> type["_AsyncRedisClusterType"]:
|
|||
from redis.exceptions import ConnectionError as _RedisConnectionError
|
||||
from redis.exceptions import TimeoutError as _RedisTimeoutError
|
||||
|
||||
node_class: Final = cluster_node_class if cluster_node_class is not None else _AsyncClusterNode
|
||||
if hasattr(node_class, "update_active_connections_for_reconnect"):
|
||||
verbose_logger.debug(
|
||||
"redis-py %s recovers a node-level connection error per-connection upstream; "
|
||||
"using the base RedisCluster without litellm's node-isolation override.",
|
||||
redis.__version__,
|
||||
)
|
||||
return _BaseAsyncRedisCluster
|
||||
|
||||
if redis.__version__ not in _VERIFIED_REDIS_VERSIONS:
|
||||
verbose_logger.warning(
|
||||
"redis-py %s is not in the set this cluster-teardown-storm fix was verified "
|
||||
|
|
|
|||
|
|
@ -49,6 +49,8 @@ LITELLM_MAX_STREAMING_DURATION_SECONDS: Final = (
|
|||
# Set to 0 to disable truncation.
|
||||
MAX_BASE64_LENGTH_FOR_LOGGING: Final = int(os.getenv("MAX_BASE64_LENGTH_FOR_LOGGING", 64))
|
||||
REDACTED_BY_LITELLM: Final = "redacted-by-litellm"
|
||||
# in-memory stand-in handed to provider converters for redacted arguments; never stored
|
||||
REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER: Final = "{}"
|
||||
|
||||
MAX_STRING_LENGTH_STDOUT_LOG: Final = get_env_int("MAX_STRING_LENGTH_STDOUT_LOG", 4096)
|
||||
|
||||
|
|
@ -147,6 +149,7 @@ LITELLM_UI_ALLOW_HEADERS: Final = [
|
|||
"x-litellm-adaptive-router-model",
|
||||
"x-litellm-applied-guardrails",
|
||||
"x-litellm-guardrail-scan-id",
|
||||
"x-litellm-cache-key",
|
||||
]
|
||||
|
||||
# Gemini model-specific minimal thinking budget constants
|
||||
|
|
@ -461,6 +464,8 @@ CONNECTION_ERROR_PATTERNS: Final[list[str]] = [
|
|||
]
|
||||
STREAM_SSE_DONE_STRING: Final[str] = "[DONE]"
|
||||
STREAM_SSE_DATA_PREFIX: Final[str] = "data: "
|
||||
STREAM_SSE_KEEPALIVE_PING_CHUNK: Final[str] = 'event: ping\ndata: {"type": "ping"}\n\n'
|
||||
STREAM_SSE_KEEPALIVE_PING_BYTES: Final[bytes] = STREAM_SSE_KEEPALIVE_PING_CHUNK.encode("utf-8")
|
||||
### SPEND TRACKING ###
|
||||
DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND: Final = float(
|
||||
os.getenv("DEFAULT_REPLICATE_GPU_PRICE_PER_SECOND", 0.001400)
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ litellm_content_retrieve tool calls server-side via the typed agentic loop plan.
|
|||
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Final, cast
|
||||
from typing import Any, ClassVar, Final, cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.compression import compress
|
||||
|
|
@ -72,6 +72,8 @@ class CompressionInterceptionLogger(CustomLogger):
|
|||
4. Build typed rerun plan with tool_result blocks from the compressed cache.
|
||||
"""
|
||||
|
||||
server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({LITELLM_CONTENT_RETRIEVE_TOOL_NAME})
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
enabled: bool = True,
|
||||
|
|
|
|||
|
|
@ -2,8 +2,8 @@
|
|||
# On success, logs events to Promptlayer
|
||||
import re
|
||||
import traceback
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional
|
||||
from collections.abc import AsyncGenerator, Mapping
|
||||
from typing import TYPE_CHECKING, Any, ClassVar, Final, Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
|
@ -60,6 +60,7 @@ _BASE64_INLINE_PATTERN: Final = re.compile(
|
|||
|
||||
class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class
|
||||
# Class variables or attributes
|
||||
server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset()
|
||||
|
||||
enforces_request_content: bool = False
|
||||
"""
|
||||
|
|
@ -292,6 +293,54 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
|
|||
Allow modifying / reviewing the response just after it's received from the deployment.
|
||||
"""
|
||||
|
||||
async def async_post_call_failure_deployment_hook(
|
||||
self,
|
||||
request_data: Mapping[str, Any],
|
||||
exception: Exception,
|
||||
call_type: CallTypes | None,
|
||||
fallback_depth: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Called once per failed deployment attempt - attempt 1, every retry, and
|
||||
every fallback chain step - because the router re-invokes the wrapped
|
||||
function on each attempt, re-entering this hook's call site fresh
|
||||
every time.
|
||||
|
||||
This is a DEPLOYMENT-LEVEL signal, distinct from the REQUEST-LEVEL
|
||||
``async_log_failure_event``, which fires once per logical client
|
||||
request behind a dedup gate. ``request_data`` is mostly this
|
||||
attempt's own kwargs, with one exception: it omits
|
||||
``attempted_targets``, the router's own bookkeeping of which fallback
|
||||
targets this request has already tried, since that one object *is*
|
||||
shared by reference across every hop of the live fallback walk.
|
||||
|
||||
Pairs with ``async_pre_call_deployment_hook`` and
|
||||
``async_post_call_success_deployment_hook`` to complete the
|
||||
pre-call/success/failure lifecycle for a single deployment attempt.
|
||||
|
||||
``fallback_depth`` is best-effort: ``None`` on the first attempt and on
|
||||
any call made without a ``Router`` (a bare SDK call has no fallback
|
||||
chain to be at a depth in), ``1`` on the first fallback hop, ``2`` on
|
||||
the second, and so on. It reflects ``Router``'s own internal fallback
|
||||
bookkeeping (``kwargs["fallback_depth"]``), not a value this hook
|
||||
computes or guarantees the shape of across versions. It tracks
|
||||
fallback hops only, not retries within the same model group - a
|
||||
retry-only failure (no fallback yet) also reports ``None``. If an
|
||||
override predates this field it's simply never passed, rather than
|
||||
raising - safe to leave off an override written before it existed.
|
||||
|
||||
``exception`` is a same-class snapshot, not the exact object about to
|
||||
be re-raised to the real caller: read it freely, but setting an
|
||||
attribute on it (e.g. ``status_code``) has no effect on what the
|
||||
caller actually receives.
|
||||
|
||||
Default: no-op. Opt in by overriding. Keep overrides fast - this
|
||||
runs on the request's exception path, so a slow implementation
|
||||
delays error propagation to the caller. The reported failure
|
||||
duration is captured before this hook runs, so a slow override
|
||||
doesn't inflate that metric, but the caller still waits for it.
|
||||
"""
|
||||
|
||||
async def async_post_call_streaming_deployment_hook(
|
||||
self,
|
||||
request_data: dict,
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from litellm.integrations.langfuse.langfuse_otel_attributes import (
|
|||
LangfuseLLMObsOTELAttributes,
|
||||
)
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
|
||||
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
|
||||
from litellm.types.integrations.langfuse_otel import (
|
||||
LangfuseSpanAttributes,
|
||||
)
|
||||
|
|
@ -197,7 +198,11 @@ class LangfuseOtelLogger(OpenTelemetry):
|
|||
)
|
||||
elif item_type == "function_call":
|
||||
arguments_str = getattr(item, "arguments", "{}")
|
||||
arguments_obj = json.loads(arguments_str) if isinstance(arguments_str, str) else arguments_str
|
||||
arguments_obj = (
|
||||
safe_json_loads(arguments_str, default={})
|
||||
if isinstance(arguments_str, str)
|
||||
else arguments_str
|
||||
)
|
||||
langfuse_tool_call = {
|
||||
"id": getattr(item, "id", ""),
|
||||
"name": getattr(item, "name", ""),
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ from litellm.types.integrations.prometheus import *
|
|||
from litellm.types.integrations.prometheus import (
|
||||
_sanitize_prometheus_label_name,
|
||||
_sanitize_prometheus_label_value,
|
||||
validate_prometheus_deployment_and_latency_caller_identity,
|
||||
)
|
||||
from litellm.types.utils import (
|
||||
StandardLoggingGuardrailInformation,
|
||||
|
|
@ -96,7 +97,10 @@ class _PaginatedPrismaTable(Protocol[_TableRowT]):
|
|||
|
||||
def _paginated_table(repository: BaseRepository[_TableRowT]) -> _PaginatedPrismaTable[_TableRowT]:
|
||||
"""View a repository's prisma table through the pagination surface budget metrics need."""
|
||||
return repository.table
|
||||
return cast(
|
||||
_PaginatedPrismaTable[_TableRowT],
|
||||
repository.table, # cast-ok: prisma rows carry the budget columns the domain model declares
|
||||
)
|
||||
|
||||
|
||||
class _OrgBudgetRow(Protocol):
|
||||
|
|
@ -172,6 +176,11 @@ class PrometheusLogger(CustomLogger):
|
|||
try:
|
||||
from prometheus_client import Counter, Gauge, Histogram
|
||||
|
||||
# Validate the caller-identity mode before any collector registers so an
|
||||
# invalid value cannot leave partially-registered metrics behind in the
|
||||
# process-global registry.
|
||||
validate_prometheus_deployment_and_latency_caller_identity()
|
||||
|
||||
# Always initialize label_filters, even for non-premium users
|
||||
self.label_filters = self._parse_prometheus_config()
|
||||
|
||||
|
|
@ -2462,6 +2471,7 @@ class PrometheusLogger(CustomLogger):
|
|||
else:
|
||||
_metadata = {
|
||||
"user_api_key_alias": getattr(_metadata_raw, "user_api_key_alias", None),
|
||||
"user_api_key_user_email": getattr(_metadata_raw, "user_api_key_user_email", None),
|
||||
"user_api_key_team_id": getattr(_metadata_raw, "user_api_key_team_id", None),
|
||||
"user_api_key_team_alias": getattr(_metadata_raw, "user_api_key_team_alias", None),
|
||||
"user_api_key_hash": getattr(_metadata_raw, "user_api_key_hash", None),
|
||||
|
|
@ -2484,6 +2494,17 @@ class PrometheusLogger(CustomLogger):
|
|||
return getattr(user_api_key_auth, "key_alias", None)
|
||||
return None
|
||||
|
||||
def _get_user_email() -> str | None:
|
||||
val = _metadata.get("user_api_key_user_email")
|
||||
if val is not None:
|
||||
return val
|
||||
val = _litellm_params_metadata.get("user_api_key_user_email")
|
||||
if val is not None:
|
||||
return val
|
||||
if user_api_key_auth is not None:
|
||||
return self._safe_get(user_api_key_auth, "user_email")
|
||||
return None
|
||||
|
||||
def _get_team_id() -> str | None:
|
||||
val = _metadata.get("user_api_key_team_id")
|
||||
if val is not None:
|
||||
|
|
@ -2519,6 +2540,7 @@ class PrometheusLogger(CustomLogger):
|
|||
|
||||
return {
|
||||
"api_key_alias": _get_api_key_alias(),
|
||||
"user_email": _get_user_email(),
|
||||
"team": _get_team_id(),
|
||||
"team_alias": _get_team_alias(),
|
||||
"hashed_api_key": _get_hashed_api_key(),
|
||||
|
|
@ -2576,6 +2598,7 @@ class PrometheusLogger(CustomLogger):
|
|||
_metadata: Final = standard_logging_payload.get("metadata", {}) or {}
|
||||
hashed_api_key: Final = fallback_values.get("hashed_api_key") or _metadata.get("user_api_key_hash")
|
||||
api_key_alias: Final = fallback_values.get("api_key_alias") or _metadata.get("user_api_key_alias")
|
||||
user_email: Final = fallback_values.get("user_email")
|
||||
team: Final = fallback_values.get("team") or _metadata.get("user_api_key_team_id")
|
||||
team_alias: Final = fallback_values.get("team_alias") or _metadata.get("user_api_key_team_alias")
|
||||
client_ip: Final = fallback_values.get("client_ip") or _metadata.get("requester_ip_address")
|
||||
|
|
@ -2616,6 +2639,7 @@ class PrometheusLogger(CustomLogger):
|
|||
requested_model=label_requested_model,
|
||||
hashed_api_key=hashed_api_key,
|
||||
api_key_alias=api_key_alias,
|
||||
user_email=user_email,
|
||||
team=team,
|
||||
team_alias=team_alias,
|
||||
tags=standard_logging_payload.get("request_tags", []),
|
||||
|
|
|
|||
|
|
@ -58,6 +58,67 @@ def safe_divide(
|
|||
return numerator / denominator
|
||||
|
||||
|
||||
def _is_litellm_limit_rejection(exception: BaseException) -> bool:
|
||||
from litellm.exceptions import RateLimitErrorCategory
|
||||
|
||||
litellm_limit_categories: Final = frozenset(
|
||||
(RateLimitErrorCategory.LITELLM_RATE_LIMIT.value, RateLimitErrorCategory.LITELLM_BATCH_RATE_LIMIT.value)
|
||||
)
|
||||
return getattr(exception, "category", None) in litellm_limit_categories
|
||||
|
||||
|
||||
def _is_proxy_rejection(exception: BaseException) -> bool:
|
||||
if _is_litellm_limit_rejection(exception):
|
||||
return True
|
||||
try:
|
||||
from starlette.exceptions import HTTPException
|
||||
except ImportError:
|
||||
return False
|
||||
return isinstance(exception, HTTPException)
|
||||
|
||||
|
||||
def _is_provider_originated(exception: BaseException) -> bool:
|
||||
if _is_proxy_rejection(exception):
|
||||
return False
|
||||
if getattr(exception, "llm_provider", None):
|
||||
return True
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
|
||||
return isinstance(exception, BaseLLMException)
|
||||
|
||||
|
||||
def is_expected_client_error(exception: BaseException | None) -> bool:
|
||||
"""
|
||||
True when the proxy itself rejected the request with an HTTP 4xx before any
|
||||
provider call (bad key, budget, unknown model, guardrail). A 4xx returned by
|
||||
a provider is an upstream or deployment problem, so it is never an expected
|
||||
client error and keeps its traceback: a mapped litellm exception carries
|
||||
``llm_provider``, and the raw ``BaseLLMException`` that provider handlers
|
||||
raise before mapping (the /v1/messages route surfaces it as-is) is one too.
|
||||
The proxy's own limiters raise ``HTTPException`` subclasses that also carry
|
||||
an ``llm_provider``, so any ``HTTPException`` stays a proxy rejection, and
|
||||
so does any exception whose unified rate-limit ``category`` names litellm's
|
||||
own limiter (``BudgetExceededError`` is a plain ``Exception`` that the auth
|
||||
handler decorates with the requested model's provider).
|
||||
|
||||
ProxyException stores the status on .code (as a str), HTTPException and
|
||||
litellm exceptions on .status_code.
|
||||
"""
|
||||
if exception is None:
|
||||
return False
|
||||
if _is_provider_originated(exception):
|
||||
return False
|
||||
code: Final[object] = getattr(exception, "code", None)
|
||||
status_code: Final[object] = code if code is not None else getattr(exception, "status_code", None)
|
||||
if status_code is None or isinstance(status_code, bool):
|
||||
return False
|
||||
try:
|
||||
status: Final = int(str(status_code))
|
||||
except ValueError:
|
||||
return False
|
||||
return 400 <= status < 500
|
||||
|
||||
|
||||
def coerce_token_limit(value: object) -> int | None:
|
||||
"""
|
||||
Coerce a max_input_tokens / max_output_tokens value to an int, treating a
|
||||
|
|
|
|||
|
|
@ -550,6 +550,13 @@ def _map_anthropic_exception(
|
|||
llm_provider="anthropic",
|
||||
model=model,
|
||||
)
|
||||
elif original_exception.status_code == 403:
|
||||
raise PermissionDeniedError(
|
||||
message=f"AnthropicException - {error_str}",
|
||||
llm_provider="anthropic",
|
||||
model=model,
|
||||
response=original_exception.response,
|
||||
)
|
||||
elif original_exception.status_code == 400 or original_exception.status_code == 413:
|
||||
raise BadRequestError(
|
||||
message=f"AnthropicException - {error_str}",
|
||||
|
|
@ -755,12 +762,19 @@ def _map_openai_like_exception(
|
|||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
)
|
||||
elif original_exception.status_code == 401 or original_exception.status_code == 403:
|
||||
elif original_exception.status_code == 401:
|
||||
raise AuthenticationError(
|
||||
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
)
|
||||
elif original_exception.status_code == 403:
|
||||
raise PermissionDeniedError(
|
||||
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=_response_or_stub(original_exception, status_code=403),
|
||||
)
|
||||
elif original_exception.status_code == 400:
|
||||
raise BadRequestError(
|
||||
message=f"{custom_llm_provider.capitalize()}Exception - {original_exception.message}",
|
||||
|
|
@ -2187,6 +2201,120 @@ def _map_openrouter_exception(
|
|||
)
|
||||
|
||||
|
||||
def _response_or_stub(original_exception: _ProviderHTTPException, status_code: int) -> httpx.Response:
|
||||
response: Final = original_exception.response if hasattr(original_exception, "response") else None
|
||||
if response is not None:
|
||||
return response
|
||||
return httpx.Response(
|
||||
status_code=status_code, request=httpx.Request(method="POST", url="https://docs.litellm.ai/docs")
|
||||
)
|
||||
|
||||
|
||||
def _map_exception_by_status(
|
||||
*,
|
||||
model: str,
|
||||
original_exception: _ProviderHTTPException,
|
||||
custom_llm_provider: str,
|
||||
error_str: str,
|
||||
exception_provider: str,
|
||||
extra_information: str,
|
||||
) -> None:
|
||||
status_code: Final = original_exception.status_code if hasattr(original_exception, "status_code") else None
|
||||
if not isinstance(status_code, int) or status_code < 400:
|
||||
return
|
||||
message: Final = f"{exception_provider} - {error_str}"
|
||||
response: Final = original_exception.response if hasattr(original_exception, "response") else None
|
||||
match status_code:
|
||||
case 401:
|
||||
raise AuthenticationError(
|
||||
message=message,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
case 403:
|
||||
raise PermissionDeniedError(
|
||||
message=message,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=_response_or_stub(original_exception, status_code=status_code),
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
case 404:
|
||||
raise NotFoundError(
|
||||
message=message,
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
case 408:
|
||||
raise Timeout(
|
||||
message=message,
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
case 429:
|
||||
raise RateLimitError(
|
||||
message=message,
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
case 500:
|
||||
raise InternalServerError(
|
||||
message=message,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
case 502:
|
||||
raise BadGatewayError(
|
||||
message=message,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
case 503:
|
||||
raise ServiceUnavailableError(
|
||||
message=message,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
case 504:
|
||||
raise Timeout(
|
||||
message=message,
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
litellm_debug_info=extra_information,
|
||||
exception_status_code=status_code,
|
||||
)
|
||||
case _ if status_code < 500:
|
||||
raise BadRequestError(
|
||||
message=message,
|
||||
model=model,
|
||||
llm_provider=custom_llm_provider,
|
||||
response=response,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
case _:
|
||||
raise APIError(
|
||||
status_code=status_code,
|
||||
message=message,
|
||||
llm_provider=custom_llm_provider,
|
||||
model=model,
|
||||
request=original_exception.request if hasattr(original_exception, "request") else None,
|
||||
litellm_debug_info=extra_information,
|
||||
)
|
||||
|
||||
|
||||
def exception_type(
|
||||
model,
|
||||
original_exception,
|
||||
|
|
@ -2501,6 +2629,14 @@ def exception_type(
|
|||
For unmapped exceptions - raise the exception with traceback - https://github.com/BerriAI/litellm/issues/4201
|
||||
"""
|
||||
exception_mapping_worked = True
|
||||
_map_exception_by_status(
|
||||
model=model,
|
||||
original_exception=mappable_exception,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
error_str=error_str,
|
||||
exception_provider=exception_provider,
|
||||
extra_information=extra_information,
|
||||
)
|
||||
if hasattr(original_exception, "request"):
|
||||
raise APIConnectionError(
|
||||
message=f"{exception_provider} - {error_str}",
|
||||
|
|
|
|||
|
|
@ -172,7 +172,7 @@ def get_supported_openai_params(
|
|||
if request_type == "embeddings":
|
||||
return litellm.JinaAIEmbeddingConfig().get_supported_openai_params(model=model)
|
||||
elif custom_llm_provider == "together_ai":
|
||||
return litellm.TogetherAIConfig().get_supported_openai_params(model=model)
|
||||
return litellm.TogetherAIChatConfig().get_supported_openai_params(model=model)
|
||||
elif custom_llm_provider == "databricks":
|
||||
if request_type == "chat_completion":
|
||||
return litellm.DatabricksConfig().get_supported_openai_params(model=model)
|
||||
|
|
|
|||
97
litellm/litellm_core_utils/json_fragment_accumulator.py
Normal file
97
litellm/litellm_core_utils/json_fragment_accumulator.py
Normal file
|
|
@ -0,0 +1,97 @@
|
|||
import json
|
||||
from typing import Final, cast # noqa: TID251 # raw_decode returns tuple[Any, int]; no cast-free unpack
|
||||
|
||||
|
||||
class JSONFragmentAccumulator:
|
||||
"""
|
||||
Buffers a JSON value that arrives piecemeal over a stream (SSE data split
|
||||
across TCP packets, one shard per network read, etc) without the O(n^2)
|
||||
cost of repeated `buffer += fragment` string concatenation, and without
|
||||
the O(n^2) cost of re-copying the unconsumed remainder on every peeled
|
||||
value when one payload holds many concatenated JSON values.
|
||||
|
||||
Fragments are appended to a list in O(1). The buffer is only rebuilt into
|
||||
a single string, and only decoded, when a caller asks for a value via
|
||||
`pop_next_value`, and `could_close_json` lets callers skip that rebuild
|
||||
entirely for fragments that plainly cannot close a JSON value yet. Once
|
||||
rebuilt, consumed values are dropped by advancing a cursor rather than
|
||||
slicing a new string, so draining N concatenated values already sitting
|
||||
in the buffer costs O(n) total, not O(n^2).
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._chunks: list[str] = [] # mutable-ok: O(1) append; string concat would copy the buffer each time
|
||||
self._buffer: str = (
|
||||
"" # mutable-ok: lazily materialized join of _chunks, rebuilt only when _chunks is non-empty
|
||||
)
|
||||
self._offset: int = 0 # mutable-ok: cursor past already-consumed values; avoids re-slicing on every pop
|
||||
self._could_close: bool = False # mutable-ok: cached heuristic; rescanning past fragments was itself O(n^2)
|
||||
|
||||
def __bool__(self) -> bool:
|
||||
return bool(self._chunks) or self._offset < len(self._buffer)
|
||||
|
||||
def append(self, fragment: str) -> None:
|
||||
self._chunks.append(fragment) # mutable-ok: see __init__
|
||||
stripped: Final = fragment.rstrip()
|
||||
if stripped:
|
||||
self._could_close = stripped[-1] in ("}", "]") # mutable-ok: see __init__
|
||||
|
||||
def could_close_json(self) -> bool:
|
||||
"""
|
||||
Whether the buffer's logical last non-whitespace byte is "}" or "]",
|
||||
i.e. whether a JSON value could plausibly be complete. Tracked
|
||||
incrementally in `append` rather than rescanned here, so a run of
|
||||
blank keepalive fragments (e.g. from a malformed upstream stream)
|
||||
can't make this, or the join+parse it gates, cost O(n^2).
|
||||
"""
|
||||
return self._could_close
|
||||
|
||||
def _materialize(self) -> None:
|
||||
if not self._chunks:
|
||||
return
|
||||
unconsumed: Final = self._buffer[self._offset :]
|
||||
self._buffer = unconsumed + "".join(self._chunks) # mutable-ok: merge pending fragments, once per append batch
|
||||
self._offset = 0 # mutable-ok: see __init__
|
||||
self._chunks = [] # mutable-ok: see __init__
|
||||
|
||||
def pop_next_value(self) -> tuple[bool, object]:
|
||||
"""
|
||||
Attempt to decode one complete JSON value from the front of the
|
||||
buffer. On success, advances a cursor past that value (keeping any
|
||||
unconsumed tail, e.g. a second concatenated value, in place rather
|
||||
than copying it) and returns (True, value). If the buffer is empty
|
||||
or holds no complete value yet, it is left untouched and this
|
||||
returns (False, None).
|
||||
"""
|
||||
self._materialize()
|
||||
length: Final = len(self._buffer)
|
||||
start = self._offset
|
||||
while start < length and self._buffer[start].isspace():
|
||||
start += 1
|
||||
if start >= length:
|
||||
self._offset = start # mutable-ok: see __init__
|
||||
return False, None
|
||||
decoder: Final = json.JSONDecoder()
|
||||
try:
|
||||
raw_value: Final = decoder.raw_decode(self._buffer, start)
|
||||
except json.JSONDecodeError:
|
||||
return False, None
|
||||
decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode returns tuple[Any, int]
|
||||
self._offset = end_index # mutable-ok: see __init__
|
||||
if self._offset >= len(self._buffer):
|
||||
self._buffer = "" # mutable-ok: see __init__
|
||||
self._offset = 0 # mutable-ok: see __init__
|
||||
self._could_close = False # mutable-ok: buffer is empty, nothing can close
|
||||
return True, decoded
|
||||
|
||||
def snapshot(self) -> str:
|
||||
self._materialize()
|
||||
return self._buffer[self._offset :]
|
||||
|
||||
def set(self, value: str) -> None:
|
||||
"""Replace the buffer's contents with a single fragment."""
|
||||
self._chunks = [] # mutable-ok: see __init__
|
||||
self._buffer = value # mutable-ok: see __init__
|
||||
self._offset = 0 # mutable-ok: see __init__
|
||||
stripped: Final = value.rstrip()
|
||||
self._could_close = bool(stripped) and stripped[-1] in ("}", "]") # mutable-ok: see __init__
|
||||
|
|
@ -62,7 +62,7 @@ from litellm.integrations.custom_logger import CustomLogger
|
|||
from litellm.integrations.deepeval.deepeval import DeepEvalLogger
|
||||
from litellm.integrations.mlflow import MlflowLogger
|
||||
from litellm.integrations.sqs import SQSLogger
|
||||
from litellm.litellm_core_utils.core_helpers import reconstruct_model_name
|
||||
from litellm.litellm_core_utils.core_helpers import is_expected_client_error, reconstruct_model_name
|
||||
from litellm.litellm_core_utils.get_litellm_params import get_litellm_params
|
||||
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
|
||||
cost_breakdown_with_guardrail,
|
||||
|
|
@ -3124,6 +3124,13 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if not hasattr(self, "model_call_details"):
|
||||
self.model_call_details = {}
|
||||
|
||||
if (
|
||||
self.model_call_details.get("log_event_type") == "failed_api_call"
|
||||
and self.model_call_details.get("exception") is exception
|
||||
and self.model_call_details.get("standard_logging_object") is not None
|
||||
):
|
||||
return start_time, self.model_call_details["end_time"]
|
||||
|
||||
self.model_call_details["log_event_type"] = "failed_api_call"
|
||||
self.model_call_details["exception"] = exception
|
||||
self.model_call_details["traceback_exception"] = (
|
||||
|
|
@ -5455,9 +5462,10 @@ class StandardLoggingPayloadSetup:
|
|||
error_class: Final[str] = str(original_exception.__class__.__name__) if original_exception else ""
|
||||
_llm_provider_in_exception: Final = getattr(original_exception, "llm_provider", "")
|
||||
|
||||
# Get traceback information (first 100 lines)
|
||||
traceback_info = traceback_str or ""
|
||||
if original_exception:
|
||||
if original_exception and (
|
||||
litellm.log_client_error_tracebacks or not is_expected_client_error(original_exception)
|
||||
):
|
||||
tb: Final[TracebackType | None] = getattr(original_exception, "__traceback__", None)
|
||||
if tb:
|
||||
tb_lines: Final = traceback.format_tb(tb)
|
||||
|
|
@ -5930,11 +5938,15 @@ def get_standard_logging_object_payload(
|
|||
response_model_name = final_response_obj.get("model")
|
||||
|
||||
# For Azure Model Router, preserve the actual model in the top-level standard
|
||||
# logging payload only when the user has opted in.
|
||||
# logging payload.
|
||||
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
|
||||
|
||||
requested_model: Final = kwargs.get("model")
|
||||
if (
|
||||
isinstance(requested_model, str)
|
||||
and ("model_router" in requested_model.lower() or "model-router" in requested_model.lower())
|
||||
stamped_selected_model: Final = AzureFoundryModelInfo.get_model_router_selected_model(hidden_params)
|
||||
if stamped_selected_model is not None:
|
||||
model_name = stamped_selected_model
|
||||
elif (
|
||||
AzureFoundryModelInfo.is_model_router_call(model=requested_model, hidden_params=hidden_params)
|
||||
and isinstance(response_model_name, str)
|
||||
and response_model_name
|
||||
):
|
||||
|
|
|
|||
|
|
@ -97,16 +97,18 @@ def _redact_function_call(function_call) -> None:
|
|||
def _redact_choice_content(choice):
|
||||
"""Helper to redact content in a choice (message or delta)."""
|
||||
if isinstance(choice, litellm.Choices):
|
||||
choice.message.content = REDACTED_BY_LITELLM
|
||||
if hasattr(choice.message, "reasoning_content"):
|
||||
if choice.message.content is not None:
|
||||
choice.message.content = REDACTED_BY_LITELLM
|
||||
if getattr(choice.message, "reasoning_content", None) is not None:
|
||||
choice.message.reasoning_content = REDACTED_BY_LITELLM
|
||||
if hasattr(choice.message, "thinking_blocks"):
|
||||
choice.message.thinking_blocks = None
|
||||
_redact_tool_calls(getattr(choice.message, "tool_calls", None))
|
||||
_redact_function_call(getattr(choice.message, "function_call", None))
|
||||
elif isinstance(choice, litellm.utils.StreamingChoices):
|
||||
choice.delta.content = REDACTED_BY_LITELLM
|
||||
if hasattr(choice.delta, "reasoning_content"):
|
||||
if choice.delta.content is not None:
|
||||
choice.delta.content = REDACTED_BY_LITELLM
|
||||
if getattr(choice.delta, "reasoning_content", None) is not None:
|
||||
choice.delta.reasoning_content = REDACTED_BY_LITELLM
|
||||
if hasattr(choice.delta, "thinking_blocks"):
|
||||
choice.delta.thinking_blocks = None
|
||||
|
|
@ -117,19 +119,19 @@ def _redact_choice_content(choice):
|
|||
def _redact_responses_api_output(output_items):
|
||||
"""Helper to redact ResponsesAPIResponse output items."""
|
||||
for output_item in output_items:
|
||||
if hasattr(output_item, "text"):
|
||||
if getattr(output_item, "text", None) is not None:
|
||||
output_item.text = REDACTED_BY_LITELLM
|
||||
|
||||
if hasattr(output_item, "content") and isinstance(output_item.content, list):
|
||||
for content_part in output_item.content:
|
||||
if hasattr(content_part, "text"):
|
||||
if getattr(content_part, "text", None) is not None:
|
||||
content_part.text = REDACTED_BY_LITELLM
|
||||
|
||||
# Redact reasoning items in output array
|
||||
if hasattr(output_item, "type") and output_item.type == "reasoning":
|
||||
if hasattr(output_item, "summary") and isinstance(output_item.summary, list):
|
||||
for summary_item in output_item.summary:
|
||||
if hasattr(summary_item, "text"):
|
||||
if getattr(summary_item, "text", None) is not None:
|
||||
summary_item.text = REDACTED_BY_LITELLM
|
||||
|
||||
if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"):
|
||||
|
|
@ -142,17 +144,17 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str):
|
|||
if not isinstance(output_item, dict):
|
||||
continue
|
||||
|
||||
if "text" in output_item:
|
||||
if output_item.get("text") is not None:
|
||||
output_item["text"] = redacted_str
|
||||
|
||||
if isinstance(output_item.get("content"), list):
|
||||
for content_item in output_item["content"]:
|
||||
if isinstance(content_item, dict) and "text" in content_item:
|
||||
if isinstance(content_item, dict) and content_item.get("text") is not None:
|
||||
content_item["text"] = redacted_str
|
||||
|
||||
if output_item.get("type") == "reasoning" and isinstance(output_item.get("summary"), list):
|
||||
for summary_item in output_item["summary"]:
|
||||
if isinstance(summary_item, dict) and "text" in summary_item:
|
||||
if isinstance(summary_item, dict) and summary_item.get("text") is not None:
|
||||
summary_item["text"] = redacted_str
|
||||
|
||||
if output_item.get("type") == "function_call" and "arguments" in output_item:
|
||||
|
|
@ -189,40 +191,42 @@ def _redact_standard_logging_object(model_call_details: dict):
|
|||
standard_logging_object["response"] = {"text": redacted_str}
|
||||
|
||||
|
||||
def _redact_tool_calls_dict(message: dict, redacted_str: str) -> None:
|
||||
def _redact_tool_calls_dict(message: dict) -> None:
|
||||
"""Redact tool call / function_call arguments in a dict-form message or delta."""
|
||||
tool_calls: Final = message.get("tool_calls")
|
||||
if isinstance(tool_calls, list):
|
||||
for tool_call in tool_calls:
|
||||
if isinstance(tool_call, dict) and isinstance(tool_call.get("function"), dict):
|
||||
tool_call["function"]["arguments"] = redacted_str
|
||||
tool_call["function"]["arguments"] = REDACTED_BY_LITELLM
|
||||
|
||||
function_call: Final = message.get("function_call")
|
||||
if isinstance(function_call, dict) and "arguments" in function_call:
|
||||
function_call["arguments"] = redacted_str
|
||||
function_call["arguments"] = REDACTED_BY_LITELLM
|
||||
|
||||
|
||||
def _redact_model_response_dict_choices(choices, redacted_str: str):
|
||||
for choice in choices:
|
||||
if isinstance(choice, dict):
|
||||
if "message" in choice and isinstance(choice["message"], dict):
|
||||
choice["message"]["content"] = redacted_str
|
||||
if "reasoning_content" in choice["message"]:
|
||||
if choice["message"].get("content") is not None:
|
||||
choice["message"]["content"] = redacted_str
|
||||
if choice["message"].get("reasoning_content") is not None:
|
||||
choice["message"]["reasoning_content"] = redacted_str
|
||||
if "thinking_blocks" in choice["message"]:
|
||||
choice["message"]["thinking_blocks"] = None
|
||||
if "audio" in choice["message"]:
|
||||
choice["message"]["audio"] = None
|
||||
_redact_tool_calls_dict(choice["message"], redacted_str)
|
||||
_redact_tool_calls_dict(choice["message"])
|
||||
elif "delta" in choice and isinstance(choice["delta"], dict):
|
||||
choice["delta"]["content"] = redacted_str
|
||||
if "reasoning_content" in choice["delta"]:
|
||||
if choice["delta"].get("content") is not None:
|
||||
choice["delta"]["content"] = redacted_str
|
||||
if choice["delta"].get("reasoning_content") is not None:
|
||||
choice["delta"]["reasoning_content"] = redacted_str
|
||||
if "thinking_blocks" in choice["delta"]:
|
||||
choice["delta"]["thinking_blocks"] = None
|
||||
if "audio" in choice["delta"]:
|
||||
choice["delta"]["audio"] = None
|
||||
_redact_tool_calls_dict(choice["delta"], redacted_str)
|
||||
_redact_tool_calls_dict(choice["delta"])
|
||||
else:
|
||||
_redact_choice_content(choice)
|
||||
|
||||
|
|
@ -263,7 +267,7 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons
|
|||
isinstance(result, (litellm.ModelResponse, litellm.ResponsesAPIResponse, litellm.EmbeddingResponse))
|
||||
or (isinstance(result, dict) and ("choices" in result or "output" in result))
|
||||
):
|
||||
return {"text": "redacted-by-litellm"}
|
||||
return {"text": REDACTED_BY_LITELLM}
|
||||
|
||||
_result: Final = copy.deepcopy(result)
|
||||
if isinstance(_result, litellm.ModelResponse):
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from litellm.anthropic_beta_headers_manager import (
|
|||
)
|
||||
from litellm.constants import RESPONSE_FORMAT_TOOL_NAME
|
||||
from litellm.litellm_core_utils.core_helpers import map_finish_reason
|
||||
from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
|
|
@ -654,7 +655,7 @@ class ModelResponseIterator:
|
|||
|
||||
# For handling partial JSON chunks from fragmentation
|
||||
# See: https://github.com/BerriAI/litellm/issues/17473
|
||||
self.accumulated_json: str = ""
|
||||
self._json_buffer = JSONFragmentAccumulator()
|
||||
self.chunk_type: Literal["valid_json", "accumulated_json"] = "valid_json"
|
||||
|
||||
# Track current content block type to avoid emitting tool calls for non-tool blocks
|
||||
|
|
@ -678,6 +679,14 @@ class ModelResponseIterator:
|
|||
self._current_server_tool_id: str | None = None
|
||||
self._container_id: str | None = None
|
||||
|
||||
@property
|
||||
def accumulated_json(self) -> str:
|
||||
return self._json_buffer.snapshot()
|
||||
|
||||
@accumulated_json.setter
|
||||
def accumulated_json(self, value: str) -> None:
|
||||
self._json_buffer.set(value)
|
||||
|
||||
def check_empty_tool_call_args(self) -> bool:
|
||||
"""
|
||||
Check if the tool call block so far has been an empty string
|
||||
|
|
@ -1149,31 +1158,39 @@ class ModelResponseIterator:
|
|||
container: Final = message_delta["delta"].get("container")
|
||||
return finish_reason, usage, container
|
||||
|
||||
def _handle_accumulated_json_chunk(self, data_str: str) -> ModelResponseStream | None:
|
||||
def _handle_accumulated_json_chunk(self, data_str: str, is_final: bool = False) -> ModelResponseStream | None:
|
||||
"""
|
||||
Handle partial JSON chunks by accumulating them until valid JSON is received.
|
||||
|
||||
This fixes network fragmentation issues where SSE data chunks may be split
|
||||
across TCP packets. See: https://github.com/BerriAI/litellm/issues/17473
|
||||
|
||||
Mid-stream, defer parsing until the buffer's last byte can close a value:
|
||||
attempting a parse after every fragment of one large object is O(n^2) and
|
||||
holds the GIL, freezing the event loop. At end of stream (is_final) no more
|
||||
data is coming, so drain whatever complete values remain regardless of the
|
||||
trailing byte.
|
||||
|
||||
Args:
|
||||
data_str: The JSON string to parse (without "data:" prefix)
|
||||
is_final: True when called from the end-of-stream drain, where the
|
||||
trailing-byte heuristic no longer applies
|
||||
|
||||
Returns:
|
||||
ModelResponseStream if JSON is complete, None if still accumulating
|
||||
"""
|
||||
# Accumulate JSON data
|
||||
self.accumulated_json += data_str
|
||||
self._json_buffer.append(data_str)
|
||||
|
||||
# Try to parse the accumulated JSON
|
||||
try:
|
||||
data_json: Final = json.loads(self.accumulated_json)
|
||||
self.accumulated_json = "" # Reset after successful parsing
|
||||
return self.chunk_parser(chunk=data_json)
|
||||
except json.JSONDecodeError:
|
||||
# If it's not valid JSON yet, continue to the next chunk
|
||||
if not is_final and not self._json_buffer.could_close_json():
|
||||
return None
|
||||
|
||||
while True:
|
||||
found, decoded = self._json_buffer.pop_next_value()
|
||||
if not found:
|
||||
return None
|
||||
if isinstance(decoded, dict):
|
||||
return self.chunk_parser(chunk=decoded)
|
||||
|
||||
def _parse_sse_data(self, str_line: str) -> ModelResponseStream | None:
|
||||
"""
|
||||
Parse SSE data line, handling both complete and partial JSON chunks.
|
||||
|
|
@ -1209,13 +1226,10 @@ class ModelResponseIterator:
|
|||
chunk = self.response_iterator.__next__()
|
||||
except StopIteration:
|
||||
# If we have accumulated JSON when stream ends, try to parse it
|
||||
if self.accumulated_json:
|
||||
try:
|
||||
data_json = json.loads(self.accumulated_json)
|
||||
self.accumulated_json = ""
|
||||
return self.chunk_parser(chunk=data_json)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
if self._json_buffer:
|
||||
result = self._handle_accumulated_json_chunk(data_str="", is_final=True)
|
||||
if result is not None:
|
||||
return result
|
||||
raise StopIteration
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"Error receiving chunk from stream: {e}")
|
||||
|
|
@ -1258,13 +1272,10 @@ class ModelResponseIterator:
|
|||
chunk = await self.async_response_iterator.__anext__()
|
||||
except StopAsyncIteration:
|
||||
# If we have accumulated JSON when stream ends, try to parse it
|
||||
if self.accumulated_json:
|
||||
try:
|
||||
data_json = json.loads(self.accumulated_json)
|
||||
self.accumulated_json = ""
|
||||
return self.chunk_parser(chunk=data_json)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
if self._json_buffer:
|
||||
result = self._handle_accumulated_json_chunk(data_str="", is_final=True)
|
||||
if result is not None:
|
||||
return result
|
||||
raise StopAsyncIteration
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"Error receiving chunk from stream: {e}")
|
||||
|
|
|
|||
|
|
@ -38,6 +38,21 @@ DROP_DISABLED_THINKING_WARNING: Final = (
|
|||
"thinking blocks, and those thinking tokens are billed as output tokens."
|
||||
)
|
||||
|
||||
# Anthropic error `type` (both the JSON error body and SSE `event: error`
|
||||
# payloads use this field) mapped to the HTTP status code it corresponds to.
|
||||
ANTHROPIC_ERROR_STATUS_CODE_MAP: Final = MappingProxyType(
|
||||
{
|
||||
"invalid_request_error": 400,
|
||||
"authentication_error": 401,
|
||||
"permission_error": 403,
|
||||
"not_found_error": 404,
|
||||
"rate_limit_error": 429,
|
||||
"api_error": 500,
|
||||
"overloaded_error": 503,
|
||||
"timeout_error": 504,
|
||||
}
|
||||
)
|
||||
|
||||
_BEDROCK_VERSION_SUFFIX_RE: Final = re.compile(r"-v\d+(?::\d+)?$")
|
||||
_INFERENCE_PROFILE_MINOR_RE: Final = re.compile(r":\d+$")
|
||||
_DATED_RELEASE_SUFFIX_RE: Final = re.compile(r"-\d{8}$")
|
||||
|
|
|
|||
|
|
@ -434,7 +434,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
content_items = list(content.get("content", []))
|
||||
|
||||
# Single-item text keeps the backward-compatible string format; a single
|
||||
# image becomes a structured image_url part
|
||||
# image or document becomes a structured image_url part
|
||||
if len(content_items) == 1:
|
||||
c = content_items[0]
|
||||
if isinstance(c, str):
|
||||
|
|
@ -454,7 +454,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
)
|
||||
self._add_cache_control_if_applicable(content, tool_result, model)
|
||||
tool_message_list.append(tool_result)
|
||||
elif c.get("type") == "image":
|
||||
elif c.get("type") in ("image", "document"):
|
||||
image_part = self._tool_result_image_part(c.get("source"))
|
||||
tool_result = ChatCompletionToolMessage(
|
||||
role="tool",
|
||||
|
|
@ -482,7 +482,7 @@ class LiteLLMAnthropicMessagesAdapter:
|
|||
text=c.get("text", ""),
|
||||
)
|
||||
)
|
||||
elif c.get("type") == "image":
|
||||
elif c.get("type") in ("image", "document"):
|
||||
image_part = self._tool_result_image_part(c.get("source"))
|
||||
if image_part:
|
||||
combined_content_parts.append(image_part)
|
||||
|
|
|
|||
|
|
@ -6,13 +6,41 @@ yields every chunk to the caller (preserving real streaming), collects
|
|||
all bytes, and on stream exhaustion rebuilds the full Anthropic response
|
||||
to run through agentic completion hooks. If an agentic hook fires, the
|
||||
follow-up response is chained as Phase 2 of the same iterator.
|
||||
|
||||
In hold-back mode (``hold_back=True``) chunks are buffered instead of yielded
|
||||
live, keepalive pings run whenever no other byte is ready, and then either the
|
||||
follow-up replaces the message or the buffer replays, except that a tool_use for
|
||||
a server-fulfilled tool fails the turn rather than reaching a client that cannot
|
||||
execute it.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any, Final, cast
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES
|
||||
|
||||
HOLD_BACK_PING_INTERVAL_SECONDS: Final = 15.0
|
||||
SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = (
|
||||
b"event: error\n"
|
||||
b'data: {"type": "error", "error": {"type": "api_error", "message": '
|
||||
b'"Server-side tool retrieval failed, so this turn could not be completed. Please retry."}}\n\n'
|
||||
)
|
||||
|
||||
|
||||
def is_server_fulfilled_tool_leak_error(chunk: object) -> bool:
|
||||
return chunk == SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES
|
||||
|
||||
|
||||
async def _anext_or_none(iterator: AsyncIterator) -> bytes | None:
|
||||
try:
|
||||
return await iterator.__anext__()
|
||||
except StopAsyncIteration:
|
||||
return None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# SSE parsing helpers (module-level to keep the class lean)
|
||||
|
|
@ -156,6 +184,9 @@ class AgenticAnthropicStreamingIterator:
|
|||
logging_obj: Any,
|
||||
custom_llm_provider: str,
|
||||
kwargs: dict,
|
||||
hold_back: bool = False,
|
||||
server_fulfilled_tool_names: frozenset[str] = frozenset(),
|
||||
ping_interval_seconds: float = HOLD_BACK_PING_INTERVAL_SECONDS,
|
||||
):
|
||||
self._inner = completion_stream.__aiter__()
|
||||
self._http_handler = http_handler
|
||||
|
|
@ -166,16 +197,32 @@ class AgenticAnthropicStreamingIterator:
|
|||
self._logging_obj = logging_obj
|
||||
self._custom_llm_provider = custom_llm_provider
|
||||
self._kwargs = kwargs
|
||||
self._hold_back = hold_back
|
||||
self._server_fulfilled_tool_names = server_fulfilled_tool_names
|
||||
self._ping_interval_seconds = ping_interval_seconds
|
||||
|
||||
self._collected_bytes: list[bytes] = []
|
||||
self._stream_exhausted = False
|
||||
self._hook_processing_done = False
|
||||
self._follow_up_iterator: AsyncIterator | None = None
|
||||
self._drain_task: asyncio.Task | None = None
|
||||
self._hook_task: asyncio.Task | None = None
|
||||
self._follow_up_chunk_task: asyncio.Task | None = None
|
||||
self._replay_index = 0
|
||||
self._error_emitted = False
|
||||
|
||||
@property
|
||||
def has_buffered_provider_output(self) -> bool:
|
||||
"""Whether provider output was received but withheld from the client behind keepalive pings."""
|
||||
return self._hold_back and bool(self._collected_bytes)
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self) -> bytes:
|
||||
if self._hold_back:
|
||||
return await self._anext_held_back()
|
||||
|
||||
# Phase 1: yield from upstream, collect bytes
|
||||
if not self._stream_exhausted:
|
||||
try:
|
||||
|
|
@ -194,11 +241,102 @@ class AgenticAnthropicStreamingIterator:
|
|||
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def _drain_upstream(self) -> None:
|
||||
try:
|
||||
while True:
|
||||
self._collected_bytes.append(await self._inner.__anext__())
|
||||
except StopAsyncIteration:
|
||||
return
|
||||
|
||||
async def _completed_within_ping_interval(self, task: asyncio.Task) -> bool:
|
||||
try:
|
||||
await asyncio.wait_for(asyncio.shield(task), timeout=self._ping_interval_seconds)
|
||||
except asyncio.TimeoutError:
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _anext_held_back(self) -> bytes:
|
||||
if self._drain_task is None:
|
||||
self._drain_task = asyncio.create_task(self._drain_upstream())
|
||||
return STREAM_SSE_KEEPALIVE_PING_BYTES
|
||||
|
||||
if not self._stream_exhausted:
|
||||
if not await self._completed_within_ping_interval(self._drain_task):
|
||||
return STREAM_SSE_KEEPALIVE_PING_BYTES
|
||||
self._stream_exhausted = True
|
||||
|
||||
if self._hook_task is None:
|
||||
self._hook_task = asyncio.create_task(self._process_agentic_hooks())
|
||||
if not await self._completed_within_ping_interval(self._hook_task):
|
||||
return STREAM_SSE_KEEPALIVE_PING_BYTES
|
||||
|
||||
if self._follow_up_iterator is not None:
|
||||
return await self._next_follow_up_chunk(self._follow_up_iterator)
|
||||
|
||||
if self._buffer_holds_server_fulfilled_tool_use():
|
||||
if self._error_emitted:
|
||||
raise StopAsyncIteration
|
||||
self._error_emitted = True
|
||||
verbose_logger.error(
|
||||
"AgenticStreamingIterator: hooks did not replace a message containing a server-fulfilled "
|
||||
"tool_use [model=%s]; emitting an SSE error instead of leaking the tool call to the client",
|
||||
self._model,
|
||||
)
|
||||
return SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES
|
||||
|
||||
if self._replay_index < len(self._collected_bytes):
|
||||
chunk: Final = self._collected_bytes[self._replay_index]
|
||||
self._replay_index += 1
|
||||
return chunk
|
||||
|
||||
raise StopAsyncIteration
|
||||
|
||||
async def _next_follow_up_chunk(self, follow_up_iterator: AsyncIterator) -> bytes:
|
||||
if self._follow_up_chunk_task is None:
|
||||
self._follow_up_chunk_task = asyncio.create_task(_anext_or_none(follow_up_iterator))
|
||||
if not await self._completed_within_ping_interval(self._follow_up_chunk_task):
|
||||
return STREAM_SSE_KEEPALIVE_PING_BYTES
|
||||
chunk: Final = self._follow_up_chunk_task.result()
|
||||
self._follow_up_chunk_task = None
|
||||
if chunk is None:
|
||||
raise StopAsyncIteration
|
||||
return chunk
|
||||
|
||||
def _buffer_holds_server_fulfilled_tool_use(self) -> bool:
|
||||
if not self._server_fulfilled_tool_names:
|
||||
return False
|
||||
started_blocks: Final = (
|
||||
data.get("content_block")
|
||||
for event_type, data in _parse_sse_events(b"".join(self._collected_bytes))
|
||||
if event_type == "content_block_start"
|
||||
)
|
||||
return any(
|
||||
isinstance(block, dict)
|
||||
and block.get("type") == "tool_use"
|
||||
and block.get("name") in self._server_fulfilled_tool_names
|
||||
for block in started_blocks
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _settle_task(task: asyncio.Task | None) -> None:
|
||||
if task is None:
|
||||
return
|
||||
if task.done():
|
||||
if not task.cancelled():
|
||||
task.exception()
|
||||
return
|
||||
task.cancel()
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
async def aclose(self) -> None:
|
||||
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
|
||||
aclose_if_supported,
|
||||
)
|
||||
|
||||
await self._settle_task(self._drain_task)
|
||||
await self._settle_task(self._hook_task)
|
||||
await self._settle_task(self._follow_up_chunk_task)
|
||||
await aclose_if_supported(self._inner)
|
||||
await aclose_if_supported(self._follow_up_iterator)
|
||||
|
||||
|
|
@ -217,11 +355,6 @@ class AgenticAnthropicStreamingIterator:
|
|||
verbose_logger.debug("AgenticStreamingIterator: Could not rebuild response from SSE bytes")
|
||||
return
|
||||
|
||||
[
|
||||
(f"{b.get('type')}({b.get('name', '')})" if b.get("type") == "tool_use" else b.get("type"))
|
||||
for b in rebuilt.get("content", [])
|
||||
]
|
||||
|
||||
result: Final = await self._http_handler._call_agentic_completion_hooks(
|
||||
response=rebuilt,
|
||||
model=self._model,
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from functools import partial
|
|||
from typing import Any, Final, cast
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.exception_mapping_utils import exception_type
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.llms.anthropic.common_utils import (
|
||||
flatten_unencrypted_web_search_results_in_anthropic_messages,
|
||||
|
|
@ -21,6 +22,7 @@ from litellm.llms.anthropic.common_utils import (
|
|||
from litellm.llms.base_llm.anthropic_messages.transformation import (
|
||||
BaseAnthropicMessagesConfig,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
||||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.types.llms.anthropic_messages.anthropic_request import AnthropicMetadata
|
||||
|
|
@ -382,13 +384,18 @@ async def anthropic_messages(
|
|||
)
|
||||
ctx: Final = contextvars.copy_context()
|
||||
func_with_context: Final = partial(ctx.run, func)
|
||||
init_response: Final = await loop.run_in_executor(None, func_with_context)
|
||||
|
||||
if asyncio.iscoroutine(init_response):
|
||||
response = await init_response
|
||||
else:
|
||||
response = init_response
|
||||
return response
|
||||
try:
|
||||
init_response: Final = await loop.run_in_executor(None, func_with_context)
|
||||
if asyncio.iscoroutine(init_response):
|
||||
return await init_response
|
||||
return init_response
|
||||
except BaseLLMException as e:
|
||||
raise exception_type(
|
||||
model=model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
original_exception=e,
|
||||
extra_kwargs=kwargs,
|
||||
)
|
||||
|
||||
|
||||
def validate_anthropic_api_metadata(metadata: dict | None = None) -> dict | None:
|
||||
|
|
|
|||
|
|
@ -46,6 +46,10 @@ class AnthropicMessagesStreamCacheWriter:
|
|||
stream._hidden_params if isinstance(stream, AnthropicMessagesStreamingResponse) else _EMPTY_MAPPING
|
||||
)
|
||||
|
||||
@property
|
||||
def has_buffered_provider_output(self) -> bool:
|
||||
return getattr(self.stream, "has_buffered_provider_output", False) is True
|
||||
|
||||
def __aiter__(self) -> "AnthropicMessagesStreamCacheWriter":
|
||||
return self
|
||||
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from datetime import datetime
|
||||
from typing import Any, Final, Protocol, runtime_checkable
|
||||
|
||||
|
|
@ -11,9 +11,11 @@ from typing_extensions import TypedDict
|
|||
from litellm.litellm_core_utils.core_helpers import process_response_headers
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.llms.anthropic.common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP
|
||||
from litellm.proxy.pass_through_endpoints.success_handler import (
|
||||
PassThroughEndpointLogging,
|
||||
)
|
||||
from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
|
||||
from litellm.types.utils import GenericStreamingChunk, ModelResponseStream
|
||||
|
||||
|
|
@ -33,26 +35,239 @@ def _is_message_stop_chunk(chunk: object) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _is_provider_error_chunk(chunk: object) -> bool:
|
||||
def is_anthropic_ping_chunk(chunk: object) -> bool:
|
||||
"""
|
||||
Whether a chunk is a pure ``ping`` keepalive frame. It carries no content
|
||||
and can recur indefinitely on a slow-starting or idle connection, so a
|
||||
mid-stream fallback wrapper drops it outright while still deciding
|
||||
whether to commit to the primary stream, rather than buffering it.
|
||||
|
||||
A physical transport chunk that coalesces a ping with any other SSE
|
||||
event (``message_start``, ``content_block_delta``, ``event: error``, ...)
|
||||
is NOT a pure ping - dropping it whole would discard those events - so
|
||||
only a chunk whose every ``event:`` line is ``event: ping`` qualifies.
|
||||
"""
|
||||
if isinstance(chunk, dict):
|
||||
return chunk.get("type") == "error"
|
||||
return chunk.get("type") == "ping"
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
return any(line == b"event: error" for line in chunk.splitlines())
|
||||
event_lines: Final = tuple(line for line in chunk.splitlines() if line.startswith(b"event:"))
|
||||
return bool(event_lines) and all(line == b"event: ping" for line in event_lines)
|
||||
return False
|
||||
|
||||
|
||||
def is_anthropic_content_delta_chunk(chunk: object) -> bool:
|
||||
"""
|
||||
Whether a chunk carries actual assistant-generated output (a
|
||||
``content_block_delta`` frame), as opposed to a lifecycle/bookkeeping
|
||||
frame (``message_start``, ``content_block_start``/``stop``,
|
||||
``message_delta``, ``message_stop``, ``ping``) that carries nothing
|
||||
worth preserving before an invisible mid-stream fallback retry.
|
||||
"""
|
||||
if isinstance(chunk, dict):
|
||||
return chunk.get("type") == "content_block_delta"
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
return any(line == b"event: content_block_delta" for line in chunk.splitlines())
|
||||
return False
|
||||
|
||||
|
||||
def _decoded_sse_data_line(line: bytes) -> object | None:
|
||||
if not line.startswith(b"data:"):
|
||||
return None
|
||||
try:
|
||||
return json.loads(line[len(b"data:") :].strip())
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
|
||||
|
||||
def _anthropic_error_event_payload(chunk: object) -> Mapping[str, object] | None:
|
||||
if isinstance(chunk, dict):
|
||||
return chunk if chunk.get("type") == "error" else None
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
decoded_lines: Final = (_decoded_sse_data_line(line) for line in chunk.splitlines())
|
||||
return next(
|
||||
(
|
||||
candidate
|
||||
for candidate in decoded_lines
|
||||
if isinstance(candidate, dict) and candidate.get("type") == "error"
|
||||
),
|
||||
None,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _anthropic_error_body(chunk: object) -> Mapping[str, object] | None:
|
||||
"""Return the ``error`` object of an Anthropic SSE ``event: error`` chunk, or None."""
|
||||
payload: Final = _anthropic_error_event_payload(chunk)
|
||||
error_body: Final = payload.get("error") if payload is not None else None
|
||||
return error_body if isinstance(error_body, dict) else None
|
||||
|
||||
|
||||
def _is_provider_error_chunk(chunk: object) -> bool:
|
||||
return _anthropic_error_body(chunk) is not None
|
||||
|
||||
|
||||
def parse_anthropic_error_event(chunk: object) -> tuple[str, str, int] | None:
|
||||
"""
|
||||
Extract ``(error_type, message, http_status_code)`` from an Anthropic SSE
|
||||
``event: error`` chunk (raw bytes or an already-decoded dict), or None if
|
||||
``chunk`` is not an error event.
|
||||
|
||||
The status code is looked up via ANTHROPIC_ERROR_STATUS_CODE_MAP,
|
||||
defaulting to 500 for an error ``type`` Anthropic hasn't documented yet.
|
||||
"""
|
||||
error_body: Final = _anthropic_error_body(chunk)
|
||||
if error_body is None:
|
||||
return None
|
||||
error_type: Final = error_body.get("type")
|
||||
if not isinstance(error_type, str):
|
||||
return None
|
||||
message: Final = error_body.get("message")
|
||||
return (
|
||||
error_type,
|
||||
message if isinstance(message, str) else error_type,
|
||||
ANTHROPIC_ERROR_STATUS_CODE_MAP.get(error_type, 500),
|
||||
)
|
||||
|
||||
|
||||
def _is_terminal_stream_chunk(chunk: object) -> bool:
|
||||
return _is_message_stop_chunk(chunk) or _is_provider_error_chunk(chunk)
|
||||
|
||||
|
||||
def _sse_event(event_type: str, payload: Mapping[str, object]) -> bytes:
|
||||
return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
|
||||
def _incomplete_stream_error_sse_event() -> bytes:
|
||||
payload: Final = json.dumps(
|
||||
{
|
||||
"type": "error",
|
||||
"error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE},
|
||||
}
|
||||
return _sse_event( # mutable-ok: one-shot JSON payload, never mutated after construction
|
||||
"error",
|
||||
{"type": "error", "error": {"type": "api_error", "message": INCOMPLETE_STREAM_ERROR_MESSAGE}},
|
||||
)
|
||||
|
||||
|
||||
def _anthropic_content_block_start_and_deltas(
|
||||
block: Mapping[str, object],
|
||||
) -> tuple[Mapping[str, object], tuple[Mapping[str, object], ...]]:
|
||||
"""
|
||||
``(content_block_start.content_block, content_block_delta.delta events)``
|
||||
for one Anthropic response content block. A thinking block emits both a
|
||||
thinking_delta and a trailing signature_delta - a real Anthropic stream
|
||||
does the same, and dropping the signature makes any replay of that
|
||||
assistant message (a follow-up turn, a tool-use continuation) fail
|
||||
Anthropic's thinking-signature verification. redacted_thinking has no
|
||||
delta at all - it is sent complete in content_block_start.
|
||||
"""
|
||||
match block.get("type"):
|
||||
case "tool_use":
|
||||
return (
|
||||
{ # mutable-ok: one-shot payload
|
||||
"id": block.get("id"),
|
||||
"name": block.get("name"),
|
||||
"input": {}, # mutable-ok: one-shot payload
|
||||
"type": "tool_use",
|
||||
},
|
||||
(
|
||||
{ # mutable-ok: one-shot payload
|
||||
"partial_json": json.dumps(block.get("input") or {}), # mutable-ok: one-shot payload
|
||||
"type": "input_json_delta",
|
||||
},
|
||||
),
|
||||
)
|
||||
case "thinking":
|
||||
signature: Final = block.get("signature")
|
||||
signature_deltas: Final = (
|
||||
({"signature": signature, "type": "signature_delta"},) # mutable-ok: one-shot payload
|
||||
if isinstance(signature, str) and signature
|
||||
else ()
|
||||
)
|
||||
return (
|
||||
{"thinking": "", "signature": "", "type": "thinking"}, # mutable-ok: one-shot payload
|
||||
(
|
||||
{"thinking": block.get("thinking") or "", "type": "thinking_delta"}, # mutable-ok: one-shot payload
|
||||
*signature_deltas,
|
||||
),
|
||||
)
|
||||
case "redacted_thinking":
|
||||
return ({"type": "redacted_thinking", "data": block.get("data")}, ()) # mutable-ok: one-shot JSON payload
|
||||
case _:
|
||||
return (
|
||||
{"type": "text", "text": ""}, # mutable-ok: one-shot JSON payload
|
||||
({"type": "text_delta", "text": block.get("text") or ""},), # mutable-ok: one-shot JSON payload
|
||||
)
|
||||
|
||||
|
||||
def anthropic_messages_response_as_sse_events(response: AnthropicMessagesResponse) -> tuple[bytes, ...]:
|
||||
"""
|
||||
Render a complete (non-streaming) AnthropicMessagesResponse as the SSE
|
||||
event sequence a real streaming request would have produced.
|
||||
|
||||
A mid-stream fallback can resolve to a non-streaming response even
|
||||
though the client asked to stream (e.g. an agentic tool-use loop that
|
||||
intercepts and returns a complete message) - yielding that dict directly
|
||||
into a `/v1/messages` SSE byte stream would produce a malformed
|
||||
response, so it's synthesized into the message_start/content_block_*/
|
||||
message_delta/message_stop lifecycle a real stream would have sent.
|
||||
"""
|
||||
content_blocks: Final = response.get("content") or ()
|
||||
content_events: Final = (
|
||||
event for index, block in enumerate(content_blocks) for event in _anthropic_content_block_events(index, block)
|
||||
)
|
||||
# A real message_start always carries a null stop_reason/stop_sequence and
|
||||
# a zero output_tokens - those are only known once generation finishes, so
|
||||
# copying the completed response's final values here would let a client
|
||||
# treat the message as already finished, or double-count output tokens.
|
||||
message_start_usage: Final = { # mutable-ok: one-shot JSON payload
|
||||
**(response.get("usage") or {}),
|
||||
"output_tokens": 0,
|
||||
}
|
||||
message_start_payload: Final = { # mutable-ok: one-shot JSON payload, never mutated after construction
|
||||
"type": "message_start",
|
||||
"message": { # mutable-ok: one-shot JSON payload
|
||||
**response,
|
||||
"content": [], # mutable-ok: one-shot JSON payload
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": message_start_usage,
|
||||
},
|
||||
}
|
||||
message_delta_payload: Final = { # mutable-ok: one-shot JSON payload, never mutated after construction
|
||||
"type": "message_delta",
|
||||
"delta": { # mutable-ok: one-shot JSON payload
|
||||
"stop_reason": response.get("stop_reason"),
|
||||
"stop_sequence": response.get("stop_sequence"),
|
||||
},
|
||||
"usage": response.get("usage") or {}, # mutable-ok: one-shot JSON payload
|
||||
}
|
||||
return (
|
||||
_sse_event("message_start", message_start_payload),
|
||||
*content_events,
|
||||
_sse_event("message_delta", message_delta_payload),
|
||||
_sse_event("message_stop", {"type": "message_stop"}), # mutable-ok: one-shot JSON payload
|
||||
)
|
||||
|
||||
|
||||
def _anthropic_content_block_events(index: int, block: Mapping[str, object]) -> tuple[bytes, ...]:
|
||||
start_block, deltas = _anthropic_content_block_start_and_deltas(block)
|
||||
start_payload: Final = { # mutable-ok: one-shot payload
|
||||
"type": "content_block_start",
|
||||
"index": index,
|
||||
"content_block": start_block,
|
||||
}
|
||||
stop_payload: Final = { # mutable-ok: one-shot payload
|
||||
"type": "content_block_stop",
|
||||
"index": index,
|
||||
}
|
||||
delta_events: Final = tuple(
|
||||
_sse_event(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": index, "delta": delta}, # mutable-ok: one-shot payload
|
||||
)
|
||||
for delta in deltas
|
||||
)
|
||||
return (
|
||||
_sse_event("content_block_start", start_payload),
|
||||
*delta_events,
|
||||
_sse_event("content_block_stop", stop_payload),
|
||||
)
|
||||
return f"event: error\ndata: {payload}\n\n".encode()
|
||||
|
||||
|
||||
class AnthropicMessagesStreamHiddenParams(TypedDict):
|
||||
|
|
@ -97,6 +312,10 @@ class AnthropicMessagesStreamingResponse:
|
|||
self.completion_stream = completion_stream
|
||||
self._hidden_params = hidden_params
|
||||
|
||||
@property
|
||||
def has_buffered_provider_output(self) -> bool:
|
||||
return getattr(self.completion_stream, "has_buffered_provider_output", False) is True
|
||||
|
||||
def __aiter__(self) -> "AnthropicMessagesStreamingResponse":
|
||||
return self
|
||||
|
||||
|
|
|
|||
|
|
@ -87,6 +87,51 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
return source.get("url")
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _translate_anthropic_document_block_to_file_part(
|
||||
block: Mapping[str, object],
|
||||
) -> dict[str, str] | None: # mutable-ok: API message payload
|
||||
"""Convert an Anthropic document block to a Responses input_file part."""
|
||||
raw_source: Final = block.get("source")
|
||||
if not isinstance(raw_source, Mapping):
|
||||
return None
|
||||
source: Final = cast(Mapping[str, object], raw_source) # cast-ok: untrusted client payload
|
||||
source_type: Final = source.get("type")
|
||||
if source_type == "base64":
|
||||
data: Final = source.get("data")
|
||||
if not isinstance(data, str) or not data:
|
||||
return None
|
||||
raw_media_type: Final = source.get("media_type")
|
||||
media_type: Final = (
|
||||
raw_media_type if isinstance(raw_media_type, str) and raw_media_type else "application/pdf"
|
||||
)
|
||||
raw_title: Final = block.get("title")
|
||||
filename: Final = raw_title if isinstance(raw_title, str) and raw_title else "document.pdf"
|
||||
return { # mutable-ok: API message payload
|
||||
"type": "input_file",
|
||||
"filename": filename,
|
||||
"file_data": f"data:{media_type};base64,{data}",
|
||||
}
|
||||
if source_type == "url":
|
||||
url: Final = source.get("url")
|
||||
if not isinstance(url, str) or not url:
|
||||
return None
|
||||
return {"type": "input_file", "file_url": url} # mutable-ok: API message payload
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _tool_result_output_value(
|
||||
output_text: str,
|
||||
file_parts: tuple[dict[str, str], ...], # mutable-ok: json content parts
|
||||
) -> str | list[dict[str, str]]: # mutable-ok: API message payload
|
||||
"""Plain string output, or a part list when document file parts are present."""
|
||||
if not file_parts:
|
||||
return output_text
|
||||
text_parts: Final = (
|
||||
[{"type": "input_text", "text": output_text}] if output_text else [] # mutable-ok: API message payload
|
||||
)
|
||||
return [*text_parts, *file_parts] # mutable-ok: API message payload
|
||||
|
||||
@staticmethod
|
||||
def _translate_midturn_system_content_to_responses(
|
||||
content: str | Iterable[AnthropicSystemMessageContent],
|
||||
|
|
@ -169,6 +214,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
system text -> message(role=system, input_text)
|
||||
user text -> message(role=user, input_text)
|
||||
user image -> message(role=user, input_image)
|
||||
user document -> message(role=user, input_file)
|
||||
user tool_result -> function_call_output
|
||||
assistant text -> message(role=assistant, output_text)
|
||||
assistant thinking -> reasoning
|
||||
|
|
@ -223,9 +269,25 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
{"type": "input_image", "image_url": url}, block.get("prompt_cache_breakpoint")
|
||||
)
|
||||
)
|
||||
elif btype == "document":
|
||||
file_part = self._translate_anthropic_document_block_to_file_part(block)
|
||||
if file_part:
|
||||
user_parts.append(
|
||||
with_prompt_cache_breakpoint(file_part, block.get("prompt_cache_breakpoint"))
|
||||
)
|
||||
elif btype == "tool_result":
|
||||
tool_use_id = block.get("tool_use_id", "")
|
||||
inner = block.get("content")
|
||||
document_candidates = (
|
||||
tuple(
|
||||
self._translate_anthropic_document_block_to_file_part(c)
|
||||
for c in inner
|
||||
if isinstance(c, dict) and c.get("type") == "document"
|
||||
)
|
||||
if isinstance(inner, list)
|
||||
else ()
|
||||
)
|
||||
tool_file_parts = tuple(part for part in document_candidates if part is not None)
|
||||
if inner is None:
|
||||
output_text = ""
|
||||
elif isinstance(inner, str):
|
||||
|
|
@ -258,7 +320,7 @@ class LiteLLMAnthropicToResponsesAPIAdapter:
|
|||
{
|
||||
"type": "function_call_output",
|
||||
"call_id": tool_use_id,
|
||||
"output": output_text,
|
||||
"output": self._tool_result_output_value(output_text, tool_file_parts),
|
||||
}
|
||||
)
|
||||
if tool_image_parts:
|
||||
|
|
|
|||
|
|
@ -22,19 +22,7 @@ from litellm.types.llms.openai import (
|
|||
from litellm.types.utils import CallTypes, LlmProviders, ModelResponse
|
||||
|
||||
from ..chat.transformation import AnthropicConfig
|
||||
from ..common_utils import AnthropicModelInfo
|
||||
|
||||
# Map Anthropic error types to HTTP status codes
|
||||
ANTHROPIC_ERROR_STATUS_CODE_MAP: Final = {
|
||||
"invalid_request_error": 400,
|
||||
"authentication_error": 401,
|
||||
"permission_error": 403,
|
||||
"not_found_error": 404,
|
||||
"rate_limit_error": 429,
|
||||
"api_error": 500,
|
||||
"overloaded_error": 503,
|
||||
"timeout_error": 504,
|
||||
}
|
||||
from ..common_utils import ANTHROPIC_ERROR_STATUS_CODE_MAP, AnthropicModelInfo
|
||||
|
||||
|
||||
class AnthropicFilesHandler:
|
||||
|
|
|
|||
|
|
@ -65,15 +65,24 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
|
|||
|
||||
Extracts the actual model used from the Azure response (e.g., gpt-5-nano-2025-08-07)
|
||||
and returns it with the azure_ai/ prefix for proper display and cost tracking.
|
||||
|
||||
Also stamps that model onto ``_hidden_params`` so downstream consumers (spend logs,
|
||||
response restamping) can read it instead of guessing the route from the model string.
|
||||
"""
|
||||
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
|
||||
from litellm.llms.azure_ai.common_utils import (
|
||||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
|
||||
AzureFoundryModelInfo,
|
||||
)
|
||||
from litellm.router_utils.add_retry_fallback_headers import (
|
||||
get_hidden_params_dict,
|
||||
)
|
||||
|
||||
# Get base model for the parent call (strips routing prefixes for API compatibility)
|
||||
base_model: Final[str] = AzureFoundryModelInfo.get_base_model(model)
|
||||
|
||||
# Call parent transform_response first - this will extract the actual model
|
||||
# from the raw response (e.g., "gpt-5-nano-2025-08-07")
|
||||
model_response = super().transform_response(
|
||||
transformed_response: Final = super().transform_response(
|
||||
model=base_model,
|
||||
raw_response=raw_response,
|
||||
model_response=model_response,
|
||||
|
|
@ -86,7 +95,15 @@ class AzureModelRouterConfig(AzureAIStudioConfig):
|
|||
api_key=api_key,
|
||||
json_mode=json_mode,
|
||||
)
|
||||
return model_response
|
||||
selected_model: Final = transformed_response.model
|
||||
if selected_model:
|
||||
# Rebuilt rather than mutated in place: ModelResponseBase declares _hidden_params as a
|
||||
# class-level dict, so an in-place write can bleed into unrelated responses.
|
||||
transformed_response._hidden_params = { # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter # mutable-ok: ModelResponse requires _hidden_params to be a plain dict
|
||||
**get_hidden_params_dict(transformed_response),
|
||||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model,
|
||||
}
|
||||
return transformed_response
|
||||
|
||||
def calculate_additional_costs(self, model: str, prompt_tokens: int, completion_tokens: int) -> dict | None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -51,6 +51,9 @@ def get_azure_ai_auth_headers(
|
|||
)
|
||||
|
||||
|
||||
AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model"
|
||||
|
||||
|
||||
class AzureFoundryModelInfo(BaseLLMModelInfo):
|
||||
"""Model info for Azure AI / Azure Foundry models."""
|
||||
|
||||
|
|
@ -82,6 +85,41 @@ class AzureFoundryModelInfo(BaseLLMModelInfo):
|
|||
return "model_router"
|
||||
return "default"
|
||||
|
||||
@staticmethod
|
||||
def get_model_router_selected_model(hidden_params: Mapping[str, object] | None) -> str | None:
|
||||
"""The model Azure Model Router actually served, stamped by ``AzureModelRouterConfig``.
|
||||
|
||||
Reading this beats re-deriving the route from a model string: the stamp is set on the
|
||||
code path that was actually taken, so it holds no matter what the caller named the model.
|
||||
"""
|
||||
if not hidden_params:
|
||||
return None
|
||||
selected: Final = hidden_params.get(AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY)
|
||||
if isinstance(selected, str) and selected:
|
||||
return selected
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def is_model_router_call(
|
||||
model: str | None = None,
|
||||
hidden_params: Mapping[str, object] | None = None,
|
||||
) -> bool:
|
||||
"""Whether a request went down the Azure Model Router route.
|
||||
|
||||
Prefers the response stamp, then the deployment's litellm model path, and only then the
|
||||
caller-supplied name. The last two go through ``get_azure_ai_route`` so the model-router
|
||||
name heuristic lives in exactly one place.
|
||||
"""
|
||||
if AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) is not None:
|
||||
return True
|
||||
deployment_model: Final = (
|
||||
hidden_params.get("litellm_model_name") or hidden_params.get("model") if hidden_params is not None else None
|
||||
)
|
||||
return any(
|
||||
isinstance(candidate, str) and AzureFoundryModelInfo.get_azure_ai_route(candidate) == "model_router"
|
||||
for candidate in (deployment_model, model)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def get_api_base(api_base: str | None = None) -> str | None:
|
||||
return api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE")
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ class BedrockRerankHandler(BaseAWSLLM):
|
|||
async def arerank(
|
||||
self,
|
||||
prepared_request: BedrockPreparedRequest,
|
||||
logging_obj: LitellmLogging,
|
||||
timeout: float | httpx.Timeout | None = None,
|
||||
client: AsyncHTTPHandler | None = None,
|
||||
):
|
||||
|
|
@ -40,6 +41,7 @@ class BedrockRerankHandler(BaseAWSLLM):
|
|||
headers=dict(prepared_request["prepped"].headers),
|
||||
data=prepared_request["body"],
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
response.raise_for_status()
|
||||
except httpx.HTTPStatusError as err:
|
||||
|
|
@ -98,6 +100,7 @@ class BedrockRerankHandler(BaseAWSLLM):
|
|||
if _is_async:
|
||||
return self.arerank(
|
||||
prepared_request,
|
||||
logging_obj=logging_obj,
|
||||
timeout=timeout,
|
||||
client=client if client is not None and isinstance(client, AsyncHTTPHandler) else None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -15,8 +15,12 @@ role / access key / profile / web identity), signed via the shared
|
|||
BaseAWSLLM._sign_request after the request body is finalized.
|
||||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from typing import Any, Final
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
|
@ -50,6 +54,33 @@ _BEDROCK_MANTLE_SUPPORTED_SERVICE_TIERS: Final = frozenset({"auto", "default"})
|
|||
|
||||
_CODEX_ADDITIONAL_TOOLS_INPUT_ITEM_TYPE: Final = "additional_tools"
|
||||
|
||||
_CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE: Final = "agent_message"
|
||||
_CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE: Final = "context_compaction"
|
||||
_CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE: Final = "local_shell_call"
|
||||
|
||||
|
||||
class _RewrittenOutputTextBlock(TypedDict):
|
||||
type: ReadOnly[str]
|
||||
text: ReadOnly[str]
|
||||
|
||||
|
||||
class _RewrittenAssistantMessageItem(TypedDict):
|
||||
type: ReadOnly[str]
|
||||
role: ReadOnly[str]
|
||||
content: ReadOnly[tuple[_RewrittenOutputTextBlock, ...]]
|
||||
|
||||
|
||||
class _RewrittenCompactionItem(TypedDict):
|
||||
type: ReadOnly[str]
|
||||
encrypted_content: ReadOnly[str]
|
||||
|
||||
|
||||
class _RewrittenFunctionCallItem(TypedDict):
|
||||
type: ReadOnly[str]
|
||||
call_id: ReadOnly[str]
|
||||
name: ReadOnly[str]
|
||||
arguments: ReadOnly[str]
|
||||
|
||||
|
||||
class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPIConfig):
|
||||
def __init__(
|
||||
|
|
@ -155,6 +186,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
|
|||
headers: dict,
|
||||
) -> dict:
|
||||
remaining_input, hoisted_tools = self._hoist_codex_additional_tools(input)
|
||||
normalized_input: Final = self._normalize_codex_input_items(remaining_input)
|
||||
request_params: Final = (
|
||||
{
|
||||
**response_api_optional_request_params,
|
||||
|
|
@ -168,7 +200,7 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
|
|||
)
|
||||
return super().transform_responses_api_request(
|
||||
model=model,
|
||||
input=remaining_input,
|
||||
input=normalized_input,
|
||||
response_api_optional_request_params=request_params,
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
|
|
@ -210,6 +242,91 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
|
|||
)
|
||||
return remaining_input, cls._filter_unsupported_tools(hoisted_tools)
|
||||
|
||||
@staticmethod
|
||||
def _agent_message_text(item: "Mapping[str, Any]") -> str:
|
||||
content: Final = item.get("content")
|
||||
if not isinstance(content, list):
|
||||
return ""
|
||||
return "".join(
|
||||
str(block.get("text") or block.get("encrypted_content") or "")
|
||||
for block in content
|
||||
if isinstance(block, dict)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _normalize_agent_message_item(cls, item: "Mapping[str, Any]") -> "_RewrittenAssistantMessageItem | None":
|
||||
text: Final = cls._agent_message_text(item)
|
||||
if not text:
|
||||
return None
|
||||
rewritten: Final[_RewrittenAssistantMessageItem] = {
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": ({"type": "output_text", "text": text},),
|
||||
}
|
||||
return rewritten
|
||||
|
||||
@staticmethod
|
||||
def _normalize_context_compaction_item(item: "Mapping[str, Any]") -> "_RewrittenCompactionItem | None":
|
||||
encrypted_content: Final = item.get("encrypted_content")
|
||||
if not isinstance(encrypted_content, str) or not encrypted_content:
|
||||
return None
|
||||
rewritten: Final[_RewrittenCompactionItem] = {"type": "compaction", "encrypted_content": encrypted_content}
|
||||
return rewritten
|
||||
|
||||
@staticmethod
|
||||
def _normalize_local_shell_call_item(item: "Mapping[str, Any]") -> "_RewrittenFunctionCallItem | None":
|
||||
call_id: Final = item.get("call_id")
|
||||
if not isinstance(call_id, str) or not call_id:
|
||||
return None
|
||||
action: Final = item.get("action")
|
||||
rewritten: Final[_RewrittenFunctionCallItem] = {
|
||||
"type": "function_call",
|
||||
"call_id": call_id,
|
||||
"name": "local_shell",
|
||||
"arguments": json.dumps(action) if isinstance(action, dict) else "{}",
|
||||
}
|
||||
return rewritten
|
||||
|
||||
@classmethod
|
||||
def _normalize_codex_input_item(cls, item: object) -> "tuple[object, str | None]":
|
||||
"""Returns (normalized item or None to drop it, original type when rewritten)."""
|
||||
if not isinstance(item, dict):
|
||||
return item, None
|
||||
item_type: Final = item.get("type")
|
||||
if item_type == _CODEX_AGENT_MESSAGE_INPUT_ITEM_TYPE:
|
||||
return cls._normalize_agent_message_item(item), item_type
|
||||
if item_type == _CODEX_CONTEXT_COMPACTION_INPUT_ITEM_TYPE:
|
||||
return cls._normalize_context_compaction_item(item), item_type
|
||||
if item_type == _CODEX_LOCAL_SHELL_CALL_INPUT_ITEM_TYPE:
|
||||
return cls._normalize_local_shell_call_item(item), item_type
|
||||
return item, None
|
||||
|
||||
@classmethod
|
||||
def _normalize_codex_input_items(
|
||||
cls,
|
||||
input: "str | ResponseInputParam",
|
||||
) -> "str | ResponseInputParam":
|
||||
"""Rewrite Codex history item types Mantle rejects with 400 "Invalid
|
||||
'input': value did not match any expected variant" into supported
|
||||
equivalents. `agent_message` (Codex multi-agent traffic; its
|
||||
encrypted_content slot carries the plaintext payload when the model
|
||||
never issued encrypted args) becomes an assistant message,
|
||||
`context_compaction` becomes the `compaction` spelling Mantle accepts,
|
||||
and `local_shell_call` becomes the function_call its recorded
|
||||
function_call_output already pairs with.
|
||||
"""
|
||||
if not isinstance(input, list):
|
||||
return input
|
||||
normalized: Final = tuple(cls._normalize_codex_input_item(item) for item in input)
|
||||
rewritten_types: Final = sorted(frozenset(item_type for _, item_type in normalized if item_type is not None))
|
||||
if rewritten_types:
|
||||
verbose_logger.warning(
|
||||
"Bedrock Mantle Responses API: rewrote Codex input item type(s) %s that Mantle rejects.",
|
||||
rewritten_types,
|
||||
)
|
||||
kept: Final = [item for item, _ in normalized if item is not None] # mutable-ok: ResponseInputParam is a list
|
||||
return kept # pyright: ignore[reportReturnType] # Codex passthrough items sit outside the OpenAI input union
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
response_api_optional_params: ResponsesAPIOptionalRequestParams,
|
||||
|
|
|
|||
|
|
@ -68,6 +68,8 @@ class CerebrasConfig(OpenAIGPTConfig):
|
|||
"tool_choice",
|
||||
"tools",
|
||||
"user",
|
||||
"max_retries",
|
||||
"extra_headers",
|
||||
]
|
||||
|
||||
# Only add reasoning_effort for models that support it
|
||||
|
|
|
|||
|
|
@ -1203,6 +1203,7 @@ class BaseLLMHTTPHandler:
|
|||
headers=headers,
|
||||
data=json.dumps(request_data),
|
||||
timeout=timeout,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
except Exception as e:
|
||||
raise self._handle_error(e=e, provider_config=provider_config)
|
||||
|
|
@ -2267,6 +2268,10 @@ class BaseLLMHTTPHandler:
|
|||
AgenticAnthropicStreamingIterator,
|
||||
)
|
||||
|
||||
held_back_tool_names: Final = self._server_fulfilled_tools_in_request(
|
||||
logging_obj=logging_obj,
|
||||
tools=anthropic_messages_optional_request_params.get("tools"),
|
||||
)
|
||||
initial_response = AgenticAnthropicStreamingIterator(
|
||||
completion_stream=completion_stream,
|
||||
http_handler=self,
|
||||
|
|
@ -2277,6 +2282,8 @@ class BaseLLMHTTPHandler:
|
|||
logging_obj=logging_obj,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
kwargs={**kwargs, "api_key": api_key} if api_key else kwargs,
|
||||
hold_back=bool(held_back_tool_names),
|
||||
server_fulfilled_tool_names=held_back_tool_names,
|
||||
)
|
||||
return AnthropicMessagesStreamingResponse(
|
||||
completion_stream=initial_response,
|
||||
|
|
@ -5124,6 +5131,20 @@ class BaseLLMHTTPHandler:
|
|||
return True
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _server_fulfilled_tools_in_request(logging_obj: LiteLLMLoggingObj, tools: object) -> frozenset[str]:
|
||||
"""The request's tools that a registered callback fulfills server-side (e.g. ``headroom_retrieve``)."""
|
||||
if not isinstance(tools, list) or not tools:
|
||||
return frozenset()
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import has_tool_with_name
|
||||
|
||||
return frozenset(
|
||||
name
|
||||
for cb in _custom_logger_callbacks(logging_obj)
|
||||
for name in getattr(cb, "server_fulfilled_tool_names", frozenset())
|
||||
if has_tool_with_name(tools, name)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _check_agentic_loop_safety(
|
||||
tool_calls: object,
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ def _normalize_reasoning_effort_for_chat_completion(
|
|||
) -> str | None:
|
||||
"""Convert reasoning_effort to the string format expected by OpenAI chat completion API.
|
||||
|
||||
The chat completion API expects a simple string: 'none', 'low', 'medium', 'high', or 'xhigh'.
|
||||
The chat completion API expects an effort string such as 'low' or 'high'.
|
||||
Config/deployments may pass the Responses API format: {'effort': 'high', 'summary': 'detailed'}.
|
||||
"""
|
||||
if value is None:
|
||||
|
|
|
|||
|
|
@ -1,58 +0,0 @@
|
|||
"""
|
||||
Support for OpenAI's `/v1/chat/completions` endpoint.
|
||||
|
||||
Calls done in OpenAI/openai.py as TogetherAI is openai-compatible.
|
||||
|
||||
Docs: https://docs.together.ai/reference/completions-1
|
||||
"""
|
||||
|
||||
from typing import Final
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.utils import supports_function_calling
|
||||
|
||||
from ..openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
|
||||
|
||||
class TogetherAIConfig(OpenAIGPTConfig):
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
"""
|
||||
Only some together models support response_format / tool calling
|
||||
|
||||
Docs: https://docs.together.ai/docs/json-mode
|
||||
"""
|
||||
# Use supports_function_calling() — which reads _get_model_info_helper
|
||||
# directly — instead of get_model_info(). get_model_info() calls
|
||||
# get_supported_openai_params() as its first step, which routes back
|
||||
# into this method for together_ai models, creating a recursion that
|
||||
# only terminates when Python's recursion limit or the "not mapped"
|
||||
# exception in _get_model_info_helper is hit (~332 deep calls).
|
||||
supports_fc: bool | None = None
|
||||
try:
|
||||
supports_fc = supports_function_calling(model, custom_llm_provider="together_ai")
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Error getting supported openai params: %s", e)
|
||||
|
||||
optional_params: Final = super().get_supported_openai_params(model)
|
||||
if supports_fc is not True:
|
||||
verbose_logger.debug(
|
||||
"Only some together models support function calling/response_format. Docs - https://docs.together.ai/docs/function-calling"
|
||||
)
|
||||
optional_params.remove("tools")
|
||||
optional_params.remove("tool_choice")
|
||||
optional_params.remove("function_call")
|
||||
optional_params.remove("response_format")
|
||||
return optional_params
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
mapped_openai_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params)
|
||||
|
||||
if "response_format" in mapped_openai_params and mapped_openai_params["response_format"] == {"type": "text"}:
|
||||
mapped_openai_params.pop("response_format")
|
||||
return mapped_openai_params
|
||||
3
litellm/llms/together_ai/chat/__init__.py
Normal file
3
litellm/llms/together_ai/chat/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
from .transformation import TogetherAIChatConfig as TogetherAIChatConfig
|
||||
|
||||
TogetherAIConfig = TogetherAIChatConfig
|
||||
168
litellm/llms/together_ai/chat/transformation.py
Normal file
168
litellm/llms/together_ai/chat/transformation.py
Normal file
|
|
@ -0,0 +1,168 @@
|
|||
"""
|
||||
Translates from OpenAI's `/v1/chat/completions` to Together AI's `/v1/chat/completions`.
|
||||
|
||||
Docs: https://docs.together.ai/docs/chat-overview
|
||||
"""
|
||||
|
||||
from collections.abc import Callable, Container, Coroutine
|
||||
from typing import (
|
||||
Final,
|
||||
Literal,
|
||||
cast, # noqa: TID251 # rebuilding a TypedDict minus keys has no checked spelling
|
||||
overload,
|
||||
)
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.exceptions import UnsupportedParamsError
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.utils import supports_function_calling, supports_response_schema
|
||||
|
||||
from ...openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
|
||||
TOOL_CALLING_PARAMS: Final = ("tools", "tool_choice", "function_call")
|
||||
LITELLM_INTERNAL_ASSISTANT_FIELDS: Final = frozenset({"thinking_blocks", "provider_specific_fields"})
|
||||
FUNCTION_CALLING_DOCS_URL: Final = "https://docs.together.ai/docs/function-calling"
|
||||
STRUCTURED_OUTPUTS_DOCS_URL: Final = "https://docs.together.ai/docs/inference/chat/structured-outputs"
|
||||
|
||||
|
||||
def _registry_verdict(model: str, flag: str, check: Callable[[str], bool]) -> bool | None:
|
||||
try:
|
||||
if check(model):
|
||||
return True
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Error checking together_ai %s for %s: %s", flag, model, e)
|
||||
registry_entry: Final = litellm.model_cost.get(f"together_ai/{model}")
|
||||
if isinstance(registry_entry, dict) and registry_entry.get(flag) is False:
|
||||
return False
|
||||
return None
|
||||
|
||||
|
||||
def _function_calling_verdict(model: str) -> bool | None:
|
||||
return _registry_verdict(
|
||||
model,
|
||||
"supports_function_calling",
|
||||
lambda checked_model: supports_function_calling(checked_model, custom_llm_provider="together_ai"),
|
||||
)
|
||||
|
||||
|
||||
def _response_schema_verdict(model: str) -> bool | None:
|
||||
return _registry_verdict(
|
||||
model,
|
||||
"supports_response_schema",
|
||||
lambda checked_model: supports_response_schema(checked_model, custom_llm_provider="together_ai"),
|
||||
)
|
||||
|
||||
|
||||
def _tool_params_to_drop(passed_params: Container[str], model: str, drop_params: bool) -> tuple[str, ...]:
|
||||
passed_tool_params: Final = tuple(param for param in TOOL_CALLING_PARAMS if param in passed_params)
|
||||
if not passed_tool_params:
|
||||
return ()
|
||||
verdict: Final = _function_calling_verdict(model)
|
||||
if verdict is True:
|
||||
return ()
|
||||
if verdict is None:
|
||||
verbose_logger.warning(
|
||||
"together_ai model %s has no function calling entry in the model registry; passing %s through for Together to validate. Docs - %s",
|
||||
model,
|
||||
", ".join(passed_tool_params),
|
||||
FUNCTION_CALLING_DOCS_URL,
|
||||
)
|
||||
return ()
|
||||
if drop_params or litellm.drop_params:
|
||||
verbose_logger.warning(
|
||||
"together_ai model %s does not support function calling per the model registry; dropping %s. Docs - %s",
|
||||
model,
|
||||
", ".join(passed_tool_params),
|
||||
FUNCTION_CALLING_DOCS_URL,
|
||||
)
|
||||
return passed_tool_params
|
||||
raise UnsupportedParamsError(
|
||||
status_code=500,
|
||||
message=f"together_ai does not support parameters: {', '.join(passed_tool_params)}, for model={model}. To drop it from the call, set `litellm.drop_params = True`.",
|
||||
)
|
||||
|
||||
|
||||
def _drop_response_format(passed_params: Container[str], model: str, drop_params: bool) -> bool:
|
||||
if "response_format" not in passed_params:
|
||||
return False
|
||||
verdict: Final = _response_schema_verdict(model)
|
||||
if verdict is True:
|
||||
return False
|
||||
if verdict is None:
|
||||
verbose_logger.warning(
|
||||
"together_ai model %s has no structured outputs entry in the model registry; passing response_format through for Together to validate. Docs - %s",
|
||||
model,
|
||||
STRUCTURED_OUTPUTS_DOCS_URL,
|
||||
)
|
||||
return False
|
||||
if drop_params or litellm.drop_params:
|
||||
verbose_logger.warning(
|
||||
"together_ai model %s does not support structured outputs per the model registry; dropping response_format. Docs - %s",
|
||||
model,
|
||||
STRUCTURED_OUTPUTS_DOCS_URL,
|
||||
)
|
||||
return True
|
||||
raise UnsupportedParamsError(
|
||||
status_code=500,
|
||||
message=f"together_ai does not support parameters: response_format, for model={model}. To drop it from the call, set `litellm.drop_params = True`.",
|
||||
)
|
||||
|
||||
|
||||
def _without_litellm_internal_fields(message: AllMessageValues) -> AllMessageValues:
|
||||
if message["role"] != "assistant" or LITELLM_INTERNAL_ASSISTANT_FIELDS.isdisjoint(message):
|
||||
return message
|
||||
return cast( # cast-ok: rebuilding the same TypedDict minus internal keys loses the narrowed type
|
||||
"AllMessageValues",
|
||||
{ # mutable-ok: TypedDict rebuild minus internal keys
|
||||
key: value for key, value in message.items() if key not in LITELLM_INTERNAL_ASSISTANT_FIELDS
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class TogetherAIChatConfig(OpenAIGPTConfig):
|
||||
@overload
|
||||
def _transform_messages(
|
||||
self,
|
||||
messages: list[AllMessageValues], # mutable-ok: inherited contract
|
||||
model: str,
|
||||
is_async: Literal[True],
|
||||
) -> Coroutine[object, object, list[AllMessageValues]]: ... # mutable-ok: inherited contract
|
||||
|
||||
@overload
|
||||
def _transform_messages(
|
||||
self,
|
||||
messages: list[AllMessageValues], # mutable-ok: inherited contract
|
||||
model: str,
|
||||
is_async: Literal[False] = False,
|
||||
) -> list[AllMessageValues]: ... # mutable-ok: inherited contract
|
||||
|
||||
def _transform_messages(
|
||||
self,
|
||||
messages: list[AllMessageValues], # mutable-ok: inherited contract
|
||||
model: str,
|
||||
is_async: bool = False,
|
||||
) -> list[AllMessageValues] | Coroutine[object, object, list[AllMessageValues]]: # mutable-ok: inherited contract
|
||||
"""Together consumes replayed assistant `reasoning_content` (preserved thinking via
|
||||
`chat_template_kwargs: {"clear_thinking": false}`), so it must stay in the payload;
|
||||
only litellm-internal fields are stripped before sending."""
|
||||
stripped: Final = [ # mutable-ok: super() requires a list
|
||||
_without_litellm_internal_fields(message) for message in messages
|
||||
]
|
||||
if is_async:
|
||||
return super()._transform_messages(stripped, model, is_async=True)
|
||||
return super()._transform_messages(stripped, model, is_async=False)
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
mapped_openai_params: Final = super().map_openai_params(non_default_params, optional_params, model, drop_params)
|
||||
for param in _tool_params_to_drop(mapped_openai_params, model, drop_params):
|
||||
mapped_openai_params.pop(param)
|
||||
if _drop_response_format(mapped_openai_params, model, drop_params):
|
||||
mapped_openai_params.pop("response_format")
|
||||
return mapped_openai_params
|
||||
|
|
@ -23,6 +23,7 @@ from litellm.constants import (
|
|||
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_FLASH_LITE,
|
||||
DEFAULT_REASONING_EFFORT_MINIMAL_THINKING_BUDGET_GEMINI_2_5_PRO,
|
||||
)
|
||||
from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAccumulator
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
_encode_tool_call_id_with_signature,
|
||||
)
|
||||
|
|
@ -3087,7 +3088,7 @@ class ModelResponseIterator:
|
|||
self.streaming_response = streaming_response
|
||||
self.response = response
|
||||
self.chunk_type: Literal["valid_json", "accumulated_json"] = "valid_json"
|
||||
self.accumulated_json = ""
|
||||
self._json_buffer = JSONFragmentAccumulator()
|
||||
self.sent_first_chunk = False
|
||||
self.logging_obj = logging_obj
|
||||
self.response_headers = response_headers or {}
|
||||
|
|
@ -3095,6 +3096,14 @@ class ModelResponseIterator:
|
|||
self.cumulative_tool_call_index: int = 0
|
||||
self.has_seen_tool_calls: bool = False
|
||||
|
||||
@property
|
||||
def accumulated_json(self) -> str:
|
||||
return self._json_buffer.snapshot()
|
||||
|
||||
@accumulated_json.setter
|
||||
def accumulated_json(self, value: str) -> None:
|
||||
self._json_buffer.set(value)
|
||||
|
||||
@staticmethod
|
||||
def _check_streaming_error(chunk: dict) -> None:
|
||||
"""Detect embedded errors (e.g. 429 RESOURCE_EXHAUSTED) in streaming chunks and raise VertexAIError."""
|
||||
|
|
@ -3298,8 +3307,8 @@ class ModelResponseIterator:
|
|||
return self.chunk_parser(chunk=json_chunk)
|
||||
|
||||
def handle_accumulated_json_chunk(self, chunk: str, is_final: bool = False) -> Optional["ModelResponseStream"]:
|
||||
message: Final = litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or ""
|
||||
self.accumulated_json = (self.accumulated_json + message.replace("\n\n", "")).strip()
|
||||
message: Final = (litellm.CustomStreamWrapper._strip_sse_data_from_chunk(chunk) or "").replace("\n\n", "")
|
||||
self._json_buffer.append(message)
|
||||
|
||||
# Mid-stream, defer parsing until the buffer's last byte can close a value:
|
||||
# attempting a parse after every fragment of one large object is O(n^2) and
|
||||
|
|
@ -3307,27 +3316,23 @@ class ModelResponseIterator:
|
|||
# data is coming, so drain whatever complete values remain regardless of the
|
||||
# trailing byte, otherwise a complete leading value sitting behind a truncated
|
||||
# trailing one would be silently dropped.
|
||||
if not is_final and (not self.accumulated_json or self.accumulated_json[-1] not in "}]"):
|
||||
if not is_final and not self._json_buffer.could_close_json():
|
||||
return None
|
||||
|
||||
# Peel one complete JSON value from the front of the buffer and keep the
|
||||
# unconsumed tail. Running json.loads over the whole buffer would fail
|
||||
# forever once it held more than one concatenated value ("Extra data") while
|
||||
# never resetting the buffer, so the buffer grew without bound and pinned the
|
||||
# core. raw_decode reports where the value ended, so concatenated values drain
|
||||
# one call at a time. A leading non-dict value (never emitted by Gemini in
|
||||
# practice) is consumed and skipped so it cannot block the dict values behind it.
|
||||
decoder: Final = json.JSONDecoder()
|
||||
while self.accumulated_json:
|
||||
try:
|
||||
raw_value = decoder.raw_decode(self.accumulated_json)
|
||||
except json.JSONDecodeError:
|
||||
# core. pop_next_value reports where the value ended, so concatenated values
|
||||
# drain one call at a time. A leading non-dict value (never emitted by Gemini
|
||||
# in practice) is consumed and skipped so it cannot block the dict values
|
||||
# behind it.
|
||||
while True:
|
||||
found, decoded = self._json_buffer.pop_next_value()
|
||||
if not found:
|
||||
return None
|
||||
decoded, end_index = cast("tuple[object, int]", raw_value) # cast-ok: raw_decode -> tuple[Any,int]
|
||||
self.accumulated_json = self.accumulated_json[end_index:].strip()
|
||||
if isinstance(decoded, dict):
|
||||
return self.chunk_parser(chunk=decoded)
|
||||
return None
|
||||
|
||||
def _common_chunk_parsing_logic(self, chunk: str) -> Optional["ModelResponseStream"]:
|
||||
try:
|
||||
|
|
@ -3351,7 +3356,7 @@ class ModelResponseIterator:
|
|||
try:
|
||||
chunk: Final = self.response_iterator.__next__()
|
||||
except StopIteration:
|
||||
if self.chunk_type == "accumulated_json" and self.accumulated_json:
|
||||
if self.chunk_type == "accumulated_json" and self._json_buffer:
|
||||
result: Final = self.handle_accumulated_json_chunk(chunk="", is_final=True)
|
||||
if result is not None:
|
||||
return result
|
||||
|
|
@ -3375,7 +3380,7 @@ class ModelResponseIterator:
|
|||
try:
|
||||
chunk: Final = await self.async_response_iterator.__anext__()
|
||||
except StopAsyncIteration:
|
||||
if self.chunk_type == "accumulated_json" and self.accumulated_json:
|
||||
if self.chunk_type == "accumulated_json" and self._json_buffer:
|
||||
result: Final = self.handle_accumulated_json_chunk(chunk="", is_final=True)
|
||||
if result is not None:
|
||||
return result
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ from concurrent import futures
|
|||
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
|
||||
from copy import deepcopy
|
||||
from functools import partial
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args
|
||||
|
||||
from litellm._logging import _redact_string
|
||||
|
|
@ -416,7 +417,7 @@ async def acompletion(
|
|||
logprobs: bool | None = None,
|
||||
top_logprobs: int | None = None,
|
||||
deployment_id=None,
|
||||
reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None,
|
||||
reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "max", "default"] | None = None,
|
||||
verbosity: Literal["low", "medium", "high"] | None = None,
|
||||
safety_identifier: str | None = None,
|
||||
service_tier: str | None = None,
|
||||
|
|
@ -1811,6 +1812,56 @@ def _complete_fireworks_ai(
|
|||
return response
|
||||
|
||||
|
||||
def _complete_together_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult:
|
||||
acompletion: Final = ctx.acompletion
|
||||
api_base: Final = ctx.api_base
|
||||
api_key: Final = ctx.api_key
|
||||
client: Final = _dispatch_client_http(ctx)
|
||||
custom_llm_provider: Final = ctx.custom_llm_provider
|
||||
headers: Final = ctx.headers
|
||||
litellm_params: Final = ctx.litellm_params
|
||||
logging: Final = ctx.logging
|
||||
messages: Final = ctx.messages
|
||||
model: Final = ctx.model
|
||||
model_response: Final = ctx.model_response
|
||||
optional_params: Final = ctx.optional_params
|
||||
provider_config: Final = ctx.provider_config
|
||||
shared_session: Final = ctx.shared_session
|
||||
stream: Final = ctx.stream
|
||||
timeout: Final = ctx.timeout
|
||||
|
||||
try:
|
||||
response: Final = base_llm_http_handler.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
headers=headers,
|
||||
model_response=model_response,
|
||||
api_key=api_key,
|
||||
api_base=api_base,
|
||||
acompletion=acompletion,
|
||||
logging_obj=logging,
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
shared_session=shared_session,
|
||||
timeout=timeout,
|
||||
client=client,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
encoding=_get_encoding(),
|
||||
stream=stream,
|
||||
provider_config=provider_config,
|
||||
)
|
||||
except Exception as e:
|
||||
logging.post_call(
|
||||
input=messages,
|
||||
api_key=api_key,
|
||||
original_response=str(e),
|
||||
additional_args=MappingProxyType({"headers": headers}),
|
||||
)
|
||||
raise
|
||||
|
||||
return response
|
||||
|
||||
|
||||
def _complete_heroku(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult:
|
||||
acompletion: Final = ctx.acompletion
|
||||
api_base: Final = ctx.api_base
|
||||
|
|
@ -4920,7 +4971,7 @@ def completion(
|
|||
logit_bias: dict | None = None,
|
||||
user: str | None = None,
|
||||
# openai v1.0+ new params
|
||||
reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None = None,
|
||||
reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "max", "default"] | None = None,
|
||||
verbosity: Literal["low", "medium", "high"] | None = None,
|
||||
response_format: dict | type[BaseModel] | None = None,
|
||||
seed: int | None = None,
|
||||
|
|
@ -5600,6 +5651,8 @@ def completion(
|
|||
elif custom_llm_provider == "fireworks_ai":
|
||||
## COMPLETION CALL
|
||||
response = _complete_fireworks_ai(_dispatch_ctx)
|
||||
elif custom_llm_provider == "together_ai":
|
||||
response = _complete_together_ai(_dispatch_ctx)
|
||||
elif custom_llm_provider == "heroku":
|
||||
response = _complete_heroku(_dispatch_ctx)
|
||||
|
||||
|
|
@ -5649,7 +5702,6 @@ def completion(
|
|||
or custom_llm_provider == "volcengine"
|
||||
or custom_llm_provider == "anyscale"
|
||||
or custom_llm_provider == "openai"
|
||||
or custom_llm_provider == "together_ai"
|
||||
or custom_llm_provider == "nebius"
|
||||
or custom_llm_provider == "wandb"
|
||||
or custom_llm_provider == "clarifai"
|
||||
|
|
@ -5699,14 +5751,6 @@ def completion(
|
|||
response = _complete_openrouter(_dispatch_ctx)
|
||||
elif custom_llm_provider == "vercel_ai_gateway":
|
||||
response = _complete_vercel_ai_gateway(_dispatch_ctx)
|
||||
elif (
|
||||
custom_llm_provider == "together_ai"
|
||||
or ("togethercomputer" in model)
|
||||
or (model in litellm.together_ai_models)
|
||||
):
|
||||
"""
|
||||
Deprecated. We now do together ai calls via the openai client - https://docs.together.ai/docs/openai-api-compatibility
|
||||
"""
|
||||
elif custom_llm_provider == "palm":
|
||||
raise ValueError(
|
||||
"Palm was decommisioned on October 2024. Please use the `gemini/` route for Gemini Google AI Studio Models. Announcement: https://ai.google.dev/palm_docs/palm?hl=en"
|
||||
|
|
|
|||
|
|
@ -37916,6 +37916,7 @@
|
|||
"output_cost_per_token": 1e-07
|
||||
},
|
||||
"together_ai/Qwen/Qwen2.5-72B-Instruct-Turbo": {
|
||||
"deprecation_date": "2026-02-06",
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37932,6 +37933,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/Qwen/Qwen3-235B-A22B-Instruct-2507-tput": {
|
||||
"deprecation_date": "2026-07-10",
|
||||
"input_cost_per_token": 2e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 262000,
|
||||
|
|
@ -37944,6 +37946,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/Qwen/Qwen3-235B-A22B-Thinking-2507": {
|
||||
"deprecation_date": "2026-04-16",
|
||||
"input_cost_per_token": 6.5e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 256000,
|
||||
|
|
@ -37956,6 +37959,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/Qwen/Qwen3-235B-A22B-fp8-tput": {
|
||||
"deprecation_date": "2026-02-06",
|
||||
"input_cost_per_token": 2e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 40000,
|
||||
|
|
@ -37967,6 +37971,7 @@
|
|||
"supports_tool_choice": false
|
||||
},
|
||||
"together_ai/Qwen/Qwen3-Coder-480B-A35B-Instruct-FP8": {
|
||||
"deprecation_date": "2026-06-04",
|
||||
"input_cost_per_token": 2e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 256000,
|
||||
|
|
@ -37979,11 +37984,15 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/deepseek-ai/DeepSeek-R1": {
|
||||
"deprecation_date": "2026-05-14",
|
||||
"input_cost_per_token": 3e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 128000,
|
||||
"max_output_tokens": 20480,
|
||||
"max_tokens": 20480,
|
||||
"metadata": {
|
||||
"successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro"
|
||||
},
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 7e-06,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -37992,6 +38001,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/deepseek-ai/DeepSeek-R1-0528-tput": {
|
||||
"deprecation_date": "2026-02-03",
|
||||
"input_cost_per_token": 5.5e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 128000,
|
||||
|
|
@ -38009,6 +38019,9 @@
|
|||
"max_input_tokens": 65536,
|
||||
"max_output_tokens": 8192,
|
||||
"max_tokens": 8192,
|
||||
"metadata": {
|
||||
"successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro"
|
||||
},
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.25e-06,
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -38017,9 +38030,13 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/deepseek-ai/DeepSeek-V3.1": {
|
||||
"deprecation_date": "2026-05-14",
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_tokens": 16384,
|
||||
"metadata": {
|
||||
"successor": "together_ai/deepseek-ai/DeepSeek-V4-Pro"
|
||||
},
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.7e-06,
|
||||
"source": "https://www.together.ai/models/deepseek-v3-1",
|
||||
|
|
@ -38031,6 +38048,7 @@
|
|||
"max_output_tokens": 16384
|
||||
},
|
||||
"together_ai/meta-llama/Llama-3.2-3B-Instruct-Turbo": {
|
||||
"deprecation_date": "2026-03-06",
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -38039,16 +38057,21 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo": {
|
||||
"input_cost_per_token": 8.8e-07,
|
||||
"input_cost_per_token": 1.04e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 8.8e-07,
|
||||
"output_cost_per_token": 1.04e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/meta-llama/Llama-3.3-70B-Instruct-Turbo-Free": {
|
||||
"deprecation_date": "2025-11-13",
|
||||
"input_cost_per_token": 0,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
|
|
@ -38059,6 +38082,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/meta-llama/Llama-4-Maverick-17B-128E-Instruct-FP8": {
|
||||
"deprecation_date": "2026-03-31",
|
||||
"input_cost_per_token": 2.7e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
|
|
@ -38069,6 +38093,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/meta-llama/Llama-4-Scout-17B-16E-Instruct": {
|
||||
"deprecation_date": "2026-02-06",
|
||||
"input_cost_per_token": 1.8e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
|
|
@ -38079,6 +38104,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo": {
|
||||
"deprecation_date": "2026-02-06",
|
||||
"input_cost_per_token": 3.5e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
|
|
@ -38089,6 +38115,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo": {
|
||||
"deprecation_date": "2026-02-25",
|
||||
"input_cost_per_token": 8.8e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
|
|
@ -38099,6 +38126,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo": {
|
||||
"deprecation_date": "2026-03-06",
|
||||
"input_cost_per_token": 1.8e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
|
|
@ -38109,6 +38137,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/mistralai/Mistral-7B-Instruct-v0.1": {
|
||||
"deprecation_date": "2025-11-13",
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -38117,6 +38146,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/mistralai/Mistral-Small-24B-Instruct-2501": {
|
||||
"deprecation_date": "2026-04-02",
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
"supports_function_calling": true,
|
||||
|
|
@ -38124,6 +38154,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/mistralai/Mixtral-8x7B-Instruct-v0.1": {
|
||||
"deprecation_date": "2026-04-16",
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"mode": "chat",
|
||||
|
|
@ -38136,6 +38167,9 @@
|
|||
"together_ai/moonshotai/Kimi-K2-Instruct": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"metadata": {
|
||||
"successor": "together_ai/moonshotai/Kimi-K3"
|
||||
},
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-06,
|
||||
"source": "https://www.together.ai/models/kimi-k2-instruct",
|
||||
|
|
@ -38179,6 +38213,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/zai-org/GLM-4.5-Air-FP8": {
|
||||
"deprecation_date": "2026-04-02",
|
||||
"input_cost_per_token": 2e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 128000,
|
||||
|
|
@ -38196,6 +38231,9 @@
|
|||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 200000,
|
||||
"max_tokens": 200000,
|
||||
"metadata": {
|
||||
"successor": "together_ai/zai-org/GLM-5.2"
|
||||
},
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.2e-06,
|
||||
"source": "https://www.together.ai/models/glm-4-6",
|
||||
|
|
@ -38205,11 +38243,15 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/zai-org/GLM-4.7": {
|
||||
"deprecation_date": "2026-04-02",
|
||||
"input_cost_per_token": 4.5e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 200000,
|
||||
"max_output_tokens": 200000,
|
||||
"max_tokens": 200000,
|
||||
"metadata": {
|
||||
"successor": "together_ai/zai-org/GLM-5.2"
|
||||
},
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2e-06,
|
||||
"source": "https://www.together.ai/models/glm-4-7",
|
||||
|
|
@ -38219,11 +38261,15 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/moonshotai/Kimi-K2.5": {
|
||||
"deprecation_date": "2026-05-21",
|
||||
"input_cost_per_token": 5e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 256000,
|
||||
"max_output_tokens": 256000,
|
||||
"max_tokens": 256000,
|
||||
"metadata": {
|
||||
"successor": "together_ai/moonshotai/Kimi-K3"
|
||||
},
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.8e-06,
|
||||
"source": "https://www.together.ai/models/kimi-k2-5",
|
||||
|
|
@ -38233,9 +38279,13 @@
|
|||
"supports_reasoning": true
|
||||
},
|
||||
"together_ai/moonshotai/Kimi-K2-Instruct-0905": {
|
||||
"deprecation_date": "2026-03-06",
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 262144,
|
||||
"metadata": {
|
||||
"successor": "together_ai/moonshotai/Kimi-K3"
|
||||
},
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-06,
|
||||
"source": "https://www.together.ai/models/kimi-k2-0905",
|
||||
|
|
@ -38244,9 +38294,13 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/Qwen/Qwen3-Next-80B-A3B-Instruct": {
|
||||
"deprecation_date": "2026-04-02",
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 262144,
|
||||
"metadata": {
|
||||
"successor": "together_ai/Qwen/Qwen3.7-Plus"
|
||||
},
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-06,
|
||||
"source": "https://www.together.ai/models/qwen3-next-80b-a3b-instruct",
|
||||
|
|
@ -38256,9 +38310,13 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/Qwen/Qwen3-Next-80B-A3B-Thinking": {
|
||||
"deprecation_date": "2026-02-25",
|
||||
"input_cost_per_token": 1.5e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 262144,
|
||||
"metadata": {
|
||||
"successor": "together_ai/Qwen/Qwen3.6-Plus"
|
||||
},
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-06,
|
||||
"source": "https://www.together.ai/models/qwen3-next-80b-a3b-thinking",
|
||||
|
|
@ -38268,6 +38326,7 @@
|
|||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/Qwen/Qwen3.5-397B-A17B": {
|
||||
"deprecation_date": "2026-06-29",
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 262144,
|
||||
|
|
@ -38279,6 +38338,292 @@
|
|||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/MiniMaxAI/MiniMax-M3": {
|
||||
"input_cost_per_token": 3e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 524288,
|
||||
"max_output_tokens": 524288,
|
||||
"max_tokens": 524288,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"together_ai/Prism-ML/Ternary-Bonsai-27B": {
|
||||
"input_cost_per_token": 0.0,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 0.0,
|
||||
"source": "https://docs.together.ai/docs/serverless-models"
|
||||
},
|
||||
"together_ai/Qwen/Qwen3.5-9B": {
|
||||
"input_cost_per_token": 1.7e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.5e-07,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"together_ai/Qwen/Qwen3.6-Plus": {
|
||||
"input_cost_per_token": 5e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 1000000,
|
||||
"max_tokens": 1000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_reasoning": true
|
||||
},
|
||||
"together_ai/Qwen/Qwen3.7-Max": {
|
||||
"input_cost_per_token": 1.25e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 1000000,
|
||||
"max_tokens": 1000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3.75e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models"
|
||||
},
|
||||
"together_ai/Qwen/Qwen3.7-Plus": {
|
||||
"input_cost_per_token": 3.2e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 1000000,
|
||||
"max_tokens": 1000000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.28e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models"
|
||||
},
|
||||
"together_ai/Qwen/Qwen3.8-2.4T-A95B": {
|
||||
"input_cost_per_token": 2.5e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1010000,
|
||||
"max_output_tokens": 1010000,
|
||||
"max_tokens": 1010000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 6.25e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models"
|
||||
},
|
||||
"together_ai/arize-ai/qwen-2-1.5b-instruct": {
|
||||
"input_cost_per_token": 1e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 32768,
|
||||
"max_output_tokens": 32768,
|
||||
"max_tokens": 32768,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1e-07,
|
||||
"source": "https://docs.together.ai/docs/serverless-models"
|
||||
},
|
||||
"together_ai/deepseek-ai/DeepSeek-V4-Flash-0731": {
|
||||
"input_cost_per_token": 1.4e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.8e-07,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/deepseek-ai/DeepSeek-V4-Pro": {
|
||||
"input_cost_per_token": 1.74e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 512000,
|
||||
"max_output_tokens": 512000,
|
||||
"max_tokens": 512000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3.48e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/deepseek-ai/DeepSeek-V4-Pro-0813": {
|
||||
"input_cost_per_token": 1.32e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3.96e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/google/gemma-3n-E4B-it": {
|
||||
"input_cost_per_token": 6e-08,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 32768,
|
||||
"max_output_tokens": 32768,
|
||||
"max_tokens": 32768,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-07,
|
||||
"source": "https://docs.together.ai/docs/serverless-models"
|
||||
},
|
||||
"together_ai/google/gemma-4-31B-it": {
|
||||
"input_cost_per_token": 3.9e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 9.7e-07,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"together_ai/intfloat/multilingual-e5-large-instruct": {
|
||||
"input_cost_per_token": 2e-08,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 514,
|
||||
"max_tokens": 514,
|
||||
"mode": "embedding",
|
||||
"output_cost_per_token": 2e-08,
|
||||
"output_vector_size": 1024,
|
||||
"source": "https://docs.together.ai/docs/serverless-models"
|
||||
},
|
||||
"together_ai/meta-llama/Llama-Guard-4-12B": {
|
||||
"input_cost_per_token": 2e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2e-07,
|
||||
"source": "https://docs.together.ai/docs/serverless-models"
|
||||
},
|
||||
"together_ai/meta-models/Muse-Glimmer-30B": {
|
||||
"input_cost_per_token": 3.5e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 131072,
|
||||
"max_output_tokens": 131072,
|
||||
"max_tokens": 131072,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models"
|
||||
},
|
||||
"together_ai/moonshotai/Kimi-K2.7-Code": {
|
||||
"input_cost_per_token": 9.5e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"together_ai/moonshotai/Kimi-K3": {
|
||||
"input_cost_per_token": 3e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1048576,
|
||||
"max_output_tokens": 1048576,
|
||||
"max_tokens": 1048576,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.5e-05,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true
|
||||
},
|
||||
"together_ai/nvidia/nemotron-3-ultra-550b-a55b": {
|
||||
"input_cost_per_token": 6e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 512288,
|
||||
"max_output_tokens": 512288,
|
||||
"max_tokens": 512288,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 3.6e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/pearl-ai/gemma-4-31b-it": {
|
||||
"input_cost_per_token": 2.8e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 262144,
|
||||
"max_output_tokens": 262144,
|
||||
"max_tokens": 262144,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 8.6e-07,
|
||||
"source": "https://docs.together.ai/docs/serverless-models"
|
||||
},
|
||||
"together_ai/thinkingmachines/Inkling": {
|
||||
"input_cost_per_token": 1e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 524288,
|
||||
"max_output_tokens": 524288,
|
||||
"max_tokens": 524288,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.05e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"together_ai/thinkingmachines/Inkling-Small": {
|
||||
"input_cost_per_token": 5e-07,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 524288,
|
||||
"max_output_tokens": 524288,
|
||||
"max_tokens": 524288,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 1.2e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models"
|
||||
},
|
||||
"together_ai/zai-org/GLM-5.2": {
|
||||
"input_cost_per_token": 1.4e-06,
|
||||
"litellm_provider": "together_ai",
|
||||
"max_input_tokens": 1048575,
|
||||
"max_output_tokens": 1048575,
|
||||
"max_tokens": 1048575,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 4.4e-06,
|
||||
"source": "https://docs.together.ai/docs/serverless-models",
|
||||
"supports_function_calling": true,
|
||||
"supports_parallel_function_calling": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true
|
||||
},
|
||||
"tts-1": {
|
||||
"input_cost_per_character": 1.5e-05,
|
||||
"litellm_provider": "openai",
|
||||
|
|
|
|||
|
|
@ -64,8 +64,8 @@ class TeamBase(LiteLLMPydanticObjectBase):
|
|||
team_alias: str | None = None
|
||||
team_id: str | None = None
|
||||
organization_id: str | None = None
|
||||
admins: list = []
|
||||
members: list = []
|
||||
admins: list[str] = []
|
||||
members: list[str] = []
|
||||
members_with_roles: list[Member] = []
|
||||
team_member_permissions: list[str] | None = None
|
||||
metadata: dict | None = None
|
||||
|
|
@ -75,7 +75,7 @@ class TeamBase(LiteLLMPydanticObjectBase):
|
|||
soft_budget: float | None = None
|
||||
budget_duration: str | None = None
|
||||
budget_limits: list[BudgetLimitEntry] | None = None
|
||||
models: list = []
|
||||
models: list[str] = []
|
||||
blocked: bool = False
|
||||
router_settings: dict | None = None
|
||||
access_group_ids: list[str] | None = None
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import hashlib
|
|||
import json
|
||||
from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, TypeVar, cast
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypedDict, cast
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -13,7 +13,6 @@ from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
|||
from litellm.proxy._experimental.mcp_server.oauth_utils import build_upstream_oauth2_token_request
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_MCPServerTable,
|
||||
LiteLLM_ObjectPermissionTable,
|
||||
MCPApprovalStatus,
|
||||
MCPEnvVar,
|
||||
MCPEnvVarScope,
|
||||
|
|
@ -30,6 +29,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import (
|
||||
MCPServerOAuthClientRepository,
|
||||
MCPServerRepository,
|
||||
|
|
@ -48,34 +48,9 @@ if TYPE_CHECKING:
|
|||
|
||||
from litellm.types.mcp_server.mcp_server_manager import MCPServer
|
||||
|
||||
_RowT = TypeVar("_RowT")
|
||||
|
||||
|
||||
class _TableActions(Protocol[_RowT]):
|
||||
async def find_unique(
|
||||
self, where: Mapping[str, object], include: Mapping[str, object] | None = None
|
||||
) -> _RowT | None: ...
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
take: int | None = None,
|
||||
where: Mapping[str, object] | None = None,
|
||||
order: Mapping[str, object] | None = None,
|
||||
) -> list[_RowT]: ...
|
||||
|
||||
async def create(self, data: Mapping[str, object]) -> _RowT: ...
|
||||
|
||||
async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT: ...
|
||||
|
||||
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT | None: ...
|
||||
|
||||
async def delete(self, where: Mapping[str, object]) -> _RowT | None: ...
|
||||
|
||||
async def delete_many(self, where: Mapping[str, object] | None = None) -> int: ...
|
||||
|
||||
|
||||
class _UserEnvVarsTransactionClient(Protocol):
|
||||
litellm_mcpuserenvvars: "_TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]"
|
||||
litellm_mcpuserenvvars: "TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]"
|
||||
|
||||
async def execute_raw(self, query: str, *args: object) -> int: ...
|
||||
|
||||
|
|
@ -473,15 +448,15 @@ def _credentials_blob_to_mutable_dict(blob: str | Mapping[str, object]) -> dict[
|
|||
|
||||
def _mcp_server_table_actions(
|
||||
prisma_client: PrismaClient,
|
||||
) -> "_TableActions[prisma_db_models.LiteLLM_MCPServerTable]":
|
||||
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPServerTable]] = MCPServerRepository(prisma_client).table
|
||||
) -> "TableActions[prisma_db_models.LiteLLM_MCPServerTable]":
|
||||
table: Final[TableActions[prisma_db_models.LiteLLM_MCPServerTable]] = MCPServerRepository(prisma_client).table
|
||||
return table
|
||||
|
||||
|
||||
def _verification_token_table_actions(
|
||||
prisma_client: PrismaClient,
|
||||
) -> "_TableActions[prisma_db_models.LiteLLM_VerificationToken]":
|
||||
table: Final[_TableActions[prisma_db_models.LiteLLM_VerificationToken]] = VerificationTokenRepository(
|
||||
) -> "TableActions[prisma_db_models.LiteLLM_VerificationToken]":
|
||||
table: Final[TableActions[prisma_db_models.LiteLLM_VerificationToken]] = VerificationTokenRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return table
|
||||
|
|
@ -489,15 +464,15 @@ def _verification_token_table_actions(
|
|||
|
||||
def _team_table_actions(
|
||||
prisma_client: PrismaClient,
|
||||
) -> "_TableActions[prisma_db_models.LiteLLM_TeamTable]":
|
||||
table: Final[_TableActions[prisma_db_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
|
||||
) -> "TableActions[prisma_db_models.LiteLLM_TeamTable]":
|
||||
table: Final[TableActions[prisma_db_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
|
||||
return table
|
||||
|
||||
|
||||
def _oauth_client_table_actions(
|
||||
prisma_client: PrismaClient,
|
||||
) -> "_TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]":
|
||||
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = MCPServerOAuthClientRepository(
|
||||
) -> "TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]":
|
||||
table: Final[TableActions[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = MCPServerOAuthClientRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return table
|
||||
|
|
@ -511,7 +486,7 @@ def _db_transaction_manager(prisma_client: PrismaClient) -> _UserEnvVarsTransact
|
|||
async def _db_find_mcp_server_rows(
|
||||
prisma_client: PrismaClient,
|
||||
where: "prisma_db_types.LiteLLM_MCPServerTableWhereInput | None" = None,
|
||||
) -> "list[prisma_db_models.LiteLLM_MCPServerTable]":
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_MCPServerTable]":
|
||||
return await _mcp_server_table_actions(prisma_client).find_many(where=where)
|
||||
|
||||
|
||||
|
|
@ -526,17 +501,19 @@ async def _db_update_mcp_server_row(
|
|||
server_id: str,
|
||||
data: "prisma_db_types.LiteLLM_MCPServerTableUpdateInput",
|
||||
) -> "prisma_db_models.LiteLLM_MCPServerTable":
|
||||
row: Final[prisma_db_models.LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.update(
|
||||
row: Final[prisma_db_models.LiteLLM_MCPServerTable | None] = await _mcp_server_table_actions(prisma_client).update(
|
||||
where={"server_id": server_id},
|
||||
data=data,
|
||||
)
|
||||
if row is None:
|
||||
raise ValueError(f"MCP server not found, passed server_id={server_id}")
|
||||
return row
|
||||
|
||||
|
||||
def _user_credential_actions(
|
||||
prisma_client: PrismaClient,
|
||||
) -> "_TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]":
|
||||
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]] = MCPUserCredentialsRepository(
|
||||
) -> "TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]":
|
||||
table: Final[TableActions[prisma_db_models.LiteLLM_MCPUserCredentials]] = MCPUserCredentialsRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return table
|
||||
|
|
@ -544,8 +521,8 @@ def _user_credential_actions(
|
|||
|
||||
def _user_env_var_actions(
|
||||
prisma_client: PrismaClient,
|
||||
) -> "_TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]":
|
||||
table: Final[_TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]] = prisma_client.db.litellm_mcpuserenvvars
|
||||
) -> "TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]":
|
||||
table: Final[TableActions[prisma_db_models.LiteLLM_MCPUserEnvVars]] = prisma_client.db.litellm_mcpuserenvvars
|
||||
return table
|
||||
|
||||
|
||||
|
|
@ -560,7 +537,7 @@ async def _db_find_user_credential_row(
|
|||
async def _db_find_user_credential_rows(
|
||||
prisma_client: PrismaClient,
|
||||
where: "prisma_db_types.LiteLLM_MCPUserCredentialsWhereInput | None" = None,
|
||||
) -> "list[prisma_db_models.LiteLLM_MCPUserCredentials]":
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserCredentials]":
|
||||
return await _user_credential_actions(prisma_client).find_many(where=where)
|
||||
|
||||
|
||||
|
|
@ -583,7 +560,7 @@ async def _db_upsert_user_credential_row(
|
|||
async def _db_find_user_env_var_rows(
|
||||
prisma_client: PrismaClient,
|
||||
where: "prisma_db_types.LiteLLM_MCPUserEnvVarsWhereInput | None" = None,
|
||||
) -> "list[prisma_db_models.LiteLLM_MCPUserEnvVars]":
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_MCPUserEnvVars]":
|
||||
return await _user_env_var_actions(prisma_client).find_many(where=where)
|
||||
|
||||
|
||||
|
|
@ -658,7 +635,7 @@ async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]
|
|||
"""
|
||||
Returns the matching mcp servers from the db with the server_ids
|
||||
"""
|
||||
_mcp_servers: Final[list[prisma_db_models.LiteLLM_MCPServerTable]] = await _mcp_server_table_actions(
|
||||
_mcp_servers: Final[Sequence[prisma_db_models.LiteLLM_MCPServerTable]] = await _mcp_server_table_actions(
|
||||
prisma_client
|
||||
).find_many(
|
||||
where={
|
||||
|
|
@ -745,13 +722,13 @@ async def get_all_mcp_servers_for_user(
|
|||
|
||||
async def get_objectpermissions_for_mcp_server(
|
||||
prisma_client: PrismaClient, mcp_server_id: str
|
||||
) -> list[LiteLLM_ObjectPermissionTable]:
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_ObjectPermissionTable]":
|
||||
"""
|
||||
Get all the object permissions records and the associated team and verficiationtoken records that have access to the mcp server
|
||||
"""
|
||||
object_permission_records: Final[list[LiteLLM_ObjectPermissionTable]] = await ObjectPermissionRepository(
|
||||
prisma_client
|
||||
).table.find_many(
|
||||
object_permission_records: Final[
|
||||
Sequence[prisma_db_models.LiteLLM_ObjectPermissionTable]
|
||||
] = await ObjectPermissionRepository(prisma_client).table.find_many(
|
||||
where={
|
||||
"mcp_servers": {"has": mcp_server_id},
|
||||
},
|
||||
|
|
@ -766,19 +743,19 @@ async def get_objectpermissions_for_mcp_server(
|
|||
|
||||
async def get_virtualkeys_for_mcp_server(
|
||||
prisma_client: PrismaClient, server_id: str
|
||||
) -> "list[prisma_db_models.LiteLLM_VerificationToken]":
|
||||
) -> "Sequence[prisma_db_models.LiteLLM_VerificationToken]":
|
||||
"""
|
||||
Get all the virtual keys that have access to the mcp server
|
||||
"""
|
||||
virtual_keys: Final[list[prisma_db_models.LiteLLM_VerificationToken] | None] = await VerificationTokenRepository(
|
||||
prisma_client
|
||||
).table.find_many(
|
||||
virtual_keys: Final[
|
||||
Sequence[prisma_db_models.LiteLLM_VerificationToken] | None
|
||||
] = await VerificationTokenRepository(prisma_client).table.find_many(
|
||||
where={
|
||||
"mcp_servers": {"has": server_id},
|
||||
},
|
||||
)
|
||||
|
||||
if virtual_keys is None:
|
||||
if virtual_keys is None: # pyright: ignore[reportUnnecessaryComparison] # unreachable per seam types; kept as-is
|
||||
return []
|
||||
return virtual_keys
|
||||
|
||||
|
|
@ -860,7 +837,7 @@ async def delete_mcp_server(
|
|||
invalidate_token_cache = global_mcp_server_manager.invalidate_user_oauth_token_cache
|
||||
for user_id in credential_user_ids:
|
||||
await invalidate_token_cache(user_id, server_id)
|
||||
return deleted_server
|
||||
return deleted_server # pyright: ignore[reportReturnType] # prisma row, not domain LiteLLM_MCPServerTable
|
||||
|
||||
|
||||
async def create_mcp_server(
|
||||
|
|
@ -880,7 +857,7 @@ async def create_mcp_server(
|
|||
data_dict["updated_by"] = touched_by
|
||||
|
||||
new_mcp_server: Final[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.create(
|
||||
data=data_dict
|
||||
data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable
|
||||
)
|
||||
|
||||
_decrypt_env_vars_on_returned_row(new_mcp_server)
|
||||
|
|
@ -982,7 +959,7 @@ async def update_mcp_server(
|
|||
data: UpdateMCPServerRequest,
|
||||
touched_by: str,
|
||||
fields_set: set[str] | None = None,
|
||||
) -> LiteLLM_MCPServerTable:
|
||||
) -> LiteLLM_MCPServerTable | None:
|
||||
"""
|
||||
Update a new mcp server record in the db
|
||||
"""
|
||||
|
|
@ -1093,9 +1070,9 @@ async def update_mcp_server(
|
|||
|
||||
data_dict["credentials"] = Json(None)
|
||||
|
||||
updated_mcp_server: Final[LiteLLM_MCPServerTable] = await MCPServerRepository(prisma_client).table.update(
|
||||
updated_mcp_server: Final[LiteLLM_MCPServerTable | None] = await MCPServerRepository(prisma_client).table.update(
|
||||
where={"server_id": data.server_id},
|
||||
data=data_dict,
|
||||
data=data_dict, # pyright: ignore[reportAssignmentType] # prisma row, not domain LiteLLM_MCPServerTable
|
||||
)
|
||||
|
||||
_decrypt_env_vars_on_returned_row(updated_mcp_server)
|
||||
|
|
@ -1181,7 +1158,7 @@ async def rotate_mcp_server_credentials_master_key(prisma_client: PrismaClient,
|
|||
)
|
||||
updated += 1
|
||||
|
||||
oauth_clients: Final[list[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = await _oauth_client_table_actions(
|
||||
oauth_clients: Final[Sequence[prisma_db_models.LiteLLM_MCPServerOAuthClient]] = await _oauth_client_table_actions(
|
||||
prisma_client
|
||||
).find_many()
|
||||
oauth_updated = 0
|
||||
|
|
@ -1914,7 +1891,7 @@ async def get_mcp_submissions(
|
|||
along with a summary count breakdown by approval_status.
|
||||
Mirrors get_guardrail_submissions() from guardrail_endpoints.py.
|
||||
"""
|
||||
rows: Final[list[prisma_db_models.LiteLLM_MCPServerTable]] = await _mcp_server_table_actions(
|
||||
rows: Final[Sequence[prisma_db_models.LiteLLM_MCPServerTable]] = await _mcp_server_table_actions(
|
||||
prisma_client
|
||||
).find_many(
|
||||
where={"submitted_at": {"not": None}},
|
||||
|
|
|
|||
|
|
@ -3546,6 +3546,8 @@ class SpendLogsMetadata(TypedDict):
|
|||
litellm_overhead_time_ms: float | None # LiteLLM overhead time in milliseconds
|
||||
attempted_retries: int | None # Number of retries attempted (0 = first attempt succeeded)
|
||||
max_retries: int | None # Max retries configured for this request
|
||||
attempted_fallbacks: ReadOnly[int | None] # Number of fallbacks attempted (0 = primary model group served)
|
||||
original_model_group: ReadOnly[str | None] # Model group requested before any fallbacks
|
||||
cost_breakdown: CostBreakdown | None # Detailed cost breakdown (input_cost, output_cost, margin, discount, etc.)
|
||||
compression_savings: CompressionSavingsMetadata | None
|
||||
autorouter_savings: ReadOnly[float | None] # stamped by the logging payload; None = not auto-routed
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ import json
|
|||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final, NamedTuple, Protocol, TypedDict
|
||||
from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
|
|
@ -12,9 +12,13 @@ from litellm.proxy.management_helpers.object_permission_utils import (
|
|||
handle_update_object_permission_common,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import AgentsRepository, ObjectPermissionRepository
|
||||
from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
|
||||
class AgentObjectPermissionRecord(Protocol):
|
||||
def model_dump(self) -> dict[str, object]: ...
|
||||
|
|
@ -42,11 +46,20 @@ class AgentRecordDump(TypedDict):
|
|||
|
||||
|
||||
class AgentRecord(Protocol):
|
||||
agent_id: str
|
||||
agent_name: str
|
||||
object_permission_id: str | None
|
||||
object_permission: AgentObjectPermissionRecord | None
|
||||
spend: float
|
||||
@property
|
||||
def agent_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def agent_name(self) -> str: ...
|
||||
|
||||
@property
|
||||
def object_permission_id(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def object_permission(self) -> AgentObjectPermissionRecord | None: ...
|
||||
|
||||
@property
|
||||
def spend(self) -> float: ...
|
||||
|
||||
def model_dump(self) -> AgentRecordDump: ...
|
||||
|
||||
|
|
@ -57,50 +70,47 @@ class AgentTableClient(Protocol):
|
|||
async def create(
|
||||
self,
|
||||
data: Mapping[str, object],
|
||||
include: Mapping[str, bool] | None = None,
|
||||
include: Mapping[str, object] | None = None,
|
||||
) -> AgentRecord: ...
|
||||
|
||||
async def find_unique(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
include: Mapping[str, bool] | None = None,
|
||||
include: Mapping[str, object] | None = None,
|
||||
) -> AgentRecord | None: ...
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
where: Mapping[str, object] | None = None,
|
||||
order: Mapping[str, str] | None = None,
|
||||
include: Mapping[str, bool] | None = None,
|
||||
include: Mapping[str, object] | None = None,
|
||||
) -> Sequence[AgentRecord]: ...
|
||||
|
||||
async def update(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
data: Mapping[str, object],
|
||||
include: Mapping[str, bool] | None = None,
|
||||
) -> AgentRecord: ...
|
||||
where: Mapping[str, object],
|
||||
include: Mapping[str, object] | None = None,
|
||||
) -> AgentRecord | None: ...
|
||||
|
||||
async def delete(self, where: Mapping[str, object]) -> AgentRecord: ...
|
||||
async def delete(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
include: Mapping[str, object] | None = None,
|
||||
) -> AgentRecord | None: ...
|
||||
|
||||
|
||||
def agents_table(prisma_client: PrismaClient) -> AgentTableClient:
|
||||
table: Final[AgentTableClient] = AgentsRepository(prisma_client).table
|
||||
table: Final[AgentTableClient] = AgentsRepository(prisma_client).table # pyright: ignore[reportAssignmentType] # prisma rows type model_dump() as dict[str, Any]
|
||||
return table
|
||||
|
||||
|
||||
class ObjectPermissionGrantRecord(Protocol):
|
||||
object_permission_id: str
|
||||
agents: list[str] | None
|
||||
|
||||
|
||||
class ObjectPermissionTableClient(Protocol):
|
||||
async def find_many(self, where: Mapping[str, object]) -> Sequence[ObjectPermissionGrantRecord]: ...
|
||||
|
||||
async def update_many(self, where: Mapping[str, object], data: Mapping[str, object]) -> int: ...
|
||||
|
||||
|
||||
def object_permission_table(prisma_client: PrismaClient) -> ObjectPermissionTableClient:
|
||||
table: Final[ObjectPermissionTableClient] = ObjectPermissionRepository(prisma_client).table
|
||||
def object_permission_table(
|
||||
prisma_client: PrismaClient,
|
||||
) -> "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]":
|
||||
table: Final[TableActions[prisma_models.LiteLLM_ObjectPermissionTable]] = ObjectPermissionRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return table
|
||||
|
||||
|
||||
|
|
@ -222,7 +232,9 @@ class AgentRegistry:
|
|||
self.load_agents_from_config(agent_config if agent_config is not None else self.config_agents)
|
||||
return self.agent_list
|
||||
|
||||
async def migrate_legacy_grant_ids(self, table: ObjectPermissionTableClient) -> GrantMigrationResult:
|
||||
async def migrate_legacy_grant_ids(
|
||||
self, table: "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]"
|
||||
) -> GrantMigrationResult:
|
||||
"""
|
||||
Rewrite object_permission.agents rows holding a legacy full-entry hash to the
|
||||
stable name-derived id.
|
||||
|
|
@ -360,6 +372,8 @@ class AgentRegistry:
|
|||
"""
|
||||
try:
|
||||
deleted_agent: Final = await agents_table(prisma_client).delete(where={"agent_id": agent_id})
|
||||
if deleted_agent is None:
|
||||
raise ValueError(f"Agent not found, passed agent_id={agent_id}")
|
||||
return dict(deleted_agent)
|
||||
except Exception as e:
|
||||
raise Exception(f"Error deleting agent from DB: {e}")
|
||||
|
|
@ -386,12 +400,12 @@ class AgentRegistry:
|
|||
The patched agent
|
||||
"""
|
||||
try:
|
||||
existing_agent = await AgentsRepository(prisma_client).table.find_unique(where={"agent_id": agent_id})
|
||||
if existing_agent is not None:
|
||||
existing_agent = dict(existing_agent)
|
||||
|
||||
if existing_agent is None:
|
||||
existing_row: Final = await AgentsRepository(prisma_client).table.find_unique(
|
||||
where={"agent_id": agent_id} # mutable-ok: prisma filters are plain dicts
|
||||
)
|
||||
if existing_row is None:
|
||||
raise Exception(f"Agent with ID {agent_id} not found")
|
||||
existing_agent: Final = dict(existing_row)
|
||||
|
||||
augment_agent: Final = {**existing_agent, **agent}
|
||||
update_data: Final[dict[str, Any]] = {}
|
||||
|
|
@ -436,6 +450,8 @@ class AgentRegistry:
|
|||
},
|
||||
include={"object_permission": True},
|
||||
)
|
||||
if patched_agent is None:
|
||||
raise ValueError(f"Agent not found, passed agent_id={agent_id}")
|
||||
patched_agent_dict: Final = patched_agent.model_dump()
|
||||
if patched_agent.object_permission is not None:
|
||||
try:
|
||||
|
|
@ -523,6 +539,8 @@ class AgentRegistry:
|
|||
include={"object_permission": True},
|
||||
)
|
||||
|
||||
if updated_agent is None:
|
||||
raise ValueError(f"Agent not found, passed agent_id={agent_id}")
|
||||
updated_agent_dict: Final = updated_agent.model_dump()
|
||||
if updated_agent.object_permission is not None:
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -543,7 +543,7 @@ async def update_plugin(
|
|||
|
||||
manifest: Final[Mapping[str, object]] = _build_plugin_manifest(plugin_name, request)
|
||||
|
||||
plugin: Final[_PluginRecord] = await ClaudeCodePluginRepository(prisma_client).table.update(
|
||||
plugin: Final[_PluginRecord | None] = await ClaudeCodePluginRepository(prisma_client).table.update(
|
||||
where={"name": plugin_name}, # mutable-ok: prisma query arguments must be plain dicts
|
||||
data={ # mutable-ok: prisma query arguments must be plain dicts
|
||||
"version": request.version,
|
||||
|
|
@ -553,6 +553,8 @@ async def update_plugin(
|
|||
"updated_at": datetime.now(timezone.utc),
|
||||
},
|
||||
)
|
||||
if plugin is None:
|
||||
raise _error_response(404, f"Plugin '{plugin_name}' not found")
|
||||
|
||||
verbose_proxy_logger.info("Plugin %s updated successfully", plugin_name)
|
||||
|
||||
|
|
|
|||
|
|
@ -71,7 +71,6 @@ from litellm.proxy.auth.budget_throttle import (
|
|||
)
|
||||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import publish_auth_cache_invalidation
|
||||
from litellm.proxy.common_utils.cache_pydantic_utils import CacheCodec
|
||||
from litellm.proxy.common_utils.http_parsing_utils import (
|
||||
_safe_get_request_headers,
|
||||
_safe_get_request_query_params,
|
||||
|
|
@ -158,7 +157,12 @@ class _PrismaVectorStoreRow(Protocol):
|
|||
|
||||
class _PrismaUserRow(Protocol):
|
||||
user_id: str
|
||||
organization_memberships: Sequence[LiteLLM_OrganizationMembershipTable | None] | None
|
||||
|
||||
@property
|
||||
def organization_memberships(self) -> Sequence[_PrismaModelDumpRow | None] | None: ...
|
||||
|
||||
@organization_memberships.setter
|
||||
def organization_memberships(self, value: Sequence[_PrismaModelDumpRow] | None) -> None: ...
|
||||
|
||||
def __iter__(self) -> Iterator[tuple[str, object]]: ...
|
||||
|
||||
|
|
@ -216,9 +220,14 @@ def _user_table(repo: _PrismaTableHolder[_PrismaUserRow]) -> _PrismaAuthTable[_P
|
|||
return repo.table
|
||||
|
||||
|
||||
class _VectorStorePermissionsRow(Protocol):
|
||||
@property
|
||||
def vector_stores(self) -> Sequence[str] | None: ...
|
||||
|
||||
|
||||
def _object_permission_table(
|
||||
repo: _PrismaTableHolder[LiteLLM_ObjectPermissionTable],
|
||||
) -> _PrismaAuthTable[LiteLLM_ObjectPermissionTable]:
|
||||
repo: _PrismaTableHolder[_VectorStorePermissionsRow],
|
||||
) -> _PrismaAuthTable[_VectorStorePermissionsRow]:
|
||||
return repo.table
|
||||
|
||||
|
||||
|
|
@ -1131,7 +1140,8 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool:
|
|||
|
||||
Parameters:
|
||||
- user_route: str - the route the user is trying to call
|
||||
- allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user.
|
||||
- allowed_routes: List[str|LiteLLMRoutes] - the list of allowed routes for the user. Entries are a route group name
|
||||
(e.g. "openai_routes"), an exact route, or a trailing-wildcard prefix (e.g. "/internal-models/*").
|
||||
"""
|
||||
from starlette.routing import compile_path
|
||||
|
||||
|
|
@ -1141,7 +1151,7 @@ def _allowed_routes_check(user_route: str, allowed_routes: list) -> bool:
|
|||
regex, _, _ = compile_path(template)
|
||||
if regex.match(user_route):
|
||||
return True
|
||||
elif allowed_route == user_route:
|
||||
elif RouteChecks.route_matches_wildcard_pattern(route=user_route, pattern=allowed_route):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
|
@ -2741,20 +2751,9 @@ async def _get_team_object_from_user_api_key_cache(
|
|||
|
||||
async def _get_team_object_from_cache(
|
||||
key: str,
|
||||
proxy_logging_obj: ProxyLogging | None,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
parent_otel_span: Span | None,
|
||||
) -> LiteLLM_TeamTableCachedObj | None:
|
||||
## INTERNAL USAGE CACHE (plain DualCache) — checked before UserApiKeyCache stores ##
|
||||
if proxy_logging_obj is not None and proxy_logging_obj.internal_usage_cache.dual_cache:
|
||||
cached_raw: Final = await proxy_logging_obj.internal_usage_cache.dual_cache.async_get_cache(
|
||||
key=key, parent_otel_span=parent_otel_span
|
||||
)
|
||||
if cached_raw is not None:
|
||||
from_internal: Final = CacheCodec.deserialize(cached_raw, LiteLLM_TeamTableCachedObj)
|
||||
if from_internal is not None:
|
||||
return from_internal
|
||||
|
||||
decoded: Final = await user_api_key_cache.async_get_cache(
|
||||
key=key,
|
||||
parent_otel_span=parent_otel_span,
|
||||
|
|
@ -2790,7 +2789,6 @@ async def get_team_object(
|
|||
if not check_db_only:
|
||||
cached_team_obj: Final = await _get_team_object_from_cache(
|
||||
key=key,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
|
|
@ -2953,7 +2951,6 @@ async def get_team_object_by_alias(
|
|||
|
||||
cached_team_obj: Final = await _get_team_object_from_cache(
|
||||
key=cache_key,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=parent_otel_span,
|
||||
)
|
||||
|
|
@ -5389,7 +5386,7 @@ async def vector_store_access_check(
|
|||
def _can_object_call_vector_stores(
|
||||
object_type: Literal["key", "team", "org"],
|
||||
vector_store_ids_to_run: list[str],
|
||||
object_permissions: LiteLLM_ObjectPermissionTable | None,
|
||||
object_permissions: _VectorStorePermissionsRow | None,
|
||||
):
|
||||
"""
|
||||
Raises ProxyException if the object (key, team, org) cannot access the specific vector store.
|
||||
|
|
|
|||
|
|
@ -11,6 +11,7 @@ import litellm
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import EMPTY_MAPPING
|
||||
from litellm.integrations.otel.runtime import seed_request_identity
|
||||
from litellm.litellm_core_utils.core_helpers import is_expected_client_error
|
||||
from litellm.proxy._types import (
|
||||
LitellmUserRoles,
|
||||
ProxyErrorTypes,
|
||||
|
|
@ -109,7 +110,12 @@ class UserAPIKeyAuthExceptionHandler:
|
|||
request=request,
|
||||
use_x_forwarded_for=general_settings.get("use_x_forwarded_for") is True,
|
||||
)
|
||||
verbose_proxy_logger.exception(
|
||||
log_fn: Final = (
|
||||
verbose_proxy_logger.error
|
||||
if is_expected_client_error(e) and not litellm.log_client_error_tracebacks
|
||||
else verbose_proxy_logger.exception
|
||||
)
|
||||
log_fn(
|
||||
"litellm.proxy.proxy_server.user_api_key_auth(): Exception occured - %s\nRequester IP Address:%s",
|
||||
e,
|
||||
requester_ip,
|
||||
|
|
|
|||
|
|
@ -608,7 +608,7 @@ def route_in_additonal_public_routes(current_route: str):
|
|||
|
||||
# Check wildcard patterns
|
||||
for route_pattern in routes_defined:
|
||||
if RouteChecks._route_matches_wildcard_pattern(route=current_route, pattern=route_pattern):
|
||||
if RouteChecks.route_matches_wildcard_pattern(route=current_route, pattern=route_pattern):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -181,7 +181,7 @@ class RouteChecks:
|
|||
|
||||
# check if wildcard pattern is allowed
|
||||
for allowed_route in valid_token.allowed_routes:
|
||||
if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route):
|
||||
if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route):
|
||||
return True
|
||||
|
||||
if denied_auth_enforced_pass_through_route:
|
||||
|
|
@ -329,7 +329,7 @@ class RouteChecks:
|
|||
route_allowed = True
|
||||
break
|
||||
|
||||
if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route):
|
||||
if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route):
|
||||
route_allowed = True
|
||||
break
|
||||
|
||||
|
|
@ -397,7 +397,7 @@ class RouteChecks:
|
|||
return True
|
||||
# Check for wildcard patterns like "/containers/*"
|
||||
if RouteChecks._is_wildcard_pattern(pattern=openai_route):
|
||||
if RouteChecks._route_matches_wildcard_pattern(route=route, pattern=openai_route):
|
||||
if RouteChecks.route_matches_wildcard_pattern(route=route, pattern=openai_route):
|
||||
return True
|
||||
|
||||
# Check for Google routes with placeholders like "/v1beta/models/{model_name}:generateContent"
|
||||
|
|
@ -517,7 +517,7 @@ class RouteChecks:
|
|||
return pattern.endswith("*")
|
||||
|
||||
@staticmethod
|
||||
def _route_matches_wildcard_pattern(route: str, pattern: str) -> bool:
|
||||
def route_matches_wildcard_pattern(route: str, pattern: str) -> bool:
|
||||
"""
|
||||
Check if route matches the wildcard pattern
|
||||
|
||||
|
|
@ -594,7 +594,7 @@ class RouteChecks:
|
|||
# e.g calling /anthropic/v1/messages is allowed if allowed_routes has /anthropic/*
|
||||
#########################################################
|
||||
if any(
|
||||
RouteChecks._route_matches_wildcard_pattern(route=route, pattern=allowed_route)
|
||||
RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route)
|
||||
for allowed_route in allowed_routes
|
||||
if RouteChecks._is_wildcard_pattern(pattern=allowed_route)
|
||||
):
|
||||
|
|
|
|||
|
|
@ -199,6 +199,15 @@ class _UserModelBudgetLimiter(Protocol):
|
|||
) -> bool: ...
|
||||
|
||||
|
||||
class _TokenTeamModels(Protocol):
|
||||
@property
|
||||
def team_models(self) -> list[str]: ...
|
||||
|
||||
|
||||
def _token_team_models(valid_token: _TokenTeamModels) -> list[str]:
|
||||
return valid_token.team_models
|
||||
|
||||
|
||||
async def _read_user_model_max_budget(
|
||||
user_id: str | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
|
|
@ -1992,7 +2001,7 @@ async def _user_api_key_auth_builder(
|
|||
include={"litellm_budget_table": True},
|
||||
)
|
||||
if _db_member is not None:
|
||||
team_member_info = LiteLLM_TeamMembership(**_db_member.dict())
|
||||
team_member_info = LiteLLM_TeamMembership(**_db_member.model_dump())
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=_cache_key,
|
||||
value=team_member_info,
|
||||
|
|
@ -2149,6 +2158,7 @@ async def _user_api_key_auth_builder(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
except HTTPException:
|
||||
token_team_models: Final = _token_team_models(valid_token)
|
||||
_team_obj = LiteLLM_TeamTableCachedObj(
|
||||
team_id=valid_token.team_id,
|
||||
max_budget=valid_token.team_max_budget,
|
||||
|
|
@ -2157,7 +2167,7 @@ async def _user_api_key_auth_builder(
|
|||
tpm_limit=valid_token.team_tpm_limit,
|
||||
rpm_limit=valid_token.team_rpm_limit,
|
||||
blocked=valid_token.team_blocked,
|
||||
models=valid_token.team_models,
|
||||
models=token_team_models,
|
||||
metadata=valid_token.team_metadata,
|
||||
object_permission_id=valid_token.team_object_permission_id,
|
||||
object_permission=await _resolve_object_permission_for_unresolvable_team(
|
||||
|
|
@ -2301,6 +2311,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached
|
|||
UserAPIKeyAuth. Only called when valid_token.team_id is known to be
|
||||
non-None (the caller gates on it)."""
|
||||
assert valid_token.team_id is not None
|
||||
token_team_models: Final = _token_team_models(valid_token)
|
||||
return LiteLLM_TeamTableCachedObj(
|
||||
team_id=valid_token.team_id,
|
||||
max_budget=valid_token.team_max_budget,
|
||||
|
|
@ -2309,7 +2320,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached
|
|||
tpm_limit=valid_token.team_tpm_limit,
|
||||
rpm_limit=valid_token.team_rpm_limit,
|
||||
blocked=valid_token.team_blocked,
|
||||
models=valid_token.team_models,
|
||||
models=token_team_models,
|
||||
metadata=valid_token.team_metadata,
|
||||
object_permission_id=valid_token.team_object_permission_id,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -30,10 +30,11 @@ from litellm.constants import (
|
|||
RETURN_RAW_MODEL_NAME_METADATA_KEY,
|
||||
ROUTER_MODEL_NAME_RESPONSE_FIELD,
|
||||
STREAM_SSE_DATA_PREFIX,
|
||||
STREAM_SSE_KEEPALIVE_PING_BYTES,
|
||||
UNSAFE_PROXY_RESPONSE_HEADERS,
|
||||
)
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
|
||||
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket, is_expected_client_error
|
||||
from litellm.litellm_core_utils.dd_tracing import NullTracer, tracer
|
||||
from litellm.litellm_core_utils.get_supported_openai_params import (
|
||||
get_supported_openai_params,
|
||||
|
|
@ -202,6 +203,10 @@ _CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInforma
|
|||
}
|
||||
|
||||
|
||||
def _withheld_provider_output(response: object) -> bool:
|
||||
return getattr(response, "has_buffered_provider_output", False) is True
|
||||
|
||||
|
||||
def _should_return_raw_model_name(request_data: dict[str, object]) -> bool:
|
||||
return any(
|
||||
isinstance(metadata, dict) and metadata.get(RETURN_RAW_MODEL_NAME_METADATA_KEY) is True
|
||||
|
|
@ -1138,24 +1143,25 @@ async def open_sse_before_first_byte(
|
|||
)
|
||||
|
||||
|
||||
def _is_azure_model_router_request(model: str) -> bool:
|
||||
def _is_azure_model_router_request(model: str, hidden_params: Mapping[str, object] | None = None) -> bool:
|
||||
"""
|
||||
Check if the requested model is an Azure Model Router.
|
||||
Check if a request went down the Azure Model Router route.
|
||||
|
||||
Azure Model Router models follow the pattern:
|
||||
- azure_ai/model_router/<deployment-name>
|
||||
- azure_ai/model-router
|
||||
- model_router/<deployment-name>
|
||||
- model-router
|
||||
``model`` here is what the *client* sent, a model group alias with no ``model_router/``
|
||||
prefix, so matching on it alone only works when the operator happened to put "model-router"
|
||||
in the alias. Where the response is in hand its stamp answers this outright, so callers
|
||||
should pass ``hidden_params``.
|
||||
|
||||
Args:
|
||||
model: The requested model name
|
||||
hidden_params: ``_hidden_params`` from the response, when the caller has it
|
||||
|
||||
Returns:
|
||||
bool: True if this is an Azure Model Router request
|
||||
"""
|
||||
model_lower: Final = model.lower()
|
||||
return "model-router" in model_lower or "model_router" in model_lower
|
||||
from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
|
||||
|
||||
return AzureFoundryModelInfo.is_model_router_call(model=model, hidden_params=hidden_params)
|
||||
|
||||
|
||||
def _override_openai_response_model(
|
||||
|
|
@ -1223,7 +1229,7 @@ def _override_openai_response_model(
|
|||
return
|
||||
|
||||
# Check if this is an Azure Model Router request - if so, preserve the actual model used
|
||||
if _is_azure_model_router_request(requested_model):
|
||||
if _is_azure_model_router_request(requested_model, hidden_params):
|
||||
verbose_proxy_logger.debug(
|
||||
"%s: Azure Model Router detected - preserving actual model used from response instead of overriding to router model.",
|
||||
log_context,
|
||||
|
|
@ -1379,7 +1385,12 @@ def _log_llm_api_exception(e: Exception) -> None:
|
|||
"litellm.proxy.proxy_server._handle_llm_api_exception(): client disconnected, upstream LLM request cancelled"
|
||||
)
|
||||
return
|
||||
verbose_proxy_logger.exception("litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - %s", e)
|
||||
log_fn: Final = (
|
||||
verbose_proxy_logger.error
|
||||
if is_expected_client_error(e) and not litellm.log_client_error_tracebacks
|
||||
else verbose_proxy_logger.exception
|
||||
)
|
||||
log_fn("litellm.proxy.proxy_server._handle_llm_api_exception(): Exception occured - %s", e)
|
||||
|
||||
|
||||
async def _cancel_llm_call_on_client_disconnect(
|
||||
|
|
@ -3442,8 +3453,9 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# so a GeneratorExit on client disconnect is raised there and any
|
||||
# statement after the yield never runs. The slow-path hook is
|
||||
# awaited above, so a cancellation during it still leaves this
|
||||
# False and refunds.
|
||||
delivered_chunk = True
|
||||
# False and refunds. A keepalive ping carries no provider output,
|
||||
# so it must not suppress that refund.
|
||||
delivered_chunk = delivered_chunk or chunk != STREAM_SSE_KEEPALIVE_PING_BYTES
|
||||
yield serialize_chunk(chunk)
|
||||
stream_completed = True
|
||||
except (asyncio.CancelledError, GeneratorExit):
|
||||
|
|
@ -3457,7 +3469,7 @@ class ProxyBaseLLMRequestProcessing:
|
|||
# only sees GeneratorExit on GC) cannot own the refund.
|
||||
if not stream_completed:
|
||||
client_disconnected = True
|
||||
if not delivered_chunk:
|
||||
if not delivered_chunk and not _withheld_provider_output(response):
|
||||
from litellm.proxy.spend_tracking.budget_reservation import (
|
||||
release_budget_reservation_on_cancel,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from dataclasses import asdict, dataclass
|
|||
from typing import TYPE_CHECKING, Final, Protocol, cast # noqa: TID251 # untyped prisma/redis boundary needs cast
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.repositories.prisma_protocols import RowT_co, TableActions
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.caching.redis_cache import RedisCache
|
||||
|
|
@ -163,13 +164,14 @@ class _PublishOnWriteActions:
|
|||
|
||||
|
||||
def wrap_table_actions_for_config_sync(
|
||||
actions: object,
|
||||
actions: "TableActions[RowT_co]",
|
||||
table_name: str,
|
||||
publish: Callable[[str], Awaitable[None]] = publish_config_change_for_object_type,
|
||||
) -> object:
|
||||
) -> "TableActions[RowT_co]":
|
||||
if table_name not in _CONFIG_SYNCED_TABLE_NAMES:
|
||||
return actions
|
||||
return _PublishOnWriteActions(actions=actions, object_type=table_name, publish=publish)
|
||||
wrapped: Final = _PublishOnWriteActions(actions=actions, object_type=table_name, publish=publish)
|
||||
return cast("TableActions[RowT_co]", wrapped) # cast-ok: dynamic write-through proxy keeps the wrapped row type
|
||||
|
||||
|
||||
class ConfigSyncSubscriber:
|
||||
|
|
|
|||
|
|
@ -4,8 +4,9 @@ Expired UI session key cleanup manager.
|
|||
Deletes expired virtual keys created for LiteLLM dashboard sessions.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Final
|
||||
from typing import Any, Final, Protocol
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
|
|
@ -14,7 +15,7 @@ from litellm.constants import (
|
|||
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
|
||||
UI_SESSION_TOKEN_TEAM_ID,
|
||||
)
|
||||
from litellm.proxy._types import KeyRequest, LiteLLM_VerificationToken, UserAPIKeyAuth
|
||||
from litellm.proxy._types import KeyRequest, UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.hooks.key_management_event_hooks import KeyManagementEventHooks
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import (
|
||||
|
|
@ -26,6 +27,11 @@ from litellm.repositories.verification_token_repository import (
|
|||
)
|
||||
|
||||
|
||||
class _ExpiredSessionKeyRow(Protocol):
|
||||
@property
|
||||
def token(self) -> str | None: ...
|
||||
|
||||
|
||||
class ExpiredUISessionKeyCleanupManager:
|
||||
"""
|
||||
Cleans up expired UI session keys.
|
||||
|
|
@ -138,7 +144,7 @@ class ExpiredUISessionKeyCleanupManager:
|
|||
|
||||
return len(tokens)
|
||||
|
||||
async def _find_expired_ui_session_keys(self) -> list[LiteLLM_VerificationToken]:
|
||||
async def _find_expired_ui_session_keys(self) -> Sequence[_ExpiredSessionKeyRow]:
|
||||
"""
|
||||
Find expired LiteLLM dashboard session keys.
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -4,8 +4,9 @@ Key Rotation Manager - Automated key rotation based on rotation schedules
|
|||
Handles finding keys that need rotation based on their individual schedules.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
|
|
@ -31,6 +32,9 @@ from litellm.repositories.verification_token_repository import (
|
|||
VerificationTokenRepository,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
|
||||
class KeyRotationManager:
|
||||
"""
|
||||
|
|
@ -106,7 +110,7 @@ class KeyRotationManager:
|
|||
cronjob_id=KEY_ROTATION_JOB_NAME,
|
||||
)
|
||||
|
||||
async def _find_keys_needing_rotation(self) -> list[LiteLLM_VerificationToken]:
|
||||
async def _find_keys_needing_rotation(self) -> "Sequence[prisma_models.LiteLLM_VerificationToken]":
|
||||
"""
|
||||
Find keys that are due for rotation based on their key_rotation_at timestamp.
|
||||
|
||||
|
|
@ -156,7 +160,7 @@ class KeyRotationManager:
|
|||
# Check if the rotation time has passed
|
||||
return now >= key.key_rotation_at
|
||||
|
||||
async def _rotate_key(self, key: LiteLLM_VerificationToken):
|
||||
async def _rotate_key(self, key: "prisma_models.LiteLLM_VerificationToken"):
|
||||
"""
|
||||
Rotate a single key using existing regenerate_key_fn and call the rotation hook
|
||||
"""
|
||||
|
|
@ -197,7 +201,7 @@ class KeyRotationManager:
|
|||
if isinstance(response, GenerateKeyResponse):
|
||||
await KeyManagementEventHooks.async_key_rotated_hook(
|
||||
data=regenerate_request,
|
||||
existing_key_row=key,
|
||||
existing_key_row=key, # pyright: ignore[reportArgumentType] # prisma row, hook wants the domain model
|
||||
response=response,
|
||||
user_api_key_dict=system_user,
|
||||
litellm_changed_by=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
|
||||
|
|
|
|||
|
|
@ -37,7 +37,7 @@ from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManage
|
|||
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
from litellm.repositories.prisma_protocols import ReadOnlyTable, SpendLinkedTable
|
||||
from litellm.repositories.prisma_protocols import SpendLinkedTable
|
||||
from litellm.repositories.table_repositories import (
|
||||
EndUserRepository,
|
||||
TagRepository,
|
||||
|
|
@ -675,7 +675,7 @@ class ResetBudgetJob:
|
|||
rely on the default budget (litellm.max_end_user_budget_id) applied
|
||||
in-memory during auth checks.
|
||||
"""
|
||||
table: Final[ReadOnlyTable] = EndUserRepository(self.prisma_client).table
|
||||
table: Final = EndUserRepository(self.prisma_client).table
|
||||
rows: Final = await self._with_db_retry(
|
||||
lambda: table.find_many(
|
||||
where={
|
||||
|
|
@ -685,7 +685,7 @@ class ResetBudgetJob:
|
|||
),
|
||||
reason="reset_budget_read_endusers_without_budget_id_failure",
|
||||
)
|
||||
return [LiteLLM_EndUserTable.model_validate(row.dict()) for row in rows]
|
||||
return [LiteLLM_EndUserTable.model_validate(row.model_dump()) for row in rows]
|
||||
|
||||
async def _write_key_reset_updates(self, updated_keys: list[LiteLLM_VerificationToken]) -> None:
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -6,7 +6,9 @@ from typing import Final
|
|||
|
||||
import anyio
|
||||
|
||||
ANTHROPIC_PING_SSE_CHUNK: Final = 'event: ping\ndata: {"type": "ping"}\n\n'
|
||||
from litellm.constants import STREAM_SSE_KEEPALIVE_PING_CHUNK
|
||||
|
||||
ANTHROPIC_PING_SSE_CHUNK: Final = STREAM_SSE_KEEPALIVE_PING_CHUNK
|
||||
SSE_COMMENT_PING: Final = ": ping\n\n"
|
||||
SSE_COMMENT_PING_BYTES: Final = SSE_COMMENT_PING.encode()
|
||||
# The byte form of proxy_server._SSE_FRAME_DELIMITERS, CR-only included: SSE
|
||||
|
|
@ -89,6 +91,17 @@ def is_sse_content_type(content_type: str | None) -> bool:
|
|||
return content_type is not None and content_type.split(";", 1)[0].strip().lower() == _SSE_MEDIA_TYPE
|
||||
|
||||
|
||||
def split_complete_sse_frames(pending: bytes) -> tuple[bytes, bytes]:
|
||||
"""Split buffered SSE bytes into ``(complete_frames, unterminated_tail)``."""
|
||||
boundary_end: Final = max(
|
||||
(pending.rfind(delimiter) + len(delimiter) for delimiter in _SSE_FRAME_DELIMITERS if delimiter in pending),
|
||||
default=0,
|
||||
)
|
||||
if boundary_end == 0:
|
||||
return b"", pending
|
||||
return pending[:boundary_end], pending[boundary_end:]
|
||||
|
||||
|
||||
def wrap_passthrough_sse_bytes_with_keepalive_pings(
|
||||
stream: AsyncGenerator[bytes, None],
|
||||
ping_interval_seconds: float | str | None,
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
|
|
@ -18,28 +18,11 @@ from litellm.repositories.table_repositories import ManagedObjectRepository
|
|||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
class _ManagedObjectRow(Protocol):
|
||||
model_object_id: str
|
||||
unified_object_id: str | None
|
||||
file_purpose: str | None
|
||||
created_by: str | None
|
||||
|
||||
|
||||
class _ManagedObjectTable(Protocol):
|
||||
async def find_unique(self, *, where: Mapping[str, str]) -> _ManagedObjectRow | None: ...
|
||||
|
||||
async def find_first(self, *, where: Mapping[str, str]) -> _ManagedObjectRow | None: ...
|
||||
|
||||
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_ManagedObjectRow]: ...
|
||||
|
||||
async def create(self, *, data: Mapping[str, str]) -> _ManagedObjectRow: ...
|
||||
|
||||
async def update(self, *, where: Mapping[str, str], data: Mapping[str, str]) -> _ManagedObjectRow | None: ...
|
||||
|
||||
|
||||
CONTAINER_OBJECT_PURPOSE: Final = "container"
|
||||
|
||||
# 60s LRU/TTL cache absorbs every container access check before it reaches
|
||||
|
|
@ -220,7 +203,7 @@ async def record_container_owner(
|
|||
verbose_proxy_logger.warning("Skipping container ownership tracking because prisma_client is None")
|
||||
return response
|
||||
|
||||
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
|
||||
table: Final = ManagedObjectRepository(prisma_client).table
|
||||
existing: Final = await table.find_unique(where={"model_object_id": model_object_id})
|
||||
if existing is not None:
|
||||
if getattr(existing, "file_purpose", None) != CONTAINER_OBJECT_PURPOSE:
|
||||
|
|
@ -272,8 +255,8 @@ async def _get_container_owner(original_container_id: str, custom_llm_provider:
|
|||
if prisma_client is None:
|
||||
return None
|
||||
|
||||
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
|
||||
row: Final[_ManagedObjectRow | None] = await table.find_first(
|
||||
table: Final = ManagedObjectRepository(prisma_client).table
|
||||
row: Final[prisma_models.LiteLLM_ManagedObjectTable | None] = await table.find_first(
|
||||
where={
|
||||
"model_object_id": model_object_id,
|
||||
"file_purpose": CONTAINER_OBJECT_PURPOSE,
|
||||
|
|
@ -309,8 +292,8 @@ async def _get_stored_container_id(original_container_id: str, custom_llm_provid
|
|||
if prisma_client is None:
|
||||
return None
|
||||
|
||||
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
|
||||
row: Final[_ManagedObjectRow | None] = await table.find_first(
|
||||
table: Final = ManagedObjectRepository(prisma_client).table
|
||||
row: Final[prisma_models.LiteLLM_ManagedObjectTable | None] = await table.find_first(
|
||||
where={
|
||||
"model_object_id": model_object_id,
|
||||
"file_purpose": CONTAINER_OBJECT_PURPOSE,
|
||||
|
|
@ -394,8 +377,8 @@ async def _get_allowed_container_ids(
|
|||
if prisma_client is None:
|
||||
return set()
|
||||
|
||||
table: Final[_ManagedObjectTable] = ManagedObjectRepository(prisma_client).table
|
||||
rows: Final[Sequence[_ManagedObjectRow]] = await table.find_many(
|
||||
table: Final = ManagedObjectRepository(prisma_client).table
|
||||
rows: Final[Sequence[prisma_models.LiteLLM_ManagedObjectTable]] = await table.find_many(
|
||||
where={
|
||||
"file_purpose": CONTAINER_OBJECT_PURPOSE,
|
||||
"created_by": {"in": owner_scopes},
|
||||
|
|
|
|||
|
|
@ -2,7 +2,10 @@
|
|||
CRUD endpoints for storing reusable credentials.
|
||||
"""
|
||||
|
||||
from typing import Final
|
||||
from typing import (
|
||||
Final,
|
||||
cast, # noqa: TID251 # jsonify_object in proxy/utils.py is annotated with a bare dict
|
||||
)
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
|
||||
|
||||
|
|
@ -88,7 +91,9 @@ async def create_credential(
|
|||
)
|
||||
encrypted_credential: Final = CredentialHelperUtils.encrypt_credential_values(processed_credential)
|
||||
credentials_dict: Final = encrypted_credential.model_dump()
|
||||
credentials_dict_jsonified: Final = jsonify_object(credentials_dict)
|
||||
credentials_dict_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
|
||||
"dict[str, object]", jsonify_object(credentials_dict)
|
||||
)
|
||||
await CredentialsRepository(prisma_client).create(
|
||||
data={
|
||||
**credentials_dict_jsonified,
|
||||
|
|
@ -310,7 +315,9 @@ async def update_credential(
|
|||
if db_credential is None:
|
||||
raise HTTPException(status_code=404, detail="Credential not found in DB.")
|
||||
merged_credential: Final = update_db_credential(db_credential, credential)
|
||||
credential_object_jsonified: Final = jsonify_object(merged_credential.model_dump())
|
||||
credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str
|
||||
"dict[str, object]", jsonify_object(merged_credential.model_dump())
|
||||
)
|
||||
await credentials_repository.update_by_name(
|
||||
credential_name,
|
||||
data={
|
||||
|
|
|
|||
|
|
@ -6,14 +6,15 @@ Admins use the management endpoints to read and update input_policy / output_pol
|
|||
"""
|
||||
|
||||
import uuid
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Final, Protocol, TypeVar
|
||||
from typing import TYPE_CHECKING, Any, Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import ToolDiscoveryQueueItem
|
||||
from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
|
||||
from litellm.repositories.object_permission_repository import ObjectPermissionRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import ToolRepository
|
||||
from litellm.types.tool_management import (
|
||||
LiteLLM_ToolTableRow,
|
||||
|
|
@ -25,33 +26,16 @@ if TYPE_CHECKING:
|
|||
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
_RowT_co: Final = TypeVar("_RowT_co", covariant=True)
|
||||
|
||||
|
||||
class _TableActions(Protocol[_RowT_co]):
|
||||
async def find_unique(self, where: Mapping[str, object]) -> _RowT_co | None: ...
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
where: Mapping[str, object] | None = None,
|
||||
order: Mapping[str, object] | None = None,
|
||||
include: Mapping[str, object] | None = None,
|
||||
) -> Sequence[_RowT_co]: ...
|
||||
|
||||
async def upsert(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT_co: ...
|
||||
|
||||
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _RowT_co | None: ...
|
||||
|
||||
|
||||
def _tool_table_actions(prisma_client: "PrismaClient") -> "_TableActions[prisma_db_models.LiteLLM_ToolTable]":
|
||||
table: Final[_TableActions[prisma_db_models.LiteLLM_ToolTable]] = ToolRepository(prisma_client).table
|
||||
def _tool_table_actions(prisma_client: "PrismaClient") -> "TableActions[prisma_db_models.LiteLLM_ToolTable]":
|
||||
table: Final[TableActions[prisma_db_models.LiteLLM_ToolTable]] = ToolRepository(prisma_client).table
|
||||
return table
|
||||
|
||||
|
||||
def _object_permission_table_actions(
|
||||
prisma_client: "PrismaClient",
|
||||
) -> "_TableActions[prisma_db_models.LiteLLM_ObjectPermissionTable]":
|
||||
table: Final[_TableActions[prisma_db_models.LiteLLM_ObjectPermissionTable]] = ObjectPermissionRepository(
|
||||
) -> "TableActions[prisma_db_models.LiteLLM_ObjectPermissionTable]":
|
||||
table: Final[TableActions[prisma_db_models.LiteLLM_ObjectPermissionTable]] = ObjectPermissionRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return table
|
||||
|
|
|
|||
|
|
@ -29,6 +29,7 @@ from litellm.proxy.guardrails.guardrail_hooks.custom_code.sandbox import (
|
|||
from litellm.proxy.guardrails.guardrail_registry import GuardrailRegistry
|
||||
from litellm.proxy.guardrails.usage_endpoints import router as guardrails_usage_router
|
||||
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import GuardrailsRepository
|
||||
from litellm.types.guardrails import (
|
||||
PII_ENTITY_CATEGORIES_MAP,
|
||||
|
|
@ -65,29 +66,12 @@ router: Final = APIRouter()
|
|||
GUARDRAIL_REGISTRY: Final = GuardrailRegistry()
|
||||
|
||||
|
||||
class _GuardrailsTableActions(Protocol):
|
||||
async def create(self, data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable": ...
|
||||
|
||||
async def delete(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ...
|
||||
|
||||
async def find_unique(self, where: Mapping[str, object]) -> "LiteLLM_GuardrailsTable | None": ...
|
||||
|
||||
async def find_many(
|
||||
self, where: Mapping[str, object], order: Mapping[str, str]
|
||||
) -> "Sequence[LiteLLM_GuardrailsTable]": ...
|
||||
|
||||
async def update(
|
||||
self, where: Mapping[str, object], data: Mapping[str, object]
|
||||
) -> "LiteLLM_GuardrailsTable | None": ...
|
||||
|
||||
|
||||
def _as_str_object_mapping(mapping: Mapping[str, object]) -> Mapping[str, object]:
|
||||
return mapping
|
||||
|
||||
|
||||
def _guardrails_table(prisma_client: "PrismaClient") -> _GuardrailsTableActions:
|
||||
table: Final[_GuardrailsTableActions] = GuardrailsRepository(prisma_client).table
|
||||
return table
|
||||
def _guardrails_table(prisma_client: "PrismaClient") -> "TableActions[LiteLLM_GuardrailsTable]":
|
||||
return GuardrailsRepository(prisma_client).table
|
||||
|
||||
|
||||
async def _create_guardrail_row(prisma_client: "PrismaClient", data: Mapping[str, object]) -> "LiteLLM_GuardrailsTable":
|
||||
|
|
|
|||
|
|
@ -339,6 +339,7 @@ def _build_responses_followup_items(
|
|||
|
||||
class HeadroomGuardrail(CustomGuardrail):
|
||||
records_own_guardrail_information: ClassVar[bool] = True
|
||||
server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset({HEADROOM_RETRIEVE_TOOL_NAME})
|
||||
|
||||
@classmethod
|
||||
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
|
||||
|
|
|
|||
|
|
@ -3,12 +3,12 @@
|
|||
import asyncio
|
||||
import importlib
|
||||
import os
|
||||
from collections.abc import Callable, Iterator, Mapping, Sequence
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from itertools import chain, count
|
||||
from typing import Final, Literal, Optional, Protocol, cast
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, cast
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from pydantic import ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
|
|
@ -39,6 +39,7 @@ from litellm.proxy.guardrails.guardrail_hooks.tool_permission import (
|
|||
)
|
||||
from litellm.proxy.types_utils.utils import get_instance_fn
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import GuardrailsRepository
|
||||
from litellm.secret_managers.main import get_secret
|
||||
from litellm.types.guardrails import (
|
||||
|
|
@ -61,6 +62,9 @@ from .guardrail_initializers import (
|
|||
initialize_tool_permission,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
|
||||
class _GuardrailRowLike(Protocol):
|
||||
@property
|
||||
|
|
@ -68,15 +72,7 @@ class _GuardrailRowLike(Protocol):
|
|||
def __iter__(self) -> Iterator[tuple[str, object]]: ...
|
||||
|
||||
|
||||
class _GuardrailTableActions(Protocol):
|
||||
async def create(self, *, data: Mapping[str, object]) -> _GuardrailRowLike: ...
|
||||
async def delete(self, *, where: Mapping[str, str]) -> object: ...
|
||||
async def update(self, *, where: Mapping[str, str], data: Mapping[str, object]) -> _GuardrailRowLike: ...
|
||||
async def find_many(self, *, where: Mapping[str, str], order: Mapping[str, str]) -> Sequence[BaseModel]: ...
|
||||
async def find_unique(self, *, where: Mapping[str, str]) -> BaseModel | None: ...
|
||||
|
||||
|
||||
def _guardrail_table(prisma_client: PrismaClient) -> _GuardrailTableActions:
|
||||
def _guardrail_table(prisma_client: PrismaClient) -> "TableActions[prisma_models.LiteLLM_GuardrailsTable]":
|
||||
"""Typed view of the guardrails table actions exposed by the Prisma repository."""
|
||||
return GuardrailsRepository(prisma_client).table
|
||||
|
||||
|
|
@ -347,7 +343,7 @@ class GuardrailRegistry:
|
|||
guardrail_info: Final[str] = safe_dumps(guardrail.get("guardrail_info", {}))
|
||||
|
||||
# Update in DB
|
||||
updated_guardrail: Final[_GuardrailRowLike] = await _guardrail_table(prisma_client).update(
|
||||
updated_guardrail: Final[_GuardrailRowLike | None] = await _guardrail_table(prisma_client).update(
|
||||
where={"guardrail_id": guardrail_id},
|
||||
data={
|
||||
"guardrail_name": guardrail_name,
|
||||
|
|
@ -356,6 +352,8 @@ class GuardrailRegistry:
|
|||
"updated_at": datetime.now(timezone.utc),
|
||||
},
|
||||
)
|
||||
if updated_guardrail is None:
|
||||
raise ValueError(f"Guardrail not found, passed guardrail_id={guardrail_id}")
|
||||
|
||||
# Convert to dict and return
|
||||
return dict(updated_guardrail)
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
|
|||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import (
|
||||
DailyGuardrailMetricsRepository,
|
||||
DailyGuardrailUsageUnitsRepository,
|
||||
|
|
@ -30,13 +31,6 @@ from litellm.repositories.table_repositories import (
|
|||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
from prisma import types as prisma_types
|
||||
from prisma.actions import (
|
||||
LiteLLM_DailyGuardrailMetricsActions,
|
||||
LiteLLM_DailyGuardrailUsageUnitsActions,
|
||||
LiteLLM_DailyPolicyMetricsActions,
|
||||
LiteLLM_GuardrailsTableActions,
|
||||
LiteLLM_PolicyTableActions,
|
||||
)
|
||||
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.types.guardrails import Guardrail
|
||||
|
|
@ -85,8 +79,8 @@ def _resolve_usage_window(start_date: str | None, end_date: str | None) -> tuple
|
|||
|
||||
def _guardrails_table(
|
||||
prisma_client: "PrismaClient",
|
||||
) -> "LiteLLM_GuardrailsTableActions[prisma_models.LiteLLM_GuardrailsTable]":
|
||||
guardrails_table: LiteLLM_GuardrailsTableActions[prisma_models.LiteLLM_GuardrailsTable] = GuardrailsRepository(
|
||||
) -> "TableActions[prisma_models.LiteLLM_GuardrailsTable]":
|
||||
guardrails_table: Final[TableActions[prisma_models.LiteLLM_GuardrailsTable]] = GuardrailsRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return guardrails_table
|
||||
|
|
@ -94,28 +88,26 @@ def _guardrails_table(
|
|||
|
||||
def _policies_table(
|
||||
prisma_client: "PrismaClient",
|
||||
) -> "LiteLLM_PolicyTableActions[prisma_models.LiteLLM_PolicyTable]":
|
||||
policies_table: Final[LiteLLM_PolicyTableActions[prisma_models.LiteLLM_PolicyTable]] = PolicyRepository(
|
||||
prisma_client
|
||||
).table
|
||||
) -> "TableActions[prisma_models.LiteLLM_PolicyTable]":
|
||||
policies_table: Final[TableActions[prisma_models.LiteLLM_PolicyTable]] = PolicyRepository(prisma_client).table
|
||||
return policies_table
|
||||
|
||||
|
||||
def _daily_guardrail_metrics_table(
|
||||
prisma_client: "PrismaClient",
|
||||
) -> "LiteLLM_DailyGuardrailMetricsActions[prisma_models.LiteLLM_DailyGuardrailMetrics]":
|
||||
metrics_table: Final[LiteLLM_DailyGuardrailMetricsActions[prisma_models.LiteLLM_DailyGuardrailMetrics]] = (
|
||||
DailyGuardrailMetricsRepository(prisma_client).table
|
||||
)
|
||||
) -> "TableActions[prisma_models.LiteLLM_DailyGuardrailMetrics]":
|
||||
metrics_table: Final[TableActions[prisma_models.LiteLLM_DailyGuardrailMetrics]] = DailyGuardrailMetricsRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return metrics_table
|
||||
|
||||
|
||||
def _daily_policy_metrics_table(
|
||||
prisma_client: "PrismaClient",
|
||||
) -> "LiteLLM_DailyPolicyMetricsActions[prisma_models.LiteLLM_DailyPolicyMetrics]":
|
||||
metrics_table: Final[LiteLLM_DailyPolicyMetricsActions[prisma_models.LiteLLM_DailyPolicyMetrics]] = (
|
||||
DailyPolicyMetricsRepository(prisma_client).table
|
||||
)
|
||||
) -> "TableActions[prisma_models.LiteLLM_DailyPolicyMetrics]":
|
||||
metrics_table: Final[TableActions[prisma_models.LiteLLM_DailyPolicyMetrics]] = DailyPolicyMetricsRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return metrics_table
|
||||
|
||||
|
||||
|
|
@ -135,8 +127,8 @@ async def _find_daily_policy_metrics(
|
|||
|
||||
def _daily_guardrail_usage_units_table(
|
||||
prisma_client: "PrismaClient",
|
||||
) -> "LiteLLM_DailyGuardrailUsageUnitsActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]":
|
||||
units_table: Final[LiteLLM_DailyGuardrailUsageUnitsActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]] = (
|
||||
) -> "TableActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]":
|
||||
units_table: Final[TableActions[prisma_models.LiteLLM_DailyGuardrailUsageUnits]] = (
|
||||
DailyGuardrailUsageUnitsRepository(prisma_client).table
|
||||
)
|
||||
return units_table
|
||||
|
|
|
|||
|
|
@ -14,6 +14,8 @@ from operator import itemgetter
|
|||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Any, Final, NamedTuple, TypeVar
|
||||
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import DB_RETRY_SAFE_ERROR_TYPES
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
|
@ -47,6 +49,18 @@ class _MetricsKey(NamedTuple):
|
|||
date: str
|
||||
|
||||
|
||||
class _UsageUnitCompoundKey(TypedDict):
|
||||
guardrail_id: ReadOnly[str]
|
||||
date: ReadOnly[str]
|
||||
team_id: ReadOnly[str]
|
||||
api_key: ReadOnly[str]
|
||||
usage_unit: ReadOnly[str]
|
||||
|
||||
|
||||
class _UsageUnitWhereUnique(TypedDict):
|
||||
guardrail_id_date_team_id_api_key_usage_unit: ReadOnly[_UsageUnitCompoundKey]
|
||||
|
||||
|
||||
class PendingRollups:
|
||||
"""Rollup rows whose connection-error retries exhausted, held for the next flush."""
|
||||
|
||||
|
|
@ -229,7 +243,7 @@ async def _upsert_usage_unit_row(prisma_client: PrismaClient, key: _UsageUnitKey
|
|||
"usage_unit": key.usage_unit,
|
||||
"units": units,
|
||||
}
|
||||
where: Final[prisma_types.LiteLLM_DailyGuardrailUsageUnitsWhereUniqueInput] = {
|
||||
where: Final[_UsageUnitWhereUnique] = {
|
||||
"guardrail_id_date_team_id_api_key_usage_unit": {
|
||||
"guardrail_id": key.guardrail_id,
|
||||
"date": key.date,
|
||||
|
|
|
|||
|
|
@ -28,6 +28,21 @@ if TYPE_CHECKING:
|
|||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
||||
_RESPONSES_API_PROVIDER_PREFIX: Final = "/openai"
|
||||
_RESPONSES_API_CREATE_ROUTES: Final = frozenset({"/v1/responses", "/responses"})
|
||||
|
||||
|
||||
def _is_responses_api_create_route(request_route: str | None) -> bool:
|
||||
if request_route is None:
|
||||
return False
|
||||
canonical: Final = (
|
||||
request_route[len(_RESPONSES_API_PROVIDER_PREFIX) :]
|
||||
if request_route.startswith(_RESPONSES_API_PROVIDER_PREFIX + "/")
|
||||
else request_route
|
||||
)
|
||||
return canonical in _RESPONSES_API_CREATE_ROUTES
|
||||
|
||||
|
||||
class ResponsesIDSecurity(CustomLogger):
|
||||
def __init__(self):
|
||||
pass
|
||||
|
|
@ -267,8 +282,7 @@ class ResponsesIDSecurity(CustomLogger):
|
|||
async for chunk in response:
|
||||
if (
|
||||
isinstance(chunk, BaseLiteLLMOpenAIResponseObject)
|
||||
and user_api_key_dict.request_route
|
||||
== "/v1/responses" # only encrypt the response id for the responses api
|
||||
and _is_responses_api_create_route(user_api_key_dict.request_route)
|
||||
and not general_settings.get("disable_responses_id_security", False)
|
||||
):
|
||||
chunk = self._encrypt_response_id(chunk, user_api_key_dict, request_encryption_cache)
|
||||
|
|
|
|||
|
|
@ -210,7 +210,6 @@ async def _patch_team_caches_add_access_group(
|
|||
for team_id in team_ids:
|
||||
cached_team = await _get_team_object_from_cache(
|
||||
key=f"team_id:{team_id}",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
|
|
@ -240,7 +239,6 @@ async def _patch_team_caches_remove_access_group(
|
|||
for team_id in team_ids:
|
||||
cached_team = await _get_team_object_from_cache(
|
||||
key=f"team_id:{team_id}",
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
parent_otel_span=None,
|
||||
)
|
||||
|
|
@ -390,7 +388,7 @@ async def list_access_groups(
|
|||
_require_admin_view(user_api_key_dict)
|
||||
prisma_client: Final = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value)
|
||||
|
||||
table: Final[_AccessGroupTable] = AccessGroupRepository(prisma_client).table
|
||||
table: Final = AccessGroupRepository(prisma_client).table
|
||||
records: Final = await table.find_many(order={"created_at": "desc"})
|
||||
return [_record_to_response(r) for r in records]
|
||||
|
||||
|
|
@ -406,7 +404,7 @@ async def get_access_group(
|
|||
_require_admin_view(user_api_key_dict)
|
||||
prisma_client: Final = get_prisma_client_or_throw(CommonProxyErrors.db_not_connected_error.value)
|
||||
|
||||
table: Final[_AccessGroupTable] = AccessGroupRepository(prisma_client).table
|
||||
table: Final = AccessGroupRepository(prisma_client).table
|
||||
record: Final = await table.find_unique(where={"access_group_id": access_group_id})
|
||||
if record is None:
|
||||
raise HTTPException(
|
||||
|
|
|
|||
|
|
@ -93,7 +93,7 @@ async def new_budget(
|
|||
budget_obj.budget_reset_at = get_budget_reset_time(budget_duration=budget_obj.budget_duration)
|
||||
|
||||
budget_obj_json: Final = budget_obj.model_dump(exclude_none=True)
|
||||
budget_obj_jsonified: Final = jsonify_object(budget_obj_json) # json dump any dictionaries
|
||||
budget_obj_jsonified: Final[dict[str, object]] = jsonify_object(budget_obj_json) # mutable-ok: prisma create input
|
||||
try:
|
||||
response: Final = await BudgetRepository(prisma_client).table.create(
|
||||
data={
|
||||
|
|
|
|||
|
|
@ -43,7 +43,8 @@ router: Final = APIRouter()
|
|||
|
||||
|
||||
class _CacheConfigRow(Protocol):
|
||||
cache_settings: str | Mapping[str, object] | None
|
||||
@property
|
||||
def cache_settings(self) -> str | Mapping[str, object] | None: ...
|
||||
|
||||
|
||||
class _CacheConfigTable(Protocol):
|
||||
|
|
|
|||
|
|
@ -441,7 +441,7 @@ async def get_api_key_metadata(
|
|||
This ensures that key_alias and team_id are preserved in historical activity logs
|
||||
even after a key is deleted or regenerated.
|
||||
"""
|
||||
key_records: list[PrismaVerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many(
|
||||
key_records: Sequence[PrismaVerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many(
|
||||
where={"token": {"in": list(api_keys)}}
|
||||
)
|
||||
result: Final[dict[str, _KeyMetadataDict]] = {
|
||||
|
|
@ -452,9 +452,9 @@ async def get_api_key_metadata(
|
|||
missing_keys: Final = api_keys - set(result.keys())
|
||||
if missing_keys:
|
||||
try:
|
||||
deleted_key_records: Final[list[PrismaDeletedVerificationToken]] = await DeletedVerificationTokenRepository(
|
||||
prisma_client
|
||||
).table.find_many(
|
||||
deleted_key_records: Final[
|
||||
Sequence[PrismaDeletedVerificationToken]
|
||||
] = await DeletedVerificationTokenRepository(prisma_client).table.find_many(
|
||||
where={"token": {"in": list(missing_keys)}},
|
||||
order={"deleted_at": "desc"},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -46,7 +46,8 @@ router: Final = APIRouter()
|
|||
|
||||
|
||||
class _ConfigOverrideRow(Protocol):
|
||||
config_value: str | Mapping[str, object] | None
|
||||
@property
|
||||
def config_value(self) -> str | Mapping[str, object] | None: ...
|
||||
|
||||
|
||||
class _ConfigOverridesTableClient(Protocol):
|
||||
|
|
|
|||
|
|
@ -15,9 +15,9 @@ These are members of a Team on LiteLLM
|
|||
import asyncio
|
||||
import json
|
||||
import traceback
|
||||
from collections.abc import Mapping, Sequence
|
||||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Final, Literal, Protocol, cast
|
||||
from typing import Any, Final, Literal, cast
|
||||
|
||||
import fastapi
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
||||
|
|
@ -58,6 +58,7 @@ from litellm.proxy.management_helpers.object_permission_utils import (
|
|||
from litellm.proxy.management_helpers.utils import management_endpoint_wrapper
|
||||
from litellm.proxy.utils import handle_exception_on_proxy, hash_password
|
||||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import (
|
||||
InvitationLinkRepository,
|
||||
OrganizationMembershipRepository,
|
||||
|
|
@ -86,15 +87,6 @@ from litellm.types.proxy.management_endpoints.scim_v2 import (
|
|||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
from prisma import types as prisma_types
|
||||
from prisma.actions import (
|
||||
LiteLLM_InvitationLinkActions,
|
||||
LiteLLM_OrganizationMembershipActions,
|
||||
LiteLLM_OrganizationTableActions,
|
||||
LiteLLM_TeamMembershipActions,
|
||||
LiteLLM_TeamTableActions,
|
||||
LiteLLM_UserTableActions,
|
||||
LiteLLM_VerificationTokenActions,
|
||||
)
|
||||
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.proxy_server import PrismaClient
|
||||
|
|
@ -105,31 +97,31 @@ router: Final = APIRouter()
|
|||
|
||||
def _user_table(
|
||||
prisma_client: "PrismaClient | None",
|
||||
) -> "LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]":
|
||||
user_table: Final[LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]] = UserRepository(prisma_client).table
|
||||
) -> "TableActions[prisma_models.LiteLLM_UserTable]":
|
||||
user_table: Final[TableActions[prisma_models.LiteLLM_UserTable]] = UserRepository(prisma_client).table
|
||||
return user_table
|
||||
|
||||
|
||||
def _team_table(
|
||||
prisma_client: "PrismaClient | None",
|
||||
) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]":
|
||||
team_table: Final[LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
|
||||
) -> "TableActions[prisma_models.LiteLLM_TeamTable]":
|
||||
team_table: Final[TableActions[prisma_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table
|
||||
return team_table
|
||||
|
||||
|
||||
def _verification_token_table(
|
||||
prisma_client: "PrismaClient | None",
|
||||
) -> "LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]":
|
||||
token_table: Final[LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]] = (
|
||||
VerificationTokenRepository(prisma_client).table
|
||||
)
|
||||
) -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
|
||||
token_table: Final[TableActions[prisma_models.LiteLLM_VerificationToken]] = VerificationTokenRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return token_table
|
||||
|
||||
|
||||
def _organization_membership_table(
|
||||
prisma_client: "PrismaClient | None",
|
||||
) -> "LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]":
|
||||
membership_table: Final[LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]] = (
|
||||
) -> "TableActions[prisma_models.LiteLLM_OrganizationMembership]":
|
||||
membership_table: Final[TableActions[prisma_models.LiteLLM_OrganizationMembership]] = (
|
||||
OrganizationMembershipRepository(prisma_client).table
|
||||
)
|
||||
return membership_table
|
||||
|
|
@ -137,8 +129,8 @@ def _organization_membership_table(
|
|||
|
||||
def _invitation_link_table(
|
||||
prisma_client: "PrismaClient | None",
|
||||
) -> "LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink]":
|
||||
invitation_table: LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink] = InvitationLinkRepository(
|
||||
) -> "TableActions[prisma_models.LiteLLM_InvitationLink]":
|
||||
invitation_table: Final[TableActions[prisma_models.LiteLLM_InvitationLink]] = InvitationLinkRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return invitation_table
|
||||
|
|
@ -146,19 +138,19 @@ def _invitation_link_table(
|
|||
|
||||
def _organization_table(
|
||||
prisma_client: "PrismaClient | None",
|
||||
) -> "LiteLLM_OrganizationTableActions[prisma_models.LiteLLM_OrganizationTable]":
|
||||
organization_table: Final[LiteLLM_OrganizationTableActions[prisma_models.LiteLLM_OrganizationTable]] = (
|
||||
OrganizationRepository(prisma_client).table
|
||||
)
|
||||
) -> "TableActions[prisma_models.LiteLLM_OrganizationTable]":
|
||||
organization_table: Final[TableActions[prisma_models.LiteLLM_OrganizationTable]] = OrganizationRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return organization_table
|
||||
|
||||
|
||||
def _team_membership_table(
|
||||
prisma_client: "PrismaClient | None",
|
||||
) -> "LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]":
|
||||
team_membership_table: Final[LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]] = (
|
||||
TeamMembershipRepository(prisma_client).table
|
||||
)
|
||||
) -> "TableActions[prisma_models.LiteLLM_TeamMembership]":
|
||||
team_membership_table: Final[TableActions[prisma_models.LiteLLM_TeamMembership]] = TeamMembershipRepository(
|
||||
prisma_client
|
||||
).table
|
||||
return team_membership_table
|
||||
|
||||
|
||||
|
|
@ -294,7 +286,7 @@ async def _add_user_to_organizations(
|
|||
organization_member_add,
|
||||
)
|
||||
|
||||
tasks: Final = []
|
||||
tasks: Final[list[Awaitable[object]]] = []
|
||||
for organization_id in organizations:
|
||||
tasks.append(
|
||||
organization_member_add(
|
||||
|
|
@ -406,7 +398,7 @@ async def add_new_user_to_default_team(
|
|||
teams: list[str] | list[NewUserRequestTeam],
|
||||
prisma_client: "PrismaClient",
|
||||
):
|
||||
tasks: Final = []
|
||||
tasks: Final[list[Awaitable[object]]] = []
|
||||
for team in teams:
|
||||
user_role: Literal["user", "admin"] = "user"
|
||||
max_budget_in_team: float | None = None
|
||||
|
|
@ -1479,7 +1471,8 @@ async def _update_single_user_helper(
|
|||
# Create new user if not found
|
||||
non_default_values["user_id"] = str(uuid.uuid4())
|
||||
non_default_values["user_email"] = user_request.user_email
|
||||
response = await prisma_client.insert_data(data=non_default_values, table_name="user")
|
||||
inserted_user_row: Final = await prisma_client.insert_data(data=non_default_values, table_name="user")
|
||||
response = inserted_user_row # pyright: ignore[reportAssignmentType] # insert_data returns a prisma row
|
||||
|
||||
if response is not None:
|
||||
await _schedule_user_update_audit_log(
|
||||
|
|
@ -1795,7 +1788,9 @@ async def bulk_user_update(
|
|||
|
||||
# Apply update transformations (reuse existing logic)
|
||||
data_json: Final[dict] = data.user_updates.model_dump(exclude_unset=True)
|
||||
non_default_values: Final = _update_internal_user_params(data_json=data_json, data=data.user_updates)
|
||||
non_default_values: Final[dict[str, object]] = _update_internal_user_params(
|
||||
data_json=data_json, data=data.user_updates
|
||||
)
|
||||
|
||||
# Remove user identification fields since we're updating by user_id
|
||||
non_default_values.pop("user_id", None)
|
||||
|
|
@ -2149,7 +2144,7 @@ async def get_users(
|
|||
_validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None
|
||||
)
|
||||
|
||||
users: Sequence[prisma_models.LiteLLM_UserTable] | None = await UserRepository(prisma_client).table.find_many(
|
||||
users: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await UserRepository(prisma_client).table.find_many(
|
||||
where=where_conditions,
|
||||
skip=skip,
|
||||
take=page_size,
|
||||
|
|
@ -2160,10 +2155,7 @@ async def get_users(
|
|||
total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where_conditions)
|
||||
|
||||
# Get key count for each user
|
||||
if users is not None:
|
||||
user_key_counts = await get_user_key_counts(prisma_client, [user.user_id for user in users])
|
||||
else:
|
||||
user_key_counts = {}
|
||||
user_key_counts: Final = await get_user_key_counts(prisma_client, [user.user_id for user in users])
|
||||
|
||||
verbose_proxy_logger.debug("Total count of users: %s", total_count)
|
||||
|
||||
|
|
@ -2172,17 +2164,14 @@ async def get_users(
|
|||
|
||||
# Prepare response
|
||||
user_list: list[LiteLLM_UserTableWithKeyCount] = []
|
||||
if users is not None:
|
||||
for user in users:
|
||||
user_dump = user.model_dump()
|
||||
user_dump["metadata"] = _redact_scim_enterprise_metadata(user_dump.get("metadata"))
|
||||
user_list.append(
|
||||
LiteLLM_UserTableWithKeyCount.model_validate(
|
||||
{**user_dump, "key_count": user_key_counts.get(user.user_id, 0)}
|
||||
)
|
||||
for user in users:
|
||||
user_dump = user.model_dump()
|
||||
user_dump["metadata"] = _redact_scim_enterprise_metadata(user_dump.get("metadata"))
|
||||
user_list.append(
|
||||
LiteLLM_UserTableWithKeyCount.model_validate(
|
||||
{**user_dump, "key_count": user_key_counts.get(user.user_id, 0)}
|
||||
)
|
||||
else:
|
||||
user_list = []
|
||||
)
|
||||
|
||||
return {
|
||||
"users": user_list,
|
||||
|
|
@ -2193,13 +2182,6 @@ async def get_users(
|
|||
}
|
||||
|
||||
|
||||
class _DeleteTeamRow(Protocol):
|
||||
team_id: str
|
||||
members_with_roles: object
|
||||
|
||||
def model_dump(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
@router.post(
|
||||
"/user/delete",
|
||||
tags=["Internal User management"],
|
||||
|
|
@ -2258,9 +2240,9 @@ async def delete_user(
|
|||
# loop an org-admin of org-A could delete users in org-B by supplying
|
||||
# {"user_ids": [victim_in_org_B], "organization_id": "org-A"}.
|
||||
caller_is_proxy_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value
|
||||
caller_admin_org_ids: set = set()
|
||||
caller_admin_org_ids: set[str] = set()
|
||||
if not caller_is_proxy_admin:
|
||||
caller_memberships: Final = (
|
||||
caller_memberships: Final[Sequence[prisma_models.LiteLLM_OrganizationMembership]] = (
|
||||
await _organization_membership_table(prisma_client).find_many(
|
||||
where={
|
||||
"user_id": user_api_key_dict.user_id,
|
||||
|
|
@ -2279,7 +2261,7 @@ async def delete_user(
|
|||
|
||||
# Batch-fetch target memberships once before the per-user loop. Avoids
|
||||
# an N+1 DB call when delete_user is called with a large user_ids list.
|
||||
target_org_ids_by_user: Final[dict[str, set]] = {}
|
||||
target_org_ids_by_user: Final[dict[str, set[str]]] = {}
|
||||
if not caller_is_proxy_admin:
|
||||
all_target_memberships: Final = await _organization_membership_table(prisma_client).find_many(
|
||||
where={"user_id": {"in": data.user_ids}}
|
||||
|
|
@ -2319,7 +2301,7 @@ async def delete_user(
|
|||
# we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes
|
||||
if is_audit_logging_enabled():
|
||||
# make an audit log for each team deleted
|
||||
_user_row = user_row.json(exclude_none=True)
|
||||
_user_row = user_row.model_dump_json(exclude_none=True)
|
||||
|
||||
asyncio.create_task(
|
||||
create_audit_log_for_update(
|
||||
|
|
@ -2342,10 +2324,10 @@ async def delete_user(
|
|||
)
|
||||
|
||||
## CLEANUP MEMBERS_WITH_ROLES
|
||||
fetch_all_teams: Sequence[_DeleteTeamRow] = await TeamRepository(prisma_client).table.find_many(
|
||||
where={"team_id": {"in": user_row.teams}}
|
||||
)
|
||||
teams_to_update = []
|
||||
fetch_all_teams: Sequence[prisma_models.LiteLLM_TeamTable] = await TeamRepository(
|
||||
prisma_client
|
||||
).table.find_many(where={"team_id": {"in": user_row.teams}})
|
||||
teams_to_update: list[tuple[str, str]] = []
|
||||
for team in fetch_all_teams:
|
||||
removed_team_members, new_team_members = _cleanup_members_with_roles(
|
||||
existing_team_row=LiteLLM_TeamTable.model_validate(team.model_dump()),
|
||||
|
|
@ -2357,15 +2339,14 @@ async def delete_user(
|
|||
)
|
||||
if removed_team_members:
|
||||
_db_new_team_members: list[dict] = [m.model_dump() for m in new_team_members]
|
||||
team.members_with_roles = json.dumps(_db_new_team_members)
|
||||
teams_to_update.append(team)
|
||||
teams_to_update.append((team.team_id, json.dumps(_db_new_team_members)))
|
||||
|
||||
## update teams
|
||||
|
||||
for team in teams_to_update:
|
||||
for team_id, members_with_roles in teams_to_update:
|
||||
await TeamRepository(prisma_client).table.update(
|
||||
where={"team_id": team.team_id},
|
||||
data={"members_with_roles": team.members_with_roles},
|
||||
where={"team_id": team_id},
|
||||
data={"members_with_roles": members_with_roles},
|
||||
)
|
||||
# End of Audit logging
|
||||
|
||||
|
|
|
|||
|
|
@ -122,6 +122,9 @@ async def update_jwt_key_mapping(
|
|||
where={"id": data.id}, data=update_data
|
||||
)
|
||||
|
||||
if updated_mapping is None:
|
||||
raise HTTPException(status_code=404, detail="Mapping not found")
|
||||
|
||||
# Invalidate new cache key if claim fields changed
|
||||
cache_key = f"jwt_key_mapping:{updated_mapping.jwt_claim_name}:{updated_mapping.jwt_claim_value}"
|
||||
await user_api_key_cache.async_delete_cache(cache_key)
|
||||
|
|
|
|||
|
|
@ -123,6 +123,7 @@ from litellm.repositories.budget_repository import BudgetRepository
|
|||
from litellm.repositories.config_repository import ConfigParam, ConfigRepository
|
||||
from litellm.repositories.credentials_repository import CredentialsRepository
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import (
|
||||
DeletedVerificationTokenRepository,
|
||||
DeprecatedVerificationTokenRepository,
|
||||
|
|
@ -151,65 +152,22 @@ from litellm.types.utils import (
|
|||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
from prisma import models as prisma_models
|
||||
|
||||
_PrismaRowT = TypeVar("_PrismaRowT")
|
||||
_RepositoryModelT = TypeVar("_RepositoryModelT", bound=BaseModel)
|
||||
|
||||
|
||||
class _PrismaTableActions(Protocol[_PrismaRowT]):
|
||||
"""Typed view of the Prisma table actions a repository exposes through its untyped ``table``."""
|
||||
|
||||
async def find_unique(
|
||||
self,
|
||||
*,
|
||||
where: Mapping[str, object],
|
||||
include: Mapping[str, object] | None = None,
|
||||
) -> _PrismaRowT | None: ...
|
||||
|
||||
async def find_first(
|
||||
self,
|
||||
*,
|
||||
where: Mapping[str, object],
|
||||
include: Mapping[str, object] | None = None,
|
||||
) -> _PrismaRowT | None: ...
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
*,
|
||||
where: Mapping[str, object] | None = None,
|
||||
include: Mapping[str, object] | None = None,
|
||||
order: Mapping[str, object] | None = None,
|
||||
skip: int | None = None,
|
||||
take: int | None = None,
|
||||
) -> list[_PrismaRowT]: ...
|
||||
|
||||
async def count(self, *, where: Mapping[str, object] | None = None) -> int: ...
|
||||
|
||||
async def create(self, *, data: Mapping[str, object]) -> _PrismaRowT: ...
|
||||
|
||||
async def create_many(self, *, data: Sequence[Mapping[str, object]]) -> int: ...
|
||||
|
||||
async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ...
|
||||
|
||||
async def update(
|
||||
self,
|
||||
*,
|
||||
where: Mapping[str, object],
|
||||
data: Mapping[str, object],
|
||||
) -> _PrismaRowT | None: ...
|
||||
|
||||
async def upsert(
|
||||
self,
|
||||
*,
|
||||
where: Mapping[str, object],
|
||||
data: Mapping[str, object],
|
||||
) -> _PrismaRowT: ...
|
||||
|
||||
|
||||
class _UserRowLike(Protocol):
|
||||
user_id: str | None
|
||||
user_email: str | None
|
||||
user_alias: str | None
|
||||
"""Read-only view of the user columns ``/key/list`` expands keys with."""
|
||||
|
||||
@property
|
||||
def user_id(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def user_email(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def user_alias(self) -> str | None: ...
|
||||
|
||||
def model_dump(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
|
@ -217,46 +175,56 @@ class _UserRowLike(Protocol):
|
|||
|
||||
|
||||
class _TxTables(Protocol):
|
||||
litellm_proxymodeltable: _PrismaTableActions[object]
|
||||
litellm_proxymodeltable: TableActions[object]
|
||||
|
||||
|
||||
class _TableSource(Protocol[_PrismaRowT]):
|
||||
"""Repository view that exposes its untyped Prisma ``table`` with a concrete row type."""
|
||||
class _ConfigTableActions(Protocol):
|
||||
"""Config table surface this module needs; the shared repository seam exposes no ``update``."""
|
||||
|
||||
@property
|
||||
def table(self) -> _PrismaTableActions[_PrismaRowT]: ...
|
||||
async def find_many(self) -> Sequence[ConfigParam]: ...
|
||||
|
||||
|
||||
def _table_of(source: _TableSource[_PrismaRowT]) -> _PrismaTableActions[_PrismaRowT]:
|
||||
return source.table
|
||||
async def update(
|
||||
self,
|
||||
*,
|
||||
where: Mapping[str, object],
|
||||
data: Mapping[str, object],
|
||||
) -> ConfigParam | None: ...
|
||||
|
||||
|
||||
def _prisma_table(
|
||||
repository: BaseRepository[_RepositoryModelT],
|
||||
) -> _PrismaTableActions[_RepositoryModelT]:
|
||||
return _table_of(repository)
|
||||
) -> TableActions[_RepositoryModelT]:
|
||||
return cast( # cast-ok: callers read only the field names the prisma row and repository model share
|
||||
"TableActions[_RepositoryModelT]", repository.table
|
||||
)
|
||||
|
||||
|
||||
def _deleted_verification_token_table(
|
||||
prisma_client: PrismaClient,
|
||||
) -> _PrismaTableActions[LiteLLM_DeletedVerificationToken]:
|
||||
return _table_of(DeletedVerificationTokenRepository(prisma_client))
|
||||
) -> "TableActions[prisma_models.LiteLLM_DeletedVerificationToken]":
|
||||
return DeletedVerificationTokenRepository(prisma_client).table
|
||||
|
||||
|
||||
def _deprecated_verification_token_table(prisma_client: PrismaClient) -> _PrismaTableActions[object]:
|
||||
return _table_of(DeprecatedVerificationTokenRepository(prisma_client))
|
||||
def _deprecated_verification_token_table(
|
||||
prisma_client: PrismaClient,
|
||||
) -> "TableActions[prisma_models.LiteLLM_DeprecatedVerificationToken]":
|
||||
return DeprecatedVerificationTokenRepository(prisma_client).table
|
||||
|
||||
|
||||
def _user_table(prisma_client: PrismaClient) -> _PrismaTableActions[_UserRowLike]:
|
||||
return _table_of(UserRepository(prisma_client))
|
||||
def _user_table(prisma_client: PrismaClient) -> TableActions[_UserRowLike]:
|
||||
return UserRepository(prisma_client).table
|
||||
|
||||
|
||||
def _credentials_table(prisma_client: PrismaClient) -> _PrismaTableActions[CredentialItem]:
|
||||
return _table_of(CredentialsRepository(prisma_client))
|
||||
def _credentials_table(prisma_client: PrismaClient) -> TableActions[CredentialItem]:
|
||||
return cast( # cast-ok: the rotation loop reads and rewrites these rows through CredentialItem names only
|
||||
"TableActions[CredentialItem]", CredentialsRepository(prisma_client).table
|
||||
)
|
||||
|
||||
|
||||
def _config_table(prisma_client: PrismaClient) -> _PrismaTableActions[ConfigParam]:
|
||||
return _table_of(ConfigRepository(prisma_client))
|
||||
def _config_table(prisma_client: PrismaClient) -> _ConfigTableActions:
|
||||
return cast( # cast-ok: ConfigRepository.table hides the write actions this module needs on that same object
|
||||
"_ConfigTableActions", ConfigRepository(prisma_client).table
|
||||
)
|
||||
|
||||
|
||||
async def _check_custom_key_allowed(custom_key_value: str | None) -> None:
|
||||
|
|
@ -1046,7 +1014,7 @@ async def _common_key_generation_helper(
|
|||
)
|
||||
new_budget: Final = prisma_client.jsonify_object(budget_row.json(exclude_none=True))
|
||||
|
||||
_budget: Final[LiteLLM_BudgetTable] = await BudgetRepository(prisma_client).table.create(
|
||||
_budget: Final[prisma_models.LiteLLM_BudgetTable] = await BudgetRepository(prisma_client).table.create(
|
||||
data={
|
||||
**new_budget,
|
||||
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
|
||||
|
|
@ -1252,7 +1220,7 @@ async def _common_key_generation_helper(
|
|||
|
||||
|
||||
def _check_key_model_specific_limits(
|
||||
keys: list[LiteLLM_VerificationToken],
|
||||
keys: Sequence[LiteLLM_VerificationToken],
|
||||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||||
entity_rpm_limit: int | None,
|
||||
entity_tpm_limit: int | None,
|
||||
|
|
@ -1323,7 +1291,7 @@ def _check_key_model_specific_limits(
|
|||
|
||||
|
||||
def _check_key_rpm_tpm_limits(
|
||||
keys: list[LiteLLM_VerificationToken],
|
||||
keys: Sequence[LiteLLM_VerificationToken],
|
||||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||||
entity_rpm_limit: int | None,
|
||||
entity_tpm_limit: int | None,
|
||||
|
|
@ -1361,7 +1329,7 @@ def _check_key_rpm_tpm_limits(
|
|||
|
||||
|
||||
def check_team_key_model_specific_limits(
|
||||
keys: list[LiteLLM_VerificationToken],
|
||||
keys: Sequence[LiteLLM_VerificationToken],
|
||||
team_table: LiteLLM_TeamTableCachedObj,
|
||||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||||
) -> None:
|
||||
|
|
@ -1386,7 +1354,7 @@ def check_team_key_model_specific_limits(
|
|||
|
||||
|
||||
def check_team_key_rpm_tpm_limits(
|
||||
keys: list[LiteLLM_VerificationToken],
|
||||
keys: Sequence[LiteLLM_VerificationToken],
|
||||
team_table: LiteLLM_TeamTableCachedObj,
|
||||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||||
) -> None:
|
||||
|
|
@ -1494,7 +1462,7 @@ async def _check_project_key_limits(
|
|||
|
||||
|
||||
def check_org_key_model_specific_limits(
|
||||
keys: list[LiteLLM_VerificationToken],
|
||||
keys: Sequence[LiteLLM_VerificationToken],
|
||||
org_table: LiteLLM_OrganizationTable,
|
||||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||||
) -> None:
|
||||
|
|
@ -1527,7 +1495,7 @@ def check_org_key_model_specific_limits(
|
|||
|
||||
|
||||
def check_org_key_rpm_tpm_limits(
|
||||
keys: list[LiteLLM_VerificationToken],
|
||||
keys: Sequence[LiteLLM_VerificationToken],
|
||||
org_table: LiteLLM_OrganizationTable,
|
||||
data: GenerateKeyRequest | UpdateKeyRequest,
|
||||
) -> None:
|
||||
|
|
@ -2242,9 +2210,9 @@ async def _get_and_validate_existing_key(
|
|||
code=status.HTTP_400_BAD_REQUEST,
|
||||
)
|
||||
|
||||
rows: list[LiteLLM_VerificationToken] = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
|
||||
where={"key_alias": key_alias}, take=2
|
||||
)
|
||||
rows: Sequence[LiteLLM_VerificationToken] = await _prisma_table(
|
||||
VerificationTokenRepository(prisma_client)
|
||||
).find_many(where={"key_alias": key_alias}, take=2)
|
||||
|
||||
if len(rows) == 0:
|
||||
raise ProxyException(
|
||||
|
|
@ -2407,7 +2375,10 @@ async def _process_single_key_update(
|
|||
)
|
||||
|
||||
_data: Final = {**non_default_values, "token": update_key_request.key}
|
||||
response: Final = await prisma_client.update_data(token=update_key_request.key, data=_data)
|
||||
response: Final[Mapping[str, object] | None] = cast( # cast-ok: every update_data branch returns a str-keyed dict
|
||||
"Mapping[str, object] | None",
|
||||
await prisma_client.update_data(token=update_key_request.key, data=_data),
|
||||
)
|
||||
|
||||
# Delete cache
|
||||
await _delete_cache_key_object(
|
||||
|
|
@ -3225,7 +3196,7 @@ async def bulk_update_team_keys(
|
|||
# `blocked` is Boolean? with no default; `/key/generate` writes NULL. Prisma's `NOT`
|
||||
# excludes NULLs, so explicitly OR `false` with `null` to include them.
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
existing_keys = await VerificationTokenRepository(prisma_client).table.find_many(
|
||||
existing_keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
|
||||
where={
|
||||
"team_id": data.team_id,
|
||||
"AND": [
|
||||
|
|
@ -3243,7 +3214,9 @@ async def bulk_update_team_keys(
|
|||
"error": f"Team {data.team_id} has more than {MAX_BATCH_SIZE} keys. Use `key_ids` to update in batches of {MAX_BATCH_SIZE}."
|
||||
},
|
||||
)
|
||||
requested_tokens = [row.token for row in existing_keys]
|
||||
requested_tokens = cast( # cast-ok: token is the table's primary key, so a row read back always carries one
|
||||
"list[str]", [row.token for row in existing_keys]
|
||||
)
|
||||
else:
|
||||
if data.key_ids is None or len(data.key_ids) == 0:
|
||||
raise HTTPException(
|
||||
|
|
@ -3261,7 +3234,7 @@ async def bulk_update_team_keys(
|
|||
seen_hashes.add(h)
|
||||
requested_tokens.append(k)
|
||||
hashed_key_ids.append(h)
|
||||
existing_keys = await VerificationTokenRepository(prisma_client).table.find_many(
|
||||
existing_keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
|
||||
where={"team_id": data.team_id, "token": {"in": hashed_key_ids}}
|
||||
)
|
||||
|
||||
|
|
@ -3698,7 +3671,7 @@ async def info_key_fn(
|
|||
hashed_key: str | None = key
|
||||
if key is not None:
|
||||
hashed_key = _hash_token_if_needed(token=key)
|
||||
key_info = await VerificationTokenRepository(prisma_client).table.find_unique(
|
||||
key_info = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique(
|
||||
where={"token": hashed_key},
|
||||
include={"litellm_budget_table": True},
|
||||
)
|
||||
|
|
@ -3727,7 +3700,7 @@ async def info_key_fn(
|
|||
key_info = key_info.model_dump()
|
||||
except Exception:
|
||||
# if using pydantic v1
|
||||
key_info = key_info.dict()
|
||||
key_info = key_info.dict() # pyright: ignore[reportDeprecated] # deliberate pydantic v1 fallback
|
||||
key_token_hash: Final = key_info.pop("token")
|
||||
|
||||
model_max_budget = key_info.get("model_max_budget") or {}
|
||||
|
|
@ -4012,7 +3985,10 @@ async def generate_key_helper_fn(
|
|||
if table_name is None or table_name == "user": # do not auto-create users for `/key/generate`
|
||||
## CREATE USER (If necessary)
|
||||
if query_type == "insert_data":
|
||||
user_row = await prisma_client.insert_data(data=user_data, table_name="user")
|
||||
user_row = cast( # cast-ok: table_name="user" is the insert_data branch returning the user row
|
||||
"prisma_models.LiteLLM_UserTable | None",
|
||||
await prisma_client.insert_data(data=user_data, table_name="user"),
|
||||
)
|
||||
|
||||
if user_row is None:
|
||||
raise Exception("Failed to create user")
|
||||
|
|
@ -4219,9 +4195,12 @@ async def delete_verification_tokens(
|
|||
if prisma_client:
|
||||
hashed_tokens: Final[list[str]] = [_hash_token_if_needed(token=key) for key in tokens]
|
||||
tokens = hashed_tokens
|
||||
_keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = await _prisma_table(
|
||||
VerificationTokenRepository(prisma_client)
|
||||
).find_many(where={"token": {"in": hashed_tokens}})
|
||||
_keys_being_deleted: Final[list[LiteLLM_VerificationToken]] = cast( # cast-ok: find_many returns a list
|
||||
"list[LiteLLM_VerificationToken]",
|
||||
await _prisma_table(VerificationTokenRepository(prisma_client)).find_many(
|
||||
where={"token": {"in": hashed_tokens}}
|
||||
),
|
||||
)
|
||||
|
||||
if len(_keys_being_deleted) == 0:
|
||||
raise HTTPException(
|
||||
|
|
@ -4297,7 +4276,7 @@ async def delete_verification_tokens(
|
|||
|
||||
|
||||
def _transform_verification_tokens_to_deleted_records(
|
||||
keys: list[LiteLLM_VerificationToken],
|
||||
keys: Sequence[LiteLLM_VerificationToken],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_changed_by: str | None = None,
|
||||
) -> list[dict[str, object]]:
|
||||
|
|
@ -4372,7 +4351,7 @@ async def _save_deleted_verification_token_records(
|
|||
|
||||
|
||||
async def _persist_deleted_verification_tokens(
|
||||
keys: list[LiteLLM_VerificationToken],
|
||||
keys: Sequence[LiteLLM_VerificationToken],
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_changed_by: str | None = None,
|
||||
|
|
@ -4435,7 +4414,9 @@ async def _rotate_master_key(
|
|||
from litellm.proxy.proxy_server import proxy_config
|
||||
|
||||
try:
|
||||
models: list | None = await _prisma_table(ModelRepository(prisma_client)).find_many()
|
||||
models: list | None = cast( # cast-ok: find_many returns a real list, which TableActions widens to Sequence
|
||||
"list[object]", await _prisma_table(ModelRepository(prisma_client)).find_many()
|
||||
)
|
||||
except Exception:
|
||||
models = None
|
||||
# 2. process model table
|
||||
|
|
@ -5361,9 +5342,9 @@ async def validate_key_list_check(
|
|||
|
||||
if key_hash:
|
||||
try:
|
||||
key_info: Final[LiteLLM_VerificationToken] = await VerificationTokenRepository(
|
||||
prisma_client
|
||||
).table.find_unique(
|
||||
key_info: Final[LiteLLM_VerificationToken | None] = await _prisma_table(
|
||||
VerificationTokenRepository(prisma_client)
|
||||
).find_unique(
|
||||
where={"token": key_hash},
|
||||
)
|
||||
except Exception:
|
||||
|
|
@ -5373,6 +5354,13 @@ async def validate_key_list_check(
|
|||
param="key_hash",
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
)
|
||||
if key_info is None:
|
||||
raise ProxyException(
|
||||
message="Key Hash not found.",
|
||||
type=ProxyErrorTypes.bad_request_error,
|
||||
param="key_hash",
|
||||
code=status.HTTP_403_FORBIDDEN,
|
||||
)
|
||||
can_user_query_key_info: Final = await _can_user_query_key_info(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
key=key_hash,
|
||||
|
|
@ -5394,8 +5382,9 @@ async def _fetch_user_team_objects(
|
|||
if complete_user_info is None or not complete_user_info.teams:
|
||||
return []
|
||||
|
||||
teams: Final[list[BaseModel] | None] = await TeamRepository(prisma_client).table.find_many(
|
||||
where={"team_id": {"in": complete_user_info.teams}}
|
||||
teams: Final[Sequence[BaseModel] | None] = cast( # cast-ok: the None guard below predates the non-optional seam
|
||||
"Sequence[BaseModel] | None",
|
||||
await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": complete_user_info.teams}}),
|
||||
)
|
||||
if teams is None:
|
||||
return []
|
||||
|
|
@ -6130,7 +6119,7 @@ async def _list_key_helper(
|
|||
key_dict = key.model_dump()
|
||||
except Exception:
|
||||
# Fallback for Pydantic v1 compatibility
|
||||
key_dict = key.dict()
|
||||
key_dict = key.dict() # pyright: ignore[reportDeprecated] # deliberate pydantic v1 fallback
|
||||
# Attach object_permission if object_permission_id is set (only for non-deleted keys)
|
||||
if not use_deleted_table:
|
||||
key_dict = await attach_object_permission_to_dict(key_dict, prisma_client)
|
||||
|
|
@ -6155,7 +6144,9 @@ async def _list_key_helper(
|
|||
# Use deleted key type to preserve deleted_at, deleted_by, etc.
|
||||
key_list.append(LiteLLM_DeletedVerificationToken.model_validate(key_dict))
|
||||
else:
|
||||
key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object
|
||||
key_list.append(
|
||||
UserAPIKeyAuth(**key_dict) # pyright: ignore[reportAny] # model_dump() is dict[str, Any]
|
||||
)
|
||||
else:
|
||||
_token = key_dict.get("token")
|
||||
key_list.append(cast(str, _token)) # Return only the token
|
||||
|
|
|
|||
|
|
@ -40,17 +40,22 @@ router: Final = APIRouter()
|
|||
|
||||
|
||||
class _DeploymentRow(Protocol):
|
||||
model_id: str
|
||||
model_name: str
|
||||
model_info: object
|
||||
@property
|
||||
def model_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def model_name(self) -> str: ...
|
||||
|
||||
@property
|
||||
def model_info(self) -> object: ...
|
||||
|
||||
|
||||
class _ModelTableClient(Protocol):
|
||||
async def find_many(self, where: Mapping[str, object] | None = None) -> Sequence[_DeploymentRow]: ...
|
||||
async def find_many(self, *, where: Mapping[str, object] | None = None) -> Sequence[_DeploymentRow]: ...
|
||||
|
||||
async def find_unique(self, where: Mapping[str, object]) -> _DeploymentRow | None: ...
|
||||
async def find_unique(self, *, where: Mapping[str, object]) -> _DeploymentRow | None: ...
|
||||
|
||||
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
|
||||
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
|
||||
|
||||
|
||||
def _model_table(prisma_client: PrismaClient) -> _ModelTableClient:
|
||||
|
|
@ -322,7 +327,9 @@ async def get_all_access_groups_from_db(
|
|||
|
||||
for deployment in deployments:
|
||||
model_info = deployment.model_info or {}
|
||||
access_groups = model_info.get("access_groups", [])
|
||||
access_groups = model_info.get( # pyright: ignore[reportAttributeAccessIssue] # Json reads back as a dict
|
||||
"access_groups", []
|
||||
)
|
||||
model_name = deployment.model_name
|
||||
|
||||
for access_group in access_groups:
|
||||
|
|
|
|||
|
|
@ -16,7 +16,7 @@ import json
|
|||
from collections.abc import Awaitable, Mapping, Sequence
|
||||
from json import JSONDecodeError
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, Protocol, cast
|
||||
from typing import TYPE_CHECKING, Final, Literal, Protocol, cast
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
||||
from pydantic import BaseModel, ConfigDict, Field, ValidationError
|
||||
|
|
@ -72,6 +72,7 @@ from litellm.proxy.spend_tracking.ptu_feature_flag import (
|
|||
)
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import ModelTableRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.router import Router
|
||||
|
|
@ -100,6 +101,9 @@ from litellm.types.router import (
|
|||
)
|
||||
from litellm.utils import get_utc_datetime
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
||||
|
|
@ -120,10 +124,14 @@ class UpdatePublicModelGroupsRequest(BaseModel):
|
|||
|
||||
|
||||
class _ProxyModelRow(Protocol):
|
||||
model_id: str
|
||||
model_name: str
|
||||
litellm_params: Mapping[str, object]
|
||||
model_info: Mapping[str, object] | None
|
||||
@property
|
||||
def model_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def model_name(self) -> str: ...
|
||||
|
||||
@property
|
||||
def model_info(self) -> object: ...
|
||||
|
||||
def model_dump_json(self, *, exclude_none: bool = False) -> str: ...
|
||||
|
||||
|
|
@ -133,7 +141,9 @@ class _ProxyModelTable(Protocol):
|
|||
|
||||
def find_many(self, *, where: Mapping[str, object]) -> Awaitable[Sequence[_ProxyModelRow]]: ...
|
||||
|
||||
def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[_ProxyModelRow]: ...
|
||||
def update(
|
||||
self, *, where: Mapping[str, object], data: Mapping[str, object]
|
||||
) -> Awaitable[_ProxyModelRow | None]: ...
|
||||
|
||||
def delete(self, *, where: Mapping[str, object]) -> Awaitable[_ProxyModelRow | None]: ...
|
||||
|
||||
|
|
@ -144,41 +154,35 @@ class _TxModelTables(Protocol):
|
|||
litellm_proxymodeltable: _ProxyModelTable
|
||||
|
||||
|
||||
class _ExistingModelRow(Protocol):
|
||||
@property
|
||||
def litellm_params(self) -> Mapping[str, object]: ...
|
||||
|
||||
def model_dump_json(self, *, exclude_none: bool = False) -> str: ...
|
||||
|
||||
|
||||
class _TeamRow(Protocol):
|
||||
models: Sequence[str]
|
||||
@property
|
||||
def models(self) -> Sequence[str]: ...
|
||||
|
||||
def model_dump(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
class _TeamTable(Protocol):
|
||||
class _TeamLookupTable(Protocol):
|
||||
def find_unique(self, *, where: Mapping[str, object]) -> Awaitable[_TeamRow | None]: ...
|
||||
|
||||
|
||||
class _TeamTable(_TeamLookupTable, Protocol):
|
||||
def update(
|
||||
self, *, where: Mapping[str, object], data: Mapping[str, object], include: Mapping[str, bool]
|
||||
) -> Awaitable[LiteLLM_TeamTable]: ...
|
||||
|
||||
|
||||
class _TeamIdRef(Protocol):
|
||||
team_id: str
|
||||
|
||||
|
||||
class _ModelAliasRow(Protocol):
|
||||
id: int
|
||||
model_aliases: dict[str, str]
|
||||
team: _TeamIdRef | None
|
||||
|
||||
|
||||
class _ModelAliasTable(Protocol):
|
||||
def find_many(self, *, include: Mapping[str, bool]) -> Awaitable[Sequence[_ModelAliasRow]]: ...
|
||||
|
||||
def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> Awaitable[object]: ...
|
||||
|
||||
|
||||
def _proxy_model_table(prisma_client: PrismaClient) -> _ProxyModelTable:
|
||||
return ModelRepository(prisma_client).table
|
||||
|
||||
|
||||
def _repo_team_table(prisma_client: PrismaClient) -> _TeamTable:
|
||||
def _repo_team_table(prisma_client: PrismaClient) -> _TeamLookupTable:
|
||||
return TeamRepository(prisma_client).table
|
||||
|
||||
|
||||
|
|
@ -186,7 +190,7 @@ def _db_team_table(prisma_client: PrismaClient) -> _TeamTable:
|
|||
return prisma_client.db.litellm_teamtable
|
||||
|
||||
|
||||
def _model_alias_table(prisma_client: PrismaClient) -> _ModelAliasTable:
|
||||
def _model_alias_table(prisma_client: PrismaClient) -> "TableActions[prisma_models.LiteLLM_ModelTable]":
|
||||
return ModelTableRepository(prisma_client).table
|
||||
|
||||
|
||||
|
|
@ -677,6 +681,14 @@ async def patch_model(
|
|||
data=update_data,
|
||||
)
|
||||
|
||||
if updated_model is None:
|
||||
raise ProxyException(
|
||||
message=f"Model {model_id} not found on proxy.",
|
||||
type=ProxyErrorTypes.not_found_error,
|
||||
code=status.HTTP_404_NOT_FOUND,
|
||||
param=None,
|
||||
)
|
||||
|
||||
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
|
||||
live_before_reload: Final = live_model_ids_snapshot()
|
||||
reload_outcome: Final = await clear_cache()
|
||||
|
|
@ -811,7 +823,7 @@ async def _set_model_blocked_status(
|
|||
live_after=reload_outcome.live_after,
|
||||
)
|
||||
|
||||
return updated_model
|
||||
return updated_model # pyright: ignore[reportReturnType] # prisma row, coerced by this route's response_model
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error in model %s: %s", action, e)
|
||||
|
|
@ -897,7 +909,7 @@ async def _add_model_to_db(
|
|||
prisma_client: PrismaClient,
|
||||
new_encryption_key: str | None = None,
|
||||
should_create_model_in_db: bool = True,
|
||||
) -> LiteLLM_ProxyModelTable | None:
|
||||
) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None":
|
||||
# encrypt litellm params #
|
||||
_litellm_params_dict: Final = model_params.litellm_params.dict(exclude_none=True)
|
||||
_original_litellm_model_name: Final = model_params.litellm_params.model
|
||||
|
|
@ -914,8 +926,9 @@ async def _add_model_to_db(
|
|||
}
|
||||
if model_params.model_info.id is not None:
|
||||
_data["model_id"] = model_params.model_info.id
|
||||
_create_data: Final = cast("Mapping[str, object]", _data) # cast-ok: str-keyed json payload built just above
|
||||
if should_create_model_in_db:
|
||||
model_response = await ModelRepository(prisma_client).table.create(data=_data)
|
||||
model_response = await ModelRepository(prisma_client).table.create(data=_create_data)
|
||||
else:
|
||||
model_response = LiteLLM_ProxyModelTable(**_data)
|
||||
return model_response
|
||||
|
|
@ -925,7 +938,7 @@ async def _add_team_model_to_db(
|
|||
model_params: Deployment,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
) -> LiteLLM_ProxyModelTable | None:
|
||||
) -> "prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None":
|
||||
"""
|
||||
If 'team_id' is provided,
|
||||
|
||||
|
|
@ -1638,7 +1651,9 @@ async def delete_team_model_alias(
|
|||
tasks: Final = []
|
||||
removed_model_aliases: Final[list[tuple[str, str]]] = []
|
||||
for team_model_alias in team_model_aliases:
|
||||
model_aliases = team_model_alias.model_aliases # {"alias": "public model name"}
|
||||
model_aliases = cast( # cast-ok: prisma types Json columns as `str`; the driver hands back the parsed dict
|
||||
"dict[str, str]", team_model_alias.model_aliases
|
||||
)
|
||||
id = team_model_alias.id
|
||||
|
||||
if public_model_name in model_aliases.values():
|
||||
|
|
@ -1733,7 +1748,7 @@ async def add_new_model(
|
|||
existing_params=None,
|
||||
)
|
||||
|
||||
model_response: LiteLLM_ProxyModelTable | None = None
|
||||
model_response: prisma_models.LiteLLM_ProxyModelTable | LiteLLM_ProxyModelTable | None = None
|
||||
# update DB
|
||||
incoming_model_info: Final = model_params.model_info.model_dump(exclude_none=True)
|
||||
_raise_if_ptu_cost_attribution_disabled(incoming_model_info)
|
||||
|
|
@ -1902,7 +1917,10 @@ async def update_model(
|
|||
|
||||
# update DB
|
||||
if store_model_in_db is True:
|
||||
_existing_litellm_params_dict: Final = dict(_existing_litellm_params.litellm_params)
|
||||
existing_model_row: Final = cast( # cast-ok: prisma types Json columns as `str`; the driver parses them
|
||||
"_ExistingModelRow", _existing_litellm_params
|
||||
)
|
||||
_existing_litellm_params_dict: Final = dict(existing_model_row.litellm_params)
|
||||
|
||||
if model_params.litellm_params is None:
|
||||
raise Exception("litellm_params not provided")
|
||||
|
|
@ -1946,8 +1964,8 @@ async def update_model(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
table_name=LitellmTableNames.PROXY_MODEL_TABLE_NAME,
|
||||
before_value=(
|
||||
_existing_litellm_params.model_dump_json(exclude_none=True)
|
||||
if isinstance(_existing_litellm_params, BaseModel)
|
||||
existing_model_row.model_dump_json(exclude_none=True)
|
||||
if isinstance(existing_model_row, BaseModel)
|
||||
else None
|
||||
),
|
||||
after_value=(
|
||||
|
|
|
|||
|
|
@ -14,7 +14,14 @@ Endpoints for /organization operations
|
|||
#### ORGANIZATION MANAGEMENT ####
|
||||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Annotated, Final, Protocol, overload
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Annotated,
|
||||
Final,
|
||||
Protocol,
|
||||
cast, # noqa: TID251 # prisma types Json columns as fields.Json but reads back plain python values
|
||||
overload,
|
||||
)
|
||||
|
||||
import fastapi
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
|
|
@ -74,6 +81,11 @@ if TYPE_CHECKING:
|
|||
router: Final = APIRouter()
|
||||
|
||||
|
||||
class _ObjectPermissionRow(Protocol):
|
||||
@property
|
||||
def object_permission_id(self) -> str | None: ...
|
||||
|
||||
|
||||
class _UserTableClient(Protocol):
|
||||
async def find_unique(self, where: Mapping[str, object]) -> "PrismaUserTable | None": ...
|
||||
|
||||
|
|
@ -681,7 +693,10 @@ async def update_organization(
|
|||
existing_metadata: Final = existing_organization_row.metadata or {}
|
||||
updated_metadata: Final = updated_organization_row_json.get("metadata", {})
|
||||
merged_metadata: Final[Mapping[str, object]] = _update_dictionary(
|
||||
existing_dict=existing_metadata.copy(), new_dict=updated_metadata
|
||||
existing_dict=cast( # cast-ok: prisma de-serializes a Json column to the plain python dict it stores
|
||||
"dict[str, object]", existing_metadata
|
||||
).copy(),
|
||||
new_dict=updated_metadata,
|
||||
)
|
||||
updated_organization_row_json["metadata"] = merged_metadata
|
||||
|
||||
|
|
@ -720,7 +735,7 @@ async def update_organization(
|
|||
|
||||
async def handle_update_object_permission(
|
||||
data_json: dict[str, object],
|
||||
existing_organization_row: LiteLLM_OrganizationTable,
|
||||
existing_organization_row: _ObjectPermissionRow,
|
||||
) -> dict[str, object]:
|
||||
"""
|
||||
Handle the update of object permission for an organization.
|
||||
|
|
@ -1276,17 +1291,20 @@ async def find_member_if_email(user_email: str, prisma_client: PrismaClient) ->
|
|||
Find a member if the user_email is in LiteLLM_UserTable
|
||||
"""
|
||||
|
||||
not_unique_user_email_error: Final = HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"Unique user not found for user_email={user_email}. Potential duplicate OR non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead."
|
||||
},
|
||||
)
|
||||
try:
|
||||
existing_user_email_row: Final[BaseModel] = await UserRepository(prisma_client).table.find_unique(
|
||||
existing_user_email_row: Final = await UserRepository(prisma_client).table.find_unique(
|
||||
where={"user_email": user_email}
|
||||
)
|
||||
except Exception:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"Unique user not found for user_email={user_email}. Potential duplicate OR non-existent user_email in LiteLLM_UserTable. Use 'user_id' instead."
|
||||
},
|
||||
)
|
||||
raise not_unique_user_email_error
|
||||
if existing_user_email_row is None:
|
||||
raise not_unique_user_email_error
|
||||
existing_user_email_row_pydantic: Final = LiteLLM_UserTable.model_validate(existing_user_email_row.model_dump())
|
||||
return existing_user_email_row_pydantic
|
||||
|
||||
|
|
@ -1537,7 +1555,10 @@ async def add_member_to_organization(
|
|||
_returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user")
|
||||
if _returned_user is not None:
|
||||
user_object = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
|
||||
elif existing_user_email_row is not None and len(existing_user_email_row) > 1:
|
||||
elif existing_user_email_row is not None and (
|
||||
len(existing_user_email_row) # pyright: ignore[reportArgumentType] # find_unique yields a row, not a list
|
||||
> 1
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "Multiple users with this email found in db. Please use 'user_id' instead."},
|
||||
|
|
|
|||
|
|
@ -33,7 +33,8 @@ class ScimTransformations:
|
|||
|
||||
# Get user's teams/groups
|
||||
groups: Final = []
|
||||
for team_id in user.teams or []:
|
||||
team_ids: Final[list[str]] = user.teams or [] # mutable-ok: scim reads the user row's team ids
|
||||
for team_id in team_ids:
|
||||
team = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id})
|
||||
if team:
|
||||
team_alias = getattr(team, "team_alias", team.team_id)
|
||||
|
|
@ -198,15 +199,8 @@ class ScimTransformations:
|
|||
|
||||
@staticmethod
|
||||
def _get_scim_member_value(member: Member) -> str:
|
||||
"""
|
||||
Get the SCIM member value. Use user_email if available, otherwise use user_id.
|
||||
SCIM member value should be the unique identifier for the user.
|
||||
"""
|
||||
if hasattr(member, "user_email") and member.user_email:
|
||||
return member.user_email
|
||||
elif hasattr(member, "user_id"):
|
||||
return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE
|
||||
return ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE
|
||||
"""The member's SCIM resource id, which LiteLLM serves as user_id (RFC 7643 §8.7.1)."""
|
||||
return member.user_id or ScimTransformations.DEFAULT_SCIM_MEMBER_VALUE
|
||||
|
||||
@staticmethod
|
||||
def _get_scim_member_display(member: Member) -> str:
|
||||
|
|
|
|||
|
|
@ -176,6 +176,10 @@ class UserProvisionerHelpers:
|
|||
is persisted too, so re-upserting an existing email demotes a user who is no
|
||||
longer in the admin group instead of leaving the stale role.
|
||||
|
||||
IdPs like Entra manage membership exclusively through /Groups and never send
|
||||
``groups`` on POST /Users, so a request without teams means "unspecified",
|
||||
not "remove from every team": existing memberships are preserved then.
|
||||
|
||||
Args:
|
||||
prisma_client: Database client
|
||||
new_user_request: New user request data
|
||||
|
|
@ -194,7 +198,8 @@ class UserProvisionerHelpers:
|
|||
if not existing_user:
|
||||
return None
|
||||
|
||||
new_teams: Final = list(dict.fromkeys(new_user_request.teams or []))
|
||||
requested_teams: Final = list(dict.fromkeys(new_user_request.teams or []))
|
||||
new_teams: Final = requested_teams if requested_teams else list(existing_user.teams or [])
|
||||
|
||||
if new_user_request.user_id != existing_user.user_id:
|
||||
verbose_proxy_logger.info(
|
||||
|
|
@ -2761,6 +2766,12 @@ async def patch_group(
|
|||
if final_team:
|
||||
updated_team = final_team
|
||||
|
||||
if updated_team is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Group not found with ID: {group_id}"}, # mutable-ok: FastAPI detail contract
|
||||
)
|
||||
|
||||
# Convert to SCIM format and return
|
||||
scim_group: Final = await ScimTransformations.transform_litellm_team_to_scim_group(
|
||||
LiteLLM_TeamTable.model_validate(updated_team.model_dump())
|
||||
|
|
|
|||
|
|
@ -369,10 +369,10 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str):
|
|||
|
||||
# Prisma returns litellm_params as dict (already parsed from JSON)
|
||||
existing_params = db_model.litellm_params
|
||||
if isinstance(existing_params, str):
|
||||
if isinstance(existing_params, str): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json stub is str
|
||||
# If it's a string, parse it
|
||||
existing_params = json.loads(existing_params)
|
||||
elif not isinstance(existing_params, dict):
|
||||
elif not isinstance(existing_params, dict): # pyright: ignore[reportUnnecessaryIsInstance] # prisma Json stub
|
||||
raise Exception(f"Unexpected litellm_params type: {type(existing_params)}")
|
||||
|
||||
# Add tag to tags array (preserve encryption of other fields)
|
||||
|
|
|
|||
|
|
@ -352,6 +352,9 @@ async def add_team_callbacks(
|
|||
include={"object_permission": True}, # mutable-ok: prisma include takes a dict literal
|
||||
)
|
||||
|
||||
if new_team_row is None:
|
||||
raise _callback_error(400, f"Team id = {team_id} does not exist. Please use a different team id.")
|
||||
|
||||
# Without this a newly registered callback stays dormant for existing keys.
|
||||
await _refresh_cached_team(
|
||||
team_row=new_team_row,
|
||||
|
|
|
|||
|
|
@ -16,11 +16,12 @@ import traceback
|
|||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Final, NamedTuple, NoReturn, Protocol, TypedDict, TypeVar, cast
|
||||
from typing import TYPE_CHECKING, Annotated, Final, NamedTuple, NoReturn, Protocol, TypeVar, cast
|
||||
|
||||
import fastapi
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException, Request, status
|
||||
from pydantic import BaseModel, JsonValue
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -33,21 +34,17 @@ from litellm.proxy._types import (
|
|||
BudgetNewRequest,
|
||||
CommonProxyErrors,
|
||||
DeleteTeamRequest,
|
||||
LiteLLM_AccessGroupTable,
|
||||
LiteLLM_AuditLogs,
|
||||
LiteLLM_BudgetTableFull,
|
||||
LiteLLM_DeletedTeamTable,
|
||||
LiteLLM_ManagementEndpoint_MetadataFields,
|
||||
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
|
||||
LiteLLM_ModelTable,
|
||||
LiteLLM_OrganizationMembershipTable,
|
||||
LiteLLM_OrganizationTable,
|
||||
LiteLLM_OrganizationTableWithMembers,
|
||||
LiteLLM_TeamMembership,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
LiteLLM_UserTable,
|
||||
LiteLLM_VerificationToken,
|
||||
LitellmTableNames,
|
||||
LitellmUserRoles,
|
||||
Member,
|
||||
|
|
@ -116,6 +113,7 @@ from litellm.proxy.management_endpoints.tag_management_endpoints import (
|
|||
get_daily_activity,
|
||||
)
|
||||
from litellm.proxy.management_helpers.access_group_team_sync import (
|
||||
TEAM_ADVISORY_LOCK_SQL,
|
||||
AccessGroupSyncTx,
|
||||
invalidate_access_group_caches,
|
||||
reconcile_team_access_group_membership,
|
||||
|
|
@ -134,12 +132,14 @@ from litellm.proxy.management_helpers.team_metadata_validation import (
|
|||
validate_team_metadata_if_configured,
|
||||
)
|
||||
from litellm.proxy.management_helpers.utils import (
|
||||
MemberWriteTx,
|
||||
add_new_member,
|
||||
management_endpoint_wrapper,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient, ProxyLogging, handle_exception_on_proxy
|
||||
from litellm.repositories.budget_repository import BudgetRepository
|
||||
from litellm.repositories.organization_repository import OrganizationRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import (
|
||||
AccessGroupRepository,
|
||||
DeletedTeamRepository,
|
||||
|
|
@ -171,6 +171,10 @@ from litellm.types.proxy.management_endpoints.team_endpoints import (
|
|||
UpdateTeamMemberPermissionsRequest,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import Prisma
|
||||
from prisma import models as prisma_models
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
_DbRecordT = TypeVar("_DbRecordT")
|
||||
|
|
@ -185,95 +189,14 @@ class _TeamIdGroupRow(TypedDict):
|
|||
_count: _TeamIdKeyCount
|
||||
|
||||
|
||||
class _PrismaTableActions(Protocol[_DbRecordT]):
|
||||
async def find_unique(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
include: Mapping[str, bool] | None = None,
|
||||
) -> _DbRecordT | None: ...
|
||||
|
||||
async def find_first(
|
||||
self,
|
||||
where: Mapping[str, object] | None = None,
|
||||
order: Mapping[str, str] | None = None,
|
||||
) -> _DbRecordT | None: ...
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
where: Mapping[str, object] | None = None,
|
||||
include: Mapping[str, bool] | None = None,
|
||||
order: Mapping[str, str] | None = None,
|
||||
skip: int | None = None,
|
||||
take: int | None = None,
|
||||
cursor: Mapping[str, object] | None = None,
|
||||
) -> list[_DbRecordT]: ...
|
||||
|
||||
async def create(
|
||||
self,
|
||||
data: Mapping[str, object],
|
||||
include: Mapping[str, bool] | None = None,
|
||||
) -> _DbRecordT: ...
|
||||
|
||||
async def create_many(
|
||||
self,
|
||||
data: Sequence[Mapping[str, object]],
|
||||
skip_duplicates: bool | None = None,
|
||||
) -> int: ...
|
||||
|
||||
async def update(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
data: Mapping[str, object],
|
||||
include: Mapping[str, bool] | None = None,
|
||||
) -> _DbRecordT: ...
|
||||
|
||||
async def update_many(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
data: Mapping[str, object],
|
||||
) -> int: ...
|
||||
|
||||
async def upsert(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
data: Mapping[str, Mapping[str, object]],
|
||||
) -> _DbRecordT: ...
|
||||
|
||||
async def delete_many(
|
||||
self,
|
||||
where: Mapping[str, object] | None = None,
|
||||
) -> int: ...
|
||||
|
||||
async def count(
|
||||
self,
|
||||
where: Mapping[str, object] | None = None,
|
||||
) -> int: ...
|
||||
|
||||
async def group_by(
|
||||
self,
|
||||
by: Sequence[str],
|
||||
where: Mapping[str, object] | None = None,
|
||||
count: Mapping[str, bool] | None = None,
|
||||
) -> Sequence[_TeamIdGroupRow]: ...
|
||||
|
||||
|
||||
class _HasTableActions(Protocol[_DbRecordT]):
|
||||
@property
|
||||
def table(self) -> "_PrismaTableActions[_DbRecordT]": ...
|
||||
|
||||
|
||||
def _typed_table(
|
||||
repo: "_HasTableActions[_DbRecordT]", record_type: type[_DbRecordT]
|
||||
) -> "_PrismaTableActions[_DbRecordT]":
|
||||
return repo.table
|
||||
|
||||
|
||||
def _as_object(value: object) -> object:
|
||||
return value
|
||||
|
||||
|
||||
def _nullable(value: _DbRecordT | None) -> _DbRecordT | None:
|
||||
return value
|
||||
def _as_list(rows: Sequence[_DbRecordT]) -> list[_DbRecordT]: # mutable-ok: pydantic list[...] fields reject Sequence
|
||||
return cast( # cast-ok: prisma-client-py find_many returns a list; TableActions only widens it to Sequence
|
||||
"list[_DbRecordT]", rows
|
||||
)
|
||||
|
||||
|
||||
class _UserIdRow(Protocol):
|
||||
|
|
@ -281,33 +204,75 @@ class _UserIdRow(Protocol):
|
|||
def user_id(self) -> str | None: ...
|
||||
|
||||
|
||||
class _HasUserIdTable(Protocol):
|
||||
@property
|
||||
def table(self) -> "_PrismaTableActions[_UserIdRow]": ...
|
||||
|
||||
|
||||
def _user_id_rows_db(repo: "_HasUserIdTable") -> "_PrismaTableActions[_UserIdRow]":
|
||||
def _user_id_rows_db(repo: UserRepository) -> "TableActions[_UserIdRow]":
|
||||
return repo.table
|
||||
|
||||
|
||||
class _RawTeamRow(Protocol):
|
||||
class _ModelDumpRow(Protocol):
|
||||
def model_dump(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
class _TeamIdRow(Protocol):
|
||||
@property
|
||||
def members_with_roles(self) -> Sequence[Mapping[str, object]] | None: ...
|
||||
def team_id(self) -> str: ...
|
||||
|
||||
|
||||
class _HasRawTeamTable(Protocol):
|
||||
class _CacheableTeamRow(_TeamIdRow, _ModelDumpRow, Protocol): ...
|
||||
|
||||
|
||||
class _ObjectPermissionRow(Protocol):
|
||||
@property
|
||||
def table(self) -> "_PrismaTableActions[_RawTeamRow]": ...
|
||||
def object_permission_id(self) -> str | None: ...
|
||||
|
||||
|
||||
def _raw_team_db(repo: "_HasRawTeamTable") -> "_PrismaTableActions[_RawTeamRow]":
|
||||
return repo.table
|
||||
class _TeamAliasBudgetRow(Protocol):
|
||||
@property
|
||||
def team_alias(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def budget_duration(self) -> str | None: ...
|
||||
|
||||
|
||||
class _TeamBudgetRow(_TeamAliasBudgetRow, Protocol):
|
||||
metadata: Mapping[str, JsonValue] | None
|
||||
|
||||
|
||||
class _AuditableTeamRow(Protocol):
|
||||
def json(self, *, exclude_none: bool = False) -> str: ...
|
||||
|
||||
|
||||
class _RawTeamRow(_TeamIdRow, _ModelDumpRow, _ObjectPermissionRow, _TeamBudgetRow, _AuditableTeamRow, Protocol):
|
||||
@property
|
||||
def members_with_roles(
|
||||
self,
|
||||
) -> Sequence[dict[str, object]] | None: ... # mutable-ok: prisma deserializes this JSON column into plain dicts
|
||||
|
||||
@property
|
||||
def organization_id(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def max_budget(self) -> float | None: ...
|
||||
|
||||
@property
|
||||
def soft_budget(self) -> float | None: ...
|
||||
|
||||
@property
|
||||
def model_id(self) -> int | None: ...
|
||||
|
||||
|
||||
def _raw_team_db(repo: TeamRepository) -> "TableActions[_RawTeamRow]":
|
||||
return cast( # cast-ok: prisma types Json columns as str; the client hands back the deserialized value
|
||||
"TableActions[_RawTeamRow]", repo.table
|
||||
)
|
||||
|
||||
|
||||
class _BudgetIdRow(Protocol):
|
||||
@property
|
||||
def budget_id(self) -> str: ...
|
||||
|
||||
|
||||
class _BudgetWriteCall(Protocol):
|
||||
async def __call__(
|
||||
self, budget_obj: BudgetNewRequest, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> LiteLLM_BudgetTableFull: ...
|
||||
async def __call__(self, budget_obj: BudgetNewRequest, user_api_key_dict: UserAPIKeyAuth) -> _BudgetIdRow: ...
|
||||
|
||||
|
||||
def _as_budget_write(fn: "_BudgetWriteCall") -> "_BudgetWriteCall":
|
||||
|
|
@ -330,9 +295,42 @@ class _TeamIdInFilter(TypedDict, total=False):
|
|||
team_id: Mapping[str, Sequence[str]]
|
||||
|
||||
|
||||
class _DeletedTeamsResult(TypedDict):
|
||||
deleted_teams: ReadOnly[Sequence[str]]
|
||||
|
||||
|
||||
class _ErrorDetail(TypedDict):
|
||||
error: ReadOnly[str]
|
||||
|
||||
|
||||
class _TeamCreateTx(AccessGroupSyncTx, Protocol):
|
||||
@property
|
||||
def litellm_teamtable(self) -> "_PrismaTableActions[LiteLLM_TeamTable]": ...
|
||||
def litellm_teamtable(self) -> "TableActions[prisma_models.LiteLLM_TeamTable]": ...
|
||||
|
||||
|
||||
class _MemberDeleteTx(Protocol):
|
||||
"""The tables `/team/member_delete` reads while it holds the team's advisory lock.
|
||||
|
||||
Reading them off the transaction keeps the whole endpoint on the one pooled connection
|
||||
it already checked out: a request that has the lock but still needs another connection
|
||||
can be starved by the lock waiters, which is a deadlock rather than a wait when enough
|
||||
of them hold the rest of the pool."""
|
||||
|
||||
@property
|
||||
def litellm_usertable(self) -> "TableActions[prisma_models.LiteLLM_UserTable]": ...
|
||||
|
||||
@property
|
||||
def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ...
|
||||
|
||||
|
||||
class _TeamDeleteTx(AccessGroupSyncTx, Protocol):
|
||||
async def execute_raw(self, query: str, *args: object) -> int: ...
|
||||
|
||||
@property
|
||||
def litellm_teamtable(self) -> "TableActions[prisma_models.LiteLLM_TeamTable]": ...
|
||||
|
||||
@property
|
||||
def litellm_teammembership(self) -> "TableActions[prisma_models.LiteLLM_TeamMembership]": ...
|
||||
|
||||
|
||||
_STRIP_DELETED_TEAM_FROM_USERS_SQL: Final = """
|
||||
|
|
@ -342,46 +340,52 @@ UPDATE "LiteLLM_UserTable" SET teams = array_remove(teams, $1) WHERE $1 = ANY(te
|
|||
_INCLUDE_MODEL_TABLE: Final = MappingProxyType({"litellm_model_table": True})
|
||||
|
||||
|
||||
def _team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamTable]":
|
||||
return _typed_table(TeamRepository(prisma_client), LiteLLM_TeamTable)
|
||||
def _team_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_TeamTable]":
|
||||
return TeamRepository(prisma_client).table
|
||||
|
||||
|
||||
def _team_membership_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_TeamMembership]":
|
||||
return _typed_table(TeamMembershipRepository(prisma_client), LiteLLM_TeamMembership)
|
||||
def _team_tx_db(tx: "Prisma") -> "TableActions[prisma_models.LiteLLM_TeamTable]":
|
||||
return cast( # cast-ok: generated actions type Json columns as str; TableActions widens inputs to Mapping
|
||||
"TableActions[prisma_models.LiteLLM_TeamTable]", tx.litellm_teamtable
|
||||
)
|
||||
|
||||
|
||||
def _user_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_UserTable]":
|
||||
return _typed_table(UserRepository(prisma_client), LiteLLM_UserTable)
|
||||
def _team_membership_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_TeamMembership]":
|
||||
return TeamMembershipRepository(prisma_client).table
|
||||
|
||||
|
||||
def _model_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_ModelTable]":
|
||||
return _typed_table(ModelTableRepository(prisma_client), LiteLLM_ModelTable)
|
||||
def _user_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_UserTable]":
|
||||
return UserRepository(prisma_client).table
|
||||
|
||||
|
||||
def _org_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_OrganizationTable]":
|
||||
return _typed_table(OrganizationRepository(prisma_client), LiteLLM_OrganizationTable)
|
||||
def _model_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_ModelTable]":
|
||||
return ModelTableRepository(prisma_client).table
|
||||
|
||||
|
||||
def _org_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_OrganizationTable]":
|
||||
return OrganizationRepository(prisma_client).table
|
||||
|
||||
|
||||
def _org_membership_db(
|
||||
prisma_client: PrismaClient | None,
|
||||
) -> "_PrismaTableActions[LiteLLM_OrganizationMembershipTable]":
|
||||
return _typed_table(OrganizationMembershipRepository(prisma_client), LiteLLM_OrganizationMembershipTable)
|
||||
) -> "TableActions[prisma_models.LiteLLM_OrganizationMembership]":
|
||||
return OrganizationMembershipRepository(prisma_client).table
|
||||
|
||||
|
||||
def _budget_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_BudgetTableFull]":
|
||||
return _typed_table(BudgetRepository(prisma_client), LiteLLM_BudgetTableFull)
|
||||
def _budget_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_BudgetTable]":
|
||||
return BudgetRepository(prisma_client).table
|
||||
|
||||
|
||||
def _deleted_team_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_DeletedTeamTable]":
|
||||
return _typed_table(DeletedTeamRepository(prisma_client), LiteLLM_DeletedTeamTable)
|
||||
def _deleted_team_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_DeletedTeamTable]":
|
||||
return DeletedTeamRepository(prisma_client).table
|
||||
|
||||
|
||||
def _access_group_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_AccessGroupTable]":
|
||||
return _typed_table(AccessGroupRepository(prisma_client), LiteLLM_AccessGroupTable)
|
||||
def _access_group_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_AccessGroupTable]":
|
||||
return AccessGroupRepository(prisma_client).table
|
||||
|
||||
|
||||
def _tokens_db(prisma_client: PrismaClient | None) -> "_PrismaTableActions[LiteLLM_VerificationToken]":
|
||||
return _typed_table(VerificationTokenRepository(prisma_client), LiteLLM_VerificationToken)
|
||||
def _tokens_db(prisma_client: PrismaClient | None) -> "TableActions[prisma_models.LiteLLM_VerificationToken]":
|
||||
return VerificationTokenRepository(prisma_client).table
|
||||
|
||||
|
||||
def _sanitize_for_log(value: object) -> str:
|
||||
|
|
@ -394,7 +398,7 @@ def _sanitize_for_log(value: object) -> str:
|
|||
|
||||
|
||||
async def _refresh_cached_team(
|
||||
team_row: LiteLLM_TeamTable,
|
||||
team_row: _CacheableTeamRow,
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> None:
|
||||
|
|
@ -483,7 +487,7 @@ class TeamMemberBudgetHandler:
|
|||
|
||||
@staticmethod
|
||||
async def create_team_member_budget_table(
|
||||
data: NewTeamRequest | LiteLLM_TeamTable,
|
||||
data: NewTeamRequest | _TeamAliasBudgetRow,
|
||||
new_team_data_json: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
team_member_budget: float | None = None,
|
||||
|
|
@ -534,7 +538,7 @@ class TeamMemberBudgetHandler:
|
|||
|
||||
@staticmethod
|
||||
async def upsert_team_member_budget_table(
|
||||
team_table: LiteLLM_TeamTable,
|
||||
team_table: _TeamBudgetRow,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
updated_kv: dict,
|
||||
team_member_budget: float | None = None,
|
||||
|
|
@ -605,7 +609,7 @@ class TeamMemberBudgetHandler:
|
|||
|
||||
@staticmethod
|
||||
async def clear_team_member_budget_fields(
|
||||
team_table: LiteLLM_TeamTable,
|
||||
team_table: _TeamBudgetRow,
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
updated_kv: dict,
|
||||
explicitly_set_fields: set,
|
||||
|
|
@ -1542,7 +1546,7 @@ async def new_team(
|
|||
|
||||
tx: _TeamCreateTx
|
||||
async with prisma_client.db.tx() as tx:
|
||||
team_row: Final[LiteLLM_TeamTable] = await tx.litellm_teamtable.create(
|
||||
team_row: Final[prisma_models.LiteLLM_TeamTable] = await tx.litellm_teamtable.create(
|
||||
data=team_creation_data,
|
||||
include=_INCLUDE_MODEL_TABLE,
|
||||
)
|
||||
|
|
@ -1597,7 +1601,7 @@ async def new_team(
|
|||
|
||||
|
||||
async def _create_team_update_audit_log(
|
||||
existing_team_row: LiteLLM_TeamTable,
|
||||
existing_team_row: _AuditableTeamRow,
|
||||
updated_kv: dict,
|
||||
team_id: str,
|
||||
litellm_changed_by: str | None,
|
||||
|
|
@ -1720,11 +1724,11 @@ async def _auto_add_team_members_to_organization(
|
|||
|
||||
async def fetch_and_validate_organization(
|
||||
organization_id: str,
|
||||
existing_team_row: LiteLLM_TeamTable,
|
||||
existing_team_row: _ModelDumpRow,
|
||||
llm_router: Router | None,
|
||||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth | None = None,
|
||||
) -> LiteLLM_OrganizationTable:
|
||||
) -> "prisma_models.LiteLLM_OrganizationTable":
|
||||
"""
|
||||
Fetch and validate an organization for team update operations.
|
||||
|
||||
|
|
@ -1998,7 +2002,9 @@ async def update_team(
|
|||
validate_budget_duration(data.budget_duration)
|
||||
validate_budget_duration(data.team_member_budget_duration)
|
||||
|
||||
existing_team_row = await _team_db(prisma_client).find_unique(where={"team_id": data.team_id})
|
||||
existing_team_row = await _raw_team_db(TeamRepository(prisma_client)).find_unique(
|
||||
where={"team_id": data.team_id}
|
||||
)
|
||||
|
||||
if existing_team_row is None:
|
||||
raise HTTPException(
|
||||
|
|
@ -2236,18 +2242,16 @@ async def update_team(
|
|||
|
||||
updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv)
|
||||
team_update_data: Final[Mapping[str, object]] = updated_kv
|
||||
team_row: Final[LiteLLM_TeamTable | None] = _nullable(
|
||||
await _team_db(prisma_client).update(
|
||||
where={"team_id": data.team_id},
|
||||
data=team_update_data,
|
||||
# `object_permission` is included so `_refresh_cached_team`
|
||||
# doesn't write a cached team with the relation nulled out —
|
||||
# see team_model_add for the full rationale.
|
||||
include={
|
||||
"litellm_model_table": True,
|
||||
"object_permission": True,
|
||||
},
|
||||
)
|
||||
team_row: Final = await _team_db(prisma_client).update(
|
||||
where={"team_id": data.team_id},
|
||||
data=team_update_data,
|
||||
# `object_permission` is included so `_refresh_cached_team`
|
||||
# doesn't write a cached team with the relation nulled out.
|
||||
# See team_model_add for the full rationale.
|
||||
include={
|
||||
"litellm_model_table": True,
|
||||
"object_permission": True,
|
||||
},
|
||||
)
|
||||
|
||||
if team_row is None or team_row.team_id is None:
|
||||
|
|
@ -2377,7 +2381,7 @@ def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None:
|
|||
updated_kv["budget_limits"] = json.dumps(initialized_windows)
|
||||
|
||||
|
||||
async def handle_update_object_permission(data_json: dict, existing_team_row: LiteLLM_TeamTable) -> dict:
|
||||
async def handle_update_object_permission(data_json: dict, existing_team_row: _ObjectPermissionRow) -> dict:
|
||||
"""
|
||||
Handle the update of object permission for a team.
|
||||
|
||||
|
|
@ -2580,8 +2584,13 @@ async def _process_team_members(
|
|||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_proxy_admin_name: str,
|
||||
tx: MemberWriteTx | None = None,
|
||||
) -> tuple[list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]:
|
||||
"""Process and add new team members."""
|
||||
"""Process and add new team members.
|
||||
|
||||
``tx`` is the caller's open transaction, when it has one, so the member writes run on the
|
||||
connection it already holds instead of checking out a second one.
|
||||
"""
|
||||
updated_users: Final[list[LiteLLM_UserTable]] = []
|
||||
updated_team_memberships: Final[list[LiteLLM_TeamMembership]] = []
|
||||
|
||||
|
|
@ -2607,6 +2616,7 @@ async def _process_team_members(
|
|||
default_team_budget_id=default_team_budget_id,
|
||||
allowed_models=member_allowed_models,
|
||||
budget_duration=data.budget_duration,
|
||||
tx=tx,
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
|
|
@ -2629,6 +2639,7 @@ async def _process_team_members(
|
|||
default_team_budget_id=default_team_budget_id,
|
||||
allowed_models=member_allowed_models,
|
||||
budget_duration=data.budget_duration,
|
||||
tx=tx,
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
|
|
@ -2707,66 +2718,39 @@ async def _add_team_members_to_team(
|
|||
prisma_client: PrismaClient,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_proxy_admin_name: str,
|
||||
) -> tuple[LiteLLM_TeamTable, list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]:
|
||||
"""Add team members to the team.
|
||||
) -> tuple["prisma_models.LiteLLM_TeamTable", list[LiteLLM_UserTable], list[LiteLLM_TeamMembership]]:
|
||||
"""Add team members to the team, under the team's advisory lock.
|
||||
|
||||
The members_with_roles reconciliation runs inside a transaction that locks
|
||||
the team row with ``SELECT ... FOR UPDATE`` before reading the current
|
||||
membership. Concurrent /team/member_add calls for the same team therefore
|
||||
serialize on the row lock and each appends onto the other's committed
|
||||
result, instead of both rewriting the whole JSON array from a stale
|
||||
snapshot (which silently drops one member on the losing write).
|
||||
The lock (``TEAM_ADVISORY_LOCK_SQL``, keyed on the team id) is taken first, and the
|
||||
team is re-read under it before any write, so a delete that already committed is
|
||||
visible here before this call writes anything: the user and membership writes only
|
||||
happen once the re-read proves the team is still live. /team/delete takes the same
|
||||
lock around its own sweep-and-delete, so the two can never interleave; whichever
|
||||
acquires the lock first runs to completion before the other's re-read can proceed.
|
||||
|
||||
The same lock serializes this against /team/delete: the delete cannot remove
|
||||
the row while the reconcile holds it, and a reconcile that finds the row
|
||||
already gone cleans up after itself rather than leaving the member pointing
|
||||
at a deleted team id.
|
||||
"""
|
||||
# Process and add new members
|
||||
updated_users, updated_team_memberships = await _process_team_members(
|
||||
data=data,
|
||||
complete_team_data=complete_team_data,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
|
||||
updated_team: Final = await _write_members_with_roles_locked(
|
||||
data=data,
|
||||
complete_team_data=complete_team_data,
|
||||
prisma_client=prisma_client,
|
||||
updated_users=updated_users,
|
||||
)
|
||||
if updated_team is None:
|
||||
await _sweep_deleted_team_references(team_ids=(data.team_id,), prisma_client=prisma_client)
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Team={data.team_id} was deleted while this member add was running"},
|
||||
)
|
||||
|
||||
return updated_team, updated_users, updated_team_memberships
|
||||
|
||||
|
||||
async def _write_members_with_roles_locked(
|
||||
data: TeamMemberAddRequest,
|
||||
complete_team_data: LiteLLM_TeamTable,
|
||||
prisma_client: PrismaClient,
|
||||
updated_users: list[LiteLLM_UserTable],
|
||||
) -> LiteLLM_TeamTable | None:
|
||||
"""Reconcile members_with_roles under the team row lock. None when the team row is gone.
|
||||
|
||||
That read is at least as recent as the user and membership writes the caller
|
||||
already made, so a missing row means /team/delete committed after them. Its
|
||||
post-delete sweep can have run before those writes landed, which is why the
|
||||
caller sweeps this team id again rather than only reporting the 404.
|
||||
The user and membership writes run on this transaction too, not on a second
|
||||
connection from the pool: a lock waiter that needs a connection it hasn't got yet is
|
||||
a waiter that can deadlock the pool, since enough concurrent adds for one team would
|
||||
hold every connection waiting on the lock while the holder waits for a free one.
|
||||
"""
|
||||
gone_detail: Final[_ErrorDetail] = {"error": f"Team={data.team_id} was deleted while this member add was running"}
|
||||
async with prisma_client.tx() as tx:
|
||||
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, data.team_id)
|
||||
|
||||
locked_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id)
|
||||
if locked_members is None:
|
||||
return None
|
||||
|
||||
raise HTTPException(status_code=404, detail=gone_detail)
|
||||
complete_team_data.members_with_roles = locked_members
|
||||
|
||||
updated_users, updated_team_memberships = await _process_team_members(
|
||||
data=data,
|
||||
complete_team_data=complete_team_data,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
tx=tx,
|
||||
)
|
||||
|
||||
await _update_team_members_list(
|
||||
data=data,
|
||||
complete_team_data=complete_team_data,
|
||||
|
|
@ -2774,10 +2758,14 @@ async def _write_members_with_roles_locked(
|
|||
)
|
||||
|
||||
_db_team_members: Final = [m.model_dump() for m in complete_team_data.members_with_roles]
|
||||
return await tx.litellm_teamtable.update(
|
||||
updated_team: Final = await _team_tx_db(tx).update(
|
||||
where={"team_id": data.team_id},
|
||||
data={"members_with_roles": json.dumps(_db_team_members)},
|
||||
)
|
||||
if updated_team is None:
|
||||
raise HTTPException(status_code=404, detail=gone_detail)
|
||||
|
||||
return updated_team, updated_users, updated_team_memberships
|
||||
|
||||
|
||||
def _emit_team_members_metric(team: LiteLLM_TeamTable) -> None:
|
||||
|
|
@ -3159,10 +3147,6 @@ async def team_member_add(
|
|||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
)
|
||||
|
||||
# Check if updated_team is None
|
||||
if updated_team is None:
|
||||
raise HTTPException(status_code=404, detail={"error": f"Team with id {data.team_id} not found"})
|
||||
|
||||
_emit_team_members_metric(complete_team_data)
|
||||
|
||||
await _create_team_member_add_audit_logs(
|
||||
|
|
@ -3276,46 +3260,62 @@ async def team_member_delete(
|
|||
)
|
||||
|
||||
## DELETE MEMBER FROM TEAM
|
||||
removed_team_members, new_team_members = _cleanup_members_with_roles(
|
||||
existing_team_row=existing_team_row,
|
||||
data=data,
|
||||
)
|
||||
|
||||
if not removed_team_members:
|
||||
raise HTTPException(status_code=400, detail={"error": "User not found in team"})
|
||||
|
||||
existing_team_row.members_with_roles = new_team_members
|
||||
|
||||
_db_new_team_members: Final[list[dict]] = [m.model_dump() for m in new_team_members]
|
||||
|
||||
## DELETE TEAM ID from USER ROW, IF EXISTS ##
|
||||
# get user row
|
||||
removed_user_ids: Final = frozenset(m.user_id for m in removed_team_members if m.user_id is not None)
|
||||
key_val: Final[Mapping[str, object]] = (
|
||||
{"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email}
|
||||
)
|
||||
existing_user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(where=key_val)
|
||||
|
||||
# Also clean up any existing team membership rows for this user and team
|
||||
user_ids_to_delete: Final = removed_user_ids.union(
|
||||
(data.user_id,) if data.user_id is not None else (),
|
||||
(user.user_id for user in existing_user_rows if user.user_id),
|
||||
)
|
||||
|
||||
## DELETE KEYS CREATED BY USER FOR THIS TEAM
|
||||
# Fetch keys before deletion so their audit records can be persisted alongside the delete.
|
||||
# An empty user_ids_to_delete still resolves cleanly: prisma's "in": [] matches no rows.
|
||||
keys_to_delete: Final[list[LiteLLM_VerificationToken]] = await _tokens_db(prisma_client).find_many(
|
||||
where={
|
||||
"user_id": {"in": sorted(user_ids_to_delete)},
|
||||
"team_id": data.team_id,
|
||||
}
|
||||
)
|
||||
|
||||
# All four cleanups run on one connection so a failure between them leaves
|
||||
# no partial removal: either every write below lands, or none of them do.
|
||||
# Everything from here on runs under the team's advisory lock, the same one
|
||||
# /team/member_add and /team/delete take: without it, this endpoint's own row-level
|
||||
# update lock used to be the only thing serializing it against a concurrent member_add,
|
||||
# and only by accident (their SELECT ... FOR UPDATE contended for the same row lock this
|
||||
# UPDATE takes). Now that member_add reads under the advisory lock instead, this has to
|
||||
# take it too, and re-read the roster under it rather than off the snapshot validated
|
||||
# above, or a member_add that commits in between can have its addition silently
|
||||
# overwritten by this delete computing from stale data.
|
||||
async with prisma_client.tx() as tx:
|
||||
await tx.litellm_teamtable.update(
|
||||
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, data.team_id)
|
||||
|
||||
fresh_members: Final = await TeamRepository(prisma_client).get_members_with_roles_locked(tx, data.team_id)
|
||||
if fresh_members is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": f"Team id={data.team_id} does not exist in db"},
|
||||
)
|
||||
|
||||
removed_team_members, new_team_members = _cleanup_members_with_roles(
|
||||
existing_team_row=LiteLLM_TeamTable(team_id=data.team_id, members_with_roles=fresh_members),
|
||||
data=data,
|
||||
)
|
||||
|
||||
if not removed_team_members:
|
||||
raise HTTPException(status_code=400, detail={"error": "User not found in team"})
|
||||
|
||||
existing_team_row.members_with_roles = new_team_members
|
||||
|
||||
_db_new_team_members: Final[list[dict]] = [m.model_dump() for m in new_team_members]
|
||||
|
||||
## DELETE TEAM ID from USER ROW, IF EXISTS ##
|
||||
# get user row
|
||||
removed_user_ids: Final = frozenset(m.user_id for m in removed_team_members if m.user_id is not None)
|
||||
key_val: Final[Mapping[str, object]] = (
|
||||
{"user_id": {"in": sorted(removed_user_ids)}} if removed_user_ids else {"user_email": data.user_email}
|
||||
)
|
||||
member_tx: Final[_MemberDeleteTx] = tx
|
||||
existing_user_rows: Final = await member_tx.litellm_usertable.find_many(where=key_val)
|
||||
|
||||
# Also clean up any existing team membership rows for this user and team
|
||||
user_ids_to_delete: Final = removed_user_ids.union(
|
||||
(data.user_id,) if data.user_id is not None else (),
|
||||
(user.user_id for user in existing_user_rows if user.user_id),
|
||||
)
|
||||
|
||||
## DELETE KEYS CREATED BY USER FOR THIS TEAM
|
||||
# Fetch keys before deletion so their audit records can be persisted alongside the delete.
|
||||
# An empty user_ids_to_delete still resolves cleanly: prisma's "in": [] matches no rows.
|
||||
keys_to_delete: Final = await member_tx.litellm_verificationtoken.find_many(
|
||||
where={
|
||||
"user_id": {"in": sorted(user_ids_to_delete)},
|
||||
"team_id": data.team_id,
|
||||
}
|
||||
)
|
||||
|
||||
await _team_tx_db(tx).update(
|
||||
where={"team_id": data.team_id},
|
||||
data={"members_with_roles": json.dumps(_db_new_team_members)},
|
||||
)
|
||||
|
|
@ -3953,9 +3953,7 @@ async def delete_team(
|
|||
_persist_deleted_verification_tokens,
|
||||
)
|
||||
|
||||
keys_to_delete: list[LiteLLM_VerificationToken] = await _tokens_db(prisma_client).find_many(
|
||||
where={"team_id": {"in": data.team_ids}}
|
||||
)
|
||||
keys_to_delete: Final = await _tokens_db(prisma_client).find_many(where={"team_id": {"in": data.team_ids}})
|
||||
|
||||
if keys_to_delete:
|
||||
await _persist_deleted_verification_tokens(
|
||||
|
|
@ -4009,7 +4007,21 @@ async def delete_team(
|
|||
await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
|
||||
|
||||
## DELETE TEAMS
|
||||
deleted_teams: Final = await prisma_client.delete_data(team_id_list=data.team_ids, table_name="team")
|
||||
# Both the delete and the reconcile sweep run under every team's advisory lock
|
||||
# (TEAM_ADVISORY_LOCK_SQL, the same one /team/member_add takes before its own writes),
|
||||
# sorted so two overlapping batch deletes always request their locks in the same order.
|
||||
# A member_add mid-flight for one of these teams either finishes its write and releases
|
||||
# the lock before this transaction starts, in which case this sweep reaches what it wrote,
|
||||
# or is still waiting on the lock, in which case its own re-read happens after this commits
|
||||
# and sees the row gone before it writes anything.
|
||||
delete_filter: Final[_TeamIdInFilter] = {"team_id": {"in": data.team_ids}}
|
||||
async with prisma_client.tx() as tx:
|
||||
for team_id in sorted(data.team_ids):
|
||||
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
|
||||
await tx.litellm_teamtable.delete_many(where=delete_filter)
|
||||
await _sweep_deleted_team_references_tx(team_ids=data.team_ids, tx=tx)
|
||||
|
||||
deleted_teams: Final[_DeletedTeamsResult] = {"deleted_teams": data.team_ids}
|
||||
|
||||
# Evict AFTER the rows are gone. Both writers of these keys (`_cache_team_object` and
|
||||
# `get_team_object_by_alias`) hydrate from the db, so evicting first leaves a window where a
|
||||
|
|
@ -4022,12 +4034,6 @@ async def delete_team(
|
|||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
# Sweep again now the team is gone. A `/team/member_add` that landed between the first sweep
|
||||
# and the delete would have re-appended the reference; an add still in flight sees the row
|
||||
# missing under its own row lock and sweeps what it wrote. Both passes are idempotent, and
|
||||
# keeping the first one means a failure here still leaves a team the admin can retry deleting.
|
||||
await _sweep_deleted_team_references(team_ids=data.team_ids, prisma_client=prisma_client)
|
||||
|
||||
for deleted_team in team_rows:
|
||||
await sync_team_access_group_membership(prisma_client=prisma_client, team_id=deleted_team.team_id)
|
||||
|
||||
|
|
@ -4056,8 +4062,18 @@ async def _sweep_deleted_team_references(team_ids: Sequence[str], prisma_client:
|
|||
_ = await _team_membership_db(prisma_client).delete_many(where=_TeamIdInFilter(team_id={"in": tuple(team_ids)}))
|
||||
|
||||
|
||||
async def _sweep_deleted_team_references_tx(team_ids: Sequence[str], tx: _TeamDeleteTx) -> None:
|
||||
"""Same sweep as `_sweep_deleted_team_references`, run on the transaction that holds
|
||||
every id's advisory lock and deletes the team rows, so it commits or rolls back with them."""
|
||||
for team_id in team_ids:
|
||||
_ = await tx.execute_raw(_STRIP_DELETED_TEAM_FROM_USERS_SQL, team_id)
|
||||
|
||||
membership_filter: Final[_TeamIdInFilter] = {"team_id": {"in": tuple(team_ids)}}
|
||||
_ = await tx.litellm_teammembership.delete_many(where=membership_filter)
|
||||
|
||||
|
||||
async def _invalidate_deleted_key_cache(
|
||||
keys: Sequence[LiteLLM_VerificationToken],
|
||||
keys: "Sequence[prisma_models.LiteLLM_VerificationToken]",
|
||||
user_api_key_cache: UserApiKeyCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> None:
|
||||
|
|
@ -4242,7 +4258,7 @@ async def _hydrate_member_emails(
|
|||
if not missing_user_ids:
|
||||
return tuple(members)
|
||||
|
||||
user_rows: Final[Sequence[LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(
|
||||
user_rows: Final[Sequence[prisma_models.LiteLLM_UserTable]] = await _user_db(prisma_client).find_many(
|
||||
where={ # mutable-ok: Prisma query filters are dict-shaped
|
||||
"user_id": { # mutable-ok: Prisma query filters are dict-shaped
|
||||
"in": sorted(missing_user_ids)
|
||||
|
|
@ -4253,7 +4269,7 @@ async def _hydrate_member_emails(
|
|||
|
||||
return tuple(
|
||||
m.model_copy(update={"user_email": email_by_user_id[m.user_id]}) # mutable-ok: pydantic update payload
|
||||
if not m.user_email and m.user_id in email_by_user_id
|
||||
if not m.user_email and m.user_id is not None and m.user_id in email_by_user_id
|
||||
else m
|
||||
for m in members
|
||||
)
|
||||
|
|
@ -4838,7 +4854,7 @@ async def _build_team_list_where_conditions(
|
|||
|
||||
async def _batch_resolve_access_group_resources(
|
||||
all_access_group_ids: list[str],
|
||||
) -> dict[str, LiteLLM_AccessGroupTable]:
|
||||
) -> "dict[str, prisma_models.LiteLLM_AccessGroupTable]":
|
||||
"""
|
||||
Batch-fetch access groups in a single DB query and return them keyed by
|
||||
access_group_id. Missing/invalid groups are silently omitted.
|
||||
|
|
@ -4856,7 +4872,7 @@ async def _batch_resolve_access_group_resources(
|
|||
|
||||
|
||||
def _convert_teams_to_response_models(
|
||||
teams: list,
|
||||
teams: Sequence,
|
||||
use_deleted_table: bool,
|
||||
keys_count_by_team: dict[str, int] | None = None,
|
||||
) -> list[TeamListItem | LiteLLM_TeamTable | LiteLLM_DeletedTeamTable]:
|
||||
|
|
@ -4890,7 +4906,7 @@ def _convert_teams_to_response_models(
|
|||
|
||||
async def _get_keys_count_by_team(
|
||||
prisma_client: PrismaClient,
|
||||
teams: Sequence[LiteLLM_TeamTable],
|
||||
teams: Sequence[_TeamIdRow],
|
||||
) -> dict[str, int]:
|
||||
"""Aggregate virtual-key counts per team for the given page of teams.
|
||||
|
||||
|
|
@ -4902,10 +4918,13 @@ async def _get_keys_count_by_team(
|
|||
if not page_team_ids:
|
||||
return {}
|
||||
|
||||
grouped: Final = await _tokens_db(prisma_client).group_by(
|
||||
by=["team_id"],
|
||||
where={"team_id": {"in": page_team_ids}},
|
||||
count={"team_id": True},
|
||||
grouped: Final = cast( # cast-ok: prisma group_by returns one row per `by` key with `count=` nested under "_count"
|
||||
"Sequence[_TeamIdGroupRow]",
|
||||
await _tokens_db(prisma_client).group_by(
|
||||
by=["team_id"],
|
||||
where={"team_id": {"in": page_team_ids}},
|
||||
count={"team_id": True},
|
||||
),
|
||||
)
|
||||
return {row["team_id"]: row.get("_count", {}).get("team_id", 0) for row in grouped if row.get("team_id")}
|
||||
|
||||
|
|
@ -5295,7 +5314,7 @@ async def list_team(
|
|||
_team_memberships.append(tm)
|
||||
|
||||
# add all keys that belong to the team
|
||||
keys = await _tokens_db(prisma_client).find_many(where={"team_id": team.team_id})
|
||||
keys = _as_list(await _tokens_db(prisma_client).find_many(where={"team_id": team.team_id}))
|
||||
|
||||
try:
|
||||
returned_responses.append(
|
||||
|
|
@ -5530,6 +5549,11 @@ async def team_model_add(
|
|||
data={"updated_at": datetime.now(timezone.utc)},
|
||||
include={"object_permission": True},
|
||||
)
|
||||
if updated_team is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Team not found, passed team_id={data.team_id}"},
|
||||
)
|
||||
|
||||
await _refresh_cached_team(
|
||||
team_row=updated_team,
|
||||
|
|
@ -5612,6 +5636,11 @@ async def team_model_delete(
|
|||
data={"models": updated_models},
|
||||
include={"object_permission": True},
|
||||
)
|
||||
if updated_team is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Team not found, passed team_id={data.team_id}"},
|
||||
)
|
||||
|
||||
await _refresh_cached_team(
|
||||
team_row=updated_team,
|
||||
|
|
@ -5746,8 +5775,13 @@ async def update_team_member_permissions(
|
|||
where={"team_id": data.team_id},
|
||||
data={"team_member_permissions": data.team_member_permissions},
|
||||
)
|
||||
if updated_team is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Team not found, passed team_id={data.team_id}"},
|
||||
)
|
||||
|
||||
return updated_team
|
||||
return updated_team # pyright: ignore[reportReturnType] # prisma row, coerced by this route's response_model
|
||||
|
||||
|
||||
@router.post(
|
||||
|
|
@ -5812,7 +5846,9 @@ async def bulk_update_team_member_permissions(
|
|||
}
|
||||
|
||||
|
||||
async def _compute_and_batch_updates(prisma_client, teams: Sequence[LiteLLM_TeamTable], permissions_to_add: set) -> int:
|
||||
async def _compute_and_batch_updates(
|
||||
prisma_client, teams: "Sequence[prisma_models.LiteLLM_TeamTable]", permissions_to_add: set
|
||||
) -> int:
|
||||
"""Compute merged permissions and batch-write updates. Returns count of teams updated."""
|
||||
updates: Final = []
|
||||
for team in teams:
|
||||
|
|
|
|||
|
|
@ -29,7 +29,6 @@ from typing import (
|
|||
NoReturn,
|
||||
Optional,
|
||||
Protocol,
|
||||
TypeVar,
|
||||
Union,
|
||||
cast,
|
||||
overload,
|
||||
|
|
@ -122,6 +121,7 @@ from litellm.proxy.utils import (
|
|||
get_custom_url,
|
||||
get_server_root_path,
|
||||
)
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import SSOConfigRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
|
|
@ -171,51 +171,16 @@ _CLI_SSO_SECRET_KEY_FRAGMENTS: Final = frozenset(
|
|||
}
|
||||
)
|
||||
|
||||
_DbRecordT: Final = TypeVar("_DbRecordT", covariant=True)
|
||||
|
||||
|
||||
class _PrismaTableActions(Protocol[_DbRecordT]):
|
||||
async def find_unique(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
) -> _DbRecordT | None: ...
|
||||
|
||||
async def find_first(
|
||||
self,
|
||||
where: Mapping[str, object] | None = None,
|
||||
) -> _DbRecordT | None: ...
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
where: Mapping[str, object] | None = None,
|
||||
include: Mapping[str, bool] | None = None,
|
||||
) -> Sequence[_DbRecordT]: ...
|
||||
|
||||
async def update(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
data: Mapping[str, object],
|
||||
) -> _DbRecordT: ...
|
||||
|
||||
async def update_many(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
data: Mapping[str, object],
|
||||
) -> int: ...
|
||||
|
||||
|
||||
class _UserMetadataRow(Protocol):
|
||||
@property
|
||||
def metadata(self) -> Mapping[str, object] | None: ...
|
||||
|
||||
|
||||
class _HasUserMetadataTable(Protocol):
|
||||
@property
|
||||
def table(self) -> "_PrismaTableActions[_UserMetadataRow]": ...
|
||||
|
||||
|
||||
def _user_meta_db(repo: "_HasUserMetadataTable") -> "_PrismaTableActions[_UserMetadataRow]":
|
||||
return repo.table
|
||||
def _user_meta_db(repo: UserRepository) -> "TableActions[_UserMetadataRow]":
|
||||
return cast( # cast-ok: prisma types Json columns as str; the client hands back the deserialized value
|
||||
"TableActions[_UserMetadataRow]", repo.table
|
||||
)
|
||||
|
||||
|
||||
class _SsoConfigRow(Protocol):
|
||||
|
|
@ -223,25 +188,17 @@ class _SsoConfigRow(Protocol):
|
|||
def sso_settings(self) -> Mapping[str, object] | None: ...
|
||||
|
||||
|
||||
class _HasSsoConfigTable(Protocol):
|
||||
@property
|
||||
def table(self) -> "_PrismaTableActions[_SsoConfigRow]": ...
|
||||
|
||||
|
||||
def _sso_config_db(repo: "_HasSsoConfigTable") -> "_PrismaTableActions[_SsoConfigRow]":
|
||||
return repo.table
|
||||
def _sso_config_db(repo: SSOConfigRepository) -> "TableActions[_SsoConfigRow]":
|
||||
return cast( # cast-ok: prisma types Json columns as str; the client hands back the deserialized value
|
||||
"TableActions[_SsoConfigRow]", repo.table
|
||||
)
|
||||
|
||||
|
||||
class _TeamDetailRow(Protocol):
|
||||
def model_dump(self) -> Mapping[str, object]: ...
|
||||
|
||||
|
||||
class _HasTeamDetailTable(Protocol):
|
||||
@property
|
||||
def table(self) -> "_PrismaTableActions[_TeamDetailRow]": ...
|
||||
|
||||
|
||||
def _team_detail_db(repo: "_HasTeamDetailTable") -> "_PrismaTableActions[_TeamDetailRow]":
|
||||
def _team_detail_db(repo: TeamRepository) -> "TableActions[_TeamDetailRow]":
|
||||
return repo.table
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -23,9 +23,11 @@ from pydantic import BaseModel, TypeAdapter
|
|||
from litellm.proxy.auth.auth_checks import _delete_cache_access_object
|
||||
|
||||
# hashtext collisions only cost two unrelated teams a little serialization, and the
|
||||
# lock is never taken by the access-group endpoints, so it cannot join their
|
||||
# access-group-then-team lock order to form a cycle.
|
||||
_LOCK_TEAM_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
|
||||
# lock is never taken by the access-group endpoints as a SELECT ... FOR UPDATE row lock,
|
||||
# so it cannot join their access-group-then-team lock order to form a cycle. team_endpoints
|
||||
# reuses this exact statement to serialize /team/member_add and /team/delete against each
|
||||
# other and against this mirror, rather than defining a second, divergent lock on the same key.
|
||||
TEAM_ADVISORY_LOCK_SQL: Final = "SELECT pg_advisory_xact_lock(hashtext($1)) IS NULL AS locked"
|
||||
|
||||
_READ_TEAM_SQL: Final = 'SELECT access_group_ids FROM "LiteLLM_TeamTable" WHERE team_id = $1'
|
||||
|
||||
|
|
@ -138,7 +140,7 @@ async def reconcile_team_access_group_membership(tx: AccessGroupSyncTx, team_id:
|
|||
concurrent write for a different team cannot be lost the way a read-modify-write of
|
||||
the whole array can, and the pair commits together or not at all.
|
||||
"""
|
||||
await tx.query_raw(_LOCK_TEAM_SQL, team_id)
|
||||
await tx.query_raw(TEAM_ADVISORY_LOCK_SQL, team_id)
|
||||
team_rows: Final = _TeamRows.validate_python(await tx.query_raw(_READ_TEAM_SQL, team_id))
|
||||
desired: Final = (team_rows[0].access_group_ids or ()) if team_rows else ()
|
||||
affected: Final = _AffectedGroups.validate_python(await tx.query_raw(_AFFECTED_SQL, team_id, desired))
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ organizations, teams, and keys.
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Any, Final, Optional
|
||||
|
||||
|
|
@ -19,6 +19,8 @@ from litellm.repositories.object_permission_repository import ObjectPermissionRe
|
|||
from litellm.repositories.table_repositories import MCPServerRepository
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_ObjectPermissionTable,
|
||||
LiteLLM_TeamTableCachedObj,
|
||||
|
|
@ -26,7 +28,7 @@ if TYPE_CHECKING:
|
|||
|
||||
|
||||
async def attach_object_permission_to_dict(
|
||||
data_dict: dict,
|
||||
data_dict: dict[str, object],
|
||||
prisma_client: PrismaClient,
|
||||
) -> dict:
|
||||
"""
|
||||
|
|
@ -61,7 +63,7 @@ async def attach_object_permission_to_dict(
|
|||
try:
|
||||
object_permission = object_permission.model_dump()
|
||||
except Exception:
|
||||
object_permission = object_permission.dict()
|
||||
object_permission = object_permission.dict() # pyright: ignore[reportDeprecated] # pydantic v1 fallback
|
||||
data_dict["object_permission"] = object_permission
|
||||
return data_dict
|
||||
|
||||
|
|
@ -188,7 +190,9 @@ async def _set_object_permission(
|
|||
return data_json
|
||||
|
||||
# Clean data: exclude None values and object_permission_id
|
||||
clean_data: Final = {k: v for k, v in permission_data.items() if v is not None and k != "object_permission_id"}
|
||||
clean_data: Final[dict[str, object]] = {
|
||||
k: v for k, v in permission_data.items() if v is not None and k != "object_permission_id"
|
||||
}
|
||||
|
||||
# Serialize mcp_tool_permissions to JSON string for GraphQL compatibility
|
||||
if "mcp_tool_permissions" in clean_data:
|
||||
|
|
@ -224,7 +228,7 @@ def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool:
|
|||
async def _get_db_mcp_servers_by_identifiers(
|
||||
identifiers: set[str],
|
||||
prisma_client: PrismaClient | None,
|
||||
) -> list[Any]:
|
||||
) -> "Sequence[prisma_models.LiteLLM_MCPServerTable]":
|
||||
if prisma_client is None or not identifiers:
|
||||
return []
|
||||
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ from litellm.proxy._types import ( # key request types; user request types; tea
|
|||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
|
||||
from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.proxy.utils import PrismaClient, jsonify_object
|
||||
from litellm.repositories.budget_repository import BudgetRepository
|
||||
from litellm.repositories.table_repositories import TeamMembershipRepository
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
|
|
@ -79,6 +79,8 @@ class _PrismaUserTable(Protocol):
|
|||
self, *, where: Mapping[str, object], data: Mapping[str, Mapping[str, object]]
|
||||
) -> _PrismaUserRecord | None: ...
|
||||
|
||||
async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_PrismaUserRecord]: ...
|
||||
|
||||
|
||||
class _PrismaTeamMembershipTable(Protocol):
|
||||
"""Team membership table actions the management helpers issue."""
|
||||
|
|
@ -86,6 +88,73 @@ class _PrismaTeamMembershipTable(Protocol):
|
|||
async def create(self, *, data: Mapping[str, object], include: Mapping[str, bool]) -> _PrismaRecord: ...
|
||||
|
||||
|
||||
class MemberWriteTx(Protocol):
|
||||
"""Transaction surface `add_new_member` writes through when the caller owns one.
|
||||
|
||||
A caller already holding a transaction, and with it a pooled connection plus that
|
||||
transaction's locks, passes it here so these writes reuse that connection rather than
|
||||
checking out another one that lock waiters may already have drained from the pool.
|
||||
"""
|
||||
|
||||
@property
|
||||
def litellm_usertable(self) -> _PrismaUserTable: ...
|
||||
|
||||
@property
|
||||
def litellm_budgettable(self) -> _PrismaBudgetTable: ...
|
||||
|
||||
@property
|
||||
def litellm_teammembership(self) -> _PrismaTeamMembershipTable: ...
|
||||
|
||||
|
||||
def _user_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaUserTable:
|
||||
return tx.litellm_usertable if tx is not None else UserRepository(prisma_client).table
|
||||
|
||||
|
||||
def _budget_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaBudgetTable:
|
||||
return tx.litellm_budgettable if tx is not None else BudgetRepository(prisma_client).table
|
||||
|
||||
|
||||
def _team_membership_table(prisma_client: PrismaClient, tx: MemberWriteTx | None) -> _PrismaTeamMembershipTable:
|
||||
return tx.litellm_teammembership if tx is not None else TeamMembershipRepository(prisma_client).table
|
||||
|
||||
|
||||
async def _find_users_by_email(
|
||||
prisma_client: PrismaClient, tx: MemberWriteTx | None, user_email: str
|
||||
) -> Sequence[_PrismaUserRecord]:
|
||||
if tx is not None:
|
||||
return await tx.litellm_usertable.find_many(where={"user_email": user_email})
|
||||
rows: Final[Sequence[_PrismaUserRecord] | None] = await prisma_client.get_data(
|
||||
key_val={"user_email": user_email},
|
||||
table_name="user",
|
||||
query_type="find_all",
|
||||
)
|
||||
return rows if rows is not None else ()
|
||||
|
||||
|
||||
async def _upsert_user_row(
|
||||
user_table: _PrismaUserTable, user_id: str, create_data: Mapping[str, object]
|
||||
) -> _PrismaUserRecord | None:
|
||||
"""Insert the user row if it is absent, leaving an existing row as it is.
|
||||
|
||||
Upserting keeps concurrent provisioning of the same new user from racing on create.
|
||||
The update branch re-states user_id rather than being empty because Prisma only
|
||||
compiles an upsert down to INSERT ... ON CONFLICT when the update is non-empty, and
|
||||
otherwise falls back to a racy SELECT-then-INSERT.
|
||||
"""
|
||||
return await user_table.upsert(
|
||||
where={"user_id": user_id},
|
||||
data={"create": create_data, "update": {"user_id": user_id}},
|
||||
)
|
||||
|
||||
|
||||
async def _create_user_row(
|
||||
prisma_client: PrismaClient, tx: MemberWriteTx | None, user_data: dict[str, object]
|
||||
) -> _PrismaUserRecord | None:
|
||||
if tx is not None:
|
||||
return await _upsert_user_row(tx.litellm_usertable, str(user_data["user_id"]), jsonify_object(user_data))
|
||||
return await prisma_client.insert_data(data=user_data, table_name="user")
|
||||
|
||||
|
||||
def get_new_internal_user_defaults(user_id: str, user_email: str | None = None) -> dict[str, object]:
|
||||
user_info: Final = litellm.default_internal_user_params or {}
|
||||
|
||||
|
|
@ -206,6 +275,7 @@ async def _clone_team_default_budget_for_member(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
litellm_proxy_admin_name: str,
|
||||
budget_duration_override: str | None = None,
|
||||
tx: MemberWriteTx | None = None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Create a new budget row that copies the values from the team's default
|
||||
|
|
@ -220,7 +290,7 @@ async def _clone_team_default_budget_for_member(
|
|||
member while keeping the default's other limits, so an admin can set a
|
||||
member's reset cadence without discarding the team default's max_budget.
|
||||
"""
|
||||
budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table
|
||||
budget_table: Final[_PrismaBudgetTable] = _budget_table(prisma_client, tx)
|
||||
default_budget: Final = await budget_table.find_unique(where={"budget_id": default_team_budget_id})
|
||||
if default_budget is None:
|
||||
return None
|
||||
|
|
@ -248,7 +318,7 @@ async def _clone_team_default_budget_for_member(
|
|||
if cloned_data.get("budget_duration"):
|
||||
cloned_data["budget_reset_at"] = get_budget_reset_time(cloned_data["budget_duration"])
|
||||
|
||||
new_budget: Final[_PrismaBudgetRecord] = await BudgetRepository(prisma_client).table.create(data=cloned_data)
|
||||
new_budget: Final[_PrismaBudgetRecord] = await budget_table.create(data=cloned_data)
|
||||
return new_budget.budget_id
|
||||
|
||||
|
||||
|
|
@ -260,6 +330,7 @@ async def _resolve_member_budget_id(
|
|||
allowed_models: list[str] | None,
|
||||
budget_duration: str | None,
|
||||
default_team_budget_id: str | None,
|
||||
tx: MemberWriteTx | None = None,
|
||||
) -> str | None:
|
||||
"""
|
||||
Resolve the budget a new team member should be linked to.
|
||||
|
|
@ -279,6 +350,7 @@ async def _resolve_member_budget_id(
|
|||
user_api_key_dict=user_api_key_dict,
|
||||
litellm_proxy_admin_name=litellm_proxy_admin_name,
|
||||
budget_duration_override=budget_duration,
|
||||
tx=tx,
|
||||
)
|
||||
|
||||
if not has_explicit_limit and budget_duration is None:
|
||||
|
|
@ -295,12 +367,14 @@ async def _resolve_member_budget_id(
|
|||
if budget_duration is not None:
|
||||
budget_data["budget_duration"] = budget_duration
|
||||
budget_data["budget_reset_at"] = get_budget_reset_time(budget_duration=budget_duration)
|
||||
budget_table: Final[_PrismaBudgetTable] = BudgetRepository(prisma_client).table
|
||||
budget_table: Final[_PrismaBudgetTable] = _budget_table(prisma_client, tx)
|
||||
response: Final = await budget_table.create(data=budget_data)
|
||||
return response.budget_id
|
||||
|
||||
|
||||
async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, team_id: str) -> None:
|
||||
async def _append_team_id_if_absent(
|
||||
prisma_client: PrismaClient, user_id: str, team_id: str, tx: MemberWriteTx | None = None
|
||||
) -> None:
|
||||
"""Append team_id to a user's teams array, only if it is not already present.
|
||||
|
||||
The row-level filter makes the append a no-op once the team is present, so
|
||||
|
|
@ -309,7 +383,7 @@ async def _append_team_id_if_absent(prisma_client: PrismaClient, user_id: str, t
|
|||
number of teams a user belongs to). Teams added concurrently for a different
|
||||
team id are unaffected, since each update filters on its own team id.
|
||||
"""
|
||||
user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
|
||||
user_table: Final[_PrismaUserTable] = _user_table(prisma_client, tx)
|
||||
await user_table.update_many(
|
||||
where={"user_id": user_id, "NOT": {"teams": {"has": team_id}}},
|
||||
data={"teams": {"push": [team_id]}},
|
||||
|
|
@ -326,6 +400,7 @@ async def add_new_member(
|
|||
default_team_budget_id: str | None = None,
|
||||
allowed_models: list[str] | None = None,
|
||||
budget_duration: str | None = None,
|
||||
tx: MemberWriteTx | None = None,
|
||||
) -> tuple[LiteLLM_UserTable, LiteLLM_TeamMembership | None]:
|
||||
"""
|
||||
Add a new member to a team
|
||||
|
|
@ -334,49 +409,41 @@ async def add_new_member(
|
|||
- add team member w/ budget to team member table
|
||||
|
||||
Returns created/existing user + team membership w/ budget id
|
||||
|
||||
Callers already inside a transaction pass it as ``tx`` so every write here runs on that
|
||||
connection instead of borrowing more from the pool while the caller's locks are held.
|
||||
"""
|
||||
returned_user: LiteLLM_UserTable | None = None
|
||||
returned_team_membership: LiteLLM_TeamMembership | None = None
|
||||
## ADD TEAM ID, to USER TABLE IF NEW ##
|
||||
if new_member.user_id is not None:
|
||||
new_user_defaults = get_new_internal_user_defaults(user_id=new_member.user_id)
|
||||
# Upsert ensures the user row exists atomically (no create race when the
|
||||
# same new user is provisioned concurrently), seeding teams on create.
|
||||
# The teams append lives in the filtered update below rather than the
|
||||
# upsert's update branch so an already-existing user does not get a
|
||||
# duplicate team id. The update branch still has to write something:
|
||||
# Prisma only compiles an upsert down to INSERT ... ON CONFLICT when it
|
||||
# is non-empty, and falls back to a racy SELECT-then-INSERT when it is
|
||||
# not, so this re-states user_id as a no-op rather than being empty.
|
||||
user_table: Final[_PrismaUserTable] = UserRepository(prisma_client).table
|
||||
_returned_user: _PrismaUserRecord | None = await user_table.upsert(
|
||||
where={"user_id": new_member.user_id},
|
||||
data={
|
||||
"create": {"teams": [team_id], **new_user_defaults},
|
||||
"update": {"user_id": new_member.user_id},
|
||||
},
|
||||
# The teams append lives in the filtered update below rather than the upsert's
|
||||
# update branch so an already-existing user does not get a duplicate team id.
|
||||
_returned_user: _PrismaUserRecord | None = await _upsert_user_row(
|
||||
_user_table(prisma_client, tx),
|
||||
new_member.user_id,
|
||||
{"teams": [team_id], **new_user_defaults},
|
||||
)
|
||||
await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id)
|
||||
await _append_team_id_if_absent(prisma_client, new_member.user_id, team_id, tx)
|
||||
if _returned_user is not None:
|
||||
returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
|
||||
elif new_member.user_email is not None:
|
||||
new_user_defaults = get_new_internal_user_defaults(user_id=str(uuid.uuid4()), user_email=new_member.user_email)
|
||||
## user email is not unique acc. to prisma schema -> future improvement
|
||||
### for now: check if it exists in db, if not - insert it
|
||||
existing_user_row: Final[list[_PrismaUserRecord] | None] = await prisma_client.get_data(
|
||||
key_val={"user_email": new_member.user_email},
|
||||
table_name="user",
|
||||
query_type="find_all",
|
||||
existing_user_row: Final[Sequence[_PrismaUserRecord]] = await _find_users_by_email(
|
||||
prisma_client, tx, new_member.user_email
|
||||
)
|
||||
if existing_user_row is None or (isinstance(existing_user_row, list) and len(existing_user_row) == 0):
|
||||
if len(existing_user_row) == 0:
|
||||
new_user_defaults["teams"] = [team_id]
|
||||
_returned_user = await prisma_client.insert_data(data=new_user_defaults, table_name="user")
|
||||
_returned_user = await _create_user_row(prisma_client, tx, new_user_defaults)
|
||||
|
||||
if _returned_user is not None:
|
||||
returned_user = LiteLLM_UserTable.model_validate(_returned_user.model_dump())
|
||||
elif len(existing_user_row) == 1:
|
||||
user_info: Final = existing_user_row[0]
|
||||
await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id)
|
||||
await _append_team_id_if_absent(prisma_client, user_info.user_id, team_id, tx)
|
||||
returned_user = LiteLLM_UserTable.model_validate(user_info.model_dump())
|
||||
elif len(existing_user_row) > 1:
|
||||
raise HTTPException(
|
||||
|
|
@ -392,10 +459,11 @@ async def add_new_member(
|
|||
allowed_models=allowed_models,
|
||||
budget_duration=budget_duration,
|
||||
default_team_budget_id=default_team_budget_id,
|
||||
tx=tx,
|
||||
)
|
||||
|
||||
if _budget_id and returned_user is not None and returned_user.user_id is not None:
|
||||
membership_table: Final[_PrismaTeamMembershipTable] = TeamMembershipRepository(prisma_client).table
|
||||
membership_table: Final[_PrismaTeamMembershipTable] = _team_membership_table(prisma_client, tx)
|
||||
_returned_team_membership: Final = await membership_table.create(
|
||||
data={
|
||||
"team_id": team_id,
|
||||
|
|
|
|||
|
|
@ -18,21 +18,20 @@ Scoping:
|
|||
"""
|
||||
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
from collections.abc import Mapping
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import (
|
||||
CommonProxyErrors,
|
||||
LiteLLM_TeamTable,
|
||||
LitellmUserRoles,
|
||||
UserAPIKeyAuth,
|
||||
user_api_key_has_admin_view,
|
||||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import MemoryRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.types.memory_management import (
|
||||
|
|
@ -44,54 +43,17 @@ from litellm.types.memory_management import (
|
|||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
||||
class _MemoryRecord(Protocol):
|
||||
memory_id: str
|
||||
key: str
|
||||
value: str
|
||||
metadata: object
|
||||
user_id: str | None
|
||||
team_id: str | None
|
||||
created_at: datetime | None
|
||||
created_by: str | None
|
||||
updated_at: datetime | None
|
||||
updated_by: str | None
|
||||
|
||||
|
||||
class _MemoryTableActions(Protocol):
|
||||
async def create(self, data: Mapping[str, object]) -> _MemoryRecord: ...
|
||||
|
||||
async def find_many(
|
||||
self,
|
||||
where: Mapping[str, object] | None = ...,
|
||||
order: Mapping[str, str] | None = ...,
|
||||
skip: int = ...,
|
||||
take: int = ...,
|
||||
) -> Sequence[_MemoryRecord]: ...
|
||||
|
||||
async def count(self, where: Mapping[str, object] | None = ...) -> int: ...
|
||||
|
||||
async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> _MemoryRecord: ...
|
||||
|
||||
async def delete(self, where: Mapping[str, object]) -> _MemoryRecord | None: ...
|
||||
|
||||
|
||||
def _memory_table(prisma_client: "PrismaClient") -> _MemoryTableActions:
|
||||
def _memory_table(prisma_client: "PrismaClient") -> TableActions["prisma_models.LiteLLM_MemoryTable"]:
|
||||
return MemoryRepository(prisma_client).table
|
||||
|
||||
|
||||
class _TeamTableActions(Protocol):
|
||||
async def find_unique(self, where: Mapping[str, str]) -> LiteLLM_TeamTable | None: ...
|
||||
|
||||
|
||||
def _team_table(prisma_client: "PrismaClient") -> _TeamTableActions:
|
||||
return TeamRepository(prisma_client).table
|
||||
|
||||
|
||||
def _serialize_metadata_for_prisma(metadata: object) -> str:
|
||||
"""
|
||||
Encode a `metadata` payload for the `Json?` column.
|
||||
|
|
@ -129,7 +91,7 @@ def _visibility_filter(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object
|
|||
return {"OR": ors}
|
||||
|
||||
|
||||
def _row_to_model(row: _MemoryRecord) -> LiteLLM_MemoryRow:
|
||||
def _row_to_model(row: "prisma_models.LiteLLM_MemoryTable") -> LiteLLM_MemoryRow:
|
||||
return LiteLLM_MemoryRow(
|
||||
memory_id=row.memory_id,
|
||||
key=row.key,
|
||||
|
|
@ -163,7 +125,7 @@ def _internal_error(log_message: str, exc: Exception, default_detail: str) -> HT
|
|||
|
||||
|
||||
async def _assert_write_access(
|
||||
prisma_client: "PrismaClient", row: _MemoryRecord, user_api_key_dict: UserAPIKeyAuth
|
||||
prisma_client: "PrismaClient", row: "prisma_models.LiteLLM_MemoryTable", user_api_key_dict: UserAPIKeyAuth
|
||||
) -> None:
|
||||
"""
|
||||
Enforce ownership for mutations (PUT/DELETE).
|
||||
|
|
@ -219,7 +181,7 @@ async def _is_team_admin_for(prisma_client: "PrismaClient", user_api_key_dict: U
|
|||
)
|
||||
|
||||
try:
|
||||
team_obj: Final = await _team_table(prisma_client).find_unique(where={"team_id": team_id})
|
||||
team_obj: Final = await TeamRepository(prisma_client).find_by_id(team_id, id_field="team_id")
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error loading team for write-auth check (team_id=%s): %s", team_id, e)
|
||||
return False
|
||||
|
|
@ -407,7 +369,7 @@ async def list_memory(
|
|||
|
||||
async def _find_memory_for_caller(
|
||||
prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> _MemoryRecord:
|
||||
) -> "prisma_models.LiteLLM_MemoryTable":
|
||||
"""Look up a memory row by key, scoped to the caller's visibility."""
|
||||
key_filter: Final[Mapping[str, object]] = {"key": key}
|
||||
vis: Final = _visibility_filter(user_api_key_dict)
|
||||
|
|
@ -418,6 +380,18 @@ async def _find_memory_for_caller(
|
|||
return rows[0]
|
||||
|
||||
|
||||
async def _find_visible_memory_or_none(
|
||||
prisma_client: "PrismaClient", key: str, user_api_key_dict: UserAPIKeyAuth
|
||||
) -> "prisma_models.LiteLLM_MemoryTable | None":
|
||||
"""The caller-visible row for `key`, or None when nothing is visible to them."""
|
||||
try:
|
||||
return await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
|
||||
except HTTPException as e:
|
||||
if e.status_code == 404:
|
||||
return None
|
||||
raise
|
||||
|
||||
|
||||
@router.get(
|
||||
"/v1/memory/{key:path}",
|
||||
tags=["memory management"],
|
||||
|
|
@ -480,17 +454,8 @@ async def upsert_memory(
|
|||
)
|
||||
data["updated_by"] = user_api_key_dict.user_id
|
||||
|
||||
async def _find_existing() -> _MemoryRecord | None:
|
||||
"""Return the caller-visible row for `key`, or None."""
|
||||
try:
|
||||
return await _find_memory_for_caller(prisma_client, key, user_api_key_dict)
|
||||
except HTTPException as e:
|
||||
if e.status_code == 404:
|
||||
return None
|
||||
raise
|
||||
|
||||
try:
|
||||
existing: Final = await _find_existing()
|
||||
existing: Final = await _find_visible_memory_or_none(prisma_client, key, user_api_key_dict)
|
||||
if existing is not None:
|
||||
# Visibility != write authority. Make sure the caller actually
|
||||
# owns this row (their user_id matches, or it's a pure team row in
|
||||
|
|
@ -530,7 +495,7 @@ async def upsert_memory(
|
|||
# instead of surfacing a 500 on a unique-violation.
|
||||
if not _is_unique_violation(e):
|
||||
raise
|
||||
existing_after_race: Final = await _find_existing()
|
||||
existing_after_race: Final = await _find_visible_memory_or_none(prisma_client, key, user_api_key_dict)
|
||||
if existing_after_race is None:
|
||||
# Row exists globally but isn't visible to this caller
|
||||
# (owned by someone else). Treat as conflict.
|
||||
|
|
@ -549,6 +514,8 @@ async def upsert_memory(
|
|||
except Exception as e:
|
||||
raise _internal_error("Error upserting memory: %s", e, "Internal error updating memory entry.")
|
||||
|
||||
if row is None:
|
||||
raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found")
|
||||
return _row_to_model(row)
|
||||
|
||||
|
||||
|
|
@ -568,8 +535,10 @@ async def delete_memory(
|
|||
# Visibility != write authority — see the upsert handler for the rationale.
|
||||
await _assert_write_access(prisma_client, row, user_api_key_dict)
|
||||
try:
|
||||
await _memory_table(prisma_client).delete(where={"memory_id": row.memory_id})
|
||||
deleted: Final = await _memory_table(prisma_client).delete(where={"memory_id": row.memory_id})
|
||||
except Exception as e:
|
||||
raise _internal_error("Error deleting memory: %s", e, "Internal error deleting memory entry.")
|
||||
|
||||
if deleted is None:
|
||||
raise HTTPException(status_code=404, detail=f"Memory with key '{key}' not found")
|
||||
return MemoryDeleteResponse(key=key, deleted=True)
|
||||
|
|
|
|||
|
|
@ -4,7 +4,16 @@ import re
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass, field
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, get_args, runtime_checkable
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
Literal,
|
||||
Optional,
|
||||
Protocol,
|
||||
cast, # noqa: TID251 # prisma types Json columns as fields.Json but de-serializes them to plain python on read
|
||||
get_args,
|
||||
runtime_checkable,
|
||||
)
|
||||
|
||||
from litellm.proxy._types import ProxyException
|
||||
from litellm.repositories.table_repositories import (
|
||||
|
|
@ -1183,7 +1192,7 @@ async def ensure_batch_response_managed_file_ids(
|
|||
prisma_client,
|
||||
verbose_proxy_logger,
|
||||
user_api_key_dict=None,
|
||||
db_batch_object=None,
|
||||
db_batch_object: "LiteLLM_ManagedObjectTable | None" = None,
|
||||
unified_batch_id: str | Literal[False] | None = None,
|
||||
) -> None:
|
||||
"""Normalize batch file IDs to managed unified IDs before DB persistence."""
|
||||
|
|
@ -1270,11 +1279,10 @@ async def get_batch_from_database(
|
|||
return None, None
|
||||
|
||||
# Parse the batch object from database
|
||||
batch_data: Final = (
|
||||
json.loads(db_batch_object.file_object)
|
||||
if isinstance(db_batch_object.file_object, str)
|
||||
else db_batch_object.file_object
|
||||
file_object: Final = cast( # cast-ok: prisma types the Json column as str; reads return the decoded value
|
||||
"Mapping[str, object] | str", db_batch_object.file_object
|
||||
)
|
||||
batch_data: Final = json.loads(file_object) if isinstance(file_object, str) else file_object
|
||||
response: Final = LiteLLMBatch.model_validate(batch_data)
|
||||
response.id = batch_id
|
||||
|
||||
|
|
@ -1360,7 +1368,7 @@ async def update_batch_in_database(
|
|||
managed_files_obj,
|
||||
prisma_client,
|
||||
verbose_proxy_logger,
|
||||
db_batch_object=None,
|
||||
db_batch_object: "LiteLLM_ManagedObjectTable | None" = None,
|
||||
operation: str = "update",
|
||||
user_api_key_dict=None,
|
||||
poller_owns_accounting: bool | None = None,
|
||||
|
|
@ -1427,7 +1435,7 @@ async def update_batch_in_database(
|
|||
# Normalize status for database storage
|
||||
db_status: Final = response.status if response.status != "completed" else "complete"
|
||||
|
||||
update_data: Final[dict] = {
|
||||
update_data: Final[dict[str, object]] = {
|
||||
"status": db_status,
|
||||
"file_object": response.model_dump_json(),
|
||||
"updated_at": litellm.utils.get_utc_datetime(),
|
||||
|
|
|
|||
|
|
@ -32,12 +32,18 @@ from __future__ import annotations
|
|||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final, TypeVar, overload
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping, Sequence
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
TypeVar,
|
||||
cast, # noqa: TID251 # prisma stubs type Json columns as fields.Json but de-serialize them on read
|
||||
overload,
|
||||
)
|
||||
from urllib.parse import quote, unquote
|
||||
|
||||
from fastapi import HTTPException
|
||||
from pydantic import JsonValue
|
||||
from pydantic import JsonValue, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.base_llm.managed_resources.isolation import (
|
||||
|
|
@ -46,6 +52,7 @@ from litellm.llms.base_llm.managed_resources.isolation import (
|
|||
)
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.batches_endpoints.common_utils import validate_batch_list_limit
|
||||
from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
|
||||
from litellm.repositories.table_repositories import (
|
||||
ManagedFileRepository,
|
||||
ManagedObjectRepository,
|
||||
|
|
@ -286,11 +293,15 @@ def _canonical_path(route: str) -> str:
|
|||
|
||||
|
||||
def _file_table(prisma_client: PrismaClient) -> ManagedFileTable:
|
||||
return ManagedFileRepository(prisma_client).table
|
||||
return cast( # cast-ok: stub-only mismatch, prisma returns real lists and de-serialized Json
|
||||
ManagedFileTable, ManagedFileRepository(prisma_client).table
|
||||
)
|
||||
|
||||
|
||||
def _object_table(prisma_client: PrismaClient) -> ManagedObjectTable:
|
||||
return ManagedObjectRepository(prisma_client).table
|
||||
return cast( # cast-ok: stub-only mismatch, prisma returns real lists and de-serialized Json
|
||||
ManagedObjectTable, ManagedObjectRepository(prisma_client).table
|
||||
)
|
||||
|
||||
|
||||
async def _resolve_one(
|
||||
|
|
@ -810,6 +821,121 @@ async def rewrite_response_ids(
|
|||
return mutated if changed else body
|
||||
|
||||
|
||||
_RESPONSE_ID_PREFIX: Final = "resp_"
|
||||
_STREAMED_RESPONSE_ID_SPEC: Final[_FieldSpec] = ("id", _RESPONSE_ID_PREFIX)
|
||||
_SSE_DATA_PREFIX: Final = "data:"
|
||||
_SSE_EVENT_ADAPTER: Final = TypeAdapter(Mapping[str, JsonValue])
|
||||
|
||||
|
||||
def _first_streamed_response(frames: bytes) -> tuple[str, Mapping[str, JsonValue]] | None:
|
||||
for line in frames.decode("utf-8", errors="replace").splitlines():
|
||||
if not line.startswith(_SSE_DATA_PREFIX):
|
||||
continue
|
||||
try:
|
||||
event = _SSE_EVENT_ADAPTER.validate_json(line[len(_SSE_DATA_PREFIX) :])
|
||||
except ValidationError:
|
||||
continue
|
||||
response = event.get("response")
|
||||
if not isinstance(response, dict):
|
||||
continue
|
||||
raw_id = response.get("id")
|
||||
if isinstance(raw_id, str) and raw_id.startswith(_RESPONSE_ID_PREFIX):
|
||||
return raw_id, response
|
||||
return None
|
||||
|
||||
|
||||
class _StreamedResponseIdRewriter:
|
||||
__slots__ = ("_is_create_route", "_pending", "_prisma_client", "_provider", "_replacement", "_user_api_key_dict")
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
provider: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
is_create_route: bool,
|
||||
) -> None:
|
||||
self._provider: Final = provider
|
||||
self._user_api_key_dict: Final = user_api_key_dict
|
||||
self._prisma_client: Final = prisma_client
|
||||
self._is_create_route: Final = is_create_route
|
||||
self._pending = b""
|
||||
self._replacement: tuple[bytes, bytes] | None = None
|
||||
|
||||
async def feed(self, chunk: bytes) -> bytes:
|
||||
complete_frames, self._pending = split_complete_sse_frames(self._pending + chunk)
|
||||
if not complete_frames:
|
||||
return b""
|
||||
if self._replacement is None:
|
||||
self._replacement = await self._mint(complete_frames)
|
||||
return self._rewrite(complete_frames)
|
||||
|
||||
def flush(self) -> bytes:
|
||||
tail: Final = self._pending
|
||||
self._pending = b""
|
||||
return self._rewrite(tail)
|
||||
|
||||
async def _mint(self, frames: bytes) -> tuple[bytes, bytes] | None:
|
||||
first: Final = _first_streamed_response(frames)
|
||||
if first is None:
|
||||
return None
|
||||
raw_id, snapshot = first
|
||||
managed_id: Final = await _mint_or_reuse_object(
|
||||
raw_id,
|
||||
self._provider,
|
||||
"response",
|
||||
snapshot,
|
||||
self._user_api_key_dict,
|
||||
self._prisma_client,
|
||||
self._is_create_route,
|
||||
)
|
||||
return raw_id.encode(), managed_id.encode()
|
||||
|
||||
def _rewrite(self, frames: bytes) -> bytes:
|
||||
if self._replacement is None:
|
||||
return frames
|
||||
raw_id, managed_id = self._replacement
|
||||
return frames.replace(raw_id, managed_id)
|
||||
|
||||
|
||||
async def rewrite_streamed_response_ids(
|
||||
stream: AsyncGenerator[bytes, None],
|
||||
provider: str,
|
||||
method: str,
|
||||
route: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
) -> AsyncGenerator[bytes, None]:
|
||||
"""
|
||||
Record ownership of the response object streamed back by a Responses API
|
||||
passthrough and swap its managed id into every SSE frame, so a streamed
|
||||
response is owned and resolved exactly like a non-streamed one.
|
||||
|
||||
Streams for any other ``(provider, method, route)`` are relayed untouched.
|
||||
"""
|
||||
from litellm.proxy.auth.auth_utils import normalize_request_route
|
||||
|
||||
canonical: Final = normalize_request_route(_canonical_path(route))
|
||||
field_specs: Final = BUILTIN_OUTPUT_ID_FIELD_MAP.get((provider, method, canonical), ())
|
||||
if _STREAMED_RESPONSE_ID_SPEC not in field_specs:
|
||||
async for chunk in stream:
|
||||
yield chunk
|
||||
return
|
||||
|
||||
rewriter: Final = _StreamedResponseIdRewriter(
|
||||
provider=provider,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
is_create_route="{" not in canonical,
|
||||
)
|
||||
async for chunk in stream:
|
||||
rewritten_frames = await rewriter.feed(chunk)
|
||||
if rewritten_frames:
|
||||
yield rewritten_frames
|
||||
tail: Final = rewriter.flush()
|
||||
if tail:
|
||||
yield tail
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# List-route interception — serve listing entirely from DB
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import json
|
|||
import posixpath
|
||||
import traceback
|
||||
from base64 import b64encode
|
||||
from collections.abc import AsyncGenerator, Callable, Iterable, Mapping
|
||||
from collections.abc import AsyncGenerator, Callable, Iterable, Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from itertools import groupby
|
||||
from typing import Any, Final, TypedDict, cast
|
||||
|
|
@ -1209,14 +1209,19 @@ async def pass_through_request(
|
|||
|
||||
return StreamingResponse(
|
||||
wrap_passthrough_sse_bytes_with_keepalive_pings(
|
||||
stream=PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
request_body=_parsed_body,
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=endpoint_type,
|
||||
start_time=start_time,
|
||||
passthrough_success_handler_obj=pass_through_endpoint_logging,
|
||||
url_route=str(url),
|
||||
stream=_own_streamed_managed_ids(
|
||||
stream=PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
request_body=_parsed_body,
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=endpoint_type,
|
||||
start_time=start_time,
|
||||
passthrough_success_handler_obj=pass_through_endpoint_logging,
|
||||
url_route=str(url),
|
||||
),
|
||||
managed_id_provider=_managed_id_provider,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
),
|
||||
ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
|
||||
upstream_headers=response.headers,
|
||||
|
|
@ -1285,14 +1290,19 @@ async def pass_through_request(
|
|||
|
||||
return StreamingResponse(
|
||||
wrap_passthrough_sse_bytes_with_keepalive_pings(
|
||||
stream=PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
request_body=_parsed_body,
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=endpoint_type,
|
||||
start_time=start_time,
|
||||
passthrough_success_handler_obj=pass_through_endpoint_logging,
|
||||
url_route=str(url),
|
||||
stream=_own_streamed_managed_ids(
|
||||
stream=PassThroughStreamingHandler.chunk_processor(
|
||||
response=response,
|
||||
request_body=_parsed_body,
|
||||
litellm_logging_obj=logging_obj,
|
||||
endpoint_type=endpoint_type,
|
||||
start_time=start_time,
|
||||
passthrough_success_handler_obj=pass_through_endpoint_logging,
|
||||
url_route=str(url),
|
||||
),
|
||||
managed_id_provider=_managed_id_provider,
|
||||
request=request,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
),
|
||||
ping_interval_seconds=litellm.sse_keepalive_ping_interval_seconds,
|
||||
upstream_headers=response.headers,
|
||||
|
|
@ -2441,6 +2451,36 @@ def _is_streaming_response(response: httpx.Response) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _own_streamed_managed_ids(
|
||||
stream: AsyncGenerator[bytes, None],
|
||||
managed_id_provider: str | None,
|
||||
request: Request,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
) -> AsyncGenerator[bytes, None]:
|
||||
from litellm.proxy.proxy_server import general_settings, prisma_client, proxy_logging_obj
|
||||
|
||||
if (
|
||||
managed_id_provider is None
|
||||
or not general_settings.get("passthrough_managed_object_ids", False)
|
||||
or prisma_client is None
|
||||
or proxy_logging_obj.get_proxy_hook("managed_files") is None
|
||||
):
|
||||
return stream
|
||||
from litellm.proxy.auth.auth_utils import get_request_route
|
||||
from litellm.proxy.pass_through_endpoints.managed_id_rewriter import (
|
||||
rewrite_streamed_response_ids,
|
||||
)
|
||||
|
||||
return rewrite_streamed_response_ids(
|
||||
stream=stream,
|
||||
provider=managed_id_provider,
|
||||
method=request.method,
|
||||
route=get_request_route(request),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
|
||||
|
||||
def _should_buffer_passthrough_response(response: httpx.Response) -> bool:
|
||||
"""
|
||||
Decide from the response headers whether the body must be read into memory.
|
||||
|
|
@ -3183,13 +3223,18 @@ async def _filter_endpoints_by_team_allowed_routes(
|
|||
)
|
||||
|
||||
# retrieve team metadata
|
||||
team_metadata: Final = team.metadata
|
||||
team_metadata: Final = cast( # cast-ok: prisma types the Json column as str; reads hand back the decoded value
|
||||
"Mapping[str, object] | None", team.metadata
|
||||
)
|
||||
if team_metadata is not None and team_metadata.get("allowed_passthrough_routes") is not None:
|
||||
## FILTER pass_through_endpoints by allowed_passthrough_routes
|
||||
pass_through_endpoints = [
|
||||
endpoint
|
||||
for endpoint in pass_through_endpoints
|
||||
if endpoint.path in team_metadata.get("allowed_passthrough_routes")
|
||||
if endpoint.path
|
||||
in cast( # cast-ok: guarded above; team metadata stores this key as a list of route paths
|
||||
"Sequence[str]", team_metadata.get("allowed_passthrough_routes")
|
||||
)
|
||||
]
|
||||
|
||||
return pass_through_endpoints
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
|||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.proxy._types import PassThroughEndpointLoggingResultValues
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.common_utils.sse_keepalive import split_complete_sse_frames
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
|
||||
from litellm.types.utils import StandardPassThroughResponseObject
|
||||
|
||||
|
|
@ -101,7 +102,7 @@ class PassThroughStreamingHandler:
|
|||
async for chunk in response.aiter_bytes():
|
||||
raw_bytes.append(chunk)
|
||||
PassThroughStreamingHandler._stamp_first_chunk_if_needed(litellm_logging_obj)
|
||||
complete_frames, pending = PassThroughStreamingHandler._split_complete_sse_frames(
|
||||
complete_frames, pending = split_complete_sse_frames(
|
||||
pending + chunk
|
||||
) # rebind-ok: SSE frame reassembly buffer across transport chunks
|
||||
if complete_frames:
|
||||
|
|
@ -139,17 +140,6 @@ class PassThroughStreamingHandler:
|
|||
except Exception as e:
|
||||
verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e)
|
||||
|
||||
@staticmethod
|
||||
def _split_complete_sse_frames(pending: bytes) -> tuple[bytes, bytes]:
|
||||
lf_boundary_end: Final = pending.rfind(b"\n\n") + 2
|
||||
crlf_boundary_end: Final = pending.rfind(b"\r\n\r\n") + 4
|
||||
boundary_end: Final = max(
|
||||
lf_boundary_end if lf_boundary_end >= 2 else 0, crlf_boundary_end if crlf_boundary_end >= 4 else 0
|
||||
)
|
||||
if boundary_end == 0:
|
||||
return b"", pending
|
||||
return pending[:boundary_end], pending[boundary_end:]
|
||||
|
||||
@staticmethod
|
||||
async def _route_streaming_logging_to_handler(
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
|
|
|
|||
|
|
@ -30,7 +30,7 @@ class PolicyMatcher:
|
|||
"""
|
||||
Check if a value matches any of the given patterns.
|
||||
|
||||
Uses the existing RouteChecks._route_matches_wildcard_pattern helper.
|
||||
Uses the existing RouteChecks.route_matches_wildcard_pattern helper.
|
||||
|
||||
Args:
|
||||
value: The value to check (e.g., team alias, key alias, model)
|
||||
|
|
@ -45,7 +45,7 @@ class PolicyMatcher:
|
|||
|
||||
for pattern in patterns:
|
||||
# Use existing wildcard pattern matching helper
|
||||
if RouteChecks._route_matches_wildcard_pattern(route=value, pattern=pattern):
|
||||
if RouteChecks.route_matches_wildcard_pattern(route=value, pattern=pattern):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -10,9 +10,20 @@ by policy_attachments (see AttachmentRegistry).
|
|||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timezone
|
||||
from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypedDict, Union
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Any,
|
||||
Final,
|
||||
Literal,
|
||||
Optional,
|
||||
Protocol,
|
||||
TypedDict,
|
||||
Union,
|
||||
cast, # noqa: TID251 # prisma types the condition/pipeline Json columns as str, but reads return decoded values
|
||||
)
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import PolicyRepository
|
||||
from litellm.types.proxy.policy_engine import (
|
||||
GuardrailPipeline,
|
||||
|
|
@ -65,15 +76,32 @@ class _PolicyRow(Protocol):
|
|||
|
||||
|
||||
class _PolicyVersionSourceRow(Protocol):
|
||||
policy_id: str
|
||||
policy_name: str
|
||||
version_number: int
|
||||
inherit: str | None
|
||||
description: str | None
|
||||
guardrails_add: Sequence[str] | None
|
||||
guardrails_remove: Sequence[str] | None
|
||||
condition: Mapping[str, object] | str | None
|
||||
pipeline: Mapping[str, object] | str | None
|
||||
@property
|
||||
def policy_id(self) -> str: ...
|
||||
|
||||
@property
|
||||
def policy_name(self) -> str: ...
|
||||
|
||||
@property
|
||||
def version_number(self) -> int: ...
|
||||
|
||||
@property
|
||||
def inherit(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def description(self) -> str | None: ...
|
||||
|
||||
@property
|
||||
def guardrails_add(self) -> Sequence[str] | None: ...
|
||||
|
||||
@property
|
||||
def guardrails_remove(self) -> Sequence[str] | None: ...
|
||||
|
||||
@property
|
||||
def condition(self) -> Mapping[str, object] | str | None: ...
|
||||
|
||||
@property
|
||||
def pipeline(self) -> Mapping[str, object] | str | None: ...
|
||||
|
||||
|
||||
class _PolicyTableClient(Protocol):
|
||||
|
|
@ -96,23 +124,15 @@ class _PolicyTableClient(Protocol):
|
|||
async def delete_many(self, where: Mapping[str, object]) -> int: ...
|
||||
|
||||
|
||||
class _PolicyVersionSourceTableClient(Protocol):
|
||||
async def find_unique(self, where: Mapping[str, object]) -> _PolicyVersionSourceRow | None: ...
|
||||
|
||||
async def find_first(
|
||||
self,
|
||||
where: Mapping[str, object],
|
||||
order: Mapping[str, str] | None = None,
|
||||
) -> _PolicyVersionSourceRow | None: ...
|
||||
|
||||
|
||||
def _policy_table(prisma_client: "PrismaClient") -> _PolicyTableClient:
|
||||
table: Final[_PolicyTableClient] = PolicyRepository(prisma_client).table
|
||||
return table
|
||||
table: Final = PolicyRepository(prisma_client).table
|
||||
return cast( # cast-ok: prisma types Json columns as str; the client hands back the decoded condition/pipeline
|
||||
"_PolicyTableClient", table
|
||||
)
|
||||
|
||||
|
||||
def _policy_version_source_table(prisma_client: "PrismaClient") -> _PolicyVersionSourceTableClient:
|
||||
table: Final[_PolicyVersionSourceTableClient] = PolicyRepository(prisma_client).table
|
||||
def _policy_version_source_table(prisma_client: "PrismaClient") -> "TableActions[_PolicyVersionSourceRow]":
|
||||
table: Final[TableActions[_PolicyVersionSourceRow]] = PolicyRepository(prisma_client).table
|
||||
return table
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -6,7 +6,8 @@ Policy resolve and attachment impact estimation endpoints.
|
|||
"""
|
||||
|
||||
import json
|
||||
from typing import Final
|
||||
from collections.abc import Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
|
||||
|
|
@ -30,25 +31,28 @@ from litellm.types.proxy.policy_engine import (
|
|||
PolicyResolveResponse,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
||||
def _build_alias_where(field: str, patterns: list) -> dict:
|
||||
def _build_alias_where(field: str, patterns: Sequence[str]) -> dict[str, object]:
|
||||
"""Build a Prisma ``where`` clause for alias patterns.
|
||||
|
||||
Supports exact matches and suffix wildcards (``prefix*``).
|
||||
Returns something like:
|
||||
{"OR": [{"field": {"in": ["a","b"]}}, {"field": {"startsWith": "dev-"}}]}
|
||||
"""
|
||||
exact: Final[list] = []
|
||||
prefix_conditions: Final[list] = []
|
||||
exact: Final[list[str]] = []
|
||||
prefix_conditions: Final[list[dict[str, object]]] = []
|
||||
for pat in patterns:
|
||||
if pat.endswith("*"):
|
||||
prefix_conditions.append({field: {"startsWith": pat[:-1]}})
|
||||
else:
|
||||
exact.append(pat)
|
||||
|
||||
conditions: Final[list] = []
|
||||
conditions: Final[list[dict[str, object]]] = []
|
||||
if exact:
|
||||
conditions.append({field: {"in": exact}})
|
||||
conditions.extend(prefix_conditions)
|
||||
|
|
@ -79,7 +83,7 @@ def _get_tags_from_metadata(metadata: object, json_metadata: object = None) -> l
|
|||
return parsed.get("tags", []) or []
|
||||
|
||||
|
||||
async def _fetch_all_teams(prisma_client: object) -> list:
|
||||
async def _fetch_all_teams(prisma_client: object) -> "Sequence[prisma_models.LiteLLM_TeamTable]":
|
||||
"""Fetch teams from DB once. Reuse the result across tag and alias lookups."""
|
||||
return await TeamRepository(prisma_client).table.find_many(
|
||||
where={},
|
||||
|
|
@ -88,19 +92,21 @@ async def _fetch_all_teams(prisma_client: object) -> list:
|
|||
)
|
||||
|
||||
|
||||
def _filter_keys_by_tags(keys: list, tag_patterns: list) -> tuple:
|
||||
def _filter_keys_by_tags(
|
||||
keys: "Sequence[prisma_models.LiteLLM_VerificationToken]", tag_patterns: Sequence[str]
|
||||
) -> tuple[list[str], int]:
|
||||
"""Filter key rows whose metadata.tags match any of the given patterns.
|
||||
|
||||
Returns (named_aliases, unnamed_count).
|
||||
"""
|
||||
|
||||
affected: Final[list] = []
|
||||
affected: Final[list[str]] = []
|
||||
unnamed_count = 0
|
||||
for key in keys:
|
||||
key_alias = key.key_alias or ""
|
||||
key_tags = _get_tags_from_metadata(key.metadata, getattr(key, "metadata_json", None))
|
||||
if key_tags and any(
|
||||
RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat)
|
||||
RouteChecks.route_matches_wildcard_pattern(route=tag, pattern=pat)
|
||||
for tag in key_tags
|
||||
for pat in tag_patterns
|
||||
):
|
||||
|
|
@ -111,19 +117,21 @@ def _filter_keys_by_tags(keys: list, tag_patterns: list) -> tuple:
|
|||
return affected, unnamed_count
|
||||
|
||||
|
||||
def _filter_teams_by_tags(teams: list, tag_patterns: list) -> tuple:
|
||||
def _filter_teams_by_tags(
|
||||
teams: "Sequence[prisma_models.LiteLLM_TeamTable]", tag_patterns: Sequence[str]
|
||||
) -> tuple[list[str], int]:
|
||||
"""Filter pre-fetched team rows whose metadata.tags match any patterns.
|
||||
|
||||
Returns (named_aliases, unnamed_count).
|
||||
"""
|
||||
|
||||
affected: Final[list] = []
|
||||
affected: Final[list[str]] = []
|
||||
unnamed_count = 0
|
||||
for team in teams:
|
||||
team_alias = team.team_alias or ""
|
||||
team_tags = _get_tags_from_metadata(team.metadata)
|
||||
if team_tags and any(
|
||||
RouteChecks._route_matches_wildcard_pattern(route=tag, pattern=pat)
|
||||
RouteChecks.route_matches_wildcard_pattern(route=tag, pattern=pat)
|
||||
for tag in team_tags
|
||||
for pat in tag_patterns
|
||||
):
|
||||
|
|
@ -136,29 +144,29 @@ def _filter_teams_by_tags(teams: list, tag_patterns: list) -> tuple:
|
|||
|
||||
async def _find_affected_by_team_patterns(
|
||||
prisma_client: object,
|
||||
all_teams: list,
|
||||
team_patterns: list,
|
||||
existing_teams: list,
|
||||
existing_keys: list,
|
||||
) -> tuple:
|
||||
all_teams: "Sequence[prisma_models.LiteLLM_TeamTable]",
|
||||
team_patterns: Sequence[str],
|
||||
existing_teams: Sequence[str],
|
||||
existing_keys: Sequence[str],
|
||||
) -> tuple[list[str], list[str], int]:
|
||||
"""Filter pre-fetched teams by alias patterns, then fetch their keys.
|
||||
|
||||
Returns (new_teams, new_keys, unnamed_keys_count).
|
||||
"""
|
||||
|
||||
new_teams: Final[list] = []
|
||||
matched_team_ids: Final[list] = []
|
||||
new_teams: Final[list[str]] = []
|
||||
matched_team_ids: Final[list[str]] = []
|
||||
|
||||
for team in all_teams:
|
||||
team_alias = team.team_alias or ""
|
||||
if team_alias and any(
|
||||
RouteChecks._route_matches_wildcard_pattern(route=team_alias, pattern=pat) for pat in team_patterns
|
||||
RouteChecks.route_matches_wildcard_pattern(route=team_alias, pattern=pat) for pat in team_patterns
|
||||
):
|
||||
if team_alias not in existing_teams:
|
||||
new_teams.append(team_alias)
|
||||
matched_team_ids.append(str(team.team_id))
|
||||
|
||||
new_keys: Final[list] = []
|
||||
new_keys: Final[list[str]] = []
|
||||
unnamed_keys_count = 0
|
||||
if matched_team_ids:
|
||||
keys: Final = await VerificationTokenRepository(prisma_client).table.find_many(
|
||||
|
|
@ -177,10 +185,12 @@ async def _find_affected_by_team_patterns(
|
|||
return new_teams, new_keys, unnamed_keys_count
|
||||
|
||||
|
||||
async def _find_affected_keys_by_alias(prisma_client: object, key_patterns: list, existing_keys: list) -> list:
|
||||
async def _find_affected_keys_by_alias(
|
||||
prisma_client: object, key_patterns: Sequence[str], existing_keys: Sequence[str]
|
||||
) -> list[str]:
|
||||
"""Find keys whose alias matches the given patterns."""
|
||||
|
||||
affected: Final[list] = []
|
||||
affected: Final[list[str]] = []
|
||||
|
||||
keys: Final = await VerificationTokenRepository(prisma_client).table.find_many(
|
||||
where=_build_alias_where("key_alias", key_patterns),
|
||||
|
|
@ -190,7 +200,7 @@ async def _find_affected_keys_by_alias(prisma_client: object, key_patterns: list
|
|||
for key in keys:
|
||||
key_alias = key.key_alias or ""
|
||||
if key_alias and any(
|
||||
RouteChecks._route_matches_wildcard_pattern(route=key_alias, pattern=pat) for pat in key_patterns
|
||||
RouteChecks.route_matches_wildcard_pattern(route=key_alias, pattern=pat) for pat in key_patterns
|
||||
):
|
||||
if key_alias not in existing_keys:
|
||||
affected.append(key_alias)
|
||||
|
|
@ -349,8 +359,8 @@ async def estimate_attachment_impact(
|
|||
sample_teams=["(global scope — affects all teams)"],
|
||||
)
|
||||
|
||||
affected_keys: list = []
|
||||
affected_teams: list = []
|
||||
affected_keys: list[str] = []
|
||||
affected_teams: list[str] = []
|
||||
unnamed_keys = 0
|
||||
unnamed_teams = 0
|
||||
|
||||
|
|
@ -358,7 +368,7 @@ async def estimate_attachment_impact(
|
|||
team_patterns: Final = request.teams or []
|
||||
|
||||
# Fetch teams once — reused by both tag-based and alias-based lookups
|
||||
all_teams: list = []
|
||||
all_teams: Sequence[prisma_models.LiteLLM_TeamTable] = []
|
||||
if tag_patterns or team_patterns:
|
||||
all_teams = await _fetch_all_teams(prisma_client)
|
||||
|
||||
|
|
|
|||
|
|
@ -93,7 +93,7 @@ class _PromptTableActions(Protocol):
|
|||
|
||||
def create(self, *, data: Mapping[str, str | int | None]) -> Awaitable[_PromptRow]: ...
|
||||
|
||||
def update(self, *, where: Mapping[str, str | int], data: Mapping[str, str]) -> Awaitable[_PromptRow]: ...
|
||||
def update(self, *, where: Mapping[str, str | int], data: Mapping[str, str]) -> Awaitable[_PromptRow | None]: ...
|
||||
|
||||
def delete_many(self, *, where: Mapping[str, str]) -> Awaitable[int]: ...
|
||||
|
||||
|
|
@ -1157,6 +1157,12 @@ async def patch_prompt(
|
|||
data=update_data,
|
||||
)
|
||||
|
||||
if updated_prompt_db_entry is None:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail=f"Prompt with ID {base_prompt_id} not found in environment {env}",
|
||||
)
|
||||
|
||||
updated_prompt_spec: Final = create_versioned_prompt_spec(db_prompt=updated_prompt_db_entry)
|
||||
|
||||
return _reload_prompt_in_registry(IN_MEMORY_PROMPT_REGISTRY, versioned_id, updated_prompt_spec)
|
||||
|
|
|
|||
|
|
@ -133,6 +133,7 @@ if TYPE_CHECKING:
|
|||
from aiohttp import ClientSession
|
||||
from fastapi.routing import APIRoute
|
||||
from opentelemetry.trace import Span as _Span
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.integrations.opentelemetry import OpenTelemetry
|
||||
|
||||
|
|
@ -634,6 +635,7 @@ from litellm.proxy.utils import (
|
|||
from litellm.proxy.video_endpoints.endpoints import router as video_router
|
||||
from litellm.repositories.base_repository import SupportsModelDump
|
||||
from litellm.repositories.credentials_repository import CredentialsRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.router import (
|
||||
AssistantsTypedDict,
|
||||
Deployment,
|
||||
|
|
@ -1642,12 +1644,21 @@ class _InvitationLinkRow(Protocol):
|
|||
class _UserTableRow(Protocol):
|
||||
user_id: str
|
||||
user_email: str | None
|
||||
user_role: str
|
||||
user_role: str | None
|
||||
|
||||
|
||||
class _ModelTableRow(Protocol):
|
||||
model_id: str | None
|
||||
created_by: str | None
|
||||
class _UserTeamsRow(Protocol):
|
||||
@property
|
||||
def teams(self) -> Sequence[str]: ...
|
||||
|
||||
|
||||
_ProxyModelRow: TypeAlias = "prisma_models.LiteLLM_ProxyModelTable"
|
||||
|
||||
|
||||
def _config_param_table(client: PrismaClient | None) -> TableActions[_ConfigParamRow]:
|
||||
return cast( # cast-ok: this is prisma's LiteLLM_Config actions object, which parses its Json column to a mapping
|
||||
"TableActions[_ConfigParamRow]", ConfigRepository(client).table
|
||||
)
|
||||
|
||||
|
||||
class _TTFTRow(TypedDict):
|
||||
|
|
@ -4376,7 +4387,7 @@ class ProxyConfig:
|
|||
if prisma_client is None or not (general_settings.get("store_model_in_db", False) is True or store_model_in_db):
|
||||
return
|
||||
|
||||
row: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
|
||||
row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "environment_variables"}
|
||||
)
|
||||
existing: Final[dict] = dict(row.param_value) if row is not None and row.param_value is not None else {}
|
||||
|
|
@ -4858,6 +4869,14 @@ class ProxyConfig:
|
|||
if litellm_settings is None:
|
||||
litellm_settings = {}
|
||||
if litellm_settings:
|
||||
# Prometheus collectors have fixed label schemas. Load and validate this
|
||||
# setting before processing callbacks so YAML key order cannot construct
|
||||
# the collectors with the default caller-identity mode, and so an invalid
|
||||
# value fails the boot instead of being swallowed by callback init.
|
||||
from litellm.types.integrations.prometheus import validate_caller_identity_settings
|
||||
|
||||
validate_caller_identity_settings(litellm_settings)
|
||||
|
||||
# ANSI escape code for blue text
|
||||
blue_color_code: Final = "\033[94m"
|
||||
reset_color_code: Final = "\033[0m"
|
||||
|
|
@ -6232,7 +6251,7 @@ class ProxyConfig:
|
|||
4. Update router settings
|
||||
"""
|
||||
if llm_router is not None and prisma_client is not None:
|
||||
db_router_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
|
||||
db_router_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "router_settings"}
|
||||
)
|
||||
|
||||
|
|
@ -6660,7 +6679,7 @@ class ProxyConfig:
|
|||
def _should_load_db_object(self, object_type: str | SupportedDBObjectType) -> bool:
|
||||
return should_load_db_object(object_type=object_type)
|
||||
|
||||
async def _get_models_from_db(self, prisma_client: PrismaClient) -> list | None:
|
||||
async def _get_models_from_db(self, prisma_client: PrismaClient) -> Sequence[_ProxyModelRow] | None:
|
||||
"""
|
||||
Fetch all model deployments from the DB.
|
||||
|
||||
|
|
@ -6670,7 +6689,7 @@ class ProxyConfig:
|
|||
as "all models deleted" and must not evict existing router deployments.
|
||||
"""
|
||||
try:
|
||||
new_models: Final[list[_ModelTableRow]] = await ModelRepository(prisma_client).table.find_many()
|
||||
new_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository(prisma_client).table.find_many()
|
||||
return new_models
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception(
|
||||
|
|
@ -6957,10 +6976,13 @@ class ProxyConfig:
|
|||
"""
|
||||
|
||||
try:
|
||||
sso_settings: Final[_SSOConfigRow | None] = await call_with_db_reconnect_retry(
|
||||
prisma_client,
|
||||
lambda: SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}),
|
||||
reason="init_sso_settings_in_db_lookup_failure",
|
||||
sso_settings: Final[_SSOConfigRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime is a dict
|
||||
"_SSOConfigRow | None",
|
||||
await call_with_db_reconnect_retry(
|
||||
prisma_client,
|
||||
lambda: SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}),
|
||||
reason="init_sso_settings_in_db_lookup_failure",
|
||||
),
|
||||
)
|
||||
if sso_settings is not None:
|
||||
sso_settings.sso_settings.pop("role_mappings", None)
|
||||
|
|
@ -6988,12 +7010,15 @@ class ProxyConfig:
|
|||
)
|
||||
|
||||
try:
|
||||
db_record: Final[_ConfigOverridesRow | None] = await call_with_db_reconnect_retry(
|
||||
prisma_client,
|
||||
lambda: ConfigOverridesRepository(prisma_client).table.find_unique(
|
||||
where={"config_type": "hashicorp_vault"}
|
||||
db_record: Final[_ConfigOverridesRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime dict
|
||||
"_ConfigOverridesRow | None",
|
||||
await call_with_db_reconnect_retry(
|
||||
prisma_client,
|
||||
lambda: ConfigOverridesRepository(prisma_client).table.find_unique(
|
||||
where={"config_type": "hashicorp_vault"}
|
||||
),
|
||||
reason="init_hashicorp_vault_config_override_lookup_failure",
|
||||
),
|
||||
reason="init_hashicorp_vault_config_override_lookup_failure",
|
||||
)
|
||||
|
||||
if db_record is None or db_record.config_value is None:
|
||||
|
|
@ -8841,8 +8866,9 @@ class ProxyStartupEvent:
|
|||
|
||||
if prisma_client is None:
|
||||
return
|
||||
db_record: Final[_UISettingsRow | None] = await UISettingsRepository(prisma_client).table.find_unique(
|
||||
where={"id": "ui_settings"}
|
||||
db_record: Final[_UISettingsRow | None] = cast( # cast-ok: prisma Json stub is `str`, runtime is a dict
|
||||
"_UISettingsRow | None",
|
||||
await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}),
|
||||
)
|
||||
if db_record and db_record.ui_settings:
|
||||
raw: Final = db_record.ui_settings
|
||||
|
|
@ -9005,7 +9031,7 @@ class ProxyStartupEvent:
|
|||
# but YAML config has False.
|
||||
if store_model_in_db is not True and prisma_client is not None:
|
||||
try:
|
||||
_db_gs_record: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
|
||||
_db_gs_record: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "general_settings"}
|
||||
)
|
||||
if _db_gs_record is not None and isinstance(_db_gs_record.param_value, dict):
|
||||
|
|
@ -12150,14 +12176,14 @@ async def _check_if_model_is_user_added(
|
|||
id = model.get("model_info", {}).get("id", None)
|
||||
if id is None:
|
||||
continue
|
||||
db_model: _ModelTableRow | None = await ModelRepository(prisma_client).table.find_unique(where={"model_id": id})
|
||||
db_model: _ProxyModelRow | None = await ModelRepository(prisma_client).table.find_unique(where={"model_id": id})
|
||||
if db_model is not None:
|
||||
if db_model.created_by == user_api_key_dict.user_id:
|
||||
filtered_models.append(model)
|
||||
return filtered_models
|
||||
|
||||
|
||||
def _check_if_model_is_team_model(models: list[DeploymentTypedDict], user_row: LiteLLM_UserTable) -> list[dict]:
|
||||
def _check_if_model_is_team_model(models: list[DeploymentTypedDict], user_row: _UserTeamsRow) -> list[dict]:
|
||||
"""
|
||||
Check if model is a team model
|
||||
|
||||
|
|
@ -12210,10 +12236,11 @@ async def non_admin_all_models(
|
|||
raise HTTPException(status_code=400, detail={"error": "User not found"})
|
||||
|
||||
# Get all models that are team models, when model team_id == user_row.teams
|
||||
all_models += _check_if_model_is_team_model(
|
||||
models=llm_router.get_model_list() or [],
|
||||
user_row=user_row,
|
||||
)
|
||||
if user_row is not None:
|
||||
all_models += _check_if_model_is_team_model(
|
||||
models=llm_router.get_model_list() or [],
|
||||
user_row=user_row,
|
||||
)
|
||||
|
||||
# de-duplicate models. Only return unique model ids
|
||||
unique_models: Final = _deduplicate_litellm_router_models(models=all_models)
|
||||
|
|
@ -12637,7 +12664,7 @@ async def _fetch_db_models_for_search(
|
|||
|
||||
db_models_total_count: Final = await ModelRepository(prisma_client).table.count(where=db_where_condition)
|
||||
|
||||
db_models_raw: list = []
|
||||
db_models_raw: Sequence[_ProxyModelRow] = []
|
||||
if take_limit > 0:
|
||||
db_models_raw = await ModelRepository(prisma_client).table.find_many(
|
||||
where=db_where_condition,
|
||||
|
|
@ -13027,7 +13054,7 @@ async def _gather_team_accessible_model_ids(
|
|||
try:
|
||||
if team_object.models and SpecialModelNames.all_proxy_models.value not in team_object.models:
|
||||
_resolved_names: Final = _team_models_resolve_to_names(team_object.models, access_groups)
|
||||
db_models: Final[Sequence[_ModelTableRow]] = await ModelRepository(prisma_client).table.find_many(
|
||||
db_models: Final[Sequence[_ProxyModelRow]] = await ModelRepository(prisma_client).table.find_many(
|
||||
where={"model_name": {"in": _resolved_names}}
|
||||
)
|
||||
for db_model in db_models:
|
||||
|
|
@ -14501,14 +14528,18 @@ async def alerting_settings(
|
|||
)
|
||||
|
||||
## get general settings from db
|
||||
db_general_settings: Final = await ConfigRepository(prisma_client).table.find_first(
|
||||
db_general_settings: Final = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "general_settings"}
|
||||
)
|
||||
|
||||
if db_general_settings is not None and db_general_settings.param_value is not None:
|
||||
db_general_settings_dict: Final = dict(db_general_settings.param_value)
|
||||
alerting_args_dict: dict = db_general_settings_dict.get("alerting_args", {})
|
||||
alerting_values: list | None = db_general_settings_dict.get("alerting")
|
||||
alerting_args_dict: dict = cast( # cast-ok: ConfigGeneralSettings validates alerting_args as a dict on write
|
||||
dict[str, JsonValue], db_general_settings_dict.get("alerting_args", {})
|
||||
)
|
||||
alerting_values: list | None = cast( # cast-ok: ConfigGeneralSettings validates alerting as a list on write
|
||||
list[JsonValue] | None, db_general_settings_dict.get("alerting")
|
||||
)
|
||||
else:
|
||||
alerting_args_dict = {}
|
||||
alerting_values = None
|
||||
|
|
@ -15059,7 +15090,7 @@ async def onboarding(invite_link: str, request: Request):
|
|||
user_id=user_obj.user_id,
|
||||
key=onboarding_token,
|
||||
user_email=user_obj.user_email,
|
||||
user_role=user_obj.user_role,
|
||||
user_role=user_obj.user_role, # pyright: ignore[reportArgumentType] # nullable DB column, no unset contract
|
||||
login_method="username_password",
|
||||
premium_user=premium_user,
|
||||
auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"),
|
||||
|
|
@ -15168,7 +15199,7 @@ async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str:
|
|||
user_id=user_obj.user_id,
|
||||
key=key,
|
||||
user_email=user_obj.user_email,
|
||||
user_role=user_obj.user_role,
|
||||
user_role=user_obj.user_role, # pyright: ignore[reportArgumentType] # nullable DB column, no unset contract
|
||||
login_method="username_password",
|
||||
premium_user=premium_user,
|
||||
auth_header_name=general_settings.get("litellm_key_header_name", "Authorization"),
|
||||
|
|
@ -15729,7 +15760,7 @@ async def update_config(
|
|||
raise Exception("No DB Connected")
|
||||
|
||||
async def _read_section(param_name: str) -> dict:
|
||||
row: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
|
||||
row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": param_name}
|
||||
)
|
||||
if row is None or row.param_value is None:
|
||||
|
|
@ -15986,7 +16017,7 @@ async def update_config_general_settings(
|
|||
)
|
||||
|
||||
## get general settings from db
|
||||
db_general_settings: Final = await ConfigRepository(prisma_client).table.find_first(
|
||||
db_general_settings: Final = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "general_settings"}
|
||||
)
|
||||
### update value
|
||||
|
|
@ -16004,7 +16035,7 @@ async def update_config_general_settings(
|
|||
if data.field_name == "plugins":
|
||||
field_value = _preserve_redacted_plugin_keys(field_value, general_settings.get("plugins"))
|
||||
|
||||
general_settings[data.field_name] = field_value
|
||||
general_settings[data.field_name] = cast(JsonValue, field_value) # cast-ok: ConfigGeneralSettings validated it
|
||||
|
||||
response: Final = await ConfigRepository(prisma_client).table.upsert(
|
||||
where={"param_name": "general_settings"},
|
||||
|
|
@ -16024,7 +16055,7 @@ async def update_config_general_settings(
|
|||
)
|
||||
|
||||
if data.field_name == "plugins":
|
||||
register_plugins_from_config(general_settings)
|
||||
register_plugins_from_config(cast(dict[str, object], general_settings)) # cast-ok: the callee only reads it
|
||||
_apply_ssrf_general_settings(general_settings)
|
||||
|
||||
return response
|
||||
|
|
@ -16200,7 +16231,7 @@ async def get_config_general_settings(
|
|||
)
|
||||
|
||||
## get general settings from db
|
||||
db_general_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
|
||||
db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "general_settings"}
|
||||
)
|
||||
### pop the value
|
||||
|
|
@ -16389,7 +16420,7 @@ async def get_config_list(
|
|||
is_full_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
|
||||
## get general settings from db
|
||||
db_general_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
|
||||
db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "general_settings"}
|
||||
)
|
||||
|
||||
|
|
@ -16485,7 +16516,7 @@ async def get_config_list(
|
|||
)
|
||||
return_val.append(_response_obj)
|
||||
|
||||
db_litellm_settings_row: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
|
||||
db_litellm_settings_row: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "litellm_settings"}
|
||||
)
|
||||
db_litellm_settings: Final[dict] = (
|
||||
|
|
@ -16562,7 +16593,7 @@ async def delete_config_general_settings(
|
|||
)
|
||||
|
||||
## get general settings from db
|
||||
db_general_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first(
|
||||
db_general_settings: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_first(
|
||||
where={"param_name": "general_settings"}
|
||||
)
|
||||
### pop the value
|
||||
|
|
@ -17129,7 +17160,7 @@ async def reload_anthropic_beta_headers(
|
|||
last_anthropic_beta_headers_reload = current_time.isoformat()
|
||||
|
||||
# Set force reload flag in database for other pods, preserving existing interval_hours
|
||||
existing_beta_config: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_unique(
|
||||
existing_beta_config: Final[_ConfigParamRow | None] = await _config_param_table(prisma_client).find_unique(
|
||||
where={"param_name": "anthropic_beta_headers_reload_config"}
|
||||
)
|
||||
existing_beta_interval = None
|
||||
|
|
@ -17307,7 +17338,7 @@ async def get_anthropic_beta_headers_reload_status(
|
|||
}
|
||||
|
||||
# Get reload configuration from database
|
||||
config_record: Final = await ConfigRepository(prisma_client).table.find_unique(
|
||||
config_record: Final = await _config_param_table(prisma_client).find_unique(
|
||||
where={"param_name": "anthropic_beta_headers_reload_config"}
|
||||
)
|
||||
|
||||
|
|
@ -17321,7 +17352,9 @@ async def get_anthropic_beta_headers_reload_status(
|
|||
}
|
||||
|
||||
config: Final = config_record.param_value
|
||||
interval_hours: Final = config.get("interval_hours")
|
||||
interval_hours: Final = cast( # cast-ok: every writer of this key stores `hours: int` or an explicit None
|
||||
int | None, config.get("interval_hours")
|
||||
)
|
||||
|
||||
if interval_hours is None:
|
||||
verbose_proxy_logger.info("No interval configured, returning not scheduled")
|
||||
|
|
|
|||
|
|
@ -90,12 +90,15 @@ async def rerank(
|
|||
fastapi_response.headers.update(
|
||||
ProxyBaseLLMRequestProcessing.get_custom_headers(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_id=hidden_params.get("litellm_call_id", None) or data.get("litellm_call_id", None),
|
||||
model_id=model_id,
|
||||
cache_key=cache_key,
|
||||
api_base=api_base,
|
||||
version=version,
|
||||
response_cost=hidden_params.get("response_cost", None),
|
||||
model_region=getattr(user_api_key_dict, "allowed_model_region", ""),
|
||||
request_data=data,
|
||||
hidden_params=hidden_params,
|
||||
**additional_headers,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,5 +1,11 @@
|
|||
import json
|
||||
from typing import Final
|
||||
from collections.abc import Mapping
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Final,
|
||||
Protocol,
|
||||
cast, # noqa: TID251 # the config repository's table protocol omits find_first
|
||||
)
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
|
||||
|
|
@ -13,6 +19,7 @@ from litellm.proxy.common_utils.encrypt_decrypt_utils import (
|
|||
)
|
||||
from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.types.proxy.cloudzero_endpoints import (
|
||||
CloudZeroExportRequest,
|
||||
CloudZeroExportResponse,
|
||||
|
|
@ -22,6 +29,9 @@ from litellm.types.proxy.cloudzero_endpoints import (
|
|||
CloudZeroSettingsView,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.proxy_server import PrismaClient
|
||||
|
||||
router: Final = APIRouter()
|
||||
|
||||
|
||||
|
|
@ -29,6 +39,18 @@ router: Final = APIRouter()
|
|||
_sensitive_masker: Final = SensitiveDataMasker()
|
||||
|
||||
|
||||
class _CloudZeroConfigRow(Protocol):
|
||||
"""The ``LiteLLM_Config`` row holding ``cloudzero_settings``, as this module reads it."""
|
||||
|
||||
@property
|
||||
def param_value(self) -> str | Mapping[str, str] | None: ...
|
||||
|
||||
|
||||
def _config_table(prisma_client: "PrismaClient") -> TableActions[_CloudZeroConfigRow]:
|
||||
repository_table: Final = ConfigRepository(prisma_client).table
|
||||
return cast(TableActions[_CloudZeroConfigRow], repository_table) # cast-ok: repo protocol omits find_first
|
||||
|
||||
|
||||
async def _set_cloudzero_settings(api_key: str, connection_id: str, timezone: str):
|
||||
"""
|
||||
Store CloudZero settings in the database with encrypted API key.
|
||||
|
|
@ -82,9 +104,7 @@ async def _get_cloudzero_settings():
|
|||
detail={"error": CommonProxyErrors.db_not_connected_error.value},
|
||||
)
|
||||
|
||||
cloudzero_config: Final = await ConfigRepository(prisma_client).table.find_first(
|
||||
where={"param_name": "cloudzero_settings"}
|
||||
)
|
||||
cloudzero_config: Final = await _config_table(prisma_client).find_first(where={"param_name": "cloudzero_settings"})
|
||||
if cloudzero_config is None or cloudzero_config.param_value is None:
|
||||
return {}
|
||||
|
||||
|
|
@ -268,7 +288,7 @@ async def is_cloudzero_setup_in_db() -> bool:
|
|||
return False
|
||||
|
||||
# Check for CloudZero settings in database
|
||||
cloudzero_config: Final = await ConfigRepository(prisma_client).table.find_first(
|
||||
cloudzero_config: Final = await _config_table(prisma_client).find_first(
|
||||
where={"param_name": "cloudzero_settings"}
|
||||
)
|
||||
|
||||
|
|
@ -530,7 +550,7 @@ async def delete_cloudzero_settings(
|
|||
)
|
||||
|
||||
# Check if CloudZero settings exist
|
||||
cloudzero_config: Final = await ConfigRepository(prisma_client).table.find_first(
|
||||
cloudzero_config: Final = await _config_table(prisma_client).find_first(
|
||||
where={"param_name": "cloudzero_settings"}
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -4,10 +4,22 @@ import json
|
|||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, NamedTuple, Protocol, TypedDict, TypeVar
|
||||
from typing import (
|
||||
TYPE_CHECKING,
|
||||
Annotated,
|
||||
Any,
|
||||
Final,
|
||||
Literal,
|
||||
NamedTuple,
|
||||
Protocol,
|
||||
TypedDict,
|
||||
TypeVar,
|
||||
cast, # noqa: TID251 # prisma group_by returns untyped aggregate mappings
|
||||
)
|
||||
|
||||
import fastapi
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from typing_extensions import ReadOnly
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -23,6 +35,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
|||
get_spend_by_team_and_customer,
|
||||
)
|
||||
from litellm.proxy.utils import handle_exception_on_proxy
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.repositories.table_repositories import SpendLogsRepository
|
||||
from litellm.repositories.team_repository import TeamRepository
|
||||
from litellm.repositories.verification_token_repository import (
|
||||
|
|
@ -30,6 +43,8 @@ from litellm.repositories.verification_token_repository import (
|
|||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from prisma import models as prisma_models
|
||||
|
||||
from litellm.proxy.proxy_server import PrismaClient
|
||||
from litellm.proxy.spend_tracking.cold_storage_handler import ColdStorageHandler
|
||||
else:
|
||||
|
|
@ -139,6 +154,18 @@ class _SessionSpendRow(TypedDict):
|
|||
mcp_tool_call_spend: float
|
||||
|
||||
|
||||
class _SpendSumAggregate(TypedDict, total=False):
|
||||
spend: ReadOnly[float]
|
||||
|
||||
|
||||
class _SpendGroupByRow(TypedDict):
|
||||
api_key: ReadOnly[str]
|
||||
user: ReadOnly[str | None]
|
||||
model: ReadOnly[str]
|
||||
startTime: ReadOnly[object]
|
||||
_sum: ReadOnly[_SpendSumAggregate]
|
||||
|
||||
|
||||
async def _query_raw(prisma_client: PrismaClient, sql_query: str, *args: object) -> Sequence[_RowT]:
|
||||
"""Run a raw read query and return its rows as the row type the caller declares."""
|
||||
return await prisma_client.db.query_raw(sql_query, *args)
|
||||
|
|
@ -149,24 +176,6 @@ async def _query_raw_or_none(prisma_client: PrismaClient, sql_query: str, *args:
|
|||
return await _query_raw(prisma_client, sql_query, *args)
|
||||
|
||||
|
||||
class _SpendLogsTable(Protocol):
|
||||
"""The subset of the Prisma spend-logs table API this module uses."""
|
||||
|
||||
async def find_many(
|
||||
self, *, where: Mapping[str, object], order: Mapping[str, str]
|
||||
) -> Sequence[_SupportsModelDump]: ...
|
||||
|
||||
async def find_unique(
|
||||
self, *, where: Mapping[str, object], include: None = None
|
||||
) -> _SpendLogOwnershipRow | None: ...
|
||||
|
||||
async def count(self, *, where: Mapping[str, object]) -> int: ...
|
||||
|
||||
async def group_by(
|
||||
self, *, by: Sequence[str], where: Mapping[str, object], count: Mapping[str, bool]
|
||||
) -> Sequence[_SessionCountRow]: ...
|
||||
|
||||
|
||||
class _TeamTable(Protocol):
|
||||
"""The subset of the Prisma team table API this module uses."""
|
||||
|
||||
|
|
@ -183,7 +192,7 @@ class _VerificationTokenTable(Protocol):
|
|||
async def update_many(self, *, data: Mapping[str, float], where: Mapping[str, object]) -> int: ...
|
||||
|
||||
|
||||
def _spend_logs_table(prisma_client: PrismaClient) -> _SpendLogsTable:
|
||||
def _spend_logs_table(prisma_client: PrismaClient) -> TableActions["prisma_models.LiteLLM_SpendLogs"]:
|
||||
return SpendLogsRepository(prisma_client).table
|
||||
|
||||
|
||||
|
|
@ -221,11 +230,12 @@ async def _count_logs_per_session(
|
|||
prisma_client: PrismaClient, session_ids: Sequence[str | None]
|
||||
) -> Sequence[_SessionCountRow]:
|
||||
"""Count spend log rows per session for the given session ids."""
|
||||
return await _spend_logs_table(prisma_client).group_by(
|
||||
rows: Final = await _spend_logs_table(prisma_client).group_by(
|
||||
by=["session_id"],
|
||||
where={"session_id": {"in": session_ids}},
|
||||
count={"session_id": True},
|
||||
)
|
||||
return cast(Sequence[_SessionCountRow], rows) # cast-ok: group_by(count=) shape is fixed by the by/count args
|
||||
|
||||
|
||||
async def _find_team_row(prisma_client: PrismaClient, team_id: str) -> _SupportsModelDump | None:
|
||||
|
|
@ -2974,8 +2984,9 @@ async def view_spend_logs(
|
|||
)
|
||||
|
||||
if isinstance(response, list) and len(response) > 0 and isinstance(response[0], dict):
|
||||
spend_rows: Final = cast(Sequence[_SpendGroupByRow], response) # cast-ok: by/sum fix the shape
|
||||
result: Final[dict] = {}
|
||||
for record in response:
|
||||
for record in spend_rows:
|
||||
dt_object = datetime.strptime(str(record["startTime"]), "%Y-%m-%dT%H:%M:%S.%fZ")
|
||||
date = dt_object.date()
|
||||
if date not in result:
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Reference in a new issue