New web interface, move llm engine to separate process
This commit is contained in:
237
test/llm_engine_test.py
Normal file
237
test/llm_engine_test.py
Normal file
@@ -0,0 +1,237 @@
|
||||
import unittest
|
||||
import os
|
||||
import tempfile
|
||||
import time
|
||||
from xml.etree import ElementTree as ET
|
||||
|
||||
from sia.llm_engine import LlmEngine
|
||||
|
||||
class TestLlmEngine(unittest.TestCase):
|
||||
def setUp(self):
|
||||
"""Set up test environment before each test."""
|
||||
# Create a temporary bash script to simulate the LLM engine
|
||||
self.temp_script = self._create_mock_llm_script()
|
||||
os.chmod(self.temp_script, 0o755) # Make it executable
|
||||
|
||||
# Initialize the LlmEngine with our mock script
|
||||
self.engine = LlmEngine(self.temp_script, "/root/sia/action_schema.xsd")
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up after the test."""
|
||||
if hasattr(self, 'engine') and self.engine:
|
||||
self.engine._terminate_process()
|
||||
|
||||
def _create_mock_llm_script(self):
|
||||
"""Create a bash script that simulates an LLM engine."""
|
||||
fd, path = tempfile.mkstemp(suffix='.sh')
|
||||
with os.fdopen(fd, 'w') as f:
|
||||
# Paste the content from the mock_llm_script_debug artifact
|
||||
# This is a placeholder - you'll replace this with the actual script content
|
||||
f.write('''#!/bin/bash
|
||||
|
||||
# Mock LLM engine for testing
|
||||
# This script simulates an LLM engine subprocess that responds to the three commands:
|
||||
# <token_limit/>, <token_count>...</token_count>, and <infer_xml>...</infer_xml>
|
||||
|
||||
# Function to read XML input until a complete closing tag is found
|
||||
read_xml_input() {
|
||||
local input=""
|
||||
local line
|
||||
local char_count=0
|
||||
|
||||
# Read until we get a complete XML command
|
||||
while IFS= read -r line; do
|
||||
input="${input}${line}"
|
||||
char_count=$((char_count + ${#line}))
|
||||
|
||||
# Debug the actual content (first 30 chars)
|
||||
|
||||
if [[ "$input" == *"<token_limit/>"* ]]; then
|
||||
printf "1024"
|
||||
printf "\004" # EOT character (hex 04)
|
||||
return
|
||||
fi
|
||||
|
||||
if [[ "$input" == *"<token_count>"*"</token_count>"* ]]; then
|
||||
printf "405"
|
||||
printf "\004" # EOT character (hex 04)
|
||||
return
|
||||
fi
|
||||
|
||||
if [[ "$input" == *"<infer_xml>"*"</infer_xml>"* ]]; then
|
||||
generate_response
|
||||
return
|
||||
fi
|
||||
|
||||
done
|
||||
}
|
||||
|
||||
# Function to generate a response token by token
|
||||
generate_response() {
|
||||
printf "<"
|
||||
sleep 0.01
|
||||
|
||||
printf "reason"
|
||||
sleep 0.01
|
||||
|
||||
printf "ing"
|
||||
sleep 0.01
|
||||
|
||||
printf ">"
|
||||
sleep 0.01
|
||||
|
||||
printf "This"
|
||||
sleep 0.01
|
||||
|
||||
printf " is"
|
||||
sleep 0.01
|
||||
|
||||
printf " a"
|
||||
sleep 0.01
|
||||
|
||||
printf " test"
|
||||
sleep 0.01
|
||||
|
||||
printf " response."
|
||||
sleep 0.01
|
||||
|
||||
printf "</"
|
||||
sleep 0.01
|
||||
|
||||
printf "reason"
|
||||
sleep 0.01
|
||||
|
||||
printf "ing"
|
||||
sleep 0.01
|
||||
|
||||
printf ">"
|
||||
sleep 0.01
|
||||
|
||||
printf "\004"
|
||||
}
|
||||
|
||||
# Main loop - keep reading input and responding
|
||||
iteration=0
|
||||
while true; do
|
||||
iteration=$((iteration + 1))
|
||||
read_xml_input
|
||||
done''')
|
||||
return path
|
||||
|
||||
def test_token_limit(self):
|
||||
"""Test retrieving the token limit from the LLM engine."""
|
||||
limit = self.engine.token_limit()
|
||||
self.assertEqual(limit, 1024, "Token limit should be 1024")
|
||||
|
||||
def test_token_count(self):
|
||||
"""Test counting tokens in a prompt."""
|
||||
system_prompt = "You are a helpful AI assistant."
|
||||
|
||||
# Create a simple context ElementTree
|
||||
context_et = ET.Element("context")
|
||||
context_et.set("time", "2024-10-18T12:00:00Z")
|
||||
context_et.text = "Test context"
|
||||
|
||||
prefix = ""
|
||||
|
||||
count = self.engine.token_count(system_prompt, context_et, prefix)
|
||||
self.assertEqual(count, 405, "Token count should be 405")
|
||||
|
||||
def test_inference_token_by_token(self):
|
||||
"""Test inference with the LLM engine, ensuring tokens come one by one."""
|
||||
system_prompt = "You are a helpful AI assistant."
|
||||
|
||||
# Create a simple context ElementTree
|
||||
context_et = ET.Element("context")
|
||||
context_et.set("time", "2024-10-18T12:00:00Z")
|
||||
context_et.text = "Test context"
|
||||
|
||||
prefix = ""
|
||||
|
||||
# Get the generator for token-by-token output
|
||||
token_generator = self.engine.infer(system_prompt, context_et, prefix)
|
||||
|
||||
# Collect tokens
|
||||
response = ""
|
||||
tokens_received = 0
|
||||
|
||||
for token in token_generator:
|
||||
response += token
|
||||
tokens_received += 1
|
||||
# Each "token" should be small (in the mock script, it's up to 10 characters)
|
||||
self.assertLessEqual(len(token), 10, "Each token should be small")
|
||||
|
||||
# Verify we received multiple tokens
|
||||
self.assertGreater(tokens_received, 10, "Should receive multiple tokens")
|
||||
|
||||
expected = "<reasoning>This is a test response.</reasoning>"
|
||||
self.assertEqual(response, expected, "Inference response should match expected output")
|
||||
|
||||
def test_multiple_inferences(self):
|
||||
"""Test running multiple inference cycles one after another."""
|
||||
system_prompt = "You are a helpful AI assistant."
|
||||
|
||||
# Create a simple context ElementTree
|
||||
context_et = ET.Element("context")
|
||||
context_et.set("time", "2024-10-18T12:00:00Z")
|
||||
context_et.text = "Test context"
|
||||
|
||||
prefix = ""
|
||||
|
||||
# Run first inference
|
||||
token_generator = self.engine.infer(system_prompt, context_et, prefix)
|
||||
response = ""
|
||||
for token in token_generator:
|
||||
response += token
|
||||
|
||||
expected = "<reasoning>This is a test response.</reasoning>"
|
||||
self.assertEqual(response, expected, "First inference response should match expected output")
|
||||
|
||||
# Add a small delay between inferences to see if it helps
|
||||
time.sleep(0.1)
|
||||
|
||||
# Run second inference
|
||||
token_generator = self.engine.infer(system_prompt, context_et, prefix)
|
||||
response = ""
|
||||
for token in token_generator:
|
||||
response += token
|
||||
|
||||
self.assertEqual(response, expected, "Second inference response should match expected output")
|
||||
|
||||
# We should be able to get the token limit again after inference
|
||||
limit = self.engine.token_limit()
|
||||
self.assertEqual(limit, 1024, "Token limit should still be 1024 after multiple inferences")
|
||||
|
||||
def test_restart(self):
|
||||
"""Test restarting the LLM engine."""
|
||||
# First get a token limit
|
||||
limit = self.engine.token_limit()
|
||||
self.assertEqual(limit, 1024, "Token limit should be 1024")
|
||||
|
||||
# Restart the engine
|
||||
self.engine.restart()
|
||||
|
||||
# Get token limit again
|
||||
limit = self.engine.token_limit()
|
||||
self.assertEqual(limit, 1024, "Token limit should still be 1024 after restart")
|
||||
|
||||
# Engine should still work for inference after restart
|
||||
system_prompt = "You are a helpful AI assistant."
|
||||
|
||||
# Create a simple context ElementTree
|
||||
context_et = ET.Element("context")
|
||||
context_et.set("time", "2024-10-18T12:00:00Z")
|
||||
context_et.text = "Test context"
|
||||
|
||||
prefix = ""
|
||||
|
||||
token_generator = self.engine.infer(system_prompt, context_et, prefix)
|
||||
response = ""
|
||||
for token in token_generator:
|
||||
response += token
|
||||
|
||||
expected = "<reasoning>This is a test response.</reasoning>"
|
||||
self.assertEqual(response, expected, "Inference should work after restart")
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user