Tracing Vercel AI SDK applications

npx create-next-app@latest --example https://github.com/braintrustdata/braintrust-cookbook/tree/main/examples/VercelAISDKTracing/complete-weather-app vercel-ai-sdk-tracing && cd vercel-ai-sdk-tracing

complete-weather-app/
  |_ app/(preview)/api/chat/route.ts
  |_ components/tools.ts
  |_ .env.local

BRAINTRUST_API_KEY="<Your Braintrust API Key>"
BRAINTRUST_PROJECT_NAME="<Your Braintrust Project Name>"
OPENAI_API_KEY="<Your OpenAI API Key>"
OTEL_EXPORTER_OTLP_ENDPOINT=https://api.braintrust.dev/otel
OTEL_EXPORTER_OTLP_HEADERS="Authorization=Bearer <Your Braintrust API Key>, x-bt-parent=project_id:<Your Braintrust Project Name>"

npm install
npm run dev

import { openai } from "@ai-sdk/openai";
import { streamText, ToolInvocation } from "ai";
import { getWeather, getFahrenheit } from "@/components/tools";
// Uncomment below to use Braintrust's tracing features
import { initLogger, wrapAISDKModel, traced, currentSpan } from "braintrust";
 
// Initialize Braintrust as the logging backend. Uncomment below
const logger = initLogger({
  apiKey: process.env.BRAINTRUST_API_KEY,
  projectName: process.env.BRAINTRUST_PROJECT_NAME,
});

// Any time this model is called, the input and output will be logged to Braintrust.
const model = wrapAISDKModel(openai("gpt-4o"));

export async function POST(request: Request) {
  // traced starts a trace span when the POST endpoint is used
  // Unlike wrapTraced, traced does not natively log inputs and outputs. Uncomment below
  return traced(
    async (span) => {
      const { messages }: { messages: Message[] } = await request.json();
 
      const stream = await streamText({
        // Our wrapped OpenAI model
        model: model,
        system: `\
        - you are an AI assistant who gives the weather. If the user gives you a location, give them the current weather in that location in Fahrenheit.
      `,
        messages: messages,
        // Important: maxSteps prevents infinite tool call loops but will stop your LLM's logic prematurely if set too low
        maxSteps: 5,
        // Register the exported tools to the LLM from @/components/tools
        tools: {
          getWeather: getWeather,
          getFahrenheit: getFahrenheit,
        },
        // Enable experimental telemetry
        experimental_telemetry: {
          isEnabled: true,
        },
        // When streamText is finished, log the input and output of the stream for the "root" span. Uncomment below
        onFinish: (result) => {
          currentSpan().log({
            input: messages,
            output: result.text,
          });
        },
      });
 
      return stream.toDataStreamResponse();
    },
    // Show this span as a function and name the span POST /api/chat. Uncomment below
    { type: "function", name: "POST /api/chat" },
  );
}

import { tool } from "ai";
import { z } from "zod";
// Uncomment below to use Braintrust's tracing features
import { wrapTraced, currentSpan } from "braintrust";
 
interface LocationInput {
  latitude: number;
  longitude: number;
}
 
// Create a simple function to note whether or not a Fahrenheit temperature is freezing
// Wrap the function with the wrapTraced function to note inputs and outputs. Uncomment wrapTraced below
const checkFreezing = wrapTraced(
  async function checkFreezing({ fahrenheit }: { fahrenheit: number }) {
    return fahrenheit < 32;
  },
  // Uncomment below
  { type: "function" },
);
 
// Create a function that takes a temperature in Celsius and returns the temperature in Fahrenheit
// Wrap the function with the wrapTraced function to note inputs and outputs. Uncomment wrapTraced
const convertToFahrenheit = wrapTraced(
  async function convertToFahrenheit({ celsius }: { celsius: number }) {
    const fahrenheit = (celsius * 9) / 5 + 32;
    const isFreezing = checkFreezing({ fahrenheit });
    return fahrenheit;
  },
  // Uncomment below
  { type: "tool" },
);
 
// Construct a tool using the tool() function in the AI package to place in the LLM call
export const getFahrenheit = tool({
  description: "Convert Celsius to Fahrenheit",
  parameters: z.object({ celsius: z.number() }),
  execute: convertToFahrenheit,
});
 
// Create a function that fetches a temperature in Celsius from open-meteo
// Wrap the function with the wrapTraced function to note inputs and outputs. Note that the function should be logged as a tool in the trace. Uncomment wrapTraced below
const weatherFunction = wrapTraced(
  async function weatherFunction({ latitude, longitude }: LocationInput) {
    const response = await fetch(
      `https://api.open-meteo.com/v1/forecast?latitude=${latitude}&longitude=${longitude}&current=temperature_2m&hourly=temperature_2m&daily=sunrise,sunset&timezone=auto`,
    );
    const weatherData = await response.json();
    // Uncomment below to add metadata to the span
    currentSpan().log({
      metadata: { foo: "bar" },
    });
    return weatherData;
  },
  // Uncomment below
  { type: "tool", name: "weatherFunction" },
);
 
// Construct a tool using the tool() function in the AI package to place in the LLM call
export const getWeather = tool({
  description: "Get the current weather at a location",
  parameters: z.object({
    latitude: z.number(),
    longitude: z.number(),
  }),
  execute: weatherFunction,
});