Add models configuration object to init()

js/init-models.test.ts

@@ -0,0 +1,103 @@
+import { expect, test, describe, beforeEach } from "vitest";
+import { init, getDefaultModel, getDefaultEmbeddingModel } from "./oai";
+import { OpenAI } from "openai";
+
+describe("init with defaultModel parameter", () => {
+  beforeEach(() => {
+    // Reset to defaults
+    init();
+  });
+
+  test("string form sets completion model (backward compatible)", () => {
+    init({
+      defaultModel: "gpt-4-turbo",
+    });
+
+    expect(getDefaultModel()).toBe("gpt-4-turbo");
+    expect(getDefaultEmbeddingModel()).toBe("text-embedding-ada-002"); // Default
+  });
+
+  test("object form can set completion model only", () => {
+    init({
+      defaultModel: {
+        completion: "gpt-4-turbo",
+      },
+    });
+
+    expect(getDefaultModel()).toBe("gpt-4-turbo");
+  });
+
+  test("object form can set embedding model only", () => {
+    init({
+      defaultModel: {
+        embedding: "text-embedding-3-large",
+      },
+    });
+
+    expect(getDefaultEmbeddingModel()).toBe("text-embedding-3-large");
+    // Completion model should remain at default since we didn't update it
+    expect(getDefaultModel()).toBe("gpt-4o");
+  });
+
+  test("object form can set both models", () => {
+    init({
+      defaultModel: {
+        completion: "claude-3-5-sonnet-20241022",
+        embedding: "text-embedding-3-large",
+      },
+    });
+
+    expect(getDefaultModel()).toBe("claude-3-5-sonnet-20241022");
+    expect(getDefaultEmbeddingModel()).toBe("text-embedding-3-large");
+  });
+
+  test("partial updates preserve unspecified models", () => {
+    // First set completion model
+    init({
+      defaultModel: {
+        completion: "gpt-4-turbo",
+      },
+    });
+
+    expect(getDefaultModel()).toBe("gpt-4-turbo");
+    expect(getDefaultEmbeddingModel()).toBe("text-embedding-ada-002");
+
+    // Then set only embedding model - completion should remain unchanged
+    init({
+      defaultModel: {
+        embedding: "text-embedding-3-large",
+      },
+    });
+
+    expect(getDefaultModel()).toBe("gpt-4-turbo"); // Should still be gpt-4-turbo
+    expect(getDefaultEmbeddingModel()).toBe("text-embedding-3-large");
+  });
+
+  test("falls back to defaults when not set", () => {
+    init();
+
+    expect(getDefaultModel()).toBe("gpt-4o");
+    expect(getDefaultEmbeddingModel()).toBe("text-embedding-ada-002");
+  });
+
+  test("string form resets embedding model to default", () => {
+    // First set both models
+    init({
+      defaultModel: {
+        completion: "gpt-4-turbo",
+        embedding: "text-embedding-3-large",
+      },
+    });
+
+    expect(getDefaultModel()).toBe("gpt-4-turbo");
+    expect(getDefaultEmbeddingModel()).toBe("text-embedding-3-large");
+
+    // Then use string form - should reset embedding to default
+    init({
+      defaultModel: "claude-3-5-sonnet-20241022",
+    });
+
+    expect(getDefaultModel()).toBe("claude-3-5-sonnet-20241022");
+    expect(getDefaultEmbeddingModel()).toBe("text-embedding-ada-002"); // Reset to default
+  });
+});

js/oai.ts

@@ -150,6 +150,7 @@ declare global {
   var __inherited_braintrust_wrap_openai: ((openai: any) => any) | undefined;
   var __client: OpenAI | undefined;
   var __defaultModel: string | undefined;
+  var __defaultEmbeddingModel: string | undefined;
 }
 
 export interface InitOptions {
@@ -160,17 +161,57 @@ export interface InitOptions {
    */
   client?: OpenAI;
   /**
-   * The default model to use for evaluations when not specified per-call.
-   * Defaults to "gpt-4o" if not set.
+   * The default model(s) to use for evaluations when not specified per-call.
+   *
+   * Can be either:
+   * - A string (for backward compatibility): Sets the default completion model only.
+   *   Defaults to "gpt-4o" if not set.
+   * - An object with `completion` and/or `embedding` properties: Allows setting
+   *   default models for different evaluation types. Only the specified models
+   *   are updated; others remain unchanged.
    *
    * When using non-OpenAI providers via the Braintrust proxy, set this to
    * the appropriate model string (e.g., "claude-3-5-sonnet-20241022").
+   *
+   * @example
+   * // String form (backward compatible)
+   * init({ defaultModel: "gpt-4-turbo" })
+   *
+   * @example
+   * // Object form: set both models
+   * init({
+   *   defaultModel: {
+   *     completion: "claude-3-5-sonnet-20241022",
+   *     embedding: "text-embedding-3-large"
+   *   }
+   * })
+   *
+   * @example
+   * // Object form: set only embedding model
+   * init({
+   *   defaultModel: {
+   *     embedding: "text-embedding-3-large"
+   *   }
+   * })
    */
-  defaultModel?: string;
+  defaultModel?:
+    | string
+    | {
+        /**
+         * Default model for LLM-as-a-judge evaluations (completion).
+         * Defaults to "gpt-4o" if not set.
+         */
+        completion?: string;
+        /**
+         * Default model for embedding-based evaluations.
+         * Defaults to "text-embedding-ada-002" if not set.
+         */
+        embedding?: string;
+      };
 }
 
 /**
- * Initialize autoevals with a custom client and/or default model.
+ * Initialize autoevals with a custom client and/or default models.
  *
  * @example
  * // Using with OpenAI (default)
@@ -189,21 +230,51 @@ export interface InitOptions {
  *     apiKey: process.env.BRAINTRUST_API_KEY,
  *     baseURL: "https://api.braintrust.dev/v1/proxy",
  *   }),
- *   defaultModel: "claude-3-5-sonnet-20241022",
+ *   defaultModel: {
+ *     completion: "claude-3-5-sonnet-20241022",
+ *     embedding: "text-embedding-3-large",
+ *   },
  * });
+ *
+ * @example
+ * // String form (backward compatible)
+ * init({ defaultModel: "gpt-4-turbo" });
  */
 export const init = ({ client, defaultModel }: InitOptions = {}) => {
   globalThis.__client = client;
-  globalThis.__defaultModel = defaultModel;
+  if (typeof defaultModel === "string") {
+    // String form: sets completion model only, resets embedding to default
+    globalThis.__defaultModel = defaultModel;
+    globalThis.__defaultEmbeddingModel = undefined;
+  } else if (defaultModel) {
+    // Object form: only update models that are explicitly provided
+    if ("completion" in defaultModel) {
+      globalThis.__defaultModel = defaultModel.completion;
+    }
+    if ("embedding" in defaultModel) {
+      globalThis.__defaultEmbeddingModel = defaultModel.embedding;
+    }
+  } else {
+    // No defaultModel: reset both to defaults
+    globalThis.__defaultModel = undefined;
+    globalThis.__defaultEmbeddingModel = undefined;
+  }
 };
 
 /**
- * Get the configured default model, or "gpt-4o" if not set.
+ * Get the configured default completion model, or "gpt-4o" if not set.
  */
 export const getDefaultModel = (): string => {
   return globalThis.__defaultModel ?? "gpt-4o";
 };
 
+/**
+ * Get the configured default embedding model, or "text-embedding-ada-002" if not set.
+ */
+export const getDefaultEmbeddingModel = (): string => {
+  return globalThis.__defaultEmbeddingModel ?? "text-embedding-ada-002";
+};
+
 export async function cachedChatCompletion(
   params: CachedLLMParams,
   options: { cache?: ChatCache } & OpenAIAuth,

js/ragas.ts

@@ -103,7 +103,7 @@ import mustache from "mustache";
 
 import { Scorer, ScorerArgs } from "./score";
 import { LLMArgs } from "./llm";
-import { getDefaultModel } from "./oai";
+import { getDefaultModel, getDefaultEmbeddingModel } from "./oai";
 import { buildOpenAIClient, extractOpenAIArgs } from "./oai";
 import OpenAI from "openai";
 import { ListContains } from "./list";
@@ -767,7 +767,7 @@ export const AnswerRelevancy: ScorerWithPartial<
         ...extractOpenAIArgs(args),
         output: question,
         expected: input,
-        model: args.embeddingModel,
+        model: args.embeddingModel ?? getDefaultEmbeddingModel(),
       });
       return { question, score };
     }),

js/string.ts

@@ -1,6 +1,6 @@
 import { Scorer, ScorerArgs } from "./score";
 import levenshtein from "js-levenshtein";
-import { OpenAIAuth, buildOpenAIClient } from "./oai";
+import { OpenAIAuth, buildOpenAIClient, getDefaultEmbeddingModel } from "./oai";
 import cossim from "compute-cosine-similarity";
 import { makePartial, ScorerWithPartial } from "./partial";
 
@@ -69,7 +69,7 @@ export const EmbeddingSimilarity: ScorerWithPartial<
     [output, expected].map((input) =>
       openai.embeddings.create({
         input,
-        model: args.model ?? "text-embedding-ada-002",
+        model: args.model ?? getDefaultEmbeddingModel(),
       }),
     ),
   );