init the project
This commit is contained in:
evyzacq
@@ -0,0 +1,753 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user