init: 初始化项目
This commit is contained in:
zwt13703
@@ -0,0 +1,181 @@
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
import httpx
|
||||
from app.core.security import decrypt_api_key
|
||||
|
||||
PROVIDER_OPENAI = "openai"
|
||||
PROVIDER_AZURE = "azure"
|
||||
PROVIDER_CUSTOM = "custom"
|
||||
|
||||
|
||||
class AIAdapter(ABC):
|
||||
@abstractmethod
|
||||
def build_request(self, model_config: dict, prompt: str, ref_content: str | None = None) -> dict:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def parse_response(self, response_data: dict) -> str:
|
||||
pass
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def provider(self) -> str:
|
||||
pass
|
||||
|
||||
|
||||
class OpenAIAdapter(AIAdapter):
|
||||
@property
|
||||
def provider(self) -> str:
|
||||
return PROVIDER_OPENAI
|
||||
|
||||
def build_request(self, model_config: dict, prompt: str, ref_content: str | None = None) -> dict:
|
||||
extra = model_config.get("extra_params", {})
|
||||
temperature = extra.get("temperature", 0.7)
|
||||
max_tokens = extra.get("max_tokens", 2000)
|
||||
|
||||
messages = [{"role": "system", "content": "你是一个专业的文档内容生成助手。"}]
|
||||
user_content = prompt
|
||||
if ref_content:
|
||||
user_content = f"参考以下内容：\n{ref_content}\n\n任务：{prompt}"
|
||||
messages.append({"role": "user", "content": user_content})
|
||||
|
||||
return {
|
||||
"url": model_config["endpoint"],
|
||||
"headers": {
|
||||
"Authorization": f"Bearer {decrypt_api_key(model_config['api_key'])}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
"json": {
|
||||
"model": extra.get("model", "gpt-4"),
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
},
|
||||
}
|
||||
|
||||
def parse_response(self, response_data: dict) -> str:
|
||||
return response_data["choices"][0]["message"]["content"]
|
||||
|
||||
|
||||
class AzureAdapter(AIAdapter):
|
||||
@property
|
||||
def provider(self) -> str:
|
||||
return PROVIDER_AZURE
|
||||
|
||||
def build_request(self, model_config: dict, prompt: str, ref_content: str | None = None) -> dict:
|
||||
extra = model_config.get("extra_params", {})
|
||||
temperature = extra.get("temperature", 0.7)
|
||||
max_tokens = extra.get("max_tokens", 2000)
|
||||
|
||||
messages = [{"role": "system", "content": "你是一个专业的文档内容生成助手。"}]
|
||||
user_content = prompt
|
||||
if ref_content:
|
||||
user_content = f"参考以下内容：\n{ref_content}\n\n任务：{prompt}"
|
||||
messages.append({"role": "user", "content": user_content})
|
||||
|
||||
api_version = extra.get("api_version", "2024-02-15-preview")
|
||||
endpoint = model_config["endpoint"]
|
||||
url = f"{endpoint}?api-version={api_version}"
|
||||
|
||||
return {
|
||||
"url": url,
|
||||
"headers": {
|
||||
"api-key": decrypt_api_key(model_config["api_key"]),
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
"json": {
|
||||
"messages": messages,
|
||||
"temperature": temperature,
|
||||
"max_tokens": max_tokens,
|
||||
},
|
||||
}
|
||||
|
||||
def parse_response(self, response_data: dict) -> str:
|
||||
return response_data["choices"][0]["message"]["content"]
|
||||
|
||||
|
||||
class CustomAdapter(AIAdapter):
|
||||
@property
|
||||
def provider(self) -> str:
|
||||
return PROVIDER_CUSTOM
|
||||
|
||||
def build_request(self, model_config: dict, prompt: str, ref_content: str | None = None) -> dict:
|
||||
extra = model_config.get("extra_params", {})
|
||||
user_content = prompt
|
||||
if ref_content:
|
||||
user_content = f"参考以下内容：\n{ref_content}\n\n任务：{prompt}"
|
||||
|
||||
messages = [{"role": "system", "content": "你是一个专业的文档内容生成助手。"}]
|
||||
messages.append({"role": "user", "content": user_content})
|
||||
|
||||
body = {
|
||||
"model": extra.get("model", "gpt-3.5-turbo"),
|
||||
"messages": messages,
|
||||
"max_tokens": extra.get("max_tokens", 2000),
|
||||
"temperature": extra.get("temperature", 0.7),
|
||||
}
|
||||
body.update({k: v for k, v in extra.items() if k not in ("model", "messages", "max_tokens", "temperature")})
|
||||
|
||||
return {
|
||||
"url": model_config["endpoint"],
|
||||
"headers": {
|
||||
"Authorization": f"Bearer {decrypt_api_key(model_config['api_key'])}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
"json": body,
|
||||
}
|
||||
|
||||
def parse_response(self, response_data: dict) -> str:
|
||||
if "choices" in response_data:
|
||||
return response_data["choices"][0]["message"]["content"]
|
||||
if "response" in response_data:
|
||||
return response_data["response"]
|
||||
if "content" in response_data:
|
||||
return response_data["content"]
|
||||
if "text" in response_data:
|
||||
return response_data["text"]
|
||||
return json.dumps(response_data)
|
||||
|
||||
|
||||
_adapters: dict[str, AIAdapter] = {
|
||||
PROVIDER_OPENAI: OpenAIAdapter(),
|
||||
PROVIDER_AZURE: AzureAdapter(),
|
||||
PROVIDER_CUSTOM: CustomAdapter(),
|
||||
}
|
||||
|
||||
|
||||
def get_adapter(provider: str) -> AIAdapter:
|
||||
adapter = _adapters.get(provider)
|
||||
if not adapter:
|
||||
raise ValueError(f"不支持的供应商: {provider}")
|
||||
return adapter
|
||||
|
||||
|
||||
async def call_ai_model(model_config: dict, prompt: str, ref_content: str | None = None) -> str:
|
||||
adapter = get_adapter(model_config["provider"])
|
||||
request = adapter.build_request(model_config, prompt, ref_content)
|
||||
|
||||
timeout = model_config.get("extra_params", {}).get("timeout", 120)
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(
|
||||
timeout=timeout,
|
||||
proxy=None,
|
||||
trust_env=False,
|
||||
) as client:
|
||||
response = await client.post(
|
||||
request["url"],
|
||||
headers=request["headers"],
|
||||
json=request["json"],
|
||||
)
|
||||
response.raise_for_status()
|
||||
return adapter.parse_response(response.json())
|
||||
except httpx.HTTPStatusError as e:
|
||||
detail = e.response.text[:500] if e.response else str(e)
|
||||
raise RuntimeError(f"AI 服务返回错误 ({e.response.status_code}): {detail}")
|
||||
except httpx.TimeoutException:
|
||||
raise RuntimeError("AI 调用超时")
|
||||
except httpx.ConnectError as e:
|
||||
raise RuntimeError(f"无法连接 AI 服务: {e}")
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"AI 调用异常: {str(e)}")
|
||||
@@ -0,0 +1,168 @@
|
||||
import mammoth
|
||||
from io import BytesIO
|
||||
from bs4 import BeautifulSoup
|
||||
from docx import Document
|
||||
from docx.shared import Pt, Inches
|
||||
from docx.enum.text import WD_ALIGN_PARAGRAPH
|
||||
from app.services.file_storage import get_file_content
|
||||
|
||||
ALIGN_MAP = {
|
||||
WD_ALIGN_PARAGRAPH.CENTER: "center",
|
||||
WD_ALIGN_PARAGRAPH.RIGHT: "right",
|
||||
WD_ALIGN_PARAGRAPH.JUSTIFY: "justify",
|
||||
}
|
||||
|
||||
|
||||
async def docx_to_html(file_content: bytes) -> str:
|
||||
result = mammoth.convert_to_html(BytesIO(file_content))
|
||||
html = result.value
|
||||
|
||||
try:
|
||||
doc = Document(BytesIO(file_content))
|
||||
soup = BeautifulSoup(html, "html.parser")
|
||||
paragraphs = doc.paragraphs
|
||||
html_blocks = soup.find_all(["p", "h1", "h2", "h3", "h4", "h5", "h6", "li"])
|
||||
|
||||
for i, para in enumerate(paragraphs):
|
||||
if i >= len(html_blocks):
|
||||
break
|
||||
if para.alignment and para.alignment in ALIGN_MAP:
|
||||
css = ALIGN_MAP[para.alignment]
|
||||
existing = html_blocks[i].get("style", "")
|
||||
styles = f"text-align:{css}"
|
||||
if existing:
|
||||
styles = existing.rstrip(";") + ";" + styles
|
||||
html_blocks[i]["style"] = styles
|
||||
|
||||
html = str(soup)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return html
|
||||
|
||||
|
||||
def html_to_docx_bytes(html_content: str) -> bytes:
|
||||
from html.parser import HTMLParser
|
||||
|
||||
doc = Document()
|
||||
|
||||
style = doc.styles["Normal"]
|
||||
font = style.font
|
||||
font.name = "Arial"
|
||||
font.size = Pt(11)
|
||||
|
||||
class RichHTMLParser(HTMLParser):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.paragraphs: list[dict] = []
|
||||
self.current = {"runs": [], "align": None}
|
||||
self.in_paragraph = False
|
||||
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
|
||||
self.tag_stack: list[str] = []
|
||||
self.heading_level = 0
|
||||
|
||||
def handle_starttag(self, tag, attrs):
|
||||
tag_lower = tag.lower()
|
||||
attrs_dict = dict(attrs)
|
||||
if tag_lower in ("p", "div", "li"):
|
||||
self.in_paragraph = True
|
||||
self.current = {"runs": [], "align": None}
|
||||
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
|
||||
style = attrs_dict.get("style", "")
|
||||
if "text-align:center" in style:
|
||||
self.current["align"] = "center"
|
||||
elif "text-align:right" in style:
|
||||
self.current["align"] = "right"
|
||||
elif "text-align:justify" in style:
|
||||
self.current["align"] = "justify"
|
||||
elif tag_lower in ("h1", "h2", "h3", "h4", "h5", "h6"):
|
||||
self.in_paragraph = True
|
||||
self.heading_level = int(tag_lower[1])
|
||||
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
|
||||
elif tag_lower in ("strong", "b"):
|
||||
self.current_run["bold"] = True
|
||||
elif tag_lower in ("em", "i"):
|
||||
self.current_run["italic"] = True
|
||||
elif tag_lower == "u":
|
||||
self.current_run["underline"] = True
|
||||
elif tag_lower in ("br",):
|
||||
if self.in_paragraph:
|
||||
self.current["runs"].append(dict(self.current_run))
|
||||
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
|
||||
self.tag_stack.append(tag_lower)
|
||||
|
||||
def handle_endtag(self, tag):
|
||||
tag_lower = tag.lower()
|
||||
if tag_lower in ("p", "div", "li", "h1", "h2", "h3", "h4", "h5", "h6"):
|
||||
if self.current_run["text"].strip():
|
||||
self.current["runs"].append(dict(self.current_run))
|
||||
if self.current["runs"]:
|
||||
p = dict(self.current)
|
||||
p["heading"] = self.heading_level
|
||||
self.paragraphs.append(p)
|
||||
self.current = {"runs": []}
|
||||
self.current_run = {"text": "", "bold": False, "italic": False, "underline": False}
|
||||
self.in_paragraph = False
|
||||
self.heading_level = 0
|
||||
if self.tag_stack:
|
||||
self.tag_stack.pop()
|
||||
|
||||
def handle_data(self, data):
|
||||
if self.in_paragraph:
|
||||
self.current_run["text"] += data
|
||||
|
||||
parser = RichHTMLParser()
|
||||
parser.feed(html_content)
|
||||
|
||||
for para_data in parser.paragraphs:
|
||||
heading = para_data.get("heading", 0)
|
||||
if heading > 0:
|
||||
p = doc.add_heading(level=min(heading, 9))
|
||||
else:
|
||||
p = doc.add_paragraph()
|
||||
|
||||
align_val = para_data.get("align")
|
||||
if align_val == "center":
|
||||
p.alignment = WD_ALIGN_PARAGRAPH.CENTER
|
||||
elif align_val == "right":
|
||||
p.alignment = WD_ALIGN_PARAGRAPH.RIGHT
|
||||
elif align_val == "justify":
|
||||
p.alignment = WD_ALIGN_PARAGRAPH.JUSTIFY
|
||||
|
||||
for run_data in para_data.get("runs", []):
|
||||
run = p.add_run(run_data["text"])
|
||||
if run_data.get("bold"):
|
||||
run.bold = True
|
||||
if run_data.get("italic"):
|
||||
run.italic = True
|
||||
if run_data.get("underline"):
|
||||
run.underline = True
|
||||
|
||||
output = BytesIO()
|
||||
doc.save(output)
|
||||
return output.getvalue()
|
||||
|
||||
|
||||
def docx_to_pdf_bytes(file_content: bytes) -> bytes:
|
||||
doc = Document(BytesIO(file_content))
|
||||
|
||||
from io import BytesIO as Bio
|
||||
from reportlab.lib.pagesizes import A4
|
||||
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
|
||||
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer
|
||||
from reportlab.lib.enums import TA_LEFT
|
||||
|
||||
buffer = Bio()
|
||||
pdf_doc = SimpleDocTemplate(buffer, pagesize=A4)
|
||||
|
||||
styles = getSampleStyleSheet()
|
||||
story = []
|
||||
|
||||
for para in doc.paragraphs:
|
||||
if para.text.strip():
|
||||
p = Paragraph(para.text, styles["Normal"])
|
||||
story.append(p)
|
||||
story.append(Spacer(1, 6))
|
||||
|
||||
pdf_doc.build(story)
|
||||
return buffer.getvalue()
|
||||
@@ -0,0 +1,43 @@
|
||||
import os
|
||||
import aiofiles
|
||||
from pathlib import Path
|
||||
from fastapi import UploadFile
|
||||
from app.core.config import get_settings
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
TEMPLATES_DIR = "templates"
|
||||
REF_FILES_DIR = "ref_files"
|
||||
RESULTS_DIR = "results"
|
||||
|
||||
|
||||
def _ensure_dir(path: str) -> str:
|
||||
os.makedirs(path, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def get_storage_dir(subdir: str) -> str:
|
||||
return _ensure_dir(os.path.join(settings.STORAGE_ROOT, subdir))
|
||||
|
||||
|
||||
async def save_upload(upload_file: UploadFile, subdir: str) -> str:
|
||||
dir_path = get_storage_dir(subdir)
|
||||
file_path = os.path.join(dir_path, upload_file.filename or "unnamed")
|
||||
async with aiofiles.open(file_path, "wb") as f:
|
||||
content = await upload_file.read()
|
||||
await f.write(content)
|
||||
return file_path
|
||||
|
||||
|
||||
async def get_file_content(file_path: str) -> bytes:
|
||||
async with aiofiles.open(file_path, "rb") as f:
|
||||
return await f.read()
|
||||
|
||||
|
||||
def delete_file(file_path: str) -> None:
|
||||
if os.path.exists(file_path):
|
||||
os.remove(file_path)
|
||||
|
||||
|
||||
def get_absolute_path(rel_path: str) -> str:
|
||||
return os.path.abspath(rel_path)
|
||||
@@ -0,0 +1,46 @@
|
||||
import os
|
||||
import json
|
||||
from docx import Document
|
||||
from PyPDF2 import PdfReader
|
||||
from app.services.file_storage import get_file_content
|
||||
|
||||
|
||||
async def parse_reference_file(file_path: str) -> str:
|
||||
ext = os.path.splitext(file_path)[1].lower()
|
||||
content = await get_file_content(file_path)
|
||||
|
||||
if ext == ".txt":
|
||||
return content.decode("utf-8", errors="ignore")
|
||||
if ext == ".docx":
|
||||
from io import BytesIO
|
||||
doc = Document(BytesIO(content))
|
||||
return "\n".join(p.text for p in doc.paragraphs if p.text.strip())
|
||||
if ext == ".pdf":
|
||||
from io import BytesIO
|
||||
reader = PdfReader(BytesIO(content))
|
||||
return "\n".join(page.extract_text() or "" for page in reader.pages)
|
||||
|
||||
raise ValueError(f"不支持的文件格式: {ext}")
|
||||
|
||||
|
||||
async def parse_reference_files(ref_file_path: str | None) -> str:
|
||||
if not ref_file_path:
|
||||
return ""
|
||||
|
||||
try:
|
||||
paths = json.loads(ref_file_path)
|
||||
if isinstance(paths, list):
|
||||
texts = []
|
||||
for path in paths:
|
||||
try:
|
||||
texts.append(await parse_reference_file(path))
|
||||
except Exception:
|
||||
texts.append(f"[无法解析文件: {path}]")
|
||||
return "\n\n".join(texts)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
try:
|
||||
return await parse_reference_file(ref_file_path)
|
||||
except Exception:
|
||||
return f"[无法解析文件: {ref_file_path}]"
|
||||
Reference in New Issue
Block a user