ai_conversation.py

import ollama
from termcolor import colored
import datetime
import tiktoken  # Used for token counting

class AIConversation:
    def __init__(
        self,
        model_1,
        model_2,
        system_prompt_1,
        system_prompt_2,
        ollama_endpoint,
        max_tokens=4000,
        limit_tokens=True
    ):
        # Initialize conversation parameters and Ollama client
        self.model_1 = model_1
        self.model_2 = model_2
        self.system_prompt_1 = system_prompt_1
        self.system_prompt_2 = system_prompt_2
        self.current_model = self.model_1
        self.messages_1 = [{"role": "system", "content": system_prompt_1}]
        self.messages_2 = [{"role": "system", "content": system_prompt_2}]
        self.client = ollama.Client(ollama_endpoint)
        self.ollama_endpoint = ollama_endpoint
        self.tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
        self.max_tokens = max_tokens
        self.limit_tokens = limit_tokens
    def count_tokens(self, messages):
        # Count the total number of tokens in the messages
        return sum(len(self.tokenizer.encode(msg["content"])) for msg in messages)

    def trim_messages(self, messages):
        # Trim messages to stay within the token limit
        if self.limit_tokens and self.count_tokens(messages) > self.max_tokens:
            print(colored(f"[SYSTEM] Max tokens reached. Sliding context window...", "magenta"))
            
            # Keep the system prompt (first message)
            system_prompt = messages[0]
            messages = messages[1:]
            
            # Remove messages from the beginning until we're under the token limit
            while self.count_tokens([system_prompt] + messages) > self.max_tokens:
                if messages:
                    messages.pop(0)  # Remove the oldest message
                else:
                    break  # Avoid removing all messages
            
            # Add the system prompt back at the beginning
            messages.insert(0, system_prompt)
        
        return messages

    def start_conversation(self, initial_message, num_exchanges=0, options=None):
        
        # Main conversation loop
        current_message = initial_message
        color_1, color_2 = "cyan", "yellow"
        conversation_log = []

        # Add initial message to system prompts
        self.messages_1[0]["content"] += f"\n\nInitial message: {current_message}"
        self.messages_2[0]["content"] += f"\n\nInitial message: {current_message}"

        print(colored(f"Starting conversation with: {current_message}", "green"))
        print(colored("Press CTRL+C to stop the conversation.", "red"))
        print()

        try:
            i = 0
            active_ai = 1  # Starting with AI 1
            while num_exchanges == 0 or i < num_exchanges:
                # Set up current AI's parameters
                name = "AI 1" if active_ai == 0 else "AI 2"
                messages = self.messages_1 if active_ai == 0 else self.messages_2
                other_messages = self.messages_2 if active_ai == 0 else self.messages_1
                color = color_1 if active_ai == 0 else color_2

                # Add user message to conversation history
                messages.append({"role": "user", "content": current_message})
                other_messages.append({"role": "assistant", "content": current_message})

                # Trim messages and get token count
                messages = self.trim_messages(messages)
                token_count = self.count_tokens(messages)
                print(colored(f"Context token count: {token_count}", "magenta"))

                # Generate AI response
                response = self.client.chat(
                    model=self.current_model,
                    messages=messages,
                    options=options,
                )
                response_content = response["message"]["content"]

                # Post-process to remove repetition
                response_content = self.remove_repetition(response_content)

                # Format and print the response with a bubble
                model_name = f"{self.current_model.upper()} ({name}):"
                formatted_response = model_name + "|:> " + response_content
                print(colored(formatted_response, color))
                conversation_log.append(
                    {"role": "assistant", "content": formatted_response}
                )

                # Update conversation history
                messages.append({"role": "assistant", "content": response_content})
                other_messages.append({"role": "user", "content": response_content})

                current_message = response_content

                # Switch to the other AI for the next turn
                self.current_model = self.model_2 if active_ai == 1 else self.model_1
                active_ai = 1 if active_ai == 0 else 0

                print(colored("---", "magenta"))
                print()

                # Check for conversation end condition
                if current_message.strip().endswith("{{end_conversation}}"):
                    print(colored("Conversation ended by the AI.", "green"))
                    break

                i += 1

        except KeyboardInterrupt:
            print(colored("\nConversation stopped by user.", "red"))

        print(colored("Conversation ended.", "green"))
        self.save_conversation_log(conversation_log)

    def save_conversation_log(self, messages, filename=None):
        # Save the conversation log to a file
        if filename is None:
            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"conversation_log_{timestamp}.txt"

        log_content = f"Conversation Log - {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
        log_content += f"Ollama Endpoint: {self.ollama_endpoint}\n"
        log_content += f"Model 1: {self.model_1}\n"
        log_content += f"Model 2: {self.model_2}\n"
        log_content += f"System Prompt 1:\n{self.system_prompt_1}\n\n"
        log_content += f"System Prompt 2:\n{self.system_prompt_2}\n\n"
        log_content += "Conversation:\n\n"

        for message in messages:
            log_content += f"{message['role'].upper()}:\n{message['content']}\n\n"

        with open(filename, "w") as f:
            f.write(log_content)

        print(f"Conversation log saved to {filename}")

    def remove_repetition(self, text):
        # Remove repeated sentences while preserving order
        split_tokens = [".", "!", "?"]
        sentences = []
        current_sentence = ""
        for char in text:
            current_sentence += char
            if char in split_tokens:
                sentences.append(current_sentence.strip())
                current_sentence = ""
        if current_sentence:  # Add any remaining text
            sentences.append(current_sentence.strip())

        # Remove duplicates while preserving order
        unique_sentences = []
        for sentence in sentences:
            if sentence not in unique_sentences:
                unique_sentences.append(sentence)

        # Join the sentences back together
        return " ".join(unique_sentences)