mirror of
https://github.com/Hopiu/llm.git
synced 2026-04-05 22:10:59 +00:00
30 lines
967 B
Python
30 lines
967 B
Python
import llm
|
|
import pytest
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_model(async_mock_model):
|
|
gathered = []
|
|
async_mock_model.enqueue(["hello world"])
|
|
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"])
|
|
response = await async_mock_model.prompt("hello")
|
|
text = await response.text()
|
|
assert text == "hello world"
|
|
assert isinstance(response, llm.AsyncResponse)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_model_conversation(async_mock_model):
|
|
async_mock_model.enqueue(["joke 1"])
|
|
conversation = async_mock_model.conversation()
|
|
response = await conversation.prompt("joke")
|
|
text = await response.text()
|
|
assert text == "joke 1"
|
|
async_mock_model.enqueue(["joke 2"])
|
|
response2 = await conversation.prompt("again")
|
|
text2 = await response2.text()
|
|
assert text2 == "joke 2"
|