50 lines
1.4 KiB
Python
50 lines
1.4 KiB
Python
import os
|
|
|
|
import pytest
|
|
|
|
from embedchain.llm.llama2 import Llama2Llm
|
|
|
|
|
|
@pytest.fixture
|
|
def llama2_llm():
|
|
os.environ["REPLICATE_API_TOKEN"] = "test_api_token"
|
|
llm = Llama2Llm()
|
|
return llm
|
|
|
|
|
|
def test_init_raises_value_error_without_api_key(mocker):
|
|
mocker.patch.dict(os.environ, clear=True)
|
|
with pytest.raises(ValueError):
|
|
Llama2Llm()
|
|
|
|
|
|
def test_get_llm_model_answer_raises_value_error_for_system_prompt(llama2_llm):
|
|
llama2_llm.config.system_prompt = "system_prompt"
|
|
with pytest.raises(ValueError):
|
|
llama2_llm.get_llm_model_answer("prompt")
|
|
|
|
|
|
def test_get_llm_model_answer(llama2_llm, mocker):
|
|
mocked_replicate = mocker.patch("embedchain.llm.llama2.Replicate")
|
|
mocked_replicate_instance = mocker.MagicMock()
|
|
mocked_replicate.return_value = mocked_replicate_instance
|
|
mocked_replicate_instance.return_value = "Test answer"
|
|
|
|
llama2_llm.config.model = "test_model"
|
|
llama2_llm.config.max_tokens = 50
|
|
llama2_llm.config.temperature = 0.7
|
|
llama2_llm.config.top_p = 0.8
|
|
|
|
answer = llama2_llm.get_llm_model_answer("Test query")
|
|
|
|
assert answer == "Test answer"
|
|
mocked_replicate.assert_called_once_with(
|
|
model="test_model",
|
|
input={
|
|
"temperature": 0.7,
|
|
"max_length": 50,
|
|
"top_p": 0.8,
|
|
},
|
|
)
|
|
mocked_replicate_instance.assert_called_once_with("Test query")
|