From 3e0646b3ddbc9244b537cb142df0388bfa0a5772 Mon Sep 17 00:00:00 2001 From: Daniel Yarmoluk <40245819+Yarmoluk@users.noreply.github.com> Date: Mon, 2 Nov 2026 16:51:57 -0600 Subject: [PATCH] fix(rerank): use a lazy per-loop Cohere AsyncClient in CohereReranker CohereReranker built its AsyncClient once in __init__ and passed it the same kwargs as the sync client. Its connection pool was bound to the first event loop that used it, so a reranker that outlived that loop failed on the next one with "Event loop is closed", and a user-supplied sync httpx_client reached the AsyncClient, so arank raised "object Response can't be used in 'await' expression". Build the async client on first use via _get_aclient(), rebuild it when the running loop changes, and keep a user-supplied async_httpx_client once built. A sync httpx_client without an async one disables the async client and arank falls back to the sync client with a warning. This is the same handling CohereTextVectorizer got in #785. --- redisvl/utils/rerank/cohere.py | 64 +++++++++- tests/unit/test_cohere_reranker.py | 181 +++++++++++++++++++++++++++++ 2 files changed, 241 insertions(+), 4 deletions(-) create mode 100644 tests/unit/test_cohere_reranker.py diff --git a/redisvl/utils/rerank/cohere.py b/redisvl/utils/rerank/cohere.py index f068551b..9eead6b3 100644 --- a/redisvl/utils/rerank/cohere.py +++ b/redisvl/utils/rerank/cohere.py @@ -1,10 +1,15 @@ +import asyncio +import logging import os +from collections.abc import Callable from typing import Any from pydantic import PrivateAttr from redisvl.utils.rerank.base import BaseReranker +logger = logging.getLogger(__name__) + class CohereReranker(BaseReranker): """ @@ -36,7 +41,10 @@ class CohereReranker(BaseReranker): """ _client: Any = PrivateAttr() - _aclient: Any = PrivateAttr() + _aclient_factory: Callable[[], Any] | None = PrivateAttr(default=None) + _aclient: Any = PrivateAttr(default=None) + _aclient_loop: asyncio.AbstractEventLoop | None = PrivateAttr(default=None) + _aclient_rebuild_on_new_loop: bool = PrivateAttr(default=True) def __init__( self, @@ -64,6 +72,11 @@ def __init__( reranked results. api_config (Optional[Dict], optional): Dictionary containing the API key. Defaults to None. + **kwargs: Passed to the Cohere ``Client`` and ``AsyncClient``, for + example ``timeout`` or ``base_url``. A custom ``httpx_client`` is + synchronous, so it goes to the sync client only and ``arank`` + falls back to it; pass ``async_httpx_client`` (an + ``httpx.AsyncClient``) to give the async client its own. Raises: ImportError: If the cohere library is not installed. @@ -77,7 +90,8 @@ def __init__( def _initialize_clients(self, api_config: dict[str, Any] | None, **kwargs): """ Setup the Cohere clients using the provided API key or an - environment variable. + environment variable; see ``__init__`` for ``httpx_client`` and + ``async_httpx_client``. """ # Dynamic import of the cohere module try: @@ -95,8 +109,42 @@ def _initialize_clients(self, api_config: dict[str, Any] | None, **kwargs): "Cohere API key is required. " "Provide it in api_config or set the COHERE_API_KEY environment variable." ) + + async_httpx_client = kwargs.pop("async_httpx_client", None) self._client = Client(api_key=api_key, client_name="redisvl", **kwargs) - self._aclient = AsyncClient(api_key=api_key, client_name="redisvl", **kwargs) + + # A user-supplied httpx_client is sync and can't be awaited, so without + # an async one arank falls back to the sync client. + self._aclient_factory = None + self._aclient = None + self._aclient_loop = None + # A user-supplied async httpx client is bound to the caller's loop, so + # it is created once and never rebuilt. + self._aclient_rebuild_on_new_loop = async_httpx_client is None + if kwargs.get("httpx_client") is None or async_httpx_client is not None: + async_kwargs = dict(kwargs) + if async_httpx_client is not None: + async_kwargs["httpx_client"] = async_httpx_client + self._aclient_factory = lambda: AsyncClient( + api_key=api_key, client_name="redisvl", **async_kwargs + ) + + def _get_aclient(self) -> Any: + """Return the async Cohere client, or ``None`` if async is disabled. + + The client's HTTP connection pool is tied to the event loop it was + created on, so a reranker that outlives a loop gets a fresh client + on the next one rather than ``Event loop is closed`` errors. + """ + if self._aclient_factory is None: + return None + loop = asyncio.get_running_loop() + if self._aclient is None or ( + self._aclient_rebuild_on_new_loop and self._aclient_loop is not loop + ): + self._aclient = self._aclient_factory() + self._aclient_loop = loop + return self._aclient def _preprocess(self, query: str, docs: list[dict[str, Any]] | list[str], **kwargs): """ @@ -186,8 +234,16 @@ async def arank( Returns: Union[Tuple[Union[List[Dict[str, Any]], List[str]], float], List[Dict[str, Any]]]: The reranked list of documents and optionally associated scores. """ + aclient = self._get_aclient() + if aclient is None: + logger.warning( + "CohereReranker was given a sync httpx_client and no " + "async_httpx_client. Falling back to sync." + ) + return self.rank(query, docs, **kwargs) + reranker_kwargs, return_score = self._preprocess(query, docs, **kwargs) - rankings = await self._aclient.rerank(**reranker_kwargs) + rankings = await aclient.rerank(**reranker_kwargs) reranked_docs, scores = self._postprocess(docs, rankings) if return_score: return reranked_docs, scores diff --git a/tests/unit/test_cohere_reranker.py b/tests/unit/test_cohere_reranker.py new file mode 100644 index 00000000..671db3fe --- /dev/null +++ b/tests/unit/test_cohere_reranker.py @@ -0,0 +1,181 @@ +import asyncio +import sys +import types + +import pytest + +from redisvl.utils.rerank.cohere import CohereReranker + +DOCS = ["doc one", "doc two", "doc three"] + + +def _response(documents, top_n): + # Rank in reverse order with descending scores so the caller can tell the + # stub's ordering apart from the input ordering. + n = min(top_n or len(documents), len(documents)) + results = [ + types.SimpleNamespace(index=i, relevance_score=1.0 - i / 10) + for i in reversed(range(len(documents))) + ][:n] + return types.SimpleNamespace(results=results) + + +@pytest.fixture +def fake_cohere_module(monkeypatch): + fake_module = types.ModuleType("cohere") + + class FakeClient: + instances = [] + + def __init__(self, api_key=None, client_name=None, **kwargs): + self.api_key = api_key + self.client_name = client_name + self.kwargs = kwargs + self.rerank_calls = [] + self.__class__.instances.append(self) + + def rerank(self, **kwargs): + self.rerank_calls.append(kwargs) + return _response(kwargs["documents"], kwargs["top_n"]) + + class FakeAsyncClient(FakeClient): + instances = [] + + async def rerank(self, **kwargs): + return FakeClient.rerank(self, **kwargs) + + fake_module.Client = FakeClient + fake_module.AsyncClient = FakeAsyncClient + monkeypatch.setitem(sys.modules, "cohere", fake_module) + return fake_module + + +@pytest.fixture +def reranker(fake_cohere_module): + return CohereReranker(api_config={"api_key": "test-key"}, limit=2, timeout=30) + + +def test_init_creates_sync_client_and_defers_async(fake_cohere_module, reranker): + sync_client = fake_cohere_module.Client.instances[-1] + + assert sync_client.api_key == "test-key" + assert sync_client.client_name == "redisvl" + assert sync_client.kwargs == {"timeout": 30} + # The async client is built on first use, inside a running event loop. + assert fake_cohere_module.AsyncClient.instances == [] + + +async def test_arank_uses_async_client(fake_cohere_module, reranker): + docs, scores = await reranker.arank("query", DOCS) + + async_client = fake_cohere_module.AsyncClient.instances[-1] + assert docs == ["doc three", "doc two"] + assert scores == [0.8, 0.9] + assert async_client.kwargs == {"timeout": 30} + assert async_client.rerank_calls == [ + { + "model": "rerank-english-v3.0", + "query": "query", + "top_n": 2, + "documents": DOCS, + "max_chunks_per_doc": None, + } + ] + + +async def test_arank_passes_rank_by_and_overrides(fake_cohere_module, reranker): + dict_docs = [{"content": d} for d in DOCS] + + docs = await reranker.arank( + "query", dict_docs, rank_by="content", limit=1, return_score=False + ) + + async_client = fake_cohere_module.AsyncClient.instances[-1] + assert docs == [{"content": "doc three"}] + call = async_client.rerank_calls[-1] + assert call["rank_fields"] == ["content"] + assert call["top_n"] == 1 + + +def test_rank_still_uses_sync_client(fake_cohere_module, reranker): + sync_client = fake_cohere_module.Client.instances[-1] + + docs, scores = reranker.rank("query", DOCS) + + assert docs == ["doc three", "doc two"] + assert scores == [0.8, 0.9] + assert sync_client.rerank_calls[-1]["documents"] == DOCS + + +def test_async_client_is_rebuilt_for_a_new_event_loop(fake_cohere_module, reranker): + # A reranker that outlives an event loop must not reuse a client whose + # connection pool is bound to the closed loop. + asyncio.run(reranker.arank("first", DOCS)) + asyncio.run(reranker.arank("second", DOCS)) + + first, second = fake_cohere_module.AsyncClient.instances + assert first is not second + assert [call["query"] for call in second.rerank_calls] == ["second"] + + +async def test_async_client_is_reused_within_one_event_loop( + fake_cohere_module, reranker +): + await reranker.arank("one", DOCS) + await reranker.arank("two", DOCS) + + assert len(fake_cohere_module.AsyncClient.instances) == 1 + + +async def test_sync_httpx_client_keeps_arank_on_sync_fallback(fake_cohere_module): + sync_httpx = object() + reranker = CohereReranker( + api_config={"api_key": "test-key"}, limit=2, httpx_client=sync_httpx + ) + sync_client = fake_cohere_module.Client.instances[-1] + + assert sync_client.kwargs == {"httpx_client": sync_httpx} + docs, scores = await reranker.arank("query", DOCS) + + assert docs == ["doc three", "doc two"] + assert scores == [0.8, 0.9] + assert fake_cohere_module.AsyncClient.instances == [] + assert sync_client.rerank_calls[-1]["query"] == "query" + + +async def test_httpx_client_none_still_uses_async_client(fake_cohere_module): + # Config dicts often carry httpx_client=None; that is not a sync client. + reranker = CohereReranker( + api_config={"api_key": "test-key"}, limit=2, httpx_client=None + ) + + await reranker.arank("query", DOCS) + + async_client = fake_cohere_module.AsyncClient.instances[-1] + assert async_client.kwargs == {"httpx_client": None} + assert async_client.rerank_calls[-1]["query"] == "query" + + +def test_async_httpx_client_goes_to_async_client_only_and_is_kept( + fake_cohere_module, +): + sync_httpx, async_httpx = object(), object() + reranker = CohereReranker( + api_config={"api_key": "test-key"}, + limit=2, + httpx_client=sync_httpx, + async_httpx_client=async_httpx, + ) + sync_client = fake_cohere_module.Client.instances[-1] + + assert sync_client.kwargs == {"httpx_client": sync_httpx} + asyncio.run(reranker.arank("first", DOCS)) + asyncio.run(reranker.arank("second", DOCS)) + + # The caller owns this client's loop, so it is never swapped out. + (async_client,) = fake_cohere_module.AsyncClient.instances + assert async_client.kwargs == {"httpx_client": async_httpx} + assert [call["query"] for call in async_client.rerank_calls] == [ + "first", + "second", + ]