wip

sourcegraph · Dec 27, 2023 · b2963db · b2963db
1 parent 250dda0
commit b2963db
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/provider/docs/src/corpus/search/embeddings.ts b/provider/docs/src/corpus/search/embeddings.ts
@@ -26,14 +26,15 @@ export async function embeddingsSearch(
 ): Promise<CorpusSearchResult[]> {
     const queryVec = await embedText(query)
 
-    const results: CorpusSearchResult[] = []
-    for (const { doc, chunks } of index.docs) {
-        for (const [i, chunk] of chunks.entries()) {
-            const chunkVec = await cachedEmbedText(chunk.text, cache)
-            const score = cos_sim(queryVec, chunkVec)
-            results.push({ doc: doc.id, chunk: i, score, excerpt: chunk.text })
-        }
-    }
+    const results: CorpusSearchResult[] = await Promise.all(
+        index.docs.flatMap(({ doc, chunks }) =>
+            chunks.map(async (chunk, i) => {
+                const chunkVec = await cachedEmbedText(chunk.text, cache)
+                const score = cos_sim(queryVec, chunkVec)
+                return { doc: doc.id, chunk: i, score, excerpt: chunk.text } satisfies CorpusSearchResult
+            })
+        )
+    )
 
     results.sort((a, b) => b.score - a.score)