✨ 优化超时配置 (#40)

Author: KomoriDev

README.md

@@ -102,6 +102,7 @@ plugins = ["nonebot_plugin_deepseek"]
 |   deepseek__enable_models    | 否 |[{ "name": "deepseek-chat" }, { "name": "deepseek-reasoner" }]|启用的模型 [配置说明](https://github.com/KomoriDev/nonebot-plugin-deepseek/wiki/%E9%85%8D%E7%BD%AE#enable_models-%E9%85%8D%E7%BD%AE%E8%AF%B4%E6%98%8E)|
 |       deepseek__prompt       | 否 |                              无                              |                模型预设                |
 |       deepseek__stream       | 否 |                             False                            |            是否启用流式传输            |
+|       deepseek__timeout      | 否 |            {"api_request": 100, "user_input": 60}            |                超时设定                |
 |     deepseek__md_to_pic      | 否 |                             False                            |        是否启用 Markdown 转图片        |
 |deepseek__enable_send_thinking| 否 |                             False                            |             是否发送思维链             |
 

nonebot_plugin_deepseek/apis/request.py

@@ -26,7 +26,7 @@ async def chat(cls, message: list[dict[str, str]], model: str = "deepseek-chat")
         prompt = model_dump(model_config, exclude_none=True).get("prompt", config.prompt)
 
         json = {
-            "messages": [{"content": prompt, "role": "system"}] + message if prompt else message,
+            "messages": ([{"content": prompt, "role": "system"}] + message if prompt else message),
             "model": model,
             **model_config.to_dict(),
         }
@@ -56,12 +56,16 @@ async def query_balance(cls, model_name: str) -> Balance:
 
 
 async def common_request(base_url: str, api_key: str, json: dict):
+    timeout_config = config.timeout
     async with httpx.AsyncClient() as client:
         response = await client.post(
             f"{base_url}/chat/completions",
-            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
+            headers={
+                "Authorization": f"Bearer {api_key}",
+                "Content-Type": "application/json",
+            },
             json=json,
-            timeout=50,
+            timeout=(timeout_config if isinstance(timeout_config, int) else timeout_config.api_request),
         )
     if error := response.json().get("error"):
         raise RequestException(error["message"])
@@ -74,7 +78,10 @@ async def stream_request(base_url: str, api_key: str, json: dict):
         async with client.stream(
             "POST",
             f"{base_url}/chat/completions",
-            headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
+            headers={
+                "Authorization": f"Bearer {api_key}",
+                "Content-Type": "application/json",
+            },
             json=json,
         ) as response:
             ret_list: Optional[StreamChoiceList] = None
@@ -107,7 +114,9 @@ async def stream_request(base_url: str, api_key: str, json: dict):
             return ret_list.transform()
 
 
-def sse_middle(line: str) -> Union[tuple[Literal["data", "event", "id", "retry", "::", "error"], str], None]:
+def sse_middle(
+    line: str,
+) -> Union[tuple[Literal["data", "event", "id", "retry", "::", "error"], str], None]:
     """单行SSE数据解析"""
     line = line.strip("\r")
     if not line:

nonebot_plugin_deepseek/config.py

@@ -139,6 +139,13 @@ def to_dict(self):
         )
 
 
+class TimeoutConfig(BaseModel):
+    api_request: int = Field(default=100)
+    """API request timeout (Not applicable for streaming)"""
+    user_input: int = Field(default=60)
+    """User input timeout"""
+
+
 class ScopedConfig(BaseModel):
     api_key: str = ""
     """Your API Key from deepseek"""
@@ -153,8 +160,8 @@ class ScopedConfig(BaseModel):
     """Text to Image"""
     enable_send_thinking: bool = False
     """Whether to send model thinking chain"""
-    context_timeout: int = Field(default=50, gt=50)
-    """Multi-round conversation timeout"""
+    timeout: Union[int, TimeoutConfig] = Field(default_factory=TimeoutConfig)
+    """Timeout"""
     stream: bool = False
     """Stream"""
 

nonebot_plugin_deepseek/utils.py

@@ -54,7 +54,8 @@ async def _handle_single_conversion(self) -> None:
             await self._send_response(message)
 
     async def _handle_multi_round_conversion(self) -> None:
-        async for resp in self.waiter(default=False, timeout=config.context_timeout):
+        timeout = config.timeout if isinstance(config.timeout, int) else config.timeout.user_input
+        async for resp in self.waiter(default=False, timeout=timeout):
             await self._process_waiter_response(resp)
 
             if resp == "rollback":
@@ -73,7 +74,12 @@ async def _handle_multi_round_conversion(self) -> None:
 
     def _setup_waiter(self) -> Waiter[Union[str, Literal[False]]]:
         permission = Permission(User.from_event(self.event, perm=self.matcher.permission))
-        waiter = Waiter(waits=["message"], handler=self._waiter_handler, matcher=self.matcher, permission=permission)
+        waiter = Waiter(
+            waits=["message"],
+            handler=self._waiter_handler,
+            matcher=self.matcher,
+            permission=permission,
+        )
         waiter.future.set_result("")
         return waiter
 
@@ -92,8 +98,14 @@ def _prompt_handler(self, msg: UniMsg) -> UniMsg:
         return msg
 
     async def _process_waiter_response(self, resp: Union[bool, str]) -> None:
+        timeout = config.timeout if isinstance(config.timeout, int) else config.timeout.user_input
+
         if resp == "" and not self.context:
-            _resp = await prompt("你想对 DeepSeek 说什么呢？", handler=self._prompt_handler, timeout=60)
+            _resp = await prompt(
+                "你想对 DeepSeek 说什么呢？",
+                handler=self._prompt_handler,
+                timeout=timeout,
+            )
             if _resp is None:
                 await UniMessage.text("等待超时").finish(reply_to=self.message_id)
             resp = self._waiter_handler(_resp, skip=True)
@@ -138,7 +150,13 @@ async def _handle_tool_calls(self, message: Message) -> bool:
             self.context.pop()
             return False
 
-        self.context.append({"role": "tool", "tool_call_id": message.tool_calls[0].id, "content": result})
+        self.context.append(
+            {
+                "role": "tool",
+                "tool_call_id": message.tool_calls[0].id,
+                "content": result,
+            }
+        )
         return True
 
     async def _get_response_message(self) -> Optional[Message]: