kaizenkodo/ChatMastermind
2
0
Fork 0
Code Issues 2 Pull Requests Packages Projects Releases Wiki Activity

Add possibility to print just the source code

Browse Source
This commit is contained in:
Oleksandr Kozachuk 2023-04-07 23:42:26 +02:00
parent 16f059920b
commit eed180e7a6
3 changed files with 20 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
chatmastermind/main.py
View File

@ -24,7 +24,9 @@ def process_and_display_chat(args: argparse.Namespace,
tags = args.tags or []
extags = args.extags or []
otags = args.output_tags or []
process_tags(tags, extags, otags)
if not args.only_source_code:
process_tags(tags, extags, otags)
question_parts = []
question_list = args.question if args.question is not None else []
@ -45,7 +47,7 @@ def process_and_display_chat(args: argparse.Namespace,
question = '\n\n'.join(question_parts)
chat = create_chat(question, tags, extags, config)
display_chat(chat, dump)
display_chat(chat, dump, args.only_source_code)
return chat, question, tags
@ -82,6 +84,7 @@ def create_parser() -> argparse.ArgumentParser:
parser.add_argument('-M', '--model', help='Model to use')
parser.add_argument('-n', '--number', help='Number of answers to produce', type=int, default=3)
parser.add_argument('-s', '--source', nargs='*', help='Source add content of a file to the query')
parser.add_argument('-S', '--only-source-code', help='Print only source code', action='store_true')
tags_arg = parser.add_argument('-t', '--tags', nargs='*', help='List of tag names', metavar='TAGS')
tags_arg.completer = tags_completer # type: ignore
extags_arg = parser.add_argument('-e', '--extags', nargs='*', help='List of tag names to exclude', metavar='EXTAGS')

19
chatmastermind/utils.py
View File

@ -40,15 +40,24 @@ def message_to_chat(message: Dict[str, str],
append_message(chat, 'assistant', message['answer'])
def display_chat(chat, dump=False) -> None:
def display_chat(chat, dump=False, source_code=False) -> None:
if dump:
pp(chat)
return
for message in chat:
if message['role'] == 'user':
if message['role'] == 'user' and not source_code:
print('-' * (terminal_width()))
if len(message['content']) > terminal_width() - len(message['role']) - 2:
print(f"{message['role'].upper()}:")
print(message['content'])
else:
if not source_code:
print(f"{message['role'].upper()}:")
if source_code:
out = 0
for line in message['content'].splitlines():
if line.strip().startswith('```'):
out += 1
elif out == 1:
print(f"{line}")
else:
print(message['content'])
elif not source_code:
print(f"{message['role'].upper()}: {message['content']}")

1
tests/test_main.py
View File

@ -95,6 +95,7 @@ class TestHandleQuestion(unittest.TestCase):
output_tags=None,
question=[self.question],
source=None,
only_source_code=False,
number=3
)
self.config = {