import os
import dotenv
from openai import OpenAI

from .Base import BaseModel

dotenv.load_dotenv()

class DeepSeekModel(BaseModel):
    def __init__(self, temperature=0):
        api_key = os.getenv("DEEPSEEK_API_KEY")
        api_base_url = os.getenv("DEEPSEEK_API_URL")
        self.model = "deepseek-chat"
        self.deepseek_client = OpenAI(api_key=api_key, base_url=api_base_url)

    def prompt(self, processed_input):
        response = self.deepseek_client.chat.completions.create(
            model=self.model,
            messages=processed_input,
            stream=False
        )
        return response.choices[0].message.content, response.usage.prompt_tokens, response.usage.completion_tokens