| 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 |
