mem0/tests/embedchain/test_embedchain.py

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

os.environ["OPENAI_API_KEY"] = "test-api-key"


@pytest.fixture
def app_instance():
    config = AppConfig(log_level="DEBUG", collect_metrics=False)
    return App(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")

    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 = {"allow_reset": True}

    app_instance = App(config=config, db_config=ChromaDbConfig(**chroma_config))
    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(ValueError):
        app_instance.add(content, data_type="json")