Getting started
This commit is contained in:
Moishe Lettvin
179
tests/test_asyncprocessor.py
Normal file
179
tests/test_asyncprocessor.py
Normal file
@@ -0,0 +1,179 @@
|
||||
import time
|
||||
import unittest
|
||||
|
||||
from queue import Queue, Empty
|
||||
from threading import Thread, Event
|
||||
from typing import Generator
|
||||
|
||||
from services.ai_services import LLMService, TTSService, ImageGenService
|
||||
from message_handler.message_handler import MessageHandler
|
||||
from async_processor.async_processor import (
|
||||
AsyncProcessor,
|
||||
AIServiceConfig,
|
||||
AsyncProcessorState,
|
||||
Response
|
||||
)
|
||||
|
||||
class MockTTSService(TTSService):
|
||||
def run_tts(self, sentence):
|
||||
for word in sentence.split(' '):
|
||||
time.sleep(0.1)
|
||||
yield bytes(word, "utf-8")
|
||||
|
||||
class MockLLMService(LLMService):
|
||||
def run_llm_async(self, messages) -> Generator[str, None, None]:
|
||||
for i in ["Hello ", "there.", "How are ", "you?", "I ", "hope ", "you ", "are ", "well."]:
|
||||
time.sleep(0.1)
|
||||
yield i
|
||||
|
||||
|
||||
class MockImageService(ImageGenService):
|
||||
def run_image_gen(self, sentence) -> None:
|
||||
return None
|
||||
|
||||
class TestResponse(unittest.TestCase):
|
||||
def test_base_state_transitions(self):
|
||||
mock_tts_service = MockTTSService()
|
||||
mock_llm_service = MockLLMService()
|
||||
mock_image_service = MockImageService()
|
||||
processor = AsyncProcessor(AIServiceConfig(tts=mock_tts_service, llm=mock_llm_service, image=mock_image_service))
|
||||
processor.prepare()
|
||||
processor.play()
|
||||
processor.finalize()
|
||||
self.assertEqual(processor.state, AsyncProcessorState.FINALIZED)
|
||||
|
||||
def test_state_transitions(self):
|
||||
output_queue = Queue()
|
||||
mock_tts_service = MockTTSService()
|
||||
mock_llm_service = MockLLMService()
|
||||
mock_image_service = MockImageService()
|
||||
message_handler = MessageHandler("Hello World")
|
||||
processor = Response(
|
||||
AIServiceConfig(
|
||||
tts=mock_tts_service, llm=mock_llm_service, image=mock_image_service
|
||||
),
|
||||
message_handler,
|
||||
output_queue,
|
||||
)
|
||||
processor.prepare()
|
||||
processor.play()
|
||||
|
||||
# Consume the output from the output queue. It's necessary to mark these tasks as done for the
|
||||
# play function to return.
|
||||
expected_words = ["Hello", "there.", "How", "are", "you?", "I", "hope", "you", "are", "well."]
|
||||
|
||||
# remove the "start_stream" message from the queue
|
||||
output_queue.get()
|
||||
output_queue.task_done()
|
||||
|
||||
while expected_words:
|
||||
# get the corresponding video frame off the queue.
|
||||
output_queue.get()
|
||||
output_queue.task_done()
|
||||
|
||||
actual_word = output_queue.get()
|
||||
word = expected_words.pop(0)
|
||||
self.assertEqual(actual_word['type'], 'audio_frame')
|
||||
self.assertEqual(actual_word['data'], bytes(word, "utf-8"))
|
||||
output_queue.task_done()
|
||||
|
||||
processor.finalize()
|
||||
|
||||
self.assertEqual(processor.state, AsyncProcessorState.FINALIZED)
|
||||
|
||||
def test_interrupt_preparation(self):
|
||||
output_queue = Queue()
|
||||
mock_tts_service = MockTTSService()
|
||||
mock_llm_service = MockLLMService()
|
||||
mock_image_service = MockImageService()
|
||||
message_handler = MessageHandler("System Message")
|
||||
processor = Response(
|
||||
AIServiceConfig(
|
||||
tts=mock_tts_service, llm=mock_llm_service, image=mock_image_service
|
||||
),
|
||||
message_handler,
|
||||
output_queue,
|
||||
)
|
||||
processor.prepare()
|
||||
interrupt_request_at = time.perf_counter()
|
||||
processor.interrupt()
|
||||
processor.finalize()
|
||||
finalized_at = time.perf_counter()
|
||||
self.assertTrue(0.1 < finalized_at - interrupt_request_at < 0.2)
|
||||
print(f"delta: {interrupt_request_at, finalized_at}")
|
||||
self.assertEqual(processor.state, AsyncProcessorState.FINALIZED)
|
||||
|
||||
def test_interrupt_play(self):
|
||||
output_queue = Queue()
|
||||
mock_tts_service = MockTTSService()
|
||||
mock_llm_service = MockLLMService()
|
||||
mock_image_service = MockImageService()
|
||||
message_handler = MessageHandler("System Message")
|
||||
processor = Response(
|
||||
AIServiceConfig(
|
||||
tts=mock_tts_service, llm=mock_llm_service, image=mock_image_service
|
||||
),
|
||||
message_handler,
|
||||
output_queue,
|
||||
)
|
||||
processor.prepare()
|
||||
processor.play()
|
||||
|
||||
stop_processing_output_queue = Event()
|
||||
def process_output_queue_async():
|
||||
# Consume the output from the output queue. It's necessary to mark these tasks as done for the
|
||||
# play function to return.
|
||||
time.sleep(0.1)
|
||||
expected_words = ["Hello", "there.", "How", "are", "you?", "I", "hope", "you", "are", "well."]
|
||||
while expected_words and not stop_processing_output_queue.is_set():
|
||||
try:
|
||||
actual_word = output_queue.get_nowait()
|
||||
if actual_word['type'] == 'audio_frame':
|
||||
time.sleep(0.1)
|
||||
word = expected_words.pop(0)
|
||||
self.assertEqual(actual_word['type'], 'audio_frame')
|
||||
self.assertEqual(actual_word['data'], bytes(word, "utf-8"))
|
||||
output_queue.task_done()
|
||||
except Empty:
|
||||
pass
|
||||
|
||||
process_output_queue = Thread(target=process_output_queue_async, daemon=True)
|
||||
process_output_queue.start()
|
||||
|
||||
time.sleep(0.5)
|
||||
processor.interrupt()
|
||||
|
||||
stop_processing_output_queue.set()
|
||||
process_output_queue.join()
|
||||
|
||||
processor.finalize()
|
||||
self.assertEqual(processor.state, AsyncProcessorState.FINALIZED)
|
||||
|
||||
def test_statechange_callback(self):
|
||||
mock_tts_service = MockTTSService()
|
||||
mock_llm_service = MockLLMService()
|
||||
mock_image_service = MockImageService()
|
||||
processor = AsyncProcessor(
|
||||
AIServiceConfig(
|
||||
tts=mock_tts_service, llm=mock_llm_service, image=mock_image_service
|
||||
)
|
||||
)
|
||||
is_finalized = False
|
||||
def set_is_finalized(async_processor:AsyncProcessor):
|
||||
nonlocal is_finalized
|
||||
is_finalized = True
|
||||
|
||||
processor.set_state_callback(
|
||||
AsyncProcessorState.FINALIZED, set_is_finalized
|
||||
)
|
||||
processor.prepare()
|
||||
self.assertFalse(is_finalized)
|
||||
processor.play()
|
||||
self.assertFalse(is_finalized)
|
||||
processor.finalize()
|
||||
self.assertTrue(is_finalized)
|
||||
self.assertEqual(processor.state, AsyncProcessorState.FINALIZED)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
129
tests/test_message_handler.py
Normal file
129
tests/test_message_handler.py
Normal file
@@ -0,0 +1,129 @@
|
||||
import time
|
||||
import unittest
|
||||
|
||||
from unittest.mock import MagicMock, call
|
||||
|
||||
from message_handler.message_handler import MessageHandler, IndexingMessageHandler
|
||||
from services.ai_services import AIService, AIServiceConfig
|
||||
from storage.search import SearchIndexer
|
||||
|
||||
|
||||
class TestMessageHandler(unittest.TestCase):
|
||||
def test_simple_intro(self):
|
||||
message_handler = MessageHandler("Hello world")
|
||||
self.assertEqual(
|
||||
message_handler.get_llm_messages(),
|
||||
[{"role": "system", "content": "Hello world"}],
|
||||
)
|
||||
|
||||
def test_simple_user_message(self):
|
||||
message_handler = MessageHandler("System prompt")
|
||||
message_handler.add_user_message("User message")
|
||||
self.assertEqual(
|
||||
message_handler.get_llm_messages(),
|
||||
[
|
||||
{"role": "system", "content": "System prompt"},
|
||||
{"role": "user", "content": "User message"},
|
||||
],
|
||||
)
|
||||
|
||||
def test_simple_user_and_assistant_message(self):
|
||||
message_handler = MessageHandler("System prompt")
|
||||
message_handler.add_user_message("User message")
|
||||
message_handler.add_assistant_message("Assistant message")
|
||||
self.assertEqual(
|
||||
message_handler.get_llm_messages(),
|
||||
[
|
||||
{"role": "system", "content": "System prompt"},
|
||||
{"role": "user", "content": "User message"},
|
||||
{"role": "assistant", "content": "Assistant message"},
|
||||
],
|
||||
)
|
||||
|
||||
def test_user_message_overwrite(self):
|
||||
message_handler = MessageHandler("System prompt")
|
||||
message_handler.add_user_message("User message")
|
||||
message_handler.add_assistant_message("Assistant message")
|
||||
message_handler.add_user_message("User message plus something else")
|
||||
self.assertEqual(
|
||||
message_handler.get_llm_messages(),
|
||||
[
|
||||
{"role": "system", "content": "System prompt"},
|
||||
{"role": "user", "content": "User message plus something else"},
|
||||
],
|
||||
)
|
||||
|
||||
def test_user_message_after_assistant(self):
|
||||
message_handler = MessageHandler("System prompt")
|
||||
message_handler.add_user_message("User message")
|
||||
message_handler.add_assistant_message("Assistant message")
|
||||
message_handler.add_user_message("other user message")
|
||||
self.assertEqual(
|
||||
message_handler.get_llm_messages(),
|
||||
[
|
||||
{"role": "system", "content": "System prompt"},
|
||||
{"role": "user", "content": "User message"},
|
||||
{"role": "assistant", "content": "Assistant message"},
|
||||
{"role": "user", "content": "other user message"},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class MockAIService(AIService):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
def run_llm(self, messages, latest_user_message=None, stream=True):
|
||||
return {"choices": [{"message": {"content": "Parsed user message."}}]}
|
||||
|
||||
|
||||
class TestIndexingMessageHandler(unittest.TestCase):
|
||||
def test_user_message_finalized(self):
|
||||
mock_ai_service = MockAIService()
|
||||
service_config = AIServiceConfig(
|
||||
mock_ai_service, mock_ai_service, mock_ai_service
|
||||
)
|
||||
|
||||
mock_indexer = MagicMock(spec=SearchIndexer)
|
||||
|
||||
message_handler = IndexingMessageHandler(
|
||||
"Hello world", "story_id", service_config, mock_indexer
|
||||
)
|
||||
message_handler.add_user_message("User message")
|
||||
message_handler.add_assistant_message("Assistant message will be ignored")
|
||||
message_handler.add_user_message("User message plus something else")
|
||||
message_handler.finalize_user_message()
|
||||
message_handler.add_assistant_message(
|
||||
"New assistant message will not be ignored"
|
||||
)
|
||||
message_handler.add_user_message("User message second time")
|
||||
message_handler.add_assistant_message("Assistant message second time")
|
||||
message_handler.write_messages_to_index()
|
||||
|
||||
time.sleep(0.5)
|
||||
|
||||
self.assertEqual(
|
||||
mock_indexer.mock_calls,
|
||||
[
|
||||
call.index_text("Parsed user message."),
|
||||
call.index_text("New assistant message will not be ignored"),
|
||||
],
|
||||
)
|
||||
|
||||
mock_indexer.reset_mock()
|
||||
|
||||
message_handler.finalize_user_message()
|
||||
|
||||
time.sleep(0.5)
|
||||
|
||||
self.assertEqual(
|
||||
mock_indexer.mock_calls,
|
||||
[
|
||||
call.index_text("Parsed user message."),
|
||||
call.index_text("Assistant message second time"),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user