10 Commits

4 changed files with 44 additions and 57 deletions
+8 -6
View File
@@ -107,7 +107,9 @@ def clear_dir(dir_path: Path,
""" """
file_iter = dir_path.glob(glob) if glob else dir_path.iterdir() file_iter = dir_path.glob(glob) if glob else dir_path.iterdir()
for file_path in file_iter: for file_path in file_iter:
if file_path.is_file() and file_path.suffix in Message.file_suffixes: if (file_path.is_file()
and file_path.name not in ignored_files # noqa: W503
and file_path.suffix in Message.file_suffixes): # noqa: W503
file_path.unlink(missing_ok=True) file_path.unlink(missing_ok=True)
@@ -494,13 +496,13 @@ class ChatDB(Chat):
else: else:
return len(self.msg_find([message], loc='db')) > 0 return len(self.msg_find([message], loc='db')) > 0
def cache_read(self) -> None: def cache_read(self, glob: Optional[str] = None, mfilter: Optional[MessageFilter] = None) -> None:
""" """
Read messages from the cache directory. New ones are added to the internal list, Read messages from the cache directory. New ones are added to the internal list,
existing ones are replaced. A message is determined as 'existing' if a message existing ones are replaced. A message is determined as 'existing' if a message
with the same base filename (i. e. 'file_path.name') is already in the list. with the same base filename (i. e. 'file_path.name') is already in the list.
""" """
new_messages = read_dir(self.cache_path, self.glob, self.mfilter) new_messages = read_dir(self.cache_path, glob, mfilter)
# remove all messages from self.messages that are in the new list # remove all messages from self.messages that are in the new list
self.messages = [m for m in self.messages if not message_in(m, new_messages)] self.messages = [m for m in self.messages if not message_in(m, new_messages)]
# copy the messages from the temporary list to self.messages and sort them # copy the messages from the temporary list to self.messages and sort them
@@ -537,11 +539,11 @@ class ChatDB(Chat):
self.messages += messages self.messages += messages
self.msg_sort() self.msg_sort()
def cache_clear(self) -> None: def cache_clear(self, glob: Optional[str] = None) -> None:
""" """
Delete all message files from the cache dir and remove them from the internal list. Delete all message files from the cache dir and remove them from the internal list.
""" """
clear_dir(self.cache_path, self.glob) clear_dir(self.cache_path, glob)
# only keep messages from DB dir (or those that have not yet been written) # only keep messages from DB dir (or those that have not yet been written)
self.messages = [m for m in self.messages if not m.file_path or m.file_path.parent.samefile(self.db_path)] self.messages = [m for m in self.messages if not m.file_path or m.file_path.parent.samefile(self.db_path)]
@@ -561,7 +563,7 @@ class ChatDB(Chat):
# (re)add it to the internal list # (re)add it to the internal list
self.msg_add([message]) self.msg_add([message])
def db_read(self) -> None: def db_read(self, glob: Optional[str] = None, mfilter: Optional[MessageFilter] = None) -> None:
""" """
Read messages from the DB directory. New ones are added to the internal list, Read messages from the DB directory. New ones are added to the internal list,
existing ones are replaced. A message is determined as 'existing' if a message existing ones are replaced. A message is determined as 'existing' if a message
+21 -47
View File
@@ -1,4 +1,3 @@
import sys
import argparse import argparse
from pathlib import Path from pathlib import Path
from itertools import zip_longest from itertools import zip_longest
@@ -74,34 +73,10 @@ def create_message(chat: ChatDB, args: argparse.Namespace) -> Message:
tags=args.output_tags, # FIXME tags=args.output_tags, # FIXME
ai=args.AI, ai=args.AI,
model=args.model) model=args.model)
# only write the message (as a backup), don't add it chat.cache_add([message])
# to the current chat history
chat.cache_write([message])
return message return message
def make_request(ai: AI, chat: ChatDB, message: Message, args: argparse.Namespace) -> None:
"""
Make an AI request with the give AI, chat history, message and CLI arguments.
Print all answers.
"""
ai.print()
chat.print(paged=False)
print(message.to_str() + '\n')
response: AIResponse = ai.request(message,
chat,
args.num_answers,
args.output_tags)
# write all answers to the cache, don't add them to the chat history
chat.cache_write(response.messages)
for idx, msg in enumerate(response.messages):
print(f"=== ANSWER {idx+1} ===")
print(msg.answer)
if response.tokens:
print("===============")
print(response.tokens)
def question_cmd(args: argparse.Namespace, config: Config) -> None: def question_cmd(args: argparse.Namespace, config: Config) -> None:
""" """
Handler for the 'question' command. Handler for the 'question' command.
@@ -120,29 +95,28 @@ def question_cmd(args: argparse.Namespace, config: Config) -> None:
# create the correct AI instance # create the correct AI instance
ai: AI = create_ai(args, config) ai: AI = create_ai(args, config)
# === ASK ===
if args.ask: if args.ask:
make_request(ai, chat, message, args) ai.print()
# === REPEAT === chat.print(paged=False)
response: AIResponse = ai.request(message,
chat,
args.num_answers, # FIXME
args.output_tags) # FIXME
chat.msg_update([response.messages[0]])
chat.cache_add(response.messages[1:])
for idx, msg in enumerate(response.messages):
print(f"=== ANSWER {idx+1} ===")
print(msg.answer)
if response.tokens:
print("===============")
print(response.tokens)
elif args.repeat is not None: elif args.repeat is not None:
lmessage = chat.msg_latest(source='cache') lmessage = chat.msg_latest()
if lmessage is None: assert lmessage
print("No message found to repeat!") # TODO: repeat either the last question or the
sys.exit(1) # one(s) given in 'args.repeat' (overwrite
else: # existing ones if 'args.overwrite' is True)
print(f"Repeating message '{lmessage.msg_id()}':") pass
# overwrite the latest message if requested or empty
if lmessage.answer is None or args.overwrite is True:
lmessage.clear_answer()
make_request(ai, chat, lmessage, args)
# otherwise create a new one
else:
args.ask = [lmessage.question]
message = create_message(chat, args)
make_request(ai, chat, message, args)
# === PROCESS ===
elif args.process is not None: elif args.process is not None:
# TODO: process either all questions without an # TODO: process either all questions without an
# answer or the one(s) given in 'args.process' # answer or the one(s) given in 'args.process'
+1 -4
View File
@@ -393,9 +393,9 @@ class Message():
try: try:
data = yaml.load(fd, Loader=yaml.FullLoader) data = yaml.load(fd, Loader=yaml.FullLoader)
data[cls.file_yaml_key] = file_path data[cls.file_yaml_key] = file_path
return cls.from_dict(data)
except Exception: except Exception:
raise MessageError(f"'{file_path}' does not contain a valid message") raise MessageError(f"'{file_path}' does not contain a valid message")
return cls.from_dict(data)
def to_str(self, with_tags: bool = False, with_file: bool = False, source_code_only: bool = False) -> str: def to_str(self, with_tags: bool = False, with_file: bool = False, source_code_only: bool = False) -> str:
""" """
@@ -540,9 +540,6 @@ class Message():
if self.tags: if self.tags:
self.tags = rename_tags(self.tags, tags_rename) self.tags = rename_tags(self.tags, tags_rename)
def clear_answer(self) -> None:
self.answer = None
def msg_id(self) -> str: def msg_id(self) -> str:
""" """
Returns an ID that is unique throughout all messages in the same (DB) directory. Returns an ID that is unique throughout all messages in the same (DB) directory.
+14
View File
@@ -2,6 +2,7 @@ import unittest
import pathlib import pathlib
import tempfile import tempfile
import time import time
import yaml
from io import StringIO from io import StringIO
from unittest.mock import patch from unittest.mock import patch
from chatmastermind.tags import TagLine from chatmastermind.tags import TagLine
@@ -184,6 +185,10 @@ class TestChatDB(unittest.TestCase):
for file in self.trash_files: for file in self.trash_files:
with open(pathlib.Path(self.db_path.name) / file, 'w') as f: with open(pathlib.Path(self.db_path.name) / file, 'w') as f:
f.write('test trash') f.write('test trash')
# also create a file with actual yaml content
with open(pathlib.Path(self.db_path.name) / 'content.yaml', 'w') as f:
yaml.dump({'key': 'value'}, f)
self.trash_files.append('content.yaml')
self.maxDiff = None self.maxDiff = None
def message_list(self, tmp_dir: tempfile.TemporaryDirectory) -> list[pathlib.Path]: def message_list(self, tmp_dir: tempfile.TemporaryDirectory) -> list[pathlib.Path]:
@@ -605,6 +610,15 @@ class TestChatDB(unittest.TestCase):
self.assertSequenceEqual(chat_db.msg_gather(loc='db'), all_messages) self.assertSequenceEqual(chat_db.msg_gather(loc='db'), all_messages)
self.assertSequenceEqual(chat_db.msg_gather(loc='disk'), all_messages) self.assertSequenceEqual(chat_db.msg_gather(loc='disk'), all_messages)
self.assertSequenceEqual(chat_db.msg_gather(loc='cache'), []) self.assertSequenceEqual(chat_db.msg_gather(loc='cache'), [])
# test with MessageFilter
self.assertSequenceEqual(chat_db.msg_gather(loc='all', mfilter=MessageFilter(tags_or={Tag('tag1')})),
[self.message1])
self.assertSequenceEqual(chat_db.msg_gather(loc='disk', mfilter=MessageFilter(tags_or={Tag('tag2')})),
[self.message2])
self.assertSequenceEqual(chat_db.msg_gather(loc='cache', mfilter=MessageFilter(tags_or={Tag('tag3')})),
[])
self.assertSequenceEqual(chat_db.msg_gather(loc='mem', mfilter=MessageFilter(question_contains="What")),
[new_message])
def test_msg_move_and_gather(self) -> None: def test_msg_move_and_gather(self) -> None:
chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name), chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),