++
This commit is contained in:
+287
-165
@@ -1,165 +1,287 @@
|
||||
from __future__ import annotations
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMClient:
|
||||
def __init__(self, api_key: str, base_url: str, model: str, max_concurrent: int = 10):
|
||||
self.client = AsyncOpenAI(api_key=api_key, base_url=base_url)
|
||||
self.model = model
|
||||
self.semaphore = asyncio.Semaphore(max_concurrent)
|
||||
self.total_tokens = 0
|
||||
self.total_calls = 0
|
||||
self._lock = asyncio.Lock()
|
||||
|
||||
async def _track(self, usage):
|
||||
if usage:
|
||||
async with self._lock:
|
||||
self.total_tokens += usage.total_tokens or 0
|
||||
self.total_calls += 1
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
system: str,
|
||||
user: str,
|
||||
json_mode: bool = False,
|
||||
max_tokens: int = 1000,
|
||||
retries: int = 3,
|
||||
) -> str:
|
||||
kwargs: dict = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": user},
|
||||
],
|
||||
"max_tokens": max_tokens,
|
||||
"temperature": 0.8,
|
||||
}
|
||||
if json_mode:
|
||||
kwargs["response_format"] = {"type": "json_object"}
|
||||
|
||||
for attempt in range(retries):
|
||||
try:
|
||||
async with self.semaphore:
|
||||
response = await self.client.chat.completions.create(**kwargs)
|
||||
await self._track(response.usage)
|
||||
content = response.choices[0].message.content or ""
|
||||
if not content:
|
||||
finish = getattr(response.choices[0], 'finish_reason', 'unknown')
|
||||
logger.warning("LLM returned empty content (finish_reason=%s)", finish)
|
||||
return content.strip()
|
||||
except Exception as e:
|
||||
if attempt == retries - 1:
|
||||
logger.error("LLM call failed after %d retries: %s", retries, e)
|
||||
raise
|
||||
wait = 2 ** attempt
|
||||
logger.warning("LLM call failed (attempt %d), retrying in %ds: %s", attempt + 1, wait, e)
|
||||
await asyncio.sleep(wait)
|
||||
return ""
|
||||
|
||||
async def generate_batch(
|
||||
self,
|
||||
prompts: list[tuple[str, str]],
|
||||
json_mode: bool = False,
|
||||
max_tokens: int = 1000,
|
||||
) -> list[str]:
|
||||
tasks = [self.generate(system, user, json_mode=json_mode, max_tokens=max_tokens) for system, user in prompts]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
return [r if isinstance(r, str) else "" for r in results]
|
||||
|
||||
async def generate_light(self, system: str, user: str, max_tokens: int = 100) -> str:
|
||||
return await self.generate(system, user, json_mode=False, max_tokens=max_tokens)
|
||||
|
||||
def get_stats(self) -> dict:
|
||||
return {
|
||||
"total_tokens": self.total_tokens,
|
||||
"total_calls": self.total_calls,
|
||||
}
|
||||
|
||||
|
||||
def parse_json(text: str) -> dict:
|
||||
text = text.strip()
|
||||
|
||||
# Strip modern LLM "thinking" blocks (Anthropic, DeepSeek, etc.) that may
|
||||
# appear either before or around the JSON payload.
|
||||
text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
|
||||
|
||||
# Markdown code fence: ```json ... ``` or ``` ... ```
|
||||
if text.startswith("```"):
|
||||
match = re.search(r"```(?:json)?\s*\n?(.*?)```", text, re.DOTALL)
|
||||
if match:
|
||||
text = match.group(1).strip()
|
||||
|
||||
# Try the whole text first.
|
||||
try:
|
||||
return json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Otherwise find the first balanced { ... } or [ ... ] substring that parses.
|
||||
def _balanced_at(start: int) -> int | None:
|
||||
open_ch, other_ch = ("{", "}") if text[start] == "{" else (("[", "]") if text[start] == "[" else (None, None))
|
||||
if open_ch is None:
|
||||
return None
|
||||
depth = 0
|
||||
in_str = False
|
||||
esc = False
|
||||
for j in range(start, len(text)):
|
||||
c = text[j]
|
||||
if in_str:
|
||||
if esc:
|
||||
esc = False
|
||||
elif c == "\\":
|
||||
esc = True
|
||||
elif c == '"':
|
||||
in_str = False
|
||||
continue
|
||||
if c == '"':
|
||||
in_str = True
|
||||
elif c == open_ch:
|
||||
depth += 1
|
||||
elif c == other_ch:
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return j
|
||||
return None
|
||||
|
||||
for opener in ("{", "["):
|
||||
i = 0
|
||||
while True:
|
||||
idx = text.find(opener, i)
|
||||
if idx == -1:
|
||||
break
|
||||
end = _balanced_at(idx)
|
||||
if end is not None:
|
||||
try:
|
||||
return json.loads(text[idx : end + 1])
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
i = idx + 1
|
||||
|
||||
# Last-ditch: trim from the end while we can find a parseable prefix.
|
||||
for i in range(len(text) - 1, -1, -1):
|
||||
if text[i] in "]}":
|
||||
try:
|
||||
return json.loads(text[: i + 1])
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
open_braces = text.count("{") - text.count("}")
|
||||
open_brackets = text.count("[") - text.count("]")
|
||||
repaired = text + ("}" * max(0, open_braces)) + ("]" * max(0, open_brackets))
|
||||
try:
|
||||
return json.loads(repaired)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
logger.error("Failed to parse JSON from LLM response: %s...", text[:200])
|
||||
return {}
|
||||
from __future__ import annotations
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from openai import AsyncOpenAI
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LLMClient:
|
||||
def __init__(self, api_key: str, base_url: str, model: str, max_concurrent: int = 10):
|
||||
self.client = AsyncOpenAI(api_key=api_key, base_url=base_url)
|
||||
self.model = model
|
||||
self._is_minimax = "minimax" in (base_url or "").lower()
|
||||
self.semaphore = asyncio.Semaphore(max_concurrent)
|
||||
self.total_tokens = 0
|
||||
self.total_calls = 0
|
||||
self._lock = asyncio.Lock()
|
||||
# Per-purpose call stats — useful when a sim stalls and you need to
|
||||
# know which LLM call category is slow. Reset by the engine at the
|
||||
# start of each sim if you want per-sim stats.
|
||||
self._purpose_stats: dict[str, dict[str, float]] = {}
|
||||
|
||||
async def _track(self, usage):
|
||||
if usage:
|
||||
async with self._lock:
|
||||
self.total_tokens += usage.total_tokens or 0
|
||||
self.total_calls += 1
|
||||
|
||||
async def _record_purpose(self, purpose: str | None, elapsed: float, ok: bool):
|
||||
if not purpose:
|
||||
return
|
||||
async with self._lock:
|
||||
st = self._purpose_stats.setdefault(
|
||||
purpose, {"calls": 0, "errors": 0, "total_ms": 0.0}
|
||||
)
|
||||
st["calls"] += 1
|
||||
st["total_ms"] += elapsed * 1000
|
||||
if not ok:
|
||||
st["errors"] += 1
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
system: str,
|
||||
user: str,
|
||||
json_mode: bool = False,
|
||||
max_tokens: int = 4096,
|
||||
retries: int = 3,
|
||||
purpose: str | None = None,
|
||||
) -> str:
|
||||
kwargs: dict = {
|
||||
"model": self.model,
|
||||
"messages": [
|
||||
{"role": "system", "content": system},
|
||||
{"role": "user", "content": user},
|
||||
],
|
||||
"temperature": 0.5 if json_mode else 1.0,
|
||||
}
|
||||
if self._is_minimax:
|
||||
# MiniMax M3 enables thinking by default. In structured-output
|
||||
# calls that can consume the completion budget before a visible
|
||||
# message.content is produced. MiniMax-specific fields must be
|
||||
# sent through extra_body when using the OpenAI Python SDK.
|
||||
kwargs["max_completion_tokens"] = max_tokens
|
||||
kwargs["extra_body"] = {"thinking": {"type": "disabled"}}
|
||||
else:
|
||||
kwargs["max_tokens"] = max_tokens
|
||||
|
||||
if json_mode and not self._is_minimax:
|
||||
# MiniMax's newer M-series models do not reliably implement the
|
||||
# OpenAI json_object response_format. They still receive the
|
||||
# explicit JSON schema in the prompt and are parsed defensively.
|
||||
kwargs["response_format"] = {"type": "json_object"}
|
||||
|
||||
t0 = time.monotonic()
|
||||
for attempt in range(retries):
|
||||
try:
|
||||
async with self.semaphore:
|
||||
response = await self.client.chat.completions.create(**kwargs)
|
||||
await self._track(response.usage)
|
||||
if not response.choices:
|
||||
logger.warning(
|
||||
"LLM 返回空 choices（用途=%s，响应ID=%s，模型=%s）",
|
||||
purpose or "-", getattr(response, "id", "-"), self.model,
|
||||
)
|
||||
content = ""
|
||||
finish = "no_choices"
|
||||
message = None
|
||||
else:
|
||||
choice = response.choices[0]
|
||||
message = choice.message
|
||||
content = getattr(message, "content", None) or ""
|
||||
finish = getattr(choice, "finish_reason", "unknown")
|
||||
|
||||
# A few OpenAI-compatible providers return content parts
|
||||
# instead of a plain string. Normalize text parts here.
|
||||
if isinstance(content, list):
|
||||
content = "".join(
|
||||
part.get("text", "") if isinstance(part, dict)
|
||||
else getattr(part, "text", "") or ""
|
||||
for part in content
|
||||
)
|
||||
|
||||
if not content:
|
||||
reasoning = getattr(message, "reasoning_content", None) if message else None
|
||||
reasoning_details = getattr(message, "reasoning_details", None) if message else None
|
||||
refusal = getattr(message, "refusal", None) if message else None
|
||||
usage = getattr(response, "usage", None)
|
||||
logger.warning(
|
||||
"LLM 返回空内容（用途=%s，结束原因=%s，完成令牌=%s，推理内容长度=%s，推理详情=%s，拒答=%s）",
|
||||
purpose or "-",
|
||||
finish,
|
||||
getattr(usage, "completion_tokens", "?"),
|
||||
len(reasoning or ""),
|
||||
bool(reasoning_details),
|
||||
bool(refusal),
|
||||
)
|
||||
if attempt < retries - 1:
|
||||
wait = 2 ** attempt
|
||||
logger.warning(
|
||||
"LLM 返回空内容，将在%d秒后重试（用途=%s，第%d/%d次）",
|
||||
wait, purpose or "-", attempt + 1, retries,
|
||||
)
|
||||
await asyncio.sleep(wait)
|
||||
continue
|
||||
await self._record_purpose(purpose, time.monotonic() - t0, ok=False)
|
||||
return ""
|
||||
await self._record_purpose(purpose, time.monotonic() - t0, ok=True)
|
||||
if purpose:
|
||||
logger.debug(
|
||||
"LLM 调用成功（用途=%s，令牌数=%s，耗时=%.2f秒）",
|
||||
purpose,
|
||||
getattr(response.usage, "total_tokens", "?"),
|
||||
time.monotonic() - t0,
|
||||
)
|
||||
return content.strip()
|
||||
except Exception as e:
|
||||
if attempt == retries - 1:
|
||||
await self._record_purpose(purpose, time.monotonic() - t0, ok=False)
|
||||
logger.error(
|
||||
"LLM 调用重试 %d 次后仍失败（用途=%s，耗时=%.2f秒）：%s",
|
||||
retries, purpose or "-", time.monotonic() - t0, e,
|
||||
)
|
||||
raise
|
||||
wait = 2 ** attempt
|
||||
logger.warning(
|
||||
"LLM 调用失败（用途=%s，第%d次尝试，%d秒后重试，已耗时=%.2f秒）：%s",
|
||||
purpose or "-", attempt + 1, wait, time.monotonic() - t0, e,
|
||||
)
|
||||
await asyncio.sleep(wait)
|
||||
return ""
|
||||
|
||||
async def generate_batch(
|
||||
self,
|
||||
prompts: list[tuple[str, str]],
|
||||
json_mode: bool = False,
|
||||
max_tokens: int = 1000,
|
||||
purpose: str | None = None,
|
||||
sub_purpose: str | None = None,
|
||||
) -> list[str]:
|
||||
# `purpose` is the batch label (e.g. "decision"); `sub_purpose` lets
|
||||
# us split stats by an inner dimension (e.g. "active/react").
|
||||
tasks = [
|
||||
self.generate(
|
||||
system, user,
|
||||
json_mode=json_mode, max_tokens=max_tokens,
|
||||
purpose=sub_purpose or purpose,
|
||||
)
|
||||
for system, user in prompts
|
||||
]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
return [r if isinstance(r, str) else "" for r in results]
|
||||
|
||||
async def generate_light(self, system: str, user: str, max_tokens: int = 100) -> str:
|
||||
return await self.generate(system, user, json_mode=False, max_tokens=max_tokens)
|
||||
|
||||
def get_stats(self) -> dict:
|
||||
return {
|
||||
"total_tokens": self.total_tokens,
|
||||
"total_calls": self.total_calls,
|
||||
"by_purpose": dict(self._purpose_stats),
|
||||
}
|
||||
|
||||
|
||||
def parse_json(text: str, purpose: str | None = None) -> Any:
|
||||
"""Parse a model response that was requested as JSON.
|
||||
|
||||
Providers occasionally return prose, truncated JSON, or a markdown code
|
||||
block even when ``response_format=json_object`` was requested. Keep the
|
||||
permissive recovery logic, but make a final failure actionable: include
|
||||
the call purpose, response length, and a bounded raw preview in the
|
||||
simulation log.
|
||||
"""
|
||||
text = (text or "").strip()
|
||||
label = purpose or "未标注用途"
|
||||
|
||||
if not text:
|
||||
logger.error("JSON 解析失败（用途=%s）：LLM 返回内容为空", label)
|
||||
return {}
|
||||
|
||||
# Strip modern LLM "thinking" blocks (Anthropic, DeepSeek, etc.) that may
|
||||
# appear either before or around the JSON payload.
|
||||
text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
|
||||
|
||||
# Markdown code fence: ```json ... ``` or ``` ... ```
|
||||
if text.startswith("```"):
|
||||
match = re.search(r"```(?:json)?\s*\n?(.*?)```", text, re.DOTALL)
|
||||
if match:
|
||||
text = match.group(1).strip()
|
||||
|
||||
# Try the whole text first.
|
||||
try:
|
||||
return json.loads(text)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
# Otherwise find the first balanced { ... } or [ ... ] substring that parses.
|
||||
def _balanced_at(start: int) -> int | None:
|
||||
open_ch, other_ch = ("{", "}") if text[start] == "{" else (("[", "]") if text[start] == "[" else (None, None))
|
||||
if open_ch is None:
|
||||
return None
|
||||
depth = 0
|
||||
in_str = False
|
||||
esc = False
|
||||
for j in range(start, len(text)):
|
||||
c = text[j]
|
||||
if in_str:
|
||||
if esc:
|
||||
esc = False
|
||||
elif c == "\\":
|
||||
esc = True
|
||||
elif c == '"':
|
||||
in_str = False
|
||||
continue
|
||||
if c == '"':
|
||||
in_str = True
|
||||
elif c == open_ch:
|
||||
depth += 1
|
||||
elif c == other_ch:
|
||||
depth -= 1
|
||||
if depth == 0:
|
||||
return j
|
||||
return None
|
||||
|
||||
for opener in ("{", "["):
|
||||
i = 0
|
||||
while True:
|
||||
idx = text.find(opener, i)
|
||||
if idx == -1:
|
||||
break
|
||||
end = _balanced_at(idx)
|
||||
if end is not None:
|
||||
try:
|
||||
return json.loads(text[idx : end + 1])
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
i = idx + 1
|
||||
|
||||
# Last-ditch: trim from the end while we can find a parseable prefix.
|
||||
for i in range(len(text) - 1, -1, -1):
|
||||
if text[i] in "]}":
|
||||
try:
|
||||
return json.loads(text[: i + 1])
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
open_braces = text.count("{") - text.count("}")
|
||||
open_brackets = text.count("[") - text.count("]")
|
||||
repaired = text + ("}" * max(0, open_braces)) + ("]" * max(0, open_brackets))
|
||||
try:
|
||||
return json.loads(repaired)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
preview = text[:1000].replace("\r", "\\r").replace("\n", "\\n")
|
||||
logger.error(
|
||||
"JSON 解析失败（用途=%s，返回长度=%d）：内容不是有效 JSON 或已被截断；原始返回片段=%s%s",
|
||||
label,
|
||||
len(text),
|
||||
preview,
|
||||
"……" if len(text) > 1000 else "",
|
||||
)
|
||||
return {}
|
||||
|
||||
Reference in New Issue
Block a user