Start work on llm_engine

tests/test_llm_engine.py

@@ -0,0 +1,66 @@
+import unittest
+from sia.llm_engine import LlmEngine, InferenceResult
+
+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.
+    """
+
+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>
+    """
+
+class TestLlmEngine(unittest.TestCase):
+    def setUp(self):
+        self.model_path = "/root/model"
+
+    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(echo_system_prompt, main_context, 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>")
+
+'''
+    def test_set_model_path(self):
+        new_model_path = "/path/to/new/model"
+        self.llm_engine.set_model_path(new_model_path)
+        # Add assertions to check if the model path was updated correctly
+
+    def test_finetune(self):
+        dataset_paths = ["/path/to/dataset1", "/path/to/dataset2"]
+        output_dir = "/path/to/output"
+
+        self.llm_engine.finetune(dataset_paths, output_dir)
+        # Add assertions to check if the fine-tuning process completed successfully
+        # For example, check if new model weights were saved in the output directory
+'''
+
+if __name__ == '__main__':
+    unittest.main()