Compare commits
2 Commits
ddbb8df5c4
...
0dbb0d3c4d
| Author | SHA1 | Date | |
|---|---|---|---|
| 0dbb0d3c4d | |||
| f07c3a58a1 |
@ -146,7 +146,7 @@ class Chat:
|
||||
return sum(m.tokens() for m in self.messages)
|
||||
|
||||
def print(self, dump: bool = False, source_code_only: bool = False,
|
||||
with_tags: bool = False, with_files: bool = False,
|
||||
with_tags: bool = False, with_file: bool = False,
|
||||
paged: bool = True) -> None:
|
||||
if dump:
|
||||
pp(self)
|
||||
@ -157,15 +157,15 @@ class Chat:
|
||||
output.extend(source_code(message.question, include_delims=True))
|
||||
continue
|
||||
output.append('-' * terminal_width())
|
||||
if with_tags:
|
||||
output.append(message.tags_str())
|
||||
if with_files:
|
||||
output.append('FILE: ' + str(message.file_path))
|
||||
output.append(Question.txt_header)
|
||||
output.append(message.question)
|
||||
if message.answer:
|
||||
output.append(Answer.txt_header)
|
||||
output.append(message.answer)
|
||||
if with_tags:
|
||||
output.append(message.tags_str())
|
||||
if with_file:
|
||||
output.append('FILE: ' + str(message.file_path))
|
||||
if paged:
|
||||
print_paged('\n'.join(output))
|
||||
else:
|
||||
|
||||
@ -12,7 +12,6 @@ from .storage import save_answers, create_chat_hist, get_tags_unique, read_file,
|
||||
from .api_client import ai, openai_api_key, print_models
|
||||
from .configuration import Config
|
||||
from .chat import ChatDB
|
||||
from .message import MessageFilter
|
||||
from itertools import zip_longest
|
||||
from typing import Any
|
||||
|
||||
@ -33,11 +32,11 @@ def create_question_with_hist(args: argparse.Namespace,
|
||||
by the specified tags.
|
||||
"""
|
||||
tags = args.tags or []
|
||||
etags = args.etags or []
|
||||
extags = args.extags or []
|
||||
otags = args.output_tags or []
|
||||
|
||||
if not args.source_code_only:
|
||||
print_tag_args(tags, etags, otags)
|
||||
if not args.only_source_code:
|
||||
print_tag_args(tags, extags, otags)
|
||||
|
||||
question_parts = []
|
||||
question_list = args.question if args.question is not None else []
|
||||
@ -54,10 +53,8 @@ def create_question_with_hist(args: argparse.Namespace,
|
||||
question_parts.append(f"```\n{r.read().strip()}\n```")
|
||||
|
||||
full_question = '\n\n'.join(question_parts)
|
||||
chat = create_chat_hist(full_question, tags, etags, config,
|
||||
match_all_tags=True if args.atags else False, # FIXME
|
||||
with_tags=False,
|
||||
with_file=False)
|
||||
chat = create_chat_hist(full_question, tags, extags, config,
|
||||
args.match_all_tags, False, False)
|
||||
return chat, full_question, tags
|
||||
|
||||
|
||||
@ -97,7 +94,7 @@ def ask_cmd(args: argparse.Namespace, config: Config) -> None:
|
||||
if args.model:
|
||||
config.openai.model = args.model
|
||||
chat, question, tags = create_question_with_hist(args, config)
|
||||
print_chat_hist(chat, False, args.source_code_only)
|
||||
print_chat_hist(chat, False, args.only_source_code)
|
||||
otags = args.output_tags or []
|
||||
answers, usage = ai(chat, config, args.number)
|
||||
save_answers(question, answers, tags, otags, config)
|
||||
@ -109,17 +106,14 @@ def hist_cmd(args: argparse.Namespace, config: Config) -> None:
|
||||
"""
|
||||
Handler for the 'hist' command.
|
||||
"""
|
||||
tags = args.tags or []
|
||||
extags = args.extags or []
|
||||
|
||||
mfilter = MessageFilter(tags_or=args.tags,
|
||||
tags_and=args.atags,
|
||||
tags_not=args.etags)
|
||||
chat = ChatDB.from_dir(pathlib.Path('.'),
|
||||
pathlib.Path(config.db),
|
||||
mfilter=mfilter)
|
||||
chat.print(args.dump,
|
||||
args.source_code_only,
|
||||
args.with_tags,
|
||||
args.with_files)
|
||||
chat = create_chat_hist(None, tags, extags, config,
|
||||
args.match_all_tags,
|
||||
args.with_tags,
|
||||
args.with_files)
|
||||
print_chat_hist(chat, args.dump, args.only_source_code)
|
||||
|
||||
|
||||
def print_cmd(args: argparse.Namespace, config: Config) -> None:
|
||||
@ -135,7 +129,7 @@ def print_cmd(args: argparse.Namespace, config: Config) -> None:
|
||||
else:
|
||||
print(f"Unknown file type: {args.file}")
|
||||
sys.exit(1)
|
||||
if args.source_code_only:
|
||||
if args.only_source_code:
|
||||
display_source_code(data['answer'])
|
||||
else:
|
||||
print(dump_data(data).strip())
|
||||
@ -155,17 +149,18 @@ def create_parser() -> argparse.ArgumentParser:
|
||||
# a parent parser for all commands that support tag selection
|
||||
tag_parser = argparse.ArgumentParser(add_help=False)
|
||||
tag_arg = tag_parser.add_argument('-t', '--tags', nargs='+',
|
||||
help='List of tag names (one must match)', metavar='TAGS')
|
||||
help='List of tag names', metavar='TAGS')
|
||||
tag_arg.completer = tags_completer # type: ignore
|
||||
atag_arg = tag_parser.add_argument('-a', '--atags', nargs='+',
|
||||
help='List of tag names (all must match)', metavar='TAGS')
|
||||
atag_arg.completer = tags_completer # type: ignore
|
||||
etag_arg = tag_parser.add_argument('-e', '--etags', nargs='+',
|
||||
help='List of tag names to exclude', metavar='ETAGS')
|
||||
etag_arg.completer = tags_completer # type: ignore
|
||||
extag_arg = tag_parser.add_argument('-e', '--extags', nargs='+',
|
||||
help='List of tag names to exclude', metavar='EXTAGS')
|
||||
extag_arg.completer = tags_completer # type: ignore
|
||||
otag_arg = tag_parser.add_argument('-o', '--output-tags', nargs='+',
|
||||
help='List of output tag names, default is input', metavar='OTAGS')
|
||||
otag_arg.completer = tags_completer # type: ignore
|
||||
tag_parser.add_argument('-a', '--match-all-tags',
|
||||
help="All given tags must match when selecting chat history entries",
|
||||
action='store_true')
|
||||
# enable autocompletion for tags
|
||||
|
||||
# 'ask' command parser
|
||||
ask_cmd_parser = cmdparser.add_parser('ask', parents=[tag_parser],
|
||||
@ -180,7 +175,7 @@ def create_parser() -> argparse.ArgumentParser:
|
||||
ask_cmd_parser.add_argument('-n', '--number', help='Number of answers to produce', type=int,
|
||||
default=1)
|
||||
ask_cmd_parser.add_argument('-s', '--source', nargs='+', help='Source add content of a file to the query')
|
||||
ask_cmd_parser.add_argument('-S', '--source-code-only', help='Add pure source code to the chat history',
|
||||
ask_cmd_parser.add_argument('-S', '--only-source-code', help='Add pure source code to the chat history',
|
||||
action='store_true')
|
||||
|
||||
# 'hist' command parser
|
||||
@ -194,7 +189,7 @@ def create_parser() -> argparse.ArgumentParser:
|
||||
action='store_true')
|
||||
hist_cmd_parser.add_argument('-W', '--with-files', help="Print chat history with filenames.",
|
||||
action='store_true')
|
||||
hist_cmd_parser.add_argument('-S', '--source-code-only', help='Print only source code',
|
||||
hist_cmd_parser.add_argument('-S', '--only-source-code', help='Print only source code',
|
||||
action='store_true')
|
||||
|
||||
# 'tag' command parser
|
||||
@ -226,7 +221,7 @@ 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_parser.add_argument('-S', '--source-code-only', help='Print only source code',
|
||||
print_cmd_parser.add_argument('-S', '--only-source-code', help='Print only source code',
|
||||
action='store_true')
|
||||
|
||||
argcomplete.autocomplete(parser)
|
||||
|
||||
@ -82,21 +82,21 @@ Answer 2
|
||||
@patch('sys.stdout', new_callable=StringIO)
|
||||
def test_print_with_tags_and_file(self, mock_stdout: StringIO) -> None:
|
||||
self.chat.add_msgs([self.message1, self.message2])
|
||||
self.chat.print(paged=False, with_tags=True, with_files=True)
|
||||
self.chat.print(paged=False, with_tags=True, with_file=True)
|
||||
expected_output = f"""{'-'*terminal_width()}
|
||||
{TagLine.prefix} atag1 btag2
|
||||
FILE: 0001.txt
|
||||
{Question.txt_header}
|
||||
Question 1
|
||||
{Answer.txt_header}
|
||||
Answer 1
|
||||
{TagLine.prefix} atag1 btag2
|
||||
FILE: 0001.txt
|
||||
{'-'*terminal_width()}
|
||||
{TagLine.prefix} btag2
|
||||
FILE: 0002.txt
|
||||
{Question.txt_header}
|
||||
Question 2
|
||||
{Answer.txt_header}
|
||||
Answer 2
|
||||
{TagLine.prefix} btag2
|
||||
FILE: 0002.txt
|
||||
"""
|
||||
self.assertEqual(mock_stdout.getvalue(), expected_output)
|
||||
|
||||
|
||||
@ -115,12 +115,11 @@ class TestHandleQuestion(CmmTestCase):
|
||||
self.question = "test question"
|
||||
self.args = argparse.Namespace(
|
||||
tags=['tag1'],
|
||||
atags=None,
|
||||
etags=['etag1'],
|
||||
extags=['extag1'],
|
||||
output_tags=None,
|
||||
question=[self.question],
|
||||
source=None,
|
||||
source_code_only=False,
|
||||
only_source_code=False,
|
||||
number=3,
|
||||
max_tokens=None,
|
||||
temperature=None,
|
||||
@ -144,18 +143,16 @@ class TestHandleQuestion(CmmTestCase):
|
||||
with patch("chatmastermind.storage.open", open_mock):
|
||||
ask_cmd(self.args, self.config)
|
||||
mock_print_tag_args.assert_called_once_with(self.args.tags,
|
||||
self.args.etags,
|
||||
self.args.extags,
|
||||
[])
|
||||
mock_create_chat_hist.assert_called_once_with(self.question,
|
||||
self.args.tags,
|
||||
self.args.etags,
|
||||
self.args.extags,
|
||||
self.config,
|
||||
match_all_tags=False,
|
||||
with_tags=False,
|
||||
with_file=False)
|
||||
False, False, False)
|
||||
mock_print_chat_hist.assert_called_once_with('test_chat',
|
||||
False,
|
||||
self.args.source_code_only)
|
||||
self.args.only_source_code)
|
||||
mock_ai.assert_called_with("test_chat",
|
||||
self.config,
|
||||
self.args.number)
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user