import asyncio import json import re from dataclasses import dataclass from urllib.parse import urlparse 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, file_summaries: list[dict] | None = None) -> 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}") if file_summaries: file_blocks = [] for item in file_summaries: file_blocks.append(f"文件名:{item['file_name']}\n文件摘要:{item['summary']}") user_parts.append("参考文件内容:\n" + "\n\n".join(file_blocks)) 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("/") parsed = urlparse(endpoint if "://" in endpoint else f"https://{endpoint}") host = parsed.netloc or parsed.path.split("/")[0] if host == "api.deepseek.com": return "https://api.deepseek.com/chat/completions" 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} - {response.text[:500]}", 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) error_message = str(last_error) if isinstance(last_error, httpx.HTTPStatusError) and last_error.response is not None: error_message = f"{error_message}\n响应内容: {last_error.response.text[:1000]}" raise RuntimeError(f"模型调用失败:{error_message}") 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, trust_env=False) 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, trust_env=False) 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, file_summaries: list[dict] | None = None) -> AiCallResult: system_prompt, user_prompt = _build_prompt(paragraph, file_summaries) if model.api_format == "anthropic": return await _call_anthropic(model, system_prompt, user_prompt) return await _call_openai_compatible(model, system_prompt, user_prompt)