fix(byllm): async streaming awaits an async generator, so every subclass but Model crashes (#8932)

jac/jaclang/byllm/llm.impl/basellm.impl.jac

@@ -1815,25 +1815,27 @@ impl BaseLLM.model_call_no_stream_async(params: dict) -> dict {
     return await asyncio.to_thread(self.model_call_no_stream, params);
 }
 
-impl BaseLLM.model_call_with_stream_async(
-    params: dict
-) -> AsyncGenerator[object, None] {
+impl BaseLLM.model_call_with_stream_async(params: dict) -> AsyncIterator[object] {
     import asyncio;
+
     def _next_chunk(gen: Generator[object, None, None]) -> tuple[bool, object] {
         try {
             return (True, next(gen));
         } except StopIteration {
             return (False, None);
         }
     }
-    gen = self.model_call_with_stream(params);
-    while True {
-        (has_value, chunk) = await asyncio.to_thread(_next_chunk, gen);
-        if not has_value {
-            break;
+    async def _iter -> AsyncGenerator[object, None] {
+        gen = iter(await asyncio.to_thread(self.model_call_with_stream, params));
+        while True {
+            (has_value, chunk) = await asyncio.to_thread(_next_chunk, gen);
+            if not has_value {
+                break;
+            }
+            yield chunk;
         }
-        yield chunk;
     }
+    return _iter();
 }
 
 impl BaseLLM._force_final_answer(mt_run: MTRuntime) -> object {

jac/jaclang/byllm/llm.impl/mockllm.impl.jac

@@ -17,7 +17,7 @@ impl MockLLM.dispatch_no_streaming(mt_run: MTRuntime) -> CompletionResult {
     if isinstance(output, MockError) {
         raise output.error;
     }
-    message = "{'role': 'assistant', 'content': {'text': 'Mock response'}}";
+    message: any = "{'role': 'assistant', 'content': {'text': 'Mock response'}}";
     mt_run.add_message(message);
 
     _usage_keys = ["prompt_tokens", "completion_tokens", "total_tokens"];
@@ -81,6 +81,31 @@ impl MockLLM.adispatch_no_streaming(mt_run: MTRuntime) -> CompletionResult {
     return self.dispatch_no_streaming(mt_run);
 }
 
+impl MockLLM.adispatch_streaming(mt_run: MTRuntime) {
+    import asyncio;
+
+    output = self.config["outputs"].pop(0);
+    usage: dict[str, object] = {};
+    if isinstance(output, tuple) and len(output) == 2 {
+        first = output[0];
+        second = output[1];
+        if isinstance(first, str) and isinstance(second, dict) {
+            output = first;
+            usage = second;
+        }
+    }
+    if mt_run.stream {
+        while output {
+            chunk_len = random.randint(3, 10);
+            yield output[:chunk_len];
+            await asyncio.sleep(random.uniform(0.01, 0.05));
+            output = output[chunk_len:];
+        }
+    }
+
+    _record_usage(mt_run, usage, self.model_name);
+}
+
 impl MockLLM.dispatch_streaming(mt_run: MTRuntime) -> Generator[str, None, None] {
     output = self.config["outputs"].pop(0);
     usage: dict[str, object] = {};

jac/jaclang/byllm/llm.impl/model.impl.jac

@@ -111,7 +111,11 @@ impl Model.ainvoke(mt_run: MTRuntime) -> object {
     logging_stream = mt_run.stream and mt_run.call_params.get("logging", False);
     no_tools = len(mt_run.tools) == 0;
     handled_directly = no_tools and not logging_stream;
+    delegate: BaseLLM | None = self._mock_delegate or self._local_delegate;
     if not handled_directly {
+        if delegate {
+            return await delegate.ainvoke(mt_run);
+        }
         return await super.ainvoke(mt_run);
     }
 
@@ -122,7 +126,9 @@ impl Model.ainvoke(mt_run: MTRuntime) -> object {
     _start = time.time();
 
     try {
-        result = await super.ainvoke(mt_run);
+        result = await delegate.ainvoke(mt_run)
+            if delegate
+            else await super.ainvoke(mt_run);
 
         if isinstance(result, _types_mod.AsyncGeneratorType) {
             _current_invocation_id.reset(_ctx_token);

jac/jaclang/byllm/llm.jac

@@ -148,10 +148,7 @@ obj BaseLLM {
     def model_call_no_stream(params: dict) -> dict;
     def model_call_with_stream(params: dict) -> Generator[object, None, None];
     async def model_call_no_stream_async(params: dict) -> dict;
-    async def model_call_with_stream_async(
-        params: dict
-    ) -> AsyncGenerator[object, None];
-
+    async def model_call_with_stream_async(params: dict) -> AsyncIterator[object];
     def _stream_final_answer(mt_run: MTRuntime) -> Generator[str, None, None];
     def _force_final_answer(mt_run: MTRuntime) -> object;
     def _attempt_recovery(mt_run: MTRuntime) -> CompletionResult | None;
@@ -186,6 +183,7 @@ obj MockLLM(BaseLLM) {
     override def dispatch_no_streaming(mt_run: MTRuntime) -> CompletionResult;
     override async def adispatch_no_streaming(mt_run: MTRuntime) -> CompletionResult;
     override def dispatch_streaming(mt_run: MTRuntime) -> Generator[str, None, None];
+    override async def adispatch_streaming(mt_run: MTRuntime);
 }
 
 obj LocalLLM(BaseLLM) {

jac/jaclang/byllm/tests/fixtures/async_stream_by_llm.jac

@@ -0,0 +1,36 @@
+"""Fixture for jac#8932: async by llm(stream=True) over a stubbed provider.
+
+acall_llm is deliberately not mocked, so the whole chain runs.
+"""
+
+import litellm;
+import from jaclang.byllm.llm { BaseLLM }
+
+
+def _chunk(text: str) -> object {
+    return litellm.ModelResponseStream(
+        **{
+            "id": "c",
+            "object": "chat.completion.chunk",
+            "created": 1,
+            "model": "stub",
+            "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}]
+        }
+    );
+}
+
+
+obj StubStreamLLM(BaseLLM) {
+    """Implements only the sync provider call, so it inherits the base async one."""
+    override def model_call_with_stream(params: dict) -> object {
+        for word in ["streamed ", "reply"] {
+            yield _chunk(word);
+        }
+    }
+}
+
+
+glob llm = StubStreamLLM(model_name="stub-stream");
+
+
+async def ask(topic: str) -> str by llm(stream=True);

jac/jaclang/byllm/tests/support_tests.jac

@@ -308,10 +308,6 @@ obj FakeLLM(BaseLLM) {
         return self.script.with_stream(params);
     }
 
-    override async def model_call_with_stream_async(params: dict) -> any {
-        return await self.script.with_stream_async(params);
-    }
-
     def sent(key: str) -> list {
         return self.script.sent(key);
     }

jac/jaclang/byllm/tests/test_byllm.jac

@@ -2,6 +2,7 @@
 
 import contextlib;
 import io;
+import time;
 import os;
 import sys;
 import warnings;
@@ -1325,6 +1326,168 @@ test "acall_llm proxy is a coroutine function" {
     );
 }
 
+# jac#8932: adispatch_streaming awaits model_call_with_stream_async, so the base
+# must be a coroutine returning an async iterator, not an async generator.
+def _chunk_8932(text: str) -> object {
+    return litellm.ModelResponseStream(
+        **{
+            "id": "c",
+            "object": "chat.completion.chunk",
+            "created": 1,
+            "model": "fake",
+            "choices": [{"index": 0, "delta": {"content": text}, "finish_reason": None}]
+        }
+    );
+}
+
+
+"""Drain an async stream, pausing between chunks to pace the producer."""
+async def _drain_8932(stream: any, pace: float = 0.0) -> str {
+    acc: list[object] = [];
+    async for chunk in stream {
+        acc.append(chunk);
+        if pace {
+            await asyncio.sleep(pace);
+        }
+    }
+    return "".join([str(c) for c in acc]);
+}
+
+
+"""Drive one LLM down the async streaming path."""
+async def _stream_8932(llm: _byllm_llm_mod.BaseLLM, pace: float = 0.0) -> str {
+    return await _drain_8932(llm.adispatch_streaming(mk_run(stream=True)), pace);
+}
+
+
+"""Builds the whole stream before returning it, like ModelPool cost-based-routing."""
+obj _EagerStream8932(_byllm_llm_mod.BaseLLM) {
+    has build_s: float = 0.0;
+
+    override def model_call_with_stream(params: dict) -> object {
+        time.sleep(self.build_s);
+        return [_chunk_8932("eager "), _chunk_8932("list")];
+    }
+}
+
+
+"""Records when each chunk is produced, so laziness can be asserted."""
+obj _LazyStream8932(_byllm_llm_mod.BaseLLM) {
+    has produced: list[float] = [];
+
+    override def model_call_with_stream(params: dict) -> object {
+        for word in ["one ", "two ", "three"] {
+            self.produced.append(time.perf_counter());
+            yield _chunk_8932(word);
+        }
+    }
+}
+
+
+"""Stream while a ticker runs, to measure whether the loop stayed free."""
+async def _ticked_stream_8932(llm: _byllm_llm_mod.BaseLLM) -> tuple[int, str] {
+    ticks: list[int] = [0];
+    stop: list[bool] = [False];
+
+    async def _tick {
+        while not stop[0] {
+            ticks[0] += 1;
+            await asyncio.sleep(0.01);
+        }
+    }
+    task = asyncio.ensure_future(_tick());
+    text = await _stream_8932(llm);
+    stop[0] = True;
+    await task;
+    return (ticks[0], text);
+}
+
+
+test "BaseLLM.model_call_with_stream_async is a coroutine, not an async generator" {
+    fn = _byllm_llm_mod.BaseLLM.model_call_with_stream_async;
+    assert not inspect.isasyncgenfunction(fn) , "awaiting one raises TypeError";
+    assert inspect.iscoroutinefunction(fn) , "must be a coroutine, matching Model";
+}
+
+
+test "async streaming works on a subclass that inherits the base method" {
+    llm = FakeLLM(model_name="fake-8932", replies=[say("hello world")]);
+    assert asyncio.run(_stream_8932(llm)) == "hello world";
+}
+
+
+test "async streaming works on MockLLM" {
+    llm = _byllm_llm_mod.MockLLM(
+        model_name="mockllm", config={"outputs": ["hello world"]}
+    );
+    assert asyncio.run(_stream_8932(llm)) == "hello world";
+}
+
+
+test "async streaming accepts a provider call that returns a list" {
+    llm = _EagerStream8932(model_name="list-8932");
+    assert asyncio.run(_stream_8932(llm)) == "eager list";
+}
+
+
+test "async streaming keeps the loop free while an eager provider builds" {
+    llm = _EagerStream8932(model_name="eager-8932", build_s=0.3);
+    (ticks, text) = asyncio.run(_ticked_stream_8932(llm));
+    assert text == "eager list";
+    assert ticks > 10 , f"loop was blocked during the eager build: {ticks} ticks";
+}
+
+
+test "async streaming stays lazy: the producer waits for the consumer" {
+    llm = _LazyStream8932(model_name="lazy-8932");
+    assert asyncio.run(_stream_8932(llm, pace=0.12)) == "one two three";
+    gaps = [llm.produced[i] - llm.produced[i - 1] for i in range(1, len(llm.produced))];
+    assert all(g > 0.08 for g in gaps) , f"stream was buffered, not lazy: {gaps}";
+}
+
+
+# Model("mockllm") and Model("local:...") build an inner delegate and hand work
+# to it. Model.invoke did; Model.ainvoke did not, so the async side reached
+# litellm with a model name no provider knows.
+test "Model with a mock delegate streams on the async path" {
+    llm = Model(model_name="mockllm", config={"outputs": ["hello world"]});
+
+    async def _run -> str {
+        return await _drain_8932(await llm.ainvoke(mk_run(stream=True)));
+    }
+    assert asyncio.run(_run()) == "hello world";
+}
+
+
+test "Model with a mock delegate answers on the async no-stream path" {
+    llm = Model(model_name="mockllm", config={"outputs": ["hello world"]});
+
+    async def _run -> object {
+        return await llm.ainvoke(mk_run());
+    }
+    assert asyncio.run(_run()) == "hello world";
+}
+
+
+# The by llm() declaration lives in a fixture: at module level here it perturbs
+# the two tests that compile a fixture with JacProgram().
+test "async def by llm(stream=True) streams end to end" {
+    code_gen = (prog := JacProgram()).compile(
+        os.path.join(FIXTURE_DIR, "async_stream_by_llm.jac")
+    );
+    assert not prog.errors_had , f"Compilation errors: {prog.errors_had}";
+    compiled = compile(code_gen.gen.py_ast[0], filename="<ast>", mode="exec");
+    module = types.ModuleType("async_stream_by_llm_fixture");
+    module.__dict__["__file__"] = code_gen.loc.mod_path;
+    exec(compiled, module.__dict__);
+
+    async def _run -> str {
+        return await _drain_8932(await module.__dict__["ask"](topic="anything"));
+    }
+    assert asyncio.run(_run()) == "streamed reply";
+}
+
+
 test "model_call_no_stream_async raises and logs AuthenticationError" {
     log_output = io.StringIO();
     _loguru_logger.remove();
@@ -1364,6 +1527,7 @@ test "model_call_no_stream_async raises and logs AuthenticationError" {
     }
 }
 
+
 test "model_call_with_stream_async raises and logs RateLimitError" {
     log_output = io.StringIO();
     _loguru_logger.remove();

release_notes/unreleased/jaclang/8955.bugfix.md

@@ -0,0 +1 @@
+- `async def ... by llm()` now works on every LLM class: the streaming path no longer raises `TypeError: 'async_generator' object can't be awaited` on `MockLLM`, `LocalLLM`, `ModelPool` or a user-defined `BaseLLM` subclass, and `Model("mockllm")` / `Model("local:...")` now reach their delegate on the async path instead of asking a provider for a model it has never heard of.