Support text = await async_mock_model.prompt("hello")

simonw · Nov 13, 2024 · 1c83a4e · 1c83a4e
1 parent 9b1e720
commit 1c83a4e
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/llm/models.py b/llm/models.py
@@ -416,6 +416,9 @@ async def datetime_utc(self) -> str:
         await self._force()
         return self._start_utcnow.isoformat() if self._start_utcnow else ""
 
+    def __await__(self):
+        return self.text().__await__()
+
     @classmethod
     def fake(
         cls,

diff --git a/tests/test_async.py b/tests/test_async.py
@@ -8,3 +8,7 @@ async def test_async_model(async_mock_model):
     async for chunk in async_mock_model.prompt("hello"):
         gathered.append(chunk)
     assert gathered == ["hello world"]
+    # Not as an iterator
+    async_mock_model.enqueue(["hello world"])
+    text = await async_mock_model.prompt("hello")
+    assert text == "hello world"