93 lines
2.6 KiB
Python
93 lines
2.6 KiB
Python
"""Tests for steward.llm.client."""
|
|
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from steward.config import Settings
|
|
from steward.llm.client import LLMClient
|
|
|
|
|
|
def _make_settings(**kwargs) -> Settings:
|
|
defaults = dict(
|
|
telegram_bot_token="test-token",
|
|
openai_api_key="test-key",
|
|
openai_model="gpt-4o",
|
|
openai_system_prompt="You are Steward.",
|
|
)
|
|
defaults.update(kwargs)
|
|
return Settings(**defaults)
|
|
|
|
|
|
@pytest.fixture
|
|
def settings():
|
|
return _make_settings()
|
|
|
|
|
|
@pytest.fixture
|
|
def llm_client(settings):
|
|
return LLMClient(settings)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_sends_correct_messages(llm_client):
|
|
"""chat() should prepend the system prompt and append the user message."""
|
|
mock_response = MagicMock()
|
|
mock_response.choices[0].message.content = "Hello from the LLM!"
|
|
|
|
with patch.object(
|
|
llm_client._client.chat.completions,
|
|
"create",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_response,
|
|
) as mock_create:
|
|
result = await llm_client.chat("Hi there")
|
|
|
|
assert result == "Hello from the LLM!"
|
|
call_kwargs = mock_create.call_args.kwargs
|
|
messages = call_kwargs["messages"]
|
|
assert messages[0]["role"] == "system"
|
|
assert messages[-1]["role"] == "user"
|
|
assert messages[-1]["content"] == "Hi there"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_with_history(llm_client):
|
|
"""chat() should include history between system and user messages."""
|
|
history = [
|
|
{"role": "user", "content": "previous question"},
|
|
{"role": "assistant", "content": "previous answer"},
|
|
]
|
|
mock_response = MagicMock()
|
|
mock_response.choices[0].message.content = "reply"
|
|
|
|
with patch.object(
|
|
llm_client._client.chat.completions,
|
|
"create",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_response,
|
|
) as mock_create:
|
|
await llm_client.chat("new question", history=history)
|
|
|
|
messages = mock_create.call_args.kwargs["messages"]
|
|
roles = [m["role"] for m in messages]
|
|
assert roles == ["system", "user", "assistant", "user"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_chat_custom_system_prompt(llm_client):
|
|
"""chat() should use a custom system prompt when provided."""
|
|
mock_response = MagicMock()
|
|
mock_response.choices[0].message.content = "ok"
|
|
|
|
with patch.object(
|
|
llm_client._client.chat.completions,
|
|
"create",
|
|
new_callable=AsyncMock,
|
|
return_value=mock_response,
|
|
) as mock_create:
|
|
await llm_client.chat("msg", system_prompt="Custom prompt")
|
|
|
|
messages = mock_create.call_args.kwargs["messages"]
|
|
assert messages[0]["content"] == "Custom prompt"
|