243 lines
9.3 KiB
Python
243 lines
9.3 KiB
Python
from __future__ import annotations
|
||
|
||
from .shared import *
|
||
|
||
class AIStewardJsonParsingMixin:
|
||
@staticmethod
|
||
def _build_project_creation_token_ladder() -> list[int]:
|
||
budgets = [
|
||
PROJECT_CREATION_STEP_INITIAL_MAX_TOKENS,
|
||
4_000,
|
||
8_000,
|
||
16_000,
|
||
32_000,
|
||
64_000,
|
||
PROJECT_CREATION_STEP_MAX_TOKENS,
|
||
]
|
||
ordered: list[int] = []
|
||
for item in budgets:
|
||
value = max(256, min(int(item), PROJECT_CREATION_STEP_MAX_TOKENS))
|
||
if value not in ordered:
|
||
ordered.append(value)
|
||
return ordered
|
||
|
||
@staticmethod
|
||
def _is_provider_budget_rejection(exc: Exception) -> bool:
|
||
lowered = str(exc or "").lower()
|
||
return "400 bad request" in lowered and any(
|
||
token in lowered
|
||
for token in ("max_tokens", "context", "length", "token", "dashscope")
|
||
)
|
||
|
||
@staticmethod
|
||
def _strip_code_fence(text: str) -> str:
|
||
raw = str(text or "").strip()
|
||
if not raw.startswith("```"):
|
||
return raw
|
||
chunks = raw.split("```")
|
||
if len(chunks) < 2:
|
||
return raw
|
||
candidate = chunks[1].strip()
|
||
if candidate.lower().startswith("json"):
|
||
candidate = candidate[4:].strip()
|
||
return candidate
|
||
|
||
@staticmethod
|
||
def _extract_balanced_json_segment(text: str, open_char: str, close_char: str) -> str:
|
||
raw = str(text or "")
|
||
start = raw.find(open_char)
|
||
if start < 0:
|
||
return ""
|
||
depth = 0
|
||
in_str = False
|
||
escaped = False
|
||
for idx in range(start, len(raw)):
|
||
ch = raw[idx]
|
||
if in_str:
|
||
if escaped:
|
||
escaped = False
|
||
elif ch == "\\":
|
||
escaped = True
|
||
elif ch == '"':
|
||
in_str = False
|
||
continue
|
||
if ch == '"':
|
||
in_str = True
|
||
continue
|
||
if ch == open_char:
|
||
depth += 1
|
||
elif ch == close_char:
|
||
depth -= 1
|
||
if depth == 0:
|
||
return raw[start : idx + 1].strip()
|
||
return ""
|
||
|
||
def _extract_json_candidate(self, text: str) -> str:
|
||
raw = self._strip_code_fence(text)
|
||
if raw.startswith("{") or raw.startswith("["):
|
||
return raw
|
||
object_candidate = self._extract_balanced_json_segment(raw, "{", "}")
|
||
array_candidate = self._extract_balanced_json_segment(raw, "[", "]")
|
||
if object_candidate and array_candidate:
|
||
return object_candidate if raw.find("{") <= raw.find("[") else array_candidate
|
||
return object_candidate or array_candidate or raw
|
||
|
||
@staticmethod
|
||
def _repair_json_text(raw: str) -> str:
|
||
text = str(raw or "")
|
||
if not text:
|
||
return text
|
||
buf: list[str] = []
|
||
in_str = False
|
||
escaped = False
|
||
for ch in text:
|
||
if in_str:
|
||
if escaped:
|
||
escaped = False
|
||
elif ch == "\\":
|
||
escaped = True
|
||
elif ch == '"':
|
||
in_str = False
|
||
elif ch == "\n":
|
||
buf.append("\\n")
|
||
continue
|
||
elif ch == "\r":
|
||
buf.append("\\r")
|
||
continue
|
||
elif ch == "\t":
|
||
buf.append("\\t")
|
||
continue
|
||
else:
|
||
if ch == '"':
|
||
in_str = True
|
||
buf.append(ch)
|
||
repaired = "".join(buf)
|
||
repaired = re.sub(r",\s*([}\]])", r"\1", repaired)
|
||
return repaired
|
||
|
||
def _parse_json_object_response(self, response_text: str, *, stage_label: str) -> Dict[str, Any]:
|
||
normalized = self._extract_json_candidate(response_text)
|
||
try:
|
||
payload = json.loads(normalized)
|
||
except Exception as exc:
|
||
repaired = self._repair_json_text(normalized)
|
||
try:
|
||
payload = json.loads(repaired)
|
||
except Exception:
|
||
try:
|
||
literal_payload = ast.literal_eval(normalized)
|
||
except Exception:
|
||
raise RuntimeError(f"{stage_label}：{exc}") from exc
|
||
payload = literal_payload
|
||
if not isinstance(payload, dict):
|
||
raise RuntimeError(f"{stage_label}：返回内容不是 JSON 对象")
|
||
return payload
|
||
|
||
async def _generate_project_creation_object(
|
||
self,
|
||
*,
|
||
prompt: str,
|
||
tenant_id: int,
|
||
function_name: str,
|
||
stage_label: str,
|
||
timeout_seconds: float | None = None,
|
||
) -> Dict[str, Any]:
|
||
max_rounds = PROJECT_CREATION_JSON_PARSE_RETRY_COUNT + 1
|
||
last_error: Exception | None = None
|
||
for round_index in range(max_rounds):
|
||
for max_tokens in self._build_project_creation_token_ladder():
|
||
try:
|
||
response = await self._generate_project_creation_text(
|
||
prompt=prompt,
|
||
tenant_id=tenant_id,
|
||
function_name=function_name,
|
||
timeout_seconds=float(timeout_seconds) if timeout_seconds is not None else PROJECT_CREATION_DRAFT_TIMEOUT_SECONDS,
|
||
max_tokens=max_tokens,
|
||
)
|
||
except Exception as exc:
|
||
last_error = exc if isinstance(exc, Exception) else RuntimeError(str(exc))
|
||
logger.warning(
|
||
"ai_steward_project_creation_token_retry",
|
||
tenant_id=int(tenant_id),
|
||
function_name=function_name,
|
||
stage_label=stage_label,
|
||
round=round_index + 1,
|
||
max_tokens=max_tokens,
|
||
error_type=type(last_error).__name__,
|
||
error=str(last_error),
|
||
)
|
||
if self._is_provider_budget_rejection(last_error):
|
||
break
|
||
continue
|
||
try:
|
||
return self._parse_json_object_response(response, stage_label=stage_label)
|
||
except Exception as exc:
|
||
last_error = exc if isinstance(exc, Exception) else RuntimeError(str(exc))
|
||
logger.warning(
|
||
"ai_steward_project_creation_json_retry",
|
||
tenant_id=int(tenant_id),
|
||
function_name=function_name,
|
||
stage_label=stage_label,
|
||
round=round_index + 1,
|
||
max_tokens=max_tokens,
|
||
error_type=type(last_error).__name__,
|
||
error=str(last_error),
|
||
)
|
||
raise RuntimeError(str(last_error or f"{stage_label}：unknown_error")) from last_error
|
||
|
||
async def _generate_project_creation_text(
|
||
self,
|
||
*,
|
||
prompt: str,
|
||
tenant_id: int,
|
||
function_name: str,
|
||
timeout_seconds: float = PROJECT_CREATION_DRAFT_TIMEOUT_SECONDS,
|
||
max_tokens: int = PROJECT_CREATION_STEP_INITIAL_MAX_TOKENS,
|
||
) -> str:
|
||
attempts_per_function = PROJECT_CREATION_STEP_RETRY_COUNT + 1
|
||
last_error: Exception | None = None
|
||
for attempt_index in range(attempts_per_function):
|
||
started_at = time.perf_counter()
|
||
try:
|
||
response = await asyncio.wait_for(
|
||
self.llm.generate(
|
||
prompt,
|
||
tenant_id=tenant_id,
|
||
function_name=function_name,
|
||
response_mode="fast",
|
||
use_prompt_profile=False,
|
||
temperature=0.0,
|
||
max_tokens=max_tokens,
|
||
),
|
||
timeout=timeout_seconds,
|
||
)
|
||
cleaned_response = str(response or "").strip()
|
||
elapsed_seconds = round(time.perf_counter() - started_at, 2)
|
||
if not cleaned_response:
|
||
raise RuntimeError("project_creation_empty_response")
|
||
logger.info(
|
||
"ai_steward_project_creation_step_ok",
|
||
tenant_id=int(tenant_id),
|
||
function_name=function_name,
|
||
attempt=attempt_index + 1,
|
||
max_tokens=max_tokens,
|
||
elapsed_seconds=elapsed_seconds,
|
||
)
|
||
return cleaned_response
|
||
except Exception as exc:
|
||
last_error = exc if isinstance(exc, Exception) else RuntimeError(str(exc))
|
||
logger.warning(
|
||
"ai_steward_project_creation_step_retry",
|
||
tenant_id=int(tenant_id),
|
||
function_name=function_name,
|
||
attempt=attempt_index + 1,
|
||
max_tokens=max_tokens,
|
||
elapsed_seconds=round(time.perf_counter() - started_at, 2),
|
||
error_type=type(exc).__name__,
|
||
error=str(exc),
|
||
)
|
||
|
||
raise RuntimeError(
|
||
f"AI 管家项目创建步骤失败：function={function_name}, reason={last_error or 'unknown_error'}"
|
||
) from last_error
|