(feat) - track org_id in SpendLogs (#8253)

* track org id in spend logs

* read org id from team table

* show user_api_key_org_id in spend logs

* test_spend_logs_payload

* test_spend_logs_with_org_id

* test_spend_logs_with_org_id

litellm/proxy/_types.py

@@ -1856,6 +1856,7 @@ class SpendLogsMetadata(TypedDict):
     user_api_key: Optional[str]
     user_api_key_alias: Optional[str]
     user_api_key_team_id: Optional[str]
+    user_api_key_org_id: Optional[str]
     user_api_key_user_id: Optional[str]
     user_api_key_team_alias: Optional[str]
     spend_logs_metadata: Optional[

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -38,6 +38,7 @@ def _get_spend_logs_metadata(metadata: Optional[dict]) -> SpendLogsMetadata:
             user_api_key=None,
             user_api_key_alias=None,
             user_api_key_team_id=None,
+            user_api_key_org_id=None,
             user_api_key_user_id=None,
             user_api_key_team_alias=None,
             spend_logs_metadata=None,

litellm/proxy/utils.py

@@ -1544,6 +1544,7 @@ class PrismaClient:
                             t.team_alias AS team_alias,
                             t.metadata AS team_metadata,
                             t.members_with_roles AS team_members_with_roles,
+                            t.organization_id as org_id,
                             tm.spend AS team_member_spend,
                             m.aliases AS team_model_aliases,
                             -- Added comma to separate b.* columns

tests/logging_callback_tests/test_spend_logs.py

@@ -66,14 +66,14 @@ def test_spend_logs_payload(model_id: Optional[str]):
                 "metadata": {
                     "tags": ["model-anthropic-claude-v2.1", "app-ishaan-prod"],
                     "user_api_key": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",
-                    "user_api_key_alias": None,
+                    "user_api_key_alias": "custom-key-alias",
                     "user_api_end_user_max_budget": None,
                     "litellm_api_version": "0.0.0",
                     "global_max_parallel_requests": None,
                     "user_api_key_user_id": "116544810872468347480",
-                    "user_api_key_org_id": None,
-                    "user_api_key_team_id": None,
-                    "user_api_key_team_alias": None,
+                    "user_api_key_org_id": "custom-org-id",
+                    "user_api_key_team_id": "custom-team-id",
+                    "user_api_key_team_alias": "custom-team-alias",
                     "user_api_key_metadata": {},
                     "requester_ip_address": "127.0.0.1",
                     "spend_logs_metadata": {"hello": "world"},
@@ -216,6 +216,10 @@ def test_spend_logs_payload(model_id: Optional[str]):
     assert (
         payload["request_tags"] == '["model-anthropic-claude-v2.1", "app-ishaan-prod"]'
     )
+    assert payload["metadata"]["user_api_key_org_id"] == "custom-org-id"
+    assert payload["metadata"]["user_api_key_team_id"] == "custom-team-id"
+    assert payload["metadata"]["user_api_key_team_alias"] == "custom-team-alias"
+    assert payload["metadata"]["user_api_key_alias"] == "custom-key-alias"
 
     assert payload["custom_llm_provider"] == "azure"
 

tests/test_spend_logs.py

@@ -1,18 +1,20 @@
 # What this tests?
 ## Tests /spend endpoints.
 
-import pytest, time, uuid
+import pytest, time, uuid, json
 import asyncio
 import aiohttp
 
 
-async def generate_key(session, models=[]):
+async def generate_key(session, models=[], team_id=None):
     url = "http://0.0.0.0:4000/key/generate"
     headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
     data = {
         "models": models,
         "duration": None,
     }
+    if team_id is not None:
+        data["team_id"] = team_id
 
     async with session.post(url, headers=headers, json=data) as response:
         status = response.status
@@ -113,6 +115,81 @@ async def test_spend_logs():
         await get_spend_logs(session=session, request_id=response["id"])
 
 
+async def generate_org(session: aiohttp.ClientSession) -> dict:
+    """
+    Generate a new organization using the API.
+
+    Args:
+        session: aiohttp client session
+
+    Returns:
+        dict: Response containing org_id
+    """
+    url = "http://0.0.0.0:4000/organization/new"
+    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+
+    request_body = {
+        "organization_alias": f"test-org-{uuid.uuid4()}",
+    }
+
+    async with session.post(url, headers=headers, json=request_body) as response:
+        return await response.json()
+
+
+async def generate_team(session: aiohttp.ClientSession, org_id: str) -> dict:
+    """
+    Generate a new team within an organization using the API.
+
+    Args:
+        session: aiohttp client session
+        org_id: Organization ID to create the team in
+
+    Returns:
+        dict: Response containing team_id
+    """
+    url = "http://0.0.0.0:4000/team/new"
+    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    data = {"organization_id": org_id}
+
+    async with session.post(url, headers=headers, json=data) as response:
+        return await response.json()
+
+
+@pytest.mark.asyncio
+async def test_spend_logs_with_org_id():
+    """
+    - Create Organization
+    - Create Team in organization
+    - Create Key in organization
+    - Make call (makes sure it's in spend logs)
+    - Get request id from logs
+    - Assert spend logs have correct org_id and team_id
+    """
+    async with aiohttp.ClientSession() as session:
+        org_gen = await generate_org(session=session)
+        print("org_gen: ", json.dumps(org_gen, indent=4, default=str))
+        org_id = org_gen["organization_id"]
+        team_gen = await generate_team(session=session, org_id=org_id)
+        print("team_gen: ", json.dumps(team_gen, indent=4, default=str))
+        team_id = team_gen["team_id"]
+        key_gen = await generate_key(session=session, team_id=team_id)
+        print("key_gen: ", json.dumps(key_gen, indent=4, default=str))
+        key = key_gen["key"]
+        response = await chat_completion(session=session, key=key)
+        await asyncio.sleep(20)
+        spend_logs_response = await get_spend_logs(
+            session=session, request_id=response["id"]
+        )
+        print(
+            "spend_logs_response: ",
+            json.dumps(spend_logs_response, indent=4, default=str),
+        )
+        spend_logs_response = spend_logs_response[0]
+        assert spend_logs_response["metadata"]["user_api_key_org_id"] == org_id
+        assert spend_logs_response["metadata"]["user_api_key_team_id"] == team_id
+        assert spend_logs_response["team_id"] == team_id
+
+
 async def get_predict_spend_logs(session):
     url = "http://0.0.0.0:4000/global/predict/spend/logs"
     headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}