braintrust 0.0.199 → 0.0.201

package/dev/authorize.ts

@@ -0,0 +1,172 @@
+import { Request, Response, NextFunction } from "express";
+import { IncomingHttpHeaders } from "http";
+import createError from "http-errors";
+
+export interface RequestContext {
+  appOrigin: string;
+  token: string | undefined;
+}
+declare module "express" {
+  interface Request {
+    ctx?: RequestContext;
+  }
+}
+
+export function authorizeRequest(
+  req: Request,
+  res: Response,
+  next: NextFunction,
+) {
+  try {
+    const ctx: RequestContext = {
+      appOrigin: extractAllowedOrigin(req.headers[ORIGIN_HEADER]),
+      token: undefined,
+    };
+
+    // Extract token and data from request
+    if (
+      req.headers.authorization ||
+      req.headers[BRAINTRUST_AUTH_TOKEN_HEADER]
+    ) {
+      const tokenText = parseBraintrustAuthHeader(req.headers);
+      if (!tokenText) {
+        return next(createError(400, "Invalid authorization token format"));
+      }
+      ctx.token = tokenText.toLowerCase() === "null" ? undefined : tokenText;
+    }
+
+    req.ctx = ctx;
+
+    next(); // Proceed to next middleware/controller
+  } catch (e) {
+    next(e);
+  }
+}
+
+export function checkAuthorized(
+  req: Request,
+  res: Response,
+  next: NextFunction,
+) {
+  if (!req.ctx?.token) {
+    return next(createError(401, "Unauthorized"));
+  }
+  next();
+}
+
+function parseBraintrustAuthHeader(
+  headers: IncomingHttpHeaders | Record<string, string>,
+): string | undefined {
+  const tokenString = parseHeader(headers, BRAINTRUST_AUTH_TOKEN_HEADER);
+  return tokenString ?? parseAuthHeader(headers) ?? undefined;
+}
+
+function parseHeader(
+  headers: IncomingHttpHeaders | Record<string, string>,
+  headerName: string,
+): string | undefined {
+  const token = headers[headerName];
+  let tokenString;
+  if (typeof token === "string") {
+    tokenString = token;
+  } else if (Array.isArray(token) && token.length > 0) {
+    tokenString = token[0];
+  }
+
+  return tokenString;
+}
+
+export type StaticOrigin =
+  | boolean
+  | string
+  | RegExp
+  | Array<boolean | string | RegExp>;
+
+export function checkOrigin(
+  requestOrigin: string | undefined,
+  callback: (err: Error | null, origin?: StaticOrigin) => void,
+) {
+  if (!requestOrigin) {
+    return callback(null, true);
+  }
+
+  // the origins can be glob patterns
+  for (const origin of WHITELISTED_ORIGINS || []) {
+    if (
+      (origin instanceof RegExp && origin.test(requestOrigin)) ||
+      origin === requestOrigin
+    ) {
+      return callback(null, requestOrigin);
+    }
+  }
+
+  return callback(null, false);
+}
+
+const BRAINTRUST_AUTH_TOKEN_HEADER = "x-bt-auth-token";
+const ORIGIN_HEADER = "origin";
+
+export function extractAllowedOrigin(originHeader: string | undefined): string {
+  let allowedOrigin: string = MAIN_ORIGIN;
+  checkOrigin(originHeader, (err, origin) => {
+    if (!err && originHeader && origin) {
+      allowedOrigin = originHeader;
+    }
+  });
+  return allowedOrigin;
+}
+
+const MAIN_ORIGIN = "https://www.braintrust.dev";
+const WHITELISTED_ORIGINS = [
+  MAIN_ORIGIN,
+  "https://www.braintrustdata.com",
+  new RegExp("https://.*.preview.braintrust.dev"),
+]
+  .concat(
+    process.env.WHITELISTED_ORIGIN ? [process.env.WHITELISTED_ORIGIN] : [],
+  )
+  .concat(
+    process.env.BRAINTRUST_APP_URL ? [process.env.BRAINTRUST_APP_URL] : [],
+  );
+
+function parseAuthHeader(
+  headers: Record<string, string | string[] | undefined>,
+) {
+  const authHeader = headers["authorization"];
+  let authValue = null;
+  if (Array.isArray(authHeader)) {
+    authValue = authHeader[authHeader.length - 1];
+  } else {
+    authValue = authHeader;
+  }
+
+  if (!authValue) {
+    return null;
+  }
+
+  const parts = authValue.split(" ");
+  if (parts.length !== 2) {
+    return null;
+  }
+  return parts[1];
+}
+
+export const baseAllowedHeaders = [
+  "Content-Type",
+  "X-Amz-Date",
+  "Authorization",
+  "X-Api-Key",
+  "X-Amz-Security-Token",
+  "x-bt-auth-token",
+  "x-bt-parent",
+  // These are eval-specific
+  "x-bt-org-name",
+  "x-bt-stream-fmt",
+  "x-bt-use-cache",
+  "x-stainless-os",
+  "x-stainless-lang",
+  "x-stainless-package-version",
+  "x-stainless-runtime",
+  "x-stainless-runtime-version",
+  "x-stainless-arch",
+];

package/dev/errorHandler.ts

@@ -0,0 +1,37 @@
+import { z } from "zod";
+import { Request, Response, ErrorRequestHandler, NextFunction } from "express";
+import { HttpError } from "http-errors";
+
+export const errorHandler: ErrorRequestHandler = (
+  err: Error | HttpError,
+  req: Request,
+  res: Response,
+  next: NextFunction,
+) => {
+  if ("status" in err) {
+    res.status(err.status).json({
+      error: {
+        message: err.message,
+        status: err.status,
+      },
+    });
+    return;
+  }
+  if (err instanceof z.ZodError) {
+    res.status(400).json({
+      error: {
+        message: "Invalid request",
+        errors: err.errors,
+      },
+    });
+    return;
+  }
+
+  console.error("Internal server error", err);
+  res.status(500).json({
+    error: {
+      message: "Internal server error",
+      status: 500,
+    },
+  });
+};

package/dev/index.ts

@@ -0,0 +1 @@
+export { runDevServer } from "./server";

package/dev/server.ts

@@ -0,0 +1,413 @@
+import express, { NextFunction, Request, Response } from "express";
+import cors from "cors";
+import {
+  callEvaluatorData,
+  Eval,
+  EvalData,
+  EvalHooks,
+  EvalScorer,
+  EvaluatorDef,
+  OneOrMoreScores,
+  scorerName,
+} from "../src/framework";
+import { errorHandler } from "./errorHandler";
+import {
+  authorizeRequest,
+  baseAllowedHeaders,
+  checkAuthorized,
+  checkOrigin,
+} from "./authorize";
+import {
+  FunctionId,
+  InvokeFunctionRequest,
+  RunEvalRequest,
+  SSEProgressEventData,
+} from "@braintrust/core/typespecs";
+import {
+  BaseMetadata,
+  BraintrustState,
+  EvalCase,
+  getSpanParentObject,
+  initDataset,
+  LoginOptions,
+  loginToState,
+} from "../src/logger";
+import { LRUCache } from "../src/prompt-cache/lru-cache";
+import {
+  BT_CURSOR_HEADER,
+  BT_FOUND_EXISTING_HEADER,
+  parseParent,
+} from "@braintrust/core";
+import { serializeSSEEvent } from "./stream";
+import {
+  evalBodySchema,
+  EvaluatorDefinitions,
+  EvaluatorManifest,
+  evalParametersSerializedSchema,
+} from "./types";
+import { EvalParameters, validateParameters } from "../src/eval-parameters";
+import { z } from "zod";
+import { promptDefinitionToPromptData } from "../src/framework2";
+import zodToJsonSchema from "zod-to-json-schema";
+export interface DevServerOpts {
+  host: string;
+  port: number;
+}
+
+export function runDevServer(
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  evaluators: EvaluatorDef<any, any, any, any, any>[],
+  opts: DevServerOpts,
+) {
+  // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+  const allEvaluators: EvaluatorManifest = Object.fromEntries(
+    evaluators.map((evaluator) => [evaluator.evalName, evaluator]),
+  ) as EvaluatorManifest;
+
+  globalThis._lazy_load = false;
+
+  const app = express();
+
+  app.use(express.json({ limit: "1gb" }));
+  console.log("Starting server");
+  app.use((req, res, next) => {
+    if (req.headers["access-control-request-private-network"]) {
+      res.setHeader("Access-Control-Allow-Private-Network", "true");
+    }
+    next();
+  });
+
+  app.use(
+    cors({
+      origin: checkOrigin,
+      methods: ["GET", "PATCH", "POST", "PUT", "DELETE", "OPTIONS"],
+      allowedHeaders: baseAllowedHeaders,
+      credentials: true,
+      exposedHeaders: [
+        BT_CURSOR_HEADER,
+        BT_FOUND_EXISTING_HEADER,
+        "x-bt-span-id",
+        "x-bt-span-export",
+      ],
+      maxAge: 86400,
+    }),
+  );
+
+  app.use(authorizeRequest);
+
+  app.get("/", (req, res) => {
+    res.send("Hello, world!");
+  });
+
+  // List endpoint - returns all available evaluators and their metadata
+  app.get("/list", (req, res) => {
+    const evalDefs: EvaluatorDefinitions = Object.fromEntries(
+      Object.entries(allEvaluators).map(([name, evaluator]) => [
+        name,
+        {
+          parameters: evaluator.parameters
+            ? makeEvalParametersSchema(evaluator.parameters)
+            : undefined,
+          scores: evaluator.scores.map((score, idx) => ({
+            name: scorerName(score, idx),
+          })),
+        },
+      ]),
+    );
+    res.json(evalDefs);
+  });
+
+  app.post(
+    "/eval",
+    checkAuthorized,
+    asyncHandler(async (req, res) => {
+      const { name, parameters, parent, data, scores, stream } =
+        evalBodySchema.parse(req.body);
+
+      const state = await cachedLogin({ apiKey: req.ctx?.token });
+
+      const evaluator = allEvaluators[name];
+      if (!evaluator) {
+        res.status(404).json({ error: `Evaluator '${name}' not found` });
+        return;
+      }
+
+      if (
+        evaluator.parameters &&
+        Object.keys(evaluator.parameters).length > 0
+      ) {
+        try {
+          if (!evaluator.parameters) {
+            res.status(400).json({
+              error: `Evaluator '${name}' does not accept parameters`,
+            });
+            return;
+          }
+
+          // This gets done again in the framework, but we do it here too to give a
+          // better error message.
+          validateParameters(parameters ?? {}, evaluator.parameters);
+        } catch (e) {
+          console.error("Error validating parameters", e);
+          if (e instanceof z.ZodError || e instanceof Error) {
+            res.status(400).json({
+              error: e.message,
+            });
+            return;
+          }
+          throw e;
+        }
+      }
+
+      const resolvedData = await getDataset(state, data);
+      const evalData = callEvaluatorData(resolvedData);
+      console.log("Starting eval", evaluator.evalName);
+
+      // Set up SSE headers
+      if (stream) {
+        res.setHeader("Content-Type", "text/event-stream");
+        res.setHeader("Cache-Control", "no-cache");
+        res.setHeader("Connection", "keep-alive");
+      } else {
+        res.setHeader("Content-Type", "application/json");
+      }
+
+      const task = async (
+        input: unknown,
+        hooks: EvalHooks<unknown, BaseMetadata, EvalParameters>,
+      ) => {
+        const result = await evaluator.task(input, hooks);
+
+        hooks.reportProgress({
+          format: "code",
+          output_type: "completion",
+          event: "json_delta",
+          data: JSON.stringify(result),
+        });
+        return result;
+      };
+
+      try {
+        const summary = await Eval(
+          "worker-thread",
+          {
+            ...evaluator,
+            data: evalData.data,
+            scores: evaluator.scores.concat(
+              scores?.map((score) =>
+                makeScorer(state, score.name, score.function_id),
+              ) ?? [],
+            ),
+            task,
+            state,
+          },
+          {
+            // Avoid printing the bar to the console.
+            progress: {
+              start: (name, total) => {},
+              stop: () => {
+                console.log("Finished running experiment");
+              },
+              increment: (name) => {},
+            },
+            stream: (data: SSEProgressEventData) => {
+              if (stream) {
+                res.write(
+                  serializeSSEEvent({
+                    event: "progress",
+                    data: JSON.stringify(data),
+                  }),
+                );
+              }
+            },
+            onStart: (metadata) => {
+              if (stream) {
+                res.write(
+                  serializeSSEEvent({
+                    event: "start",
+                    data: JSON.stringify(metadata),
+                  }),
+                );
+              }
+            },
+            parent: parseParent(parent),
+            parameters: parameters ?? {},
+          },
+        );
+
+        if (stream) {
+          res.write(
+            serializeSSEEvent({
+              event: "summary",
+              data: JSON.stringify(summary.summary),
+            }),
+          );
+          res.write(
+            serializeSSEEvent({
+              event: "done",
+              data: "",
+            }),
+          );
+        } else {
+          res.json(summary.summary);
+        }
+      } catch (e) {
+        console.error("Error running eval", e);
+        if (stream) {
+          res.write(
+            serializeSSEEvent({
+              event: "error",
+              data: JSON.stringify(e),
+            }),
+          );
+        } else {
+          res.status(500).json({ error: e });
+        }
+      } finally {
+        res.end();
+      }
+    }),
+  );
+
+  app.use(errorHandler);
+
+  // Start the server
+  app.listen(opts.port, opts.host, () => {
+    console.log(`Dev server running at http://${opts.host}:${opts.port}`);
+  });
+}
+const asyncHandler =
+  (fn: (req: Request, res: Response, next: NextFunction) => Promise<void>) =>
+  (req: Request, res: Response, next: NextFunction) => {
+    Promise.resolve(fn(req, res, next)).catch(next);
+  };
+
+const loginCache = new LRUCache<string, BraintrustState>({
+  max: 32, // TODO: Make this configurable
+});
+
+async function cachedLogin(options: LoginOptions): Promise<BraintrustState> {
+  const key = JSON.stringify(options);
+  const cached = loginCache.get(key);
+  if (cached) {
+    return cached;
+  }
+
+  const state = await loginToState(options);
+  loginCache.set(key, state);
+  return state;
+}
+
+async function getDataset(
+  state: BraintrustState,
+  data: RunEvalRequest["data"],
+): Promise<EvalData<unknown, unknown, BaseMetadata>> {
+  if ("project_name" in data) {
+    return initDataset({
+      state,
+      project: data.project_name,
+      dataset: data.dataset_name,
+      _internal_btql: data._internal_btql ?? undefined,
+    });
+  } else if ("dataset_id" in data) {
+    const datasetInfo = await getDatasetById({
+      state,
+      datasetId: data.dataset_id,
+    });
+    return initDataset({
+      state,
+      projectId: datasetInfo.projectId,
+      dataset: datasetInfo.dataset,
+      _internal_btql: data._internal_btql ?? undefined,
+    });
+  } else {
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    return data.data as EvalCase<unknown, unknown, BaseMetadata>[];
+  }
+}
+
+const datasetFetchSchema = z.object({
+  project_id: z.string(),
+  name: z.string(),
+});
+async function getDatasetById({
+  state,
+  datasetId,
+}: {
+  state: BraintrustState;
+  datasetId: string;
+}): Promise<{ projectId: string; dataset: string }> {
+  const dataset = await state.appConn().post_json("api/dataset/get", {
+    id: datasetId,
+  });
+  const parsed = z.array(datasetFetchSchema).parse(dataset);
+  if (parsed.length === 0) {
+    throw new Error(`Dataset '${datasetId}' not found`);
+  }
+  return { projectId: parsed[0].project_id, dataset: parsed[0].name };
+}
+
+function makeScorer(
+  state: BraintrustState,
+  name: string,
+  score: FunctionId,
+): EvalScorer<unknown, unknown, unknown, BaseMetadata> {
+  const ret = async (input: EvalCase<unknown, unknown, BaseMetadata>) => {
+    const request: InvokeFunctionRequest = {
+      ...score,
+      input,
+      parent: await getSpanParentObject().export(),
+      stream: false,
+      mode: "auto",
+      strict: true,
+    };
+    const result = await state.proxyConn().post(`function/invoke`, request, {
+      headers: {
+        Accept: "application/json",
+      },
+    });
+    const data = await result.json();
+    // NOTE: Ideally we can parse this value with a zod schema.
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    return data as OneOrMoreScores;
+  };
+
+  Object.defineProperties(ret, {
+    name: { value: `Remote eval scorer (${name})` },
+  });
+
+  return ret;
+}
+
+function makeEvalParametersSchema(
+  parameters: EvalParameters,
+): z.infer<typeof evalParametersSerializedSchema> {
+  return Object.fromEntries(
+    Object.entries(parameters).map(([name, value]) => {
+      if ("type" in value && value.type === "prompt") {
+        return [
+          name,
+          {
+            type: "prompt",
+            default: value.default
+              ? promptDefinitionToPromptData(value.default)
+              : undefined,
+            description: value.description,
+          },
+        ];
+      } else {
+        return [
+          name,
+          {
+            type: "data",
+            // Since this schema is bundled, it won't pass an instanceof check.
+            // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+            schema: zodToJsonSchema(value as z.ZodSchema),
+            default: value.default,
+            description: value.description,
+          },
+        ];
+      }
+    }),
+  );
+}

package/dev/stream.ts

@@ -0,0 +1,14 @@
+export interface RawSSEEvent {
+  id?: string;
+  event?: string;
+  data: string;
+}
+
+export function serializeSSEEvent(event: RawSSEEvent): string {
+  return (
+    Object.entries(event)
+      .filter(([_key, value]) => value !== undefined)
+      .map(([key, value]) => `${key}: ${value}`)
+      .join("\n") + "\n\n"
+  );
+}

package/dev/types.ts

@@ -0,0 +1,63 @@
+import {
+  functionIdSchema,
+  invokeParent,
+  runEvalSchema,
+  promptDataSchema,
+} from "@braintrust/core/typespecs";
+import { z } from "zod";
+import { EvaluatorDef } from "../src/framework";
+import { BaseMetadata } from "../src/logger";
+
+export const evalBodySchema = z.object({
+  name: z.string(),
+  parameters: z.record(z.string(), z.unknown()).nullish(),
+  data: runEvalSchema.shape.data,
+  scores: z
+    .array(
+      z.object({
+        function_id: functionIdSchema,
+        name: z.string(),
+      }),
+    )
+    .nullish(),
+  parent: invokeParent.optional(),
+  stream: z.boolean().optional(),
+});
+
+export type EvaluatorManifest = Record<
+  string,
+  EvaluatorDef<unknown, unknown, unknown, BaseMetadata>
+>;
+
+export const evalParametersSerializedSchema = z.record(
+  z.string(),
+  z.union([
+    z.object({
+      type: z.literal("prompt"),
+      default: promptDataSchema.optional(),
+      description: z.string().optional(),
+    }),
+    z.object({
+      type: z.literal("data"),
+      schema: z.record(z.unknown()), // JSON Schema
+      default: z.unknown().optional(),
+      description: z.string().optional(),
+    }),
+  ]),
+);
+
+export type EvalParameterSerializedSchema = z.infer<
+  typeof evalParametersSerializedSchema
+>;
+
+export const evaluatorDefinitionSchema = z.object({
+  parameters: evalParametersSerializedSchema.optional(),
+});
+export type EvaluatorDefinition = z.infer<typeof evaluatorDefinitionSchema>;
+
+export const evaluatorDefinitionsSchema = z.record(
+  z.string(),
+  evaluatorDefinitionSchema,
+);
+
+export type EvaluatorDefinitions = z.infer<typeof evaluatorDefinitionsSchema>;