Private instance vars, public get properties

test/llm_engine_test.py

@@ -0,0 +1,25 @@
+import unittest
+
+from itertools import tee
+
+from . import test_data
+
+from sia.llm_engine import LlmEngine
+
+class LlmEngineTest(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)
+        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>")