scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

feat: enforce x-litellm-trace-id in header, if required

Browse source
This commit is contained in:
Krrish Dholakia 2026-03-04 17:37:37 -08:00
parent 8541272629
commit b979830529
9 changed files with 883 additions and 758 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

252
litellm/model_prices_and_context_window_backup.json
View file

@ -9779,6 +9779,122 @@
}
]
},
"dashscope/qwen3-max-2026-01-23": {
"litellm_provider": "dashscope",
"max_input_tokens": 258048,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"source": "https://www.alibabacloud.com/help/en/model-studio/models",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"tiered_pricing": [
{
"input_cost_per_token": 1.2e-06,
"output_cost_per_token": 6e-06,
"range": [
0,
32000.0
]
},
{
"input_cost_per_token": 2.4e-06,
"output_cost_per_token": 1.2e-05,
"range": [
32000.0,
128000.0
]
},
{
"input_cost_per_token": 3e-06,
"output_cost_per_token": 1.5e-05,
"range": [
128000.0,
252000.0
]
}
]
},
"dashscope/qwen3-next-80b-a3b-instruct": {
"input_cost_per_token": 1.5e-07,
"litellm_provider": "dashscope",
"max_input_tokens": 262144,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
"supports_function_calling": true,
"supports_tool_choice": true
},
"dashscope/qwen3-next-80b-a3b-thinking": {
"input_cost_per_token": 1.5e-07,
"litellm_provider": "dashscope",
"max_input_tokens": 262144,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 1.2e-06,
"source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"dashscope/qwen3-vl-235b-a22b-instruct": {
"input_cost_per_token": 4e-07,
"litellm_provider": "dashscope",
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 1.6e-06,
"source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_vision": true
},
"dashscope/qwen3-vl-235b-a22b-thinking": {
"input_cost_per_token": 4e-07,
"litellm_provider": "dashscope",
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 4e-06,
"source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"dashscope/qwen3-vl-32b-instruct": {
"input_cost_per_token": 1.6e-07,
"litellm_provider": "dashscope",
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 6.4e-07,
"source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
"supports_function_calling": true,
"supports_tool_choice": true,
"supports_vision": true
},
"dashscope/qwen3-vl-32b-thinking": {
"input_cost_per_token": 1.6e-07,
"litellm_provider": "dashscope",
"max_input_tokens": 131072,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "chat",
"output_cost_per_token": 2.87e-06,
"source": "https://www.alibabacloud.com/help/en/model-studio/model-pricing",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"dashscope/qwen3-vl-plus": {
"litellm_provider": "dashscope",
"max_input_tokens": 260096,
@ -25806,6 +25922,30 @@
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
},
"openrouter/anthropic/claude-sonnet-4.6": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_creation_input_token_cost_above_200k_tokens": 7.5e-06,
"cache_read_input_token_cost": 3e-07,
"cache_read_input_token_cost_above_200k_tokens": 6e-07,
"input_cost_per_token": 3e-06,
"input_cost_per_token_above_200k_tokens": 6e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1000000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"output_cost_per_token_above_200k_tokens": 2.25e-05,
"source": "https://openrouter.ai/anthropic/claude-sonnet-4.6",
"supports_assistant_prefill": true,
"supports_computer_use": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true,
"tool_use_system_prompt_tokens": 159
},
"openrouter/anthropic/claude-opus-4.5": {
"cache_creation_input_token_cost": 6.25e-06,
"cache_read_input_token_cost": 5e-07,
@ -26156,6 +26296,39 @@
"supports_web_search": true,
"tpm": 800000
},
"openrouter/google/gemini-3.1-pro-preview": {
"cache_read_input_token_cost": 2e-07,
"cache_read_input_token_cost_above_200k_tokens": 4e-07,
"cache_creation_input_token_cost_above_200k_tokens": 2.5e-07,
"input_cost_per_token": 2e-06,
"input_cost_per_token_above_200k_tokens": 4e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 1048576,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 1.2e-05,
"output_cost_per_token_above_200k_tokens": 1.8e-05,
"source": "https://openrouter.ai/google/gemini-3.1-pro-preview",
"supported_modalities": [
"text",
"image",
"audio",
"video"
],
"supported_output_modalities": [
"text"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_pdf_input": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"openrouter/gryphe/mythomax-l2-13b": {
"input_cost_per_token": 1.875e-06,
"litellm_provider": "openrouter",
@ -26533,6 +26706,29 @@
"supports_reasoning": true,
"supports_tool_choice": true
},
"openrouter/openai/gpt-5.1-codex-max": {
"cache_read_input_token_cost": 1.25e-07,
"input_cost_per_token": 1.25e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 400000,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 1e-05,
"source": "https://openrouter.ai/openai/gpt-5.1-codex-max",
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text"
],
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"supports_vision": true
},
"openrouter/openai/gpt-5.2": {
"input_cost_per_image": 0,
"cache_read_input_token_cost": 1.75e-07,
@ -26687,6 +26883,19 @@
"supports_tool_choice": true,
"supports_function_calling": true
},
"openrouter/qwen/qwen3-coder-plus": {
"input_cost_per_token": 1e-06,
"litellm_provider": "openrouter",
"max_input_tokens": 997952,
"max_output_tokens": 65536,
"max_tokens": 65536,
"mode": "chat",
"output_cost_per_token": 5e-06,
"source": "https://openrouter.ai/qwen/qwen3-coder-plus",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"openrouter/qwen/qwen3-235b-a22b-2507": {
"input_cost_per_token": 7.1e-08,
"litellm_provider": "openrouter",
@ -26822,6 +27031,19 @@
"supports_vision": true,
"supports_prompt_caching": false
},
"openrouter/z-ai/glm-5": {
"input_cost_per_token": 8e-07,
"litellm_provider": "openrouter",
"max_input_tokens": 202752,
"max_output_tokens": 128000,
"max_tokens": 128000,
"mode": "chat",
"output_cost_per_token": 2.56e-06,
"source": "https://openrouter.ai/z-ai/glm-5",
"supports_function_calling": true,
"supports_reasoning": true,
"supports_tool_choice": true
},
"openrouter/minimax/minimax-m2.1": {
"input_cost_per_token": 2.7e-07,
"output_cost_per_token": 1.2e-06,
@ -34327,6 +34549,36 @@
"supports_tool_choice": true,
"source": "https://aws.amazon.com/bedrock/pricing/"
},
"zai/glm-5": {
"cache_creation_input_token_cost": 0,
"cache_read_input_token_cost": 2e-07,
"input_cost_per_token": 1e-06,
"output_cost_per_token": 3.2e-06,
"litellm_provider": "zai",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"mode": "chat",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"source": "https://docs.z.ai/guides/overview/pricing"
},
"zai/glm-5-code": {
"cache_creation_input_token_cost": 0,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 1.2e-06,
"output_cost_per_token": 5e-06,
"litellm_provider": "zai",
"max_input_tokens": 200000,
"max_output_tokens": 128000,
"mode": "chat",
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_tool_choice": true,
"source": "https://docs.z.ai/guides/overview/pricing"
},
"zai/glm-4.7": {
"cache_creation_input_token_cost": 0,
"cache_read_input_token_cost": 1.1e-07,

76
litellm/proxy/_types.py
View file

@ -1,59 +1,40 @@
import enum
import json
from datetime import datetime
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union
from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal,
Optional, Union)
import httpx
from pydantic import (
BaseModel,
ConfigDict,
Field,
Json,
field_validator,
model_validator,
)
from pydantic import (BaseModel, ConfigDict, Field, Json, field_validator,
model_validator)
from typing_extensions import Required, TypedDict
from litellm._uuid import uuid
from litellm.types.integrations.slack_alerting import AlertType
from litellm.types.llms.openai import (
AllMessageValues,
OpenAIFileObject,
ResponsesAPIResponse,
)
from litellm.types.mcp import (
MCPAuthType,
MCPCredentials,
MCPTransport,
MCPTransportType,
)
from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject,
ResponsesAPIResponse)
from litellm.types.mcp import (MCPAuthType, MCPCredentials, MCPTransport,
MCPTransportType)
from litellm.types.mcp_server.mcp_server_manager import MCPInfo
from litellm.types.router import RouterErrors, UpdateRouterConfig
from litellm.types.secret_managers.main import KeyManagementSystem
from litellm.types.utils import (
CallTypes,
CostBreakdown,
EmbeddingResponse,
GenericBudgetConfigType,
ImageResponse,
LiteLLMBatch,
LiteLLMFineTuningJob,
LiteLLMPydanticObjectBase,
ModelResponse,
ProviderField,
StandardCallbackDynamicParams,
StandardLoggingGuardrailInformation,
StandardLoggingMCPToolCall,
StandardLoggingModelInformation,
StandardLoggingPayloadErrorInformation,
StandardLoggingPayloadStatus,
StandardLoggingVectorStoreRequest,
StandardPassThroughResponseObject,
TextCompletionResponse,
)
from litellm.types.utils import (CallTypes, CostBreakdown, EmbeddingResponse,
GenericBudgetConfigType, ImageResponse,
LiteLLMBatch, LiteLLMFineTuningJob,
LiteLLMPydanticObjectBase, ModelResponse,
ProviderField, StandardCallbackDynamicParams,
StandardLoggingGuardrailInformation,
StandardLoggingMCPToolCall,
StandardLoggingModelInformation,
StandardLoggingPayloadErrorInformation,
StandardLoggingPayloadStatus,
StandardLoggingVectorStoreRequest,
StandardPassThroughResponseObject,
TextCompletionResponse)
from litellm.types.videos.main import VideoObject
from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type
from .types_utils.utils import (get_instance_fn,
validate_custom_validate_return_type)
if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span
@ -200,6 +181,7 @@ class Litellm_EntityType(enum.Enum):
ORGANIZATION = "organization"
PROJECT = "project"
TAG = "tag"
AGENT = "agent"
# global proxy level entity
PROXY = "proxy"
@ -2385,7 +2367,8 @@ class UserAPIKeyAuth(
This is used to track number of requests/spend for health check calls.
"""
from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
from litellm.constants import \
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
@ -2417,7 +2400,8 @@ class UserAPIKeyAuth(
This is used to track actions performed by automated system jobs.
"""
from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
from litellm.constants import \
LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME
return cls(
api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME,
@ -2823,7 +2807,8 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase):
@model_validator(mode="after")
def mask_api_keys(self):
from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker
from litellm.litellm_core_utils.sensitive_data_masker import \
SensitiveDataMasker
masker = SensitiveDataMasker(sensitive_patterns={"key"})
@ -4146,6 +4131,7 @@ class DBSpendUpdateTransactions(TypedDict):
team_member_list_transactions: Optional[Dict[str, float]]
org_list_transactions: Optional[Dict[str, float]]
tag_list_transactions: Optional[Dict[str, float]]
agent_list_transactions: Optional[Dict[str, float]]
class SpendUpdateQueueItem(TypedDict, total=False):

128
litellm/proxy/auth/auth_checks.py
View file

@ -11,7 +11,8 @@ Run checks for:
import asyncio
import re
import time
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union,
cast)
from fastapi import HTTPException, Request, status
from pydantic import BaseModel
@ -20,42 +21,29 @@ import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.caching.dual_cache import LimitedSizeOrderedDict
from litellm.constants import (
CLI_JWT_EXPIRATION_HOURS,
CLI_JWT_TOKEN_NAME,
DEFAULT_ACCESS_GROUP_CACHE_TTL,
DEFAULT_IN_MEMORY_TTL,
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
DEFAULT_MAX_RECURSE_DEPTH,
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
)
from litellm.constants import (CLI_JWT_EXPIRATION_HOURS, CLI_JWT_TOKEN_NAME,
DEFAULT_ACCESS_GROUP_CACHE_TTL,
DEFAULT_IN_MEMORY_TTL,
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
DEFAULT_MAX_RECURSE_DEPTH,
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE)
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.proxy._types import (
RBAC_ROLES,
CallInfo,
LiteLLM_AccessGroupTable,
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
Litellm_EntityType,
LiteLLM_JWTAuth,
LiteLLM_ObjectPermissionTable,
LiteLLM_OrganizationMembershipTable,
LiteLLM_OrganizationTable,
LiteLLM_ProjectTableCachedObj,
LiteLLM_TagTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable,
LiteLLMRoutes,
LitellmUserRoles,
NewTeamRequest,
ProxyErrorTypes,
ProxyException,
RoleBasedPermissions,
SpecialModelNames,
UserAPIKeyAuth,
)
from litellm.proxy._types import (RBAC_ROLES, CallInfo,
LiteLLM_AccessGroupTable,
LiteLLM_BudgetTable, LiteLLM_EndUserTable,
Litellm_EntityType, LiteLLM_JWTAuth,
LiteLLM_ObjectPermissionTable,
LiteLLM_OrganizationMembershipTable,
LiteLLM_OrganizationTable,
LiteLLM_ProjectTableCachedObj,
LiteLLM_TagTable, LiteLLM_TeamMembership,
LiteLLM_TeamTable,
LiteLLM_TeamTableCachedObj,
LiteLLM_UserTable, LiteLLMRoutes,
LitellmUserRoles, NewTeamRequest,
ProxyErrorTypes, ProxyException,
RoleBasedPermissions, SpecialModelNames,
UserAPIKeyAuth)
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.route_llm_request import route_request
@ -282,6 +270,29 @@ async def common_checks(
code=status.HTTP_401_UNAUTHORIZED,
)
# Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent
if valid_token is not None and valid_token.agent_id:
from litellm.proxy.agent_endpoints.agent_registry import \
global_agent_registry
from litellm.proxy.litellm_pre_call_utils import \
get_chain_id_from_headers
agent = global_agent_registry.get_agent_by_id(agent_id=valid_token.agent_id)
if agent is not None:
require_trace_id = (agent.litellm_params or {}).get(
"require_trace_id_on_calls_by_agent"
)
if require_trace_id:
headers_dict = dict(request.headers)
trace_id = get_chain_id_from_headers(headers_dict)
if not trace_id:
raise ProxyException(
message="Requests made with this agent's key must include the x-litellm-trace-id header.",
type=ProxyErrorTypes.bad_request_error,
param=None,
code=status.HTTP_400_BAD_REQUEST,
)
## 2.1 If user can call model (if personal key)
if _model and team_object is None and user_object is not None:
await can_user_call_model(
@ -436,7 +447,8 @@ async def common_checks(
_request_metadata: dict = request_body.get("metadata", {}) or {}
if _request_metadata.get("guardrails"):
# check if team allowed to modify guardrails
from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails
from litellm.proxy.guardrails.guardrail_helpers import \
can_modify_guardrails
can_modify: bool = can_modify_guardrails(team_object)
if can_modify is False:
@ -1881,9 +1893,8 @@ class ExperimentalUIJWTToken:
def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str:
from datetime import timedelta
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
encrypt_value_helper,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import \
encrypt_value_helper
if user_info.user_role is None:
raise Exception("User role is required for experimental UI login")
@ -1929,9 +1940,8 @@ class ExperimentalUIJWTToken:
"""
from datetime import timedelta
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
encrypt_value_helper,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import \
encrypt_value_helper
if user_info.user_role is None:
raise Exception("User role is required for CLI JWT login")
@ -1970,9 +1980,8 @@ class ExperimentalUIJWTToken:
import json
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_utils.encrypt_decrypt_utils import (
decrypt_value_helper,
)
from litellm.proxy.common_utils.encrypt_decrypt_utils import \
decrypt_value_helper
decrypted_token = decrypt_value_helper(
hashed_token, key="ui_hash_key", exception_type="debug"
@ -2070,13 +2079,11 @@ async def get_key_object(
)
# else, check db
_valid_token: Optional[BaseModel] = (
await _fetch_key_object_from_db_with_reconnect(
hashed_token=hashed_token,
prisma_client=prisma_client,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
_valid_token: Optional[BaseModel] = await _fetch_key_object_from_db_with_reconnect(
hashed_token=hashed_token,
prisma_client=prisma_client,
parent_otel_span=parent_otel_span,
proxy_logging_obj=proxy_logging_obj,
)
if _valid_token is None:
@ -2222,9 +2229,9 @@ async def get_org_object(
# Cache the result
await user_api_key_cache.async_set_cache(
key=cache_key,
value=response.model_dump()
if hasattr(response, "model_dump")
else response,
value=(
response.model_dump() if hasattr(response, "model_dump") else response
),
ttl=DEFAULT_IN_MEMORY_TTL,
)
@ -2267,8 +2274,10 @@ async def _get_resources_from_access_groups(
# Lazy import to avoid circular imports
if prisma_client is None or user_api_key_cache is None:
from litellm.proxy.proxy_server import prisma_client as _prisma_client
from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj
from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache
from litellm.proxy.proxy_server import \
proxy_logging_obj as _proxy_logging_obj
from litellm.proxy.proxy_server import \
user_api_key_cache as _user_api_key_cache
prisma_client = prisma_client or _prisma_client
user_api_key_cache = user_api_key_cache or _user_api_key_cache
@ -3224,7 +3233,8 @@ async def _tag_max_budget_check(
BudgetExceededError if any tag is over its max budget.
Triggers a budget alert if any tag is over its max budget.
"""
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
from litellm.proxy.common_utils.http_parsing_utils import \
get_tags_from_request_body
if prisma_client is None:
return

140
litellm/proxy/db/db_spend_update_writer.py
View file

@ -13,49 +13,36 @@ import random
import time
import traceback
from datetime import datetime, timedelta, timezone
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Literal,
Optional,
Union,
cast,
overload,
)
from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union,
cast, overload)
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching import DualCache, RedisCache
from litellm.constants import DB_SPEND_UPDATE_JOB_NAME
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.proxy._types import (
DB_CONNECTION_ERROR_TYPES,
BaseDailySpendTransaction,
DailyAgentSpendTransaction,
DailyEndUserSpendTransaction,
DailyOrganizationSpendTransaction,
DailyTagSpendTransaction,
DailyTeamSpendTransaction,
DailyUserSpendTransaction,
DBSpendUpdateTransactions,
Litellm_EntityType,
LiteLLM_UserTable,
SpendLogsMetadata,
SpendLogsPayload,
SpendUpdateQueueItem,
ToolDiscoveryQueueItem,
)
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
DailySpendUpdateQueue,
)
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import (
ToolDiscoveryQueue,
)
from litellm.proxy._types import (DB_CONNECTION_ERROR_TYPES,
BaseDailySpendTransaction,
DailyAgentSpendTransaction,
DailyEndUserSpendTransaction,
DailyOrganizationSpendTransaction,
DailyTagSpendTransaction,
DailyTeamSpendTransaction,
DailyUserSpendTransaction,
DBSpendUpdateTransactions,
Litellm_EntityType, LiteLLM_UserTable,
SpendLogsMetadata, SpendLogsPayload,
SpendUpdateQueueItem, ToolDiscoveryQueueItem)
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import \
DailySpendUpdateQueue
from litellm.proxy.db.db_transaction_queue.pod_lock_manager import \
PodLockManager
from litellm.proxy.db.db_transaction_queue.redis_update_buffer import \
RedisUpdateBuffer
from litellm.proxy.db.db_transaction_queue.spend_update_queue import \
SpendUpdateQueue
from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import \
ToolDiscoveryQueue
from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING
if TYPE_CHECKING:
@ -104,12 +91,10 @@ class DBSpendUpdateWriter:
end_time: Optional[datetime],
response_cost: Optional[float],
):
from litellm.proxy.proxy_server import (
disable_spend_logs,
litellm_proxy_budget_name,
prisma_client,
user_api_key_cache,
)
from litellm.proxy.proxy_server import (disable_spend_logs,
litellm_proxy_budget_name,
prisma_client,
user_api_key_cache)
from litellm.proxy.utils import ProxyUpdateSpend, hash_token
try:
@ -124,9 +109,8 @@ class DBSpendUpdateWriter:
hashed_token = token
## CREATE SPEND LOG PAYLOAD ##
from litellm.proxy.spend_tracking.spend_tracking_utils import (
get_logging_payload,
)
from litellm.proxy.spend_tracking.spend_tracking_utils import \
get_logging_payload
payload = get_logging_payload(
kwargs=kwargs,
@ -384,6 +368,18 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
try:
await self._update_agent_db(
response_cost=response_cost,
agent_id=payload_copy.get("agent_id"),
prisma_client=prisma_client,
)
except Exception:
verbose_proxy_logger.debug(
"_batch_database_updates: _update_agent_db failed: %s",
traceback.format_exc(),
)
try:
await self.add_spend_log_transaction_to_daily_user_transaction(
payload=payload_copy,
@ -614,6 +610,37 @@ class DBSpendUpdateWriter:
)
raise e
async def _update_agent_db(
self,
response_cost: Optional[float],
agent_id: Optional[str],
prisma_client: Optional[PrismaClient],
):
try:
if agent_id is None or prisma_client is None:
verbose_proxy_logger.debug(
"track_cost_callback: agent_id is None or prisma_client is None. Not tracking spend for agent"
)
return
await self.spend_update_queue.add_update(
update=SpendUpdateQueueItem(
entity_type=Litellm_EntityType.AGENT,
entity_id=agent_id,
response_cost=response_cost,
)
)
except Exception as e:
verbose_proxy_logger.error(
"Spend tracking - failed to enqueue agent spend update. "
"agent_id=%s, response_cost=%s - %s\n%s",
agent_id,
response_cost,
str(e),
traceback.format_exc(),
)
raise e
async def _update_tag_db(
self,
response_cost: Optional[float],
@ -773,7 +800,7 @@ class DBSpendUpdateWriter:
if db_spend_update_transactions is not None:
verbose_proxy_logger.info(
"Spend tracking - committing spend updates from Redis to DB: "
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d",
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d",
len(db_spend_update_transactions.get("key_list_transactions") or {}),
len(db_spend_update_transactions.get("user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_list_transactions") or {}),
@ -781,6 +808,7 @@ class DBSpendUpdateWriter:
len(db_spend_update_transactions.get("end_user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_member_list_transactions") or {}),
len(db_spend_update_transactions.get("tag_list_transactions") or {}),
len(db_spend_update_transactions.get("agent_list_transactions") or {}),
)
await self._commit_spend_updates_to_db(
prisma_client=prisma_client,
@ -985,10 +1013,8 @@ class DBSpendUpdateWriter:
Commits all the spend `UPDATE` transactions to the Database
"""
from litellm.proxy.utils import (
ProxyUpdateSpend,
_raise_failed_update_spend_exception,
)
from litellm.proxy.utils import (ProxyUpdateSpend,
_raise_failed_update_spend_exception)
### UPDATE USER TABLE ###
user_list_transactions = db_spend_update_transactions["user_list_transactions"]
@ -1264,6 +1290,18 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE AGENT TABLE ###
agent_list_transactions = db_spend_update_transactions["agent_list_transactions"]
await DBSpendUpdateWriter._update_entity_spend_in_db(
entity_name="Agent",
transactions=agent_list_transactions,
table_accessor="litellm_agentstable",
where_field="agent_id",
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
@staticmethod
async def _update_entity_spend_in_db(
entity_name: str,

50
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -10,33 +10,31 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import (
MAX_REDIS_BUFFER_DEQUEUE_COUNT,
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
REDIS_UPDATE_BUFFER_KEY,
)
from litellm.constants import (MAX_REDIS_BUFFER_DEQUEUE_COUNT,
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
REDIS_UPDATE_BUFFER_KEY)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import (
DailyTagSpendTransaction,
DailyTeamSpendTransaction,
DailyUserSpendTransaction,
DailyOrganizationSpendTransaction,
DailyEndUserSpendTransaction,
DBSpendUpdateTransactions,
DailyAgentSpendTransaction,
)
from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
DailySpendUpdateQueue,
)
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
from litellm.proxy._types import (DailyAgentSpendTransaction,
DailyEndUserSpendTransaction,
DailyOrganizationSpendTransaction,
DailyTagSpendTransaction,
DailyTeamSpendTransaction,
DailyUserSpendTransaction,
DBSpendUpdateTransactions)
from litellm.proxy.db.db_transaction_queue.base_update_queue import \
service_logger_obj
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import \
DailySpendUpdateQueue
from litellm.proxy.db.db_transaction_queue.spend_update_queue import \
SpendUpdateQueue
from litellm.secret_managers.main import str_to_bool
from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation
from litellm.types.caching import (RedisPipelineLpopOperation,
RedisPipelineRpushOperation)
from litellm.types.services import ServiceTypes
if TYPE_CHECKING:
@ -579,6 +577,7 @@ class RedisUpdateBuffer:
team_member_list_transactions={},
org_list_transactions={},
tag_list_transactions={},
agent_list_transactions={},
)
# Define the transaction fields to process
@ -590,6 +589,7 @@ class RedisUpdateBuffer:
"team_member_list_transactions",
"org_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
]
# Loop through each transaction and combine the values

17
litellm/proxy/db/db_transaction_queue/spend_update_queue.py
View file

@ -3,15 +3,10 @@ from typing import Dict, List, Optional
from litellm._logging import verbose_proxy_logger
from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE
from litellm.proxy._types import (
DBSpendUpdateTransactions,
Litellm_EntityType,
SpendUpdateQueueItem,
)
from litellm.proxy._types import (DBSpendUpdateTransactions,
Litellm_EntityType, SpendUpdateQueueItem)
from litellm.proxy.db.db_transaction_queue.base_update_queue import (
BaseUpdateQueue,
service_logger_obj,
)
BaseUpdateQueue, service_logger_obj)
from litellm.types.services import ServiceTypes
@ -145,6 +140,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
team_member_list_transactions={},
org_list_transactions={},
tag_list_transactions={},
agent_list_transactions={},
)
# Map entity types to their corresponding transaction dictionary keys
@ -156,6 +152,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions",
Litellm_EntityType.ORGANIZATION: "org_list_transactions",
Litellm_EntityType.TAG: "tag_list_transactions",
Litellm_EntityType.AGENT: "agent_list_transactions",
}
for update in updates:
@ -207,6 +204,10 @@ class SpendUpdateQueue(BaseUpdateQueue):
transactions_dict = db_spend_update_transactions[
"tag_list_transactions"
]
elif dict_key == "agent_list_transactions":
transactions_dict = db_spend_update_transactions[
"agent_list_transactions"
]
else:
continue

56
litellm/proxy/litellm_pre_call_utils.py
View file

@ -10,16 +10,12 @@ import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm._service_logger import ServiceLogging
from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
from litellm.proxy._types import (
AddTeamCallback,
CommonProxyErrors,
LitellmDataForBackendLLMCall,
LitellmUserRoles,
SpecialHeaders,
TeamCallbackMetadata,
UserAPIKeyAuth,
)
from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers
from litellm.proxy._types import (AddTeamCallback, CommonProxyErrors,
LitellmDataForBackendLLMCall,
LitellmUserRoles, SpecialHeaders,
TeamCallbackMetadata, UserAPIKeyAuth)
from litellm.proxy.common_utils.http_parsing_utils import \
_safe_get_request_headers
# Cache special headers as a frozenset for O(1) lookup performance
_SPECIAL_HEADERS_CACHE = frozenset(
@ -28,12 +24,9 @@ _SPECIAL_HEADERS_CACHE = frozenset(
from litellm.router import Router
from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS
from litellm.types.services import ServiceTypes
from litellm.types.utils import (
LlmProviders,
ProviderSpecificHeader,
StandardLoggingUserAPIKeyMetadata,
SupportedCacheControls,
)
from litellm.types.utils import (LlmProviders, ProviderSpecificHeader,
StandardLoggingUserAPIKeyMetadata,
SupportedCacheControls)
service_logger_obj = ServiceLogging() # used for tracking latency on OTEL
@ -667,6 +660,12 @@ class LiteLLMProxyRequestSetup:
"user_api_key"
] = user_api_key_dict.api_key # this is just the hashed token
# Key-owned agent_id for spend attribution; keep existing (e.g. from header) if key has none
data[_metadata_variable_name]["agent_id"] = (
getattr(user_api_key_dict, "agent_id", None)
or data[_metadata_variable_name].get("agent_id")
)
data[_metadata_variable_name]["user_api_end_user_max_budget"] = getattr(
user_api_key_dict, "end_user_max_budget", None
)
@ -689,8 +688,7 @@ class LiteLLMProxyRequestSetup:
return data
from litellm.proxy._types import (
LiteLLM_ManagementEndpoint_MetadataFields,
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
)
LiteLLM_ManagementEndpoint_MetadataFields_Premium)
# ignore any special fields
added_metadata = {}
@ -859,7 +857,8 @@ async def add_litellm_data_to_request( # noqa: PLR0915
"""
from litellm.proxy.proxy_server import llm_router, premium_user
from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields
from litellm.types.proxy.litellm_pre_call_utils import (RedactedDict,
SecretFields)
_raw_headers: Dict[str, str] = RedactedDict(_safe_get_request_headers(request))
@ -1534,7 +1533,8 @@ async def move_guardrails_to_metadata(
# Only check policy engine if no local config (avoid import + registry lookup)
if not (has_key_config or has_team_config or has_request_config):
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.proxy.policy_engine.policy_registry import \
get_policy_registry
if not get_policy_registry().is_initialized():
# Nothing configured anywhere - clean up request body fields and return
@ -1598,14 +1598,16 @@ async def move_guardrails_to_metadata(
def _is_policy_version_id(s: str) -> bool:
"""Return True if string is a policy version ID (starts with policy_<uuid> prefix)."""
from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX
from litellm.proxy.policy_engine.policy_registry import \
POLICY_VERSION_ID_PREFIX
return isinstance(s, str) and s.startswith(POLICY_VERSION_ID_PREFIX)
def _extract_policy_id(s: str) -> Optional[str]:
"""Extract raw UUID from policy_<uuid> string, or None if not a valid version ID."""
from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX
from litellm.proxy.policy_engine.policy_registry import \
POLICY_VERSION_ID_PREFIX
if not _is_policy_version_id(s):
return None
@ -1626,10 +1628,9 @@ def _match_and_track_policies(
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.callback_utils import (
add_policy_sources_to_metadata,
add_policy_to_applied_policies_header,
)
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
add_policy_sources_to_metadata, add_policy_to_applied_policies_header)
from litellm.proxy.policy_engine.attachment_registry import \
get_attachment_registry
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
# Get matching policies via attachments (with match reasons for attribution)
@ -1774,7 +1775,8 @@ async def add_guardrails_from_policy_engine(
user_api_key_dict: The user's API key authentication info
"""
from litellm._logging import verbose_proxy_logger
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
from litellm.proxy.common_utils.http_parsing_utils import \
get_tags_from_request_body
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
from litellm.types.proxy.policy_engine import PolicyMatchContext

921
litellm/proxy/proxy_server.py
View file

File diff suppressed because it is too large Load diff

1
litellm/proxy/schema.prisma
View file

@ -66,6 +66,7 @@ model LiteLLM_AgentsTable {
agent_access_groups String[] @default([])
object_permission_id String?
object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id])
spend Float @default(0.0)
created_at DateTime @default(now()) @map("created_at")
created_by String
updated_at DateTime @default(now()) @updatedAt @map("updated_at")