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

1"""In-memory quota backend for LLM routing. 

2 

3Suitable for development, testing, and single-process deployments. 

4All state is lost on process restart; use ``DatabaseQuotaBackend`` for 

5production deployments that require persistence across restarts. 

6""" 

7 

8from __future__ import annotations 

9 

10import asyncio 

11from datetime import UTC, date, datetime 

12 

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

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

15from lexigram.logging import ( 

16 get_logger, 

17) 

18 

19logger = get_logger(__name__) 

20 

21__all__ = ["InMemoryQuotaBackend"] 

22 

23 

24class InMemoryQuotaBackend: 

25 """Thread-safe in-memory quota backend using asyncio locks. 

26 

27 Quota state is keyed on ``(provider, date)`` and resets automatically 

28 across UTC midnight boundaries — no cron job required. 

29 

30 Example: 

31 >>> backend = InMemoryQuotaBackend() 

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

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

34 False 

35 >>> await backend.mark_exhausted("groq") 

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

37 True 

38 """ 

39 

40 def __init__(self) -> None: 

41 """Initialise an empty in-memory quota backend.""" 

42 self._lock = asyncio.Lock() 

43 # Keys: (provider, date_str); values: ProviderUsage 

44 self._store: dict[tuple[str, str], ProviderUsage] = {} 

45 

46 def _today(self) -> str: 

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

48 return date.today().isoformat() 

49 

50 def _get_or_create(self, provider: str) -> ProviderUsage: 

51 """Return (or create) today's usage record for *provider*. 

52 

53 Must be called while holding ``self._lock``. 

54 

55 Args: 

56 provider: Provider name. 

57 

58 Returns: 

59 Mutable :class:`ProviderUsage` for today. 

60 """ 

61 today = self._today() 

62 key = (provider, today) 

63 if key not in self._store: 

64 self._store[key] = ProviderUsage( 

65 provider=provider, 

66 usage_date=today, 

67 ) 

68 return self._store[key] 

69 

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

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

72 

73 Args: 

74 provider: Provider name. 

75 

76 Returns: 

77 Whether the provider is currently exhausted. 

78 """ 

79 async with self._lock: 

80 record = self._store.get((provider, self._today())) 

81 if record is None or record.exhausted_until is None: 

82 return False 

83 return datetime.now(UTC) < record.exhausted_until 

84 

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

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

87 

88 Args: 

89 provider: Provider name. 

90 """ 

91 async with self._lock: 

92 record = self._get_or_create(provider) 

93 record.success_count += 1 

94 

95 async def mark_exhausted( 

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

97 ) -> None: 

98 """Mark *provider* exhausted until *until*. 

99 

100 Args: 

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

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

103 """ 

104 expiry = until or end_of_utc_day() 

105 async with self._lock: 

106 record = self._get_or_create(provider) 

107 record.is_exhausted = True 

108 record.exhausted_until = expiry 

109 logger.info( 

110 "llm.quota.memory: %s marked exhausted until %s", 

111 provider, 

112 expiry.isoformat(), 

113 ) 

114 

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

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

117 

118 Args: 

119 provider: Provider name. 

120 """ 

121 async with self._lock: 

122 record = self._get_or_create(provider) 

123 record.error_count += 1 

124 

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

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

127 

128 Args: 

129 provider: Provider name. 

130 

131 Returns: 

132 :class:`ProviderUsage` or ``None`` if no activity recorded today. 

133 """ 

134 async with self._lock: 

135 return self._store.get((provider, self._today())) 

136 

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

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

139 

140 Returns: 

141 List of :class:`ProviderUsage` for every provider active today. 

142 """ 

143 today = self._today() 

144 async with self._lock: 

145 return [ 

146 record 

147 for (provider, date_str), record in self._store.items() 

148 if date_str == today 

149 ]