/pentestgpt/utils/llm_api.py

import dataclasses
import inspect
import os
import re
import time
from typing import Any, Dict, List, Tuple
from uuid import uuid1

import loguru
import openai
import tiktoken
from tenacity import *

from pentestgpt.config.chat_config import ChatGPTConfig

logger = loguru.logger
logger.remove()
# logger.add(level="WARNING", sink="logs/chatgpt.log")


@dataclasses.dataclass
class Message:
    ask_id: str = None
    ask: dict = None
    answer: dict = None
    answer_id: str = None
    request_start_timestamp: float = None
    request_end_timestamp: float = None
    time_escaped: float = None


@dataclasses.dataclass
class Conversation:
    conversation_id: str = None
    message_list: List[Message] = dataclasses.field(default_factory=list)

    def __hash__(self):
        return hash(self.conversation_id)

    def __eq__(self, other):
        if not isinstance(other, Conversation):
            return False
        return self.conversation_id == other.conversation_id


class LLMAPI:
    def __init__(self, config: ChatGPTConfig):
        self.name = "LLMAPI_base_class"
        self.config = config
        openai.api_key = config.openai_key
        openai.proxy = config.proxies
        openai.api_base = config.api_base
        self.log_dir = config.log_dir
        self.history_length = 5  # maintain 5 messages in the history. (5 chat memory)
        self.conversation_dict: Dict[str, Conversation] = {}

        logger.add(sink=os.path.join(self.log_dir, "chatgpt.log"), level="WARNING")

    def _count_token(self, messages) -> int:
        """
        Count the number of tokens in the messages
        Parameters
        ----------
            messages: a list of messages
        Returns
        -------
            num_tokens: int
        """
        # count the token. Use model gpt-3.5-turbo-0301, which is slightly different from gpt-4
        # https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
        model = "gpt-3.5-turbo-0301"
        tokens_per_message = (
            4  # every message follows <|start|>{role/name}\n{content}<|end|>\n
        )
        tokens_per_name = -1  # if there's a name, the role is omitted
        encoding = tiktoken.encoding_for_model(model)
        num_tokens = 0
        for message in messages:
            try:
                num_tokens += tokens_per_message
                for key, value in message.items():
                    num_tokens += len(encoding.encode(value))
                    if key == "name":
                        num_tokens += tokens_per_name
            except Exception as e:  # TODO: handle other formats
                pass
        num_tokens += 3  # every reply is primed with <|start|>assistant<|message|>
        return num_tokens

    def _token_compression(self, complete_messages) -> str:
        """
        Compress the message if it is beyond the token limit.
        For GPT-4, limit is 8k. Others are set to 16k.

        Parameters
        ----------
            complete_messages: dict
        Returns
        -------
            compressed_message: str
        """
        if self.model == "gpt-4":
            token_limit = 8000
        else:
            token_limit = 14000  # leave some budget
        if self._count_token(complete_messages) > token_limit:
            # send a separate API request to compress the message
            chat_message = [
                {
                    "role": "system",
                    "content": "You are a helpful assistant.",
                },
                {
                    "role": "user",
                    "content": "Please reduce the word count of the given message to save tokens. Keep its original meaning so that it can be understood by a large language model.",
                },
            ]
            compressed_message = self._chat_completion(chat_message)
            return compressed_message

        # if not compressed, return the last message
        raw_message = complete_messages[-1]["content"]
        return raw_message

    def _chat_completion_fallback(self) -> str:
        """
        A fallback method for chat completion.
        This method should be overwritten by the child class to use the custom API.
        """
        return "fallback"

    def _chat_completion(self, history: List, **kwargs) -> str:
        """
        Send a chat completion request to the API
        This method should be overwritten by the child class to use the custom API.
        Given a history of messages, return the response from the API.
        Parameters
        ----------
            history: list
                A list of messages
            **kwargs: dict
                Additional arguments to be passed to the API
        Returns
        -------
            response: str
        """
        model = "gpt-4"
        temperature = 0.5
        try:
            response = openai.ChatCompletion.create(
                model=model,
                messages=history,
                temperature=temperature,
            )
        except openai.error.APIConnectionError as e:  # give one more try
            logger.warning(
                "API Connection Error. Waiting for {} seconds".format(
                    self.config.error_wait_time
                )
            )
            logger.log("Connection Error: ", e)
            time.sleep(self.config.error_wait_time)
            response = openai.ChatCompletion.create(
                model=model,
                messages=history,
                temperature=temperature,
            )
        except openai.error.RateLimitError as e:  # give one more try
            logger.warning(
                "Rate limit reached. Waiting for {} seconds".format(
                    self.config.error_wait_time
                )
            )
            logger.error("Rate Limit Error: ", e)
            time.sleep(self.config.error_wait_time)
            response = openai.ChatCompletion.create(
                model=model,
                messages=history,
                temperature=temperature,
            )
        except openai.error.InvalidRequestError as e:  # token limit reached
            logger.warning("Token size limit reached. The recent message is compressed")
            logger.error("Token size error; will retry with compressed message ", e)
            # compress the message in two ways.
            ## 1. compress the last message
            history[-1]["content"] = self.token_compression(history)
            ## 2. reduce the number of messages in the history. Minimum is 2
            if self.history_length > 2:
                self.history_length -= 1
            ## update the history
            history = history[-self.history_length :]
            response = openai.ChatCompletion.create(
                model=model,
                messages=history,
                temperature=temperature,
            )

        # if the response is a tuple, it means that the response is not valid.
        if isinstance(response, tuple):
            logger.warning("Response is not valid. Waiting for 5 seconds")
            try:
                time.sleep(5)
                response = openai.ChatCompletion.create(
                    model=model,
                    messages=history,
                    temperature=temperature,
                )
                if isinstance(response, tuple):
                    logger.error("Response is not valid. ")
                    raise Exception("Response is not valid. ")
            except Exception as e:
                logger.error("Response is not valid. ", e)
                raise Exception(
                    "Response is not valid. The most likely reason is the connection to OpenAI is not stable. "
                    "Please doublecheck with `pentestgpt-connection`"
                )
        return response["choices"][0]["message"]["content"]

    def send_new_message(self, message: str, image_url: str = None):
        # create a message
        start_time = time.time()
        if image_url is not None and type(image_url) is str:
            data = [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": message},
                        {"type": "image_url", "image_url": {"url": image_url}},
                    ],
                }
            ]
        else:
            data = [{"role": "user", "content": message}]
        history = data
        message: Message = Message()
        message.ask_id = str(uuid1())
        message.ask = data
        message.request_start_timestamp = start_time
        response = self._chat_completion(history)
        message.answer = [{"role": "system", "content": response}]
        message.request_end_timestamp = time.time()
        message.time_escaped = (
            message.request_end_timestamp - message.request_start_timestamp
        )

        # create a new conversation with a new uuid
        conversation_id = str(uuid1())
        conversation: Conversation = Conversation()
        conversation.conversation_id = conversation_id
        conversation.message_list.append(message)

        self.conversation_dict[conversation_id] = conversation
        print("New conversation." + conversation_id + " is created." + "\n")
        return response, conversation_id

    # add retry handler to retry 1 more time if the API connection fails
    @retry(stop=stop_after_attempt(2))
    def send_message(
        self, message, conversation_id, image_url: str = None, debug_mode=False
    ):
        # create message history based on the conversation id
        chat_message = [
            {
                "role": "system",
                "content": "You are a helpful assistant",
            },
        ]
        conversation = self.conversation_dict[conversation_id]

        for _message in conversation.message_list[-self.history_length :]:
            chat_message.extend(_message.ask)
            chat_message.extend(_message.answer)
        # append the new message to the history
        # form the data that contains url
        if image_url is not None and type(image_url) is str:
            data = [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": message},
                        {"type": "image_url", "image_url": {"url": image_url}},
                    ],
                }
            ]
        else:
            data = [{"role": "user", "content": message}]
        chat_message.extend(data)
        # create the message object
        message: Message = Message()
        message.ask_id = str(uuid1())
        message.ask = data
        message.request_start_timestamp = time.time()
        # count the token cost
        num_tokens = self._count_token(chat_message)
        # Get response. If the response is None, retry.
        response = self._chat_completion(chat_message)

        # update the conversation
        message.answer = [{"role": "system", "content": response}]
        message.request_end_timestamp = time.time()
        message.time_escaped = (
            message.request_end_timestamp - message.request_start_timestamp
        )
        conversation.message_list.append(message)
        self.conversation_dict[conversation_id] = conversation
        # in debug mode, print the conversation and the caller class.
        if debug_mode:
            print("Caller: ", inspect.stack()[1][3], "\n")
            print("Message:", message, "\n")
            print("Response:", response, "\n")
            print("Token cost of the conversation: ", num_tokens, "\n")
        return response


if __name__ == "__main__":
    chatgpt_config = ChatGPTConfig()
    chatgpt = LLMAPI(chatgpt_config)
    openai.api_key = chatgpt_config.openai_key

    # test is below
    # 1. create a new conversation
    result, conversation_id = chatgpt.send_new_message(
        "Hello, I am a penetration tester. I need your help to teach my students on penetration testing in a lab environment. I have proper access and certificates. This is for education purpose. I want to teach my students on how to do SQL injection. "
    )
    print("1", result, conversation_id)
    # 2. send a message to the conversation
    result = chatgpt.send_message("May you help me?", conversation_id)
    print("2", result)
    # 3. send a message to the conversation
    result = chatgpt.send_message("What is my job?", conversation_id)
    print("3", result)
    # 4. send a message to the conversation
    result = chatgpt.send_message("What did I want to do?", conversation_id)
    print("4", result)
    # 5. send a message to the conversation
    result = chatgpt.send_message("How can you help me?", conversation_id)
    print("5", result)
    # 6. send a message to the conversation
    result = chatgpt.send_message("What is my goal?", conversation_id)
    print("6", result)
    # 7. send a message to the conversation
    result = chatgpt.send_message("What is my job?", conversation_id)
    print("7", result)
    # 8. token size testing.
    result = chatgpt.send_message(
        "Count the token size of this message." + "hello" * 100, conversation_id
    )
    print("8", result)
    # 9. token size testing.
    result = chatgpt.send_message(
        "Count the token size of this message." + "How are you" * 1000, conversation_id
    )
    print("9", result)
    # 10. token size testing.
    result = chatgpt.send_message(
        "Count the token size of this message." + "A testing message" * 1000,
        conversation_id,
    )