fix(spend-tracking): hand plain dict rows to polars in the CloudZero and Focus exports

litellm/integrations/cloudzero/database.py

@@ -106,6 +106,6 @@ class LiteLLMDatabase:
                 else []
             )
             recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
-            return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None)
+            return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None)
         except Exception as e:
             raise Exception(f"Error retrieving usage data: {e}")

litellm/integrations/focus/database.py

@@ -108,7 +108,7 @@ class FocusLiteLLMDatabase:
                 else []
             )
             recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
-            return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None)
+            return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None)
         except Exception as exc:
             raise RuntimeError(f"Error retrieving usage data: {exc}") from exc
 

tests/test_litellm/integrations/cloudzero/test_cloudzero.py

@@ -1,3 +1,4 @@
+import hashlib
 from datetime import datetime
 from unittest.mock import AsyncMock, MagicMock, patch
 
@@ -165,3 +166,26 @@ class TestCloudZeroHourlyExport:
             logger = CloudZeroLogger(api_key="test", connection_id="test")
 
             await logger._hourly_usage_data_export()
+
+
+class TestLiteLLMDatabaseUsageData:
+    @pytest.mark.asyncio
+    async def test_builds_frame_from_rows_recovered_for_double_hashed_keys(self, monkeypatch: pytest.MonkeyPatch):
+        double_hashed = hashlib.sha256(b"sk-hashed-token").hexdigest()
+        joined_row = {"api_key": "sk-joined", "api_key_alias": "joined", "team_id": "team-0", "user_email": None, "spend": 0.1}
+        dirty_row = {"api_key": double_hashed, "api_key_alias": None, "team_id": None, "user_email": None, "spend": 0.5}
+
+        async def query_raw(query: str, *params):
+            if "sha256(" in query:
+                return [{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": None}]
+            return [joined_row, dirty_row]
+
+        fake_client = MagicMock()
+        fake_client.db.query_raw = AsyncMock(side_effect=query_raw)
+        db = LiteLLMDatabase()
+        monkeypatch.setattr(db, "_ensure_prisma_client", lambda: fake_client)
+
+        result = await db.get_usage_data()
+
+        assert result["api_key_alias"].to_list() == ["joined", "batch-worker"]
+        assert result["team_id"].to_list() == ["team-0", "team-1"]

tests/test_litellm/integrations/focus/test_focus_database.py

@@ -1,5 +1,6 @@
 """Tests for FocusLiteLLMDatabase query construction."""
 
+import hashlib
 from datetime import datetime, timezone
 from types import SimpleNamespace
 from unittest.mock import AsyncMock
@@ -87,3 +88,24 @@ async def test_should_join_organization_table(monkeypatch: pytest.MonkeyPatch):
     )
     assert "ot.organization_alias as organization_alias" in query_text
     assert 'LEFT JOIN "LiteLLM_OrganizationTable" ot' in query_text
+
+
+@pytest.mark.asyncio
+async def test_should_build_frame_from_rows_recovered_for_double_hashed_keys(monkeypatch: pytest.MonkeyPatch):
+    double_hashed = hashlib.sha256(b"sk-hashed-token").hexdigest()
+    joined_row = {"api_key": "sk-joined", "api_key_alias": "joined", "team_id": "team-0", "user_email": None, "spend": 0.1}
+    dirty_row = {"api_key": double_hashed, "api_key_alias": None, "team_id": None, "user_email": None, "spend": 0.5}
+
+    async def query_raw(query: str, *params):
+        if "sha256(" in query:
+            return [{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": None}]
+        return [joined_row, dirty_row]
+
+    mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=AsyncMock(side_effect=query_raw)))
+    db = FocusLiteLLMDatabase()
+    monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client)
+
+    result = await db.get_usage_data()
+
+    assert result["api_key_alias"].to_list() == ["joined", "batch-worker"]
+    assert result["team_id"].to_list() == ["team-0", "team-1"]