Intial commit. Switch from OpenAI to Ollama module

assistant.py

@@ -0,0 +1,140 @@
+#!/bin/python3
+# Chat with an intelligent assistant in your terminal
+from ollama import Client
+import re
+import pyperclip
+import sys
+import argparse
+
+model = 'mistral-nemo'
+
+pattern = r'```[a-z]*\n[\s\S]*?\n```'
+line_pattern = r'`[a-z]*[\s\S]*?`'
+
+def extract_code_block(markdown_text):
+    # Use the regular expression pattern to find all matches in the markdown text
+    matches = re.finditer(pattern, markdown_text)
+
+    # Iterate over the matches and extract the code blocks
+    code_blocks = []
+    for match in matches:
+        code_block = match.group(0)
+
+        # Add the code block to the list of code blocks
+        code_blocks.append('\n'.join(code_block.split('\n')[1:-1]))
+
+    if len(code_blocks) == 0:
+        line_matches = re.finditer(line_pattern, markdown_text)
+        for match in line_matches:
+            code_block = match.group(0)
+            code_blocks.append(code_block[1:-1])
+    return code_blocks
+
+def copy_string_to_clipboard(string):
+    try:
+        pyperclip.copy(string)
+    except:
+        return
+
+# Point to the local server
+client = Client(host='hayden-desktop:11434')
+
+code_history = []
+
+history = [
+        {"role": "system", "content": "You are a helpful, smart, kind, and efficient AI assistant. You always fulfill the user's requests accurately and concisely."},
+        ]
+
+def chat(message, stream=True):
+    history.append({"role": "user", "content": message})
+    completion = client.chat(
+            model=model,
+            messages=history,
+            stream=stream
+            )
+    result = ''
+    for chunk in completion:
+        if stream:
+            print(chunk['message']['content'], end='', flush=True)
+        result += chunk['message']['content']
+    if stream:
+        print()
+    history.append({"role": 'assistant', 'content': result})
+    return result
+
+def parse_args():
+    # Create the parser
+    parser = argparse.ArgumentParser(description='Copy and open a source file in TextEdit')
+    # Add the --follow-up (-f) argument
+    parser.add_argument('--follow-up', '-f', nargs='?', const=True, default=False, help='Ask a follow up question when piping in context')
+    # Add the --copy (-c) argument
+    parser.add_argument('--copy', '-c', action='store_true', help='copy a codeblock if it appears')
+    # Add the --shell (-s) argument
+    parser.add_argument('--shell', '-s', nargs='?', const=True, default=False, help='output a shell command that does as described')
+    # Add the --model (-m) argument
+    parser.add_argument('--model', '-m', nargs='?', const=True, default=False, help='Specify model')
+    # Parse the arguments
+    return parser.parse_args()
+
+def arg_follow_up(args):
+    sys.stdin = open('/dev/tty')
+    if args.follow_up != True:
+        second_input = args.follow_up
+    else:
+        second_input = input('> ')
+    return second_input
+
+def arg_shell(args):
+    query = 'Form a shell command based on the following description. Only output a working shell command and then <|eot_id|>.\nDescription: '
+    if args.shell != True:
+        query += args.shell
+    else:
+        query += input('> ')
+    result = chat(query, False)
+    result = blocks[0] if len(blocks := extract_code_block(result)) else result
+    print(result)
+    copy_string_to_clipboard(result)
+
+def handle_piped_input(args):
+    all_input = sys.stdin.read()
+    query = '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:\nSTART CONTEXT\n' + all_input + '\nEND CONTEXT\n'
+    if args.copy:
+        query += 'Answer the question using a codeblock for any code or shell scripts\n'
+    if args.follow_up:
+        query += arg_follow_up(args)
+    query += '\n'
+
+    result = chat(query)
+    blocks = extract_code_block(result)
+    if args.copy and len(blocks):
+        copy_string_to_clipboard(blocks[0])
+
+def handle_non_piped_input(args):
+    if args.shell:
+        arg_shell(args)
+        exit()
+    if args.follow_up:
+        user_input = arg_follow_up(args)
+        chat(user_input)
+        exit()
+    while True:
+        try:
+            user_input = input('> ')
+        except (EOFError, KeyboardInterrupt):
+            print()
+            exit()
+        else:
+            chat(user_input)
+
+def main():
+    args = parse_args()
+    if args.model:
+        global model 
+        model = args.model
+    if not sys.stdin.isatty():
+        handle_piped_input(args)
+    else:
+        handle_non_piped_input(args)
+
+if __name__ == '__main__':
+    main()