/pentestgpt/utils/chat_utils/chatbot.py

import os
import traceback
import uuid
from typing import Any, Dict, List

import tiktoken
from litellm import completion

from pentestgpt.utils.chat_utils.message import (
    AssistantMessage,
    Message,
    SystemMessage,
    UserMessage,
)


def process_message(message: Message) -> Dict:
    MAX_MESSAGE_LENGTH = 1048576
    MAX_TOKEN_COUNT = 120000

    raw_content = message.content[:MAX_MESSAGE_LENGTH]

    # If the message is too long, truncate it
    tokenizer = tiktoken.get_encoding("o200k_base")
    token_count = len(tokenizer.encode(raw_content))
    if token_count > MAX_TOKEN_COUNT:
        raw_content = tokenizer.decode(tokenizer.encode(raw_content)[:MAX_TOKEN_COUNT])

    if isinstance(message, UserMessage):
        user_dict: Dict = {"role": "user", "content": raw_content}
        return user_dict
    elif isinstance(message, SystemMessage):
        system_dict: Dict = {"role": "system", "content": raw_content}
        return system_dict
    elif isinstance(message, AssistantMessage):
        assistant_dict: Dict = {"role": "assistant", "content": raw_content}
        return assistant_dict
    else:
        raise Exception("Unknown message type")


class LLMLiteChatbot:
    def __init__(self, chatbot_name: str, model_name: str):
        self.chatbot_id = str(uuid.uuid4())
        self.chatbot_name = chatbot_name
        self.model_name = model_name

    def chat(self, messages: List[Message]) -> AssistantMessage:
        message_list = []
        for message in messages:
            message_list.append(process_message(message))
        response = completion(model=self.model_name, messages=message_list)
        # convert response to AssistantMessage
        return AssistantMessage(response.choices[0].message.content)


if __name__ == "__main__":
    chatbot = LLMLiteChatbot("test", "gpt-4o", "you're a kind asssistant")
    model_response = chatbot.chat([UserMessage("Hello")])
    print(model_response.get_content())