From 9ebf480b71ce036517d28dfb5cc878690d730c01 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 24 Jul 2026 18:35:20 +0000 Subject: [PATCH] ci(spend-logs): put the spend logs query budget under the budget ratchet Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- scripts/budget_ratchet_check.py | 7 +- spend-logs-query-budget.json | 100 +++++++++++++----- .../check_spend_logs_query_budget.py | 7 +- .../test_check_spend_logs_query_budget.py | 28 ++++- 4 files changed, 107 insertions(+), 35 deletions(-) diff --git a/scripts/budget_ratchet_check.py b/scripts/budget_ratchet_check.py index 10a78483643..c175f98b29f 100644 --- a/scripts/budget_ratchet_check.py +++ b/scripts/budget_ratchet_check.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 """Non-gating ratchet guard: budget limits may only fall, never rise. -Every `*-budget.json` file (ruff-strict, type-discipline, basedpyright-code) is a -one-way ratchet: each rule's ceiling is its `limit`, and that limit is meant to be -driven DOWN over time. This check compares every budget file against its own +Every `*-budget.json` file (ruff-strict, type-discipline, basedpyright-code, +spend-logs-query) is a one-way ratchet: each rule's ceiling is its `limit`, and that +limit is meant to be driven DOWN over time. This check compares every budget file against its own content at the merge-base with the target branch and fails (exits 1, red) if: * a rule's `limit` went up, @@ -39,6 +39,7 @@ DEFAULT_BUDGETS: tuple[str, ...] = ( "ruff-strict-budget.json", "type-discipline-budget.json", "basedpyright-code-budget.json", + "spend-logs-query-budget.json", ) diff --git a/spend-logs-query-budget.json b/spend-logs-query-budget.json index 31cab56ddba..4ec23f5192c 100644 --- a/spend-logs-query-budget.json +++ b/spend-logs-query-budget.json @@ -1,27 +1,77 @@ { - "litellm/integrations/cloudzero/database.py": 1, - "litellm/integrations/focus/database.py": 2, - "litellm/models/__init__.py": 1, - "litellm/proxy/_types.py": 1, - "litellm/proxy/agent_endpoints/endpoints.py": 1, - "litellm/proxy/analytics_endpoints/analytics_endpoints.py": 1, - "litellm/proxy/db/create_views.py": 7, - "litellm/proxy/db/db_spend_update_writer.py": 12, - "litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py": 1, - "litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py": 1, - "litellm/proxy/management_endpoints/common_daily_activity.py": 12, - "litellm/proxy/management_endpoints/customer_endpoints.py": 1, - "litellm/proxy/management_endpoints/internal_user_endpoints.py": 2, - "litellm/proxy/management_endpoints/organization_endpoints.py": 1, - "litellm/proxy/management_endpoints/tag_management_endpoints.py": 1, - "litellm/proxy/management_endpoints/team_endpoints.py": 1, - "litellm/proxy/management_endpoints/tool_management_endpoints.py": 2, - "litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py": 3, - "litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py": 5, - "litellm/proxy/proxy_server.py": 4, - "litellm/proxy/spend_tracking/spend_management_endpoints.py": 25, - "litellm/proxy/spend_tracking/spend_tracking_utils.py": 2, - "litellm/proxy/utils.py": 1, - "litellm/repositories/table_repositories.py": 6, - "litellm/responses/litellm_completion_transformation/session_handler.py": 1 + "litellm/integrations/cloudzero/database.py": { + "limit": 1 + }, + "litellm/integrations/focus/database.py": { + "limit": 2 + }, + "litellm/models/__init__.py": { + "limit": 1 + }, + "litellm/proxy/_types.py": { + "limit": 1 + }, + "litellm/proxy/agent_endpoints/endpoints.py": { + "limit": 1 + }, + "litellm/proxy/analytics_endpoints/analytics_endpoints.py": { + "limit": 1 + }, + "litellm/proxy/db/create_views.py": { + "limit": 7 + }, + "litellm/proxy/db/db_spend_update_writer.py": { + "limit": 12 + }, + "litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py": { + "limit": 1 + }, + "litellm/proxy/db/db_transaction_queue/spend_logs_partition_manager.py": { + "limit": 1 + }, + "litellm/proxy/management_endpoints/common_daily_activity.py": { + "limit": 12 + }, + "litellm/proxy/management_endpoints/customer_endpoints.py": { + "limit": 1 + }, + "litellm/proxy/management_endpoints/internal_user_endpoints.py": { + "limit": 2 + }, + "litellm/proxy/management_endpoints/organization_endpoints.py": { + "limit": 1 + }, + "litellm/proxy/management_endpoints/tag_management_endpoints.py": { + "limit": 1 + }, + "litellm/proxy/management_endpoints/team_endpoints.py": { + "limit": 1 + }, + "litellm/proxy/management_endpoints/tool_management_endpoints.py": { + "limit": 2 + }, + "litellm/proxy/management_endpoints/usage_endpoints/ai_usage_chat.py": { + "limit": 3 + }, + "litellm/proxy/management_endpoints/user_agent_analytics_endpoints.py": { + "limit": 5 + }, + "litellm/proxy/proxy_server.py": { + "limit": 4 + }, + "litellm/proxy/spend_tracking/spend_management_endpoints.py": { + "limit": 25 + }, + "litellm/proxy/spend_tracking/spend_tracking_utils.py": { + "limit": 2 + }, + "litellm/proxy/utils.py": { + "limit": 1 + }, + "litellm/repositories/table_repositories.py": { + "limit": 6 + }, + "litellm/responses/litellm_completion_transformation/session_handler.py": { + "limit": 1 + } } diff --git a/tests/code_coverage_tests/check_spend_logs_query_budget.py b/tests/code_coverage_tests/check_spend_logs_query_budget.py index 0d6b50f400d..f5cdb32a56c 100644 --- a/tests/code_coverage_tests/check_spend_logs_query_budget.py +++ b/tests/code_coverage_tests/check_spend_logs_query_budget.py @@ -96,12 +96,13 @@ def over_budget(counts: dict[str, int], budget: dict[str, int]) -> dict[str, int def load_budget(budget_path: Path) -> dict[str, int]: parsed: object = json.loads(budget_path.read_text(encoding="utf-8")) if not isinstance(parsed, dict): - raise ValueError(f"{budget_path.name} must be an object mapping file paths to query counts") - return {str(path): int(limit) for path, limit in parsed.items()} + raise ValueError(f"{budget_path.name} must be an object mapping file paths to {{'limit': int}}") + return {str(path): int(spec["limit"]) for path, spec in parsed.items()} def write_budget(budget_path: Path, counts: dict[str, int]) -> None: - _ = budget_path.write_text(json.dumps(counts, indent=2, sort_keys=True) + "\n", encoding="utf-8") + specs = {path: {"limit": count} for path, count in counts.items()} + _ = budget_path.write_text(json.dumps(specs, indent=2, sort_keys=True) + "\n", encoding="utf-8") def _report(sites: tuple[QuerySite, ...], violations: dict[str, int], budget: dict[str, int]) -> None: diff --git a/tests/test_litellm/test_check_spend_logs_query_budget.py b/tests/test_litellm/test_check_spend_logs_query_budget.py index b12517be92a..d53395ff7a3 100644 --- a/tests/test_litellm/test_check_spend_logs_query_budget.py +++ b/tests/test_litellm/test_check_spend_logs_query_budget.py @@ -8,16 +8,18 @@ the two things that would silently defeat it: detecting every shape of query we prose mentions, and failing the budget comparison for files that grow or appear. """ +import json import os import sys from pathlib import Path _CODE_COVERAGE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "code_coverage_tests") -sys.path.insert(0, _CODE_COVERAGE_DIR) - -import check_spend_logs_query_budget as guard # noqa: E402 - _REPO_ROOT = Path(_CODE_COVERAGE_DIR).resolve().parents[1] +sys.path.insert(0, _CODE_COVERAGE_DIR) +sys.path.insert(0, str(_REPO_ROOT / "scripts")) + +import budget_ratchet_check as ratchet # noqa: E402 +import check_spend_logs_query_budget as guard # noqa: E402 def _scan(tmp_path: Path, source: str) -> tuple[guard.QuerySite, ...]: @@ -79,6 +81,24 @@ def test_over_budget_flags_new_file_and_growth_but_not_shrinkage(): assert guard.over_budget(counts, budget) == {"a.py": 2, "b.py": 1} +def test_budget_file_uses_the_ratchet_schema(tmp_path): + budget_path = tmp_path / "spend-logs-query-budget.json" + guard.write_budget(budget_path, {"litellm/module.py": 3}) + assert json.loads(budget_path.read_text(encoding="utf-8")) == {"litellm/module.py": {"limit": 3}} + assert guard.load_budget(budget_path) == {"litellm/module.py": 3} + + +def test_committed_budget_is_watched_by_the_ratchet(): + assert "spend-logs-query-budget.json" in ratchet.DEFAULT_BUDGETS + budget = json.loads((_REPO_ROOT / "spend-logs-query-budget.json").read_text(encoding="utf-8")) + raised = ratchet.regressions_for( + "spend-logs-query-budget.json", + budget, + {path: {"limit": spec["limit"] + 1} for path, spec in budget.items()}, + ) + assert len(raised) == len(budget) + + def test_committed_budget_matches_the_repository(): counts = guard.counts_by_file(guard.scan_sites(_REPO_ROOT)) budget = guard.load_budget(_REPO_ROOT / "spend-logs-query-budget.json")