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

1""" 

2Distributed tracing for Lexigram Intelligence operations. 

3 

4This module provides distributed tracing capabilities for: 

5- LLM API calls with token and cost tracking 

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

7- RAG pipeline stages (retrieval, ranking, synthesis) 

8- Embedding operations 

9 

10All traces use lexigram-monitor's Tracer for OpenTelemetry compatibility. 

11""" 

12 

13from __future__ import annotations 

14 

15from functools import wraps 

16from typing import TYPE_CHECKING, Any 

17 

18if TYPE_CHECKING: 

19 from collections.abc import Callable 

20 

21 from lexigram.ai.observability.tracing.core import AITracer 

22 

23 

24# The local TracerProtocol and SpanProtocol definitions are removed as per instruction. 

25# The type alias Span = SpanProtocol is also removed. 

26 

27# TracerProtocol is an alias for Tracer in contracts 

28 

29 

30def trace_llm( 

31 provider: str, 

32 model: str, 

33 tracer: AITracer, 

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

35 """Decorator to automatically trace LLM calls. 

36 

37 Args: 

38 provider: LLM provider name 

39 model: Model name 

40 tracer: AITracer instance to use for tracing 

41 

42 Returns: 

43 Decorator function 

44 

45 Example: 

46 >>> tracer = AITracer(some_tracer) 

47 >>> @trace_llm(provider="openai", model="gpt-4", tracer=tracer) 

48 ... async def complete(messages): 

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

50 ... return response 

51 """ 

52 

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

54 @wraps(func) 

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

56 nonlocal tracer 

57 if tracer is None: 

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

59 

60 with tracer.trace_llm_call(provider, model) as span: 

61 try: 

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

63 

64 # Extract and record metrics from result 

65 if hasattr(result, "usage"): 

66 usage = result.usage 

67 if hasattr(usage, "total_tokens"): 

68 span.set_attribute("llm.tokens.total", usage.total_tokens) 

69 if hasattr(usage, "prompt_tokens"): 

70 span.set_attribute("llm.tokens.prompt", usage.prompt_tokens) 

71 if hasattr(usage, "completion_tokens"): 

72 span.set_attribute( 

73 "llm.tokens.completion", 

74 usage.completion_tokens, 

75 ) 

76 

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

78 span.set_attribute("llm.cost", result.cost) 

79 

80 if hasattr(result, "content"): 

81 span.set_attribute("llm.response.length", len(result.content)) 

82 

83 span.set_attribute("status", "success") 

84 

85 except ( 

86 Exception 

87 ) as e: # tracing decorator must capture all exception types 

88 span.set_attribute("status", "error") 

89 span.set_attribute("error.type", type(e).__name__) 

90 span.set_attribute("error.message", str(e)) 

91 span.add_event("exception", {"exception": str(e)}) 

92 raise 

93 else: 

94 return result 

95 

96 return wrapper 

97 

98 return decorator 

99 

100 

101def trace_vector( 

102 operation: str, 

103 provider: str, 

104 tracer: AITracer, 

105 collection: str | None = None, 

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

107 """Decorator to automatically trace vector operations. 

108 

109 Args: 

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

111 provider: Vector store provider 

112 tracer: AITracer instance to use for tracing 

113 collection: Optional collection name 

114 

115 Returns: 

116 Decorator function 

117 

118 Example: 

119 >>> tracer = AITracer(some_tracer) 

120 >>> @trace_vector(operation="search", provider="pgvector", tracer=tracer, collection="docs") 

121 ... async def search(query, limit=10): 

122 ... results = await store.search(query, limit) 

123 ... return results 

124 """ 

125 

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

127 @wraps(func) 

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

129 nonlocal tracer 

130 if tracer is None: 

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

132 

133 with tracer.trace_vector_operation(operation, provider, collection) as span: 

134 try: 

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

136 

137 # Record result metrics 

138 if hasattr(result, "__len__"): 

139 span.set_attribute("results.count", len(result)) 

140 

141 if operation == "search" and isinstance(result, list): 

142 if result and hasattr(result[0], "score"): 

143 span.set_attribute("results.top_score", result[0].score) 

144 

145 span.set_attribute("status", "success") 

146 

147 except ( 

148 Exception 

149 ) as e: # tracing decorator must capture all exception types 

150 span.set_attribute("status", "error") 

151 span.set_attribute("error.type", type(e).__name__) 

152 span.set_attribute("error.message", str(e)) 

153 span.add_event("exception", {"exception": str(e)}) 

154 raise 

155 else: 

156 return result 

157 

158 return wrapper 

159 

160 return decorator 

161 

162 

163def trace_rag( 

164 stage: str, 

165 tracer: AITracer, 

166 pipeline: str = "default", 

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

168 """Decorator to automatically trace RAG pipeline stages. 

169 

170 Args: 

171 stage: Stage name (e.g., "retrieval", "ranking", "synthesis") 

172 tracer: AITracer instance to use for tracing 

173 pipeline: Pipeline name 

174 

175 Returns: 

176 Decorator function 

177 

178 Example: 

179 >>> tracer = AITracer(some_tracer) 

180 >>> @trace_rag(stage="retrieval", tracer=tracer, pipeline="default") 

181 ... async def retrieve(query): 

182 ... documents = await retriever.retrieve(query) 

183 ... return documents 

184 """ 

185 

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

187 @wraps(func) 

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

189 nonlocal tracer 

190 if tracer is None: 

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

192 

193 with tracer.trace_rag_stage(stage, pipeline) as span: 

194 try: 

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

196 

197 # Record stage-specific metrics 

198 if stage == "retrieval" and hasattr(result, "__len__"): 

199 span.set_attribute("documents.retrieved", len(result)) 

200 

201 if stage == "synthesis" and hasattr(result, "answer"): 

202 span.set_attribute("answer.length", len(result.answer)) 

203 

204 span.set_attribute("status", "success") 

205 

206 except ( 

207 Exception 

208 ) as e: # tracing decorator must capture all exception types 

209 span.set_attribute("status", "error") 

210 span.set_attribute("error.type", type(e).__name__) 

211 span.set_attribute("error.message", str(e)) 

212 span.add_event("exception", {"exception": str(e)}) 

213 raise 

214 else: 

215 return result 

216 

217 return wrapper 

218 

219 return decorator