Send OpenTelemetry Traces to Weave

import { trace, context } from "@opentelemetry/api";
import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node";
import {
  BatchSpanProcessor,
  ConsoleSpanExporter,
} from "@opentelemetry/sdk-trace-base";
import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-proto";
import { Resource } from "@opentelemetry/resources";

// Configuration
const ENTITY = "<your-team-name>";
const PROJECT = "<your-project-name>";
const WANDB_API_KEY = process.env.WANDB_API_KEY;

if (!WANDB_API_KEY) {
  console.error("Error: WANDB_API_KEY environment variable is not set");
  console.error("Run: export WANDB_API_KEY=your_api_key_here");
  process.exit(1);
}

// OTel Setup
const OTEL_EXPORTER_OTLP_ENDPOINT = "https://trace.wandb.ai/otel/v1/traces";

const exporter = new OTLPTraceExporter({
  url: OTEL_EXPORTER_OTLP_ENDPOINT,
  headers: { "wandb-api-key": WANDB_API_KEY },
});

// Initialize tracer provider with span processors
const provider = new NodeTracerProvider({
  resource: new Resource({
    "wandb.entity": ENTITY,
    "wandb.project": PROJECT,
  }),
  spanProcessors: [
    new BatchSpanProcessor(exporter),
    new BatchSpanProcessor(new ConsoleSpanExporter()),
  ],
});

// Register the tracer provider
provider.register();

// Create a tracer from the global tracer provider
const tracer = trace.getTracer("threads-examples");

function example_1_basic_thread_and_turn() {
  console.log("\n=== Example 1: Basic Thread and Turn ===");

  // Create a thread context
  const threadId = "thread_example_1";

  // This span represents a turn (direct child of thread)
  tracer.startActiveSpan("process_user_message", (turnSpan) => {
    // Set thread attributes
    turnSpan.setAttribute("wandb.thread_id", threadId);
    turnSpan.setAttribute("wandb.is_turn", true);

    // Add some example attributes
    turnSpan.setAttribute("input.value", "Hello, help me with setup");

    let response: string;
    
    // Simulate some work with nested spans
    tracer.startActiveSpan("generate_response", (nestedSpan) => {
      // This is a nested call within the turn, so is_turn should be false or unset
      nestedSpan.setAttribute("wandb.thread_id", threadId);
      // wandb.is_turn is not set or set to false for nested calls

      response = "I'll help you get started with the setup process.";
      nestedSpan.setAttribute("output.value", response);
      nestedSpan.end();
    });
    
    turnSpan.setAttribute("output.value", response!);
    console.log(`Turn completed in thread: ${threadId}`);
    turnSpan.end();
  });
}

function main() {
  example_1_basic_thread_and_turn();
}

main();

function example_2_multiple_turns() {
  console.log("\n=== Example 2: Multiple Turns in Thread ===");

  const threadId = "thread_conversation_123";

  // Turn 1
  tracer.startActiveSpan("process_message_turn1", (turn1Span) => {
    turn1Span.setAttribute("wandb.thread_id", threadId);
    turn1Span.setAttribute("wandb.is_turn", true);
    turn1Span.setAttribute(
      "input.value",
      "What programming languages do you recommend?"
    );

    // Nested operations
    tracer.startActiveSpan("analyze_query", (analyzeSpan) => {
      analyzeSpan.setAttribute("wandb.thread_id", threadId);
      // No is_turn attribute or set to false for nested spans
      analyzeSpan.end();
    });

    const response1 =
      "I recommend Python for beginners and JavaScript for web development.";
    turn1Span.setAttribute("output.value", response1);
    console.log(`Turn 1 completed in thread: ${threadId}`);
    turn1Span.end();
  });

  // Turn 2
  tracer.startActiveSpan("process_message_turn2", (turn2Span) => {
    turn2Span.setAttribute("wandb.thread_id", threadId);
    turn2Span.setAttribute("wandb.is_turn", true);
    turn2Span.setAttribute("input.value", "Can you explain Python vs JavaScript?");

    // Nested operations
    tracer.startActiveSpan("comparison_analysis", (compareSpan) => {
      compareSpan.setAttribute("wandb.thread_id", threadId);
      compareSpan.setAttribute("wandb.is_turn", false); // Explicitly false for nested
      compareSpan.end();
    });

    const response2 =
      "Python excels at data science while JavaScript dominates web development.";
    turn2Span.setAttribute("output.value", response2);
    console.log(`Turn 2 completed in thread: ${threadId}`);
    turn2Span.end();
  });
}

function main() {
  example_2_multiple_turns();
}

main();

function example_3_complex_nested_structure() {
  console.log("\n=== Example 3: Complex Nested Structure ===");

  const threadId = "thread_complex_456";

  // Turn with multiple levels of nesting
  tracer.startActiveSpan("handle_complex_request", (turnSpan) => {
    turnSpan.setAttribute("wandb.thread_id", threadId);
    turnSpan.setAttribute("wandb.is_turn", true);
    turnSpan.setAttribute(
      "input.value",
      "Analyze this code and suggest improvements"
    );

    // Level 1 nested operation
    tracer.startActiveSpan("code_analysis", (analysisSpan) => {
      analysisSpan.setAttribute("wandb.thread_id", threadId);
      // No is_turn for nested operations

      // Level 2 nested operation
      tracer.startActiveSpan("syntax_check", (syntaxSpan) => {
        syntaxSpan.setAttribute("wandb.thread_id", threadId);
        syntaxSpan.setAttribute("result", "No syntax errors found");
        syntaxSpan.end();
      });

      // Another Level 2 nested operation
      tracer.startActiveSpan("performance_check", (perfSpan) => {
        perfSpan.setAttribute("wandb.thread_id", threadId);
        perfSpan.setAttribute("result", "Found 2 optimization opportunities");
        perfSpan.end();
      });

      analysisSpan.end();
    });

    // Another Level 1 nested operation
    tracer.startActiveSpan("generate_suggestions", (suggestSpan) => {
      suggestSpan.setAttribute("wandb.thread_id", threadId);
      const suggestions = ["Use list comprehension", "Consider caching results"];
      suggestSpan.setAttribute("suggestions", JSON.stringify(suggestions));
      suggestSpan.end();
    });

    turnSpan.setAttribute(
      "output.value",
      "Analysis complete with 2 improvement suggestions"
    );
    console.log(`Complex turn completed in thread: ${threadId}`);
    turnSpan.end();
  });
}

function main() {
  example_3_complex_nested_structure();
}

main();

function example_4_non_turn_operations() {
  console.log("\n=== Example 4: Non-Turn Thread Operations ===");

  const threadId = "thread_background_789";

  // Background operation that's part of thread but not a turn
  tracer.startActiveSpan("background_indexing", (bgSpan) => {
    bgSpan.setAttribute("wandb.thread_id", threadId);
    // wandb.is_turn is unset or false - this is not a turn
    bgSpan.setAttribute("wandb.is_turn", false);
    bgSpan.setAttribute("operation", "Indexing conversation history");
    console.log(`Background operation in thread: ${threadId}`);
    bgSpan.end();
  });

  // Actual turn in the same thread
  tracer.startActiveSpan("user_query", (turnSpan) => {
    turnSpan.setAttribute("wandb.thread_id", threadId);
    turnSpan.setAttribute("wandb.is_turn", true);
    turnSpan.setAttribute("input.value", "Search my previous conversations");
    turnSpan.setAttribute("output.value", "Found 5 relevant conversations");
    console.log(`Turn completed in thread: ${threadId}`);
    turnSpan.end();
  });
}

function main() {
  example_4_non_turn_operations();
}

main();