mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
feat(proxy): expose lifetime total_spend on virtual keys
Adds a persistent total_spend column to LiteLLM_VerificationToken and LiteLLM_DeletedVerificationToken, incremented in the same write as spend and left alone by budget resets. Surfaces it on /key/info, /key/list and the Admin UI Virtual Keys table and key detail view Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
a8979fe054
commit
6cf35ed71b
18 changed files with 213 additions and 1 deletions
|
|
@ -0,0 +1,5 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
||||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "total_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0;
|
||||
|
|
@ -426,6 +426,7 @@ model LiteLLM_VerificationToken {
|
|||
key_alias String?
|
||||
soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
expires DateTime?
|
||||
models String[]
|
||||
aliases Json @default("{}")
|
||||
|
|
@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
key_alias String?
|
||||
soft_budget_cooldown Boolean @default(false)
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
expires DateTime?
|
||||
models String[]
|
||||
aliases Json @default("{}")
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase):
|
|||
key_name: str | None = None
|
||||
key_alias: str | None = None
|
||||
spend: float = 0.0
|
||||
total_spend: float = 0.0
|
||||
max_budget: float | None = None
|
||||
expires: str | datetime | None = None
|
||||
models: list = []
|
||||
|
|
|
|||
|
|
@ -1619,6 +1619,7 @@ class DBSpendUpdateWriter:
|
|||
where={"token": token},
|
||||
data={
|
||||
"spend": {"increment": response_cost},
|
||||
"total_spend": {"increment": response_cost},
|
||||
"last_active": datetime.now(timezone.utc),
|
||||
},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -426,6 +426,7 @@ model LiteLLM_VerificationToken {
|
|||
key_alias String?
|
||||
soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
expires DateTime?
|
||||
models String[]
|
||||
aliases Json @default("{}")
|
||||
|
|
@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
key_alias String?
|
||||
soft_budget_cooldown Boolean @default(false)
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
expires DateTime?
|
||||
models String[]
|
||||
aliases Json @default("{}")
|
||||
|
|
|
|||
|
|
@ -426,6 +426,7 @@ model LiteLLM_VerificationToken {
|
|||
key_alias String?
|
||||
soft_budget_cooldown Boolean @default(false) // key-level state on if budget alerts need to be cooled down
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
expires DateTime?
|
||||
models String[]
|
||||
aliases Json @default("{}")
|
||||
|
|
@ -528,6 +529,7 @@ model LiteLLM_DeletedVerificationToken {
|
|||
key_alias String?
|
||||
soft_budget_cooldown Boolean @default(false)
|
||||
spend Float @default(0.0)
|
||||
total_spend Float @default(0.0)
|
||||
expires DateTime?
|
||||
models String[]
|
||||
aliases Json @default("{}")
|
||||
|
|
|
|||
|
|
@ -362,6 +362,14 @@ class TestVerificationToken:
|
|||
assert deleted.deleted_at is not None
|
||||
assert deleted.token == "t1"
|
||||
|
||||
def test_total_spend_is_carried_separately_from_resettable_spend(self):
|
||||
token = LiteLLM_VerificationToken(token="t1", spend=0.0, total_spend=12.5)
|
||||
assert token.model_dump()["total_spend"] == 12.5
|
||||
assert token.model_dump()["spend"] == 0.0
|
||||
|
||||
deleted = LiteLLM_DeletedVerificationToken.model_validate({**token.model_dump(), "deleted_by": "admin"})
|
||||
assert deleted.total_spend == 12.5
|
||||
|
||||
|
||||
class TestConfigTable:
|
||||
def test_config_creation(self):
|
||||
|
|
|
|||
|
|
@ -291,6 +291,23 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client):
|
|||
assert set(write["data"].keys()) == {"spend", "budget_reset_at"}
|
||||
|
||||
|
||||
def test_reset_budget_for_key_leaves_lifetime_total_spend_alone(reset_budget_job, mock_prisma_client):
|
||||
"""A period reset zeroes spend but must neither write nor touch the lifetime total_spend."""
|
||||
now = datetime.now(timezone.utc)
|
||||
key = LiteLLM_VerificationToken(
|
||||
token="tok-key-1", spend=100.0, total_spend=340.0, budget_duration="30d", budget_reset_at=now
|
||||
)
|
||||
mock_prisma_client.data["key"] = [key]
|
||||
|
||||
asyncio.run(reset_budget_job.reset_budget_for_litellm_keys())
|
||||
|
||||
(write,) = _batch_writes(mock_prisma_client, "key")
|
||||
assert write["data"]["spend"] == {"decrement": 100.0}
|
||||
assert "total_spend" not in write["data"]
|
||||
assert key.spend == 0.0
|
||||
assert key.total_spend == 340.0
|
||||
|
||||
|
||||
def test_reset_budget_for_key_honors_injected_reset_time(mock_prisma_client, mock_proxy_logging):
|
||||
"""Injected BudgetResetSettings drives the written reset time end to end (DI, no globals).
|
||||
|
||||
|
|
|
|||
|
|
@ -1658,6 +1658,58 @@ async def test_commit_key_spend_updates_includes_last_active():
|
|||
assert before_call <= last_active <= after_call
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_commit_spend_updates_to_db_increments_key_total_spend_alongside_spend():
|
||||
"""
|
||||
The key table write must increment the lifetime total_spend by the same amount as the
|
||||
resettable spend, in the same update so the two cannot drift.
|
||||
"""
|
||||
db_writer = DBSpendUpdateWriter()
|
||||
|
||||
mock_batcher = MagicMock()
|
||||
mock_batcher.litellm_verificationtoken = MagicMock()
|
||||
mock_batcher.litellm_verificationtoken.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)
|
||||
|
||||
db_spend_update_transactions = {
|
||||
"user_list_transactions": {},
|
||||
"end_user_list_transactions": {},
|
||||
"key_list_transactions": {"hashed_token_abc": 0.05, "hashed_token_def": 1.25},
|
||||
"team_list_transactions": {},
|
||||
"team_member_list_transactions": {},
|
||||
"org_list_transactions": {},
|
||||
"tag_list_transactions": {},
|
||||
"agent_list_transactions": {},
|
||||
}
|
||||
|
||||
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=MagicMock(),
|
||||
db_spend_update_transactions=db_spend_update_transactions,
|
||||
)
|
||||
|
||||
calls = mock_batcher.litellm_verificationtoken.update_many.call_args_list
|
||||
assert [c.kwargs["where"] for c in calls] == [{"token": "hashed_token_abc"}, {"token": "hashed_token_def"}]
|
||||
for call, expected_cost in zip(calls, (0.05, 1.25)):
|
||||
assert call.kwargs["data"]["spend"] == {"increment": expected_cost}
|
||||
assert call.kwargs["data"]["total_spend"] == call.kwargs["data"]["spend"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_update_database_creates_single_task():
|
||||
"""
|
||||
|
|
@ -2813,7 +2865,7 @@ async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at
|
|||
mock_batcher.litellm_verificationtoken.update_many.assert_called_once()
|
||||
call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1]
|
||||
assert call_kwargs["where"] == {"token": token}
|
||||
assert set(call_kwargs["data"]) == {"spend", "last_active"}
|
||||
assert set(call_kwargs["data"]) == {"spend", "total_spend", "last_active"}
|
||||
assert call_kwargs["data"]["spend"] == {"increment": response_cost}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1431,6 +1431,60 @@ async def test_key_info_returns_object_permission(monkeypatch):
|
|||
)
|
||||
|
||||
|
||||
def _stored_key_with_lifetime_spend(token: str, spend: float, total_spend: float) -> LiteLLM_VerificationToken:
|
||||
return LiteLLM_VerificationToken.model_validate(
|
||||
{"token": token, "user_id": "user123", "spend": spend, "total_spend": total_spend}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_info_returns_lifetime_total_spend_next_to_resettable_spend(monkeypatch):
|
||||
"""After a budget reset the period spend is 0 while total_spend keeps the lifetime figure."""
|
||||
from litellm.proxy.management_endpoints.key_management_endpoints import info_key_fn
|
||||
|
||||
mock_prisma_client = AsyncMock()
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(
|
||||
return_value=_stored_key_with_lifetime_spend(token="hashed_key", spend=0.0, total_spend=3.75)
|
||||
)
|
||||
|
||||
result = await info_key_fn(
|
||||
key="sk-test-key-456",
|
||||
user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-test-key-456"),
|
||||
)
|
||||
|
||||
assert result["info"]["spend"] == 0.0
|
||||
assert result["info"]["total_spend"] == 3.75
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_keys_full_object_returns_lifetime_total_spend():
|
||||
mock_prisma_client = AsyncMock()
|
||||
mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(
|
||||
return_value=[_stored_key_with_lifetime_spend(token="hashed_key", spend=0.0, total_spend=3.75)]
|
||||
)
|
||||
mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=1)
|
||||
|
||||
result = await _list_key_helper(
|
||||
prisma_client=mock_prisma_client,
|
||||
page=1,
|
||||
size=50,
|
||||
user_id=None,
|
||||
team_id=None,
|
||||
organization_id=None,
|
||||
key_alias=None,
|
||||
key_hash=None,
|
||||
exclude_team_id=None,
|
||||
return_full_object=True,
|
||||
admin_team_ids=None,
|
||||
)
|
||||
|
||||
listed_key = result["keys"][0]
|
||||
assert isinstance(listed_key, UserAPIKeyAuth)
|
||||
assert listed_key.spend == 0.0
|
||||
assert listed_key.total_spend == 3.75
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_new_token_with_valid_key(monkeypatch):
|
||||
"""Test get_new_token function when provided with a valid key that starts with 'sk-'"""
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ const mockDeletedKey: DeletedKeyResponse = {
|
|||
key_name: "test-key",
|
||||
key_alias: "Test Key Alias",
|
||||
spend: 5.5,
|
||||
total_spend: 5.5,
|
||||
max_budget: 100,
|
||||
expires: "2024-12-31T23:59:59Z",
|
||||
models: ["gpt-3.5-turbo"],
|
||||
|
|
|
|||
|
|
@ -79,6 +79,7 @@ const mockKey: KeyResponse = {
|
|||
key_name: "test-key",
|
||||
key_alias: "Test Key Alias",
|
||||
spend: 5.5,
|
||||
total_spend: 42.25,
|
||||
max_budget: 100,
|
||||
expires: "2999-12-31T23:59:59Z",
|
||||
models: ["gpt-3.5-turbo", "gpt-4"],
|
||||
|
|
@ -236,6 +237,14 @@ it("should display key information correctly", async () => {
|
|||
});
|
||||
});
|
||||
|
||||
it("shows lifetime spend in its own column next to the period spend meter", async () => {
|
||||
renderWithProviders(<VirtualKeysTable />);
|
||||
|
||||
expect(await screen.findByText("Lifetime Spend")).toBeInTheDocument();
|
||||
expect(screen.getByText("$42.2500")).toBeInTheDocument();
|
||||
expect(screen.getByText("$5.5000")).toBeInTheDocument();
|
||||
});
|
||||
|
||||
it("should display user email correctly", async () => {
|
||||
renderWithProviders(<VirtualKeysTable />);
|
||||
|
||||
|
|
|
|||
|
|
@ -13,6 +13,7 @@ import {
|
|||
IdCell,
|
||||
IdentityCell,
|
||||
ModelsCell,
|
||||
MoneyCell,
|
||||
SpendBudgetCell,
|
||||
StatusBadge,
|
||||
UserPopoverCell,
|
||||
|
|
@ -274,6 +275,20 @@ export const getKeyTableColumns = ({
|
|||
);
|
||||
},
|
||||
},
|
||||
{
|
||||
id: "total_spend",
|
||||
accessorKey: "total_spend",
|
||||
meta: { title: "Lifetime Spend" },
|
||||
header: () => (
|
||||
<InfoHeader
|
||||
label="Lifetime Spend"
|
||||
tooltip="Cumulative spend across every budget period. Budget resets do not touch this value. Keys created before this field existed only count spend from then on."
|
||||
/>
|
||||
),
|
||||
size: 130,
|
||||
enableSorting: false,
|
||||
cell: (info) => <MoneyCell value={info.getValue() as number | null | undefined} showZero />,
|
||||
},
|
||||
{
|
||||
id: "budget_reset_at",
|
||||
accessorKey: "budget_reset_at",
|
||||
|
|
|
|||
|
|
@ -39,6 +39,7 @@ export interface KeyResponse {
|
|||
key_name: string;
|
||||
key_alias: string;
|
||||
spend: number;
|
||||
total_spend: number;
|
||||
max_budget: number;
|
||||
expires: string;
|
||||
models: string[];
|
||||
|
|
|
|||
|
|
@ -174,6 +174,7 @@ describe("KeyEditView", () => {
|
|||
key_name: "sk-...TUuw",
|
||||
key_alias: "asdasdas",
|
||||
spend: 0,
|
||||
total_spend: 0,
|
||||
max_budget: 0,
|
||||
expires: "null",
|
||||
models: [],
|
||||
|
|
|
|||
|
|
@ -119,6 +119,7 @@ describe("KeyInfoView", () => {
|
|||
key_name: "sk-...TUuw",
|
||||
key_alias: "asdasdas",
|
||||
spend: 0,
|
||||
total_spend: 0,
|
||||
max_budget: 0,
|
||||
expires: "null",
|
||||
models: [],
|
||||
|
|
@ -272,6 +273,23 @@ describe("KeyInfoView", () => {
|
|||
});
|
||||
});
|
||||
|
||||
it("shows lifetime spend separately from the resettable period spend", async () => {
|
||||
vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock);
|
||||
|
||||
renderWithProviders(
|
||||
<KeyInfoView
|
||||
keyData={{ ...MOCK_KEY_DATA, spend: 0.25, total_spend: 340.5 }}
|
||||
onClose={() => {}}
|
||||
keyId={"test-key-id"}
|
||||
onKeyDataUpdate={() => {}}
|
||||
teams={[]}
|
||||
/>,
|
||||
);
|
||||
|
||||
expect(await screen.findByText("$0.2500")).toBeInTheDocument();
|
||||
expect(screen.getByTestId("key-lifetime-spend")).toHaveTextContent("Lifetime spend: $340.5000");
|
||||
});
|
||||
|
||||
it("should render the key's saved router fallbacks", async () => {
|
||||
vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock);
|
||||
|
||||
|
|
|
|||
|
|
@ -677,6 +677,9 @@ export default function KeyInfoView({
|
|||
{currentKeyData.budget_reset_at && (
|
||||
<p className="text-sm">Resets {formatTimestamp(currentKeyData.budget_reset_at)}</p>
|
||||
)}
|
||||
<p className="text-sm mt-2" data-testid="key-lifetime-spend">
|
||||
Lifetime spend: ${formatNumberWithCommas(currentKeyData.total_spend ?? 0, 4)}
|
||||
</p>
|
||||
</div>
|
||||
</Card>
|
||||
|
||||
|
|
@ -935,6 +938,11 @@ export default function KeyInfoView({
|
|||
<p className="text-sm">${formatNumberWithCommas(currentKeyData.spend, 4)} USD</p>
|
||||
</div>
|
||||
|
||||
<div>
|
||||
<p className="text-sm font-medium">Lifetime Spend</p>
|
||||
<p className="text-sm">${formatNumberWithCommas(currentKeyData.total_spend ?? 0, 4)} USD</p>
|
||||
</div>
|
||||
|
||||
<div>
|
||||
<p className="text-sm font-medium">Budget</p>
|
||||
<p className="text-sm">
|
||||
|
|
|
|||
15
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
15
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -29686,6 +29686,11 @@ export interface components {
|
|||
team_id?: string | null;
|
||||
/** Token */
|
||||
token?: string | null;
|
||||
/**
|
||||
* Total Spend
|
||||
* @default 0
|
||||
*/
|
||||
total_spend: number;
|
||||
/** Tpd Limit */
|
||||
tpd_limit?: number | null;
|
||||
/** Tpm Limit */
|
||||
|
|
@ -31259,6 +31264,11 @@ export interface components {
|
|||
team_id?: string | null;
|
||||
/** Token */
|
||||
token?: string | null;
|
||||
/**
|
||||
* Total Spend
|
||||
* @default 0
|
||||
*/
|
||||
total_spend: number;
|
||||
/** Tpd Limit */
|
||||
tpd_limit?: number | null;
|
||||
/** Tpm Limit */
|
||||
|
|
@ -40057,6 +40067,11 @@ export interface components {
|
|||
team_tpm_limit?: number | null;
|
||||
/** Token */
|
||||
token?: string | null;
|
||||
/**
|
||||
* Total Spend
|
||||
* @default 0
|
||||
*/
|
||||
total_spend: number;
|
||||
/** Tpd Limit */
|
||||
tpd_limit?: number | null;
|
||||
/** Tpm Limit */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue