Merge branch 'main' of forgejo.zaydenlab.com:Chewt/ai-assistant

This commit introduces several improvements and new features to the AI Assistant:

- **Refactored code:** Improved code structure and readability.
- **Added clipboard functionality:** Implemented a `/copy` command to copy code blocks to the clipboard.
- **Improved command parsing:** Added a `/copy` command to copy code blocks to the clipboard.
- **Enhanced error handling:** Added error handling for clipboard operations.
- **Improved documentation:** Updated documentation to reflect changes.

refactored.py

@@ -1,4 +1,4 @@
-#/bin/python
+#!/bin/python
 # AI Assistant in the terminal
 import argparse
 import os
@@ -16,6 +16,7 @@ from pygments.formatters import TerminalFormatter
 from prompt_toolkit.key_binding import KeyBindings
 from prompt_toolkit import PromptSession
 
+
 class AIAssistant:
     def __init__(self, server="http://localhost:11434", model="gemma3:12b"):
         self.server = server
@@ -81,12 +82,15 @@ class AIAssistant:
             stream=stream
         )
         result = ''
+        all_chunks = []
         large_chunk = []
         language = None
 
+        if stream:
             for chunk in completion:
                 text = chunk['message']['content']
                 large_chunk.append(text)
+                all_chunks.append(text)
                 large_text = ''.join(large_chunk)
 
                 if ('```' in large_text) and ('\n' in large_text.split('```')[1]):
@@ -96,22 +100,20 @@ class AIAssistant:
                         print(large_text, end='', flush=True)
                         language = None
 
-            if stream:
                 if language and ('\n' in large_text) and large_chunk:
                     output = self.highlight_code(language, large_text)
                     print(output, end='', flush=True)
                     large_chunk = []
                 elif not language or not large_chunk:
                     print(text, end='', flush=True)
-
+            result = ''.join(all_chunks)
-        if not stream:
-            result = completion['message']['content']
         else:
-            result = ''.join(large_chunk)
+            result = completion['message']['content']
         self.history.append({"role": 'assistant', 'content': result})
         self.save_history()
         return result
 
+
 class CommandLineParser:
     def __init__(self):
         self.parser = argparse.ArgumentParser(description='Chat with an intelligent assistant')
@@ -132,16 +134,22 @@ class CommandLineParser:
     def parse(self):
         return self.parser.parse_args()
 
+
 # Keybindings
 bindings = KeyBindings()
 
+
 @bindings.add('c-d')
 def _(event):
     event.current_buffer.validate_and_handle()
 
+
 class InputHandler:
-    def __init__(self, assistant):
+    def __init__(self, assistant, command_parser):
         self.assistant = assistant
+        self.command_parser = command_parser
+        self.command_parser.assistant = assistant
+        self.command_parser.input_handler = self
         self.session = PromptSession(multiline=True, prompt_continuation='', key_bindings=bindings)
 
     def handle_input(self, args):
@@ -150,33 +158,39 @@ class InputHandler:
         else:
             self.handle_interactive_input(args)
 
+    def copy_string_to_clipboard(self, s):
+        try:
+            pyperclip.copy(s)
+        except:
+            return
+
     def handle_piped_input(self, args):
         all_input = sys.stdin.read()
         query = f'Use the following context to answer the question. There will be no follow up questions from the user so make sure your answer is complete:\n{all_input}\n'
         if args.copy:
             query += 'Answer the question using a codeblock for any code or shell scripts\n'
         if args.follow_up:
-            second_input = improved_input()
+            second_input = self.improved_input()
             query += f'\n{second_input}'
         result = self.assistant.chat(query, stream=False)
         blocks = self.extract_code_block(result)
         if args.copy and len(blocks):
-            pyperclip.copy(blocks[0])
+            self.copy_string_to_clipboard(blocks[0])
 
-    def arg_shell(args):
+    def arg_shell(self, args):
         query = '''
 Form a shell command based on the following description. Only output a working shell command. Format the command like this: `command`
 
-Description: 
+Description:\n
         '''
-        if args.shell != True:
+        if type(args.shell) is str:
             query += args.shell
         else:
             query += self.improved_input()
         result = self.assistant.chat(query, stream=False)
         result = blocks[0] if len(blocks := self.extract_code_block(result)) else result
         print(result)
-        copy_string_to_clipboard(result)
+        self.copy_string_to_clipboard(result)
 
     def handle_interactive_input(self, args):
         if args.shell:
@@ -191,7 +205,9 @@ Description:
                     break
                 if full_input.strip() == '':
                     continue
-                result = self.assistant.chat(full_input)
+                command_result = self.command_parser.parse_commands(full_input)
+                if type(command_result) is str:
+                    self.assistant.chat(command_result)
                     print()
             except (EOFError, KeyboardInterrupt):
                 print("\nExiting...")
@@ -209,15 +225,18 @@ Description:
             print("\nUser aborted input")
             return None
 
-    def extract_code_block(self, text):
+    def extract_code_block(self, text, highlight=True):
         pattern = r'```[a-z]*\n[\s\S]*?\n```'
         code_blocks = []
         matches = re.finditer(pattern, text)
         for match in matches:
             code_block = match.group(0)
+            if highlight:
                 lexer_name = self.assistant.determine_lexer(code_block)
                 highlighted_code = self.assistant.highlight_code(lexer_name, code_block)
                 code_blocks.append(highlighted_code)
+            else:
+                code_blocks.append(code_block)
         if not code_blocks:
             line_pattern = r'`[a-z]*[\s\S]*?`'
             line_matches = re.finditer(line_pattern, text)
@@ -226,16 +245,32 @@ Description:
                 code_blocks.append(code_block[1:-1])
         return code_blocks
 
+
 class CommandParser:
     def __init__(self):
         self.commands = {
             '/save': self.handle_save,
             '/clear': self.handle_clear,
-            '/clipboard': self.handle_clipboard,
+            '/clipboard': None,
-            '/exit': self.handle_exit
+            '/exit': self.handle_exit,
+            '/copy': self.handle_copy
         }
 
     def parse_commands(self, text):
+        """
+        Parses the given text to check if it contains a recognized command.
+
+        If the command is standalone returns True.
+        If the command requires passing the text through a string handler (e.g., /clipboard), returns the result of that handler as a string.
+        Otherwise, returns text if no command is recognized.
+
+        Args:
+            text: The input text to parse.
+
+        Returns:
+            True if the command is standalone
+            A string containing the result of the command handler if the command is not standalone.
+        """
         tokens = text.split(' ')
         if not tokens:
             return False
@@ -246,18 +281,19 @@ class CommandParser:
                 context_query = '\n\nThe following is context provided by the user:\n'
                 clipboard_content = pyperclip.paste()
                 if clipboard_content:
-                    context_query += clipboard_content + '\n'
+                    context_query += clipboard_content + '\n' + text
-                return handler(context_query)
+                return context_query
             else:
-                handler()
+                return handler()
-            return True
+        return text
-        return False
 
     def handle_save(self):
         filename = input('Enter filename to save conversation: ')
         self.save_conversation(filename)
+        return True
 
     def save_conversation(self, filename='conversation.md'):
+        # TODO finish this
         if not filename.endswith('.md'):
             filename += '.md'
         base, extension = os.path.splitext(filename)
@@ -271,14 +307,19 @@ class CommandParser:
     def handle_clear(self):
         self.assistant.history = [self.assistant.system_prompt()]
         self.assistant.save_history()
+        return True
 
-    def handle_clipboard(self, context_query):
+    def handle_copy(self):
-        # Implementation for clipboard command
+        blocks = self.input_handler.extract_code_block(self.assistant.history[-1]['content'], highlight=False)
-        pass
+        if len(blocks):
+            block = '\n'.join(blocks[0].split('\n')[1:-1])
+            self.input_handler.copy_string_to_clipboard(block)
+        return True
 
     def handle_exit(self):
         sys.exit(0)
 
+
 def main():
     args = CommandLineParser().parse()
     assistant = AIAssistant()
@@ -296,8 +337,10 @@ def main():
     else:
         assistant.load_history()
 
-    input_handler = InputHandler(assistant)
+    command_parser = CommandParser()
+    input_handler = InputHandler(assistant, command_parser)
     input_handler.handle_input(args)
 
+
 if __name__ == '__main__':
     main()