fix(embeddings): forward output dimensions to Gemini for consistent embedding dims

CHANGELOG.md

@@ -14,6 +14,7 @@ _In development — bullets added per PR; finalized at release._
 
 ### 🐛 Fixed
 
+- **fix(embeddings):** forward output dimensions to Gemini for consistent embedding dims. (thanks @nguyenha935)
 - **fix(combo): round-robin members fail over faster under concurrency saturation via a configurable queue depth** — when a round-robin combo member was saturated, requests sat in the per-model semaphore's **unbounded** queue and only failed over to the next member after the full `queueTimeoutMs` (default 30s) elapsed — so a burst of agentic requests deep-queued one hot member instead of spilling to healthy ones. The per-model semaphore now accepts a bounded queue depth and emits `SEMAPHORE_QUEUE_FULL` once it is full (the round-robin loop already cascades on that code), so a configured low depth fails over immediately. A new `queueDepth` combo-config knob (global default / provider override / per-combo, default **20** for backward compatibility; **0** = never queue → fail over now) is exposed in Settings → Combo Defaults. ([#3872](https://github.com/diegosouzapw/OmniRoute/issues/3872) — thanks @KooshaPari)
 
 ---

open-sse/handlers/embeddings.ts

@@ -144,6 +144,20 @@ export async function handleEmbedding({
     }
   }
 
+  // Gemini embedding models (gemini-embedding-001 / -2-preview / text-embedding-004)
+  // default to 3072-dim vectors. Clients targeting pgvector-style schemas typically
+  // request a smaller size (e.g. 1536) via OpenAI's `dimensions` field, but Google's
+  // OpenAI-compatibility shim at /v1beta/openai/embeddings does not document the
+  // `dimensions` → `outputDimensionality` translation. Mirror the request value into
+  // the Gemini-native `outputDimensionality` field so the upstream actually returns
+  // the requested vector size. Ported from upstream decolua/9router#1366.
+  if (provider === "gemini" && upstreamBody.outputDimensionality === undefined) {
+    const outputDimensionality = Number(body.dimensions);
+    if (Number.isFinite(outputDimensionality) && outputDimensionality > 0) {
+      upstreamBody.outputDimensionality = outputDimensionality;
+    }
+  }
+
   // Inject model-level default params (e.g. NVIDIA NIM asymmetric models require
   // `input_type`) only for keys the client did not already supply, so a
   // client-sent value is never overwritten. Symmetric models carry no defaults

tests/unit/embeddings-gemini-dimensions.test.ts

@@ -0,0 +1,163 @@
+import test from "node:test";
+import assert from "node:assert/strict";
+import { mkdtempSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+process.env.DATA_DIR = mkdtempSync(join(tmpdir(), "omniroute-embeddings-gemini-dim-"));
+
+const { handleEmbedding } = await import("../../open-sse/handlers/embeddings.ts");
+
+// Ported from upstream decolua/9router#1366 (author @nguyenha935).
+// Gemini embedding models can return 3072 dimensions by default. OpenAI-compatible
+// clients may request a smaller embedding (e.g. 1536 for pgvector schemas) via the
+// `dimensions` field. The Gemini native API uses `outputDimensionality` instead;
+// Google's OpenAI-compatibility shim does not document the `dimensions` translation,
+// so OmniRoute must forward `outputDimensionality` alongside `dimensions` for Gemini
+// embedding requests to guarantee the requested vector size lands at the model.
+
+function captureFetch(captured: { body?: Record<string, unknown> }) {
+  return async (_url: unknown, options: { headers?: unknown; body?: unknown } = {}) => {
+    captured.body = JSON.parse(String(options.body || "{}"));
+    return new Response(
+      JSON.stringify({
+        data: [{ object: "embedding", embedding: new Array(1536).fill(0.1), index: 0 }],
+        usage: { prompt_tokens: 4, total_tokens: 4 },
+      }),
+      { status: 200, headers: { "content-type": "application/json" } }
+    );
+  };
+}
+
+test("handleEmbedding forwards Gemini dimensions as outputDimensionality (single input)", async () => {
+  const originalFetch = globalThis.fetch;
+  const captured: { body?: Record<string, unknown> } = {};
+  globalThis.fetch = captureFetch(captured) as typeof fetch;
+
+  try {
+    const result = await handleEmbedding({
+      body: {
+        model: "gemini/text-embedding-004",
+        input: "test",
+        dimensions: 1536,
+      },
+      credentials: { apiKey: "gemini-key" },
+      log: null,
+    });
+
+    assert.equal(result.success, true);
+    // OpenAI-style `dimensions` must still be forwarded (back-compat).
+    assert.equal(captured.body?.dimensions, 1536);
+    // Gemini-native `outputDimensionality` must also be present so the upstream
+    // returns the requested vector size regardless of the OpenAI-shim behavior.
+    assert.equal(captured.body?.outputDimensionality, 1536);
+  } finally {
+    globalThis.fetch = originalFetch;
+  }
+});
+
+test("handleEmbedding forwards Gemini dimensions as outputDimensionality (batch input)", async () => {
+  const originalFetch = globalThis.fetch;
+  const captured: { body?: Record<string, unknown> } = {};
+  globalThis.fetch = captureFetch(captured) as typeof fetch;
+
+  try {
+    const result = await handleEmbedding({
+      body: {
+        model: "gemini/text-embedding-004",
+        input: ["hello", "world"],
+        dimensions: 1536,
+      },
+      credentials: { apiKey: "gemini-key" },
+      log: null,
+    });
+
+    assert.equal(result.success, true);
+    assert.equal(captured.body?.dimensions, 1536);
+    assert.equal(captured.body?.outputDimensionality, 1536);
+  } finally {
+    globalThis.fetch = originalFetch;
+  }
+});
+
+test("handleEmbedding does not inject outputDimensionality when dimensions is omitted (Gemini)", async () => {
+  const originalFetch = globalThis.fetch;
+  const captured: { body?: Record<string, unknown> } = {};
+  globalThis.fetch = captureFetch(captured) as typeof fetch;
+
+  try {
+    const result = await handleEmbedding({
+      body: {
+        model: "gemini/text-embedding-004",
+        input: "test",
+      },
+      credentials: { apiKey: "gemini-key" },
+      log: null,
+    });
+
+    assert.equal(result.success, true);
+    assert.equal(
+      "outputDimensionality" in (captured.body || {}),
+      false,
+      "outputDimensionality must not be injected when the client did not request a specific size"
+    );
+  } finally {
+    globalThis.fetch = originalFetch;
+  }
+});
+
+test("handleEmbedding does not inject outputDimensionality for non-Gemini providers", async () => {
+  const originalFetch = globalThis.fetch;
+  const captured: { body?: Record<string, unknown> } = {};
+  globalThis.fetch = captureFetch(captured) as typeof fetch;
+
+  try {
+    const result = await handleEmbedding({
+      body: {
+        model: "openai/text-embedding-3-small",
+        input: "test",
+        dimensions: 1536,
+      },
+      credentials: { apiKey: "openai-key" },
+      log: null,
+    });
+
+    assert.equal(result.success, true);
+    // OpenAI gets the standard `dimensions` field — not `outputDimensionality`.
+    assert.equal(captured.body?.dimensions, 1536);
+    assert.equal(
+      "outputDimensionality" in (captured.body || {}),
+      false,
+      "outputDimensionality is Gemini-specific and must not leak into other providers"
+    );
+  } finally {
+    globalThis.fetch = originalFetch;
+  }
+});
+
+test("handleEmbedding ignores non-finite/non-positive dimensions for Gemini", async () => {
+  const originalFetch = globalThis.fetch;
+  const captured: { body?: Record<string, unknown> } = {};
+  globalThis.fetch = captureFetch(captured) as typeof fetch;
+
+  try {
+    const result = await handleEmbedding({
+      body: {
+        model: "gemini/text-embedding-004",
+        input: "test",
+        dimensions: 0,
+      },
+      credentials: { apiKey: "gemini-key" },
+      log: null,
+    });
+
+    assert.equal(result.success, true);
+    assert.equal(
+      "outputDimensionality" in (captured.body || {}),
+      false,
+      "0/NaN/negative dimensions must not map to outputDimensionality"
+    );
+  } finally {
+    globalThis.fetch = originalFetch;
+  }
+});