doc-forge/backend/services/ai_service.py

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)