⚡️ Speed up function global_view_all_end_users by 24%

[codeflash-ai]

📄 24% (0.24x) speedup for global_view_all_end_users in litellm/proxy/spend_tracking/spend_management_endpoints.py

⏱️ Runtime : 5.42 milliseconds → 4.36 milliseconds (best of 150 runs)

📝 Explanation and details

The optimization replaces an explicit loop with .append() calls with a list comprehension, achieving a 24% runtime improvement and 2% throughput increase.

Key optimization applied:

• Original approach: Creates empty list _end_users = [], then iterates through db_response with a for loop, calling _end_users.append(row["end_user"]) for each item
• Optimized approach: Uses list comprehension [row["end_user"] for row in db_response] to build the list directly in the return statement

Why this is faster:

1. Reduced Python bytecode overhead - List comprehensions are implemented at the C level in CPython, avoiding the interpreter overhead of explicit loops
2. Eliminated method lookups - Removes repeated .append() method lookups (61,118 calls in the profiler), which require attribute resolution on each iteration
3. Better memory allocation - List comprehensions can pre-allocate memory more efficiently than incremental appends

Performance impact by workload:
The line profiler shows the optimization eliminates the two most expensive operations:

• 30.4% of runtime spent on the for loop iteration (now eliminated)
• 45.6% of runtime spent on .append() calls (now eliminated)

Test case performance:
The optimization is particularly effective for:

• Large-scale scenarios (500-1000 users) where the loop overhead compounds
• High-throughput concurrent calls where the per-request savings accumulate across many simultaneous operations
• Medium load tests (50-100 concurrent calls) benefit from reduced CPU cycles per operation

This optimization maintains identical functionality while leveraging Python's optimized list comprehension implementation for better performance across all data sizes.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 354 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	100.0%

🌀 Generated Regression Tests and Runtime

import asyncio # used to run async functions
from typing import TYPE_CHECKING, Any
from unittest.mock import AsyncMock, MagicMock, patch

import pytest # used for our unit tests

function to test

(EXACT COPY, DO NOT MODIFY)

from fastapi import APIRouter, Depends, HTTPException
from litellm.proxy._types import TYPE_CHECKING, Any
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.spend_tracking.spend_management_endpoints import
global_view_all_end_users

router = APIRouter()
from litellm.proxy.spend_tracking.spend_management_endpoints import
global_view_all_end_users

---------------------- UNIT TESTS ----------------------

Helper: Patch the prisma_client in the correct module

PRISMA_CLIENT_PATH = "litellm.proxy.proxy_server.prisma_client"

@pytest.mark.asyncio
async def test_global_view_all_end_users_basic_nonempty():
"""
Basic: Test that the function returns the correct list of end users when db returns multiple users.
"""
# Prepare fake db response
fake_db_response = [
{"end_user": "alice"},
{"end_user": "bob"},
{"end_user": "carol"}
]
# Patch prisma_client and its db.query_raw
mock_query_raw = AsyncMock(return_value=fake_db_response)
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    result = await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_basic_empty():
"""
Basic: Test that the function returns an empty list if db returns an empty list.
"""
mock_query_raw = AsyncMock(return_value=[])
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    result = await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_basic_single_user():
"""
Basic: Test that the function returns a single end user if only one is present.
"""
mock_query_raw = AsyncMock(return_value=[{"end_user": "alice"}])
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    result = await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_db_none():
"""
Edge: Test that the function returns [] if db.query_raw returns None.
"""
mock_query_raw = AsyncMock(return_value=None)
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    result = await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_prisma_client_none():
"""
Edge: Test that the function raises HTTPException if prisma_client is None.
"""
with patch(PRISMA_CLIENT_PATH, None):
with pytest.raises(HTTPException) as exc_info:
await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_concurrent_calls():
"""
Edge: Test concurrent execution of the function.
"""
fake_db_response = [{"end_user": "alice"}, {"end_user": "bob"}]
mock_query_raw = AsyncMock(return_value=fake_db_response)
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    # Run 5 concurrent calls
    results = await asyncio.gather(*[global_view_all_end_users() for _ in range(5)])
    for result in results:
        pass

@pytest.mark.asyncio
async def test_global_view_all_end_users_handles_nonstring_users():
"""
Edge: Test that the function can handle non-string end_user values.
"""
fake_db_response = [
{"end_user": 123},
{"end_user": None},
{"end_user": "bob"}
]
mock_query_raw = AsyncMock(return_value=fake_db_response)
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    result = await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_db_row_missing_key():
"""
Edge: Test that the function raises KeyError if a row is missing 'end_user'.
"""
fake_db_response = [
{"not_end_user": "foo"},
]
mock_query_raw = AsyncMock(return_value=fake_db_response)
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    with pytest.raises(KeyError):
        await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_large_scale():
"""
Large Scale: Test with a large number of end users (up to 1000).
"""
fake_db_response = [{"end_user": f"user{i}"} for i in range(1000)]
mock_query_raw = AsyncMock(return_value=fake_db_response)
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    result = await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_large_scale_concurrent():
"""
Large Scale: Test concurrent calls with large user lists.
"""
fake_db_response = [{"end_user": f"user{i}"} for i in range(500)]
mock_query_raw = AsyncMock(return_value=fake_db_response)
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    results = await asyncio.gather(*[global_view_all_end_users() for _ in range(10)])
    for result in results:
        pass

@pytest.mark.asyncio
async def test_global_view_all_end_users_throughput_small_load():
"""
Throughput: Test performance with small load (10 concurrent calls, 10 users each).
"""
fake_db_response = [{"end_user": f"user{i}"} for i in range(10)]
mock_query_raw = AsyncMock(return_value=fake_db_response)
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    results = await asyncio.gather(*[global_view_all_end_users() for _ in range(10)])
    for result in results:
        pass

@pytest.mark.asyncio
async def test_global_view_all_end_users_throughput_medium_load():
"""
Throughput: Test performance with medium load (50 concurrent calls, 100 users each).
"""
fake_db_response = [{"end_user": f"user{i}"} for i in range(100)]
mock_query_raw = AsyncMock(return_value=fake_db_response)
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    results = await asyncio.gather(*[global_view_all_end_users() for _ in range(50)])
    for result in results:
        pass

@pytest.mark.asyncio
async def test_global_view_all_end_users_throughput_high_load():
"""
Throughput: Test performance with high load (100 concurrent calls, 500 users each).
"""
fake_db_response = [{"end_user": f"user{i}"} for i in range(500)]
mock_query_raw = AsyncMock(return_value=fake_db_response)
mock_db = MagicMock()
mock_db.query_raw = mock_query_raw
mock_prisma_client = MagicMock()
mock_prisma_client.db = mock_db

with patch(PRISMA_CLIENT_PATH, mock_prisma_client):
    results = await asyncio.gather(*[global_view_all_end_users() for _ in range(100)])
    for result in results:
        pass

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

#------------------------------------------------
import asyncio # used to run async functions
import sys

Helper: Patch prisma_client in litellm.proxy.proxy_server

import types
from typing import TYPE_CHECKING, Any

import pytest # used for our unit tests
from fastapi import APIRouter, Depends, HTTPException
from litellm.proxy._types import TYPE_CHECKING, Any
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.spend_tracking.spend_management_endpoints import
global_view_all_end_users

router = APIRouter()
from litellm.proxy.spend_tracking.spend_management_endpoints import
global_view_all_end_users

Helper: Mock db.query_raw coroutine

class MockDB:
def init(self, response):
self._response = response
self._calls = []

async def query_raw(self, query):
    self._calls.append(query)
    return self._response

class MockPrismaClient:
def init(self, db_response):
self.db = MockDB(db_response)

--- Basic Test Cases ---

@pytest.mark.asyncio
async def test_global_view_all_end_users_basic_nonempty(monkeypatch):
"""
Test basic async/await behavior with a non-empty db_response.
"""
# Setup mock prisma_client with sample end_users
sample_rows = [
{"end_user": "alice"},
{"end_user": "bob"},
{"end_user": "carol"},
]
mock_client = MockPrismaClient(db_response=sample_rows)
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", mock_client
)

result = await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_basic_empty(monkeypatch):
"""
Test with db_response as empty list (no end_users).
"""
mock_client = MockPrismaClient(db_response=[])
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", mock_client
)

result = await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_basic_none(monkeypatch):
"""
Test with db_response as None (should return empty list).
"""
mock_client = MockPrismaClient(db_response=None)
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", mock_client
)

result = await global_view_all_end_users()

--- Edge Test Cases ---

@pytest.mark.asyncio
async def test_global_view_all_end_users_no_db(monkeypatch):
"""
Test when prisma_client is None (should raise HTTPException).
"""
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", None
)
with pytest.raises(HTTPException) as excinfo:
await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_concurrent(monkeypatch):
"""
Test concurrent execution of the function.
"""
sample_rows = [
{"end_user": "alice"},
{"end_user": "bob"},
]
mock_client = MockPrismaClient(db_response=sample_rows)
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", mock_client
)

# Run multiple coroutines concurrently
results = await asyncio.gather(
    global_view_all_end_users(),
    global_view_all_end_users(),
    global_view_all_end_users(),
)
for result in results:
    pass

@pytest.mark.asyncio
async def test_global_view_all_end_users_row_missing_end_user(monkeypatch):
"""
Test when a row in db_response is missing 'end_user' key.
Should raise KeyError.
"""
sample_rows = [
{"end_user": "alice"},
{},
]
mock_client = MockPrismaClient(db_response=sample_rows)
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", mock_client
)

with pytest.raises(KeyError):
    await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_row_end_user_none(monkeypatch):
"""
Test when a row has 'end_user' as None.
Should include None in the result.
"""
sample_rows = [
{"end_user": "alice"},
{"end_user": None},
]
mock_client = MockPrismaClient(db_response=sample_rows)
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", mock_client
)

result = await global_view_all_end_users()

--- Large Scale Test Cases ---

@pytest.mark.asyncio
async def test_global_view_all_end_users_large_scale(monkeypatch):
"""
Test with a large number of unique end_users.
"""
num_users = 500
sample_rows = [{"end_user": f"user_{i}"} for i in range(num_users)]
mock_client = MockPrismaClient(db_response=sample_rows)
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", mock_client
)

result = await global_view_all_end_users()

@pytest.mark.asyncio
async def test_global_view_all_end_users_large_scale_duplicates(monkeypatch):
"""
Test with duplicate end_users in db_response.
Should return all as given (function does not deduplicate).
"""
sample_rows = [
{"end_user": "alice"},
{"end_user": "bob"},
{"end_user": "alice"},
{"end_user": "bob"},
{"end_user": "carol"},
]
mock_client = MockPrismaClient(db_response=sample_rows)
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", mock_client
)

result = await global_view_all_end_users()

--- Throughput Test Cases ---

@pytest.mark.asyncio
async def test_global_view_all_end_users_throughput_small_load(monkeypatch):
"""
Throughput test: Small load, multiple concurrent calls.
"""
sample_rows = [
{"end_user": "alice"},
{"end_user": "bob"},
]
mock_client = MockPrismaClient(db_response=sample_rows)
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", mock_client
)

# Run 10 concurrent calls
results = await asyncio.gather(
    *[global_view_all_end_users() for _ in range(10)]
)
for result in results:
    pass

@pytest.mark.asyncio
async def test_global_view_all_end_users_throughput_medium_load(monkeypatch):
"""
Throughput test: Medium load, moderate concurrent calls.
"""
sample_rows = [{"end_user": f"user_{i}"} for i in range(50)]
mock_client = MockPrismaClient(db_response=sample_rows)
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", mock_client
)

# Run 50 concurrent calls
results = await asyncio.gather(
    *[global_view_all_end_users() for _ in range(50)]
)
for result in results:
    pass

@pytest.mark.asyncio
async def test_global_view_all_end_users_throughput_high_volume(monkeypatch):
"""
Throughput test: High volume, stress test with many concurrent calls.
"""
num_users = 200
sample_rows = [{"end_user": f"user_{i}"} for i in range(num_users)]
mock_client = MockPrismaClient(db_response=sample_rows)
monkeypatch.setattr(
sys.modules["litellm.proxy.proxy_server"], "prisma_client", mock_client
)

# Run 100 concurrent calls
results = await asyncio.gather(
    *[global_view_all_end_users() for _ in range(100)]
)
for result in results:
    pass

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-global_view_all_end_users-mhu07mug and push.