From b737f0a6112ed049855419d1b0c8e74b51f0c3be Mon Sep 17 00:00:00 2001 From: Nitjsefnie Date: Sat, 18 Jul 2026 19:58:18 +0000 Subject: [PATCH] fix(providers): use AsyncElasticsearch in elasticsearch retriever (#82) The retriever constructed a synchronous Elasticsearch client but called self._client.search(...) without await from an async def retrieve(), blocking the event loop during network I/O. Mirrors the fix shape from #182 (pgvector): swap the sync client for the library's async client and await the I/O call. Unlike pgvector, AsyncElasticsearch's constructor performs no I/O, so no lazy-connection deferral is needed. - Use AsyncElasticsearch instead of Elasticsearch. - await self._client.search(...) in retrieve(). - Add close() to terminate the async client. Adds tests/unit/test_providers/test_elasticsearch_retriever.py (4 tests, elasticsearch async surface mocked via sys.modules patch, no real Elasticsearch cluster required). Co-Authored-By: Claude Sonnet 4.6 --- .../providers/retrievers/elasticsearch.py | 10 +- .../test_elasticsearch_retriever.py | 113 ++++++++++++++++++ 2 files changed, 120 insertions(+), 3 deletions(-) create mode 100644 tests/unit/test_providers/test_elasticsearch_retriever.py diff --git a/openagent_eval/providers/retrievers/elasticsearch.py b/openagent_eval/providers/retrievers/elasticsearch.py index d8383f2..78346a5 100644 --- a/openagent_eval/providers/retrievers/elasticsearch.py +++ b/openagent_eval/providers/retrievers/elasticsearch.py @@ -61,7 +61,7 @@ def __init__( ) try: - from elasticsearch import Elasticsearch + from elasticsearch import AsyncElasticsearch except ImportError as exc: # pragma: no cover - depends on installed dep raise ImportError( "elasticsearch is required for the elasticsearch retriever. " @@ -69,7 +69,7 @@ def __init__( ) from exc try: - self._client = Elasticsearch(hosts=hosts, api_key=api_key) + self._client = AsyncElasticsearch(hosts=hosts, api_key=api_key) except Exception as exc: raise ProviderConnectionError( message=f"Failed to connect to Elasticsearch: {exc}", @@ -103,7 +103,7 @@ async def retrieve(self, query: str, k: int = 5) -> list[Document]: "size": k, "_source": [self._content_field], } - resp = self._client.search(index=self._index, **body) + resp = await self._client.search(index=self._index, **body) except Exception as exc: raise ProviderExecutionError( message=f"Elasticsearch query failed: {exc}", @@ -130,3 +130,7 @@ async def retrieve(self, query: str, k: int = 5) -> list[Document]: ) ) return documents + + async def close(self) -> None: + """Close the underlying async Elasticsearch client.""" + await self._client.close() diff --git a/tests/unit/test_providers/test_elasticsearch_retriever.py b/tests/unit/test_providers/test_elasticsearch_retriever.py new file mode 100644 index 0000000..97fe9ea --- /dev/null +++ b/tests/unit/test_providers/test_elasticsearch_retriever.py @@ -0,0 +1,113 @@ +"""Tests for the Elasticsearch retriever (issue #82). + +The retriever must use ``AsyncElasticsearch`` and ``await`` the ``search`` +call. A synchronous ``Elasticsearch`` client's ``search()`` blocks the event +loop when driven from an ``async def retrieve()``. These tests mock the +``elasticsearch`` async surface so no real Elasticsearch cluster is required. +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from openagent_eval.exceptions.provider import ProviderConnectionError +from openagent_eval.providers.embedders.base import Embedder +from openagent_eval.providers.retrievers.elasticsearch import ElasticsearchRetriever + + +class _FakeEmbedder(Embedder): + name = "fake" + description = "fake embedder for tests" + + async def embed(self, texts: list[str]) -> list[list[float]]: + return [[0.1, 0.2, 0.3] for _ in texts] + + +def _make_es_mock(resp: dict | None = None) -> tuple[MagicMock, AsyncMock]: + """Build a fake ``elasticsearch`` module exposing an async client. + + ``AsyncElasticsearch`` is constructed synchronously (no ``await``, matching + real ``elasticsearch-py`` behaviour), but its ``search`` method is a + coroutine that must be ``await``-ed. + """ + if resp is None: + resp = { + "hits": { + "hits": [ + {"_score": 5.0, "_source": {"content": "doc text"}, "_id": "1"} + ] + } + } + + client = AsyncMock() + client.search = AsyncMock(return_value=resp) + client.close = AsyncMock() + + es_module = MagicMock() + es_module.AsyncElasticsearch = MagicMock(return_value=client) + # A sync Elasticsearch must NOT exist on the fake; the retriever must not + # construct or call it. + del es_module.Elasticsearch + return es_module, client + + +def _build_retriever(**kwargs) -> ElasticsearchRetriever: + return ElasticsearchRetriever( + index="docs", + hosts="http://localhost:9200", + **kwargs, + ) + + +class TestElasticsearchAsyncClient: + async def test_retrieve_uses_async_client_and_awaits_search(self): + es_module, client = _make_es_mock() + + with patch.dict("sys.modules", {"elasticsearch": es_module}): + retriever = _build_retriever() + docs = await retriever.retrieve("hello", k=3) + + # Async client class was used (not a sync blocking client). + es_module.AsyncElasticsearch.assert_called_once() + # The search call was awaited, not just invoked synchronously. + client.search.assert_awaited_once() + assert len(docs) == 1 + assert docs[0].content == "doc text" + assert docs[0].id == "1" + + async def test_retrieve_knn_mode_awaits_search(self): + es_module, client = _make_es_mock() + + with patch.dict("sys.modules", {"elasticsearch": es_module}): + retriever = _build_retriever( + mode="knn", embedder=_FakeEmbedder(), vector_field="embedding" + ) + docs = await retriever.retrieve("hello", k=3) + + client.search.assert_awaited_once() + _, kwargs = client.search.call_args + assert kwargs["knn"]["query_vector"] == [0.1, 0.2, 0.3] + assert len(docs) == 1 + + async def test_construction_failure_raises_typed_error(self): + es_module = MagicMock() + es_module.AsyncElasticsearch = MagicMock(side_effect=RuntimeError("bad config")) + del es_module.Elasticsearch + + with ( + patch.dict("sys.modules", {"elasticsearch": es_module}), + pytest.raises(ProviderConnectionError, match="Failed to connect"), + ): + _build_retriever() + + async def test_close_terminates_client(self): + es_module, client = _make_es_mock() + + with patch.dict("sys.modules", {"elasticsearch": es_module}): + retriever = _build_retriever() + await retriever.retrieve("q") + await retriever.close() + + client.close.assert_awaited_once()