API walkthrough

import os
from uuid import uuid4
 
import requests
 
API_URL = "https://api.braintrustdata.com/v1"
headers = {"Authorization": "Bearer " + os.environ["BRAINTRUST_API_KEY"]}
 
if __name__ == "__main__":
    # Create a project, if it does not already exist
    project = requests.post(f"{API_URL}/project", headers=headers, json={"name": "rest_test"}).json()
    print(project)
 
    # Create an experiment. This should always be new
    experiment = requests.post(
        f"{API_URL}/experiment", headers=headers, json={"name": "rest_test", "project_id": project["id"]}
    ).json()
    print(experiment)
 
    # Log some stuff
    for i in range(10):
        resp = requests.post(
            f"{API_URL}/experiment/{experiment['id']}/insert",
            headers=headers,
            json={"events": [{"id": uuid4().hex, "input": 1, "output": 2, "scores": {"accuracy": 0.5}}]},
        )
        if not resp.ok:
            raise Exception(f"Error: {resp.status_code} {resp.text}: {resp.content}")

// If you're self-hosting Braintrust, then use your stack's API URL, e.g.
//  https://0o71finhla.execute-api.us-east-1.amazonaws.com/api
export const BRAINTRUST_API_URL = "https://api.braintrustdata.com";
export const API_KEY = process.env.BRAINTRUST_API_KEY;
 
export async function* paginateDataset(args: {
  project: string;
  dataset: string;
  version?: string;
  // Number of rows to fetch per request. You can adjust this to be a lower number
  // if your rows are very large (e.g. several MB each).
  perRequestLimit?: number;
}) {
  const { project, dataset, version, perRequestLimit } = args;
  const headers = {
    Accept: "application/json",
    "Accept-Encoding": "gzip",
    Authorization: `Bearer ${API_KEY}`,
  };
  const fullURL = `${BRAINTRUST_API_URL}/v1/dataset?project_name=${encodeURIComponent(
    project,
  )}&dataset_name=${encodeURIComponent(dataset)}`;
  const ds = await fetch(fullURL, {
    method: "GET",
    headers,
  });
  if (!ds.ok) {
    throw new Error(
      `Error fetching dataset metadata: ${ds.status}: ${await ds.text()}`,
    );
  }
  const dsJSON = await ds.json();
  const dsMetadata = dsJSON.objects[0];
  if (!dsMetadata?.id) {
    throw new Error(`Dataset not found: ${project}/${dataset}`);
  }
 
  let cursor: string | null = null;
  while (true) {
    const body: string = JSON.stringify({
      query: {
        from: {
          op: "function",
          name: { op: "ident", name: ["dataset"] },
          args: [{ op: "literal", value: dsMetadata.id }],
        },
        select: [{ op: "star" }],
        limit: perRequestLimit,
        cursor,
      },
      fmt: "jsonl",
      version,
    });
    const response = await fetch(`${BRAINTRUST_API_URL}/btql`, {
      method: "POST",
      headers,
      body,
    });
    if (!response.ok) {
      throw new Error(
        `Error fetching rows for ${dataset}: ${
          response.status
        }: ${await response.text()}`,
      );
    }
 
    cursor =
      response.headers.get("x-bt-cursor") ??
      response.headers.get("x-amz-meta-bt-cursor");
 
    // Parse jsonl line-by-line
    const allRows = await response.text();
    const rows = allRows.split("\n");
    let rowCount = 0;
    for (const row of rows) {
      if (!row.trim()) {
        continue;
      }
      yield JSON.parse(row);
      rowCount++;
    }
 
    if (rowCount === 0) {
      break;
    }
  }
}
 
async function main() {
  for await (const row of paginateDataset({
    project: "Your project name", // Replace with your project name
    dataset: "Your dataset name", // Replace with your dataset name
    perRequestLimit: 100,
  })) {
    console.log(row);
  }
}
 
main();