From 603ec9b35329ffbca2d036a8dd16eee63b0083ad Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Sat, 26 Sep 2026 12:09:01 -0700 Subject: [PATCH] feat(agents): define persisted budget windows and accounting keys --- .../migration.sql | 15 ++++++++++ .../litellm_proxy_extras/schema.prisma | 4 +++ litellm/proxy/schema.prisma | 4 +++ litellm/types/agents.py | 24 +++++++++++++++- litellm/types/proxy/agent_identity.py | 16 +++++++++++ schema.prisma | 4 +++ tests/unit/types/proxy/test_agent_identity.py | 20 +++++++++++++ tests/unit/types/test_agents.py | 28 +++++++++++++++++++ 8 files changed, 114 insertions(+), 1 deletion(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260925220000_agent_budgets/migration.sql create mode 100644 tests/unit/types/proxy/test_agent_identity.py create mode 100644 tests/unit/types/test_agents.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260925220000_agent_budgets/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260925220000_agent_budgets/migration.sql new file mode 100644 index 00000000000..6f1db719855 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260925220000_agent_budgets/migration.sql @@ -0,0 +1,15 @@ +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "budget_id" TEXT; + +-- CreateIndex +CREATE UNIQUE INDEX IF NOT EXISTS "LiteLLM_AgentsTable_budget_id_key" ON "LiteLLM_AgentsTable"("budget_id"); + +-- AddForeignKey +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'LiteLLM_AgentsTable_budget_id_fkey') THEN + ALTER TABLE "LiteLLM_AgentsTable" ADD CONSTRAINT "LiteLLM_AgentsTable_budget_id_fkey" FOREIGN KEY ("budget_id") REFERENCES "LiteLLM_BudgetTable"("budget_id") ON DELETE SET NULL ON UPDATE CASCADE; + END IF; +END $$; + + +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "spend_window" TIMESTAMP(3); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 8535d1004f7..92dfe53e29e 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -36,6 +36,7 @@ model LiteLLM_BudgetTable { model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization + agents LiteLLM_AgentsTable[] } // Models on proxy @@ -83,6 +84,9 @@ model LiteLLM_AgentsTable { execution_mode String @default("autonomous") identity LiteLLM_AgentIdentity? retired_identities LiteLLM_RetiredAgentIdentity[] + budget_id String? @unique + spend_window DateTime? + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) tpm_limit Int? rpm_limit Int? session_tpm_limit Int? diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 8535d1004f7..92dfe53e29e 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -36,6 +36,7 @@ model LiteLLM_BudgetTable { model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization + agents LiteLLM_AgentsTable[] } // Models on proxy @@ -83,6 +84,9 @@ model LiteLLM_AgentsTable { execution_mode String @default("autonomous") identity LiteLLM_AgentIdentity? retired_identities LiteLLM_RetiredAgentIdentity[] + budget_id String? @unique + spend_window DateTime? + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) tpm_limit Int? rpm_limit Int? session_tpm_limit Int? diff --git a/litellm/types/agents.py b/litellm/types/agents.py index 94adb9f7c4a..5564105de7d 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -1,5 +1,5 @@ from collections.abc import Mapping, Sequence -from datetime import datetime +from datetime import datetime, timezone from typing import TYPE_CHECKING, Annotated, Any, Final, Literal, TypeAlias from urllib.parse import urlsplit @@ -15,6 +15,7 @@ from litellm.types.proxy.agent_identity import ( if TYPE_CHECKING: from a2a.types import SendMessageResponse + from prisma.types import LiteLLM_AgentsTableWhereInput # AgentProvider @@ -311,6 +312,27 @@ class AgentKeySummary(BaseModel): key_name: str | None = None +def agent_budget_counter_key(agent_id: str, reset_at: datetime | None) -> str: + if reset_at is None: + return f"spend:agent:{agent_id}" + aware: Final = reset_at if reset_at.tzinfo is not None else reset_at.replace(tzinfo=timezone.utc) + window: Final = aware.astimezone(timezone.utc).strftime("%Y%m%dT%H%M%S.%fZ") + return f"spend:agent_window:{window}:{agent_id}" + + +def agent_spend_filter(counter_key: str) -> "LiteLLM_AgentsTableWhereInput": + if counter_key.startswith("spend:agent_window:"): + _, _, raw_window, agent_id = counter_key.split(":", 3) + window: Final = datetime.strptime(raw_window, "%Y%m%dT%H%M%S.%fZ").replace(tzinfo=timezone.utc) + windowed: Final[LiteLLM_AgentsTableWhereInput] = {"agent_id": agent_id, "spend_window": window} + return windowed + cumulative: Final[LiteLLM_AgentsTableWhereInput] = { + "agent_id": counter_key.removeprefix("spend:agent:"), + "spend_window": None, + } + return cumulative + + class AgentResponse(BaseModel): identity: AgentIdentityBinding | None = None identity_managed: bool = False diff --git a/litellm/types/proxy/agent_identity.py b/litellm/types/proxy/agent_identity.py index a7fe0be37e1..2dc43976e0f 100644 --- a/litellm/types/proxy/agent_identity.py +++ b/litellm/types/proxy/agent_identity.py @@ -46,6 +46,22 @@ class AgentIdentityBinding(BaseModel): last_authenticated_at: datetime | None = None +class AgentBudgetConfig(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + max_budget: float = Field(ge=0, allow_inf_nan=False) + budget_duration: str | None = None + + +class AgentBudgetState(BaseModel): + model_config = ConfigDict(frozen=True) + + budget_id: str + max_budget: float | None = None + budget_duration: str | None = None + budget_reset_at: datetime | None = None + + class AgentSubject(BaseModel): model_config = ConfigDict(frozen=True) diff --git a/schema.prisma b/schema.prisma index 8535d1004f7..92dfe53e29e 100644 --- a/schema.prisma +++ b/schema.prisma @@ -36,6 +36,7 @@ model LiteLLM_BudgetTable { model_access_groups LiteLLM_ModelAccessGroupBudgetTable[] // multiple model access groups can have the same budget team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization + agents LiteLLM_AgentsTable[] } // Models on proxy @@ -83,6 +84,9 @@ model LiteLLM_AgentsTable { execution_mode String @default("autonomous") identity LiteLLM_AgentIdentity? retired_identities LiteLLM_RetiredAgentIdentity[] + budget_id String? @unique + spend_window DateTime? + litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) tpm_limit Int? rpm_limit Int? session_tpm_limit Int? diff --git a/tests/unit/types/proxy/test_agent_identity.py b/tests/unit/types/proxy/test_agent_identity.py new file mode 100644 index 00000000000..ed4b0901da2 --- /dev/null +++ b/tests/unit/types/proxy/test_agent_identity.py @@ -0,0 +1,20 @@ +import pytest +from pydantic import ValidationError + +from litellm.types.proxy.agent_identity import AgentBudgetConfig + + +@pytest.mark.parametrize("amount", [-1, float("inf"), float("-inf"), float("nan")]) +def test_agent_budget_rejects_negative_or_nonfinite_caps(amount: float) -> None: + with pytest.raises(ValidationError): + AgentBudgetConfig(max_budget=amount) + + +@pytest.mark.parametrize("amount", [0, 0.01, 100]) +def test_agent_budget_preserves_a_finite_nonnegative_cap(amount: float) -> None: + assert AgentBudgetConfig(max_budget=amount).max_budget == amount + + +def test_agent_budget_rejects_unknown_policy_fields() -> None: + with pytest.raises(ValidationError): + AgentBudgetConfig.model_validate({"max_budget": 1, "unknown_control": True}) diff --git a/tests/unit/types/test_agents.py b/tests/unit/types/test_agents.py new file mode 100644 index 00000000000..a583ab9e8a2 --- /dev/null +++ b/tests/unit/types/test_agents.py @@ -0,0 +1,28 @@ +from datetime import datetime, timedelta, timezone +from typing import Final + +import pytest + +from litellm.types.agents import agent_budget_counter_key, agent_spend_filter + + +@pytest.mark.parametrize("offset", [None, timezone.utc, timezone(timedelta(hours=5, minutes=30))]) +def test_agent_window_key_round_trip_preserves_the_admitted_instant(offset) -> None: + instant: Final = datetime(2030, 1, 1, 12, 30, 1, 123000, tzinfo=offset) + expected: Final = instant.replace(tzinfo=timezone.utc) if offset is None else instant.astimezone(timezone.utc) + key: Final = agent_budget_counter_key("agent:with:colons", instant) + assert agent_spend_filter(key) == {"agent_id": "agent:with:colons", "spend_window": expected} + assert key == agent_budget_counter_key("agent:with:colons", expected) + + +def test_unbudgeted_key_is_filtered_to_an_unbudgeted_row() -> None: + key: Final = agent_budget_counter_key("agent-one", None) + assert agent_spend_filter(key) == {"agent_id": "agent-one", "spend_window": None} + + +def test_different_budget_windows_never_share_a_settlement_filter() -> None: + first: Final = datetime(2030, 1, 1, tzinfo=timezone.utc) + second: Final = first + timedelta(days=1) + assert agent_spend_filter(agent_budget_counter_key("agent-one", first)) != agent_spend_filter( + agent_budget_counter_key("agent-one", second) + )