Update llm_adapters.py

This commit is contained in:
CNlaojing
2025-03-13 06:47:27 +08:00
committed by GitHub
parent d2576ddaed
commit 33944c0266
+35
View File
@@ -9,6 +9,7 @@ from azure.ai.inference import ChatCompletionsClient
from azure.core.credentials import AzureKeyCredential from azure.core.credentials import AzureKeyCredential
from azure.ai.inference.models import SystemMessage, UserMessage from azure.ai.inference.models import SystemMessage, UserMessage
from openai import OpenAI from openai import OpenAI
import requests
def check_base_url(url: str) -> str: def check_base_url(url: str) -> str:
@@ -306,6 +307,38 @@ class VolcanoEngineAIAdapter(BaseLLMAdapter):
logging.error(f"火山引擎API调用超时或失败: {e}") logging.error(f"火山引擎API调用超时或失败: {e}")
return "" return ""
class SiliconFlowAdapter(BaseLLMAdapter):
def __init__(self, api_key: str, base_url: str, model_name: str, max_tokens: int, temperature: float = 0.7, timeout: Optional[int] = 600):
self.base_url = check_base_url(base_url)
self.api_key = api_key
self.model_name = model_name
self.max_tokens = max_tokens
self.temperature = temperature
self.timeout = timeout
self._client = OpenAI(
base_url=base_url,
api_key=api_key,
timeout=timeout # 添加超时配置
)
def invoke(self, prompt: str) -> str:
try:
response = self._client.chat.completions.create(
model=self.model_name,
messages=[
{"role": "system", "content": "你是DeepSeek,是一个 AI 人工智能助手"},
{"role": "user", "content": prompt},
],
timeout=self.timeout # 添加超时参数
)
if not response:
logging.warning("No response from DeepSeekAdapter.")
return ""
return response.choices[0].message.content
except Exception as e:
logging.error(f"硅基流动API调用超时或失败: {e}")
return ""
def create_llm_adapter( def create_llm_adapter(
interface_format: str, interface_format: str,
base_url: str, base_url: str,
@@ -338,5 +371,7 @@ def create_llm_adapter(
return OpenAIAdapter(api_key, base_url, model_name, max_tokens, temperature, timeout) return OpenAIAdapter(api_key, base_url, model_name, max_tokens, temperature, timeout)
elif fmt == "火山引擎": elif fmt == "火山引擎":
return VolcanoEngineAIAdapter(api_key, base_url, model_name, max_tokens, temperature, timeout) return VolcanoEngineAIAdapter(api_key, base_url, model_name, max_tokens, temperature, timeout)
elif fmt == "硅基流动":
return SiliconFlowAdapter(api_key, base_url, model_name, max_tokens, temperature, timeout)
else: else:
raise ValueError(f"Unknown interface_format: {interface_format}") raise ValueError(f"Unknown interface_format: {interface_format}")