External/roboco
1
0
Fork 0
mirror of https://github.com/rennf93/roboco.git synced 2026-08-03 07:23:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
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.
This commit is contained in:
Renn F
2026-06-15 06:12:00 +02:00
parent dfd0d1f188
commit d7aee91b39
3 changed files with 163 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
roboco/services/optimal.py
+65 -26
View File
@@ -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",
roboco/services/optimal_brain/indexes/base.py
+44 -15
View File
@@ -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.
tests/unit/services/test_optimal_search_fanout.py
+54
View File
@@ -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()