perf(rag): embed the query once and search indexes concurrently

OptimalService.search / query (via _aggregate_citations) ran each index's
plugin.search() sequentially, and every plugin.search re-ran HyDE + embed — so an
N-index query made N LLM+embed round-trips in series (~28s across all indexes,
even though the SQL is fast). Embed the query ONCE
(BaseIndexPlugin.compute_query_embedding) and run every index's vector search
concurrently against that single embedding (search_with_embedding +
asyncio.gather). The search/query signatures and return contract are unchanged;
behavior is identical, just ~Nx fewer embed calls and parallel fetch.

Adds a regression test asserting one embed + per-index fan-out.

roboco/services/optimal.py

@@ -1034,6 +1034,30 @@ class OptimalService:
     # SEARCH OPERATIONS
     # =========================================================================
 
+    @staticmethod
+    def _collect_outcomes(
+        plugins: list[tuple[IndexType, BaseIndexPlugin]],
+        outcomes: list[Any],
+    ) -> list[SearchResult]:
+        """Flatten gathered per-index search outcomes; log and skip failures."""
+        results: list[SearchResult] = []
+        for (index_type, _), outcome in zip(plugins, outcomes, strict=True):
+            if isinstance(outcome, BaseException):
+                logger.warning(
+                    "Search failed for index",
+                    index_type=index_type.value,
+                    error=str(outcome),
+                )
+            elif outcome.success:
+                results.extend(outcome.results)
+            else:
+                logger.warning(
+                    "Search failed for index",
+                    index_type=index_type.value,
+                    error=outcome.error_message,
+                )
+        return results
+
     async def search(
         self,
         query: str,
@@ -1054,46 +1078,47 @@ class OptimalService:
         if not self._initialized:
             raise RuntimeError("OptimalService not initialized")
 
-        results: list[SearchResult] = []
         index_types = (
             context.index_types if context and context.index_types else list(IndexType)
         )
+        plugins = [(it, self._plugins[it]) for it in index_types if it in self._plugins]
+        if not plugins:
+            return []
 
-        for index_type in index_types:
+        # Embed the query ONCE (HyDE + embed), then run every index's vector
-            plugin = self._plugins.get(index_type)
+        # search concurrently — instead of each index re-embedding in series.
-            if plugin:
+        try:
-                outcome = await plugin.search(query=query, top_k=top_k)
+            query_embedding = await plugins[0][1].compute_query_embedding(query)
-                if outcome.success:
+        except Exception as e:
-                    results.extend(outcome.results)
+            logger.warning("Query embedding failed", error=str(e))
-                else:
+            return []
-                    logger.warning(
-                        "Search failed for index",
-                        index_type=index_type.value,
-                        error=outcome.error_message,
-                    )
 
-        # Sort by score descending
+        outcomes = await asyncio.gather(
+            *(
+                plugin.search_with_embedding(query_embedding, top_k=top_k)
+                for _, plugin in plugins
+            ),
+            return_exceptions=True,
+        )
+
+        results = self._collect_outcomes(plugins, outcomes)
         results.sort(key=lambda r: r.score, reverse=True)
         return results[: top_k * len(index_types)]
 
     async def _search_single_index(
         self,
         index_type: IndexType,
-        query: str,
+        plugin: BaseIndexPlugin,
+        query_embedding: list[float],
         top_k: int,
         buf: _QueryAggregationBuffer,
     ) -> None:
-        """Search one index and update aggregate buffers in place."""
+        """Search one index with a pre-computed embedding; update buf in place."""
-        plugin = self._plugins.get(index_type)
+        if await plugin.count() == 0:
-        if not plugin:
-            return
-
-        count = await plugin.count()
-        if count == 0:
             logger.debug("Skipping empty index", index_type=index_type.value)
             return
 
-        outcome = await plugin.search(query=query, top_k=top_k)
+        outcome = await plugin.search_with_embedding(query_embedding, top_k=top_k)
         if outcome.success:
             buf.stats[index_type.value] = len(outcome.results)
             buf.citations.extend(outcome.results)
@@ -1109,11 +1134,25 @@ class OptimalService:
     async def _aggregate_citations(
         self, index_types: list[IndexType], query: str, top_k: int
     ) -> tuple[list[SearchResult], dict[str, int], dict[str, str]]:
-        """Run search across the requested indexes and aggregate citations."""
+        """Embed once, then search the requested indexes concurrently."""
         buf = _QueryAggregationBuffer()
+        plugins = [(it, self._plugins[it]) for it in index_types if it in self._plugins]
+        if not plugins:
+            return buf.citations, buf.stats, buf.errors
 
-        for index_type in index_types:
+        try:
-            await self._search_single_index(index_type, query, top_k, buf)
+            query_embedding = await plugins[0][1].compute_query_embedding(query)
+        except Exception as e:
+            logger.warning("RAG query embedding failed", error=str(e))
+            return buf.citations, buf.stats, buf.errors
+
+        await asyncio.gather(
+            *(
+                self._search_single_index(it, plugin, query_embedding, top_k, buf)
+                for it, plugin in plugins
+            ),
+            return_exceptions=True,
+        )
 
         logger.info(
             "RAG search complete",

roboco/services/optimal_brain/indexes/base.py

@@ -854,34 +854,33 @@ class BaseIndexPlugin(ABC):
                 break
         return results
 
-    async def search(
+    async def compute_query_embedding(self, query: str) -> list[float]:
+        """Embed a query (incl. HyDE) once.
+
+        Lets the service compute one embedding and fan it out across every
+        index, instead of each index re-running HyDE + embed sequentially.
+        """
+        return await self._compute_query_embedding(query)
+
+    async def search_with_embedding(
         self,
-        query: str,
+        query_embedding: list[float],
         top_k: int = 5,
         filters: dict[str, Any] | None = None,
     ) -> SearchOutcome:
-        """
+        """Search using a pre-computed query embedding (skips HyDE/embed).
-        Search the index.
 
-        Args:
+        Splitting embed from fetch lets the service embed once and run every
-            query: Natural language search query
+        index's vector search concurrently.
-            top_k: Number of results to return
-            filters: Optional metadata filters
-
-        Returns:
-            SearchOutcome with results and success status
         """
         import time
 
         start_time = time.time()
-
         try:
-            query_embedding = await self._compute_query_embedding(query)
             chunks = await self._fetch_citations(
                 query_embedding, top_k, has_filters=bool(filters)
             )
             results = self._citations_to_results(chunks, top_k, filters)
-
             elapsed_ms = (time.time() - start_time) * 1000
             return SearchOutcome(
                 results=results,
@@ -889,7 +888,6 @@ class BaseIndexPlugin(ABC):
                 index_type=self.index_type,
                 search_time_ms=elapsed_ms,
             )
-
         except Exception as e:
             elapsed_ms = (time.time() - start_time) * 1000
             logger.warning(
@@ -906,6 +904,37 @@ class BaseIndexPlugin(ABC):
                 search_time_ms=elapsed_ms,
             )
 
+    async def search(
+        self,
+        query: str,
+        top_k: int = 5,
+        filters: dict[str, Any] | None = None,
+    ) -> SearchOutcome:
+        """Search the index: embed the query, then run the vector search."""
+        import time
+
+        start_time = time.time()
+        try:
+            query_embedding = await self._compute_query_embedding(query)
+        except Exception as e:
+            elapsed_ms = (time.time() - start_time) * 1000
+            logger.warning(
+                "Search failed",
+                index_type=self.index_type.value,
+                error=str(e),
+                search_time_ms=elapsed_ms,
+            )
+            return SearchOutcome(
+                results=[],
+                success=False,
+                error_message=str(e),
+                index_type=self.index_type,
+                search_time_ms=elapsed_ms,
+            )
+        return await self.search_with_embedding(
+            query_embedding, top_k=top_k, filters=filters
+        )
+
     def _fallback_answer(self, _search_results: list[SearchResult]) -> str:
         """
         Return empty to let OptimalService continue searching other indexes.

tests/unit/services/test_optimal_search_fanout.py

@@ -0,0 +1,54 @@
+"""OptimalService.search embeds the query once and fans out across indexes.
+
+Guards the latency fix: previously each index re-ran HyDE + embed sequentially
+(N LLM calls, serial), so an all-index search took ~28s. Now the query is
+embedded once and every index's vector search runs concurrently with that single
+embedding.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+from roboco.models.optimal import IndexType, SearchOutcome
+from roboco.services.optimal import OptimalService
+
+
+def _fake_plugin(index_type: IndexType) -> MagicMock:
+    plugin = MagicMock()
+    plugin.compute_query_embedding = AsyncMock(return_value=[0.1, 0.2, 0.3])
+    plugin.search_with_embedding = AsyncMock(
+        return_value=SearchOutcome(results=[], success=True, index_type=index_type)
+    )
+    plugin.count = AsyncMock(return_value=1)
+    return plugin
+
+
+@pytest.mark.asyncio
+async def test_search_embeds_once_and_fans_out(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    svc = OptimalService.__new__(OptimalService)
+    p_docs = _fake_plugin(IndexType.DOCUMENTATION)
+    p_journals = _fake_plugin(IndexType.JOURNALS)
+    monkeypatch.setattr(svc, "_initialized", True, raising=False)
+    monkeypatch.setattr(
+        svc,
+        "_plugins",
+        {IndexType.DOCUMENTATION: p_docs, IndexType.JOURNALS: p_journals},
+        raising=False,
+    )
+
+    await svc.search("anything")
+
+    # Embedded exactly once total (on the first plugin), reused across indexes —
+    # not once per index.
+    embed_calls = (
+        p_docs.compute_query_embedding.await_count
+        + p_journals.compute_query_embedding.await_count
+    )
+    assert embed_calls == 1
+    # Every index ran a vector search with the pre-computed embedding.
+    p_docs.search_with_embedding.assert_awaited_once()
+    p_journals.search_with_embedding.assert_awaited_once()