diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.7-py3-none-any.whl b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.7-py3-none-any.whl new file mode 100644 index 00000000000..376c1e0d070 Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.7-py3-none-any.whl differ diff --git a/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.7.tar.gz b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.7.tar.gz new file mode 100644 index 00000000000..0bb0fd9c74a Binary files /dev/null and b/litellm-proxy-extras/dist/litellm_proxy_extras-0.4.7.tar.gz differ diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114180624_Add_org_usage_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114180624_Add_org_usage_table/migration.sql new file mode 100644 index 00000000000..74e0eea3134 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251114180624_Add_org_usage_table/migration.sql @@ -0,0 +1,42 @@ +-- CreateTable +CREATE TABLE "LiteLLM_DailyOrganizationSpend" ( + "id" TEXT NOT NULL, + "organization_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_DailyOrganizationSpend_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyOrganizationSpend_date_idx" ON "LiteLLM_DailyOrganizationSpend"("date"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyOrganizationSpend_organization_id_idx" ON "LiteLLM_DailyOrganizationSpend"("organization_id"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyOrganizationSpend_api_key_idx" ON "LiteLLM_DailyOrganizationSpend"("api_key"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyOrganizationSpend_model_idx" ON "LiteLLM_DailyOrganizationSpend"("model"); + +-- CreateIndex +CREATE INDEX "LiteLLM_DailyOrganizationSpend_mcp_namespaced_tool_name_idx" ON "LiteLLM_DailyOrganizationSpend"("mcp_namespaced_tool_name"); + +-- CreateIndex +CREATE UNIQUE INDEX "LiteLLM_DailyOrganizationSpend_organization_id_date_api_key_key" ON "LiteLLM_DailyOrganizationSpend"("organization_id", "date", "api_key", "model", "custom_llm_provider", "mcp_namespaced_tool_name"); + diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20251122125322_Add organization_id to spend logs/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251122125322_Add organization_id to spend logs/migration.sql new file mode 100644 index 00000000000..4ea082f2750 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20251122125322_Add organization_id to spend logs/migration.sql @@ -0,0 +1,3 @@ +-- AlterTable +ALTER TABLE "LiteLLM_SpendLogs" ADD COLUMN "organization_id" TEXT; + diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 6cfbb90c362..2883dfc4b82 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -304,6 +304,7 @@ model LiteLLM_SpendLogs { cache_key String? @default("") request_tags Json? @default("[]") team_id String? + organization_id String? end_user String? requester_ip_address String? messages Json? @default("{}") @@ -432,6 +433,35 @@ model LiteLLM_DailyUserSpend { @@index([mcp_namespaced_tool_name]) } +// Track daily organization spend metrics per model and key +model LiteLLM_DailyOrganizationSpend { + id String @id @default(uuid()) + organization_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([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name]) + @@index([date]) + @@index([organization_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()) diff --git a/litellm/constants.py b/litellm/constants.py index cf3d4c6e742..9235916dd43 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -139,6 +139,7 @@ DEFAULT_SSL_CIPHERS = os.getenv( REDIS_UPDATE_BUFFER_KEY = "litellm_spend_update_buffer" 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_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", 10000)) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fd2b12e063e..d05454c5784 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2647,6 +2647,7 @@ class SpendLogsPayload(TypedDict): cache_key: str request_tags: str # json str team_id: Optional[str] + organization_id: Optional[str] end_user: Optional[str] requester_ip_address: Optional[str] custom_llm_provider: Optional[str] @@ -3583,6 +3584,10 @@ class DailyTeamSpendTransaction(BaseDailySpendTransaction): team_id: str +class DailyOrganizationSpendTransaction(BaseDailySpendTransaction): + organization_id: str + + class DailyUserSpendTransaction(BaseDailySpendTransaction): user_id: str diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 6c4f3bdd458..6c9289e3ff6 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -23,6 +23,7 @@ from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, BaseDailySpendTransaction, DailyTagSpendTransaction, + DailyOrganizationSpendTransaction, DailyTeamSpendTransaction, DailyUserSpendTransaction, DBSpendUpdateTransactions, @@ -64,6 +65,7 @@ class DBSpendUpdateWriter: self.spend_update_queue = SpendUpdateQueue() self.daily_spend_update_queue = DailySpendUpdateQueue() self.daily_team_spend_update_queue = DailySpendUpdateQueue() + self.daily_org_spend_update_queue = DailySpendUpdateQueue() self.daily_tag_spend_update_queue = DailySpendUpdateQueue() async def update_database( @@ -116,6 +118,12 @@ class DBSpendUpdateWriter: payload["startTime"] = payload["startTime"].isoformat() if isinstance(payload["endTime"], datetime): payload["endTime"] = payload["endTime"].isoformat() + + if org_id is not None and org_id != "": + payload["organization_id"] = org_id + + if team_id is not None and team_id != "": + payload["team_id"] = team_id asyncio.create_task( self._update_user_db( @@ -180,7 +188,13 @@ class DBSpendUpdateWriter: prisma_client=prisma_client, ) ) - + asyncio.create_task( + self.add_spend_log_transaction_to_daily_org_transaction( + payload=payload, + org_id=org_id, + prisma_client=prisma_client, + ) + ) asyncio.create_task( self.add_spend_log_transaction_to_daily_tag_transaction( payload=payload, @@ -460,6 +474,7 @@ class DBSpendUpdateWriter: spend_update_queue=self.spend_update_queue, daily_spend_update_queue=self.daily_spend_update_queue, daily_team_spend_update_queue=self.daily_team_spend_update_queue, + daily_org_spend_update_queue=self.daily_org_spend_update_queue, daily_tag_spend_update_queue=self.daily_tag_spend_update_queue, ) @@ -502,6 +517,17 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_team_spend_update_transactions, ) + daily_org_spend_update_transactions = ( + await self.redis_update_buffer.get_all_daily_org_spend_update_transactions_from_redis_buffer() + ) + if daily_org_spend_update_transactions is not None: + await DBSpendUpdateWriter.update_daily_org_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_org_spend_update_transactions, + ) + daily_tag_spend_update_transactions = ( await self.redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer() ) @@ -573,6 +599,20 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_team_spend_update_transactions, ) + ################## Daily Organization Spend Update Transactions ################## + # Aggregate all in memory daily org spend transactions and commit to db + daily_org_spend_update_transactions = cast( + Dict[str, DailyOrganizationSpendTransaction], + await self.daily_org_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions(), + ) + + await DBSpendUpdateWriter.update_daily_org_spend( + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + daily_spend_transactions=daily_org_spend_update_transactions, + ) + ################## Daily Tag Spend Update Transactions ################## # Aggregate all in memory daily tag spend transactions and commit to db daily_tag_spend_update_transactions = cast( @@ -936,6 +976,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, DailyOrganizationSpendTransaction], + entity_type: Literal["org"], + entity_id_field: str, + table_name: str, + unique_constraint_name: str, + ) -> None: + ... + @overload @staticmethod async def _update_daily_spend( @@ -960,14 +1014,15 @@ class DBSpendUpdateWriter: Dict[str, DailyUserSpendTransaction], Dict[str, DailyTeamSpendTransaction], Dict[str, DailyTagSpendTransaction], + Dict[str, DailyOrganizationSpendTransaction], ], - entity_type: Literal["user", "team", "tag"], + entity_type: Literal["user", "team", "org", "tag"], entity_id_field: str, table_name: str, unique_constraint_name: str, ) -> None: """ - Generic function to update daily spend for any entity type (user, team, tag) + Generic function to update daily spend for any entity type (user, team, org, tag) """ from litellm.proxy.utils import _raise_failed_update_spend_exception @@ -1191,6 +1246,27 @@ class DBSpendUpdateWriter: unique_constraint_name="team_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name", ) + @staticmethod + async def update_daily_org_spend( + n_retry_times: int, + prisma_client: PrismaClient, + proxy_logging_obj: ProxyLogging, + daily_spend_transactions: Dict[str, DailyOrganizationSpendTransaction], + ): + """ + Batch job to update LiteLLM_DailyOrganizationSpend 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="org", + entity_id_field="organization_id", + table_name="litellm_dailyorganizationspend", + unique_constraint_name="organization_id_date_api_key_model_custom_llm_provider_mcp_namespaced_tool_name", + ) + @staticmethod async def update_daily_tag_spend( n_retry_times: int, @@ -1216,13 +1292,15 @@ class DBSpendUpdateWriter: self, payload: Union[dict, SpendLogsPayload], prisma_client: PrismaClient, - type: Literal["user", "team", "request_tags"] = "user", + type: Literal["user", "team", "org", "request_tags"] = "user", ) -> Optional[BaseDailySpendTransaction]: common_expected_keys = ["startTime", "api_key"] if type == "user": expected_keys = ["user", *common_expected_keys] elif type == "team": expected_keys = ["team_id", *common_expected_keys] + elif type == "org": + expected_keys = ["organization_id", *common_expected_keys] elif type == "request_tags": expected_keys = ["request_tags", *common_expected_keys] else: @@ -1354,6 +1432,48 @@ class DBSpendUpdateWriter: update={daily_transaction_key: daily_transaction} ) + async def add_spend_log_transaction_to_daily_org_transaction( + self, + payload: SpendLogsPayload, + prisma_client: Optional[PrismaClient] = None, + org_id: Optional[str] = None, + ) -> None: + if prisma_client is None: + verbose_proxy_logger.debug( + "prisma_client is None. Skipping writing spend logs to db." + ) + return + + if org_id is None: + verbose_proxy_logger.debug( + "organization_id is None for request. Skipping incrementing organization spend." + ) + return + + payload_with_org = cast( + SpendLogsPayload, + { + **payload, + "organization_id": org_id, + }, + ) + + base_daily_transaction = ( + await self._common_add_spend_log_transaction_to_daily_transaction( + payload_with_org, prisma_client, "org" + ) + ) + if base_daily_transaction is None: + return + + daily_transaction_key = f"{org_id}_{base_daily_transaction['date']}_{payload_with_org['api_key']}_{payload_with_org['model']}_{payload_with_org['custom_llm_provider']}" + daily_transaction = DailyOrganizationSpendTransaction( + organization_id=org_id, **base_daily_transaction + ) + await self.daily_org_spend_update_queue.add_update( + update={daily_transaction_key: daily_transaction} + ) + async def add_spend_log_transaction_to_daily_tag_transaction( self, payload: 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 91d0bee1d3a..921fd9701bd 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -15,6 +15,7 @@ from litellm.constants import ( 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_UPDATE_BUFFER_KEY, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -22,6 +23,7 @@ from litellm.proxy._types import ( DailyTagSpendTransaction, DailyTeamSpendTransaction, DailyUserSpendTransaction, + DailyOrganizationSpendTransaction, DBSpendUpdateTransactions, ) from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj @@ -104,6 +106,7 @@ class RedisUpdateBuffer: spend_update_queue: SpendUpdateQueue, daily_spend_update_queue: DailySpendUpdateQueue, daily_team_spend_update_queue: DailySpendUpdateQueue, + daily_org_spend_update_queue: DailySpendUpdateQueue, daily_tag_spend_update_queue: DailySpendUpdateQueue, ): """ @@ -166,6 +169,9 @@ class RedisUpdateBuffer: daily_team_spend_update_transactions = ( await daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() ) + daily_org_spend_update_transactions = ( + await daily_org_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() ) @@ -195,6 +201,12 @@ class RedisUpdateBuffer: service_type=ServiceTypes.REDIS_DAILY_TEAM_SPEND_UPDATE_QUEUE, ) + await self._store_transactions_in_redis( + transactions=daily_org_spend_update_transactions, + redis_key=REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + service_type=ServiceTypes.REDIS_DAILY_SPEND_UPDATE_QUEUE, + ) + await self._store_transactions_in_redis( transactions=daily_tag_spend_update_transactions, redis_key=REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, @@ -329,6 +341,30 @@ class RedisUpdateBuffer: ), ) + async def get_all_daily_org_spend_update_transactions_from_redis_buffer( + self, + ) -> Optional[Dict[str, DailyOrganizationSpendTransaction]]: + """ + Gets all the daily organization 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_ORG_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, DailyOrganizationSpendTransaction], + 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]]: diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 402dffaaa7d..99b37c765a8 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -25,6 +25,7 @@ from litellm.proxy.management_endpoints.budget_management_endpoints import ( update_budget, ) from litellm.proxy.management_endpoints.common_utils import _set_object_metadata_field +from litellm.proxy.management_endpoints.common_utils import _user_has_admin_view from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, ) @@ -34,6 +35,10 @@ from litellm.proxy.management_helpers.utils import ( ) from litellm.proxy.utils import PrismaClient from litellm.utils import _update_dictionary +from litellm.types.proxy.management_endpoints.common_daily_activity import ( + SpendAnalyticsPaginatedResponse, +) +from litellm.proxy.management_endpoints.common_daily_activity import get_daily_activity router = APIRouter() @@ -255,6 +260,98 @@ async def new_organization( return response +@router.get( + "/organization/daily/activity", + response_model=SpendAnalyticsPaginatedResponse, + tags=["organization management"], +) +async def get_organization_daily_activity( + organization_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_organization_ids: Optional[str] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get daily activity for specific organizations or all accessible organizations. + """ + 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}, + ) + + # Parse comma-separated ids + org_ids_list = organization_ids.split(",") if organization_ids else None + exclude_org_ids_list: Optional[List[str]] = None + if exclude_organization_ids: + exclude_org_ids_list = ( + exclude_organization_ids.split(",") if exclude_organization_ids else None + ) + + # Restrict non-proxy-admins to only organizations where they are org_admin + if not _user_has_admin_view(user_api_key_dict): + memberships = await prisma_client.db.litellm_organizationmembership.find_many( + where={"user_id": user_api_key_dict.user_id} + ) + admin_org_ids = [ + m.organization_id + for m in memberships + if m.user_role == LitellmUserRoles.ORG_ADMIN.value + ] + if org_ids_list is None: + # Default to orgs where user is org_admin + org_ids_list = admin_org_ids + else: + # Ensure user is org_admin for all requested orgs + for org_id in org_ids_list: + if org_id not in admin_org_ids: + raise HTTPException( + status_code=403, + detail={ + "error": "User is not org_admin for Organization= {}.".format( + org_id + ) + }, + ) + + # Fetch organization aliases for metadata + where_condition = {} + if org_ids_list: + where_condition["organization_id"] = {"in": list(org_ids_list)} + org_aliases = await prisma_client.db.litellm_organizationtable.find_many( + where=where_condition + ) + org_alias_metadata = { + o.organization_id: {"organization_alias": o.organization_alias} + for o in org_aliases + } + + # Query daily activity for organizations + return await get_daily_activity( + prisma_client=prisma_client, + table_name="litellm_dailyorganizationspend", + entity_id_field="organization_id", + entity_id=org_ids_list, + entity_metadata_field=org_alias_metadata, + exclude_entity_ids=exclude_org_ids_list, + start_date=start_date, + end_date=end_date, + model=model, + api_key=api_key, + page=page, + page_size=page_size, + ) + + async def _set_object_permission( data: NewOrganizationRequest, prisma_client: Optional[PrismaClient], diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 6cfbb90c362..2883dfc4b82 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -304,6 +304,7 @@ model LiteLLM_SpendLogs { cache_key String? @default("") request_tags Json? @default("[]") team_id String? + organization_id String? end_user String? requester_ip_address String? messages Json? @default("{}") @@ -432,6 +433,35 @@ model LiteLLM_DailyUserSpend { @@index([mcp_namespaced_tool_name]) } +// Track daily organization spend metrics per model and key +model LiteLLM_DailyOrganizationSpend { + id String @id @default(uuid()) + organization_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([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name]) + @@index([date]) + @@index([organization_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()) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 59f712e5b6f..a5ad105c3ba 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -381,6 +381,7 @@ def get_logging_payload( # noqa: PLR0915 model=kwargs.get("model", "") or "", user=metadata.get("user_api_key_user_id", "") or "", team_id=metadata.get("user_api_key_team_id", "") or "", + organization_id=metadata.get("user_api_key_org_id") or "", metadata=safe_dumps(clean_metadata), cache_key=cache_key, spend=kwargs.get("response_cost", 0), diff --git a/schema.prisma b/schema.prisma index 6cfbb90c362..2883dfc4b82 100644 --- a/schema.prisma +++ b/schema.prisma @@ -304,6 +304,7 @@ model LiteLLM_SpendLogs { cache_key String? @default("") request_tags Json? @default("[]") team_id String? + organization_id String? end_user String? requester_ip_address String? messages Json? @default("{}") @@ -432,6 +433,35 @@ model LiteLLM_DailyUserSpend { @@index([mcp_namespaced_tool_name]) } +// Track daily organization spend metrics per model and key +model LiteLLM_DailyOrganizationSpend { + id String @id @default(uuid()) + organization_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([organization_id, date, api_key, model, custom_llm_provider, mcp_namespaced_tool_name]) + @@index([date]) + @@index([organization_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()) diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index b25080df0d0..8b2941672b3 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -10,6 +10,7 @@ "model": "gpt-4o", "user": "", "team_id": "", + "organization_id": "", "metadata": "{\"applied_guardrails\": [], \"batch_models\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"guardrail_information\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, 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 43dbc63689e..181d21b44f6 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 @@ -490,4 +490,86 @@ async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_i transaction_dict = call[1]["update"] # Each transaction should have one key with the format tag_date_api_key_model_provider for key, transaction in transaction_dict.items(): - assert transaction["request_id"] == request_id, f"request_id should be {request_id} but got {transaction.get('request_id')}" \ No newline at end of file + assert transaction["request_id"] == request_id, f"request_id should be {request_id} but got {transaction.get('request_id')}" + + +@pytest.mark.asyncio +async def test_add_spend_log_transaction_to_daily_org_transaction_injects_org_id_and_queues_update(): + """ + Verify org_id is injected into payload for daily aggregation and the update is queued. + """ + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + mock_prisma.get_request_status = MagicMock(return_value="success") + + org_id = "org-xyz" + payload = { + "request_id": "req-1", + "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": 10, + "completion_tokens": 5, + "spend": 0.2, + "metadata": '{"usage_object": {}}', + } + + writer.daily_org_spend_update_queue.add_update = AsyncMock() + + await writer.add_spend_log_transaction_to_daily_org_transaction( + payload=payload, + prisma_client=mock_prisma, + org_id=org_id, + ) + + # Should enqueue one org spend update + writer.daily_org_spend_update_queue.add_update.assert_called_once() + + # Validate key and injected fields + call_args = writer.daily_org_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"{org_id}_2024-01-01_test-key_gpt-4_openai" + assert transaction["organization_id"] == org_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_org_transaction_skips_when_org_id_missing(): + """ + Ensure no update is queued when org_id is None. + """ + writer = DBSpendUpdateWriter() + mock_prisma = MagicMock() + mock_prisma.get_request_status = MagicMock(return_value="success") + + payload = { + "request_id": "req-2", + "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": 10, + "completion_tokens": 5, + "spend": 0.2, + "metadata": '{"usage_object": {}}', + } + + writer.daily_org_spend_update_queue.add_update = AsyncMock() + + await writer.add_spend_log_transaction_to_daily_org_transaction( + payload=payload, + prisma_client=mock_prisma, + org_id=None, + ) + + writer.daily_org_spend_update_queue.add_update.assert_not_called() \ No newline at end of file diff --git a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py index 981c26d6db9..c02db727bf0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py @@ -94,6 +94,181 @@ async def test_organization_update_object_permissions_existing_permission(monkey mock_prisma_client.db.litellm_objectpermissiontable.upsert.assert_called_once() +@pytest.mark.asyncio +async def test_get_organization_daily_activity_admin_param_passing(monkeypatch): + """ + As admin, ensure parsed params are forwarded to get_daily_activity with correct values. + """ + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints import organization_endpoints + from litellm.proxy.management_endpoints.organization_endpoints import ( + get_organization_daily_activity, + ) + + # Mock prisma client + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( + return_value=[] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Admin view -> skip membership restriction + monkeypatch.setattr( + "litellm.proxy.management_endpoints.organization_endpoints._user_has_admin_view", + lambda _: True, + ) + + # Patch downstream common function and verify call args + mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") + get_daily_activity_mock = AsyncMock(return_value=mocked_response) + monkeypatch.setattr( + organization_endpoints, "get_daily_activity", get_daily_activity_mock + ) + + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin1") + result = await get_organization_daily_activity( + organization_ids="org1,org2", + start_date="2024-01-01", + end_date="2024-01-31", + model="gpt-4", + api_key="test-key", + page=2, + page_size=5, + exclude_organization_ids="org3", + user_api_key_dict=auth, + ) + + # Ensure passthrough to common method with correct args + get_daily_activity_mock.assert_awaited_once() + kwargs = get_daily_activity_mock.call_args.kwargs + assert kwargs["table_name"] == "litellm_dailyorganizationspend" + assert kwargs["entity_id_field"] == "organization_id" + assert kwargs["entity_id"] == ["org1", "org2"] + assert kwargs["exclude_entity_ids"] == ["org3"] + 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_organization_daily_activity_non_admin_defaults_to_admin_orgs(monkeypatch): + """ + Non-admin with no explicit organization_ids should default to orgs they are ORG_ADMIN of. + """ + from types import SimpleNamespace + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints import organization_endpoints + from litellm.proxy.management_endpoints.organization_endpoints import ( + get_organization_daily_activity, + ) + + # Mock prisma client and memberships + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( + return_value=[] + ) + mock_prisma_client.db.litellm_organizationmembership.find_many = AsyncMock( + return_value=[ + SimpleNamespace( + organization_id="orgA", user_role=LitellmUserRoles.ORG_ADMIN.value + ), + SimpleNamespace( + organization_id="orgB", user_role=LitellmUserRoles.ORG_ADMIN.value + ), + ] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Non-admin view + monkeypatch.setattr( + "litellm.proxy.management_endpoints.organization_endpoints._user_has_admin_view", + lambda _: False, + ) + + # Patch downstream aggregator + mocked_response = MagicMock(name="SpendAnalyticsPaginatedResponse") + get_daily_activity_mock = AsyncMock(return_value=mocked_response) + monkeypatch.setattr( + organization_endpoints, "get_daily_activity", get_daily_activity_mock + ) + + auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="regular-user" + ) + await get_organization_daily_activity( + organization_ids=None, + start_date="2024-02-01", + end_date="2024-02-28", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_organization_ids=None, + user_api_key_dict=auth, + ) + + kwargs = get_daily_activity_mock.call_args.kwargs + assert kwargs["entity_id"] == ["orgA", "orgB"] + assert kwargs["start_date"] == "2024-02-01" + assert kwargs["end_date"] == "2024-02-28" + + +@pytest.mark.asyncio +async def test_get_organization_daily_activity_non_admin_unauthorized_org_raises(monkeypatch): + """ + Non-admin requesting an org they aren't ORG_ADMIN for should raise 403. + """ + from types import SimpleNamespace + + from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth + from litellm.proxy.management_endpoints.organization_endpoints import ( + get_organization_daily_activity, + ) + + # Mock prisma client and memberships (only orgA is admin) + mock_prisma_client = AsyncMock() + mock_prisma_client.db.litellm_organizationmembership.find_many = AsyncMock( + return_value=[ + SimpleNamespace( + organization_id="orgA", user_role=LitellmUserRoles.ORG_ADMIN.value + ) + ] + ) + mock_prisma_client.db.litellm_organizationtable.find_many = AsyncMock( + return_value=[] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + # Non-admin view + monkeypatch.setattr( + "litellm.proxy.management_endpoints.organization_endpoints._user_has_admin_view", + lambda _: False, + ) + + auth = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, user_id="regular-user" + ) + + with pytest.raises(HTTPException) as exc: + await get_organization_daily_activity( + organization_ids="orgA,orgX", # orgX is unauthorized + start_date="2024-03-01", + end_date="2024-03-31", + model=None, + api_key=None, + page=1, + page_size=10, + exclude_organization_ids=None, + user_api_key_dict=auth, + ) + assert exc.value.status_code == 403 + @pytest.mark.asyncio async def test_organization_update_object_permissions_no_existing_permission( monkeypatch, diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 65b245b9e7e..33715eb461a 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -196,6 +196,7 @@ ignored_keys = [ "endTime", "completionStartTime", "endTime", + "organization_id", "metadata.model_map_information", "metadata.usage_object", "metadata.cold_storage_object_key",