From c1841981a90fbed8966074da17b8c9943d889680 Mon Sep 17 00:00:00 2001 From: pm25coder Date: Mon, 19 Oct 2026 21:54:00 +0800 Subject: [PATCH] fix: stream Ollama responses through AsyncClient instead of a sync generator _ollama_stream() is an async def that iterated ollama.Client.chat(stream=True), which returns a synchronous generator: every next() blocked the event loop for the length of that network read, so a streaming Ollama reply froze every other coroutine for the whole response. Await AsyncClient.chat() instead. Both clients are now built from one _ollama_connection_options() helper, so the streaming and non-streaming paths cannot drift onto different endpoints. tests/test_ollama_stream_nonblocking.py drives the stream with fake clients and measures the loop: a ticker coroutine gets 0 turns on the old code and double-digit turns on the new one, while the chunk contract (text in order, empty content dropped) is pinned in the same file. Closes #103 --- devagent/core/llm.py | 48 ++++++-- tests/test_ollama_stream_nonblocking.py | 144 ++++++++++++++++++++++++ 2 files changed, 180 insertions(+), 12 deletions(-) create mode 100644 tests/test_ollama_stream_nonblocking.py diff --git a/devagent/core/llm.py b/devagent/core/llm.py index 3e30c6a..87eb4b0 100644 --- a/devagent/core/llm.py +++ b/devagent/core/llm.py @@ -575,15 +575,17 @@ def list_ollama_models(self) -> list[str]: except Exception: return [] - def _ollama_client(self) -> ollama.Client: - """Return an ollama.Client configured for local or Ollama Cloud. + def _ollama_connection_options(self) -> dict[str, Any]: + """Connection options shared by the sync and async Ollama clients. + + One builder for both: the two clients serve the same config from + different call sites, so anything that made only one of them differ + would send part of a conversation to a different endpoint. Cloud is detected when api_key is set, or when base_url is not the default localhost address. Cloud requests use https://ollama.com as the host and pass the key as an Authorization: Bearer header. """ - import ollama - _LOCAL_DEFAULT = "http://localhost:11434" base_url = self.cfg.base_url or _LOCAL_DEFAULT api_key = self.cfg.api_key or "" @@ -591,13 +593,27 @@ def _ollama_client(self) -> ollama.Client: # Cloud: explicit key always wins; non-default URL is also cloud-like. if api_key: host = base_url if base_url != _LOCAL_DEFAULT else "https://ollama.com" - return ollama.Client( - host=host, - headers={"Authorization": f"Bearer {api_key}"}, - ) + return {"host": host, "headers": {"Authorization": f"Bearer {api_key}"}} if base_url != _LOCAL_DEFAULT: - return ollama.Client(host=base_url) - return ollama.Client() + return {"host": base_url} + return {} + + def _ollama_client(self) -> ollama.Client: + """Return an ollama.Client configured for local or Ollama Cloud.""" + import ollama + + return ollama.Client(**self._ollama_connection_options()) + + def _ollama_async_client(self) -> ollama.AsyncClient: + """Return an ollama.AsyncClient with the same configuration. + + ``AsyncClient`` performs its reads on the running event loop's IO + layer, which is what the streaming path needs: the sync client returns + a generator whose every ``next()`` blocks. + """ + import ollama + + return ollama.AsyncClient(**self._ollama_connection_options()) def _ollama(self, messages: list[AgentMessage], tools: list[ToolDef] | None) -> LLMResponse: @@ -686,7 +702,15 @@ async def _ollama_stream( messages: list[AgentMessage], tools: list[ToolDef] | None = None, ) -> AsyncIterator[str]: - client = self._ollama_client() + """Stream assistant text from Ollama without blocking the event loop. + + ``ollama.Client.chat(stream=True)`` returns a *synchronous* generator: + each ``next()`` blocks until the next network read returns. Iterating + it inside an ``async def`` therefore froze the loop for the length of + the stream (issue #103). ``AsyncClient`` returns an async iterator for + the same request, so every chunk is awaited instead. + """ + client = self._ollama_async_client() kwargs: dict[str, Any] = { "model": self.cfg.model, "messages": _to_ollama_messages(messages), @@ -696,7 +720,7 @@ async def _ollama_stream( if tools: kwargs["tools"] = _to_openai_tools(tools) - for chunk in client.chat(**kwargs): + async for chunk in await client.chat(**kwargs): text = chunk.message.content if text: yield text diff --git a/tests/test_ollama_stream_nonblocking.py b/tests/test_ollama_stream_nonblocking.py new file mode 100644 index 0000000..dc8d672 --- /dev/null +++ b/tests/test_ollama_stream_nonblocking.py @@ -0,0 +1,144 @@ +"""Issue #103: ``_ollama_stream()`` must not block the event loop. + +The stream used to iterate a *synchronous* ``ollama.Client.chat()`` generator +inside an ``async def``, so every network read froze the loop for the length of +that read. An ``async def`` wrapper makes a generator awaitable; it does not +make a blocking read concurrent. + +The first test drives the stream with a fake client whose per-chunk delay +stands in for a network read, and measures whether a coroutine sharing the +loop still gets to run. The second pins the chunk contract so the fix cannot +change what callers receive. No Ollama server is contacted. +""" + +from __future__ import annotations + +import asyncio +import time +from types import SimpleNamespace + +import pytest + +from devagent.core.config import LLMConfig +from devagent.core.llm import AgentMessage, LLMClient + +CHUNKS = ("part0", "part1", "part2") +CHUNK_DELAY = 0.05 +TICK = 0.01 + + +def _chunk(text: str) -> SimpleNamespace: + return SimpleNamespace(message=SimpleNamespace(content=text, tool_calls=[])) + + +class _SlowSyncClient: + """A stand-in for ``ollama.Client``: blocking reads between chunks.""" + + def __init__(self, delay: float = CHUNK_DELAY) -> None: + self.delay = delay + + def chat(self, **_kwargs): + for text in CHUNKS: + time.sleep(self.delay) + yield _chunk(text) + + +class _SlowAsyncClient: + """A stand-in for ``ollama.AsyncClient``: the same reads, awaited.""" + + def __init__(self, delay: float = CHUNK_DELAY) -> None: + self.delay = delay + + async def chat(self, **_kwargs): + async def _gen(): + for text in CHUNKS: + await asyncio.sleep(self.delay) + yield _chunk(text) + + return _gen() + + +class _MixedSyncClient: + """Sync client yielding empty content, to pin the chunk contract.""" + + def chat(self, **_kwargs): + for text in ("a", "", "b"): + yield _chunk(text) + + +class _MixedAsyncClient: + """Async twin of :class:`_MixedSyncClient`.""" + + async def chat(self, **_kwargs): + async def _gen(): + for text in ("a", "", "b"): + await asyncio.sleep(0) + yield _chunk(text) + + return _gen() + + +def _client() -> LLMClient: + return LLMClient(LLMConfig(provider="ollama", model="qwen2.5-coder:7b")) + + +def _patch_clients( + monkeypatch: pytest.MonkeyPatch, sync_client, async_client +) -> None: + """Point both client factories at fakes, whichever one the code reaches for.""" + monkeypatch.setattr(LLMClient, "_ollama_client", lambda self: sync_client) + monkeypatch.setattr( + LLMClient, "_ollama_async_client", lambda self: async_client, raising=False + ) + + +async def _consume_with_ticker(llm: LLMClient) -> tuple[list[str], int]: + """Consume the stream while a second coroutine counts its turns.""" + ticks = 0 + done = False + + async def ticker() -> None: + nonlocal ticks + while not done: + ticks += 1 + await asyncio.sleep(TICK) + + task = asyncio.create_task(ticker()) + parts = [ + text + async for text in llm._ollama_stream([AgentMessage(role="user", content="hi")]) + ] + done = True + await task + return parts, ticks + + +async def test_ollama_stream_leaves_the_event_loop_responsive( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A coroutine on the same loop keeps running while the stream is consumed.""" + _patch_clients(monkeypatch, _SlowSyncClient(), _SlowAsyncClient()) + + parts, ticks = await _consume_with_ticker(_client()) + + assert parts == list(CHUNKS) + # Streaming takes ~0.15s and the ticker runs every 0.01s, so a free loop + # gives double-digit turns. A loop that is blocked by the sync reads gives + # zero -- the ticker is created but never scheduled. + assert ticks >= 3, f"event loop was blocked during the stream ({ticks} ticks)" + + +async def test_ollama_stream_yields_every_non_empty_chunk( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The chunk contract is unchanged: text in order, empty content dropped.""" + _patch_clients(monkeypatch, _MixedSyncClient(), _MixedAsyncClient()) + + parts = [ + text + async for text in _client()._ollama_stream( + [AgentMessage(role="user", content="hi")] + ) + ] + + assert parts == ["a", "b"]