my-project-public

#1

import os

#2

import unittest

#3

from unittest.mock import MagicMock, patch

#4

#5

from embedchain import App

#6

from embedchain.config import AppConfig, BaseLlmConfig

#7

from embedchain.llm.base import BaseLlm

#8

from embedchain.memory.base import ChatHistory

#9

from embedchain.memory.message import ChatMessage

#10

#11

#12

class TestApp(unittest.TestCase):

#13

def setUp(self):

#14

os.environ["OPENAI_API_KEY"] = "test_key"

#15

self.app = App(config=AppConfig(collect_metrics=False))

#16

#17

@patch.object(App, "_retrieve_from_database", return_value=["Test context"])

#18

@patch.object(BaseLlm, "get_answer_from_llm", return_value="Test answer")

#19

def test_chat_with_memory(self, mock_get_answer, mock_retrieve):

#20

"""

#21

This test checks the functionality of the 'chat' method in the App class with respect to the chat history

#22

memory.

#23

The 'chat' method is called twice. The first call initializes the chat history memory.

#24

The second call is expected to use the chat history from the first call.

#25

#26

Key assumptions tested:

#27

called with correct arguments, adding the correct chat history.

#28

- After the first call, 'memory.chat_memory.add_user_message' and 'memory.chat_memory.add_ai_message' are

#29

- During the second call, the 'chat' method uses the chat history from the first call.

#30

#31

The test isolates the 'chat' method behavior by mocking out '_retrieve_from_database', 'get_answer_from_llm' and

#32

'memory' methods.

#33

"""

#34

config = AppConfig(collect_metrics=False)

#35

app = App(config=config)

#36

with patch.object(BaseLlm, "add_history") as mock_history:

#37

first_answer = app.chat("Test query 1")

#38

self.assertEqual(first_answer, "Test answer")

#39

mock_history.assert_called_with(app.config.id, "Test query 1", "Test answer", session_id="default")

#40

#41

second_answer = app.chat("Test query 2", session_id="test_session")

#42

self.assertEqual(second_answer, "Test answer")

#43

mock_history.assert_called_with(app.config.id, "Test query 2", "Test answer", session_id="test_session")

#44

#45

@patch.object(App, "_retrieve_from_database", return_value=["Test context"])

#46

@patch.object(BaseLlm, "get_answer_from_llm", return_value="Test answer")

#47

def test_template_replacement(self, mock_get_answer, mock_retrieve):

#48

"""

#49

Tests that if a default template is used and it doesn't contain history,

#50

the default template is swapped in.

#51

#52

Also tests that a dry run does not change the history

#53

"""

#54

with patch.object(ChatHistory, "get") as mock_memory:

#55

mock_message = ChatMessage()

#56

mock_message.add_user_message("Test query 1")

#57

mock_message.add_ai_message("Test answer")

#58

mock_memory.return_value = [mock_message]

#59

#60

config = AppConfig(collect_metrics=False)

#61

app = App(config=config)

#62

first_answer = app.chat("Test query 1")

#63

self.assertEqual(first_answer, "Test answer")

#64

self.assertEqual(len(app.llm.history), 1)

#65

history = app.llm.history

#66

dry_run = app.chat("Test query 2", dry_run=True)

#67

self.assertIn("Conversation history:", dry_run)

#68

self.assertEqual(history, app.llm.history)

#69

self.assertEqual(len(app.llm.history), 1)

#70

#71

@patch("chromadb.api.models.Collection.Collection.add", MagicMock)

#72

def test_chat_with_where_in_params(self):

"""

Test where filter

"""

with patch.object(self.app, "_retrieve_from_database") as mock_retrieve:

#77

mock_retrieve.return_value = ["Test context"]

#78

with patch.object(self.app.llm, "get_llm_model_answer") as mock_answer:

#79

mock_answer.return_value = "Test answer"

#80

answer = self.app.chat("Test query", where={"attribute": "value"})

#81

#82

self.assertEqual(answer, "Test answer")

#83

_args, kwargs = mock_retrieve.call_args

#84

self.assertEqual(kwargs.get("input_query"), "Test query")

#85

self.assertEqual(kwargs.get("where"), {"attribute": "value"})

#86

mock_answer.assert_called_once()

#87

#88

@patch("chromadb.api.models.Collection.Collection.add", MagicMock)

#89

def test_chat_with_where_in_chat_config(self):

#90

"""

#91

This test checks the functionality of the 'chat' method in the App class.

#92

It simulates a scenario where the '_retrieve_from_database' method returns a context list based on

#93

a where filter and 'get_llm_model_answer' returns an expected answer string.

#94

#95

The 'chat' method is expected to call '_retrieve_from_database' with the where filter specified

#96

in the BaseLlmConfig and 'get_llm_model_answer' methods appropriately and return the right answer.

#97

#98

Key assumptions tested:

#99

- '_retrieve_from_database' method is called exactly once with arguments: "Test query" and an instance of

#100

BaseLlmConfig.

#101

- 'get_llm_model_answer' is called exactly once. The specific arguments are not checked in this test.

#102

- 'chat' method returns the value it received from 'get_llm_model_answer'.

#103

#104

The test isolates the 'chat' method behavior by mocking out '_retrieve_from_database' and

#105

'get_llm_model_answer' methods.

#106

"""

#107

with patch.object(self.app.llm, "get_llm_model_answer") as mock_answer:

#108

mock_answer.return_value = "Test answer"

#109

with patch.object(self.app.db, "query") as mock_database_query:

#110

mock_database_query.return_value = ["Test context"]

#111

llm_config = BaseLlmConfig(where={"attribute": "value"})

#112

answer = self.app.chat("Test query", llm_config)

#113

#114

self.assertEqual(answer, "Test answer")

#115

_args, kwargs = mock_database_query.call_args

#116

self.assertEqual(kwargs.get("input_query"), "Test query")

#117

where = kwargs.get("where")

#118

assert "app_id" in where

#119

assert "attribute" in where

#120

mock_answer.assert_called_once()

#121

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public