From c189f4ee4a62d71a7dd89df7489d36e77aadb471 Mon Sep 17 00:00:00 2001 From: Anshul Jain Date: Sat, 1 Aug 2026 12:38:57 +0530 Subject: [PATCH] feat(#159): add RAG data lineage tracing with OpenTelemetry Implement OpenTelemetry-based tracing for complete RAG pipeline visibility: - app/observability/rag_tracing.py: Core tracing module with: * RAGTracer context manager for end-to-end request tracing * RetrievalSpan for vector search operations and latency tracking * RerankerSpan for chunk reranking metrics * LLMSpan for language model generation with token tracking * Document lineage tracking with source attribution * Automatic latency metrics for all operations * Token usage tracking per request * Cache hit/miss tracing events - tests/test_rag_tracing.py: Comprehensive test suite with: * Span creation and context management * Attribute setting and query text handling * Token usage tracking * Error handling and status reporting * Integration scenario testing * Latency metric recording Features: - Request lifecycle tracing with user/tenant/session correlation - Retrieved document metadata capture (source, chunk index, similarity) - Vector search latency metrics - LLM execution time tracking - Token consumption metrics per request - Reranking operation tracing - Cache operation visibility - Structured spans for audit trails Observability benefits: - Full visibility into RAG execution pipeline - Performance bottleneck identification via latency metrics - Token usage tracking for cost analysis - Document lineage for debugging and validation - Request correlation across distributed systems - Production audit trails for enterprise compliance Fixes #159 --- app/observability/rag_tracing.py | 324 +++++++++++++++++++++++++++++++ tests/test_rag_tracing.py | 267 +++++++++++++++++++++++++ 2 files changed, 591 insertions(+) create mode 100644 app/observability/rag_tracing.py create mode 100644 tests/test_rag_tracing.py diff --git a/app/observability/rag_tracing.py b/app/observability/rag_tracing.py new file mode 100644 index 0000000..5adcc22 --- /dev/null +++ b/app/observability/rag_tracing.py @@ -0,0 +1,324 @@ +"""OpenTelemetry-based tracing for RAG pipeline execution and data lineage. + +Issue #159: Provides complete visibility into RAG query execution including +retrieval latency, document lineage, LLM execution time, and token consumption. + +Implements structured spans for debugging, performance analysis, and audit trails +in enterprise production environments. +""" + +import time +from typing import Any, Optional + +from opentelemetry import trace, metrics +from opentelemetry.trace import Status, StatusCode + +# Get tracer and meter instances +tracer = trace.get_tracer(__name__) +meter = metrics.get_meter(__name__) + +# Create metrics +retrieval_latency = meter.create_histogram( + "rag.retrieval.latency_ms", + unit="ms", + description="Time taken for retrieval operation", +) + +llm_latency = meter.create_histogram( + "rag.llm.latency_ms", + unit="ms", + description="Time taken for LLM generation", +) + +reranking_latency = meter.create_histogram( + "rag.reranking.latency_ms", + unit="ms", + description="Time taken for chunk reranking", +) + +retrieval_chunk_count = meter.create_histogram( + "rag.retrieval.chunk_count", + description="Number of chunks retrieved", +) + +token_usage = meter.create_histogram( + "rag.token_usage", + description="Tokens used per request", +) + + +class RAGSpanAttributes: + """Standard attributes for RAG-related spans.""" + + # Query attributes + QUERY_TEXT = "rag.query.text" + QUERY_LENGTH = "rag.query.length" + USER_ID = "rag.user.id" + TENANT_ID = "rag.tenant.id" + SESSION_ID = "rag.session.id" + + # Retrieval attributes + RETRIEVAL_TOP_K = "rag.retrieval.top_k" + RETRIEVAL_CHUNK_COUNT = "rag.retrieval.chunk_count" + RETRIEVAL_METHOD = "rag.retrieval.method" + VECTOR_STORE_TYPE = "rag.vector_store.type" + + # Retrieved document attributes + DOCUMENT_ID = "rag.document.id" + DOCUMENT_SOURCE = "rag.document.source" + DOCUMENT_CHUNK_INDEX = "rag.document.chunk_index" + DOCUMENT_SIMILARITY_SCORE = "rag.document.similarity_score" + + # LLM attributes + LLM_MODEL = "rag.llm.model" + LLM_TEMPERATURE = "rag.llm.temperature" + INPUT_TOKENS = "rag.llm.input_tokens" + OUTPUT_TOKENS = "rag.llm.output_tokens" + TOTAL_TOKENS = "rag.llm.total_tokens" + + # Reranking attributes + RERANKING_MODEL = "rag.reranking.model" + RERANKING_TOP_N = "rag.reranking.top_n" + + # Status attributes + CACHE_HIT = "rag.cache.hit" + ERROR_TYPE = "rag.error.type" + ERROR_MESSAGE = "rag.error.message" + + +class RAGTracer: + """Context manager for tracing RAG pipeline execution.""" + + def __init__(self, query: str, user_id: str, tenant_id: str, session_id: Optional[str] = None): + """Initialize RAG tracer. + + Args: + query: The user's query string + user_id: Unique user identifier + tenant_id: Tenant identifier + session_id: Optional session identifier for correlation + """ + self.query = query + self.user_id = user_id + self.tenant_id = tenant_id + self.session_id = session_id + self.span = None + self.start_time = None + + def __enter__(self): + """Start the main RAG request span.""" + self.start_time = time.time() + self.span = tracer.start_span("rag.request") + + # Set query attributes + self.span.set_attribute(RAGSpanAttributes.QUERY_TEXT, self.query[:256]) # Limit to 256 chars + self.span.set_attribute(RAGSpanAttributes.QUERY_LENGTH, len(self.query)) + self.span.set_attribute(RAGSpanAttributes.USER_ID, self.user_id) + self.span.set_attribute(RAGSpanAttributes.TENANT_ID, self.tenant_id) + + if self.session_id: + self.span.set_attribute(RAGSpanAttributes.SESSION_ID, self.session_id) + + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """End the main RAG request span.""" + if exc_type is not None: + self.span.set_status(Status(StatusCode.ERROR)) + self.span.set_attribute(RAGSpanAttributes.ERROR_TYPE, exc_type.__name__) + self.span.set_attribute(RAGSpanAttributes.ERROR_MESSAGE, str(exc_val)) + else: + self.span.set_status(Status(StatusCode.OK)) + + self.span.end() + + def trace_retrieval(self, top_k: int = 10, vector_store: str = "faiss"): + """Create a span for vector retrieval operation. + + Args: + top_k: Number of top results to retrieve + vector_store: Type of vector store being used + + Returns: + RetrievalSpan context manager + """ + return RetrievalSpan(self.span, top_k, vector_store) + + def trace_reranking(self, model: str = "cross-encoder", top_n: int = 3): + """Create a span for reranking operation. + + Args: + model: Reranking model being used + top_n: Number of top results after reranking + + Returns: + RerankerSpan context manager + """ + return RerankerSpan(self.span, model, top_n) + + def trace_llm_generation(self, model: str, temperature: float = 0.7): + """Create a span for LLM generation. + + Args: + model: LLM model being used + temperature: Temperature parameter for generation + + Returns: + LLMSpan context manager + """ + return LLMSpan(self.span, model, temperature) + + def trace_document_retrieval(self, doc_id: str, source: str, chunk_index: int, similarity: float): + """Add event for each retrieved document. + + Args: + doc_id: Document identifier + source: Document source/path + chunk_index: Index of the chunk within document + similarity: Similarity score from retrieval + """ + with tracer.start_as_current_span("rag.document.retrieved") as span: + span.set_attribute(RAGSpanAttributes.DOCUMENT_ID, doc_id) + span.set_attribute(RAGSpanAttributes.DOCUMENT_SOURCE, source) + span.set_attribute(RAGSpanAttributes.DOCUMENT_CHUNK_INDEX, chunk_index) + span.set_attribute(RAGSpanAttributes.DOCUMENT_SIMILARITY_SCORE, similarity) + + +class RetrievalSpan: + """Context manager for vector retrieval operations.""" + + def __init__(self, parent_span: Any, top_k: int, vector_store: str): + self.parent_span = parent_span + self.top_k = top_k + self.vector_store = vector_store + self.span = None + self.start_time = None + + def __enter__(self): + self.start_time = time.time() + self.span = tracer.start_span("rag.retrieval", attributes={ + RAGSpanAttributes.RETRIEVAL_TOP_K: self.top_k, + RAGSpanAttributes.VECTOR_STORE_TYPE: self.vector_store, + RAGSpanAttributes.RETRIEVAL_METHOD: "semantic_search", + }) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if exc_type is not None: + self.span.set_status(Status(StatusCode.ERROR)) + self.span.set_attribute(RAGSpanAttributes.ERROR_TYPE, exc_type.__name__) + else: + self.span.set_status(Status(StatusCode.OK)) + + # Record latency metric + latency_ms = (time.time() - self.start_time) * 1000 + retrieval_latency.record(latency_ms, {"vector_store": self.vector_store}) + + self.span.end() + + def set_chunk_count(self, count: int): + """Record number of chunks retrieved.""" + self.span.set_attribute(RAGSpanAttributes.RETRIEVAL_CHUNK_COUNT, count) + retrieval_chunk_count.record(count) + + +class RerankerSpan: + """Context manager for chunk reranking operations.""" + + def __init__(self, parent_span: Any, model: str, top_n: int): + self.parent_span = parent_span + self.model = model + self.top_n = top_n + self.span = None + self.start_time = None + + def __enter__(self): + self.start_time = time.time() + self.span = tracer.start_span("rag.reranking", attributes={ + RAGSpanAttributes.RERANKING_MODEL: self.model, + RAGSpanAttributes.RERANKING_TOP_N: self.top_n, + }) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if exc_type is not None: + self.span.set_status(Status(StatusCode.ERROR)) + else: + self.span.set_status(Status(StatusCode.OK)) + + # Record latency metric + latency_ms = (time.time() - self.start_time) * 1000 + reranking_latency.record(latency_ms, {"model": self.model}) + + self.span.end() + + +class LLMSpan: + """Context manager for LLM generation operations.""" + + def __init__(self, parent_span: Any, model: str, temperature: float): + self.parent_span = parent_span + self.model = model + self.temperature = temperature + self.span = None + self.start_time = None + + def __enter__(self): + self.start_time = time.time() + self.span = tracer.start_span("rag.llm.generation", attributes={ + RAGSpanAttributes.LLM_MODEL: self.model, + RAGSpanAttributes.LLM_TEMPERATURE: self.temperature, + }) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + if exc_type is not None: + self.span.set_status(Status(StatusCode.ERROR)) + else: + self.span.set_status(Status(StatusCode.OK)) + + # Record latency metric + latency_ms = (time.time() - self.start_time) * 1000 + llm_latency.record(latency_ms, {"model": self.model}) + + self.span.end() + + def set_token_usage(self, input_tokens: int, output_tokens: int): + """Record token usage for the LLM generation. + + Args: + input_tokens: Number of input tokens + output_tokens: Number of output tokens + """ + total_tokens = input_tokens + output_tokens + + self.span.set_attribute(RAGSpanAttributes.INPUT_TOKENS, input_tokens) + self.span.set_attribute(RAGSpanAttributes.OUTPUT_TOKENS, output_tokens) + self.span.set_attribute(RAGSpanAttributes.TOTAL_TOKENS, total_tokens) + + token_usage.record(total_tokens, {"model": self.model}) + + +def trace_cache_hit(query_id: str): + """Record a cache hit event. + + Args: + query_id: Identifier for the cached query + """ + with tracer.start_as_current_span("rag.cache.hit") as span: + span.set_attribute("cache.query_id", query_id) + + +def trace_cache_miss(query_id: str): + """Record a cache miss event. + + Args: + query_id: Identifier for the query + """ + with tracer.start_as_current_span("rag.cache.miss") as span: + span.set_attribute("cache.query_id", query_id) + + +def get_current_span(): + """Get the current active span for adding events.""" + return trace.get_current_span() diff --git a/tests/test_rag_tracing.py b/tests/test_rag_tracing.py new file mode 100644 index 0000000..3f32ac5 --- /dev/null +++ b/tests/test_rag_tracing.py @@ -0,0 +1,267 @@ +"""Tests for OpenTelemetry-based RAG data lineage tracing (Issue #159).""" + +from unittest.mock import MagicMock, patch + +import pytest + +from app.observability.rag_tracing import ( + LLMSpan, + RAGSpanAttributes, + RAGTracer, + RerankerSpan, + RetrievalSpan, + trace_cache_hit, + trace_cache_miss, +) + + +class TestRAGTracer: + """Tests for main RAG tracer.""" + + def test_rag_tracer_initialization(self): + """Should initialize RAG tracer with query and user info.""" + tracer = RAGTracer( + query="What is machine learning?", + user_id="user123", + tenant_id="tenant456", + session_id="session789", + ) + + assert tracer.query == "What is machine learning?" + assert tracer.user_id == "user123" + assert tracer.tenant_id == "tenant456" + assert tracer.session_id == "session789" + + @patch("app.observability.rag_tracing.tracer") + def test_rag_tracer_context_manager(self, mock_tracer): + """Should create and manage span context correctly.""" + mock_span = MagicMock() + mock_tracer.start_span.return_value = mock_span + + with RAGTracer( + query="test query", + user_id="user123", + tenant_id="tenant456", + ) as rag_tracer: + assert rag_tracer.span is not None + + # Verify span was ended + mock_span.end.assert_called_once() + # Verify attributes were set + assert mock_span.set_attribute.called + + @patch("app.observability.rag_tracing.tracer") + def test_rag_tracer_error_handling(self, mock_tracer): + """Should set error status on exception.""" + mock_span = MagicMock() + mock_tracer.start_span.return_value = mock_span + + try: + with RAGTracer( + query="test query", + user_id="user123", + tenant_id="tenant456", + ): + raise ValueError("Test error") + except ValueError: + pass + + # Verify error status was set + mock_span.set_status.assert_called() + + @patch("app.observability.rag_tracing.tracer") + def test_rag_tracer_query_text_truncation(self, mock_tracer): + """Should truncate long query text for span attributes.""" + mock_span = MagicMock() + mock_tracer.start_span.return_value = mock_span + + long_query = "a" * 500 # 500 characters + + with RAGTracer( + query=long_query, + user_id="user123", + tenant_id="tenant456", + ): + pass + + # Find the call that sets query text + for call in mock_span.set_attribute.call_args_list: + if call[0][0] == RAGSpanAttributes.QUERY_TEXT: + # Should be truncated to 256 chars + assert len(call[0][1]) == 256 + + +class TestRetrievalSpan: + """Tests for vector retrieval spans.""" + + @patch("app.observability.rag_tracing.tracer") + @patch("app.observability.rag_tracing.retrieval_latency") + def test_retrieval_span_creation(self, mock_latency, mock_tracer): + """Should create retrieval span with correct attributes.""" + mock_parent_span = MagicMock() + mock_span = MagicMock() + mock_tracer.start_span.return_value = mock_span + + with RetrievalSpan(mock_parent_span, top_k=10, vector_store="faiss"): + pass + + # Verify span was created with attributes + mock_tracer.start_span.assert_called() + + @patch("app.observability.rag_tracing.tracer") + @patch("app.observability.rag_tracing.retrieval_latency") + def test_retrieval_span_chunk_count(self, mock_latency, mock_tracer): + """Should record chunk count in retrieval span.""" + mock_parent_span = MagicMock() + mock_span = MagicMock() + mock_tracer.start_span.return_value = mock_span + + with RetrievalSpan(mock_parent_span, top_k=10, vector_store="faiss") as span: + span.set_chunk_count(5) + + # Verify chunk count was set + assert mock_span.set_attribute.called + + +class TestRerankerSpan: + """Tests for chunk reranking spans.""" + + @patch("app.observability.rag_tracing.tracer") + @patch("app.observability.rag_tracing.reranking_latency") + def test_reranker_span_creation(self, mock_latency, mock_tracer): + """Should create reranker span with correct attributes.""" + mock_parent_span = MagicMock() + mock_span = MagicMock() + mock_tracer.start_span.return_value = mock_span + + with RerankerSpan(mock_parent_span, model="cross-encoder", top_n=3): + pass + + mock_tracer.start_span.assert_called() + + +class TestLLMSpan: + """Tests for LLM generation spans.""" + + @patch("app.observability.rag_tracing.tracer") + @patch("app.observability.rag_tracing.llm_latency") + def test_llm_span_creation(self, mock_latency, mock_tracer): + """Should create LLM span with correct attributes.""" + mock_parent_span = MagicMock() + mock_span = MagicMock() + mock_tracer.start_span.return_value = mock_span + + with LLMSpan(mock_parent_span, model="gpt-3.5-turbo", temperature=0.7): + pass + + mock_tracer.start_span.assert_called() + + @patch("app.observability.rag_tracing.tracer") + @patch("app.observability.rag_tracing.llm_latency") + @patch("app.observability.rag_tracing.token_usage") + def test_llm_span_token_usage(self, mock_token_usage, mock_latency, mock_tracer): + """Should record token usage in LLM span.""" + mock_parent_span = MagicMock() + mock_span = MagicMock() + mock_tracer.start_span.return_value = mock_span + + with LLMSpan(mock_parent_span, model="gpt-3.5-turbo", temperature=0.7) as span: + span.set_token_usage(input_tokens=100, output_tokens=50) + + # Verify tokens were set + mock_span.set_attribute.assert_called() + # Verify metric was recorded + mock_token_usage.record.assert_called() + + +class TestCacheTracing: + """Tests for cache hit/miss tracing.""" + + @patch("app.observability.rag_tracing.tracer") + def test_trace_cache_hit(self, mock_tracer): + """Should trace cache hit event.""" + mock_span = MagicMock() + mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span + + trace_cache_hit("query123") + + mock_tracer.start_as_current_span.assert_called() + + @patch("app.observability.rag_tracing.tracer") + def test_trace_cache_miss(self, mock_tracer): + """Should trace cache miss event.""" + mock_span = MagicMock() + mock_tracer.start_as_current_span.return_value.__enter__.return_value = mock_span + + trace_cache_miss("query123") + + mock_tracer.start_as_current_span.assert_called() + + +class TestSpanAttributes: + """Tests for RAG span attribute constants.""" + + def test_span_attributes_exist(self): + """Should have all required span attributes defined.""" + # Query attributes + assert hasattr(RAGSpanAttributes, "QUERY_TEXT") + assert hasattr(RAGSpanAttributes, "QUERY_LENGTH") + assert hasattr(RAGSpanAttributes, "USER_ID") + assert hasattr(RAGSpanAttributes, "TENANT_ID") + + # Retrieval attributes + assert hasattr(RAGSpanAttributes, "RETRIEVAL_TOP_K") + assert hasattr(RAGSpanAttributes, "RETRIEVAL_CHUNK_COUNT") + assert hasattr(RAGSpanAttributes, "VECTOR_STORE_TYPE") + + # Document attributes + assert hasattr(RAGSpanAttributes, "DOCUMENT_ID") + assert hasattr(RAGSpanAttributes, "DOCUMENT_SOURCE") + assert hasattr(RAGSpanAttributes, "DOCUMENT_SIMILARITY_SCORE") + + # LLM attributes + assert hasattr(RAGSpanAttributes, "LLM_MODEL") + assert hasattr(RAGSpanAttributes, "INPUT_TOKENS") + assert hasattr(RAGSpanAttributes, "OUTPUT_TOKENS") + assert hasattr(RAGSpanAttributes, "TOTAL_TOKENS") + + +class TestIntegrationScenario: + """Integration tests for complete RAG tracing scenario.""" + + @patch("app.observability.rag_tracing.tracer") + def test_full_rag_tracing_scenario(self, mock_tracer): + """Should trace complete RAG pipeline execution.""" + mock_span = MagicMock() + mock_retrieval_span = MagicMock() + mock_reranker_span = MagicMock() + mock_llm_span = MagicMock() + + # Setup mock return values + mock_tracer.start_span.side_effect = [ + mock_span, # Main RAG span + mock_retrieval_span, # Retrieval span + mock_reranker_span, # Reranker span + mock_llm_span, # LLM span + ] + + with RAGTracer( + query="What is AI?", + user_id="user123", + tenant_id="tenant456", + session_id="session789", + ) as rag_tracer: + # Simulate retrieval + with rag_tracer.trace_retrieval(top_k=10, vector_store="faiss") as ret: + ret.set_chunk_count(8) + + # Simulate reranking + with rag_tracer.trace_reranking(model="cross-encoder", top_n=3): + pass + + # Simulate LLM generation + with rag_tracer.trace_llm_generation(model="gpt-3.5-turbo") as llm: + llm.set_token_usage(input_tokens=200, output_tokens=100) + + # Verify all spans were created + assert mock_tracer.start_span.called