From 78c0c1a0807a8043318b9487249173e0c9b5c19e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 4 Mar 2026 18:24:36 -0800 Subject: [PATCH] feat: update spend for agent --- .../migration.sql | 3 + .../migration.sql | 8 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/a2a_protocol/main.py | 104 +++---- litellm/batches/batch_utils.py | 7 +- litellm/fine_tuning/main.py | 57 ++-- litellm/proxy/auth/auth_checks.py | 254 ++++++++++-------- litellm/proxy/db/db_spend_update_writer.py | 9 +- .../generic_guardrail_api.py | 103 +++---- litellm/proxy/litellm_pre_call_utils.py | 58 ++-- litellm/router.py | 6 +- schema.prisma | 1 + .../proxy/db/test_db_spend_update_writer.py | 133 ++++++++- .../components/agents/agent_form_fields.tsx | 17 ++ 14 files changed, 457 insertions(+), 304 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175123_baseline_diff/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql new file mode 100644 index 00000000000..01f3936a6fc --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175123_baseline_diff/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175123_baseline_diff/migration.sql new file mode 100644 index 00000000000..8af167950ec --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175123_baseline_diff/migration.sql @@ -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"); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 6d20a71f039..3fb523524d1 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -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") diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 485b57e311b..a3951bd4276 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -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 = ( diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index a55e30ebeb9..c752e84b967 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -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) diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py index db77fa32919..ec8267e007f 100644 --- a/litellm/fine_tuning/main.py +++ b/litellm/fine_tuning/main.py @@ -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, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 6de9e6b1993..61d442a6eb6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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 diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 679ee200f7b..4419ede64ff 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -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, diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 990e7b3ede6..9c6ea46c77b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -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)}") diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9e06a987721..a30d8f1a3c8 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -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_ 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_ 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 diff --git a/litellm/router.py b/litellm/router.py index 8d44882cdf8..a858ca6cc6a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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, diff --git a/schema.prisma b/schema.prisma index f7c07112417..a6882003030 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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") diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index abd59a66a36..b5f82ef04c5 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -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( diff --git a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx index d5429d2a3b5..8e2574efec2 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_form_fields.tsx @@ -188,6 +188,23 @@ const AgentFormFields: React.FC = ({ showAgentName = true, ))} )} + + {/* Tracing */} + {shouldShow(AGENT_FORM_CONFIG.tracing.key) && ( + + {AGENT_FORM_CONFIG.tracing.fields.map((field) => ( + + + + ))} + + )} );