A collector on your machine

"""An OpenTelemetry collector and a fake OpenAI API, both on this machine.

AgentOps sends traces to otlp.agentops.ai. Pointing it here instead lets you read
every span it would have sent. The fake chat endpoint answers like OpenAI's API
with a keyword rule, so LLM calls can be traced without a key. start() also
blocks connections to other machines for the rest of the program.
"""
import json
import socket
import threading
from http.server import BaseHTTPRequestHandler, HTTPServer

from opentelemetry.proto.collector.trace.v1.trace_service_pb2 import ExportTraceServiceRequest

SPANS = []


def _value(v):
    kind = v.WhichOneof("value")
    if kind == "array_value":
        return [_value(item) for item in v.array_value.values]
    return getattr(v, kind) if kind else None


class Handler(BaseHTTPRequestHandler):
    def do_POST(self):
        body = self.rfile.read(int(self.headers.get("Content-Length", 0)))
        if self.path.endswith("/v1/traces"):
            request = ExportTraceServiceRequest()
            request.ParseFromString(body)
            for resource in request.resource_spans:
                for scope in resource.scope_spans:
                    for span in scope.spans:
                        SPANS.append({
                            "name": span.name, "id": span.span_id.hex(), "parent": span.parent_span_id.hex(),
                            "attributes": {a.key: _value(a.value) for a in span.attributes},
                            "error": span.status.message if span.status.code == 2 else None,
                        })
            return self._reply(200, b"{}")
        if self.path.endswith("/chat/completions"):
            asked = json.loads(body)["messages"][-1]["content"]
            text = "Our billing team will help." if "charged" in asked else "Our support team will help."
            answer = {"id": "chatcmpl-1", "object": "chat.completion", "created": 0, "model": "gpt-4.1-mini",
                      "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": text}}],
                      "usage": {"prompt_tokens": len(asked.split()), "completion_tokens": len(text.split()), "total_tokens": len(asked.split()) + len(text.split())}}
            return self._reply(200, json.dumps(answer).encode())
        return self._reply(200, b"{}")

    def _reply(self, code, data):
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        self.wfile.write(data)

    def log_message(self, *args):
        pass


def _local_only():
    """Refuse network connections to anything but this machine, so nothing can leave by mistake."""
    connect = socket.socket.connect

    def guarded(sock, address):
        if isinstance(address, tuple) and address[0] not in ("127.0.0.1", "::1"):
            raise ConnectionRefusedError(f"blocked a connection to {address[0]}: this lab stays on your machine")
        return connect(sock, address)

    socket.socket.connect = guarded


def start():
    """Start the collector in the background and return its address."""
    _local_only()
    SPANS.clear()
    server = HTTPServer(("127.0.0.1", 0), Handler)
    threading.Thread(target=server.serve_forever, daemon=True).start()
    return f"http://127.0.0.1:{server.server_address[1]}"


def flush():
    """Send every finished span now, instead of waiting for AgentOps' next batch."""
    from opentelemetry import trace

    trace.get_tracer_provider().force_flush()


def tree():
    """Print the spans received so far as a tree."""
    ids = {span["id"] for span in SPANS}

    def show(span, depth):
        print("  " * depth + span["name"])
        for child in SPANS:
            if child["parent"] == span["id"]:
                show(child, depth + 1)

    for root in [span for span in SPANS if span["parent"] not in ids]:
        show(root, 0)

import agentops

import local_collector

url = local_collector.start()
agentops.init(endpoint=url, exporter_endpoint=f"{url}/v1/traces", auto_start_session=False, log_level="ERROR")


@agentops.trace(name="support-ticket")
def handle(ticket):
    return f"Received: {ticket}"

print(handle("Where is my order A-1001?"))
local_collector.flush()
for span in local_collector.SPANS:
    print(span["name"], span["attributes"]["agentops.span.kind"])

python first_trace.py