From 33f47db17574254b41b51289ef8a758f59700045 Mon Sep 17 00:00:00 2001
From: kayo09 <68217041+kayo09@users.noreply.github.com>
Date: Sat, 4 Apr 2026 18:14:11 +0530
Subject: [PATCH] fix(lib): handle embedding dimension mismatch in Unlimiformer
 retrieval

Skip mismatched key/query lengths instead of throwing from cosineSimilarity;
attentionScore returns NaN for length mismatch; extend tests.
---
 packages/lib/unlimiformer.test.ts | 40 ++++++++++++++++++++++++++++++
 packages/lib/unlimiformer.ts      | 41 +++++++++++++++++++++++--------
 2 files changed, 71 insertions(+), 10 deletions(-)

diff --git a/packages/lib/unlimiformer.test.ts b/packages/lib/unlimiformer.test.ts
index ae988a22..cd72921e 100644
--- a/packages/lib/unlimiformer.test.ts
+++ b/packages/lib/unlimiformer.test.ts
@@ -39,6 +39,27 @@ describe("topKAttentionKeys", () => {
 		const top = topKAttentionKeys(q, keys, 10)
 		expect(top).toHaveLength(2)
 	})
+
+	test("skips keys whose dimension does not match the query (no throw)", () => {
+		const q = unit(1, 0, 0)
+		const keys = [unit(1, 0, 0), [1, 0], unit(0, 1, 0)]
+		const top = topKAttentionKeys(q, keys, 5)
+		expect(top.map((t) => t.index)).toEqual([0, 2])
+	})
+
+	test("returns empty when no key matches query dimension", () => {
+		const q = unit(1, 0, 0)
+		expect(
+			topKAttentionKeys(
+				q,
+				[
+					[1, 0],
+					[0, 1],
+				],
+				3,
+			),
+		).toEqual([])
+	})
 })
 
 describe("attentionScores", () => {
@@ -48,6 +69,14 @@ describe("attentionScores", () => {
 		const scores = attentionScores(q, keys)
 		expect(scores).toHaveLength(2)
 	})
+
+	test("uses NaN when key dimension mismatches query", () => {
+		const q = unit(1, 0, 0)
+		const scores = attentionScores(q, [unit(1, 0, 0), [1, 0]])
+		expect(scores).toHaveLength(2)
+		expect(Number.isFinite(scores[0] ?? Number.NaN)).toBe(true)
+		expect(scores[1]).toBeNaN()
+	})
 })
 
 describe("topKAttentionKeysMultiHead", () => {
@@ -73,4 +102,15 @@ describe("rankItemsByAttentionTopK", () => {
 		expect(ranked[0]?.item.id).toBe("c")
 		expect(ranked[0]?.originalIndex).toBe(2)
 	})
+
+	test("skips items whose embedding length does not match the query", () => {
+		const items = [
+			{ id: "wide", e: [0.1, 0.2, 0.3, 0.4] },
+			{ id: "ok", e: unit(0, 1, 0) },
+		]
+		const q = unit(0, 1, 0)
+		const ranked = rankItemsByAttentionTopK(q, items, (x) => x.e, 2)
+		expect(ranked).toHaveLength(1)
+		expect(ranked[0]?.item.id).toBe("ok")
+	})
 })
diff --git a/packages/lib/unlimiformer.ts b/packages/lib/unlimiformer.ts
index a72a0c10..882fae9c 100644
--- a/packages/lib/unlimiformer.ts
+++ b/packages/lib/unlimiformer.ts
@@ -17,12 +17,21 @@ export type AttentionTopK = {
 /**
  * Dot-product attention score between one query vector and one key vector.
  * For normalized embeddings this matches cosine similarity.
+ *
+ * Returns `NaN` when `query` and `key` have different lengths (e.g. mixed embedding
+ * models) so callers can avoid throwing from `cosineSimilarity`.
  */
-export const attentionScore = (query: number[], key: number[]): number =>
-	cosineSimilarity(query, key)
+export const attentionScore = (query: number[], key: number[]): number => {
+	if (query.length !== key.length) {
+		return Number.NaN
+	}
+	return cosineSimilarity(query, key)
+}
 
 /**
- * Attention scores for `query` against every row in `keys` (same dimension as query).
+ * Attention scores for `query` against every row in `keys`, aligned by index.
+ * Entries are `NaN` when a key length does not match the query (same embedding model
+ * is required for a meaningful score).
  */
 export const attentionScores = (query: number[], keys: number[][]): number[] =>
 	keys.map((key) => attentionScore(query, key))
@@ -40,12 +49,19 @@ export const topKAttentionKeys = (
 		return []
 	}
 
-	const effectiveK = Math.min(k, keys.length)
-	const scored: AttentionTopK[] = keys.map((key, index) => ({
-		index,
-		score: attentionScore(query, key),
-	}))
+	const scored: AttentionTopK[] = keys.flatMap((key, index) => {
+		const score = attentionScore(query, key)
+		if (!Number.isFinite(score)) {
+			return []
+		}
+		return [{ index, score }]
+	})
 
+	if (scored.length === 0) {
+		return []
+	}
+
+	const effectiveK = Math.min(k, scored.length)
 	scored.sort((a, b) => b.score - a.score)
 	return scored.slice(0, effectiveK)
 }
@@ -68,7 +84,8 @@ export type RankedItem<T> = {
 
 /**
  * Rank arbitrary items that carry embeddings, returning the top-k by attention score.
- * Items with missing or empty embeddings are skipped.
+ * Items with missing or empty embeddings, or embeddings whose length does not match
+ * `queryEmbedding`, are skipped.
  */
 export const rankItemsByAttentionTopK = <T>(
 	queryEmbedding: number[],
@@ -87,7 +104,11 @@ export const rankItemsByAttentionTopK = <T>(
 		const item = items[i]
 		if (item === undefined) continue
 		const embedding = getEmbedding(item)
-		if (embedding && embedding.length > 0) {
+		if (
+			embedding &&
+			embedding.length > 0 &&
+			embedding.length === queryEmbedding.length
+		) {
 			packed.push({ item, originalIndex: i, embedding })
 		}
 	}