scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 137

feat: update spend for agent

Browse source
This commit is contained in:
Krrish Dholakia 2026-03-04 18:24:36 -08:00
parent aa7ef0802f
commit 78c0c1a080
14 changed files with 457 additions and 304 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql Normal file
View file

@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;

8
litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175123_baseline_diff/migration.sql Normal file
View file

@ -0,0 +1,8 @@
-- AlterTable
ALTER TABLE "LiteLLM_GuardrailsTable" ADD COLUMN "reviewed_at" TIMESTAMP(3),
ADD COLUMN "status" TEXT NOT NULL DEFAULT 'active',
ADD COLUMN "submitted_at" TIMESTAMP(3);
-- CreateIndex
CREATE INDEX "LiteLLM_GuardrailsTable_status_idx" ON "LiteLLM_GuardrailsTable"("status");

1
litellm-proxy-extras/litellm_proxy_extras/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")

104
litellm/a2a_protocol/main.py
View file

@ -7,7 +7,8 @@ Provides standalone functions with @client decorator for LiteLLM logging integra
import asyncio
import datetime
import uuid
from typing import TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict, Optional, Union
from typing import (TYPE_CHECKING, Any, AsyncIterator, Coroutine, Dict,
Optional, Union)
import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger
@ -15,16 +16,15 @@ from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator
from litellm.a2a_protocol.utils import A2ARequestUtils
from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.llms.custom_httpx.http_handler import (get_async_httpx_client,
httpxSpecialProvider)
from litellm.types.agents import LiteLLMSendMessageResponse
from litellm.utils import client
if TYPE_CHECKING:
from a2a.client import A2AClient as A2AClientType
from a2a.types import AgentCard, SendMessageRequest, SendStreamingMessageRequest
from a2a.types import (AgentCard, SendMessageRequest,
SendStreamingMessageRequest)
# Runtime imports with availability check
A2A_SDK_AVAILABLE = False
@ -41,9 +41,7 @@ except ImportError:
# Import our custom card resolver that supports multiple well-known paths
from litellm.a2a_protocol.card_resolver import LiteLLMA2ACardResolver
from litellm.a2a_protocol.exception_mapping_utils import (
handle_a2a_localhost_retry,
map_a2a_exception,
)
handle_a2a_localhost_retry, map_a2a_exception)
from litellm.a2a_protocol.exceptions import A2ALocalhostURLError
# Use our custom resolver instead of the default A2A SDK resolver
@ -142,9 +140,8 @@ async def _send_message_via_completion_bridge(
f"A2A using completion bridge: provider={custom_llm_provider}, api_base={api_base}"
)
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
from litellm.a2a_protocol.litellm_completion_bridge.handler import \
A2ACompletionBridgeHandler
params = (
request.params.model_dump(mode="json")
@ -162,6 +159,44 @@ async def _send_message_via_completion_bridge(
return LiteLLMSendMessageResponse.from_dict(response_dict)
async def _send_message_with_retry(
a2a_client: "A2AClientType",
request: "SendMessageRequest",
agent_card: Any,
card_url: Optional[str],
agent_name: str,
api_base: Optional[str],
) -> tuple:
a2a_response = None
for _ in range(2):
try:
a2a_response = await a2a_client.send_message(request)
break
except A2ALocalhostURLError as e:
a2a_client = handle_a2a_localhost_retry(
error=e,
agent_card=agent_card,
a2a_client=a2a_client,
is_streaming=False,
)
card_url = agent_card.url if agent_card else None
except Exception as e:
try:
map_a2a_exception(e, card_url, api_base, model=agent_name)
except A2ALocalhostURLError as localhost_err:
a2a_client = handle_a2a_localhost_retry(
error=localhost_err,
agent_card=agent_card,
a2a_client=a2a_client,
is_streaming=False,
)
card_url = agent_card.url if agent_card else None
continue
except Exception:
raise
return a2a_response, a2a_client
@client
async def asend_message(
a2a_client: Optional["A2AClientType"] = None,
@ -279,38 +314,14 @@ async def asend_message(
if getattr(message, "context_id", None) is None:
message.context_id = context_id
# Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL
a2a_response = None
for _ in range(2): # max 2 attempts: original + 1 retry
try:
a2a_response = await a2a_client.send_message(request)
break # success, exit retry loop
except A2ALocalhostURLError as e:
# Localhost URL error - fix and retry
a2a_client = handle_a2a_localhost_retry(
error=e,
agent_card=agent_card,
a2a_client=a2a_client,
is_streaming=False,
)
card_url = agent_card.url if agent_card else None
except Exception as e:
# Map exception - will raise A2ALocalhostURLError if applicable
try:
map_a2a_exception(e, card_url, api_base, model=agent_name)
except A2ALocalhostURLError as localhost_err:
# Localhost URL error - fix and retry
a2a_client = handle_a2a_localhost_retry(
error=localhost_err,
agent_card=agent_card,
a2a_client=a2a_client,
is_streaming=False,
)
card_url = agent_card.url if agent_card else None
continue
except Exception:
# Re-raise the mapped exception
raise
a2a_response, a2a_client = await _send_message_with_retry(
a2a_client=a2a_client,
request=request,
agent_card=agent_card,
card_url=card_url,
agent_name=agent_name,
api_base=api_base,
)
verbose_logger.info(f"A2A send_message completed, request_id={request.id}")
@ -477,9 +488,8 @@ async def asend_message_streaming(
f"A2A streaming using completion bridge: provider={custom_llm_provider}"
)
from litellm.a2a_protocol.litellm_completion_bridge.handler import (
A2ACompletionBridgeHandler,
)
from litellm.a2a_protocol.litellm_completion_bridge.handler import \
A2ACompletionBridgeHandler
# Extract params from request
params = (

7
litellm/batches/batch_utils.py
View file

@ -198,9 +198,8 @@ async def _get_batch_output_file_content_as_dictionary(
Required for Azure and other providers that need authentication
"""
from litellm.files.main import afile_content
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
from litellm.proxy.openai_files_endpoints.common_utils import \
_is_base64_encoded_unified_file_id
if custom_llm_provider == "vertex_ai":
raise ValueError("Vertex AI does not support file content retrieval")
@ -227,7 +226,7 @@ async def _get_batch_output_file_content_as_dictionary(
credentials = _extract_file_access_credentials(litellm_params)
file_content_kwargs.update(credentials)
_file_content = await afile_content(**file_content_kwargs)
_file_content = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType]
return _get_file_content_as_dictionary(_file_content.content)

57
litellm/fine_tuning/main.py
View file

@ -126,6 +126,18 @@ async def acreate_fine_tuning_job(
raise e
def _build_fine_tuning_job_data(model, training_file, hyperparameters, suffix, validation_file, integrations, seed):
return FineTuningJobCreate(
model=model,
training_file=training_file,
hyperparameters=hyperparameters,
suffix=suffix,
validation_file=validation_file,
integrations=integrations,
seed=seed,
)
@client
def create_fine_tuning_job(
model: str,
@ -204,19 +216,9 @@ def create_fine_tuning_job(
or os.getenv("OPENAI_API_KEY")
)
create_fine_tuning_job_data = FineTuningJobCreate(
model=model,
training_file=training_file,
hyperparameters=_oai_hyperparameters,
suffix=suffix,
validation_file=validation_file,
integrations=integrations,
seed=seed,
)
create_fine_tuning_job_data_dict = create_fine_tuning_job_data.model_dump(
exclude_none=True
)
create_fine_tuning_job_data_dict = _build_fine_tuning_job_data(
model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed,
).model_dump(exclude_none=True)
response = openai_fine_tuning_apis_instance.create_fine_tuning_job(
api_base=api_base,
@ -258,20 +260,10 @@ def create_fine_tuning_job(
# Prepare Azure-specific parameters for extra_body
extra_body = _prepare_azure_extra_body(extra_body, kwargs, azure_specific_hyperparams)
create_fine_tuning_job_data = FineTuningJobCreate(
model=model,
training_file=training_file,
hyperparameters=_oai_hyperparameters,
suffix=suffix,
validation_file=validation_file,
integrations=integrations,
seed=seed,
)
create_fine_tuning_job_data_dict = _build_fine_tuning_job_data(
model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed,
).model_dump(exclude_none=True)
create_fine_tuning_job_data_dict = create_fine_tuning_job_data.model_dump(
exclude_none=True
)
# Add extra_body if it has Azure-specific parameters
if extra_body:
create_fine_tuning_job_data_dict["extra_body"] = extra_body
@ -301,18 +293,11 @@ def create_fine_tuning_job(
vertex_credentials = optional_params.vertex_credentials or get_secret_str(
"VERTEXAI_CREDENTIALS"
)
create_fine_tuning_job_data = FineTuningJobCreate(
model=model,
training_file=training_file,
hyperparameters=_oai_hyperparameters,
suffix=suffix,
validation_file=validation_file,
integrations=integrations,
seed=seed,
)
response = vertex_fine_tuning_apis_instance.create_fine_tuning_job(
_is_async=_is_async,
create_fine_tuning_job_data=create_fine_tuning_job_data,
create_fine_tuning_job_data=_build_fine_tuning_job_data(
model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed,
),
vertex_credentials=vertex_credentials,
vertex_project=vertex_ai_project,
vertex_location=vertex_ai_location,

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

@ -11,8 +11,7 @@ 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
@ -21,29 +20,42 @@ 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.guardrails.tool_name_extraction import (
@ -253,6 +265,88 @@ async def check_tools_allowlist(
)
def _enforce_user_param_check(
general_settings: dict, request: Request, request_body: dict, route: str
) -> None:
if not general_settings.get("enforce_user_param", False):
return
http_method = request.method if hasattr(request, "method") else None
is_post_method = http_method and http_method.upper() == "POST"
is_openai_route = RouteChecks.is_llm_api_route(route=route)
is_mcp_route = (
route in LiteLLMRoutes.mcp_routes.value
or RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
)
)
if (
is_post_method
and is_openai_route
and not is_mcp_route
and "user" not in request_body
):
raise Exception(
f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}"
)
def _reject_clientside_metadata_tags_check(
general_settings: dict, request_body: dict, route: str
) -> None:
if not general_settings.get("reject_clientside_metadata_tags", False):
return
if (
RouteChecks.is_llm_api_route(route=route)
and "metadata" in request_body
and isinstance(request_body["metadata"], dict)
and "tags" in request_body["metadata"]
):
raise ProxyException(
message=f"Client-side 'metadata.tags' not allowed in request. 'reject_clientside_metadata_tags'={general_settings['reject_clientside_metadata_tags']}. Tags can only be set via API key metadata.",
type=ProxyErrorTypes.bad_request_error,
param="metadata.tags",
code=status.HTTP_400_BAD_REQUEST,
)
def _global_proxy_budget_check(
global_proxy_spend: Optional[float], skip_budget_checks: bool, route: str
) -> None:
if (
litellm.max_budget > 0
and not skip_budget_checks
and global_proxy_spend is not None
and RouteChecks.is_llm_api_route(route=route)
and route != "/v1/models"
and route != "/models"
):
if global_proxy_spend > litellm.max_budget:
raise litellm.BudgetExceededError(
current_cost=global_proxy_spend, max_budget=litellm.max_budget
)
def _guardrail_modification_check(
request_body: dict, team_object: Optional[LiteLLM_TeamTable]
) -> None:
_request_metadata: dict = request_body.get("metadata", {}) or {}
if not _request_metadata.get("guardrails"):
return
from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails
if not can_modify_guardrails(team_object):
raise HTTPException(
status_code=403,
detail={
"error": "Your team does not have permission to modify guardrails."
},
)
async def common_checks( # noqa: PLR0915
request_body: dict,
team_object: Optional[LiteLLM_TeamTable],
@ -316,10 +410,8 @@ async def common_checks( # noqa: PLR0915
# 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
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:
@ -426,84 +518,10 @@ async def common_checks( # noqa: PLR0915
message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}",
)
# 6. [OPTIONAL] If 'enforce_user_param' enabled - did developer pass in 'user' param for openai endpoints
if (
general_settings.get("enforce_user_param", None) is not None
and general_settings["enforce_user_param"] is True
):
# Get HTTP method from request
http_method = request.method if hasattr(request, "method") else None
# Check if it's a POST request and if it's an OpenAI route but not MCP
is_post_method = http_method and http_method.upper() == "POST"
is_openai_route = RouteChecks.is_llm_api_route(route=route)
is_mcp_route = (
route in LiteLLMRoutes.mcp_routes.value
or RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
)
)
# Enforce user param only for POST requests on OpenAI routes (excluding MCP routes)
if (
is_post_method
and is_openai_route
and not is_mcp_route
and "user" not in request_body
):
raise Exception(
f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}"
)
# 6.1 [OPTIONAL] If 'reject_clientside_metadata_tags' enabled - reject request if it has client-side 'metadata.tags'
if (
general_settings.get("reject_clientside_metadata_tags", None) is not None
and general_settings["reject_clientside_metadata_tags"] is True
):
if (
RouteChecks.is_llm_api_route(route=route)
and "metadata" in request_body
and isinstance(request_body["metadata"], dict)
and "tags" in request_body["metadata"]
):
raise ProxyException(
message=f"Client-side 'metadata.tags' not allowed in request. 'reject_clientside_metadata_tags'={general_settings['reject_clientside_metadata_tags']}. Tags can only be set via API key metadata.",
type=ProxyErrorTypes.bad_request_error,
param="metadata.tags",
code=status.HTTP_400_BAD_REQUEST,
)
# 7. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget
if (
litellm.max_budget > 0
and not skip_budget_checks
and global_proxy_spend is not None
# only run global budget checks for OpenAI routes
# Reason - the Admin UI should continue working if the proxy crosses it's global budget
and RouteChecks.is_llm_api_route(route=route)
and route != "/v1/models"
and route != "/models"
):
if global_proxy_spend > litellm.max_budget:
raise litellm.BudgetExceededError(
current_cost=global_proxy_spend, max_budget=litellm.max_budget
)
_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
can_modify: bool = can_modify_guardrails(team_object)
if can_modify is False:
from fastapi import HTTPException
raise HTTPException(
status_code=403,
detail={
"error": "Your team does not have permission to modify guardrails."
},
)
_enforce_user_param_check(general_settings, request, request_body, route)
_reject_clientside_metadata_tags_check(general_settings, request_body, route)
_global_proxy_budget_check(global_proxy_spend, skip_budget_checks, route)
_guardrail_modification_check(request_body, team_object)
# 10 [OPTIONAL] Organization RBAC checks
organization_role_based_access_check(
@ -1942,8 +1960,9 @@ 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")
@ -1989,8 +2008,9 @@ 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")
@ -2029,8 +2049,9 @@ 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"
@ -2346,10 +2367,8 @@ 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
@ -3305,8 +3324,7 @@ 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

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

@ -374,10 +374,11 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
_agent_id_for_spend = payload_copy.get("agent_id")
try:
await self._update_agent_db(
response_cost=response_cost,
agent_id=payload_copy.get("agent_id"),
agent_id=_agent_id_for_spend,
prisma_client=prisma_client,
)
except Exception:
@ -624,9 +625,6 @@ class DBSpendUpdateWriter:
):
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(
@ -2092,9 +2090,6 @@ class DBSpendUpdateWriter:
)
return
if payload["agent_id"] is None:
verbose_proxy_logger.debug(
"agent_id is None for request. Skipping incrementing agent spend."
)
return
payload_with_agent_id = cast(
SpendLogsPayload,

103
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py
View file

@ -14,24 +14,19 @@ import httpx
from litellm._logging import verbose_proxy_logger
from litellm._version import version as litellm_version
from litellm.exceptions import GuardrailRaisedException, Timeout
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.integrations.custom_guardrail import (CustomGuardrail,
log_guardrail_information)
from litellm.llms.custom_httpx.http_handler import (get_async_httpx_client,
httpxSpecialProvider)
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.proxy.guardrails.guardrail_hooks.generic_guardrail_api import (
GenericGuardrailAPIMetadata,
GenericGuardrailAPIRequest,
GenericGuardrailAPIResponse,
)
GenericGuardrailAPIMetadata, GenericGuardrailAPIRequest,
GenericGuardrailAPIResponse)
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.litellm_logging import \
Logging as LiteLLMLoggingObj
GUARDRAIL_NAME = "generic_guardrail_api"
@ -334,6 +329,30 @@ class GenericGuardrailAPI(CustomGuardrail):
return_inputs["tools"] = tools
return return_inputs
def _handle_guardrail_request_error(
self,
error: Exception,
inputs: GenericGuardrailAPIInputs,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"],
is_unreachable: bool = True,
) -> GenericGuardrailAPIInputs:
if is_unreachable and self.unreachable_fallback == "fail_open":
http_status_code = getattr(
getattr(error, "response", None), "status_code", None
)
return self._fail_open_passthrough(
inputs=inputs,
input_type=input_type,
logging_obj=logging_obj,
error=error,
**({"http_status_code": http_status_code} if http_status_code else {}),
)
verbose_proxy_logger.error(
"Generic Guardrail API: failed to make request: %s", str(error)
)
raise Exception(f"Generic Guardrail API failed: {str(error)}")
@log_guardrail_information
async def apply_guardrail(
self,
@ -462,58 +481,24 @@ class GenericGuardrailAPI(CustomGuardrail):
)
except GuardrailRaisedException:
# Re-raise guardrail exceptions as-is
raise
except Timeout as e:
# AsyncHTTPHandler wraps httpx.TimeoutException into litellm.Timeout
if self.unreachable_fallback == "fail_open":
return self._fail_open_passthrough(
inputs=inputs,
input_type=input_type,
logging_obj=logging_obj,
error=e,
)
verbose_proxy_logger.error(
"Generic Guardrail API: failed to make request: %s", str(e)
return self._handle_guardrail_request_error(
e, inputs, input_type, logging_obj
)
raise Exception(f"Generic Guardrail API failed: {str(e)}")
except httpx.HTTPStatusError as e:
# Common reverse-proxy/LB failures can present as HTTP errors even when the backend is unreachable.
status_code = getattr(getattr(e, "response", None), "status_code", None)
if self.unreachable_fallback == "fail_open" and status_code in (
502,
503,
504,
):
return self._fail_open_passthrough(
inputs=inputs,
input_type=input_type,
logging_obj=logging_obj,
error=e,
http_status_code=status_code,
)
verbose_proxy_logger.error(
"Generic Guardrail API: failed to make request: %s", str(e)
status_code = getattr(
getattr(e, "response", None), "status_code", None
)
is_unreachable = status_code in (502, 503, 504)
return self._handle_guardrail_request_error(
e, inputs, input_type, logging_obj, is_unreachable=is_unreachable
)
raise Exception(f"Generic Guardrail API failed: {str(e)}")
except httpx.RequestError as e:
# Guardrail endpoint is unreachable (DNS/connect/timeout/etc)
if self.unreachable_fallback == "fail_open":
return self._fail_open_passthrough(
inputs=inputs,
input_type=input_type,
logging_obj=logging_obj,
error=e,
)
verbose_proxy_logger.error(
"Generic Guardrail API: failed to make request: %s", str(e)
return self._handle_guardrail_request_error(
e, inputs, input_type, logging_obj
)
raise Exception(f"Generic Guardrail API failed: {str(e)}")
except Exception as e:
verbose_proxy_logger.error(
"Generic Guardrail API: failed to make request: %s", str(e)
return self._handle_guardrail_request_error(
e, inputs, input_type, logging_obj, is_unreachable=False
)
raise Exception(f"Generic Guardrail API failed: {str(e)}")

58
litellm/proxy/litellm_pre_call_utils.py
View file

@ -10,12 +10,16 @@ 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(
@ -24,9 +28,12 @@ _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
@ -661,10 +668,10 @@ class LiteLLMProxyRequestSetup:
] = 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")
)
_key_agent_id = getattr(user_api_key_dict, "agent_id", None)
_existing_agent_id = data[_metadata_variable_name].get("agent_id")
_resolved_agent_id = _key_agent_id or _existing_agent_id
data[_metadata_variable_name]["agent_id"] = _resolved_agent_id
data[_metadata_variable_name]["user_api_end_user_max_budget"] = getattr(
user_api_key_dict, "end_user_max_budget", None
@ -688,7 +695,8 @@ 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 = {}
@ -857,8 +865,7 @@ 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))
@ -1542,8 +1549,7 @@ 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
@ -1607,16 +1613,14 @@ 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
@ -1637,9 +1641,10 @@ 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)
@ -1784,8 +1789,7 @@ 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

6
litellm/router.py
View file

@ -164,11 +164,7 @@ from litellm.types.utils import (
)
from litellm.types.utils import ModelInfo
from litellm.types.utils import ModelInfo as ModelMapInfo
from litellm.types.utils import (
ModelResponseStream,
StandardLoggingPayload,
Usage,
)
from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage
from litellm.utils import (
CustomStreamWrapper,
EmbeddingResponse,

1
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")

133
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -9,7 +9,7 @@ sys.path.insert(
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch, call
from unittest.mock import AsyncMock, MagicMock, call, patch
import pytest
@ -516,6 +516,131 @@ async def test_update_tag_db_without_prisma_client():
assert writer.spend_update_queue.add_update.call_count == 0
@pytest.mark.asyncio
async def test_update_agent_db_enqueues_agent_spend():
"""
Test that _update_agent_db enqueues a SpendUpdateQueueItem with entity_type=AGENT.
"""
from litellm.proxy._types import Litellm_EntityType
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
agent_id = "agent-123"
response_cost = 0.1
writer.spend_update_queue.add_update = AsyncMock()
await writer._update_agent_db(
response_cost=response_cost,
agent_id=agent_id,
prisma_client=mock_prisma,
)
writer.spend_update_queue.add_update.assert_called_once()
call_args = writer.spend_update_queue.add_update.call_args[1]
assert call_args["update"]["entity_type"] == Litellm_EntityType.AGENT
assert call_args["update"]["entity_id"] == agent_id
assert call_args["update"]["response_cost"] == response_cost
@pytest.mark.asyncio
async def test_update_agent_db_skips_when_agent_id_none():
"""_update_agent_db does not enqueue when agent_id is None."""
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
writer.spend_update_queue.add_update = AsyncMock()
await writer._update_agent_db(
response_cost=0.05,
agent_id=None,
prisma_client=mock_prisma,
)
writer.spend_update_queue.add_update.assert_not_called()
@pytest.mark.asyncio
async def test_update_agent_db_skips_when_prisma_client_none():
"""_update_agent_db does not enqueue when prisma_client is None."""
writer = DBSpendUpdateWriter()
writer.spend_update_queue.add_update = AsyncMock()
await writer._update_agent_db(
response_cost=0.05,
agent_id="agent-456",
prisma_client=None,
)
writer.spend_update_queue.add_update.assert_not_called()
@pytest.mark.asyncio
async def test_commit_spend_updates_to_db_increments_agent_spend():
"""
Test that _commit_spend_updates_to_db calls litellm_agentstable.update_many
with spend increment when agent_list_transactions is present.
"""
db_writer = DBSpendUpdateWriter()
mock_batcher = MagicMock()
mock_batcher.litellm_verificationtoken = MagicMock()
mock_batcher.litellm_verificationtoken.update_many = MagicMock()
mock_batcher.litellm_usertable = MagicMock()
mock_batcher.litellm_usertable.update_many = MagicMock()
mock_batcher.litellm_teamtable = MagicMock()
mock_batcher.litellm_teamtable.update_many = MagicMock()
mock_batcher.litellm_teammembership = MagicMock()
mock_batcher.litellm_teammembership.update_many = MagicMock()
mock_batcher.litellm_organizationtable = MagicMock()
mock_batcher.litellm_organizationtable.update_many = MagicMock()
mock_batcher.litellm_tagtable = MagicMock()
mock_batcher.litellm_tagtable.update_many = MagicMock()
mock_batcher.litellm_agentstable = MagicMock()
mock_batcher.litellm_agentstable.update_many = MagicMock()
mock_transaction = AsyncMock()
mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
mock_transaction.__aexit__ = AsyncMock(return_value=False)
mock_transaction.batch_ = MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=mock_batcher),
__aexit__=AsyncMock(return_value=False),
)
)
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
mock_proxy_logging = MagicMock()
agent_id = "agent-789"
response_cost = 0.25
db_spend_update_transactions = {
"user_list_transactions": {},
"end_user_list_transactions": {},
"key_list_transactions": {},
"team_list_transactions": {},
"team_member_list_transactions": {},
"org_list_transactions": {},
"tag_list_transactions": {},
"agent_list_transactions": {agent_id: response_cost},
}
with patch("litellm.proxy.utils._raise_failed_update_spend_exception"):
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=mock_proxy_logging,
db_spend_update_transactions=db_spend_update_transactions,
)
mock_batcher.litellm_agentstable.update_many.assert_called_once()
call_kwargs = mock_batcher.litellm_agentstable.update_many.call_args[1]
assert call_kwargs["where"] == {"agent_id": agent_id}
assert call_kwargs["data"] == {"spend": {"increment": response_cost}}
@pytest.mark.asyncio
async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
"""
@ -1048,6 +1173,8 @@ async def test_commit_key_spend_updates_includes_last_active():
mock_batcher.litellm_teamtable.update_many = MagicMock()
mock_batcher.litellm_organizationtable = MagicMock()
mock_batcher.litellm_organizationtable.update_many = MagicMock()
mock_batcher.litellm_agentstable = MagicMock()
mock_batcher.litellm_agentstable.update_many = MagicMock()
mock_proxy_logging = MagicMock()
@ -1059,6 +1186,7 @@ async def test_commit_key_spend_updates_includes_last_active():
"team_member_list_transactions": {},
"org_list_transactions": {},
"tag_list_transactions": {},
"agent_list_transactions": {},
}
before_call = datetime.now(timezone.utc)
@ -1142,6 +1270,7 @@ async def test_batch_database_updates_isolation_on_failure():
db_writer._update_team_db = AsyncMock()
db_writer._update_org_db = AsyncMock()
db_writer._update_tag_db = AsyncMock()
db_writer._update_agent_db = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock()
@ -1169,6 +1298,7 @@ async def test_batch_database_updates_isolation_on_failure():
db_writer._update_team_db.assert_awaited_once()
db_writer._update_org_db.assert_awaited_once()
db_writer._update_tag_db.assert_awaited_once()
db_writer._update_agent_db.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_user_transaction.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_end_user_transaction.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_agent_transaction.assert_awaited_once()
@ -1203,6 +1333,7 @@ async def test_daily_agent_receives_deepcopied_payload():
db_writer._update_team_db = AsyncMock()
db_writer._update_org_db = AsyncMock()
db_writer._update_tag_db = AsyncMock()
db_writer._update_agent_db = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock(

17
ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx
View file

@ -188,6 +188,23 @@ const AgentFormFields: React.FC<AgentFormFieldsProps> = ({ showAgentName = true,
))}
</Panel>
)}
{/* Tracing */}
{shouldShow(AGENT_FORM_CONFIG.tracing.key) && (
<Panel header={AGENT_FORM_CONFIG.tracing.title} key={AGENT_FORM_CONFIG.tracing.key}>
{AGENT_FORM_CONFIG.tracing.fields.map((field) => (
<Form.Item
key={field.name}
label={field.label}
name={field.name}
valuePropName="checked"
tooltip={field.tooltip}
>
<Switch />
</Form.Item>
))}
</Panel>
)}
</Collapse>
</>
);