Standardize LLM clients and expose configurable base URLs.

QueryEngine/agent.py

@@ -60,7 +60,8 @@ class DeepSearchAgent:
         if self.config.default_llm_provider == "deepseek":
             return DeepSeekLLM(
                 api_key=self.config.deepseek_api_key,
-                model_name=self.config.deepseek_model
+                model_name=self.config.deepseek_model,
+                base_url=self.config.deepseek_base_url
             )
         elif self.config.default_llm_provider == "openai":
             return OpenAILLM(

QueryEngine/llms/deepseek.py

@@ -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
         }

QueryEngine/utils/config.py

@@ -15,6 +15,8 @@ class Config:
     deepseek_api_key: Optional[str] = None
     openai_api_key: Optional[str] = None
     tavily_api_key: Optional[str] = None
+    deepseek_base_url: str = "https://api.deepseek.com"
+    openai_base_url: Optional[str] = None
     
     # 模型配置
     default_llm_provider: str = "deepseek"  # deepseek 或 openai
@@ -66,6 +68,8 @@ class Config:
                 deepseek_api_key=getattr(config_module, "DEEPSEEK_API_KEY", None),
                 openai_api_key=getattr(config_module, "OPENAI_API_KEY", None),
                 tavily_api_key=getattr(config_module, "TAVILY_API_KEY", None),
+                deepseek_base_url=getattr(config_module, "DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
+                openai_base_url=getattr(config_module, "OPENAI_BASE_URL", None),
                 default_llm_provider=getattr(config_module, "DEFAULT_LLM_PROVIDER", "deepseek"),
                 deepseek_model=getattr(config_module, "DEEPSEEK_MODEL", "deepseek-chat"),
                 openai_model=getattr(config_module, "OPENAI_MODEL", "gpt-4o-mini"),
@@ -93,6 +97,8 @@ class Config:
                 deepseek_api_key=config_dict.get("DEEPSEEK_API_KEY"),
                 openai_api_key=config_dict.get("OPENAI_API_KEY"),
                 tavily_api_key=config_dict.get("TAVILY_API_KEY"),
+                deepseek_base_url=config_dict.get("DEEPSEEK_BASE_URL", "https://api.deepseek.com"),
+                openai_base_url=config_dict.get("OPENAI_BASE_URL"),
                 default_llm_provider=config_dict.get("DEFAULT_LLM_PROVIDER", "deepseek"),
                 deepseek_model=config_dict.get("DEEPSEEK_MODEL", "deepseek-chat"),
                 openai_model=config_dict.get("OPENAI_MODEL", "gpt-4o-mini"),