| import os |
| import requests |
| import openai |
| from dotenv import load_dotenv |
| import json |
|
|
| load_dotenv() |
|
|
| class LLM: |
|
|
| usages = [] |
|
|
| def __init__(self, model_name, logger=None, user_id=None): |
| self.model_name = model_name |
| self.logger = logger |
| self.user_id = user_id |
| self.api_key = None |
| self.api_base = None |
|
|
| def reset(self, api_key=None, api_base=None, model_name=None): |
| if api_key: |
| self.api_key = api_key |
| if api_base: |
| self.api_base = api_base |
| if model_name: |
| self.model_name = model_name |
| self.client = openai.Client(api_key=self.api_key, base_url=self.api_base) |
|
|
| def generate(self, prompt, system='', usage=True): |
| try: |
| if self.model_name in ['deepseek-chat', 'deepseek-reasoner']: |
| response = self.client.chat.completions.create( |
| model=self.model_name, |
| messages=[ |
| {'role': 'system', 'content': system}, |
| {'role': 'user', 'content': prompt} |
| ], |
| temperature=0.7, |
| top_p=1.0, |
| frequency_penalty=0.0, |
| presence_penalty=0.0 |
| ) |
| answer = response.choices[0].message.content |
| usage = { |
| 'completion_tokens': response.usage.completion_tokens, |
| 'prompt_tokens': response.usage.prompt_tokens, |
| 'total_tokens': response.usage.total_tokens |
| } |
| else: |
| messages=[{'role': 'system', 'content': system}, |
| {'role': 'user', 'content': prompt}] |
| raw_request = { |
| "model": self.model_name, |
| "temperature": 0.7, |
| "top_p": 1.0, |
| "frequency_penalty": 0.0, |
| "presence_penalty": 0.0 |
| } |
| |
| HEADERS = { |
| "Content-Type": "application/json", |
| "Authorization": "Bearer {}".format(self.api_key) |
| } |
| response = requests.post(self.api_base, headers=HEADERS, data=json.dumps({"messages": messages, **raw_request})) |
| response.raise_for_status() |
| response = response.json() |
| answer = response["choices"][0]["message"]["content"] |
| if self.model_name == 'DeepSeek-R1-671B': |
| answer = answer.split('\n</think>\n\n')[1] |
| usage = { |
| 'completion_tokens': response['usage']['completion_tokens'], |
| 'prompt_tokens': response['usage']['prompt_tokens'], |
| 'total_tokens': response['usage']['total_tokens'] |
| } |
| if self.logger: |
| self.logger.info(f"[LLM] UserID: {self.user_id} Key: {self.api_key}, Model: {self.model_name}, Usage: {usage}") |
| if usage: |
| self.usages.append(usage) |
| return answer |
|
|
| except Exception as e: |
| return f'An error occurred: {e}' |
|
|
| def get_total_usage(self): |
| total_usage = { |
| 'completion_tokens': 0, |
| 'prompt_tokens': 0, |
| 'total_tokens': 0 |
| } |
| for usage in self.usages: |
| for key, value in usage.items(): |
| total_usage[key] += value |
| return total_usage |
| |
| def clear_usage(self): |
| self.usages = [] |
|
|
|
|
| if __name__ == '__main__': |
| llm = LLM('deepseek-chat') |
| print(llm.generate('Hello, how are you?', 'You are a helpful assistant.', usage=True)) |
|
|
| |
| |
| |
| |
|
|
|
|
|
|
|
|
|
|
|
|
