@vellumai/assistant 0.11.10 → 0.11.11-staging.1

package/src/persistence/embeddings/embedding-backend.test.ts

@@ -3,15 +3,18 @@ import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
 import type { AssistantConfig } from "../../config/types.js";
 import {
   clearEmbeddingBackendCache,
+  customEmbeddingSpaceIdentity,
   embedWithBackend,
   isEmbeddingDimensionAvailable,
   resetLocalEmbeddingFailureState,
+  resolveBackendDimension,
   selectEmbeddingBackend,
 } from "./embedding-backend.js";
 import {
   _resetEmbeddingBillingBreaker,
   recordBillingBlock,
 } from "./embedding-billing-breaker.js";
+import { OpenAIEmbeddingBackend } from "./embedding-openai.js";
 
 const getProviderKeyAsyncMock = mock(
   async (_provider: string): Promise<string | undefined> => undefined,
@@ -361,3 +364,162 @@ describe("managed-proxy Gemini fallback to direct key", () => {
     expect(directUrl).toContain("key=direct-key");
   });
 });
+
+function customConfig(overrides: {
+  baseUrl?: string;
+  customModel?: string;
+  customDimensions?: number;
+}): AssistantConfig {
+  return {
+    memory: {
+      embeddings: {
+        provider: "custom",
+        customModel: overrides.customModel ?? "text-embedding-3-small",
+        baseUrl: overrides.baseUrl,
+        customDimensions: overrides.customDimensions,
+      },
+    },
+  } as unknown as AssistantConfig;
+}
+
+describe("custom OpenAI-compatible embedding backend selection", () => {
+  afterEach(() => {
+    clearEmbeddingBackendCache();
+    getProviderKeyAsyncMock.mockReset();
+    getProviderKeyAsyncMock.mockImplementation(async () => undefined);
+  });
+
+  test("selects a custom backend with a trimmed baseUrl", async () => {
+    getProviderKeyAsyncMock.mockImplementation(async (provider) =>
+      provider === "custom" ? "gw-key" : undefined,
+    );
+    const { backend, reason } = await selectEmbeddingBackend(
+      customConfig({
+        baseUrl: "http://127.0.0.1:4000/v1/",
+        customModel: "embed-mistral",
+        customDimensions: 1024,
+      }),
+    );
+    expect(reason).toBeNull();
+    expect(backend).toBeInstanceOf(OpenAIEmbeddingBackend);
+    expect(backend?.provider).toBe("custom");
+    expect(backend?.model).toBe("embed-mistral");
+    expect((backend as OpenAIEmbeddingBackend).baseURL).toBe(
+      "http://127.0.0.1:4000/v1",
+    );
+    expect((backend as OpenAIEmbeddingBackend).dimensions).toBe(1024);
+  });
+
+  test("selects a custom backend without an API key", async () => {
+    const { backend, reason } = await selectEmbeddingBackend(
+      customConfig({ baseUrl: "http://127.0.0.1:1234/v1" }),
+    );
+    expect(reason).toBeNull();
+    expect(backend?.provider).toBe("custom");
+  });
+
+  test("is not configured when baseUrl is missing", async () => {
+    const { backend, reason } = await selectEmbeddingBackend(customConfig({}));
+    expect(backend).toBeNull();
+    expect(reason).toMatch(/baseUrl/);
+  });
+
+  test("is not configured when baseUrl is not http(s)", async () => {
+    const { backend, reason } = await selectEmbeddingBackend(
+      customConfig({ baseUrl: "not-a-url" }),
+    );
+    expect(backend).toBeNull();
+    expect(reason).toMatch(/baseUrl/);
+  });
+
+  test("is not selected by auto even when baseUrl is set", async () => {
+    const { backend } = await selectEmbeddingBackend({
+      memory: {
+        embeddings: {
+          provider: "auto",
+          localModel: "BAAI/bge-small-en-v1.5",
+          baseUrl: "http://127.0.0.1:4000/v1",
+          customModel: "embed-mistral",
+        },
+      },
+    } as unknown as AssistantConfig);
+    expect(backend?.provider).toBe("local");
+  });
+
+  test("caches custom backends by baseUrl", async () => {
+    const first = await selectEmbeddingBackend(
+      customConfig({ baseUrl: "http://127.0.0.1:4000/v1" }),
+    );
+    const second = await selectEmbeddingBackend(
+      customConfig({ baseUrl: "http://127.0.0.1:4001/v1" }),
+    );
+    const firstAgain = await selectEmbeddingBackend(
+      customConfig({ baseUrl: "http://127.0.0.1:4000/v1" }),
+    );
+    expect(first.backend).not.toBe(second.backend);
+    expect(firstAgain.backend).toBe(first.backend);
+  });
+
+  test("caches custom backends by requested dimensions", async () => {
+    const first = await selectEmbeddingBackend(
+      customConfig({
+        baseUrl: "http://127.0.0.1:4000/v1",
+        customDimensions: 1024,
+      }),
+    );
+    const second = await selectEmbeddingBackend(
+      customConfig({
+        baseUrl: "http://127.0.0.1:4000/v1",
+        customDimensions: 1536,
+      }),
+    );
+    expect(first.backend).not.toBe(second.backend);
+  });
+
+  test("dimension probes for custom backends are keyed by baseUrl", async () => {
+    const { backend: first } = await selectEmbeddingBackend(
+      customConfig({ baseUrl: "http://127.0.0.1:4000/v1" }),
+    );
+    const { backend: second } = await selectEmbeddingBackend(
+      customConfig({ baseUrl: "http://127.0.0.1:4001/v1" }),
+    );
+    if (!first || !second) {
+      throw new Error("expected custom backends");
+    }
+    (first as { embed: typeof first.embed }).embed = mock(async () => [
+      new Array(1536).fill(0),
+    ]) as unknown as typeof first.embed;
+    (second as { embed: typeof second.embed }).embed = mock(async () => [
+      new Array(1024).fill(0),
+    ]) as unknown as typeof second.embed;
+
+    expect(await resolveBackendDimension(first)).toBe(1536);
+    expect(await resolveBackendDimension(second)).toBe(1024);
+  });
+
+  test("customEmbeddingSpaceIdentity includes model, url, and dimensions", () => {
+    expect(
+      customEmbeddingSpaceIdentity(
+        customConfig({
+          baseUrl: "http://127.0.0.1:4000/v1/",
+          customModel: "embed-mistral",
+          customDimensions: 1024,
+        }),
+      ),
+    ).toBe("custom\0embed-mistral\0url=http://127.0.0.1:4000/v1\0dim=1024");
+  });
+
+  test("customEmbeddingSpaceIdentity is null when the provider is not custom", () => {
+    expect(
+      customEmbeddingSpaceIdentity({
+        memory: {
+          embeddings: {
+            provider: "openai",
+            openaiModel: "text-embedding-3-small",
+            baseUrl: "http://127.0.0.1:4000/v1",
+          },
+        },
+      } as unknown as AssistantConfig),
+    ).toBeNull();
+  });
+});

package/src/persistence/embeddings/embedding-backend.ts

@@ -17,7 +17,10 @@ import {
 } from "./embedding-billing-breaker.js";
 import { GeminiEmbeddingBackend } from "./embedding-gemini.js";
 import { OllamaEmbeddingBackend } from "./embedding-ollama.js";
-import { OpenAIEmbeddingBackend } from "./embedding-openai.js";
+import {
+  OpenAIEmbeddingBackend,
+  resolveOpenAICompatibleBaseUrl,
+} from "./embedding-openai.js";
 import { EmbeddingRuntimeManager } from "./embedding-runtime-manager.js";
 import {
   EMBEDDING_DIMENSION_PROBE_TEXT,
@@ -428,6 +431,70 @@ export function geminiCacheExtras(config: AssistantConfig): string[] {
   return extras;
 }
 
+/**
+ * Cache-key fragments for a custom OpenAI-compatible embeddings endpoint.
+ * Base URL and requested output dimensionality both change the vectors a
+ * given model returns, so they belong in backend, in-memory, and durable
+ * cache identity.
+ */
+export function customCacheExtras(config: AssistantConfig): string[] {
+  const extras: string[] = [];
+  const baseUrl = resolveOpenAICompatibleBaseUrl(
+    config.memory.embeddings.baseUrl,
+  );
+  if (baseUrl) {
+    extras.push(`url=${baseUrl}`);
+  }
+  if (config.memory.embeddings.customDimensions != null) {
+    extras.push(`dim=${config.memory.embeddings.customDimensions}`);
+  }
+  return extras;
+}
+
+/**
+ * Provider extras that belong in durable `memory_embeddings` identity.
+ * Gemini task type / dimensions and custom endpoint URL / dimensions both
+ * change the vector for identical text, so cache readers must fold them in.
+ */
+export function durableEmbeddingCacheExtras(
+  config: AssistantConfig,
+  provider: string | null,
+): string[] {
+  if (provider === "gemini") {
+    return geminiCacheExtras(config);
+  }
+  if (provider === "custom") {
+    return customCacheExtras(config);
+  }
+  return [];
+}
+
+/**
+ * Checkpoint key for the custom embedding space last used to populate
+ * Qdrant. Compared at reconcile time so a same-dimension endpoint change
+ * still enqueues a reembed instead of mixing vector spaces.
+ */
+export const CUSTOM_EMBEDDING_SPACE_CHECKPOINT =
+  "memory:custom_embedding_space";
+
+/**
+ * Stable identity of the configured custom embeddings endpoint. Null when
+ * the provider is not `custom`. Includes model, normalized base URL, and
+ * requested dimensions.
+ */
+export function customEmbeddingSpaceIdentity(
+  config: AssistantConfig,
+): string | null {
+  if (config.memory.embeddings.provider !== "custom") {
+    return null;
+  }
+  return [
+    "custom",
+    config.memory.embeddings.customModel,
+    ...customCacheExtras(config),
+  ].join("\0");
+}
+
 /** Build (or reuse) the direct-API Gemini backend for the given key. */
 function getDirectGeminiBackend(
   config: AssistantConfig,
@@ -449,6 +516,29 @@ function getDirectGeminiBackend(
   );
 }
 
+/** Build (or reuse) a custom OpenAI-compatible embeddings backend. */
+function getCustomEmbeddingBackend(
+  config: AssistantConfig,
+  baseUrl: string,
+  apiKey: string | undefined,
+): EmbeddingBackend {
+  return getCachedOrCreate(
+    "custom",
+    config.memory.embeddings.customModel,
+    () =>
+      new OpenAIEmbeddingBackend(
+        apiKey ?? "",
+        config.memory.embeddings.customModel,
+        {
+          provider: "custom",
+          baseURL: baseUrl,
+          dimensions: config.memory.embeddings.customDimensions,
+        },
+      ),
+    customCacheExtras(config),
+  );
+}
+
 /**
  * Build (or reuse) the managed-proxy Gemini backend, or return undefined
  * when the managed proxy prerequisites (platform URL + assistant API key)
@@ -517,6 +607,23 @@ export async function selectEmbeddingBackend(
       reason: null,
     };
   }
+  if (requested === "custom") {
+    const baseUrl = resolveOpenAICompatibleBaseUrl(
+      config.memory.embeddings.baseUrl,
+    );
+    if (!baseUrl) {
+      return {
+        backend: null,
+        reason:
+          'Embedding backend "custom" requires memory.embeddings.baseUrl to be an http(s) URL',
+      };
+    }
+    const customKey = (await getProviderKeyAsync("custom")) ?? undefined;
+    return {
+      backend: getCustomEmbeddingBackend(config, baseUrl, customKey),
+      reason: null,
+    };
+  }
 
   // When managed proxy prerequisites are satisfied, insert managed-proxy Gemini
   // at the front of the auto chain so platform assistants use Vellum-managed
@@ -620,6 +727,9 @@ export async function selectEmbeddingBackend(
           reason: null,
         };
       }
+      case "custom":
+        // Custom is explicit-only and is handled before the auto chain.
+        continue;
     }
   }
 
@@ -666,14 +776,30 @@ export async function getMemoryBackendStatus(config: AssistantConfig): Promise<{
 }
 
 /**
- * Memoized output dimension per "provider:model". A backend's vector dimension
- * is fixed for the life of a (provider, model) pair, so a single probe answers
- * every subsequent {@link isEmbeddingDimensionAvailable} call without another
- * backend round-trip. Cleared alongside the backend cache so a credential
- * change or explicit reset re-probes the (possibly different) backend.
+ * Memoized output dimension per backend identity. A backend's vector
+ * dimension is fixed for the life of that identity (provider, model, and for
+ * custom OpenAI-compatible endpoints the base URL and requested dimensions),
+ * so a single probe answers every subsequent
+ * {@link isEmbeddingDimensionAvailable} call without another backend
+ * round-trip. Cleared alongside the backend cache so a credential change or
+ * explicit reset re-probes the (possibly different) backend.
  */
 const backendDimCache = new Map<string, number>();
 
+function backendDimCacheKey(backend: EmbeddingBackend): string {
+  if (backend instanceof OpenAIEmbeddingBackend) {
+    const extras: string[] = [];
+    if (backend.baseURL) {
+      extras.push(`url=${backend.baseURL}`);
+    }
+    if (backend.dimensions != null) {
+      extras.push(`dim=${backend.dimensions}`);
+    }
+    return cacheKey(backend.provider, backend.model, extras);
+  }
+  return `${backend.provider}:${backend.model}`;
+}
+
 /**
  * Whether the currently-reachable embedding backend can produce vectors of the
  * committed Qdrant collection dimension (`config.memory.qdrant.vectorSize`).
@@ -692,7 +818,7 @@ const backendDimCache = new Map<string, number>();
  *
  * The selected backend's output dimension is not statically known from
  * provider/model alone, so it is measured with a fixed-string probe and
- * memoized per (provider, model) — steady-state calls resolve from
+ * memoized per backend identity — steady-state calls resolve from
  * {@link backendDimCache} without a backend round-trip.
  */
 export async function isEmbeddingDimensionAvailable(
@@ -713,7 +839,7 @@ export async function isEmbeddingDimensionAvailable(
   // assembler, so this preflight can never be more pessimistic than the real
   // embed path. Dense recall stays available if ANY backend in the chain
   // produces the committed dimension. `resolveBackendDimension` memoizes per
-  // provider:model, so each backend is probed at most once.
+  // backend identity, so each backend is probed at most once.
   const expected = config.memory.qdrant.vectorSize;
   for (const candidate of await assembleEmbeddingBackends(config, backend)) {
     if ((await resolveBackendDimension(candidate)) === expected) {
@@ -736,7 +862,7 @@ export async function isEmbeddingDimensionAvailable(
 export async function resolveBackendDimension(
   backend: EmbeddingBackend,
 ): Promise<number | null> {
-  const key = `${backend.provider}:${backend.model}`;
+  const key = backendDimCacheKey(backend);
   const cached = backendDimCache.get(key);
   if (cached != null) {
     return cached;
@@ -849,7 +975,11 @@ export async function embedWithBackend(
 
   // ── Compute provider-specific vector cache extras ───────────────
   const vectorExtras =
-    primaryProvider === "gemini" ? geminiCacheExtras(config) : undefined;
+    primaryProvider === "gemini"
+      ? geminiCacheExtras(config)
+      : primaryProvider === "custom"
+        ? customCacheExtras(config)
+        : undefined;
 
   // ── In-memory cache check (primary provider only) ──────────────
   const cached: (number[] | null)[] = inputs.map((input) => {
@@ -918,7 +1048,11 @@ export async function embedWithBackend(
 
       // Populate cache with freshly embedded vectors
       const backendExtras =
-        backend.provider === "gemini" ? geminiCacheExtras(config) : undefined;
+        backend.provider === "gemini"
+          ? geminiCacheExtras(config)
+          : backend.provider === "custom"
+            ? customCacheExtras(config)
+            : undefined;
       for (let i = 0; i < inputsToEmbed.length; i++) {
         putInVectorCache(
           backend.provider,
@@ -1060,6 +1194,9 @@ async function selectFallbackBackends(
         }
         break;
       }
+      case "custom":
+      case "local":
+        break;
     }
   }
   return backends;

package/src/persistence/embeddings/embedding-openai.test.ts

@@ -0,0 +1,113 @@
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test";
+
+const createMock = mock(async (_body: unknown, _opts?: unknown) => ({
+  data: [{ embedding: [0.1, 0.2, 0.3] }],
+}));
+
+let lastCtorArgs: { apiKey?: string; baseURL?: string } | undefined;
+
+mock.module("openai", () => ({
+  default: class OpenAI {
+    embeddings = { create: createMock };
+    constructor(opts: { apiKey: string; baseURL?: string }) {
+      lastCtorArgs = opts;
+    }
+  },
+}));
+
+import {
+  OpenAIEmbeddingBackend,
+  resolveOpenAICompatibleBaseUrl,
+} from "./embedding-openai.js";
+
+describe("resolveOpenAICompatibleBaseUrl", () => {
+  test("trims and strips a trailing slash", () => {
+    expect(
+      resolveOpenAICompatibleBaseUrl(" https://gateway.example.com/v1/ "),
+    ).toBe("https://gateway.example.com/v1");
+  });
+
+  test("accepts http localhost URLs", () => {
+    expect(resolveOpenAICompatibleBaseUrl("http://127.0.0.1:4000/v1")).toBe(
+      "http://127.0.0.1:4000/v1",
+    );
+  });
+
+  test("rejects empty, non-http, and invalid values", () => {
+    expect(resolveOpenAICompatibleBaseUrl(undefined)).toBeNull();
+    expect(resolveOpenAICompatibleBaseUrl("")).toBeNull();
+    expect(resolveOpenAICompatibleBaseUrl("   ")).toBeNull();
+    expect(resolveOpenAICompatibleBaseUrl("not-a-url")).toBeNull();
+    expect(resolveOpenAICompatibleBaseUrl("ftp://gateway.example.com/v1")).toBe(
+      null,
+    );
+  });
+});
+
+describe("OpenAIEmbeddingBackend", () => {
+  beforeEach(() => {
+    createMock.mockClear();
+    lastCtorArgs = undefined;
+  });
+
+  afterEach(() => {
+    createMock.mockClear();
+  });
+
+  test("constructs the official OpenAI client without a baseURL", () => {
+    const backend = new OpenAIEmbeddingBackend(
+      "sk-test",
+      "text-embedding-3-small",
+    );
+    expect(backend.provider).toBe("openai");
+    expect(lastCtorArgs).toEqual({ apiKey: "sk-test" });
+  });
+
+  test("routes custom embeddings through the supplied baseURL", async () => {
+    const backend = new OpenAIEmbeddingBackend("gw-key", "embed-mistral", {
+      provider: "custom",
+      baseURL: "http://127.0.0.1:4000/v1",
+      dimensions: 1024,
+    });
+    expect(backend.provider).toBe("custom");
+    expect(backend.baseURL).toBe("http://127.0.0.1:4000/v1");
+    expect(lastCtorArgs).toEqual({
+      apiKey: "gw-key",
+      baseURL: "http://127.0.0.1:4000/v1",
+    });
+
+    const vectors = await backend.embed(["hello"]);
+    expect(vectors).toEqual([[0.1, 0.2, 0.3]]);
+    expect(createMock).toHaveBeenCalledTimes(1);
+    const [body] = createMock.mock.calls[0] as [
+      {
+        model: string;
+        input: string[];
+        encoding_format: string;
+        dimensions?: number;
+      },
+    ];
+    expect(body.model).toBe("embed-mistral");
+    expect(body.input).toEqual(["hello"]);
+    expect(body.encoding_format).toBe("float");
+    expect(body.dimensions).toBe(1024);
+  });
+
+  test("uses a placeholder API key when the custom endpoint is keyless", () => {
+    new OpenAIEmbeddingBackend("", "text-embedding-3-small", {
+      provider: "custom",
+      baseURL: "http://127.0.0.1:1234/v1",
+    });
+    expect(lastCtorArgs?.apiKey).toBe("not-needed");
+  });
+
+  test("omits dimensions when they are not configured", async () => {
+    const backend = new OpenAIEmbeddingBackend(
+      "sk-test",
+      "text-embedding-3-small",
+    );
+    await backend.embed(["hello"]);
+    const [body] = createMock.mock.calls[0] as [{ dimensions?: number }];
+    expect(body.dimensions).toBeUndefined();
+  });
+});

package/src/persistence/embeddings/embedding-openai.ts

@@ -3,18 +3,40 @@ import OpenAI from "openai";
 import {
   type EmbeddingBackend,
   type EmbeddingInput,
+  type EmbeddingProviderName,
   type EmbeddingRequestOptions,
   normalizeEmbeddingInput,
 } from "./embedding-types.js";
 
+export type OpenAICompatibleEmbeddingProvider = Extract<
+  EmbeddingProviderName,
+  "openai" | "custom"
+>;
+
+export interface OpenAIEmbeddingOptions {
+  baseURL?: string;
+  dimensions?: number;
+  provider?: OpenAICompatibleEmbeddingProvider;
+}
+
+const PLACEHOLDER_API_KEY = "not-needed";
+
 export class OpenAIEmbeddingBackend implements EmbeddingBackend {
-  readonly provider = "openai" as const;
+  readonly provider: OpenAICompatibleEmbeddingProvider;
   readonly model: string;
+  readonly baseURL?: string;
+  readonly dimensions?: number;
   private readonly client: OpenAI;
 
-  constructor(apiKey: string, model: string) {
+  constructor(apiKey: string, model: string, options?: OpenAIEmbeddingOptions) {
     this.model = model;
-    this.client = new OpenAI({ apiKey });
+    this.provider = options?.provider ?? "openai";
+    this.baseURL = options?.baseURL;
+    this.dimensions = options?.dimensions;
+    this.client = new OpenAI({
+      apiKey: apiKey || PLACEHOLDER_API_KEY,
+      ...(options?.baseURL ? { baseURL: options.baseURL } : {}),
+    });
   }
 
   async embed(
@@ -28,7 +50,9 @@ export class OpenAIEmbeddingBackend implements EmbeddingBackend {
     const texts = inputs.map((i) => {
       const n = normalizeEmbeddingInput(i);
       if (n.type !== "text") {
-        throw new Error("OpenAI embedding backend only supports text inputs");
+        throw new Error(
+          `${this.provider === "custom" ? "Custom" : "OpenAI"} embedding backend only supports text inputs`,
+        );
       }
       return n.text;
     });
@@ -38,6 +62,7 @@ export class OpenAIEmbeddingBackend implements EmbeddingBackend {
         model: this.model,
         input: texts,
         encoding_format: "float",
+        ...(this.dimensions != null ? { dimensions: this.dimensions } : {}),
       },
       {
         signal: options?.signal,
@@ -46,3 +71,29 @@ export class OpenAIEmbeddingBackend implements EmbeddingBackend {
     return response.data.map((item) => item.embedding);
   }
 }
+
+/**
+ * Trim and strip a trailing slash from a user-supplied OpenAI-compatible
+ * embeddings base URL. Returns null when the value is empty, not http(s),
+ * or not a valid absolute URL.
+ */
+export function resolveOpenAICompatibleBaseUrl(
+  value: string | undefined,
+): string | null {
+  if (value == null) {
+    return null;
+  }
+  const trimmed = value.trim().replace(/\/+$/, "");
+  if (!trimmed) {
+    return null;
+  }
+  try {
+    const parsed = new URL(trimmed);
+    if (parsed.protocol !== "http:" && parsed.protocol !== "https:") {
+      return null;
+    }
+  } catch {
+    return null;
+  }
+  return trimmed;
+}

package/src/persistence/embeddings/embedding-types.test.ts

@@ -1,6 +1,7 @@
 import { describe, expect, test } from "bun:test";
 
 import {
+  embeddingContentHashWithExtras,
   embeddingInputContentHash,
   type MultimodalEmbeddingInput,
   normalizeEmbeddingInput,
@@ -114,3 +115,26 @@ describe("embeddingInputContentHash", () => {
     expect(hash).toMatch(/^[0-9a-f]{64}$/);
   });
 });
+
+describe("embeddingContentHashWithExtras", () => {
+  test("matches the bare content hash when extras are empty", () => {
+    expect(embeddingContentHashWithExtras("hello")).toBe(
+      embeddingInputContentHash("hello"),
+    );
+    expect(embeddingContentHashWithExtras("hello", [])).toBe(
+      embeddingInputContentHash("hello"),
+    );
+  });
+
+  test("changes when extras change for identical text", () => {
+    const text = "hello";
+    const a = embeddingContentHashWithExtras(text, [
+      "url=http://127.0.0.1:4000/v1",
+    ]);
+    const b = embeddingContentHashWithExtras(text, [
+      "url=http://127.0.0.1:4001/v1",
+    ]);
+    expect(a).not.toBe(b);
+    expect(a).not.toBe(embeddingInputContentHash(text));
+  });
+});