await model.prompt() now returns AsyncResponse

Refs #613 (comment)

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

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",

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(

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)