You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
49 lines
1.8 KiB
49 lines
1.8 KiB
import unittest
|
|
import os
|
|
from unittest.mock import patch, MagicMock
|
|
from langchain import PromptTemplate, HuggingFaceHub, ChatOpenAI, LLMChain
|
|
|
|
from swarms.utils.llm import LLM
|
|
|
|
class TestLLM(unittest.TestCase):
|
|
@patch.object(HuggingFaceHub, '__init__', return_value=None)
|
|
@patch.object(ChatOpenAI, '__init__', return_value=None)
|
|
def setUp(self, mock_hf_init, mock_openai_init):
|
|
self.llm_openai = LLM(openai_api_key='mock_openai_key')
|
|
self.llm_hf = LLM(hf_repo_id='mock_repo_id', hf_api_token='mock_hf_token')
|
|
self.prompt = "Who won the FIFA World Cup in 1998?"
|
|
|
|
def test_init(self):
|
|
self.assertEqual(self.llm_openai.openai_api_key, 'mock_openai_key')
|
|
self.assertEqual(self.llm_hf.hf_repo_id, 'mock_repo_id')
|
|
self.assertEqual(self.llm_hf.hf_api_token, 'mock_hf_token')
|
|
|
|
@patch.object(HuggingFaceHub, 'run', return_value="France")
|
|
@patch.object(ChatOpenAI, 'run', return_value="France")
|
|
def test_run(self, mock_hf_run, mock_openai_run):
|
|
result_openai = self.llm_openai.run(self.prompt)
|
|
mock_openai_run.assert_called_once()
|
|
self.assertEqual(result_openai, "France")
|
|
|
|
result_hf = self.llm_hf.run(self.prompt)
|
|
mock_hf_run.assert_called_once()
|
|
self.assertEqual(result_hf, "France")
|
|
|
|
def test_error_on_no_keys(self):
|
|
with self.assertRaises(ValueError):
|
|
LLM()
|
|
|
|
@patch.object(os, 'environ', {})
|
|
def test_error_on_missing_hf_token(self):
|
|
with self.assertRaises(ValueError):
|
|
LLM(hf_repo_id='mock_repo_id')
|
|
|
|
@patch.dict(os.environ, {"HUGGINGFACEHUB_API_TOKEN": "mock_hf_token"})
|
|
def test_hf_token_from_env(self):
|
|
llm = LLM(hf_repo_id='mock_repo_id')
|
|
self.assertEqual(llm.hf_api_token, "mock_hf_token")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|