79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
import os
|
|
|
|
import pytest
|
|
from chromadb.api.models.Collection import Collection
|
|
|
|
from embedchain import App
|
|
from embedchain.config import AppConfig, ChromaDbConfig
|
|
from embedchain.embedchain import EmbedChain
|
|
from embedchain.llm.base import BaseLlm
|
|
from embedchain.memory.base import ChatHistory
|
|
from embedchain.vectordb.chroma import ChromaDB
|
|
|
|
os.environ["OPENAI_API_KEY"] = "test-api-key"
|
|
|
|
|
|
@pytest.fixture
|
|
def app_instance():
|
|
config = AppConfig(log_level="DEBUG", collect_metrics=False)
|
|
return App(config=config)
|
|
|
|
|
|
def test_whole_app(app_instance, mocker):
|
|
knowledge = "lorem ipsum dolor sit amet, consectetur adipiscing"
|
|
|
|
mocker.patch.object(EmbedChain, "add")
|
|
mocker.patch.object(EmbedChain, "_retrieve_from_database")
|
|
mocker.patch.object(BaseLlm, "get_answer_from_llm", return_value=knowledge)
|
|
mocker.patch.object(BaseLlm, "get_llm_model_answer", return_value=knowledge)
|
|
mocker.patch.object(BaseLlm, "generate_prompt")
|
|
mocker.patch.object(
|
|
BaseLlm,
|
|
"add_history",
|
|
)
|
|
mocker.patch.object(ChatHistory, "delete", autospec=True)
|
|
|
|
app_instance.add(knowledge, data_type="text")
|
|
app_instance.query("What text did I give you?")
|
|
app_instance.chat("What text did I give you?")
|
|
|
|
assert BaseLlm.generate_prompt.call_count == 2
|
|
app_instance.reset()
|
|
|
|
|
|
def test_add_after_reset(app_instance, mocker):
|
|
mocker.patch("embedchain.vectordb.chroma.chromadb.Client")
|
|
|
|
config = AppConfig(log_level="DEBUG", collect_metrics=False)
|
|
chroma_config = ChromaDbConfig(allow_reset=True)
|
|
db = ChromaDB(config=chroma_config)
|
|
app_instance = App(config=config, db=db)
|
|
|
|
# mock delete chat history
|
|
mocker.patch.object(ChatHistory, "delete", autospec=True)
|
|
|
|
app_instance.reset()
|
|
|
|
app_instance.db.client.heartbeat()
|
|
|
|
mocker.patch.object(Collection, "add")
|
|
|
|
app_instance.db.collection.add(
|
|
embeddings=[[1.1, 2.3, 3.2], [4.5, 6.9, 4.4], [1.1, 2.3, 3.2]],
|
|
metadatas=[
|
|
{"chapter": "3", "verse": "16"},
|
|
{"chapter": "3", "verse": "5"},
|
|
{"chapter": "29", "verse": "11"},
|
|
],
|
|
ids=["id1", "id2", "id3"],
|
|
)
|
|
|
|
app_instance.reset()
|
|
|
|
|
|
def test_add_with_incorrect_content(app_instance, mocker):
|
|
content = [{"foo": "bar"}]
|
|
|
|
with pytest.raises(TypeError):
|
|
app_instance.add(content, data_type="json")
|