diff options
Diffstat (limited to '')
| -rw-r--r-- | style.py | 42 |
1 files changed, 37 insertions, 5 deletions
@@ -1,14 +1,47 @@ import sys import time -from prompt_toolkit import prompt +from prompt_toolkit import PromptSession, print_formatted_text, prompt +from prompt_toolkit.formatted_text import HTML +from prompt_toolkit.lexers import PygmentsLexer from prompt_toolkit.styles import Style +from pygments.lexers import PythonLexer +from pygments.styles.native import NativeStyle +from prompt_toolkit.auto_suggest import AutoSuggestFromHistory +from prompt_toolkit.history import FileHistory +from prompt_toolkit.completion import WordCompleter class StyleLog: - def prompt(role: str): + styler = None + + style = Style.from_dict({ + 'input': 'bg:#000000 #00ff00', + 'assistant': 'bg:#000000 #7777ff', + 'system': 'bg:#000000 #ff00ff', + }) + + def __init__(self): + self.styler = PromptSession(lexer=PygmentsLexer(PythonLexer), auto_suggest=AutoSuggestFromHistory(), history=FileHistory('history.txt')) + + def prompt(self, role: str, message: str): + if role == 'assistant': + print_formatted_text(HTML(f"<assistant>Assistant: </assistant>%s") % (message, ), style = self.style) + elif role == 'user': + user_input = prompt( + [ + ('class:input', "\nInput: "), + ('', '') + ], + style = self.style + ) + return user_input + elif role == 'system': + print_formatted_text(HTML(f'<system>System:</system> {message}'), style = self.style) + elif role == 'none': + print_formatted_text(HTML(f'{message}'), style = self.style) return - + def show_ellipsis(self, api_call_done): loop = True while loop: @@ -21,5 +54,4 @@ class StyleLog: sys.stdout.write('.') sys.stdout.flush() sys.stdout.write('\r' + ' ' * 4 + '\r') - sys.stdout.flush() - + sys.stdout.flush() |