#!/bin/python # AI Assistant in the terminal import argparse import os import sys import json from ollama import Client import re import pyperclip import pygments from pygments.lexers import get_lexer_by_name, guess_lexer 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 self.model = model self.client = Client(host=self.server) self.temperature = 0.2 self.num_ctx = 4096 self.history = [self.system_prompt()] def set_host(self, host): self.server = host self.client = Client(host=host) def system_prompt(self): return {"role": "system", "content": "You are a helpful, smart, kind, and efficient AI assistant. You always fulfill the user's requests accurately and concisely."} def load_history(self): path = os.environ.get('HOME') + '/.cache/ai-assistant.history' try: with open(path, 'r') as f: self.history = json.load(f) except FileNotFoundError: pass def save_history(self): path = os.environ.get('HOME') + '/.cache/ai-assistant.history' with open(path, 'w+') as f: json.dump(self.history, f) def determine_lexer(self, code_block): lexer_name = None lines = code_block.split('\n') for line in lines: if line.strip().startswith('```'): lexer_part = line.strip().split('```')[1].strip() if lexer_part: lexer_name = lexer_part break elif line.strip().startswith('lang:'): lexer_part = line.strip().split(':')[1].strip() if lexer_part: lexer_name = lexer_part break return lexer_name def highlight_code(self, lexer_name, code): try: lexer = get_lexer_by_name(lexer_name) if lexer_name else guess_lexer(code) except ValueError: lexer = guess_lexer('\n'.join(code.split('\n')[1:-1])) if not lexer: lexer = get_lexer_by_name('bash') formatter = TerminalFormatter() highlighted_code = pygments.highlight(code, lexer, formatter) return highlighted_code def chat(self, message, stream=True): self.history.append({"role": "user", "content": message}) completion = self.client.chat( model=self.model, options={"temperature": self.temperature, "num_ctx": self.num_ctx}, messages=self.history, 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]): language = large_text.split('```')[1].split('\n')[0] large_chunk = [] if language == '': print(large_text, end='', flush=True) language = None 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) else: 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') self.add_arguments() def add_arguments(self): parser = self.parser parser.add_argument('--host', nargs='?', const=True, default=False, help='Specify host of Ollama server') parser.add_argument('--model', '-m', nargs='?', const=True, default=False, help='Specify model') parser.add_argument('--temp', '-t', nargs='?', type=float, const=0.2, default=False, help='Specify temperature') parser.add_argument('--context', type=int, default=4096, help='Specify context size') parser.add_argument('--reasoning', '-r', action='store_true', help='Use the default reasoning model deepseek-r1:14b') parser.add_argument('--new', '-n', action='store_true', help='Start a chat with a fresh history') parser.add_argument('--follow-up', '-f', nargs='?', const=True, default=False, help='Ask a follow up question when piping in context') parser.add_argument('--copy', '-c', action='store_true', help='Copy a codeblock if it appears') parser.add_argument('--shell', '-s', nargs='?', const=True, default=False, help='Output a shell command that does as described') 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, 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): if not sys.stdin.isatty(): self.handle_piped_input(args) 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 = 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): self.copy_string_to_clipboard(blocks[0]) 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:\n ''' 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) self.copy_string_to_clipboard(result) def handle_interactive_input(self, args): if args.shell: self.arg_shell(args) exit() print("\033[91massistant\033[0m: Type your message (press Ctrl+D to send):") while True: try: full_input = self.improved_input() if full_input is None: break if full_input.strip() == '': continue 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 def improved_input(self, prompt="> "): """ Returns the full text (including embedded newlines) when you press Ctrl-D. Arrow keys edit within or across lines automatically. """ try: text = self.session.prompt(prompt) return text except KeyboardInterrupt: print("\nUser aborted input") return None 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) for match in line_matches: code_block = match.group(0) 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': 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 command = tokens[0] if command in self.commands: handler = self.commands[command] if len(tokens) > 1 and command == '/clipboard': context_query = '\n\nThe following is context provided by the user:\n' clipboard_content = pyperclip.paste() if clipboard_content: context_query += clipboard_content + '\n' + text return context_query else: 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) i = 1 while os.path.exists(filename): filename = f"{base}_{i}{extension}" i += 1 with open(filename, 'w') as f: f.write(conversation) def handle_clear(self): self.assistant.history = [self.assistant.system_prompt()] self.assistant.save_history() return True 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() if args.host: assistant.set_host(args.host) if args.model: assistant.model = args.model if args.temp: assistant.temperature = args.temp if args.context: assistant.num_ctx = args.context if args.new: assistant.history = [assistant.system_prompt()] assistant.save_history() else: assistant.load_history() command_parser = CommandParser() input_handler = InputHandler(assistant, command_parser) input_handler.handle_input(args) if __name__ == '__main__': main()