feat: instrument plain-agents with opensearch-genai-observability-sdk-py (opensearch-project#85)

Replace manual OTel TracerProvider/exporter setup and span.set_attribute()
calls with the SDK's register(), observe(), and enrich() APIs.

- register() replaces ~20-30 lines of TracerProvider + exporter config
- observe() as context manager replaces tracer.start_as_current_span()
- enrich() replaces dozens of span.set_attribute("gen_ai.*", ...) calls
- Net reduction of ~112 lines while preserving identical trace output
- Metrics and logging setup unchanged (SDK handles tracing only)
- Fault injection, ASGI middleware, MCP integration all preserved

Requires opensearch-genai-observability-sdk-py >= 0.2.6.

Signed-off-by: Vamshi Vijay Nakkirtha <vamsimanohar@gmail.com>
Signed-off-by: Vamsi Manohar <reddyvam@amazon.com>

Author: vamsimanohar

examples/plain-agents/multi-agent-planner/events-agent/Dockerfile

@@ -7,8 +7,7 @@ RUN pip install --no-cache-dir \
     fastapi \
     uvicorn \
     httpx \
-    opentelemetry-api \
-    opentelemetry-sdk \
+    opensearch-genai-observability-sdk-py>=0.2.7 \
     opentelemetry-exporter-otlp-proto-grpc \
     opentelemetry-instrumentation-asgi
 

examples/plain-agents/multi-agent-planner/events-agent/main.py

@@ -2,6 +2,10 @@
 """
 Events Agent - Fetches local events for a destination.
 Supports fault injection for testing observability.
+
+Instrumented with opensearch-genai-observability-sdk-py:
+- register() replaces ~20 lines of manual TracerProvider/exporter setup
+- observe() + enrich() replace manual span creation + set_attribute() calls
 """
 
 import json
@@ -15,10 +19,7 @@
 import httpx
 from fastapi import FastAPI
 from opentelemetry import trace, metrics
-from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
 from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter
-from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.sdk.metrics import MeterProvider
 from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
 from opentelemetry.sdk.resources import Resource
@@ -27,6 +28,8 @@
 from opentelemetry.propagate import inject
 from pydantic import BaseModel, Field
 
+from opensearch_genai_observability_sdk_py import Op, enrich, observe, register
+
 
 # MCP Server configuration
 MCP_SERVER_URL = os.getenv("MCP_SERVER_URL", "http://mcp-server:8003")
@@ -141,29 +144,25 @@ class ErrorResponse(BaseModel):
     agent_id: str
 
 
-def setup_telemetry(service_name: str, otlp_endpoint: str):
-    resource = Resource.create({
-        "service.name": service_name,
-        "service.version": "1.0.0",
-        "gen_ai.agent.id": AGENT_ID,
-        "gen_ai.agent.name": AGENT_NAME,
-    })
-    tracer_provider = TracerProvider(resource=resource)
-    tracer_provider.add_span_processor(
-        BatchSpanProcessor(OTLPSpanExporter(endpoint=otlp_endpoint, insecure=True))
-    )
-    trace.set_tracer_provider(tracer_provider)
-    metric_reader = PeriodicExportingMetricReader(
-        OTLPMetricExporter(endpoint=otlp_endpoint, insecure=True),
-        export_interval_millis=10000,
-    )
-    meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
-    metrics.set_meter_provider(meter_provider)
-    return trace.get_tracer(service_name), metrics.get_meter(service_name)
-
-
+# --- Telemetry setup ---
 otlp_endpoint = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4317")
-tracer, meter = setup_telemetry("events-agent", otlp_endpoint)
+
+# One line replaces ~20 lines of TracerProvider + exporter setup
+register(
+    endpoint=f"grpc://{otlp_endpoint.replace('http://', '').replace('https://', '')}",
+    service_name="events-agent",
+    service_version="1.0.0",
+)
+
+# Metrics (SDK handles tracing only)
+# TODO: unify Resource with register() when SDK supports metrics
+resource = Resource.create({"service.name": "events-agent"})
+metric_reader = PeriodicExportingMetricReader(
+    OTLPMetricExporter(endpoint=otlp_endpoint, insecure=True),
+    export_interval_millis=10000,
+)
+meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
+metrics.set_meter_provider(meter_provider)
 
 inner_app = FastAPI(title="Events Agent", version="1.0.0")
 
@@ -182,41 +181,40 @@ async def health():
 @inner_app.post("/events")
 async def get_events(request: EventsRequest):
     model = random.choice(MODELS)
-    
+    provider = SYSTEMS[model]
+
     # Promote gen_ai attributes to the root HTTP span so the UI can read them
+    enrich(
+        model=model,
+        provider=provider,
+        agent_id=AGENT_ID,
+        input_messages=[{"role": "user", "parts": [{"type": "text", "content": f"Find events in {request.destination}"}]}],
+    )
     root_span = trace.get_current_span()
-    root_span.set_attribute("gen_ai.system", SYSTEMS[model])
     root_span.set_attribute("gen_ai.agent.name", AGENT_NAME)
-    root_span.set_attribute("gen_ai.request.model", model)
     root_span.set_attribute("gen_ai.operation.name", "invoke_agent")
-    root_span.set_attribute("gen_ai.input.messages", json.dumps(
-        [{"role": "user", "parts": [{"type": "text", "content": f"Find events in {request.destination}"}]}]
-    ))
-    
-    with tracer.start_as_current_span(
-        "invoke_agent",
-        kind=SpanKind.INTERNAL,
-        attributes={
-            "gen_ai.operation.name": "invoke_agent",
-            "gen_ai.agent.id": AGENT_ID,
-            "gen_ai.agent.name": AGENT_NAME,
-            "gen_ai.system": SYSTEMS[model],
-            "gen_ai.request.model": model,
-            "gen_ai.tool.definitions": json.dumps(TOOL_DEFINITIONS),
-        },
-    ) as span:
+
+    with observe(AGENT_NAME, op=Op.INVOKE_AGENT) as span:
+        enrich(
+            model=model,
+            provider=provider,
+            agent_id=AGENT_ID,
+            tool_definitions=TOOL_DEFINITIONS,
+        )
+
         destination = request.destination.lower()
         date = request.date or datetime.now().strftime("%Y-%m-%d")
         fault = request.fault
 
         # Synthetic "thinking" LLM call
-        with tracer.start_as_current_span("chat", kind=SpanKind.INTERNAL) as chat_span:
-            chat_span.set_attribute("gen_ai.operation.name", "chat")
-            chat_span.set_attribute("gen_ai.system", SYSTEMS[model])
-            chat_span.set_attribute("gen_ai.request.model", model)
-            chat_span.set_attribute("gen_ai.usage.input_tokens", random.randint(100, 500))
-            chat_span.set_attribute("gen_ai.usage.output_tokens", random.randint(50, 200))
-            chat_span.set_attribute("gen_ai.response.finish_reasons", ["tool_calls"])
+        with observe("events-reasoning", op=Op.CHAT):
+            enrich(
+                model=model,
+                provider=provider,
+                input_tokens=random.randint(100, 500),
+                output_tokens=random.randint(50, 200),
+                finish_reason="tool_calls",
+            )
             time.sleep(random.uniform(0.05, 0.15))
 
         # Check for fault injection
@@ -267,20 +265,16 @@ async def get_events(request: EventsRequest):
         # Tool execution via MCP server
         session_id = uuid4().hex
         request_id = uuid4().hex[:8]
-        with tracer.start_as_current_span(
-            "tools/call fetch_events_api",
-            kind=SpanKind.CLIENT,
-            attributes={
-                "mcp.method.name": "tools/call",
-                "mcp.session.id": session_id,
-                "mcp.protocol.version": MCP_PROTOCOL_VERSION,
-                "jsonrpc.request.id": request_id,
-                "gen_ai.operation.name": "execute_tool",
-                "gen_ai.tool.name": "fetch_events_api",
-                "network.transport": "tcp",
-                "network.protocol.name": "http",
-            },
-        ):
+
+        # MCP tool call — uses observe() for the span, with MCP-specific attributes set manually
+        with observe("fetch_events_api", op=Op.EXECUTE_TOOL, kind=SpanKind.CLIENT) as tool_span:
+            tool_span.set_attribute("mcp.method.name", "tools/call")
+            tool_span.set_attribute("mcp.session.id", session_id)
+            tool_span.set_attribute("mcp.protocol.version", MCP_PROTOCOL_VERSION)
+            tool_span.set_attribute("jsonrpc.request.id", request_id)
+            tool_span.set_attribute("network.transport", "tcp")
+            tool_span.set_attribute("network.protocol.name", "http")
+
             headers = {"mcp-session-id": session_id}
             inject(headers)
             payload = {
@@ -293,12 +287,15 @@ async def get_events(request: EventsRequest):
             events = [Event(name=e["name"], type=e["type"], venue=e.get("venue", "TBD"), date=e.get("date", date)) for e in events_list]
 
         span.set_attribute("events.count", len(events))
-        
-        root_span.set_attribute("gen_ai.output.messages", json.dumps(
-            [{"role": "assistant", "parts": [{"type": "text", "content": json.dumps([e.model_dump() for e in events])}]}]
-        ))
-        
-        return EventsResponse(destination=request.destination, events=events, agent_id=AGENT_ID)
+
+    # Set output on the parent HTTP request span. This enrich() is intentionally
+    # outside the observe() block — exiting observe() restores the parent span
+    # context, so enrich() here targets the HTTP request span, not the agent span.
+    enrich(
+        output_messages=[{"role": "assistant", "parts": [{"type": "text", "content": json.dumps([e.model_dump() for e in events])}]}],
+    )
+
+    return EventsResponse(destination=request.destination, events=events, agent_id=AGENT_ID)
 
 
 app = OpenTelemetryMiddleware(inner_app)

examples/plain-agents/multi-agent-planner/orchestrator/Dockerfile

@@ -7,8 +7,7 @@ RUN pip install --no-cache-dir \
     fastapi \
     uvicorn \
     httpx \
-    opentelemetry-api \
-    opentelemetry-sdk \
+    opensearch-genai-observability-sdk-py>=0.2.7 \
     opentelemetry-exporter-otlp-proto-grpc \
     opentelemetry-instrumentation-httpx \
     opentelemetry-instrumentation-asgi