Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/health_check_latest.py: 86%

1""" 

2Latest health-check row per model, deduplicated by Postgres. 

3 

4prisma-client-py's ``find_many(distinct=...)`` dedups client-side: the emitted 

5SQL carries no DISTINCT, so the whole append-only history table streams to the 

6worker on every call. ``SELECT DISTINCT ON`` keeps the transfer at one row per 

7(model_id, model_name) and is served by the matching descending index. 

8""" 

9 

10from __future__ import annotations 

11 

12import json 

13from collections.abc import Sequence 

14from datetime import datetime, timezone 

15from typing import TYPE_CHECKING, Final 

16 

17from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, field_validator 

18 

19from litellm._logging import verbose_proxy_logger 

20 

21if TYPE_CHECKING: 21 ↛ 22line 21 didn't jump to line 22 because the condition on line 21 was never true

22 from litellm.proxy.utils import PrismaClient 

23 

24LATEST_HEALTH_CHECKS_SQL: Final = """ 

25SELECT DISTINCT ON ("model_id", "model_name") 

26 "health_check_id", "model_name", "model_id", "status", 

27 "healthy_count", "unhealthy_count", "error_message", 

28 "response_time_ms", "details", "checked_by", 

29 "checked_at", "created_at", "updated_at" 

30FROM "LiteLLM_HealthCheckTable" 

31ORDER BY "model_id" ASC, "model_name" ASC, "checked_at" DESC 

32""" 

33 

34LATEST_HEALTH_CHECKS_FOR_MODELS_SQL: Final = """ 

35SELECT DISTINCT ON ("model_id", "model_name") 

36 "health_check_id", "model_name", "model_id", "status", 

37 "healthy_count", "unhealthy_count", "error_message", 

38 "response_time_ms", "details", "checked_by", 

39 "checked_at", "created_at", "updated_at" 

40FROM "LiteLLM_HealthCheckTable" 

41WHERE "model_name" = ANY($1) 

42ORDER BY "model_id" ASC, "model_name" ASC, "checked_at" DESC 

43""" 

44 

45 

46class LatestHealthCheckRow(BaseModel): 

47 model_config = ConfigDict(frozen=True, protected_namespaces=()) 

48 

49 health_check_id: str 

50 model_name: str 

51 model_id: str | None = None 

52 status: str 

53 healthy_count: int = 0 

54 unhealthy_count: int = 0 

55 error_message: str | None = None 

56 response_time_ms: float | None = None 

57 details: JsonValue | None = None 

58 checked_by: str | None = None 

59 checked_at: datetime 

60 created_at: datetime 

61 updated_at: datetime 

62 

63 @field_validator("details", mode="before") 

64 @classmethod 

65 def _decode_json_text(cls, value: object) -> object: 

66 return json.loads(value) if isinstance(value, str) else value 

67 

68 @field_validator("checked_at", "created_at", "updated_at") 

69 @classmethod 

70 def _assume_utc(cls, value: datetime) -> datetime: 

71 return value.replace(tzinfo=timezone.utc) if value.tzinfo is None else value 

72 

73 

74_ROWS_ADAPTER: Final = TypeAdapter(tuple[LatestHealthCheckRow, ...]) 

75 

76 

77async def query_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]: 

78 rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_SQL) 

79 return _ROWS_ADAPTER.validate_python(rows) 

80 

81 

82async def fetch_latest_health_checks(prisma_client: PrismaClient) -> tuple[LatestHealthCheckRow, ...]: 

83 try: 

84 return await query_latest_health_checks(prisma_client) 

85 except Exception as query_err: # noqa: BLE001 # health decorates other reads; a driver error must not fail them 

86 verbose_proxy_logger.error("Error getting all latest health checks: %s", query_err) 

87 return () 

88 

89 

90async def fetch_latest_health_checks_for_models( 

91 prisma_client: PrismaClient, model_names: Sequence[str] 

92) -> tuple[LatestHealthCheckRow, ...]: 

93 if not model_names: 

94 return () 

95 try: 

96 rows: Final = await prisma_client.db.query_raw(LATEST_HEALTH_CHECKS_FOR_MODELS_SQL, list(model_names)) 

97 return _ROWS_ADAPTER.validate_python(rows) 

98 except Exception as query_err: # noqa: BLE001 # a paged model list must not fail on its health decoration 

99 verbose_proxy_logger.error("Error getting latest health checks for models: %s", query_err) 

100 return ()