Context propagation with Traces for Agents

[nadheesh]

To maintain correlation between traces generated by different agents for the same request, it is necessary to propagate a unique identifier (e.g., test-case-id) through the trace context. This ensures that traces from all agent executions related to a single test case can be linked and analyzed together.

OpenTelemetry (OTEL) provides support for context propagation, but leveraging it in practice requires additional setup and configuration. This capability becomes particularly critical when evaluating agents across multiple test cases: once an agent is invoked, we need a way to associate all resulting traces with the current test case. Passing this identifier as baggage to the agent enables downstream spans to carry this metadata, which can later be used to filter or correlate traces.

User stories:

• Distributed tracing for distributed multi-agents: Observe and analyse distributed agent's behaviors
• Simplified Agent evaluation: Need a correlation id to filter traces for each test-case accurately and efficiently.
• Session/Task Id support: Tracking sessions and task through a correlation id.

Approaches for Achieving Context Propagation

1) Framework Instrumentation (e.g., FastAPI)

Automatically extracts context from incoming requests. Propagates baggage to all child spans, including agent traces.

from fastapi import FastAPI
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.requests import RequestsInstrumentor

app = FastAPI()

# These two lines fix propagation for the entire app
FastAPIInstrumentor.instrument_app(app) # Extracts context from incoming headers
RequestsInstrumentor().instrument()     # Injects context into outgoing requests

Notes:

• Works only with supported frameworks (e.g., FastAPI, Flask).
• Requires modifying the application code to register instrumentation for each app instance.

2) Manual Propagation

Explicitly extract incoming headers and attach them to the OpenTelemetry context in the code. Baggage can then be propagated to all spans spawned during agent execution.

from fastapi import Request
from opentelemetry import propagate
import requests

@app.get("/chat")
async def chat_agent(request: Request):
    # 1. Extract context from incoming FastAPI headers
    ctx = propagate.extract(request.headers)
    
    # 2. Inject context into a headers dictionary for outgoing 'requests'
    headers = {}
    propagate.inject(headers, context=ctx)
    
    # 3. Use the headers in your call
    response = requests.get("https://other-service.com", headers=headers)
    return response.json()

Notes:

• Works with any framework or service type.
• Requires changes in the code wherever context propagation is needed.

3) Auto-Instrumentation with OpenTelemetry

Uses OTEL’s auto-instrumentation capabilities to capture context across HTTP, RPC, or other supported protocols.

a. Installation

pip install opentelemetry-distro opentelemetry-exporter-otlp
opentelemetry-bootstrap -a install

b. Run

export OTEL_SERVICE_NAME="your-service-name"
export OTEL_TRACES_EXPORTER="otlp"
export OTEL_EXPORTER_OTLP_ENDPOINT="http://your-collector:4317"

# Wrap the existing startup command
opentelemetry-instrument uvicorn main:app --host 0.0.0.0 --port 8000

Notes:

• Provides broad coverage without manual instrumentation.
• Limited by supported frameworks and languages.

4) Hybrid approach to extend amp-instumentation (Recommended)

To support any framework without modifying user code, I propose extending sitecustomization.py to handle context propagation centrally.

from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor

# 1. Initialize the instrumentor FIRST
# Repeat for other supported frameworks 
FastAPIInstrumentor().instrument()

# 2. ONLY NOW import and use FastAPI
from fastapi import FastAPI
app = FastAPI()

This approach enables us to enhance our existing amp-instrumentation package to provide context propagation not only for FastAPI but for any supported framework. By implementing this centralized mechanism, we can ensure consistent trace correlation across distributed agents, streamline agent evaluations, and minimize the need for framework-specific instrumentation or manual code changes.