diff --git a/docs/python-api.md b/docs/python-api.md index 8d437a0d..0450031a 100644 --- a/docs/python-api.md +++ b/docs/python-api.md @@ -127,16 +127,17 @@ model = llm.get_async_model("gpt-4o") You can then run a prompt using `await model.prompt(...)`: ```python -result = await model.prompt( +response = await model.prompt( "Five surprising names for a pet pelican" ) +print(await response.text()) ``` Or use `async for chunk in ...` to stream the response as it is generated: ```python async for chunk in model.prompt( "Five surprising names for a pet pelican" ): - print(chunk, end="") + print(chunk, end="", flush=True) ``` ## Conversations diff --git a/llm/__init__.py b/llm/__init__.py index 16ff9aed..d6df280f 100644 --- a/llm/__init__.py +++ b/llm/__init__.py @@ -5,6 +5,7 @@ ) from .models import ( AsyncModel, + AsyncResponse, Attachment, Conversation, Model, @@ -31,6 +32,7 @@ "get_model", "get_key", "user_dir", + "AsyncResponse", "Attachment", "Collection", "Conversation", diff --git a/llm/models.py b/llm/models.py index 82777edb..cb9c7ab3 100644 --- a/llm/models.py +++ b/llm/models.py @@ -391,6 +391,7 @@ async def _force(self): if not self._done: async for _ in self: pass + return self async def text(self) -> str: await self._force() @@ -409,7 +410,7 @@ async def datetime_utc(self) -> str: return self._start_utcnow.isoformat() if self._start_utcnow else "" def __await__(self): - return self.text().__await__() + return self._force().__await__() @classmethod def fake( diff --git a/tests/test_async.py b/tests/test_async.py index db4cf529..a84dd97d 100644 --- a/tests/test_async.py +++ b/tests/test_async.py @@ -1,3 +1,4 @@ +import llm import pytest @@ -10,5 +11,7 @@ async def test_async_model(async_mock_model): assert gathered == ["hello world"] # Not as an iterator async_mock_model.enqueue(["hello world"]) - text = await async_mock_model.prompt("hello") + response = await async_mock_model.prompt("hello") + text = await response.text() assert text == "hello world" + assert isinstance(response, llm.AsyncResponse)