scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 82

Merge pull request #17781 from BerriAI/litellm_daily_agent_table

Browse source 
[Feature] Daily Agent Usage Backend
This commit is contained in:
yuneng-jiang 2025-12-11 15:01:50 -08:00 • committed by GitHub
commit de9a195d8c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 700 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.13-py3-none-any.whl vendored Normal file
View file

Binary file not shown.

BIN
litellm-proxy-extras/dist/litellm_proxy_extras-0.4.13.tar.gz vendored Normal file
View file

Binary file not shown.

45
litellm-proxy-extras/litellm_proxy_extras/migrations/20251210205007_add_daily_agent_spend_table/migration.sql Normal file
View file

@ -0,0 +1,45 @@
-- AlterTable
ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "agent_id" TEXT;
-- CreateTable
CREATE TABLE "LiteLLM_DailyAgentSpend" (
"id" TEXT NOT NULL,
"agent_id" TEXT,
"date" TEXT NOT NULL,
"api_key" TEXT NOT NULL,
"model" TEXT,
"model_group" TEXT,
"custom_llm_provider" TEXT,
"mcp_namespaced_tool_name" TEXT,
"prompt_tokens" BIGINT NOT NULL DEFAULT 0,
"completion_tokens" BIGINT NOT NULL DEFAULT 0,
"cache_read_input_tokens" BIGINT NOT NULL DEFAULT 0,
"cache_creation_input_tokens" BIGINT NOT NULL DEFAULT 0,
"spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0,
"api_requests" BIGINT NOT NULL DEFAULT 0,
"successful_requests" BIGINT NOT NULL DEFAULT 0,
"failed_requests" BIGINT NOT NULL DEFAULT 0,
"created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP,
"updated_at" TIMESTAMP(3) NOT NULL,
CONSTRAINT "LiteLLM_DailyAgentSpend_pkey" PRIMARY KEY ("id")
);
-- CreateIndex
CREATE INDEX "LiteLLM_DailyAgentSpend_date_idx" ON "LiteLLM_DailyAgentSpend"("date");
-- CreateIndex
CREATE INDEX "LiteLLM_DailyAgentSpend_agent_id_idx" ON "LiteLLM_DailyAgentSpend"("agent_id");
-- CreateIndex
CREATE INDEX "LiteLLM_DailyAgentSpend_api_key_idx" ON "LiteLLM_DailyAgentSpend"("api_key");
-- CreateIndex
CREATE INDEX "LiteLLM_DailyAgentSpend_model_idx" ON "LiteLLM_DailyAgentSpend"("model");
-- CreateIndex
CREATE INDEX "LiteLLM_DailyAgentSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyAgentSpend"("mcp_namespaced_tool_name");
-- CreateIndex
CREATE UNIQUE INDEX "LiteLLM_DailyAgentSpend_agent_id_date_api_key_model_custom__key" ON "LiteLLM_DailyAgentSpend"("agent_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name");

28
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -494,6 +494,34 @@ model LiteLLM_DailyEndUserSpend {
@@index([mcp_namespaced_tool_name])
}
// Track daily agent spend metrics per model and key
model LiteLLM_DailyAgentSpend {
id String @id @default(uuid())
agent_id String?
date String
api_key String
model String?
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@index([date])
@@index([agent_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
}
// Track daily team spend metrics per model and key
model LiteLLM_DailyTeamSpend {
id String @id @default(uuid())

4
litellm-proxy-extras/pyproject.toml
View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-proxy-extras"
version = "0.4.12"
version = "0.4.13"
description = "Additional files for the LiteLLM Proxy. Reduces the size of the main litellm package."
authors = ["BerriAI"]
readme = "README.md"
@ -22,7 +22,7 @@ requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.commitizen]
version = "0.4.12"
version = "0.4.13"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-proxy-extras==",

1
litellm/constants.py
View file

@ -150,6 +150,7 @@ REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_spend_update_buffer"
REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_team_spend_update_buffer"
REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_org_spend_update_buffer"
REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_end_user_spend_update_buffer"
REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_agent_spend_update_buffer"
REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer"
MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100))
MAX_SIZE_IN_MEMORY_QUEUE = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", 2000))

3
litellm/proxy/_types.py
View file

@ -3654,6 +3654,9 @@ class DailyTagSpendTransaction(BaseDailySpendTransaction):
request_id: Optional[str]
tag: str
class DailyAgentSpendTransaction(BaseDailySpendTransaction):
agent_id: str
class DBSpendUpdateTransactions(TypedDict):
"""

68
litellm/proxy/agent_endpoints/endpoints.py
View file

@ -8,7 +8,7 @@ Follows the A2A Spec.
3. Get specific agent via GET `/v1/agents/{agent_id}`
"""
from typing import Any, List
from typing import Any, List, Optional
from fastapi import APIRouter, Depends, HTTPException, Request
@ -24,6 +24,11 @@ from litellm.types.agents import (
PatchAgentRequest,
)
from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity
from litellm.types.proxy.management_endpoints.common_daily_activity import (
SpendAnalyticsPaginatedResponse,
)
router = APIRouter()
@ -703,3 +708,64 @@ async def make_agents_public(
except Exception as e:
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"],
dependencies=[Depends(user_api_key_auth)],
response_model=SpendAnalyticsPaginatedResponse,
)
async def get_agent_daily_activity(
agent_ids: Optional[str] = None,
start_date: Optional[str] = None,
end_date: Optional[str] = None,
model: Optional[str] = None,
api_key: Optional[str] = None,
page: int = 1,
page_size: int = 10,
exclude_agent_ids: Optional[str] = None,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Get daily activity for specific agents or all accessible agents.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
agent_ids_list = agent_ids.split(",") if agent_ids else None
exclude_agent_ids_list: Optional[List[str]] = None
if exclude_agent_ids:
exclude_agent_ids_list = (
exclude_agent_ids.split(",") if exclude_agent_ids else None
)
where_condition = {}
if agent_ids_list:
where_condition["agent_id"] = {"in": list(agent_ids_list)}
agent_records = await prisma_client.db.litellm_agentstable.find_many(
where=where_condition
)
agent_metadata = {
agent.agent_id: {"agent_name": agent.agent_name} for agent in agent_records
}
return await get_daily_activity(
prisma_client=prisma_client,
table_name="litellm_dailyagentspend",
entity_id_field="agent_id",
entity_id=agent_ids_list,
entity_metadata_field=agent_metadata,
exclude_entity_ids=exclude_agent_ids_list,
start_date=start_date,
end_date=end_date,
model=model,
api_key=api_key,
page=page,
page_size=page_size,
)

122
litellm/proxy/db/db_spend_update_writer.py
View file

@ -28,6 +28,7 @@ from litellm.proxy._types import (
DailyTeamSpendTransaction,
DailyEndUserSpendTransaction,
DailyUserSpendTransaction,
DailyAgentSpendTransaction,
DBSpendUpdateTransactions,
Litellm_EntityType,
LiteLLM_UserTable,
@ -68,6 +69,7 @@ class DBSpendUpdateWriter:
self.daily_spend_update_queue = DailySpendUpdateQueue()
self.daily_team_spend_update_queue = DailySpendUpdateQueue()
self.daily_end_user_spend_update_queue = DailySpendUpdateQueue()
self.daily_agent_spend_update_queue = DailySpendUpdateQueue()
self.daily_org_spend_update_queue = DailySpendUpdateQueue()
self.daily_tag_spend_update_queue = DailySpendUpdateQueue()
@ -192,6 +194,13 @@ class DBSpendUpdateWriter:
)
)
asyncio.create_task(
self.add_spend_log_transaction_to_daily_agent_transaction(
payload=payload,
prisma_client=prisma_client,
)
)
asyncio.create_task(
self.add_spend_log_transaction_to_daily_team_transaction(
payload=copy.deepcopy(payload),
@ -488,6 +497,7 @@ class DBSpendUpdateWriter:
daily_team_spend_update_queue=self.daily_team_spend_update_queue,
daily_org_spend_update_queue=self.daily_org_spend_update_queue,
daily_end_user_spend_update_queue=self.daily_end_user_spend_update_queue,
daily_agent_spend_update_queue=self.daily_agent_spend_update_queue,
daily_tag_spend_update_queue=self.daily_tag_spend_update_queue,
)
@ -561,6 +571,16 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_end_user_spend_update_transactions,
)
daily_agent_spend_update_transactions = (
await self.redis_update_buffer.get_all_daily_agent_spend_update_transactions_from_redis_buffer()
)
if daily_agent_spend_update_transactions is not None:
await DBSpendUpdateWriter.update_daily_agent_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_agent_spend_update_transactions,
)
except Exception as e:
verbose_proxy_logger.error(f"Error committing spend updates: {e}")
finally:
@ -664,6 +684,20 @@ class DBSpendUpdateWriter:
daily_spend_transactions=daily_end_user_spend_update_transactions,
)
################## Daily Agent Spend Update Transactions ##################
# Aggregate all in memory daily agent spend transactions and commit to db
daily_agent_spend_update_transactions = cast(
Dict[str, DailyAgentSpendTransaction],
await self.daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(),
)
await DBSpendUpdateWriter.update_daily_agent_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_agent_spend_update_transactions,
)
async def _commit_spend_updates_to_db( # noqa: PLR0915
self,
prisma_client: PrismaClient,
@ -1041,6 +1075,20 @@ class DBSpendUpdateWriter:
) -> None:
...
@overload
@staticmethod
async def _update_daily_spend(
n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
daily_spend_transactions: Dict[str, DailyAgentSpendTransaction],
entity_type: Literal["agent"],
entity_id_field: str,
table_name: str,
unique_constraint_name: str,
) -> None:
...
@overload
@staticmethod
async def _update_daily_spend(
@ -1067,14 +1115,15 @@ class DBSpendUpdateWriter:
Dict[str, DailyTagSpendTransaction],
Dict[str, DailyOrganizationSpendTransaction],
Dict[str, DailyEndUserSpendTransaction],
Dict[str, DailyAgentSpendTransaction],
],
entity_type: Literal["user", "team", "org", "tag", "end_user"],
entity_type: Literal["user", "team", "org", "tag", "end_user", "agent"],
entity_id_field: str,
table_name: str,
unique_constraint_name: str,
) -> None:
"""
Generic function to update daily spend for any entity type (user, team, org, tag, end_user)
Generic function to update daily spend for any entity type (user, team, org, tag, end_user, agent)
"""
from litellm.proxy.utils import _raise_failed_update_spend_exception
@ -1340,6 +1389,27 @@ class DBSpendUpdateWriter:
unique_constraint_name="end_user_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name",
)
@staticmethod
async def update_daily_agent_spend(
n_retry_times: int,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
daily_spend_transactions: Dict[str, DailyAgentSpendTransaction],
):
"""
Batch job to update LiteLLM_DailyAgentSpend table using in-memory daily_spend_transactions
"""
await DBSpendUpdateWriter._update_daily_spend(
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
daily_spend_transactions=daily_spend_transactions,
entity_type="agent",
entity_id_field="agent_id",
table_name="litellm_dailyagentspend",
unique_constraint_name="agent_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name",
)
@staticmethod
async def update_daily_tag_spend(
n_retry_times: int,
@ -1365,7 +1435,7 @@ class DBSpendUpdateWriter:
self,
payload: Union[dict, SpendLogsPayload],
prisma_client: PrismaClient,
type: Literal["user", "team", "org", "request_tags", "end_user"] = "user",
type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user",
) -> Optional[BaseDailySpendTransaction]:
common_expected_keys = ["startTime", "api_key"]
if type == "user":
@ -1378,6 +1448,8 @@ class DBSpendUpdateWriter:
expected_keys = ["request_tags", *common_expected_keys]
elif type == "end_user":
expected_keys = ["end_user_id", *common_expected_keys]
elif type == "agent":
expected_keys = ["agent_id", *common_expected_keys]
else:
raise ValueError(f"Invalid type: {type}")
if not all(key in payload for key in expected_keys):
@ -1591,6 +1663,50 @@ class DBSpendUpdateWriter:
update={daily_transaction_key: daily_transaction}
)
async def add_spend_log_transaction_to_daily_agent_transaction(
self,
payload: SpendLogsPayload,
prisma_client: Optional[PrismaClient] = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug(
"prisma_client is None. Skipping writing spend logs to db."
)
return
base_daily_transaction = (
await self._common_add_spend_log_transaction_to_daily_transaction(
payload, prisma_client, "agent"
)
)
if base_daily_transaction is None:
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,
{
**payload,
"agent_id": payload["agent_id"],
},
)
base_daily_transaction = (
await self._common_add_spend_log_transaction_to_daily_transaction(
payload_with_agent_id, prisma_client, "agent"
)
)
if base_daily_transaction is None:
return
daily_transaction_key = f"{payload['agent_id']}_{base_daily_transaction['date']}_{payload_with_agent_id['api_key']}_{payload_with_agent_id['model']}_{payload_with_agent_id['custom_llm_provider']}"
daily_transaction = DailyAgentSpendTransaction(
agent_id=payload['agent_id'], **base_daily_transaction
)
await self.daily_agent_spend_update_queue.add_update(
update={daily_transaction_key: daily_transaction}
)
async def add_spend_log_transaction_to_daily_tag_transaction(
self,
payload: SpendLogsPayload,

36
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -17,6 +17,7 @@ from litellm.constants import (
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.litellm_core_utils.safe_json_dumps import safe_dumps
@ -27,6 +28,7 @@ from litellm.proxy._types import (
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 (
@ -110,6 +112,7 @@ class RedisUpdateBuffer:
daily_team_spend_update_queue: DailySpendUpdateQueue,
daily_org_spend_update_queue: DailySpendUpdateQueue,
daily_end_user_spend_update_queue: DailySpendUpdateQueue,
daily_agent_spend_update_queue: DailySpendUpdateQueue,
daily_tag_spend_update_queue: DailySpendUpdateQueue,
):
"""
@ -178,6 +181,9 @@ class RedisUpdateBuffer:
daily_end_user_spend_update_transactions = (
await daily_end_user_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
daily_agent_spend_update_transactions = (
await daily_agent_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
daily_tag_spend_update_transactions = (
await daily_tag_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions()
)
@ -219,6 +225,12 @@ class RedisUpdateBuffer:
service_type=ServiceTypes.REDIS_DAILY_END_USER_SPEND_UPDATE_QUEUE,
)
await self._store_transactions_in_redis(
transactions=daily_agent_spend_update_transactions,
redis_key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
service_type=ServiceTypes.REDIS_DAILY_AGENT_SPEND_UPDATE_QUEUE,
)
await self._store_transactions_in_redis(
transactions=daily_tag_spend_update_transactions,
redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY,
@ -401,6 +413,30 @@ class RedisUpdateBuffer:
),
)
async def get_all_daily_agent_spend_update_transactions_from_redis_buffer(
self,
) -> Optional[Dict[str, DailyAgentSpendTransaction]]:
"""
Gets all the daily agent spend update transactions from Redis
"""
if self.redis_cache is None:
return None
list_of_transactions = await self.redis_cache.async_lpop(
key=REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY,
count=MAX_REDIS_BUFFER_DEQUEUE_COUNT,
)
if list_of_transactions is None:
return None
list_of_daily_spend_update_transactions = [
json.loads(transaction) for transaction in list_of_transactions
]
return cast(
Dict[str, DailyAgentSpendTransaction],
DailySpendUpdateQueue.get_aggregated_daily_spend_update_transactions(
list_of_daily_spend_update_transactions
),
)
async def get_all_daily_tag_spend_update_transactions_from_redis_buffer(
self,
) -> Optional[Dict[str, DailyTagSpendTransaction]]:

28
litellm/proxy/schema.prisma
View file

@ -494,6 +494,34 @@ model LiteLLM_DailyEndUserSpend {
@@index([mcp_namespaced_tool_name])
}
// Track daily agent spend metrics per model and key
model LiteLLM_DailyAgentSpend {
id String @id @default(uuid())
agent_id String?
date String
api_key String
model String?
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@index([date])
@@index([agent_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
}
// Track daily team spend metrics per model and key
model LiteLLM_DailyTeamSpend {
id String @id @default(uuid())

2
pyproject.toml
View file

@ -59,7 +59,7 @@ websockets = {version = "^15.0.1", optional = true}
boto3 = {version = "1.36.0", optional = true}
redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.9' and python_version < '3.14'"}
mcp = {version = "^1.21.2", optional = true, python = ">=3.10"}
litellm-proxy-extras = {version = "0.4.12", optional = true}
litellm-proxy-extras = {version = "0.4.13", optional = true}
rich = {version = "13.7.1", optional = true}
litellm-enterprise = {version = "0.1.25", optional = true}
diskcache = {version = "^5.6.1", optional = true}

2
requirements.txt
View file

@ -44,7 +44,7 @@ sentry_sdk==2.21.0 # for sentry error handling
detect-secrets==1.5.0 # Enterprise - secret detection / masking in LLM requests
cryptography==44.0.1
tzdata==2025.1 # IANA time zone database
litellm-proxy-extras==0.4.12 # for proxy extras - e.g. prisma migrations
litellm-proxy-extras==0.4.13 # for proxy extras - e.g. prisma migrations
### LITELLM PACKAGE DEPENDENCIES
python-dotenv==1.0.1 # for env
tiktoken==0.8.0 # for calculating usage

28
schema.prisma
View file

@ -494,6 +494,34 @@ model LiteLLM_DailyEndUserSpend {
@@index([mcp_namespaced_tool_name])
}
// Track daily agent spend metrics per model and key
model LiteLLM_DailyAgentSpend {
id String @id @default(uuid())
agent_id String?
date String
api_key String
model String?
model_group String?
custom_llm_provider String?
mcp_namespaced_tool_name String?
prompt_tokens BigInt @default(0)
completion_tokens BigInt @default(0)
cache_read_input_tokens BigInt @default(0)
cache_creation_input_tokens BigInt @default(0)
spend Float @default(0.0)
api_requests BigInt @default(0)
successful_requests BigInt @default(0)
failed_requests BigInt @default(0)
created_at DateTime @default(now())
updated_at DateTime @updatedAt
@@unique([agent_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name])
@@index([date])
@@index([agent_id])
@@index([api_key])
@@index([model])
@@index([mcp_namespaced_tool_name])
}
// Track daily team spend metrics per model and key
model LiteLLM_DailyTeamSpend {
id String @id @default(uuid())

260
tests/test_litellm/proxy/agent_endpoints/test_endpoints.py Normal file
View file

@ -0,0 +1,260 @@
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.agent_endpoints import endpoints as agent_endpoints
from litellm.proxy.agent_endpoints.endpoints import (
get_agent_daily_activity,
router,
user_api_key_auth,
)
from litellm.types.agents import AgentResponse
def _sample_agent_card_params() -> dict:
return {
"protocolVersion": "1.0",
"name": "Test Agent",
"description": "desc",
"url": "http://localhost",
"version": "1.0.0",
"capabilities": {"streaming": True},
"defaultInputModes": ["text"],
"defaultOutputModes": ["text"],
"skills": [],
}
def _sample_agent_config() -> dict:
return {
"agent_name": "Test Agent",
"agent_card_params": _sample_agent_card_params(),
"litellm_params": {"make_public": False},
}
def _sample_agent_response(
agent_id: str = "agent-123", agent_name: str = "Test Agent"
) -> AgentResponse:
return AgentResponse(
agent_id=agent_id,
agent_name=agent_name,
agent_card_params=_sample_agent_card_params(),
litellm_params={"make_public": False},
)
app = FastAPI()
app.include_router(router)
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(
user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN
)
client = TestClient(app)
@pytest.fixture
def mock_prisma_client():
with patch("litellm.proxy.proxy_server.prisma_client") as mock:
yield mock
@pytest.fixture
def mock_user_api_key_auth():
with patch("litellm.proxy.agent_endpoints.endpoints.user_api_key_auth") as mock:
mock.return_value = UserAPIKeyAuth(
user_id="test-user", user_role=LitellmUserRoles.PROXY_ADMIN
)
yield mock
def test_update_agent_success(mock_prisma_client, mock_user_api_key_auth, monkeypatch):
existing_agent = {
"agent_id": "agent-123",
"agent_name": "Existing Agent",
"agent_card_params": _sample_agent_card_params(),
}
mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(
return_value=existing_agent
)
mock_registry = MagicMock()
mock_registry.update_agent_in_db = AsyncMock(
return_value=_sample_agent_response(agent_id="agent-123")
)
mock_registry.deregister_agent = MagicMock()
mock_registry.register_agent = MagicMock()
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)
response = client.put(
"/v1/agents/agent-123",
json=_sample_agent_config(),
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 200
assert response.json()["agent_id"] == "agent-123"
assert response.json()["agent_name"] == "Test Agent"
def test_update_agent_not_found(
mock_prisma_client, mock_user_api_key_auth, monkeypatch
):
mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
mock_registry = MagicMock()
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)
response = client.put(
"/v1/agents/missing-agent",
json=_sample_agent_config(),
headers={"Authorization": "Bearer test-key"},
)
assert response.status_code == 404
assert "Agent with ID missing-agent not found" in response.json()["detail"]
def test_get_agent_by_id_not_found(
mock_prisma_client, mock_user_api_key_auth, monkeypatch
):
mock_registry = MagicMock()
mock_registry.get_agent_by_id = MagicMock(return_value=None)
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)
mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
response = client.get(
"/v1/agents/missing-agent", headers={"Authorization": "Bearer test-key"}
)
assert response.status_code == 404
assert "Agent with ID missing-agent not found" in response.json()["detail"]
def test_delete_agent_not_found(
mock_prisma_client, mock_user_api_key_auth, monkeypatch
):
mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
mock_registry = MagicMock()
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)
response = client.delete(
"/v1/agents/missing-agent", headers={"Authorization": "Bearer test-key"}
)
assert response.status_code == 404
assert "Agent with ID missing-agent not found in DB." in response.json()["detail"]
def test_agent_error_schema_consistency(
mock_prisma_client, mock_user_api_key_auth, monkeypatch
):
mock_registry = MagicMock()
mock_registry.get_agent_by_id = MagicMock(return_value=None)
mock_registry.update_agent_in_db = AsyncMock(side_effect=Exception("should not run"))
mock_registry.delete_agent_from_db = AsyncMock(side_effect=Exception("should not run"))
monkeypatch.setattr(agent_endpoints, "AGENT_REGISTRY", mock_registry)
mock_prisma_client.db.litellm_agentstable.find_unique = AsyncMock(return_value=None)
missing_agent_id = "missing-agent"
responses = [
client.get(
f"/v1/agents/{missing_agent_id}",
headers={"Authorization": "Bearer test-key"},
),
client.put(
f"/v1/agents/{missing_agent_id}",
json=_sample_agent_config(),
headers={"Authorization": "Bearer test-key"},
),
client.delete(
f"/v1/agents/{missing_agent_id}",
headers={"Authorization": "Bearer test-key"},
),
]
for resp in responses:
assert resp.status_code == 404
detail = resp.json()["detail"]
assert isinstance(detail, str)
assert missing_agent_id in detail
@pytest.mark.asyncio
async def test_get_agent_daily_activity_admin_param_passing(monkeypatch):
mock_prisma = AsyncMock()
mock_prisma.db.litellm_agentstable.find_many = AsyncMock(return_value=[])
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse")
get_daily_activity_mock = AsyncMock(return_value=mocked_response)
monkeypatch.setattr(agent_endpoints, "get_daily_activity", get_daily_activity_mock)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1")
result = await get_agent_daily_activity(
agent_ids="agent-1,agent-2",
start_date="2024-01-01",
end_date="2024-01-31",
model="gpt-4",
api_key="test-key",
page=2,
page_size=5,
exclude_agent_ids="agent-3",
user_api_key_dict=auth,
)
get_daily_activity_mock.assert_awaited_once()
kwargs = get_daily_activity_mock.call_args.kwargs
assert kwargs["table_name"] == "litellm_dailyagentspend"
assert kwargs["entity_id_field"] == "agent_id"
assert kwargs["entity_id"] == ["agent-1", "agent-2"]
assert kwargs["exclude_entity_ids"] == ["agent-3"]
assert kwargs["start_date"] == "2024-01-01"
assert kwargs["end_date"] == "2024-01-31"
assert kwargs["model"] == "gpt-4"
assert kwargs["api_key"] == "test-key"
assert kwargs["page"] == 2
assert kwargs["page_size"] == 5
assert result is mocked_response
@pytest.mark.asyncio
async def test_get_agent_daily_activity_with_agent_names(monkeypatch):
mock_prisma = AsyncMock()
mock_agent1 = MagicMock()
mock_agent1.agent_id = "agent-1"
mock_agent1.agent_name = "First Agent"
mock_agent2 = MagicMock()
mock_agent2.agent_id = "agent-2"
mock_agent2.agent_name = "Second Agent"
mock_prisma.db.litellm_agentstable.find_many = AsyncMock(
return_value=[mock_agent1, mock_agent2]
)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma)
mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse")
get_daily_activity_mock = AsyncMock(return_value=mocked_response)
monkeypatch.setattr(agent_endpoints, "get_daily_activity", get_daily_activity_mock)
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1")
await get_agent_daily_activity(
agent_ids="agent-1,agent-2",
start_date="2024-01-01",
end_date="2024-01-31",
model=None,
api_key=None,
page=1,
page_size=10,
exclude_agent_ids=None,
user_api_key_dict=auth,
)
kwargs = get_daily_activity_mock.call_args.kwargs
assert kwargs["entity_metadata_field"] == {
"agent-1": {"agent_name": "First Agent"},
"agent-2": {"agent_name": "Second Agent"},
}

82
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -645,4 +645,84 @@ async def test_add_spend_log_transaction_to_daily_end_user_transaction_skips_whe
prisma_client=mock_prisma,
)
writer.daily_end_user_spend_update_queue.add_update.assert_not_called()
writer.daily_end_user_spend_update_queue.add_update.assert_not_called()
@pytest.mark.asyncio
async def test_add_spend_log_transaction_to_daily_agent_transaction_injects_agent_id_and_queues_update():
"""
Ensure agent_id is injected and queued for daily aggregation.
"""
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
mock_prisma.get_request_status = MagicMock(return_value="success")
agent_id = "agent-123"
payload = {
"request_id": "req-123",
"agent_id": agent_id,
"user": "test-user",
"startTime": "2024-01-01T12:00:00",
"api_key": "test-key",
"model": "gpt-4",
"custom_llm_provider": "openai",
"model_group": "gpt-4-group",
"prompt_tokens": 20,
"completion_tokens": 10,
"spend": 0.3,
"metadata": '{"usage_object": {}}',
}
writer.daily_agent_spend_update_queue.add_update = AsyncMock()
await writer.add_spend_log_transaction_to_daily_agent_transaction(
payload=payload,
prisma_client=mock_prisma,
)
writer.daily_agent_spend_update_queue.add_update.assert_called_once()
call_args = writer.daily_agent_spend_update_queue.add_update.call_args[1]
update_dict = call_args["update"]
assert len(update_dict) == 1
for key, transaction in update_dict.items():
assert key == f"{agent_id}_2024-01-01_test-key_gpt-4_openai"
assert transaction["agent_id"] == agent_id
assert transaction["date"] == "2024-01-01"
assert transaction["api_key"] == "test-key"
assert transaction["model"] == "gpt-4"
assert transaction["custom_llm_provider"] == "openai"
@pytest.mark.asyncio
async def test_add_spend_log_transaction_to_daily_agent_transaction_skips_when_agent_id_missing():
"""
Do not queue agent spend updates when agent_id is None.
"""
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
mock_prisma.get_request_status = MagicMock(return_value="success")
payload = {
"request_id": "req-456",
"agent_id": None,
"user": "test-user",
"startTime": "2024-01-01T12:00:00",
"api_key": "test-key",
"model": "gpt-4",
"custom_llm_provider": "openai",
"model_group": "gpt-4-group",
"prompt_tokens": 15,
"completion_tokens": 5,
"spend": 0.1,
"metadata": '{"usage_object": {}}',
}
writer.daily_agent_spend_update_queue.add_update = AsyncMock()
await writer.add_spend_log_transaction_to_daily_agent_transaction(
payload=payload,
prisma_client=mock_prisma,
)
writer.daily_agent_spend_update_queue.add_update.assert_not_called()