Allow to customize glob and location in command question and hist, convert location type to enum.

This commit is contained in:
Oleksandr Kozachuk
2023-10-19 14:45:05 +02:00
parent 5f29f60168
commit 234cca394f
8 changed files with 140 additions and 107 deletions
+16 -14
View File
@@ -9,7 +9,7 @@ from chatmastermind.configuration import Config
from chatmastermind.commands.question import create_message, question_cmd
from chatmastermind.tags import Tag
from chatmastermind.message import Message, Question, Answer
from chatmastermind.chat import Chat, ChatDB
from chatmastermind.chat import Chat, ChatDB, msg_location
from chatmastermind.ai import AIError
from .test_common import TestWithFakeAI
@@ -234,6 +234,8 @@ class TestQuestionCmd(TestWithFakeAI):
# create a mock argparse.Namespace
self.args = argparse.Namespace(
ask=['What is the meaning of life?'],
glob=None,
location='db',
num_answers=1,
output_tags=['science'],
AI='FakeAI',
@@ -279,7 +281,7 @@ class TestQuestionCmdAsk(TestQuestionCmd):
# check for the expected message files
chat = ChatDB.from_dir(Path(self.cache_dir.name),
Path(self.db_dir.name))
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_msgs_equal_except_file_path(cached_msg, expected_responses)
@@ -337,7 +339,7 @@ class TestQuestionCmdAsk(TestQuestionCmd):
# check for the expected message files
chat = ChatDB.from_dir(Path(self.cache_dir.name),
Path(self.db_dir.name))
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_msgs_equal_except_file_path(cached_msg, [expected_question])
@@ -375,7 +377,7 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
# we expect the original message + the one with the new response
expected_responses = [message] + [expected_response]
question_cmd(self.args, self.config)
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
print(self.message_list(self.cache_dir))
self.assertEqual(len(self.message_list(self.cache_dir)), 2)
self.assert_msgs_equal_except_file_path(cached_msg, expected_responses)
@@ -396,7 +398,7 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
model=self.args.model,
file_path=Path(self.cache_dir.name) / f'0001{msg_suffix}')
chat.msg_write([message])
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
assert cached_msg[0].file_path
cached_msg_file_id = cached_msg[0].file_path.stem
@@ -412,7 +414,7 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
tags=message.tags,
file_path=Path('<NOT COMPARED>'))
question_cmd(self.args, self.config)
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_msgs_equal_except_file_path(cached_msg, [expected_response])
# also check that the file ID has not been changed
@@ -435,7 +437,7 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
model=self.args.model,
file_path=Path(self.cache_dir.name) / f'0001{msg_suffix}')
chat.msg_write([message])
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
assert cached_msg[0].file_path
cached_msg_file_id = cached_msg[0].file_path.stem
@@ -452,7 +454,7 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
tags=message.tags,
file_path=Path('<NOT COMPARED>'))
question_cmd(self.args, self.config)
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_msgs_equal_except_file_path(cached_msg, [expected_response])
# also check that the file ID has not been changed
@@ -475,7 +477,7 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
model=self.args.model,
file_path=Path(self.cache_dir.name) / f'0001{msg_suffix}')
chat.msg_write([message])
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
assert cached_msg[0].file_path
# repeat the last question with new arguments (without overwriting)
@@ -493,7 +495,7 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
tags={Tag('newtag')},
file_path=Path('<NOT COMPARED>'))
question_cmd(self.args, self.config)
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
self.assertEqual(len(self.message_list(self.cache_dir)), 2)
self.assert_msgs_equal_except_file_path(cached_msg, [message] + [new_expected_response])
@@ -513,7 +515,7 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
model=self.args.model,
file_path=Path(self.cache_dir.name) / f'0001{msg_suffix}')
chat.msg_write([message])
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
assert cached_msg[0].file_path
# repeat the last question with new arguments
@@ -530,7 +532,7 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
tags={Tag('newtag')},
file_path=Path('<NOT COMPARED>'))
question_cmd(self.args, self.config)
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_msgs_equal_except_file_path(cached_msg, [new_expected_response])
@@ -586,8 +588,8 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
self.assertEqual(len(self.message_list(self.cache_dir)), 4)
self.assertEqual(len(self.message_list(self.db_dir)), 1)
expected_cache_messages = [expected_responses[0], message2, expected_responses[1], expected_responses[2]]
cached_msg = chat.msg_gather(loc='cache')
cached_msg = chat.msg_gather(loc=msg_location.CACHE)
self.assert_msgs_equal_except_file_path(cached_msg, expected_cache_messages)
# check that the DB message has not been modified at all
db_msg = chat.msg_gather(loc='db')
db_msg = chat.msg_gather(loc=msg_location.DB)
self.assert_msgs_all_equal(db_msg, [message3])