Coverage for merco/core/llm/_client.py: 57%
284 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-07 14:04 +0800
1"""LLM 客户端 - OpenAI 兼容接口"""
2import asyncio
3import json
4import logging
5import re
6import tempfile
7import time
8from abc import ABC, abstractmethod
9from pathlib import Path
10from typing import Any, Optional, AsyncIterator
12logger = logging.getLogger("merco.llm")
15# 清理字符串中的代理对字符（surrogates），防止 API 序列化崩溃
16_SURROGATE_RE = re.compile(r'[\ud800-\udfff]')
19def _clean_surrogates(obj):
20 """递归清理数据结构中的代理对字符"""
21 if isinstance(obj, str):
22 return _SURROGATE_RE.sub('', obj)
23 if isinstance(obj, list):
24 return [_clean_surrogates(x) for x in obj]
25 if isinstance(obj, dict):
26 return {k: _clean_surrogates(v) for k, v in obj.items()}
27 return obj
30# ── Thinking 提取策略体系 ─────────────────────────────────────────
32# 统一的 think 标签对配置（开标签 → 闭标签）
33# 新增 provider 只需在此添加一对标签
34THINK_TAG_PAIRS: tuple[tuple[str, str], ...] = (
35 ("<think>", "[/think]"), # MiniMax 等
36 ("<think>", "</think>"), # 标准格式
37 ("<thinking>", "</thinking>"), # XML 格式
38)
41def _build_think_block_re() -> re.Pattern:
42 """根据 THINK_TAG_PAIRS 构建匹配所有 think 块的正则。
44 构造策略：
45 1. 收集所有模式（块 + 单标签）
46 2. 按"长度降序"排序——块（长）必须先于裸标签（短）
47 3. 用非捕获组组装，避免 sub() 因空捕获组错位
49 排序是关键：Python re 交替是"左优先"，短模式排在前面会"吞掉"块首字符，
50 导致长块模式失配。例如块模式 <think>...[/think] 在前时，遇到 <think>hi</think>
51 会先匹配 <think> 裸标签，块模式再也匹配不上。
52 """
53 raw_patterns = []
54 for open_tag, close_tag in THINK_TAG_PAIRS:
55 eo = re.escape(open_tag)
56 ec = re.escape(close_tag)
57 raw_patterns.append(f"{eo}.*?{ec}") # 完整块
58 raw_patterns.append(eo) # 孤儿开标签
59 raw_patterns.append(ec) # 孤儿闭标签
60 # 长度降序：块（最长）必须先于单标签
61 raw_patterns.sort(key=len, reverse=True)
62 return re.compile("|".join(f"(?:{p})" for p in raw_patterns), re.IGNORECASE | re.DOTALL)
65_THINK_BLOCK_RE = _build_think_block_re()
68def _strip_think_tags(text: str) -> str:
69 """chunk 安全：只去 think 标签，不动空白。
70 流式场景专用——每 chunk 调一次，.strip() 会破坏词边界。"""
71 return _THINK_BLOCK_RE.sub("", text)
74def _clean_content(text: str) -> str:
75 """完整 content 终态处理：去 think 标签 + 去前后空白。
76 非流式专用——完整响应只需调一次，strip 收尾合理。"""
77 return _THINK_BLOCK_RE.sub("", text).strip()
80def _extract_usage(response) -> dict:
81 """从 API 响应中提取 token 用量，包括缓存命中"""
82 usage = response.usage
83 if not usage:
84 return {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
86 result = {
87 "prompt_tokens": getattr(usage, "prompt_tokens", 0) or 0,
88 "completion_tokens": getattr(usage, "completion_tokens", 0) or 0,
89 "total_tokens": getattr(usage, "total_tokens", 0) or 0,
90 }
92 # 缓存命中 — 多 provider 兼容采集
93 # Anthropic: usage.cache_read_input_tokens / cache_creation_input_tokens
94 cached = getattr(usage, "cache_read_input_tokens", None)
95 if cached is not None:
96 result["cache_read_tokens"] = cached
97 result["cache_write_tokens"] = getattr(usage, "cache_creation_input_tokens", 0) or 0
99 # OpenAI: usage.prompt_tokens_details.cached_tokens
100 details = getattr(usage, "prompt_tokens_details", None)
101 if details is not None:
102 cd = getattr(details, "cached_tokens", None)
103 if cd is not None:
104 result["cached_tokens"] = cd
106 return result
109class ThinkingStrategy(ABC):
110 """思考内容提取策略基类。子类注册到 ThinkingExtractor 后按优先级调用。"""
112 @abstractmethod
113 def extract_from_delta(self, delta: Any) -> dict | None:
114 """从流式 delta 中提取。返回 {'content'?: str, 'reasoning'?: str} 或 None。"""
115 ...
117 @abstractmethod
118 def extract_from_message(self, message: Any) -> dict | None:
119 """从完整 message（非流式）中提取。"""
120 ...
122 def reset(self) -> None:
123 """重置跨 chunk 状态。"""
124 pass
127class DirectFieldStrategy(ThinkingStrategy):
128 """直接从对象属性检查 reasoning_content / reasoning（scnet 等代理放在顶层字段）。"""
130 def extract_from_delta(self, delta: Any) -> dict | None:
131 return self._check(delta)
133 def extract_from_message(self, message: Any) -> dict | None:
134 return self._check(message)
136 @staticmethod
137 def _check(obj: Any) -> dict | None:
138 try:
139 for attr in ("reasoning_content", "reasoning"):
140 val = getattr(obj, attr, None)
141 if val and isinstance(val, str):
142 return {"reasoning": val}
143 except Exception:
144 pass
145 return None
148class ModelExtraStrategy(ThinkingStrategy):
149 """从 model_extra 提取 reasoning_content / reasoning（OpenAI o1 / DeepSeek R1 等）。"""
151 def extract_from_delta(self, delta: Any) -> dict | None:
152 return self._extract_from(delta)
154 def extract_from_message(self, message: Any) -> dict | None:
155 return self._extract_from(message)
157 @staticmethod
158 def _extract_from(obj: Any) -> dict | None:
159 try:
160 extra = getattr(obj, "model_extra", None)
161 if isinstance(extra, dict):
162 rc = extra.get("reasoning_content") or extra.get("reasoning") or ""
163 if rc:
164 return {"reasoning": str(rc)}
165 except Exception:
166 pass
167 return None
170class ThinkTagStrategy(ThinkingStrategy):
171 """从 think 标签中提取思考内容。
173 使用统一的 THINK_TAG_PAIRS 配置，流式场景用状态机处理标签跨 chunk 的情况。
174 新增标签格式只需修改 THINK_TAG_PAIRS。
175 """
177 def __init__(self):
178 self._in_thinking = False
179 self._open_tag = ""
180 self._close_tag = ""
182 def reset(self) -> None:
183 self._in_thinking = False
184 self._open_tag = ""
185 self._close_tag = ""
187 def extract_from_delta(self, delta: Any) -> dict | None:
188 content = getattr(delta, "content", None) or ""
189 if not content:
190 return None
192 if self._in_thinking:
193 # 继续处理跨 chunk 的 think 块
194 if self._close_tag in content:
195 before_close, after_close = content.split(self._close_tag, 1)
196 self._in_thinking = False
197 result: dict = {}
198 if before_close:
199 result["reasoning"] = before_close
200 result["content"] = after_close
201 return result
202 else:
203 return {"reasoning": content, "content": ""}
204 else:
205 # 检测开标签（使用 THINK_TAG_PAIRS，优先匹配较长的）
206 for ot, ct in THINK_TAG_PAIRS:
207 if ot in content:
208 before_open, rest = content.split(ot, 1)
209 if ct in rest:
210 thinking, after_close = rest.split(ct, 1)
211 result: dict = {}
212 if thinking:
213 result["reasoning"] = thinking
214 result["content"] = before_open + after_close
215 return result
216 else:
217 # think 块跨 chunk，进入状态机
218 self._in_thinking = True
219 self._open_tag = ot
220 self._close_tag = ct
221 result: dict = {"reasoning": rest, "content": before_open}
222 return result
223 return {"content": content}
225 def extract_from_message(self, message: Any) -> dict | None:
226 """非流式：完整 content 一次性提取所有 think 块。"""
227 content = getattr(message, "content", None) or ""
228 for open_tag, close_tag in THINK_TAG_PAIRS:
229 if open_tag in content:
230 pattern = re.compile(
231 re.escape(open_tag) + r"(.*?)" + re.escape(close_tag),
232 re.DOTALL
233 )
234 thinking_parts = pattern.findall(content)
235 if thinking_parts:
236 cleaned = pattern.sub("", content).strip()
237 thinking = "\n\n".join(t.strip() for t in thinking_parts)
238 result: dict = {"reasoning": thinking}
239 if cleaned:
240 result["content"] = cleaned
241 return result
242 return None
245class ThinkingExtractor:
246 """策略链式思考提取器。按注册顺序尝试各策略，
247 首个返回非空 reasoning 的结果生效。"""
249 def __init__(self):
250 self._strategies: list[ThinkingStrategy] = [
251 DirectFieldStrategy(),
252 ModelExtraStrategy(),
253 ThinkTagStrategy(),
254 ]
256 def register(self, strategy: ThinkingStrategy) -> None:
257 """扩展点：注册新策略，插到链首。"""
258 self._strategies.insert(0, strategy)
260 def extract_from_delta(self, delta: Any) -> dict:
261 """流式块提取，返回 dict 可能含 content / reasoning。"""
262 for s in self._strategies:
263 result = s.extract_from_delta(delta)
264 if result is not None:
265 if "content" not in result:
266 raw = getattr(delta, "content", None) or ""
267 result["content"] = _strip_think_tags(raw)
268 return result
269 raw = getattr(delta, "content", None) or ""
270 return {"content": _strip_think_tags(raw)}
272 def extract_from_message(self, message: Any) -> dict:
273 """非流式响应提取。"""
274 for s in self._strategies:
275 result = s.extract_from_message(message)
276 if result is not None:
277 return result
278 return {}
280 def reset(self) -> None:
281 for s in self._strategies:
282 s.reset()
285class LLMClient:
286 """OpenAI 兼容的 LLM 客户端 — 纯传输层，重试/恢复由 Agent RecoveryPipeline 统一处理"""
288 def __init__(
289 self,
290 api_key: str,
291 model: str = "gpt-4",
292 base_url: Optional[str] = None,
293 temperature: float = 0.7,
294 max_tokens: int = 4096,
295 cooldown: float = 0, # 请求最小间隔（秒），0=不禁用
296 extra_params: Optional[dict] = None, # 额外 API 参数（top_p / seed 等）
297 headers: Optional[dict] = None, # 自定义 HTTP header（X-Title 等）
298 stream_options: Optional[dict] = None, # 流式额外参数（如 {"include_usage": True}）
299 ):
300 self.model = model
301 self.temperature = temperature
302 self.max_tokens = max_tokens
303 self.cooldown = cooldown
304 self.extra_params = extra_params or {}
305 self.stream_options = stream_options
306 self._last_request_time = 0.0
307 self._client_ready = False
309 import httpx
310 from openai import AsyncOpenAI
312 client_kwargs: dict = {
313 "api_key": api_key,
314 "max_retries": 0,
315 "timeout": httpx.Timeout(connect=10.0, read=None, write=None, pool=10.0),
316 }
317 if base_url:
318 client_kwargs["base_url"] = base_url
319 if headers:
320 client_kwargs["http_client"] = httpx.AsyncClient(headers=headers)
322 self.client = AsyncOpenAI(**client_kwargs)
324 # ── 公共方法 ─────────────────────────────────────────────────
326 async def chat(
327 self,
328 messages: list[dict],
329 tools: list[dict] = None,
330 tool_choice: str | dict[str, Any] = "auto",
331 ) -> dict:
332 """非流式聊天"""
333 params = self._build_params(messages, tools, tool_choice)
334 response = await self._request(params)
335 result = self._parse_response(response)
336 logger.debug("← API 响应: finish=%s content_len=%d tool_calls=%d",
337 result.get("finish_reason"), len(result.get("content", "")),
338 len(result.get("tool_calls", [])))
339 return result
341 async def chat_stream(
342 self,
343 messages: list[dict],
344 tools: list[dict] = None,
345 tool_choice: str | dict[str, Any] = "auto",
346 ) -> AsyncIterator[dict]:
347 """流式聊天"""
348 extractor = ThinkingExtractor()
349 params = self._build_params(messages, tools, tool_choice, stream=True)
350 stream = await self._request(params)
351 assert stream is not None
352 async for chunk in stream:
353 if parsed := self._parse_chunk(chunk, extractor):
354 yield parsed
356 # ── 私有方法 ─────────────────────────────────────────────────
358 def _build_params(
359 self, messages: list[dict], tools: list[dict] | None,
360 tool_choice: str | dict[str, Any], stream: bool = False,
361 ) -> dict:
362 """构造 API 请求参数"""
363 params = {
364 "model": self.model,
365 "messages": _clean_surrogates(messages),
366 "temperature": self.temperature,
367 "max_tokens": self.max_tokens,
368 }
369 if stream:
370 params["stream"] = True
371 options = {"include_usage": True}
372 if self.stream_options:
373 options.update(self.stream_options)
374 params["stream_options"] = options
375 if tools:
376 params["tools"] = tools
377 params["tool_choice"] = tool_choice
378 params.update(self.extra_params)
379 return params
381 async def _ensure_client_ready(self):
382 """确保 httpx 异步连接池初始化完成（首次请求前执行一次）。"""
383 if self._client_ready:
384 return
385 await asyncio.sleep(0)
386 self._client_ready = True
388 async def _request(self, params: dict):
389 """发送请求（含 cooldown），不重试。错误原样上抛交 Agent RecoveryPipeline。"""
390 logger.debug("→ API 请求: model=%s messages=%d tools=%d",
391 self.model, len(params.get("messages", [])),
392 len(params.get("tools", [])))
394 if self.cooldown > 0:
395 elapsed = time.monotonic() - self._last_request_time
396 if elapsed < self.cooldown:
397 wait = self.cooldown - elapsed
398 logger.debug("⏳ 冷却 %.1fs（距上次请求 %.1fs）", wait, elapsed)
399 await asyncio.sleep(wait)
401 try:
402 await self._ensure_client_ready()
403 response = await self.client.chat.completions.create(**params)
404 except asyncio.CancelledError:
405 logger.debug("⚠ 请求被取消")
406 raise
407 self._last_request_time = time.monotonic()
408 return response
410 @staticmethod
411 def _normalize_tool_calls(tool_calls: list) -> list[dict]:
412 """统一将 SDK tool_call 列表归一为安全 dict 列表，避免 None 穿透。"""
413 result = []
414 for tc in tool_calls:
415 entry: dict = {
416 "id": tc.id or "",
417 "name": tc.function.name or "",
418 "arguments": tc.function.arguments or "",
419 }
420 idx = getattr(tc, "index", None)
421 if idx is not None:
422 entry["index"] = idx
423 result.append(entry)
424 return result
426 def _parse_response(self, response) -> dict:
427 """解析完整响应"""
428 if not response.choices:
429 return {"content": "", "finish_reason": None, "usage": _extract_usage(response)}
431 choice = response.choices[0]
432 message = choice.message
434 content = message.content or ""
435 usage_data = _extract_usage(response)
436 result = {
437 "role": message.role,
438 "content": content,
439 "reasoning": "",
440 "finish_reason": choice.finish_reason,
441 "usage": usage_data,
442 }
444 extracted = ThinkingExtractor().extract_from_message(message)
445 if extracted:
446 if extracted.get("reasoning"):
447 result["reasoning"] = extracted["reasoning"]
448 if "content" in extracted:
449 result["content"] = extracted["content"] if extracted["content"] else ""
451 # 兜底：无论哪个策略命中，都清理 content 中残留的 <thinking> 标签
452 # 防止 DirectFieldStrategy 命中后 ThinkTagStrategy 没跑导致标签泄漏
453 result["content"] = _clean_content(result["content"])
455 if message.tool_calls:
456 normalized = self._normalize_tool_calls(message.tool_calls)
457 logger.debug("_parse_response: 检测到 %d 个 tool_calls: %s",
458 len(normalized),
459 [f"{tc['name']}({tc['id'][:8]})" for tc in normalized])
460 result["tool_calls"] = [
461 {**tc, "arguments": json.loads(tc["arguments"]) if tc["arguments"] else {}}
462 for tc in normalized
463 ]
465 return result
467 def _parse_chunk(self, chunk, extractor: Optional["ThinkingExtractor"] = None) -> Optional[dict]:
468 """解析流式块"""
469 choice = chunk.choices[0] if chunk.choices else None
470 if not choice:
471 return None
473 delta = choice.delta
474 if delta is None:
475 return None
476 result = {}
477 ex = extractor if extractor is not None else ThinkingExtractor()
478 extracted = ex.extract_from_delta(delta)
479 if extracted.get("content"):
480 result["content"] = _strip_think_tags(extracted["content"])
481 if extracted.get("reasoning"):
482 result["reasoning"] = _strip_think_tags(extracted["reasoning"])
483 if delta.tool_calls:
484 result["tool_calls"] = self._normalize_tool_calls(delta.tool_calls)
485 tcs = result["tool_calls"]
486 logger.debug("_parse_chunk: 流式收到 %d 个 tool_calls delta: index=%s name=%s",
487 len(tcs),
488 [tc.get("index", "?") for tc in tcs],
489 [tc.get("name", "") or "?" for tc in tcs])
491 if choice.finish_reason:
492 result["finish_reason"] = choice.finish_reason
493 if hasattr(chunk, "usage") and chunk.usage:
494 result["usage"] = {
495 "prompt_tokens": chunk.usage.prompt_tokens,
496 "completion_tokens": chunk.usage.completion_tokens,
497 "total_tokens": chunk.usage.total_tokens,
498 }
499 return result if result else None