import asyncio import json import re from dataclasses import dataclass import httpx from config import settings from models.ai_model import AiModel from models.paragraph import Paragraph from services.security import decrypt_text @dataclass class AiCallResult: content: dict raw_text: str used_model: str def _ensure_json_content(text: str) -> dict: stripped = text.strip() if not stripped: return {"content": [{"type": "text", "text": ""}]} try: parsed = json.loads(stripped) if isinstance(parsed, dict) and "content" in parsed: return parsed except json.JSONDecodeError: pass code_block_match = re.search(r"```json\s*(.*?)\s*```", stripped, re.S) if code_block_match: try: parsed = json.loads(code_block_match.group(1)) if isinstance(parsed, dict) and "content" in parsed: return parsed except json.JSONDecodeError: pass return {"content": [{"type": "text", "text": stripped}]} def _build_prompt(paragraph: Paragraph) -> tuple[str, str]: system_prompt = ( "你是一个企业文档撰写助手。" "请严格输出 JSON,不要输出 JSON 之外的说明。" '格式为:{"content":[{"type":"text","text":"..."},{"type":"table","title":"...","headers":["..."],"rows":[["..."]]}]}。' ) user_parts = [f"段落标题:{paragraph.title}"] if paragraph.content: user_parts.append(f"模板上下文:{paragraph.content}") if paragraph.need_prompt and paragraph.prompt_text: user_parts.append(f"附加要求:{paragraph.prompt_text}") user_parts.append(f"输出格式:{paragraph.output_format}") return system_prompt, "\n\n".join(user_parts) def _normalize_openai_endpoint(api_endpoint: str) -> str: endpoint = api_endpoint.rstrip("/") if endpoint.endswith("/chat/completions"): return endpoint if endpoint.endswith("/v1"): return f"{endpoint}/chat/completions" return f"{endpoint}/v1/chat/completions" def _normalize_anthropic_endpoint(api_endpoint: str) -> str: endpoint = api_endpoint.rstrip("/") if endpoint.endswith("/messages"): return endpoint if endpoint.endswith("/v1"): return f"{endpoint}/messages" return f"{endpoint}/v1/messages" async def _post_with_retry( client: httpx.AsyncClient, url: str, headers: dict, payload: dict, ) -> httpx.Response: last_error: Exception | None = None for attempt in range(settings.AI_MAX_RETRIES): try: response = await client.post(url, headers=headers, json=payload) if response.status_code in (429, 500, 502, 503, 504): raise httpx.HTTPStatusError( f"上游模型响应异常: {response.status_code}", request=response.request, response=response, ) response.raise_for_status() return response except (httpx.TimeoutException, httpx.HTTPStatusError, httpx.NetworkError) as error: last_error = error if attempt == settings.AI_MAX_RETRIES - 1: break await asyncio.sleep(2 ** attempt) raise RuntimeError(f"模型调用失败:{last_error}") async def _call_openai_compatible(model: AiModel, system_prompt: str, user_prompt: str) -> AiCallResult: api_key = decrypt_text(model.api_key_encrypted) if not api_key: raise RuntimeError("模型 API Key 不可用") payload = { "model": model.name, "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ], "temperature": 0.3, } headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} async with httpx.AsyncClient(timeout=settings.AI_REQUEST_TIMEOUT) as client: response = await _post_with_retry(client, _normalize_openai_endpoint(model.api_endpoint), headers, payload) body = response.json() text = body["choices"][0]["message"]["content"] return AiCallResult(content=_ensure_json_content(text), raw_text=text, used_model=model.name) async def _call_anthropic(model: AiModel, system_prompt: str, user_prompt: str) -> AiCallResult: api_key = decrypt_text(model.api_key_encrypted) if not api_key: raise RuntimeError("模型 API Key 不可用") payload = { "model": model.name, "max_tokens": 2048, "system": system_prompt, "messages": [{"role": "user", "content": user_prompt}], } headers = { "x-api-key": api_key, "anthropic-version": "2023-06-01", "content-type": "application/json", } async with httpx.AsyncClient(timeout=settings.AI_REQUEST_TIMEOUT) as client: response = await _post_with_retry(client, _normalize_anthropic_endpoint(model.api_endpoint), headers, payload) body = response.json() text = "" for item in body.get("content", []): if item.get("type") == "text": text += item.get("text", "") return AiCallResult(content=_ensure_json_content(text), raw_text=text, used_model=model.name) async def call_ai(paragraph: Paragraph, model: AiModel) -> AiCallResult: system_prompt, user_prompt = _build_prompt(paragraph) if model.api_format == "anthropic": return await _call_anthropic(model, system_prompt, user_prompt) return await _call_openai_compatible(model, system_prompt, user_prompt)