« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""
2Distributed tracing for Lexigram Intelligence operations.
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
10All traces use lexigram-monitor's Tracer for OpenTelemetry compatibility.
11"""
13from __future__ import annotations
15from typing import TYPE_CHECKING, Any
17from lexigram.contracts.ai.callbacks import CallbackHandlerProtocol
18from lexigram.contracts.ai.llm import ChatMessage, Completion
19from lexigram.contracts.core.logging import RedactorProtocol
20from lexigram.contracts.observability.tracing import SpanProtocol as Span
21from lexigram.contracts.observability.tracing import TracerProtocol as Tracer
22from lexigram.di.decorators import inject
23from lexigram.observability.core import NoOpTracer
25if TYPE_CHECKING:
26 from contextlib import AbstractContextManager as ContextManager
29# TracerProtocol is an alias for Tracer in contracts
30TracerProtocol = Tracer
33@inject
34class AITracer(CallbackHandlerProtocol):
35 """Distributed tracer for intelligence operations.
37 Provides span management and context propagation for:
38 - LLM completions and streaming
39 - Vector store operations
40 - RAG pipeline execution
41 - Embedding generation
43 Also implements CallbackHandlerProtocol for event-driven tracing.
45 Example:
46 >>> tracer = AITracer()
47 >>> async with tracer.trace_llm_call("openai", "gpt-4") as span:
48 ... response = await client.complete(messages)
49 ... span.set_attribute("tokens.total", response.usage.total_tokens)
50 ... span.set_attribute("cost", response.cost)
51 """
53 def __init__(
54 self,
55 tracer: Tracer | None = None,
56 *,
57 redaction_policy: RedactorProtocol | None = None,
58 max_attribute_length: int | None = None,
59 ) -> None:
60 """Initialize intelligence tracer.
62 Args:
63 tracer: Tracer instance to use for tracing. Defaults to a
64 no-op tracer when ``None``.
65 redaction_policy: Optional policy applied to every payload
66 dict before it reaches a span boundary (``start_span``
67 attributes and ``add_event``). Defaults to ``None``,
68 which passes payloads through unchanged.
69 max_attribute_length: Optional cap on string attribute
70 values, applied recursively to dicts, lists, and
71 tuples, independent of the redaction policy. Defaults
72 to ``None`` (no truncation).
73 """
74 self.tracer = tracer if tracer is not None else NoOpTracer()
75 self._redaction_policy = redaction_policy
76 self._max_attribute_length = max_attribute_length
78 def trace_llm_call(
79 self,
80 provider: str,
81 model: str,
82 **attributes: Any,
83 ) -> ContextManager[Span]:
84 """Create a span for LLM API call.
86 Args:
87 provider: LLM provider name (e.g., "openai", "anthropic")
88 model: Model name (e.g., "gpt-4", "claude-3-opus")
89 **attributes: Additional span attributes
91 Returns:
92 Span context manager
94 Example:
95 >>> tracer = AITracer()
96 >>> with tracer.trace_llm_call("openai", "gpt-4") as span:
97 ... response = await client.complete(messages)
98 ... span.set_attribute("tokens.total", response.usage.total_tokens)
99 """
100 span_attributes = {
101 "llm.provider": provider,
102 "llm.model": model,
103 "operation.type": "llm.completion",
104 **attributes,
105 }
106 return self.tracer.start_span(
107 name=f"llm.{provider}.{model}",
108 attributes=span_attributes,
109 )
111 def trace_operation(
112 self,
113 name: str,
114 **attrs: Any,
115 ) -> ContextManager[Span]:
116 """Generic operation tracing helper.
118 This mirrors the `Tracer.trace_operation` API and is used by
119 worker code that needs a generic operation span (e.g., document
120 parsing/chunking).
121 """
122 return self.tracer.start_span(name=name, attributes=attrs or {})
124 def trace_vector_operation(
125 self,
126 operation: str,
127 provider: str,
128 collection: str | None = None,
129 **attributes: Any,
130 ) -> ContextManager[Span]:
131 """Create a span for vector store operation.
133 Args:
134 operation: Operation type (e.g., "add", "search", "delete")
135 provider: Vector store provider (e.g., "pgvector", "chroma")
136 collection: Optional collection/table name
137 **attributes: Additional span attributes
139 Returns:
140 Span context manager
142 Example:
143 >>> with tracer.trace_vector_operation("search", "pgvector", "documents") as span:
144 ... results = await store.search(query, limit=10)
145 ... span.set_attribute("results.count", len(results))
146 """
147 span_attributes = {
148 "vector.operation": operation,
149 "vector.provider": provider,
150 "operation.type": "vector.operation",
151 **attributes,
152 }
153 if collection:
154 span_attributes["vector.collection"] = collection
156 return self.tracer.start_span(
157 name=f"vector.{operation}.{provider}",
158 attributes=span_attributes,
159 )
161 def trace_embedding_operation(
162 self,
163 model: str,
164 batch_size: int | None = None,
165 **attributes: Any,
166 ) -> ContextManager[Span]:
167 """Create a span for embedding generation.
169 Args:
170 model: Embedding model name
171 batch_size: Optional number of texts being embedded
172 **attributes: Additional span attributes
174 Returns:
175 Span context manager
177 Example:
178 >>> with tracer.trace_embedding_operation("text-embedding-ada-002", 5) as span:
179 ... embeddings = await embedder.embed(texts)
180 ... span.set_attribute("embeddings.dimensions", len(embeddings[0]))
181 """
182 span_attributes = {
183 "embedding.model": model,
184 "operation.type": "embedding.generation",
185 **attributes,
186 }
187 if batch_size is not None:
188 span_attributes["embedding.batch_size"] = batch_size
190 return self.tracer.start_span(
191 name=f"embedding.{model}",
192 attributes=span_attributes,
193 )
195 def trace_rag_stage(
196 self,
197 stage: str,
198 pipeline: str = "default",
199 **attributes: Any,
200 ) -> ContextManager[Span]:
201 """Create a span for RAG pipeline stage.
203 Args:
204 stage: Stage name (e.g., "retrieval", "ranking", "synthesis")
205 pipeline: Pipeline name
206 **attributes: Additional span attributes
208 Returns:
209 Span context manager
211 Example:
212 >>> with tracer.trace_rag_stage("retrieval", "default") as span:
213 ... documents = await retriever.retrieve(query)
214 ... span.set_attribute("documents.count", len(documents))
215 """
216 span_attributes = {
217 "rag.stage": stage,
218 "rag.pipeline": pipeline,
219 "operation.type": "rag.stage",
220 **attributes,
221 }
223 return self.tracer.start_span(
224 name=f"rag.{stage}",
225 attributes=span_attributes,
226 )
228 def trace_rag_query(
229 self,
230 query: str,
231 pipeline: str = "default",
232 **attributes: Any,
233 ) -> ContextManager[Span]:
234 """Create a span for complete RAG query.
236 Args:
237 query: Query text
238 pipeline: Pipeline name
239 **attributes: Additional span attributes
241 Returns:
242 Span context manager
244 Example:
245 >>> with tracer.trace_rag_query("What is Python?") as span:
246 ... result = await rag_pipeline.query(query)
247 ... span.set_attribute("answer.length", len(result.answer))
248 """
249 span_attributes = {
250 "rag.query": query[:100], # Truncate long queries
251 "rag.pipeline": pipeline,
252 "operation.type": "rag.query",
253 **attributes,
254 }
256 return self.tracer.start_span(
257 name="rag.query",
258 attributes=span_attributes,
259 )
261 def get_current_span(self) -> Span | None:
262 """Get the currently active span.
264 Returns:
265 Current span or None
266 """
267 return self.tracer.get_current_span()
269 def _sanitize_attributes(self, attributes: dict[str, Any]) -> dict[str, Any]:
270 """Redact and/or truncate *attributes* before a span boundary.
272 When neither ``redaction_policy`` nor ``max_attribute_length``
273 is configured, *attributes* is returned unchanged so trace
274 output stays byte-identical to an unsanitized tracer.
276 Args:
277 attributes: The payload dict destined for ``start_span``
278 attributes or ``add_event``.
280 Returns:
281 The sanitized payload dict.
282 """
283 if self._redaction_policy is None and self._max_attribute_length is None:
284 return attributes
285 sanitized = attributes
286 if self._redaction_policy is not None:
287 sanitized = self._redaction_policy.redact_dict(sanitized)
288 if self._max_attribute_length is not None:
289 sanitized = self._truncate_attributes(sanitized)
290 return sanitized
292 def _truncate_attributes(self, attributes: dict[str, Any]) -> dict[str, Any]:
293 """Cap *attributes* string values at ``max_attribute_length``."""
294 return {key: self._truncate_value(value) for key, value in attributes.items()}
296 def _truncate_value(self, value: Any) -> Any:
297 """Truncate *value* if it is an oversized string, recursing containers."""
298 max_length = self._max_attribute_length
299 if max_length is None:
300 return value
301 if isinstance(value, str):
302 if len(value) > max_length:
303 return value[:max_length]
304 return value
305 if isinstance(value, dict):
306 return self._truncate_attributes(value)
307 if isinstance(value, list):
308 return [self._truncate_value(item) for item in value]
309 if isinstance(value, tuple):
310 return tuple(self._truncate_value(item) for item in value)
311 return value
313 async def on_llm_start(
314 self,
315 messages: list[ChatMessage],
316 model: str,
317 **kwargs: Any,
318 ) -> None:
319 """Called when an LLM call starts."""
320 span_attributes = {
321 "llm.model": model,
322 "operation.type": "llm.start",
323 **kwargs,
324 }
325 self.tracer.start_span(name=f"llm.{model}", attributes=span_attributes)
327 async def on_llm_new_token(
328 self,
329 token: str,
330 **kwargs: Any,
331 ) -> None:
332 """Called for each new token in a streaming LLM response."""
333 span = self.tracer.get_current_span()
334 if span:
335 span.add_event("llm.token", {"token": token})
337 async def on_llm_end(
338 self,
339 response: Completion,
340 **kwargs: Any,
341 ) -> None:
342 """Called when an LLM call completes successfully."""
343 span = self.tracer.get_current_span()
344 if span:
345 span.add_event("llm.end", {"model": response.model})
346 span.end()
348 async def on_llm_error(
349 self,
350 error: Exception,
351 **kwargs: Any,
352 ) -> None:
353 """Called when an LLM call fails."""
354 span = self.tracer.get_current_span()
355 if span:
356 span.add_event("llm.error", {"error": str(error)})
357 span.set_status("error")
358 span.end()
360 async def on_chain_start(
361 self,
362 name: str,
363 inputs: dict[str, Any],
364 **kwargs: Any,
365 ) -> None:
366 """Called when a chain/pipeline starts executing."""
367 self.tracer.start_span(
368 name=f"chain.{name}", attributes={"chain.name": name, **kwargs}
369 )
371 async def on_chain_end(
372 self,
373 name: str,
374 outputs: dict[str, Any],
375 **kwargs: Any,
376 ) -> None:
377 """Called when a chain/pipeline completes."""
378 span = self.tracer.get_current_span()
379 if span:
380 span.add_event("chain.end", {"chain.name": name})
381 span.end()
383 async def on_tool_start(
384 self,
385 tool_name: str,
386 arguments: dict[str, Any],
387 **kwargs: Any,
388 ) -> None:
389 """Called when a tool starts executing."""
390 self.tracer.start_span(
391 name=f"tool.{tool_name}",
392 attributes=self._sanitize_attributes(
393 {"tool.name": tool_name, "tool.args": arguments, **kwargs}
394 ),
395 )
397 async def on_tool_end(
398 self,
399 tool_name: str,
400 result: Any,
401 **kwargs: Any,
402 ) -> None:
403 """Called when a tool finishes executing."""
404 span = self.tracer.get_current_span()
405 if span:
406 span.add_event("tool.end", {"tool.name": tool_name})
407 span.end()
409 async def on_agent_action(
410 self,
411 action: dict[str, Any],
412 **kwargs: Any,
413 ) -> None:
414 """Called when an agent takes an action."""
415 span = self.tracer.get_current_span()
416 if span:
417 span.add_event("agent.action", self._sanitize_attributes(action))
419 async def on_agent_finish(
420 self,
421 response: dict[str, Any],
422 **kwargs: Any,
423 ) -> None:
424 """Called when an agent finishes executing."""
425 span = self.tracer.get_current_span()
426 if span:
427 span.add_event("agent.finish", self._sanitize_attributes(response))
428 span.end()
430 async def on_retriever_start(
431 self,
432 query: str,
433 **kwargs: Any,
434 ) -> None:
435 """Called when a retriever starts a search."""
436 self.tracer.start_span(
437 name="retriever.search",
438 attributes=self._sanitize_attributes({"retriever.query": query, **kwargs}),
439 )
441 async def on_retriever_end(
442 self,
443 documents: list[Any],
444 **kwargs: Any,
445 ) -> None:
446 """Called when a retriever completes a search."""
447 span = self.tracer.get_current_span()
448 if span:
449 span.add_event("retriever.end", {"documents.count": len(documents)})
450 span.end()