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>")