diff --git a/docutranslate/agents/agent.py b/docutranslate/agents/agent.py index 32b11f5..1c1cfef 100644 --- a/docutranslate/agents/agent.py +++ b/docutranslate/agents/agent.py @@ -13,6 +13,7 @@ from urllib.parse import urlparse import httpx +from docutranslate.agents.thinking.thinking_factory import get_thinking_mode from docutranslate.logger import global_logger from docutranslate.utils.utils import get_httpx_proxies @@ -197,33 +198,6 @@ ErrorResultHandlerType = Callable[[str, logging.Logger], Any] class Agent: - _think_factory = { - "open.bigmodel.cn": ("thinking", {"type": "enabled"}, {"type": "disabled"}), - "dashscope.aliyuncs.com": ( - "extra_body", - {"enable_thinking": True}, - {"enable_thinking": False}, - ), - "ark.cn-beijing.volces.com": ( - "thinking", - {"type": "enabled"}, - {"type": "disabled"}, - ), - "generativelanguage.googleapis.com": ( - "extra_body", - { - "google": { - "thinking_config": {"thinking_budget": -1, "include_thoughts": True} - } - }, - { - "google": { - "thinking_config": {"thinking_budget": 0, "include_thoughts": False} - } - }, - ), - "api.siliconflow.cn": ("enable_thinking", True, False), - } def __init__(self, config: AgentConfig): @@ -251,9 +225,10 @@ class Agent: self.system_proxy_enable = config.system_proxy_enable def _add_thinking_mode(self, data: dict): - if self.domain not in self._think_factory: + thinking_mode_result=get_thinking_mode(self.domain,data.get("model")) + if thinking_mode_result is None: return - field_thinking, val_enable, val_disable = self._think_factory[self.domain] + field_thinking, val_enable, val_disable = thinking_mode_result if self.thinking == "enable": data[field_thinking] = val_enable elif self.thinking == "disable": @@ -313,6 +288,7 @@ class Agent: headers=headers, timeout=self.timeout, ) + # print(f"【测试】json:\n{data}") response.raise_for_status() # print(f"【测试】resp:\n{response.json()}") result = response.json()["choices"][0]["message"]["content"] diff --git a/docutranslate/agents/thinking/__init__.py b/docutranslate/agents/thinking/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/docutranslate/agents/thinking/thinking_factory.py b/docutranslate/agents/thinking/thinking_factory.py new file mode 100644 index 0000000..e1753fe --- /dev/null +++ b/docutranslate/agents/thinking/thinking_factory.py @@ -0,0 +1,69 @@ +from typing import TypeAlias + +mode_type:TypeAlias=str +thinking_field: TypeAlias=str +enable_value: TypeAlias= str | dict +disable_value: TypeAlias= str | dict + + +thinking_mode:dict[mode_type,tuple[thinking_field, enable_value, disable_value]]={ + "bigmodel": ("thinking", {"type": "enabled"}, {"type": "disabled"}), + "aliyun": ( + "extra_body", + {"enable_thinking": True}, + {"enable_thinking": False}, + ), + "volces": ( + "thinking", + {"type": "enabled"}, + {"type": "disabled"}, + ), + "google": ( + "extra_body", + { + "google": { + "thinking_config": {"thinking_budget": -1, "include_thoughts": True} + } + }, + { + "google": { + "thinking_config": {"thinking_budget": 0, "include_thoughts": False} + } + }, + ), + "siliconflow": ("enable_thinking", True, False), + } + + +def get_thinking_mode_by_model_id(model_id: str) -> tuple[str, str | dict, str | dict] | None: + model_id = model_id.strip().lower() + if "glm-4.5" in model_id: + return thinking_mode["bigmodel"] + elif "qwen/qwen3" in model_id: + return thinking_mode["aliyun"] + elif "seed-1-6" in model_id: + return thinking_mode["volces"] + elif "gemini" in model_id: + return thinking_mode["google"] + return None + + +def get_thinking_mode(provider: str, model_id: str) -> tuple[str, str | dict, str | dict] | None: + provider = provider.strip() + if provider == "open.bigmodel.cn": + return thinking_mode["bigmodel"] + elif provider == "dashscope.aliyuncs.com": + return thinking_mode["aliyun"] + elif provider == "ark.cn-beijing.volces.com": + return thinking_mode["volces"] + elif provider == "generativelanguage.googleapis.com": + return thinking_mode["google"] + elif provider == "api.siliconflow.cn": + return thinking_mode["siliconflow"] + elif provider == "api.302.ai": + return get_thinking_mode_by_model_id(model_id) + return None + + +# def add_thinking_mode(data: dict, provider: str, model_id: str, think_enable: bool): +# pass