my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import pytest
#2	from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
#3
#4	from embedchain.config import BaseLlmConfig
#5	from embedchain.llm.aws_bedrock import AWSBedrockLlm
#6
#7
#8	@pytest.fixture
#9	def config(monkeypatch):
#10	monkeypatch.setenv("AWS_ACCESS_KEY_ID", "test_access_key_id")
#11	monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test_secret_access_key")
#12	config = BaseLlmConfig(
#13	model="amazon.titan-text-express-v1",
#14	model_kwargs={
#15	"temperature": 0.5,
#16	"topP": 1,
#17	"maxTokenCount": 1000,
#18	},
#19	)
#20	yield config
#21	monkeypatch.delenv("AWS_ACCESS_KEY_ID")
#22	monkeypatch.delenv("AWS_SECRET_ACCESS_KEY")
#23
#24
#25	def test_get_llm_model_answer(config, mocker):
#26	mocked_get_answer = mocker.patch("embedchain.llm.aws_bedrock.AWSBedrockLlm._get_answer", return_value="Test answer")
#27
#28	llm = AWSBedrockLlm(config)
#29	answer = llm.get_llm_model_answer("Test query")
#30
#31	assert answer == "Test answer"
#32	mocked_get_answer.assert_called_once_with("Test query", config)
#33
#34
#35	def test_get_llm_model_answer_empty_prompt(config, mocker):
#36	mocked_get_answer = mocker.patch("embedchain.llm.aws_bedrock.AWSBedrockLlm._get_answer", return_value="Test answer")
#37
#38	llm = AWSBedrockLlm(config)
#39	answer = llm.get_llm_model_answer("")
#40
#41	assert answer == "Test answer"
#42	mocked_get_answer.assert_called_once_with("", config)
#43
#44
#45	def test_get_llm_model_answer_with_streaming(config, mocker):
#46	config.stream = True
#47	mocked_bedrock_chat = mocker.patch("embedchain.llm.aws_bedrock.BedrockLLM")
#48
#49	llm = AWSBedrockLlm(config)
#50	llm.get_llm_model_answer("Test query")
#51
#52	mocked_bedrock_chat.assert_called_once()
#53	callbacks = [callback[1]["callbacks"] for callback in mocked_bedrock_chat.call_args_list]
#54	assert any(isinstance(callback[0], StreamingStdOutCallbackHandler) for callback in callbacks)
#55

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public