Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-observability/src/lexigram/ai/observability/metrics/decorators.py: 11%

1""" 

2Metrics collection for Lexigram Intelligence operations. 

3 

4This module provides comprehensive metrics collection for: 

5- LLM operations (requests, tokens, duration, costs) 

6- Vector store operations (add, search, delete) 

7- Cache operations (hits, misses) 

8- RAG pipeline operations (queries, retrievals, latency) 

9 

10All metrics use lexigram-monitor's MetricsCollectorProtocol for consistent observability. 

11""" 

12 

13from __future__ import annotations 

14 

15from functools import wraps 

16import time 

17from typing import TYPE_CHECKING, Any 

18 

19if TYPE_CHECKING: 

20 from collections.abc import Callable 

21 

22 from lexigram.ai.observability.metrics.core import AIMetrics 

23 

24 

25from lexigram.contracts.observability.metrics import ( 

26 MetricsCollectorProtocol as MetricsCollectorProtocol, 

27) 

28 

29 

30def track_llm_call( 

31 provider: str, 

32 model: str, 

33 metrics: AIMetrics | None = None, 

34) -> Callable[[Callable[..., Any]], Callable[..., Any]]: 

35 """Decorator to automatically track LLM call metrics. 

36 

37 Args: 

38 provider: LLM provider name (e.g., "openai", "anthropic") 

39 model: Model name (e.g., "gpt-4", "claude-3-opus") 

40 metrics: AIMetrics instance to use. If None, creates a new one. 

41 

42 Returns: 

43 Decorator function 

44 

45 Example: 

46 >>> @track_llm_call(provider="openai", model="gpt-4") 

47 ... async def complete(messages): 

48 ... response = await client.complete(messages) 

49 ... return response 

50 """ 

51 

52 def decorator(func: Callable[..., Any]) -> Callable[..., Any]: 

53 @wraps(func) 

54 async def wrapper(*args: Any, **kwargs: Any) -> Any: 

55 nonlocal metrics 

56 

57 if metrics is None: 

58 return await func(*args, **kwargs) 

59 

60 labels = {"provider": provider, "model": model} 

61 

62 # Track active requests 

63 metrics.llm_active_requests.increment(labels=labels) 

64 

65 start_time = time.time() 

66 status = "success" 

67 

68 try: 

69 result = await func(*args, **kwargs) 

70 

71 # Extract metrics from result if available 

72 if hasattr(result, "usage"): 

73 usage = result.usage 

74 if hasattr(usage, "total_tokens"): 

75 metrics.llm_tokens_total.increment( 

76 amount=usage.total_tokens, 

77 labels={**labels, "type": "total"}, 

78 ) 

79 if hasattr(usage, "prompt_tokens"): 

80 metrics.llm_tokens_total.increment( 

81 amount=usage.prompt_tokens, 

82 labels={**labels, "type": "prompt"}, 

83 ) 

84 if hasattr(usage, "completion_tokens"): 

85 metrics.llm_tokens_total.increment( 

86 amount=usage.completion_tokens, 

87 labels={**labels, "type": "completion"}, 

88 ) 

89 

90 if hasattr(result, "cost") and result.cost is not None: 

91 metrics.llm_cost_dollars.increment( 

92 amount=int(result.cost * 1_000_000), # Store as micro-dollars 

93 labels=labels, 

94 ) 

95 

96 except ( 

97 Exception 

98 ): # metrics decorator must observe all exceptions without narrowing 

99 status = "error" 

100 raise 

101 

102 else: 

103 return result 

104 

105 finally: 

106 # Track request completion 

107 duration = time.time() - start_time 

108 metrics.llm_duration_seconds.observe(value=duration, labels=labels) 

109 metrics.llm_requests_total.increment( 

110 labels={**labels, "status": status}, 

111 ) 

112 metrics.llm_active_requests.decrement(labels=labels) 

113 

114 return wrapper 

115 

116 return decorator 

117 

118 

119def track_vector_operation( 

120 operation: str, 

121 provider: str, 

122 metrics: AIMetrics | None = None, 

123) -> Callable[[Callable[..., Any]], Callable[..., Any]]: 

124 """Decorator to automatically track vector store operation metrics. 

125 

126 Args: 

127 operation: Operation type (e.g., "add", "search", "delete") 

128 provider: Vector store provider (e.g., "pgvector", "chroma", "qdrant") 

129 metrics: AIMetrics instance to use. If None, creates a new one. 

130 

131 Returns: 

132 Decorator function 

133 

134 Example: 

135 >>> @track_vector_operation(operation="search", provider="pgvector") 

136 ... async def search(query_embedding, limit=10): 

137 ... results = await store.search(query_embedding, limit) 

138 ... return results 

139 """ 

140 

141 def decorator(func: Callable[..., Any]) -> Callable[..., Any]: 

142 @wraps(func) 

143 async def wrapper(*args: Any, **kwargs: Any) -> Any: 

144 nonlocal metrics 

145 

146 if metrics is None: 

147 return await func(*args, **kwargs) 

148 

149 labels = {"operation": operation, "provider": provider} 

150 

151 start_time = time.time() 

152 

153 try: 

154 result = await func(*args, **kwargs) 

155 

156 # Track document counts for add/search operations 

157 if operation == "add" and hasattr(result, "__len__"): 

158 metrics.vector_documents_total.increment( 

159 amount=len(result), 

160 labels={**labels, "type": "added"}, 

161 ) 

162 elif operation == "search" and hasattr(result, "__len__"): 

163 metrics.vector_documents_total.increment( 

164 amount=len(result), 

165 labels={**labels, "type": "retrieved"}, 

166 ) 

167 

168 return result 

169 

170 finally: 

171 duration = time.time() - start_time 

172 metrics.vector_duration_seconds.observe(value=duration, labels=labels) 

173 metrics.vector_operations_total.increment(labels=labels) 

174 

175 return wrapper 

176 

177 return decorator 

178 

179 

180def track_embedding_operation( 

181 model: str, 

182 metrics: AIMetrics | None = None, 

183) -> Callable[[Callable[..., Any]], Callable[..., Any]]: 

184 """Decorator to automatically track embedding operation metrics. 

185 

186 Args: 

187 model: Embedding model name (e.g., "text-embedding-ada-002") 

188 metrics: AIMetrics instance to use. If None, creates a new one. 

189 

190 Returns: 

191 Decorator function 

192 

193 Example: 

194 >>> @track_embedding_operation(model="text-embedding-ada-002") 

195 ... async def embed_batch(texts): 

196 ... embeddings = await embedder.embed(texts) 

197 ... return embeddings 

198 """ 

199 

200 def decorator(func: Callable[..., Any]) -> Callable[..., Any]: 

201 @wraps(func) 

202 async def wrapper(*args: Any, **kwargs: Any) -> Any: 

203 nonlocal metrics 

204 

205 if metrics is None: 

206 return await func(*args, **kwargs) 

207 

208 labels = {"model": model} 

209 

210 start_time = time.time() 

211 

212 try: 

213 result = await func(*args, **kwargs) 

214 

215 # Track batch size 

216 if hasattr(result, "__len__"): 

217 metrics.embedding_batch_size.observe( 

218 value=len(result), 

219 labels=labels, 

220 ) 

221 

222 return result 

223 

224 finally: 

225 duration = time.time() - start_time 

226 metrics.embedding_duration_seconds.observe( 

227 value=duration, 

228 labels=labels, 

229 ) 

230 metrics.embedding_operations_total.increment(labels=labels) 

231 

232 return wrapper 

233 

234 return decorator