litellm-chat-bot-v0.1

from litellm import completion
from datetime import datetime
import os
import glob
# importlib.reload(chatbot); from chatbot import Chatbot; hacs = Chatbot()

class Chatbot:
    def __init__(self, model="ollama/llama3.2:latest", api_base="http://localhost:11434"):
        self.model = model
        self.api_base = api_base
        self.context = []
        self.saved_contexts = {}  # Dictionary to hold named contexts
        print ('hacs2 loaded')


    def set_model(self, new_model):
        self.model = new_model
        print(f"Model set to: {self.model}")

    def set_context(self, context_string):
        self.context = [{"role": "user", "content": context_string}]
        print("Context updated.")

    def add_context(self, context_string):
        """Append a new context entry to the main context."""
        self.context.append({"role": "user", "content": context_string})
        print("Context added.")

    def save_context(self, context_name):
        """Save the current context under a specified name."""
        self.saved_contexts[context_name] = self.context.copy()
        print(f"Context saved as '{context_name}'.")

    def get_context(self, context_name):
        """Retrieve and display a saved context."""
        if context_name in self.saved_contexts:
            print(f"Context '{context_name}':")
            for message in self.saved_contexts[context_name]:
                role = message["role"].capitalize()
                content = message["content"]
                print(f"{role}: {content}")
        else:
            print(f"No context found with the name '{context_name}'.")

    def clear_context(self):
        """Clear the main context."""
        self.context = []
        print("Context cleared.")


    def stream_response(self, response):
        for chunk in response:
            chunk_content = chunk['choices'][0]['delta'].get('content', '')
            if chunk_content:
                yield chunk_content

    def send(self, message):
        self.context.append({"role": "user", "content": message})
        response = completion(
            model=self.model,
            messages=self.context,
            api_base=self.api_base,
            stream=True
        )

        ai_response = ""
        print("AI: ", end='', flush=True)
        for chunk in self.stream_response(response):
            print(chunk, end='', flush=True)
            ai_response += chunk
        print()  # Move to the next line after finishing

        self.context.append({"role": "assistant", "content": ai_response})
        return ai_response

    def save_chat(self):
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"chat_history_{timestamp}.txt"
        with open(filename, "w") as file:
            for message in self.context:
                role = message["role"].capitalize()
                content = message["content"]
                file.write(f"{role}: {content}\n\n")
        print(f"Chat history saved to {filename}")

    def load_chat(self):
        chat_files = glob.glob("chat_history_*.txt")
        if not chat_files:
            print("No saved chat history found.")
            return None

        latest_file = max(chat_files, key=os.path.getctime)
        context = []
        with open(latest_file, "r") as file:
            current_role = None
            current_content = ""
            for line in file:
                line = line.strip()
                if line.startswith("User:") or line.startswith("Assistant:"):
                    if current_role:
                        context.append({"role": current_role.lower(), "content": current_content.strip()})
                    current_role = line.split(":")[0]
                    current_content = line.split(":", 1)[1].strip()
                elif line:
                    current_content += " " + line
            if current_role:
                context.append({"role": current_role.lower(), "content": current_content.strip()})

        print(f"Loaded chat history from {latest_file}")
        self.context = context  # Update the context of the chatbot

def chat_cli():
    chatbot = Chatbot()
    print("Welcome to the CLI Chat with Memory!")
    print("Commands: 'exit' to end, 'save chat' to save, 'load chat' to load previous chat, 'set model [model_name]' to change model.")

    while True:
        user_input = input("You: ").strip()
        if user_input.lower() == 'exit':
            print("Goodbye!")
            break
        elif user_input.lower() == 'save chat':
            chatbot.save_chat()
            continue
        elif user_input.lower() == 'load chat':
            chatbot.load_chat()
            continue
        elif user_input.lower().startswith('set model'):
            _, new_model = user_input.split(maxsplit=1)
            chatbot.set_model(new_model)
            continue

        chatbot.send(user_input)

if __name__ == "__main__":
    chat_cli()