diff --git a/roboco/services/optimal.py b/roboco/services/optimal.py index 1391ed0a..2311662a 100644 --- a/roboco/services/optimal.py +++ b/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: - plugin = self._plugins.get(index_type) - if plugin: - outcome = await plugin.search(query=query, top_k=top_k) - if outcome.success: - results.extend(outcome.results) - else: - logger.warning( - "Search failed for index", - index_type=index_type.value, - error=outcome.error_message, - ) + # Embed the query ONCE (HyDE + embed), then run every index's vector + # search concurrently — instead of each index re-embedding in series. + try: + query_embedding = await plugins[0][1].compute_query_embedding(query) + except Exception as e: + logger.warning("Query embedding failed", error=str(e)) + return [] - # 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.""" - plugin = self._plugins.get(index_type) - if not plugin: - return - - count = await plugin.count() - if count == 0: + """Search one index with a pre-computed embedding; update buf in place.""" + if await plugin.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: - await self._search_single_index(index_type, query, top_k, buf) + try: + 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", diff --git a/roboco/services/optimal_brain/indexes/base.py b/roboco/services/optimal_brain/indexes/base.py index e9bbc122..7ca1726f 100644 --- a/roboco/services/optimal_brain/indexes/base.py +++ b/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 the index. + """Search using a pre-computed query embedding (skips HyDE/embed). - Args: - query: Natural language search query - top_k: Number of results to return - filters: Optional metadata filters - - Returns: - SearchOutcome with results and success status + Splitting embed from fetch lets the service embed once and run every + index's vector search concurrently. """ 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. diff --git a/tests/unit/services/test_optimal_search_fanout.py b/tests/unit/services/test_optimal_search_fanout.py new file mode 100644 index 00000000..9d7580be --- /dev/null +++ b/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()