From cf439c269c77909a02fba3cdcd827c2e4bed1eaa Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Sat, 7 Mar 2026 19:12:42 -0800 Subject: [PATCH] Agents - add max budget + tpm/rpm limiting per agent AND per agent session (#22849) * feat: enforce x-litellm-trace-id in header, if required * feat: update spend for agent * refactor: update agent table to follow similar format as other entities - also add a spend column - allows us to see spend of an agent * fix: cleanup ui * feat: return spend on agent endpoints * feat: scope pr * feat(agents/): support budgets + rate limiting on agents + agent sessions * fix: address PR review feedback - Add missing tpm_limit, rpm_limit, session_tpm_limit, session_rpm_limit columns to root schema.prisma to match proxy and extras schemas - Add backwards-compatible fallback to key metadata for max_iterations so existing users don't silently lose enforcement Co-Authored-By: Claude Opus 4.6 * fix: qa'ed RPM limiting on agents --------- Co-authored-by: Claude Opus 4.6 --- docs/my-website/docs/a2a.md | 1 + docs/my-website/docs/a2a_iteration_budgets.md | 188 +++++++ docs/my-website/docs/proxy/users.md | 130 +++++ docs/my-website/sidebars.js | 3 +- .../migration.sql | 3 + .../migration.sql | 5 + .../litellm_proxy_extras/schema.prisma | 5 + litellm/batches/batch_utils.py | 7 +- litellm/fine_tuning/main.py | 57 +- litellm/proxy/_types.py | 2 + .../proxy/agent_endpoints/a2a_endpoints.py | 57 +- .../proxy/agent_endpoints/agent_registry.py | 20 +- litellm/proxy/agent_endpoints/endpoints.py | 54 +- litellm/proxy/auth/auth_checks.py | 318 ++++++----- litellm/proxy/auth/user_api_key_auth.py | 148 +++-- litellm/proxy/db/db_spend_update_writer.py | 146 +++-- .../redis_update_buffer.py | 50 +- .../spend_update_queue.py | 17 +- .../generic_guardrail_api.py | 82 ++- litellm/proxy/hooks/__init__.py | 4 + .../hooks/max_budget_per_session_limiter.py | 271 +++++++++ litellm/proxy/hooks/max_iterations_limiter.py | 27 +- .../hooks/parallel_request_limiter_v3.py | 177 ++++-- litellm/proxy/litellm_pre_call_utils.py | 6 + litellm/proxy/proxy_server.py | 28 +- litellm/proxy/schema.prisma | 5 + litellm/proxy/utils.py | 2 +- litellm/router.py | 6 +- litellm/types/agents.py | 13 + schema.prisma | 5 + .../proxy/db/test_db_spend_update_writer.py | 133 ++++- .../test_max_budget_per_session_limiter.py | 165 ++++++ .../hooks/test_max_iterations_limiter.py | 158 ++++-- .../hooks/test_parallel_request_limiter_v3.py | 521 ++++++++++++++++++ .../src/components/agents.tsx | 115 +++- .../src/components/agents/add_agent_form.tsx | 144 ++++- .../src/components/agents/agent_config.ts | 22 + .../src/components/agents/agent_info.tsx | 25 +- .../agents/dynamic_agent_form_fields.tsx | 9 +- .../src/components/agents/types.ts | 5 + .../src/components/networking.tsx | 19 +- 41 files changed, 2596 insertions(+), 557 deletions(-) create mode 100644 docs/my-website/docs/a2a_iteration_budgets.md 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/20260305000000_add_rate_limits_to_agents/migration.sql create mode 100644 litellm/proxy/hooks/max_budget_per_session_limiter.py create mode 100644 tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py diff --git a/docs/my-website/docs/a2a.md b/docs/my-website/docs/a2a.md index b1166a7809..9c86d0de38 100644 --- a/docs/my-website/docs/a2a.md +++ b/docs/my-website/docs/a2a.md @@ -20,6 +20,7 @@ Add A2A Agents on LiteLLM AI Gateway, Invoke agents in A2A Protocol, track reque | Logging | ✅ | | Load Balancing | ✅ | | Streaming | ✅ | +| [Iteration Budgets](a2a_iteration_budgets) | ✅ | :::tip diff --git a/docs/my-website/docs/a2a_iteration_budgets.md b/docs/my-website/docs/a2a_iteration_budgets.md new file mode 100644 index 0000000000..47beca3470 --- /dev/null +++ b/docs/my-website/docs/a2a_iteration_budgets.md @@ -0,0 +1,188 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Agent Iteration Budgets + +Control runaway costs from agentic loops with per-session iteration and budget caps. + +## Overview + +When agents run agentic loops, they can make unbounded LLM calls, causing unexpected costs. LiteLLM provides two controls: + +| Control | Description | +|---------|-------------| +| **Max Iterations** | Hard cap on the number of LLM calls per session | +| **Max Budget Per Session** | Dollar cap per session (identified by `x-litellm-trace-id`) | + +Both controls require a `session_id` (sent via `x-litellm-trace-id` header or `metadata.session_id`) to track calls within a session. + +## Trace-ID Enforcement + +LiteLLM supports two independent trace-id flags, configured in `litellm_params` on the agent: + +| Flag | Description | +|------|-------------| +| `require_trace_id_on_calls_to_agent` | Requires callers invoking this agent to include `x-litellm-trace-id`. Use when the agent should only be called as a sub-agent with a trace context. Returns **400** if missing. | +| `require_trace_id_on_calls_by_agent` | Requires all LLM/MCP calls made **by** this agent (via its virtual key) to include `x-litellm-trace-id`. This is what enables `max_iterations` and `max_budget_per_session` tracking. Returns **400** if missing. | + +## Configuring via UI + +When creating an agent in the LiteLLM Admin UI: + +1. Navigate to the **Agents** tab and click **Add Agent** +2. In the **Agent Settings** step, expand the **Tracing** section +3. Toggle **Require x-litellm-trace-id on calls BY this agent** to enable session tracking +4. Set **Max Iterations** to cap the number of LLM calls per session +5. Set **Max Budget Per Session ($)** to cap spend per session + +The trace-id flags are stored on the agent's `litellm_params`. Budget controls (`max_iterations`, `max_budget_per_session`) are stored in the virtual key's metadata. + +## Configuring via API + +Set trace-id enforcement on the agent itself: + +```bash +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent with budget controls", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "litellm_params": { + "require_trace_id_on_calls_to_agent": true, + "require_trace_id_on_calls_by_agent": true + } + }' +``` + +Budget controls are set on the agent's `litellm_params` (not on individual keys), so they apply across all keys for the agent: + +```bash +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent with budget controls", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "litellm_params": { + "require_trace_id_on_calls_by_agent": true, + "max_iterations": 25, + "max_budget_per_session": 5.00 + } + }' +``` + +## How It Works + +### Session Tracking + +Callers identify their session by including a `session_id` in one of these ways: +- **Header**: `x-litellm-trace-id: my-session-123` +- **Metadata**: `{"metadata": {"session_id": "my-session-123"}}` + +### Max Iterations + +When `max_iterations` is set in agent `litellm_params`: +- Each LLM call for a session increments a counter +- When the counter exceeds `max_iterations`, the request receives a **429 Too Many Requests** +- Counters expire after 1 hour by default (configurable via `LITELLM_MAX_ITERATIONS_TTL` env var) + +### Max Budget Per Session + +When `max_budget_per_session` is set in agent `litellm_params`: +- After each successful LLM call, the response cost is accumulated for the session +- Before each call, the accumulated spend is checked against the budget +- When spend exceeds the budget, the request receives a **429 Too Many Requests** +- Session spend counters expire after 1 hour by default (configurable via `LITELLM_MAX_BUDGET_PER_SESSION_TTL` env var) + +## Example + +Create an agent with max 25 iterations and a $5 budget cap: + + + + +1. Go to **Agents** → **Add Agent** +2. Configure your agent (name, model, etc.) +3. In **Agent Settings**, expand the **Tracing** section +4. Toggle on **Require x-litellm-trace-id on calls BY this agent** +5. Set **Max Iterations** to `25` +6. Set **Max Budget Per Session** to `5.00` +7. Proceed to create a new key for the agent +8. Click **Create Agent** + + + + +```bash +# 1. Create the agent with trace-id enforcement +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent with budget controls", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "litellm_params": { + "require_trace_id_on_calls_by_agent": true + } + }' + +# 2. Create a key for the agent +curl -X POST 'http://localhost:4000/key/generate' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_id": "", + "key_alias": "my-research-agent-key" + }' +``` + + + + +### Making Calls with Session Tracking + +```bash +curl -X POST 'http://localhost:4000/chat/completions' \ + -H 'Authorization: Bearer sk-agent-key-xxx' \ + -H 'x-litellm-trace-id: session-abc-123' \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Hello"}] + }' +``` + +After 25 calls or $5 spent within this session, subsequent requests will receive: + +```json +{ + "error": { + "message": "Session budget exceeded for session session-abc-123. Current spend: $5.0032, max_budget_per_session: $5.00.", + "type": "budget_exceeded", + "code": 429 + } +} +``` + +## Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `LITELLM_MAX_ITERATIONS_TTL` | `3600` (1 hour) | TTL in seconds for session iteration counters | +| `LITELLM_MAX_BUDGET_PER_SESSION_TTL` | `3600` (1 hour) | TTL in seconds for session budget counters | diff --git a/docs/my-website/docs/proxy/users.md b/docs/my-website/docs/proxy/users.md index 8517db51a8..58813eaf49 100644 --- a/docs/my-website/docs/proxy/users.md +++ b/docs/my-website/docs/proxy/users.md @@ -10,6 +10,8 @@ import TabItem from '@theme/TabItem'; **Team member budgets**: Set individual spending limits within the team's shared budget +**Agent budgets**: Set rate limits (tpm/rpm) and session-level caps (iterations, dollar budget) on agents [**Jump**](#agents) + ***If a key belongs to a team, the team budget is applied, not the user's personal budget.*** ::: @@ -420,6 +422,109 @@ Expected response on failure +### Agents + +Set budgets and rate limits on agents registered with LiteLLM's [Agent Gateway](../a2a.md). You can control: +- **Per-agent rate limits**: `tpm_limit` and `rpm_limit` on the agent itself +- **Per-session rate limits**: `session_tpm_limit` and `session_rpm_limit` applied per session +- **Per-session iteration cap**: `max_iterations` in agent `litellm_params` +- **Per-session budget cap**: `max_budget_per_session` in agent `litellm_params` + + + + +Set `tpm_limit` and `rpm_limit` on the agent to cap total throughput across all sessions. + +```bash +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "tpm_limit": 100000, + "rpm_limit": 100 + }' +``` + + + + +Set `session_tpm_limit` and `session_rpm_limit` to cap throughput per individual session. + +```bash +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "session_tpm_limit": 50000, + "session_rpm_limit": 50 + }' +``` + + + + +Set `max_iterations` and `max_budget_per_session` in agent `litellm_params` to cap individual sessions. Requires `require_trace_id_on_calls_by_agent` so LiteLLM can track calls per session. + +```bash +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "litellm_params": { + "require_trace_id_on_calls_by_agent": true, + "max_iterations": 25, + "max_budget_per_session": 5.00 + } + }' +``` + +When a session exceeds the limit, requests receive a **429 Too Many Requests** response. + +See the [Agent Iteration Budgets](../a2a_iteration_budgets) guide for full details. + + + + +:::info + +You can also update rate limits on existing agents using `PATCH /v1/agents/{agent_id}`: + +```bash +curl -X PATCH 'http://localhost:4000/v1/agents/' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "tpm_limit": 200000, + "rpm_limit": 200, + "session_tpm_limit": 50000, + "session_rpm_limit": 50 + }' +``` + +::: + + ### Customers Use this to budget `user` passed to `/chat/completions`, **without needing to create a key for every user** @@ -685,6 +790,31 @@ These headers indicate: - 1 request remaining for the GPT-4 model for key=`sk-ulGNRXWtv7M0lFnnsQk0wQ` - 179 tokens remaining for the GPT-4 model for key=`sk-ulGNRXWtv7M0lFnnsQk0wQ` + + + +Set rate limits on agents registered with the [Agent Gateway](../a2a.md). + +**Agent-level limits** cap total throughput across all sessions: + +```shell +curl -X POST 'http://0.0.0.0:4000/v1/agents' \ +--header 'Authorization: Bearer sk-1234' \ +--header 'Content-Type: application/json' \ +--data '{"agent_name": "my-agent", "agent_card_params": {"name": "my-agent", "description": "My agent", "url": "http://my-agent:8080", "version": "1.0.0"}, "tpm_limit": 100000, "rpm_limit": 100}' +``` + +**Session-level limits** cap throughput per individual session: + +```shell +curl -X POST 'http://0.0.0.0:4000/v1/agents' \ +--header 'Authorization: Bearer sk-1234' \ +--header 'Content-Type: application/json' \ +--data '{"agent_name": "my-agent", "agent_card_params": {"name": "my-agent", "description": "My agent", "url": "http://my-agent:8080", "version": "1.0.0"}, "session_tpm_limit": 50000, "session_rpm_limit": 50}' +``` + +You can also set **max_iterations** (call count cap) and **max_budget_per_session** (dollar cap) per session via `litellm_params`. See [Agent Iteration Budgets](../a2a_iteration_budgets) for details. + diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index a77e905724..ef2df2d8ad 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -542,7 +542,8 @@ const sidebars = { "a2a_invoking_agents", "a2a_agent_headers", "a2a_cost_tracking", - "a2a_agent_permissions" + "a2a_agent_permissions", + "a2a_iteration_budgets" ], }, "assistants", 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 0000000000..01f3936a6f --- /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/20260305000000_add_rate_limits_to_agents/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql new file mode 100644 index 0000000000..3cd8ca638a --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "tpm_limit" INTEGER; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "rpm_limit" INTEGER; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "session_tpm_limit" INTEGER; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "session_rpm_limit" INTEGER; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 329ff80933..8d4bdffb2d 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -68,6 +68,11 @@ 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) + tpm_limit Int? + rpm_limit Int? + session_tpm_limit Int? + session_rpm_limit Int? created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index a55e30ebeb..c752e84b96 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 e6f2a6f86d..93fa56ff97 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, + ) + + def _resolve_fine_tuning_timeout( timeout: Any, custom_llm_provider: str, @@ -206,19 +218,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, @@ -260,20 +262,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 @@ -303,18 +295,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/_types.py b/litellm/proxy/_types.py index e6b7b7d285..b3e9a8ea24 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -202,6 +202,7 @@ class Litellm_EntityType(enum.Enum): ORGANIZATION = "organization" PROJECT = "project" TAG = "tag" + AGENT = "agent" # global proxy level entity PROXY = "proxy" @@ -4228,6 +4229,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): diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 3ffced7a12..63e0dad332 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -39,7 +39,8 @@ def _jsonrpc_error( def _get_agent(agent_id: str): """Look up an agent by ID or name. Returns None if not found.""" - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry agent = global_agent_registry.get_agent_by_id(agent_id=agent_id) if agent is None: @@ -47,6 +48,26 @@ def _get_agent(agent_id: str): return agent +def _enforce_inbound_trace_id(agent: Any, request: Request) -> None: + """Raise 400 if agent requires x-litellm-trace-id on inbound calls and it is missing.""" + agent_litellm_params = agent.litellm_params or {} + if not agent_litellm_params.get("require_trace_id_on_calls_to_agent"): + return + + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + headers_dict = dict(request.headers) + trace_id = get_chain_id_from_headers(headers_dict) + if not trace_id: + raise HTTPException( + status_code=400, + detail=( + f"Agent '{agent.agent_id}' requires x-litellm-trace-id header " + "on all inbound requests." + ), + ) + + async def _handle_stream_message( api_base: Optional[str], request_id: str, @@ -116,9 +137,8 @@ async def _handle_stream_message( and request_data is not None and proxy_logging_obj is not None ): - from litellm.proxy.common_request_processing import ( - ProxyBaseLLMRequestProcessing, - ) + from litellm.proxy.common_request_processing import \ + ProxyBaseLLMRequestProcessing def _ndjson_chunk(chunk: Any) -> str: if hasattr(chunk, "model_dump"): @@ -218,9 +238,8 @@ async def get_agent_card( The URL in the agent card is rewritten to point to the LiteLLM proxy, so all subsequent A2A calls go through LiteLLM for logging and cost tracking. """ - from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( - AgentRequestHandler, - ) + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import \ + AgentRequestHandler try: agent = _get_agent(agent_id) @@ -284,15 +303,10 @@ async def invoke_agent_a2a( # noqa: PLR0915 """ from litellm.a2a_protocol import asend_message from litellm.a2a_protocol.main import A2A_SDK_AVAILABLE - from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( - AgentRequestHandler, - ) - from litellm.proxy.proxy_server import ( - general_settings, - proxy_config, - proxy_logging_obj, - version, - ) + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import \ + AgentRequestHandler + from litellm.proxy.proxy_server import (general_settings, proxy_config, + proxy_logging_obj, version) body = {} try: @@ -345,6 +359,8 @@ async def invoke_agent_a2a( # noqa: PLR0915 detail=f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access.", ) + _enforce_inbound_trace_id(agent, request) + # Get backend URL and agent name agent_url = agent.agent_card_params.get("url") agent_name = agent.agent_card_params.get("name", agent_id) @@ -365,6 +381,10 @@ async def invoke_agent_a2a( # noqa: PLR0915 ) # Set up data dict for litellm processing + if "metadata" not in body: + body["metadata"] = {} + body["metadata"]["agent_id"] = agent.agent_id + body.update( { "model": f"a2a_agent/{agent_name}", @@ -373,9 +393,8 @@ async def invoke_agent_a2a( # noqa: PLR0915 ) # Add litellm data (user_api_key, user_id, team_id, etc.) - from litellm.proxy.common_request_processing import ( - ProxyBaseLLMRequestProcessing, - ) + from litellm.proxy.common_request_processing import \ + ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=body) data, logging_obj = await processor.common_processing_pre_call_logic( diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 91bbbd73d1..ce6b1055ee 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -5,9 +5,8 @@ from typing import Any, Dict, List, Optional import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy.management_helpers.object_permission_utils import ( - handle_update_object_permission_common, -) +from litellm.proxy.management_helpers.object_permission_utils import \ + handle_update_object_permission_common from litellm.proxy.utils import PrismaClient from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -152,6 +151,11 @@ class AgentRegistry: if object_permission_id is not None: create_data["object_permission_id"] = object_permission_id + for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): + _val = agent.get(rate_field) + if _val is not None: + create_data[rate_field] = _val + # Create agent in DB created_agent = await prisma_client.db.litellm_agentstable.create( data=create_data, @@ -226,6 +230,10 @@ class AgentRegistry: update_data["agent_card_params"] = safe_dumps( augment_agent.get("agent_card_params") ) + + for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): + if rate_field in agent: + update_data[rate_field] = agent.get(rate_field) if "static_headers" in agent: headers_value = agent.get("static_headers") update_data["static_headers"] = safe_dumps( @@ -321,6 +329,12 @@ class AgentRegistry: "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), } + + for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): + _val = agent.get(rate_field) + if _val is not None: + update_data[rate_field] = _val + if agent.get("object_permission") is not None: existing_agent = await prisma_client.db.litellm_agentstable.find_unique( where={"agent_id": agent_id} diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index eef969951a..252db29a9d 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -135,7 +135,7 @@ async def get_agents( try: returned_agents: List[AgentResponse] = [] - + # Admin users get all agents if ( user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN @@ -147,7 +147,7 @@ async def get_agents( allowed_agent_ids = await AgentRequestHandler.get_allowed_agents( user_api_key_auth=user_api_key_dict ) - + # If no restrictions (empty list), return all agents if len(allowed_agent_ids) == 0: returned_agents = global_agent_registry.get_agent_list() @@ -155,10 +155,23 @@ async def get_agents( # Filter agents by allowed IDs all_agents = global_agent_registry.get_agent_list() returned_agents = [ - agent for agent in all_agents - if agent.agent_id in allowed_agent_ids + agent for agent in all_agents if agent.agent_id in allowed_agent_ids ] + # Fetch current spend from DB for all returned agents + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is not None: + agent_ids = [agent.agent_id for agent in returned_agents] + if agent_ids: + db_agents = await prisma_client.db.litellm_agentstable.find_many( + where={"agent_id": {"in": agent_ids}}, + ) + spend_map = {a.agent_id: a.spend for a in db_agents} + for agent in returned_agents: + if agent.agent_id in spend_map: + agent.spend = spend_map[agent.agent_id] + # add is_public field to each agent - we do it this way, to allow setting config agents as public for agent in returned_agents: if agent.litellm_params is None: @@ -222,9 +235,8 @@ async def get_agents( #### CRUD ENDPOINTS FOR AGENTS #### -from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry as AGENT_REGISTRY, -) +from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry as AGENT_REGISTRY @router.post( @@ -363,10 +375,21 @@ async def get_agent_by_id( agent_dict = agent_row.model_dump() if agent_row.object_permission is not None: try: - agent_dict["object_permission"] = agent_row.object_permission.model_dump() + agent_dict["object_permission"] = ( + agent_row.object_permission.model_dump() + ) except Exception: - agent_dict["object_permission"] = agent_row.object_permission.dict() + agent_dict["object_permission"] = ( + agent_row.object_permission.dict() + ) agent = AgentResponse(**agent_dict) # type: ignore + else: + # Agent found in memory — refresh spend from DB + db_row = await prisma_client.db.litellm_agentstable.find_unique( + where={"agent_id": agent_id} + ) + if db_row is not None: + agent.spend = db_row.spend if agent is None: raise HTTPException( @@ -674,9 +697,8 @@ async def make_agent_public( try: # Update the public model groups import litellm - from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry as AGENT_REGISTRY, - ) + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry as AGENT_REGISTRY from litellm.proxy.proxy_server import proxy_config # Check if user has admin permissions @@ -791,9 +813,8 @@ async def make_agents_public( try: # Update the public model groups import litellm - from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry as AGENT_REGISTRY, - ) + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry as AGENT_REGISTRY from litellm.proxy.proxy_server import proxy_config # Load existing config @@ -853,6 +874,7 @@ async def make_agents_public( verbose_proxy_logger.exception(f"Error making agent public: {e}") raise HTTPException(status_code=500, detail=str(e)) + @router.get( "/agent/daily/activity", tags=["Agent Management"], @@ -914,4 +936,4 @@ async def get_agent_daily_activity( api_key=api_key, page=page, page_size=page_size, - ) \ No newline at end of file + ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index eb690931fb..db794c5ac3 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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,48 +21,33 @@ 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 ( - TOOL_CAPABLE_CALL_TYPES, - extract_request_tool_names, -) + TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names) from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.router import Router @@ -224,6 +210,89 @@ async def _run_project_checks( ) +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 check_tools_allowlist( request_body: dict, valid_token: Optional[UserAPIKeyAuth], @@ -235,23 +304,34 @@ async def check_tools_allowlist( effective allowlist is read from valid_token.metadata and valid_token.team_metadata. Raises ProxyException with tool_access_denied if a tool is not allowed. """ - from litellm.litellm_core_utils.api_route_to_call_types import ( - get_call_types_for_route, - ) + from litellm.litellm_core_utils.api_route_to_call_types import \ + get_call_types_for_route if valid_token is None: return call_types = get_call_types_for_route(route) - if not call_types or not any(ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types): + if not call_types or not any( + ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types + ): return tool_names = extract_request_tool_names(route, request_body) if not tool_names: return - key_meta = (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} - team_meta = (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {} + key_meta = ( + (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} + ) + team_meta = ( + (valid_token.team_metadata or {}) + if isinstance(valid_token.team_metadata, dict) + else {} + ) key_allowed = key_meta.get("allowed_tools") team_allowed = team_meta.get("allowed_tools") - effective = key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed + effective = ( + key_allowed + if (isinstance(key_allowed, list) and len(key_allowed) > 0) + else team_allowed + ) if not isinstance(effective, list) or len(effective) == 0: return allowed_set = {str(t) for t in effective} @@ -326,6 +406,29 @@ async def common_checks( # noqa: PLR0915 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( @@ -415,83 +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( @@ -1932,9 +1962,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") @@ -1980,9 +2009,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") @@ -2021,9 +2049,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" @@ -2144,13 +2171,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: @@ -2296,9 +2321,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, ) @@ -2341,8 +2366,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 @@ -3298,7 +3325,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 diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 82341c9a70..94cb7510a5 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -25,53 +25,38 @@ from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( - ExperimentalUIJWTToken, - _cache_key_object, - _delete_cache_key_object, - _get_user_role, - _is_user_proxy_admin, - _virtual_key_max_budget_alert_check, - _virtual_key_max_budget_check, - _virtual_key_soft_budget_check, - can_key_call_model, - common_checks, - get_end_user_object, - get_jwt_key_mapping_object, - get_key_object, - get_project_object, - get_team_object, - get_user_object, - is_valid_fallback_model, -) -from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler -from litellm.proxy.auth.auth_utils import ( - abbreviate_api_key, - get_end_user_id_from_request_body, - get_model_from_request, - get_request_route, - normalize_request_route, - pre_db_read_auth_checks, - route_in_additonal_public_routes, -) + ExperimentalUIJWTToken, _cache_key_object, _delete_cache_key_object, + _get_user_role, _is_user_proxy_admin, _virtual_key_max_budget_alert_check, + _virtual_key_max_budget_check, _virtual_key_soft_budget_check, + can_key_call_model, common_checks, get_end_user_object, + get_jwt_key_mapping_object, get_key_object, get_project_object, + get_team_object, get_user_object, is_valid_fallback_model) +from litellm.proxy.auth.auth_exception_handler import \ + UserAPIKeyAuthExceptionHandler +from litellm.proxy.auth.auth_utils import (abbreviate_api_key, + get_end_user_id_from_request_body, + get_model_from_request, + get_request_route, + normalize_request_route, + pre_db_read_auth_checks, + route_in_additonal_public_routes) from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator +from litellm.proxy.common_utils.cache_coordinator import \ + EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( - _read_request_body, - _safe_get_request_headers, - populate_request_with_path_params, -) + _read_request_body, _safe_get_request_headers, + populate_request_with_path_params) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes try: - from litellm_enterprise.proxy.auth.user_api_key_auth import ( - enterprise_custom_auth as _enterprise_custom_auth, - ) + from litellm_enterprise.proxy.auth.user_api_key_auth import \ + enterprise_custom_auth as _enterprise_custom_auth enterprise_custom_auth: Optional[Callable] = _enterprise_custom_auth except ImportError as e: @@ -351,9 +336,8 @@ def get_api_key( Tuple[Optional[str], Optional[str]]: Tuple of the api_key and the passed_in_key """ from litellm.proxy.auth.route_checks import RouteChecks - from litellm.proxy.common_utils.http_parsing_utils import ( - _safe_get_request_query_params, - ) + from litellm.proxy.common_utils.http_parsing_utils import \ + _safe_get_request_query_params api_key = api_key passed_in_key: Optional[str] = None @@ -519,20 +503,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request_data: dict, custom_litellm_key_header: Optional[str] = None, ) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import ( - general_settings, - jwt_handler, - litellm_proxy_admin_name, - llm_model_list, - llm_router, - master_key, - model_max_budget_limiter, - open_telemetry_logger, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - user_custom_auth, - ) + from litellm.proxy.proxy_server import (general_settings, jwt_handler, + litellm_proxy_admin_name, + llm_model_list, llm_router, + master_key, + model_max_budget_limiter, + open_telemetry_logger, + prisma_client, proxy_logging_obj, + user_api_key_cache, + user_custom_auth) parent_otel_span: Optional[Span] = None start_time = datetime.now() @@ -730,9 +709,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_object is not None else None ), - team_metadata=team_object.metadata - if team_object is not None - else None, + team_metadata=( + team_object.metadata + if team_object is not None + else None + ), org_id=org_id, end_user_id=end_user_id, parent_otel_span=parent_otel_span, @@ -750,9 +731,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 team_rpm_limit=( team_object.rpm_limit if team_object is not None else None ), - team_models=team_object.models - if team_object is not None - else [], + team_models=( + team_object.models if team_object is not None else [] + ), user_role=( LitellmUserRoles(user_object.user_role) if user_object is not None @@ -779,16 +760,17 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_membership is not None else None ), - team_metadata=team_object.metadata - if team_object is not None - else None, + team_metadata=( + team_object.metadata if team_object is not None else None + ), ) # Check if model has zero cost - if so, skip all budget checks model = get_model_from_request(request_data, route) skip_budget_checks = False if model is not None and llm_router is not None: - from litellm.proxy.auth.auth_checks import _is_model_cost_zero + from litellm.proxy.auth.auth_checks import \ + _is_model_cost_zero skip_budget_checks = _is_model_cost_zero( model=model, llm_router=llm_router @@ -893,9 +875,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, ) if _end_user_object is not None: - end_user_params[ - "allowed_model_region" - ] = _end_user_object.allowed_model_region + end_user_params["allowed_model_region"] = ( + _end_user_object.allowed_model_region + ) if _end_user_object.litellm_budget_table is not None: _apply_budget_limits_to_end_user_params( end_user_params=end_user_params, @@ -904,9 +886,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) elif litellm.max_end_user_budget_id is not None: # End user doesn't exist yet, but apply default budget limits if configured - from litellm.proxy.auth.auth_checks import ( - get_default_end_user_budget, - ) + from litellm.proxy.auth.auth_checks import \ + get_default_end_user_budget default_budget = await get_default_end_user_budget( prisma_client=prisma_client, @@ -1463,9 +1444,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if _end_user_object is not None: valid_token_dict.update(end_user_params) - valid_token_dict[ - "end_user_object_permission" - ] = _end_user_object.object_permission + valid_token_dict["end_user_object_permission"] = ( + _end_user_object.object_permission + ) # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions # sso/login, ui/login, /key functions and /user functions @@ -1687,7 +1668,8 @@ async def _lookup_end_user_and_apply_budget( valid_token=valid_token, end_user_params=end_user_params ) elif litellm.max_end_user_budget_id is not None: - from litellm.proxy.auth.auth_checks import get_default_end_user_budget + from litellm.proxy.auth.auth_checks import \ + get_default_end_user_budget default_budget = await get_default_end_user_budget( prisma_client=prisma_client, @@ -1718,14 +1700,10 @@ async def _run_post_custom_auth_checks( route: str, parent_otel_span: Optional[Span], ) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import ( - prisma_client, - user_api_key_cache, - proxy_logging_obj, - general_settings, - llm_router, - model_max_budget_limiter, - ) + from litellm.proxy.proxy_server import (general_settings, llm_router, + model_max_budget_limiter, + prisma_client, proxy_logging_obj, + user_api_key_cache) # 1. Look up end_user object from DB if end_user_id is set end_user_object = None @@ -1756,9 +1734,11 @@ async def _run_post_custom_auth_checks( message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", type=ProxyErrorTypes.expired_key, code=400, - param=abbreviate_api_key(api_key=valid_token.token) - if valid_token.token - else "", + param=( + abbreviate_api_key(api_key=valid_token.token) + if valid_token.token + else "" + ), ) current_model = request_data.get("model", None) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4c96e079c9..28b1e6601b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -13,36 +13,49 @@ 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: @@ -91,10 +104,12 @@ 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: @@ -109,8 +124,9 @@ 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, @@ -374,6 +390,19 @@ 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=_agent_id_for_spend, + 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, @@ -604,6 +633,34 @@ 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: + 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], @@ -765,7 +822,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 {} @@ -798,6 +855,12 @@ class DBSpendUpdateWriter: 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, @@ -1002,8 +1065,10 @@ 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"] @@ -1279,6 +1344,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, @@ -2031,9 +2108,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/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 51201f96d7..4f38e71bbf 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -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 diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 3e059cf8c1..b7cd06a64f 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -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 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 feea3023d4..2d0ce040a6 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 @@ -341,6 +341,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, @@ -466,58 +490,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/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index 1d1e559d4b..790ebcd879 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -5,6 +5,8 @@ from . import * from .cache_control_check import _PROXY_CacheControlCheck from .litellm_skills import SkillsInjectionHook from .max_budget_limiter import _PROXY_MaxBudgetLimiter +from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler +from .max_iterations_limiter import _PROXY_MaxIterationsHandler from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler from .parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 from .responses_id_security import ResponsesIDSecurity @@ -23,6 +25,8 @@ PROXY_HOOKS = { "cache_control_check": _PROXY_CacheControlCheck, "responses_id_security": ResponsesIDSecurity, "litellm_skills": SkillsInjectionHook, + "max_iterations_limiter": _PROXY_MaxIterationsHandler, + "max_budget_per_session_limiter": _PROXY_MaxBudgetPerSessionHandler, } ## FEATURE FLAG HOOKS ## diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py new file mode 100644 index 0000000000..a981207f00 --- /dev/null +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -0,0 +1,271 @@ +""" +Per-Session Budget Limiter for LiteLLM Proxy. + +Enforces a dollar-amount cap per session (identified by `session_id` / +`x-litellm-trace-id`). After each successful LLM call the response cost is +accumulated against the session. When the accumulated spend exceeds +`max_budget_per_session` (configured in agent litellm_params), subsequent +requests for that session receive a 429. + +Note: trace-id enforcement (require_trace_id_on_calls_by_agent) is handled +separately in auth_checks.py at the agent level, not in this hook. + +Works across multiple proxy instances via DualCache (in-memory + Redis). +Follows the same pattern as max_iterations_limiter.py. +""" + +import os +from typing import TYPE_CHECKING, Any, Optional, Union + +from fastapi import HTTPException + +from litellm import DualCache +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth + +if TYPE_CHECKING: + from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache + + InternalUsageCache = _InternalUsageCache +else: + InternalUsageCache = Any + + +# Redis Lua script for atomic float increment with TTL. +# INCRBYFLOAT returns the new value as a string. +# Only sets EXPIRE on first call (when prior value was nil). +MAX_BUDGET_SESSION_INCREMENT_SCRIPT = """ +local key = KEYS[1] +local amount = ARGV[1] +local ttl = tonumber(ARGV[2]) + +local existed = redis.call('EXISTS', key) +local new_val = redis.call('INCRBYFLOAT', key, amount) +if existed == 0 then + redis.call('EXPIRE', key, ttl) +end + +return new_val +""" + +# Default TTL for session budget counters (1 hour) +DEFAULT_MAX_BUDGET_PER_SESSION_TTL = 3600 + + +class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): + """ + Pre-call hook that enforces max_budget_per_session. + + Configuration (set in agent litellm_params): + - max_budget_per_session: dollar cap per session_id + + Cache key pattern: + {session_budget:}:spend + """ + + def __init__(self, internal_usage_cache: InternalUsageCache): + self.internal_usage_cache = internal_usage_cache + self.ttl = int( + os.getenv( + "LITELLM_MAX_BUDGET_PER_SESSION_TTL", + DEFAULT_MAX_BUDGET_PER_SESSION_TTL, + ) + ) + + if self.internal_usage_cache.dual_cache.redis_cache is not None: + self.increment_script = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + MAX_BUDGET_SESSION_INCREMENT_SCRIPT + ) + ) + else: + self.increment_script = None + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: str, + ) -> Optional[Union[Exception, str, dict]]: + """ + Before each LLM call, check if max_budget_per_session is set and + whether accumulated spend exceeds the budget (429 if so). + """ + max_budget = self._get_max_budget_per_session(user_api_key_dict) + + session_id = self._get_session_id(data) + + if max_budget is None or session_id is None: + return None + + max_budget = float(max_budget) + cache_key = self._make_cache_key(session_id) + current_spend = await self._get_current_spend(cache_key) + + verbose_proxy_logger.debug( + "MaxBudgetPerSessionHandler: session_id=%s, spend=%.4f, max=%.2f", + session_id, + current_spend, + max_budget, + ) + + if current_spend >= max_budget: + raise HTTPException( + status_code=429, + detail=( + f"Session budget exceeded for session {session_id}. " + f"Current spend: ${current_spend:.4f}, " + f"max_budget_per_session: ${max_budget:.2f}." + ), + ) + + return None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + After a successful LLM call, increment the session spend by the response cost. + """ + try: + litellm_params = kwargs.get("litellm_params") or {} + metadata = litellm_params.get("metadata") or {} + session_id = metadata.get("session_id") + if session_id is None: + return + + agent_id = metadata.get("agent_id") + if agent_id is None: + return + + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + ) + + agent = global_agent_registry.get_agent_by_id(agent_id=str(agent_id)) + if agent is None: + return + + agent_litellm_params = agent.litellm_params or {} + max_budget = agent_litellm_params.get("max_budget_per_session") + if max_budget is None: + return + + response_cost = kwargs.get("response_cost") or 0.0 + if response_cost <= 0: + return + + cache_key = self._make_cache_key(str(session_id)) + await self._increment_spend(cache_key, float(response_cost)) + + verbose_proxy_logger.debug( + "MaxBudgetPerSessionHandler: incremented session %s spend by %.6f", + session_id, + response_cost, + ) + except Exception as e: + verbose_proxy_logger.warning( + "MaxBudgetPerSessionHandler: error in async_log_success_event: %s", + str(e), + ) + + def _get_session_id(self, data: dict) -> Optional[str]: + """Extract session_id from request metadata.""" + metadata = data.get("metadata") or {} + session_id = metadata.get("session_id") + if session_id is not None: + return str(session_id) + + litellm_metadata = data.get("litellm_metadata") or {} + session_id = litellm_metadata.get("session_id") + if session_id is not None: + return str(session_id) + + return None + + def _get_max_budget_per_session( + self, user_api_key_dict: UserAPIKeyAuth + ) -> Optional[float]: + """Extract max_budget_per_session from agent litellm_params.""" + agent_id = user_api_key_dict.agent_id + if agent_id is None: + return None + + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + agent = global_agent_registry.get_agent_by_id(agent_id=agent_id) + if agent is None: + return None + + litellm_params = agent.litellm_params or {} + max_budget = litellm_params.get("max_budget_per_session") + if max_budget is not None: + return float(max_budget) + return None + + def _make_cache_key(self, session_id: str) -> str: + return f"{{session_budget:{session_id}}}:spend" + + async def _get_current_spend(self, cache_key: str) -> float: + """Read current accumulated spend for a session.""" + if ( + self.internal_usage_cache.dual_cache.redis_cache is not None + ): + try: + result = await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache( + key=cache_key + ) + if result is not None: + return float(result) + return 0.0 + except Exception as e: + verbose_proxy_logger.warning( + "MaxBudgetPerSessionHandler: Redis GET failed, " + "falling back to in-memory: %s", + str(e), + ) + + result = await self.internal_usage_cache.async_get_cache( + key=cache_key, + litellm_parent_otel_span=None, + local_only=True, + ) + if result is not None: + return float(result) + return 0.0 + + async def _increment_spend(self, cache_key: str, amount: float) -> float: + """Atomically increment the session spend and return the new value.""" + if self.increment_script is not None: + try: + result = await self.increment_script( + keys=[cache_key], + args=[str(amount), self.ttl], + ) + return float(result) + except Exception as e: + verbose_proxy_logger.warning( + "MaxBudgetPerSessionHandler: Redis INCRBYFLOAT failed, " + "falling back to in-memory: %s", + str(e), + ) + + return await self._in_memory_increment_spend(cache_key, amount) + + async def _in_memory_increment_spend( + self, cache_key: str, amount: float + ) -> float: + current = await self.internal_usage_cache.async_get_cache( + key=cache_key, + litellm_parent_otel_span=None, + local_only=True, + ) + new_value = (float(current) if current is not None else 0.0) + amount + await self.internal_usage_cache.async_set_cache( + key=cache_key, + value=new_value, + ttl=self.ttl, + litellm_parent_otel_span=None, + local_only=True, + ) + return new_value diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index 8d481f6b26..b6fde2b178 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -4,7 +4,7 @@ Max Iterations Limiter for LiteLLM Proxy. Enforces a per-session cap on the number of LLM calls an agentic loop can make. Callers send a `session_id` with each request (via `x-litellm-session-id` header or `metadata.session_id`), and this hook counts calls per session. When the count -exceeds `max_iterations` (configured in key/team metadata), returns 429. +exceeds `max_iterations` (configured in agent litellm_params or key metadata), returns 429. Works across multiple proxy instances via DualCache (in-memory + Redis). Follows the same pattern as parallel_request_limiter_v3.py. @@ -52,8 +52,9 @@ class _PROXY_MaxIterationsHandler(CustomLogger): Pre-call hook that enforces max_iterations per session. Configuration: - - max_iterations: set in key metadata via /key/generate or /key/update - e.g. metadata={"max_iterations": 25} + - max_iterations: set in agent litellm_params (preferred) + e.g. litellm_params={"max_iterations": 25} + Falls back to key metadata max_iterations for backwards compatibility. - session_id: sent by caller via x-litellm-session-id header or metadata.session_id in request body @@ -93,14 +94,13 @@ class _PROXY_MaxIterationsHandler(CustomLogger): Check session iteration count before making the API call. Extracts session_id from request metadata and max_iterations from - key metadata. If the session has exceeded max_iterations, raises 429. + agent litellm_params. If the session has exceeded max_iterations, raises 429. """ # Extract session_id from request data session_id = self._get_session_id(data) if session_id is None: return None - # Extract max_iterations from key metadata max_iterations = self._get_max_iterations(user_api_key_dict) if max_iterations is None: return None @@ -151,7 +151,22 @@ class _PROXY_MaxIterationsHandler(CustomLogger): def _get_max_iterations( self, user_api_key_dict: UserAPIKeyAuth ) -> Optional[int]: - """Extract max_iterations from key metadata.""" + """Extract max_iterations from agent litellm_params, with fallback to key metadata.""" + # Try agent litellm_params first + agent_id = user_api_key_dict.agent_id + if agent_id is not None: + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + ) + + agent = global_agent_registry.get_agent_by_id(agent_id=agent_id) + if agent is not None: + litellm_params = agent.litellm_params or {} + max_iterations = litellm_params.get("max_iterations") + if max_iterations is not None: + return int(max_iterations) + + # Fallback to key metadata for backwards compatibility metadata = user_api_key_dict.metadata or {} max_iterations = metadata.get("max_iterations") if max_iterations is not None: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b5bbb4237c..856975ea09 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -7,18 +7,8 @@ This is currently in development and not yet ready for production. import binascii import os from datetime import datetime -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - List, - Literal, - Optional, - TypedDict, - Union, - cast, -) +from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, + Optional, TypedDict, Union, cast) from fastapi import HTTPException @@ -175,9 +165,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """Get or lazy-load the batch rate limiter.""" if self._batch_rate_limiter is None: try: - from litellm.proxy.hooks.batch_rate_limiter import ( - _PROXY_BatchRateLimiter, - ) + from litellm.proxy.hooks.batch_rate_limiter import \ + _PROXY_BatchRateLimiter self._batch_rate_limiter = _PROXY_BatchRateLimiter( internal_usage_cache=self.internal_usage_cache, @@ -679,10 +668,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model: The model being requested descriptors: List of rate limit descriptors to append to """ - from litellm.proxy.auth.auth_utils import ( - get_key_model_rpm_limit, - get_key_model_tpm_limit, - ) + from litellm.proxy.auth.auth_utils import (get_key_model_rpm_limit, + get_key_model_tpm_limit) if not requested_model: return @@ -791,6 +778,92 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ return rpm_limit_type == "dynamic" or tpm_limit_type == "dynamic" + def _get_agent_from_registry(self, agent_id: str) -> Optional[Any]: + """Look up an agent from the in-memory registry by ID.""" + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry + + return global_agent_registry.get_agent_by_id(agent_id=agent_id) + + def _get_resolved_agent_id( + self, user_api_key_dict: UserAPIKeyAuth, data: dict + ) -> Optional[str]: + """ + Resolve the agent_id from either the API key or request metadata. + Key-level agent_id takes precedence over metadata/header-supplied agent_id. + """ + key_agent_id = getattr(user_api_key_dict, "agent_id", None) + if key_agent_id: + return key_agent_id + metadata = data.get("metadata") or {} + return metadata.get("agent_id") + + def _get_session_id_from_data(self, data: dict) -> Optional[str]: + """Extract session_id from request metadata or litellm_session_id.""" + session_id = data.get("litellm_session_id") + if session_id: + return str(session_id) + metadata = data.get("metadata") or {} + session_id = metadata.get("session_id") + if session_id: + return str(session_id) + litellm_metadata = data.get("litellm_metadata") or {} + session_id = litellm_metadata.get("session_id") + if session_id: + return str(session_id) + return None + + def _create_agent_rate_limit_descriptors( + self, + agent_id: str, + data: dict, + ) -> List[RateLimitDescriptor]: + """ + Create rate limit descriptors for agent-level and session-level limits. + + Agent-level: caps total RPM/TPM across all sessions for a given agent. + Session-level: caps RPM/TPM within a single session (identified by session_id). + """ + descriptors: List[RateLimitDescriptor] = [] + + agent = self._get_agent_from_registry(agent_id) + if agent is None: + return descriptors + + agent_rpm = getattr(agent, "rpm_limit", None) + agent_tpm = getattr(agent, "tpm_limit", None) + if agent_rpm is not None or agent_tpm is not None: + descriptors.append( + RateLimitDescriptor( + key="agent", + value=agent_id, + rate_limit={ + "requests_per_unit": agent_rpm, + "tokens_per_unit": agent_tpm, + "window_size": self.window_size, + }, + ) + ) + + session_rpm = getattr(agent, "session_rpm_limit", None) + session_tpm = getattr(agent, "session_tpm_limit", None) + if session_rpm is not None or session_tpm is not None: + session_id = self._get_session_id_from_data(data) + if session_id is not None: + descriptors.append( + RateLimitDescriptor( + key="agent_session", + value=f"{agent_id}:{session_id}", + rate_limit={ + "requests_per_unit": session_rpm, + "tokens_per_unit": session_tpm, + "window_size": self.window_size, + }, + ) + ) + + return descriptors + def _create_rate_limit_descriptors( self, user_api_key_dict: UserAPIKeyAuth, @@ -802,12 +875,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Create all rate limit descriptors for the request. - Returns list of descriptors for API key, user, team, team member, end user, and model-specific limits. + Returns list of descriptors for API key, user, team, team member, end user, + model-specific, agent, and agent-session limits. """ - from litellm.proxy.auth.auth_utils import ( - get_team_model_rpm_limit, - get_team_model_tpm_limit, - ) + from litellm.proxy.auth.auth_utils import (get_team_model_rpm_limit, + get_team_model_tpm_limit) descriptors = [] @@ -956,6 +1028,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + # Agent-level and session-level rate limits + resolved_agent_id = self._get_resolved_agent_id(user_api_key_dict, data) + + if resolved_agent_id: + descriptors.extend( + self._create_agent_rate_limit_descriptors( + agent_id=resolved_agent_id, + data=data, + ) + ) + return descriptors async def _check_model_has_recent_failures( @@ -970,9 +1053,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Returns True if any deployment has failures in the current minute. """ from litellm.proxy.proxy_server import llm_router - from litellm.router_utils.router_callbacks.track_deployment_metrics import ( - get_deployment_failures_for_current_minute, - ) + from litellm.router_utils.router_callbacks.track_deployment_metrics import \ + get_deployment_failures_for_current_minute if llm_router is None: return False @@ -1386,12 +1468,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Update TPM usage on successful API calls by incrementing counters using pipeline """ - from litellm.litellm_core_utils.core_helpers import ( - _get_parent_otel_span_from_kwargs, - ) - from litellm.proxy.common_utils.callback_utils import ( - get_model_group_from_litellm_kwargs, - ) + from litellm.litellm_core_utils.core_helpers import \ + _get_parent_otel_span_from_kwargs + from litellm.proxy.common_utils.callback_utils import \ + get_model_group_from_litellm_kwargs from litellm.types.caching import RedisPipelineIncrementOperation rate_limit_type = self.get_rate_limit_type() @@ -1533,6 +1613,32 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + # Agent TPM + agent_id = standard_logging_metadata.get("agent_id") + if agent_id: + pipeline_operations.extend( + self._create_pipeline_operations( + key="agent", + value=agent_id, + rate_limit_type="tokens", + total_tokens=total_tokens, + ) + ) + + # Agent Session TPM + session_id = standard_logging_metadata.get( + "session_id" + ) or standard_logging_metadata.get("trace_id") + if session_id: + pipeline_operations.extend( + self._create_pipeline_operations( + key="agent_session", + value=f"{agent_id}:{session_id}", + rate_limit_type="tokens", + total_tokens=total_tokens, + ) + ) + # Execute all increments in a single pipeline if pipeline_operations: await self.async_increment_tokens_with_ttl_preservation( @@ -1549,9 +1655,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Decrement max parallel requests counter for the API Key """ - from litellm.litellm_core_utils.core_helpers import ( - _get_parent_otel_span_from_kwargs, - ) + from litellm.litellm_core_utils.core_helpers import \ + _get_parent_otel_span_from_kwargs from litellm.types.caching import RedisPipelineIncrementOperation try: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d0fabe9010..cf4729db94 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -690,6 +690,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 + _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 ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index db0bb735ba..bd7b21c3b5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -373,8 +373,9 @@ from litellm.proxy.management_endpoints.fallback_management_endpoints import ( from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) -from litellm.proxy.management_endpoints.internal_user_endpoints import ( - user_update, +from litellm.proxy.management_endpoints.internal_user_endpoints import user_update +from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( + router as jwt_key_mapping_router, ) from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( router as jwt_key_mapping_router, @@ -444,9 +445,7 @@ from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_route from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import ( - set_files_config, -) +from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -545,9 +544,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) from litellm.types.realtime import RealtimeQueryParams -from litellm.types.router import ( - DeploymentTypedDict, -) +from litellm.types.router import DeploymentTypedDict from litellm.types.router import ModelInfo as RouterModelInfo from litellm.types.router import ( RouterGeneralSettings, @@ -6682,6 +6679,11 @@ async def chat_completion( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + if ( + hasattr(user_api_key_dict, "agent_id") + and user_api_key_dict.agent_id is not None + ): + data["metadata"]["agent_id"] = user_api_key_dict.agent_id base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) try: result = await base_llm_response_processor.base_process_llm_request( @@ -6851,6 +6853,11 @@ async def completion( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + if ( + hasattr(user_api_key_dict, "agent_id") + and user_api_key_dict.agent_id is not None + ): + data["metadata"]["agent_id"] = user_api_key_dict.agent_id base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) return await base_llm_response_processor.base_process_llm_request( request=request, @@ -7088,6 +7095,11 @@ async def embeddings( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + if ( + hasattr(user_api_key_dict, "agent_id") + and user_api_key_dict.agent_id is not None + ): + data["metadata"]["agent_id"] = user_api_key_dict.agent_id # Use unified request processor (same as chat/completions and responses) base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 329ff80933..8d4bdffb2d 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -68,6 +68,11 @@ 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) + tpm_limit Int? + rpm_limit Int? + session_tpm_limit Int? + session_rpm_limit Int? created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c5f399e3ad..44c02f8b16 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -325,7 +325,7 @@ class ProxyLogging: if email_logger_class is not None: # All email logger classes now accept internal_usage_cache self.email_logging_instance = email_logger_class( - internal_usage_cache=self.internal_usage_cache.dual_cache, + internal_usage_cache=self.internal_usage_cache.dual_cache, # type: ignore[call-arg] ) self.premium_user = premium_user self.service_logging_obj = ServiceLogging() diff --git a/litellm/router.py b/litellm/router.py index 7119d2e850..43b53d14d7 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/litellm/types/agents.py b/litellm/types/agents.py index 7879cae9ff..951fbfcabd 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -179,6 +179,10 @@ class AgentConfig(TypedDict, total=False): agent_card_params: Required[AgentCard] litellm_params: Dict[str, Any] # allow for any future litellm params object_permission: AgentObjectPermission + tpm_limit: Optional[int] + rpm_limit: Optional[int] + session_tpm_limit: Optional[int] + session_rpm_limit: Optional[int] static_headers: Optional[Dict[str, str]] extra_headers: Optional[List[str]] @@ -188,6 +192,10 @@ class PatchAgentRequest(TypedDict, total=False): agent_card_params: AgentCard litellm_params: Dict[str, Any] object_permission: AgentObjectPermission + tpm_limit: Optional[int] + rpm_limit: Optional[int] + session_tpm_limit: Optional[int] + session_rpm_limit: Optional[int] static_headers: Optional[Dict[str, str]] extra_headers: Optional[List[str]] @@ -201,6 +209,11 @@ class AgentResponse(BaseModel): litellm_params: Optional[Dict[str, Any]] = None agent_card_params: Dict[str, Any] object_permission: Optional[Dict[str, Any]] = None + spend: Optional[float] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + session_tpm_limit: Optional[int] = None + session_rpm_limit: Optional[int] = None static_headers: Optional[Dict[str, str]] = None extra_headers: Optional[List[str]] = None created_at: Optional[datetime] = None diff --git a/schema.prisma b/schema.prisma index 329ff80933..8d4bdffb2d 100644 --- a/schema.prisma +++ b/schema.prisma @@ -68,6 +68,11 @@ 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) + tpm_limit Int? + rpm_limit Int? + session_tpm_limit Int? + session_rpm_limit Int? 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 abd59a66a3..b5f82ef04c 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/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py new file mode 100644 index 0000000000..879e2d65c7 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py @@ -0,0 +1,165 @@ +""" +Unit Tests for the per-session budget limiter for the proxy. + +Tests that session-scoped budget tracking works correctly: +- Enforces max_budget_per_session per session_id (read from agent litellm_params) +- Different sessions have independent budgets +- Requests under budget pass through +- Requests without agent_id pass through +""" + +from unittest.mock import patch + +import pytest +from fastapi import HTTPException + +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.max_budget_per_session_limiter import ( + _PROXY_MaxBudgetPerSessionHandler, +) +from litellm.proxy.utils import InternalUsageCache +from litellm.types.agents import AgentResponse + + +def _make_mock_agent(max_budget_per_session: float) -> AgentResponse: + return AgentResponse( + agent_id="agent-budget-123", + agent_name="budget-agent", + litellm_params={"max_budget_per_session": max_budget_per_session}, + agent_card_params={"name": "budget-agent", "version": "1.0.0"}, + ) + + +@pytest.mark.asyncio +async def test_budget_per_session_under_budget_passes(): + """ + Requests under budget should pass through without error. + """ + local_cache = DualCache() + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-budget", + agent_id="agent-budget-123", + ) + + mock_agent = _make_mock_agent(max_budget_per_session=5.0) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = mock_agent + + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-budget-1"}}, + call_type="", + ) + assert result is None + + +@pytest.mark.asyncio +async def test_budget_per_session_exceeds_budget(): + """ + After accumulating spend beyond max_budget_per_session, the next + pre-call check should raise 429. + """ + local_cache = DualCache() + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-budget", + agent_id="agent-budget-123", + ) + + session_id = "session-over-budget" + cache_key = handler._make_cache_key(session_id) + await handler._increment_spend(cache_key, 1.50) + + mock_agent = _make_mock_agent(max_budget_per_session=1.0) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = mock_agent + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": session_id}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + assert "budget exceeded" in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_budget_per_session_independent_sessions(): + """ + Different session_ids have independent budget counters. + Exhausting session A does not affect session B. + """ + local_cache = DualCache() + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-budget", + agent_id="agent-budget-123", + ) + + cache_key_a = handler._make_cache_key("session-A") + await handler._increment_spend(cache_key_a, 3.0) + + mock_agent = _make_mock_agent(max_budget_per_session=2.0) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = mock_agent + + # Session A should be blocked + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-A"}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + + # Session B should still pass + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-B"}}, + call_type="", + ) + assert result is None + + +@pytest.mark.asyncio +async def test_no_agent_id_passes(): + """ + When no agent_id is set on the key, all requests pass through. + """ + local_cache = DualCache() + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-no-agent", + ) + + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "any-session"}}, + call_type="", + ) + assert result is None diff --git a/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py b/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py index deb1c483b8..20928ef46d 100644 --- a/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py @@ -2,10 +2,12 @@ Unit Tests for the max iterations limiter for the proxy. Tests that session-scoped iteration counting works correctly: -- Enforces max_iterations per session_id +- Enforces max_iterations per session_id (read from agent litellm_params) - Different sessions have independent counters """ +from unittest.mock import MagicMock, patch + import pytest from fastapi import HTTPException @@ -13,6 +15,16 @@ from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler from litellm.proxy.utils import InternalUsageCache +from litellm.types.agents import AgentResponse + + +def _make_mock_agent(max_iterations: int) -> AgentResponse: + return AgentResponse( + agent_id="agent-test-123", + agent_name="test-agent", + litellm_params={"max_iterations": max_iterations}, + agent_card_params={"name": "test-agent", "version": "1.0.0"}, + ) @pytest.mark.asyncio @@ -28,28 +40,36 @@ async def test_max_iterations_basic_enforcement(): internal_usage_cache=InternalUsageCache(local_cache), ) user_api_key_dict = UserAPIKeyAuth( - api_key="sk-test-key-1234", metadata={"max_iterations": 3} + api_key="sk-test-key-1234", + agent_id="agent-test-123", ) - # First 3 requests should succeed - for i in range(3): - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-abc"}}, - call_type="", - ) + mock_agent = _make_mock_agent(max_iterations=3) - # 4th request should fail with 429 - with pytest.raises(HTTPException) as exc_info: - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-abc"}}, - call_type="", - ) - assert exc_info.value.status_code == 429 - assert "max_iterations" in str(exc_info.value.detail).lower() + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = mock_agent + + # First 3 requests should succeed + for i in range(3): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-abc"}}, + call_type="", + ) + + # 4th request should fail with 429 + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-abc"}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + assert "max_iterations" in str(exc_info.value.detail).lower() @pytest.mark.asyncio @@ -65,42 +85,72 @@ async def test_max_iterations_different_sessions_independent(): internal_usage_cache=InternalUsageCache(local_cache), ) user_api_key_dict = UserAPIKeyAuth( - api_key="sk-test-key-5678", metadata={"max_iterations": 2} + api_key="sk-test-key-5678", + agent_id="agent-test-123", ) - # Session A: 2 calls succeed - for _ in range(2): - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-A"}}, - call_type="", - ) + mock_agent = _make_mock_agent(max_iterations=2) - # Session B: 2 calls succeed (independent counter) - for _ in range(2): - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-B"}}, - call_type="", - ) + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = mock_agent - # Session A: 3rd call fails - with pytest.raises(HTTPException) as exc_info: - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-A"}}, - call_type="", - ) - assert exc_info.value.status_code == 429 + # Session A: 2 calls succeed + for _ in range(2): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-A"}}, + call_type="", + ) - # Session B: 3rd call also fails - with pytest.raises(HTTPException): - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-B"}}, - call_type="", - ) + # Session B: 2 calls succeed (independent counter) + for _ in range(2): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-B"}}, + call_type="", + ) + + # Session A: 3rd call fails + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-A"}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + + # Session B: 3rd call also fails + with pytest.raises(HTTPException): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-B"}}, + call_type="", + ) + + +@pytest.mark.asyncio +async def test_max_iterations_no_agent_id_passes(): + """ + When no agent_id is set on the key, all requests pass through. + """ + local_cache = DualCache() + handler = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-no-agent", + ) + + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-any"}}, + call_type="", + ) + assert result is None diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index d92e152d89..3eb481991f 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -1981,6 +1981,527 @@ async def test_execute_token_increment_script_cluster_compatibility(): ), f"Each key should have 2 args, got {len(args)} args for {len(keys)} keys" +@pytest.mark.asyncio +async def test_agent_level_rate_limit_descriptors(): + """ + Test that agent-level rate limit descriptors are created when + an agent has rpm_limit and/or tpm_limit configured. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_abc123" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="test-agent", + agent_card_params={"name": "Test Agent"}, + rpm_limit=50, + tpm_limit=5000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + assert captured_descriptors is not None + + agent_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent": + agent_descriptor = d + break + + assert agent_descriptor is not None, "Agent descriptor should be present" + assert agent_descriptor["value"] == _agent_id + assert agent_descriptor["rate_limit"]["requests_per_unit"] == 50 + assert agent_descriptor["rate_limit"]["tokens_per_unit"] == 5000 + + +@pytest.mark.asyncio +async def test_agent_session_rate_limit_descriptors(): + """ + Test that session-level rate limit descriptors are created when + an agent has session_rpm_limit/session_tpm_limit and a session_id is present. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_abc123" + _session_id = "sess_xyz789" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="test-agent", + agent_card_params={"name": "Test Agent"}, + session_rpm_limit=10, + session_tpm_limit=1000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4", + "metadata": {"session_id": _session_id}, + }, + call_type="", + ) + + assert captured_descriptors is not None + + session_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent_session": + session_descriptor = d + break + + assert session_descriptor is not None, "Agent session descriptor should be present" + assert session_descriptor["value"] == f"{_agent_id}:{_session_id}" + assert session_descriptor["rate_limit"]["requests_per_unit"] == 10 + assert session_descriptor["rate_limit"]["tokens_per_unit"] == 1000 + + +@pytest.mark.asyncio +async def test_agent_session_rate_limit_skipped_without_session_id(): + """ + Test that session-level rate limit descriptors are NOT created + when no session_id is available in the request. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_abc123" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="test-agent", + agent_card_params={"name": "Test Agent"}, + session_rpm_limit=10, + session_tpm_limit=1000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + # should_rate_limit should not have been called (no agent-level limits, only session limits + # but no session_id) + assert captured_descriptors is None, ( + "No descriptors should be created when agent has only session limits " + "but no session_id in request" + ) + + +@pytest.mark.asyncio +async def test_agent_rate_limit_from_metadata_agent_id(): + """ + Test that agent rate limits work when agent_id comes from + request metadata (header) rather than from the API key. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_from_header" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="header-agent", + agent_card_params={"name": "Header Agent"}, + rpm_limit=25, + tpm_limit=2500, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4", + "metadata": {"agent_id": _agent_id}, + }, + call_type="", + ) + + assert captured_descriptors is not None + + agent_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent": + agent_descriptor = d + break + + assert agent_descriptor is not None, "Agent descriptor should be created from metadata agent_id" + assert agent_descriptor["value"] == _agent_id + assert agent_descriptor["rate_limit"]["requests_per_unit"] == 25 + + +@pytest.mark.asyncio +async def test_agent_both_agent_and_session_rate_limits(): + """ + Test that both agent-level and session-level descriptors are created + when both types of limits are configured on the agent. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_dual" + _session_id = "sess_dual" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="dual-agent", + agent_card_params={"name": "Dual Agent"}, + rpm_limit=100, + tpm_limit=10000, + session_rpm_limit=20, + session_tpm_limit=2000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4", + "metadata": {"session_id": _session_id}, + }, + call_type="", + ) + + assert captured_descriptors is not None + + agent_descriptor = None + session_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent": + agent_descriptor = d + elif d["key"] == "agent_session": + session_descriptor = d + + assert agent_descriptor is not None, "Agent-level descriptor should be present" + assert agent_descriptor["rate_limit"]["requests_per_unit"] == 100 + assert agent_descriptor["rate_limit"]["tokens_per_unit"] == 10000 + + assert session_descriptor is not None, "Session-level descriptor should be present" + assert session_descriptor["value"] == f"{_agent_id}:{_session_id}" + assert session_descriptor["rate_limit"]["requests_per_unit"] == 20 + assert session_descriptor["rate_limit"]["tokens_per_unit"] == 2000 + + +@pytest.mark.asyncio +async def test_agent_rate_limit_tpm_increment_on_success(monkeypatch): + """ + Test that async_log_success_event increments agent and session + TPM counters when agent_id and session_id are in metadata. + """ + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_tpm_test" + _session_id = "sess_tpm_test" + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + def mock_get_rate_limit_type(): + return "total" + + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type + ) + + mock_usage = Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50) + mock_response = ModelResponse( + id="mock-response", + object="chat.completion", + created=int(datetime.now().timestamp()), + model="gpt-4", + usage=mock_usage, + choices=[], + ) + + mock_kwargs = { + "standard_logging_object": { + "metadata": { + "user_api_key_hash": _api_key, + "user_api_key_user_id": None, + "user_api_key_team_id": None, + "user_api_key_end_user_id": None, + "agent_id": _agent_id, + "session_id": _session_id, + } + }, + "model": "gpt-4", + } + + captured_operations = [] + + async def mock_increment_pipeline(increment_list, **kwargs): + captured_operations.extend(increment_list) + return True + + monkeypatch.setattr( + parallel_request_handler.internal_usage_cache.dual_cache, + "async_increment_cache_pipeline", + mock_increment_pipeline, + ) + + await parallel_request_handler.async_log_success_event( + kwargs=mock_kwargs, + response_obj=mock_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + agent_tpm_op = None + session_tpm_op = None + for op in captured_operations: + if op["key"] == f"{{agent:{_agent_id}}}:tokens": + agent_tpm_op = op + elif op["key"] == f"{{agent_session:{_agent_id}:{_session_id}}}:tokens": + session_tpm_op = op + + assert agent_tpm_op is not None, "Agent TPM increment should be present" + assert agent_tpm_op["increment_value"] == 50 + + assert session_tpm_op is not None, "Session TPM increment should be present" + assert session_tpm_op["increment_value"] == 50 + + +@pytest.mark.asyncio +async def test_agent_rate_limit_429_on_over_limit(monkeypatch, time_controller): + """ + Test end-to-end that agent rate limiting returns 429 when the agent + RPM limit is exceeded. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "2") + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_429_test" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=time_controller.now, + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="rate-limited-agent", + agent_card_params={"name": "Rate Limited Agent"}, + rpm_limit=2, + ) + + window_starts: Dict[str, int] = {} + request_counts: Dict[str, int] = {} + + async def mock_batch_rate_limiter(*args, **kwargs): + keys = kwargs.get("keys") if kwargs else args[0] + args_list = kwargs.get("args") if kwargs else args[1] + now = args_list[0] + window_size = args_list[1] + results = [] + for i in range(0, len(keys), 2): + window_key = keys[i] + counter_key = keys[i + 1] + prev_window = window_starts.get(window_key) + prev_counter = request_counts.get(counter_key, 0) + if prev_window is None or (now - prev_window) >= window_size: + window_starts[window_key] = now + new_counter = 1 + request_counts[counter_key] = new_counter + await local_cache.async_set_cache( + key=window_key, value=now, ttl=window_size + ) + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) + else: + new_counter = prev_counter + 1 + request_counts[counter_key] = new_counter + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) + results.append(now) + results.append(new_counter) + return results + + parallel_request_handler.batch_rate_limiter_script = mock_batch_rate_limiter + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + with pytest.raises(HTTPException) as exc_info: + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + assert exc_info.value.status_code == 429 + assert "agent" in exc_info.value.detail + + class TestGetTotalTokensFromUsageCacheExclusion: """ Tests for _get_total_tokens_from_usage cache token exclusion. diff --git a/ui/litellm-dashboard/src/components/agents.tsx b/ui/litellm-dashboard/src/components/agents.tsx index de544acb9f..b60af25413 100644 --- a/ui/litellm-dashboard/src/components/agents.tsx +++ b/ui/litellm-dashboard/src/components/agents.tsx @@ -1,14 +1,26 @@ import React, { useState, useEffect } from "react"; -import { Button } from "@tremor/react"; -import { Modal, Alert, Switch, Tooltip } from "antd"; +import { + Button, + Card, + Table, + TableBody, + TableCell, + TableHead, + TableHeaderCell, + TableRow, + Badge, + Text, +} from "@tremor/react"; +import { Modal, Alert, Tooltip, Skeleton } from "antd"; import { CheckCircleOutlined } from "@ant-design/icons"; import { getAgentsList, deleteAgentCall, keyListCall } from "./networking"; import AddAgentForm from "./agents/add_agent_form"; -import AgentCardGrid from "./agents/agent_card_grid"; import { isAdminRole } from "@/utils/roles"; import AgentInfoView from "./agents/agent_info"; import NotificationsManager from "./molecules/notifications_manager"; import { Agent, AgentKeyInfo } from "./agents/types"; +import { formatNumberWithCommas } from "@/utils/dataUtils"; +import TableIconActionButton from "./common_components/IconActionButton/TableIconActionButtons/TableIconActionButton"; interface AgentsPanelProps { accessToken: string | null; @@ -136,6 +148,14 @@ const AgentsPanel: React.FC = ({ accessToken, userRole }) => { setAgentToDelete(null); }; + const sortedAgents = [...agentsList].sort((a, b) => { + const dateA = a.created_at ? new Date(a.created_at).getTime() : 0; + const dateB = b.created_at ? new Date(b.created_at).getTime() : 0; + return dateB - dateA; + }); + + const columnCount = isAdmin ? 7 : 6; + return (
@@ -177,16 +197,84 @@ const AgentsPanel: React.FC = ({ accessToken, userRole }) => { isAdmin={isAdmin} /> ) : ( - setSelectedAgentId(id)} - /> + + {isLoading ? ( + + ) : ( + + + + Agent Name + Agent ID + Spend (USD) + Model + Created + Status + {isAdmin && Actions} + + + + {sortedAgents.length === 0 ? ( + + + No agents found. Click "+ Add New Agent" to create one. + + + ) : ( + sortedAgents.map((agent) => ( + + + {agent.agent_name} + + + + + + + + {formatNumberWithCommas(agent.spend, 4)} + + + + {agent.litellm_params?.model || "N/A"} + + + + + {agent.created_at + ? new Date(agent.created_at).toLocaleDateString() + : "N/A"} + + + + {keyInfoMap[agent.agent_id]?.has_key ? ( + Active + ) : ( + Needs Setup + )} + + {isAdmin && ( + + handleDeleteClick(agent.agent_id, agent.agent_name)} + /> + + )} + + )) + )} + +
+ )} + )} = ({ accessToken, userRole }) => { }; export default AgentsPanel; - diff --git a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx index 10fbfa4615..0cec0331f4 100644 --- a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Modal, Form, message, Select, Input, Steps, Radio, Tag, Divider } from "antd"; +import { Modal, Form, message, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd"; import { Button } from "@tremor/react"; import { CheckCircleFilled, KeyOutlined, RobotOutlined, AppstoreOutlined, InfoCircleOutlined } from "@ant-design/icons"; import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; @@ -60,6 +60,12 @@ const AddAgentForm: React.FC = ({ const [createdKeyValue, setCreatedKeyValue] = useState(null); const [assignedKeyAlias, setAssignedKeyAlias] = useState(null); + // Tracing & guardrails state + const [requireTraceIdInbound, setRequireTraceIdInbound] = useState(false); + const [requireTraceIdOutbound, setRequireTraceIdOutbound] = useState(false); + const [maxIterations, setMaxIterations] = useState(null); + const [maxBudgetPerSession, setMaxBudgetPerSession] = useState(null); + // Fetch agent type metadata on mount useEffect(() => { const fetchMetadata = async () => { @@ -218,6 +224,19 @@ const AddAgentForm: React.FC = ({ } } + // Wire trace-id flags and budget controls into agent litellm_params (before create call) + if (requireTraceIdInbound || requireTraceIdOutbound) { + if (!agentData.litellm_params) agentData.litellm_params = {}; + if (requireTraceIdInbound) { + agentData.litellm_params.require_trace_id_on_calls_to_agent = true; + } + if (requireTraceIdOutbound) { + agentData.litellm_params.require_trace_id_on_calls_by_agent = true; + if (maxIterations) agentData.litellm_params.max_iterations = maxIterations; + if (maxBudgetPerSession) agentData.litellm_params.max_budget_per_session = maxBudgetPerSession; + } + } + const agentResponse = await createAgentCall(accessToken, agentData); const agentId: string = agentResponse.agent_id; const agentName: string = agentResponse.agent_name || values.agent_name || agentId; @@ -267,6 +286,10 @@ const AddAgentForm: React.FC = ({ setCreatedAgentName(""); setCreatedKeyValue(null); setAssignedKeyAlias(null); + setRequireTraceIdInbound(false); + setRequireTraceIdOutbound(false); + setMaxIterations(null); + setMaxBudgetPerSession(null); onClose(); }; @@ -315,6 +338,122 @@ const AddAgentForm: React.FC = ({
)} + + Tracing, + children: ( +
+
+
+ + Require x-litellm-trace-id on calls TO this agent + +

+ Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent). +

+
+ +
+ +
+
+ + Require x-litellm-trace-id on calls BY this agent + +

+ Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking. +

+
+ { + setRequireTraceIdOutbound(checked); + if (!checked) { + setMaxIterations(null); + setMaxBudgetPerSession(null); + } + }} + /> +
+
+ ), + }, + { + key: "budgets_and_rate_limits", + label: Budgets & Rate Limits, + children: ( +
+ {!requireTraceIdOutbound && ( +
+ Enable "Require x-litellm-trace-id on calls BY this agent" in Tracing to configure budgets and rate limits. +
+ )} + +
Session Budgets
+
+
+ + setMaxIterations(val)} + /> +

Hard cap on LLM calls per session

+
+
+ + setMaxBudgetPerSession(val)} + /> +

Max spend per trace before returning 429

+
+
+ + + +
Agent Rate Limits
+

+ Global rate limits applied across all callers of this agent. +

+
+ + + + + + +
+ +
Per-Session Rate Limits
+

+ Rate limits per session (x-litellm-trace-id). Each session gets its own counters. +

+
+ + + + + + +
+
+ ), + }, + ]} />
); @@ -456,6 +595,7 @@ const AddAgentForm: React.FC = ({ ) : null} + ); @@ -643,7 +783,7 @@ const AddAgentForm: React.FC = ({ {/* Step indicator */} - + diff --git a/ui/litellm-dashboard/src/components/agents/agent_config.ts b/ui/litellm-dashboard/src/components/agents/agent_config.ts index 01041c5cee..e87b191b19 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_config.ts +++ b/ui/litellm-dashboard/src/components/agents/agent_config.ts @@ -29,6 +29,7 @@ export const AGENT_FORM_CONFIG: { optional: SectionConfig; litellm: SectionConfig; cost: SectionConfig; + tracing: SectionConfig; } = { basic: { key: "basic", @@ -174,6 +175,19 @@ export const AGENT_FORM_CONFIG: { }, ], }, + tracing: { + key: "tracing", + title: "Tracing", + fields: [ + { + name: "enable_tracing", + label: "Enable Tracing", + type: "switch", + defaultValue: false, + tooltip: "Enable request tracing for this agent", + }, + ], + }, }; export const SKILL_FIELD_CONFIG = { @@ -269,6 +283,10 @@ export const buildAgentDataFromForm = (values: any, existingAgent?: any) => { agentData.litellm_params = params; } + if (values.tpm_limit != null) agentData.tpm_limit = values.tpm_limit; + if (values.rpm_limit != null) agentData.rpm_limit = values.rpm_limit; + if (values.session_tpm_limit != null) agentData.session_tpm_limit = values.session_tpm_limit; + if (values.session_rpm_limit != null) agentData.session_rpm_limit = values.session_rpm_limit; // static_headers: convert [{header, value}, ...] → {header: value, ...} if (Array.isArray(values.static_headers) && values.static_headers.length > 0) { const staticHeaders: Record = {}; @@ -319,6 +337,10 @@ export const parseAgentForForm = (agent: any) => { cost_per_query: agent.litellm_params?.cost_per_query, input_cost_per_token: agent.litellm_params?.input_cost_per_token, output_cost_per_token: agent.litellm_params?.output_cost_per_token, + tpm_limit: agent.tpm_limit, + rpm_limit: agent.rpm_limit, + session_tpm_limit: agent.session_tpm_limit, + session_rpm_limit: agent.session_rpm_limit, // static_headers: {key: value} → [{header, value}, ...] static_headers: agent.static_headers ? Object.entries(agent.static_headers as Record).map(([header, value]) => ({ diff --git a/ui/litellm-dashboard/src/components/agents/agent_info.tsx b/ui/litellm-dashboard/src/components/agents/agent_info.tsx index deb4f90037..b41e318a76 100644 --- a/ui/litellm-dashboard/src/components/agents/agent_info.tsx +++ b/ui/litellm-dashboard/src/components/agents/agent_info.tsx @@ -1,6 +1,6 @@ import React, { useState, useEffect } from "react"; import { Card, Title, Text, Button as TremorButton, Tab, TabGroup, TabList, TabPanel, TabPanels} from "@tremor/react"; -import { Form, Input, Button as AntButton, message, Spin, Descriptions } from "antd"; +import { Form, Input, InputNumber, Button as AntButton, message, Spin, Descriptions, Divider } from "antd"; import { ArrowLeftIcon } from "@heroicons/react/outline"; import { getAgentInfo, patchAgentCall, getAgentCreateMetadata, AgentCreateInfo } from "../networking"; import { Agent } from "./types"; @@ -201,6 +201,10 @@ const AgentInfoView: React.FC = ({ {agent.agent_card_params?.documentationUrl && ( {agent.agent_card_params.documentationUrl} )} + {agent.tpm_limit ?? "Unlimited"} + {agent.rpm_limit ?? "Unlimited"} + {agent.session_tpm_limit ?? "Unlimited"} + {agent.session_rpm_limit ?? "Unlimited"} {formatDate(agent.created_at)} {formatDate(agent.updated_at)} @@ -295,6 +299,25 @@ const AgentInfoView: React.FC = ({ )} + + Rate Limits +
+ + + + + + +
+
+ + + + + + +
+
{ setIsEditing(false); diff --git a/ui/litellm-dashboard/src/components/agents/dynamic_agent_form_fields.tsx b/ui/litellm-dashboard/src/components/agents/dynamic_agent_form_fields.tsx index 55a4a62953..1138b0de73 100644 --- a/ui/litellm-dashboard/src/components/agents/dynamic_agent_form_fields.tsx +++ b/ui/litellm-dashboard/src/components/agents/dynamic_agent_form_fields.tsx @@ -118,7 +118,7 @@ export const buildDynamicAgentData = ( litellmParams.model = model; } - return { + const agentData: Record = { agent_name: values.agent_name, agent_card_params: { protocolVersion: "1.0", @@ -140,6 +140,13 @@ export const buildDynamicAgentData = ( }, litellm_params: litellmParams, }; + + if (values.tpm_limit != null) agentData.tpm_limit = values.tpm_limit; + if (values.rpm_limit != null) agentData.rpm_limit = values.rpm_limit; + if (values.session_tpm_limit != null) agentData.session_tpm_limit = values.session_tpm_limit; + if (values.session_rpm_limit != null) agentData.session_rpm_limit = values.session_rpm_limit; + + return agentData; }; export default DynamicAgentFormFields; diff --git a/ui/litellm-dashboard/src/components/agents/types.ts b/ui/litellm-dashboard/src/components/agents/types.ts index 2e903b1402..3e27177c81 100644 --- a/ui/litellm-dashboard/src/components/agents/types.ts +++ b/ui/litellm-dashboard/src/components/agents/types.ts @@ -23,6 +23,11 @@ export interface Agent { [key: string]: any; }; object_permission?: AgentObjectPermission; + spend?: number; + tpm_limit?: number | null; + rpm_limit?: number | null; + session_tpm_limit?: number | null; + session_rpm_limit?: number | null; created_at?: string; updated_at?: string; created_by?: string; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 244144f56b..0cd94d9895 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -898,19 +898,24 @@ export const keyCreateForAgentCall = async ( agentId: string, keyAlias: string, models: string[], + metadata?: Record, ) => { const url = proxyBaseUrl ? `${proxyBaseUrl}/key/generate` : `/key/generate`; + const body: Record = { + agent_id: agentId, + key_alias: keyAlias, + models: models.length > 0 ? models : [], + }; + if (metadata && Object.keys(metadata).length > 0) { + body.metadata = metadata; + } const response = await fetch(url, { method: "POST", headers: { [globalLitellmHeaderName]: `Bearer ${accessToken}`, "Content-Type": "application/json", }, - body: JSON.stringify({ - agent_id: agentId, - key_alias: keyAlias, - models: models.length > 0 ? models : [], - }), + body: JSON.stringify(body), }); if (!response.ok) { @@ -7696,6 +7701,10 @@ export const patchAgentCall = async ( agent_name?: string; litellm_params?: Record; agent_card_params?: Record; + tpm_limit?: number | null; + rpm_limit?: number | null; + session_tpm_limit?: number | null; + session_rpm_limit?: number | null; }, ) => { try {