Gpt4 Commandline Chat application.

OpenAiCliChat.py

#!/usr/bin/python
import sys
import re
import openai
import pyperclip
from dotenv import load_dotenv
import json
import os
from pathlib import Path
import hashlib
import re
import os
import textwrap

from pygments import highlight
from pygments.lexers import get_lexer_by_name
from pygments.formatters import TerminalFormatter

class CLI_Color:
    # Foreground colors
    RED = '\033[31m'
    GREEN = '\033[32m'
    YELLOW = '\033[33m'
    BLUE = '\033[34m'
    MAGENTA = '\033[35m'
    CYAN = '\033[36m'
    WHITE = '\033[37m'
    # Background colors
    BG_BLACK = '\033[40m'
    BG_RED = '\033[41m'
    BG_GREEN = '\033[42m'
    BG_YELLOW = '\033[43m'
    BG_BLUE = '\033[44m'
    BG_MAGENTA = '\033[45m'
    BG_CYAN = '\033[46m'
    BG_WHITE = '\033[47m'
    # Darker background (using bright black which is often dark gray)
    BG_DARK_GRAY = '\033[100m' # You might need to adjust this if it's not dark enough
    
    BLACK_BG = '\033[40m'
    # Reset
    ENDC = '\033[0m'   
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'
    GOLD = '\033[33m'
    GREEN = '\033[32m'
    RED='\033[31m'
    GRAY='\033[90m'


# apply color to code blocks in answer
def colorize_all_code_blocks(answer):
    code_blocks = re.findall(r"```([^\n]*)\n?(.*?)```", answer, re.DOTALL)
    
    for i, block in enumerate(code_blocks, 1):
        answer = answer.replace(f"```{block[0]}\n{block[1]}```", f"```{block[0]}\n{highlight_code(block[1], block[0])}```\n")
    #return answer
    return answer


def highlight_code(code, language):
    lexer = get_lexer_by_name(language)
    highlighted_code = highlight(code, lexer, TerminalFormatter())
    return highlighted_code

# Correctly load ~/.env file
load_dotenv(os.path.expanduser("~/.env"))
# Load environment 

class OpenAIChatbot:
    is_new: bool = True
    messages: list = []
    def __init__(self,title, model_id='gpt-4',system=None,max_tokens=1500):
        self.title = title
        self.max_tokens = max_tokens
        if system is None:
            self.system = "You are a helpful assistant. Output code only in code blocks. do not say 'example code here' or psudo code. only valid bash,python."
        else:
            self.system = system
        self.api_key = os.getenv('OPENAI_API_KEY')
        if not self.api_key:
            print(CLI_Color.WARNING + "OpenAI API key not found. Please check your .env file." + CLI_Color.ENDC)
            raise ValueError("OpenAI API key not found. Please check your .env file.")
        self.model_id = model_id
        openai.api_key = self.api_key
        if self.load_context():
            self.is_new = False
        # Ensure the directory exists(system)
        else:
            self.is_new = True
            self.init_context()
    def init_context(self):
        self.messages=[
            {"role": "system", "content": self.system},
        ]
        self.query(self.title)
    def load_context(self):
        title_md5 = hashlib.md5(self.title.encode()).hexdigest()
        f = Path.home() / f"openai_context/{title_md5}.json"
        if f.exists():
            self.messages = json.loads(f.read_text())
            return True
        return False
    
    def save_context(self):
        first_message_md5 = hashlib.md5(self.title.encode()).hexdigest()
        f = Path.home() / f"openai_context/{first_message_md5}.json"
        f.parent.mkdir(parents=True, exist_ok=True)
        f.write_text(json.dumps(self.messages))
    
    def list_models(self):

        try:
            models = openai.Model.list()
            for model in models.data:
                print(f"Model ID: {model.id} - Family: {model.family}")
        except Exception as e:
            print("Failed to list models:", e)
    
   
    def extract_code_blocks(self, text):
        code_blocks = re.findall(r"```([^\n]*)\n?(.*?)```", text, re.DOTALL)
        return code_blocks
    
    def copy_to_clipboard(self, text):
        pyperclip.copy(text)
        print("Code block copied to clipboard.")
    
    def query(self, question):
        try:
            self.messages.append({"role": "user", "content": question})
            response = openai.ChatCompletion.create(
                model=self.model_id,
                messages=self.messages,
                max_tokens=self.max_tokens
            )
            self.messages.append(response.choices[0]['message'])
            return self.messages[-1]['content']
        except Exception as e:
            print("Failed to get an answer:", e)
            return None

class ChatInterface:
    def __init__(self, chatbot,system=None):
        self.chatbot = chatbot
    def input(self, prompt):
        print(CLI_Color.OKGREEN + prompt + CLI_Color.ENDC)
        return input("")
    def mlinput(self, prompt):
        # multi-line input end in >end, inform user
        print(CLI_Color.OKGREEN + prompt + CLI_Color.ENDC)
        print("Multi-line input ends in '>end'.")
        lines = []
        while True:
            line = input("")
            if line == '>end':
                break
            lines.append(line)
        return '\n'.join(lines)

    def query(self, question):
        answer =  self.chatbot.query(question)
        if answer:
            colored_answer = colorize_all_code_blocks(answer)
            print(f"{CLI_Color.OKGREEN}Chatbot:{CLI_Color.ENDC}\n{colored_answer}")
            #highlight code blocks
            
            code_blocks = self.chatbot.extract_code_blocks(answer)
            if code_blocks:
                for i, block in enumerate(code_blocks, 1):
                    print(f"\n[{i}]")
                    print(highlight_code(block[1], block[0]))
                choice = input("Enter the number of the code block to copy to clipboard (or press Enter to skip): ")
                if choice.isdigit():
                    choice = int(choice) - 1
                    if 0 <= choice < len(code_blocks):
                        self.chatbot.copy_to_clipboard(code_blocks[choice])[1]
                    else:
                        print(f"{CLI_Color.RED}Invalid selection.{CLI_Color.ENDC}")
                elif choice:\
                    print(CLI_Color.WARNING + "Invalid input. Skipping clipboard copy." + CLI_Color.ENDC)

            else:
                print(f"{CLI_Color.GRAY}No code blocks found.{CLI_Color.ENDC}")
        print("#" * 50)

    def run(self):
        print("Type 'exit' or 'quit' to end the chat.")
        while True:
            question = self.mlinput("You: ")
            if question.lower() in ['exit', 'quit']:
                print("Exiting chat.")
                break
            if question.strip() == ">save":
                self.chatbot.save_context()
                print("Chat context saved.")
                continue
            if question.strip() == ">branch":
                self.chatbot.title = input("Enter the new title for branch: ")
                self.chatbot.is_new = True
                self.chatbot.save_context()
                print("Chat context branched.")
                continue
            print(f"{CLI_Color.GOLD}Chatbot is thinking...{CLI_Color.ENDC}")
            self.query(question)
    def print_history(self):
        print("Chat history:")
        for message in self.chatbot.messages:
            role = message['role']
            content = message['content']
            content = colorize_all_code_blocks(content)
            print(f"{CLI_Color.OKGREEN}{role.capitalize()}:{CLI_Color.ENDC}\n {content}" )
            print("#" * 50)
        print("End of chat history.")

if __name__ == "__main__":
    try:

        if len(sys.argv) > 1 and sys.argv[1]:
            title = sys.argv[1]
        chatbot = OpenAIChatbot(title=title) 
        chat_interface = ChatInterface(chatbot)
        chat_interface.print_history()
        chat_interface.run()
    except ValueError as e:
        print(e)