Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-llm/src/lexigram/ai/llm/routing/backends/database.py: 29%

1"""Database-backed quota backend for LLM routing. 

2 

3Persists per-provider daily quota state to the ``provider_daily_usage`` 

4table via ``DatabaseProviderProtocol``. Suitable for multi-process and 

5production deployments. 

6 

7All SQL errors are absorbed so that quota tracking failures never interrupt 

8inference. If the DB is unreachable, methods return safe defaults 

9(is_exhausted → False, increment → no-op). 

10""" 

11 

12from __future__ import annotations 

13 

14from datetime import UTC, date, datetime 

15from typing import TYPE_CHECKING 

16 

17from lexigram.ai.llm.routing.backends._time import end_of_utc_day 

18from lexigram.ai.llm.routing.types import ProviderUsage 

19from lexigram.logging import ( 

20 get_logger, 

21) 

22 

23if TYPE_CHECKING: 

24 from lexigram.contracts.data import DatabaseProviderProtocol 

25 

26logger = get_logger(__name__) 

27 

28__all__ = ["DatabaseQuotaBackend"] 

29 

30_INCREMENT_SQL = """ 

31INSERT INTO provider_daily_usage (provider, usage_date, success_count) 

32VALUES ($1, $2, 1) 

33ON CONFLICT (provider, usage_date) 

34DO UPDATE SET 

35 success_count = provider_daily_usage.success_count + 1, 

36 updated_at = NOW() 

37""" 

38 

39_MARK_EXHAUSTED_SQL = """ 

40INSERT INTO provider_daily_usage (provider, usage_date, is_exhausted, exhausted_until) 

41VALUES ($1, $2, TRUE, $3) 

42ON CONFLICT (provider, usage_date) 

43DO UPDATE SET 

44 is_exhausted = TRUE, 

45 exhausted_until = $3, 

46 updated_at = NOW() 

47""" 

48 

49_RECORD_ERROR_SQL = """ 

50INSERT INTO provider_daily_usage (provider, usage_date, error_count) 

51VALUES ($1, $2, 1) 

52ON CONFLICT (provider, usage_date) 

53DO UPDATE SET 

54 error_count = provider_daily_usage.error_count + 1, 

55 updated_at = NOW() 

56""" 

57 

58_GET_TODAY_SQL = """ 

59SELECT provider, usage_date::text, success_count, error_count, is_exhausted, 

60 exhausted_until 

61FROM provider_daily_usage 

62WHERE provider = $1 AND usage_date = $2 

63""" 

64 

65_GET_ALL_TODAY_SQL = """ 

66SELECT provider, usage_date::text, success_count, error_count, is_exhausted, 

67 exhausted_until 

68FROM provider_daily_usage 

69WHERE usage_date = $1 

70""" 

71 

72 

73class DatabaseQuotaBackend: 

74 """PostgreSQL-backed quota backend using ``DatabaseProviderProtocol``. 

75 

76 Uses UPSERT statements for thread-safe and process-safe concurrent 

77 updates without application-level locking. 

78 

79 Example: 

80 >>> backend = DatabaseQuotaBackend(db=db_provider) 

81 >>> await backend.increment("groq") 

82 >>> await backend.is_exhausted("groq") 

83 False 

84 """ 

85 

86 def __init__(self, db: DatabaseProviderProtocol) -> None: 

87 """Initialise the database quota backend. 

88 

89 Args: 

90 db: Framework database provider (injected from the DI container). 

91 """ 

92 self._db = db 

93 

94 def _today(self) -> str: 

95 """Return today's UTC date as an ISO 8601 string.""" 

96 return date.today().isoformat() 

97 

98 async def is_exhausted(self, provider: str) -> bool: 

99 """Return ``True`` when *provider* is quota-exhausted today. 

100 

101 Args: 

102 provider: Provider name. 

103 

104 Returns: 

105 Whether the provider is currently exhausted; defaults to ``False`` 

106 when the database call fails. 

107 """ 

108 try: 

109 async with self._db.scoped_context(): 

110 conn = await self._db.get_scoped_connection() 

111 row = await conn.fetchrow(_GET_TODAY_SQL, provider, self._today()) 

112 if row is None: 

113 return False 

114 exhausted_until = row["exhausted_until"] 

115 if exhausted_until is None: 

116 return False 

117 return bool(datetime.now(UTC) < exhausted_until) 

118 except Exception as e: 

119 logger.exception( 

120 "llm.quota.db: is_exhausted query failed for %s", provider, error=str(e) 

121 ) 

122 return False 

123 

124 async def increment(self, provider: str) -> None: 

125 """Record one successful completion for *provider* today. 

126 

127 Args: 

128 provider: Provider name. 

129 """ 

130 try: 

131 async with self._db.scoped_context(): 

132 conn = await self._db.get_scoped_connection() 

133 await conn.execute(_INCREMENT_SQL, provider, self._today()) 

134 except Exception as e: 

135 logger.exception( 

136 "llm.quota.db: increment failed for %s", provider, error=str(e) 

137 ) 

138 

139 async def mark_exhausted( 

140 self, provider: str, *, until: datetime | None = None 

141 ) -> None: 

142 """Mark *provider* exhausted until *until*. 

143 

144 Args: 

145 provider: Cascade-entry key (``name:model``) or provider name. 

146 until: Exhaustion expiry; ``None`` means the rest of today (UTC). 

147 """ 

148 expiry = until or end_of_utc_day() 

149 try: 

150 async with self._db.scoped_context(): 

151 conn = await self._db.get_scoped_connection() 

152 await conn.execute(_MARK_EXHAUSTED_SQL, provider, self._today(), expiry) 

153 logger.info( 

154 "llm.quota.db: %s marked exhausted until %s", 

155 provider, 

156 expiry.isoformat(), 

157 ) 

158 except Exception as e: 

159 logger.exception( 

160 "llm.quota.db: mark_exhausted failed for %s", provider, error=str(e) 

161 ) 

162 

163 async def record_error(self, provider: str) -> None: 

164 """Record a non-exhaustion error for *provider* today. 

165 

166 Args: 

167 provider: Provider name. 

168 """ 

169 try: 

170 async with self._db.scoped_context(): 

171 conn = await self._db.get_scoped_connection() 

172 await conn.execute(_RECORD_ERROR_SQL, provider, self._today()) 

173 except Exception as e: 

174 logger.exception( 

175 "llm.quota.db: record_error failed for %s", provider, error=str(e) 

176 ) 

177 

178 async def get_usage(self, provider: str) -> ProviderUsage | None: 

179 """Return today's usage record for *provider*. 

180 

181 Args: 

182 provider: Provider name. 

183 

184 Returns: 

185 :class:`ProviderUsage` or ``None`` when not found or on DB error. 

186 """ 

187 try: 

188 async with self._db.scoped_context(): 

189 conn = await self._db.get_scoped_connection() 

190 row = await conn.fetchrow(_GET_TODAY_SQL, provider, self._today()) 

191 if row is None: 

192 return None 

193 return ProviderUsage( 

194 provider=row["provider"], 

195 usage_date=row["usage_date"], 

196 success_count=row["success_count"], 

197 error_count=row["error_count"], 

198 is_exhausted=row["is_exhausted"], 

199 exhausted_until=row["exhausted_until"], 

200 ) 

201 except Exception as e: 

202 logger.exception( 

203 "llm.quota.db: get_usage failed for %s", provider, error=str(e) 

204 ) 

205 return None 

206 

207 async def get_all_usage(self) -> list[ProviderUsage]: 

208 """Return all today's usage records. 

209 

210 Returns: 

211 List of :class:`ProviderUsage`; empty list on DB error. 

212 """ 

213 try: 

214 async with self._db.scoped_context(): 

215 conn = await self._db.get_scoped_connection() 

216 rows = await conn.fetch(_GET_ALL_TODAY_SQL, self._today()) 

217 return [ 

218 ProviderUsage( 

219 provider=row["provider"], 

220 usage_date=row["usage_date"], 

221 success_count=row["success_count"], 

222 error_count=row["error_count"], 

223 is_exhausted=row["is_exhausted"], 

224 ) 

225 for row in rows 

226 ] 

227 except Exception as e: 

228 logger.exception("llm.quota.db: get_all_usage failed", error=str(e)) 

229 return []