From f41488e250b6a93131262d459dabf6f7f7dafb96 Mon Sep 17 00:00:00 2001
From: Stephen Adamson <popertots@gmail.com>
Date: Sat, 15 Nov 2025 01:21:23 +0000
Subject: [PATCH] redo WeightedEmbeddingFunction

---
 mcp_codebase.py | 54 ++++++++++++++++++++++++++++++-------------------
 1 file changed, 33 insertions(+), 21 deletions(-)

diff --git a/mcp_codebase.py b/mcp_codebase.py
index c9e9f04..c919d1e 100644
--- a/mcp_codebase.py
+++ b/mcp_codebase.py
@@ -2133,32 +2133,44 @@ def build_indexes():
 
     # Create a wrapper that will return our pre-computed weighted embeddings
     class WeightedEmbeddingFunction:
-        def __init__(self, weighted_embs):
-            self.weighted_embs = weighted_embs
-            self.idx = 0
+        def __init__(self, texts, weighted_embs):
+            """
+            texts:          list[str]
+            weighted_embs:  list[list[float]]
+            """
 
-        def embed_documents(self, texts):
-            # Return the pre-computed weighted embeddings
-            start = self.idx
-            end = start + len(texts)
-            result = self.weighted_embs[start:end]
-            self.idx = end
+            # Map each text to a queue of its embeddings (duplicate-safe)
+            self.lookup = defaultdict(deque)
+
+            for t, emb in zip(texts, weighted_embs):
+                self.lookup[t].append(emb)
+
+        def embed_documents(self, docs):
+            """
+            Chroma may pass docs in any size or order.
+            Always return correct per-document embeddings.
+            """
+            result = []
+            for d in docs:
+                if not self.lookup[d]:
+                    raise ValueError(
+                        f"No embedding left for text {repr(d[:80])} — "
+                        f"likely mismatch between texts and embeddings."
+                    )
+                result.append(self.lookup[d].popleft())
             return result
 
-        def batch_embed_documents(texts: List[str]) -> List[List[float]]:
-            # Embed documents in batches
-            all_embeddings = []
-            for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
-                batch = texts[i:i + EMBEDDING_BATCH_SIZE]
-                embeddings = embeddings.embed_documents(batch)
-                all_embeddings.extend(embeddings)
-            return all_embeddings
+        def embed_query(self, q):
+            """
+            Queries should use base model, not weighted context embeddings.
+            """
+            return embeddings.embed_query(q)
 
-        def embed_query(self, text):
-            # For queries, use the base embedding model (no weighting)
-            return embeddings.embed_query(text)
+    weighted_emb_func = WeightedEmbeddingFunction(
+        texts=texts,
+        weighted_embs=weighted_embeddings
+    )
 
-    weighted_emb_func = WeightedEmbeddingFunction(weighted_embeddings)
     vectorstore = Chroma.from_texts(
         texts=texts,
         embedding=weighted_emb_func,