redo WeightedEmbeddingFunction

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):
+        def __init__(self, texts, weighted_embs):
-            self.weighted_embs = weighted_embs
+            """
-            self.idx = 0
+            texts:          list[str]
+            weighted_embs:  list[list[float]]
+            """
 
-        def embed_documents(self, texts):
+            # Map each text to a queue of its embeddings (duplicate-safe)
-            # Return the pre-computed weighted embeddings
+            self.lookup = defaultdict(deque)
-            start = self.idx
+
-            end = start + len(texts)
+            for t, emb in zip(texts, weighted_embs):
-            result = self.weighted_embs[start:end]
+                self.lookup[t].append(emb)
-            self.idx = end
+
+        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]]:
+        def embed_query(self, q):
-            # Embed documents in batches
+            """
-            all_embeddings = []
+            Queries should use base model, not weighted context embeddings.
-            for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
+            """
-                batch = texts[i:i + EMBEDDING_BATCH_SIZE]
+            return embeddings.embed_query(q)
-                embeddings = embeddings.embed_documents(batch)
-                all_embeddings.extend(embeddings)
-            return all_embeddings
 
-        def embed_query(self, text):
+    weighted_emb_func = WeightedEmbeddingFunction(
-            # For queries, use the base embedding model (no weighting)
+        texts=texts,
-            return embeddings.embed_query(text)
+        weighted_embs=weighted_embeddings
+    )
 
-    weighted_emb_func = WeightedEmbeddingFunction(weighted_embeddings)
     vectorstore = Chroma.from_texts(
         texts=texts,
         embedding=weighted_emb_func,