chat GPT api

chatgpt_api.py

import contextlib
import json
import logging
import time
import uuid
from abc import ABCMeta
from abc import abstractmethod
from functools import wraps
from os import environ
from os import getenv

import requests

BASE_URL = environ.get("CHATGPT_BASE_URL") or "https://bypass.churchless.tech/api/"

log = logging.getLogger(__name__)


class colors:
    """
    Colors for printing
    """

    HEADER = "\033[95m"
    OKBLUE = "\033[94m"
    OKCYAN = "\033[96m"
    OKGREEN = "\033[92m"
    WARNING = "\033[93m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"
    BOLD = "\033[1m"
    UNDERLINE = "\033[4m"

    def __init__(self) -> None:
        if getenv("NO_COLOR"):
            self.HEADER = ""
            self.OKBLUE = ""
            self.OKCYAN = ""
            self.OKGREEN = ""
            self.WARNING = ""
            self.FAIL = ""
            self.ENDC = ""
            self.BOLD = ""
            self.UNDERLINE = ""


def logger(is_timed: bool):
    """Logger decorator

    Args:
        is_timed (bool): Whether to include function running time in exit log

    Returns:
        _type_: decorated function
    """

    def decorator(func):
        wraps(func)

        def wrapper(*args, **kwargs):
            log.debug(
                "Entering %s with args %s and kwargs %s",
                func.__name__,
                args,
                kwargs,
            )
            start = time.time()
            out = func(*args, **kwargs)
            end = time.time()
            if is_timed:
                log.debug(
                    "Exiting %s with return value %s. Took %s seconds.",
                    func.__name__,
                    out,
                    end - start,
                )
            else:
                log.debug("Exiting %s with return value %s", func.__name__, out)

            return out

        return wrapper

    return decorator


bcolors = colors()


class ChatbotError(Exception, metaclass=ABCMeta):
    """
    Base class for all ChatGPT errors in this Project
    """

    @abstractmethod
    def __init__(self, *args: object) -> None:
        super().__init__(*args)


class MetaNotAllowInstance(type):
    """
    Metaclasses that do not allow classes to be instantiated
    """

    def __call__(self, *args, **kwds):
        error = RuntimeError("This class is not allowed to be instantiated")
        raise error


class Error(ChatbotError):
    """
    Base class for exceptions in V1 module.
    Error codes:
    -1: User error
    0: Unknown error
    1: Server error
    2: Rate limit error
    3: Invalid request error
    4: Expired access token error
    5: Invalid access token error
    6: Prohibited concurrent query error
    """

    def __init__(self, source: str, message: str, code=0, *args) -> None:
        self.source: str = source
        self.message: str = message
        self.code = code
        super().__init__(*args)

    def __str__(self) -> str:
        return f"{self.source}: {self.message} (code: {self.code})"

    def __repr__(self) -> str:
        return f"{self.source}: {self.message} (code: {self.code})"


class ErrorType(metaclass=MetaNotAllowInstance):
    # define consts for the error codes
    USER_ERROR = -1
    UNKNOWN_ERROR = 0
    SERVER_ERROR = 1
    RATE_LIMIT_ERROR = 2
    INVALID_REQUEST_ERROR = 3
    EXPIRED_ACCESS_TOKEN_ERROR = 4
    INVALID_ACCESS_TOKEN_ERROR = 5
    PROHIBITED_CONCURRENT_QUERY_ERROR = 6
    AUTHENTICATION_ERROR = 7
    CLOUDFLARE_ERROR = 8


class AuthenticationError(ChatbotError):
    """
    Subclass of ChatbotError

    The object of the error thrown by a validation failure or exception
    """

    def __init__(self, *args: object) -> None:
        super().__init__(*args)


class ChatGPT:
    """
    ChatGPT class for ChatGPT
    """

    @logger(is_timed=True)
    def __init__(
            self,
            config,
            conversation_id=None,
            parent_id=None,
            session_client=None,
            lazy_loading: bool = True,
    ) -> None:
        """Initialize a chatbot

        Args:
            config (dict[str, str]): Login and proxy info. Example:
                {
                    "email": "OpenAI account email",
                    "password": "OpenAI account password",
                    "session_token": "<session_token>"
                    "access_token": "<access_token>"
                    "proxy": "<proxy_url_string>",
                    "paid": True/False, # whether this is a plus account
                }
                More details on these are available at https://github.com/acheong08/ChatGPT#configuration
            conversation_id (str | None, optional): Id of the conversation to continue on. Defaults to None.
            parent_id (str | None, optional): Id of the previous response message to continue on. Defaults to None.
            session_client (_type_, optional): _description_. Defaults to None.

        Raises:
            Exception: _description_
        """

        self.config = config
        self.session = session_client() if session_client else requests.Session()
        if "proxy" in config:
            if not isinstance(config["proxy"], str):
                error = TypeError("Proxy must be a string!")
                raise error
            proxies = {
                "http": config["proxy"],
                "https": config["proxy"],
            }
            self.session.proxies.update(proxies)
        self.conversation_id = conversation_id
        self.parent_id = parent_id
        self.conversation_mapping = {}
        self.conversation_id_prev_queue = []
        self.parent_id_prev_queue = []
        self.lazy_loading = lazy_loading

        self.__check_credentials()

    @logger(is_timed=True)
    def __check_credentials(self) -> None:
        """Check login info and perform login

        Any one of the following is sufficient for login. Multiple login info can be provided at the same time and they will be used in the order listed below.
            - access_token
            - session_token
            - email + password

        Raises:
            Exception: _description_
            AuthError: _description_
        """
        if "access_token" in self.config:
            self.set_access_token(self.config["access_token"])
        else:
            raise AuthenticationError("Insufficient login details provided!")

    @logger(is_timed=False)
    def set_access_token(self, access_token: str) -> None:
        """Set access token in request header and self.config, then cache it to file.

        Args:
            access_token (str): access_token
        """
        self.session.headers.clear()
        self.session.headers.update(
            {
                "Accept": "text/event-stream",
                "Authorization": f"Bearer {access_token}",
                "Content-Type": "application/json",
                "X-Openai-Assistant-App-Id": "",
                "Connection": "close",
                "Accept-Language": "en-US,en;q=0.9",
                "Referer": "https://chat.openai.com/chat",
            },
        )

        self.config["access_token"] = access_token

    @logger(is_timed=True)
    def ask(
            self,
            prompt: str,
            conversation_id=None,
            parent_id=None,
            timeout: float = 360,
    ) -> str:
        """Ask a question to the chatbot
        Args:
            prompt (str): The question
            conversation_id (str | None, optional): UUID for the conversation to continue on. Defaults to None.
            parent_id (str | None, optional): UUID for the message to continue on. Defaults to None.
            timeout (float, optional): Timeout for getting the full response, unit is second. Defaults to 360.

        Raises:
            Error: _description_
            Exception: _description_
            Error: _description_
            Error: _description_
            Error: _description_

        Yields:
            _type_: _description_
        """

        if parent_id is not None and conversation_id is None:
            log.error("conversation_id must be set once parent_id is set")
            error = Error(
                source="User",
                message="conversation_id must be set once parent_id is set",
                code=ErrorType.USER_ERROR,
            )
            raise error

        if conversation_id is not None and conversation_id != self.conversation_id:
            log.debug("Updating to new conversation by setting parent_id to None")
            self.parent_id = None

        conversation_id = conversation_id or self.conversation_id
        parent_id = parent_id or self.parent_id
        if conversation_id is None and parent_id is None:
            parent_id = str(uuid.uuid4())
            log.debug("New conversation, setting parent_id to new UUID4: %s", parent_id)

        if conversation_id is not None and parent_id is None:
            if conversation_id not in self.conversation_mapping:
                if self.lazy_loading:
                    log.debug(
                        "Conversation ID %s not found in conversation mapping, try to get conversation history for the given ID",
                        conversation_id,
                    )
                    with contextlib.suppress(Exception):
                        history = self.get_msg_history(conversation_id)
                        self.conversation_mapping[conversation_id] = history[
                            "current_node"
                        ]
                else:
                    log.debug(
                        "Conversation ID %s not found in conversation mapping, mapping conversations",
                        conversation_id,
                    )

                    self.__map_conversations()

            if conversation_id in self.conversation_mapping:
                log.debug(
                    "Conversation ID %s found in conversation mapping, setting parent_id to %s",
                    conversation_id,
                    self.conversation_mapping[conversation_id],
                )
                parent_id = self.conversation_mapping[conversation_id]
            else:  # invalid conversation_id provided, treat as a new conversation
                conversation_id = None
                parent_id = str(uuid.uuid4())
        data = {
            "action": "next",
            "messages": [
                {
                    "id": str(uuid.uuid4()),
                    "role": "user",
                    "author": {"role": "user"},
                    "content": {"content_type": "text", "parts": [prompt]},
                },
            ],
            "conversation_id": conversation_id,
            "parent_message_id": parent_id,
            "model": self.config.get("model")
                     or (
                         "text-davinci-002-render-paid"
                         if self.config.get("paid")
                         else "text-davinci-002-render-sha"
                     ),
        }
        log.debug("Sending the payload")
        log.debug(json.dumps(data, indent=2))

        self.conversation_id_prev_queue.append(
            data["conversation_id"],
        )
        self.parent_id_prev_queue.append(data["parent_message_id"])
        response = self.session.post(
            url=f"{BASE_URL}conversation",
            data=json.dumps(data),
            timeout=timeout,
            stream=True,
        )
        self.__check_response(response)
        done: bool = False
        for line in response.iter_lines():
            # remove b' and ' at the beginning and end and ignore case
            line = str(line)[2:-1]
            if line.lower() == "internal server error":
                log.error("Internal Server Error: %s", line)
                error = Error(
                    source="ask",
                    message="Internal Server Error",
                    code=ErrorType.SERVER_ERROR,
                )
                raise error
            if not line or line is None:
                continue
            if "data: " in line:
                line = line[6:]
            if line == "[DONE]":
                done = True
                break

            line = line.replace('\\"', '"')
            line = line.replace("\\'", "'")
            line = line.replace("\\\\", "\\")

            try:
                line = json.loads(line)
            except json.decoder.JSONDecodeError:
                continue
            if not self.__check_fields(line) or response.status_code != 200:
                log.error("Field missing", exc_info=True)
                log.error(response.text)
                if response.status_code == 401:
                    error = Error(
                        source="ask",
                        message="Permission denied",
                        code=ErrorType.AUTHENTICATION_ERROR,
                    )
                elif response.status_code == 403:
                    error = Error(
                        source="ask",
                        message="Cloudflare triggered a 403 error",
                        code=ErrorType.CLOUDFLARE_ERROR,
                    )
                elif response.status_code == 429:
                    error = Error(
                        source="ask",
                        message="Rate limit exceeded",
                        code=ErrorType.RATE_LIMIT_ERROR,
                    )
                else:
                    error = Error(
                        source="ask",
                        message=line,
                        code=ErrorType.SERVER_ERROR,
                    )
                raise error
            message: str = line["message"]["content"]["parts"][0]
            if message == prompt:
                continue
            conversation_id = line["conversation_id"]
            parent_id = line["message"]["id"]
            try:
                model = line["message"]["metadata"]["model_slug"]
            except KeyError:
                model = None
            log.debug("Received message: %s", message)
            log.debug("Received conversation_id: %s", conversation_id)
            log.debug("Received parent_id: %s", parent_id)
            yield {
                "message": message.strip("\n"),
                "conversation_id": conversation_id,
                "parent_id": parent_id,
                "model": model,
            }
        if not done:
            pass
        self.conversation_mapping[conversation_id] = parent_id
        if parent_id is not None:
            self.parent_id = parent_id
        if conversation_id is not None:
            self.conversation_id = conversation_id

    @logger(is_timed=False)
    def __check_fields(self, data: dict) -> bool:
        try:
            data["message"]["content"]
        except (TypeError, KeyError):
            return False
        return True

    @logger(is_timed=False)
    def __check_response(self, response: requests.Response) -> None:
        """Make sure response is success

        Args:
            response (_type_): _description_

        Raises:
            Error: _description_
        """
        if response.status_code != 200:
            print(response.text)
            error = Error(
                source="OpenAI",
                message=response.text,
                code=response.status_code,
            )
            raise error

    @logger(is_timed=True)
    def get_conversations(
            self,
            offset=0,
            limit=20,
            encoding=None,
    ) -> list:
        """
        Get conversations
        :param offset: Integer
        :param limit: Integer
        """
        url = f"{BASE_URL}conversations?offset={offset}&limit={limit}"
        response = self.session.get(url)
        self.__check_response(response)
        if encoding is not None:
            response.encoding = encoding
        data = json.loads(response.text)
        return data["items"]

    @logger(is_timed=True)
    def get_msg_history(self, convo_id: str, encoding=None) -> list:
        """
        Get message history
        :param id: UUID of conversation
        :param encoding: String
        """
        url = f"{BASE_URL}conversation/{convo_id}"
        response = self.session.get(url)
        self.__check_response(response)
        if encoding is not None:
            response.encoding = encoding
        return json.loads(response.text)

    @logger(is_timed=True)
    def gen_title(self, convo_id: str, message_id: str) -> str:
        """
        Generate title for conversation
        """
        response = self.session.post(
            f"{BASE_URL}conversation/gen_title/{convo_id}",
            data=json.dumps(
                {"message_id": message_id, "model": "text-davinci-002-render"},
            ),
        )
        self.__check_response(response)
        return response.json().get("title", "Error generating title")

    @logger(is_timed=True)
    def change_title(self, convo_id: str, title: str) -> None:
        """
        Change title of conversation
        :param id: UUID of conversation
        :param title: String
        """
        url = f"{BASE_URL}conversation/{convo_id}"
        response = self.session.patch(url, data=json.dumps({"title": title}))
        self.__check_response(response)

    @logger(is_timed=True)
    def delete_conversation(self, convo_id: str) -> None:
        """
        Delete conversation
        :param id: UUID of conversation
        """
        url = f"{BASE_URL}conversation/{convo_id}"
        response = self.session.patch(url, data='{"is_visible": false}')
        self.__check_response(response)

    @logger(is_timed=True)
    def clear_conversations(self) -> None:
        """
        Delete all conversations
        """
        url = f"{BASE_URL}conversations"
        response = self.session.patch(url, data='{"is_visible": false}')
        self.__check_response(response)

    @logger(is_timed=False)
    def __map_conversations(self) -> None:
        conversations = self.get_conversations()
        histories = [self.get_msg_history(x["id"]) for x in conversations]
        for x, y in zip(conversations, histories):
            self.conversation_mapping[x["id"]] = y["current_node"]

    @logger(is_timed=False)
    def reset_chat(self) -> None:
        """
        Reset the conversation ID and parent ID.

        :return: None
        """
        self.conversation_id = None
        self.parent_id = str(uuid.uuid4())

    @logger(is_timed=False)
    def rollback_conversation(self, num=1) -> None:
        """
        Rollback the conversation.
        :param num: Integer. The number of messages to rollback
        :return: None
        """
        for _ in range(num):
            self.conversation_id = self.conversation_id_prev_queue.pop()
            self.parent_id = self.parent_id_prev_queue.pop()


@logger(is_timed=False)
def main(config: dict):
    """
    Main function for the chatGPT program.
    """
    chatbot = ChatGPT(
        config,
        conversation_id=config.get("conversation_id"),
        parent_id=config.get("parent_id"),
    )

    print()
    try:
        while True:
            print(f"{bcolors.OKBLUE + bcolors.BOLD}You: {bcolors.ENDC}")

            prompt = input("> ")

            print()
            print(f"{bcolors.OKGREEN + bcolors.BOLD}ChatGPT: {bcolors.ENDC}")
            prev_text = ""
            for data in chatbot.ask(prompt):
                message = data["message"][len(prev_text):]
                print(message, end="", flush=True)
                prev_text = data["message"]
            print(bcolors.ENDC)
            print()
    except (KeyboardInterrupt, EOFError):
        exit()


if __name__ == "__main__":
    logging.basicConfig(
        format="%(asctime)s - %(name)s - %(levelname)s - %(funcName)s - %(message)s",
    )

    print(
        """
        ChatGPT - A command-line interface to OpenAI's ChatGPT (https://chat.openai.com/chat)
        """,
    )
    # first login to chatgpt, then go to https://chat.openai.com/api/auth/session
    # copy the access_token, valid for about 8 hours
    cfg = {
        "access_token": "eyJhbGc....."}

    main(cfg)