evyzacq
754 lines
31 KiB
Python
754 lines
31 KiB
Python
# 3-Stage IR Generation Pipeline
|
||
#
|
||
# Stage 1: Semantic Index — full document → function_units + concepts
|
||
# Stage 2: Per-Unit IR Extraction — precision context → detailed IR rules
|
||
# Stage 3: Deterministic Merge & Audit — dedup, assign IDs, coverage report
|
||
|
||
import concurrent.futures
|
||
import hashlib
|
||
import json
|
||
import logging
|
||
import re
|
||
import time
|
||
import uuid
|
||
from collections import defaultdict
|
||
from datetime import datetime, timezone
|
||
from pathlib import Path
|
||
|
||
from server.config import settings
|
||
from server.core.llm_provider.router import router
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
PROMPTS_DIR = Path(__file__).parent.parent / "prompts" # services/prompts/
|
||
MAX_CONCURRENCY = 5
|
||
|
||
PASS, WARN, FAIL = "[PASS]", "[WARN]", "[FAIL]"
|
||
|
||
|
||
# =============================================================================
|
||
# Stage 1: Semantic Index Generation
|
||
# =============================================================================
|
||
|
||
def format_document_for_prompt(doc: dict) -> str:
|
||
"""Render the full parsed document as a readable string for the LLM prompt."""
|
||
lines = ["=== SECTIONS ==="]
|
||
|
||
for i, section in enumerate(doc.get("sections", [])):
|
||
source = section.get("source", f"(无标题-章节{i})")
|
||
lines.append(f"\n--- Section: {source} ---")
|
||
|
||
for block in section.get("blocks", []):
|
||
if block["type"] == "para":
|
||
lines.append(f"[段落 {block['index']}] {block['text']}")
|
||
elif block["type"] == "table":
|
||
lines.append(f"[表格 {block.get('table', '?')}]")
|
||
headers = block.get("headers", [])
|
||
lines.append(f" 表头: {' | '.join(headers)}")
|
||
for row in block.get("rows", []):
|
||
cols = row.get("columns", [])
|
||
cell_texts = []
|
||
for c in cols:
|
||
cell_texts.append(f"[行{c.get('row','?')}]{c.get('name','')}: {c.get('text','')}")
|
||
lines.append(f" {'; '.join(cell_texts)}")
|
||
|
||
images = section.get("images", [])
|
||
if images:
|
||
lines.append(f" 图片引用: {', '.join(images)}")
|
||
|
||
# Image Analysis
|
||
lines.append("\n\n=== IMAGE_ANALYSIS ===")
|
||
for img in doc.get("image_analysis", []):
|
||
rid = img.get("rid", "?")
|
||
img_type = img.get("type", "?")
|
||
lines.append(f"\n--- Image: {rid} (type={img_type}) ---")
|
||
lines.append(f" 描述: {img.get('description', '')[:500]}")
|
||
|
||
# Resolved Conflicts
|
||
conflicts = doc.get("resolved_conflicts", [])
|
||
if conflicts:
|
||
lines.append("\n\n=== RESOLVED_CONFLICTS ===")
|
||
for c in conflicts:
|
||
lines.append(
|
||
f" [{c.get('conflict_type','?')}] {c.get('section','?')}: "
|
||
f"以{c.get('source','?')}为准 — {c.get('correction','')}"
|
||
)
|
||
|
||
return "\n".join(lines)
|
||
|
||
|
||
def extract_json_from_response(text: str, expect_array: bool = False) -> str:
|
||
"""Robustly extract JSON object or array from LLM response."""
|
||
open_char = "[" if expect_array else "{"
|
||
close_char = "]" if expect_array else "}"
|
||
|
||
# Try code fence first
|
||
pattern = rf"```(?:json)?\s*(\[{open_char}[\s\S]*?\{close_char}\])\s*```" if expect_array else rf"```(?:json)?\s*(\{{[\s\S]*?\}})\s*```"
|
||
m = re.search(pattern, text)
|
||
if m:
|
||
return m.group(1).strip()
|
||
|
||
# Find outermost bracket
|
||
start = text.find(open_char)
|
||
if start == -1:
|
||
raise ValueError(f"No JSON {'array' if expect_array else 'object'} found in LLM response")
|
||
|
||
depth = 0
|
||
for i in range(start, len(text)):
|
||
if text[i] == open_char:
|
||
depth += 1
|
||
elif text[i] == close_char:
|
||
depth -= 1
|
||
if depth == 0:
|
||
return text[start: i + 1]
|
||
|
||
raise ValueError(f"Unclosed JSON {'array' if expect_array else 'object'} in LLM response")
|
||
|
||
|
||
async def stage1_semantic_index(doc: dict) -> dict:
|
||
"""Stage 1: Generate semantic index from full document."""
|
||
logger.info("[PIPELINE] ═══ Stage 1/3: 语义索引 ═══")
|
||
|
||
template_path = PROMPTS_DIR / "ir_step1_semantic_index.txt"
|
||
if not template_path.exists():
|
||
logger.warning("[PIPELINE] Stage 1 prompt template not found, using fallback")
|
||
return _fallback_semantic_index(doc)
|
||
|
||
template = template_path.read_text(encoding="utf-8")
|
||
formatted_doc = format_document_for_prompt(doc)
|
||
prompt = template.replace("{document_json}", formatted_doc)
|
||
|
||
logger.info("[PIPELINE] Prompt: %d chars (~%d tokens)", len(prompt), len(prompt) // 3)
|
||
|
||
client = router.get_text_client()
|
||
for attempt in range(3):
|
||
try:
|
||
raw = client.chat(
|
||
model=router.text_model,
|
||
messages=[
|
||
{"role": "system", "content": "你是一个精确的 JSON 输出引擎。只输出合法的 JSON，不输出任何其他文字。"},
|
||
{"role": "user", "content": prompt},
|
||
],
|
||
temperature=0.1,
|
||
)
|
||
json_str = extract_json_from_response(raw)
|
||
result = json.loads(json_str)
|
||
|
||
# Validate
|
||
if "function_units" not in result or not result["function_units"]:
|
||
raise ValueError("Missing function_units")
|
||
logger.info("[PIPELINE] ✅ Stage 1 完成: %d function_units, %d concepts",
|
||
len(result["function_units"]), len(result.get("concepts", [])))
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.warning("[PIPELINE] Stage 1 attempt %d failed: %s", attempt + 1, e)
|
||
if attempt < 2:
|
||
time.sleep(2)
|
||
|
||
# Fallback
|
||
logger.warning("[PIPELINE] Stage 1 全部重试失败 exhausted, using fallback")
|
||
return _fallback_semantic_index(doc)
|
||
|
||
|
||
def _fallback_semantic_index(doc: dict) -> dict:
|
||
"""Generate a basic semantic index without LLM (for mock/demo)."""
|
||
units = []
|
||
unit_idx = 1
|
||
|
||
for section in doc.get("sections", []):
|
||
source = section.get("source", "")
|
||
for block in section.get("blocks", []):
|
||
if block["type"] == "para":
|
||
# Create a function unit from each substantial paragraph
|
||
text = block.get("text", "")
|
||
if len(text) > 30:
|
||
units.append({
|
||
"unit_id": f"FU-{unit_idx:03d}",
|
||
"name": text[:60],
|
||
"description": text[:200],
|
||
"sources": [{"section": source, "type": "para", "text_snippet": text[:200]}],
|
||
})
|
||
unit_idx += 1
|
||
elif block["type"] == "table":
|
||
for row in block.get("rows", []):
|
||
cols = row.get("columns", [])
|
||
texts = [f"{c.get('name','')}: {c.get('text','')}" for c in cols if c.get('text')]
|
||
if texts:
|
||
combined = "; ".join(texts)
|
||
units.append({
|
||
"unit_id": f"FU-{unit_idx:03d}",
|
||
"name": combined[:60],
|
||
"description": combined[:200],
|
||
"sources": [{"section": source, "type": "table", "row": row.get("columns", [{}])[0].get("row", 1) if cols else 1, "text_snippet": combined[:200]}],
|
||
})
|
||
unit_idx += 1
|
||
|
||
return {
|
||
"feature_name": doc.get("sections", [{}])[0].get("source", "未命名功能"),
|
||
"concepts": [],
|
||
"function_units": units,
|
||
}
|
||
|
||
|
||
# =============================================================================
|
||
# Stage 2: Per-Unit IR Extraction
|
||
# =============================================================================
|
||
|
||
def build_context_package(fu: dict, doc: dict, image_sources: dict) -> dict:
|
||
"""Build a precision context package for one function unit."""
|
||
texts = []
|
||
tables = []
|
||
images = []
|
||
seen_sections = set()
|
||
|
||
img_by_rid = {img.get("rid", ""): img for img in doc.get("image_analysis", [])}
|
||
|
||
for src in fu.get("sources", []):
|
||
section_key = src.get("section", "")
|
||
src_type = src.get("type", "")
|
||
|
||
if src_type in ("para", "table") and section_key:
|
||
if section_key in seen_sections:
|
||
continue
|
||
seen_sections.add(section_key)
|
||
|
||
# Find matching section
|
||
for section in doc.get("sections", []):
|
||
if section.get("source", "") == section_key:
|
||
for block in section.get("blocks", []):
|
||
if block["type"] == "para":
|
||
texts.append({
|
||
"section": section_key,
|
||
"text": block["text"],
|
||
})
|
||
elif block["type"] == "table":
|
||
all_rows = []
|
||
for row in block.get("rows", []):
|
||
for col in row.get("columns", []):
|
||
all_rows.append({
|
||
"row": col.get("row"),
|
||
"name": col.get("name"),
|
||
"text": col.get("text"),
|
||
})
|
||
# Trim: max 20 cells per table to reduce prompt size
|
||
tables.append({
|
||
"section": section_key,
|
||
"headers": block.get("headers", []),
|
||
"all_rows": all_rows[:20],
|
||
"_truncated": len(all_rows) > 20,
|
||
})
|
||
break
|
||
|
||
if src_type in ("image", "logic_tree"):
|
||
img_id = src.get("image_id", "")
|
||
if img_id and img_id in img_by_rid:
|
||
img = img_by_rid[img_id]
|
||
images.append({
|
||
"image_id": img_id,
|
||
"type": img.get("type", "other"),
|
||
"description": img.get("description", ""),
|
||
})
|
||
|
||
return {
|
||
"unit_id": fu["unit_id"],
|
||
"unit_name": fu.get("name", ""),
|
||
"unit_description": fu.get("description", ""),
|
||
"texts": texts,
|
||
"tables": tables,
|
||
"images": images,
|
||
"resolved_conflicts": doc.get("resolved_conflicts", []),
|
||
}
|
||
|
||
|
||
async def stage2_extract_rules(semantic_index: dict, doc: dict) -> list[dict]:
|
||
"""Stage 2: Extract detailed IR rules for each function unit."""
|
||
logger.info("=" * 50)
|
||
logger.info("[PIPELINE] ═══ Stage 2/3: Per-Unit IR Extraction")
|
||
|
||
template_path = PROMPTS_DIR / "ir_step2_extraction.txt"
|
||
if not template_path.exists():
|
||
return _fallback_rules(semantic_index)
|
||
|
||
template = template_path.read_text(encoding="utf-8")
|
||
function_units = semantic_index.get("function_units", [])
|
||
image_sources = doc.get("image_sources", {})
|
||
|
||
logger.info("Processing %d function units (max concurrency: %d)", len(function_units), MAX_CONCURRENCY)
|
||
|
||
# Build context packages
|
||
packages = []
|
||
for fu in function_units:
|
||
pkg = build_context_package(fu, doc, image_sources)
|
||
packages.append(pkg)
|
||
|
||
# Parallel LLM calls
|
||
fragments = []
|
||
client = router.get_text_client()
|
||
|
||
def extract_one(pkg: dict) -> dict:
|
||
"""Extract rules for one function unit."""
|
||
prompt = _build_stage2_prompt(template, pkg)
|
||
unit_id = pkg["unit_id"]
|
||
|
||
for attempt in range(3):
|
||
try:
|
||
raw = client.chat(
|
||
model=router.text_model,
|
||
messages=[
|
||
{"role": "system", "content": "你是一个精确的 JSON 输出引擎。只输出合法的 JSON 数组，不输出任何其他文字。"},
|
||
{"role": "user", "content": prompt},
|
||
],
|
||
temperature=0.1,
|
||
)
|
||
json_str = extract_json_from_response(raw, expect_array=True)
|
||
rules = json.loads(json_str)
|
||
if not isinstance(rules, list):
|
||
raise ValueError(f"Expected array, got {type(rules).__name__}")
|
||
return {"unit_id": unit_id, "unit_name": pkg["unit_name"], "rules": rules}
|
||
except Exception as e:
|
||
if attempt < 2:
|
||
time.sleep(2)
|
||
else:
|
||
logger.warning("Unit %s failed after 3 attempts: %s", unit_id, e)
|
||
return {"unit_id": unit_id, "unit_name": pkg["unit_name"], "rules": [], "error": str(e)}
|
||
return {"unit_id": unit_id, "unit_name": pkg["unit_name"], "rules": [], "error": "unknown"}
|
||
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_CONCURRENCY) as executor:
|
||
futures = {executor.submit(extract_one, pkg): pkg for pkg in packages}
|
||
for future in concurrent.futures.as_completed(futures):
|
||
result = future.result()
|
||
fragments.append(result)
|
||
n_rules = len(result.get("rules", []))
|
||
logger.info("[PIPELINE] %s: %d rules %s", result["unit_id"], n_rules,
|
||
"❌" if result.get("error") else "✅")
|
||
|
||
fragments.sort(key=lambda f: f["unit_id"])
|
||
total = sum(len(f.get("rules", [])) for f in fragments)
|
||
logger.info("[PIPELINE] ✅ Stage 2 完成: %d fragments, %d total rules", len(fragments), total)
|
||
return fragments
|
||
|
||
|
||
def _build_stage2_prompt(template: str, pkg: dict) -> str:
|
||
"""Render the stage 2 prompt template with context package data."""
|
||
# Simple str.replace (avoids format issues with curly braces in JSON)
|
||
prompt = template
|
||
prompt = prompt.replace("{unit_id}", pkg["unit_id"])
|
||
prompt = prompt.replace("{unit_name}", pkg["unit_name"])
|
||
prompt = prompt.replace("{unit_description}", pkg["unit_description"])
|
||
prompt = prompt.replace("{texts}", json.dumps(pkg.get("texts", []), ensure_ascii=False, indent=2))
|
||
prompt = prompt.replace("{tables}", json.dumps(pkg.get("tables", []), ensure_ascii=False, indent=2))
|
||
prompt = prompt.replace("{images}", json.dumps(pkg.get("images", []), ensure_ascii=False, indent=2))
|
||
prompt = prompt.replace("{resolved_conflicts}", json.dumps(pkg.get("resolved_conflicts", []), ensure_ascii=False, indent=2))
|
||
return prompt
|
||
|
||
|
||
def _fallback_rules(semantic_index: dict) -> list[dict]:
|
||
"""Generate basic rules without LLM."""
|
||
fragments = []
|
||
for fu in semantic_index.get("function_units", []):
|
||
fragments.append({
|
||
"unit_id": fu["unit_id"],
|
||
"unit_name": fu.get("name", ""),
|
||
"rules": [{
|
||
"description": fu.get("description", ""),
|
||
"priority": "P2",
|
||
"sources": fu.get("sources", []),
|
||
"precondition": {},
|
||
"trigger": {"operator": "AND", "conditions": []},
|
||
"actions": [],
|
||
}],
|
||
})
|
||
return fragments
|
||
|
||
|
||
# =============================================================================
|
||
# Stage 3: Merge & Audit
|
||
# =============================================================================
|
||
|
||
def rule_signature(rule: dict) -> str:
|
||
"""SHA256 signature from trigger + actions for dedup."""
|
||
trigger = rule.get("trigger", {})
|
||
actions = rule.get("actions", [])
|
||
conditions = sorted(trigger.get("conditions", []), key=lambda c: json.dumps(c, sort_keys=True))
|
||
sorted_actions = sorted(actions, key=lambda a: json.dumps(a, sort_keys=True))
|
||
sig = json.dumps({"conditions": conditions, "actions": sorted_actions}, ensure_ascii=False, sort_keys=True)
|
||
return hashlib.sha256(sig.encode()).hexdigest()[:16]
|
||
|
||
|
||
def stage3_merge_and_audit(fragments: list[dict], semantic_index: dict, doc: dict) -> dict:
|
||
"""Stage 3: Merge rules, assign IDs, and generate audit report."""
|
||
logger.info("=" * 50)
|
||
logger.info("[PIPELINE] ═══ Stage 3/3: Merge & Audit")
|
||
|
||
# Merge and dedup
|
||
sig_map: dict[str, dict] = {}
|
||
order = []
|
||
for f in fragments:
|
||
for rule in f.get("rules", []):
|
||
sig = rule_signature(rule)
|
||
if sig in sig_map:
|
||
existing = sig_map[sig]
|
||
existing_sources = existing.setdefault("sources", [])
|
||
for src in rule.get("sources", []):
|
||
if src not in existing_sources:
|
||
existing_sources.append(src)
|
||
if len(rule.get("description", "")) > len(existing.get("description", "")):
|
||
existing["description"] = rule["description"]
|
||
else:
|
||
sig_map[sig] = dict(rule)
|
||
order.append(sig)
|
||
|
||
merged = [sig_map[sig] for sig in order]
|
||
logger.info("[PIPELINE] 合并: %d → %d rules (dedup)", sum(len(f.get("rules", [])) for f in fragments), len(merged))
|
||
|
||
# Assign rule IDs: IR-{feature_id}-{section}-{seq}
|
||
section_counters = defaultdict(int)
|
||
feature_name = semantic_index.get("feature_name", "FEATURE")
|
||
feature_id = re.sub(r'[^A-Z]', '', feature_name.upper())[:6] or "FEAT"
|
||
if len(feature_id) < 2:
|
||
feature_id = (feature_id + "FEATURE")[:4]
|
||
|
||
for rule in merged:
|
||
# Extract section from first para/table source, fallback to "0"
|
||
section = "0"
|
||
for src in rule.get("sources", []):
|
||
s = src.get("section", "")
|
||
if s:
|
||
# Extract first number pattern like "4.2.1" or "3.1"
|
||
m = re.search(r'[\d]+(?:\.[\d]+)*', s)
|
||
if m:
|
||
section = m.group()
|
||
break
|
||
section_counters[section] += 1
|
||
rule["rule_id"] = f"IR-{feature_id}-{section}-{section_counters[section]:03d}"
|
||
|
||
# Build final IR
|
||
ir_final = {
|
||
"feature": feature_name,
|
||
"feature_id": feature_id,
|
||
"rules": merged,
|
||
"rule_count": len(merged),
|
||
"generated_at": datetime.now(timezone.utc).isoformat(),
|
||
}
|
||
|
||
# Generate audit report
|
||
audit_results = _generate_audit(merged, doc)
|
||
report = _format_audit_report(feature_name, len(merged), audit_results)
|
||
|
||
logger.info("[PIPELINE] ✅ Stage 3 完成: %d rules, %d audit issues", len(merged), audit_results.get("issue_count", 0))
|
||
return {"ir": ir_final, "audit": audit_results, "audit_report": report}
|
||
|
||
|
||
def _generate_audit(rules: list[dict], doc: dict) -> dict:
|
||
"""Run coverage audits and return structured results."""
|
||
results = []
|
||
|
||
# 1. Source coverage: check that all sections with content are referenced
|
||
all_sources = set()
|
||
for rule in rules:
|
||
for src in rule.get("sources", []):
|
||
section = src.get("section", "")
|
||
if section:
|
||
all_sources.add(section)
|
||
|
||
doc_sections = set()
|
||
for section in doc.get("sections", []):
|
||
source = section.get("source", "")
|
||
if source and section.get("blocks"):
|
||
doc_sections.add(source)
|
||
|
||
uncovered = doc_sections - all_sources
|
||
results.append({
|
||
"check": "文档章节覆盖",
|
||
"status": PASS if not uncovered else WARN,
|
||
"detail": f"已覆盖 {len(all_sources)}/{len(doc_sections)} 章节"
|
||
+ (f"; 未覆盖: {list(uncovered)[:5]}" if uncovered else ""),
|
||
})
|
||
|
||
# 2. Image coverage
|
||
img_count = len(doc.get("image_analysis", []))
|
||
img_referenced = set()
|
||
for rule in rules:
|
||
for src in rule.get("sources", []):
|
||
if src.get("type") in ("image", "logic_tree"):
|
||
img_referenced.add(src.get("image_id", ""))
|
||
results.append({
|
||
"check": "图片覆盖率",
|
||
"status": PASS if len(img_referenced) >= img_count else (WARN if img_referenced else FAIL),
|
||
"detail": f"已引用 {len(img_referenced)}/{img_count} 张图片",
|
||
})
|
||
|
||
# 3. Priority distribution
|
||
p0 = sum(1 for r in rules if r.get("priority") == "P0")
|
||
p1 = sum(1 for r in rules if r.get("priority") == "P1")
|
||
p2 = sum(1 for r in rules if r.get("priority") == "P2")
|
||
results.append({
|
||
"check": "优先级分布",
|
||
"status": PASS if p0 > 0 else WARN,
|
||
"detail": f"P0={p0}, P1={p1}, P2={p2}",
|
||
})
|
||
|
||
issue_count = sum(1 for r in results if r["status"] in (WARN, FAIL))
|
||
return {"checks": results, "issue_count": issue_count}
|
||
|
||
|
||
def _sync_call(async_fn, *args, **kwargs):
|
||
"""Run an async function synchronously using the existing event loop if available."""
|
||
import asyncio
|
||
try:
|
||
loop = asyncio.get_running_loop()
|
||
except RuntimeError:
|
||
# No event loop running — just use asyncio.run
|
||
return asyncio.run(async_fn(*args, **kwargs))
|
||
else:
|
||
# Event loop is running (e.g., inside FastAPI) — must use run_until_complete
|
||
# Create a new loop in a thread to avoid conflicts
|
||
import concurrent.futures
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
|
||
future = ex.submit(asyncio.run, async_fn(*args, **kwargs))
|
||
return future.result()
|
||
|
||
|
||
def _format_audit_report(feature_name: str, rule_count: int, audit: dict) -> str:
|
||
"""Generate Markdown audit report."""
|
||
lines = [
|
||
f"# IR 完整性审计报告",
|
||
f"",
|
||
f"**功能**: {feature_name}",
|
||
f"**规则总数**: {rule_count}",
|
||
f"**生成时间**: {datetime.now(timezone.utc).isoformat()}",
|
||
f"",
|
||
f"> ⚠️ 请人工审查以下 ⚠️ 和 ❌ 项。",
|
||
f"",
|
||
f"## 审计结果",
|
||
f"",
|
||
f"| 检查项 | 状态 | 详情 |",
|
||
f"|--------|------|------|",
|
||
]
|
||
for r in audit.get("checks", []):
|
||
lines.append(f"| {r['check']} | {r['status']} | {r['detail']} |")
|
||
lines.append("")
|
||
lines.append(f"**审计问题总计**: {audit.get('issue_count', 0)} 个需要人工审查")
|
||
return "\n".join(lines)
|
||
|
||
|
||
# =============================================================================
|
||
# Pipeline Orchestrator
|
||
# =============================================================================
|
||
|
||
class IRPipeline:
|
||
"""Orchestrates the 3-stage IR generation pipeline."""
|
||
|
||
def __init__(self):
|
||
self._store: dict[str, dict] = {}
|
||
|
||
def run_streaming(self, prd_id: str, parsed_doc: dict):
|
||
"""Run pipeline as a sync generator yielding progress events.
|
||
|
||
This is a plain generator (NOT async) so it never blocks the event loop
|
||
between yields. The SSE endpoint wraps it in a thread + Queue bridge.
|
||
"""
|
||
t_total_start = time.time()
|
||
n_sections = len(parsed_doc.get("sections", []))
|
||
est_units = max(n_sections * 2, 3)
|
||
est_total = 20 + (est_units * 20 / 3) + 1
|
||
|
||
# --- Stage 1 ---
|
||
yield {"stage": 1, "status": "running", "message": "Stage 1: 分析文档结构，识别功能单元...",
|
||
"elapsed": 0, "stage_total": 3, "detail": f"文档含 {n_sections} 个章节"}
|
||
t0 = time.time()
|
||
try:
|
||
semantic_index = _sync_call(stage1_semantic_index, parsed_doc)
|
||
except Exception as e:
|
||
yield {"stage": 1, "status": "error", "message": f"Stage 1 失败: {e}", "stage_total": 3}
|
||
return
|
||
t1 = time.time()
|
||
n_units = len(semantic_index.get("function_units", []))
|
||
est_total = 20 + (max(n_units, 1) * 25 / 3) + 1
|
||
yield {"stage": 1, "status": "done", "message": f"Stage 1 完成: 识别 {n_units} 个功能单元",
|
||
"elapsed": round(t1 - t0, 1), "stage_total": 3, "detail": f"识别 {n_units} 个功能单元",
|
||
"n_units": n_units, "estimated_total": round(est_total)}
|
||
|
||
# --- Stage 2 ---
|
||
est_stage2 = max(n_units, 1) * 25 / min(MAX_CONCURRENCY, max(n_units, 1))
|
||
yield {"stage": 2, "status": "running", "message": f"Stage 2: 逐单元提取 IR 规则 ({n_units} 个单元)...",
|
||
"elapsed": 0, "stage_total": 3, "detail": f"预计 {est_stage2:.0f}s，{n_units} 个单元 × {min(MAX_CONCURRENCY, n_units)} 并发",
|
||
"n_units": n_units, "estimated_total": round(est_total)}
|
||
t1b = time.time()
|
||
completed = 0
|
||
total_rules = 0
|
||
|
||
fragments = []
|
||
function_units = semantic_index.get("function_units", [])
|
||
|
||
# Optimization: filter out trivial units (very short desc, no text sources)
|
||
filtered_units = []
|
||
for fu in function_units:
|
||
sources = fu.get("sources", [])
|
||
desc = fu.get("description", "")
|
||
has_text = any(s.get("type") in ("para", "table") for s in sources)
|
||
if len(desc) < 15 and not has_text:
|
||
logger.info("[PIPELINE] 跳过空单元: %s", fu.get("unit_id", "?"))
|
||
continue
|
||
filtered_units.append(fu)
|
||
|
||
skipped = len(function_units) - len(filtered_units)
|
||
if skipped:
|
||
logger.info("[PIPELINE] Stage 2 过滤掉 %d 个空单元，剩余 %d", skipped, len(filtered_units))
|
||
function_units = filtered_units
|
||
|
||
image_sources = parsed_doc.get("image_sources", {})
|
||
template_path = PROMPTS_DIR / "ir_step2_extraction.txt"
|
||
template = template_path.read_text(encoding="utf-8") if template_path.exists() else ""
|
||
client = router.get_text_client()
|
||
packages = [build_context_package(fu, parsed_doc, image_sources) for fu in function_units]
|
||
|
||
with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_CONCURRENCY) as executor:
|
||
futures = {}
|
||
for i, pkg in enumerate(packages):
|
||
future = executor.submit(self._extract_one_sync, template, pkg, client)
|
||
futures[future] = (i, pkg["unit_id"], pkg["unit_name"])
|
||
|
||
for future in concurrent.futures.as_completed(futures):
|
||
i, uid, uname = futures[future]
|
||
try:
|
||
rules = future.result()
|
||
fragments.append({"unit_id": uid, "unit_name": uname, "rules": rules})
|
||
total_rules += len(rules)
|
||
except Exception as e:
|
||
logger.warning("Unit %s failed: %s", uid, e)
|
||
fragments.append({"unit_id": uid, "unit_name": uname, "rules": [], "error": str(e)})
|
||
|
||
completed += 1
|
||
elapsed = round(time.time() - t1b, 1)
|
||
yield {"stage": 2, "status": "running",
|
||
"message": f"Stage 2: {completed}/{n_units} 单元完成 ({total_rules} 条规则)",
|
||
"elapsed": elapsed, "stage_total": 3,
|
||
"detail": f"已完成 {completed}/{n_units}，累计 {total_rules} 条规则",
|
||
"completed": completed, "total": n_units, "rules_so_far": total_rules,
|
||
"estimated_total": round(est_total)}
|
||
|
||
fragments.sort(key=lambda f: f["unit_id"])
|
||
t2 = time.time()
|
||
logger.info("[PIPELINE] ✅ Stage 2 完成: %d fragments, %d total rules", len(fragments), total_rules)
|
||
yield {"stage": 2, "status": "done",
|
||
"message": f"Stage 2 完成: {len(fragments)} 个单元 → {total_rules} 条规则",
|
||
"elapsed": round(t2 - t1b, 1), "stage_total": 3,
|
||
"detail": f"{len(fragments)} 个单元 → {total_rules} 条规则",
|
||
"total_rules": total_rules}
|
||
|
||
# --- Stage 3 ---
|
||
yield {"stage": 3, "status": "running", "message": "Stage 3: 合并去重 + 审计报告...",
|
||
"elapsed": 0, "stage_total": 3, "detail": "去重 + 分配 ID + 生成审计报告"}
|
||
result = stage3_merge_and_audit(fragments, semantic_index, parsed_doc)
|
||
t3 = time.time()
|
||
n_rules = len(result["ir"]["rules"])
|
||
|
||
ir_id = str(uuid.uuid4())[:8]
|
||
output = {
|
||
"ir_id": ir_id, "prd_id": prd_id,
|
||
"yaml_content": json.dumps(result["ir"], ensure_ascii=False, indent=2),
|
||
"ir_json": result["ir"],
|
||
"audit": result["audit"],
|
||
"audit_report": result["audit_report"],
|
||
"skill_used": "default",
|
||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||
"pipeline_stats": {
|
||
"stage1_seconds": round(t1 - t0, 1),
|
||
"stage2_seconds": round(t2 - t1b, 1),
|
||
"stage3_seconds": round(t3 - t2, 1),
|
||
"total_seconds": round(t3 - t_total_start, 1),
|
||
},
|
||
}
|
||
self._store[ir_id] = output
|
||
|
||
# Save to disk
|
||
output_dir = settings.OUTPUT_DIR / prd_id
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
with open(output_dir / f"ir_{ir_id}.json", "w", encoding="utf-8") as f:
|
||
json.dump(output, f, ensure_ascii=False, indent=2, default=str)
|
||
|
||
yield {"stage": 3, "status": "done", "message": f"流水线完成！共 {n_rules} 条规则",
|
||
"elapsed": round(t3 - t2, 1), "stage_total": 3,
|
||
"detail": f"合并后 {n_rules} 条规则，审计问题 {result['audit'].get('issue_count', 0)} 个",
|
||
"done": True, "result": output}
|
||
|
||
def _extract_one_sync(self, template: str, pkg: dict, client) -> list[dict]:
|
||
"""Synchronous rule extraction for one unit (runs in thread pool)."""
|
||
prompt = _build_stage2_prompt(template, pkg)
|
||
for attempt in range(3):
|
||
try:
|
||
raw = client.chat(
|
||
model=router.text_model,
|
||
messages=[
|
||
{"role": "system", "content": "你是一个精确的 JSON 输出引擎。只输出合法的 JSON 数组，不输出任何其他文字。"},
|
||
{"role": "user", "content": prompt},
|
||
],
|
||
temperature=0.1,
|
||
)
|
||
json_str = extract_json_from_response(raw, expect_array=True)
|
||
rules = json.loads(json_str)
|
||
if isinstance(rules, list):
|
||
return rules
|
||
except Exception as e:
|
||
if attempt == 2:
|
||
raise
|
||
time.sleep(2)
|
||
return []
|
||
|
||
async def run(self, prd_id: str, parsed_doc: dict) -> dict:
|
||
"""Run the full 3-stage pipeline and return the final IR + audit."""
|
||
logger.info("[PIPELINE] ═══ 流水线启动 for PRD %s", prd_id)
|
||
|
||
# Stage 1
|
||
t0 = time.time()
|
||
semantic_index = await stage1_semantic_index(parsed_doc)
|
||
t1 = time.time()
|
||
logger.info("[PIPELINE] Stage 1 耗时: %.1fs", t1 - t0)
|
||
|
||
# Stage 2
|
||
fragments = await stage2_extract_rules(semantic_index, parsed_doc)
|
||
t2 = time.time()
|
||
logger.info("[PIPELINE] Stage 2 耗时 %.1fs", t2 - t1)
|
||
|
||
# Stage 3
|
||
result = stage3_merge_and_audit(fragments, semantic_index, parsed_doc)
|
||
t3 = time.time()
|
||
logger.info("[PIPELINE] Stage 3 耗时 %.1fs", t3 - t2)
|
||
|
||
ir_id = str(uuid.uuid4())[:8]
|
||
output = {
|
||
"ir_id": ir_id,
|
||
"prd_id": prd_id,
|
||
"yaml_content": json.dumps(result["ir"], ensure_ascii=False, indent=2),
|
||
"ir_json": result["ir"],
|
||
"audit": result["audit"],
|
||
"audit_report": result["audit_report"],
|
||
"skill_used": "default",
|
||
"created_at": datetime.now(timezone.utc).isoformat(),
|
||
"pipeline_stats": {
|
||
"stage1_seconds": round(t1 - t0, 1),
|
||
"stage2_seconds": round(t2 - t1, 1),
|
||
"stage3_seconds": round(t3 - t2, 1),
|
||
"total_seconds": round(t3 - t0, 1),
|
||
},
|
||
}
|
||
self._store[ir_id] = output
|
||
|
||
# Save output
|
||
output_dir = settings.OUTPUT_DIR / prd_id
|
||
output_dir.mkdir(parents=True, exist_ok=True)
|
||
with open(output_dir / f"ir_{ir_id}.json", "w", encoding="utf-8") as f:
|
||
json.dump(output, f, ensure_ascii=False, indent=2, default=str)
|
||
|
||
logger.info("[PIPELINE] ✅ 流水线完成: ir_id=%s, rules=%d, total=%.1fs",
|
||
ir_id, len(result["ir"]["rules"]), t3 - t0)
|
||
return output
|
||
|
||
async def get_ir(self, ir_id: str) -> dict | None:
|
||
return self._store.get(ir_id)
|
||
|
||
|
||
ir_pipeline = IRPipeline()
|