diff --git a/Dockerfile b/Dockerfile index 605e702d2ae..75ccff29663 100644 --- a/Dockerfile +++ b/Dockerfile @@ -49,7 +49,7 @@ USER root # Install runtime dependencies (libsndfile needed for audio processing on ARM64) RUN apk add --no-cache bash openssl tzdata nodejs npm python3 py3-pip libsndfile && \ - npm install -g npm@latest tar@7.5.8 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \ + npm install -g npm@latest tar@7.5.10 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \ # SECURITY FIX: npm bundles tar, glob, and brace-expansion at multiple nested # levels inside its dependency tree. `npm install -g ` only creates a # SEPARATE global package, it does NOT replace npm's internal copies. @@ -70,7 +70,15 @@ RUN apk add --no-cache bash openssl tzdata nodejs npm python3 py3-pip libsndfile find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ done && \ - npm cache clean --force + # SECURITY FIX: patch npm's own package.json metadata so scanners see the + # actual installed versions instead of the stale declared dependencies. + find /usr/local/lib /usr/lib -path "*/node_modules/npm/package.json" -exec \ + sed -i 's/"tar": "\^7\.5\.[0-9]*"/"tar": "^7.5.10"/g; s/"minimatch": "\^10\.[0-9.]*"/"minimatch": "^10.2.4"/g' {} + 2>/dev/null && \ + npm cache clean --force && \ + # Remove the apk-tracked npm so its stale SBOM metadata (tar 7.5.9) is + # no longer visible to image scanners. The globally installed npm@latest + # at /usr/local/lib/node_modules/npm/ remains fully functional. + { apk del --no-cache npm 2>/dev/null || true; } WORKDIR /app # Copy the current directory contents into the container at /app @@ -96,6 +104,7 @@ RUN find /usr/lib -type f -path "*/tornado/test/*" -delete && \ # npm with old vulnerable deps at /usr/lib/python3.*/site-packages/nodejs_wheel/. # Patch every copy of tar, glob, and brace-expansion inside that tree. RUN GLOBAL="$(npm root -g)" && \ + [ -n "$GLOBAL" ] || { echo "ERROR: npm root -g returned empty; aborting"; exit 1; } && \ find /usr/lib -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ done && \ diff --git a/docker/Dockerfile.custom_ui b/docker/Dockerfile.custom_ui index fb98846a6cc..4052c7a51bc 100644 --- a/docker/Dockerfile.custom_ui +++ b/docker/Dockerfile.custom_ui @@ -19,7 +19,7 @@ RUN apt-get update && apt-get upgrade -y \ libgnutls30 \ libc6 && \ apt-get install -y nodejs npm && \ - npm install -g npm@latest tar@7.5.8 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.1 diff@8.0.3 && \ + npm install -g npm@latest tar@7.5.10 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \ GLOBAL="$(npm root -g)" && \ find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ @@ -36,7 +36,10 @@ RUN apt-get update && apt-get upgrade -y \ find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ done && \ - npm cache clean --force + find /usr/local/lib /usr/lib -path "*/node_modules/npm/package.json" -exec \ + sed -i 's/"tar": "\^7\.5\.[0-9]*"/"tar": "^7.5.10"/g; s/"minimatch": "\^10\.[0-9.]*"/"minimatch": "^10.2.4"/g' {} + 2>/dev/null && \ + npm cache clean --force && \ + apt-get purge -y npm # Copy the UI source into the container COPY ./ui/litellm-dashboard /app/ui/litellm-dashboard diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index 371766bd9db..962d129e57f 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -50,7 +50,7 @@ USER root # Install runtime dependencies RUN apk add --no-cache bash openssl tzdata nodejs npm python3 py3-pip libsndfile && \ - npm install -g npm@latest tar@7.5.8 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.1 diff@8.0.3 && \ + npm install -g npm@latest tar@7.5.10 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 && \ GLOBAL="$(npm root -g)" && \ find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ @@ -67,7 +67,10 @@ RUN apk add --no-cache bash openssl tzdata nodejs npm python3 py3-pip libsndfile find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ done && \ - npm cache clean --force + find /usr/local/lib /usr/lib -path "*/node_modules/npm/package.json" -exec \ + sed -i 's/"tar": "\^7\.5\.[0-9]*"/"tar": "^7.5.10"/g; s/"minimatch": "\^10\.[0-9.]*"/"minimatch": "^10.2.4"/g' {} + 2>/dev/null && \ + npm cache clean --force && \ + { apk del --no-cache npm 2>/dev/null || true; } WORKDIR /app # Copy the current directory contents into the container at /app @@ -85,6 +88,7 @@ RUN pip install *.whl /wheels/* --no-index --find-links=/wheels/ && rm -f *.whl # npm with old vulnerable deps at /usr/lib/python3.*/site-packages/nodejs_wheel/. # Patch every copy of tar, glob, and brace-expansion inside that tree. RUN GLOBAL="$(npm root -g)" && \ + [ -n "$GLOBAL" ] || { echo "ERROR: npm root -g returned empty; aborting"; exit 1; } && \ find /usr/lib -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ done && \ diff --git a/docker/Dockerfile.dev b/docker/Dockerfile.dev index a5312dec9e3..cfc4c646ba2 100644 --- a/docker/Dockerfile.dev +++ b/docker/Dockerfile.dev @@ -75,7 +75,7 @@ RUN apt-get update && apt-get upgrade -y \ nodejs \ npm \ && rm -rf /var/lib/apt/lists/* \ - && npm install -g npm@latest tar@7.5.8 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.1 diff@8.0.3 \ + && npm install -g npm@latest tar@7.5.10 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 \ && GLOBAL="$(npm root -g)" \ && find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ @@ -92,7 +92,10 @@ RUN apt-get update && apt-get upgrade -y \ && find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ done \ - && npm cache clean --force + && find /usr/local/lib /usr/lib -path "*/node_modules/npm/package.json" -exec \ + sed -i 's/"tar": "\^7\.5\.[0-9]*"/"tar": "^7.5.10"/g; s/"minimatch": "\^10\.[0-9.]*"/"minimatch": "^10.2.4"/g' {} + 2>/dev/null \ + && npm cache clean --force \ + && apt-get purge -y npm WORKDIR /app @@ -114,6 +117,7 @@ RUN pip install --no-cache-dir *.whl /wheels/* --no-index --find-links=/wheels/ # npm with old vulnerable deps at /usr/lib/python3.*/site-packages/nodejs_wheel/. # Patch every copy of tar, glob, and brace-expansion inside that tree. RUN GLOBAL="$(npm root -g)" && \ + [ -n "$GLOBAL" ] || { echo "ERROR: npm root -g returned empty; aborting"; exit 1; } && \ find /usr/lib -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ done && \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index fda591df083..fbc16e4f876 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -106,7 +106,7 @@ RUN for i in 1 2 3; do \ apk add --no-cache python3 py3-pip bash openssl tzdata nodejs npm supervisor && break || sleep 5; \ done \ && apk upgrade --no-cache nodejs \ - && npm install -g npm@latest tar@7.5.8 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.1 diff@8.0.3 \ + && npm install -g npm@latest tar@7.5.10 glob@11.1.0 @isaacs/brace-expansion@5.0.1 minimatch@10.2.4 diff@8.0.3 \ && GLOBAL="$(npm root -g)" \ && find "$GLOBAL/npm" -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ @@ -123,7 +123,10 @@ RUN for i in 1 2 3; do \ && find "$GLOBAL/npm" -type d -name "diff" -path "*/node_modules/diff" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/diff" "$d"; \ done \ - && npm cache clean --force + && find /usr/local/lib /usr/lib -path "*/node_modules/npm/package.json" -exec \ + sed -i 's/"tar": "\^7\.5\.[0-9]*"/"tar": "^7.5.10"/g; s/"minimatch": "\^10\.[0-9.]*"/"minimatch": "^10.2.4"/g' {} + 2>/dev/null \ + && npm cache clean --force \ + && { apk del --no-cache npm 2>/dev/null || true; } # Copy artifacts from builder COPY --from=builder /app/requirements.txt /app/requirements.txt @@ -169,6 +172,7 @@ RUN pip install --no-index --find-links=/wheels/ -r requirements.txt && \ # npm with old vulnerable deps at /usr/lib/python3.*/site-packages/nodejs_wheel/. # Patch every copy of tar, glob, and brace-expansion inside that tree. RUN GLOBAL="$(npm root -g)" && \ + [ -n "$GLOBAL" ] || { echo "ERROR: npm root -g returned empty; aborting"; exit 1; } && \ find /usr/lib -type d -name "tar" -path "*/node_modules/tar" | while read d; do \ rm -rf "$d" && cp -rL "$GLOBAL/tar" "$d"; \ done && \ diff --git a/docs/my-website/docs/a2a.md b/docs/my-website/docs/a2a.md index b1166a7809c..9c86d0de383 100644 --- a/docs/my-website/docs/a2a.md +++ b/docs/my-website/docs/a2a.md @@ -20,6 +20,7 @@ Add A2A Agents on LiteLLM AI Gateway, Invoke agents in A2A Protocol, track reque | Logging | ✅ | | Load Balancing | ✅ | | Streaming | ✅ | +| [Iteration Budgets](a2a_iteration_budgets) | ✅ | :::tip diff --git a/docs/my-website/docs/a2a_iteration_budgets.md b/docs/my-website/docs/a2a_iteration_budgets.md new file mode 100644 index 00000000000..47beca3470f --- /dev/null +++ b/docs/my-website/docs/a2a_iteration_budgets.md @@ -0,0 +1,188 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Agent Iteration Budgets + +Control runaway costs from agentic loops with per-session iteration and budget caps. + +## Overview + +When agents run agentic loops, they can make unbounded LLM calls, causing unexpected costs. LiteLLM provides two controls: + +| Control | Description | +|---------|-------------| +| **Max Iterations** | Hard cap on the number of LLM calls per session | +| **Max Budget Per Session** | Dollar cap per session (identified by `x-litellm-trace-id`) | + +Both controls require a `session_id` (sent via `x-litellm-trace-id` header or `metadata.session_id`) to track calls within a session. + +## Trace-ID Enforcement + +LiteLLM supports two independent trace-id flags, configured in `litellm_params` on the agent: + +| Flag | Description | +|------|-------------| +| `require_trace_id_on_calls_to_agent` | Requires callers invoking this agent to include `x-litellm-trace-id`. Use when the agent should only be called as a sub-agent with a trace context. Returns **400** if missing. | +| `require_trace_id_on_calls_by_agent` | Requires all LLM/MCP calls made **by** this agent (via its virtual key) to include `x-litellm-trace-id`. This is what enables `max_iterations` and `max_budget_per_session` tracking. Returns **400** if missing. | + +## Configuring via UI + +When creating an agent in the LiteLLM Admin UI: + +1. Navigate to the **Agents** tab and click **Add Agent** +2. In the **Agent Settings** step, expand the **Tracing** section +3. Toggle **Require x-litellm-trace-id on calls BY this agent** to enable session tracking +4. Set **Max Iterations** to cap the number of LLM calls per session +5. Set **Max Budget Per Session ($)** to cap spend per session + +The trace-id flags are stored on the agent's `litellm_params`. Budget controls (`max_iterations`, `max_budget_per_session`) are stored in the virtual key's metadata. + +## Configuring via API + +Set trace-id enforcement on the agent itself: + +```bash +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent with budget controls", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "litellm_params": { + "require_trace_id_on_calls_to_agent": true, + "require_trace_id_on_calls_by_agent": true + } + }' +``` + +Budget controls are set on the agent's `litellm_params` (not on individual keys), so they apply across all keys for the agent: + +```bash +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent with budget controls", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "litellm_params": { + "require_trace_id_on_calls_by_agent": true, + "max_iterations": 25, + "max_budget_per_session": 5.00 + } + }' +``` + +## How It Works + +### Session Tracking + +Callers identify their session by including a `session_id` in one of these ways: +- **Header**: `x-litellm-trace-id: my-session-123` +- **Metadata**: `{"metadata": {"session_id": "my-session-123"}}` + +### Max Iterations + +When `max_iterations` is set in agent `litellm_params`: +- Each LLM call for a session increments a counter +- When the counter exceeds `max_iterations`, the request receives a **429 Too Many Requests** +- Counters expire after 1 hour by default (configurable via `LITELLM_MAX_ITERATIONS_TTL` env var) + +### Max Budget Per Session + +When `max_budget_per_session` is set in agent `litellm_params`: +- After each successful LLM call, the response cost is accumulated for the session +- Before each call, the accumulated spend is checked against the budget +- When spend exceeds the budget, the request receives a **429 Too Many Requests** +- Session spend counters expire after 1 hour by default (configurable via `LITELLM_MAX_BUDGET_PER_SESSION_TTL` env var) + +## Example + +Create an agent with max 25 iterations and a $5 budget cap: + + + + +1. Go to **Agents** → **Add Agent** +2. Configure your agent (name, model, etc.) +3. In **Agent Settings**, expand the **Tracing** section +4. Toggle on **Require x-litellm-trace-id on calls BY this agent** +5. Set **Max Iterations** to `25` +6. Set **Max Budget Per Session** to `5.00` +7. Proceed to create a new key for the agent +8. Click **Create Agent** + + + + +```bash +# 1. Create the agent with trace-id enforcement +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent with budget controls", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "litellm_params": { + "require_trace_id_on_calls_by_agent": true + } + }' + +# 2. Create a key for the agent +curl -X POST 'http://localhost:4000/key/generate' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_id": "", + "key_alias": "my-research-agent-key" + }' +``` + + + + +### Making Calls with Session Tracking + +```bash +curl -X POST 'http://localhost:4000/chat/completions' \ + -H 'Authorization: Bearer sk-agent-key-xxx' \ + -H 'x-litellm-trace-id: session-abc-123' \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Hello"}] + }' +``` + +After 25 calls or $5 spent within this session, subsequent requests will receive: + +```json +{ + "error": { + "message": "Session budget exceeded for session session-abc-123. Current spend: $5.0032, max_budget_per_session: $5.00.", + "type": "budget_exceeded", + "code": 429 + } +} +``` + +## Environment Variables + +| Variable | Default | Description | +|----------|---------|-------------| +| `LITELLM_MAX_ITERATIONS_TTL` | `3600` (1 hour) | TTL in seconds for session iteration counters | +| `LITELLM_MAX_BUDGET_PER_SESSION_TTL` | `3600` (1 hour) | TTL in seconds for session budget counters | diff --git a/docs/my-website/docs/proxy/dynamic_rate_limit.md b/docs/my-website/docs/proxy/dynamic_rate_limit.md index 3c3500f8a6c..09a111f7297 100644 --- a/docs/my-website/docs/proxy/dynamic_rate_limit.md +++ b/docs/my-website/docs/proxy/dynamic_rate_limit.md @@ -3,6 +3,8 @@ Prevent projects from gobbling too much tpm/rpm. +**See Also:** [Request Prioritization](../scheduler.md) - Prioritize LLM API requests in high-traffic by adding them to a priority queue. + Dynamically allocate TPM/RPM quota to api keys, based on active keys in that minute. [**See Code**](https://github.com/BerriAI/litellm/blob/9bffa9a48e610cc6886fc2dce5c1815aeae2ad46/litellm/proxy/hooks/dynamic_rate_limiter.py#L125) ## Quick Start Usage diff --git a/docs/my-website/docs/proxy/users.md b/docs/my-website/docs/proxy/users.md index 8517db51a8f..58813eaf49e 100644 --- a/docs/my-website/docs/proxy/users.md +++ b/docs/my-website/docs/proxy/users.md @@ -10,6 +10,8 @@ import TabItem from '@theme/TabItem'; **Team member budgets**: Set individual spending limits within the team's shared budget +**Agent budgets**: Set rate limits (tpm/rpm) and session-level caps (iterations, dollar budget) on agents [**Jump**](#agents) + ***If a key belongs to a team, the team budget is applied, not the user's personal budget.*** ::: @@ -420,6 +422,109 @@ Expected response on failure +### Agents + +Set budgets and rate limits on agents registered with LiteLLM's [Agent Gateway](../a2a.md). You can control: +- **Per-agent rate limits**: `tpm_limit` and `rpm_limit` on the agent itself +- **Per-session rate limits**: `session_tpm_limit` and `session_rpm_limit` applied per session +- **Per-session iteration cap**: `max_iterations` in agent `litellm_params` +- **Per-session budget cap**: `max_budget_per_session` in agent `litellm_params` + + + + +Set `tpm_limit` and `rpm_limit` on the agent to cap total throughput across all sessions. + +```bash +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "tpm_limit": 100000, + "rpm_limit": 100 + }' +``` + + + + +Set `session_tpm_limit` and `session_rpm_limit` to cap throughput per individual session. + +```bash +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "session_tpm_limit": 50000, + "session_rpm_limit": 50 + }' +``` + + + + +Set `max_iterations` and `max_budget_per_session` in agent `litellm_params` to cap individual sessions. Requires `require_trace_id_on_calls_by_agent` so LiteLLM can track calls per session. + +```bash +curl -X POST 'http://localhost:4000/v1/agents' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "agent_name": "my-research-agent", + "agent_card_params": { + "name": "my-research-agent", + "description": "A research agent", + "url": "http://my-agent:8080", + "version": "1.0.0" + }, + "litellm_params": { + "require_trace_id_on_calls_by_agent": true, + "max_iterations": 25, + "max_budget_per_session": 5.00 + } + }' +``` + +When a session exceeds the limit, requests receive a **429 Too Many Requests** response. + +See the [Agent Iteration Budgets](../a2a_iteration_budgets) guide for full details. + + + + +:::info + +You can also update rate limits on existing agents using `PATCH /v1/agents/{agent_id}`: + +```bash +curl -X PATCH 'http://localhost:4000/v1/agents/' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' \ + -d '{ + "tpm_limit": 200000, + "rpm_limit": 200, + "session_tpm_limit": 50000, + "session_rpm_limit": 50 + }' +``` + +::: + + ### Customers Use this to budget `user` passed to `/chat/completions`, **without needing to create a key for every user** @@ -685,6 +790,31 @@ These headers indicate: - 1 request remaining for the GPT-4 model for key=`sk-ulGNRXWtv7M0lFnnsQk0wQ` - 179 tokens remaining for the GPT-4 model for key=`sk-ulGNRXWtv7M0lFnnsQk0wQ` + + + +Set rate limits on agents registered with the [Agent Gateway](../a2a.md). + +**Agent-level limits** cap total throughput across all sessions: + +```shell +curl -X POST 'http://0.0.0.0:4000/v1/agents' \ +--header 'Authorization: Bearer sk-1234' \ +--header 'Content-Type: application/json' \ +--data '{"agent_name": "my-agent", "agent_card_params": {"name": "my-agent", "description": "My agent", "url": "http://my-agent:8080", "version": "1.0.0"}, "tpm_limit": 100000, "rpm_limit": 100}' +``` + +**Session-level limits** cap throughput per individual session: + +```shell +curl -X POST 'http://0.0.0.0:4000/v1/agents' \ +--header 'Authorization: Bearer sk-1234' \ +--header 'Content-Type: application/json' \ +--data '{"agent_name": "my-agent", "agent_card_params": {"name": "my-agent", "description": "My agent", "url": "http://my-agent:8080", "version": "1.0.0"}, "session_tpm_limit": 50000, "session_rpm_limit": 50}' +``` + +You can also set **max_iterations** (call count cap) and **max_budget_per_session** (dollar cap) per session via `litellm_params`. See [Agent Iteration Budgets](../a2a_iteration_budgets) for details. + diff --git a/docs/my-website/package.json b/docs/my-website/package.json index 758e53ebf59..20462de2dd7 100644 --- a/docs/my-website/package.json +++ b/docs/my-website/package.json @@ -61,7 +61,7 @@ "mermaid": ">=11.10.0", "gray-matter": "4.0.3", "glob": ">=11.1.0", - "tar": ">=7.5.8", + "tar": ">=7.5.10", "minimatch": ">=10.2.4", "diff": ">=8.0.3", "@isaacs/brace-expansion": ">=5.0.1", diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index a77e905724e..ef2df2d8ad4 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -542,7 +542,8 @@ const sidebars = { "a2a_invoking_agents", "a2a_agent_headers", "a2a_cost_tracking", - "a2a_agent_permissions" + "a2a_agent_permissions", + "a2a_iteration_budgets" ], }, "assistants", diff --git a/litellm-js/spend-logs/package.json b/litellm-js/spend-logs/package.json index 5a7a08cb9ef..adfe49017d1 100644 --- a/litellm-js/spend-logs/package.json +++ b/litellm-js/spend-logs/package.json @@ -12,8 +12,8 @@ }, "overrides": { "glob": ">=11.1.0", - "tar": ">=7.5.8", - "minimatch": ">=10.2.1", + "tar": ">=7.5.10", + "minimatch": ">=10.2.4", "diff": ">=8.0.3", "@isaacs/brace-expansion": ">=5.0.1", "@babel/traverse": ">=7.23.2", diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql new file mode 100644 index 00000000000..01f3936a6fc --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260304175016_add_spend_to_agent_table/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql new file mode 100644 index 00000000000..3cd8ca638a4 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260305000000_add_rate_limits_to_agents/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "tpm_limit" INTEGER; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "rpm_limit" INTEGER; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "session_tpm_limit" INTEGER; +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN "session_rpm_limit" INTEGER; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 329ff80933f..8d4bdffb2dd 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -68,6 +68,11 @@ model LiteLLM_AgentsTable { agent_access_groups String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + spend Float @default(0.0) + tpm_limit Int? + rpm_limit Int? + session_tpm_limit Int? + session_rpm_limit Int? created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index a55e30ebeb9..c752e84b967 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -198,9 +198,8 @@ async def _get_batch_output_file_content_as_dictionary( Required for Azure and other providers that need authentication """ from litellm.files.main import afile_content - from litellm.proxy.openai_files_endpoints.common_utils import ( - _is_base64_encoded_unified_file_id, - ) + from litellm.proxy.openai_files_endpoints.common_utils import \ + _is_base64_encoded_unified_file_id if custom_llm_provider == "vertex_ai": raise ValueError("Vertex AI does not support file content retrieval") @@ -227,7 +226,7 @@ async def _get_batch_output_file_content_as_dictionary( credentials = _extract_file_access_credentials(litellm_params) file_content_kwargs.update(credentials) - _file_content = await afile_content(**file_content_kwargs) + _file_content = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType] return _get_file_content_as_dictionary(_file_content.content) diff --git a/litellm/fine_tuning/main.py b/litellm/fine_tuning/main.py index e6f2a6f86db..93fa56ff971 100644 --- a/litellm/fine_tuning/main.py +++ b/litellm/fine_tuning/main.py @@ -126,6 +126,18 @@ async def acreate_fine_tuning_job( raise e +def _build_fine_tuning_job_data(model, training_file, hyperparameters, suffix, validation_file, integrations, seed): + return FineTuningJobCreate( + model=model, + training_file=training_file, + hyperparameters=hyperparameters, + suffix=suffix, + validation_file=validation_file, + integrations=integrations, + seed=seed, + ) + + def _resolve_fine_tuning_timeout( timeout: Any, custom_llm_provider: str, @@ -206,19 +218,9 @@ def create_fine_tuning_job( or os.getenv("OPENAI_API_KEY") ) - create_fine_tuning_job_data = FineTuningJobCreate( - model=model, - training_file=training_file, - hyperparameters=_oai_hyperparameters, - suffix=suffix, - validation_file=validation_file, - integrations=integrations, - seed=seed, - ) - - create_fine_tuning_job_data_dict = create_fine_tuning_job_data.model_dump( - exclude_none=True - ) + create_fine_tuning_job_data_dict = _build_fine_tuning_job_data( + model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed, + ).model_dump(exclude_none=True) response = openai_fine_tuning_apis_instance.create_fine_tuning_job( api_base=api_base, @@ -260,20 +262,10 @@ def create_fine_tuning_job( # Prepare Azure-specific parameters for extra_body extra_body = _prepare_azure_extra_body(extra_body, kwargs, azure_specific_hyperparams) - create_fine_tuning_job_data = FineTuningJobCreate( - model=model, - training_file=training_file, - hyperparameters=_oai_hyperparameters, - suffix=suffix, - validation_file=validation_file, - integrations=integrations, - seed=seed, - ) + create_fine_tuning_job_data_dict = _build_fine_tuning_job_data( + model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed, + ).model_dump(exclude_none=True) - create_fine_tuning_job_data_dict = create_fine_tuning_job_data.model_dump( - exclude_none=True - ) - # Add extra_body if it has Azure-specific parameters if extra_body: create_fine_tuning_job_data_dict["extra_body"] = extra_body @@ -303,18 +295,11 @@ def create_fine_tuning_job( vertex_credentials = optional_params.vertex_credentials or get_secret_str( "VERTEXAI_CREDENTIALS" ) - create_fine_tuning_job_data = FineTuningJobCreate( - model=model, - training_file=training_file, - hyperparameters=_oai_hyperparameters, - suffix=suffix, - validation_file=validation_file, - integrations=integrations, - seed=seed, - ) response = vertex_fine_tuning_apis_instance.create_fine_tuning_job( _is_async=_is_async, - create_fine_tuning_job_data=create_fine_tuning_job_data, + create_fine_tuning_job_data=_build_fine_tuning_job_data( + model, training_file, _oai_hyperparameters, suffix, validation_file, integrations, seed, + ), vertex_credentials=vertex_credentials, vertex_project=vertex_ai_project, vertex_location=vertex_ai_location, diff --git a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py index e5698843e8f..328c3a0b977 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude3_transformation.py @@ -108,6 +108,9 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig): _anthropic_request.pop("stream", None) # Bedrock Invoke doesn't support output_format parameter _anthropic_request.pop("output_format", None) + # Bedrock Invoke doesn't support output_config parameter + # Fixes: https://github.com/BerriAI/litellm/issues/22797 + _anthropic_request.pop("output_config", None) if "anthropic_version" not in _anthropic_request: _anthropic_request["anthropic_version"] = self.anthropic_version diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 9fae5fd2a17..b11215e7f6b 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -419,6 +419,10 @@ class AmazonAnthropicClaudeMessagesConfig( anthropic_messages_request=anthropic_messages_request, ) + # 5b. Strip `output_config` — Bedrock Invoke doesn't support it + # Fixes: https://github.com/BerriAI/litellm/issues/22797 + anthropic_messages_request.pop("output_config", None) + # 5a. Remove `custom` field from tools (Bedrock doesn't support it) # Claude Code sends `custom: {defer_loading: true}` on tool definitions, # which causes Bedrock to reject the request with "Extra inputs are not permitted" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 12b0a69ce77..d797d9c7e0a 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -202,6 +202,7 @@ class Litellm_EntityType(enum.Enum): ORGANIZATION = "organization" PROJECT = "project" TAG = "tag" + AGENT = "agent" # global proxy level entity PROXY = "proxy" @@ -4230,6 +4231,7 @@ class DBSpendUpdateTransactions(TypedDict): team_member_list_transactions: Optional[Dict[str, float]] org_list_transactions: Optional[Dict[str, float]] tag_list_transactions: Optional[Dict[str, float]] + agent_list_transactions: Optional[Dict[str, float]] class SpendUpdateQueueItem(TypedDict, total=False): diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 3ffced7a12c..63e0dad3322 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -39,7 +39,8 @@ def _jsonrpc_error( def _get_agent(agent_id: str): """Look up an agent by ID or name. Returns None if not found.""" - from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry agent = global_agent_registry.get_agent_by_id(agent_id=agent_id) if agent is None: @@ -47,6 +48,26 @@ def _get_agent(agent_id: str): return agent +def _enforce_inbound_trace_id(agent: Any, request: Request) -> None: + """Raise 400 if agent requires x-litellm-trace-id on inbound calls and it is missing.""" + agent_litellm_params = agent.litellm_params or {} + if not agent_litellm_params.get("require_trace_id_on_calls_to_agent"): + return + + from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers + + headers_dict = dict(request.headers) + trace_id = get_chain_id_from_headers(headers_dict) + if not trace_id: + raise HTTPException( + status_code=400, + detail=( + f"Agent '{agent.agent_id}' requires x-litellm-trace-id header " + "on all inbound requests." + ), + ) + + async def _handle_stream_message( api_base: Optional[str], request_id: str, @@ -116,9 +137,8 @@ async def _handle_stream_message( and request_data is not None and proxy_logging_obj is not None ): - from litellm.proxy.common_request_processing import ( - ProxyBaseLLMRequestProcessing, - ) + from litellm.proxy.common_request_processing import \ + ProxyBaseLLMRequestProcessing def _ndjson_chunk(chunk: Any) -> str: if hasattr(chunk, "model_dump"): @@ -218,9 +238,8 @@ async def get_agent_card( The URL in the agent card is rewritten to point to the LiteLLM proxy, so all subsequent A2A calls go through LiteLLM for logging and cost tracking. """ - from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( - AgentRequestHandler, - ) + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import \ + AgentRequestHandler try: agent = _get_agent(agent_id) @@ -284,15 +303,10 @@ async def invoke_agent_a2a( # noqa: PLR0915 """ from litellm.a2a_protocol import asend_message from litellm.a2a_protocol.main import A2A_SDK_AVAILABLE - from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( - AgentRequestHandler, - ) - from litellm.proxy.proxy_server import ( - general_settings, - proxy_config, - proxy_logging_obj, - version, - ) + from litellm.proxy.agent_endpoints.auth.agent_permission_handler import \ + AgentRequestHandler + from litellm.proxy.proxy_server import (general_settings, proxy_config, + proxy_logging_obj, version) body = {} try: @@ -345,6 +359,8 @@ async def invoke_agent_a2a( # noqa: PLR0915 detail=f"Agent '{agent_id}' is not allowed for your key/team. Contact proxy admin for access.", ) + _enforce_inbound_trace_id(agent, request) + # Get backend URL and agent name agent_url = agent.agent_card_params.get("url") agent_name = agent.agent_card_params.get("name", agent_id) @@ -365,6 +381,10 @@ async def invoke_agent_a2a( # noqa: PLR0915 ) # Set up data dict for litellm processing + if "metadata" not in body: + body["metadata"] = {} + body["metadata"]["agent_id"] = agent.agent_id + body.update( { "model": f"a2a_agent/{agent_name}", @@ -373,9 +393,8 @@ async def invoke_agent_a2a( # noqa: PLR0915 ) # Add litellm data (user_api_key, user_id, team_id, etc.) - from litellm.proxy.common_request_processing import ( - ProxyBaseLLMRequestProcessing, - ) + from litellm.proxy.common_request_processing import \ + ProxyBaseLLMRequestProcessing processor = ProxyBaseLLMRequestProcessing(data=body) data, logging_obj = await processor.common_processing_pre_call_logic( diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index 91bbbd73d11..ce6b1055ee1 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -5,9 +5,8 @@ from typing import Any, Dict, List, Optional import litellm from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy.management_helpers.object_permission_utils import ( - handle_update_object_permission_common, -) +from litellm.proxy.management_helpers.object_permission_utils import \ + handle_update_object_permission_common from litellm.proxy.utils import PrismaClient from litellm.types.agents import AgentConfig, AgentResponse, PatchAgentRequest @@ -152,6 +151,11 @@ class AgentRegistry: if object_permission_id is not None: create_data["object_permission_id"] = object_permission_id + for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): + _val = agent.get(rate_field) + if _val is not None: + create_data[rate_field] = _val + # Create agent in DB created_agent = await prisma_client.db.litellm_agentstable.create( data=create_data, @@ -226,6 +230,10 @@ class AgentRegistry: update_data["agent_card_params"] = safe_dumps( augment_agent.get("agent_card_params") ) + + for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): + if rate_field in agent: + update_data[rate_field] = agent.get(rate_field) if "static_headers" in agent: headers_value = agent.get("static_headers") update_data["static_headers"] = safe_dumps( @@ -321,6 +329,12 @@ class AgentRegistry: "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), } + + for rate_field in ("tpm_limit", "rpm_limit", "session_tpm_limit", "session_rpm_limit"): + _val = agent.get(rate_field) + if _val is not None: + update_data[rate_field] = _val + if agent.get("object_permission") is not None: existing_agent = await prisma_client.db.litellm_agentstable.find_unique( where={"agent_id": agent_id} diff --git a/litellm/proxy/agent_endpoints/endpoints.py b/litellm/proxy/agent_endpoints/endpoints.py index 80c55f634f7..252db29a9da 100644 --- a/litellm/proxy/agent_endpoints/endpoints.py +++ b/litellm/proxy/agent_endpoints/endpoints.py @@ -8,12 +8,15 @@ Follows the A2A Spec. 3. Get specific agent via GET `/v1/agents/{agent_id}` """ -from typing import Any, List, Optional +import asyncio +import os +from typing import Any, Dict, List, Optional -from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi import APIRouter, Depends, HTTPException, Query, Request import litellm from litellm._logging import verbose_proxy_logger +from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.proxy._types import CommonProxyErrors, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.rbac_utils import check_feature_access_for_user @@ -25,6 +28,7 @@ from litellm.types.agents import ( MakeAgentsPublicRequest, PatchAgentRequest, ) +from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.proxy.management_endpoints.common_daily_activity import ( SpendAnalyticsPaginatedResponse, ) @@ -49,6 +53,48 @@ def _check_agent_management_permission(user_api_key_dict: UserAPIKeyAuth) -> Non ) +AGENT_HEALTH_CHECK_TIMEOUT_SECONDS = float( + os.environ.get("LITELLM_AGENT_HEALTH_CHECK_TIMEOUT", "5.0") +) +AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS = float( + os.environ.get("LITELLM_AGENT_HEALTH_CHECK_GATHER_TIMEOUT", "30.0") +) + + +async def _check_agent_url_health( + agent: AgentResponse, +) -> Dict[str, Any]: + """ + Perform a GET request against the agent's URL and return the health result. + + Returns a dict with ``agent_id``, ``healthy`` (bool), and an optional + ``error`` message. + """ + url = (agent.agent_card_params or {}).get("url") + if not url: + return {"agent_id": agent.agent_id, "healthy": True} + + try: + client = get_async_httpx_client( + llm_provider=httpxSpecialProvider.AgentHealthCheck, + params={"timeout": AGENT_HEALTH_CHECK_TIMEOUT_SECONDS}, + ) + response = await client.get(url) + if response.status_code >= 500: + return { + "agent_id": agent.agent_id, + "healthy": False, + "error": f"HTTP {response.status_code}", + } + return {"agent_id": agent.agent_id, "healthy": True} + except Exception as exc: + return { + "agent_id": agent.agent_id, + "healthy": False, + "error": str(exc), + } + + @router.get( "/v1/agents", tags=["[beta] A2A Agents"], @@ -57,6 +103,9 @@ def _check_agent_management_permission(user_api_key_dict: UserAPIKeyAuth) -> Non ) async def get_agents( request: Request, + health_check: bool = Query( + False, + description="When true, performs a GET request to each agent's URL. Agents with reachable URLs (HTTP status < 500) and agents without a URL are returned; unreachable agents are filtered out.", user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), # Used for auth ): """ @@ -67,6 +116,13 @@ async def get_agents( -H "Authorization: Bearer your-key" \ ``` + Pass `?health_check=true` to filter out agents whose URL is unreachable: + ``` + curl -X GET "http://localhost:4000/v1/agents?health_check=true" \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer your-key" \ + ``` + Returns: List[AgentResponse] """ @@ -79,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 @@ -91,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() @@ -99,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: @@ -112,6 +181,44 @@ async def get_agents( and (agent.agent_id in litellm.public_agent_groups) ) + if health_check: + agents_with_url = [ + agent + for agent in returned_agents + if (agent.agent_card_params or {}).get("url") + ] + agents_without_url = [ + agent + for agent in returned_agents + if not (agent.agent_card_params or {}).get("url") + ] + try: + health_results = await asyncio.wait_for( + asyncio.gather( + *[_check_agent_url_health(agent) for agent in agents_with_url] + ), + timeout=AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + verbose_proxy_logger.warning( + "Agent health check gather timed out after %s seconds", + AGENT_HEALTH_CHECK_GATHER_TIMEOUT_SECONDS, + ) + health_results = [ + {"agent_id": agent.agent_id, "healthy": False, "error": "Health check timed out"} + for agent in agents_with_url + ] + healthy_ids = { + result["agent_id"] + for result in health_results + if result["healthy"] + } + returned_agents = [ + agent + for agent in agents_with_url + if agent.agent_id in healthy_ids + ] + agents_without_url + return returned_agents except HTTPException: raise @@ -128,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( @@ -269,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( @@ -580,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 @@ -697,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 @@ -759,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"], @@ -820,4 +936,4 @@ async def get_agent_daily_activity( api_key=api_key, page=page, page_size=page_size, - ) \ No newline at end of file + ) diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index eb690931fb9..db794c5ac3d 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -11,7 +11,8 @@ Run checks for: import asyncio import re import time -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast +from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, + cast) from fastapi import HTTPException, Request, status from pydantic import BaseModel @@ -20,48 +21,33 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.caching.dual_cache import LimitedSizeOrderedDict -from litellm.constants import ( - CLI_JWT_EXPIRATION_HOURS, - CLI_JWT_TOKEN_NAME, - DEFAULT_ACCESS_GROUP_CACHE_TTL, - DEFAULT_IN_MEMORY_TTL, - DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, - DEFAULT_MAX_RECURSE_DEPTH, - EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, -) +from litellm.constants import (CLI_JWT_EXPIRATION_HOURS, CLI_JWT_TOKEN_NAME, + DEFAULT_ACCESS_GROUP_CACHE_TTL, + DEFAULT_IN_MEMORY_TTL, + DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + DEFAULT_MAX_RECURSE_DEPTH, + EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.proxy._types import ( - RBAC_ROLES, - CallInfo, - LiteLLM_AccessGroupTable, - LiteLLM_BudgetTable, - LiteLLM_EndUserTable, - Litellm_EntityType, - LiteLLM_JWTAuth, - LiteLLM_ObjectPermissionTable, - LiteLLM_OrganizationMembershipTable, - LiteLLM_OrganizationTable, - LiteLLM_ProjectTableCachedObj, - LiteLLM_TagTable, - LiteLLM_TeamMembership, - LiteLLM_TeamTable, - LiteLLM_TeamTableCachedObj, - LiteLLM_UserTable, - LiteLLMRoutes, - LitellmUserRoles, - NewTeamRequest, - ProxyErrorTypes, - ProxyException, - RoleBasedPermissions, - SpecialModelNames, - UserAPIKeyAuth, -) +from litellm.proxy._types import (RBAC_ROLES, CallInfo, + LiteLLM_AccessGroupTable, + LiteLLM_BudgetTable, LiteLLM_EndUserTable, + Litellm_EntityType, LiteLLM_JWTAuth, + LiteLLM_ObjectPermissionTable, + LiteLLM_OrganizationMembershipTable, + LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TagTable, LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, LiteLLMRoutes, + LitellmUserRoles, NewTeamRequest, + ProxyErrorTypes, ProxyException, + RoleBasedPermissions, SpecialModelNames, + UserAPIKeyAuth) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( - TOOL_CAPABLE_CALL_TYPES, - extract_request_tool_names, -) + TOOL_CAPABLE_CALL_TYPES, extract_request_tool_names) from litellm.proxy.route_llm_request import route_request from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.router import Router @@ -224,6 +210,89 @@ async def _run_project_checks( ) +def _enforce_user_param_check( + general_settings: dict, request: Request, request_body: dict, route: str +) -> None: + if not general_settings.get("enforce_user_param", False): + return + + http_method = request.method if hasattr(request, "method") else None + is_post_method = http_method and http_method.upper() == "POST" + is_openai_route = RouteChecks.is_llm_api_route(route=route) + is_mcp_route = ( + route in LiteLLMRoutes.mcp_routes.value + or RouteChecks.check_route_access( + route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value + ) + ) + + if ( + is_post_method + and is_openai_route + and not is_mcp_route + and "user" not in request_body + ): + raise Exception( + f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}" + ) + + +def _reject_clientside_metadata_tags_check( + general_settings: dict, request_body: dict, route: str +) -> None: + if not general_settings.get("reject_clientside_metadata_tags", False): + return + + if ( + RouteChecks.is_llm_api_route(route=route) + and "metadata" in request_body + and isinstance(request_body["metadata"], dict) + and "tags" in request_body["metadata"] + ): + raise ProxyException( + message=f"Client-side 'metadata.tags' not allowed in request. 'reject_clientside_metadata_tags'={general_settings['reject_clientside_metadata_tags']}. Tags can only be set via API key metadata.", + type=ProxyErrorTypes.bad_request_error, + param="metadata.tags", + code=status.HTTP_400_BAD_REQUEST, + ) + + +def _global_proxy_budget_check( + global_proxy_spend: Optional[float], skip_budget_checks: bool, route: str +) -> None: + if ( + litellm.max_budget > 0 + and not skip_budget_checks + and global_proxy_spend is not None + and RouteChecks.is_llm_api_route(route=route) + and route != "/v1/models" + and route != "/models" + ): + if global_proxy_spend > litellm.max_budget: + raise litellm.BudgetExceededError( + current_cost=global_proxy_spend, max_budget=litellm.max_budget + ) + + +def _guardrail_modification_check( + request_body: dict, team_object: Optional[LiteLLM_TeamTable] +) -> None: + _request_metadata: dict = request_body.get("metadata", {}) or {} + if not _request_metadata.get("guardrails"): + return + + from litellm.proxy.guardrails.guardrail_helpers import \ + can_modify_guardrails + + if not can_modify_guardrails(team_object): + raise HTTPException( + status_code=403, + detail={ + "error": "Your team does not have permission to modify guardrails." + }, + ) + + async def check_tools_allowlist( request_body: dict, valid_token: Optional[UserAPIKeyAuth], @@ -235,23 +304,34 @@ async def check_tools_allowlist( effective allowlist is read from valid_token.metadata and valid_token.team_metadata. Raises ProxyException with tool_access_denied if a tool is not allowed. """ - from litellm.litellm_core_utils.api_route_to_call_types import ( - get_call_types_for_route, - ) + from litellm.litellm_core_utils.api_route_to_call_types import \ + get_call_types_for_route if valid_token is None: return call_types = get_call_types_for_route(route) - if not call_types or not any(ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types): + if not call_types or not any( + ct.value in TOOL_CAPABLE_CALL_TYPES for ct in call_types + ): return tool_names = extract_request_tool_names(route, request_body) if not tool_names: return - key_meta = (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} - team_meta = (valid_token.team_metadata or {}) if isinstance(valid_token.team_metadata, dict) else {} + key_meta = ( + (valid_token.metadata or {}) if isinstance(valid_token.metadata, dict) else {} + ) + team_meta = ( + (valid_token.team_metadata or {}) + if isinstance(valid_token.team_metadata, dict) + else {} + ) key_allowed = key_meta.get("allowed_tools") team_allowed = team_meta.get("allowed_tools") - effective = key_allowed if (isinstance(key_allowed, list) and len(key_allowed) > 0) else team_allowed + effective = ( + key_allowed + if (isinstance(key_allowed, list) and len(key_allowed) > 0) + else team_allowed + ) if not isinstance(effective, list) or len(effective) == 0: return allowed_set = {str(t) for t in effective} @@ -326,6 +406,29 @@ async def common_checks( # noqa: PLR0915 code=status.HTTP_401_UNAUTHORIZED, ) + # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent + if valid_token is not None and valid_token.agent_id: + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry + from litellm.proxy.litellm_pre_call_utils import \ + get_chain_id_from_headers + + agent = global_agent_registry.get_agent_by_id(agent_id=valid_token.agent_id) + if agent is not None: + require_trace_id = (agent.litellm_params or {}).get( + "require_trace_id_on_calls_by_agent" + ) + if require_trace_id: + headers_dict = dict(request.headers) + trace_id = get_chain_id_from_headers(headers_dict) + if not trace_id: + raise ProxyException( + message="Requests made with this agent's key must include the x-litellm-trace-id header.", + type=ProxyErrorTypes.bad_request_error, + param=None, + code=status.HTTP_400_BAD_REQUEST, + ) + ## 2.1 If user can call model (if personal key) if _model and team_object is None and user_object is not None: await can_user_call_model( @@ -415,83 +518,10 @@ async def common_checks( # noqa: PLR0915 message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}", ) - # 6. [OPTIONAL] If 'enforce_user_param' enabled - did developer pass in 'user' param for openai endpoints - if ( - general_settings.get("enforce_user_param", None) is not None - and general_settings["enforce_user_param"] is True - ): - # Get HTTP method from request - http_method = request.method if hasattr(request, "method") else None - - # Check if it's a POST request and if it's an OpenAI route but not MCP - is_post_method = http_method and http_method.upper() == "POST" - is_openai_route = RouteChecks.is_llm_api_route(route=route) - is_mcp_route = ( - route in LiteLLMRoutes.mcp_routes.value - or RouteChecks.check_route_access( - route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value - ) - ) - - # Enforce user param only for POST requests on OpenAI routes (excluding MCP routes) - if ( - is_post_method - and is_openai_route - and not is_mcp_route - and "user" not in request_body - ): - raise Exception( - f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}" - ) - - # 6.1 [OPTIONAL] If 'reject_clientside_metadata_tags' enabled - reject request if it has client-side 'metadata.tags' - if ( - general_settings.get("reject_clientside_metadata_tags", None) is not None - and general_settings["reject_clientside_metadata_tags"] is True - ): - if ( - RouteChecks.is_llm_api_route(route=route) - and "metadata" in request_body - and isinstance(request_body["metadata"], dict) - and "tags" in request_body["metadata"] - ): - raise ProxyException( - message=f"Client-side 'metadata.tags' not allowed in request. 'reject_clientside_metadata_tags'={general_settings['reject_clientside_metadata_tags']}. Tags can only be set via API key metadata.", - type=ProxyErrorTypes.bad_request_error, - param="metadata.tags", - code=status.HTTP_400_BAD_REQUEST, - ) - # 7. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget - if ( - litellm.max_budget > 0 - and not skip_budget_checks - and global_proxy_spend is not None - # only run global budget checks for OpenAI routes - # Reason - the Admin UI should continue working if the proxy crosses it's global budget - and RouteChecks.is_llm_api_route(route=route) - and route != "/v1/models" - and route != "/models" - ): - if global_proxy_spend > litellm.max_budget: - raise litellm.BudgetExceededError( - current_cost=global_proxy_spend, max_budget=litellm.max_budget - ) - - _request_metadata: dict = request_body.get("metadata", {}) or {} - if _request_metadata.get("guardrails"): - # check if team allowed to modify guardrails - from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails - - can_modify: bool = can_modify_guardrails(team_object) - if can_modify is False: - from fastapi import HTTPException - - raise HTTPException( - status_code=403, - detail={ - "error": "Your team does not have permission to modify guardrails." - }, - ) + _enforce_user_param_check(general_settings, request, request_body, route) + _reject_clientside_metadata_tags_check(general_settings, request_body, route) + _global_proxy_budget_check(global_proxy_spend, skip_budget_checks, route) + _guardrail_modification_check(request_body, team_object) # 10 [OPTIONAL] Organization RBAC checks organization_role_based_access_check( @@ -1932,9 +1962,8 @@ class ExperimentalUIJWTToken: def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str: from datetime import timedelta - from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - encrypt_value_helper, - ) + from litellm.proxy.common_utils.encrypt_decrypt_utils import \ + encrypt_value_helper if user_info.user_role is None: raise Exception("User role is required for experimental UI login") @@ -1980,9 +2009,8 @@ class ExperimentalUIJWTToken: """ from datetime import timedelta - from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - encrypt_value_helper, - ) + from litellm.proxy.common_utils.encrypt_decrypt_utils import \ + encrypt_value_helper if user_info.user_role is None: raise Exception("User role is required for CLI JWT login") @@ -2021,9 +2049,8 @@ class ExperimentalUIJWTToken: import json from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth - from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - decrypt_value_helper, - ) + from litellm.proxy.common_utils.encrypt_decrypt_utils import \ + decrypt_value_helper decrypted_token = decrypt_value_helper( hashed_token, key="ui_hash_key", exception_type="debug" @@ -2144,13 +2171,11 @@ async def get_key_object( ) # else, check db - _valid_token: Optional[BaseModel] = ( - await _fetch_key_object_from_db_with_reconnect( - hashed_token=hashed_token, - prisma_client=prisma_client, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) + _valid_token: Optional[BaseModel] = await _fetch_key_object_from_db_with_reconnect( + hashed_token=hashed_token, + prisma_client=prisma_client, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) if _valid_token is None: @@ -2296,9 +2321,9 @@ async def get_org_object( # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, - value=response.model_dump() - if hasattr(response, "model_dump") - else response, + value=( + response.model_dump() if hasattr(response, "model_dump") else response + ), ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -2341,8 +2366,10 @@ async def _get_resources_from_access_groups( # Lazy import to avoid circular imports if prisma_client is None or user_api_key_cache is None: from litellm.proxy.proxy_server import prisma_client as _prisma_client - from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj - from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache + from litellm.proxy.proxy_server import \ + proxy_logging_obj as _proxy_logging_obj + from litellm.proxy.proxy_server import \ + user_api_key_cache as _user_api_key_cache prisma_client = prisma_client or _prisma_client user_api_key_cache = user_api_key_cache or _user_api_key_cache @@ -3298,7 +3325,8 @@ async def _tag_max_budget_check( BudgetExceededError if any tag is over its max budget. Triggers a budget alert if any tag is over its max budget. """ - from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body + from litellm.proxy.common_utils.http_parsing_utils import \ + get_tags_from_request_body if prisma_client is None: return diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 82341c9a704..94cb7510a54 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -25,53 +25,38 @@ from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( - ExperimentalUIJWTToken, - _cache_key_object, - _delete_cache_key_object, - _get_user_role, - _is_user_proxy_admin, - _virtual_key_max_budget_alert_check, - _virtual_key_max_budget_check, - _virtual_key_soft_budget_check, - can_key_call_model, - common_checks, - get_end_user_object, - get_jwt_key_mapping_object, - get_key_object, - get_project_object, - get_team_object, - get_user_object, - is_valid_fallback_model, -) -from litellm.proxy.auth.auth_exception_handler import UserAPIKeyAuthExceptionHandler -from litellm.proxy.auth.auth_utils import ( - abbreviate_api_key, - get_end_user_id_from_request_body, - get_model_from_request, - get_request_route, - normalize_request_route, - pre_db_read_auth_checks, - route_in_additonal_public_routes, -) + ExperimentalUIJWTToken, _cache_key_object, _delete_cache_key_object, + _get_user_role, _is_user_proxy_admin, _virtual_key_max_budget_alert_check, + _virtual_key_max_budget_check, _virtual_key_soft_budget_check, + can_key_call_model, common_checks, get_end_user_object, + get_jwt_key_mapping_object, get_key_object, get_project_object, + get_team_object, get_user_object, is_valid_fallback_model) +from litellm.proxy.auth.auth_exception_handler import \ + UserAPIKeyAuthExceptionHandler +from litellm.proxy.auth.auth_utils import (abbreviate_api_key, + get_end_user_id_from_request_body, + get_model_from_request, + get_request_route, + normalize_request_route, + pre_db_read_auth_checks, + route_in_additonal_public_routes) from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler from litellm.proxy.auth.oauth2_check import Oauth2Handler from litellm.proxy.auth.oauth2_proxy_hook import handle_oauth2_proxy_request from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.common_utils.cache_coordinator import EventDrivenCacheCoordinator +from litellm.proxy.common_utils.cache_coordinator import \ + EventDrivenCacheCoordinator from litellm.proxy.common_utils.http_parsing_utils import ( - _read_request_body, - _safe_get_request_headers, - populate_request_with_path_params, -) + _read_request_body, _safe_get_request_headers, + populate_request_with_path_params) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.secret_managers.main import get_secret_bool from litellm.types.services import ServiceTypes try: - from litellm_enterprise.proxy.auth.user_api_key_auth import ( - enterprise_custom_auth as _enterprise_custom_auth, - ) + from litellm_enterprise.proxy.auth.user_api_key_auth import \ + enterprise_custom_auth as _enterprise_custom_auth enterprise_custom_auth: Optional[Callable] = _enterprise_custom_auth except ImportError as e: @@ -351,9 +336,8 @@ def get_api_key( Tuple[Optional[str], Optional[str]]: Tuple of the api_key and the passed_in_key """ from litellm.proxy.auth.route_checks import RouteChecks - from litellm.proxy.common_utils.http_parsing_utils import ( - _safe_get_request_query_params, - ) + from litellm.proxy.common_utils.http_parsing_utils import \ + _safe_get_request_query_params api_key = api_key passed_in_key: Optional[str] = None @@ -519,20 +503,15 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 request_data: dict, custom_litellm_key_header: Optional[str] = None, ) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import ( - general_settings, - jwt_handler, - litellm_proxy_admin_name, - llm_model_list, - llm_router, - master_key, - model_max_budget_limiter, - open_telemetry_logger, - prisma_client, - proxy_logging_obj, - user_api_key_cache, - user_custom_auth, - ) + from litellm.proxy.proxy_server import (general_settings, jwt_handler, + litellm_proxy_admin_name, + llm_model_list, llm_router, + master_key, + model_max_budget_limiter, + open_telemetry_logger, + prisma_client, proxy_logging_obj, + user_api_key_cache, + user_custom_auth) parent_otel_span: Optional[Span] = None start_time = datetime.now() @@ -730,9 +709,11 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_object is not None else None ), - team_metadata=team_object.metadata - if team_object is not None - else None, + team_metadata=( + team_object.metadata + if team_object is not None + else None + ), org_id=org_id, end_user_id=end_user_id, parent_otel_span=parent_otel_span, @@ -750,9 +731,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 team_rpm_limit=( team_object.rpm_limit if team_object is not None else None ), - team_models=team_object.models - if team_object is not None - else [], + team_models=( + team_object.models if team_object is not None else [] + ), user_role=( LitellmUserRoles(user_object.user_role) if user_object is not None @@ -779,16 +760,17 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if team_membership is not None else None ), - team_metadata=team_object.metadata - if team_object is not None - else None, + team_metadata=( + team_object.metadata if team_object is not None else None + ), ) # Check if model has zero cost - if so, skip all budget checks model = get_model_from_request(request_data, route) skip_budget_checks = False if model is not None and llm_router is not None: - from litellm.proxy.auth.auth_checks import _is_model_cost_zero + from litellm.proxy.auth.auth_checks import \ + _is_model_cost_zero skip_budget_checks = _is_model_cost_zero( model=model, llm_router=llm_router @@ -893,9 +875,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, ) if _end_user_object is not None: - end_user_params[ - "allowed_model_region" - ] = _end_user_object.allowed_model_region + end_user_params["allowed_model_region"] = ( + _end_user_object.allowed_model_region + ) if _end_user_object.litellm_budget_table is not None: _apply_budget_limits_to_end_user_params( end_user_params=end_user_params, @@ -904,9 +886,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) elif litellm.max_end_user_budget_id is not None: # End user doesn't exist yet, but apply default budget limits if configured - from litellm.proxy.auth.auth_checks import ( - get_default_end_user_budget, - ) + from litellm.proxy.auth.auth_checks import \ + get_default_end_user_budget default_budget = await get_default_end_user_budget( prisma_client=prisma_client, @@ -1463,9 +1444,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if _end_user_object is not None: valid_token_dict.update(end_user_params) - valid_token_dict[ - "end_user_object_permission" - ] = _end_user_object.object_permission + valid_token_dict["end_user_object_permission"] = ( + _end_user_object.object_permission + ) # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions # sso/login, ui/login, /key functions and /user functions @@ -1687,7 +1668,8 @@ async def _lookup_end_user_and_apply_budget( valid_token=valid_token, end_user_params=end_user_params ) elif litellm.max_end_user_budget_id is not None: - from litellm.proxy.auth.auth_checks import get_default_end_user_budget + from litellm.proxy.auth.auth_checks import \ + get_default_end_user_budget default_budget = await get_default_end_user_budget( prisma_client=prisma_client, @@ -1718,14 +1700,10 @@ async def _run_post_custom_auth_checks( route: str, parent_otel_span: Optional[Span], ) -> UserAPIKeyAuth: - from litellm.proxy.proxy_server import ( - prisma_client, - user_api_key_cache, - proxy_logging_obj, - general_settings, - llm_router, - model_max_budget_limiter, - ) + from litellm.proxy.proxy_server import (general_settings, llm_router, + model_max_budget_limiter, + prisma_client, proxy_logging_obj, + user_api_key_cache) # 1. Look up end_user object from DB if end_user_id is set end_user_object = None @@ -1756,9 +1734,11 @@ async def _run_post_custom_auth_checks( message=f"Authentication Error - Expired Key. Key Expiry time {expiry_time} and current time {current_time}", type=ProxyErrorTypes.expired_key, code=400, - param=abbreviate_api_key(api_key=valid_token.token) - if valid_token.token - else "", + param=( + abbreviate_api_key(api_key=valid_token.token) + if valid_token.token + else "" + ), ) current_model = request_data.get("model", None) diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4c96e079c9e..28b1e6601b1 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -13,36 +13,49 @@ import random import time import traceback from datetime import datetime, timedelta, timezone -from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, - cast, overload) +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Optional, + Union, + cast, + overload, +) import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache, RedisCache from litellm.constants import DB_SPEND_UPDATE_JOB_NAME from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.proxy._types import (DB_CONNECTION_ERROR_TYPES, - BaseDailySpendTransaction, - DailyAgentSpendTransaction, - DailyEndUserSpendTransaction, - DailyOrganizationSpendTransaction, - DailyTagSpendTransaction, - DailyTeamSpendTransaction, - DailyUserSpendTransaction, - DBSpendUpdateTransactions, - Litellm_EntityType, LiteLLM_UserTable, - SpendLogsMetadata, SpendLogsPayload, - SpendUpdateQueueItem, ToolDiscoveryQueueItem) -from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import \ - DailySpendUpdateQueue -from litellm.proxy.db.db_transaction_queue.pod_lock_manager import \ - PodLockManager -from litellm.proxy.db.db_transaction_queue.redis_update_buffer import \ - RedisUpdateBuffer -from litellm.proxy.db.db_transaction_queue.spend_update_queue import \ - SpendUpdateQueue -from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import \ - ToolDiscoveryQueue +from litellm.proxy._types import ( + DB_CONNECTION_ERROR_TYPES, + BaseDailySpendTransaction, + DailyAgentSpendTransaction, + DailyEndUserSpendTransaction, + DailyOrganizationSpendTransaction, + DailyTagSpendTransaction, + DailyTeamSpendTransaction, + DailyUserSpendTransaction, + DBSpendUpdateTransactions, + Litellm_EntityType, + LiteLLM_UserTable, + SpendLogsMetadata, + SpendLogsPayload, + SpendUpdateQueueItem, + ToolDiscoveryQueueItem, +) +from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( + DailySpendUpdateQueue, +) +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, +) from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING if TYPE_CHECKING: @@ -91,10 +104,12 @@ class DBSpendUpdateWriter: end_time: Optional[datetime], response_cost: Optional[float], ): - from litellm.proxy.proxy_server import (disable_spend_logs, - litellm_proxy_budget_name, - prisma_client, - user_api_key_cache) + from litellm.proxy.proxy_server import ( + disable_spend_logs, + litellm_proxy_budget_name, + prisma_client, + user_api_key_cache, + ) from litellm.proxy.utils import ProxyUpdateSpend, hash_token try: @@ -109,8 +124,9 @@ class DBSpendUpdateWriter: hashed_token = token ## CREATE SPEND LOG PAYLOAD ## - from litellm.proxy.spend_tracking.spend_tracking_utils import \ - get_logging_payload + from litellm.proxy.spend_tracking.spend_tracking_utils import ( + get_logging_payload, + ) payload = get_logging_payload( kwargs=kwargs, @@ -374,6 +390,19 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + _agent_id_for_spend = payload_copy.get("agent_id") + try: + await self._update_agent_db( + response_cost=response_cost, + agent_id=_agent_id_for_spend, + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_agent_db failed: %s", + traceback.format_exc(), + ) + try: await self.add_spend_log_transaction_to_daily_user_transaction( payload=payload_copy, @@ -604,6 +633,34 @@ class DBSpendUpdateWriter: ) raise e + async def _update_agent_db( + self, + response_cost: Optional[float], + agent_id: Optional[str], + prisma_client: Optional[PrismaClient], + ): + try: + if agent_id is None or prisma_client is None: + return + + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.AGENT, + entity_id=agent_id, + response_cost=response_cost, + ) + ) + except Exception as e: + verbose_proxy_logger.error( + "Spend tracking - failed to enqueue agent spend update. " + "agent_id=%s, response_cost=%s - %s\n%s", + agent_id, + response_cost, + str(e), + traceback.format_exc(), + ) + raise e + async def _update_tag_db( self, response_cost: Optional[float], @@ -765,7 +822,7 @@ class DBSpendUpdateWriter: if db_spend_update_transactions is not None: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " - "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d", + "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d", len( db_spend_update_transactions.get("key_list_transactions") or {} @@ -798,6 +855,12 @@ class DBSpendUpdateWriter: db_spend_update_transactions.get("tag_list_transactions") or {} ), + len( + db_spend_update_transactions.get( + "agent_list_transactions" + ) + or {} + ), ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -1002,8 +1065,10 @@ class DBSpendUpdateWriter: Commits all the spend `UPDATE` transactions to the Database """ - from litellm.proxy.utils import (ProxyUpdateSpend, - _raise_failed_update_spend_exception) + from litellm.proxy.utils import ( + ProxyUpdateSpend, + _raise_failed_update_spend_exception, + ) ### UPDATE USER TABLE ### user_list_transactions = db_spend_update_transactions["user_list_transactions"] @@ -1279,6 +1344,18 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + ### UPDATE AGENT TABLE ### + agent_list_transactions = db_spend_update_transactions["agent_list_transactions"] + await DBSpendUpdateWriter._update_entity_spend_in_db( + entity_name="Agent", + transactions=agent_list_transactions, + table_accessor="litellm_agentstable", + where_field="agent_id", + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + @staticmethod async def _update_entity_spend_in_db( entity_name: str, @@ -2031,9 +2108,6 @@ class DBSpendUpdateWriter: ) return if payload["agent_id"] is None: - verbose_proxy_logger.debug( - "agent_id is None for request. Skipping incrementing agent spend." - ) return payload_with_agent_id = cast( SpendLogsPayload, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 51201f96d77..4f38e71bbfa 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -10,33 +10,31 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache -from litellm.constants import ( - MAX_REDIS_BUFFER_DEQUEUE_COUNT, - REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, - REDIS_UPDATE_BUFFER_KEY, -) +from litellm.constants import (MAX_REDIS_BUFFER_DEQUEUE_COUNT, + REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, + REDIS_UPDATE_BUFFER_KEY) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._types import ( - DailyTagSpendTransaction, - DailyTeamSpendTransaction, - DailyUserSpendTransaction, - DailyOrganizationSpendTransaction, - DailyEndUserSpendTransaction, - DBSpendUpdateTransactions, - DailyAgentSpendTransaction, -) -from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj -from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( - DailySpendUpdateQueue, -) -from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy._types import (DailyAgentSpendTransaction, + DailyEndUserSpendTransaction, + DailyOrganizationSpendTransaction, + DailyTagSpendTransaction, + DailyTeamSpendTransaction, + DailyUserSpendTransaction, + DBSpendUpdateTransactions) +from litellm.proxy.db.db_transaction_queue.base_update_queue import \ + service_logger_obj +from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import \ + DailySpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.spend_update_queue import \ + SpendUpdateQueue from litellm.secret_managers.main import str_to_bool -from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation +from litellm.types.caching import (RedisPipelineLpopOperation, + RedisPipelineRpushOperation) from litellm.types.services import ServiceTypes if TYPE_CHECKING: @@ -579,6 +577,7 @@ class RedisUpdateBuffer: team_member_list_transactions={}, org_list_transactions={}, tag_list_transactions={}, + agent_list_transactions={}, ) # Define the transaction fields to process @@ -590,6 +589,7 @@ class RedisUpdateBuffer: "team_member_list_transactions", "org_list_transactions", "tag_list_transactions", + "agent_list_transactions", ] # Loop through each transaction and combine the values diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 3e059cf8c1f..b7cd06a64f3 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -3,15 +3,10 @@ from typing import Dict, List, Optional from litellm._logging import verbose_proxy_logger from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE -from litellm.proxy._types import ( - DBSpendUpdateTransactions, - Litellm_EntityType, - SpendUpdateQueueItem, -) +from litellm.proxy._types import (DBSpendUpdateTransactions, + Litellm_EntityType, SpendUpdateQueueItem) from litellm.proxy.db.db_transaction_queue.base_update_queue import ( - BaseUpdateQueue, - service_logger_obj, -) + BaseUpdateQueue, service_logger_obj) from litellm.types.services import ServiceTypes @@ -145,6 +140,7 @@ class SpendUpdateQueue(BaseUpdateQueue): team_member_list_transactions={}, org_list_transactions={}, tag_list_transactions={}, + agent_list_transactions={}, ) # Map entity types to their corresponding transaction dictionary keys @@ -156,6 +152,7 @@ class SpendUpdateQueue(BaseUpdateQueue): Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions", Litellm_EntityType.ORGANIZATION: "org_list_transactions", Litellm_EntityType.TAG: "tag_list_transactions", + Litellm_EntityType.AGENT: "agent_list_transactions", } for update in updates: @@ -207,6 +204,10 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict = db_spend_update_transactions[ "tag_list_transactions" ] + elif dict_key == "agent_list_transactions": + transactions_dict = db_spend_update_transactions[ + "agent_list_transactions" + ] else: continue diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index feea3023d46..2d0ce040a6c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -341,6 +341,30 @@ class GenericGuardrailAPI(CustomGuardrail): return_inputs["tools"] = tools return return_inputs + def _handle_guardrail_request_error( + self, + error: Exception, + inputs: GenericGuardrailAPIInputs, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"], + is_unreachable: bool = True, + ) -> GenericGuardrailAPIInputs: + if is_unreachable and self.unreachable_fallback == "fail_open": + http_status_code = getattr( + getattr(error, "response", None), "status_code", None + ) + return self._fail_open_passthrough( + inputs=inputs, + input_type=input_type, + logging_obj=logging_obj, + error=error, + **({"http_status_code": http_status_code} if http_status_code else {}), + ) + verbose_proxy_logger.error( + "Generic Guardrail API: failed to make request: %s", str(error) + ) + raise Exception(f"Generic Guardrail API failed: {str(error)}") + @log_guardrail_information async def apply_guardrail( self, @@ -466,58 +490,24 @@ class GenericGuardrailAPI(CustomGuardrail): ) except GuardrailRaisedException: - # Re-raise guardrail exceptions as-is raise except Timeout as e: - # AsyncHTTPHandler wraps httpx.TimeoutException into litellm.Timeout - if self.unreachable_fallback == "fail_open": - return self._fail_open_passthrough( - inputs=inputs, - input_type=input_type, - logging_obj=logging_obj, - error=e, - ) - - verbose_proxy_logger.error( - "Generic Guardrail API: failed to make request: %s", str(e) + return self._handle_guardrail_request_error( + e, inputs, input_type, logging_obj ) - raise Exception(f"Generic Guardrail API failed: {str(e)}") except httpx.HTTPStatusError as e: - # Common reverse-proxy/LB failures can present as HTTP errors even when the backend is unreachable. - status_code = getattr(getattr(e, "response", None), "status_code", None) - if self.unreachable_fallback == "fail_open" and status_code in ( - 502, - 503, - 504, - ): - return self._fail_open_passthrough( - inputs=inputs, - input_type=input_type, - logging_obj=logging_obj, - error=e, - http_status_code=status_code, - ) - - verbose_proxy_logger.error( - "Generic Guardrail API: failed to make request: %s", str(e) + status_code = getattr( + getattr(e, "response", None), "status_code", None + ) + is_unreachable = status_code in (502, 503, 504) + return self._handle_guardrail_request_error( + e, inputs, input_type, logging_obj, is_unreachable=is_unreachable ) - raise Exception(f"Generic Guardrail API failed: {str(e)}") except httpx.RequestError as e: - # Guardrail endpoint is unreachable (DNS/connect/timeout/etc) - if self.unreachable_fallback == "fail_open": - return self._fail_open_passthrough( - inputs=inputs, - input_type=input_type, - logging_obj=logging_obj, - error=e, - ) - - verbose_proxy_logger.error( - "Generic Guardrail API: failed to make request: %s", str(e) + return self._handle_guardrail_request_error( + e, inputs, input_type, logging_obj ) - raise Exception(f"Generic Guardrail API failed: {str(e)}") except Exception as e: - verbose_proxy_logger.error( - "Generic Guardrail API: failed to make request: %s", str(e) + return self._handle_guardrail_request_error( + e, inputs, input_type, logging_obj, is_unreachable=False ) - raise Exception(f"Generic Guardrail API failed: {str(e)}") diff --git a/litellm/proxy/hooks/__init__.py b/litellm/proxy/hooks/__init__.py index 1d1e559d4be..790ebcd8791 100644 --- a/litellm/proxy/hooks/__init__.py +++ b/litellm/proxy/hooks/__init__.py @@ -5,6 +5,8 @@ from . import * from .cache_control_check import _PROXY_CacheControlCheck from .litellm_skills import SkillsInjectionHook from .max_budget_limiter import _PROXY_MaxBudgetLimiter +from .max_budget_per_session_limiter import _PROXY_MaxBudgetPerSessionHandler +from .max_iterations_limiter import _PROXY_MaxIterationsHandler from .parallel_request_limiter import _PROXY_MaxParallelRequestsHandler from .parallel_request_limiter_v3 import _PROXY_MaxParallelRequestsHandler_v3 from .responses_id_security import ResponsesIDSecurity @@ -23,6 +25,8 @@ PROXY_HOOKS = { "cache_control_check": _PROXY_CacheControlCheck, "responses_id_security": ResponsesIDSecurity, "litellm_skills": SkillsInjectionHook, + "max_iterations_limiter": _PROXY_MaxIterationsHandler, + "max_budget_per_session_limiter": _PROXY_MaxBudgetPerSessionHandler, } ## FEATURE FLAG HOOKS ## diff --git a/litellm/proxy/hooks/max_budget_per_session_limiter.py b/litellm/proxy/hooks/max_budget_per_session_limiter.py new file mode 100644 index 00000000000..a981207f000 --- /dev/null +++ b/litellm/proxy/hooks/max_budget_per_session_limiter.py @@ -0,0 +1,271 @@ +""" +Per-Session Budget Limiter for LiteLLM Proxy. + +Enforces a dollar-amount cap per session (identified by `session_id` / +`x-litellm-trace-id`). After each successful LLM call the response cost is +accumulated against the session. When the accumulated spend exceeds +`max_budget_per_session` (configured in agent litellm_params), subsequent +requests for that session receive a 429. + +Note: trace-id enforcement (require_trace_id_on_calls_by_agent) is handled +separately in auth_checks.py at the agent level, not in this hook. + +Works across multiple proxy instances via DualCache (in-memory + Redis). +Follows the same pattern as max_iterations_limiter.py. +""" + +import os +from typing import TYPE_CHECKING, Any, Optional, Union + +from fastapi import HTTPException + +from litellm import DualCache +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_logger import CustomLogger +from litellm.proxy._types import UserAPIKeyAuth + +if TYPE_CHECKING: + from litellm.proxy.utils import InternalUsageCache as _InternalUsageCache + + InternalUsageCache = _InternalUsageCache +else: + InternalUsageCache = Any + + +# Redis Lua script for atomic float increment with TTL. +# INCRBYFLOAT returns the new value as a string. +# Only sets EXPIRE on first call (when prior value was nil). +MAX_BUDGET_SESSION_INCREMENT_SCRIPT = """ +local key = KEYS[1] +local amount = ARGV[1] +local ttl = tonumber(ARGV[2]) + +local existed = redis.call('EXISTS', key) +local new_val = redis.call('INCRBYFLOAT', key, amount) +if existed == 0 then + redis.call('EXPIRE', key, ttl) +end + +return new_val +""" + +# Default TTL for session budget counters (1 hour) +DEFAULT_MAX_BUDGET_PER_SESSION_TTL = 3600 + + +class _PROXY_MaxBudgetPerSessionHandler(CustomLogger): + """ + Pre-call hook that enforces max_budget_per_session. + + Configuration (set in agent litellm_params): + - max_budget_per_session: dollar cap per session_id + + Cache key pattern: + {session_budget:}:spend + """ + + def __init__(self, internal_usage_cache: InternalUsageCache): + self.internal_usage_cache = internal_usage_cache + self.ttl = int( + os.getenv( + "LITELLM_MAX_BUDGET_PER_SESSION_TTL", + DEFAULT_MAX_BUDGET_PER_SESSION_TTL, + ) + ) + + if self.internal_usage_cache.dual_cache.redis_cache is not None: + self.increment_script = ( + self.internal_usage_cache.dual_cache.redis_cache.async_register_script( + MAX_BUDGET_SESSION_INCREMENT_SCRIPT + ) + ) + else: + self.increment_script = None + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict, + call_type: str, + ) -> Optional[Union[Exception, str, dict]]: + """ + Before each LLM call, check if max_budget_per_session is set and + whether accumulated spend exceeds the budget (429 if so). + """ + max_budget = self._get_max_budget_per_session(user_api_key_dict) + + session_id = self._get_session_id(data) + + if max_budget is None or session_id is None: + return None + + max_budget = float(max_budget) + cache_key = self._make_cache_key(session_id) + current_spend = await self._get_current_spend(cache_key) + + verbose_proxy_logger.debug( + "MaxBudgetPerSessionHandler: session_id=%s, spend=%.4f, max=%.2f", + session_id, + current_spend, + max_budget, + ) + + if current_spend >= max_budget: + raise HTTPException( + status_code=429, + detail=( + f"Session budget exceeded for session {session_id}. " + f"Current spend: ${current_spend:.4f}, " + f"max_budget_per_session: ${max_budget:.2f}." + ), + ) + + return None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + """ + After a successful LLM call, increment the session spend by the response cost. + """ + try: + litellm_params = kwargs.get("litellm_params") or {} + metadata = litellm_params.get("metadata") or {} + session_id = metadata.get("session_id") + if session_id is None: + return + + agent_id = metadata.get("agent_id") + if agent_id is None: + return + + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + ) + + agent = global_agent_registry.get_agent_by_id(agent_id=str(agent_id)) + if agent is None: + return + + agent_litellm_params = agent.litellm_params or {} + max_budget = agent_litellm_params.get("max_budget_per_session") + if max_budget is None: + return + + response_cost = kwargs.get("response_cost") or 0.0 + if response_cost <= 0: + return + + cache_key = self._make_cache_key(str(session_id)) + await self._increment_spend(cache_key, float(response_cost)) + + verbose_proxy_logger.debug( + "MaxBudgetPerSessionHandler: incremented session %s spend by %.6f", + session_id, + response_cost, + ) + except Exception as e: + verbose_proxy_logger.warning( + "MaxBudgetPerSessionHandler: error in async_log_success_event: %s", + str(e), + ) + + def _get_session_id(self, data: dict) -> Optional[str]: + """Extract session_id from request metadata.""" + metadata = data.get("metadata") or {} + session_id = metadata.get("session_id") + if session_id is not None: + return str(session_id) + + litellm_metadata = data.get("litellm_metadata") or {} + session_id = litellm_metadata.get("session_id") + if session_id is not None: + return str(session_id) + + return None + + def _get_max_budget_per_session( + self, user_api_key_dict: UserAPIKeyAuth + ) -> Optional[float]: + """Extract max_budget_per_session from agent litellm_params.""" + agent_id = user_api_key_dict.agent_id + if agent_id is None: + return None + + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + agent = global_agent_registry.get_agent_by_id(agent_id=agent_id) + if agent is None: + return None + + litellm_params = agent.litellm_params or {} + max_budget = litellm_params.get("max_budget_per_session") + if max_budget is not None: + return float(max_budget) + return None + + def _make_cache_key(self, session_id: str) -> str: + return f"{{session_budget:{session_id}}}:spend" + + async def _get_current_spend(self, cache_key: str) -> float: + """Read current accumulated spend for a session.""" + if ( + self.internal_usage_cache.dual_cache.redis_cache is not None + ): + try: + result = await self.internal_usage_cache.dual_cache.redis_cache.async_get_cache( + key=cache_key + ) + if result is not None: + return float(result) + return 0.0 + except Exception as e: + verbose_proxy_logger.warning( + "MaxBudgetPerSessionHandler: Redis GET failed, " + "falling back to in-memory: %s", + str(e), + ) + + result = await self.internal_usage_cache.async_get_cache( + key=cache_key, + litellm_parent_otel_span=None, + local_only=True, + ) + if result is not None: + return float(result) + return 0.0 + + async def _increment_spend(self, cache_key: str, amount: float) -> float: + """Atomically increment the session spend and return the new value.""" + if self.increment_script is not None: + try: + result = await self.increment_script( + keys=[cache_key], + args=[str(amount), self.ttl], + ) + return float(result) + except Exception as e: + verbose_proxy_logger.warning( + "MaxBudgetPerSessionHandler: Redis INCRBYFLOAT failed, " + "falling back to in-memory: %s", + str(e), + ) + + return await self._in_memory_increment_spend(cache_key, amount) + + async def _in_memory_increment_spend( + self, cache_key: str, amount: float + ) -> float: + current = await self.internal_usage_cache.async_get_cache( + key=cache_key, + litellm_parent_otel_span=None, + local_only=True, + ) + new_value = (float(current) if current is not None else 0.0) + amount + await self.internal_usage_cache.async_set_cache( + key=cache_key, + value=new_value, + ttl=self.ttl, + litellm_parent_otel_span=None, + local_only=True, + ) + return new_value diff --git a/litellm/proxy/hooks/max_iterations_limiter.py b/litellm/proxy/hooks/max_iterations_limiter.py index 8d481f6b261..b6fde2b1780 100644 --- a/litellm/proxy/hooks/max_iterations_limiter.py +++ b/litellm/proxy/hooks/max_iterations_limiter.py @@ -4,7 +4,7 @@ Max Iterations Limiter for LiteLLM Proxy. Enforces a per-session cap on the number of LLM calls an agentic loop can make. Callers send a `session_id` with each request (via `x-litellm-session-id` header or `metadata.session_id`), and this hook counts calls per session. When the count -exceeds `max_iterations` (configured in key/team metadata), returns 429. +exceeds `max_iterations` (configured in agent litellm_params or key metadata), returns 429. Works across multiple proxy instances via DualCache (in-memory + Redis). Follows the same pattern as parallel_request_limiter_v3.py. @@ -52,8 +52,9 @@ class _PROXY_MaxIterationsHandler(CustomLogger): Pre-call hook that enforces max_iterations per session. Configuration: - - max_iterations: set in key metadata via /key/generate or /key/update - e.g. metadata={"max_iterations": 25} + - max_iterations: set in agent litellm_params (preferred) + e.g. litellm_params={"max_iterations": 25} + Falls back to key metadata max_iterations for backwards compatibility. - session_id: sent by caller via x-litellm-session-id header or metadata.session_id in request body @@ -93,14 +94,13 @@ class _PROXY_MaxIterationsHandler(CustomLogger): Check session iteration count before making the API call. Extracts session_id from request metadata and max_iterations from - key metadata. If the session has exceeded max_iterations, raises 429. + agent litellm_params. If the session has exceeded max_iterations, raises 429. """ # Extract session_id from request data session_id = self._get_session_id(data) if session_id is None: return None - # Extract max_iterations from key metadata max_iterations = self._get_max_iterations(user_api_key_dict) if max_iterations is None: return None @@ -151,7 +151,22 @@ class _PROXY_MaxIterationsHandler(CustomLogger): def _get_max_iterations( self, user_api_key_dict: UserAPIKeyAuth ) -> Optional[int]: - """Extract max_iterations from key metadata.""" + """Extract max_iterations from agent litellm_params, with fallback to key metadata.""" + # Try agent litellm_params first + agent_id = user_api_key_dict.agent_id + if agent_id is not None: + from litellm.proxy.agent_endpoints.agent_registry import ( + global_agent_registry, + ) + + agent = global_agent_registry.get_agent_by_id(agent_id=agent_id) + if agent is not None: + litellm_params = agent.litellm_params or {} + max_iterations = litellm_params.get("max_iterations") + if max_iterations is not None: + return int(max_iterations) + + # Fallback to key metadata for backwards compatibility metadata = user_api_key_dict.metadata or {} max_iterations = metadata.get("max_iterations") if max_iterations is not None: diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index b5bbb4237c1..856975ea093 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -7,18 +7,8 @@ This is currently in development and not yet ready for production. import binascii import os from datetime import datetime -from typing import ( - TYPE_CHECKING, - Any, - Callable, - Dict, - List, - Literal, - Optional, - TypedDict, - Union, - cast, -) +from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, + Optional, TypedDict, Union, cast) from fastapi import HTTPException @@ -175,9 +165,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """Get or lazy-load the batch rate limiter.""" if self._batch_rate_limiter is None: try: - from litellm.proxy.hooks.batch_rate_limiter import ( - _PROXY_BatchRateLimiter, - ) + from litellm.proxy.hooks.batch_rate_limiter import \ + _PROXY_BatchRateLimiter self._batch_rate_limiter = _PROXY_BatchRateLimiter( internal_usage_cache=self.internal_usage_cache, @@ -679,10 +668,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): requested_model: The model being requested descriptors: List of rate limit descriptors to append to """ - from litellm.proxy.auth.auth_utils import ( - get_key_model_rpm_limit, - get_key_model_tpm_limit, - ) + from litellm.proxy.auth.auth_utils import (get_key_model_rpm_limit, + get_key_model_tpm_limit) if not requested_model: return @@ -791,6 +778,92 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ return rpm_limit_type == "dynamic" or tpm_limit_type == "dynamic" + def _get_agent_from_registry(self, agent_id: str) -> Optional[Any]: + """Look up an agent from the in-memory registry by ID.""" + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry + + return global_agent_registry.get_agent_by_id(agent_id=agent_id) + + def _get_resolved_agent_id( + self, user_api_key_dict: UserAPIKeyAuth, data: dict + ) -> Optional[str]: + """ + Resolve the agent_id from either the API key or request metadata. + Key-level agent_id takes precedence over metadata/header-supplied agent_id. + """ + key_agent_id = getattr(user_api_key_dict, "agent_id", None) + if key_agent_id: + return key_agent_id + metadata = data.get("metadata") or {} + return metadata.get("agent_id") + + def _get_session_id_from_data(self, data: dict) -> Optional[str]: + """Extract session_id from request metadata or litellm_session_id.""" + session_id = data.get("litellm_session_id") + if session_id: + return str(session_id) + metadata = data.get("metadata") or {} + session_id = metadata.get("session_id") + if session_id: + return str(session_id) + litellm_metadata = data.get("litellm_metadata") or {} + session_id = litellm_metadata.get("session_id") + if session_id: + return str(session_id) + return None + + def _create_agent_rate_limit_descriptors( + self, + agent_id: str, + data: dict, + ) -> List[RateLimitDescriptor]: + """ + Create rate limit descriptors for agent-level and session-level limits. + + Agent-level: caps total RPM/TPM across all sessions for a given agent. + Session-level: caps RPM/TPM within a single session (identified by session_id). + """ + descriptors: List[RateLimitDescriptor] = [] + + agent = self._get_agent_from_registry(agent_id) + if agent is None: + return descriptors + + agent_rpm = getattr(agent, "rpm_limit", None) + agent_tpm = getattr(agent, "tpm_limit", None) + if agent_rpm is not None or agent_tpm is not None: + descriptors.append( + RateLimitDescriptor( + key="agent", + value=agent_id, + rate_limit={ + "requests_per_unit": agent_rpm, + "tokens_per_unit": agent_tpm, + "window_size": self.window_size, + }, + ) + ) + + session_rpm = getattr(agent, "session_rpm_limit", None) + session_tpm = getattr(agent, "session_tpm_limit", None) + if session_rpm is not None or session_tpm is not None: + session_id = self._get_session_id_from_data(data) + if session_id is not None: + descriptors.append( + RateLimitDescriptor( + key="agent_session", + value=f"{agent_id}:{session_id}", + rate_limit={ + "requests_per_unit": session_rpm, + "tokens_per_unit": session_tpm, + "window_size": self.window_size, + }, + ) + ) + + return descriptors + def _create_rate_limit_descriptors( self, user_api_key_dict: UserAPIKeyAuth, @@ -802,12 +875,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Create all rate limit descriptors for the request. - Returns list of descriptors for API key, user, team, team member, end user, and model-specific limits. + Returns list of descriptors for API key, user, team, team member, end user, + model-specific, agent, and agent-session limits. """ - from litellm.proxy.auth.auth_utils import ( - get_team_model_rpm_limit, - get_team_model_tpm_limit, - ) + from litellm.proxy.auth.auth_utils import (get_team_model_rpm_limit, + get_team_model_tpm_limit) descriptors = [] @@ -956,6 +1028,17 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + # Agent-level and session-level rate limits + resolved_agent_id = self._get_resolved_agent_id(user_api_key_dict, data) + + if resolved_agent_id: + descriptors.extend( + self._create_agent_rate_limit_descriptors( + agent_id=resolved_agent_id, + data=data, + ) + ) + return descriptors async def _check_model_has_recent_failures( @@ -970,9 +1053,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): Returns True if any deployment has failures in the current minute. """ from litellm.proxy.proxy_server import llm_router - from litellm.router_utils.router_callbacks.track_deployment_metrics import ( - get_deployment_failures_for_current_minute, - ) + from litellm.router_utils.router_callbacks.track_deployment_metrics import \ + get_deployment_failures_for_current_minute if llm_router is None: return False @@ -1386,12 +1468,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Update TPM usage on successful API calls by incrementing counters using pipeline """ - from litellm.litellm_core_utils.core_helpers import ( - _get_parent_otel_span_from_kwargs, - ) - from litellm.proxy.common_utils.callback_utils import ( - get_model_group_from_litellm_kwargs, - ) + from litellm.litellm_core_utils.core_helpers import \ + _get_parent_otel_span_from_kwargs + from litellm.proxy.common_utils.callback_utils import \ + get_model_group_from_litellm_kwargs from litellm.types.caching import RedisPipelineIncrementOperation rate_limit_type = self.get_rate_limit_type() @@ -1533,6 +1613,32 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) ) + # Agent TPM + agent_id = standard_logging_metadata.get("agent_id") + if agent_id: + pipeline_operations.extend( + self._create_pipeline_operations( + key="agent", + value=agent_id, + rate_limit_type="tokens", + total_tokens=total_tokens, + ) + ) + + # Agent Session TPM + session_id = standard_logging_metadata.get( + "session_id" + ) or standard_logging_metadata.get("trace_id") + if session_id: + pipeline_operations.extend( + self._create_pipeline_operations( + key="agent_session", + value=f"{agent_id}:{session_id}", + rate_limit_type="tokens", + total_tokens=total_tokens, + ) + ) + # Execute all increments in a single pipeline if pipeline_operations: await self.async_increment_tokens_with_ttl_preservation( @@ -1549,9 +1655,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): """ Decrement max parallel requests counter for the API Key """ - from litellm.litellm_core_utils.core_helpers import ( - _get_parent_otel_span_from_kwargs, - ) + from litellm.litellm_core_utils.core_helpers import \ + _get_parent_otel_span_from_kwargs from litellm.types.caching import RedisPipelineIncrementOperation try: diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d0fabe90103..cf4729db94b 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -690,6 +690,12 @@ class LiteLLMProxyRequestSetup: "user_api_key" ] = user_api_key_dict.api_key # this is just the hashed token + # Key-owned agent_id for spend attribution; keep existing (e.g. from header) if key has none + _key_agent_id = getattr(user_api_key_dict, "agent_id", None) + _existing_agent_id = data[_metadata_variable_name].get("agent_id") + _resolved_agent_id = _key_agent_id or _existing_agent_id + data[_metadata_variable_name]["agent_id"] = _resolved_agent_id + data[_metadata_variable_name]["user_api_end_user_max_budget"] = getattr( user_api_key_dict, "end_user_max_budget", None ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index db0bb735bab..bd7b21c3b55 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -373,8 +373,9 @@ from litellm.proxy.management_endpoints.fallback_management_endpoints import ( from litellm.proxy.management_endpoints.internal_user_endpoints import ( router as internal_user_router, ) -from litellm.proxy.management_endpoints.internal_user_endpoints import ( - user_update, +from litellm.proxy.management_endpoints.internal_user_endpoints import user_update +from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( + router as jwt_key_mapping_router, ) from litellm.proxy.management_endpoints.jwt_key_mapping_endpoints import ( router as jwt_key_mapping_router, @@ -444,9 +445,7 @@ from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_route from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, ) -from litellm.proxy.openai_files_endpoints.files_endpoints import ( - set_files_config, -) +from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( passthrough_endpoint_router, ) @@ -545,9 +544,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import ( LiteLLM_UpperboundKeyGenerateParams, ) from litellm.types.realtime import RealtimeQueryParams -from litellm.types.router import ( - DeploymentTypedDict, -) +from litellm.types.router import DeploymentTypedDict from litellm.types.router import ModelInfo as RouterModelInfo from litellm.types.router import ( RouterGeneralSettings, @@ -6682,6 +6679,11 @@ async def chat_completion( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + if ( + hasattr(user_api_key_dict, "agent_id") + and user_api_key_dict.agent_id is not None + ): + data["metadata"]["agent_id"] = user_api_key_dict.agent_id base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) try: result = await base_llm_response_processor.base_process_llm_request( @@ -6851,6 +6853,11 @@ async def completion( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + if ( + hasattr(user_api_key_dict, "agent_id") + and user_api_key_dict.agent_id is not None + ): + data["metadata"]["agent_id"] = user_api_key_dict.agent_id base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) return await base_llm_response_processor.base_process_llm_request( request=request, @@ -7088,6 +7095,11 @@ async def embeddings( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + if ( + hasattr(user_api_key_dict, "agent_id") + and user_api_key_dict.agent_id is not None + ): + data["metadata"]["agent_id"] = user_api_key_dict.agent_id # Use unified request processor (same as chat/completions and responses) base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 329ff80933f..8d4bdffb2dd 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -68,6 +68,11 @@ model LiteLLM_AgentsTable { agent_access_groups String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + spend Float @default(0.0) + tpm_limit Int? + rpm_limit Int? + session_tpm_limit Int? + session_rpm_limit Int? created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 65d90454ebd..2f9d27568e3 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -325,7 +325,7 @@ class ProxyLogging: if email_logger_class is not None: # All email logger classes now accept internal_usage_cache self.email_logging_instance = email_logger_class( - internal_usage_cache=self.internal_usage_cache.dual_cache, + internal_usage_cache=self.internal_usage_cache.dual_cache, # type: ignore[call-arg] ) self.premium_user = premium_user self.service_logging_obj = ServiceLogging() diff --git a/litellm/router.py b/litellm/router.py index 7119d2e850d..43b53d14d79 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -164,11 +164,7 @@ from litellm.types.utils import ( ) from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.utils import ( - ModelResponseStream, - StandardLoggingPayload, - Usage, -) +from litellm.types.utils import ModelResponseStream, StandardLoggingPayload, Usage from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 7879cae9ff6..951fbfcabd1 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -179,6 +179,10 @@ class AgentConfig(TypedDict, total=False): agent_card_params: Required[AgentCard] litellm_params: Dict[str, Any] # allow for any future litellm params object_permission: AgentObjectPermission + tpm_limit: Optional[int] + rpm_limit: Optional[int] + session_tpm_limit: Optional[int] + session_rpm_limit: Optional[int] static_headers: Optional[Dict[str, str]] extra_headers: Optional[List[str]] @@ -188,6 +192,10 @@ class PatchAgentRequest(TypedDict, total=False): agent_card_params: AgentCard litellm_params: Dict[str, Any] object_permission: AgentObjectPermission + tpm_limit: Optional[int] + rpm_limit: Optional[int] + session_tpm_limit: Optional[int] + session_rpm_limit: Optional[int] static_headers: Optional[Dict[str, str]] extra_headers: Optional[List[str]] @@ -201,6 +209,11 @@ class AgentResponse(BaseModel): litellm_params: Optional[Dict[str, Any]] = None agent_card_params: Dict[str, Any] object_permission: Optional[Dict[str, Any]] = None + spend: Optional[float] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + session_tpm_limit: Optional[int] = None + session_rpm_limit: Optional[int] = None static_headers: Optional[Dict[str, str]] = None extra_headers: Optional[List[str]] = None created_at: Optional[datetime] = None diff --git a/litellm/types/llms/custom_http.py b/litellm/types/llms/custom_http.py index 8f192d876c4..792adb4182d 100644 --- a/litellm/types/llms/custom_http.py +++ b/litellm/types/llms/custom_http.py @@ -25,6 +25,7 @@ class httpxSpecialProvider(str, Enum): MCP = "mcp" RAG = "rag" A2AProvider = "a2a_provider" + AgentHealthCheck = "agent_health_check" A2A = "a2a" PromptManagement = "prompt_management" UI = "ui" diff --git a/package.json b/package.json index a45e116b277..b5be819a451 100644 --- a/package.json +++ b/package.json @@ -12,8 +12,8 @@ }, "overrides": { "glob": ">=11.1.0", - "tar": ">=7.5.8", - "minimatch": ">=10.2.1", + "tar": ">=7.5.10", + "minimatch": ">=10.2.4", "diff": ">=8.0.3", "@isaacs/brace-expansion": ">=5.0.1", "@babel/traverse": ">=7.23.2", diff --git a/poetry.lock b/poetry.lock index 2c9f9802b1f..bad6a75b6e7 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,4 +1,4 @@ -# This file is automatically @generated by Poetry 2.2.0 and should not be changed by hand. +# This file is automatically @generated by Poetry 2.3.2 and should not be changed by hand. [[package]] name = "a2a-sdk" @@ -7,11 +7,11 @@ description = "A2A Python SDK" optional = false python-versions = ">=3.10" groups = ["main", "proxy-dev"] -markers = "python_version >= \"3.10\"" files = [ {file = "a2a_sdk-0.3.22-py3-none-any.whl", hash = "sha256:b98701135bb90b0ff85d35f31533b6b7a299bf810658c1c65f3814a6c15ea385"}, {file = "a2a_sdk-0.3.22.tar.gz", hash = "sha256:77a5694bfc4f26679c11b70c7f1062522206d430b34bc1215cfbb1eba67b7e7d"}, ] +markers = {main = "python_version >= \"3.10\" and extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [package.dependencies] google-api-core = ">=1.26.0" @@ -385,6 +385,7 @@ files = [ {file = "azure_core-1.36.0-py3-none-any.whl", hash = "sha256:fee9923a3a753e94a259563429f3644aaf05c486d45b1215d098115102d91d3b"}, {file = "azure_core-1.36.0.tar.gz", hash = "sha256:22e5605e6d0bf1d229726af56d9e92bc37b6e726b141a18be0b4d424131741b7"}, ] +markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""} [package.dependencies] requests = ">=2.21.0" @@ -405,6 +406,7 @@ files = [ {file = "azure_identity-1.25.1-py3-none-any.whl", hash = "sha256:e9edd720af03dff020223cd269fa3a61e8f345ea75443858273bcb44844ab651"}, {file = "azure_identity-1.25.1.tar.gz", hash = "sha256:87ca8328883de6036443e1c37b40e8dc8fb74898240f61071e09d2e369361456"}, ] +markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""} [package.dependencies] azure-core = ">=1.31.0" @@ -598,7 +600,7 @@ files = [ {file = "cachetools-6.2.2-py3-none-any.whl", hash = "sha256:6c09c98183bf58560c97b2abfcedcbaf6a896a490f534b031b661d3723b45ace"}, {file = "cachetools-6.2.2.tar.gz", hash = "sha256:8e6d266b25e539df852251cfd6f990b4bc3a141db73b939058d809ebd2590fc6"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [[package]] name = "certifi" @@ -705,7 +707,7 @@ files = [ {file = "cffi-2.0.0-cp39-cp39-win_amd64.whl", hash = "sha256:b882b3df248017dba09d6b16defe9b5c407fe32fc7c65a9c69798e6175601be9"}, {file = "cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529"}, ] -markers = {main = "platform_python_implementation != \"PyPy\" or extra == \"proxy\"", dev = "platform_python_implementation != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\""} +markers = {main = "(platform_python_implementation != \"PyPy\" or extra == \"proxy\") and (python_version >= \"3.10\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\")", dev = "platform_python_implementation != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\""} [package.dependencies] pycparser = {version = "*", markers = "implementation_name != \"PyPy\""} @@ -1055,6 +1057,7 @@ files = [ {file = "cryptography-43.0.3-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:2ce6fae5bdad59577b44e4dfed356944fbf1d925269114c28be377692643b4ff"}, {file = "cryptography-43.0.3.tar.gz", hash = "sha256:315b9001266a492a6ff443b61238f956b214dbec9910a081ba5b6646a055a805"}, ] +markers = {main = "python_version >= \"3.10\" and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\") or extra == \"proxy\" or extra == \"extra-proxy\""} [package.dependencies] cffi = {version = ">=1.12", markers = "platform_python_implementation != \"PyPy\""} @@ -1837,11 +1840,11 @@ description = "Google API client core library" optional = false python-versions = ">=3.7" groups = ["main", "proxy-dev"] -markers = "python_version >= \"3.14\"" files = [ {file = "google_api_core-2.25.2-py3-none-any.whl", hash = "sha256:e9a8f62d363dc8424a8497f4c2a47d6bcda6c16514c935629c257ab5d10210e7"}, {file = "google_api_core-2.25.2.tar.gz", hash = "sha256:1c63aa6af0d0d5e37966f157a77f9396d820fba59f9e43e9415bc3dc5baff300"}, ] +markers = {main = "python_version >= \"3.14\" and (extra == \"extra-proxy\" or extra == \"google\")", proxy-dev = "python_version >= \"3.14\""} [package.dependencies] google-auth = ">=2.14.1,<3.0.0" @@ -1869,7 +1872,7 @@ files = [ {file = "google_api_core-2.28.1-py3-none-any.whl", hash = "sha256:4021b0f8ceb77a6fb4de6fde4502cecab45062e66ff4f2895169e0b35bc9466c"}, {file = "google_api_core-2.28.1.tar.gz", hash = "sha256:2b405df02d68e68ce0fbc138559e6036559e685159d148ae5861013dc201baf8"}, ] -markers = {main = "(python_version >= \"3.10\" or extra == \"google\" or extra == \"extra-proxy\") and python_version < \"3.14\"", proxy-dev = "python_version >= \"3.10\" and python_version < \"3.14\""} +markers = {main = "python_version < \"3.14\" and (extra == \"extra-proxy\" or extra == \"google\")", proxy-dev = "python_version >= \"3.10\" and python_version < \"3.14\""} [package.dependencies] google-auth = ">=2.14.1,<3.0.0" @@ -1906,7 +1909,7 @@ files = [ {file = "google_auth-2.43.0-py2.py3-none-any.whl", hash = "sha256:af628ba6fa493f75c7e9dbe9373d148ca9f4399b5ea29976519e0a3848eddd16"}, {file = "google_auth-2.43.0.tar.gz", hash = "sha256:88228eee5fc21b62a1b5fe773ca15e67778cb07dc8363adcb4a8827b52d81483"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [package.dependencies] cachetools = ">=2.0.0,<7.0" @@ -2078,11 +2081,11 @@ files = [ ] [package.dependencies] -google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0dev", extras = ["grpc"]} -google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0dev" -grpc-google-iam-v1 = ">=0.12.4,<1.0.0dev" -proto-plus = ">=1.22.3,<2.0.0dev" -protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0dev" +google-api-core = {version = ">=1.34.1,<2.0.dev0 || >=2.11.dev0,<3.0.0.dev0", extras = ["grpc"]} +google-auth = ">=2.14.1,<2.24.0 || >2.24.0,<2.25.0 || >2.25.0,<3.0.0.dev0" +grpc-google-iam-v1 = ">=0.12.4,<1.0.0.dev0" +proto-plus = ">=1.22.3,<2.0.0.dev0" +protobuf = ">=3.20.2,<4.21.0 || >4.21.0,<4.21.1 || >4.21.1,<4.21.2 || >4.21.2,<4.21.3 || >4.21.3,<4.21.4 || >4.21.4,<4.21.5 || >4.21.5,<6.0.0.dev0" [[package]] name = "google-cloud-resource-manager" @@ -2264,7 +2267,7 @@ files = [ {file = "googleapis_common_protos-1.72.0-py3-none-any.whl", hash = "sha256:4299c5a82d5ae1a9702ada957347726b167f9f8d1fc352477702a1e851ff4038"}, {file = "googleapis_common_protos-1.72.0.tar.gz", hash = "sha256:e55a601c1b32b52d7a3e65f43563e2aa61bcd737998ee672ac9b951cd49319f5"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\") or extra == \"google\" or extra == \"extra-proxy\""} [package.dependencies] grpcio = {version = ">=1.44.0,<2.0.0", optional = true, markers = "extra == \"grpc\""} @@ -2673,11 +2676,11 @@ description = "Consume Server-Sent Event (SSE) messages with HTTPX." optional = false python-versions = ">=3.9" groups = ["main", "proxy-dev"] -markers = "python_version >= \"3.10\"" files = [ {file = "httpx_sse-0.4.3-py3-none-any.whl", hash = "sha256:0ac1c9fe3c0afad2e0ebb25a934a59f4c7823b60792691f779fad2c5568830fc"}, {file = "httpx_sse-0.4.3.tar.gz", hash = "sha256:9b1ed0127459a66014aec3c56bebd93da3c1bc8bb6618c8082039a44889a755d"}, ] +markers = {main = "python_version >= \"3.10\" and (extra == \"proxy\" or extra == \"extra-proxy\")", proxy-dev = "python_version >= \"3.10\""} [[package]] name = "huey" @@ -3042,7 +3045,7 @@ files = [ [package.dependencies] attrs = ">=22.2.0" -jsonschema-specifications = ">=2023.03.6" +jsonschema-specifications = ">=2023.3.6" referencing = ">=0.28.4" rpds-py = ">=0.7.1" @@ -3713,6 +3716,7 @@ files = [ {file = "msal-1.34.0-py3-none-any.whl", hash = "sha256:f669b1644e4950115da7a176441b0e13ec2975c29528d8b9e81316023676d6e1"}, {file = "msal-1.34.0.tar.gz", hash = "sha256:76ba83b716ea5a6d75b0279c0ac353a0e05b820ca1f6682c0eb7f45190c43c2f"}, ] +markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""} [package.dependencies] cryptography = ">=2.5,<49" @@ -3733,6 +3737,7 @@ files = [ {file = "msal_extensions-1.3.1-py3-none-any.whl", hash = "sha256:96d3de4d034504e969ac5e85bae8106c8373b5c6568e4c8fa7af2eca9dbe6bca"}, {file = "msal_extensions-1.3.1.tar.gz", hash = "sha256:c5b0fd10f65ef62b5f1d62f4251d51cbcaf003fcedae8c91b040a488614be1a4"}, ] +markers = {main = "extra == \"proxy\" or extra == \"extra-proxy\""} [package.dependencies] msal = ">=1.29,<2" @@ -3983,6 +3988,7 @@ files = [ {file = "nodeenv-1.9.1-py2.py3-none-any.whl", hash = "sha256:ba11c9782d29c27c70ffbdda2d7415098754709be8a7056d79a737cd901155c9"}, {file = "nodeenv-1.9.1.tar.gz", hash = "sha256:6ec12890a2dab7946721edbfbcd91f3319c6ccc9aec47be7c7e6b7011ee6645f"}, ] +markers = {main = "extra == \"extra-proxy\""} [[package]] name = "numpy" @@ -4105,7 +4111,7 @@ files = [ {file = "opentelemetry_api-1.39.1-py3-none-any.whl", hash = "sha256:2edd8463432a7f8443edce90972169b195e7d6a05500cd29e6d13898187c9950"}, {file = "opentelemetry_api-1.39.1.tar.gz", hash = "sha256:fbde8c80e1b937a2c61f20347e91c0c18a1940cecf012d62e65a7caf08967c9c"}, ] -markers = {main = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""} [package.dependencies] importlib-metadata = ">=6.0,<8.8.0" @@ -4220,7 +4226,7 @@ files = [ {file = "opentelemetry_sdk-1.39.1-py3-none-any.whl", hash = "sha256:4d5482c478513ecb0a5d938dcc61394e647066e0cc2676bee9f3af3f3f45f01c"}, {file = "opentelemetry_sdk-1.39.1.tar.gz", hash = "sha256:cf4d4563caf7bff906c9f7967e2be22d0d6b349b908be0d90fb21c8e9c995cc6"}, ] -markers = {main = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""} [package.dependencies] opentelemetry-api = "1.39.1" @@ -4238,7 +4244,7 @@ files = [ {file = "opentelemetry_semantic_conventions-0.60b1-py3-none-any.whl", hash = "sha256:9fa8c8b0c110da289809292b0591220d3a7b53c1526a23021e977d68597893fb"}, {file = "opentelemetry_semantic_conventions-0.60b1.tar.gz", hash = "sha256:87c228b5a0669b748c76d76df6c364c369c28f1c465e50f661e39737e84bc953"}, ] -markers = {main = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and extra == \"mlflow\""} [package.dependencies] opentelemetry-api = "1.39.1" @@ -4722,6 +4728,7 @@ files = [ {file = "prisma-0.11.0-py3-none-any.whl", hash = "sha256:22bb869e59a2968b99f3483bb417717273ffbc569fd1e9ceed95e5614cbaf53a"}, {file = "prisma-0.11.0.tar.gz", hash = "sha256:3f2f2fd2361e1ec5ff655f2a04c7860c2f2a5bc4c91f78ca9c5c6349735bf693"}, ] +markers = {main = "extra == \"extra-proxy\""} [package.dependencies] click = ">=7.1.2" @@ -4895,7 +4902,7 @@ files = [ {file = "proto_plus-1.26.1-py3-none-any.whl", hash = "sha256:13285478c2dcf2abb829db158e1047e2f1e8d63a077d94263c2b88b043c75a66"}, {file = "proto_plus-1.26.1.tar.gz", hash = "sha256:21a515a4c4c0088a773899e23c7bbade3d18f9c66c73edd4c7ee3816bc96a012"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""} +markers = {main = "extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [package.dependencies] protobuf = ">=3.19.0,<7.0.0" @@ -4923,7 +4930,7 @@ files = [ {file = "protobuf-5.29.5-py3-none-any.whl", hash = "sha256:6cf42630262c59b2d8de33954443d94b746c952b01434fc58a417fdbd2e84bd5"}, {file = "protobuf-5.29.5.tar.gz", hash = "sha256:bc1463bafd4b0929216c35f437a8e28731a2b7fe3d98bb77a600efced5a15c84"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\""} [[package]] name = "psutil" @@ -5083,7 +5090,7 @@ files = [ {file = "pyasn1-0.6.1-py3-none-any.whl", hash = "sha256:0d632f46f2ba09143da3a8afe9e33fb6f92fa2320ab7e886e2d0f7672af84629"}, {file = "pyasn1-0.6.1.tar.gz", hash = "sha256:6f580d2bdd84365380830acf45550f2511469f673cb4a5ae3857a3170128b034"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [[package]] name = "pyasn1-modules" @@ -5096,7 +5103,7 @@ files = [ {file = "pyasn1_modules-0.4.2-py3-none-any.whl", hash = "sha256:29253a9207ce32b64c3ac6600edc75368f98473906e8fd1043bd6b5b1de2c14a"}, {file = "pyasn1_modules-0.4.2.tar.gz", hash = "sha256:677091de870a80aae844b1ca6134f54652fa2c8c5a52aa396440ac3106e941e6"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [package.dependencies] pyasn1 = ">=0.6.1,<0.7.0" @@ -5124,7 +5131,7 @@ files = [ {file = "pycparser-2.23-py3-none-any.whl", hash = "sha256:e5c6e8d3fbad53479cab09ac03729e0a9faf2bee3db8208a550daf5af81a5934"}, {file = "pycparser-2.23.tar.gz", hash = "sha256:78816d4f24add8f10a06d6f05b4d424ad9e96cfebf68a4ddc99c65c0720d00c2"}, ] -markers = {main = "implementation_name != \"PyPy\" and (platform_python_implementation != \"PyPy\" or extra == \"proxy\")", dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\""} +markers = {main = "implementation_name != \"PyPy\" and (platform_python_implementation != \"PyPy\" or extra == \"proxy\") and (python_version >= \"3.10\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"proxy\" or extra == \"extra-proxy\" or extra == \"mlflow\")", dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\"", proxy-dev = "platform_python_implementation != \"PyPy\" and implementation_name != \"PyPy\""} [[package]] name = "pydantic" @@ -5347,6 +5354,7 @@ files = [ {file = "PyJWT-2.10.1-py3-none-any.whl", hash = "sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb"}, {file = "pyjwt-2.10.1.tar.gz", hash = "sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953"}, ] +markers = {main = "(python_version <= \"3.13\" or extra == \"proxy\" or extra == \"extra-proxy\") and (extra == \"extra-proxy\" or extra == \"proxy\")"} [package.dependencies] cryptography = {version = ">=3.4.0", optional = true, markers = "extra == \"crypto\""} @@ -6289,7 +6297,7 @@ files = [ {file = "rsa-4.9.1-py3-none-any.whl", hash = "sha256:68635866661c6836b8d39430f97a996acbd61bfa49406748ea243539fe239762"}, {file = "rsa-4.9.1.tar.gz", hash = "sha256:e7bdbfdb5497da4c07dfd35530e1a902659db6ff241e39d9953cad06ebd0ae75"}, ] -markers = {main = "extra == \"google\" or extra == \"extra-proxy\" or python_version >= \"3.10\"", proxy-dev = "python_version >= \"3.10\""} +markers = {main = "python_version >= \"3.10\" and (extra == \"extra-proxy\" or extra == \"google\" or extra == \"mlflow\") or extra == \"google\" or extra == \"extra-proxy\"", proxy-dev = "python_version >= \"3.10\""} [package.dependencies] pyasn1 = ">=0.1.3" @@ -6335,10 +6343,10 @@ files = [ ] [package.dependencies] -botocore = ">=1.37.4,<2.0a.0" +botocore = ">=1.37.4,<2.0a0" [package.extras] -crt = ["botocore[crt] (>=1.37.4,<2.0a.0)"] +crt = ["botocore[crt] (>=1.37.4,<2.0a0)"] [[package]] name = "scikit-learn" @@ -6491,9 +6499,9 @@ tornado = ">=6.4.2,<7" urllib3 = ">=1.26,<3" [package.extras] -all = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)", "cohere (>=5.9.4,<6.00)", "dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\"", "google-cloud-aiplatform (>=1.45.0,<2)", "ipykernel (>=6.25.0,<7)", "llama-cpp-python (>=0.2.28,<0.2.86) ; python_version < \"3.13\"", "mistralai (>=0.0.12,<0.1.0)", "mypy (>=1.7.1,<2)", "ollama (>=0.1.7)", "pillow (>=10.2.0,<11.0.0) ; python_version < \"3.13\"", "pinecone[asyncio] (>=7.0.0,<8.0.0)", "psycopg[binary] (>=3.1.0,<4)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "qdrant-client (>=1.11.1,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "sentence-transformers (>=5.0.0) ; python_version < \"3.13\"", "tokenizers (>=0.19) ; python_version < \"3.13\"", "torch (>=2.6.0) ; python_version < \"3.13\"", "torchvision (>=0.17.0) ; python_version < \"3.13\"", "transformers (>=4.36.2) ; python_version < \"3.13\"", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"] +all = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)", "cohere (>=5.9.4,<6.0)", "dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\"", "google-cloud-aiplatform (>=1.45.0,<2)", "ipykernel (>=6.25.0,<7)", "llama-cpp-python (>=0.2.28,<0.2.86) ; python_version < \"3.13\"", "mistralai (>=0.0.12,<0.1.0)", "mypy (>=1.7.1,<2)", "ollama (>=0.1.7)", "pillow (>=10.2.0,<11.0.0) ; python_version < \"3.13\"", "pinecone[asyncio] (>=7.0.0,<8.0.0)", "psycopg[binary] (>=3.1.0,<4)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "qdrant-client (>=1.11.1,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "sentence-transformers (>=5.0.0) ; python_version < \"3.13\"", "tokenizers (>=0.19) ; python_version < \"3.13\"", "torch (>=2.6.0) ; python_version < \"3.13\"", "torchvision (>=0.17.0) ; python_version < \"3.13\"", "transformers (>=4.36.2) ; python_version < \"3.13\"", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"] bedrock = ["boto3 (>=1.34.98,<2)", "botocore (>=1.34.110,<2)"] -cohere = ["cohere (>=5.9.4,<6.00)"] +cohere = ["cohere (>=5.9.4,<6.0)"] dev = ["dagger-io (>=0.1.1) ; python_version >= \"3.11\"", "ipykernel (>=6.25.0,<7)", "mypy (>=1.7.1,<2)", "pytest (>=8.2,<9.0)", "pytest-asyncio (>=0.24.0,<0.25)", "pytest-cov (>=4.1.0,<5)", "pytest-mock (>=3.12.0,<4)", "pytest-timeout", "pytest-xdist (>=3.5.0,<4)", "python-dotenv (>=1.0.0,<2)", "requests-mock (>=1.12.1,<2)", "ruff (>=0.11.2,<0.12)", "types-pyyaml (>=6.0.12.12,<7)", "types-requests (>=2.31.0,<3)"] docs = ["pydoc-markdown (>=4.8.2) ; python_version < \"3.12\""] fastembed = ["fastembed (>=0.3.0,<0.4) ; python_version < \"3.13\""] @@ -7221,6 +7229,7 @@ files = [ {file = "tomlkit-0.13.3-py3-none-any.whl", hash = "sha256:c89c649d79ee40629a9fda55f8ace8c6a1b42deb912b2a8fd8d942ddadb606b0"}, {file = "tomlkit-0.13.3.tar.gz", hash = "sha256:430cf247ee57df2b94ee3fbe588e71d362a941ebb545dec29b53961d61add2a1"}, ] +markers = {main = "extra == \"extra-proxy\""} [[package]] name = "tornado" @@ -7993,4 +8002,4 @@ utils = ["numpydoc"] [metadata] lock-version = "2.1" python-versions = ">=3.9,<4.0" -content-hash = "1066f8f705d83a46c40f64b0fa4e1ffd0a88817c2f1f4761bbf565bed91588f9" +content-hash = "fa110a048c30d0ad4e66414290ec103dba7707d99474827ea0cf3e4a2058d165" diff --git a/pyproject.toml b/pyproject.toml index 0c42c18c70d..346e9114640 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,7 +47,7 @@ fastapi-sso = { version = "^0.16.0", optional = true } PyJWT = { version = "^2.10.1", optional = true, python = ">=3.9" } python-multipart = { version = ">=0.0.20", optional = true} cryptography = {version = "*", optional = true} -prisma = {version = "0.11.0", optional = true} +prisma = {version = "^0.11.0", optional = true} azure-identity = {version = "^1.15.0", optional = true, python = ">=3.9"} azure-keyvault-secrets = {version = "^4.8.0", optional = true} azure-storage-blob = {version="^12.25.1", optional=true} @@ -57,13 +57,13 @@ google-cloud-aiplatform = {version = ">=1.38.0", optional = true} resend = {version = ">=0.8.0", optional = true} pynacl = {version = "^1.5.0", optional = true} websockets = {version = "^15.0.1", optional = true} -boto3 = { version = "1.40.76", optional = true } +boto3 = { version = "^1.40.76", optional = true } redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"} mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"} a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"} -litellm-proxy-extras = {version = "0.4.52", optional = true} -rich = {version = "13.7.1", optional = true} -litellm-enterprise = {version = "0.1.33", optional = true} +litellm-proxy-extras = {version = "^0.4.52", optional = true} +rich = {version = "^13.7.1", optional = true} +litellm-enterprise = {version = "^0.1.33", optional = true} diskcache = {version = "^5.6.1", optional = true} polars = {version = "^1.31.0", optional = true, python = ">=3.10"} semantic-router = {version = ">=0.1.12", optional = true, python = ">=3.9,<3.14"} diff --git a/requirements.txt b/requirements.txt index a6446dc0cf1..103b2981458 100644 --- a/requirements.txt +++ b/requirements.txt @@ -75,7 +75,7 @@ jsonschema>=4.23.0,<5.0.0 # validating json schema - aligned with openapi-core + websockets==15.0.1 # for realtime API soundfile==0.12.1 # for audio file processing openapi-core==0.21.0 # for OpenAPI compliance tests -pypdf>=6.6.2 # for PDF text extraction in RAG ingestion +pypdf>=6.7.3 # for PDF text extraction in RAG ingestion (CVE-2026-27888) ######################## # LITELLM ENTERPRISE DEPENDENCIES diff --git a/schema.prisma b/schema.prisma index 329ff80933f..8d4bdffb2dd 100644 --- a/schema.prisma +++ b/schema.prisma @@ -68,6 +68,11 @@ model LiteLLM_AgentsTable { agent_access_groups String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + spend Float @default(0.0) + tpm_limit Int? + rpm_limit Int? + session_tpm_limit Int? + session_rpm_limit Int? created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") diff --git a/tests/proxy_admin_ui_tests/package.json b/tests/proxy_admin_ui_tests/package.json index 037ec7082a7..ac726b64b76 100644 --- a/tests/proxy_admin_ui_tests/package.json +++ b/tests/proxy_admin_ui_tests/package.json @@ -13,8 +13,8 @@ }, "overrides": { "glob": ">=11.1.0", - "tar": ">=7.5.8", - "minimatch": ">=10.2.1", + "tar": ">=7.5.10", + "minimatch": ">=10.2.4", "diff": ">=8.0.3", "@isaacs/brace-expansion": ">=5.0.1", "@babel/traverse": ">=7.23.2", diff --git a/tests/proxy_admin_ui_tests/ui_unit_tests/package.json b/tests/proxy_admin_ui_tests/ui_unit_tests/package.json index eb9c7473a5b..9f1c689721c 100644 --- a/tests/proxy_admin_ui_tests/ui_unit_tests/package.json +++ b/tests/proxy_admin_ui_tests/ui_unit_tests/package.json @@ -25,8 +25,8 @@ }, "overrides": { "glob": ">=11.1.0", - "tar": ">=7.5.8", - "minimatch": ">=10.2.1", + "tar": ">=7.5.10", + "minimatch": ">=10.2.4", "diff": ">=8.0.3", "@isaacs/brace-expansion": ">=5.0.1", "@babel/traverse": ">=7.23.2", diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py index d2fb45643de..cb05531c2f8 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_bedrock_chat_invoke_transformations_anthropic_claude3_transformation.py @@ -386,6 +386,40 @@ def test_opus_4_5_model_detection(): # f"computer-use beta should be kept, got: {anthropic_beta}" +def test_output_config_removed_from_bedrock_chat_invoke_request(): + """ + Test that output_config parameter is stripped from Bedrock Chat Invoke requests. + + Bedrock Invoke API doesn't support the output_config parameter (Anthropic-only). + Ensures the chat/invoke path mirrors the messages/invoke path fix. + + Fixes: https://github.com/BerriAI/litellm/issues/22797 + """ + config = AmazonAnthropicClaudeConfig() + + messages = [{"role": "user", "content": "test"}] + + # Inject output_config into optional_params (simulates Anthropic SDK forwarding it) + optional_params = { + "max_tokens": 100, + "output_config": {"effort": "high"}, + } + + result = config.transform_request( + model="anthropic.claude-sonnet-4-20250514-v1:0", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + assert "output_config" not in result, ( + f"output_config should be stripped for Bedrock Chat Invoke, got keys: {list(result.keys())}" + ) + # Verify normal params survive + assert result["max_tokens"] == 100 + + def test_output_format_removed_from_bedrock_invoke_request(): """ Test that output_format parameter is removed from Bedrock Invoke requests. diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index e0b06ced172..f69f478278f 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -275,3 +275,70 @@ def test_remove_scope_from_cache_control(): # Verify scope is removed from messages assert "scope" not in request["messages"][0]["content"][0]["cache_control"] assert request["messages"][0]["content"][0]["cache_control"]["type"] == "ephemeral" + + +def test_bedrock_messages_strips_output_config(): + """ + Ensure output_config is stripped from the request before sending to + Bedrock Invoke, which doesn't support this Anthropic-specific parameter. + + Regression test for: https://github.com/BerriAI/litellm/issues/22797 + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "output_config": { + "effort": "high", + }, + } + + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "output_config" not in result, ( + "output_config should be stripped — Bedrock Invoke rejects it" + ) + # Other params should be preserved + assert result.get("max_tokens") == 4096 + + +def test_bedrock_messages_strips_output_config_with_output_format(): + """ + When both output_config and output_format are present, both should be + stripped (output_format is converted to inline schema, output_config + is simply dropped). + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}] + optional_params = { + "max_tokens": 4096, + "output_config": {"effort": "low"}, + "output_format": { + "type": "json_schema", + "schema": { + "type": "object", + "properties": {"answer": {"type": "string"}}, + }, + }, + } + + result = cfg.transform_anthropic_messages_request( + model="anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "output_config" not in result + assert "output_format" not in result diff --git a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py index fcf8f048190..00d08504fe5 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_endpoints.py @@ -439,3 +439,200 @@ class TestAgentRoutesIncludesAgentIdPattern: from litellm.proxy._types import LiteLLMRoutes assert "/v1/agents/{agent_id}" in LiteLLMRoutes.agent_routes.value + + +class TestAgentHealthCheck: + """Tests for the health_check query parameter on GET /v1/agents.""" + + @pytest.fixture(autouse=True) + def _setup(self, monkeypatch): + from litellm.proxy.agent_endpoints import agent_registry as ar_mod + + self.admin_client = _make_app_with_role(LitellmUserRoles.PROXY_ADMIN) + self.mock_registry = MagicMock() + monkeypatch.setattr(ar_mod, "global_agent_registry", self.mock_registry) + + def _make_agent(self, agent_id: str, url: str | None = None) -> AgentResponse: + card = _sample_agent_card_params() + if url is not None: + card["url"] = url + else: + card.pop("url", None) + return AgentResponse( + agent_id=agent_id, + agent_name=f"Agent {agent_id}", + agent_card_params=card, + litellm_params={}, + ) + + def test_should_return_all_agents_when_health_check_disabled(self): + agents = [self._make_agent("a1", "http://reachable"), self._make_agent("a2", "http://unreachable")] + self.mock_registry.get_agent_list = MagicMock(return_value=agents) + + resp = self.admin_client.get( + "/v1/agents", headers={"Authorization": "Bearer k"} + ) + assert resp.status_code == 200 + assert len(resp.json()) == 2 + + def test_should_filter_unhealthy_agents_when_health_check_enabled(self, monkeypatch): + agents = [ + self._make_agent("a1", "http://reachable"), + self._make_agent("a2", "http://unreachable"), + ] + self.mock_registry.get_agent_list = MagicMock(return_value=agents) + + results = iter([ + {"agent_id": "a1", "healthy": True}, + {"agent_id": "a2", "healthy": False, "error": "Connection refused"}, + ]) + monkeypatch.setattr( + agent_endpoints, + "_check_agent_url_health", + AsyncMock(side_effect=lambda agent: next(results)), + ) + + resp = self.admin_client.get( + "/v1/agents?health_check=true", + headers={"Authorization": "Bearer k"}, + ) + assert resp.status_code == 200 + data = resp.json() + assert len(data) == 1 + assert data[0]["agent_id"] == "a1" + + def test_should_return_empty_list_when_all_agents_unhealthy(self, monkeypatch): + agents = [self._make_agent("a1", "http://down")] + self.mock_registry.get_agent_list = MagicMock(return_value=agents) + monkeypatch.setattr( + agent_endpoints, + "_check_agent_url_health", + AsyncMock(return_value={"agent_id": "a1", "healthy": False, "error": "timeout"}), + ) + + resp = self.admin_client.get( + "/v1/agents?health_check=true", + headers={"Authorization": "Bearer k"}, + ) + assert resp.status_code == 200 + assert len(resp.json()) == 0 + + def test_should_return_all_agents_when_all_healthy(self, monkeypatch): + agents = [self._make_agent("a1", "http://ok1"), self._make_agent("a2", "http://ok2")] + self.mock_registry.get_agent_list = MagicMock(return_value=agents) + + results = iter([ + {"agent_id": "a1", "healthy": True}, + {"agent_id": "a2", "healthy": True}, + ]) + monkeypatch.setattr( + agent_endpoints, + "_check_agent_url_health", + AsyncMock(side_effect=lambda agent: next(results)), + ) + + resp = self.admin_client.get( + "/v1/agents?health_check=true", + headers={"Authorization": "Bearer k"}, + ) + assert resp.status_code == 200 + assert len(resp.json()) == 2 + + +class TestCheckAgentUrlHealth: + """Unit tests for the _check_agent_url_health helper.""" + + @pytest.mark.asyncio + async def test_should_return_healthy_when_no_url(self): + from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health + + agent = AgentResponse( + agent_id="no-url", + agent_name="No URL Agent", + agent_card_params={"name": "test"}, + litellm_params={}, + ) + result = await _check_agent_url_health(agent) + assert result["healthy"] is True + assert "error" not in result + + @pytest.mark.asyncio + @patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client") + async def test_should_return_healthy_for_200(self, mock_get_client): + from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health + + mock_response = MagicMock() + mock_response.status_code = 200 + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + agent = AgentResponse( + agent_id="ok", + agent_name="OK Agent", + agent_card_params={"url": "http://example.com"}, + litellm_params={}, + ) + result = await _check_agent_url_health(agent) + assert result["healthy"] is True + + @pytest.mark.asyncio + @patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client") + async def test_should_return_unhealthy_for_500(self, mock_get_client): + from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health + + mock_response = MagicMock() + mock_response.status_code = 500 + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + agent = AgentResponse( + agent_id="err", + agent_name="Error Agent", + agent_card_params={"url": "http://failing.com"}, + litellm_params={}, + ) + result = await _check_agent_url_health(agent) + assert result["healthy"] is False + assert "HTTP 500" in result["error"] + + @pytest.mark.asyncio + @patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client") + async def test_should_return_unhealthy_on_connection_error(self, mock_get_client): + from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health + + mock_client = AsyncMock() + mock_client.get = AsyncMock(side_effect=Exception("Connection refused")) + mock_get_client.return_value = mock_client + + agent = AgentResponse( + agent_id="down", + agent_name="Down Agent", + agent_card_params={"url": "http://down.com"}, + litellm_params={}, + ) + result = await _check_agent_url_health(agent) + assert result["healthy"] is False + assert "Connection refused" in result["error"] + + @pytest.mark.asyncio + @patch("litellm.proxy.agent_endpoints.endpoints.get_async_httpx_client") + async def test_should_treat_404_as_healthy(self, mock_get_client): + """A 404 means the server is reachable, just not the specific path.""" + from litellm.proxy.agent_endpoints.endpoints import _check_agent_url_health + + mock_response = MagicMock() + mock_response.status_code = 404 + mock_client = AsyncMock() + mock_client.get = AsyncMock(return_value=mock_response) + mock_get_client.return_value = mock_client + + agent = AgentResponse( + agent_id="notfound", + agent_name="NotFound Agent", + agent_card_params={"url": "http://example.com/missing"}, + litellm_params={}, + ) + result = await _check_agent_url_health(agent) + assert result["healthy"] is True diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index abd59a66a36..b5f82ef04c5 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -9,7 +9,7 @@ sys.path.insert( from datetime import datetime, timezone -from unittest.mock import AsyncMock, MagicMock, patch, call +from unittest.mock import AsyncMock, MagicMock, call, patch import pytest @@ -516,6 +516,131 @@ async def test_update_tag_db_without_prisma_client(): assert writer.spend_update_queue.add_update.call_count == 0 +@pytest.mark.asyncio +async def test_update_agent_db_enqueues_agent_spend(): + """ + Test that _update_agent_db enqueues a SpendUpdateQueueItem with entity_type=AGENT. + """ + from litellm.proxy._types import Litellm_EntityType + + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + agent_id = "agent-123" + response_cost = 0.1 + + writer.spend_update_queue.add_update = AsyncMock() + + await writer._update_agent_db( + response_cost=response_cost, + agent_id=agent_id, + prisma_client=mock_prisma, + ) + + writer.spend_update_queue.add_update.assert_called_once() + call_args = writer.spend_update_queue.add_update.call_args[1] + assert call_args["update"]["entity_type"] == Litellm_EntityType.AGENT + assert call_args["update"]["entity_id"] == agent_id + assert call_args["update"]["response_cost"] == response_cost + + +@pytest.mark.asyncio +async def test_update_agent_db_skips_when_agent_id_none(): + """_update_agent_db does not enqueue when agent_id is None.""" + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + writer.spend_update_queue.add_update = AsyncMock() + + await writer._update_agent_db( + response_cost=0.05, + agent_id=None, + prisma_client=mock_prisma, + ) + + writer.spend_update_queue.add_update.assert_not_called() + + +@pytest.mark.asyncio +async def test_update_agent_db_skips_when_prisma_client_none(): + """_update_agent_db does not enqueue when prisma_client is None.""" + writer = DBSpendUpdateWriter() + writer.spend_update_queue.add_update = AsyncMock() + + await writer._update_agent_db( + response_cost=0.05, + agent_id="agent-456", + prisma_client=None, + ) + + writer.spend_update_queue.add_update.assert_not_called() + + +@pytest.mark.asyncio +async def test_commit_spend_updates_to_db_increments_agent_spend(): + """ + Test that _commit_spend_updates_to_db calls litellm_agentstable.update_many + with spend increment when agent_list_transactions is present. + """ + db_writer = DBSpendUpdateWriter() + + mock_batcher = MagicMock() + mock_batcher.litellm_verificationtoken = MagicMock() + mock_batcher.litellm_verificationtoken.update_many = MagicMock() + mock_batcher.litellm_usertable = MagicMock() + mock_batcher.litellm_usertable.update_many = MagicMock() + mock_batcher.litellm_teamtable = MagicMock() + mock_batcher.litellm_teamtable.update_many = MagicMock() + mock_batcher.litellm_teammembership = MagicMock() + mock_batcher.litellm_teammembership.update_many = MagicMock() + mock_batcher.litellm_organizationtable = MagicMock() + mock_batcher.litellm_organizationtable.update_many = MagicMock() + mock_batcher.litellm_tagtable = MagicMock() + mock_batcher.litellm_tagtable.update_many = MagicMock() + mock_batcher.litellm_agentstable = MagicMock() + mock_batcher.litellm_agentstable.update_many = MagicMock() + + mock_transaction = AsyncMock() + mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction) + mock_transaction.__aexit__ = AsyncMock(return_value=False) + mock_transaction.batch_ = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + ) + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction) + + mock_proxy_logging = MagicMock() + + agent_id = "agent-789" + response_cost = 0.25 + db_spend_update_transactions = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": {}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {agent_id: response_cost}, + } + + with patch("litellm.proxy.utils._raise_failed_update_spend_exception"): + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=mock_proxy_logging, + db_spend_update_transactions=db_spend_update_transactions, + ) + + mock_batcher.litellm_agentstable.update_many.assert_called_once() + call_kwargs = mock_batcher.litellm_agentstable.update_many.call_args[1] + assert call_kwargs["where"] == {"agent_id": agent_id} + assert call_kwargs["data"] == {"spend": {"increment": response_cost}} + + @pytest.mark.asyncio async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id(): """ @@ -1048,6 +1173,8 @@ async def test_commit_key_spend_updates_includes_last_active(): mock_batcher.litellm_teamtable.update_many = MagicMock() mock_batcher.litellm_organizationtable = MagicMock() mock_batcher.litellm_organizationtable.update_many = MagicMock() + mock_batcher.litellm_agentstable = MagicMock() + mock_batcher.litellm_agentstable.update_many = MagicMock() mock_proxy_logging = MagicMock() @@ -1059,6 +1186,7 @@ async def test_commit_key_spend_updates_includes_last_active(): "team_member_list_transactions": {}, "org_list_transactions": {}, "tag_list_transactions": {}, + "agent_list_transactions": {}, } before_call = datetime.now(timezone.utc) @@ -1142,6 +1270,7 @@ async def test_batch_database_updates_isolation_on_failure(): db_writer._update_team_db = AsyncMock() db_writer._update_org_db = AsyncMock() db_writer._update_tag_db = AsyncMock() + db_writer._update_agent_db = AsyncMock() db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock() @@ -1169,6 +1298,7 @@ async def test_batch_database_updates_isolation_on_failure(): db_writer._update_team_db.assert_awaited_once() db_writer._update_org_db.assert_awaited_once() db_writer._update_tag_db.assert_awaited_once() + db_writer._update_agent_db.assert_awaited_once() db_writer.add_spend_log_transaction_to_daily_user_transaction.assert_awaited_once() db_writer.add_spend_log_transaction_to_daily_end_user_transaction.assert_awaited_once() db_writer.add_spend_log_transaction_to_daily_agent_transaction.assert_awaited_once() @@ -1203,6 +1333,7 @@ async def test_daily_agent_receives_deepcopied_payload(): db_writer._update_team_db = AsyncMock() db_writer._update_org_db = AsyncMock() db_writer._update_tag_db = AsyncMock() + db_writer._update_agent_db = AsyncMock() db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock() db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock( diff --git a/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py b/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py new file mode 100644 index 00000000000..879e2d65c7a --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_max_budget_per_session_limiter.py @@ -0,0 +1,165 @@ +""" +Unit Tests for the per-session budget limiter for the proxy. + +Tests that session-scoped budget tracking works correctly: +- Enforces max_budget_per_session per session_id (read from agent litellm_params) +- Different sessions have independent budgets +- Requests under budget pass through +- Requests without agent_id pass through +""" + +from unittest.mock import patch + +import pytest +from fastapi import HTTPException + +from litellm.caching.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.max_budget_per_session_limiter import ( + _PROXY_MaxBudgetPerSessionHandler, +) +from litellm.proxy.utils import InternalUsageCache +from litellm.types.agents import AgentResponse + + +def _make_mock_agent(max_budget_per_session: float) -> AgentResponse: + return AgentResponse( + agent_id="agent-budget-123", + agent_name="budget-agent", + litellm_params={"max_budget_per_session": max_budget_per_session}, + agent_card_params={"name": "budget-agent", "version": "1.0.0"}, + ) + + +@pytest.mark.asyncio +async def test_budget_per_session_under_budget_passes(): + """ + Requests under budget should pass through without error. + """ + local_cache = DualCache() + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-budget", + agent_id="agent-budget-123", + ) + + mock_agent = _make_mock_agent(max_budget_per_session=5.0) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = mock_agent + + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-budget-1"}}, + call_type="", + ) + assert result is None + + +@pytest.mark.asyncio +async def test_budget_per_session_exceeds_budget(): + """ + After accumulating spend beyond max_budget_per_session, the next + pre-call check should raise 429. + """ + local_cache = DualCache() + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-budget", + agent_id="agent-budget-123", + ) + + session_id = "session-over-budget" + cache_key = handler._make_cache_key(session_id) + await handler._increment_spend(cache_key, 1.50) + + mock_agent = _make_mock_agent(max_budget_per_session=1.0) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = mock_agent + + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": session_id}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + assert "budget exceeded" in str(exc_info.value.detail).lower() + + +@pytest.mark.asyncio +async def test_budget_per_session_independent_sessions(): + """ + Different session_ids have independent budget counters. + Exhausting session A does not affect session B. + """ + local_cache = DualCache() + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-budget", + agent_id="agent-budget-123", + ) + + cache_key_a = handler._make_cache_key("session-A") + await handler._increment_spend(cache_key_a, 3.0) + + mock_agent = _make_mock_agent(max_budget_per_session=2.0) + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = mock_agent + + # Session A should be blocked + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-A"}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + + # Session B should still pass + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-B"}}, + call_type="", + ) + assert result is None + + +@pytest.mark.asyncio +async def test_no_agent_id_passes(): + """ + When no agent_id is set on the key, all requests pass through. + """ + local_cache = DualCache() + handler = _PROXY_MaxBudgetPerSessionHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-no-agent", + ) + + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "any-session"}}, + call_type="", + ) + assert result is None diff --git a/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py b/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py index deb1c483b87..20928ef46d5 100644 --- a/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py +++ b/tests/test_litellm/proxy/hooks/test_max_iterations_limiter.py @@ -2,10 +2,12 @@ Unit Tests for the max iterations limiter for the proxy. Tests that session-scoped iteration counting works correctly: -- Enforces max_iterations per session_id +- Enforces max_iterations per session_id (read from agent litellm_params) - Different sessions have independent counters """ +from unittest.mock import MagicMock, patch + import pytest from fastapi import HTTPException @@ -13,6 +15,16 @@ from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.hooks.max_iterations_limiter import _PROXY_MaxIterationsHandler from litellm.proxy.utils import InternalUsageCache +from litellm.types.agents import AgentResponse + + +def _make_mock_agent(max_iterations: int) -> AgentResponse: + return AgentResponse( + agent_id="agent-test-123", + agent_name="test-agent", + litellm_params={"max_iterations": max_iterations}, + agent_card_params={"name": "test-agent", "version": "1.0.0"}, + ) @pytest.mark.asyncio @@ -28,28 +40,36 @@ async def test_max_iterations_basic_enforcement(): internal_usage_cache=InternalUsageCache(local_cache), ) user_api_key_dict = UserAPIKeyAuth( - api_key="sk-test-key-1234", metadata={"max_iterations": 3} + api_key="sk-test-key-1234", + agent_id="agent-test-123", ) - # First 3 requests should succeed - for i in range(3): - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-abc"}}, - call_type="", - ) + mock_agent = _make_mock_agent(max_iterations=3) - # 4th request should fail with 429 - with pytest.raises(HTTPException) as exc_info: - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-abc"}}, - call_type="", - ) - assert exc_info.value.status_code == 429 - assert "max_iterations" in str(exc_info.value.detail).lower() + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = mock_agent + + # First 3 requests should succeed + for i in range(3): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-abc"}}, + call_type="", + ) + + # 4th request should fail with 429 + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-abc"}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + assert "max_iterations" in str(exc_info.value.detail).lower() @pytest.mark.asyncio @@ -65,42 +85,72 @@ async def test_max_iterations_different_sessions_independent(): internal_usage_cache=InternalUsageCache(local_cache), ) user_api_key_dict = UserAPIKeyAuth( - api_key="sk-test-key-5678", metadata={"max_iterations": 2} + api_key="sk-test-key-5678", + agent_id="agent-test-123", ) - # Session A: 2 calls succeed - for _ in range(2): - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-A"}}, - call_type="", - ) + mock_agent = _make_mock_agent(max_iterations=2) - # Session B: 2 calls succeed (independent counter) - for _ in range(2): - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-B"}}, - call_type="", - ) + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry" + ) as mock_registry: + mock_registry.get_agent_by_id.return_value = mock_agent - # Session A: 3rd call fails - with pytest.raises(HTTPException) as exc_info: - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-A"}}, - call_type="", - ) - assert exc_info.value.status_code == 429 + # Session A: 2 calls succeed + for _ in range(2): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-A"}}, + call_type="", + ) - # Session B: 3rd call also fails - with pytest.raises(HTTPException): - await handler.async_pre_call_hook( - user_api_key_dict=user_api_key_dict, - cache=local_cache, - data={"metadata": {"session_id": "session-B"}}, - call_type="", - ) + # Session B: 2 calls succeed (independent counter) + for _ in range(2): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-B"}}, + call_type="", + ) + + # Session A: 3rd call fails + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-A"}}, + call_type="", + ) + assert exc_info.value.status_code == 429 + + # Session B: 3rd call also fails + with pytest.raises(HTTPException): + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-B"}}, + call_type="", + ) + + +@pytest.mark.asyncio +async def test_max_iterations_no_agent_id_passes(): + """ + When no agent_id is set on the key, all requests pass through. + """ + local_cache = DualCache() + handler = _PROXY_MaxIterationsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + ) + user_api_key_dict = UserAPIKeyAuth( + api_key="sk-test-key-no-agent", + ) + + result = await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"metadata": {"session_id": "session-any"}}, + call_type="", + ) + assert result is None diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index d92e152d89a..3eb481991f7 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -1981,6 +1981,527 @@ async def test_execute_token_increment_script_cluster_compatibility(): ), f"Each key should have 2 args, got {len(args)} args for {len(keys)} keys" +@pytest.mark.asyncio +async def test_agent_level_rate_limit_descriptors(): + """ + Test that agent-level rate limit descriptors are created when + an agent has rpm_limit and/or tpm_limit configured. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_abc123" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="test-agent", + agent_card_params={"name": "Test Agent"}, + rpm_limit=50, + tpm_limit=5000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + assert captured_descriptors is not None + + agent_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent": + agent_descriptor = d + break + + assert agent_descriptor is not None, "Agent descriptor should be present" + assert agent_descriptor["value"] == _agent_id + assert agent_descriptor["rate_limit"]["requests_per_unit"] == 50 + assert agent_descriptor["rate_limit"]["tokens_per_unit"] == 5000 + + +@pytest.mark.asyncio +async def test_agent_session_rate_limit_descriptors(): + """ + Test that session-level rate limit descriptors are created when + an agent has session_rpm_limit/session_tpm_limit and a session_id is present. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_abc123" + _session_id = "sess_xyz789" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="test-agent", + agent_card_params={"name": "Test Agent"}, + session_rpm_limit=10, + session_tpm_limit=1000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4", + "metadata": {"session_id": _session_id}, + }, + call_type="", + ) + + assert captured_descriptors is not None + + session_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent_session": + session_descriptor = d + break + + assert session_descriptor is not None, "Agent session descriptor should be present" + assert session_descriptor["value"] == f"{_agent_id}:{_session_id}" + assert session_descriptor["rate_limit"]["requests_per_unit"] == 10 + assert session_descriptor["rate_limit"]["tokens_per_unit"] == 1000 + + +@pytest.mark.asyncio +async def test_agent_session_rate_limit_skipped_without_session_id(): + """ + Test that session-level rate limit descriptors are NOT created + when no session_id is available in the request. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_abc123" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="test-agent", + agent_card_params={"name": "Test Agent"}, + session_rpm_limit=10, + session_tpm_limit=1000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + # should_rate_limit should not have been called (no agent-level limits, only session limits + # but no session_id) + assert captured_descriptors is None, ( + "No descriptors should be created when agent has only session limits " + "but no session_id in request" + ) + + +@pytest.mark.asyncio +async def test_agent_rate_limit_from_metadata_agent_id(): + """ + Test that agent rate limits work when agent_id comes from + request metadata (header) rather than from the API key. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_from_header" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="header-agent", + agent_card_params={"name": "Header Agent"}, + rpm_limit=25, + tpm_limit=2500, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4", + "metadata": {"agent_id": _agent_id}, + }, + call_type="", + ) + + assert captured_descriptors is not None + + agent_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent": + agent_descriptor = d + break + + assert agent_descriptor is not None, "Agent descriptor should be created from metadata agent_id" + assert agent_descriptor["value"] == _agent_id + assert agent_descriptor["rate_limit"]["requests_per_unit"] == 25 + + +@pytest.mark.asyncio +async def test_agent_both_agent_and_session_rate_limits(): + """ + Test that both agent-level and session-level descriptors are created + when both types of limits are configured on the agent. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_dual" + _session_id = "sess_dual" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="dual-agent", + agent_card_params={"name": "Dual Agent"}, + rpm_limit=100, + tpm_limit=10000, + session_rpm_limit=20, + session_tpm_limit=2000, + ) + + captured_descriptors = None + + async def mock_should_rate_limit(descriptors, **kwargs): + nonlocal captured_descriptors + captured_descriptors = descriptors + return {"overall_code": "OK", "statuses": []} + + parallel_request_handler.should_rate_limit = mock_should_rate_limit + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={ + "model": "gpt-4", + "metadata": {"session_id": _session_id}, + }, + call_type="", + ) + + assert captured_descriptors is not None + + agent_descriptor = None + session_descriptor = None + for d in captured_descriptors: + if d["key"] == "agent": + agent_descriptor = d + elif d["key"] == "agent_session": + session_descriptor = d + + assert agent_descriptor is not None, "Agent-level descriptor should be present" + assert agent_descriptor["rate_limit"]["requests_per_unit"] == 100 + assert agent_descriptor["rate_limit"]["tokens_per_unit"] == 10000 + + assert session_descriptor is not None, "Session-level descriptor should be present" + assert session_descriptor["value"] == f"{_agent_id}:{_session_id}" + assert session_descriptor["rate_limit"]["requests_per_unit"] == 20 + assert session_descriptor["rate_limit"]["tokens_per_unit"] == 2000 + + +@pytest.mark.asyncio +async def test_agent_rate_limit_tpm_increment_on_success(monkeypatch): + """ + Test that async_log_success_event increments agent and session + TPM counters when agent_id and session_id are in metadata. + """ + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_tpm_test" + _session_id = "sess_tpm_test" + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + def mock_get_rate_limit_type(): + return "total" + + monkeypatch.setattr( + parallel_request_handler, "get_rate_limit_type", mock_get_rate_limit_type + ) + + mock_usage = Usage(prompt_tokens=20, completion_tokens=30, total_tokens=50) + mock_response = ModelResponse( + id="mock-response", + object="chat.completion", + created=int(datetime.now().timestamp()), + model="gpt-4", + usage=mock_usage, + choices=[], + ) + + mock_kwargs = { + "standard_logging_object": { + "metadata": { + "user_api_key_hash": _api_key, + "user_api_key_user_id": None, + "user_api_key_team_id": None, + "user_api_key_end_user_id": None, + "agent_id": _agent_id, + "session_id": _session_id, + } + }, + "model": "gpt-4", + } + + captured_operations = [] + + async def mock_increment_pipeline(increment_list, **kwargs): + captured_operations.extend(increment_list) + return True + + monkeypatch.setattr( + parallel_request_handler.internal_usage_cache.dual_cache, + "async_increment_cache_pipeline", + mock_increment_pipeline, + ) + + await parallel_request_handler.async_log_success_event( + kwargs=mock_kwargs, + response_obj=mock_response, + start_time=datetime.now(), + end_time=datetime.now(), + ) + + agent_tpm_op = None + session_tpm_op = None + for op in captured_operations: + if op["key"] == f"{{agent:{_agent_id}}}:tokens": + agent_tpm_op = op + elif op["key"] == f"{{agent_session:{_agent_id}:{_session_id}}}:tokens": + session_tpm_op = op + + assert agent_tpm_op is not None, "Agent TPM increment should be present" + assert agent_tpm_op["increment_value"] == 50 + + assert session_tpm_op is not None, "Session TPM increment should be present" + assert session_tpm_op["increment_value"] == 50 + + +@pytest.mark.asyncio +async def test_agent_rate_limit_429_on_over_limit(monkeypatch, time_controller): + """ + Test end-to-end that agent rate limiting returns 429 when the agent + RPM limit is exceeded. + """ + from unittest.mock import patch + + from litellm.types.agents import AgentResponse + + monkeypatch.setenv("LITELLM_RATE_LIMIT_WINDOW_SIZE", "2") + _api_key = "sk-12345" + _api_key = hash_token(_api_key) + _agent_id = "agent_429_test" + + user_api_key_dict = UserAPIKeyAuth( + api_key=_api_key, + agent_id=_agent_id, + ) + + local_cache = DualCache() + parallel_request_handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache), + time_provider=time_controller.now, + ) + + mock_agent = AgentResponse( + agent_id=_agent_id, + agent_name="rate-limited-agent", + agent_card_params={"name": "Rate Limited Agent"}, + rpm_limit=2, + ) + + window_starts: Dict[str, int] = {} + request_counts: Dict[str, int] = {} + + async def mock_batch_rate_limiter(*args, **kwargs): + keys = kwargs.get("keys") if kwargs else args[0] + args_list = kwargs.get("args") if kwargs else args[1] + now = args_list[0] + window_size = args_list[1] + results = [] + for i in range(0, len(keys), 2): + window_key = keys[i] + counter_key = keys[i + 1] + prev_window = window_starts.get(window_key) + prev_counter = request_counts.get(counter_key, 0) + if prev_window is None or (now - prev_window) >= window_size: + window_starts[window_key] = now + new_counter = 1 + request_counts[counter_key] = new_counter + await local_cache.async_set_cache( + key=window_key, value=now, ttl=window_size + ) + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) + else: + new_counter = prev_counter + 1 + request_counts[counter_key] = new_counter + await local_cache.async_set_cache( + key=counter_key, value=new_counter, ttl=window_size + ) + results.append(now) + results.append(new_counter) + return results + + parallel_request_handler.batch_rate_limiter_script = mock_batch_rate_limiter + + with patch( + "litellm.proxy.agent_endpoints.agent_registry.global_agent_registry.get_agent_by_id", + return_value=mock_agent, + ): + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + with pytest.raises(HTTPException) as exc_info: + await parallel_request_handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4"}, + call_type="", + ) + + assert exc_info.value.status_code == 429 + assert "agent" in exc_info.value.detail + + class TestGetTotalTokensFromUsageCacheExclusion: """ Tests for _get_total_tokens_from_usage cache token exclusion. diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 69efbf19c38..d4f60ff0d7f 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -13279,6 +13279,21 @@ "type": "github", "url": "https://github.com/sponsors/wooorm" } + }, + "node_modules/@next/swc-win32-ia32-msvc": { + "version": "14.2.33", + "resolved": "https://registry.npmjs.org/@next/swc-win32-ia32-msvc/-/swc-win32-ia32-msvc-14.2.33.tgz", + "integrity": "sha512-pc9LpGNKhJ0dXQhZ5QMmYxtARwwmWLpeocFmVG5Z0DzWq5Uf0izcI8tLc+qOpqxO1PWqZ5A7J1blrUIKrIFc7Q==", + "cpu": [ + "ia32" + ], + "optional": true, + "os": [ + "win32" + ], + "engines": { + "node": ">= 10" + } } } } diff --git a/ui/litellm-dashboard/package.json b/ui/litellm-dashboard/package.json index ea84ea6f401..06aff85e3f2 100644 --- a/ui/litellm-dashboard/package.json +++ b/ui/litellm-dashboard/package.json @@ -88,8 +88,8 @@ "mermaid": ">=11.10.0", "js-yaml": ">=4.1.1", "glob": ">=11.1.0", - "tar": ">=7.5.8", - "minimatch": ">=10.2.1", + "tar": ">=7.5.10", + "minimatch": ">=10.2.4", "@isaacs/brace-expansion": ">=5.0.1", "node-forge": ">=1.3.2", "lodash-es": ">=4.17.23", diff --git a/ui/litellm-dashboard/src/components/agents.test.tsx b/ui/litellm-dashboard/src/components/agents.test.tsx index 2d4c879dec0..8914f8ad6bc 100644 --- a/ui/litellm-dashboard/src/components/agents.test.tsx +++ b/ui/litellm-dashboard/src/components/agents.test.tsx @@ -1,7 +1,8 @@ import React from "react"; -import { render, screen, waitFor } from "@testing-library/react"; +import { render, screen, waitFor, act, fireEvent } from "@testing-library/react"; import { describe, it, expect, vi, beforeEach } from "vitest"; import AgentsPanel from "./agents"; +import * as networking from "./networking"; vi.mock("./networking", () => ({ getAgentsList: vi.fn().mockResolvedValue({ agents: [] }), @@ -68,4 +69,37 @@ describe("AgentsPanel", () => { expect(grid).toHaveAttribute("data-is-admin", "false"); }); }); + + it("should render the Health Check toggle", async () => { + render(); + expect(screen.getByText("Health Check")).toBeInTheDocument(); + }); + + it("should render the Health Check toggle for non-admin users too", async () => { + render(); + expect(screen.getByText("Health Check")).toBeInTheDocument(); + }); + + it("should call getAgentsList with health_check=false on initial load", async () => { + render(); + await waitFor(() => { + expect(networking.getAgentsList).toHaveBeenCalledWith("test-token", false); + }); + }); + + it("should call getAgentsList with health_check=true when toggle is enabled", async () => { + render(); + await waitFor(() => { + expect(networking.getAgentsList).toHaveBeenCalledWith("test-token", false); + }); + + const toggle = screen.getByRole("switch"); + await act(async () => { + fireEvent.click(toggle); + }); + + await waitFor(() => { + expect(networking.getAgentsList).toHaveBeenCalledWith("test-token", true); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/agents.tsx b/ui/litellm-dashboard/src/components/agents.tsx index 8dd9bc7d01c..b60af254133 100644 --- a/ui/litellm-dashboard/src/components/agents.tsx +++ b/ui/litellm-dashboard/src/components/agents.tsx @@ -1,13 +1,26 @@ import React, { useState, useEffect } from "react"; -import { Button } from "@tremor/react"; -import { Modal, Alert } 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; @@ -26,17 +39,18 @@ const AgentsPanel: React.FC = ({ accessToken, userRole }) => { const [isDeleting, setIsDeleting] = useState(false); const [agentToDelete, setAgentToDelete] = useState<{ id: string; name: string } | null>(null); const [selectedAgentId, setSelectedAgentId] = useState(null); + const [healthCheckEnabled, setHealthCheckEnabled] = useState(false); const isAdmin = userRole ? isAdminRole(userRole) : false; - const fetchAgents = async () => { + const fetchAgents = async (healthCheck?: boolean) => { if (!accessToken) { return; } setIsLoading(true); try { - const response: AgentsResponse = await getAgentsList(accessToken); + const response: AgentsResponse = await getAgentsList(accessToken, healthCheck ?? healthCheckEnabled); setAgentsList(response.agents || []); } catch (error) { console.error("Error fetching agents:", error); @@ -89,6 +103,11 @@ const AgentsPanel: React.FC = ({ accessToken, userRole }) => { } }, [accessToken, agentsList.length]); + const handleHealthCheckToggle = (checked: boolean) => { + setHealthCheckEnabled(checked); + fetchAgents(checked); + }; + const handleAddAgent = () => { if (selectedAgentId) { setSelectedAgentId(null); @@ -129,6 +148,14 @@ const AgentsPanel: React.FC = ({ accessToken, userRole }) => { setAgentToDelete(null); }; + const sortedAgents = [...agentsList].sort((a, b) => { + const dateA = a.created_at ? new Date(a.created_at).getTime() : 0; + const dateB = b.created_at ? new Date(b.created_at).getTime() : 0; + return dateB - dateA; + }); + + const columnCount = isAdmin ? 7 : 6; + return (
@@ -141,13 +168,25 @@ const AgentsPanel: React.FC = ({ accessToken, userRole }) => { showIcon className="mb-3" /> - {isAdmin && ( -
+
+ {isAdmin && ( -
- )} + )} + +
+ + Health Check + +
+ +
{selectedAgentId ? ( @@ -158,16 +197,84 @@ const AgentsPanel: React.FC = ({ accessToken, userRole }) => { isAdmin={isAdmin} /> ) : ( - setSelectedAgentId(id)} - /> + + {isLoading ? ( + + ) : ( + + + + Agent Name + Agent ID + Spend (USD) + Model + Created + Status + {isAdmin && Actions} + + + + {sortedAgents.length === 0 ? ( + + + No agents found. Click "+ Add New Agent" to create one. + + + ) : ( + sortedAgents.map((agent) => ( + + + {agent.agent_name} + + + + + + + + {formatNumberWithCommas(agent.spend, 4)} + + + + {agent.litellm_params?.model || "N/A"} + + + + + {agent.created_at + ? new Date(agent.created_at).toLocaleDateString() + : "N/A"} + + + + {keyInfoMap[agent.agent_id]?.has_key ? ( + Active + ) : ( + Needs Setup + )} + + {isAdmin && ( + + handleDeleteClick(agent.agent_id, agent.agent_name)} + /> + + )} + + )) + )} + +
+ )} + )} = ({ accessToken, userRole }) => { }; export default AgentsPanel; - diff --git a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx index 10fbfa4615a..0cec0331f43 100644 --- a/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/components/agents/add_agent_form.tsx @@ -1,5 +1,5 @@ import React, { useState, useEffect } from "react"; -import { Modal, Form, message, Select, Input, Steps, Radio, Tag, Divider } from "antd"; +import { Modal, Form, message, Select, Input, Steps, Radio, Tag, Divider, Switch, InputNumber, Collapse } from "antd"; import { Button } from "@tremor/react"; import { CheckCircleFilled, KeyOutlined, RobotOutlined, AppstoreOutlined, InfoCircleOutlined } from "@ant-design/icons"; import CreatedKeyDisplay from "../shared/CreatedKeyDisplay"; @@ -60,6 +60,12 @@ const AddAgentForm: React.FC = ({ const [createdKeyValue, setCreatedKeyValue] = useState(null); const [assignedKeyAlias, setAssignedKeyAlias] = useState(null); + // Tracing & guardrails state + const [requireTraceIdInbound, setRequireTraceIdInbound] = useState(false); + const [requireTraceIdOutbound, setRequireTraceIdOutbound] = useState(false); + const [maxIterations, setMaxIterations] = useState(null); + const [maxBudgetPerSession, setMaxBudgetPerSession] = useState(null); + // Fetch agent type metadata on mount useEffect(() => { const fetchMetadata = async () => { @@ -218,6 +224,19 @@ const AddAgentForm: React.FC = ({ } } + // Wire trace-id flags and budget controls into agent litellm_params (before create call) + if (requireTraceIdInbound || requireTraceIdOutbound) { + if (!agentData.litellm_params) agentData.litellm_params = {}; + if (requireTraceIdInbound) { + agentData.litellm_params.require_trace_id_on_calls_to_agent = true; + } + if (requireTraceIdOutbound) { + agentData.litellm_params.require_trace_id_on_calls_by_agent = true; + if (maxIterations) agentData.litellm_params.max_iterations = maxIterations; + if (maxBudgetPerSession) agentData.litellm_params.max_budget_per_session = maxBudgetPerSession; + } + } + const agentResponse = await createAgentCall(accessToken, agentData); const agentId: string = agentResponse.agent_id; const agentName: string = agentResponse.agent_name || values.agent_name || agentId; @@ -267,6 +286,10 @@ const AddAgentForm: React.FC = ({ setCreatedAgentName(""); setCreatedKeyValue(null); setAssignedKeyAlias(null); + setRequireTraceIdInbound(false); + setRequireTraceIdOutbound(false); + setMaxIterations(null); + setMaxBudgetPerSession(null); onClose(); }; @@ -315,6 +338,122 @@ const AddAgentForm: React.FC = ({
)} + + Tracing, + children: ( +
+
+
+ + Require x-litellm-trace-id on calls TO this agent + +

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

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

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

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

Hard cap on LLM calls per session

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

Max spend per trace before returning 429

+
+
+ + + +
Agent Rate Limits
+

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

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

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

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