my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2
#3	import pytest
#4	from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
#5
#6	from embedchain.config import BaseLlmConfig
#7	from embedchain.llm.jina import JinaLlm
#8
#9
#10	@pytest.fixture
#11	def config():
#12	os.environ["JINACHAT_API_KEY"] = "test_api_key"
#13	config = BaseLlmConfig(temperature=0.7, max_tokens=50, top_p=0.8, stream=False, system_prompt="System prompt")
#14	yield config
#15	os.environ.pop("JINACHAT_API_KEY")
#16
#17
#18	def test_init_raises_value_error_without_api_key(mocker):
#19	mocker.patch.dict(os.environ, clear=True)
#20	with pytest.raises(ValueError):
#21	JinaLlm()
#22
#23
#24	def test_get_llm_model_answer(config, mocker):
#25	mocked_get_answer = mocker.patch("embedchain.llm.jina.JinaLlm._get_answer", return_value="Test answer")
#26
#27	llm = JinaLlm(config)
#28	answer = llm.get_llm_model_answer("Test query")
#29
#30	assert answer == "Test answer"
#31	mocked_get_answer.assert_called_once_with("Test query", config)
#32
#33
#34	def test_get_llm_model_answer_with_system_prompt(config, mocker):
#35	config.system_prompt = "Custom system prompt"
#36	mocked_get_answer = mocker.patch("embedchain.llm.jina.JinaLlm._get_answer", return_value="Test answer")
#37
#38	llm = JinaLlm(config)
#39	answer = llm.get_llm_model_answer("Test query")
#40
#41	assert answer == "Test answer"
#42	mocked_get_answer.assert_called_once_with("Test query", config)
#43
#44
#45	def test_get_llm_model_answer_empty_prompt(config, mocker):
#46	mocked_get_answer = mocker.patch("embedchain.llm.jina.JinaLlm._get_answer", return_value="Test answer")
#47
#48	llm = JinaLlm(config)
#49	answer = llm.get_llm_model_answer("")
#50
#51	assert answer == "Test answer"
#52	mocked_get_answer.assert_called_once_with("", config)
#53
#54
#55	def test_get_llm_model_answer_with_streaming(config, mocker):
#56	config.stream = True
#57	mocked_jinachat = mocker.patch("embedchain.llm.jina.JinaChat")
#58
#59	llm = JinaLlm(config)
#60	llm.get_llm_model_answer("Test query")
#61
#62	mocked_jinachat.assert_called_once()
#63	callbacks = [callback[1]["callbacks"] for callback in mocked_jinachat.call_args_list]
#64	assert any(isinstance(callback[0], StreamingStdOutCallbackHandler) for callback in callbacks)
#65
#66
#67	def test_get_llm_model_answer_without_system_prompt(config, mocker):
#68	config.system_prompt = None
#69	mocked_jinachat = mocker.patch("embedchain.llm.jina.JinaChat")
#70
#71	llm = JinaLlm(config)
#72	llm.get_llm_model_answer("Test query")
#73
#74	mocked_jinachat.assert_called_once_with(
#75	temperature=config.temperature,
#76	max_tokens=config.max_tokens,
#77	jinachat_api_key=os.environ["JINACHAT_API_KEY"],
#78	model_kwargs={"top_p": config.top_p},
#79	)
#80

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public