130 lines
3.7 KiB
Python
130 lines
3.7 KiB
Python
"""DeepSeek LLM客户端 - OpenAI兼容API"""
|
||
|
||
import time
|
||
from typing import List, Dict, Optional
|
||
from loguru import logger
|
||
import httpx
|
||
|
||
from .base_client import BaseLLMClient, LLMResponse, LLMConfig, LLMProvider
|
||
|
||
|
||
class DeepSeekClient(BaseLLMClient):
|
||
"""
|
||
DeepSeek API客户端（OpenAI兼容格式）
|
||
|
||
支持模型：
|
||
- deepseek-chat
|
||
- deepseek-coder
|
||
- deepseek-reasoner
|
||
- deepseek-v3
|
||
- deepseek-v3.2
|
||
- deepseek-v4-flash
|
||
"""
|
||
|
||
SUPPORTED_MODELS = [
|
||
"deepseek-chat",
|
||
"deepseek-coder",
|
||
"deepseek-reasoner",
|
||
"deepseek-v3",
|
||
"deepseek-v3.2",
|
||
"deepseek-v4-flash"
|
||
]
|
||
|
||
def __init__(self, config: LLMConfig):
|
||
if config.provider != LLMProvider.DEEPSEEK:
|
||
raise ValueError(f"配置provider应为DEEPSEEK，实际为{config.provider}")
|
||
super().__init__(config)
|
||
self._init_client()
|
||
|
||
def _init_client(self):
|
||
"""初始化HTTP客户端"""
|
||
self._client = httpx.Client(
|
||
base_url=self.config.base_url,
|
||
headers={
|
||
"Authorization": f"Bearer {self.config.api_key}",
|
||
"Content-Type": "application/json"
|
||
},
|
||
timeout=self.config.timeout
|
||
)
|
||
logger.info(f"DeepSeek客户端初始化完成: {self.config.base_url} - {self.config.model}")
|
||
|
||
def chat(
|
||
self,
|
||
messages: List[Dict[str, str]],
|
||
max_tokens: Optional[int] = None,
|
||
temperature: Optional[float] = None,
|
||
**kwargs
|
||
) -> LLMResponse:
|
||
"""对话补全"""
|
||
start_time = time.time()
|
||
|
||
try:
|
||
payload = {
|
||
"model": self.config.model,
|
||
"messages": messages,
|
||
"max_tokens": max_tokens or self.config.max_tokens,
|
||
"temperature": temperature or self.config.temperature,
|
||
"top_p": kwargs.get("top_p", self.config.top_p),
|
||
"stream": False
|
||
}
|
||
|
||
response = self._client.post("/chat/completions", json=payload)
|
||
response.raise_for_status()
|
||
|
||
data = response.json()
|
||
|
||
latency_ms = int((time.time() - start_time) * 1000)
|
||
|
||
choices = data.get("choices", [{}])
|
||
message = choices[0].get("message", {})
|
||
|
||
return LLMResponse(
|
||
content=message.get("content", ""),
|
||
model=data.get("model", self.config.model),
|
||
usage=data.get("usage", {}),
|
||
finish_reason=choices[0].get("finish_reason", "stop"),
|
||
latency_ms=latency_ms
|
||
)
|
||
|
||
except httpx.HTTPStatusError as e:
|
||
logger.error(f"DeepSeek API错误: {e.response.status_code} - {e.response.text}")
|
||
return LLMResponse(
|
||
content="",
|
||
model=self.config.model,
|
||
error=f"API错误: {e.response.status_code} - {e.response.text[:200]}"
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"DeepSeek调用失败: {e}")
|
||
return LLMResponse(
|
||
content="",
|
||
model=self.config.model,
|
||
error=str(e)
|
||
)
|
||
|
||
def get_available_models(self) -> List[str]:
|
||
"""获取可用模型列表"""
|
||
return self.SUPPORTED_MODELS
|
||
|
||
def close(self):
|
||
"""关闭客户端"""
|
||
if self._client:
|
||
self._client.close()
|
||
|
||
|
||
def create_deepseek_client(
|
||
api_key: str,
|
||
model: str = "deepseek-v4-flash",
|
||
base_url: str = "http://6.86.80.4:30080/v1",
|
||
**kwargs
|
||
) -> DeepSeekClient:
|
||
"""便捷函数：创建DeepSeek客户端"""
|
||
config = LLMConfig(
|
||
provider=LLMProvider.DEEPSEEK,
|
||
model=model,
|
||
api_key=api_key,
|
||
base_url=base_url,
|
||
**kwargs
|
||
)
|
||
return DeepSeekClient(config)
|