diff --git a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py index 2a928db50ec..c6f58ab0b2a 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py +++ b/litellm-proxy-extras/litellm_proxy_extras/request_log_indexes.py @@ -95,8 +95,10 @@ _CREATE_INDEX_STATEMENT: Final = re.compile( ) _TABLE_KIND_SQL: Final = "SELECT c.relkind = 'p' AS partitioned FROM pg_class c WHERE c.oid = to_regclass(%s)" _TABLE_BYTES_SQL: Final = ( - "SELECT COALESCE(SUM(pg_table_size(c.oid)), 0)::bigint AS total FROM pg_class c WHERE c.relkind = 'r' " - "AND (c.oid = to_regclass(%(table)s) OR c.oid IN (SELECT relid FROM pg_partition_tree(to_regclass(%(table)s))))" + "WITH RECURSIVE tree(oid) AS (SELECT to_regclass(%(table)s)::oid " + "UNION ALL SELECT i.inhrelid FROM pg_inherits i JOIN tree ON i.inhparent = tree.oid) " + "SELECT COALESCE(SUM(pg_table_size(c.oid)), 0)::bigint AS total " + "FROM tree JOIN pg_class c ON c.oid = tree.oid WHERE c.relkind = 'r'" ) _CHILDREN_WITHOUT_THE_INDEX_SQL: Final = ( "SELECT child.relname AS name, n.nspname AS schema, child.relkind = 'p' AS partitioned " @@ -290,12 +292,16 @@ def _too_large_to_build( return False logger.warning( "Not building index %s on %s: the table takes %s on disk, more than the %s it builds by itself. Build it " - "off-peak by running each of these statements on its own, and the next index build finds it in place: %s", + "off-peak by running each of these statements on its own, for example with psql -f, and the next index build " + "finds it in place:\n%s", index.name, index.table, _gib(size), _gib(index.max_bytes), - "; ".join(_operator_build_statements(connection, schema, index, index.table, index.name, partitioned)), + "".join( + f"{statement};\n" + for statement in _operator_build_statements(connection, schema, index, index.table, index.name, partitioned) + ), ) return True diff --git a/litellm/proxy/spend_tracking/spend_management_endpoints.py b/litellm/proxy/spend_tracking/spend_management_endpoints.py index 09fc2285f51..59ccfe6ea12 100644 --- a/litellm/proxy/spend_tracking/spend_management_endpoints.py +++ b/litellm/proxy/spend_tracking/spend_management_endpoints.py @@ -99,6 +99,10 @@ _BATCH_CALL_TYPES_SQL: Final = "('acreate_batch', 'create_batch', 'aretrieve_bat _ERROR_CODE_SQL: Final = "metadata->'error_information'->>'error_code'" _INDEXED_ERROR_CODE_PREFIX_LENGTH: Final = 128 _INDEXED_ERROR_CODE_PREFIX_SQL: Final = f"left({_ERROR_CODE_SQL}, {_INDEXED_ERROR_CODE_PREFIX_LENGTH})" +_ERROR_CODE_INDEX_READY_SQL: Final = ( + "SELECT EXISTS (SELECT 1 FROM pg_index " + "WHERE indexrelid = to_regclass('\"LiteLLM_SpendLogs_error_code_startTime_idx\"') AND indisvalid) AS ready" +) _SPAN_TYPE_SQL_CONDITIONS: Final[Mapping[str, str]] = MappingProxyType( { "mcp": f"call_type IN {_MCP_CALL_TYPES_SQL}", @@ -229,7 +233,18 @@ class _SessionSpendRow(TypedDict): session_models: ReadOnly[Sequence[str]] -def _error_code_condition(placeholder: str, error_code: str) -> str: +class _IndexReadyRow(TypedDict): + ready: ReadOnly[bool] + + +async def _error_code_index_is_ready(prisma_client: PrismaClient) -> bool: + rows: Final[Sequence[_IndexReadyRow]] = await _query_raw(prisma_client, _ERROR_CODE_INDEX_READY_SQL) + return bool(rows) and rows[0]["ready"] is True + + +def _error_code_condition(placeholder: str, error_code: str, index_ready: bool) -> str: + if not index_ready: + return f"{_ERROR_CODE_SQL} = {placeholder}" if len(error_code) < _INDEXED_ERROR_CODE_PREFIX_LENGTH: return f"{_INDEXED_ERROR_CODE_PREFIX_SQL} = {placeholder}" return ( @@ -2947,7 +2962,9 @@ async def ui_view_spend_logs( sql_params.append(f"%{key_alias}%") p += 1 if error_code is not None: - sql_conditions.append(_error_code_condition(f"${p}", error_code)) + sql_conditions.append( + _error_code_condition(f"${p}", error_code, await _error_code_index_is_ready(prisma_client)) + ) sql_params.append(error_code) p += 1 if error_message is not None: diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index 7e489a35706..b21873fbb52 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -935,7 +935,7 @@ def _seed_spend_logs(database_url: str, partitioned: bool) -> None: def _run_the_logged_statements(database_url: str, messages: list[str]) -> None: logged: Final = next(message for message in messages if "Not building index" in message) with psycopg.connect(database_url, autocommit=True) as conn: - for statement in logged.split("finds it in place: ", 1)[1].split("; "): + for statement in logged.split("finds it in place:\n", 1)[1].splitlines(): conn.execute(statement.encode()) diff --git a/tests/integration/spend/test_spend_log_error_filters.py b/tests/integration/spend/test_spend_log_error_filters.py index a1e3b97e5db..8a8a7377240 100644 --- a/tests/integration/spend/test_spend_log_error_filters.py +++ b/tests/integration/spend/test_spend_log_error_filters.py @@ -1,6 +1,6 @@ import secrets import time -from collections.abc import Mapping +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import datetime, timedelta, timezone from itertools import product @@ -19,6 +19,7 @@ from psycopg.types.json import Jsonb pytestmark: Final = pytest.mark.timeout(300) ERROR_INDEX: Final = "LiteLLM_SpendLogs_error_code_startTime_idx" +HAND_BUILT_INDEX: Final = "spend_logs_error_code_by_hand_idx" INDEXES_IN_PLACE: Final = "Request-log indexes are all in place" SUCCESS_ROWS: Final = 5000 STATS_FLUSH_WINDOW_SECONDS: Final = 11.0 @@ -58,11 +59,11 @@ ERROR_CODES: Final = MappingProxyType( LONG_CODE_PREFIX[:-1]: frozenset(), } ) +SHORT_ERROR_CODES: Final = ("500", "499") SESSION_VIEWS: Final = ("false", "true") -FILTERS: Final = tuple(product(SESSION_VIEWS, ERROR_CODES)) -def _seed(database_url: str, started: datetime) -> None: +def _seed(database_url: str, started: datetime, error_rows: tuple[ErrorRow, ...]) -> None: utc: Final = started.astimezone(timezone.utc).replace(tzinfo=None) with psycopg.connect(database_url, autocommit=True) as connection: connection.execute( @@ -92,31 +93,41 @@ def _seed(database_url: str, started: datetime) -> None: ), row.status, ) - for row in ERROR_ROWS + for row in error_rows ), ) connection.execute('ANALYZE "LiteLLM_SpendLogs"') -def _error_index_scans(database_url: str) -> int: +def _index_scans(database_url: str, index: str) -> int: with psycopg.connect(database_url) as connection, connection.cursor(row_factory=class_row(_ScanRow)) as cursor: row: Final = cursor.execute( - "SELECT idx_scan AS scans FROM pg_stat_user_indexes WHERE indexrelname = %s", (ERROR_INDEX,) + "SELECT idx_scan AS scans FROM pg_stat_user_indexes WHERE indexrelname = %s", (index,) ).fetchone() - assert row is not None, f"{ERROR_INDEX} was not built" + assert row is not None, f"{index} was not built" return row.scans -def _settled_error_index_scans(database_url: str, previous: int, unchanged_since: float) -> int: +def _settled_index_scans(database_url: str, index: str, previous: int, unchanged_since: float) -> int: changed: Final = eventually( - lambda: _error_index_scans(database_url) != previous, + lambda: _index_scans(database_url, index) != previous, lambda drifted: drifted, seconds=STATS_FLUSH_WINDOW_SECONDS - (time.monotonic() - unchanged_since), return_last_on_timeout=True, ) if not changed: return previous - return _settled_error_index_scans(database_url, _error_index_scans(database_url), time.monotonic()) + return _settled_index_scans(database_url, index, _index_scans(database_url, index), time.monotonic()) + + +def _replace_the_error_index_with_one_on_the_full_error_code(database_url: str) -> None: + with psycopg.connect(database_url, autocommit=True) as connection: + connection.execute(f'DROP INDEX "{ERROR_INDEX}"') + connection.execute( + f'CREATE INDEX "{HAND_BUILT_INDEX}" ON "LiteLLM_SpendLogs" ' + "((metadata->'error_information'->>'error_code'), \"startTime\")" + ) + connection.execute('ANALYZE "LiteLLM_SpendLogs"') def _listed(gateway: Gateway, window: Mapping[str, str], query: Mapping[str, str]) -> frozenset[str]: @@ -127,6 +138,7 @@ def _listed(gateway: Gateway, window: Mapping[str, str], query: Mapping[str, str def _list_each_filter( gateway: Gateway, database_url: str, + index: str, window: Mapping[str, str], filters: tuple[tuple[str, str], ...], scans: int, @@ -135,14 +147,39 @@ def _list_each_filter( return () session, code = filters[0] listed: Final = _listed(gateway, window, {"error_code": code, "group_by_session": session}) - settled: Final = _settled_error_index_scans(database_url, scans, time.monotonic()) + settled: Final = _settled_index_scans(database_url, index, scans, time.monotonic()) return ( ((session, code), listed, settled - scans), - *_list_each_filter(gateway, database_url, window, filters[1:], settled), + *_list_each_filter(gateway, database_url, index, window, filters[1:], settled), ) -def test_error_code_filters_read_only_the_rows_that_carry_an_error_code(gateway: Gateway, tmp_path: Path) -> None: +def _no_change(_: str) -> None: + return None + + +@dataclass(frozen=True, slots=True) +class IndexCase: + index: str + error_rows: tuple[ErrorRow, ...] + codes: tuple[str, ...] + prepare: Callable[[str], None] + + +INDEX_CASES: Final = ( + IndexCase(ERROR_INDEX, ERROR_ROWS, tuple(ERROR_CODES), _no_change), + IndexCase( + HAND_BUILT_INDEX, + tuple(row for row in ERROR_ROWS if row.error_code in ("500", "429", "499")), + SHORT_ERROR_CODES, + _replace_the_error_index_with_one_on_the_full_error_code, + ), +) + + +def _filter_through_index( + gateway: Gateway, tmp_path: Path, case: IndexCase, filters: tuple[tuple[str, str], ...] +) -> tuple[tuple[tuple[str, str], frozenset[str], int], ...]: now: Final = datetime.now(timezone.utc) window: Final = { "start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d %H:%M:%S"), @@ -155,14 +192,23 @@ def test_error_code_filters_read_only_the_rows_that_carry_an_error_code(gateway: eventually( lambda: booted.log.read_text(errors="replace"), lambda text: INDEXES_IN_PLACE in text, seconds=60 ) - _seed(database_url, now - timedelta(hours=1)) - scans: Final = _settled_error_index_scans( - database_url, _error_index_scans(database_url), time.monotonic() + _seed(database_url, now - timedelta(hours=1), case.error_rows) + case.prepare(database_url) + scans: Final = _settled_index_scans( + database_url, case.index, _index_scans(database_url, case.index), time.monotonic() ) + return _list_each_filter(booted.gateway, database_url, case.index, window, filters, scans) - results: Final = _list_each_filter(booted.gateway, database_url, window, FILTERS, scans) - assert {combo: listed for combo, listed, _ in results} == { - (session, code): ERROR_CODES[code] for session, code in FILTERS - } - assert {combo: index_scans > 0 for combo, _, index_scans in results} == {combo: True for combo in FILTERS} +@pytest.mark.parametrize("case", INDEX_CASES, ids=("built_by_the_index_job", "built_by_hand_on_the_full_error_code")) +def test_error_code_filters_read_only_the_rows_that_carry_an_error_code( + gateway: Gateway, tmp_path: Path, case: IndexCase +) -> None: + filters: Final = tuple(product(SESSION_VIEWS, case.codes)) + + results: Final = _filter_through_index(gateway, tmp_path, case, filters) + + assert {combo: listed for combo, listed, _ in results} == { + (session, code): ERROR_CODES[code] for session, code in filters + } + assert {combo: index_scans > 0 for combo, _, index_scans in results} == {combo: True for combo in filters} diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 985dadb9557..49a30019b4c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -208,7 +208,9 @@ def _reconstruct_ui_where_from_sql(sql_query, params): return where -def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=None, query_observer=None): +def make_ui_spend_logs_mock_prisma( + mock_spend_logs, filter_fn, team_lookup_fn=None, query_observer=None, error_code_index_ready=False +): """ Create a MockPrismaClient for /spend/logs/ui endpoint tests. @@ -239,6 +241,8 @@ def make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_fn, team_lookup_fn=No query_observer(sql_query, params) if "mcp_tool_call_count" in sql_query: return [] + if "FROM pg_index" in sql_query: + return [{"ready": error_code_index_ready}] if 'SELECT DISTINCT "user", team_id' in sql_query: return _emulate_spend_log_owner_lookup(mock_spend_logs, sql_query, params) filtered = filter_fn(_reconstruct_ui_where_from_sql(sql_query, params)) @@ -4698,7 +4702,8 @@ async def test_view_spend_logs_summarize_unhashed_api_key_without_padding(client @pytest.mark.asyncio -async def test_ui_view_spend_logs_with_error_code(client): +@pytest.mark.parametrize("error_code_index_ready", [True, False], ids=["with_the_index", "without_the_index"]) +async def test_ui_view_spend_logs_with_error_code(client, error_code_index_ready): """Test filtering spend logs by error code""" mock_spend_logs = [ { @@ -4744,7 +4749,9 @@ async def test_ui_view_spend_logs_with_error_code(client): with patch.object( ps, "prisma_client", - make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_error_code), + make_ui_spend_logs_mock_prisma( + mock_spend_logs, filter_by_error_code, error_code_index_ready=error_code_index_ready + ), ): start_date, end_date = _default_date_range() @@ -4773,7 +4780,10 @@ async def test_ui_view_spend_logs_with_error_code(client): @pytest.mark.asyncio -async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact_code(client: TestClient) -> None: +@pytest.mark.parametrize("error_code_index_ready", [True, False], ids=["with_the_index", "without_the_index"]) +async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact_code( + client: TestClient, error_code_index_ready: bool +) -> None: prefix: Final = "E" * 128 codes: Final = MappingProxyType({f"log-{suffix}": prefix + suffix for suffix in ("first", "second")}) mock_spend_logs: Final = [ @@ -4798,7 +4808,11 @@ async def test_ui_view_spend_logs_with_a_long_error_code_returns_only_that_exact ) try: with patch.object( - ps, "prisma_client", make_ui_spend_logs_mock_prisma(mock_spend_logs, filter_by_exact_error_code) + ps, + "prisma_client", + make_ui_spend_logs_mock_prisma( + mock_spend_logs, filter_by_exact_error_code, error_code_index_ready=error_code_index_ready + ), ): start_date, end_date = _default_date_range() response = client.get(