237 lines
7.4 KiB
Python
237 lines
7.4 KiB
Python
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() |