LlamaIndexRetrieval and FilesRetrieval block the event loop on every lookup #7674
Description
🔴 Required Information
Describe the Bug:
LlamaIndexRetrieval.run_async calls the retriever's synchronous retrieve() directly (src/google/adk/tools/retrieval/llama_index_retrieval.py:41), so the whole lookup runs on the event loop. Retrieval blocks on I/O: a VectorStoreIndex retriever embeds the query before it searches, and FilesRetrieval's default GoogleGenAIEmbedding makes a synchronous embed_content request to do it. Nothing else on the loop runs until the lookup returns, so other sessions on the same server, streamed responses and the tool's sibling calls all wait.
When the model asks for several lookups in one turn, ADK runs the calls in parallel, but each call holds the loop, so they run one after another.
RunConfig.tool_thread_pool_config does not help outside live mode, because it only moves synchronous FunctionTool and FunctionNode callables to the pool. VertexAiRagRetrieval had the same problem in #5033, which was fixed by running the call in asyncio.to_thread.
Steps to Reproduce:
- Install
google-adk2.12.0 withllama-index-core(also reproduced onmainatfacd7dda). - Run the script under "Minimal Reproduction Code". It needs no network or API key: its embedding model sleeps for one second to stand in for the request a remote embedding model makes.
Expected Behavior:
The event loop keeps running other tasks while the retriever works, and two lookups started together finish in about the time of one.
Observed Behavior:
Two retrievals took 2.01s.
The event loop stalled for up to 2.01s.
Environment Details:
- ADK Library Version (pip show google-adk): 2.12.0,
mainatfacd7dda - Desktop OS: Windows 11
- Python Version (python -V): Python 3.11.3
Model Information:
- Are you using LiteLLM: No
- Which model is being used: N/A, the script calls the tool directly (the
Runnerrun under "Logs" uses a scripted model)
🟡 Optional Information
Regression:
No. run_async has called retrieve() on the event loop since the tool was added.
Logs:
The same stall through Runner: an Agent with this tool and a scripted model that asks for two lookups in one turn, with a 10 ms heartbeat task measuring the longest stall.
Runner turn with two parallel search_docs calls -> ['search_docs', 'search_docs']
took 2.14s, event loop stalled for 2.01s
Screenshots / Video:
N/A
Additional Context:
Switching to aretrieve would not be enough. BaseRetriever only requires _retrieve, and its default _aretrieve calls _retrieve on the loop, so retrievers without an async path of their own still block. Several in llama-index-core are like that, such as SummaryIndexEmbeddingRetriever, SQLRetriever and AutoMergingRetriever. Running retrieve in asyncio.to_thread covers every retriever.
Minimal Reproduction Code:
import asyncio
import time
from google.adk.tools.retrieval.llama_index_retrieval import LlamaIndexRetrieval
from llama_index.core import Document
from llama_index.core import VectorStoreIndex
from llama_index.core.base.embeddings.base import BaseEmbedding
class SlowEmbedding(BaseEmbedding):
"""Takes one second to embed a query, like a request to a remote model."""
def _get_query_embedding(self, query):
time.sleep(1)
return [1.0, 0.0]
async def _aget_query_embedding(self, query):
await asyncio.sleep(1)
return [1.0, 0.0]
def _get_text_embedding(self, text):
return [1.0, 0.0]
async def _aget_text_embedding(self, text):
return [1.0, 0.0]
async def heartbeat(gaps):
"""Ticks every 10 ms and records how long each tick actually took."""
last = time.perf_counter()
while True:
await asyncio.sleep(0.01)
now = time.perf_counter()
gaps.append(now - last)
last = now
async def main():
index = VectorStoreIndex.from_documents(
[Document(text="ADK runs tools on the asyncio event loop.")],
embed_model=SlowEmbedding(),
)
tool = LlamaIndexRetrieval(
name="search_docs",
description="Searches the product docs.",
retriever=index.as_retriever(),
)
gaps = []
beat = asyncio.create_task(heartbeat(gaps))
await asyncio.sleep(0.05)
# Two lookups at once, the way ADK runs parallel function calls.
start = time.perf_counter()
await asyncio.gather(
tool.run_async(args={"query": "event loop"}, tool_context=None),
tool.run_async(args={"query": "tool calls"}, tool_context=None),
)
elapsed = time.perf_counter() - start
beat.cancel()
print(f"Two retrievals took {elapsed:.2f}s.")
print(f"The event loop stalled for up to {max(gaps):.2f}s.")
asyncio.run(main())How often has this issue occurred?:
- Always (100%)
Activity