mirror of
https://github.com/ggozad/oterm.git
synced 2026-10-21 22:51:26 +02:00
61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
import pytest
|
|
from ollama import ResponseError
|
|
|
|
from oterm.ollamaclient import OllamaLLM
|
|
from oterm.tools.date_time import DateTimeTool
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generate(default_model):
|
|
llm = OllamaLLM(model=default_model)
|
|
res = ""
|
|
async for _, text in llm.stream(prompt="Please add 42 and 42"):
|
|
res = text
|
|
assert "84" in res
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_llm_context(default_model):
|
|
llm = OllamaLLM(model=default_model)
|
|
async for _, _ in llm.stream("I am testing oterm, a python client for Ollama."):
|
|
pass
|
|
# There should now be a context saved for the conversation.
|
|
res = ""
|
|
async for _, text in llm.stream("Do you remember what I am testing?"):
|
|
res = text
|
|
assert "oterm" in res.lower()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_multi_modal_llm(llama_image):
|
|
llm = OllamaLLM(model="llava")
|
|
res = ""
|
|
async for _, text in llm.stream("Describe this image", images=[llama_image]):
|
|
res = text
|
|
assert "llama" in res or "animal" in res
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_errors():
|
|
llm = OllamaLLM(model="non-existent-model")
|
|
try:
|
|
async for _, _ in llm.stream("This should fail."):
|
|
pass
|
|
except ResponseError as e:
|
|
assert 'model "non-existent-model" not found' in str(e)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_tool_streaming(default_model):
|
|
llm = OllamaLLM(
|
|
model=default_model,
|
|
tool_defs=[
|
|
{"tool": DateTimeTool, "callable": lambda: "2025-01-01"},
|
|
],
|
|
)
|
|
response = ""
|
|
async for _, text in llm.stream(
|
|
"What is the current date in YYYY-MM-DD format?. Use the date_time tool to answer."
|
|
):
|
|
response = text
|
|
assert "2025-01-01" in response or ("January" in response and "2025" in response)
|