zeekerWatchmen/ref/ir_generation_new_skill/step1_semantic_index.py

"""
Stage 1: Semantic Index Generation.

Reads the full parsed PRD JSON, calls the LLM once to produce a semantic index
identifying all function units, concepts, and their document sources.

Output: output/semantic_index.json
"""

import json
import re
import sys
import time
from pathlib import Path

import config


def format_document_for_prompt(doc: dict) -> str:
    """Render the full parsed document as a readable string for the LLM prompt."""
    lines = []

    # ---- Sections ----
    lines.append("=== 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', '')[:300]}")

        lt = img.get("logic_tree")
        if lt:
            lines.append(f"  逻辑树根节点: {lt.get('root', '?')}")
            lines.append("  节点详情:")
            for node in lt.get("nodes", []):
                nid = node.get("id", "?")
                ntype = node.get("type", "?")
                desc = node.get("description", "") or node.get("condition", "")
                lines.append(f"    [{ntype}] {nid}: {desc}")
                branches = node.get("branches", [])
                if branches:
                    for br in branches:
                        lines.append(f"      → {br['value']} → {br['target']}")

    # ---- 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 build_prompt(doc: dict) -> str:
    """Load the prompt template and inject the formatted document."""
    template_path = Path(config.PROMPTS_DIR) / "step1_semantic_index.txt"
    template = template_path.read_text(encoding="utf-8")

    formatted_doc = format_document_for_prompt(doc)
    prompt = template.replace("{document_json}", formatted_doc)
    return prompt


def extract_json_from_response(text: str) -> str:
    """Robustly extract JSON from LLM response, handling markdown fences."""
    # Try ```json ... ``` first
    m = re.search(r"```(?:json)?\s*([\s\S]*?)```", text)
    if m:
        return m.group(1).strip()

    # Try to find the outermost { ... }
    start = text.find("{")
    if start == -1:
        raise ValueError("No JSON object found in LLM response")

    depth = 0
    for i in range(start, len(text)):
        if text[i] == "{":
            depth += 1
        elif text[i] == "}":
            depth -= 1
            if depth == 0:
                return text[start : i + 1]

    raise ValueError("Unclosed JSON object in LLM response")


def call_llm(prompt: str, max_retries: int = 2) -> dict:
    """Send the prompt to the LLM and return the parsed semantic index JSON."""
    client = config.llm_client()

    for attempt in range(max_retries + 1):
        print(f"  LLM 调用 (尝试 {attempt + 1}/{max_retries + 1})...", flush=True)
        try:
            resp = client.chat.completions.create(
                model=config.MODEL_NAME,
                messages=[
                    {"role": "system", "content": "你是一个精确的 JSON 输出引擎。只输出合法的 JSON，不输出任何其他文字。"},
                    {"role": "user", "content": prompt},
                ],
                temperature=config.TEMPERATURE,
                max_tokens=config.MAX_TOKENS,
            )
            content = resp.choices[0].message.content
            if content is None:
                raise RuntimeError("LLM returned empty response")

            json_str = extract_json_from_response(content)
            result = json.loads(json_str)

            # Quick structural validation
            _validate_schema(result)
            return result

        except (json.JSONDecodeError, ValueError) as e:
            print(f"  JSON 解析失败: {e}")
            if attempt < max_retries:
                time.sleep(2)
            else:
                raise RuntimeError(f"无法从 LLM 响应中解析 JSON: {e}") from e


def _validate_schema(data: dict) -> None:
    """Validate the top-level schema of the semantic index."""
    if "feature_name" not in data:
        raise ValueError("semantic_index 缺少 'feature_name' 字段")
    if "function_units" not in data:
        raise ValueError("semantic_index 缺少 'function_units' 字段")

    units = data["function_units"]
    if not isinstance(units, list) or len(units) == 0:
        raise ValueError("function_units 必须是非空数组")

    for i, fu in enumerate(units):
        if not fu.get("unit_id"):
            raise ValueError(f"function_unit[{i}] 缺少 unit_id")
        if not fu.get("name"):
            raise ValueError(f"function_unit[{i}] 缺少 name")


def main():
    print("=" * 60)
    print("阶段一：宏观语义索引 (Semantic Index)")
    print("=" * 60)

    # 1. Load input
    print(f"\n[1/4] 加载输入文档: {config.INPUT_JSON}")
    doc = config.load_input_document()
    print(f"  已加载 {len(doc.get('sections', []))} 个 section, "
          f"{len(doc.get('image_analysis', []))} 张图片分析")

    # 2. Build prompt
    print(f"\n[2/4] 构造 Prompt...")
    prompt = build_prompt(doc)
    print(f"  Prompt 长度: {len(prompt)} 字符 (~{len(prompt)//3} tokens)")

    # 3. Call LLM
    print(f"\n[3/4] 调用 LLM ({config.MODEL_NAME})...")
    semantic_index = call_llm(prompt)

    # 4. Save result
    print(f"\n[4/4] 保存语义索引: {config.SEMANTIC_INDEX_JSON}")
    config.save_json(semantic_index, config.SEMANTIC_INDEX_JSON)

    n_concepts = len(semantic_index.get("concepts", []))
    n_units = len(semantic_index.get("function_units", []))
    print(f"\n完成! 提取了 {n_concepts} 个概念, {n_units} 个功能单元.")
    print(f"输出: {config.SEMANTIC_INDEX_JSON}")


if __name__ == "__main__":
    main()