index.ts

import { initLogger, wrapAISDKModel, traced } from "braintrust";
import { createGoogleGenerativeAI } from "@ai-sdk/google";
import { generateText } from "ai";

export type GoogleGenerativeAI =
  | "gemini-2.0-pro-exp-02-05"
  | "gemini-2.0-flash";

export const initBraintrust = (braintrustApiKey: string) => {
  const logger = initLogger({
    apiKey: braintrustApiKey,
    projectName: "Parrot",
    asyncFlush: true,
  });
  return logger;
};

export const createGoogleModel = async (
  modelName: GoogleGenerativeAI,
  apiKey: string,
  braintrustApiKey: string
) => {
  await initBraintrust(braintrustApiKey);

  const google = createGoogleGenerativeAI({
    apiKey,
  });

  return wrapAISDKModel(google(modelName));
};

export const run = async () => {
  const model = await createGoogleModel(
    "gemini-2.0-pro-exp-02-05",
    process.env.GOOGLE_AI_STUDIO_GEMINI_API_KEY!,
    process.env.BRAINTRUST_API_KEY!
  );

  const result = await generateText({
    model,
    prompt: "Hello, world!",
  });

  // you can also wrap arbitrary functions
  const dbResult = await traced(async (span) => {
    const result = await someDatabaseCall();
    span.log({
      output: result,
      input: "some input",
      metadata: {
        query: "select * from 123",
      },
    });
    return result;
  });

  console.log(result);
  console.log(dbResult);
};