Customize traces

import { initLogger, wrapOpenAI, wrapTraced } from "braintrust";
import OpenAI from "openai";
import { ChatCompletionMessageParam } from "openai/resources";

const logger = initLogger({
  projectName: "My Project",
  apiKey: process.env.BRAINTRUST_API_KEY,
});

const client = wrapOpenAI(new OpenAI({ apiKey: process.env.OPENAI_API_KEY }));

// wrapTraced() automatically logs the input (args) and output (return value)
// of this function to a span. To ensure the span is named `answerQuestion`,
// you should name the inline function definition (inside of wrapTraced).
const answerQuestion = wrapTraced(async function answerQuestion(
  body: string,
): Promise<string> {
  const prompt: ChatCompletionMessageParam[] = [
    { role: "system", content: "You are a helpful assistant." },
    { role: "user", content: body },
  ];

  const result = await client.chat.completions.create({
    model: "gpt-4o",
    messages: prompt,
  });

  const content = result.choices[0].message.content;
  if (!content) {
    throw new Error("The LLM response content is empty or undefined.");
  }

  return content;
});

async function main() {
  const input = "How can I improve my productivity?";
  const result = await answerQuestion(input);
  console.log(result);
}

// In TypeScript, you can use the `name` parameter in either `wrapTraced` or `traced()`.
import { wrapTraced } from "braintrust";

const myFunction = wrapTraced(
  async function (input) {
    return process(input);
  },
  { name: "Custom Span Name", type: "task" },
);

await logger.traced(
  async (span) => {
    span.log({ input: data });
    // Do work
    span.log({ output: result });
  },
  { name: "Custom Span Name", type: "task" },
);

import { initLogger, traced, wrapTraced } from "braintrust";

const logger = initLogger({
  projectName: "My Project",
  apiKey: process.env.BRAINTRUST_API_KEY,
});

interface LLMCompletion {
  completion: string;
  metrics: {
    prompt_tokens: number;
    completion_tokens: number;
  };
}

async function callMyLLM(
  input: string,
  params: { temperature: number },
): Promise<LLMCompletion> {
  // Replace with your custom LLM implementation
  return {
    completion: "Hello, world!",
    metrics: {
      prompt_tokens: input.length,
      completion_tokens: 10,
    },
  };
}

export const invokeCustomLLM = wrapTraced(
  async function invokeCustomLLM(
    llmInput: string,
    params: { temperature: number },
  ) {
    return traced(async (span) => {
      const result = await callMyLLM(llmInput, params);
      const content = result.completion;
      span.log({
        input: [{ role: "user", content: llmInput }],
        output: content,
        metrics: {
          prompt_tokens: result.metrics.prompt_tokens,
          completion_tokens: result.metrics.completion_tokens,
          tokens:
            result.metrics.prompt_tokens + result.metrics.completion_tokens,
        },
        metadata: params,
      });
      return content;
    });
  },
  {
    type: "llm",
    name: "Custom LLM",
  },
);

export async function POST(req: Request) {
  return traced(async (span) => {
    const result = await invokeCustomLLM(await req.text(), {
      temperature: 0.1,
    });
    span.log({ input: req.body, output: result });
    return result;
  });
}

import { Attachment, initLogger } from "braintrust";

const logger = initLogger({ projectName: "Attachment Example" });

logger.log({
  input: {
    question: "What is this?",
    context: new Attachment({
      data: "path/to/input_image.jpg",
      filename: "user_input.jpg",
      contentType: "image/jpeg",
    }),
  },
  output: "Example response.",
});

import { JSONAttachment, initLogger } from "braintrust";

const logger = initLogger({ projectName: "Large Trace Example" });

// Example: Large document collection
const documents = Array.from({ length: 500 }, (_, i) => ({
  id: `doc_${i}`,
  title: `Document ${i}`,
  content: `This is a long document with lots of text content...`.repeat(100),
  metadata: {
    author: `Author ${i % 20}`,
    created_at: new Date(Date.now() - i * 3600000).toISOString(),
    tags: [`tag_${i % 10}`, `category_${i % 5}`],
    embeddings: Array.from({ length: 768 }, () => Math.random()),
  },
}));

logger.log({
  input: {
    query: "Find documents about machine learning",
    search_context: new JSONAttachment(documents, {
      filename: "document_collection.json",
      pretty: true, // Optional: pretty-print the JSON
    }),
    search_params: {
      limit: 10,
      similarity_threshold: 0.8,
    },
  },
  output: {
    results: documents.slice(0, 10).map((d) => ({ id: d.id, title: d.title })),
    total_searched: documents.length,
  },
  metrics: {
    search_duration_ms: 1250,
    documents_processed: documents.length,
  },
});

// Example: Complex nested configuration
const systemConfig = {
  models: Array.from({ length: 50 }, (_, i) => ({
    id: `model_${i}`,
    name: `Model ${i}`,
    parameters: {
      temperature: Math.random(),
      max_tokens: 1000 + i * 100,
      top_p: 0.9,
      frequency_penalty: Math.random() * 0.5,
      presence_penalty: Math.random() * 0.5,
    },
    performance_metrics: {
      latency_p50: Math.random() * 1000,
      latency_p95: Math.random() * 2000,
      latency_p99: Math.random() * 3000,
      success_rate: 0.95 + Math.random() * 0.05,
    },
  })),
  prompts: Array.from({ length: 100 }, (_, i) => ({
    id: `prompt_${i}`,
    template: `System prompt template ${i} with lots of instructions...`.repeat(
      50,
    ),
    version: `v${i}.0.0`,
    test_cases: Array.from({ length: 20 }, (_, j) => ({
      input: `Test input ${j}`,
      expected: `Expected output ${j}`,
    })),
  })),
};

logger.log({
  input: {
    experiment_name: "model_comparison",
    config: new JSONAttachment(systemConfig, {
      filename: "experiment_config.json",
    }),
  },
  output: {
    best_model: "model_42",
    summary: "Completed comparison of 50 models across 100 prompts",
  },
});

import { wrapTraced, currentSpan } from "braintrust";

async function processRequest(input: string) {
  return input.length > 10
    ? { error: "Input too long" }
    : { data: "Hello, world!" };
}

const requestHandler = wrapTraced(async function requestHandler(req: Request) {
  const body = await req.text();
  const result = await processRequest(body);
  if (result.error) {
    currentSpan().log({ error: result.error });
  } else {
    currentSpan().log({ input: req.body, output: result.data });
  }
  return result;
});

import {
  currentSpan,
  initLogger,
  traced,
  wrapOpenAI,
  wrapTraced,
} from "braintrust";
import OpenAI from "openai";

const logger = initLogger();
const client = wrapOpenAI(new OpenAI({ apiKey: process.env.OPENAI_API_KEY }));

export const runLLM = wrapTraced(async function runLLM(input) {
  const model = Math.random() > 0.5 ? "gpt-4o" : "gpt-4o-mini";
  const result = await client.chat.completions.create({
    model,
    messages: [{ role: "user", content: input }],
  });
  const output = result.choices[0].message.content;
  currentSpan().log({
    metadata: {
      randomModel: model,
    },
  });
  return output;
});

export const someLogic = wrapTraced(async function someLogic(input: string) {
  return await runLLM(
    "You are a magical wizard. Answer the following question: " + input,
  );
});

export async function POST(req: Request) {
  return await traced(async () => {
    const body = await req.json();
    const result = await someLogic(body.text);
    currentSpan().log({
      input: body.text,
      output: result,
      metadata: { user_id: body.userId },
    });
    return result;
  });
}

import { setMaskingFunction } from "braintrust";

// Set a global masking function
setMaskingFunction((data: any): any => {
  // Your masking logic here
  const maskedData = data; // Process data as needed
  return maskedData;
});

// To disable masking
setMaskingFunction(null);

import { setMaskingFunction, initLogger } from "braintrust";

const maskFunction = (data: any): any => {
  if (typeof data === "string") {
    // Mask strings containing sensitive patterns
    return data.replace(
      /\b(api[_-]?key|password|token)[\s:=]+\S+/gi,
      "$1: [REDACTED]",
    );
  }

  if (typeof data === "object" && data !== null) {
    if (Array.isArray(data)) {
      return data.map((item) => maskFunction(item));
    }

    const masked: any = {};
    for (const [key, value] of Object.entries(data)) {
      // Mask values for sensitive keys
      if (/^(api[_-]?key|password|secret|token|auth|credential)$/i.test(key)) {
        masked[key] = "[REDACTED]";
      } else {
        masked[key] = maskFunction(value);
      }
    }
    return masked;
  }

  return data;
};

setMaskingFunction(maskFunction);

// Usage example
const logger = initLogger({ projectName: "My Project" });

logger.log({
  input: {
    query: "Process payment",
    api_key: "sk-1234567890",
    user: "[email protected]",
  },
  output: "Payment processed with token: abc123xyz",
  metadata: {
    password: "super-secret",
    request_id: "req-456",
  },
});

// This will log:
// input: { query: "Process payment", api_key: "[REDACTED]", user: "[email protected]" }
// output: "Payment processed with token: [REDACTED]"
// metadata: { password: "[REDACTED]", request_id: "req-456" }

import { setMaskingFunction } from "braintrust";

const maskPII = (data: any): any => {
  if (typeof data === "string") {
    let masked = data;
    // Mask email addresses
    masked = masked.replace(
      /\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b/g,
      "[EMAIL]",
    );
    // Mask phone numbers (US format)
    masked = masked.replace(/\b\d{3}[-.]?\d{3}[-.]?\d{4}\b/g, "[PHONE]");
    // Mask SSN
    masked = masked.replace(/\b\d{3}-\d{2}-\d{4}\b/g, "[SSN]");
    return masked;
  }

  if (typeof data === "object" && data !== null) {
    if (Array.isArray(data)) {
      return data.map((item) => maskPII(item));
    }

    const masked: any = {};
    for (const [key, value] of Object.entries(data)) {
      if (
        ["email", "phone", "ssn", "phone_number"].includes(key.toLowerCase())
      ) {
        masked[key] = `[${key.toUpperCase()}]`;
      } else {
        masked[key] = maskPII(value);
      }
    }
    return masked;
  }

  return data;
};

setMaskingFunction(maskPII);

// Usage example
import { initLogger } from "braintrust";
const logger = initLogger({ projectName: "My Project" });

logger.log({
  input: {
    message: "Contact [email protected] or call 555-123-4567",
    user: {
      name: "John Doe",
      email: "[email protected]",
      phone: "555-123-4567",
      ssn: "123-45-6789",
    },
  },
});

// This will log:
// input: {
//   message: "Contact [EMAIL] or call [PHONE]",
//   user: {
//     name: "John Doe",
//     email: "[EMAIL]",
//     phone: "[PHONE]",
//     ssn: "[SSN]"
//   }
// }

import { setMaskingFunction } from "braintrust";

const customMask = (data: any): any => {
  // Handle different data types
  if (typeof data === "string") {
    // Only mask if the string contains certain keywords
    if (data.toLowerCase().includes("confidential")) {
      return "[CONFIDENTIAL DATA REMOVED]";
    }
    return data;
  }

  if (typeof data === "number") {
    // Example: mask large monetary values
    if (data > 10000) {
      return -1; // Sentinel value for masked numbers
    }
    return data;
  }

  if (typeof data === "object" && data !== null) {
    // Handle special data structures
    if ("credit_card" in data && "cvv" in data) {
      // Mask credit card info but keep last 4 digits
      return {
        ...data,
        credit_card: (data as any).credit_card?.replace(/\d(?=\d{4})/g, "X"),
        cvv: "XXX",
        amount: (data as any).amount, // Keep amount unmasked
      };
    }

    if (Array.isArray(data)) {
      return data.map((item) => customMask(item));
    }

    // Default object handling
    const masked: any = {};
    for (const [key, value] of Object.entries(data)) {
      // Skip masking for specific fields
      if (["timestamp", "request_id", "trace_id"].includes(key)) {
        masked[key] = value;
      } else {
        masked[key] = customMask(value);
      }
    }
    return masked;
  }

  return data;
};

setMaskingFunction(customMask);

// Usage example
import { initLogger } from "braintrust";
const logger = initLogger({ projectName: "My Project" });

logger.log({
  input: {
    transaction: {
      credit_card: "4532-1234-5678-9012",
      cvv: "123",
      amount: 15000,
      timestamp: "2024-01-01T00:00:00Z",
    },
    internal_note: "Confidential: Premium customer",
  },
  metadata: {
    trace_id: "trace-123",
    debug_info: "Processing large transaction",
  },
});

// This will log:
// input: {
//   transaction: {
//     credit_card: "XXXX-XXXX-XXXX-9012",
//     cvv: "XXX",
//     amount: -1,
//     timestamp: "2024-01-01T00:00:00Z"
//   },
//   internal_note: "[CONFIDENTIAL DATA REMOVED]"
// },
// metadata: {
//   trace_id: "trace-123",
//   debug_info: "Processing large transaction"
// }

import { currentSpan, initLogger, wrapTraced } from "braintrust";
import { ChatCompletionMessageParam } from "openai/resources";

const logger = initLogger({ projectName: "my-project" });

async function remoteChatCompletion(args: {
  model: string;
  messages: ChatCompletionMessageParam[];
  extraHeaders?: Record<string, string>;
}) {
  // This is a placeholder for code that would call a remote service
}

const bedTimeStory = wrapTraced(async function bedtimeStory(input: {
  summary: string;
  length: number;
}) {
  return await remoteChatCompletion({
    model: "gpt-3.5-turbo",
    messages: [
      {
        role: "system",
        content:
          "Come up with a bedtime story with the following summary and approximate length (in words)",
      },
      {
        role: "user",
        content: `summary: ${input.summary}\nlength: ${input.length}`,
      },
    ],
    extraHeaders: {
      request_id: await currentSpan().export(),
    },
  });
});

import { traced, wrapOpenAI } from "braintrust";
import OpenAI from "openai";
import { ChatCompletionMessageParam } from "openai/resources";

const client = wrapOpenAI(new OpenAI({ apiKey: process.env.OPENAI_API_KEY }));

async function serverSideChatCompletion(request: {
  model: string;
  messages: ChatCompletionMessageParam[];
  headers?: Record<string, string>;
}) {
  return await traced(
    async (span) => {
      const output = await client.chat.completions.create({
        model: request.model,
        messages: request.messages,
      });
      return output.choices[0].message.content;
    },
    {
      name: "text_generator_server",
      type: "llm",
      // This will be a fresh, root-level trace if headers or request_id are undefined,
      // or will create sub-spans under the parent trace if they are defined.
      parent: request.headers?.request_id,
    },
  );
}

import { initLogger, wrapTraced, currentSpan } from "braintrust";

const logger = initLogger({
  projectName: "my-project", // Replace with your project name
  apiKey: process.env.BRAINTRUST_API_KEY, // Replace with your API key
});

const startRequest = wrapTraced(async function startRequest(request) {
  const handle = startSomething(request.body);
  return {
    result: handle,
    spanId: currentSpan().id,
  };
});

const finishRequest = wrapTraced(async function finishRequest(handle, spanId) {
  const result = await finishSomething(handle);
  logger.updateSpan({
    id: spanId,
    output: result,
  });
  return result;
});

import { initLogger, wrapTraced, currentSpan, updateSpan } from "braintrust";

const logger = initLogger({
  projectName: "my-project", // Replace with your project name
  apiKey: process.env.BRAINTRUST_API_KEY, // Replace with your API key
});

const startRequest = wrapTraced(async function startRequest(request) {
  const handle = startSomething(request.body);
  return {
    result: handle,
    exported: currentSpan().export(),
  };
});

const finishRequest = wrapTraced(
  async function finishRequest(handle, exported) {
    const result = await finishSomething(handle);
    updateSpan({
      exported,
      output: result,
    });
    return result;
  },
);

import {
  Span,
  initLogger,
  startSpan,
  wrapOpenAI,
  wrapTraced,
} from "braintrust";
import { OpenAI } from "openai";

const client = wrapOpenAI(new OpenAI({ apiKey: process.env.OPENAI_API_KEY }));
const logger = initLogger({ projectName: "My long-running project" });

const computeOutput = wrapTraced(async function computeOutput(
  systemPrompt: string,
  userInput: string,
  parentSpan: Span,
) {
  return await client.chat.completions.create({
    model: "gpt-3.5-turbo",
    messages: [
      { role: "system", content: systemPrompt },
      { role: "user", content: userInput },
    ],
  });
});

class MyHandler {
  private liveSpans: Record<string, { span: Span; input: string }>;

  constructor() {
    this.liveSpans = {};
  }

  async onRequestStart(requestId: string, input: string, expected: string) {
    const span = startSpan({ name: requestId, event: { input, expected } });
    this.liveSpans[requestId] = { span, input };
  }

  async onGetOutput(requestId: string, systemPrompt: string) {
    const { span, input } = this.liveSpans[requestId];
    const output = await computeOutput(systemPrompt, input, span);
    span.log({ output });
  }

  async onRequestEnd(requestId: string, metadata: Record<string, string>) {
    const { span } = this.liveSpans[requestId];
    delete this.liveSpans[requestId];
    span.log({ metadata });
    span.end();
  }
}