25 lines
931 B
Python
25 lines
931 B
Python
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>") |