diff --git a/litellm/repositories/daily_activity_repository.py b/litellm/repositories/daily_activity_repository.py index e9d8c3bd309..2e34582091a 100644 --- a/litellm/repositories/daily_activity_repository.py +++ b/litellm/repositories/daily_activity_repository.py @@ -282,13 +282,20 @@ class DailyActivityRepository: scope.timezone_offset_minutes, include_current_utc_day=scope.include_current_utc_day, ) - entity_filter: Final = { - **({"in": list(scope.entity_ids)} if scope.entity_ids is not None else {}), - **({"not": {"in": list(scope.exclude_entity_ids)}} if scope.exclude_entity_ids else {}), - } + exclusion_filter: Final = ( + { + "OR": [ + {scope.entity_id_field: None}, + {scope.entity_id_field: {"not": {"in": list(scope.exclude_entity_ids)}}}, + ] + } + if scope.exclude_entity_ids + else {} + ) conditions: Final = { "date": {"gte": adjusted_start, "lte": adjusted_end}, - **({scope.entity_id_field: entity_filter} if entity_filter else {}), + **({scope.entity_id_field: {"in": list(scope.entity_ids)}} if scope.entity_ids is not None else {}), + **exclusion_filter, **({"model": scope.model} if scope.model else {}), **({"api_key": {"in": list(scope.api_keys)}} if scope.api_keys is not None else {}), } diff --git a/litellm/repositories/daily_activity_sql.py b/litellm/repositories/daily_activity_sql.py index 96920b55eb0..f12dc1be5ae 100644 --- a/litellm/repositories/daily_activity_sql.py +++ b/litellm/repositories/daily_activity_sql.py @@ -78,7 +78,14 @@ def build_where_clause(scope: DailyActivityScope, *, start_index: int = 1) -> tu if has_entity_array else () ), - *((f'NOT ("{scope.entity_id_field}" = ANY(${exclusion_index}::text[]))',) if scope.exclude_entity_ids else ()), + *( + ( + f'("{scope.entity_id_field}" IS NULL ' + f'OR NOT ("{scope.entity_id_field}" = ANY(${exclusion_index}::text[])))', + ) + if scope.exclude_entity_ids + else () + ), *((f"model = ${model_index}",) if scope.model else ()), *( ("FALSE",) diff --git a/tests/code_coverage_tests/unbounded_in_baseline.txt b/tests/code_coverage_tests/unbounded_in_baseline.txt index 0a95ebfc01f..c42a6b0ddf5 100644 --- a/tests/code_coverage_tests/unbounded_in_baseline.txt +++ b/tests/code_coverage_tests/unbounded_in_baseline.txt @@ -138,7 +138,7 @@ litellm/proxy/utils.py PrismaClient.get_data prisma budget_id.in `budget_id_list litellm/proxy/utils.py PrismaClient.get_data prisma team_id.in `team_id_list` 0 litellm/proxy/utils.py PrismaClient.get_data prisma user_id.in `user_id_list` 0 litellm/proxy/utils.py prefetch_config_params prisma param_name.in `param_names` 0 -litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma ?.in `list(scope.entity_ids)` 0 +litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma [scope.entity_id_field].in `list(scope.entity_ids)` 0 litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma api_key.in `list(scope.api_keys)` 0 litellm/repositories/daily_activity_repository.py DailyActivityRepository.daily_rows prisma not.in `list(scope.exclude_entity_ids)` 0 litellm/router_utils/auto_router_model_naming.py raw-sql classifier_type.IN `IN ({_LLM_CLASSIFIER_TYPES_SQL})` 0 diff --git a/tests/integration/spend/_daily_activity_fixtures.py b/tests/integration/spend/_daily_activity_fixtures.py index d6e819edca6..9cde3ed0fb9 100644 --- a/tests/integration/spend/_daily_activity_fixtures.py +++ b/tests/integration/spend/_daily_activity_fixtures.py @@ -317,3 +317,25 @@ def seed_daily_team_unassigned_fixture( rows, ) connection.commit() + + +def seed_daily_team_exclusion_fixture(connection: psycopg.Connection, *, schema: str) -> None: + team_table: Final = sql.Identifier(schema, "LiteLLM_DailyTeamSpend") + rows: Final = ( + ("exclusion-null", None, "key-excluded-null", 3.0), + ("exclusion-empty", "", "key-excluded-empty", 7.0), + ("exclusion-dashboard", "litellm-dashboard", "key-excluded-dashboard", 11.0), + ("exclusion-normal", "team-normal", "key-excluded-normal", 13.0), + ) + with connection.cursor() as cursor: + cursor.executemany( + sql.SQL(""" + INSERT INTO {} + (id, team_id, date, api_key, model, model_group, custom_llm_provider, + mcp_namespaced_tool_name, endpoint, prompt_tokens, spend, api_requests, updated_at) + VALUES (%s, %s, '2026-06-04', %s, 'model-a', '', 'provider-a', NULL, '/v1/chat/completions', + 1, %s, 1, '2026-06-04 12:00:00') + """).format(team_table), + rows, + ) + connection.commit() diff --git a/tests/integration/spend/test_daily_activity_repository.py b/tests/integration/spend/test_daily_activity_repository.py index c347e18bf64..c6ffeef2eba 100644 --- a/tests/integration/spend/test_daily_activity_repository.py +++ b/tests/integration/spend/test_daily_activity_repository.py @@ -15,6 +15,7 @@ from integration.spend._daily_activity_fixtures import ( seed_daily_activity_fixture, seed_daily_tag_activity_fixture, seed_daily_tag_float_tie_fixture, + seed_daily_team_exclusion_fixture, seed_daily_team_unassigned_fixture, ) from prisma import Prisma @@ -91,6 +92,7 @@ async def _daily_activity_database( include_tag_activity: bool = False, include_tag_float_tie_activity: bool = False, include_team_unassigned_activity: bool = False, + include_team_exclusion_activity: bool = False, ) -> AsyncIterator[Prisma]: schema: Final = f"integration_{uuid.uuid4().hex}" url: Final = os.environ["DATABASE_URL"] @@ -111,6 +113,8 @@ async def _daily_activity_database( seed_daily_team_unassigned_fixture( connection, schema=schema, ptu_sentinel_api_key=constants.PTU_SENTINEL_API_KEY ) + if include_team_exclusion_activity: + seed_daily_team_exclusion_fixture(connection, schema=schema) database: Final = Prisma(datasource={"url": _scoped_url(url, schema)}) await database.connect() try: @@ -622,3 +626,38 @@ async def test_team_entity_rollups_merge_null_and_empty_entity_ids() -> None: keyed_rows: Final = tuple(row for row in aggregate.entity_rows or () if not row.api_key_rolled) assert {row.entity_id for row in keyed_rows} == {""} assert {row.api_key for row in keyed_rows} == {"key-unassigned-null", "key-unassigned-empty"} + + +@pytest.mark.asyncio +async def test_team_exclusion_keeps_null_and_empty_entity_rows() -> None: + async with _daily_activity_database(include_team_exclusion_activity=True) as database: + repository: Final = _repository(database) + scope: Final = DailyActivityScope( + table=DailyActivityTable.TEAM, + entity_id_field="team_id", + entity_ids=None, + exclude_entity_ids=("litellm-dashboard",), + api_keys=None, + start_date="2026-06-04", + end_date="2026-06-04", + model=None, + timezone_offset_minutes=None, + ) + aggregate: Final = await repository.aggregated(scope, include_entity_breakdown=True, api_key_limit=10) + + totals: Final = tuple(row for row in aggregate.grouping_rows if row.group_level == 127) + assert len(totals) == 1 + assert totals[0].spend == 23.0 + assert aggregate.distinct_api_keys == 3 + + keyed_rows: Final = tuple(row for row in aggregate.entity_rows or () if not row.api_key_rolled) + assert {row.api_key for row in keyed_rows} == {"key-excluded-null", "key-excluded-empty", "key-excluded-normal"} + assert {row.entity_id for row in keyed_rows} == {"", "team-normal"} + + page: Final = await repository.key_page(scope, offset=0, limit=10) + assert page.total_api_keys == 3 + assert {row.api_key for row in page.rows} == {"key-excluded-null", "key-excluded-empty", "key-excluded-normal"} + + daily: Final = await repository.daily_rows(scope, page=1, page_size=10) + assert daily.total_count == 3 + assert {row.api_key for row in daily.rows} == {"key-excluded-null", "key-excluded-empty", "key-excluded-normal"} diff --git a/tests/integration/spend/test_daily_activity_routes.py b/tests/integration/spend/test_daily_activity_routes.py index 7c8db96966e..9e4c16e573c 100644 --- a/tests/integration/spend/test_daily_activity_routes.py +++ b/tests/integration/spend/test_daily_activity_routes.py @@ -512,3 +512,37 @@ async def test_user_key_pages_and_details_respect_caller_scope() -> None: other_body: Final = JSON_OBJECT.validate_json(other_details.content) assert object_value(other_body["metadata"])["total_api_keys"] == 0 assert _aggregate_top_keys(other_body["results"]) == frozenset() + + +@pytest.mark.asyncio +async def test_team_routes_exclusion_keeps_unassigned_keys() -> None: + async with _daily_activity_database(include_team_exclusion_activity=True) as database: + repository: Final = _repository(database) + app: Final = FastAPI() + app.include_router(daily_activity_router) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth( + user_id="integration-admin", user_role=LitellmUserRoles.PROXY_ADMIN + ) + app.dependency_overrides[get_daily_activity_prisma_client] = lambda: _PrismaDatabase(database) + app.dependency_overrides[get_daily_activity_repository] = lambda: repository + + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver") as client: + params: Final = { + "start_date": "2026-06-04", + "end_date": "2026-06-04", + "exclude_team_ids": "litellm-dashboard", + } + surviving_keys: Final = frozenset(("key-excluded-null", "key-excluded-empty", "key-excluded-normal")) + + aggregated: Final = await client.get("/team/daily/activity/aggregated", params=params) + assert aggregated.status_code == 200, aggregated.text + aggregated_body: Final = JSON_OBJECT.validate_json(aggregated.content) + assert object_value(aggregated_body["metadata"])["total_spend"] == 23.0 + assert object_value(aggregated_body["metadata"])["total_api_keys"] == 3 + assert _aggregate_top_keys(aggregated_body["results"]) == surviving_keys + + page: Final = await client.get("/team/daily/activity/aggregated/keys", params={**params, "limit": 10}) + assert page.status_code == 200, page.text + page_body: Final = DailyActivityKeyPageResponse.model_validate_json(page.content) + assert page_body.total_api_keys == 3 + assert frozenset(row.api_key for row in page_body.api_keys) == surviving_keys diff --git a/tests/unit/repositories/test_daily_activity_repository.py b/tests/unit/repositories/test_daily_activity_repository.py index f0bbfd32d2c..4bb833f2bc2 100644 --- a/tests/unit/repositories/test_daily_activity_repository.py +++ b/tests/unit/repositories/test_daily_activity_repository.py @@ -494,7 +494,8 @@ async def test_daily_rows_selects_the_table_and_applies_filters_and_pagination( expected_where: Final = { "date": {"gte": "2026-01-01", "lte": "2026-01-31"}, - entity_field: {"in": ["entity-1"], "not": {"in": ["excluded-1"]}}, + entity_field: {"in": ["entity-1"]}, + "OR": [{entity_field: None}, {entity_field: {"not": {"in": ["excluded-1"]}}}], "model": "model-1", "api_key": {"in": ["key-1"]}, } @@ -516,6 +517,22 @@ async def test_daily_rows_selects_the_table_and_applies_filters_and_pagination( assert sum(len(daily_table.find_many_calls) for daily_table in tables.values()) == 1 +@pytest.mark.asyncio +async def test_daily_rows_exclusion_without_entity_filter_keeps_null_entity_rows() -> None: + database = _FakeDatabase() + repository, _ = _repository(database) + scope = _scope(table=DailyActivityTable.TEAM, entity_ids=None, exclude_entity_ids=("litellm-dashboard",)) + + await repository.daily_rows(scope, page=1, page_size=10) + + expected_where: Final = { + "date": {"gte": "2026-01-01", "lte": "2026-01-31"}, + "OR": [{"team_id": None}, {"team_id": {"not": {"in": ["litellm-dashboard"]}}}], + } + assert database.litellm_dailyteamspend.count_calls == [expected_where] + assert database.litellm_dailyteamspend.find_many_calls == [expected_where] + + @pytest.mark.asyncio async def test_export_is_lazy_and_uses_the_last_row_as_the_next_cursor(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(constants, "USAGE_EXPORT_BATCH_SIZE", 2) diff --git a/tests/unit/repositories/test_daily_activity_sql.py b/tests/unit/repositories/test_daily_activity_sql.py index 67f15712b11..c775dfd1f88 100644 --- a/tests/unit/repositories/test_daily_activity_sql.py +++ b/tests/unit/repositories/test_daily_activity_sql.py @@ -67,7 +67,7 @@ def test_where_clause_binds_each_filter_as_a_single_array_parameter() -> None: assert sql == ( 'date >= $1 AND date <= $2 AND "user_id" = ANY($3::text[]) ' - 'AND NOT ("user_id" = ANY($4::text[])) AND model = $5 AND api_key = ANY($6::text[])' + 'AND ("user_id" IS NULL OR NOT ("user_id" = ANY($4::text[]))) AND model = $5 AND api_key = ANY($6::text[])' ) assert params == ( "2026-01-01", @@ -79,6 +79,15 @@ def test_where_clause_binds_each_filter_as_a_single_array_parameter() -> None: ) +def test_where_clause_exclusion_keeps_null_entity_rows() -> None: + scope = _scope(table=DailyActivityTable.TEAM, entity_ids=None, exclude_entity_ids=("litellm-dashboard",)) + + sql, params = build_where_clause(scope) + + assert sql == 'date >= $1 AND date <= $2 AND ("team_id" IS NULL OR NOT ("team_id" = ANY($3::text[])))' + assert params == ("2026-01-01", "2026-01-31", ["litellm-dashboard"]) + + @pytest.mark.parametrize( ("entity_ids", "api_keys", "expected_sql", "expected_params"), [