import os
import dotenv
from openai import OpenAI
from tenacity import retry, wait_random_exponential, stop_after_attempt

from .Base import BaseModel

dotenv.load_dotenv()

class Pangu72B(BaseModel):
    def __init__(self, temperature=0):
        api_key = os.getenv("PANGU_API_KEY_72b")
        api_base_url = os.getenv("PANGU_API_URL_72b")
        self.model = "pangu-pro"
        self.deepseek_client = OpenAI(api_key=api_key, base_url=api_base_url)


    @retry(wait=wait_random_exponential(min=1, max=120), stop=stop_after_attempt(30))
    def prompt(self, processed_input):
        response = self.deepseek_client.chat.completions.create(
            model=self.model,
            messages=processed_input,
            stream=False,
            temperature=0.7
        )
        return response.choices[0].message.content, response.usage.prompt_tokens, response.usage.completion_tokens