Coverage for agentos/health/__init__.py: 32%

1"""AgentOS health checks — readiness, liveness, and dependency probes. 

2 

3Provides standard health-check endpoints for Kubernetes, Docker, and load balancers. 

4""" 

5 

6from __future__ import annotations 

7 

8import time 

9from collections.abc import Callable 

10from dataclasses import dataclass 

11from enum import Enum 

12 

13 

14class HealthStatus(Enum): 

15 """健康状态枚举。""" 

16 

17 HEALTHY = "healthy" 

18 DEGRADED = "degraded" 

19 UNHEALTHY = "unhealthy" 

20 

21 

22@dataclass 

23class HealthCheck: 

24 """健康检查定义。""" 

25 

26 name: str 

27 check_fn: Callable[[], bool] 

28 timeout_seconds: float = 5.0 

29 description: str = "" 

30 

31 

32@dataclass 

33class CheckResult: 

34 """检查结果。""" 

35 

36 name: str 

37 status: HealthStatus 

38 latency_ms: float 

39 message: str = "" 

40 

41 def to_dict(self) -> dict: 

42 return { 

43 "name": self.name, 

44 "status": self.status.value, 

45 "latency_ms": round(self.latency_ms, 2), 

46 "message": self.message, 

47 } 

48 

49 

50class HealthChecker: 

51 """Aggregate readiness and liveness checks.""" 

52 

53 def __init__(self): 

54 self._readiness_checks: list[HealthCheck] = [] 

55 self._liveness_checks: list[HealthCheck] = [] 

56 

57 def add_readiness(self, check: HealthCheck): 

58 self._readiness_checks.append(check) 

59 

60 def add_liveness(self, check: HealthCheck): 

61 self._liveness_checks.append(check) 

62 

63 def _run_checks(self, checks: list[HealthCheck]) -> tuple[HealthStatus, list[CheckResult]]: 

64 results: list[CheckResult] = [] 

65 overall = HealthStatus.HEALTHY 

66 for chk in checks: 

67 start = time.monotonic() 

68 try: 

69 ok = chk.check_fn() 

70 except Exception as e: 

71 ok = False 

72 msg = str(e) 

73 else: 

74 msg = "ok" if ok else "check returned False" 

75 latency = (time.monotonic() - start) * 1000 

76 status = HealthStatus.HEALTHY if ok else HealthStatus.UNHEALTHY 

77 if status == HealthStatus.UNHEALTHY and overall != HealthStatus.UNHEALTHY: 

78 overall = HealthStatus.DEGRADED 

79 if status == HealthStatus.UNHEALTHY: 

80 overall = HealthStatus.UNHEALTHY 

81 results.append( 

82 CheckResult(name=chk.name, status=status, latency_ms=latency, message=msg) 

83 ) 

84 return (overall, results) 

85 

86 def readiness(self) -> dict: 

87 """Run all readiness checks. Returns a dict suitable for a /health/ready endpoint.""" 

88 overall, results = self._run_checks(self._readiness_checks) 

89 return { 

90 "status": overall.value, 

91 "timestamp": time.time(), 

92 "checks": [r.to_dict() for r in results], 

93 } 

94 

95 def liveness(self) -> dict: 

96 """Run all liveness checks. Returns a dict suitable for a /health/live endpoint.""" 

97 overall, results = self._run_checks(self._liveness_checks) 

98 return { 

99 "status": overall.value, 

100 "timestamp": time.time(), 

101 "checks": [r.to_dict() for r in results], 

102 } 

103 

104 def all(self) -> dict: 

105 """Combined readiness + liveness report, suitable for /health.""" 

106 r = self.readiness() 

107 lv = self.liveness() 

108 combined_status = HealthStatus.HEALTHY 

109 for s in (r["status"], lv["status"]): 

110 if s == HealthStatus.UNHEALTHY.value: 

111 combined_status = HealthStatus.UNHEALTHY 

112 break 

113 if s == HealthStatus.DEGRADED.value: 

114 combined_status = HealthStatus.DEGRADED 

115 return { 

116 "status": combined_status.value, 

117 "timestamp": time.time(), 

118 "readiness": r, 

119 "liveness": lv, 

120 } 

121 

122 

123# ── Built-in checks ─────────────────────────────────────────────────────────── 

124 

125 

126def check_openai_connectivity(api_key: str | None = None) -> HealthCheck: 

127 """Verify connectivity to the OpenAI API.""" 

128 

129 def _check() -> bool: 

130 try: 

131 import urllib.request 

132 

133 req = urllib.request.Request("https://api.openai.com/v1/models", method="HEAD") 

134 if api_key: 

135 req.add_header("Authorization", f"Bearer {api_key}") 

136 urllib.request.urlopen(req, timeout=5) 

137 return True 

138 except Exception: 

139 return False 

140 

141 return HealthCheck( 

142 name="openai-connectivity", 

143 check_fn=_check, 

144 timeout_seconds=5.0, 

145 description="Check OpenAI API reachability", 

146 ) 

147 

148 

149def check_vectorstore_health(db_instance=None) -> HealthCheck: 

150 """Check vector store connection health.""" 

151 

152 def _check() -> bool: 

153 if db_instance is None: 

154 return False 

155 try: 

156 return hasattr(db_instance, "is_healthy") and db_instance.is_healthy() 

157 except Exception: 

158 return False 

159 

160 return HealthCheck( 

161 name="vectorstore-health", 

162 check_fn=_check, 

163 timeout_seconds=5.0, 

164 description="Check vector store connection", 

165 ) 

166 

167 

168def check_disk_space(threshold_bytes: int = 100 * 1024 * 1024) -> HealthCheck: 

169 """Check available disk space exceeds threshold (default 100MB).""" 

170 

171 def _check() -> bool: 

172 import shutil 

173 

174 usage = shutil.disk_usage("/") 

175 return usage.free >= threshold_bytes 

176 

177 return HealthCheck( 

178 name="disk-space", 

179 check_fn=_check, 

180 timeout_seconds=1.0, 

181 description=f"Free disk space >= {threshold_bytes/1024/1024:.0f}MB", 

182 ) 

183 

184 

185def check_memory(threshold_bytes: int = 50 * 1024 * 1024) -> HealthCheck: 

186 """Check available system memory exceeds threshold (default 50MB).""" 

187 

188 def _check() -> bool: 

189 try: 

190 with open("/proc/meminfo") as f: 

191 for line in f: 

192 if line.startswith("MemAvailable:"): 

193 kb = int(line.split()[1]) 

194 return kb * 1024 >= threshold_bytes 

195 except Exception: 

196 return True # can't check, assume OK 

197 return True 

198 

199 return HealthCheck( 

200 name="memory", 

201 check_fn=_check, 

202 timeout_seconds=1.0, 

203 description=f"Available memory >= {threshold_bytes/1024/1024:.0f}MB", 

204 ) 

205 

206 

207# ── Default health checker factory ──────────────────────────────────────────── 

208 

209 

210def create_default_health_checker() -> HealthChecker: 

211 """Return a HealthChecker pre-loaded with sensible built-in checks.""" 

212 hc = HealthChecker() 

213 hc.add_liveness(check_memory()) 

214 hc.add_readiness(check_disk_space()) 

215 return hc