Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| f3007111e8 | |||
| 29fd722699 | |||
| 810332d140 |
@@ -66,9 +66,3 @@ class AI(Protocol):
|
||||
and is not implemented for all AIs.
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def print(self) -> None:
|
||||
"""
|
||||
Print some info about the current AI, like system message.
|
||||
"""
|
||||
pass
|
||||
|
||||
@@ -4,31 +4,25 @@ Creates different AI instances, based on the given configuration.
|
||||
|
||||
import argparse
|
||||
from typing import cast
|
||||
from .configuration import Config, AIConfig, OpenAIConfig
|
||||
from .configuration import Config, OpenAIConfig, default_ai_ID
|
||||
from .ai import AI, AIError
|
||||
from .ais.openai import OpenAI
|
||||
|
||||
|
||||
def create_ai(args: argparse.Namespace, config: Config) -> AI: # noqa: 11
|
||||
def create_ai(args: argparse.Namespace, config: Config) -> AI:
|
||||
"""
|
||||
Creates an AI subclass instance from the given arguments
|
||||
and configuration file. If AI has not been set in the
|
||||
arguments, it searches for the ID 'default'. If that
|
||||
is not found, it uses the first AI in the list.
|
||||
and configuration file.
|
||||
"""
|
||||
ai_conf: AIConfig
|
||||
if args.AI:
|
||||
if args.ai:
|
||||
try:
|
||||
ai_conf = config.ais[args.AI]
|
||||
ai_conf = config.ais[args.ai]
|
||||
except KeyError:
|
||||
raise AIError(f"AI ID '{args.AI}' does not exist in this configuration")
|
||||
elif 'default' in config.ais:
|
||||
ai_conf = config.ais['default']
|
||||
raise AIError(f"AI ID '{args.ai}' does not exist in this configuration")
|
||||
elif default_ai_ID in config.ais:
|
||||
ai_conf = config.ais[default_ai_ID]
|
||||
else:
|
||||
try:
|
||||
ai_conf = next(iter(config.ais.values()))
|
||||
except StopIteration:
|
||||
raise AIError("No AI found in this configuration")
|
||||
raise AIError("No AI name given and no default exists")
|
||||
|
||||
if ai_conf.name == 'openai':
|
||||
ai = OpenAI(cast(OpenAIConfig, ai_conf))
|
||||
@@ -40,4 +34,4 @@ def create_ai(args: argparse.Namespace, config: Config) -> AI: # noqa: 11
|
||||
ai.config.temperature = args.temperature
|
||||
return ai
|
||||
else:
|
||||
raise AIError(f"AI '{args.AI}' is not supported")
|
||||
raise AIError(f"AI '{args.ai}' is not supported")
|
||||
|
||||
@@ -43,20 +43,16 @@ class OpenAI(AI):
|
||||
n=num_answers,
|
||||
frequency_penalty=self.config.frequency_penalty,
|
||||
presence_penalty=self.config.presence_penalty)
|
||||
question.answer = Answer(response['choices'][0]['message']['content'])
|
||||
question.tags = otags
|
||||
question.ai = self.ID
|
||||
question.model = self.config.model
|
||||
answers: list[Message] = [question]
|
||||
for choice in response['choices'][1:]: # type: ignore
|
||||
answers: list[Message] = []
|
||||
for choice in response['choices']: # type: ignore
|
||||
answers.append(Message(question=question.question,
|
||||
answer=Answer(choice['message']['content']),
|
||||
tags=otags,
|
||||
ai=self.ID,
|
||||
ai=self.name,
|
||||
model=self.config.model))
|
||||
return AIResponse(answers, Tokens(response['usage']['prompt_tokens'],
|
||||
response['usage']['completion_tokens'],
|
||||
response['usage']['total_tokens']))
|
||||
return AIResponse(answers, Tokens(response['usage']['prompt'],
|
||||
response['usage']['completion'],
|
||||
response['usage']['total']))
|
||||
|
||||
def models(self) -> list[str]:
|
||||
"""
|
||||
@@ -99,8 +95,3 @@ class OpenAI(AI):
|
||||
|
||||
def tokens(self, data: Union[Message, Chat]) -> int:
|
||||
raise NotImplementedError
|
||||
|
||||
def print(self) -> None:
|
||||
print(f"MODEL: {self.config.model}")
|
||||
print("=== SYSTEM ===")
|
||||
print(self.config.system)
|
||||
|
||||
+1
-20
@@ -62,10 +62,7 @@ def make_file_path(dir_path: Path,
|
||||
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}{file_suffix}"
|
||||
while file_path.exists():
|
||||
file_path = dir_path / f"{next_fid():04d}{file_suffix}"
|
||||
return file_path
|
||||
return dir_path / f"{next_fid():04d}{file_suffix}"
|
||||
|
||||
|
||||
def write_dir(dir_path: Path,
|
||||
@@ -389,19 +386,3 @@ class ChatDB(Chat):
|
||||
msgs = iter(messages if messages else self.messages)
|
||||
while (m := next(msgs, None)):
|
||||
m.to_file()
|
||||
|
||||
def update_messages(self, messages: list[Message], write: bool = True) -> None:
|
||||
"""
|
||||
Update existing messages. A message is determined as 'existing' if a message with
|
||||
the same base filename (i. e. 'file_path.name') is already in the list. Only accepts
|
||||
existing messages.
|
||||
"""
|
||||
if any(not message_in(m, self.messages) for m in messages):
|
||||
raise ChatError("Can't update messages that are not in the internal list")
|
||||
# remove old versions and add new ones
|
||||
self.messages = [m for m in self.messages if not message_in(m, messages)]
|
||||
self.messages += messages
|
||||
self.sort()
|
||||
# write the UPDATED messages if requested
|
||||
if write:
|
||||
self.write_messages(messages)
|
||||
|
||||
@@ -13,15 +13,7 @@ def print_cmd(args: argparse.Namespace, config: Config) -> None:
|
||||
try:
|
||||
message = Message.from_file(fname)
|
||||
if message:
|
||||
if args.question:
|
||||
print(message.question)
|
||||
elif args.answer:
|
||||
print(message.answer)
|
||||
elif message.answer and args.only_source_code:
|
||||
for code in message.answer.source_code():
|
||||
print(code)
|
||||
else:
|
||||
print(message.to_str())
|
||||
print(message.to_str(source_code_only=args.source_code_only))
|
||||
except MessageError:
|
||||
print(f"File is not a valid message: {args.file}")
|
||||
sys.exit(1)
|
||||
|
||||
@@ -1,9 +1,8 @@
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from itertools import zip_longest
|
||||
from ..configuration import Config
|
||||
from ..chat import ChatDB
|
||||
from ..message import Message, MessageFilter, Question, source_code
|
||||
from ..message import Message, Question
|
||||
from ..ai_factory import create_ai
|
||||
from ..ai import AI, AIResponse
|
||||
|
||||
@@ -12,37 +11,10 @@ def create_message(chat: ChatDB, args: argparse.Namespace) -> Message:
|
||||
"""
|
||||
Creates (and writes) a new message from the given arguments.
|
||||
"""
|
||||
question_parts = []
|
||||
question_list = args.ask if args.ask 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 []
|
||||
|
||||
for question, source, code in zip_longest(question_list, text_files, code_files, fillvalue=None):
|
||||
if question is not None and len(question.strip()) > 0:
|
||||
question_parts.append(question)
|
||||
if source is not None and len(source) > 0:
|
||||
with open(source) as r:
|
||||
content = r.read().strip()
|
||||
if len(content) > 0:
|
||||
question_parts.append(content)
|
||||
if code is not None and len(code) > 0:
|
||||
with open(code) as r:
|
||||
content = r.read().strip()
|
||||
if len(content) == 0:
|
||||
continue
|
||||
# try to extract and add source code
|
||||
code_parts = source_code(content, include_delims=True)
|
||||
if len(code_parts) > 0:
|
||||
question_parts += code_parts
|
||||
# if there's none, add the whole file
|
||||
else:
|
||||
question_parts.append(f"```\n{content}\n```")
|
||||
|
||||
full_question = '\n\n'.join(question_parts)
|
||||
|
||||
message = Message(question=Question(full_question),
|
||||
tags=args.output_tags, # FIXME
|
||||
ai=args.AI,
|
||||
# FIXME: add sources to the question
|
||||
message = Message(question=Question(args.question),
|
||||
tags=args.ouput_tags, # FIXME
|
||||
ai=args.ai,
|
||||
model=args.model)
|
||||
chat.add_to_cache([message])
|
||||
return message
|
||||
@@ -52,12 +24,8 @@ def question_cmd(args: argparse.Namespace, config: Config) -> None:
|
||||
"""
|
||||
Handler for the 'question' command.
|
||||
"""
|
||||
mfilter = MessageFilter(tags_or=args.or_tags if args.or_tags is not None else set(),
|
||||
tags_and=args.and_tags if args.and_tags is not None else set(),
|
||||
tags_not=args.exclude_tags if args.exclude_tags is not None else set())
|
||||
chat = ChatDB.from_dir(cache_path=Path('.'),
|
||||
db_path=Path(config.db),
|
||||
mfilter=mfilter)
|
||||
db_path=Path(config.db))
|
||||
# if it's a new question, create and store it immediately
|
||||
if args.ask or args.create:
|
||||
message = create_message(chat, args)
|
||||
@@ -67,28 +35,23 @@ def question_cmd(args: argparse.Namespace, config: Config) -> None:
|
||||
# create the correct AI instance
|
||||
ai: AI = create_ai(args, config)
|
||||
if args.ask:
|
||||
ai.print()
|
||||
chat.print(paged=False)
|
||||
response: AIResponse = ai.request(message,
|
||||
chat,
|
||||
args.num_answers, # FIXME
|
||||
args.output_tags) # FIXME
|
||||
chat.update_messages([response.messages[0]])
|
||||
chat.add_to_cache(response.messages[1:])
|
||||
for idx, msg in enumerate(response.messages):
|
||||
print(f"=== ANSWER {idx+1} ===")
|
||||
print(msg.answer)
|
||||
if response.tokens:
|
||||
print("===============")
|
||||
print(response.tokens)
|
||||
elif args.repeat is not None:
|
||||
args.otags) # FIXME
|
||||
assert response
|
||||
# TODO:
|
||||
# * add answer to the message above (and create
|
||||
# more messages for any additional answers)
|
||||
pass
|
||||
elif args.repeat:
|
||||
lmessage = chat.latest_message()
|
||||
assert lmessage
|
||||
# TODO: repeat either the last question or the
|
||||
# one(s) given in 'args.repeat' (overwrite
|
||||
# existing ones if 'args.overwrite' is True)
|
||||
pass
|
||||
elif args.process is not None:
|
||||
elif args.process:
|
||||
# TODO: process either all questions without an
|
||||
# answer or the one(s) given in 'args.process'
|
||||
pass
|
||||
|
||||
@@ -9,6 +9,7 @@ OpenAIConfigInst = TypeVar('OpenAIConfigInst', bound='OpenAIConfig')
|
||||
|
||||
|
||||
supported_ais: list[str] = ['openai']
|
||||
default_ai_ID: str = 'default'
|
||||
default_config_path = '.config.yaml'
|
||||
|
||||
|
||||
@@ -16,18 +17,6 @@ class ConfigError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def str_presenter(dumper: yaml.Dumper, data: str) -> yaml.ScalarNode:
|
||||
"""
|
||||
Changes the YAML dump style to multiline syntax for multiline strings.
|
||||
"""
|
||||
if len(data.splitlines()) > 1:
|
||||
return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')
|
||||
return dumper.represent_scalar('tag:yaml.org,2002:str', data)
|
||||
|
||||
|
||||
yaml.add_representer(str, str_presenter)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AIConfig:
|
||||
"""
|
||||
@@ -57,15 +46,15 @@ class OpenAIConfig(AIConfig):
|
||||
|
||||
# all members have default values, so we can easily create
|
||||
# a default configuration
|
||||
ID: str = 'myopenai'
|
||||
ID: str = 'default'
|
||||
api_key: str = '0123456789'
|
||||
system: str = 'You are an assistant'
|
||||
model: str = 'gpt-3.5-turbo-16k'
|
||||
temperature: float = 1.0
|
||||
max_tokens: int = 4000
|
||||
top_p: float = 1.0
|
||||
frequency_penalty: float = 0.0
|
||||
presence_penalty: float = 0.0
|
||||
system: str = 'You are an assistant'
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls: Type[OpenAIConfigInst], source: dict[str, Any]) -> OpenAIConfigInst:
|
||||
@@ -73,14 +62,14 @@ class OpenAIConfig(AIConfig):
|
||||
Create OpenAIConfig from a dict.
|
||||
"""
|
||||
res = cls(
|
||||
system=str(source['system']),
|
||||
api_key=str(source['api_key']),
|
||||
model=str(source['model']),
|
||||
max_tokens=int(source['max_tokens']),
|
||||
temperature=float(source['temperature']),
|
||||
top_p=float(source['top_p']),
|
||||
frequency_penalty=float(source['frequency_penalty']),
|
||||
presence_penalty=float(source['presence_penalty']),
|
||||
system=str(source['system'])
|
||||
presence_penalty=float(source['presence_penalty'])
|
||||
)
|
||||
# overwrite default ID if provided
|
||||
if 'ID' in source:
|
||||
@@ -159,8 +148,6 @@ class Config:
|
||||
|
||||
def as_dict(self) -> dict[str, Any]:
|
||||
res = asdict(self)
|
||||
# add the AI name manually (as first element)
|
||||
# (not done by 'asdict' because it's a class variable)
|
||||
for ID, conf in res['ais'].items():
|
||||
res['ais'][ID] = {**{'name': self.ais[ID].name}, **conf}
|
||||
conf.update({'name': self.ais[ID].name})
|
||||
return res
|
||||
|
||||
@@ -67,8 +67,9 @@ def create_parser() -> argparse.ArgumentParser:
|
||||
question_group.add_argument('-p', '--process', nargs='*', help='Process existing questions')
|
||||
question_cmd_parser.add_argument('-O', '--overwrite', help='Overwrite existing messages when repeating them',
|
||||
action='store_true')
|
||||
question_cmd_parser.add_argument('-s', '--source-text', nargs='+', help='Add content of a file to the query')
|
||||
question_cmd_parser.add_argument('-S', '--source-code', nargs='+', help='Add source code file content to the chat history')
|
||||
question_cmd_parser.add_argument('-s', '--source', nargs='+', help='Source add content of a file to the query')
|
||||
question_cmd_parser.add_argument('-S', '--source-code-only', help='Add pure source code to the chat history',
|
||||
action='store_true')
|
||||
|
||||
# 'hist' command parser
|
||||
hist_cmd_parser = cmdparser.add_parser('hist', parents=[tag_parser],
|
||||
@@ -113,10 +114,8 @@ def create_parser() -> argparse.ArgumentParser:
|
||||
aliases=['p'])
|
||||
print_cmd_parser.set_defaults(func=print_cmd)
|
||||
print_cmd_parser.add_argument('-f', '--file', help='File to print', required=True)
|
||||
print_cmd_modes = print_cmd_parser.add_mutually_exclusive_group()
|
||||
print_cmd_modes.add_argument('-q', '--question', help='Print only question', action='store_true')
|
||||
print_cmd_modes.add_argument('-a', '--answer', help='Print only answer', action='store_true')
|
||||
print_cmd_modes.add_argument('-S', '--only-source-code', help='Print only source code', action='store_true')
|
||||
print_cmd_parser.add_argument('-S', '--source-code-only', help='Print source code only (from the answer, if available)',
|
||||
action='store_true')
|
||||
|
||||
argcomplete.autocomplete(parser)
|
||||
return parser
|
||||
|
||||
+11
-17
@@ -3,8 +3,6 @@ Module implementing message related functions and classes.
|
||||
"""
|
||||
import pathlib
|
||||
import yaml
|
||||
import tempfile
|
||||
import shutil
|
||||
from typing import Type, TypeVar, ClassVar, Optional, Any, Union, Final, Literal, Iterable
|
||||
from dataclasses import dataclass, asdict, field
|
||||
from .tags import Tag, TagLine, TagError, match_tags, rename_tags
|
||||
@@ -314,7 +312,7 @@ class Message():
|
||||
mfilter.tags_not if mfilter else None)
|
||||
else:
|
||||
message = cls.__from_file_yaml(file_path)
|
||||
if message and (mfilter is None or message.match(mfilter)):
|
||||
if message and (not mfilter or (mfilter and message.match(mfilter))):
|
||||
return message
|
||||
else:
|
||||
return None
|
||||
@@ -416,7 +414,7 @@ class Message():
|
||||
return '\n'.join(output)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return self.to_str(True, True, False)
|
||||
return self.to_str(False, False, False)
|
||||
|
||||
def to_file(self, file_path: Optional[pathlib.Path]=None) -> None: # noqa: 11
|
||||
"""
|
||||
@@ -447,18 +445,16 @@ class Message():
|
||||
* Answer.txt_header
|
||||
* Answer
|
||||
"""
|
||||
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)
|
||||
with open(file_path, "w") as fd:
|
||||
if self.tags:
|
||||
temp_fd.write(f'{TagLine.from_set(self.tags)}\n')
|
||||
fd.write(f'{TagLine.from_set(self.tags)}\n')
|
||||
if self.ai:
|
||||
temp_fd.write(f'{AILine.from_ai(self.ai)}\n')
|
||||
fd.write(f'{AILine.from_ai(self.ai)}\n')
|
||||
if self.model:
|
||||
temp_fd.write(f'{ModelLine.from_model(self.model)}\n')
|
||||
temp_fd.write(f'{Question.txt_header}\n{self.question}\n')
|
||||
fd.write(f'{ModelLine.from_model(self.model)}\n')
|
||||
fd.write(f'{Question.txt_header}\n{self.question}\n')
|
||||
if self.answer:
|
||||
temp_fd.write(f'{Answer.txt_header}\n{self.answer}\n')
|
||||
shutil.move(temp_file_path, file_path)
|
||||
fd.write(f'{Answer.txt_header}\n{self.answer}\n')
|
||||
|
||||
def __to_file_yaml(self, file_path: pathlib.Path) -> None:
|
||||
"""
|
||||
@@ -470,8 +466,7 @@ class Message():
|
||||
* Message.ai_yaml_key: str [Optional]
|
||||
* Message.model_yaml_key: str [Optional]
|
||||
"""
|
||||
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)
|
||||
with open(file_path, "w") as fd:
|
||||
data: YamlDict = {Question.yaml_key: str(self.question)}
|
||||
if self.answer:
|
||||
data[Answer.yaml_key] = str(self.answer)
|
||||
@@ -481,8 +476,7 @@ class Message():
|
||||
data[self.model_yaml_key] = self.model
|
||||
if self.tags:
|
||||
data[self.tags_yaml_key] = sorted([str(tag) for tag in self.tags])
|
||||
yaml.dump(data, temp_fd, sort_keys=False)
|
||||
shutil.move(temp_file_path, file_path)
|
||||
yaml.dump(data, fd, sort_keys=False)
|
||||
|
||||
def filter_tags(self, prefix: Optional[str] = None, contain: Optional[str] = None) -> set[Tag]:
|
||||
"""
|
||||
@@ -514,7 +508,7 @@ class Message():
|
||||
Return True if all attributes match, else False.
|
||||
"""
|
||||
mytags = self.tags or set()
|
||||
if (((mfilter.tags_or is not None or mfilter.tags_and is not None or mfilter.tags_not is not None)
|
||||
if (((mfilter.tags_or or mfilter.tags_and or mfilter.tags_not)
|
||||
and not match_tags(mytags, mfilter.tags_or, mfilter.tags_and, mfilter.tags_not)) # noqa: W503
|
||||
or (mfilter.ai and (not self.ai or mfilter.ai != self.ai)) # noqa: W503
|
||||
or (mfilter.model and (not self.model or mfilter.model != self.model)) # noqa: W503
|
||||
|
||||
@@ -10,7 +10,7 @@ from chatmastermind.ais.openai import OpenAI
|
||||
class TestCreateAI(unittest.TestCase):
|
||||
def setUp(self) -> None:
|
||||
self.args = MagicMock(spec=argparse.Namespace)
|
||||
self.args.AI = 'myopenai'
|
||||
self.args.ai = 'default'
|
||||
self.args.model = None
|
||||
self.args.max_tokens = None
|
||||
self.args.temperature = None
|
||||
@@ -18,19 +18,19 @@ class TestCreateAI(unittest.TestCase):
|
||||
def test_create_ai_from_args(self) -> None:
|
||||
# Create an AI with the default configuration
|
||||
config = Config()
|
||||
self.args.AI = 'myopenai'
|
||||
self.args.ai = 'default'
|
||||
ai = create_ai(self.args, config)
|
||||
self.assertIsInstance(ai, OpenAI)
|
||||
|
||||
def test_create_ai_from_default(self) -> None:
|
||||
self.args.AI = None
|
||||
self.args.ai = None
|
||||
# Create an AI with the default configuration
|
||||
config = Config()
|
||||
ai = create_ai(self.args, config)
|
||||
self.assertIsInstance(ai, OpenAI)
|
||||
|
||||
def test_create_empty_ai_error(self) -> None:
|
||||
self.args.AI = None
|
||||
self.args.ai = None
|
||||
# Create Config with empty AIs
|
||||
config = Config()
|
||||
config.ais = {}
|
||||
@@ -40,7 +40,7 @@ class TestCreateAI(unittest.TestCase):
|
||||
|
||||
def test_create_unsupported_ai_error(self) -> None:
|
||||
# Mock argparse.Namespace with ai='invalid_ai'
|
||||
self.args.AI = 'invalid_ai'
|
||||
self.args.ai = 'invalid_ai'
|
||||
# Create default Config
|
||||
config = Config()
|
||||
# Call create_ai function and assert that it raises AIError
|
||||
|
||||
+2
-48
@@ -202,25 +202,7 @@ class TestChatDB(unittest.TestCase):
|
||||
self.assertEqual(chat_db.messages[1].file_path,
|
||||
pathlib.Path(self.db_path.name, '0003.txt'))
|
||||
|
||||
def test_chat_db_from_dir_filter_tags(self) -> None:
|
||||
chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),
|
||||
pathlib.Path(self.db_path.name),
|
||||
mfilter=MessageFilter(tags_or={Tag('tag1')}))
|
||||
self.assertEqual(len(chat_db.messages), 1)
|
||||
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.messages[0].file_path,
|
||||
pathlib.Path(self.db_path.name, '0001.txt'))
|
||||
|
||||
def test_chat_db_from_dir_filter_tags_empty(self) -> None:
|
||||
chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),
|
||||
pathlib.Path(self.db_path.name),
|
||||
mfilter=MessageFilter(tags_or=set(),
|
||||
tags_and=set(),
|
||||
tags_not=set()))
|
||||
self.assertEqual(len(chat_db.messages), 0)
|
||||
|
||||
def test_chat_db_from_dir_filter_answer(self) -> None:
|
||||
def test_chat_db_filter(self) -> None:
|
||||
chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),
|
||||
pathlib.Path(self.db_path.name),
|
||||
mfilter=MessageFilter(answer_contains='Answer 2'))
|
||||
@@ -231,7 +213,7 @@ class TestChatDB(unittest.TestCase):
|
||||
pathlib.Path(self.db_path.name, '0002.yaml'))
|
||||
self.assertEqual(chat_db.messages[0].answer, 'Answer 2')
|
||||
|
||||
def test_chat_db_from_messages(self) -> None:
|
||||
def test_chat_db_from_messges(self) -> None:
|
||||
chat_db = ChatDB.from_messages(pathlib.Path(self.cache_path.name),
|
||||
pathlib.Path(self.db_path.name),
|
||||
messages=[self.message1, self.message2,
|
||||
@@ -458,31 +440,3 @@ class TestChatDB(unittest.TestCase):
|
||||
cache_dir_files = self.message_list(self.cache_path)
|
||||
self.assertEqual(len(cache_dir_files), 1)
|
||||
self.assertIn(pathlib.Path(self.cache_path.name, '123456.txt'), cache_dir_files)
|
||||
|
||||
def test_chat_db_update_messages(self) -> None:
|
||||
# create a new ChatDB instance
|
||||
chat_db = ChatDB.from_dir(pathlib.Path(self.cache_path.name),
|
||||
pathlib.Path(self.db_path.name))
|
||||
|
||||
db_dir_files = self.message_list(self.db_path)
|
||||
self.assertEqual(len(db_dir_files), 4)
|
||||
cache_dir_files = self.message_list(self.cache_path)
|
||||
self.assertEqual(len(cache_dir_files), 0)
|
||||
|
||||
message = chat_db.messages[0]
|
||||
message.answer = Answer("New answer")
|
||||
# update message without writing
|
||||
chat_db.update_messages([message], write=False)
|
||||
self.assertEqual(chat_db.messages[0].answer, Answer("New answer"))
|
||||
# re-read the message and check for old content
|
||||
chat_db.read_db()
|
||||
self.assertEqual(chat_db.messages[0].answer, Answer("Answer 1"))
|
||||
# now check with writing (message should be overwritten)
|
||||
chat_db.update_messages([message], write=True)
|
||||
chat_db.read_db()
|
||||
self.assertEqual(chat_db.messages[0].answer, Answer("New answer"))
|
||||
# test without file_path -> expect error
|
||||
message1 = Message(question=Question("Question 1"),
|
||||
answer=Answer("Answer 1"))
|
||||
with self.assertRaises(ChatError):
|
||||
chat_db.update_messages([message1])
|
||||
|
||||
@@ -59,7 +59,7 @@ class TestConfig(unittest.TestCase):
|
||||
source_dict = {
|
||||
'db': './test_db/',
|
||||
'ais': {
|
||||
'myopenai': {
|
||||
'default': {
|
||||
'name': 'openai',
|
||||
'system': 'Custom system',
|
||||
'api_key': '9876543210',
|
||||
@@ -75,10 +75,10 @@ class TestConfig(unittest.TestCase):
|
||||
config = Config.from_dict(source_dict)
|
||||
self.assertEqual(config.db, './test_db/')
|
||||
self.assertEqual(len(config.ais), 1)
|
||||
self.assertEqual(config.ais['myopenai'].name, 'openai')
|
||||
self.assertEqual(cast(OpenAIConfig, config.ais['myopenai']).system, 'Custom system')
|
||||
self.assertEqual(config.ais['default'].name, 'openai')
|
||||
self.assertEqual(cast(OpenAIConfig, config.ais['default']).system, 'Custom system')
|
||||
# check that 'ID' has been added
|
||||
self.assertEqual(config.ais['myopenai'].ID, 'myopenai')
|
||||
self.assertEqual(config.ais['default'].ID, 'default')
|
||||
|
||||
def test_create_default_should_create_default_config(self) -> None:
|
||||
Config.create_default(Path(self.test_file.name))
|
||||
@@ -117,8 +117,8 @@ class TestConfig(unittest.TestCase):
|
||||
config = Config(
|
||||
db='./test_db/',
|
||||
ais={
|
||||
'myopenai': OpenAIConfig(
|
||||
ID='myopenai',
|
||||
'default': OpenAIConfig(
|
||||
ID='default',
|
||||
system='Custom system',
|
||||
api_key='9876543210',
|
||||
model='custom_model',
|
||||
@@ -135,7 +135,7 @@ class TestConfig(unittest.TestCase):
|
||||
saved_config = yaml.load(f, Loader=yaml.FullLoader)
|
||||
self.assertEqual(saved_config['db'], './test_db/')
|
||||
self.assertEqual(len(saved_config['ais']), 1)
|
||||
self.assertEqual(saved_config['ais']['myopenai']['system'], 'Custom system')
|
||||
self.assertEqual(saved_config['ais']['default']['system'], 'Custom system')
|
||||
|
||||
def test_from_file_error_unknown_ai(self) -> None:
|
||||
source_dict = {
|
||||
|
||||
@@ -300,12 +300,6 @@ This is a question.
|
||||
MessageFilter(tags_or={Tag('tag1')}))
|
||||
self.assertIsNone(message)
|
||||
|
||||
def test_from_file_txt_empty_tags_dont_match(self) -> None:
|
||||
message = Message.from_file(self.file_path_min,
|
||||
MessageFilter(tags_or=set(),
|
||||
tags_and=set()))
|
||||
self.assertIsNone(message)
|
||||
|
||||
def test_from_file_txt_no_tags_match_tags_not(self) -> None:
|
||||
message = Message.from_file(self.file_path_min,
|
||||
MessageFilter(tags_not={Tag('tag1')}))
|
||||
|
||||
@@ -1,162 +0,0 @@
|
||||
import os
|
||||
import unittest
|
||||
import argparse
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
from chatmastermind.commands.question import create_message
|
||||
from chatmastermind.message import Message, Question
|
||||
from chatmastermind.chat import ChatDB
|
||||
|
||||
|
||||
class TestMessageCreate(unittest.TestCase):
|
||||
"""
|
||||
Test if messages created by the 'question' command have
|
||||
the correct format.
|
||||
"""
|
||||
def setUp(self) -> None:
|
||||
# create ChatDB structure
|
||||
self.db_path = tempfile.TemporaryDirectory()
|
||||
self.cache_path = tempfile.TemporaryDirectory()
|
||||
self.chat = ChatDB.from_dir(cache_path=Path(self.cache_path.name),
|
||||
db_path=Path(self.db_path.name))
|
||||
# create arguments mock
|
||||
self.args = MagicMock(spec=argparse.Namespace)
|
||||
self.args.source_text = None
|
||||
self.args.source_code = None
|
||||
self.args.AI = None
|
||||
self.args.model = None
|
||||
self.args.output_tags = None
|
||||
# File 1 : no source code block, only text
|
||||
self.source_file1 = tempfile.NamedTemporaryFile(delete=False)
|
||||
self.source_file1_content = """This is just text.
|
||||
No source code.
|
||||
Nope. Go look elsewhere!"""
|
||||
with open(self.source_file1.name, 'w') as f:
|
||||
f.write(self.source_file1_content)
|
||||
# File 2 : one embedded source code block
|
||||
self.source_file2 = tempfile.NamedTemporaryFile(delete=False)
|
||||
self.source_file2_content = """This is just text.
|
||||
```
|
||||
This is embedded source code.
|
||||
```
|
||||
And some text again."""
|
||||
with open(self.source_file2.name, 'w') as f:
|
||||
f.write(self.source_file2_content)
|
||||
# File 3 : all source code
|
||||
self.source_file3 = tempfile.NamedTemporaryFile(delete=False)
|
||||
self.source_file3_content = """This is all source code.
|
||||
Yes, really.
|
||||
Language is called 'brainfart'."""
|
||||
with open(self.source_file3.name, 'w') as f:
|
||||
f.write(self.source_file3_content)
|
||||
# File 4 : two source code blocks
|
||||
self.source_file4 = tempfile.NamedTemporaryFile(delete=False)
|
||||
self.source_file4_content = """This is just text.
|
||||
```
|
||||
This is embedded source code.
|
||||
```
|
||||
And some text again.
|
||||
```
|
||||
This is embedded source code.
|
||||
```
|
||||
Aaaand again some text."""
|
||||
with open(self.source_file4.name, 'w') as f:
|
||||
f.write(self.source_file4_content)
|
||||
|
||||
def tearDown(self) -> None:
|
||||
os.remove(self.source_file1.name)
|
||||
os.remove(self.source_file2.name)
|
||||
os.remove(self.source_file3.name)
|
||||
|
||||
def message_list(self, tmp_dir: tempfile.TemporaryDirectory) -> list[Path]:
|
||||
# exclude '.next'
|
||||
return list(Path(tmp_dir.name).glob('*.[ty]*'))
|
||||
|
||||
def test_message_file_created(self) -> None:
|
||||
self.args.ask = ["What is this?"]
|
||||
cache_dir_files = self.message_list(self.cache_path)
|
||||
self.assertEqual(len(cache_dir_files), 0)
|
||||
create_message(self.chat, self.args)
|
||||
cache_dir_files = self.message_list(self.cache_path)
|
||||
self.assertEqual(len(cache_dir_files), 1)
|
||||
message = Message.from_file(cache_dir_files[0])
|
||||
self.assertIsInstance(message, Message)
|
||||
self.assertEqual(message.question, Question("What is this?")) # type: ignore [union-attr]
|
||||
|
||||
def test_single_question(self) -> None:
|
||||
self.args.ask = ["What is this?"]
|
||||
message = create_message(self.chat, self.args)
|
||||
self.assertIsInstance(message, Message)
|
||||
self.assertEqual(message.question, Question("What is this?"))
|
||||
self.assertEqual(len(message.question.source_code()), 0)
|
||||
|
||||
def test_multipart_question(self) -> None:
|
||||
self.args.ask = ["What is this", "'bard' thing?", "Is it good?"]
|
||||
message = create_message(self.chat, self.args)
|
||||
self.assertIsInstance(message, Message)
|
||||
self.assertEqual(message.question, Question("""What is this
|
||||
|
||||
'bard' thing?
|
||||
|
||||
Is it good?"""))
|
||||
|
||||
def test_single_question_with_text_only_file(self) -> None:
|
||||
self.args.ask = ["What is this?"]
|
||||
self.args.source_text = [f"{self.source_file1.name}"]
|
||||
message = create_message(self.chat, self.args)
|
||||
self.assertIsInstance(message, Message)
|
||||
# file contains no source code (only text)
|
||||
# -> don't expect any in the question
|
||||
self.assertEqual(len(message.question.source_code()), 0)
|
||||
self.assertEqual(message.question, Question(f"""What is this?
|
||||
|
||||
{self.source_file1_content}"""))
|
||||
|
||||
def test_single_question_with_text_file_and_embedded_code(self) -> None:
|
||||
self.args.ask = ["What is this?"]
|
||||
self.args.source_code = [f"{self.source_file2.name}"]
|
||||
message = create_message(self.chat, self.args)
|
||||
self.assertIsInstance(message, Message)
|
||||
# file contains 1 source code block
|
||||
# -> expect it in the question
|
||||
self.assertEqual(len(message.question.source_code()), 1)
|
||||
self.assertEqual(message.question, Question("""What is this?
|
||||
|
||||
```
|
||||
This is embedded source code.
|
||||
```
|
||||
"""))
|
||||
|
||||
def test_single_question_with_code_only_file(self) -> None:
|
||||
self.args.ask = ["What is this?"]
|
||||
self.args.source_code = [f"{self.source_file3.name}"]
|
||||
message = create_message(self.chat, self.args)
|
||||
self.assertIsInstance(message, Message)
|
||||
# file is complete source code
|
||||
self.assertEqual(len(message.question.source_code()), 1)
|
||||
self.assertEqual(message.question, Question(f"""What is this?
|
||||
|
||||
```
|
||||
{self.source_file3_content}
|
||||
```"""))
|
||||
|
||||
def test_single_question_with_text_file_and_multi_embedded_code(self) -> None:
|
||||
self.args.ask = ["What is this?"]
|
||||
self.args.source_code = [f"{self.source_file4.name}"]
|
||||
message = create_message(self.chat, self.args)
|
||||
self.assertIsInstance(message, Message)
|
||||
# file contains 2 source code blocks
|
||||
# -> expect them in the question
|
||||
self.assertEqual(len(message.question.source_code()), 2)
|
||||
self.assertEqual(message.question, Question("""What is this?
|
||||
|
||||
```
|
||||
This is embedded source code.
|
||||
```
|
||||
|
||||
|
||||
```
|
||||
This is embedded source code.
|
||||
```
|
||||
"""))
|
||||
Reference in New Issue
Block a user