zwt13703
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
|