Exporting observability traces with OpenTelemetry

#!/usr/bin/env bash
set -euo pipefail

cd "$(dirname "$0")"
source .venv/bin/activate

export TENANT_ID="<your-tenant-id>"
export AGENT_ID="<your-agent-id>"
export ENVIRONMENT_NAME="draft"
export OTEL_RESOURCE_ATTRIBUTES="tenant.id=${TENANT_ID},deployment.environment=${ENVIRONMENT_NAME}"
export API_KEY="<your-api-key>"
export TOKEN_URL="https://iam.platform.saas.ibm.com/siusermgr/api/1.0/apikeys/token"
export OTEL_EXPORT_URL="https://api.watson-orchestrate.ibm.com/instances/<instance-id>/v1/orchestrate/inject/traces"

import uuid
import os

RESOURCE_ATTRIBUTES = {
    "service.name": "my-external-agent",
    "service.version": "1.0.0",
    "application": "my-agent-server",
}

def get_required_env(name: str) -> str:
    value = os.environ.get(name)
    if not value:
        raise RuntimeError(f"Missing required environment variable: {name}")
    return value

def build_span_attributes() -> dict[str, str]:
    return {
        "agent.id": get_required_env("AGENT_ID"),
        "langfuse.user.id": f"usr_{uuid.uuid4().hex[:12]}",
        "langfuse.session.id": str(uuid.uuid4()),
    }

import json
import ssl
from urllib.request import Request, urlopen

from opentelemetry import trace
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import BatchSpanProcessor

def get_bearer_token() -> str:
    """Exchange API_KEY for a bearer token via the IAM/MCSP token endpoint."""
    api_key = get_required_env("API_KEY")
    token_url = get_required_env("TOKEN_URL")

    payload = json.dumps({"apikey": api_key}).encode("utf-8")
    request = Request(
        token_url,
        data=payload,
        headers={
            "Content-Type": "application/json",
            "Accept": "application/json",
        },
        method="POST",
    )
    with urlopen(request, timeout=30, context=ssl.create_default_context()) as response:
        token_response = json.loads(response.read().decode("utf-8"))

    bearer_token = token_response.get("token") or token_response.get("access_token")
    if not bearer_token:
        raise RuntimeError(f"Token response missing token/access_token: {token_response}")
    return str(bearer_token)

def setup_tracer() -> tuple[trace.Tracer, TracerProvider]:
    otel_export_url = get_required_env("OTEL_EXPORT_URL")
    resource = Resource.create(RESOURCE_ATTRIBUTES)
    provider = TracerProvider(resource=resource)

    bearer_token = get_bearer_token()
    exporter = OTLPSpanExporter(
        endpoint=otel_export_url,
        headers={"Authorization": f"Bearer {bearer_token}"},
    )

    processor = BatchSpanProcessor(exporter)
    provider.add_span_processor(processor)
    trace.set_tracer_provider(provider)
    return trace.get_tracer("wxo-server"), provider

from typing import cast

def main() -> None:
    tracer, provider = setup_tracer()
    span_attributes = build_span_attributes()

    with tracer.start_as_current_span("sample-operation") as span:
        for key, value in span_attributes.items():
            span.set_attribute(key, value)

        span.set_attribute("input", "User asked Agent to greet and start an orchestrate run")
        span.set_attribute("output", "Agent prepared a greeting response and submitted the orchestrate run")

        # Child span: load conversation context
        with tracer.start_as_current_span("load-conversation-context") as child_span:
            child_span.set_attribute("operation.type", "context_lookup")

        # Child span: resolve and call the agent
        with tracer.start_as_current_span("resolve-agent") as child_span:
            child_span.set_attribute("operation.type", "agent_resolution")

            with tracer.start_as_current_span("prepare-agent-response") as nested_span:
                nested_span.set_attribute("operation.type", "response_preparation")

                with tracer.start_as_current_span("calling-agent") as nested_span:
                    nested_span.set_attribute("operation.type", "response_preparation")

        # Child span: submit the orchestrate run
        with tracer.start_as_current_span("submit-orchestrate-run") as child_span:
            child_span.set_attribute("operation.type", "http_request")
            child_span.set_attribute("http.method", "POST")
            child_span.set_attribute("http.route", "/orchestrate/runs")

    # Flush all buffered spans before the process exits
    cast(TracerProvider, provider).force_flush()

if __name__ == "__main__":
    main()