my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	from unittest.mock import MagicMock, patch
#2
#3	import pytest
#4	from langchain.schema import HumanMessage, SystemMessage
#5
#6	from embedchain.config import BaseLlmConfig
#7	from embedchain.core.db.database import database_manager
#8	from embedchain.llm.vertex_ai import VertexAILlm
#9
#10
#11	@pytest.fixture(autouse=True)
#12	def setup_database():
#13	database_manager.setup_engine()
#14
#15
#16	@pytest.fixture
#17	def vertexai_llm():
#18	config = BaseLlmConfig(temperature=0.6, model="chat-bison")
#19	return VertexAILlm(config)
#20
#21
#22	def test_get_llm_model_answer(vertexai_llm):
#23	with patch.object(VertexAILlm, "_get_answer", return_value="Test Response") as mock_method:
#24	prompt = "Test Prompt"
#25	response = vertexai_llm.get_llm_model_answer(prompt)
#26	assert response == "Test Response"
#27	mock_method.assert_called_once_with(prompt, vertexai_llm.config)
#28
#29
#30	def test_get_llm_model_answer_with_token_usage(vertexai_llm):
#31	test_config = BaseLlmConfig(
#32	temperature=vertexai_llm.config.temperature,
#33	max_tokens=vertexai_llm.config.max_tokens,
#34	top_p=vertexai_llm.config.top_p,
#35	model=vertexai_llm.config.model,
#36	token_usage=True,
#37	)
#38	vertexai_llm.config = test_config
#39	with patch.object(
#40	VertexAILlm,
#41	"_get_answer",
#42	return_value=("Test Response", {"prompt_token_count": 1, "candidates_token_count": 2}),
#43	):
#44	response, token_info = vertexai_llm.get_llm_model_answer("Test Query")
#45	assert response == "Test Response"
#46	assert token_info == {
#47	"prompt_tokens": 1,
#48	"completion_tokens": 2,
#49	"total_tokens": 3,
#50	"total_cost": 3.75e-07,
#51	"cost_currency": "USD",
#52	}
#53
#54
#55	@patch("embedchain.llm.vertex_ai.ChatVertexAI")
#56	def test_get_answer(mock_chat_vertexai, vertexai_llm, caplog):
#57	mock_chat_vertexai.return_value.invoke.return_value = MagicMock(content="Test Response")
#58
#59	config = vertexai_llm.config
#60	prompt = "Test Prompt"
#61	messages = vertexai_llm._get_messages(prompt)
#62	response = vertexai_llm._get_answer(prompt, config)
#63	mock_chat_vertexai.return_value.invoke.assert_called_once_with(messages)
#64
#65	assert response == "Test Response" # Assertion corrected
#66	assert "Config option `top_p` is not supported by this model." not in caplog.text
#67
#68
#69	def test_get_messages(vertexai_llm):
#70	prompt = "Test Prompt"
#71	system_prompt = "Test System Prompt"
#72	messages = vertexai_llm._get_messages(prompt, system_prompt)
#73	assert messages == [
#74	SystemMessage(content="Test System Prompt", additional_kwargs={}),
#75	HumanMessage(content="Test Prompt", additional_kwargs={}, example=False),
#76	]
#77

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public