diff --git a/refactored.py b/refactored.py index 8d176e1..5e2ddd0 100644 --- a/refactored.py +++ b/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,37 +82,38 @@ class AIAssistant: stream=stream ) result = '' + all_chunks = [] large_chunk = [] language = None - for chunk in completion: - text = chunk['message']['content'] - large_chunk.append(text) - large_text = ''.join(large_chunk) + 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]): - language = large_text.split('```')[1].split('\n')[0] - large_chunk = [] - if language == '': - print(large_text, end='', flush=True) - language = None + if ('```' in large_text) and ('\n' in large_text.split('```')[1]): + language = large_text.split('```')[1].split('\n')[0] + large_chunk = [] + if language == '': + 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) - - if not stream: - result = completion['message']['content'] + result = ''.join(all_chunks) 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,8 +205,10 @@ Description: break if full_input.strip() == '': continue - result = self.assistant.chat(full_input) - print() + 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...") break @@ -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) - lexer_name = self.assistant.determine_lexer(code_block) - highlighted_code = self.assistant.highlight_code(lexer_name, code_block) - code_blocks.append(highlighted_code) + 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, - '/exit': self.handle_exit + '/clipboard': None, + '/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' - return handler(context_query) + context_query += clipboard_content + '\n' + text + return context_query else: - handler() - return True - return False + return handler() + return text 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): - # Implementation for clipboard command - pass + def handle_copy(self): + blocks = self.input_handler.extract_code_block(self.assistant.history[-1]['content'], highlight=False) + 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()