@baimingtao/lbs-agent 1.0.0

package/runtime/core/llm/think_scrubber.py

@@ -0,0 +1,232 @@
+"""
+Think Scrubber —— 流式 thinking/reasoning block 过滤器。
+
+对标 LBS think_scrubber.py，精简实现：
+  1. 流式过滤：逐 chunk 处理，状态机驱动
+  2. 非流式过滤：对完整文本做一次性清除
+  3. 支持多种 tag：<think>/<thinking>/<reasoning>/<thought>
+  4. 边界安全：partial tag 缓存，避免误删正文
+
+用法：
+  # 非流式
+  clean = ThinkScrubber.strip(text)
+
+  # 流式
+  scrubber = StreamingThinkScrubber()
+  for delta in stream:
+      visible = scrubber.feed(delta)
+      if visible:
+          emit(visible)
+  tail = scrubber.flush()
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Tuple
+
+__all__ = ["ThinkScrubber", "StreamingThinkScrubber"]
+
+
+# 匹配的 tag 名称（大小写不敏感）
+_THINK_TAGS = ("think", "thinking", "reasoning", "thought", "REASONING_SCRATCHPAD")
+
+
+class ThinkScrubber:
+    """非流式 thinking block 过滤器。"""
+
+    # 匹配完整的 <think>...</think> 对（含变体）
+    _PATTERN = re.compile(
+        r"<(?:" + "|".join(_THINK_TAGS) + r")\b[^>]*>.*?</(?:" + "|".join(_THINK_TAGS) + r")>",
+        re.IGNORECASE | re.DOTALL,
+    )
+
+    @staticmethod
+    def strip(text: str) -> str:
+        """从完整文本中移除所有 thinking block。"""
+        if not text:
+            return text
+        # 移除闭合对
+        result = ThinkScrubber._PATTERN.sub("", text)
+        # 移除未闭合的开启标签（<think> 后面没有 </think>）
+        result = re.sub(
+            r"<(?:" + "|".join(_THINK_TAGS) + r")\b[^>]*>.*$",
+            "",
+            result,
+            flags=re.IGNORECASE | re.DOTALL,
+        )
+        # 清理多余空行
+        result = re.sub(r"\n{3,}", "\n\n", result).strip()
+        return result
+
+    @staticmethod
+    def has_think_blocks(text: str) -> bool:
+        """检查文本是否包含 thinking block。"""
+        if not text:
+            return False
+        return bool(re.search(r"<(?:" + "|".join(_THINK_TAGS) + r")\b", text, re.IGNORECASE))
+
+
+class StreamingThinkScrubber:
+    """流式 thinking block 过滤器。
+
+    状态机：
+      - _in_block: True 时在 block 内部，丢弃所有内容
+      - _buf: 持有 partial tag 的尾部，下个 delta 决定是 tag 还是正文
+
+    每个 chunk 通过 feed() 输入，返回应该展示给用户的部分。
+    流结束时调用 flush() 获取缓冲区剩余。
+    """
+
+    _MAX_TAG_LEN = 30  # 最长 tag 名 + 尖括号
+
+    def __init__(self):
+        self._in_block = False
+        self._buf = ""
+        self._last_emit_newline = True  # 流开始等价于行首
+
+    def reset(self):
+        """重置状态（每个新 turn 开头调用）。"""
+        self._in_block = False
+        self._buf = ""
+        self._last_emit_newline = True
+
+    def feed(self, delta: str) -> str:
+        """处理一个 chunk，返回应该展示的内容。"""
+        if not delta:
+            return ""
+
+        # 合并缓冲区和新 delta
+        text = self._buf + delta
+        self._buf = ""
+        output = []
+
+        i = 0
+        while i < len(text):
+            if self._in_block:
+                # 在 block 内部：寻找关闭标签
+                close_idx = self._find_close_tag(text, i)
+                if close_idx >= 0:
+                    # 找到关闭标签
+                    tag_end = text.index(">", close_idx) + 1
+                    self._in_block = False
+                    i = tag_end
+                    # 关闭标签后面可能有换行，吃掉它
+                    if i < len(text) and text[i] == "\n":
+                        i += 1
+                else:
+                    # 没找到关闭标签，保留剩余到缓冲区
+                    # 但检查是否是 partial close tag
+                    remaining = text[i:]
+                    partial = self._extract_partial_close(remaining)
+                    if partial:
+                        self._buf = partial
+                        output_text = text[i:len(text) - len(partial)]
+                        # 这部分在 block 内，丢弃
+                    # 全部在 block 内，丢弃
+                    break  # 剩余都在 block 内
+            else:
+                # 在 block 外部：寻找开启标签
+                open_idx = self._find_open_tag(text, i)
+                if open_idx >= 0:
+                    # 找到开启标签
+                    # 先输出 open_idx 之前的文本
+                    before = text[i:open_idx]
+                    if before:
+                        output.append(before)
+                        self._last_emit_newline = before.endswith("\n")
+                    # 进入 block
+                    tag_end = text.index(">", open_idx) + 1
+                    self._in_block = True
+                    i = tag_end
+                else:
+                    # 没找到开启标签，检查尾部是否是 partial open tag
+                    remaining = text[i:]
+                    partial = self._extract_partial_open(remaining)
+                    if partial:
+                        # 输出 partial 之前的部分
+                        before_partial = text[i:len(text) - len(partial)]
+                        if before_partial:
+                            output.append(before_partial)
+                            self._last_emit_newline = before_partial.endswith("\n")
+                        self._buf = partial
+                    else:
+                        # 没有 partial，全部输出
+                        output.append(remaining)
+                        self._last_emit_newline = remaining.endswith("\n")
+                    break
+
+        result = "".join(output)
+        return result
+
+    def flush(self) -> str:
+        """流结束时调用，返回缓冲区剩余内容。"""
+        result = ""
+        if self._buf:
+            if self._in_block:
+                # 在 block 内部的缓冲，丢弃（不完整的 block）
+                pass
+            else:
+                # 不在 block 内，缓冲是 partial tag 但最终不是 tag
+                result = self._buf
+        self._buf = ""
+        self._in_block = False
+        return result
+
+    # ── 内部方法 ─────────────────────────────────────
+
+    def _find_open_tag(self, text: str, start: int) -> int:
+        """在 text[start:] 中寻找 thinking block 开启标签。
+
+        只有在行首（或文本开头）才认为是 block 开启。
+        """
+        for tag in _THINK_TAGS:
+            pattern = f"<{tag}"
+            idx = text.lower().find(pattern.lower(), start)
+            if idx >= 0:
+                # 检查是否在行首（前面只有空白或文本开头）
+                if idx == start and self._last_emit_newline:
+                    return idx
+                # 检查前面是否有换行
+                line_start = text.rfind("\n", start, idx)
+                if line_start >= 0:
+                    between = text[line_start + 1:idx]
+                    if between.strip() == "":
+                        return idx
+                elif idx == 0:
+                    return idx
+        return -1
+
+    def _find_close_tag(self, text: str, start: int) -> int:
+        """在 text[start:] 中寻找关闭标签。"""
+        for tag in _THINK_TAGS:
+            pattern = f"</{tag}"
+            idx = text.lower().find(pattern.lower(), start)
+            if idx >= 0:
+                return idx
+        return -1
+
+    def _extract_partial_open(self, text: str) -> str:
+        """检查文本尾部是否是 partial open tag（如 "<thi"）。"""
+        if not text or text[-1] != ">" and text[-1] != "\n":
+            # 检查尾部是否是 < 后跟部分 tag 名
+            for length in range(min(len(text), self._MAX_TAG_LEN), 0, -1):
+                tail = text[-length:]
+                if tail.startswith("<"):
+                    # 可能是 partial tag
+                    tail_lower = tail.lower()
+                    for tag in _THINK_TAGS:
+                        if tag.startswith(tail_lower[1:]) or tail_lower[1:] == tag[:len(tail_lower) - 1]:
+                            return tail
+        return ""
+
+    def _extract_partial_close(self, text: str) -> str:
+        """检查文本尾部是否是 partial close tag。"""
+        for length in range(min(len(text), self._MAX_TAG_LEN), 0, -1):
+            tail = text[-length:]
+            if tail.startswith("</"):
+                tail_lower = tail.lower()
+                for tag in _THINK_TAGS:
+                    if tag.startswith(tail_lower[2:]) or tail_lower[2:] == tag[:len(tail_lower) - 2]:
+                        return tail
+        return ""

package/runtime/core/llm/tokenizer.py

@@ -0,0 +1,150 @@
+"""Token 精确计数。
+
+封装 tiktoken，按模型选择最佳编码器。fallback 到近似计数（1 char approx 0.25 token）。
+集成到 ContextManager 用于精确窗口管理 + 成本预估。
+"""
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+_MODEL_TO_ENCODING = {
+    "gpt-4": "cl100k_base",
+    "gpt-4-32k": "cl100k_base",
+    "gpt-4-turbo": "cl100k_base",
+    "gpt-4o": "o200k_base",
+    "gpt-4o-mini": "o200k_base",
+    "gpt-4.1": "o200k_base",
+    "gpt-4.1-mini": "o200k_base",
+    "gpt-4.1-nano": "o200k_base",
+    "gpt-5": "o200k_base",
+    "gpt-3.5-turbo": "cl100k_base",
+    "gpt-3.5": "cl100k_base",
+    "text-davinci-003": "p50k_base",
+    "claude": "cl100k_base",
+    "claude-3": "cl100k_base",
+    "claude-3.5": "cl100k_base",
+    "claude-4": "cl100k_base",
+    "deepseek": "cl100k_base",
+    "deepseek-v2": "cl100k_base",
+    "deepseek-v3": "cl100k_base",
+    "deepseek-r1": "cl100k_base",
+    "llama": "cl100k_base",
+    "llama-3": "cl100k_base",
+    "mistral": "cl100k_base",
+    "mixtral": "cl100k_base",
+    "qwen": "cl100k_base",
+    "qwen2": "cl100k_base",
+    "glm": "cl100k_base",
+    "gemini": "cl100k_base",
+    "gemma": "cl100k_base",
+    "yi": "cl100k_base",
+    "minmax": "cl100k_base",
+}
+
+
+class TokenCounter:
+    """Token 精确计数。"""
+
+    def __init__(self):
+        self._encoders = {}
+        self._has_tiktoken = False
+        self._warned = False
+        self._try_import()
+
+    def _try_import(self):
+        try:
+            import tiktoken as _tiktoken_mod
+            self._tiktoken = _tiktoken_mod
+            self._has_tiktoken = True
+        except ImportError:
+            if not self._warned:
+                logger.warning("tiktoken 未安装，使用近似计数")
+                self._warned = True
+
+    def _resolve_encoding(self, model):
+        ml = model.lower()
+        if ml in _MODEL_TO_ENCODING:
+            return _MODEL_TO_ENCODING[ml]
+        for key, enc in _MODEL_TO_ENCODING.items():
+            if ml.startswith(key):
+                return enc
+        return "cl100k_base"
+
+    def _get_encoder(self, model):
+        enc_name = self._resolve_encoding(model)
+        if not self._has_tiktoken:
+            return None
+        if enc_name not in self._encoders:
+            try:
+                self._encoders[enc_name] = self._tiktoken.get_encoding(enc_name)
+            except Exception:
+                logger.warning("获取编码 %s 失败", enc_name)
+                return None
+        return self._encoders[enc_name]
+
+    def count(self, text: str, model: str = "claude-4") -> int:
+        if not text:
+            return 0
+        enc = self._get_encoder(model)
+        if enc is not None:
+            try:
+                return len(enc.encode(text))
+            except Exception:
+                pass
+        return int(len(text) * 0.25) + 1
+
+    def count_messages(self, messages: list, model: str = "claude-4") -> int:
+        total = 0
+        for msg in messages:
+            total += 4
+            content = msg.get("content", "")
+            if isinstance(content, list):
+                for part in content:
+                    if isinstance(part, dict):
+                        total += self.count(part.get("text", "") or "", model)
+                    elif isinstance(part, str):
+                        total += self.count(part, model)
+            elif isinstance(content, str):
+                total += self.count(content, model)
+            total += len(msg.get("role", "user"))
+        total += 2
+        return total
+
+    def truncate(self, text: str, max_tokens: int, model: str = "claude-4") -> str:
+        if not text:
+            return text
+        enc = self._get_encoder(model)
+        if enc is not None:
+            try:
+                tokens = enc.encode(text)
+                if len(tokens) <= max_tokens:
+                    return text
+                return enc.decode(tokens[:max_tokens])
+            except Exception:
+                pass
+        return text[:int(max_tokens * 4)]
+
+    def estimate_cost(self, text: str, model: str = "claude-4",
+                      cost_input_per_m: float = 3.0,
+                      cost_output_per_m: float = 15.0):
+        tokens = self.count(text, model)
+        in_cost = tokens / 1_000_000 * cost_input_per_m
+        out_cost = tokens / 1_000_000 * cost_output_per_m
+        return {
+            "tokens": tokens,
+            "input_cost_est": round(in_cost, 6),
+            "output_cost_est": round(out_cost, 6),
+            "total_est": round(in_cost + out_cost, 6),
+        }
+
+
+_counter = None
+
+
+def get_counter():
+    global _counter
+    if _counter is None:
+        _counter = TokenCounter()
+    return _counter