evyzacq
261 lines
11 KiB
Python
261 lines
11 KiB
Python
# PRD 管理服务：上传、解析、版本快照
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
import time
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
|
|
from server.config import settings
|
|
from server.services.prd_manager.word_parser import WordParser
|
|
from server.services.prd_manager.image_parser import ImageParser
|
|
|
|
logger = logging.getLogger("testflow")
|
|
RATE_LIMIT_DELAY = 0.5
|
|
|
|
|
|
class PRDService:
|
|
"""PRD management: handle file upload, text extraction, and version snapshots."""
|
|
|
|
def __init__(self):
|
|
self.output_dir = settings.OUTPUT_DIR
|
|
os.makedirs(self.output_dir, exist_ok=True)
|
|
self._prds: dict[str, dict] = {}
|
|
|
|
async def upload_and_parse(self, filename: str, content: bytes) -> dict:
|
|
prd_id = str(uuid.uuid4())[:8]
|
|
basename = Path(filename).stem
|
|
ext = Path(filename).suffix.lower()
|
|
|
|
upload_dir = self.output_dir / prd_id
|
|
os.makedirs(upload_dir, exist_ok=True)
|
|
filepath = upload_dir / filename
|
|
|
|
t0 = time.time()
|
|
with open(filepath, "wb") as f:
|
|
f.write(content)
|
|
|
|
logger.info("[UPLOAD] 收到文件: %s (%.1f KB, %s)", filename, len(content) / 1024, ext)
|
|
|
|
if ext in (".txt", ".md"):
|
|
text = content.decode("utf-8", errors="replace")
|
|
sections = self._parse_markdown_sections(text)
|
|
parsed = {
|
|
"source": str(filepath), "sections": sections,
|
|
"image_sources": {}, "image_analysis": [], "full_text": text,
|
|
}
|
|
logger.info("[PARSE] %.4s %s → %d 章节, %d 字符, %.1fs",
|
|
prd_id, ext, len(sections), len(text), time.time() - t0)
|
|
|
|
elif ext == ".docx":
|
|
logger.info("[PARSE] %.4s → 开始解析 .docx...", prd_id)
|
|
parsed = await self._parse_docx(str(filepath), upload_dir)
|
|
logger.info("[PARSE] %.4s .docx 完成: %d 章节, %.1fs",
|
|
prd_id, len(parsed.get("sections", [])), time.time() - t0)
|
|
|
|
elif ext == ".pdf":
|
|
logger.info("[PARSE] %.4s → 开始解析 .pdf...", prd_id)
|
|
parsed = await self._parse_pdf(str(filepath), upload_dir)
|
|
logger.info("[PARSE] %.4s .pdf 完成: %d 页, %.1fs",
|
|
prd_id, len(parsed.get("sections", [])), time.time() - t0)
|
|
else:
|
|
raise ValueError(f"Unsupported file format: {ext}")
|
|
|
|
parsed_path = upload_dir / f"{basename}_parsed.json"
|
|
with open(parsed_path, "w", encoding="utf-8") as f:
|
|
json.dump(parsed, f, ensure_ascii=False, indent=2)
|
|
|
|
prd_version = {
|
|
"prd_id": prd_id, "filename": filename,
|
|
"uploaded_at": datetime.now(timezone.utc).isoformat(), "status": "ready",
|
|
"parsed_path": str(parsed_path), "full_text": parsed.get("full_text", ""),
|
|
"sections": parsed.get("sections", []), "images": parsed.get("image_analysis", []),
|
|
}
|
|
self._prds[prd_id] = prd_version
|
|
|
|
n_paras = sum(1 for s in parsed.get("sections", []) for b in s.get("blocks", []) if b["type"] == "para")
|
|
n_tables = sum(1 for s in parsed.get("sections", []) for b in s.get("blocks", []) if b["type"] == "table")
|
|
logger.info("[UPLOAD] %.4s 解析完成: %d章节 %d段落 %d表格 %d图片 → %s",
|
|
prd_id, len(parsed.get("sections", [])), n_paras, n_tables,
|
|
len(parsed.get("image_analysis", [])), parsed_path)
|
|
return prd_version
|
|
|
|
async def upload_quick(self, filename: str, content: bytes) -> dict:
|
|
"""Quick save — returns immediately, parsing happens in background."""
|
|
prd_id = str(uuid.uuid4())[:8]
|
|
upload_dir = self.output_dir / prd_id
|
|
os.makedirs(upload_dir, exist_ok=True)
|
|
filepath = upload_dir / filename
|
|
with open(filepath, "wb") as f:
|
|
f.write(content)
|
|
|
|
logger.info("[UPLOAD] %.4s 快速保存: %s (%.1f KB)", prd_id, filename, len(content) / 1024)
|
|
|
|
prd_version = {
|
|
"prd_id": prd_id, "filename": filename,
|
|
"uploaded_at": datetime.now(timezone.utc).isoformat(),
|
|
"status": "parsing", "full_text": "", "sections": [], "images": [],
|
|
"parsed_path": "", "filepath": str(filepath),
|
|
}
|
|
self._prds[prd_id] = prd_version
|
|
return prd_version
|
|
|
|
async def parse_async(self, prd_id: str, filename: str, content: bytes):
|
|
"""Full parse (text + images + vision) — called in background thread."""
|
|
basename = Path(filename).stem
|
|
ext = Path(filename).suffix.lower()
|
|
upload_dir = self.output_dir / prd_id
|
|
filepath = upload_dir / filename
|
|
t0 = time.time()
|
|
|
|
# Parse
|
|
if ext in (".txt", ".md"):
|
|
text = content.decode("utf-8", errors="replace")
|
|
sections = self._parse_markdown_sections(text)
|
|
parsed = {"source": str(filepath), "sections": sections,
|
|
"image_sources": {}, "image_analysis": [], "full_text": text}
|
|
elif ext == ".docx":
|
|
parsed = await self._parse_docx(str(filepath), upload_dir)
|
|
elif ext == ".pdf":
|
|
parsed = await self._parse_pdf(str(filepath), upload_dir)
|
|
else:
|
|
self.set_status(prd_id, "error", f"Unsupported format: {ext}")
|
|
return
|
|
|
|
# Save parsed
|
|
parsed_path = upload_dir / f"{basename}_parsed.json"
|
|
with open(parsed_path, "w", encoding="utf-8") as f:
|
|
json.dump(parsed, f, ensure_ascii=False, indent=2)
|
|
|
|
# Update in-memory
|
|
prd = self._prds.get(prd_id, {})
|
|
prd.update({
|
|
"status": "ready", "parsed_path": str(parsed_path),
|
|
"full_text": parsed.get("full_text", ""),
|
|
"sections": parsed.get("sections", []),
|
|
"images": parsed.get("image_analysis", []),
|
|
})
|
|
|
|
n_paras = sum(1 for s in parsed.get("sections", []) for b in s.get("blocks", []) if b["type"] == "para")
|
|
n_tables = sum(1 for s in parsed.get("sections", []) for b in s.get("blocks", []) if b["type"] == "table")
|
|
logger.info("[UPLOAD] %.4s 后台解析完成: %d章节 %d段落 %d表格 %d图片, %.1fs",
|
|
prd_id, len(parsed.get("sections", [])), n_paras, n_tables,
|
|
len(parsed.get("image_analysis", [])), time.time() - t0)
|
|
|
|
def set_status(self, prd_id: str, status: str, error: str = ""):
|
|
if prd_id in self._prds:
|
|
self._prds[prd_id]["status"] = status
|
|
if error:
|
|
self._prds[prd_id]["error"] = error
|
|
|
|
async def get_status(self, prd_id: str) -> dict | None:
|
|
prd = self._prds.get(prd_id)
|
|
if not prd:
|
|
return None
|
|
return {"prd_id": prd_id, "status": prd.get("status", "ready"),
|
|
"error": prd.get("error", ""),
|
|
"sections_count": len(prd.get("sections", [])),
|
|
"images_count": len(prd.get("images", []))}
|
|
|
|
async def get_prd(self, prd_id: str) -> dict | None:
|
|
return self._prds.get(prd_id)
|
|
|
|
# ---- internal parsers ----
|
|
|
|
async def _parse_docx(self, filepath: str, output_dir: Path) -> dict:
|
|
word = WordParser(filepath)
|
|
sections, image_sources = word.extract_sections()
|
|
full_text = word.extract_full_text()
|
|
|
|
images_dir = output_dir / "images"
|
|
images = word.extract_images(str(images_dir))
|
|
image_analysis = []
|
|
n_chart_images = 0
|
|
|
|
if images:
|
|
parser = ImageParser()
|
|
logger.info("[IMAGE] %.4s 发现 %d 张图片，开始视觉分析...", output_dir.parent.name, len(images))
|
|
for i, img in enumerate(images):
|
|
t_img = time.time()
|
|
logger.info("[IMAGE] [%d/%d] rid=%s, 调用 Qwen VL...", i + 1, len(images), img["rid"])
|
|
result = parser.parse_image(img["path"])
|
|
if result is None:
|
|
result = {"type": "other", "description": ""}
|
|
result["rid"] = img["rid"]
|
|
result["path"] = img["path"]
|
|
if img["rid"] in image_sources:
|
|
result["context"] = image_sources[img["rid"]]
|
|
image_analysis.append(result)
|
|
if result.get("type") in ("flowchart", "architecture", "state", "sequence", "activity"):
|
|
n_chart_images += 1
|
|
logger.info("[IMAGE] [%d/%d] rid=%s 完成，类型=%s (%.1fs)",
|
|
i + 1, len(images), img["rid"], result.get("type", "?"), time.time() - t_img)
|
|
if i < len(images) - 1:
|
|
time.sleep(RATE_LIMIT_DELAY)
|
|
|
|
usg = parser.usage
|
|
logger.info("[IMAGE] %.4s 图片分析完成: %d张, 图表类%d张, tokens(p=%d c=%d t=%d)",
|
|
output_dir.parent.name, len(images), n_chart_images,
|
|
usg["prompt_tokens"], usg["completion_tokens"], usg["total_tokens"])
|
|
else:
|
|
logger.info("[IMAGE] %.4s 文档无图片", output_dir.parent.name)
|
|
|
|
return {
|
|
"source": filepath, "sections": sections,
|
|
"image_sources": image_sources, "image_analysis": image_analysis,
|
|
"full_text": full_text,
|
|
}
|
|
|
|
def _parse_markdown_sections(self, text: str) -> list[dict]:
|
|
import re
|
|
sections: list[dict] = []
|
|
current_source = "正文"
|
|
current_blocks: list[dict] = []
|
|
para_idx = 0
|
|
|
|
for line in text.split("\n"):
|
|
stripped = line.strip()
|
|
if not stripped:
|
|
continue
|
|
heading_match = re.match(r"^(#{1,6})\s+(.+)", stripped)
|
|
if heading_match:
|
|
if current_blocks:
|
|
sections.append({"source": current_source, "blocks": current_blocks, "images": []})
|
|
current_blocks = []
|
|
level = len(heading_match.group(1))
|
|
current_source = f"{'#' * level} {heading_match.group(2)}"
|
|
continue
|
|
if "|" in stripped and stripped.count("|") >= 2:
|
|
cells = [c.strip() for c in stripped.split("|") if c.strip()]
|
|
if cells:
|
|
current_blocks.append({
|
|
"type": "table",
|
|
"table": len([b for b in current_blocks if b["type"] == "table"]) + 1,
|
|
"headers": cells,
|
|
"rows": [{"columns": [{"name": c, "row": 1, "col": i + 1, "text": c} for i, c in enumerate(cells)]}],
|
|
})
|
|
continue
|
|
para_idx += 1
|
|
current_blocks.append({"type": "para", "index": para_idx, "text": stripped})
|
|
|
|
if current_blocks:
|
|
sections.append({"source": current_source, "blocks": current_blocks, "images": []})
|
|
return sections
|
|
|
|
async def _parse_pdf(self, filepath: str, output_dir: Path) -> dict:
|
|
from PyPDF2 import PdfReader
|
|
reader = PdfReader(filepath)
|
|
full_text_parts = []
|
|
for i, page in enumerate(reader.pages):
|
|
text = page.extract_text()
|
|
if text:
|
|
full_text_parts.append(text)
|
|
full_text = "\n".join(full_text_parts)
|
|
sections = [{"source": "PDF 内容", "blocks": [{"type": "para", "index": i + 1, "text": t} for i, t in enumerate(full_text_parts)], "images": []}]
|
|
return {"source": filepath, "sections": sections, "image_sources": {}, "image_analysis": [], "full_text": full_text}
|
|
|
|
|
|
prd_service = PRDService()
|