Standardize LLM clients and expose configurable base URLs.

This commit is contained in:
666ghj
2025-10-09 02:23:08 +08:00
parent 960608cee9
commit 49d4893780
22 changed files with 189 additions and 62 deletions
+8 -3
View File
@@ -9,6 +9,8 @@ from typing import Optional, Dict, Any
from openai import OpenAI
from .base import BaseLLM
DEFAULT_DEEPSEEK_BASE_URL = "https://api.deepseek.com"
# 添加utils目录到Python路径并导入重试模块
try:
current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -29,13 +31,14 @@ except ImportError:
class DeepSeekLLM(BaseLLM):
"""DeepSeek LLM实现类"""
def __init__(self, api_key: Optional[str] = None, model_name: Optional[str] = None):
def __init__(self, api_key: Optional[str] = None, model_name: Optional[str] = None, base_url: Optional[str] = None):
"""
初始化DeepSeek客户端
Args:
api_key: DeepSeek API密钥,如果不提供则从环境变量读取
model_name: 模型名称,默认使用deepseek-chat
base_url: DeepSeek API基础地址
"""
if api_key is None:
api_key = os.getenv("DEEPSEEK_API_KEY")
@@ -44,10 +47,12 @@ class DeepSeekLLM(BaseLLM):
super().__init__(api_key, model_name)
self.base_url = base_url or os.getenv("DEEPSEEK_BASE_URL") or DEFAULT_DEEPSEEK_BASE_URL
# 初始化OpenAI客户端,使用DeepSeek的endpoint
self.client = OpenAI(
api_key=self.api_key,
base_url="https://api.deepseek.com"
base_url=self.base_url
)
self.default_model = model_name or self.get_default_model()
@@ -109,5 +114,5 @@ class DeepSeekLLM(BaseLLM):
return {
"provider": "DeepSeek",
"model": self.default_model,
"api_base": "https://api.deepseek.com"
"api_base": self.base_url
}
+8 -3
View File
@@ -9,6 +9,8 @@ from typing import Optional, Dict, Any
from openai import OpenAI
from .base import BaseLLM
DEFAULT_GEMINI_BASE_URL = "https://www.chataiapi.com/v1"
# 添加utils目录到Python路径并导入重试模块
try:
current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -29,13 +31,14 @@ except ImportError:
class GeminiLLM(BaseLLM):
"""Gemini LLM实现类"""
def __init__(self, api_key: Optional[str] = None, model_name: Optional[str] = None):
def __init__(self, api_key: Optional[str] = None, model_name: Optional[str] = None, base_url: Optional[str] = None):
"""
初始化Gemini客户端
Args:
api_key: Gemini API密钥,如果不提供则从环境变量读取
model_name: 模型名称,默认使用gemini-2.5-pro
base_url: Gemini API基础地址
"""
if api_key is None:
api_key = os.getenv("GEMINI_API_KEY")
@@ -44,10 +47,12 @@ class GeminiLLM(BaseLLM):
super().__init__(api_key, model_name)
self.base_url = base_url or os.getenv("GEMINI_BASE_URL") or DEFAULT_GEMINI_BASE_URL
# 初始化OpenAI客户端,使用Gemini的中转endpoint
self.client = OpenAI(
api_key=self.api_key,
base_url="https://www.chataiapi.com/v1"
base_url=self.base_url
)
self.default_model = model_name or self.get_default_model()
@@ -109,5 +114,5 @@ class GeminiLLM(BaseLLM):
return {
"provider": "Gemini",
"model": self.default_model,
"api_base": "https://www.chataiapi.com/v1"
"api_base": self.base_url
}