308 lines
13 KiB
Python
308 lines
13 KiB
Python
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 _format_file_context(file_summaries: list[dict]) -> str:
|
||
file_blocks: list[str] = []
|
||
for item in file_summaries:
|
||
file_name = item.get("file_name") or "未命名文件"
|
||
summary = item.get("summary") or "文件内容为空。"
|
||
file_blocks.append(f"文件：{file_name}\n内容：\n{summary}")
|
||
return "\n\n".join(file_blocks)
|
||
|
||
|
||
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":[["..."]]}]}。'
|
||
)
|
||
if getattr(paragraph, "enable_reasoning", False):
|
||
system_prompt += "你可以先进行充分思考，再给出最终答案，但最终只输出要求的结果内容。"
|
||
|
||
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:
|
||
user_parts.append("参考文件内容：\n" + _format_file_context(file_summaries))
|
||
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 _stream_openai_compatible(model: AiModel, system_prompt: str, user_prompt: str):
|
||
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,
|
||
"stream": True,
|
||
}
|
||
headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
|
||
last_error: Exception | None = None
|
||
for attempt in range(settings.AI_MAX_RETRIES):
|
||
try:
|
||
async with httpx.AsyncClient(timeout=settings.AI_REQUEST_TIMEOUT, trust_env=False) as client:
|
||
async with client.stream("POST", _normalize_openai_endpoint(model.api_endpoint), headers=headers, json=payload) as response:
|
||
if response.status_code in (429, 500, 502, 503, 504):
|
||
body = await response.aread()
|
||
raise httpx.HTTPStatusError(
|
||
f"上游模型流式响应异常: {response.status_code} - {body.decode('utf-8', errors='ignore')[:500]}",
|
||
request=response.request,
|
||
response=response,
|
||
)
|
||
response.raise_for_status()
|
||
async for line in response.aiter_lines():
|
||
if not line or not line.startswith("data:"):
|
||
continue
|
||
payload_line = line[5:].strip()
|
||
if payload_line == "[DONE]":
|
||
break
|
||
try:
|
||
chunk = json.loads(payload_line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
delta_payload = chunk.get("choices", [{}])[0].get("delta", {})
|
||
delta = delta_payload.get("content", "")
|
||
reasoning = delta_payload.get("reasoning_content", "")
|
||
if isinstance(delta, list):
|
||
delta = "".join(
|
||
item.get("text", "") if isinstance(item, dict) else str(item)
|
||
for item in delta
|
||
)
|
||
if delta:
|
||
yield delta
|
||
if reasoning:
|
||
yield reasoning
|
||
return
|
||
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_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 _stream_anthropic(model: AiModel, system_prompt: str, user_prompt: str):
|
||
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}],
|
||
"stream": True,
|
||
}
|
||
headers = {
|
||
"x-api-key": api_key,
|
||
"anthropic-version": "2023-06-01",
|
||
"content-type": "application/json",
|
||
}
|
||
last_error: Exception | None = None
|
||
for attempt in range(settings.AI_MAX_RETRIES):
|
||
try:
|
||
async with httpx.AsyncClient(timeout=settings.AI_REQUEST_TIMEOUT, trust_env=False) as client:
|
||
async with client.stream("POST", _normalize_anthropic_endpoint(model.api_endpoint), headers=headers, json=payload) as response:
|
||
if response.status_code in (429, 500, 502, 503, 504):
|
||
body = await response.aread()
|
||
raise httpx.HTTPStatusError(
|
||
f"上游模型流式响应异常: {response.status_code} - {body.decode('utf-8', errors='ignore')[:500]}",
|
||
request=response.request,
|
||
response=response,
|
||
)
|
||
response.raise_for_status()
|
||
async for line in response.aiter_lines():
|
||
if not line or not line.startswith("data:"):
|
||
continue
|
||
payload_line = line[5:].strip()
|
||
if payload_line == "[DONE]":
|
||
break
|
||
try:
|
||
chunk = json.loads(payload_line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
if chunk.get("type") == "content_block_delta":
|
||
delta = chunk.get("delta", {}).get("text", "")
|
||
if delta:
|
||
yield delta
|
||
return
|
||
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_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)
|
||
|
||
|
||
def build_test_stream_prompt(paragraph: Paragraph, file_summaries: list[dict] | None = None) -> tuple[str, str]:
|
||
system_prompt = "你是一个企业文档撰写助手。请直接输出适合预览的正文内容或 Markdown 表格，不要输出 JSON。"
|
||
if getattr(paragraph, "enable_reasoning", False):
|
||
system_prompt += "你可以先进行充分思考，再持续输出最终可展示的内容。"
|
||
_, user_prompt = _build_prompt(paragraph, file_summaries or [])
|
||
return system_prompt, user_prompt
|
||
|
||
|
||
async def stream_ai_preview(paragraph: Paragraph, model: AiModel, file_summaries: list[dict] | None = None):
|
||
system_prompt, user_prompt = build_test_stream_prompt(paragraph, file_summaries)
|
||
if model.api_format == "anthropic":
|
||
async for chunk in _stream_anthropic(model, system_prompt, user_prompt):
|
||
yield chunk
|
||
return
|
||
async for chunk in _stream_openai_compatible(model, system_prompt, user_prompt):
|
||
yield chunk
|