Skip to content

add support for braintrust ai proxy #546

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 9 commits into
base: main
Choose a base branch
from
Draft

add support for braintrust ai proxy #546

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
85a33d9
working
kamath Mar 1, 2025
8020598
support braintrust
kamath Mar 1, 2025
a866680
example
kamath Mar 1, 2025
a67eae7
add braintrust models to evals
kamath Mar 2, 2025
f8c1394
changeset
kamath Mar 2, 2025
69e02b4
evals baby
kamath Mar 2, 2025
b795865
pls
kamath Mar 3, 2025
9cc46ee
only use braintrust api key
kamath Mar 3, 2025
5b83692
working locally
kamath Mar 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/shy-poets-run.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
"@browserbasehq/stagehand": minor
---

add support for braintrust ai proxy
15 changes: 2 additions & 13 deletions .github/workflows/ci.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ on:
- unlabeled

env:
EVAL_MODELS: "gpt-4o,gpt-4o-mini,claude-3-5-sonnet-latest"
EVAL_MODELS: "braintrust-gpt-4o,braintrust-gpt-4.5-preview,braintrust-gpt-4o-mini,braintrust-claude-3-5-sonnet-latest,braintrust-claude-3-7-sonnet-latest,braintrust-gemini-2.0-flash,braintrust-llama-3.3-70b-versatile,braintrust-llama-3.1-8b-instant,braintrust-deepseek-r1-distill-llama-70b"
EVAL_CATEGORIES: "observe,act,combination,extract,text_extract"

concurrency:
Expand Down Expand Up @@ -151,8 +151,7 @@ jobs:
github.event_name == 'push' ||
(github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name == github.repository)
env:
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
BRAINTRUST_API_KEY: ${{ secrets.BRAINTRUST_API_KEY }}
BROWSERBASE_API_KEY: ${{ secrets.BROWSERBASE_API_KEY }}
BROWSERBASE_PROJECT_ID: ${{ secrets.BROWSERBASE_PROJECT_ID }}
HEADLESS: true
Expand Down Expand Up @@ -185,8 +184,6 @@ jobs:
runs-on: ubuntu-latest
timeout-minutes: 40
env:
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
BRAINTRUST_API_KEY: ${{ secrets.BRAINTRUST_API_KEY }}
BROWSERBASE_API_KEY: ${{ secrets.BROWSERBASE_API_KEY }}
BROWSERBASE_PROJECT_ID: ${{ secrets.BROWSERBASE_PROJECT_ID }}
Expand Down Expand Up @@ -234,8 +231,6 @@ jobs:
runs-on: ubuntu-latest
timeout-minutes: 25
env:
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
BRAINTRUST_API_KEY: ${{ secrets.BRAINTRUST_API_KEY }}
BROWSERBASE_API_KEY: ${{ secrets.BROWSERBASE_API_KEY }}
BROWSERBASE_PROJECT_ID: ${{ secrets.BROWSERBASE_PROJECT_ID }}
Expand Down Expand Up @@ -302,8 +297,6 @@ jobs:
runs-on: ubuntu-latest
timeout-minutes: 50
env:
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
BRAINTRUST_API_KEY: ${{ secrets.BRAINTRUST_API_KEY }}
BROWSERBASE_API_KEY: ${{ secrets.BROWSERBASE_API_KEY }}
BROWSERBASE_PROJECT_ID: ${{ secrets.BROWSERBASE_PROJECT_ID }}
Expand Down Expand Up @@ -387,8 +380,6 @@ jobs:
runs-on: ubuntu-latest
timeout-minutes: 120
env:
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
BRAINTRUST_API_KEY: ${{ secrets.BRAINTRUST_API_KEY }}
BROWSERBASE_API_KEY: ${{ secrets.BROWSERBASE_API_KEY }}
BROWSERBASE_PROJECT_ID: ${{ secrets.BROWSERBASE_PROJECT_ID }}
Expand Down Expand Up @@ -456,8 +447,6 @@ jobs:
runs-on: ubuntu-latest
timeout-minutes: 60
env:
OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
BRAINTRUST_API_KEY: ${{ secrets.BRAINTRUST_API_KEY }}
BROWSERBASE_API_KEY: ${{ secrets.BROWSERBASE_API_KEY }}
BROWSERBASE_PROJECT_ID: ${{ secrets.BROWSERBASE_PROJECT_ID }}
Expand Down
34 changes: 17 additions & 17 deletions evals/deterministic/tests/local/create.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,9 +14,9 @@ test.describe("Local browser launch options", () => {
debugDom: true,
domSettleTimeoutMs: 30_000,
enableCaching: true,
modelName: "gpt-4o",
modelName: "braintrust-gpt-4o",
modelClientOptions: {
apiKey: process.env.OPENAI_API_KEY,
apiKey: process.env.BRAINTRUST_API_KEY,
},
});
await stagehand.init();
Expand All @@ -38,9 +38,9 @@ test.describe("Local browser launch options", () => {
debugDom: true,
domSettleTimeoutMs: 30_000,
enableCaching: true,
modelName: "gpt-4o",
modelName: "braintrust-gpt-4o",
modelClientOptions: {
apiKey: process.env.OPENAI_API_KEY,
apiKey: process.env.BRAINTRUST_API_KEY,
},
localBrowserLaunchOptions: {
userDataDir: customUserDataDir,
Expand All @@ -66,9 +66,9 @@ test.describe("Local browser launch options", () => {
debugDom: true,
domSettleTimeoutMs: 30_000,
enableCaching: true,
modelName: "gpt-4o",
modelName: "braintrust-gpt-4o",
modelClientOptions: {
apiKey: process.env.OPENAI_API_KEY,
apiKey: process.env.BRAINTRUST_API_KEY,
},
localBrowserLaunchOptions: {
viewport: customViewport,
Expand Down Expand Up @@ -105,9 +105,9 @@ test.describe("Local browser launch options", () => {
debugDom: true,
domSettleTimeoutMs: 30_000,
enableCaching: true,
modelName: "gpt-4o",
modelName: "braintrust-gpt-4o",
modelClientOptions: {
apiKey: process.env.OPENAI_API_KEY,
apiKey: process.env.BRAINTRUST_API_KEY,
},
localBrowserLaunchOptions: {
cookies: testCookies,
Expand Down Expand Up @@ -139,9 +139,9 @@ test.describe("Local browser launch options", () => {
debugDom: true,
domSettleTimeoutMs: 30_000,
enableCaching: true,
modelName: "gpt-4o",
modelName: "braintrust-gpt-4o",
modelClientOptions: {
apiKey: process.env.OPENAI_API_KEY,
apiKey: process.env.BRAINTRUST_API_KEY,
},
localBrowserLaunchOptions: {
geolocation: customGeolocation,
Expand All @@ -151,7 +151,7 @@ test.describe("Local browser launch options", () => {
await stagehand.init();

const page = await stagehand.context.newPage();
await page.goto("https://example.com");
await page.goto("https://docs.stagehand.dev");

const location = await page.evaluate(() => {
return new Promise((resolve) => {
Expand Down Expand Up @@ -180,9 +180,9 @@ test.describe("Local browser launch options", () => {
debugDom: true,
domSettleTimeoutMs: 30_000,
enableCaching: true,
modelName: "gpt-4o",
modelName: "braintrust-gpt-4o",
modelClientOptions: {
apiKey: process.env.OPENAI_API_KEY,
apiKey: process.env.BRAINTRUST_API_KEY,
},
localBrowserLaunchOptions: {
locale: "ja-JP",
Expand All @@ -192,7 +192,7 @@ test.describe("Local browser launch options", () => {
await stagehand.init();

const page = await stagehand.context.newPage();
await page.goto("https://example.com");
await page.goto("https://docs.stagehand.dev");

const { locale, timezone } = await page.evaluate(() => ({
locale: navigator.language,
Expand All @@ -216,9 +216,9 @@ test.describe("Local browser launch options", () => {
debugDom: true,
domSettleTimeoutMs: 30_000,
enableCaching: true,
modelName: "gpt-4o",
modelName: "braintrust-gpt-4o",
modelClientOptions: {
apiKey: process.env.OPENAI_API_KEY,
apiKey: process.env.BRAINTRUST_API_KEY,
},
localBrowserLaunchOptions: {
recordVideo: {
Expand All @@ -230,7 +230,7 @@ test.describe("Local browser launch options", () => {
await stagehand.init();

const page = await stagehand.context.newPage();
await page.goto("https://example.com");
await page.goto("https://docs.stagehand.dev");
await stagehand.close();

const videos = fs.readdirSync(videoDir);
Expand Down
9 changes: 3 additions & 6 deletions evals/initStagehand.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,9 @@ const StagehandConfig = {
headless: false,
enableCaching,
domSettleTimeoutMs: 30_000,
modelName: "gpt-4o", // default model, can be overridden by initStagehand arguments
modelName: "braintrust-gpt-4o", // default model, can be overridden by initStagehand arguments
modelClientOptions: {
apiKey: process.env.OPENAI_API_KEY,
apiKey: process.env.BRAINTRUST_API_KEY,
},
logger: (logLine: LogLine) =>
console.log(`[stagehand::${logLine.category}] ${logLine.message}`),
Expand Down Expand Up @@ -61,10 +61,7 @@ export const initStagehand = async ({
logger: EvalLogger;
configOverrides?: Partial<ConstructorParams>;
}) => {
let chosenApiKey: string | undefined = process.env.OPENAI_API_KEY;
if (modelName.startsWith("claude")) {
chosenApiKey = process.env.ANTHROPIC_API_KEY;
}
const chosenApiKey = process.env.BRAINTRUST_API_KEY;

const config = {
...StagehandConfig,
Expand Down
12 changes: 11 additions & 1 deletion evals/taskConfig.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,17 @@ if (filterByEvalName && !tasksByName[filterByEvalName]) {
*/
const DEFAULT_EVAL_MODELS = process.env.EVAL_MODELS
? process.env.EVAL_MODELS.split(",")
: ["gpt-4o", "claude-3-5-sonnet-latest"];
: [
"braintrust-gpt-4o",
"braintrust-gpt-4.5-preview",
"braintrust-gpt-4o-mini",
"braintrust-claude-3-5-sonnet-latest",
"braintrust-claude-3-7-sonnet-latest",
"braintrust-gemini-2.0-flash",
"braintrust-llama-3.3-70b-versatile",
"braintrust-llama-3.1-8b-instant",
"braintrust-deepseek-r1-distill-llama-70b",
];

/**
* getModelList:
Expand Down
2 changes: 1 addition & 1 deletion evals/tasks/allrecipes.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ export const allrecipes: EvalFunction = async ({
action: 'Type "chocolate chip cookies" in the search bar',
});
await stagehand.page.act({
action: "hit enter",
action: "click search",
});

const recipeDetails = await stagehand.page.extract({
Expand Down
17 changes: 16 additions & 1 deletion examples/example.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,15 +5,30 @@
* npx create-browser-app@latest my-browser-app
*/

import { Stagehand } from "@/dist";
import { AvailableModel, Stagehand } from "@/dist";
import StagehandConfig from "@/stagehand.config";
import { z } from "zod";

async function example() {
const stagehand = new Stagehand({
...StagehandConfig,
modelName: "braintrust-claude-3-7-sonnet-latest" as AvailableModel,
modelClientOptions: {
apiKey: process.env.BRAINTRUST_API_KEY,
},
env: "LOCAL",
});
await stagehand.init();
await stagehand.page.goto("https://docs.stagehand.dev");
const result = await stagehand.page.extract({
instruction: "get the heading",
schema: z.object({
heading: z.string().describe("The heading of the page"),
}),
useTextExtract: true,
});
console.log(result);
await stagehand.close();
}

(async () => {
Expand Down
Loading