From c5d082654f20961e1982bbc203c5a2b19d993af5 Mon Sep 17 00:00:00 2001
From: Schark <jordan@schark.online>
Date: Mon, 4 Dec 2023 13:27:00 -0800
Subject: Adding some styling

---
 style.py | 42 +++++++++++++++++++++++++++++++++++++-----
 1 file changed, 37 insertions(+), 5 deletions(-)

(limited to 'style.py')

diff --git a/style.py b/style.py
index 65fecf1..ad616d8 100644
--- a/style.py
+++ b/style.py
@@ -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() 
-- 
cgit v1.2.3-18-g5258