tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-17 02:17:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Agents - add max budget + tpm/rpm limiting per agent AND per agent session (#22849)

Browse Source 
* 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 <noreply@anthropic.com>

* fix: qa'ed RPM limiting on agents

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Krish Dholakia
2026-03-07 19:12:42 -08:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 03ca98123f
commit cf439c269c
41 changed files with 2596 additions and 557 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/my-website/docs/a2a.md
+1
View File
@@ -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
docs/my-website/docs/a2a_iteration_budgets.md
+188
View File
@@ -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:
<Tabs>
<TabItem value="ui" label="Via UI">
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**
</TabItem>
<TabItem value="api" label="Via API">
```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": "<agent_id_from_step_1>",
"key_alias": "my-research-agent-key"
}'
```
</TabItem>
</Tabs>
### 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 |
docs/my-website/docs/proxy/users.md
+130
View File
@@ -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
</Tabs>
### 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`
<Tabs>
<TabItem value="agent-rate-limits" label="Agent Rate Limits">
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
}'
```
</TabItem>
<TabItem value="session-rate-limits" label="Session Rate Limits">
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
}'
```
</TabItem>
<TabItem value="session-budgets" label="Session Budgets">
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.
</TabItem>
</Tabs>
:::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/<agent_id>' \
-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`
</TabItem>
<TabItem value="per-agent" label="Per Agent">
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.
</TabItem>
<TabItem value="per-end-user" label="For customers">
docs/my-website/sidebars.js
+2 -1
View File
@@ -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",
litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql
+3
View File
@@ -0,0 +1,3 @@
-- AlterTable
ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql
+5
View File
@@ -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;
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
+5
View File
@@ -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")
litellm/batches/batch_utils.py
+3 -4
View File
@@ -198,9 +198,8 @@ async def _get_batch_output_file_content_as_dictionary(
Required for Azure and other providers that need authentication
"""
from litellm.files.main import afile_content
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
from litellm.proxy.openai_files_endpoints.common_utils import \
_is_base64_encoded_unified_file_id
if custom_llm_provider == "vertex_ai":
raise ValueError("Vertex AI does not support file content retrieval")
@@ -227,7 +226,7 @@ async def _get_batch_output_file_content_as_dictionary(
credentials = _extract_file_access_credentials(litellm_params)
file_content_kwargs.update(credentials)
_file_content = await afile_content(**file_content_kwargs)
_file_content = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType]
return _get_file_content_as_dictionary(_file_content.content)
litellm/fine_tuning/main.py
+21 -36
View File
@@ -126,6 +126,18 @@ async def acreate_fine_tuning_job(
raise e
def _build_fine_tuning_job_data(model, training_file, hyperparameters, suffix, validation_file, integrations, seed):
return FineTuningJobCreate(
model=model,
training_file=training_file,
hyperparameters=hyperparameters,
suffix=suffix,
validation_file=validation_file,
integrations=integrations,
seed=seed,
)
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,
litellm/proxy/_types.py
+2
View File
@@ -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):
litellm/proxy/agent_endpoints/a2a_endpoints.py
+38 -19
View File
@@ -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(
litellm/proxy/agent_endpoints/agent_registry.py
+17 -3
View File
@@ -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}
litellm/proxy/agent_endpoints/endpoints.py
+38 -16
View File
@@ -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,
)
)
litellm/proxy/auth/auth_checks.py
+173 -145
View File
@@ -11,7 +11,8 @@ Run checks for:
import asyncio
import re
import time
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union,
cast)
from fastapi import HTTPException, Request, status
from pydantic import BaseModel
@@ -20,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
litellm/proxy/auth/user_api_key_auth.py
+64 -84
View File
@@ -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)
litellm/proxy/db/db_spend_update_writer.py
+110 -36
View File
@@ -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,
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
+25 -25
View File
@@ -10,33 +10,31 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
from litellm._logging import verbose_proxy_logger
from litellm.caching import RedisCache
from litellm.constants import (
MAX_REDIS_BUFFER_DEQUEUE_COUNT,
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
REDIS_UPDATE_BUFFER_KEY,
)
from litellm.constants import (MAX_REDIS_BUFFER_DEQUEUE_COUNT,
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY,
REDIS_UPDATE_BUFFER_KEY)
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
from litellm.proxy._types import (
DailyTagSpendTransaction,
DailyTeamSpendTransaction,
DailyUserSpendTransaction,
DailyOrganizationSpendTransaction,
DailyEndUserSpendTransaction,
DBSpendUpdateTransactions,
DailyAgentSpendTransaction,
)
from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import (
DailySpendUpdateQueue,
)
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
from litellm.proxy._types import (DailyAgentSpendTransaction,
DailyEndUserSpendTransaction,
DailyOrganizationSpendTransaction,
DailyTagSpendTransaction,
DailyTeamSpendTransaction,
DailyUserSpendTransaction,
DBSpendUpdateTransactions)
from litellm.proxy.db.db_transaction_queue.base_update_queue import \
service_logger_obj
from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import \
DailySpendUpdateQueue
from litellm.proxy.db.db_transaction_queue.spend_update_queue import \
SpendUpdateQueue
from litellm.secret_managers.main import str_to_bool
from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation
from litellm.types.caching import (RedisPipelineLpopOperation,
RedisPipelineRpushOperation)
from litellm.types.services import ServiceTypes
if TYPE_CHECKING:
@@ -579,6 +577,7 @@ class RedisUpdateBuffer:
team_member_list_transactions={},
org_list_transactions={},
tag_list_transactions={},
agent_list_transactions={},
)
# Define the transaction fields to process
@@ -590,6 +589,7 @@ class RedisUpdateBuffer:
"team_member_list_transactions",
"org_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
]
# Loop through each transaction and combine the values
litellm/proxy/db/db_transaction_queue/spend_update_queue.py
+9 -8
View File
@@ -3,15 +3,10 @@ from typing import Dict, List, Optional
from litellm._logging import verbose_proxy_logger
from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE
from litellm.proxy._types import (
DBSpendUpdateTransactions,
Litellm_EntityType,
SpendUpdateQueueItem,
)
from litellm.proxy._types import (DBSpendUpdateTransactions,
Litellm_EntityType, SpendUpdateQueueItem)
from litellm.proxy.db.db_transaction_queue.base_update_queue import (
BaseUpdateQueue,
service_logger_obj,
)
BaseUpdateQueue, service_logger_obj)
from litellm.types.services import ServiceTypes
@@ -145,6 +140,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
team_member_list_transactions={},
org_list_transactions={},
tag_list_transactions={},
agent_list_transactions={},
)
# Map entity types to their corresponding transaction dictionary keys
@@ -156,6 +152,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions",
Litellm_EntityType.ORGANIZATION: "org_list_transactions",
Litellm_EntityType.TAG: "tag_list_transactions",
Litellm_EntityType.AGENT: "agent_list_transactions",
}
for update in updates:
@@ -207,6 +204,10 @@ class SpendUpdateQueue(BaseUpdateQueue):
transactions_dict = db_spend_update_transactions[
"tag_list_transactions"
]
elif dict_key == "agent_list_transactions":
transactions_dict = db_spend_update_transactions[
"agent_list_transactions"
]
else:
continue
litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py
+36 -46
View File
@@ -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)}")
litellm/proxy/hooks/__init__.py
+4
View File
@@ -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 ##
litellm/proxy/hooks/max_budget_per_session_limiter.py
+271
View File
@@ -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:<session_id>}: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
litellm/proxy/hooks/max_iterations_limiter.py
+21 -6
View File
@@ -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:
litellm/proxy/hooks/parallel_request_limiter_v3.py
+141 -36
View File
@@ -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:
litellm/proxy/litellm_pre_call_utils.py
+6
View File
@@ -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
)
litellm/proxy/proxy_server.py
+20 -8
View File
@@ -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)
litellm/proxy/schema.prisma
+5
View File
@@ -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")
litellm/proxy/utils.py
+1 -1
View File
@@ -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()
litellm/router.py
+1 -5
View File
@@ -164,11 +164,7 @@ from litellm.types.utils import (
)
from litellm.types.utils import ModelInfo
from litellm.types.utils import ModelInfo as ModelMapInfo
from litellm.types.utils import (
ModelResponseStream,
StandardLoggingPayload,
Usage,
)
from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage
from litellm.utils import (
CustomStreamWrapper,
EmbeddingResponse,
litellm/types/agents.py
+13
View File
@@ -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
schema.prisma
+5
View File
@@ -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")
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
+132 -1
View File
@@ -9,7 +9,7 @@ sys.path.insert(
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch, call
from unittest.mock import AsyncMock, MagicMock, call, patch
import pytest
@@ -516,6 +516,131 @@ async def test_update_tag_db_without_prisma_client():
assert writer.spend_update_queue.add_update.call_count == 0
@pytest.mark.asyncio
async def test_update_agent_db_enqueues_agent_spend():
"""
Test that _update_agent_db enqueues a SpendUpdateQueueItem with entity_type=AGENT.
"""
from litellm.proxy._types import Litellm_EntityType
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
agent_id = "agent-123"
response_cost = 0.1
writer.spend_update_queue.add_update = AsyncMock()
await writer._update_agent_db(
response_cost=response_cost,
agent_id=agent_id,
prisma_client=mock_prisma,
)
writer.spend_update_queue.add_update.assert_called_once()
call_args = writer.spend_update_queue.add_update.call_args[1]
assert call_args["update"]["entity_type"] == Litellm_EntityType.AGENT
assert call_args["update"]["entity_id"] == agent_id
assert call_args["update"]["response_cost"] == response_cost
@pytest.mark.asyncio
async def test_update_agent_db_skips_when_agent_id_none():
"""_update_agent_db does not enqueue when agent_id is None."""
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
writer.spend_update_queue.add_update = AsyncMock()
await writer._update_agent_db(
response_cost=0.05,
agent_id=None,
prisma_client=mock_prisma,
)
writer.spend_update_queue.add_update.assert_not_called()
@pytest.mark.asyncio
async def test_update_agent_db_skips_when_prisma_client_none():
"""_update_agent_db does not enqueue when prisma_client is None."""
writer = DBSpendUpdateWriter()
writer.spend_update_queue.add_update = AsyncMock()
await writer._update_agent_db(
response_cost=0.05,
agent_id="agent-456",
prisma_client=None,
)
writer.spend_update_queue.add_update.assert_not_called()
@pytest.mark.asyncio
async def test_commit_spend_updates_to_db_increments_agent_spend():
"""
Test that _commit_spend_updates_to_db calls litellm_agentstable.update_many
with spend increment when agent_list_transactions is present.
"""
db_writer = DBSpendUpdateWriter()
mock_batcher = MagicMock()
mock_batcher.litellm_verificationtoken = MagicMock()
mock_batcher.litellm_verificationtoken.update_many = MagicMock()
mock_batcher.litellm_usertable = MagicMock()
mock_batcher.litellm_usertable.update_many = MagicMock()
mock_batcher.litellm_teamtable = MagicMock()
mock_batcher.litellm_teamtable.update_many = MagicMock()
mock_batcher.litellm_teammembership = MagicMock()
mock_batcher.litellm_teammembership.update_many = MagicMock()
mock_batcher.litellm_organizationtable = MagicMock()
mock_batcher.litellm_organizationtable.update_many = MagicMock()
mock_batcher.litellm_tagtable = MagicMock()
mock_batcher.litellm_tagtable.update_many = MagicMock()
mock_batcher.litellm_agentstable = MagicMock()
mock_batcher.litellm_agentstable.update_many = MagicMock()
mock_transaction = AsyncMock()
mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
mock_transaction.__aexit__ = AsyncMock(return_value=False)
mock_transaction.batch_ = MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=mock_batcher),
__aexit__=AsyncMock(return_value=False),
)
)
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
mock_proxy_logging = MagicMock()
agent_id = "agent-789"
response_cost = 0.25
db_spend_update_transactions = {
"user_list_transactions": {},
"end_user_list_transactions": {},
"key_list_transactions": {},
"team_list_transactions": {},
"team_member_list_transactions": {},
"org_list_transactions": {},
"tag_list_transactions": {},
"agent_list_transactions": {agent_id: response_cost},
}
with patch("litellm.proxy.utils._raise_failed_update_spend_exception"):
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=mock_proxy_logging,
db_spend_update_transactions=db_spend_update_transactions,
)
mock_batcher.litellm_agentstable.update_many.assert_called_once()
call_kwargs = mock_batcher.litellm_agentstable.update_many.call_args[1]
assert call_kwargs["where"] == {"agent_id": agent_id}
assert call_kwargs["data"] == {"spend": {"increment": response_cost}}
@pytest.mark.asyncio
async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
"""
@@ -1048,6 +1173,8 @@ async def test_commit_key_spend_updates_includes_last_active():
mock_batcher.litellm_teamtable.update_many = MagicMock()
mock_batcher.litellm_organizationtable = MagicMock()
mock_batcher.litellm_organizationtable.update_many = MagicMock()
mock_batcher.litellm_agentstable = MagicMock()
mock_batcher.litellm_agentstable.update_many = MagicMock()
mock_proxy_logging = MagicMock()
@@ -1059,6 +1186,7 @@ async def test_commit_key_spend_updates_includes_last_active():
"team_member_list_transactions": {},
"org_list_transactions": {},
"tag_list_transactions": {},
"agent_list_transactions": {},
}
before_call = datetime.now(timezone.utc)
@@ -1142,6 +1270,7 @@ async def test_batch_database_updates_isolation_on_failure():
db_writer._update_team_db = AsyncMock()
db_writer._update_org_db = AsyncMock()
db_writer._update_tag_db = AsyncMock()
db_writer._update_agent_db = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock()
@@ -1169,6 +1298,7 @@ async def test_batch_database_updates_isolation_on_failure():
db_writer._update_team_db.assert_awaited_once()
db_writer._update_org_db.assert_awaited_once()
db_writer._update_tag_db.assert_awaited_once()
db_writer._update_agent_db.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_user_transaction.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_end_user_transaction.assert_awaited_once()
db_writer.add_spend_log_transaction_to_daily_agent_transaction.assert_awaited_once()
@@ -1203,6 +1333,7 @@ async def test_daily_agent_receives_deepcopied_payload():
db_writer._update_team_db = AsyncMock()
db_writer._update_org_db = AsyncMock()
db_writer._update_tag_db = AsyncMock()
db_writer._update_agent_db = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock(
tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py
+165
View File
@@ -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
tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py
+104 -54
View File
@@ -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
tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py
+521
View File
@@ -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.
ui/litellm-dashboard/src/components/agents.tsx
+101 -14
View File
@@ -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<AgentsPanelProps> = ({ 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 (
<div className="w-full mx-auto flex-auto overflow-y-auto m-8 p-2">
<div className="flex flex-col gap-2 mb-4">
@@ -177,16 +197,84 @@ const AgentsPanel: React.FC<AgentsPanelProps> = ({ accessToken, userRole }) => {
isAdmin={isAdmin}
/>
) : (
<AgentCardGrid
agentsList={agentsList}
keyInfoMap={keyInfoMap}
isLoading={isLoading}
onDeleteClick={handleDeleteClick}
accessToken={accessToken}
onAgentUpdated={fetchAgents}
isAdmin={isAdmin}
onAgentClick={(id) => setSelectedAgentId(id)}
/>
<Card>
{isLoading ? (
<Skeleton active paragraph={{ rows: 3 }} />
) : (
<Table>
<TableHead>
<TableRow>
<TableHeaderCell>Agent Name</TableHeaderCell>
<TableHeaderCell>Agent ID</TableHeaderCell>
<TableHeaderCell>Spend (USD)</TableHeaderCell>
<TableHeaderCell>Model</TableHeaderCell>
<TableHeaderCell>Created</TableHeaderCell>
<TableHeaderCell>Status</TableHeaderCell>
{isAdmin && <TableHeaderCell>Actions</TableHeaderCell>}
</TableRow>
</TableHead>
<TableBody>
{sortedAgents.length === 0 ? (
<TableRow>
<TableCell colSpan={columnCount}>
<Text className="text-center">No agents found. Click &quot;+ Add New Agent&quot; to create one.</Text>
</TableCell>
</TableRow>
) : (
sortedAgents.map((agent) => (
<TableRow key={agent.agent_id}>
<TableCell>
<Text>{agent.agent_name}</Text>
</TableCell>
<TableCell>
<Tooltip title={agent.agent_id}>
<Button
size="xs"
variant="light"
className="font-mono text-blue-500 bg-blue-50 hover:bg-blue-100 text-xs font-normal px-2 py-0.5 text-left overflow-hidden truncate max-w-[200px]"
onClick={() => setSelectedAgentId(agent.agent_id)}
>
{agent.agent_id.slice(0, 7)}...
</Button>
</Tooltip>
</TableCell>
<TableCell>
<Text>{formatNumberWithCommas(agent.spend, 4)}</Text>
</TableCell>
<TableCell>
<Badge size="xs" color="blue">
{agent.litellm_params?.model || "N/A"}
</Badge>
</TableCell>
<TableCell>
<Text>
{agent.created_at
? new Date(agent.created_at).toLocaleDateString()
: "N/A"}
</Text>
</TableCell>
<TableCell>
{keyInfoMap[agent.agent_id]?.has_key ? (
<Badge color="green">Active</Badge>
) : (
<Badge color="yellow">Needs Setup</Badge>
)}
</TableCell>
{isAdmin && (
<TableCell>
<TableIconActionButton
variant="Delete"
onClick={() => handleDeleteClick(agent.agent_id, agent.agent_name)}
/>
</TableCell>
)}
</TableRow>
))
)}
</TableBody>
</Table>
)}
</Card>
)}
<AddAgentForm
@@ -215,4 +303,3 @@ const AgentsPanel: React.FC<AgentsPanelProps> = ({ accessToken, userRole }) => {
};
export default AgentsPanel;
ui/litellm-dashboard/src/components/agents/add_agent_form.tsx
+142 -2
View File
@@ -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<AddAgentFormProps> = ({
const [createdKeyValue, setCreatedKeyValue] = useState<string | null>(null);
const [assignedKeyAlias, setAssignedKeyAlias] = useState<string | null>(null);
// Tracing & guardrails state
const [requireTraceIdInbound, setRequireTraceIdInbound] = useState(false);
const [requireTraceIdOutbound, setRequireTraceIdOutbound] = useState(false);
const [maxIterations, setMaxIterations] = useState<number | null>(null);
const [maxBudgetPerSession, setMaxBudgetPerSession] = useState<number | null>(null);
// Fetch agent type metadata on mount
useEffect(() => {
const fetchMetadata = async () => {
@@ -218,6 +224,19 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
}
}
// 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<AddAgentFormProps> = ({
setCreatedAgentName("");
setCreatedKeyValue(null);
setAssignedKeyAlias(null);
setRequireTraceIdInbound(false);
setRequireTraceIdOutbound(false);
setMaxIterations(null);
setMaxBudgetPerSession(null);
onClose();
};
@@ -315,6 +338,122 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
</div>
)}
</Form.Item>
<Collapse ghost className="mt-6" items={[
{
key: "tracing",
label: <span className="text-sm font-medium text-gray-700">Tracing</span>,
children: (
<div className="space-y-4">
<div className="flex items-center justify-between">
<div>
<span className="text-sm font-medium text-gray-700">
Require x-litellm-trace-id on calls TO this agent
</span>
<p className="text-xs text-gray-500 mt-1">
Only accept this agent being invoked with a trace-id (e.g. when used as a sub-agent).
</p>
</div>
<Switch
checked={requireTraceIdInbound}
onChange={setRequireTraceIdInbound}
/>
</div>
<div className="flex items-center justify-between">
<div>
<span className="text-sm font-medium text-gray-700">
Require x-litellm-trace-id on calls BY this agent
</span>
<p className="text-xs text-gray-500 mt-1">
Requires LLM/MCP calls made by this agent to include x-litellm-trace-id for session tracking.
</p>
</div>
<Switch
checked={requireTraceIdOutbound}
onChange={(checked) => {
setRequireTraceIdOutbound(checked);
if (!checked) {
setMaxIterations(null);
setMaxBudgetPerSession(null);
}
}}
/>
</div>
</div>
),
},
{
key: "budgets_and_rate_limits",
label: <span className="text-sm font-medium text-gray-700">Budgets &amp; Rate Limits</span>,
children: (
<div className="space-y-4">
{!requireTraceIdOutbound && (
<div className="p-3 bg-yellow-50 border border-yellow-200 rounded-lg text-sm text-yellow-800">
Enable &quot;Require x-litellm-trace-id on calls BY this agent&quot; in Tracing to configure budgets and rate limits.
</div>
)}
<div className="text-sm font-medium text-gray-700">Session Budgets</div>
<div className="grid grid-cols-2 gap-4">
<div>
<label className="text-sm text-gray-600 block mb-1">Max Iterations</label>
<InputNumber
className="w-full"
min={1}
placeholder="e.g. 25"
disabled={!requireTraceIdOutbound}
value={maxIterations}
onChange={(val) => setMaxIterations(val)}
/>
<p className="text-xs text-gray-400 mt-1">Hard cap on LLM calls per session</p>
</div>
<div>
<label className="text-sm text-gray-600 block mb-1">Max Budget Per Session ($)</label>
<InputNumber
className="w-full"
min={0.01}
step={0.5}
placeholder="e.g. 5.00"
disabled={!requireTraceIdOutbound}
value={maxBudgetPerSession}
onChange={(val) => setMaxBudgetPerSession(val)}
/>
<p className="text-xs text-gray-400 mt-1">Max spend per trace before returning 429</p>
</div>
</div>
<Divider className="my-2" />
<div className="text-sm font-medium text-gray-700">Agent Rate Limits</div>
<p className="text-xs text-gray-500">
Global rate limits applied across all callers of this agent.
</p>
<div className="grid grid-cols-2 gap-4">
<Form.Item label="TPM Limit" name="tpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 100000" disabled={!requireTraceIdOutbound} />
</Form.Item>
<Form.Item label="RPM Limit" name="rpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 100" disabled={!requireTraceIdOutbound} />
</Form.Item>
</div>
<div className="text-sm font-medium text-gray-700 mt-4">Per-Session Rate Limits</div>
<p className="text-xs text-gray-500">
Rate limits per session (x-litellm-trace-id). Each session gets its own counters.
</p>
<div className="grid grid-cols-2 gap-4">
<Form.Item label="Session TPM Limit" name="session_tpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 10000" disabled={!requireTraceIdOutbound} />
</Form.Item>
<Form.Item label="Session RPM Limit" name="session_rpm_limit" className="mb-0">
<InputNumber className="w-full" min={0} placeholder="e.g. 20" disabled={!requireTraceIdOutbound} />
</Form.Item>
</div>
</div>
),
},
]} />
</div>
);
@@ -456,6 +595,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
<DynamicAgentFormFields agentTypeInfo={selectedAgentTypeInfo} />
) : null}
</div>
</>
);
@@ -643,7 +783,7 @@ const AddAgentForm: React.FC<AddAgentFormProps> = ({
{/* Step indicator */}
<Steps current={currentStep} size="small" className="mb-8">
<Step title="Configure" />
<Step title="MCP Tools" />
<Step title="Agent Settings" />
<Step title="Assign Key" />
<Step title="Ready" />
</Steps>
ui/litellm-dashboard/src/components/agents/agent_config.ts
+22
View File
@@ -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<string, string> = {};
@@ -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<string, string>).map(([header, value]) => ({
ui/litellm-dashboard/src/components/agents/agent_info.tsx
+24 -1
View File
@@ -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<AgentInfoViewProps> = ({
{agent.agent_card_params?.documentationUrl && (
<Descriptions.Item label="Documentation URL">{agent.agent_card_params.documentationUrl}</Descriptions.Item>
)}
<Descriptions.Item label="TPM Limit">{agent.tpm_limit ?? "Unlimited"}</Descriptions.Item>
<Descriptions.Item label="RPM Limit">{agent.rpm_limit ?? "Unlimited"}</Descriptions.Item>
<Descriptions.Item label="Session TPM Limit">{agent.session_tpm_limit ?? "Unlimited"}</Descriptions.Item>
<Descriptions.Item label="Session RPM Limit">{agent.session_rpm_limit ?? "Unlimited"}</Descriptions.Item>
<Descriptions.Item label="Created At">{formatDate(agent.created_at)}</Descriptions.Item>
<Descriptions.Item label="Updated At">{formatDate(agent.updated_at)}</Descriptions.Item>
</Descriptions>
@@ -295,6 +299,25 @@ const AgentInfoView: React.FC<AgentInfoViewProps> = ({
<AgentFormFields showAgentName={true} />
)}
<Divider />
<Title className="mb-4">Rate Limits</Title>
<div className="grid grid-cols-2 gap-4">
<Form.Item label="TPM Limit" name="tpm_limit">
<InputNumber className="w-full" min={0} placeholder="Unlimited" />
</Form.Item>
<Form.Item label="RPM Limit" name="rpm_limit">
<InputNumber className="w-full" min={0} placeholder="Unlimited" />
</Form.Item>
</div>
<div className="grid grid-cols-2 gap-4">
<Form.Item label="Session TPM Limit" name="session_tpm_limit">
<InputNumber className="w-full" min={0} placeholder="Unlimited" />
</Form.Item>
<Form.Item label="Session RPM Limit" name="session_rpm_limit">
<InputNumber className="w-full" min={0} placeholder="Unlimited" />
</Form.Item>
</div>
<div className="flex justify-end gap-2 mt-6">
<AntButton onClick={() => {
setIsEditing(false);
ui/litellm-dashboard/src/components/agents/dynamic_agent_form_fields.tsx
+8 -1
View File
@@ -118,7 +118,7 @@ export const buildDynamicAgentData = (
litellmParams.model = model;
}
return {
const agentData: Record<string, any> = {
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;
ui/litellm-dashboard/src/components/agents/types.ts
+5
View File
@@ -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;
ui/litellm-dashboard/src/components/networking.tsx
+14 -5
View File
@@ -898,19 +898,24 @@ export const keyCreateForAgentCall = async (
agentId: string,
keyAlias: string,
models: string[],
metadata?: Record<string, any>,
) => {
const url = proxyBaseUrl ? `${proxyBaseUrl}/key/generate` : `/key/generate`;
const body: Record<string, any> = {
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<string, any>;
agent_card_params?: Record<string, any>;
tpm_limit?: number | null;
rpm_limit?: number | null;
session_tpm_limit?: number | null;
session_rpm_limit?: number | null;
},
) => {
try {