9 Commits

16 changed files with 464 additions and 954 deletions
-1
View File
@@ -106,7 +106,6 @@ celerybeat.pid
.venv .venv
env/ env/
venv/ venv/
.old/
ENV/ ENV/
env.bak/ env.bak/
venv.bak/ venv.bak/
+6 -7
View File
@@ -68,21 +68,20 @@ cmm question [-t OTAGS]... [-k ATAGS]... [-x XTAGS]... [-o OUTTAGS]... [-A AI_ID
#### Hist #### Hist
The `hist` command is used to print and manage the chat history. The `hist` command is used to print the chat history.
```bash ```bash
cmm hist [--print | --convert FORMAT] [-t OTAGS]... [-k ATAGS]... [-x XTAGS]... [-w] [-W] [-S] [-A SUBSTRING] [-Q SUBSTRING] cmm hist [-t OTAGS]... [-k ATAGS]... [-x XTAGS]... [-w] [-W] [-S] [-A SUBSTRING] [-Q SUBSTRING]
``` ```
* `-p, --print`: Print the DB chat history
* `-c, --convert FORMAT`: Convert all messages to the given format
* `-t, --or-tags OTAGS`: List of tags (one must match) * `-t, --or-tags OTAGS`: List of tags (one must match)
* `-k, --and-tags ATAGS`: List of tags (all must match) * `-k, --and-tags ATAGS`: List of tags (all must match)
* `-x, --exclude-tags XTAGS`: List of tags to exclude * `-x, --exclude-tags XTAGS`: List of tags to exclude
* `-w, --with-metadata`: Print chat history with metadata (tags, filenames, AI, etc.) * `-w, --with-tags`: Print chat history with tags
* `-W, --with-files`: Print chat history with filenames
* `-S, --source-code-only`: Only print embedded source code * `-S, --source-code-only`: Only print embedded source code
* `-A, --answer SUBSTRING`: Filter for answer substring * `-A, --answer SUBSTRING`: Search for answer substring
* `-Q, --question SUBSTRING`: Filter for question substring * `-Q, --question SUBSTRING`: Search for question substring
#### Tags #### Tags
+6 -12
View File
@@ -3,20 +3,18 @@ Creates different AI instances, based on the given configuration.
""" """
import argparse import argparse
from typing import cast, Optional from typing import cast
from .configuration import Config, AIConfig, OpenAIConfig from .configuration import Config, AIConfig, OpenAIConfig
from .ai import AI, AIError from .ai import AI, AIError
from .ais.openai import OpenAI from .ais.openai import OpenAI
def create_ai(args: argparse.Namespace, config: Config, # noqa: 11 def create_ai(args: argparse.Namespace, config: Config) -> AI: # noqa: 11
def_ai: Optional[str] = None,
def_model: Optional[str] = None) -> AI:
""" """
Creates an AI subclass instance from the given arguments and configuration file. Creates an AI subclass instance from the given arguments
If AI has not been set in the arguments, it searches for the ID 'default'. If and configuration file. If AI has not been set in the
that is not found, it uses the first AI in the list. It's also possible to arguments, it searches for the ID 'default'. If that
specify a default AI and model using 'def_ai' and 'def_model'. is not found, it uses the first AI in the list.
""" """
ai_conf: AIConfig ai_conf: AIConfig
if hasattr(args, 'AI') and args.AI: if hasattr(args, 'AI') and args.AI:
@@ -24,8 +22,6 @@ def create_ai(args: argparse.Namespace, config: Config, # noqa: 11
ai_conf = config.ais[args.AI] ai_conf = config.ais[args.AI]
except KeyError: except KeyError:
raise AIError(f"AI ID '{args.AI}' does not exist in this configuration") raise AIError(f"AI ID '{args.AI}' does not exist in this configuration")
elif def_ai:
ai_conf = config.ais[def_ai]
elif 'default' in config.ais: elif 'default' in config.ais:
ai_conf = config.ais['default'] ai_conf = config.ais['default']
else: else:
@@ -38,8 +34,6 @@ def create_ai(args: argparse.Namespace, config: Config, # noqa: 11
ai = OpenAI(cast(OpenAIConfig, ai_conf)) ai = OpenAI(cast(OpenAIConfig, ai_conf))
if hasattr(args, 'model') and args.model: if hasattr(args, 'model') and args.model:
ai.config.model = args.model ai.config.model = args.model
elif def_model:
ai.config.model = def_model
if hasattr(args, 'max_tokens') and args.max_tokens: if hasattr(args, 'max_tokens') and args.max_tokens:
ai.config.max_tokens = args.max_tokens ai.config.max_tokens = args.max_tokens
if hasattr(args, 'temperature') and args.temperature: if hasattr(args, 'temperature') and args.temperature:
+1 -1
View File
@@ -44,7 +44,7 @@ class OpenAI(AI):
frequency_penalty=self.config.frequency_penalty, frequency_penalty=self.config.frequency_penalty,
presence_penalty=self.config.presence_penalty) presence_penalty=self.config.presence_penalty)
question.answer = Answer(response['choices'][0]['message']['content']) question.answer = Answer(response['choices'][0]['message']['content'])
question.tags = set(otags) if otags is not None else None question.tags = otags
question.ai = self.ID question.ai = self.ID
question.model = self.config.model question.model = self.config.model
answers: list[Message] = [question] answers: list[Message] = [question]
+37 -47
View File
@@ -6,9 +6,9 @@ from pathlib import Path
from pprint import PrettyPrinter from pprint import PrettyPrinter
from pydoc import pager from pydoc import pager
from dataclasses import dataclass from dataclasses import dataclass
from typing import TypeVar, Type, Optional, Any, Callable, Literal, Union from typing import TypeVar, Type, Optional, ClassVar, Any, Callable, Literal, Union
from .configuration import default_config_file from .configuration import default_config_file
from .message import Message, MessageFilter, MessageError, MessageFormat, message_in, message_valid_formats from .message import Message, MessageFilter, MessageError, message_in
from .tags import Tag from .tags import Tag
ChatInst = TypeVar('ChatInst', bound='Chat') ChatInst = TypeVar('ChatInst', bound='Chat')
@@ -17,7 +17,6 @@ ChatDBInst = TypeVar('ChatDBInst', bound='ChatDB')
db_next_file = '.next' db_next_file = '.next'
ignored_files = [db_next_file, default_config_file] ignored_files = [db_next_file, default_config_file]
msg_location = Literal['mem', 'disk', 'cache', 'db', 'all'] msg_location = Literal['mem', 'disk', 'cache', 'db', 'all']
msg_suffix = Message.file_suffix_write
class ChatError(Exception): class ChatError(Exception):
@@ -53,7 +52,7 @@ def read_dir(dir_path: Path,
for file_path in sorted(file_iter): for file_path in sorted(file_iter):
if (file_path.is_file() if (file_path.is_file()
and file_path.name not in ignored_files # noqa: W503 and file_path.name not in ignored_files # noqa: W503
and file_path.suffix in Message.file_suffixes_read): # noqa: W503 and file_path.suffix in Message.file_suffixes): # noqa: W503
try: try:
message = Message.from_file(file_path, mfilter) message = Message.from_file(file_path, mfilter)
if message: if message:
@@ -64,20 +63,22 @@ def read_dir(dir_path: Path,
def make_file_path(dir_path: Path, def make_file_path(dir_path: Path,
file_suffix: str,
next_fid: Callable[[], int]) -> Path: next_fid: Callable[[], int]) -> Path:
""" """
Create a file_path for the given directory using the given ID generator function. Create a file_path for the given directory using the
given file_suffix and ID generator function.
""" """
file_path = dir_path / f"{next_fid():04d}{msg_suffix}" file_path = dir_path / f"{next_fid():04d}{file_suffix}"
while file_path.exists(): while file_path.exists():
file_path = dir_path / f"{next_fid():04d}{msg_suffix}" file_path = dir_path / f"{next_fid():04d}{file_suffix}"
return file_path return file_path
def write_dir(dir_path: Path, def write_dir(dir_path: Path,
messages: list[Message], messages: list[Message],
next_fid: Callable[[], int], file_suffix: str,
mformat: MessageFormat = Message.default_format) -> None: next_fid: Callable[[], int]) -> None:
""" """
Write all messages to the given directory. If a message has no file_path, Write all messages to the given directory. If a message has no file_path,
a new one will be created. If message.file_path exists, it will be modified a new one will be created. If message.file_path exists, it will be modified
@@ -85,17 +86,18 @@ def write_dir(dir_path: Path,
Parameters: Parameters:
* 'dir_path': destination directory * 'dir_path': destination directory
* 'messages': list of messages to write * 'messages': list of messages to write
* 'file_suffix': suffix for the message files ['.txt'|'.yaml']
* 'next_fid': callable that returns the next file ID * 'next_fid': callable that returns the next file ID
""" """
for message in messages: for message in messages:
file_path = message.file_path file_path = message.file_path
# message has no file_path: create one # message has no file_path: create one
if not file_path: if not file_path:
file_path = make_file_path(dir_path, next_fid) file_path = make_file_path(dir_path, file_suffix, next_fid)
# file_path does not point to given directory: modify it # file_path does not point to given directory: modify it
elif not file_path.parent.samefile(dir_path): elif not file_path.parent.samefile(dir_path):
file_path = dir_path / file_path.name file_path = dir_path / file_path.name
message.to_file(file_path, mformat=mformat) message.to_file(file_path)
def clear_dir(dir_path: Path, def clear_dir(dir_path: Path,
@@ -107,7 +109,7 @@ def clear_dir(dir_path: Path,
for file_path in file_iter: for file_path in file_iter:
if (file_path.is_file() if (file_path.is_file()
and file_path.name not in ignored_files # noqa: W503 and file_path.name not in ignored_files # noqa: W503
and file_path.suffix in Message.file_suffixes_read): # noqa: W503 and file_path.suffix in Message.file_suffixes): # noqa: W503
file_path.unlink(missing_ok=True) file_path.unlink(missing_ok=True)
@@ -144,7 +146,7 @@ class Chat:
Matching is True if: Matching is True if:
* 'name' matches the full 'file_path' * 'name' matches the full 'file_path'
* 'name' matches 'file_path.name' (i. e. including the suffix) * 'name' matches 'file_path.name' (i. e. including the suffix)
* 'name' matches 'file_path.stem' (i. e. without the suffix) * 'name' matches 'file_path.stem' (i. e. without a suffix)
""" """
return Path(name) == file_path or name == file_path.name or name == file_path.stem return Path(name) == file_path or name == file_path.name or name == file_path.stem
@@ -255,17 +257,14 @@ class Chat:
return sum(m.tokens() for m in self.messages) return sum(m.tokens() for m in self.messages)
def print(self, source_code_only: bool = False, def print(self, source_code_only: bool = False,
with_metadata: bool = False, with_tags: bool = False, with_files: bool = False,
paged: bool = True, paged: bool = True) -> None:
tight: bool = False) -> None:
output: list[str] = [] output: list[str] = []
for message in self.messages: for message in self.messages:
if source_code_only: if source_code_only:
output.append(message.to_str(source_code_only=True)) output.append(message.to_str(source_code_only=True))
continue continue
output.append(message.to_str(with_metadata)) output.append(message.to_str(with_tags, with_files))
if not tight:
output.append('\n' + ('-' * terminal_width()) + '\n')
if paged: if paged:
print_paged('\n'.join(output)) print_paged('\n'.join(output))
else: else:
@@ -282,14 +281,15 @@ class ChatDB(Chat):
persistently. persistently.
""" """
default_file_suffix: ClassVar[str] = '.txt'
cache_path: Path cache_path: Path
db_path: Path db_path: Path
# a MessageFilter that all messages must match (if given) # a MessageFilter that all messages must match (if given)
mfilter: Optional[MessageFilter] = None mfilter: Optional[MessageFilter] = None
file_suffix: str = default_file_suffix
# the glob pattern for all messages # the glob pattern for all messages
glob: Optional[str] = None glob: Optional[str] = None
# message format (for writing)
mformat: MessageFormat = Message.default_format
def __post_init__(self) -> None: def __post_init__(self) -> None:
# contains the latest message ID # contains the latest message ID
@@ -317,7 +317,8 @@ class ChatDB(Chat):
when reading them. when reading them.
""" """
messages = read_dir(db_path, glob, mfilter) messages = read_dir(db_path, glob, mfilter)
return cls(messages, cache_path, db_path, mfilter, glob) return cls(messages, cache_path, db_path, mfilter,
cls.default_file_suffix, glob)
@classmethod @classmethod
def from_messages(cls: Type[ChatDBInst], def from_messages(cls: Type[ChatDBInst],
@@ -344,17 +345,7 @@ class ChatDB(Chat):
with open(self.next_path, 'w') as f: with open(self.next_path, 'w') as f:
f.write(f'{fid}') f.write(f'{fid}')
def set_msg_format(self, mformat: MessageFormat) -> None: def msg_write(self, messages: Optional[list[Message]] = None) -> None:
"""
Set message format for writing messages.
"""
if mformat not in message_valid_formats:
raise ChatError(f"Message format '{mformat}' is not supported")
self.mformat = mformat
def msg_write(self,
messages: Optional[list[Message]] = None,
mformat: Optional[MessageFormat] = None) -> None:
""" """
Write either the given messages or the internal ones to their CURRENT file_path. Write either the given messages or the internal ones to their CURRENT file_path.
If messages are given, they all must have a valid file_path. When writing the If messages are given, they all must have a valid file_path. When writing the
@@ -365,7 +356,7 @@ class ChatDB(Chat):
raise ChatError("Can't write files without a valid file_path") raise ChatError("Can't write files without a valid file_path")
msgs = iter(messages if messages else self.messages) msgs = iter(messages if messages else self.messages)
while (m := next(msgs, None)): while (m := next(msgs, None)):
m.to_file(mformat=mformat if mformat else self.mformat) m.to_file()
def msg_update(self, messages: list[Message], write: bool = True) -> None: def msg_update(self, messages: list[Message], write: bool = True) -> None:
""" """
@@ -386,7 +377,6 @@ class ChatDB(Chat):
def msg_gather(self, def msg_gather(self,
loc: msg_location, loc: msg_location,
require_file_path: bool = False, require_file_path: bool = False,
glob: Optional[str] = None,
mfilter: Optional[MessageFilter] = None) -> list[Message]: mfilter: Optional[MessageFilter] = None) -> list[Message]:
""" """
Gather and return messages from the given locations: Gather and return messages from the given locations:
@@ -405,9 +395,9 @@ class ChatDB(Chat):
else: else:
loc_messages += [m for m in self.messages if (mfilter is None or m.match(mfilter))] loc_messages += [m for m in self.messages if (mfilter is None or m.match(mfilter))]
if loc in ['cache', 'disk', 'all']: if loc in ['cache', 'disk', 'all']:
loc_messages += read_dir(self.cache_path, glob=glob, mfilter=mfilter) loc_messages += read_dir(self.cache_path, mfilter=mfilter)
if loc in ['db', 'disk', 'all']: if loc in ['db', 'disk', 'all']:
loc_messages += read_dir(self.db_path, glob=glob, mfilter=mfilter) loc_messages += read_dir(self.db_path, mfilter=mfilter)
# remove_duplicates and sort the list # remove_duplicates and sort the list
unique_messages: list[Message] = [] unique_messages: list[Message] = []
for m in loc_messages: for m in loc_messages:
@@ -528,8 +518,8 @@ class ChatDB(Chat):
""" """
write_dir(self.cache_path, write_dir(self.cache_path,
messages if messages else self.messages, messages if messages else self.messages,
self.get_next_fid, self.file_suffix,
self.mformat) self.get_next_fid)
def cache_add(self, messages: list[Message], write: bool = True) -> None: def cache_add(self, messages: list[Message], write: bool = True) -> None:
""" """
@@ -541,11 +531,11 @@ class ChatDB(Chat):
if write: if write:
write_dir(self.cache_path, write_dir(self.cache_path,
messages, messages,
self.get_next_fid, self.file_suffix,
self.mformat) self.get_next_fid)
else: else:
for m in messages: for m in messages:
m.file_path = make_file_path(self.cache_path, self.get_next_fid) m.file_path = make_file_path(self.cache_path, self.default_file_suffix, self.get_next_fid)
self.messages += messages self.messages += messages
self.msg_sort() self.msg_sort()
@@ -595,8 +585,8 @@ class ChatDB(Chat):
""" """
write_dir(self.db_path, write_dir(self.db_path,
messages if messages else self.messages, messages if messages else self.messages,
self.get_next_fid, self.file_suffix,
self.mformat) self.get_next_fid)
def db_add(self, messages: list[Message], write: bool = True) -> None: def db_add(self, messages: list[Message], write: bool = True) -> None:
""" """
@@ -608,11 +598,11 @@ class ChatDB(Chat):
if write: if write:
write_dir(self.db_path, write_dir(self.db_path,
messages, messages,
self.get_next_fid, self.file_suffix,
self.mformat) self.get_next_fid)
else: else:
for m in messages: for m in messages:
m.file_path = make_file_path(self.db_path, self.get_next_fid) m.file_path = make_file_path(self.db_path, self.default_file_suffix, self.get_next_fid)
self.messages += messages self.messages += messages
self.msg_sort() self.msg_sort()
+5 -54
View File
@@ -1,51 +1,13 @@
import sys
import argparse import argparse
from pathlib import Path from pathlib import Path
from ..configuration import Config from ..configuration import Config
from ..chat import ChatDB from ..chat import ChatDB
from ..message import MessageFilter, Message from ..message import MessageFilter
msg_suffix = Message.file_suffix_write # currently '.msg' def hist_cmd(args: argparse.Namespace, config: Config) -> None:
def convert_messages(args: argparse.Namespace, config: Config) -> None:
""" """
Convert messages to a new format. Also used to change old suffixes Handler for the 'hist' command.
('.txt', '.yaml') to the latest default message file suffix ('.msg').
"""
chat = ChatDB.from_dir(Path(config.cache),
Path(config.db))
# read all known message files
msgs = chat.msg_gather(loc='disk', glob='*.*')
# make a set of all message IDs
msg_ids = set([m.msg_id() for m in msgs])
# set requested format and write all messages
chat.set_msg_format(args.convert)
# delete the current suffix
# -> a new one will automatically be created
for m in msgs:
if m.file_path:
m.file_path = m.file_path.with_suffix('')
chat.msg_write(msgs)
# read all messages with the current default suffix
msgs = chat.msg_gather(loc='disk', glob=f'*{msg_suffix}')
# make sure we converted all of the original messages
for mid in msg_ids:
if not any(mid == m.msg_id() for m in msgs):
print(f"Message '{mid}' has not been found after conversion. Aborting.")
sys.exit(1)
# delete messages with old suffixes
msgs = chat.msg_gather(loc='disk', glob='*.*')
for m in msgs:
if m.file_path and m.file_path.suffix != msg_suffix:
m.rm_file()
print(f"Successfully converted {len(msg_ids)} messages.")
def print_chat(args: argparse.Namespace, config: Config) -> None:
"""
Print the DB chat history.
""" """
mfilter = MessageFilter(tags_or=args.or_tags, mfilter = MessageFilter(tags_or=args.or_tags,
@@ -57,16 +19,5 @@ def print_chat(args: argparse.Namespace, config: Config) -> None:
Path(config.db), Path(config.db),
mfilter=mfilter) mfilter=mfilter)
chat.print(args.source_code_only, chat.print(args.source_code_only,
args.with_metadata, args.with_tags,
paged=not args.no_paging, args.with_files)
tight=args.tight)
def hist_cmd(args: argparse.Namespace, config: Config) -> None:
"""
Handler for the 'hist' command.
"""
if args.print:
print_chat(args, config)
elif args.convert:
convert_messages(args, config)
+18 -84
View File
@@ -2,7 +2,6 @@ import sys
import argparse import argparse
from pathlib import Path from pathlib import Path
from itertools import zip_longest from itertools import zip_longest
from copy import deepcopy
from ..configuration import Config from ..configuration import Config
from ..chat import ChatDB from ..chat import ChatDB
from ..message import Message, MessageFilter, MessageError, Question, source_code from ..message import Message, MessageFilter, MessageError, Question, source_code
@@ -10,10 +9,6 @@ from ..ai_factory import create_ai
from ..ai import AI, AIResponse from ..ai import AI, AIResponse
class QuestionCmdError(Exception):
pass
def add_file_as_text(question_parts: list[str], file: str) -> None: def add_file_as_text(question_parts: list[str], file: str) -> None:
""" """
Add the given file as plain text to the question part list. Add the given file as plain text to the question part list.
@@ -55,41 +50,13 @@ def add_file_as_code(question_parts: list[str], file: str) -> None:
question_parts.append(f"```\n{content}\n```") question_parts.append(f"```\n{content}\n```")
def create_msg_args(msg: Message, args: argparse.Namespace) -> argparse.Namespace:
"""
Takes an existing message and CLI arguments, and returns modified args based
on the members of the given message. Used e.g. when repeating messages, where
it's necessary to determine the correct AI, module and output tags to use
(either from the existing message or the given args).
"""
msg_args = args
# if AI, model or output tags have not been specified,
# use those from the original message
if (args.AI is None
or args.model is None # noqa: W503
or args.output_tags is None): # noqa: W503
msg_args = deepcopy(args)
if args.AI is None and msg.ai is not None:
msg_args.AI = msg.ai
if args.model is None and msg.model is not None:
msg_args.model = msg.model
if args.output_tags is None and msg.tags is not None:
msg_args.output_tags = msg.tags
return msg_args
def create_message(chat: ChatDB, args: argparse.Namespace) -> Message: def create_message(chat: ChatDB, args: argparse.Namespace) -> Message:
""" """
Create a new message from the given arguments and write it Create a new message from the given arguments and write it
to the cache directory. to the cache directory.
""" """
question_parts = [] question_parts = []
if args.create is not None: question_list = args.ask if args.ask is not None else []
question_list = args.create
elif args.ask is not None:
question_list = args.ask
else:
raise QuestionCmdError("No question found")
text_files = args.source_text if args.source_text is not None else [] text_files = args.source_text if args.source_text is not None else []
code_files = args.source_code if args.source_code is not None else [] code_files = args.source_code if args.source_code is not None else []
@@ -138,52 +105,13 @@ def make_request(ai: AI, chat: ChatDB, message: Message, args: argparse.Namespac
print(response.tokens) print(response.tokens)
def repeat_messages(messages: list[Message], chat: ChatDB, args: argparse.Namespace, config: Config) -> None:
"""
Repeat the given messages using the given arguments.
"""
ai: AI
for msg in messages:
msg_args = create_msg_args(msg, args)
ai = create_ai(msg_args, config)
print(f"--------- Repeating message '{msg.msg_id()}': ---------")
# overwrite the latest message if requested or empty
# -> but not if it's in the DB!
if ((msg.answer is None or msg_args.overwrite is True)
and (not chat.msg_in_db(msg))): # noqa: W503
msg.clear_answer()
make_request(ai, chat, msg, msg_args)
# otherwise create a new one
else:
msg_args.ask = [msg.question]
message = create_message(chat, msg_args)
make_request(ai, chat, message, msg_args)
def invert_input_tag_args(args: argparse.Namespace) -> None:
"""
Changes the semantics of the INPUT tags for this command:
* not tags specified on the CLI -> no tags are selected
* empty tags specified on the CLI -> all tags are selected
"""
if args.or_tags is None:
args.or_tags = set()
elif len(args.or_tags) == 0:
args.or_tags = None
if args.and_tags is None:
args.and_tags = set()
elif len(args.and_tags) == 0:
args.and_tags = None
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.
""" """
invert_input_tag_args(args) mfilter = MessageFilter(tags_or=args.or_tags if args.or_tags is not None else set(),
mfilter = MessageFilter(tags_or=args.or_tags, tags_and=args.and_tags if args.and_tags is not None else set(),
tags_and=args.and_tags, tags_not=args.exclude_tags if args.exclude_tags is not None else set())
tags_not=args.exclude_tags)
chat = ChatDB.from_dir(cache_path=Path(config.cache), chat = ChatDB.from_dir(cache_path=Path(config.cache),
db_path=Path(config.db), db_path=Path(config.db),
mfilter=mfilter) mfilter=mfilter)
@@ -193,24 +121,30 @@ def question_cmd(args: argparse.Namespace, config: Config) -> None:
if args.create: if args.create:
return return
# create the correct AI instance
ai: AI = create_ai(args, config)
# === ASK === # === ASK ===
if args.ask: if args.ask:
ai: AI = create_ai(args, config)
make_request(ai, chat, message, args) make_request(ai, chat, message, args)
# === REPEAT === # === REPEAT ===
elif args.repeat is not None: elif args.repeat is not None:
repeat_msgs: list[Message] = []
# repeat latest message
if len(args.repeat) == 0:
lmessage = chat.msg_latest(loc='cache') lmessage = chat.msg_latest(loc='cache')
if lmessage is None: if lmessage is None:
print("No message found to repeat!") print("No message found to repeat!")
sys.exit(1) sys.exit(1)
repeat_msgs.append(lmessage)
# repeat given message(s)
else: else:
repeat_msgs = chat.msg_find(args.repeat, loc='disk') print(f"Repeating message '{lmessage.msg_id()}':")
repeat_messages(repeat_msgs, chat, args, config) # 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 === # === 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
+9 -12
View File
@@ -34,13 +34,13 @@ def create_parser() -> argparse.ArgumentParser:
# a parent parser for all commands that support tag selection # a parent parser for all commands that support tag selection
tag_parser = argparse.ArgumentParser(add_help=False) tag_parser = argparse.ArgumentParser(add_help=False)
tag_arg = tag_parser.add_argument('-t', '--or-tags', nargs='*', tag_arg = tag_parser.add_argument('-t', '--or-tags', nargs='+',
help='List of tags (one must match)', metavar='OTAGS') help='List of tags (one must match)', metavar='OTAGS')
tag_arg.completer = tags_completer # type: ignore tag_arg.completer = tags_completer # type: ignore
atag_arg = tag_parser.add_argument('-k', '--and-tags', nargs='*', atag_arg = tag_parser.add_argument('-k', '--and-tags', nargs='+',
help='List of tags (all must match)', metavar='ATAGS') help='List of tags (all must match)', metavar='ATAGS')
atag_arg.completer = tags_completer # type: ignore atag_arg.completer = tags_completer # type: ignore
etag_arg = tag_parser.add_argument('-x', '--exclude-tags', nargs='*', etag_arg = tag_parser.add_argument('-x', '--exclude-tags', nargs='+',
help='List of tags to exclude', metavar='XTAGS') help='List of tags to exclude', metavar='XTAGS')
etag_arg.completer = tags_completer # type: ignore etag_arg.completer = tags_completer # type: ignore
otag_arg = tag_parser.add_argument('-o', '--output-tags', nargs='+', otag_arg = tag_parser.add_argument('-o', '--output-tags', nargs='+',
@@ -73,20 +73,17 @@ def create_parser() -> argparse.ArgumentParser:
# 'hist' command parser # 'hist' command parser
hist_cmd_parser = cmdparser.add_parser('hist', parents=[tag_parser], hist_cmd_parser = cmdparser.add_parser('hist', parents=[tag_parser],
help="Print and manage chat history.", help="Print chat history.",
aliases=['h']) aliases=['h'])
hist_cmd_parser.set_defaults(func=hist_cmd) hist_cmd_parser.set_defaults(func=hist_cmd)
hist_group = hist_cmd_parser.add_mutually_exclusive_group(required=True) hist_cmd_parser.add_argument('-w', '--with-tags', help="Print chat history with tags.",
hist_group.add_argument('-p', '--print', help='Print the DB chat history', action='store_true') action='store_true')
hist_group.add_argument('-c', '--convert', help='Convert all message files to the given format [txt|yaml]', metavar='FORMAT') hist_cmd_parser.add_argument('-W', '--with-files', help="Print chat history with filenames.",
hist_cmd_parser.add_argument('-w', '--with-metadata', help="Print chat history with metadata (tags, filename, AI, etc.).",
action='store_true') action='store_true')
hist_cmd_parser.add_argument('-S', '--source-code-only', help='Only print embedded source code', hist_cmd_parser.add_argument('-S', '--source-code-only', help='Only print embedded source code',
action='store_true') action='store_true')
hist_cmd_parser.add_argument('-A', '--answer', help='Print only answers with given substring', metavar='SUBSTRING') hist_cmd_parser.add_argument('-A', '--answer', help='Search for answer substring', metavar='SUBSTRING')
hist_cmd_parser.add_argument('-Q', '--question', help='Print only questions with given substring', metavar='SUBSTRING') hist_cmd_parser.add_argument('-Q', '--question', help='Search for question substring', metavar='SUBSTRING')
hist_cmd_parser.add_argument('-d', '--tight', help='Print without message separators', action='store_true')
hist_cmd_parser.add_argument('-P', '--no-paging', help='Print without paging', action='store_true')
# 'tags' command parser # 'tags' command parser
tags_cmd_parser = cmdparser.add_parser('tags', tags_cmd_parser = cmdparser.add_parser('tags',
+33 -50
View File
@@ -5,8 +5,7 @@ import pathlib
import yaml import yaml
import tempfile import tempfile
import shutil import shutil
from typing import Type, TypeVar, ClassVar, Optional, Any, Union, Final, Literal, Iterable, Tuple from typing import Type, TypeVar, ClassVar, Optional, Any, Union, Final, Literal, Iterable
from typing import get_args as typing_get_args
from dataclasses import dataclass, asdict, field from dataclasses import dataclass, asdict, field
from .tags import Tag, TagLine, TagError, match_tags, rename_tags from .tags import Tag, TagLine, TagError, match_tags, rename_tags
@@ -16,9 +15,6 @@ MessageInst = TypeVar('MessageInst', bound='Message')
AILineInst = TypeVar('AILineInst', bound='AILine') AILineInst = TypeVar('AILineInst', bound='AILine')
ModelLineInst = TypeVar('ModelLineInst', bound='ModelLine') ModelLineInst = TypeVar('ModelLineInst', bound='ModelLine')
YamlDict = dict[str, Union[QuestionInst, AnswerInst, set[Tag]]] YamlDict = dict[str, Union[QuestionInst, AnswerInst, set[Tag]]]
MessageFormat = Literal['txt', 'yaml']
message_valid_formats: Final[Tuple[MessageFormat, ...]] = typing_get_args(MessageFormat)
message_default_format: Final[MessageFormat] = 'txt'
class MessageError(Exception): class MessageError(Exception):
@@ -96,7 +92,7 @@ class MessageFilter:
class AILine(str): class AILine(str):
""" """
A line that represents the AI name in the 'txt' format. A line that represents the AI name in a '.txt' file..
""" """
prefix: Final[str] = 'AI:' prefix: Final[str] = 'AI:'
@@ -116,7 +112,7 @@ class AILine(str):
class ModelLine(str): class ModelLine(str):
""" """
A line that represents the model name in the 'txt' format. A line that represents the model name in a '.txt' file..
""" """
prefix: Final[str] = 'MODEL:' prefix: Final[str] = 'MODEL:'
@@ -220,9 +216,7 @@ class Message():
model: Optional[str] = field(default=None, compare=False) model: Optional[str] = field(default=None, compare=False)
file_path: Optional[pathlib.Path] = field(default=None, compare=False) file_path: Optional[pathlib.Path] = field(default=None, compare=False)
# class variables # class variables
file_suffixes_read: ClassVar[list[str]] = ['.msg', '.txt', '.yaml'] file_suffixes: ClassVar[list[str]] = ['.txt', '.yaml']
file_suffix_write: ClassVar[str] = '.msg'
default_format: ClassVar[MessageFormat] = message_default_format
tags_yaml_key: ClassVar[str] = 'tags' tags_yaml_key: ClassVar[str] = 'tags'
file_yaml_key: ClassVar[str] = 'file_path' file_yaml_key: ClassVar[str] = 'file_path'
ai_yaml_key: ClassVar[str] = 'ai' ai_yaml_key: ClassVar[str] = 'ai'
@@ -282,8 +276,16 @@ class Message():
tags: set[Tag] = set() tags: set[Tag] = set()
if not file_path.exists(): if not file_path.exists():
raise MessageError(f"Message file '{file_path}' does not exist") raise MessageError(f"Message file '{file_path}' does not exist")
if file_path.suffix not in cls.file_suffixes_read: if file_path.suffix not in cls.file_suffixes:
raise MessageError(f"File type '{file_path.suffix}' is not supported") raise MessageError(f"File type '{file_path.suffix}' is not supported")
# for TXT, it's enough to read the TagLine
if file_path.suffix == '.txt':
with open(file_path, "r") as fd:
try:
tags = TagLine(fd.readline()).tags(prefix, contain)
except TagError:
pass # message without tags
else: # '.yaml'
try: try:
message = cls.from_file(file_path) message = cls.from_file(file_path)
if message: if message:
@@ -326,16 +328,15 @@ class Message():
""" """
if not file_path.exists(): if not file_path.exists():
raise MessageError(f"Message file '{file_path}' does not exist") raise MessageError(f"Message file '{file_path}' does not exist")
if file_path.suffix not in cls.file_suffixes_read: if file_path.suffix not in cls.file_suffixes:
raise MessageError(f"File type '{file_path.suffix}' is not supported") raise MessageError(f"File type '{file_path.suffix}' is not supported")
# try TXT first
try: if file_path.suffix == '.txt':
message = cls.__from_file_txt(file_path, message = cls.__from_file_txt(file_path,
mfilter.tags_or if mfilter else None, mfilter.tags_or if mfilter else None,
mfilter.tags_and if mfilter else None, mfilter.tags_and if mfilter else None,
mfilter.tags_not if mfilter else None) mfilter.tags_not if mfilter else None)
# then YAML else:
except MessageError:
message = cls.__from_file_yaml(file_path) message = cls.__from_file_yaml(file_path)
if message and (mfilter is None or message.match(mfilter)): if message and (mfilter is None or message.match(mfilter)):
return message return message
@@ -372,6 +373,10 @@ class Message():
tags = TagLine(fd.readline()).tags() tags = TagLine(fd.readline()).tags()
except TagError: except TagError:
fd.seek(pos) fd.seek(pos)
if tags_or or tags_and or tags_not:
# match with an empty set if the file has no tags
if not match_tags(tags, tags_or, tags_and, tags_not):
return None
# AILine (Optional) # AILine (Optional)
try: try:
pos = fd.tell() pos = fd.tell()
@@ -396,12 +401,6 @@ class Message():
answer = Answer.from_list(text[answer_idx + 1:]) answer = Answer.from_list(text[answer_idx + 1:])
except ValueError: except ValueError:
question = Question.from_list(text[question_idx:]) question = Question.from_list(text[question_idx:])
# match tags AFTER reading the whole file
# -> make sure it's a valid 'txt' file format
if tags_or or tags_and or tags_not:
# match with an empty set if the file has no tags
if not match_tags(tags, tags_or, tags_and, tags_not):
return None
return cls(question, answer, tags, ai, model, file_path) return cls(question, answer, tags, ai, model, file_path)
@classmethod @classmethod
@@ -422,7 +421,7 @@ class Message():
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")
def to_str(self, with_metadata: 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:
""" """
Return the current Message as a string. Return the current Message as a string.
""" """
@@ -432,11 +431,10 @@ class Message():
if self.answer: if self.answer:
output.extend(self.answer.source_code(include_delims=True)) output.extend(self.answer.source_code(include_delims=True))
return '\n'.join(output) if len(output) > 0 else '' return '\n'.join(output) if len(output) > 0 else ''
if with_metadata: if with_tags:
output.append(self.tags_str()) output.append(self.tags_str())
if with_file:
output.append('FILE: ' + str(self.file_path)) output.append('FILE: ' + str(self.file_path))
output.append('AI: ' + str(self.ai))
output.append('MODEL: ' + str(self.model))
output.append(Question.txt_header) output.append(Question.txt_header)
output.append(self.question) output.append(self.question)
if self.answer: if self.answer:
@@ -444,29 +442,21 @@ class Message():
output.append(self.answer) output.append(self.answer)
return '\n'.join(output) return '\n'.join(output)
def to_file(self, file_path: Optional[pathlib.Path]=None, mformat: MessageFormat = message_default_format) -> None: # noqa: 11 def to_file(self, file_path: Optional[pathlib.Path]=None) -> None: # noqa: 11
""" """
Write a Message to the given file. Supported message file formats are 'txt' and 'yaml'. Write a Message to the given file. Type is determined based on the suffix.
Suffix is always '.msg'. Currently supported suffixes: ['.txt', '.yaml']
""" """
if file_path: if file_path:
self.file_path = file_path self.file_path = file_path
if not self.file_path: if not self.file_path:
raise MessageError("Got no valid path to write message") raise MessageError("Got no valid path to write message")
if mformat not in message_valid_formats: if self.file_path.suffix not in self.file_suffixes:
raise MessageError(f"File format '{mformat}' is not supported") raise MessageError(f"File type '{self.file_path.suffix}' is not supported")
# check for valid suffix
# -> add one if it's empty
# -> refuse old or otherwise unsupported suffixes
if not self.file_path.suffix:
self.file_path = self.file_path.with_suffix(self.file_suffix_write)
elif self.file_path.suffix != self.file_suffix_write:
raise MessageError(f"File suffix '{self.file_path.suffix}' is not supported")
# TXT # TXT
if mformat == 'txt': if self.file_path.suffix == '.txt':
return self.__to_file_txt(self.file_path) return self.__to_file_txt(self.file_path)
# YAML elif self.file_path.suffix == '.yaml':
elif mformat == 'yaml':
return self.__to_file_yaml(self.file_path) return self.__to_file_yaml(self.file_path)
def __to_file_txt(self, file_path: pathlib.Path) -> None: def __to_file_txt(self, file_path: pathlib.Path) -> None:
@@ -478,8 +468,8 @@ class Message():
* Model [Optional] * Model [Optional]
* Question.txt_header * Question.txt_header
* Question * Question
* Answer.txt_header [Optional] * Answer.txt_header
* Answer [Optional] * Answer
""" """
with tempfile.NamedTemporaryFile(dir=file_path.parent, prefix=file_path.name, mode="w", delete=False) as temp_fd: with tempfile.NamedTemporaryFile(dir=file_path.parent, prefix=file_path.name, mode="w", delete=False) as temp_fd:
temp_file_path = pathlib.Path(temp_fd.name) temp_file_path = pathlib.Path(temp_fd.name)
@@ -518,13 +508,6 @@ class Message():
yaml.dump(data, temp_fd, sort_keys=False) yaml.dump(data, temp_fd, sort_keys=False)
shutil.move(temp_file_path, file_path) shutil.move(temp_file_path, file_path)
def rm_file(self) -> None:
"""
Delete the message file. Ignore empty file_path and not existing files.
"""
if self.file_path is not None:
self.file_path.unlink(missing_ok=True)
def filter_tags(self, prefix: Optional[str] = None, contain: Optional[str] = None) -> set[Tag]: def filter_tags(self, prefix: Optional[str] = None, contain: Optional[str] = None) -> set[Tag]:
""" """
Filter tags based on their prefix (i. e. the tag starts with a given string) Filter tags based on their prefix (i. e. the tag starts with a given string)
-1
View File
@@ -2,4 +2,3 @@ openai
PyYAML PyYAML
argcomplete argcomplete
pytest pytest
Jinja2
+6 -3
View File
@@ -2,8 +2,6 @@ from setuptools import setup, find_packages
with open("README.md", "r", encoding="utf-8") as fh: with open("README.md", "r", encoding="utf-8") as fh:
long_description = fh.read() long_description = fh.read()
with open("requirements.txt", "r") as fh:
install_requirements = [line.strip() for line in fh]
setup( setup(
name="ChatMastermind", name="ChatMastermind",
@@ -30,7 +28,12 @@ setup(
"Topic :: Utilities", "Topic :: Utilities",
"Topic :: Text Processing", "Topic :: Text Processing",
], ],
install_requires=install_requirements, install_requires=[
"openai",
"PyYAML",
"argcomplete",
"pytest",
],
python_requires=">=3.9", python_requires=">=3.9",
test_suite="tests", test_suite="tests",
entry_points={ entry_points={
+95 -129
View File
@@ -10,20 +10,6 @@ from chatmastermind.message import Message, Question, Answer, Tag, MessageFilter
from chatmastermind.chat import Chat, ChatDB, ChatError from chatmastermind.chat import Chat, ChatDB, ChatError
msg_suffix: str = Message.file_suffix_write
def msg_to_file_force_suffix(msg: Message) -> None:
"""
Force writing a message file with illegal suffixes.
"""
def_suffix = Message.file_suffix_write
assert msg.file_path
Message.file_suffix_write = msg.file_path.suffix
msg.to_file()
Message.file_suffix_write = def_suffix
class TestChatBase(unittest.TestCase): class TestChatBase(unittest.TestCase):
def assert_messages_equal(self, msg1: list[Message], msg2: list[Message]) -> None: def assert_messages_equal(self, msg1: list[Message], msg2: list[Message]) -> None:
""" """
@@ -41,15 +27,11 @@ class TestChat(TestChatBase):
self.message1 = Message(Question('Question 1'), self.message1 = Message(Question('Question 1'),
Answer('Answer 1'), Answer('Answer 1'),
{Tag('atag1'), Tag('btag2')}, {Tag('atag1'), Tag('btag2')},
ai='FakeAI', file_path=pathlib.Path('0001.txt'))
model='FakeModel',
file_path=pathlib.Path(f'0001{msg_suffix}'))
self.message2 = Message(Question('Question 2'), self.message2 = Message(Question('Question 2'),
Answer('Answer 2'), Answer('Answer 2'),
{Tag('btag2')}, {Tag('btag2')},
ai='FakeAI', file_path=pathlib.Path('0002.txt'))
model='FakeModel',
file_path=pathlib.Path(f'0002{msg_suffix}'))
self.maxDiff = None self.maxDiff = None
def test_unique_id(self) -> None: def test_unique_id(self) -> None:
@@ -117,24 +99,24 @@ class TestChat(TestChatBase):
def test_find_remove_messages(self) -> None: def test_find_remove_messages(self) -> None:
self.chat.msg_add([self.message1, self.message2]) self.chat.msg_add([self.message1, self.message2])
msgs = self.chat.msg_find(['0001']) msgs = self.chat.msg_find(['0001.txt'])
self.assertListEqual(msgs, [self.message1]) self.assertListEqual(msgs, [self.message1])
msgs = self.chat.msg_find(['0001', '0002']) msgs = self.chat.msg_find(['0001.txt', '0002.txt'])
self.assertListEqual(msgs, [self.message1, self.message2]) self.assertListEqual(msgs, [self.message1, self.message2])
# add new Message with full path # add new Message with full path
message3 = Message(Question('Question 2'), message3 = Message(Question('Question 2'),
Answer('Answer 2'), Answer('Answer 2'),
{Tag('btag2')}, {Tag('btag2')},
file_path=pathlib.Path(f'/foo/bla/0003{msg_suffix}')) file_path=pathlib.Path('/foo/bla/0003.txt'))
self.chat.msg_add([message3]) self.chat.msg_add([message3])
# find new Message by full path # find new Message by full path
msgs = self.chat.msg_find([f'/foo/bla/0003{msg_suffix}']) msgs = self.chat.msg_find(['/foo/bla/0003.txt'])
self.assertListEqual(msgs, [message3]) self.assertListEqual(msgs, [message3])
# find Message with full path only by filename # find Message with full path only by filename
msgs = self.chat.msg_find([f'0003{msg_suffix}']) msgs = self.chat.msg_find(['0003.txt'])
self.assertListEqual(msgs, [message3]) self.assertListEqual(msgs, [message3])
# remove last message # remove last message
self.chat.msg_remove(['0003']) self.chat.msg_remove(['0003.txt'])
self.assertListEqual(self.chat.messages, [self.message1, self.message2]) self.assertListEqual(self.chat.messages, [self.message1, self.message2])
def test_latest_message(self) -> None: def test_latest_message(self) -> None:
@@ -147,7 +129,7 @@ class TestChat(TestChatBase):
@patch('sys.stdout', new_callable=StringIO) @patch('sys.stdout', new_callable=StringIO)
def test_print(self, mock_stdout: StringIO) -> None: def test_print(self, mock_stdout: StringIO) -> None:
self.chat.msg_add([self.message1, self.message2]) self.chat.msg_add([self.message1, self.message2])
self.chat.print(paged=False, tight=True) self.chat.print(paged=False)
expected_output = f"""{Question.txt_header} expected_output = f"""{Question.txt_header}
Question 1 Question 1
{Answer.txt_header} {Answer.txt_header}
@@ -160,21 +142,17 @@ Answer 2
self.assertEqual(mock_stdout.getvalue(), expected_output) self.assertEqual(mock_stdout.getvalue(), expected_output)
@patch('sys.stdout', new_callable=StringIO) @patch('sys.stdout', new_callable=StringIO)
def test_print_with_metadata(self, mock_stdout: StringIO) -> None: def test_print_with_tags_and_file(self, mock_stdout: StringIO) -> None:
self.chat.msg_add([self.message1, self.message2]) self.chat.msg_add([self.message1, self.message2])
self.chat.print(paged=False, with_metadata=True, tight=True) self.chat.print(paged=False, with_tags=True, with_files=True)
expected_output = f"""{TagLine.prefix} atag1 btag2 expected_output = f"""{TagLine.prefix} atag1 btag2
FILE: 0001{msg_suffix} FILE: 0001.txt
AI: FakeAI
MODEL: FakeModel
{Question.txt_header} {Question.txt_header}
Question 1 Question 1
{Answer.txt_header} {Answer.txt_header}
Answer 1 Answer 1
{TagLine.prefix} btag2 {TagLine.prefix} btag2
FILE: 0002{msg_suffix} FILE: 0002.txt
AI: FakeAI
MODEL: FakeModel
{Question.txt_header} {Question.txt_header}
Question 2 Question 2
{Answer.txt_header} {Answer.txt_header}
@@ -190,27 +168,31 @@ class TestChatDB(TestChatBase):
self.message1 = Message(Question('Question 1'), self.message1 = Message(Question('Question 1'),
Answer('Answer 1'), Answer('Answer 1'),
{Tag('tag1')}) {Tag('tag1')},
file_path=pathlib.Path('0001.txt'))
self.message2 = Message(Question('Question 2'), self.message2 = Message(Question('Question 2'),
Answer('Answer 2'), Answer('Answer 2'),
{Tag('tag2')}) {Tag('tag2')},
file_path=pathlib.Path('0002.yaml'))
self.message3 = Message(Question('Question 3'), self.message3 = Message(Question('Question 3'),
Answer('Answer 3'), Answer('Answer 3'),
{Tag('tag3')}) {Tag('tag3')},
file_path=pathlib.Path('0003.txt'))
self.message4 = Message(Question('Question 4'), self.message4 = Message(Question('Question 4'),
Answer('Answer 4'), Answer('Answer 4'),
{Tag('tag4')}) {Tag('tag4')},
file_path=pathlib.Path('0004.yaml'))
self.message1.to_file(pathlib.Path(self.db_path.name, '0001'), mformat='txt') self.message1.to_file(pathlib.Path(self.db_path.name, '0001.txt'))
self.message2.to_file(pathlib.Path(self.db_path.name, '0002'), mformat='yaml') self.message2.to_file(pathlib.Path(self.db_path.name, '0002.yaml'))
self.message3.to_file(pathlib.Path(self.db_path.name, '0003'), mformat='txt') self.message3.to_file(pathlib.Path(self.db_path.name, '0003.txt'))
self.message4.to_file(pathlib.Path(self.db_path.name, '0004'), mformat='yaml') self.message4.to_file(pathlib.Path(self.db_path.name, '0004.yaml'))
# make the next FID match the current state # make the next FID match the current state
next_fname = pathlib.Path(self.db_path.name) / '.next' next_fname = pathlib.Path(self.db_path.name) / '.next'
with open(next_fname, 'w') as f: with open(next_fname, 'w') as f:
f.write('4') f.write('4')
# add some "trash" in order to test if it's correctly handled / ignored # add some "trash" in order to test if it's correctly handled / ignored
self.trash_files = ['.config.yaml', 'foo.yaml', 'bla.txt', 'fubar.msg'] self.trash_files = ['.config.yaml', 'foo.yaml', 'bla.txt']
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')
@@ -225,7 +207,7 @@ class TestChatDB(TestChatBase):
List all Message files in the given TemporaryDirectory. List all Message files in the given TemporaryDirectory.
""" """
# exclude '.next' # exclude '.next'
return [f for f in pathlib.Path(tmp_dir.name).glob('*.[tym]*') if f.name not in self.trash_files] return [f for f in pathlib.Path(tmp_dir.name).glob('*.[ty]*') if f.name not in self.trash_files]
def tearDown(self) -> None: def tearDown(self) -> None:
self.db_path.cleanup() self.db_path.cleanup()
@@ -236,31 +218,13 @@ class TestChatDB(TestChatBase):
duplicate_message = Message(Question('Question 4'), duplicate_message = Message(Question('Question 4'),
Answer('Answer 4'), Answer('Answer 4'),
{Tag('tag4')}, {Tag('tag4')},
file_path=pathlib.Path(self.db_path.name, '0004.txt')) file_path=pathlib.Path('0004.txt'))
msg_to_file_force_suffix(duplicate_message) duplicate_message.to_file(pathlib.Path(self.db_path.name, '0004.txt'))
with self.assertRaises(ChatError) as cm: with self.assertRaises(ChatError) as cm:
ChatDB.from_dir(pathlib.Path(self.cache_path.name), ChatDB.from_dir(pathlib.Path(self.cache_path.name),
pathlib.Path(self.db_path.name)) pathlib.Path(self.db_path.name))
self.assertEqual(str(cm.exception), "Validation failed") self.assertEqual(str(cm.exception), "Validation failed")
def test_file_path_ID_exists(self) -> None:
"""
Tests if the CacheDB chooses another ID if a file path with
the given one exists.
"""
# create a new and empty CacheDB
db_path = tempfile.TemporaryDirectory()
cache_path = tempfile.TemporaryDirectory()
chat_db = ChatDB.from_dir(pathlib.Path(cache_path.name),
pathlib.Path(db_path.name))
# add a message file
message = Message(Question('What?'),
file_path=pathlib.Path(cache_path.name) / f'0001{msg_suffix}')
message.to_file()
message1 = Message(Question('Where?'))
chat_db.cache_write([message1])
self.assertEqual(message1.msg_id(), '0002')
def test_from_dir(self) -> None: def test_from_dir(self) -> None:
chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name), chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),
pathlib.Path(self.db_path.name)) pathlib.Path(self.db_path.name))
@@ -269,23 +233,25 @@ class TestChatDB(TestChatBase):
self.assertEqual(chat_db.db_path, pathlib.Path(self.db_path.name)) self.assertEqual(chat_db.db_path, pathlib.Path(self.db_path.name))
# check that the files are sorted # check that the files are sorted
self.assertEqual(chat_db.messages[0].file_path, self.assertEqual(chat_db.messages[0].file_path,
pathlib.Path(self.db_path.name, f'0001{msg_suffix}')) pathlib.Path(self.db_path.name, '0001.txt'))
self.assertEqual(chat_db.messages[1].file_path, self.assertEqual(chat_db.messages[1].file_path,
pathlib.Path(self.db_path.name, f'0002{msg_suffix}')) pathlib.Path(self.db_path.name, '0002.yaml'))
self.assertEqual(chat_db.messages[2].file_path, self.assertEqual(chat_db.messages[2].file_path,
pathlib.Path(self.db_path.name, f'0003{msg_suffix}')) pathlib.Path(self.db_path.name, '0003.txt'))
self.assertEqual(chat_db.messages[3].file_path, self.assertEqual(chat_db.messages[3].file_path,
pathlib.Path(self.db_path.name, f'0004{msg_suffix}')) pathlib.Path(self.db_path.name, '0004.yaml'))
def test_from_dir_glob(self) -> None: def test_from_dir_glob(self) -> None:
chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name), chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),
pathlib.Path(self.db_path.name), pathlib.Path(self.db_path.name),
glob='*1.*') glob='*.txt')
self.assertEqual(len(chat_db.messages), 1) self.assertEqual(len(chat_db.messages), 2)
self.assertEqual(chat_db.cache_path, pathlib.Path(self.cache_path.name)) self.assertEqual(chat_db.cache_path, pathlib.Path(self.cache_path.name))
self.assertEqual(chat_db.db_path, pathlib.Path(self.db_path.name)) self.assertEqual(chat_db.db_path, pathlib.Path(self.db_path.name))
self.assertEqual(chat_db.messages[0].file_path, self.assertEqual(chat_db.messages[0].file_path,
pathlib.Path(self.db_path.name, f'0001{msg_suffix}')) pathlib.Path(self.db_path.name, '0001.txt'))
self.assertEqual(chat_db.messages[1].file_path,
pathlib.Path(self.db_path.name, '0003.txt'))
def test_from_dir_filter_tags(self) -> None: def test_from_dir_filter_tags(self) -> None:
chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name), chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),
@@ -295,7 +261,7 @@ class TestChatDB(TestChatBase):
self.assertEqual(chat_db.cache_path, pathlib.Path(self.cache_path.name)) self.assertEqual(chat_db.cache_path, pathlib.Path(self.cache_path.name))
self.assertEqual(chat_db.db_path, pathlib.Path(self.db_path.name)) self.assertEqual(chat_db.db_path, pathlib.Path(self.db_path.name))
self.assertEqual(chat_db.messages[0].file_path, self.assertEqual(chat_db.messages[0].file_path,
pathlib.Path(self.db_path.name, f'0001{msg_suffix}')) pathlib.Path(self.db_path.name, '0001.txt'))
def test_from_dir_filter_tags_empty(self) -> None: def test_from_dir_filter_tags_empty(self) -> None:
chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name), chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),
@@ -313,7 +279,7 @@ class TestChatDB(TestChatBase):
self.assertEqual(chat_db.cache_path, pathlib.Path(self.cache_path.name)) self.assertEqual(chat_db.cache_path, pathlib.Path(self.cache_path.name))
self.assertEqual(chat_db.db_path, pathlib.Path(self.db_path.name)) self.assertEqual(chat_db.db_path, pathlib.Path(self.db_path.name))
self.assertEqual(chat_db.messages[0].file_path, self.assertEqual(chat_db.messages[0].file_path,
pathlib.Path(self.db_path.name, f'0002{msg_suffix}')) pathlib.Path(self.db_path.name, '0002.yaml'))
self.assertEqual(chat_db.messages[0].answer, 'Answer 2') self.assertEqual(chat_db.messages[0].answer, 'Answer 2')
def test_from_messages(self) -> None: def test_from_messages(self) -> None:
@@ -358,25 +324,25 @@ class TestChatDB(TestChatBase):
chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name), chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),
pathlib.Path(self.db_path.name)) pathlib.Path(self.db_path.name))
# check that Message.file_path is correct # check that Message.file_path is correct
self.assertEqual(chat_db.messages[0].file_path, pathlib.Path(self.db_path.name, f'0001{msg_suffix}')) self.assertEqual(chat_db.messages[0].file_path, pathlib.Path(self.db_path.name, '0001.txt'))
self.assertEqual(chat_db.messages[1].file_path, pathlib.Path(self.db_path.name, f'0002{msg_suffix}')) self.assertEqual(chat_db.messages[1].file_path, pathlib.Path(self.db_path.name, '0002.yaml'))
self.assertEqual(chat_db.messages[2].file_path, pathlib.Path(self.db_path.name, f'0003{msg_suffix}')) self.assertEqual(chat_db.messages[2].file_path, pathlib.Path(self.db_path.name, '0003.txt'))
self.assertEqual(chat_db.messages[3].file_path, pathlib.Path(self.db_path.name, f'0004{msg_suffix}')) self.assertEqual(chat_db.messages[3].file_path, pathlib.Path(self.db_path.name, '0004.yaml'))
# write the messages to the cache directory # write the messages to the cache directory
chat_db.cache_write() chat_db.cache_write()
# check if the written files are in the cache directory # check if the written files are in the cache directory
cache_dir_files = self.message_list(self.cache_path) cache_dir_files = self.message_list(self.cache_path)
self.assertEqual(len(cache_dir_files), 4) self.assertEqual(len(cache_dir_files), 4)
self.assertIn(pathlib.Path(self.cache_path.name, f'0001{msg_suffix}'), cache_dir_files) self.assertIn(pathlib.Path(self.cache_path.name, '0001.txt'), cache_dir_files)
self.assertIn(pathlib.Path(self.cache_path.name, f'0002{msg_suffix}'), cache_dir_files) self.assertIn(pathlib.Path(self.cache_path.name, '0002.yaml'), cache_dir_files)
self.assertIn(pathlib.Path(self.cache_path.name, f'0003{msg_suffix}'), cache_dir_files) self.assertIn(pathlib.Path(self.cache_path.name, '0003.txt'), cache_dir_files)
self.assertIn(pathlib.Path(self.cache_path.name, f'0004{msg_suffix}'), cache_dir_files) self.assertIn(pathlib.Path(self.cache_path.name, '0004.yaml'), cache_dir_files)
# check that Message.file_path has been correctly updated # check that Message.file_path has been correctly updated
self.assertEqual(chat_db.messages[0].file_path, pathlib.Path(self.cache_path.name, f'0001{msg_suffix}')) self.assertEqual(chat_db.messages[0].file_path, pathlib.Path(self.cache_path.name, '0001.txt'))
self.assertEqual(chat_db.messages[1].file_path, pathlib.Path(self.cache_path.name, f'0002{msg_suffix}')) self.assertEqual(chat_db.messages[1].file_path, pathlib.Path(self.cache_path.name, '0002.yaml'))
self.assertEqual(chat_db.messages[2].file_path, pathlib.Path(self.cache_path.name, f'0003{msg_suffix}')) self.assertEqual(chat_db.messages[2].file_path, pathlib.Path(self.cache_path.name, '0003.txt'))
self.assertEqual(chat_db.messages[3].file_path, pathlib.Path(self.cache_path.name, f'0004{msg_suffix}')) self.assertEqual(chat_db.messages[3].file_path, pathlib.Path(self.cache_path.name, '0004.yaml'))
# check the timestamp of the files in the DB directory # check the timestamp of the files in the DB directory
db_dir_files = self.message_list(self.db_path) db_dir_files = self.message_list(self.db_path)
@@ -388,18 +354,18 @@ class TestChatDB(TestChatBase):
# check if the written files are in the DB directory # check if the written files are in the DB directory
db_dir_files = self.message_list(self.db_path) db_dir_files = self.message_list(self.db_path)
self.assertEqual(len(db_dir_files), 4) self.assertEqual(len(db_dir_files), 4)
self.assertIn(pathlib.Path(self.db_path.name, f'0001{msg_suffix}'), db_dir_files) self.assertIn(pathlib.Path(self.db_path.name, '0001.txt'), db_dir_files)
self.assertIn(pathlib.Path(self.db_path.name, f'0002{msg_suffix}'), db_dir_files) self.assertIn(pathlib.Path(self.db_path.name, '0002.yaml'), db_dir_files)
self.assertIn(pathlib.Path(self.db_path.name, f'0003{msg_suffix}'), db_dir_files) self.assertIn(pathlib.Path(self.db_path.name, '0003.txt'), db_dir_files)
self.assertIn(pathlib.Path(self.db_path.name, f'0004{msg_suffix}'), db_dir_files) self.assertIn(pathlib.Path(self.db_path.name, '0004.yaml'), db_dir_files)
# check if all files in the DB dir have actually been overwritten # check if all files in the DB dir have actually been overwritten
for file in db_dir_files: for file in db_dir_files:
self.assertGreater(file.stat().st_mtime, old_timestamps[file]) self.assertGreater(file.stat().st_mtime, old_timestamps[file])
# check that Message.file_path has been correctly updated (again) # check that Message.file_path has been correctly updated (again)
self.assertEqual(chat_db.messages[0].file_path, pathlib.Path(self.db_path.name, f'0001{msg_suffix}')) self.assertEqual(chat_db.messages[0].file_path, pathlib.Path(self.db_path.name, '0001.txt'))
self.assertEqual(chat_db.messages[1].file_path, pathlib.Path(self.db_path.name, f'0002{msg_suffix}')) self.assertEqual(chat_db.messages[1].file_path, pathlib.Path(self.db_path.name, '0002.yaml'))
self.assertEqual(chat_db.messages[2].file_path, pathlib.Path(self.db_path.name, f'0003{msg_suffix}')) self.assertEqual(chat_db.messages[2].file_path, pathlib.Path(self.db_path.name, '0003.txt'))
self.assertEqual(chat_db.messages[3].file_path, pathlib.Path(self.db_path.name, f'0004{msg_suffix}')) self.assertEqual(chat_db.messages[3].file_path, pathlib.Path(self.db_path.name, '0004.yaml'))
def test_db_read(self) -> None: def test_db_read(self) -> None:
# create a new ChatDB instance # create a new ChatDB instance
@@ -414,65 +380,65 @@ class TestChatDB(TestChatBase):
new_message2 = Message(Question('Question 6'), new_message2 = Message(Question('Question 6'),
Answer('Answer 6'), Answer('Answer 6'),
{Tag('tag6')}) {Tag('tag6')})
new_message1.to_file(pathlib.Path(self.db_path.name, f'0005{msg_suffix}'), mformat='txt') new_message1.to_file(pathlib.Path(self.db_path.name, '0005.txt'))
new_message2.to_file(pathlib.Path(self.db_path.name, f'0006{msg_suffix}'), mformat='yaml') new_message2.to_file(pathlib.Path(self.db_path.name, '0006.yaml'))
# read and check them # read and check them
chat_db.db_read() chat_db.db_read()
self.assertEqual(len(chat_db.messages), 6) self.assertEqual(len(chat_db.messages), 6)
self.assertEqual(chat_db.messages[4].file_path, pathlib.Path(self.db_path.name, f'0005{msg_suffix}')) self.assertEqual(chat_db.messages[4].file_path, pathlib.Path(self.db_path.name, '0005.txt'))
self.assertEqual(chat_db.messages[5].file_path, pathlib.Path(self.db_path.name, f'0006{msg_suffix}')) self.assertEqual(chat_db.messages[5].file_path, pathlib.Path(self.db_path.name, '0006.yaml'))
# create 2 new files in the cache directory # create 2 new files in the cache directory
new_message3 = Message(Question('Question 7'), new_message3 = Message(Question('Question 7'),
Answer('Answer 7'), Answer('Answer 5'),
{Tag('tag7')}) {Tag('tag7')})
new_message4 = Message(Question('Question 8'), new_message4 = Message(Question('Question 8'),
Answer('Answer 8'), Answer('Answer 6'),
{Tag('tag8')}) {Tag('tag8')})
new_message3.to_file(pathlib.Path(self.cache_path.name, f'0007{msg_suffix}'), mformat='txt') new_message3.to_file(pathlib.Path(self.cache_path.name, '0007.txt'))
new_message4.to_file(pathlib.Path(self.cache_path.name, f'0008{msg_suffix}'), mformat='yaml') new_message4.to_file(pathlib.Path(self.cache_path.name, '0008.yaml'))
# read and check them # read and check them
chat_db.cache_read() chat_db.cache_read()
self.assertEqual(len(chat_db.messages), 8) self.assertEqual(len(chat_db.messages), 8)
# check that the new message have the cache dir path # check that the new message have the cache dir path
self.assertEqual(chat_db.messages[6].file_path, pathlib.Path(self.cache_path.name, f'0007{msg_suffix}')) self.assertEqual(chat_db.messages[6].file_path, pathlib.Path(self.cache_path.name, '0007.txt'))
self.assertEqual(chat_db.messages[7].file_path, pathlib.Path(self.cache_path.name, f'0008{msg_suffix}')) self.assertEqual(chat_db.messages[7].file_path, pathlib.Path(self.cache_path.name, '0008.yaml'))
# an the old ones keep their path (since they have not been replaced) # an the old ones keep their path (since they have not been replaced)
self.assertEqual(chat_db.messages[4].file_path, pathlib.Path(self.db_path.name, f'0005{msg_suffix}')) self.assertEqual(chat_db.messages[4].file_path, pathlib.Path(self.db_path.name, '0005.txt'))
self.assertEqual(chat_db.messages[5].file_path, pathlib.Path(self.db_path.name, f'0006{msg_suffix}')) self.assertEqual(chat_db.messages[5].file_path, pathlib.Path(self.db_path.name, '0006.yaml'))
# now overwrite two messages in the DB directory # now overwrite two messages in the DB directory
new_message1.question = Question('New Question 1') new_message1.question = Question('New Question 1')
new_message2.question = Question('New Question 2') new_message2.question = Question('New Question 2')
new_message1.to_file(pathlib.Path(self.db_path.name, f'0005{msg_suffix}'), mformat='txt') new_message1.to_file(pathlib.Path(self.db_path.name, '0005.txt'))
new_message2.to_file(pathlib.Path(self.db_path.name, f'0006{msg_suffix}'), mformat='yaml') new_message2.to_file(pathlib.Path(self.db_path.name, '0006.yaml'))
# read from the DB dir and check if the modified messages have been updated # read from the DB dir and check if the modified messages have been updated
chat_db.db_read() chat_db.db_read()
self.assertEqual(len(chat_db.messages), 8) self.assertEqual(len(chat_db.messages), 8)
self.assertEqual(chat_db.messages[4].question, 'New Question 1') self.assertEqual(chat_db.messages[4].question, 'New Question 1')
self.assertEqual(chat_db.messages[5].question, 'New Question 2') self.assertEqual(chat_db.messages[5].question, 'New Question 2')
self.assertEqual(chat_db.messages[4].file_path, pathlib.Path(self.db_path.name, f'0005{msg_suffix}')) self.assertEqual(chat_db.messages[4].file_path, pathlib.Path(self.db_path.name, '0005.txt'))
self.assertEqual(chat_db.messages[5].file_path, pathlib.Path(self.db_path.name, f'0006{msg_suffix}')) self.assertEqual(chat_db.messages[5].file_path, pathlib.Path(self.db_path.name, '0006.yaml'))
# now write the messages from the cache to the DB directory # now write the messages from the cache to the DB directory
new_message3.to_file(pathlib.Path(self.db_path.name, f'0007{msg_suffix}')) new_message3.to_file(pathlib.Path(self.db_path.name, '0007.txt'))
new_message4.to_file(pathlib.Path(self.db_path.name, f'0008{msg_suffix}')) new_message4.to_file(pathlib.Path(self.db_path.name, '0008.yaml'))
# read and check them # read and check them
chat_db.db_read() chat_db.db_read()
self.assertEqual(len(chat_db.messages), 8) self.assertEqual(len(chat_db.messages), 8)
# check that they now have the DB path # check that they now have the DB path
self.assertEqual(chat_db.messages[6].file_path, pathlib.Path(self.db_path.name, f'0007{msg_suffix}')) self.assertEqual(chat_db.messages[6].file_path, pathlib.Path(self.db_path.name, '0007.txt'))
self.assertEqual(chat_db.messages[7].file_path, pathlib.Path(self.db_path.name, f'0008{msg_suffix}')) self.assertEqual(chat_db.messages[7].file_path, pathlib.Path(self.db_path.name, '0008.yaml'))
def test_cache_clear(self) -> None: def test_cache_clear(self) -> None:
# create a new ChatDB instance # create a new ChatDB instance
chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name), chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),
pathlib.Path(self.db_path.name)) pathlib.Path(self.db_path.name))
# check that Message.file_path is correct # check that Message.file_path is correct
self.assertEqual(chat_db.messages[0].file_path, pathlib.Path(self.db_path.name, f'0001{msg_suffix}')) self.assertEqual(chat_db.messages[0].file_path, pathlib.Path(self.db_path.name, '0001.txt'))
self.assertEqual(chat_db.messages[1].file_path, pathlib.Path(self.db_path.name, f'0002{msg_suffix}')) self.assertEqual(chat_db.messages[1].file_path, pathlib.Path(self.db_path.name, '0002.yaml'))
self.assertEqual(chat_db.messages[2].file_path, pathlib.Path(self.db_path.name, f'0003{msg_suffix}')) self.assertEqual(chat_db.messages[2].file_path, pathlib.Path(self.db_path.name, '0003.txt'))
self.assertEqual(chat_db.messages[3].file_path, pathlib.Path(self.db_path.name, f'0004{msg_suffix}')) self.assertEqual(chat_db.messages[3].file_path, pathlib.Path(self.db_path.name, '0004.yaml'))
# write the messages to the cache directory # write the messages to the cache directory
chat_db.cache_write() chat_db.cache_write()
@@ -484,10 +450,10 @@ class TestChatDB(TestChatBase):
chat_db.db_write() chat_db.db_write()
db_dir_files = self.message_list(self.db_path) db_dir_files = self.message_list(self.db_path)
self.assertEqual(len(db_dir_files), 4) self.assertEqual(len(db_dir_files), 4)
self.assertIn(pathlib.Path(self.db_path.name, f'0001{msg_suffix}'), db_dir_files) self.assertIn(pathlib.Path(self.db_path.name, '0001.txt'), db_dir_files)
self.assertIn(pathlib.Path(self.db_path.name, f'0002{msg_suffix}'), db_dir_files) self.assertIn(pathlib.Path(self.db_path.name, '0002.yaml'), db_dir_files)
self.assertIn(pathlib.Path(self.db_path.name, f'0003{msg_suffix}'), db_dir_files) self.assertIn(pathlib.Path(self.db_path.name, '0003.txt'), db_dir_files)
self.assertIn(pathlib.Path(self.db_path.name, f'0004{msg_suffix}'), db_dir_files) self.assertIn(pathlib.Path(self.db_path.name, '0004.yaml'), db_dir_files)
# add a new message with empty file_path # add a new message with empty file_path
message_empty = Message(question=Question("What the hell am I doing here?"), message_empty = Message(question=Question("What the hell am I doing here?"),
@@ -495,7 +461,7 @@ class TestChatDB(TestChatBase):
# and one for the cache dir # and one for the cache dir
message_cache = Message(question=Question("What the hell am I doing here?"), message_cache = Message(question=Question("What the hell am I doing here?"),
answer=Answer("You're a creep!"), answer=Answer("You're a creep!"),
file_path=pathlib.Path(self.cache_path.name, '0005')) file_path=pathlib.Path(self.cache_path.name, '0005.txt'))
chat_db.msg_add([message_empty, message_cache]) chat_db.msg_add([message_empty, message_cache])
# clear the cache and check the cache dir # clear the cache and check the cache dir
@@ -557,11 +523,11 @@ class TestChatDB(TestChatBase):
chat_db.msg_write([message]) chat_db.msg_write([message])
# write a message with a valid file_path # write a message with a valid file_path
message.file_path = pathlib.Path(self.cache_path.name) / '123456' message.file_path = pathlib.Path(self.cache_path.name) / '123456.txt'
chat_db.msg_write([message]) chat_db.msg_write([message])
cache_dir_files = self.message_list(self.cache_path) cache_dir_files = self.message_list(self.cache_path)
self.assertEqual(len(cache_dir_files), 1) self.assertEqual(len(cache_dir_files), 1)
self.assertIn(pathlib.Path(self.cache_path.name, f'123456{msg_suffix}'), cache_dir_files) self.assertIn(pathlib.Path(self.cache_path.name, '123456.txt'), cache_dir_files)
def test_msg_update(self) -> None: def test_msg_update(self) -> None:
# create a new ChatDB instance # create a new ChatDB instance
@@ -597,21 +563,21 @@ class TestChatDB(TestChatBase):
# search for a DB file in memory # search for a DB file in memory
self.assertEqual(chat_db.msg_find([str(self.message1.file_path)], loc='mem'), [self.message1]) self.assertEqual(chat_db.msg_find([str(self.message1.file_path)], loc='mem'), [self.message1])
self.assertEqual(chat_db.msg_find([self.message1.msg_id()], loc='mem'), [self.message1]) self.assertEqual(chat_db.msg_find([self.message1.msg_id()], loc='mem'), [self.message1])
self.assertEqual(chat_db.msg_find(['0001.msg'], loc='mem'), [self.message1]) self.assertEqual(chat_db.msg_find(['0001.txt'], loc='mem'), [self.message1])
self.assertEqual(chat_db.msg_find(['0001'], loc='mem'), [self.message1]) self.assertEqual(chat_db.msg_find(['0001'], loc='mem'), [self.message1])
# and on disk # and on disk
self.assertEqual(chat_db.msg_find([str(self.message2.file_path)], loc='db'), [self.message2]) self.assertEqual(chat_db.msg_find([str(self.message2.file_path)], loc='db'), [self.message2])
self.assertEqual(chat_db.msg_find([self.message2.msg_id()], loc='db'), [self.message2]) self.assertEqual(chat_db.msg_find([self.message2.msg_id()], loc='db'), [self.message2])
self.assertEqual(chat_db.msg_find(['0002.msg'], loc='db'), [self.message2]) self.assertEqual(chat_db.msg_find(['0002.yaml'], loc='db'), [self.message2])
self.assertEqual(chat_db.msg_find(['0002'], loc='db'), [self.message2]) self.assertEqual(chat_db.msg_find(['0002'], loc='db'), [self.message2])
# now search the cache -> expect empty result # now search the cache -> expect empty result
self.assertEqual(chat_db.msg_find([str(self.message3.file_path)], loc='cache'), []) self.assertEqual(chat_db.msg_find([str(self.message3.file_path)], loc='cache'), [])
self.assertEqual(chat_db.msg_find([self.message3.msg_id()], loc='cache'), []) self.assertEqual(chat_db.msg_find([self.message3.msg_id()], loc='cache'), [])
self.assertEqual(chat_db.msg_find(['0003.msg'], loc='cache'), []) self.assertEqual(chat_db.msg_find(['0003.txt'], loc='cache'), [])
self.assertEqual(chat_db.msg_find(['0003'], loc='cache'), []) self.assertEqual(chat_db.msg_find(['0003'], loc='cache'), [])
# search for multiple messages # search for multiple messages
# -> search one twice, expect result to be unique # -> search one twice, expect result to be unique
search_names = ['0001', '0002.msg', self.message3.msg_id(), str(self.message3.file_path)] search_names = ['0001', '0002.yaml', self.message3.msg_id(), str(self.message3.file_path)]
expected_result = [self.message1, self.message2, self.message3] expected_result = [self.message1, self.message2, self.message3]
result = chat_db.msg_find(search_names, loc='all') result = chat_db.msg_find(search_names, loc='all')
self.assert_messages_equal(result, expected_result) self.assert_messages_equal(result, expected_result)
-100
View File
@@ -1,100 +0,0 @@
import unittest
import argparse
from typing import Union, Optional
from chatmastermind.configuration import Config, AIConfig
from chatmastermind.tags import Tag
from chatmastermind.message import Message, Answer
from chatmastermind.chat import Chat
from chatmastermind.ai import AI, AIResponse, Tokens, AIError
class FakeAI(AI):
"""
A mocked version of the 'AI' class.
"""
ID: str
name: str
config: AIConfig
def models(self) -> list[str]:
raise NotImplementedError
def tokens(self, data: Union[Message, Chat]) -> int:
return 123
def print(self) -> None:
pass
def print_models(self) -> None:
pass
def __init__(self, ID: str, model: str, error: bool = False):
self.ID = ID
self.model = model
self.error = error
def request(self,
question: Message,
chat: Chat,
num_answers: int = 1,
otags: Optional[set[Tag]] = None) -> AIResponse:
"""
Mock the 'ai.request()' function by either returning fake
answers or raising an exception.
"""
if self.error:
raise AIError
question.answer = Answer("Answer 0")
question.tags = set(otags) if otags is not None else None
question.ai = self.ID
question.model = self.model
answers: list[Message] = [question]
for n in range(1, num_answers):
answers.append(Message(question=question.question,
answer=Answer(f"Answer {n}"),
tags=otags,
ai=self.ID,
model=self.model))
return AIResponse(answers, Tokens(10, 10, 20))
class TestWithFakeAI(unittest.TestCase):
"""
Base class for all tests that need to use the FakeAI.
"""
def assert_msgs_equal_except_file_path(self, msg1: list[Message], msg2: list[Message]) -> None:
"""
Compare messages using Question, Answer and all metadata excecot for the file_path.
"""
self.assertEqual(len(msg1), len(msg2))
for m1, m2 in zip(msg1, msg2):
# exclude the file_path, compare only Q, A and metadata
self.assertTrue(m1.equals(m2, file_path=False, verbose=True))
def assert_msgs_all_equal(self, msg1: list[Message], msg2: list[Message]) -> None:
"""
Compare messages using Question, Answer and ALL metadata.
"""
self.assertEqual(len(msg1), len(msg2))
for m1, m2 in zip(msg1, msg2):
self.assertTrue(m1.equals(m2, verbose=True))
def assert_msgs_content_equal(self, msg1: list[Message], msg2: list[Message]) -> None:
"""
Compare messages using only Question and Answer.
"""
self.assertEqual(len(msg1), len(msg2))
for m1, m2 in zip(msg1, msg2):
self.assertEqual(m1, m2)
def mock_create_ai(self, args: argparse.Namespace, config: Config) -> AI:
"""
Mocked 'create_ai' that returns a 'FakeAI' instance.
"""
return FakeAI(args.AI, args.model)
def mock_create_ai_with_error(self, args: argparse.Namespace, config: Config) -> AI:
"""
Mocked 'create_ai' that returns a 'FakeAI' instance.
"""
return FakeAI(args.AI, args.model, error=True)
-62
View File
@@ -1,62 +0,0 @@
import unittest
import argparse
import tempfile
import yaml
from pathlib import Path
from chatmastermind.message import Message, Question
from chatmastermind.chat import ChatDB, ChatError
from chatmastermind.configuration import Config
from chatmastermind.commands.hist import convert_messages
msg_suffix = Message.file_suffix_write
class TestConvertMessages(unittest.TestCase):
def setUp(self) -> None:
self.db_dir = tempfile.TemporaryDirectory()
self.cache_dir = tempfile.TemporaryDirectory()
self.db_path = Path(self.db_dir.name)
self.cache_path = Path(self.cache_dir.name)
self.args = argparse.Namespace()
self.config = Config()
self.config.cache = self.cache_dir.name
self.config.db = self.db_dir.name
# Prepare some messages
self.chat = ChatDB.from_dir(Path(self.cache_path),
Path(self.db_path))
self.messages = [Message(Question(f'Question {i}')) for i in range(0, 6)]
self.chat.db_write(self.messages[0:2])
self.chat.cache_write(self.messages[2:])
# Change some of the suffixes
assert self.messages[0].file_path
assert self.messages[1].file_path
self.messages[0].file_path.rename(self.messages[0].file_path.with_suffix('.txt'))
self.messages[1].file_path.rename(self.messages[1].file_path.with_suffix('.yaml'))
def tearDown(self) -> None:
self.db_dir.cleanup()
self.cache_dir.cleanup()
def test_convert_messages(self) -> None:
self.args.convert = 'yaml'
convert_messages(self.args, self.config)
msgs = self.chat.msg_gather(loc='disk', glob='*.*')
# Check if the number of messages is the same as before
self.assertEqual(len(msgs), len(self.messages))
# Check if all messages have the requested suffix
for msg in msgs:
assert msg.file_path
self.assertEqual(msg.file_path.suffix, msg_suffix)
# Check if the message IDs are correctly maintained
for m_new, m_old in zip(msgs, self.messages):
self.assertEqual(m_new.msg_id(), m_old.msg_id())
# check if all messages have the new format
for m in msgs:
with open(str(m.file_path), "r") as fd:
yaml.load(fd, Loader=yaml.FullLoader)
def test_convert_messages_wrong_format(self) -> None:
self.args.convert = 'foo'
with self.assertRaises(ChatError):
convert_messages(self.args, self.config)
+39 -102
View File
@@ -1,16 +1,11 @@
import unittest import unittest
import pathlib import pathlib
import tempfile import tempfile
import itertools
from typing import cast from typing import cast
from chatmastermind.message import source_code, Message, MessageError, Question, Answer, AILine, ModelLine,\ from chatmastermind.message import source_code, Message, MessageError, Question, Answer, AILine, ModelLine, MessageFilter, message_in
MessageFilter, message_in, message_valid_formats
from chatmastermind.tags import Tag, TagLine from chatmastermind.tags import Tag, TagLine
msg_suffix: str = Message.file_suffix_write
class SourceCodeTestCase(unittest.TestCase): class SourceCodeTestCase(unittest.TestCase):
def test_source_code_with_include_delims(self) -> None: def test_source_code_with_include_delims(self) -> None:
text = """ text = """
@@ -106,7 +101,7 @@ class AnswerTestCase(unittest.TestCase):
class MessageToFileTxtTestCase(unittest.TestCase): class MessageToFileTxtTestCase(unittest.TestCase):
def setUp(self) -> None: def setUp(self) -> None:
self.file = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
self.file_path = pathlib.Path(self.file.name) self.file_path = pathlib.Path(self.file.name)
self.message_complete = Message(Question('This is a question.'), self.message_complete = Message(Question('This is a question.'),
Answer('This is an answer.'), Answer('This is an answer.'),
@@ -122,7 +117,7 @@ class MessageToFileTxtTestCase(unittest.TestCase):
self.file_path.unlink() self.file_path.unlink()
def test_to_file_txt_complete(self) -> None: def test_to_file_txt_complete(self) -> None:
self.message_complete.to_file(self.file_path, mformat='txt') self.message_complete.to_file(self.file_path)
with open(self.file_path, "r") as fd: with open(self.file_path, "r") as fd:
content = fd.read() content = fd.read()
@@ -137,7 +132,7 @@ This is an answer.
self.assertEqual(content, expected_content) self.assertEqual(content, expected_content)
def test_to_file_txt_min(self) -> None: def test_to_file_txt_min(self) -> None:
self.message_min.to_file(self.file_path, mformat='txt') self.message_min.to_file(self.file_path)
with open(self.file_path, "r") as fd: with open(self.file_path, "r") as fd:
content = fd.read() content = fd.read()
@@ -146,17 +141,11 @@ This is a question.
""" """
self.assertEqual(content, expected_content) self.assertEqual(content, expected_content)
def test_to_file_unsupported_file_suffix(self) -> None: def test_to_file_unsupported_file_type(self) -> None:
unsupported_file_path = pathlib.Path("example.doc") unsupported_file_path = pathlib.Path("example.doc")
with self.assertRaises(MessageError) as cm: with self.assertRaises(MessageError) as cm:
self.message_complete.to_file(unsupported_file_path) self.message_complete.to_file(unsupported_file_path)
self.assertEqual(str(cm.exception), "File suffix '.doc' is not supported") self.assertEqual(str(cm.exception), "File type '.doc' is not supported")
def test_to_file_unsupported_file_format(self) -> None:
unsupported_file_format = pathlib.Path(f"example{msg_suffix}")
with self.assertRaises(MessageError) as cm:
self.message_complete.to_file(unsupported_file_format, mformat='doc') # type: ignore [arg-type]
self.assertEqual(str(cm.exception), "File format 'doc' is not supported")
def test_to_file_no_file_path(self) -> None: def test_to_file_no_file_path(self) -> None:
""" """
@@ -170,24 +159,10 @@ This is a question.
# reset the internal file_path # reset the internal file_path
self.message_complete.file_path = self.file_path self.message_complete.file_path = self.file_path
def test_to_file_txt_auto_suffix(self) -> None:
"""
Test if suffix is auto-generated if omitted.
"""
file_path_no_suffix = self.file_path.with_suffix('')
# test with file_path member
self.message_min.file_path = file_path_no_suffix
self.message_min.to_file(mformat='txt')
self.assertEqual(self.message_min.file_path.suffix, msg_suffix)
# test with explicit file_path
self.message_min.file_path = file_path_no_suffix
self.message_min.to_file(file_path=file_path_no_suffix, mformat='txt')
self.assertEqual(self.message_min.file_path.suffix, msg_suffix)
class MessageToFileYamlTestCase(unittest.TestCase): class MessageToFileYamlTestCase(unittest.TestCase):
def setUp(self) -> None: def setUp(self) -> None:
self.file = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file = tempfile.NamedTemporaryFile(delete=False, suffix='.yaml')
self.file_path = pathlib.Path(self.file.name) self.file_path = pathlib.Path(self.file.name)
self.message_complete = Message(Question('This is a question.'), self.message_complete = Message(Question('This is a question.'),
Answer('This is an answer.'), Answer('This is an answer.'),
@@ -209,7 +184,7 @@ class MessageToFileYamlTestCase(unittest.TestCase):
self.file_path.unlink() self.file_path.unlink()
def test_to_file_yaml_complete(self) -> None: def test_to_file_yaml_complete(self) -> None:
self.message_complete.to_file(self.file_path, mformat='yaml') self.message_complete.to_file(self.file_path)
with open(self.file_path, "r") as fd: with open(self.file_path, "r") as fd:
content = fd.read() content = fd.read()
@@ -224,7 +199,7 @@ class MessageToFileYamlTestCase(unittest.TestCase):
self.assertEqual(content, expected_content) self.assertEqual(content, expected_content)
def test_to_file_yaml_multiline(self) -> None: def test_to_file_yaml_multiline(self) -> None:
self.message_multiline.to_file(self.file_path, mformat='yaml') self.message_multiline.to_file(self.file_path)
with open(self.file_path, "r") as fd: with open(self.file_path, "r") as fd:
content = fd.read() content = fd.read()
@@ -243,31 +218,17 @@ class MessageToFileYamlTestCase(unittest.TestCase):
self.assertEqual(content, expected_content) self.assertEqual(content, expected_content)
def test_to_file_yaml_min(self) -> None: def test_to_file_yaml_min(self) -> None:
self.message_min.to_file(self.file_path, mformat='yaml') self.message_min.to_file(self.file_path)
with open(self.file_path, "r") as fd: with open(self.file_path, "r") as fd:
content = fd.read() content = fd.read()
expected_content = f"{Question.yaml_key}: This is a question.\n" expected_content = f"{Question.yaml_key}: This is a question.\n"
self.assertEqual(content, expected_content) self.assertEqual(content, expected_content)
def test_to_file_yaml_auto_suffix(self) -> None:
"""
Test if suffix is auto-generated if omitted.
"""
file_path_no_suffix = self.file_path.with_suffix('')
# test with file_path member
self.message_min.file_path = file_path_no_suffix
self.message_min.to_file(mformat='yaml')
self.assertEqual(self.message_min.file_path.suffix, msg_suffix)
# test with explicit file_path
self.message_min.file_path = file_path_no_suffix
self.message_min.to_file(file_path=file_path_no_suffix, mformat='yaml')
self.assertEqual(self.message_min.file_path.suffix, msg_suffix)
class MessageFromFileTxtTestCase(unittest.TestCase): class MessageFromFileTxtTestCase(unittest.TestCase):
def setUp(self) -> None: def setUp(self) -> None:
self.file = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
self.file_path = pathlib.Path(self.file.name) self.file_path = pathlib.Path(self.file.name)
with open(self.file_path, "w") as fd: with open(self.file_path, "w") as fd:
fd.write(f"""{TagLine.prefix} tag1 tag2 fd.write(f"""{TagLine.prefix} tag1 tag2
@@ -278,7 +239,7 @@ This is a question.
{Answer.txt_header} {Answer.txt_header}
This is an answer. This is an answer.
""") """)
self.file_min = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file_min = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
self.file_path_min = pathlib.Path(self.file_min.name) self.file_path_min = pathlib.Path(self.file_min.name)
with open(self.file_path_min, "w") as fd: with open(self.file_path_min, "w") as fd:
fd.write(f"""{Question.txt_header} fd.write(f"""{Question.txt_header}
@@ -298,7 +259,7 @@ This is a question.
message = Message.from_file(self.file_path) message = Message.from_file(self.file_path)
self.assertIsNotNone(message) self.assertIsNotNone(message)
self.assertIsInstance(message, Message) self.assertIsInstance(message, Message)
assert message if message: # mypy bug
self.assertEqual(message.question, 'This is a question.') self.assertEqual(message.question, 'This is a question.')
self.assertEqual(message.answer, 'This is an answer.') self.assertEqual(message.answer, 'This is an answer.')
self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')}) self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')})
@@ -313,7 +274,7 @@ This is a question.
message = Message.from_file(self.file_path_min) message = Message.from_file(self.file_path_min)
self.assertIsNotNone(message) self.assertIsNotNone(message)
self.assertIsInstance(message, Message) self.assertIsInstance(message, Message)
assert message if message: # mypy bug
self.assertEqual(message.question, 'This is a question.') self.assertEqual(message.question, 'This is a question.')
self.assertEqual(message.file_path, self.file_path_min) self.assertEqual(message.file_path, self.file_path_min)
self.assertIsNone(message.answer) self.assertIsNone(message.answer)
@@ -323,7 +284,7 @@ This is a question.
MessageFilter(tags_or={Tag('tag1')})) MessageFilter(tags_or={Tag('tag1')}))
self.assertIsNotNone(message) self.assertIsNotNone(message)
self.assertIsInstance(message, Message) self.assertIsInstance(message, Message)
assert message if message: # mypy bug
self.assertEqual(message.question, 'This is a question.') self.assertEqual(message.question, 'This is a question.')
self.assertEqual(message.answer, 'This is an answer.') self.assertEqual(message.answer, 'This is an answer.')
self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')}) self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')})
@@ -350,13 +311,13 @@ This is a question.
MessageFilter(tags_not={Tag('tag1')})) MessageFilter(tags_not={Tag('tag1')}))
self.assertIsNotNone(message) self.assertIsNotNone(message)
self.assertIsInstance(message, Message) self.assertIsInstance(message, Message)
assert message if message: # mypy bug
self.assertEqual(message.question, 'This is a question.') self.assertEqual(message.question, 'This is a question.')
self.assertSetEqual(cast(set[Tag], message.tags), set()) self.assertSetEqual(cast(set[Tag], message.tags), set())
self.assertEqual(message.file_path, self.file_path_min) self.assertEqual(message.file_path, self.file_path_min)
def test_from_file_not_exists(self) -> None: def test_from_file_not_exists(self) -> None:
file_not_exists = pathlib.Path(f"example{msg_suffix}") file_not_exists = pathlib.Path("example.txt")
with self.assertRaises(MessageError) as cm: with self.assertRaises(MessageError) as cm:
Message.from_file(file_not_exists) Message.from_file(file_not_exists)
self.assertEqual(str(cm.exception), f"Message file '{file_not_exists}' does not exist") self.assertEqual(str(cm.exception), f"Message file '{file_not_exists}' does not exist")
@@ -435,7 +396,7 @@ This is a question.
class MessageFromFileYamlTestCase(unittest.TestCase): class MessageFromFileYamlTestCase(unittest.TestCase):
def setUp(self) -> None: def setUp(self) -> None:
self.file = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file = tempfile.NamedTemporaryFile(delete=False, suffix='.yaml')
self.file_path = pathlib.Path(self.file.name) self.file_path = pathlib.Path(self.file.name)
with open(self.file_path, "w") as fd: with open(self.file_path, "w") as fd:
fd.write(f""" fd.write(f"""
@@ -449,7 +410,7 @@ class MessageFromFileYamlTestCase(unittest.TestCase):
- tag1 - tag1
- tag2 - tag2
""") """)
self.file_min = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file_min = tempfile.NamedTemporaryFile(delete=False, suffix='.yaml')
self.file_path_min = pathlib.Path(self.file_min.name) self.file_path_min = pathlib.Path(self.file_min.name)
with open(self.file_path_min, "w") as fd: with open(self.file_path_min, "w") as fd:
fd.write(f""" fd.write(f"""
@@ -470,7 +431,7 @@ class MessageFromFileYamlTestCase(unittest.TestCase):
message = Message.from_file(self.file_path) message = Message.from_file(self.file_path)
self.assertIsInstance(message, Message) self.assertIsInstance(message, Message)
self.assertIsNotNone(message) self.assertIsNotNone(message)
assert message if message: # mypy bug
self.assertEqual(message.question, 'This is a question.') self.assertEqual(message.question, 'This is a question.')
self.assertEqual(message.answer, 'This is an answer.') self.assertEqual(message.answer, 'This is an answer.')
self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')}) self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')})
@@ -485,14 +446,14 @@ class MessageFromFileYamlTestCase(unittest.TestCase):
message = Message.from_file(self.file_path_min) message = Message.from_file(self.file_path_min)
self.assertIsInstance(message, Message) self.assertIsInstance(message, Message)
self.assertIsNotNone(message) self.assertIsNotNone(message)
assert message if message: # mypy bug
self.assertEqual(message.question, 'This is a question.') self.assertEqual(message.question, 'This is a question.')
self.assertSetEqual(cast(set[Tag], message.tags), set()) self.assertSetEqual(cast(set[Tag], message.tags), set())
self.assertEqual(message.file_path, self.file_path_min) self.assertEqual(message.file_path, self.file_path_min)
self.assertIsNone(message.answer) self.assertIsNone(message.answer)
def test_from_file_not_exists(self) -> None: def test_from_file_not_exists(self) -> None:
file_not_exists = pathlib.Path(f"example{msg_suffix}") file_not_exists = pathlib.Path("example.yaml")
with self.assertRaises(MessageError) as cm: with self.assertRaises(MessageError) as cm:
Message.from_file(file_not_exists) Message.from_file(file_not_exists)
self.assertEqual(str(cm.exception), f"Message file '{file_not_exists}' does not exist") self.assertEqual(str(cm.exception), f"Message file '{file_not_exists}' does not exist")
@@ -502,7 +463,7 @@ class MessageFromFileYamlTestCase(unittest.TestCase):
MessageFilter(tags_or={Tag('tag1')})) MessageFilter(tags_or={Tag('tag1')}))
self.assertIsNotNone(message) self.assertIsNotNone(message)
self.assertIsInstance(message, Message) self.assertIsInstance(message, Message)
assert message if message: # mypy bug
self.assertEqual(message.question, 'This is a question.') self.assertEqual(message.question, 'This is a question.')
self.assertEqual(message.answer, 'This is an answer.') self.assertEqual(message.answer, 'This is an answer.')
self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')}) self.assertSetEqual(cast(set[Tag], message.tags), {Tag('tag1'), Tag('tag2')})
@@ -523,7 +484,7 @@ class MessageFromFileYamlTestCase(unittest.TestCase):
MessageFilter(tags_not={Tag('tag1')})) MessageFilter(tags_not={Tag('tag1')}))
self.assertIsNotNone(message) self.assertIsNotNone(message)
self.assertIsInstance(message, Message) self.assertIsInstance(message, Message)
assert message if message: # mypy bug
self.assertEqual(message.question, 'This is a question.') self.assertEqual(message.question, 'This is a question.')
self.assertSetEqual(cast(set[Tag], message.tags), set()) self.assertSetEqual(cast(set[Tag], message.tags), set())
self.assertEqual(message.file_path, self.file_path_min) self.assertEqual(message.file_path, self.file_path_min)
@@ -602,7 +563,7 @@ class MessageFromFileYamlTestCase(unittest.TestCase):
class TagsFromFileTestCase(unittest.TestCase): class TagsFromFileTestCase(unittest.TestCase):
def setUp(self) -> None: def setUp(self) -> None:
self.file_txt = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file_txt = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
self.file_path_txt = pathlib.Path(self.file_txt.name) self.file_path_txt = pathlib.Path(self.file_txt.name)
with open(self.file_path_txt, "w") as fd: with open(self.file_path_txt, "w") as fd:
fd.write(f"""{TagLine.prefix} tag1 tag2 ptag3 fd.write(f"""{TagLine.prefix} tag1 tag2 ptag3
@@ -611,7 +572,7 @@ This is a question.
{Answer.txt_header} {Answer.txt_header}
This is an answer. This is an answer.
""") """)
self.file_txt_no_tags = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file_txt_no_tags = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
self.file_path_txt_no_tags = pathlib.Path(self.file_txt_no_tags.name) self.file_path_txt_no_tags = pathlib.Path(self.file_txt_no_tags.name)
with open(self.file_path_txt_no_tags, "w") as fd: with open(self.file_path_txt_no_tags, "w") as fd:
fd.write(f"""{Question.txt_header} fd.write(f"""{Question.txt_header}
@@ -619,7 +580,7 @@ This is a question.
{Answer.txt_header} {Answer.txt_header}
This is an answer. This is an answer.
""") """)
self.file_txt_tags_empty = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file_txt_tags_empty = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
self.file_path_txt_tags_empty = pathlib.Path(self.file_txt_tags_empty.name) self.file_path_txt_tags_empty = pathlib.Path(self.file_txt_tags_empty.name)
with open(self.file_path_txt_tags_empty, "w") as fd: with open(self.file_path_txt_tags_empty, "w") as fd:
fd.write(f"""TAGS: fd.write(f"""TAGS:
@@ -628,7 +589,7 @@ This is a question.
{Answer.txt_header} {Answer.txt_header}
This is an answer. This is an answer.
""") """)
self.file_yaml = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file_yaml = tempfile.NamedTemporaryFile(delete=False, suffix='.yaml')
self.file_path_yaml = pathlib.Path(self.file_yaml.name) self.file_path_yaml = pathlib.Path(self.file_yaml.name)
with open(self.file_path_yaml, "w") as fd: with open(self.file_path_yaml, "w") as fd:
fd.write(f""" fd.write(f"""
@@ -641,7 +602,7 @@ This is an answer.
- tag2 - tag2
- ptag3 - ptag3
""") """)
self.file_yaml_no_tags = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file_yaml_no_tags = tempfile.NamedTemporaryFile(delete=False, suffix='.yaml')
self.file_path_yaml_no_tags = pathlib.Path(self.file_yaml_no_tags.name) self.file_path_yaml_no_tags = pathlib.Path(self.file_yaml_no_tags.name)
with open(self.file_path_yaml_no_tags, "w") as fd: with open(self.file_path_yaml_no_tags, "w") as fd:
fd.write(f""" fd.write(f"""
@@ -718,25 +679,24 @@ class TagsFromDirTestCase(unittest.TestCase):
{Tag('ctag5'), Tag('ctag6')} {Tag('ctag5'), Tag('ctag6')}
] ]
self.files = [ self.files = [
pathlib.Path(self.temp_dir.name, f'file1{msg_suffix}'), pathlib.Path(self.temp_dir.name, 'file1.txt'),
pathlib.Path(self.temp_dir.name, f'file2{msg_suffix}'), pathlib.Path(self.temp_dir.name, 'file2.yaml'),
pathlib.Path(self.temp_dir.name, f'file3{msg_suffix}') pathlib.Path(self.temp_dir.name, 'file3.txt')
] ]
self.files_no_tags = [ self.files_no_tags = [
pathlib.Path(self.temp_dir_no_tags.name, f'file4{msg_suffix}'), pathlib.Path(self.temp_dir_no_tags.name, 'file4.txt'),
pathlib.Path(self.temp_dir_no_tags.name, f'file5{msg_suffix}'), pathlib.Path(self.temp_dir_no_tags.name, 'file5.yaml'),
pathlib.Path(self.temp_dir_no_tags.name, f'file6{msg_suffix}') pathlib.Path(self.temp_dir_no_tags.name, 'file6.txt')
] ]
mformats = itertools.cycle(message_valid_formats)
for file, tags in zip(self.files, self.tag_sets): for file, tags in zip(self.files, self.tag_sets):
message = Message(Question('This is a question.'), message = Message(Question('This is a question.'),
Answer('This is an answer.'), Answer('This is an answer.'),
tags) tags)
message.to_file(file, next(mformats)) message.to_file(file)
for file in self.files_no_tags: for file in self.files_no_tags:
message = Message(Question('This is a question.'), message = Message(Question('This is a question.'),
Answer('This is an answer.')) Answer('This is an answer.'))
message.to_file(file, next(mformats)) message.to_file(file)
def tearDown(self) -> None: def tearDown(self) -> None:
self.temp_dir.cleanup() self.temp_dir.cleanup()
@@ -759,7 +719,7 @@ class TagsFromDirTestCase(unittest.TestCase):
class MessageIDTestCase(unittest.TestCase): class MessageIDTestCase(unittest.TestCase):
def setUp(self) -> None: def setUp(self) -> None:
self.file = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix) self.file = tempfile.NamedTemporaryFile(delete=False, suffix='.txt')
self.file_path = pathlib.Path(self.file.name) self.file_path = pathlib.Path(self.file.name)
self.message = Message(Question('This is a question.'), self.message = Message(Question('This is a question.'),
file_path=self.file_path) file_path=self.file_path)
@@ -856,8 +816,6 @@ class MessageToStrTestCase(unittest.TestCase):
def setUp(self) -> None: def setUp(self) -> None:
self.message = Message(Question('This is a question.'), self.message = Message(Question('This is a question.'),
Answer('This is an answer.'), Answer('This is an answer.'),
ai=('FakeAI'),
model=('FakeModel'),
tags={Tag('atag1'), Tag('btag2')}, tags={Tag('atag1'), Tag('btag2')},
file_path=pathlib.Path('/tmp/foo/bla')) file_path=pathlib.Path('/tmp/foo/bla'))
@@ -871,29 +829,8 @@ This is an answer."""
def test_to_str_with_tags_and_file(self) -> None: def test_to_str_with_tags_and_file(self) -> None:
expected_output = f"""{TagLine.prefix} atag1 btag2 expected_output = f"""{TagLine.prefix} atag1 btag2
FILE: /tmp/foo/bla FILE: /tmp/foo/bla
AI: FakeAI
MODEL: FakeModel
{Question.txt_header} {Question.txt_header}
This is a question. This is a question.
{Answer.txt_header} {Answer.txt_header}
This is an answer.""" This is an answer."""
self.assertEqual(self.message.to_str(with_metadata=True), expected_output) self.assertEqual(self.message.to_str(with_tags=True, with_file=True), expected_output)
class MessageRmFileTestCase(unittest.TestCase):
def setUp(self) -> None:
self.file = tempfile.NamedTemporaryFile(delete=False, suffix=msg_suffix)
self.file_path = pathlib.Path(self.file.name)
self.message = Message(Question('This is a question.'),
file_path=self.file_path)
self.message.to_file()
def tearDown(self) -> None:
self.file.close()
self.file_path.unlink(missing_ok=True)
def test_rm_file(self) -> None:
assert self.message.file_path
self.assertTrue(self.message.file_path.exists())
self.message.rm_file()
self.assertFalse(self.message.file_path.exists())
+143 -223
View File
@@ -1,23 +1,31 @@
import os import os
import unittest
import argparse import argparse
import tempfile import tempfile
from copy import copy
from pathlib import Path from pathlib import Path
from unittest import mock from unittest import mock
from unittest.mock import MagicMock, call from unittest.mock import MagicMock, call, ANY
from typing import Optional
from chatmastermind.configuration import Config from chatmastermind.configuration import Config
from chatmastermind.commands.question import create_message, question_cmd from chatmastermind.commands.question import create_message, question_cmd
from chatmastermind.tags import Tag from chatmastermind.tags import Tag
from chatmastermind.message import Message, Question, Answer from chatmastermind.message import Message, Question, Answer
from chatmastermind.chat import Chat, ChatDB from chatmastermind.chat import Chat, ChatDB
from chatmastermind.ai import AIError from chatmastermind.ai import AI, AIResponse, Tokens, AIError
from .test_common import TestWithFakeAI
msg_suffix = Message.file_suffix_write class TestQuestionCmdBase(unittest.TestCase):
def assert_messages_equal(self, msg1: list[Message], msg2: list[Message]) -> None:
"""
Compare messages using more than just Question and Answer.
"""
self.assertEqual(len(msg1), len(msg2))
for m1, m2 in zip(msg1, msg2):
# exclude the file_path, compare only Q, A and metadata
self.assertTrue(m1.equals(m2, file_path=False, verbose=True))
class TestMessageCreate(TestWithFakeAI): class TestMessageCreate(TestQuestionCmdBase):
""" """
Test if messages created by the 'question' command have Test if messages created by the 'question' command have
the correct format. the correct format.
@@ -41,8 +49,6 @@ class TestMessageCreate(TestWithFakeAI):
self.args.AI = None self.args.AI = None
self.args.model = None self.args.model = None
self.args.output_tags = None self.args.output_tags = None
self.args.ask = None
self.args.create = None
# File 1 : no source code block, only text # File 1 : no source code block, only text
self.source_file1 = tempfile.NamedTemporaryFile(delete=False) self.source_file1 = tempfile.NamedTemporaryFile(delete=False)
self.source_file1_content = """This is just text. self.source_file1_content = """This is just text.
@@ -88,7 +94,7 @@ Aaaand again some text."""
def message_list(self, tmp_dir: tempfile.TemporaryDirectory) -> list[Path]: def message_list(self, tmp_dir: tempfile.TemporaryDirectory) -> list[Path]:
# exclude '.next' # exclude '.next'
return list(Path(tmp_dir.name).glob(f'*{msg_suffix}')) return list(Path(tmp_dir.name).glob('*.[ty]*'))
def test_message_file_created(self) -> None: def test_message_file_created(self) -> None:
self.args.ask = ["What is this?"] self.args.ask = ["What is this?"]
@@ -206,22 +212,7 @@ It is embedded code
""")) """))
class TestCreateOption(TestMessageCreate): class TestQuestionCmd(TestQuestionCmdBase):
def test_message_file_created(self) -> None:
self.args.create = ["How does question --create work?"]
self.args.ask = None
cache_dir_files = self.message_list(self.cache_dir)
self.assertEqual(len(cache_dir_files), 0)
create_message(self.chat, self.args)
cache_dir_files = self.message_list(self.cache_dir)
self.assertEqual(len(cache_dir_files), 1)
message = Message.from_file(cache_dir_files[0])
self.assertIsInstance(message, Message)
self.assertEqual(message.question, Question("How does question --create work?")) # type: ignore [union-attr]
class TestQuestionCmd(TestWithFakeAI):
def setUp(self) -> None: def setUp(self) -> None:
# create DB and cache # create DB and cache
@@ -236,8 +227,8 @@ class TestQuestionCmd(TestWithFakeAI):
ask=['What is the meaning of life?'], ask=['What is the meaning of life?'],
num_answers=1, num_answers=1,
output_tags=['science'], output_tags=['science'],
AI='FakeAI', AI='openai',
model='FakeModel', model='gpt-3.5-turbo',
or_tags=None, or_tags=None,
and_tags=None, and_tags=None,
exclude_tags=None, exclude_tags=None,
@@ -248,27 +239,57 @@ class TestQuestionCmd(TestWithFakeAI):
process=None, process=None,
overwrite=None overwrite=None
) )
# create a mock AI instance
self.ai = MagicMock(spec=AI)
self.ai.request.side_effect = self.mock_request
def input_message(self, args: argparse.Namespace) -> Message:
"""
Create the expected input message for a question using the
given arguments.
"""
# NOTE: we only use the first question from the "ask" list
# -> message creation using "question.create_message()" is
# tested above
# the answer is always empty for the input message
return Message(Question(args.ask[0]),
tags=args.output_tags,
ai=args.AI,
model=args.model)
def mock_request(self,
question: Message,
chat: Chat,
num_answers: int = 1,
otags: Optional[set[Tag]] = None) -> AIResponse:
"""
Mock the 'ai.request()' function
"""
question.answer = Answer("Answer 0")
question.tags = set(otags) if otags else None
question.ai = 'FakeAI'
question.model = 'FakeModel'
answers: list[Message] = [question]
for n in range(1, num_answers):
answers.append(Message(question=question.question,
answer=Answer(f"Answer {n}"),
tags=otags,
ai='FakeAI',
model='FakeModel'))
return AIResponse(answers, Tokens(10, 10, 20))
def message_list(self, tmp_dir: tempfile.TemporaryDirectory) -> list[Path]: def message_list(self, tmp_dir: tempfile.TemporaryDirectory) -> list[Path]:
# exclude '.next' # exclude '.next'
return sorted([f for f in Path(tmp_dir.name).glob(f'*{msg_suffix}')]) return sorted([f for f in Path(tmp_dir.name).glob('*.[ty]*')])
class TestQuestionCmdAsk(TestQuestionCmd):
@mock.patch('chatmastermind.commands.question.create_ai') @mock.patch('chatmastermind.commands.question.create_ai')
def test_ask_single_answer(self, mock_create_ai: MagicMock) -> None: def test_ask_single_answer(self, mock_create_ai: MagicMock) -> None:
""" """
Test single answer with no errors. Test single answer with no errors.
""" """
mock_create_ai.side_effect = self.mock_create_ai mock_create_ai.return_value = self.ai
expected_question = Message(Question(self.args.ask[0]), expected_question = self.input_message(self.args)
tags=set(self.args.output_tags), expected_responses = self.mock_request(expected_question,
ai=self.args.AI,
model=self.args.model,
file_path=Path('<NOT COMPARED>'))
fake_ai = self.mock_create_ai(self.args, self.config)
expected_responses = fake_ai.request(expected_question,
Chat([]), Chat([]),
self.args.num_answers, self.args.num_answers,
self.args.output_tags).messages self.args.output_tags).messages
@@ -276,12 +297,17 @@ class TestQuestionCmdAsk(TestQuestionCmd):
# execute the command # execute the command
question_cmd(self.args, self.config) question_cmd(self.args, self.config)
# check for correct request call
self.ai.request.assert_called_once_with(expected_question,
ANY,
self.args.num_answers,
self.args.output_tags)
# check for the expected message files # check for the expected message files
chat = ChatDB.from_dir(Path(self.cache_dir.name), chat = ChatDB.from_dir(Path(self.cache_dir.name),
Path(self.db_dir.name)) Path(self.db_dir.name))
cached_msg = chat.msg_gather(loc='cache') cached_msg = chat.msg_gather(loc='cache')
self.assertEqual(len(self.message_list(self.cache_dir)), 1) self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_msgs_equal_except_file_path(cached_msg, expected_responses) self.assert_messages_equal(cached_msg, expected_responses)
@mock.patch('chatmastermind.commands.question.ChatDB.from_dir') @mock.patch('chatmastermind.commands.question.ChatDB.from_dir')
@mock.patch('chatmastermind.commands.question.create_ai') @mock.patch('chatmastermind.commands.question.create_ai')
@@ -292,14 +318,9 @@ class TestQuestionCmdAsk(TestQuestionCmd):
chat = MagicMock(spec=ChatDB) chat = MagicMock(spec=ChatDB)
mock_from_dir.return_value = chat mock_from_dir.return_value = chat
mock_create_ai.side_effect = self.mock_create_ai mock_create_ai.return_value = self.ai
expected_question = Message(Question(self.args.ask[0]), expected_question = self.input_message(self.args)
tags=set(self.args.output_tags), expected_responses = self.mock_request(expected_question,
ai=self.args.AI,
model=self.args.model,
file_path=Path('<NOT COMPARED>'))
fake_ai = self.mock_create_ai(self.args, self.config)
expected_responses = fake_ai.request(expected_question,
Chat([]), Chat([]),
self.args.num_answers, self.args.num_answers,
self.args.output_tags).messages self.args.output_tags).messages
@@ -307,6 +328,12 @@ class TestQuestionCmdAsk(TestQuestionCmd):
# execute the command # execute the command
question_cmd(self.args, self.config) question_cmd(self.args, self.config)
# check for correct request call
self.ai.request.assert_called_once_with(expected_question,
chat,
self.args.num_answers,
self.args.output_tags)
# check for the correct ChatDB calls: # check for the correct ChatDB calls:
# - initial question has been written (prior to the actual request) # - initial question has been written (prior to the actual request)
# - responses have been written (after the request) # - responses have been written (after the request)
@@ -323,98 +350,86 @@ class TestQuestionCmdAsk(TestQuestionCmd):
Provoke an error during the AI request and verify that the question Provoke an error during the AI request and verify that the question
has been correctly stored in the cache. has been correctly stored in the cache.
""" """
mock_create_ai.side_effect = self.mock_create_ai_with_error mock_create_ai.return_value = self.ai
expected_question = Message(Question(self.args.ask[0]), expected_question = self.input_message(self.args)
tags=set(self.args.output_tags), self.ai.request.side_effect = AIError
ai=self.args.AI,
model=self.args.model,
file_path=Path('<NOT COMPARED>'))
# execute the command # execute the command
with self.assertRaises(AIError): with self.assertRaises(AIError):
question_cmd(self.args, self.config) question_cmd(self.args, self.config)
# check for correct request call
self.ai.request.assert_called_once_with(expected_question,
ANY,
self.args.num_answers,
self.args.output_tags)
# check for the expected message files # check for the expected message files
chat = ChatDB.from_dir(Path(self.cache_dir.name), chat = ChatDB.from_dir(Path(self.cache_dir.name),
Path(self.db_dir.name)) Path(self.db_dir.name))
cached_msg = chat.msg_gather(loc='cache') cached_msg = chat.msg_gather(loc='cache')
self.assertEqual(len(self.message_list(self.cache_dir)), 1) self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_msgs_equal_except_file_path(cached_msg, [expected_question]) self.assert_messages_equal(cached_msg, [expected_question])
class TestQuestionCmdRepeat(TestQuestionCmd):
@mock.patch('chatmastermind.commands.question.create_ai') @mock.patch('chatmastermind.commands.question.create_ai')
def test_repeat_single_question(self, mock_create_ai: MagicMock) -> None: def test_repeat_single_question(self, mock_create_ai: MagicMock) -> None:
""" """
Repeat a single question. Repeat a single question.
""" """
mock_create_ai.side_effect = self.mock_create_ai # 1. ask a question
mock_create_ai.return_value = self.ai
expected_question = self.input_message(self.args)
expected_responses = self.mock_request(expected_question,
Chat([]),
self.args.num_answers,
self.args.output_tags).messages
question_cmd(self.args, self.config)
chat = ChatDB.from_dir(Path(self.cache_dir.name), chat = ChatDB.from_dir(Path(self.cache_dir.name),
Path(self.db_dir.name)) Path(self.db_dir.name))
# create a message cached_msg = chat.msg_gather(loc='cache')
message = Message(Question(self.args.ask[0]), self.assertEqual(len(self.message_list(self.cache_dir)), 1)
Answer('Old Answer'), self.assert_messages_equal(cached_msg, expected_responses)
tags=set(self.args.output_tags),
ai=self.args.AI,
model=self.args.model,
file_path=Path(self.cache_dir.name) / f'0001{msg_suffix}')
chat.msg_write([message])
# repeat the last question (without overwriting) # 2. repeat the last question (without overwriting)
# -> expect two identical messages (except for the file_path) # -> expect two identical messages (except for the file_path)
self.args.ask = None self.args.ask = None
self.args.repeat = [] self.args.repeat = []
self.args.overwrite = False self.args.overwrite = False
expected_response = Message(Question(message.question), expected_responses += expected_responses
Answer('Answer 0'),
ai=message.ai,
model=message.model,
tags=message.tags,
file_path=Path('<NOT COMPARED>'))
# we expect the original message + the one with the new response
expected_responses = [message] + [expected_response]
question_cmd(self.args, self.config) question_cmd(self.args, self.config)
cached_msg = chat.msg_gather(loc='cache') cached_msg = chat.msg_gather(loc='cache')
print(self.message_list(self.cache_dir))
self.assertEqual(len(self.message_list(self.cache_dir)), 2) self.assertEqual(len(self.message_list(self.cache_dir)), 2)
self.assert_msgs_equal_except_file_path(cached_msg, expected_responses) self.assert_messages_equal(cached_msg, expected_responses)
@mock.patch('chatmastermind.commands.question.create_ai') @mock.patch('chatmastermind.commands.question.create_ai')
def test_repeat_single_question_overwrite(self, mock_create_ai: MagicMock) -> None: def test_repeat_single_question_overwrite(self, mock_create_ai: MagicMock) -> None:
""" """
Repeat a single question and overwrite the old one. Repeat a single question and overwrite the old one.
""" """
mock_create_ai.side_effect = self.mock_create_ai # 1. ask a question
mock_create_ai.return_value = self.ai
expected_question = self.input_message(self.args)
expected_responses = self.mock_request(expected_question,
Chat([]),
self.args.num_answers,
self.args.output_tags).messages
question_cmd(self.args, self.config)
chat = ChatDB.from_dir(Path(self.cache_dir.name), chat = ChatDB.from_dir(Path(self.cache_dir.name),
Path(self.db_dir.name)) Path(self.db_dir.name))
# create a message
message = Message(Question(self.args.ask[0]),
Answer('Old Answer'),
tags=set(self.args.output_tags),
ai=self.args.AI,
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='cache')
assert cached_msg[0].file_path assert cached_msg[0].file_path
cached_msg_file_id = cached_msg[0].file_path.stem cached_msg_file_id = cached_msg[0].file_path.stem
self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_messages_equal(cached_msg, expected_responses)
# repeat the last question (WITH overwriting) # 2. repeat the last question (WITH overwriting)
# -> expect a single message afterwards (with a new answer) # -> expect a single message afterwards
self.args.ask = None self.args.ask = None
self.args.repeat = [] self.args.repeat = []
self.args.overwrite = True self.args.overwrite = True
expected_response = Message(Question(message.question),
Answer('Answer 0'),
ai=message.ai,
model=message.model,
tags=message.tags,
file_path=Path('<NOT COMPARED>'))
question_cmd(self.args, self.config) question_cmd(self.args, self.config)
cached_msg = chat.msg_gather(loc='cache') cached_msg = chat.msg_gather(loc='cache')
self.assertEqual(len(self.message_list(self.cache_dir)), 1) self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_msgs_equal_except_file_path(cached_msg, [expected_response]) self.assert_messages_equal(cached_msg, expected_responses)
# also check that the file ID has not been changed # also check that the file ID has not been changed
assert cached_msg[0].file_path assert cached_msg[0].file_path
self.assertEqual(cached_msg_file_id, cached_msg[0].file_path.stem) self.assertEqual(cached_msg_file_id, cached_msg[0].file_path.stem)
@@ -424,37 +439,35 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
""" """
Repeat a single question after an error. Repeat a single question after an error.
""" """
mock_create_ai.side_effect = self.mock_create_ai # 1. ask a question and provoke an error
mock_create_ai.return_value = self.ai
expected_question = self.input_message(self.args)
self.ai.request.side_effect = AIError
with self.assertRaises(AIError):
question_cmd(self.args, self.config)
chat = ChatDB.from_dir(Path(self.cache_dir.name), chat = ChatDB.from_dir(Path(self.cache_dir.name),
Path(self.db_dir.name)) Path(self.db_dir.name))
# create a question WITHOUT an answer
# -> just like after an error, which is tested above
message = Message(Question(self.args.ask[0]),
tags=set(self.args.output_tags),
ai=self.args.AI,
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='cache')
assert cached_msg[0].file_path assert cached_msg[0].file_path
cached_msg_file_id = cached_msg[0].file_path.stem cached_msg_file_id = cached_msg[0].file_path.stem
self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_messages_equal(cached_msg, [expected_question])
# repeat the last question (without overwriting) # 2. repeat the last question (without overwriting)
# -> expect a single message because if the original has # -> expect a single message because if the original has
# no answer, it should be overwritten by default # no answer, it should be overwritten by default
self.args.ask = None self.args.ask = None
self.args.repeat = [] self.args.repeat = []
self.args.overwrite = False self.args.overwrite = False
expected_response = Message(Question(message.question), self.ai.request.side_effect = self.mock_request
Answer('Answer 0'), expected_responses = self.mock_request(expected_question,
ai=message.ai, Chat([]),
model=message.model, self.args.num_answers,
tags=message.tags, self.args.output_tags).messages
file_path=Path('<NOT COMPARED>'))
question_cmd(self.args, self.config) question_cmd(self.args, self.config)
cached_msg = chat.msg_gather(loc='cache') cached_msg = chat.msg_gather(loc='cache')
self.assertEqual(len(self.message_list(self.cache_dir)), 1) self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_msgs_equal_except_file_path(cached_msg, [expected_response]) self.assert_messages_equal(cached_msg, expected_responses)
# also check that the file ID has not been changed # also check that the file ID has not been changed
assert cached_msg[0].file_path assert cached_msg[0].file_path
self.assertEqual(cached_msg_file_id, cached_msg[0].file_path.stem) self.assertEqual(cached_msg_file_id, cached_msg[0].file_path.stem)
@@ -464,130 +477,37 @@ class TestQuestionCmdRepeat(TestQuestionCmd):
""" """
Repeat a single question with new arguments. Repeat a single question with new arguments.
""" """
mock_create_ai.side_effect = self.mock_create_ai # 1. ask a question
mock_create_ai.return_value = self.ai
expected_question = self.input_message(self.args)
expected_responses = self.mock_request(expected_question,
Chat([]),
self.args.num_answers,
self.args.output_tags).messages
question_cmd(self.args, self.config)
chat = ChatDB.from_dir(Path(self.cache_dir.name), chat = ChatDB.from_dir(Path(self.cache_dir.name),
Path(self.db_dir.name)) Path(self.db_dir.name))
# create a message
message = Message(Question(self.args.ask[0]),
Answer('Old Answer'),
tags=set(self.args.output_tags),
ai=self.args.AI,
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='cache')
assert cached_msg[0].file_path self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_messages_equal(cached_msg, expected_responses)
# repeat the last question with new arguments (without overwriting) # 2. repeat the last question with new arguments (without overwriting)
# -> expect two messages with identical question but different metadata and new answer # -> expect two messages with identical question and answer, but different metadata
self.args.ask = None self.args.ask = None
self.args.repeat = [] self.args.repeat = []
self.args.overwrite = False self.args.overwrite = False
self.args.output_tags = ['newtag'] self.args.output_tags = ['newtag']
self.args.AI = 'newai' self.args.AI = 'newai'
self.args.model = 'newmodel' self.args.model = 'newmodel'
new_expected_response = Message(Question(message.question), new_expected_question = Message(question=Question(expected_question.question),
Answer('Answer 0'),
ai='newai',
model='newmodel',
tags={Tag('newtag')},
file_path=Path('<NOT COMPARED>'))
question_cmd(self.args, self.config)
cached_msg = chat.msg_gather(loc='cache')
self.assertEqual(len(self.message_list(self.cache_dir)), 2)
self.assert_msgs_equal_except_file_path(cached_msg, [message] + [new_expected_response])
@mock.patch('chatmastermind.commands.question.create_ai')
def test_repeat_single_question_new_args_overwrite(self, mock_create_ai: MagicMock) -> None:
"""
Repeat a single question with new arguments, overwriting the old one.
"""
mock_create_ai.side_effect = self.mock_create_ai
chat = ChatDB.from_dir(Path(self.cache_dir.name),
Path(self.db_dir.name))
# create a message
message = Message(Question(self.args.ask[0]),
Answer('Old Answer'),
tags=set(self.args.output_tags), tags=set(self.args.output_tags),
ai=self.args.AI, ai=self.args.AI,
model=self.args.model, model=self.args.model)
file_path=Path(self.cache_dir.name) / f'0001{msg_suffix}') expected_responses += self.mock_request(new_expected_question,
chat.msg_write([message])
cached_msg = chat.msg_gather(loc='cache')
assert cached_msg[0].file_path
# repeat the last question with new arguments
self.args.ask = None
self.args.repeat = []
self.args.overwrite = True
self.args.output_tags = ['newtag']
self.args.AI = 'newai'
self.args.model = 'newmodel'
new_expected_response = Message(Question(message.question),
Answer('Answer 0'),
ai='newai',
model='newmodel',
tags={Tag('newtag')},
file_path=Path('<NOT COMPARED>'))
question_cmd(self.args, self.config)
cached_msg = chat.msg_gather(loc='cache')
self.assertEqual(len(self.message_list(self.cache_dir)), 1)
self.assert_msgs_equal_except_file_path(cached_msg, [new_expected_response])
@mock.patch('chatmastermind.commands.question.create_ai')
def test_repeat_multiple_questions(self, mock_create_ai: MagicMock) -> None:
"""
Repeat multiple questions.
"""
mock_create_ai.side_effect = self.mock_create_ai
chat = ChatDB.from_dir(Path(self.cache_dir.name),
Path(self.db_dir.name))
# 1. === create three questions ===
# cached message without an answer
message1 = Message(Question(self.args.ask[0]),
tags=self.args.output_tags,
ai=self.args.AI,
model=self.args.model,
file_path=Path(self.cache_dir.name) / f'0001{msg_suffix}')
# cached message with an answer
message2 = Message(Question(self.args.ask[0]),
Answer('Old Answer'),
tags=self.args.output_tags,
ai=self.args.AI,
model=self.args.model,
file_path=Path(self.cache_dir.name) / f'0002{msg_suffix}')
# DB message without an answer
message3 = Message(Question(self.args.ask[0]),
tags=self.args.output_tags,
ai=self.args.AI,
model=self.args.model,
file_path=Path(self.db_dir.name) / f'0003{msg_suffix}')
chat.msg_write([message1, message2, message3])
questions = [message1, message2, message3]
expected_responses: list[Message] = []
fake_ai = self.mock_create_ai(self.args, self.config)
for question in questions:
# since the message's answer is modified, we use a copy
# -> the original is used for comparison below
expected_responses += fake_ai.request(copy(question),
Chat([]), Chat([]),
self.args.num_answers, self.args.num_answers,
set(self.args.output_tags)).messages set(self.args.output_tags)).messages
# 2. === repeat all three questions (without overwriting) ===
self.args.ask = None
self.args.repeat = ['0001', '0002', '0003']
self.args.overwrite = False
question_cmd(self.args, self.config) question_cmd(self.args, self.config)
# two new files should be in the cache directory
# * the repeated cached message with answer
# * the repeated DB message
# -> the cached message without answer should be overwritten
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='cache')
self.assert_msgs_equal_except_file_path(cached_msg, expected_cache_messages) self.assertEqual(len(self.message_list(self.cache_dir)), 2)
# check that the DB message has not been modified at all self.assert_messages_equal(cached_msg, expected_responses)
db_msg = chat.msg_gather(loc='db')
self.assert_msgs_all_equal(db_msg, [message3])