Spaces:
Sleeping
Sleeping
File size: 3,487 Bytes
d12a6df |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
try:
from openai import OpenAI
except ImportError:
raise ImportError("Please install the openai package by running `pip install openai`, and add 'DEEPSEEK_API_KEY' to your environment variables.")
import os
import platformdirs
from tenacity import (
retry,
stop_after_attempt,
wait_random_exponential,
)
from typing import List, Union
from .base import EngineLM, CachedEngine
class ChatDeepseek(EngineLM, CachedEngine):
DEFAULT_SYSTEM_PROMPT = "You are a helpful assistant."
def __init__(
self,
model_string="deepseek-chat",
use_cache: bool=False,
system_prompt=DEFAULT_SYSTEM_PROMPT,
is_multimodal: bool=False):
self.model_string = model_string
self.use_cache = use_cache
self.system_prompt = system_prompt
self.is_multimodal = is_multimodal
self.is_chat_model = any(x in model_string for x in ["deepseek-chat"])
self.is_reasoning_model = any(x in model_string for x in ["deepseek-reasoner"])
if self.use_cache:
root = platformdirs.user_cache_dir("agentflow")
cache_path = os.path.join(root, f"cache_deepseek_{model_string}.db")
super().__init__(cache_path=cache_path)
if os.getenv("DEEPSEEK_API_KEY") is None:
raise ValueError("Please set the DEEPSEEK_API_KEY environment variable.")
self.client = OpenAI(
api_key=os.getenv("DEEPSEEK_API_KEY"),
base_url="https://api.deepseek.com"
)
@retry(wait=wait_random_exponential(min=1, max=5), stop=stop_after_attempt(5))
def generate(self, content: Union[str, List[Union[str, bytes]]], system_prompt=None, **kwargs):
if isinstance(content, list) and len(content) == 1:
content = content[0]
if isinstance(content, str):
return self._generate_text(content, system_prompt=system_prompt, **kwargs)
def _generate_text(
self, prompt, system_prompt=None, temperature=0, max_tokens=4000, top_p=0.99, response_format=None
):
sys_prompt_arg = system_prompt if system_prompt else self.system_prompt
if self.use_cache:
cache_key = sys_prompt_arg + prompt
cache_or_none = self._check_cache(cache_key)
if cache_or_none is not None:
return cache_or_none
if self.is_chat_model:
response = self.client.chat.completions.create(
model=self.model_string,
messages=[
{"role": "system", "content": sys_prompt_arg},
{"role": "user", "content": prompt},
],
temperature=temperature,
max_tokens=max_tokens,
top_p=top_p,
)
response = response.choices[0].message.content
elif self.is_reasoning_model:
response = self.client.chat.completions.create(
model=self.model_string,
messages=[
{"role": "system", "content": sys_prompt_arg},
{"role": "user", "content": prompt},
],
temperature=temperature,
top_p=top_p,
)
response = response.choices[0].message.content
if self.use_cache:
self._save_cache(cache_key, response)
return response
def __call__(self, prompt, **kwargs):
return self.generate(prompt, **kwargs)
|