接入真实模型调用与参考文件上传

backend/services/ai_service.py

@@ -0,0 +1,156 @@
+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)