Add configurable LM Studio flushing and LLM retries.
This commit is contained in:
@@ -1,5 +1,7 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from photools3.llm_session import LlmSession
|
||||
|
||||
|
||||
@@ -84,3 +86,70 @@ def test_llm_session_captions_and_keywords_mocked(test_photo):
|
||||
|
||||
# Verify that the mock was called
|
||||
assert mock_client.chat.completions.create.called
|
||||
|
||||
|
||||
def test_llm_session_retries_from_original_request_without_accumulating_messages():
|
||||
with patch("photools3.llm_session.OpenAI") as mocked_openai:
|
||||
mock_client = MagicMock()
|
||||
mocked_openai.return_value = mock_client
|
||||
mock_model = MagicMock()
|
||||
mock_model.id = "gpt-4o"
|
||||
mock_client.models.list.return_value.data = [mock_model]
|
||||
|
||||
def response(content):
|
||||
message = MagicMock()
|
||||
message.content = content
|
||||
choice = MagicMock()
|
||||
choice.message = message
|
||||
result = MagicMock()
|
||||
result.choices = [choice]
|
||||
return result
|
||||
|
||||
mock_client.chat.completions.create.side_effect = [
|
||||
response('{"wrong": "first"}'),
|
||||
response('{"wrong": "second"}'),
|
||||
response('{"text": "correct"}'),
|
||||
]
|
||||
|
||||
session = LlmSession(system_prompt="You are a helper.")
|
||||
result = session.call(photos=[], prompt="Return JSON")
|
||||
|
||||
assert result.text == "correct"
|
||||
requests = mock_client.chat.completions.create.call_args_list
|
||||
messages_per_attempt = [call.kwargs["messages"] for call in requests]
|
||||
assert [len(messages) for messages in messages_per_attempt] == [2, 3, 3]
|
||||
final_correction = messages_per_attempt[2][-1]["content"]
|
||||
assert '"wrong": "second"' in final_correction
|
||||
assert '"wrong": "first"' not in final_correction
|
||||
|
||||
|
||||
def test_llm_session_uses_configured_retry_count(monkeypatch):
|
||||
monkeypatch.setattr("photools3.llm_session.settings.llm_max_retries", 1)
|
||||
|
||||
with patch("photools3.llm_session.OpenAI") as mocked_openai:
|
||||
mock_client = MagicMock()
|
||||
mocked_openai.return_value = mock_client
|
||||
mock_model = MagicMock()
|
||||
mock_model.id = "gpt-4o"
|
||||
mock_client.models.list.return_value.data = [mock_model]
|
||||
|
||||
def response(content):
|
||||
message = MagicMock()
|
||||
message.content = content
|
||||
choice = MagicMock()
|
||||
choice.message = message
|
||||
result = MagicMock()
|
||||
result.choices = [choice]
|
||||
return result
|
||||
|
||||
mock_client.chat.completions.create.side_effect = [
|
||||
response('{"wrong": "first"}'),
|
||||
response('{"wrong": "second"}'),
|
||||
response('{"text": "would be a third call"}'),
|
||||
]
|
||||
session = LlmSession(system_prompt="You are a helper.")
|
||||
|
||||
with pytest.raises(ValueError, match="after 1 retries"):
|
||||
session.call(photos=[])
|
||||
|
||||
assert mock_client.chat.completions.create.call_count == 2
|
||||
|
||||
Reference in New Issue
Block a user