Begin implementation, basic inferenece

test/llm_engine_test.py

@@ -0,0 +1,25 @@
+import unittest
+
+from . import test_data
+from . import test_util
+
+from sia.llm_engine import LlmEngine, InferenceResult
+
+class LlmEngineTest(unittest.TestCase):
+    def setUp(self):
+        self.model_path = "/root/model"
+        self.llm_engine = LlmEngine(self.model_path)
+
+    def test_initialization(self):
+        llm_engine = LlmEngine(self.model_path)
+        self.assertIsInstance(llm_engine, LlmEngine)
+
+    def test_infer(self):
+        main_context = "This is a test"
+        llm_engine = LlmEngine(self.model_path)
+        result = llm_engine.infer(test_data.echo_system_prompt, main_context, test_data.echo_action_schema)
+        self.assertIsInstance(result, InferenceResult)
+        self.assertIsInstance(result.reasoning, str)
+        self.assertIsInstance(result.actions, str)
+        self.assertEqual(result.reasoning, main_context)
+        self.assertEqual(result.actions, f"<test_tag>{main_context}</test_tag>")

test/test_data.py

@@ -0,0 +1,32 @@
+echo_system_prompt = """
+Your answer always consists of 4 parts:
+    - The original request
+    - <test_tag>
+    - The original request
+    - </test_tag>
+    You never provide an answer.
+    Only the entered text, the xml open tag, the same text again and the xml close tag.
+    Don't add whitespace or newlines.
+    Don't modify the text or change casing.
+    Be exact, this is for testing purposes.
+
+    example input:
+    hello
+    example output:
+    hello<test_tag>hello</test_tag>
+""".strip()
+
+echo_action_schema = """
+<?xml version="1.0" encoding="UTF-8"?>
+<xs:schema xmlns:xs="http://www.w3.org/2001/XMLSchema">
+<xs:element name="test_tag">
+    <xs:complexType>
+        <xs:simpleContent>
+            <xs:extension base="xs:string">
+                <xs:attribute name="id" type="xs:string" use="required"/>
+            </xs:extension>
+        </xs:simpleContent>
+    </xs:complexType>
+</xs:element>
+</xs:schema>
+""".strip()

test/test_util.py

@@ -0,0 +1,8 @@
+import unittest
+
+class SequentialTestLoader(unittest.TestLoader):
+    def getTestCaseNames(self, testCaseClass):
+        test_names = super().getTestCaseNames(testCaseClass)
+        testcase_methods = list(testCaseClass.__dict__.keys())
+        test_names.sort(key=testcase_methods.index)
+        return test_names

test/util_test.py

@@ -0,0 +1,17 @@
+import unittest
+
+from . import test_data
+
+from sia import util
+
+class UtilTest(unittest.TestCase):
+    def test_get_valid_root_elements_single(self):
+        valid_elements = util.get_valid_root_elements(test_data.echo_action_schema)
+        self.assertEqual(valid_elements, {'test_tag'})
+
+    def test_split_response_single_element(self):
+        response = "Some reasoning here\n<test_tag>content</test_tag>"
+        valid_elements = {'test_tag'}
+        result = util.split_response(response, valid_elements)
+        self.assertEqual(result.reasoning, "Some reasoning here")
+        self.assertEqual(result.actions, "<test_tag>content</test_tag>")