feat: Refactor code and add clipboard functionality
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.
This commit is contained in:
parent
c0d83584c7
commit
4d6e431563
123
refactored.py
123
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()
|
||||
|
|
|
|||
Loading…
Reference in a new issue