Demo of different trace hierarchies for spans

ai-sdk-tool-span-hierarchy.ts

import { InMemorySpanExporter, NodeTracerProvider, ReadableSpan, SimpleSpanProcessor } from "@opentelemetry/sdk-trace-node";
import { gateway, streamText, stepCountIs, tool } from "ai";
import { anthropic } from "@ai-sdk/anthropic";
import { z } from "zod";

const spanExporter = new InMemorySpanExporter();
const provider = new NodeTracerProvider({
  spanProcessors: [new SimpleSpanProcessor(spanExporter)],
});
provider.register();

const getWeather = tool({
  description: "Get the weather in a given location",
  inputSchema: z.object({
    location: z.string().describe("The city and state, e.g. San Francisco, CA")
  }),
  execute: async ({ location }) => {
    await new Promise(resolve => setTimeout(resolve, 300));
    return { location, weather: "Sunny as always!" };
  },
});

const callLlm = async (useGateway: boolean = false) => {
  const result = streamText({
    model: useGateway ? gateway("anthropic/claude-haiku-4-5") : anthropic("claude-haiku-4-5"),
    messages: [
      {
        role: "user",
        content: "What is the weather in SF?",
      },
    ],
    experimental_telemetry: {
      isEnabled: true,
      tracer: provider.getTracer('defaultTracer'),
    },
    tools: {
      getWeather,
    },
    stopWhen: stepCountIs(5),
  });
  for await (const chunk of result.textStream) {
    console.log(chunk);
  }
};

const visualizeSpanStructure = (spans: ReadableSpan[]) => {
  // Build a map of spanId to span for quick lookup
  const spanMap = new Map<string, ReadableSpan>();
  spans.forEach(span => {
    const spanId = span.spanContext().spanId;
    spanMap.set(spanId, span);
  });

  // Find root spans (those without a parent)
  const rootSpans = spans.filter(span => !span.parentSpanContext?.spanId);

  // Build children map
  const childrenMap = new Map<string, ReadableSpan[]>();
  spans.forEach(span => {
    const parentId = span.parentSpanContext?.spanId;
    if (parentId) {
      if (!childrenMap.has(parentId)) {
        childrenMap.set(parentId, []);
      }
      childrenMap.get(parentId)!.push(span);
    }
  });

  // Recursive function to print tree
  const printTree = (span: ReadableSpan, prefix: string = '', isLast: boolean = true) => {
    const connector = isLast ? '└── ' : '├── ';
    console.log(prefix + connector + span.name);

    const spanId = span.spanContext().spanId;
    const children = childrenMap.get(spanId) || [];

    children.forEach((child, index) => {
      const isLastChild = index === children.length - 1;
      const newPrefix = prefix + (isLast ? '    ' : '│   ');
      printTree(child, newPrefix, isLastChild);
    });
  };

  // Print each root span and its children
  rootSpans.forEach((rootSpan, index) => {
    if (index === 0) {
      console.log(rootSpan.name);
    } else {
      console.log(rootSpan.name);
    }

    const spanId = rootSpan.spanContext().spanId;
    const children = childrenMap.get(spanId) || [];
    children.forEach((child, childIndex) => {
      const isLastChild = childIndex === children.length - 1;
      printTree(child, '', isLastChild);
    });
  });
}

const main = async () => {
  await callLlm(false);
  console.log("Span structure without using gateway:");
  const spans = spanExporter.getFinishedSpans();
  console.log(visualizeSpanStructure(spans));

  spanExporter.reset();

  await callLlm(true);
  console.log("Span structure using gateway:");
  const spans2 = spanExporter.getFinishedSpans();
  console.log(visualizeSpanStructure(spans2));
};

main().catch((e) => {
  console.error(e);
  process.exit(1);
});