llm/SIA
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fixed context usage calculation

Browse Source
This commit is contained in:
Niels Geens
2024-11-14 18:23:33 +01:00
parent ede0a642d3
commit 4ce421bbce
13 changed files with 249 additions and 119 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
test/base_agent_test.py
View File

@@ -49,13 +49,14 @@ class MockEntry(Entry):
class TestBaseAgent(BaseAgent):
"""Concrete implementation of BaseAgent for testing."""
def run(self) -> None:
pass
pass
class BaseAgentTest(unittest.TestCase):
def setUp(self):
"""Set up test cases with mocked components."""
self.mock_llm = Mock(spec=LlmEngine)
self.mock_llm.token_count.return_value = 100
self.mock_llm.token_limit.return_value = 1000
self.mock_validator = Mock(spec=XMLValidator)
self.io_buffer = WebIOBuffer()
self.working_memory = WorkingMemory()
@@ -75,9 +76,10 @@ class BaseAgentTest(unittest.TestCase):
# Create test agent
self.agent = TestBaseAgent(
system_prompt="test prompt",
action_schema="test schema",
working_memory=self.working_memory,
system_metrics=self.mock_metrics,
metrics=self.mock_metrics,
llm=self.mock_llm,
validator=self.mock_validator,
parser=self.parser
@@ -130,7 +132,7 @@ class BaseAgentTest(unittest.TestCase):
self.assertEqual(root.get("disk_total"), "10000")
# Check context size is 0 (empty memory)
self.assertEqual(root.get("context"), "0")
self.assertEqual(root.get("context"), "10.0")
# Check no memory entries
self.assertEqual(len(list(root)), 0)
@@ -144,7 +146,7 @@ class BaseAgentTest(unittest.TestCase):
root = ET.fromstring(context)
# Check context size reflects one entry
self.assertEqual(root.get("context"), "1")
self.assertEqual(root.get("context"), "10.0")
# Verify entry content
reasoning_elem = root.find("reasoning")
@@ -166,7 +168,7 @@ class BaseAgentTest(unittest.TestCase):
root = ET.fromstring(context)
# Check context size reflects three entries
self.assertEqual(root.get("context"), "3")
self.assertEqual(root.get("context"), "10.0")
# Verify entry order maintained
reasoning_elems = root.findall("reasoning")

29
test/local_llm_engine_test.py
View File

@@ -7,20 +7,43 @@ from . import test_data
from sia.llm_engine import LlmEngine
from sia.local_llm_engine import LocalLlmEngine
class LlmEngineTest(unittest.TestCase):
class LocalLlmEngineTest(unittest.TestCase):
def setUp(self):
self.model_path = "/root/model"
def test_initialization(self):
llm_engine = LocalLlmEngine(self.model_path)
llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
self.assertIsInstance(llm_engine, LlmEngine)
def test_infer(self):
main_context = "This is a test"
llm_engine = LocalLlmEngine(self.model_path)
llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
tokens = llm_engine.infer(test_data.echo_system_prompt, main_context)
print_tokens, result_tokens = tee(tokens)
for token in print_tokens:
print(token, end="", flush=True)
result = ''.join(result_tokens)
self.assertEqual(result, f"{main_context}<test_tag>{main_context}</test_tag>")
def test_token_count(self):
"""Test token counting returns reasonable numbers"""
llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
# Test with short inputs
count = llm_engine.token_count("Short prompt", "Brief context")
self.assertGreater(count, 5)
self.assertLess(count, 50)
# Test with longer inputs
long_prompt = "A detailed system prompt with multiple sentences. " * 5
long_context = "An extensive context containing various details. " * 10
count = llm_engine.token_count(long_prompt, long_context)
self.assertGreater(count, 100)
self.assertLess(count, 1000)
def test_token_limit(self):
"""Test token limit is within expected range for modern LLMs"""
llm_engine = LocalLlmEngine(self.model_path, 0.2, 1024)
limit = llm_engine.token_limit()
self.assertGreaterEqual(limit, 2048)
self.assertLessEqual(limit, 1048576)

4
test/web_agent_test.py
View File

@@ -60,13 +60,15 @@ class WebAgentTest(unittest.TestCase):
yield "<reasoning>test reasoning</reasoning>"
self.mock_llm.infer.side_effect = mock_infer
self.mock_llm.token_count.return_value = 100
self.mock_llm.token_limit.return_value = 1000
# Create agent with all components
self.agent = WebAgent(
system_prompt="test prompt",
action_schema="test schema",
working_memory=self.working_memory,
system_metrics=self.mock_metrics,
metrics=self.mock_metrics,
llm=self.mock_llm,
validator=self.mock_validator,
parser=self.parser

8
test/web_socket_manager_test.py
View File

@@ -31,7 +31,9 @@ class WebSocketManagerTest(AioHTTPTestCase):
"memory_used": 1000,
"memory_total": 2000,
"disk_used": 5000,
"disk_total": 10000
"disk_total": 10000,
"token_count": 100,
"token_limit": 1000
}
# Create minimal mocks for other components
@@ -41,6 +43,8 @@ class WebSocketManagerTest(AioHTTPTestCase):
self.mock_llm.infer.side_effect = mock_infer
self.mock_validator = Mock(spec=XMLValidator)
self.mock_validator.validate.return_value = None
self.mock_llm.token_count.return_value = 100
self.mock_llm.token_limit.return_value = 1000
# Create parser with real IO buffer
self.parser = ResponseParser(self.io_buffer)
@@ -50,7 +54,7 @@ class WebSocketManagerTest(AioHTTPTestCase):
system_prompt="test prompt",
action_schema="test schema",
working_memory=self.working_memory,
system_metrics=self.mock_metrics,
metrics=self.mock_metrics,
llm=self.mock_llm,
validator=self.mock_validator,
parser=self.parser