Stop Flying Blind: Adding OpenTelemetry Tracing to Your AI Gateway

Building an AI Gateway has become standard practice for organizations moving from experimental LLM playgrounds to production-grade applications. The gateway sits between your application code and multiple model providers (OpenAI, Anthropic, Local LLMs), handling routing, caching, rate limiting, and cost control. However, without deep observability, your team is essentially flying blind. When a user complains that the chat response was "slow" or "weird," can you distinguish between a network timeout, a specific model hallucination, or a cache miss that triggered a redundant expensive inference?

Standard APM tools struggle with AI workloads because they lack semantic understanding of Large Language Model (LLM) interactions. They see HTTP requests, but they miss the semantic "spans" that matter: the actual prompt sent, the token count consumed, the finish reason, and the semantic relationship between a retrieval step (RAG) and the final generation. This tutorial walks you through implementing a robust OpenTelemetry (OTel) tracing pipeline for an AI Gateway. We will cover semantic attribute conventions, context propagation across asynchronous boundaries, and how to export this data to a backend like Jaeger or Grafana Tempo.

Table of Contents

• 1. Why Standard Tracing Fails for AI Gateways
• 2. Architecture: The AI Gateway Observability Pipeline
• 3. Semantic Conventions: What to Record
• 4. Implementation: Instrumenting the Gateway with Python
• 5. Handling Context Propagation in Async Python
• 6. Redacting PII and Sensitive Data
• 7. Visualizing the Data: Jaeger and Grafana
• 8. Frequently Asked Questions

1. Why Standard Tracing Fails for AI Gateways

Traditional distributed tracing focuses on "where" the latency is (DB query, API call). In AI systems, the question is often "why" the model behaved a certain way. A standard HTTP trace will tell you that /api/chat took 3.2 seconds. It will not tell you that 2.5 seconds of that was spent in the LLM inference engine generating tokens, or that the prompt was 4,000 tokens long (hence the high cost).

Furthermore, AI workflows are often multi-step. A typical RAG (Retrieval-Augmented Generation) flow looks like this:

1. Embed the user query.
2. Search the vector database.
3. Rerank the results.
4. Send the context + query to the LLM.

If you don't link these operations into a single trace, you cannot diagnose issues like "the vector DB returned irrelevant results, causing the LLM to hallucinate." OpenTelemetry provides the primitive Span that allows us to nest these operations, creating a hierarchical view of the request lifecycle.

2. Architecture: The AI Gateway Observability Pipeline

Before writing code, we must define the architecture. An AI Gateway typically acts as a reverse proxy or a SDK wrapper. For this tutorial, we will assume a Python-based gateway service that receives requests and forwards them to LLM providers.

The Data Flow:

1. Ingestion: The Gateway receives an inbound HTTP request.
2. Root Span: The Gateway creates a root span representing the entire request.
3. Child Spans: As the Gateway processes the request, it creates child spans for specific operations:
   • retrieval.search: For fetching context from a vector DB.
   • llm.inference: For the actual call to the model provider.
   • llm.generate: For streaming chunks (if applicable).
4. Enrichment: Each span is enriched with attributes defined by the OpenTelemetry GenAI Semantic Conventions (e.g., gen_ai.model.name, gen_ai.usage.prompt_tokens).
5. Export: The spans are batched and exported via OTLP to a collector.
6. Storage & Visualization: The Collector routes data to Jaeger/Tempo for viewing.

*Note: For a deeper dive into the underlying collector architecture, you might find the technical breakdown on Tamiz's Insights on high-throughput telemetry pipelines useful.

3. Semantic Conventions: What to Record

To make traces readable by standard tools, you should adhere to the emerging OpenTelemetry GenAI Semantic Conventions. While these are still evolving, defining a consistent set of attributes now ensures you won't have to rewrite your instrumentation later.

Key attributes to record:

Important: Do not record raw prompt/completion content in spans unless you have a strict PII redaction pipeline. Spans are often sampled at a lower rate than metrics, but they are still logged. We will address this in Section 6.

4. Implementation: Instrumenting the Gateway with Python

We will build a minimal AI Gateway using FastAPI and httpx. We will instrument it using the opentelemetry-sdk and opentelemetry-exporter-otlp.

Prerequisites

Ensure you have Python 3.10+ installed. Create a virtual environment and install dependencies:

pip install fastapi uvicorn httpx opentelemetry-sdk opentelemetry-exporter-otlp-proto-http opentelemetry-instrumentation-fastapi opentelemetry-instrumentation-httpx

4.1 Setting up the Tracer Provider

First, we need to initialize the OpenTelemetry SDK. This sets up the environment where spans will be generated and exported.

# app/telemetry.py
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import Resource, SERVICE_NAME
from opentelemetry.sdk.trace.sampling import TraceIdRatioBased
import os

OTEL_ENDPOINT = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317")
OTEL_HEADERS = os.getenv("OTEL_EXPORTER_OTLP_HEADERS", "")

def get_tracer_provider() -> TracerProvider:
    resource = Resource.create({
        SERVICE_NAME: "ai-gateway",
        "service.version": "1.0.0",
    })
    
    provider = TracerProvider(resource=resource)
    
    # Configure sampling. In production, you might use parent-based sampling.
    # Here we use a 50% ratio for demonstration.
    sampler = TraceIdRatioBased(0.5)
    
    exporter = OTLPSpanExporter(
        endpoint=OTEL_ENDPOINT,
        headers=OTEL_HEADERS,
        insecure=True # For local dev; use HTTPS in prod
    )
    
    processor = BatchSpanProcessor(exporter)
    provider.add_span_processor(processor)
    
    trace.set_tracer_provider(provider)
    return provider

4.2 The Gateway Logic and Instrumentation

Here is the core of the AI Gateway. We will create a service that wraps a hypothetical LLM provider. Notice how we create nested spans for the "Retrieval" step and the "Generation" step.

# app/gateway.py
import asyncio
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import httpx
from opentelemetry import trace
from app.telemetry import get_tracer_provider

tracer = trace.get_tracer("ai_gateway")

app = FastAPI()

# Initialize provider
get_tracer_provider()

class ChatRequest(BaseModel):
    message: str
    user_id: str

class ChatResponse(BaseModel):
    response: str
    usage: dict

# Mock LLM and Vector DB calls for demonstration
async def mock_vector_search(query: str):
    """Simulates a vector database lookup."""
    async with tracer.start_as_current_span("vector_db.search") as span:
        span.set_attribute("db.system", "pinecone")
        await asyncio.sleep(0.1) # Simulate latency
        return ["Context chunk 1", "Context chunk 2"]

async def mock_llm_call(prompt: str, model: str = "gpt-4"):
    """Simulates an LLM API call."""
    async with tracer.start_as_current_span("llm.generate") as span:
        span.set_attribute("gen_ai.system", "openai")
        span.set_attribute("gen_ai.request.model", model)
        span.set_attribute("gen_ai.request.temperature", 0.7)
        
        # Simulate token counting and generation latency
        await asyncio.sleep(0.5)
        
        # Mock response
        response_text = f"Response based on: {prompt[:20]}..."
        
        span.set_attribute("gen_ai.response.model", f"{model}-2024-01")
        span.set_attribute("gen_ai.usage.prompt_tokens", 150)
        span.set_attribute("gen_ai.usage.completion_tokens", 1024)
        span.set_attribute("gen_ai.finish_reason", "stop")
        
        return response_text

@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    # 1. Create Root Span for the entire request
    with tracer.start_as_current_span("gateway.handle_chat") as root_span:
        root_span.set_attribute("user.id", request.user_id)
        
        # 2. Step 1: Retrieval
        context = await mock_vector_search(request.message)
        
        # 3. Step 2: Construction of Prompt
        final_prompt = f"Context: {str(context)}\nQuery: {request.message}"
        
        # 4. Step 3: Generation
        response_text = await mock_llm_call(final_prompt)
        
        return ChatResponse(
            response=response_text,
            usage={"prompt_tokens": 150, "completion_tokens": 1024}
        )

Key Takeaway: Notice the use of tracer.start_as_current_span. This context manager automatically handles the parent-child relationship. The vector_db.search and llm.generate spans will be automatically nested inside gateway.handle_chat.

5. Handling Context Propagation in Async Python

A common pitfall when instrumenting Python AI gateways is context loss in asynchronous code. If you spawn a new asyncio.Task or thread, the OpenTelemetry context (which holds the current trace ID) does not propagate automatically unless you handle it correctly.

In the example above, we awaited everything in the same task, so propagation worked. However, if your gateway uses concurrent.futures or separate worker processes, you must explicitly capture the SpanContext and inject it into the child environment.

Here is how to propagate context to an httpx request (which many gateways use to call LLM providers). The opentelemetry-instrumentation-httpx package handles this mostly, but you must ensure the traceparent header is sent.

import httpx
from opentelemetry.trace import get_current_span, SpanKind
from opentelemetry.trace.context import Context
from opentelemetry.propagate import inject

async def call_external_llm_with_trace(url: str, payload: dict):
    tracer = trace.get_tracer("http_client")
    
    with tracer.start_as_current_span("http.client.llm_call", kind=SpanKind.CLIENT) as span:
        span.set_attribute("http.method", "POST")
        span.set_attribute("http.url", url)
        
        # We need to inject the current trace context into the headers
        headers = {}
        inject(headers) # This adds 'traceparent' header
        
        async with httpx.AsyncClient() as client:
            try:
                response = await client.post(url, json=payload, headers=headers)
                response.raise_for_status()
                
                span.set_attribute("http.status_code", response.status_code)
                return response.json()
            except httpx.HTTPError as exc:
                span.set_status(trace.StatusCode.ERROR, str(exc))
                span.record_exception(exc)
                raise

By injecting the headers, you ensure that if the LLM provider also supports OpenTelemetry (or if you are tracing across microservices), the trace remains contiguous. For local LLMs that don't support OTel, at least the client-side duration and status are recorded.

6. Redacting PII and Sensitive Data

This is the most critical security aspect of AI tracing. LLM prompts often contain user PII, API keys, or proprietary data. Never record raw prompts/completions in spans without a redaction layer.

We can add a middleware or a utility function to sanitize attributes before they are attached to the span.

import re
from opentelemetry.trace import Span

def redact_pii(text: str) -> str:
    """Simple regex-based redaction for emails and phone numbers."""
    if not text:
        return text
    
    # Email addresses
    text = re.sub(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}\b', '[EMAIL_REDACTED]', text)
    
    # US Phone Numbers (example)
    text = re.sub(r'\b\d{3}[-.]?\d{3}[-.]?\d{4}\b', '[PHONE_REDACTED]', text)
    
    # API Keys (generic pattern)
    text = re.sub(r'sk-[A-Za-z0-9]{20,}', '[API_KEY_REDACTED]', text)
    
    return text

# Modify the mock_llm_call to use this
async def mock_llm_call_safe(prompt: str, model: str = "gpt-4"):
    async with tracer.start_as_current_span("llm.generate") as span:
        # ... other attributes ...
        
        # Instead of span.set_attribute("gen_ai.prompt", prompt)
        span.set_attribute("gen_ai.prompt", redact_pii(prompt))
        
        # In production, log the full prompt to a secure S3 bucket with IAM controls,
        # and store only the hash or reference in the trace span.
        span.set_attribute("gen_ai.prompt.hash", get_sha256_hash(prompt))

Best Practice: Store the full prompt/completion in a separate, secure log store (like CloudWatch or Datadog Logs) with strict access controls. In the OpenTelemetry span, store only the hash or a reference ID. This allows developers to debug by looking up the ID in the secure logs, while keeping the trace backend clean and compliant.

7. Visualizing the Data: Jaeger and Grafana

To see the results of our work, we need a backend. For local development, the easiest way is to run Jaeger in All-In-One mode.

docker run -p 16686:16686 -p 4317:4317 jaegertracing/all-in-one

Visit http://localhost:16686 in your browser. Search for the service ai-gateway. You will see a trace waterfall that looks like this:

gateway.handle_chat (200ms)
├── vector_db.search (50ms)
└── llm.generate (150ms)

Interpreting the Waterfall

1. Latency Breakdown: If llm.generate takes 80% of the total time, you know the issue is inference, not retrieval. If vector_db.search is slow, you need to optimize your indexing.
2. Error Propagation: If llm.generate fails (e.g., 429 Too Many Requests), the root span gateway.handle_chat will also be marked as errored because of the context propagation. This makes it easy to filter for "All Errored Requests" in Jaeger.
3. Cost Correlation: By exporting gen_ai.usage.* attributes, you can build dashboards in Grafana that correlate Trace IDs with Cost. For example, you can create a panel that shows the average cost per trace ID, helping you identify expensive queries that might be failing or returning low-quality results.

8. Frequently Asked Questions

Should I trace every single LLM call, or only sampled ones?

In high-throughput systems, tracing every call can be expensive. OpenTelemetry's BatchSpanProcessor and sampling strategies (like TraceIdRatioBased(0.1)) allow you to sample only 10% of traces. However, you should use Head-Based Sampling with a higher rate for error spans. This ensures that if an LLM call fails (which is more interesting), you always capture that trace, even if successful calls are dropped.

How do I handle streaming responses?

Streaming responses (SSE) create a challenge because the completion_tokens are not known until the stream ends. You can update the span attributes after the stream finishes. Use span.set_attribute after the async for loop completes. Do not finalize the span (end it) until the entire stream has been processed, otherwise, you will have an incomplete trace.

Does OpenTelemetry work with local LLMs like Llama 3?

Yes, but you must instrument the client library (e.g., transformers, ollama, vllm). Since local LLMs run on your infrastructure, the "network" latency is replaced by CPU/GPU compute latency. You can use opentelemetry-instrumentation-nvidia-smi (or similar) to capture GPU utilization metrics alongside the trace spans. This provides a holistic view: the trace tells you what happened, and the metrics tell you how the hardware performed.

By implementing this observability layer, you transform your AI Gateway from a "black box" into a transparent, debuggable system. You gain the ability to answer complex questions about performance, cost, and quality, ultimately enabling safer and more reliable deployment of AI applications in production.