Tracing Vercel AI SDK Stream Response with Braintrust

[meleksomai]

We're running into an issue with Braintrust logging and streaming responses. We're using streamText and createDataStreamResponse from Vercel AI SDK with Braintrust's wrapTraced and traced approach, but only the first stream is being captured in the logs.

Here's a simplified example of what we're trying to do:

const generateStream = wrapTraced(
  async (messages) => {
    return createDataStreamResponse({
      execute: (dataStream) => {
        const result = streamText({
          model: provider.languageModel("generic"),
          messages: messages,
          experimental_telemetry: {
            isEnabled: true,
            // tracer: logger,  // We've tried enabling this but no luck
          },
          onFinish: async ({ response }) => {
            // Only the first stream is captured in Braintrust logs
            // Subsequent streams aren't being tracked
          }
        });
        result.mergeIntoDataStream(dataStream);
      }
    });
  },
  {
    name: "generate-chat-stream",
    type: "function"
  }
);

// Main POST handler
export async function POST(request: Request) {
  return traced(
    async (span) => {
      const { messages, model } = await initalizeRequest(request);
      return generateStream(messages);
    },
    {
      name: "POST /api/chat",
      type: "function",
      setCurrent: true,
    }
  );
}

The issue seems to be that while we're wrapping the main function with wrapTraced, the individual streams within the streamText function aren't being properly traced. We've tried:

• Enabling experimental_telemetry with the logger
• Using traced at different levels of the function
• Wrapping the stream handling in additional wrapTraced calls

Has anyone else encountered this issue or have suggestions on how to properly trace streaming responses with Braintrust?

Thanks in advance for any help!

[lgrammel]

Have you reported this w/ braintrust? Most likely the issue is on their end.

[meleksomai]

I did reach out on their Discord but did not receive an answer yet.

[cacoos]

Hey! That didn't work for me, so I ended up implementing our own spans. Something like this:

import { streamWithContinuation } from "@/app/api/websites/utils/stream-with-continuation";
import { anthropic } from "@ai-sdk/anthropic";
import { type CoreMessage, createDataStreamResponse } from "ai";
import { initLogger, wrapAISDKModel } from "braintrust";

const logger = initLogger({
  projectName: env.NEXT_PUBLIC_ENVIRONMENT_NAME,
  apiKey: env.BRAINTRUST_API_KEY,
});

export const maxDuration = 250;

const model = wrapAISDKModel(anthropic("claude-3-7-sonnet-latest"));

export async function POST(req: Request) {
  try {
    const requestData = (await req.json()) as {
      ...
    };

    return await logger.traced(
      async (span) => {
        const input = {
          ...
        };

        span.log({
          input,
        });

        return createDataStreamResponse({
          onError: (error) => {
            span.log({
              error: {
                errorMessage:
                  error instanceof Error ? error.message : String(error),
                errorStack: error instanceof Error ? error.stack : undefined,
                errorName: error instanceof Error ? error.name : undefined,
              },
            });

            span.end();

            return "An error occurred.";
          },
          execute: async (dataStream) => {
            await streamWithContinuation({
              model,
              initialMessages: [...],
              dataStream,
              maxTokens: 8192,
              maxContinuations: 3,
              abortSignal: req.signal,
            });
          },
        });
      },
      {
        name: "...",
      },
    );
  } catch (error) {
    console.error("Error in POST /api/your-route:", error);
    throw error;
  }
}

streamWithContinuation is a custom function that allows you to continue generating even if you reach maxTokens. But is basically a while loop wrapper over streamText. Inside streamText you can use onChunk, onFinish, etc. to call span.log with output, metadata, whatever you want.

Hope that helps!