✨ 自定义模型 API Key 和 prompt (#18)

Co-authored-by: Zaxpris <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: KomoriDev <[email protected]>

Author: 3 people

README.md

@@ -112,6 +112,8 @@ plugins = ["nonebot_plugin_deepseek"]
 
 - `name`: 模型名称（必填）
 - `base_url`: 接口地址（默认为：<https://api.deepseek.com>）（自建模型必填）
+- `api_key`: API Key（默认使用 `deepseek__api_key`）
+- `prompt`: 模型预设（默认使用 `deepseek__prompt`）
 - `max_tokens`: 模型生成补全的最大 token 数
   - `deepseek-chat`: 介于 1 到 8192 间的整数，默认使用 4096
   - `deepseek-reasoner`: 默认为 4K，最大为 8K

nonebot_plugin_deepseek/__init__.py

@@ -35,6 +35,7 @@
 from .apis import API
 from . import hook as hook
 from .utils import DeepSeekHandler
+from .exception import RequestException
 from .extension import CleanDocExtension
 from .config import Config, config, model_config
 
@@ -114,7 +115,7 @@ async def _(is_superuser: bool = Depends(SuperUser())):
     if not is_superuser:
         return
     try:
-        balances = await API.query_balance()
+        balances = await API.query_balance(model_config.default_model)
 
         await deepseek.finish(
             "".join(
@@ -129,6 +130,8 @@ async def _(is_superuser: bool = Depends(SuperUser())):
         )
     except ValueError as e:
         await deepseek.finish(str(e))
+    except RequestException as e:
+        await deepseek.finish(str(e))
 
 
 @deepseek.assign("model.list")

nonebot_plugin_deepseek/apis/request.py

@@ -2,6 +2,7 @@
 from nonebot.log import logger
 
 from ..config import config
+from ..compat import model_dump
 
 # from ..function_call import registry
 from ..exception import RequestException
@@ -11,17 +12,18 @@
 class API:
     _headers = {
         "Accept": "application/json",
-        "Authorization": f"Bearer {config.api_key}",
     }
 
     @classmethod
     async def chat(cls, message: list[dict[str, str]], model: str = "deepseek-chat") -> ChatCompletions:
         """普通对话"""
         model_config = config.get_model_config(model)
+
+        api_key = model_config.api_key or config.api_key
+        prompt = model_dump(model_config, exclude_none=True).get("prompt", config.prompt)
+
         json = {
-            "messages": [{"content": config.prompt, "role": "system"}] + message
-            if config.prompt and model == "deepseek-chat"
-            else message,
+            "messages": [{"content": prompt, "role": "system"}] + message if prompt else message,
             "model": model,
             **model_config.to_dict(),
         }
@@ -31,7 +33,7 @@ async def chat(cls, message: list[dict[str, str]], model: str = "deepseek-chat")
         async with httpx.AsyncClient() as client:
             response = await client.post(
                 f"{model_config.base_url}/chat/completions",
-                headers={**cls._headers, "Content-Type": "application/json"},
+                headers={**cls._headers, "Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
                 json=json,
                 timeout=50,
             )
@@ -40,12 +42,15 @@ async def chat(cls, message: list[dict[str, str]], model: str = "deepseek-chat")
         return ChatCompletions(**response.json())
 
     @classmethod
-    async def query_balance(cls) -> Balance:
-        """查询账号余额"""
+    async def query_balance(cls, model_name: str) -> Balance:
+        model_config = config.get_model_config(model_name)
+        api_key = model_config.api_key or config.api_key
+
         async with httpx.AsyncClient() as client:
             response = await client.get(
-                f"{config.get_model_url('deepseek-chat')}/user/balance",
-                headers=cls._headers,
+                f"{model_config.base_url}/user/balance",
+                headers={**cls._headers, "Authorization": f"Bearer {api_key}"},
             )
-
+        if response.status_code == 404:
+            raise RequestException("本地模型不支持查询余额，请更换默认模型")
         return Balance(**response.json())

nonebot_plugin_deepseek/config.py

@@ -44,6 +44,10 @@ class CustomModel(BaseModel):
     """Model Name"""
     base_url: str = "https://api.deepseek.com"
     """Custom base URL for this model (optional)"""
+    api_key: Optional[str] = None
+    """Custom API Key for the model (optional)"""
+    prompt: Optional[str] = None
+    """Custom character preset for the model (optional)"""
     max_tokens: int = Field(default=4090, gt=1, lt=8192)
     """
     限制一次请求中模型生成 completion 的最大 token 数
@@ -123,7 +127,9 @@ def check_max_token(cls, data: Any) -> Any:
         return data
 
     def to_dict(self):
-        return model_dump(self, exclude_unset=True, exclude_none=True, exclude={"name", "base_url"})
+        return model_dump(
+            self, exclude_unset=True, exclude_none=True, exclude={"name", "base_url", "api_key", "prompt"}
+        )
 
 
 class ScopedConfig(BaseModel):

tutorial.md

@@ -31,14 +31,21 @@ ollama run deepseek-r1:7b
 配置示例
 
 ```bash
+deepseek__prompt = "You are a helpful assistant."
 deepseek__api_key = "sk-xxx"
 deepseek__enable_models='
 [
   { "name": "deepseek-chat" },
   { "name": "deepseek-reasoner" },
   {
     "name": "deepseek-r1:7b",
-    "base_url": "http://localhost:11434/v1"
+    "base_url": "http://localhost:11434/v1",
+    "prompt": None
+  },
+  {
+    "name": "deepseek-r1:free",
+    "base_url": "https://openrouter.ai/api/v1",
+    "api_key": "sk-openrouter-xxx"
   }
 ]
 '