From dc772722d661fd0fdd162be3e3988e18cb16ee70 Mon Sep 17 00:00:00 2001
From: Hayden Johnson <hajohn100@gmail.com>
Date: Tue, 1 Oct 2024 12:50:55 -0700
Subject: [PATCH] Add /save /clear and /exit commands

The save command takes an optional filename to save the conversation to.
the clear command clears the chat history.
---
 assistant.py | 47 ++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 38 insertions(+), 9 deletions(-)

diff --git a/assistant.py b/assistant.py
index 9b47ce8..6c2835e 100755
--- a/assistant.py
+++ b/assistant.py
@@ -8,7 +8,7 @@ import argparse
 import pygments
 from pygments.lexers import get_lexer_by_name
 from pygments.formatters import TerminalFormatter
-
+import os
 
 server = 'localhost:11434'
 model = 'llama3.1:8b-instruct-q8_0'
@@ -17,16 +17,41 @@ temp = 0.2
 pattern = r'```[a-z]*\n[\s\S]*?\n```'
 line_pattern = r'`[a-z]*[\s\S]*?`'
 
+def save_conversation(filename='conversation.md'):
+    # check if filename already exists and increment filename if so
+    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
+    # save conversation to filename
+    global conversation
+    with open(filename, 'w') as f:
+        f.write(conversation)
+
 def parse_commands(text):
     # See if user wrote any commands here
+    # returns bool: True if command was executed, False if not
+    # importantly, the command doesn't need to execute succesfully for it to return True
     tokens = text.split(' ')
     match tokens[0]:
         case '/save':
-            print('saving')
+            if len(tokens) > 1:
+                save_conversation(tokens[1])
+            else:
+                save_conversation()
+            return True
         case '/clear':
-            print('clearing context')
+            global history
+            history = [ system_prompt ]
+            return True
         case '/exit':
-            print('exiting')
+            exit()
+    return False
+
 
 def highlight_code(language_name, code):
     # Check if the language is specified in the first line
@@ -102,9 +127,9 @@ def copy_string_to_clipboard(string):
 
 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."},
-        ]
+system_prompt = {"role": "system", "content": "You are a helpful, smart, kind, and efficient AI assistant. You always fulfill the user's requests accurately and concisely."}
+
+history = [ system_prompt ]
 
 conversation = ""
 
@@ -150,11 +175,15 @@ def chat(message, stream=True):
     return result
 
 def chat2(args, user_input, stream=True):
-    conversation += 'user: ' + user_input + '\n'
-    if args.reflect:
+    global conversation
+    if parse_commands(user_input):
+        result = ''
+    elif args.reflect:
         result = reflection_mode(user_input, stream)
     else:
         result = chat(user_input, stream)
+
+    conversation += 'user: ' + user_input + '\n'
     conversation += 'assistant: ' + result + '\n'
     return result