From 5fe1a7c868c8cef2a29f94dbbb901e62211e0192 Mon Sep 17 00:00:00 2001 From: shuiyu67 Date: Thu, 6 Aug 2026 22:50:18 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E7=BB=9F=E4=B8=80=20max=5Finput=5Ftoke?= =?UTF-8?q?ns=20API=EF=BC=8C=E6=89=80=E6=9C=89=E5=BC=95=E6=93=8E=E5=BF=85?= =?UTF-8?q?=E9=A1=BB=E8=BF=94=E5=9B=9E=E6=9C=80=E5=A4=A7=E8=BE=93=E5=85=A5?= =?UTF-8?q?=20token=20=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - openai/ollama/manual 三引擎均暴露 max_input_tokens 属性（统一契约） - ollama 额外上报 prompt_eval_count 到 last_prompt_tokens，使本地模型也走 token 感知压缩 - should_compress 改用 max_input_tokens，系统依当前引擎上限的 70%（compress_ratio）自动触发 - 压缩始终调用当前引擎（current_engine）做总结，符合"调用当前引擎总结压缩" - 新增 tests/test_max_input_tokens.py 覆盖契约 + 70% 触发 + 消息数兜底 Co-Authored-By: WorkBuddy --- ai_engines/manual_engine.py | 7 ++++ ai_engines/ollama_engine.py | 20 ++++++++++- ai_engines/openai_engine.py | 5 +++ tests/test_max_input_tokens.py | 66 ++++++++++++++++++++++++++++++++++ xcli_core/cli_core.py | 5 +-- xcli_core/memory.py | 8 ++--- 6 files changed, 104 insertions(+), 7 deletions(-) create mode 100644 tests/test_max_input_tokens.py diff --git a/ai_engines/manual_engine.py b/ai_engines/manual_engine.py index 08bd3d2..3cdd691 100644 --- a/ai_engines/manual_engine.py +++ b/ai_engines/manual_engine.py @@ -44,6 +44,8 @@ class ManualAI: self.conversation_history = [] self.max_history = 10 self._subprocess = None + # 最近一次请求的 prompt token 数（手动引擎无真实模型，恒为 None） + self.last_prompt_tokens = None self._session_id = str(int(time.time())) self._started = False @@ -240,6 +242,11 @@ while True: time.sleep(0.3) ''' + @property + def max_input_tokens(self): + """当前引擎支持的最大输入 token 数（统一 API）。手动引擎无真实模型，给名义上限。""" + return 128000 + def generate_response(self, user_input, tool_results=None, system_prompt=None): """生成响应 — 转发给另一个终端的人工""" paths = self._get_paths() diff --git a/ai_engines/ollama_engine.py b/ai_engines/ollama_engine.py index 676aa33..55d449e 100644 --- a/ai_engines/ollama_engine.py +++ b/ai_engines/ollama_engine.py @@ -113,6 +113,9 @@ Ollama AI引擎插件帮助信息 # 共享对话历史引用（由CLI设置） self.shared_conversation_history = None + + # 最近一次请求的 prompt token 数（供 token 感知压缩使用） + self.last_prompt_tokens = None # 直接启用思考标记，由主程序自动处理深度思考模型的响应 self.thinking_start_marker = "" @@ -156,7 +159,13 @@ Ollama AI引擎插件帮助信息 except Exception as e: # 如果出现异常,说明服务未运行 return False - + + @property + def max_input_tokens(self): + """当前引擎支持的最大输入 token 数（统一 API，供上下文压缩使用）。 + 本地 Ollama 由配置 max_token_k 换算的 num_ctx 决定。""" + return self.num_ctx + def generate_response(self, user_input, tool_results=None, system_prompt=None, tools=None): """ 生成AI响应 - 使用Ollama API @@ -203,6 +212,15 @@ Ollama AI引擎插件帮助信息 # ── Function Calling: 检查 tool_calls ── message = response.get("message", {}) + # 上报本轮 prompt token 数（Ollama 返回 prompt_eval_count），供 token 感知压缩 + try: + if isinstance(response, dict): + _pet = response.get("prompt_eval_count") + else: + _pet = getattr(response, "prompt_eval_count", None) + self.last_prompt_tokens = int(_pet) if _pet else None + except Exception: + self.last_prompt_tokens = None tool_calls_raw = message.get("tool_calls") if tool_calls_raw: # Ollama 的 tool_calls 格式: [{"function": {"name": ..., "arguments": {...}}}] diff --git a/ai_engines/openai_engine.py b/ai_engines/openai_engine.py index 7867aac..ffe1d1c 100644 --- a/ai_engines/openai_engine.py +++ b/ai_engines/openai_engine.py @@ -582,6 +582,11 @@ OpenAI 兼容格式引擎帮助信息 return win return 128000 + @property + def max_input_tokens(self): + """当前引擎支持的最大输入 token 数（统一 API，供上下文压缩使用）。""" + return self.context_window() + def list_models(self): """列出可用模型（调用 API 的 /models 端点）""" if not self.base_url: diff --git a/tests/test_max_input_tokens.py b/tests/test_max_input_tokens.py new file mode 100644 index 0000000..e857b7e --- /dev/null +++ b/tests/test_max_input_tokens.py @@ -0,0 +1,66 @@ +"""max_input_tokens 统一 API + token 感知压缩 回归测试 + +验证: +1. 所有引擎（openai / ollama / manual）都暴露 max_input_tokens 且返回正整数（API 契约）。 +2. 系统依据 max_input_tokens 的 70% 自动判定压缩（should_compress）。 +""" +import os +import sys + +import pytest + +PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if PROJECT_ROOT not in sys.path: + sys.path.insert(0, PROJECT_ROOT) + + +def _engine_classes(): + from ai_engines.openai_engine import OpenaiAI + from ai_engines.ollama_engine import OllamaAI + from ai_engines.manual_engine import ManualAI + return {"openai": OpenaiAI, "ollama": OllamaAI, "manual": ManualAI} + + +def test_all_engines_expose_max_input_tokens(): + """所有引擎必须返回 max_input_tokens（统一 API 契约）""" + for name, cls in _engine_classes().items(): + assert hasattr(cls, "max_input_tokens"), f"{name} 引擎缺少 max_input_tokens 属性" + + +def test_max_input_tokens_positive_when_instantiated(): + """实例化后 max_input_tokens 应为正整数""" + for name, cls in _engine_classes().items(): + try: + inst = cls() + except Exception as e: + pytest.skip(f"{name} 引擎实例化失败（环境/网络），跳过实例检查: {e}") + val = inst.max_input_tokens + assert isinstance(val, int) and val > 0, f"{name} 的 max_input_tokens 非法: {val!r}" + + +def _mgr(): + from xcli_core.memory import MemoryManager + return MemoryManager(project_dir=PROJECT_ROOT) + + +def _conv(n): + return [{"role": "user", "content": f"msg {i}"} for i in range(n)] + + +def test_should_compress_token_aware_triggers_at_70pct(): + mgr = _mgr() + # 7000 / 10000 = 0.7 >= 0.7 → 触发 + assert mgr.should_compress(_conv(5), prompt_tokens=7000, max_input_tokens=10000, ratio=0.7) is True + + +def test_should_compress_token_aware_below_threshold(): + mgr = _mgr() + # 6000 / 10000 = 0.6 < 0.7 → 不触发（消息数也未超阈值） + assert mgr.should_compress(_conv(5), prompt_tokens=6000, max_input_tokens=10000, ratio=0.7) is False + + +def test_should_compress_message_count_fallback(): + mgr = _mgr() + # 无 token 信息，仅消息数 > 50 → 兜底触发 + assert mgr.should_compress(_conv(60)) is True + assert mgr.should_compress(_conv(10)) is False diff --git a/xcli_core/cli_core.py b/xcli_core/cli_core.py index 7200b9c..cdfb320 100644 --- a/xcli_core/cli_core.py +++ b/xcli_core/cli_core.py @@ -760,10 +760,11 @@ multi 操作支持一次修改多处： if hasattr(self, 'memory_manager') and self.memory_manager: engine = self.current_engine prompt_tokens = getattr(engine, 'last_prompt_tokens', None) if engine else None + # 统一 API：所有引擎必须返回 max_input_tokens（最大输入 token 数） ctx = None - if engine and hasattr(engine, 'context_window'): + if engine: try: - ctx = engine.context_window() + ctx = engine.max_input_tokens except Exception: ctx = None ratio = get_system_config('compress_ratio', 0.7) diff --git a/xcli_core/memory.py b/xcli_core/memory.py index 6f27ced..a9abc9c 100644 --- a/xcli_core/memory.py +++ b/xcli_core/memory.py @@ -307,18 +307,18 @@ class MemoryManager: except (TypeError, ValueError): pass - def should_compress(self, conversation: list, prompt_tokens=None, context_window=None, ratio=None) -> bool: + def should_compress(self, conversation: list, prompt_tokens=None, max_input_tokens=None, ratio=None) -> bool: """ 判断是否需要压缩上下文。 优先级: - 1. token 感知: prompt_tokens/context_window >= ratio 触发（OpenAI 格式可拿到真实 token 数） + 1. token 感知: prompt_tokens/max_input_tokens >= ratio 触发（引擎需返回 max_input_tokens） 2. 兜底: 消息数超过 compress_threshold """ if ratio is None: ratio = self.compress_ratio - if prompt_tokens and context_window: + if prompt_tokens and max_input_tokens: try: - if int(context_window) > 0 and (int(prompt_tokens) / int(context_window)) >= ratio: + if int(max_input_tokens) > 0 and (int(prompt_tokens) / int(max_input_tokens)) >= ratio: return True except (TypeError, ValueError): pass