chat.py

from abc import ABC, abstractmethod
from utils import extract_all_blocks

import time
import os

from openai import OpenAI

# Set OPENAI_API_KEY and OPENAI_BASE_URL environment variables before running.
# Example:
#   export OPENAI_API_KEY="sk-..."
#   export OPENAI_BASE_URL="https://api.openai.com/v1"
class BaseChat(ABC):
    def __init__(self, model: str, temperature: float = 1.0):
        self.model = model
        self.temperature = float(temperature)
        self.messages = []

    @abstractmethod
    def get_response(self, prompt, system_prompt=None, return_token_usage=False) -> str:
        pass

    def get_model_response(self, prompt, code_format=None, tag="unknown", return_token_usage=False, system_prompt=None) -> list:
        code_blocks = []
        cur_token_usage = -1
        max_try = 3
        latest_response = None
        msg = ""
        while code_blocks == [] and max_try > 0:
            max_try -= 1
            try:
                if return_token_usage:
                    response, cur_token_usage = self.get_response(prompt, system_prompt=system_prompt, return_token_usage=return_token_usage)
                else:
                    response = self.get_response(prompt, system_prompt=system_prompt)
                latest_response = response
                if response is None:
                    raise Exception(f"Tag: {tag}, API returned None response")
            except Exception as e:
                response = getattr(e, "response", None)
                handled_rate_limit = False

                if response is not None and getattr(response, "status_code", 0) == 429:
                    try:
                        error_data = response.json()
                        inner_error = error_data.get("error", {})
                        if inner_error.get("type") == "rate_limit_exceeded":
                            wait_seconds = inner_error.get("retry_after", 0)
                            if wait_seconds > 0:
                                print(f"Tag: {tag}. Rate limit exceeded. Waiting {wait_seconds} seconds...")
                                time.sleep(wait_seconds)
                                handled_rate_limit = True
                    except Exception as parse_err:
                        print(f"[WARNING] Tag: {tag}, failed to parse 429 body: {parse_err}")

                if not handled_rate_limit:
                    print(f"Tag: {tag}, max_try in get_model_response: {max_try}, exception: {e}")
                    msg += f"{e}\n"

            code_blocks = extract_all_blocks(response, code_format)

        if code_blocks == []:
            error_msg = f"[ERROR] Tag: {tag}, get_model_response() failed after all retries. Latest response: {latest_response}"
            print(error_msg)
            raise Exception(error_msg)
        if return_token_usage:
            return code_blocks, cur_token_usage
        else:
            return code_blocks

    def get_model_response_txt(self, prompt, system_prompt=None, tag="unknown", return_token_usage=False) -> str:
        max_try = 3
        msg = ""
        while max_try > 0:
            max_try -= 1
            try:
                response = self.get_response(prompt, system_prompt=system_prompt, return_token_usage=return_token_usage)
                return response
            except Exception as e:
                response = getattr(e, "response", None)
                handled_rate_limit = False

                if response is not None and getattr(response, "status_code", 0) == 429:
                    try:
                        error_data = response.json()
                        inner_error = error_data.get("error", {})
                        if inner_error.get("type") == "rate_limit_exceeded":
                            wait_seconds = inner_error.get("retry_after", 0)
                            if wait_seconds > 0:
                                print(f"Tag: {tag}. Rate limit exceeded. Waiting {wait_seconds} seconds...")
                                time.sleep(wait_seconds)
                                handled_rate_limit = True
                    except Exception as parse_err:
                        print(f"[WARNING] Tag: {tag}, failed to parse 429 body: {parse_err}")

                if not handled_rate_limit:
                    print(f"Tag: {tag}, max_try in get_model_response_txt: {max_try}, exception: {e}")
                    msg += f"{e}\n"

        error_msg = f"[ERROR] Tag: {tag}, get_model_response_txt() failed after all retries."
        raise Exception(error_msg)

    def get_message_len(self):
        return {
            "prompt_len": sum(len(item["content"]) for item in self.messages if item["role"] == "user"),
            "response_len": sum(len(item["content"]) for item in self.messages if item["role"] == "assistant"),
            "num_calls": len(self.messages) // 2
        }

    def init_messages(self):
        self.messages = []


class GPTChat(BaseChat):
    def __init__(self, model="gpt-4o", temperature=0.6):
        super().__init__(model, temperature)

    def init_system_prompt(self, system_prompt):
        if len(self.messages) == 0:
            self.messages.append({"role": "system", "content": system_prompt})
        else:
            if self.messages[0]["role"] != "system":
                raise Exception(f"Conflicts: The current first message is not a system message, but you provided a system prompt.")
            else:
                self.messages[0]["content"] = system_prompt

    def get_response(self, prompt, system_prompt=None, return_token_usage=False) -> str:
        if system_prompt:
            if len(self.messages) == 0:
                self.messages.append({"role": "system", "content": system_prompt})
            else:
                if self.messages[0]["role"] != "system":
                    raise Exception(f"Conflicts: The current first message is not a system message, but you provided a system prompt.")
                else:
                    self.messages[0]["content"] = system_prompt
        return self.get_response_openai(prompt, return_token_usage=return_token_usage)

    def get_response_openai(self, prompt, return_token_usage=False) -> str:
        """
        Generic OpenAI-compatible API call.
        Set the following environment variables before running:
          - OPENAI_API_KEY:  your API key (e.g. sk-...)
          - OPENAI_BASE_URL: API base URL (default: https://api.openai.com/v1)
        """
        api_key = os.environ.get("OPENAI_API_KEY", "YOUR_API_KEY_HERE")
        base_url = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")

        client = OpenAI(api_key=api_key, base_url=base_url)

        response = client.chat.completions.create(
            model=self.model,
            messages=self.messages + [{"role": "user", "content": prompt}],
            temperature=self.temperature,
        )

        cur_token_usage = response.usage.total_tokens
        main_content = response.choices[0].message.content.strip()

        self.messages.append({"role": "user", "content": prompt})
        self.messages.append({"role": "assistant", "content": main_content})

        if return_token_usage:
            return main_content, cur_token_usage
        else:
            return main_content


if __name__ == "__main__":
    chat = GPTChat(model="gpt-4o", temperature=1.0)
    print(chat.get_model_response_txt("What is the capital of France?", system_prompt="You are a helpful assistant.", tag="test", return_token_usage=True))
    print(chat.messages)
    chat.init_messages()