Blame for core/generation_runner.py

# -*- coding: utf-8 -*-
"""R3-Free headless SmartPaper generation runner.

The runner contains no Streamlit calls. It is safe to execute in a subprocess
whose lifetime is independent from Streamlit ScriptRunner reruns and browser
connections. Checkpoints are task-scoped, immutable-by-sequence JSON snapshots.
"""
from __future__ import annotations

import hashlib
import json
import logging
import math
import os
import re
from datetime import datetime, timezone, timedelta
from pathlib import Path
from typing import Callable, Any

from config import APP_CONFIG, AMJ_PROMPTS, ASQ_PROMPTS
from core.checkpoint_store import TaskCheckpointStore
from core.meta_extractor import extract_golden_metadata
from core.pipeline_agent import SSciPipelineAgent
from core.production_guard import build_validated_package, replace_abstract
from core.qualitative_engine import get_qualitative_engine
from core.quantitative_engine import get_quantitative_engine
from core.rag_engine import initialize_rag
from utils.file_parser import extract_text_from_files

ProgressCallback = Callable[..., None]


def _clean_trailing_references(text, keywords=("reference", "references", "reference list", "参考文献")):
    lines = str(text or "").split("\n")
    for i in range(len(lines) - 1, max(-1, len(lines) - 100), -1):
        if any(re.match(rf"^##?\s*{kw}", lines[i], re.IGNORECASE) for kw in keywords):
            return "\n".join(lines[:i]).strip()
    return str(text or "").strip()


def _clean_markdown_headings(text):
    return re.sub(r"(?m)^#+\s+(.*)$", r"**\1**", str(text or ""))


def _canonical_request_hash(request_data: dict[str, Any]) -> str:
    material = dict(request_data or {})
    material.pop("request_hash", None)
    encoded = json.dumps(material, sort_keys=True, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


class GenerationRunner:
    def __init__(
        self,
        *,
        task_dir: str | os.PathLike,
        request_data: dict[str, Any],
        execution_epoch: int,
        progress_callback: ProgressCallback | None = None,
        cancel_checker: Callable[[], bool] | None = None,
    ):
        self.task_dir = Path(task_dir)
        self.request = dict(request_data)
        self.execution_epoch = int(execution_epoch)
        self.progress_callback = progress_callback or (lambda **_: None)
        self.cancel_checker = cancel_checker or (lambda: False)
        expected_hash = str(self.request.get("request_hash") or "")
        actual_hash = _canonical_request_hash(self.request)
        if not expected_hash or expected_hash != actual_hash:
            raise RuntimeError("request_hash_mismatch")
        self.request_hash = expected_hash
        self.checkpoints = TaskCheckpointStore(self.task_dir, self.request_hash, self.execution_epoch)
        self.output_dir = self.task_dir / "outputs"
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.agent = None

    def _progress(self, pct: int, node: str, *, checkpoint_seq=None, checkpoint_stage=None) -> None:
        agent = self.agent
        self.progress_callback(
            progress_pct=max(0, min(100, int(pct))),
            current_node=str(node or "")[:240],
            checkpoint_seq=checkpoint_seq,
            checkpoint_stage=checkpoint_stage,
            ds_calls=int(getattr(agent, "ds_calls", 0) or 0),
            s2_calls=int(getattr(agent, "s2_calls", 0) or 0),
            oa_calls=int(getattr(agent, "oa_calls", 0) or 0),
        )

    def _checkpoint(self, stage: str, state: dict[str, Any], pct: int) -> None:
        seq, normalized_stage = self.checkpoints.save(stage, state)
        self._progress(pct, normalized_stage, checkpoint_seq=seq, checkpoint_stage=normalized_stage)

    @staticmethod
    def _stream_text(generator) -> str:
        return "".join(str(chunk or "") for chunk in generator)

    @staticmethod
    def _bounded_target(base_target: int, accumulated_deviation: float, remaining_including_current: int, safe_chunk: int) -> int:
        base = max(80, int(base_target))
        correction = -float(accumulated_deviation) / max(1, int(remaining_including_current))
        lower = max(80, int(round(base * 0.75)))
        upper = min(int(safe_chunk * 1.5), max(lower, int(round(base * 1.25))))
        return max(lower, min(upper, int(round(base + correction))))

    @staticmethod
    def _english_count_for_target(manuscript: str) -> int:
        body = re.split(r"(?m)^#\s+References\s*$", str(manuscript or ""), maxsplit=1)[0]
        # Markdown punctuation is ignored; contractions/hyphenated terms stay one token.
        return len(re.findall(r"\b[A-Za-z0-9]+(?:[-'][A-Za-z0-9]+)*\b", body))

    def _validate_target_length(self, manuscript: str, target_words: int) -> None:
        quality = APP_CONFIG.get("manuscript_validation", {})
        minimum = int(target_words * float(quality.get("minimum_total_ratio", 0.85)))
        maximum = int(target_words * float(quality.get("maximum_total_ratio", 1.15)))
        actual = self._english_count_for_target(manuscript)
        if actual < minimum or actual > maximum:
            raise RuntimeError(
                f"English manuscript length outside validated range: actual={actual}, "
                f"required={minimum}-{maximum}, target={target_words}."
            )

    def _resolve_paths(self, items) -> list[str]:
        output = []
        for item in items or []:
            path = (self.task_dir / str(item)).resolve()
            if self.task_dir.resolve() not in path.parents:
                raise RuntimeError("input_path_escape")
            if path.is_file():
                output.append(str(path))
        return output

    def run(self) -> dict[str, Any]:
        req = self.request
        mode = str(req.get("mode") or "").lower()
        if mode not in {"amj", "asq"}:
            raise RuntimeError("invalid_generation_mode")
        target_words = int(req.get("target_words") or 0)
        if not APP_CONFIG["target_words_min"] <= target_words <= APP_CONFIG["target_words_max"]:
            raise RuntimeError("target_words_out_of_range")

        ds_key = str(os.environ.get("DEEPSEEK_API_KEY") or "").strip()
        s2_key = str(os.environ.get("S2_API_KEY") or "").strip()
        if not ds_key:
            raise RuntimeError("DEEPSEEK_API_KEY_missing")

        if mode == "amj":
            APP_CONFIG.update(AMJ_PROMPTS)
        else:
            APP_CONFIG.update(ASQ_PROMPTS)

        agent = SSciPipelineAgent(ds_key, s2_key)
        self.agent = agent
        agent.set_cancel_checker(self.cancel_checker)
        # DeepSeek/search activity marks the executor as genuinely alive; the
        # heartbeat thread decides separately when to renew D1's lease.
        agent.set_execution_progress_hook(lambda: self._progress_activity())
        safe_chunk_words = int(agent.chunk_words_default)
        quant_engine = get_quantitative_engine()
        qual_engine = get_qualitative_engine(agent)

        state = self.checkpoints.load_latest() or {}
        if state:
            logging.info(
                "[SP_EXECUTOR] resume checkpoint seq=%s stage=%s",
                state.get("checkpoint_seq"), state.get("checkpoint_stage"),
            )

        topic = str(req.get("topic") or "").strip()
        keywords = str(req.get("keywords") or "").strip()
        target_lang = str(req.get("target_lang") or "").strip()
        output_suffix = str(req.get("output_suffix") or "Target").strip()
        interface_language = str(req.get("interface_language") or "zh").strip()
        reference_paths = self._resolve_paths(req.get("reference_files") or [])
        data_paths = self._resolve_paths(req.get("data_files") or [])
        interview_paths = self._resolve_paths(req.get("interview_files") or [])

        # Restore agent state before continuing any dependent step.
        if isinstance(state.get("agent_state"), dict):
            agent.load_state(state["agent_state"])

        user_lit_context = str(state.get("user_lit_context") or "")
        user_data_md = str(state.get("user_data_md") or "")
        data_metadata = dict(state.get("data_metadata") or {})

        if not state.get("evidence_ready"):
            self._progress(3, "inputs_and_evidence")
            if reference_paths:
                user_lit_context = extract_text_from_files(reference_paths[:20])
                all_paper_texts = [extract_text_from_files([path]) for path in reference_paths[:20]]
                agent.global_memory["core_golden_metadata"] = extract_golden_metadata(agent, all_paper_texts)
            if mode == "amj":
                data_file = data_paths[0] if data_paths else None
                _, user_data_md, _, data_metadata = quant_engine.resolve_empirical_data(data_file, interface_language)
            else:
                raw_interviews = extract_text_from_files(interview_paths[:1]) if interview_paths else ""
                user_data_md, _, data_metadata = qual_engine.resolve_qualitative_data(topic, keywords, raw_interviews)
            state.update(
                {
                    "topic": topic,
                    "keywords": keywords,
                    "target_lang": target_lang,
                    "target_words": target_words,
                    "mode": mode,
                    "interface_language": interface_language,
                    "user_lit_context": user_lit_context,
                    "user_data_md": user_data_md,
                    "data_metadata": data_metadata,
                    "agent_state": agent.export_state(),
                    "evidence_ready": True,
                }
            )
            self._checkpoint("evidence_ready", state, 6)

        refs_text = str(state.get("refs_text") or "")
        if not refs_text:
            self._progress(8, "S2 + OpenAlex")
            refs_text = agent.fetch_references_parallel(
                topic,
                keywords,
                min_required=50,
                progress_callback=lambda: self._progress(9, "S2 + OpenAlex"),
            )
            if len(agent.citation_map) < 50:
                raise RuntimeError(f"Insufficient eligible references: {len(agent.citation_map)}/50")
            if user_lit_context:
                core_refs = agent.parse_user_core_literature(user_lit_context)
                if core_refs:
                    refs_text = (
                        "【USER PROVIDED CORE LITERATURE】\n" + core_refs
                        + "\n\n【RETRIEVED LITERATURE】\n" + refs_text
                    )
            state.update({"refs_text": refs_text, "agent_state": agent.export_state()})
            self._checkpoint("literature_ready", state, 11)

        rag_engine = initialize_rag([value for value in (user_lit_context, refs_text) if value])

        synthesis_result = str(state.get("synthesis_result") or "")
        if not synthesis_result:
            self._progress(12, "literature_synthesis")
            synthesis_result = self._stream_text(
                agent.stream_synthesize_literature(topic, refs_text, user_core_literature=user_lit_context)
            )
            synthesis_result = re.sub(r"<think>.*?</think>", "", synthesis_result, flags=re.DOTALL).strip()
            state.update({"synthesis_result": synthesis_result, "agent_state": agent.export_state()})
            self._checkpoint("synthesis_ready", state, 13)

        adaptive_structure = state.get("adaptive_structure")
        if not isinstance(adaptive_structure, list) or not adaptive_structure:
            self._progress(14, "anchors_and_outline")
            anchor_response = self._stream_text(
                agent.stream_generate_anchors(
                    topic, keywords, synthesis_result, refs_text,
                    data_evidence=user_data_md, mode=mode,
                )
            )
            anchor_clean = re.sub(r"<think>.*?</think>", "", anchor_response, flags=re.DOTALL).strip()
            abstract_match = re.search(
                r"-+\s*ABSTRACT\s*-+\s*(.*?)(?=\s*-+\s*CONCLUSION\s*-+|$)",
                anchor_clean, re.DOTALL | re.IGNORECASE,
            )
            conclusion_match = re.search(
                r"-+\s*CONCLUSION\s*-+\s*(.*?)$", anchor_clean, re.DOTALL | re.IGNORECASE,
            )
            abstract_anchor = abstract_match.group(1).strip() if abstract_match else anchor_clean
            conclusion_anchor = conclusion_match.group(1).strip() if conclusion_match else ""
            agent.global_memory["abstract_anchor"] = abstract_anchor
            agent.global_memory["conclusion_anchor"] = conclusion_anchor

            plan = agent.generate_outline(topic, keywords, mode=mode, data_metadata=data_metadata)
            en_topic = str(plan.get("en_topic", topic))
            structure = plan.get("structure", [])
            en_abstract_clean = _clean_markdown_headings(
                re.sub(r"(?i)^abstract[:\s]*", "", str(abstract_anchor)).strip()
            )
            full_draft_en = f"# {en_topic}\n\n## Abstract\n{en_abstract_clean}\n\n## Keywords\n{keywords}\n\n"
            adaptive_structure = agent.build_generation_tasks(structure, target_words, safe_chunk_words)
            for task in adaptive_structure:
                task["is_first"] = bool(task.get("section_first"))
            jit_texts = agent.prefetch_jit_literature(topic, keywords, adaptive_structure)
            if jit_texts:
                rag_engine._initialize_chunks([jit_texts])
            state.update(
                {
                    "adaptive_structure": adaptive_structure,
                    "full_draft_en": full_draft_en,
                    "full_draft_translated": "",
                    "completed_sections": [],
                    "context_tail": "",
                    "accumulated_deviation": 0.0,
                    "agent_state": agent.export_state(),
                    "outline_ready": True,
                }
            )
            self._checkpoint("outline_ready", state, 15)
        else:
            full_draft_en = str(state.get("full_draft_en") or "")
            # JIT text itself is already represented in agent citation state;
            # refs + user literature rebuild a safe RAG baseline on resume.

        completed_sections = list(state.get("completed_sections") or [])
        context_tail = str(state.get("context_tail") or "")
        accumulated_deviation = float(state.get("accumulated_deviation") or 0.0)
        all_cite_keys = list(agent.citation_map.keys())
        valid_chunks = [
            item for item in adaptive_structure
            if not any(word in item["section"].lower() for word in ["abstract", "keyword", "appendix"])
        ]
        allocation_map: dict[str, list[str]] = {}
        current_cite_idx = 0
        for idx, item in enumerate(valid_chunks):
            count = max(1, math.ceil(len(all_cite_keys) / max(1, len(valid_chunks))))
            allocation_map[item["display_name"]] = (
                all_cite_keys[current_cite_idx:]
                if idx == len(valid_chunks) - 1
                else all_cite_keys[current_cite_idx:current_cite_idx + count]
            )
            current_cite_idx += count

        total_steps = len(adaptive_structure)
        for sec in adaptive_structure:
            display_name = sec["display_name"]
            if display_name in completed_sections:
                continue
            if self.cancel_checker():
                raise RuntimeError("task_cancel_requested")
            remaining = sum(1 for item in adaptive_structure if item["display_name"] not in completed_sections)
            adjusted_target = self._bounded_target(
                int(sec["target_words"]), accumulated_deviation, remaining, safe_chunk_words
            )
            sec_title = sec["section"]
            subsection = sec.get("subsection")
            is_quarantined = any(
                word in sec_title.lower()
                for word in ["abstract", "keyword", "conclusion", "references", "appendix"]
            )
            assigned_tags = ", ".join(allocation_map.get(display_name, [])) or "the single most relevant source"
            citation_rule = (
                "NO IN-TEXT CITATIONS ALLOWED IN THIS SECTION."
                if is_quarantined
                else "Use claim-specific citation. Each sentence may contain at most one source tag. "
                     f"Candidate tags: {assigned_tags}. <AVAILABLE_CITATIONS>{refs_text}</AVAILABLE_CITATIONS>"
            )
            memory_context = agent._build_memory_context(sec_title, 2500)
            memory_block = f"\n<GLOBAL_MEMORY>\n{memory_context}\n</GLOBAL_MEMORY>\n" if memory_context else ""
            snippets = rag_engine.get_relevant_context(display_name) if rag_engine is not None else ""
            if snippets:
                memory_block += f"\n<TARGETED_RAG>\n{snippets}\n</TARGETED_RAG>\n"
            heading_directive = "Do not output the section heading; it will be inserted by the system."
            if subsection:
                heading_directive += (
                    f" Do not output the subsection heading '{subsection}'; it will also be inserted by the system. "
                    "Use level-3 headings only when analytically necessary."
                )
            elif sec_title.lower() in {"introduction", "conclusion"}:
                heading_directive += " Do not use additional headings."
            else:
                heading_directive += (
                    " Use level-2 or level-3 headings only when they clarify the argument; "
                    "never use level-4 headings or numbered headings."
                )
            table_rule = (
                "Present statistical evidence in standard Markdown tables with plain-text captions."
                if sec["is_method_or_result"] and mode == "amj" else ""
            )
            concepts = ", ".join(sec.get("concepts") or [])
            qualitative_rule = (
                f"Assigned qualitative concepts: {concepts}. Use only their assigned representative quotations "
                "and do not repeat a quotation already used elsewhere."
                if concepts else ""
            )
            data_injection = (
                f"\n<RESEARCH_EVIDENCE>\n{user_data_md}\n</RESEARCH_EVIDENCE>\n"
                if user_data_md and sec["is_method_or_result"] else ""
            )
            completed_label = ", ".join(completed_sections)
            anti_repeat = (
                f"Already completed: {completed_label}. Do not restate their background, gap, or contribution."
                if completed_label else ""
            )
            sys_prompt = APP_CONFIG.get("sys_academic_writer", "") + "\n" + APP_CONFIG.get("deep_refinement_directive", "")
            user_prompt = f"""Write the assigned English manuscript unit.
Section: {sec_title}
Subsection: {subsection or 'None'}
Target length: {int(adjusted_target * 0.9)}-{int(adjusted_target * 1.1)} words.
{heading_directive}
{citation_rule}
{table_rule}
{str(sec.get('required_content') or '')}
{qualitative_rule}
{anti_repeat}
{memory_block}
Literature synthesis: {synthesis_result}
{data_injection}
Previous tail: {context_tail[-2500:]}
Output only final academic prose. Do not report the word count."""
            if agent._count_tokens(sys_prompt) + agent._count_tokens(user_prompt) > 60000:
                user_prompt = agent._truncate_context(user_prompt, 60000 - agent._count_tokens(sys_prompt))

            current_index = len(completed_sections) + 1
            pct = min(88, 15 + int((current_index / max(1, total_steps)) * 72))
            self._progress(pct, display_name)
            en_raw = self._stream_text(
                agent.stream_deepseek(
                    sys_prompt, user_prompt, max_tokens=7000,
                    model_role="writing", thinking=False, allow_fallback=False,
                )
            )
            en_raw = re.sub(r"<think>.*?</think>", "", en_raw, flags=re.S).strip()
            en_raw = re.sub(rf"(?im)^#\s*{re.escape(sec_title)}\s*\n+", "", en_raw).strip()
            if subsection:
                en_raw = re.sub(rf"(?im)^##\s*{re.escape(subsection)}\s*\n+", "", en_raw).strip()
            en_raw = re.sub(r"(?m)^(#{2,3})\s+[\d\.\s一二三四五六七八九十、]+\s*(.*)", r"\1 \2", en_raw)
            en_raw = re.sub(r"(?m)^#{4,}\s+(.*)", r"**\1**", en_raw)
            en_raw = re.sub(r"(?m)^#\s+(.+)$", r"## \1", en_raw)
            if mode == "asq" and user_data_md and "appendix" not in sec_title.lower():
                quote_dict = {tag: quote for quote, tag in re.findall(r'"([^"]+)"\s*\[(QUOTE-\d+)\]', user_data_md)}
                for tag, quote in quote_dict.items():
                    en_raw = en_raw.replace(tag, f'"{quote}"')
            en_raw = agent.polish_academic_section(en_raw, display_name, abstract=False)
            prefix = ""
            if sec.get("section_first"):
                prefix += f"# {sec_title}\n\n"
            if sec.get("subsection_first") and subsection:
                prefix += f"## {subsection}\n\n"
            en_resolved = agent.audit_and_resolve_citations(
                prefix + en_raw, section_name=display_name, quarantine=is_quarantined
            )
            full_draft_en += en_resolved + "\n\n"
            agent._extract_memory_from_text(en_raw, sec_title)
            context_tail = " ".join(en_raw.split()[-200:])
            actual_words = len(en_raw.split())
            accumulated_deviation += actual_words - adjusted_target
            completed_sections.append(display_name)
            state.update(
                {
                    "completed_sections": completed_sections,
                    "full_draft_en": full_draft_en,
                    "full_draft_translated": "",
                    "context_tail": context_tail,
                    "accumulated_deviation": accumulated_deviation,
                    "agent_state": agent.export_state(),
                    "data_metadata": data_metadata,
                }
            )
            self._checkpoint(f"section_{len(completed_sections):03d}", state, pct)

        if not state.get("english_finalized"):
            self._progress(91, "english_validation")
            full_draft_en = _clean_trailing_references(full_draft_en, ["Reference", "References", "Reference List"])
            final_abstract = agent.generate_final_abstract(
                topic, keywords, full_draft_en, data_evidence=user_data_md, mode=mode
            )
            full_draft_en = replace_abstract(full_draft_en, final_abstract)
            agent.validate_manuscript_structure(full_draft_en, adaptive_structure, translated=False)
            if mode == "asq":
                repeated_quotes = agent.find_repeated_qualitative_quotes(full_draft_en, data_metadata)
                if repeated_quotes:
                    raise RuntimeError(f"Repeated qualitative quotations: {repeated_quotes[:5]}")
            references = []
            seen_refs = set()
            for tag in sorted(agent.used_citations, key=agent._citation_tag_number):
                ref = str((agent.citation_map.get(tag) or {}).get("ref") or "").strip()
                if ref and ref not in seen_refs:
                    references.append(f"- {ref}")
                    seen_refs.add(ref)
            if not references:
                raise RuntimeError("No cited references remained after citation validation.")
            full_draft_en = full_draft_en.rstrip() + "\n\n# References\n" + "\n\n".join(references)
            self._validate_target_length(full_draft_en, target_words)
            state.update(
                {
                    "full_draft_en": full_draft_en,
                    "english_finalized": True,
                    "agent_state": agent.export_state(),
                }
            )
            self._checkpoint("english_complete", state, 93)
        else:
            full_draft_en = str(state.get("full_draft_en") or "")

        full_draft_translated = str(state.get("full_draft_translated") or "")
        if not state.get("translation_complete") or not full_draft_translated:
            self._progress(95, "translation")
            full_draft_translated = agent.translate_locked_manuscript(full_draft_en, target_lang)
            agent.validate_translation_alignment(full_draft_en, full_draft_translated)
            state.update(
                {
                    "full_draft_translated": full_draft_translated,
                    "translation_complete": True,
                    "agent_state": agent.export_state(),
                }
            )
            self._checkpoint("translation_complete", state, 96)

        existing_output = str(state.get("output_filename") or "").strip()
        if state.get("package_validated") and existing_output:
            existing_path = self.output_dir / os.path.basename(existing_output)
            if existing_path.is_file() and existing_path.stat().st_size > 0:
                self._progress(99, "package_validated")
                return {"output_filename": existing_path.name, "output_path": str(existing_path)}

        self._progress(97, "packaging")
        bj_time = datetime.now(timezone(timedelta(hours=8))).strftime("%Y%m%d_%H%M")
        package = build_validated_package(
            english_markdown=full_draft_en,
            translated_markdown=full_draft_translated,
            english_filename=f"Research_EN_{bj_time}.docx",
            translated_filename=f"Research_{output_suffix}_{bj_time}.docx",
            data_metadata=data_metadata,
            mode=mode,
            citation_audit=agent.citation_audit_log,
            literature_metrics=agent.get_source_metrics(),
        )
        output_filename = f"SmartPaper_{bj_time}.zip"
        pending = self.output_dir / (".pending_" + output_filename)
        pending.write_bytes(package.zip_bytes)
        final = self.output_dir / output_filename
        os.replace(pending, final)
        state.update(
            {
                "output_filename": output_filename,
                "package_validated": True,
                "agent_state": agent.export_state(),
            }
        )
        self._checkpoint("package_validated", state, 99)
        return {"output_filename": output_filename, "output_path": str(final)}

    def _progress_activity(self) -> None:
        # This call intentionally does not update D1 directly. The executor's
        # heartbeat thread uses it only to prove that real backend work is still
        # occurring, which avoids browser-driven or timer-only zombie liveness.
        try:
            self.progress_callback(activity_only=True)
        except TypeError:
            pass