my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import pytest
#2
#3	from embedchain.config import BaseLlmConfig
#4	from embedchain.llm.mistralai import MistralAILlm
#5
#6
#7	@pytest.fixture
#8	def mistralai_llm_config(monkeypatch):
#9	monkeypatch.setenv("MISTRAL_API_KEY", "fake_api_key")
#10	yield BaseLlmConfig(model="mistral-tiny", max_tokens=100, temperature=0.7, top_p=0.5, stream=False)
#11	monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
#12
#13
#14	def test_mistralai_llm_init_missing_api_key(monkeypatch):
#15	monkeypatch.delenv("MISTRAL_API_KEY", raising=False)
#16	with pytest.raises(ValueError, match="Please set the MISTRAL_API_KEY environment variable."):
#17	MistralAILlm()
#18
#19
#20	def test_mistralai_llm_init(monkeypatch):
#21	monkeypatch.setenv("MISTRAL_API_KEY", "fake_api_key")
#22	llm = MistralAILlm()
#23	assert llm is not None
#24
#25
#26	def test_get_llm_model_answer(monkeypatch, mistralai_llm_config):
#27	def mock_get_answer(self, prompt, config):
#28	return "Generated Text"
#29
#30	monkeypatch.setattr(MistralAILlm, "_get_answer", mock_get_answer)
#31	llm = MistralAILlm(config=mistralai_llm_config)
#32	result = llm.get_llm_model_answer("test prompt")
#33
#34	assert result == "Generated Text"
#35
#36
#37	def test_get_llm_model_answer_with_system_prompt(monkeypatch, mistralai_llm_config):
#38	mistralai_llm_config.system_prompt = "Test system prompt"
#39	monkeypatch.setattr(MistralAILlm, "_get_answer", lambda self, prompt, config: "Generated Text")
#40	llm = MistralAILlm(config=mistralai_llm_config)
#41	result = llm.get_llm_model_answer("test prompt")
#42
#43	assert result == "Generated Text"
#44
#45
#46	def test_get_llm_model_answer_empty_prompt(monkeypatch, mistralai_llm_config):
#47	monkeypatch.setattr(MistralAILlm, "_get_answer", lambda self, prompt, config: "Generated Text")
#48	llm = MistralAILlm(config=mistralai_llm_config)
#49	result = llm.get_llm_model_answer("")
#50
#51	assert result == "Generated Text"
#52
#53
#54	def test_get_llm_model_answer_without_system_prompt(monkeypatch, mistralai_llm_config):
#55	mistralai_llm_config.system_prompt = None
#56	monkeypatch.setattr(MistralAILlm, "_get_answer", lambda self, prompt, config: "Generated Text")
#57	llm = MistralAILlm(config=mistralai_llm_config)
#58	result = llm.get_llm_model_answer("test prompt")
#59
#60	assert result == "Generated Text"
#61
#62
#63	def test_get_llm_model_answer_with_token_usage(monkeypatch, mistralai_llm_config):
#64	test_config = BaseLlmConfig(
#65	temperature=mistralai_llm_config.temperature,
#66	max_tokens=mistralai_llm_config.max_tokens,
#67	top_p=mistralai_llm_config.top_p,
#68	model=mistralai_llm_config.model,
#69	token_usage=True,
#70	)
#71	monkeypatch.setattr(
#72	MistralAILlm,
#73	"_get_answer",
#74	lambda self, prompt, config: ("Generated Text", {"prompt_tokens": 1, "completion_tokens": 2}),
#75	)
#76
#77	llm = MistralAILlm(test_config)
#78	answer, token_info = llm.get_llm_model_answer("Test query")
#79
#80	assert answer == "Generated Text"
#81	assert token_info == {
#82	"prompt_tokens": 1,
#83	"completion_tokens": 2,
#84	"total_tokens": 3,
#85	"total_cost": 7.5e-07,
#86	"cost_currency": "USD",
#87	}
#88

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public