my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	import os
#2	from unittest.mock import MagicMock, patch
#3
#4	import pytest
#5
#6	from embedchain import App
#7	from embedchain.config import AppConfig, BaseLlmConfig
#8	from embedchain.llm.openai import OpenAILlm
#9
#10
#11	@pytest.fixture
#12	def app():
#13	os.environ["OPENAI_API_KEY"] = "test_api_key"
#14	app = App(config=AppConfig(collect_metrics=False))
#15	return app
#16
#17
#18	@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
#19	def test_query(app):
#20	with patch.object(app, "_retrieve_from_database") as mock_retrieve:
#21	mock_retrieve.return_value = ["Test context"]
#22	with patch.object(app.llm, "get_llm_model_answer") as mock_answer:
#23	mock_answer.return_value = "Test answer"
#24	answer = app.query(input_query="Test query")
#25	assert answer == "Test answer"
#26
#27	mock_retrieve.assert_called_once()
#28	_, kwargs = mock_retrieve.call_args
#29	input_query_arg = kwargs.get("input_query")
#30	assert input_query_arg == "Test query"
#31	mock_answer.assert_called_once()
#32
#33
#34	@patch("embedchain.llm.openai.OpenAILlm._get_answer")
#35	def test_query_config_app_passing(mock_get_answer):
#36	mock_get_answer.return_value = MagicMock()
#37	mock_get_answer.return_value = "Test answer"
#38
#39	config = AppConfig(collect_metrics=False)
#40	chat_config = BaseLlmConfig(system_prompt="Test system prompt")
#41	llm = OpenAILlm(config=chat_config)
#42	app = App(config=config, llm=llm)
#43	answer = app.llm.get_llm_model_answer("Test query")
#44
#45	assert app.llm.config.system_prompt == "Test system prompt"
#46	assert answer == "Test answer"
#47
#48
#49	@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
#50	def test_query_with_where_in_params(app):
#51	with patch.object(app, "_retrieve_from_database") as mock_retrieve:
#52	mock_retrieve.return_value = ["Test context"]
#53	with patch.object(app.llm, "get_llm_model_answer") as mock_answer:
#54	mock_answer.return_value = "Test answer"
#55	answer = app.query("Test query", where={"attribute": "value"})
#56
#57	assert answer == "Test answer"
#58	_, kwargs = mock_retrieve.call_args
#59	assert kwargs.get("input_query") == "Test query"
#60	assert kwargs.get("where") == {"attribute": "value"}
#61	mock_answer.assert_called_once()
#62
#63
#64	@patch("chromadb.api.models.Collection.Collection.add", MagicMock)
#65	def test_query_with_where_in_query_config(app):
#66	with patch.object(app.llm, "get_llm_model_answer") as mock_answer:
#67	mock_answer.return_value = "Test answer"
#68	with patch.object(app.db, "query") as mock_database_query:
#69	mock_database_query.return_value = ["Test context"]
#70	llm_config = BaseLlmConfig(where={"attribute": "value"})
#71	answer = app.query("Test query", llm_config)
#72
#73	assert answer == "Test answer"
#74	_, kwargs = mock_database_query.call_args
#75	assert kwargs.get("input_query") == "Test query"
#76	where = kwargs.get("where")
#77	assert "app_id" in where
#78	assert "attribute" in where
#79	mock_answer.assert_called_once()
#80

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public