source-kb 0.2.2__py3-none-any.whl

engine/pipeline/recovery.py

@@ -0,0 +1,250 @@
+"""Checkpoint and recovery module for pipeline execution.
+
+Provides:
+- Checkpoint dataclass for recording step completion state
+- save()/load() methods for persistence
+- --resume breakpoint resume logic
+- Corruption detection and rollback
+
+Requirements: 19.1, 19.2, 19.3, 19.4
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+from dataclasses import dataclass, field, asdict
+from pathlib import Path
+from typing import Any
+
+from core.interfaces import LlmStrategy
+from engine.runner import SubagentTask, SubagentResult, run_single
+
+logger = logging.getLogger(__name__)
+
+CHECKPOINT_FILENAME = ".pipeline-checkpoint.json"
+CHECKPOINT_VERSION = 2
+
+
+# ---------------------------------------------------------------------------
+# Checkpoint dataclass
+# ---------------------------------------------------------------------------
+
+@dataclass
+class StepState:
+    """State of a single pipeline step."""
+    name: str
+    status: str  # "pending", "running", "done", "failed", "skipped"
+    started_at: float = 0.0
+    completed_at: float = 0.0
+    message: str = ""
+    details: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class Checkpoint:
+    """Pipeline execution checkpoint for resume support."""
+    pipeline_name: str
+    kb_name: str
+    version: int = CHECKPOINT_VERSION
+    created_at: float = field(default_factory=time.time)
+    updated_at: float = field(default_factory=time.time)
+    steps: list[StepState] = field(default_factory=list)
+    state_snapshot: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def completed_steps(self) -> list[str]:
+        """Names of steps that completed successfully."""
+        return [s.name for s in self.steps if s.status in ("done", "skipped")]
+
+    @property
+    def last_completed(self) -> str | None:
+        """Name of the last completed step."""
+        completed = self.completed_steps
+        return completed[-1] if completed else None
+
+    @property
+    def next_step(self) -> str | None:
+        """Name of the next step to execute (first pending/failed)."""
+        for s in self.steps:
+            if s.status in ("pending", "failed", "running"):
+                return s.name
+        return None
+
+    def mark_step(self, name: str, status: str, message: str = "", **details) -> None:
+        """Update a step's status."""
+        for s in self.steps:
+            if s.name == name:
+                s.status = status
+                s.message = message
+                s.details.update(details)
+                if status == "running":
+                    s.started_at = time.time()
+                elif status in ("done", "failed", "skipped"):
+                    s.completed_at = time.time()
+                break
+        self.updated_at = time.time()
+
+
+# ---------------------------------------------------------------------------
+# Save / Load
+# ---------------------------------------------------------------------------
+
+def save_checkpoint(checkpoint: Checkpoint, knowledge_dir: Path) -> Path:
+    """Save checkpoint to disk.
+
+    Args:
+        checkpoint: Checkpoint to save
+        knowledge_dir: Knowledge directory (checkpoint stored in .meta/)
+
+    Returns:
+        Path to saved checkpoint file
+    """
+    checkpoint.updated_at = time.time()
+
+    cp_dir = knowledge_dir / ".meta"
+    cp_dir.mkdir(parents=True, exist_ok=True)
+    cp_path = cp_dir / CHECKPOINT_FILENAME
+
+    data = {
+        "version": checkpoint.version,
+        "pipeline_name": checkpoint.pipeline_name,
+        "kb_name": checkpoint.kb_name,
+        "created_at": checkpoint.created_at,
+        "updated_at": checkpoint.updated_at,
+        "steps": [asdict(s) for s in checkpoint.steps],
+        "state_snapshot": _sanitize_state(checkpoint.state_snapshot),
+    }
+
+    # Write atomically (write to temp then rename)
+    tmp_path = cp_path.with_suffix(".tmp")
+    tmp_path.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
+    tmp_path.replace(cp_path)
+
+    logger.debug("[checkpoint] Saved: %s (%d steps)", cp_path, len(checkpoint.steps))
+    return cp_path
+
+
+def load_checkpoint(knowledge_dir: Path) -> Checkpoint | None:
+    """Load checkpoint from disk.
+
+    Returns None if no checkpoint exists or if JSON is corrupted.
+    """
+    cp_path = knowledge_dir / ".meta" / CHECKPOINT_FILENAME
+    if not cp_path.exists():
+        return None
+
+    try:
+        data = json.loads(cp_path.read_text(encoding="utf-8"))
+    except (json.JSONDecodeError, OSError) as e:
+        logger.warning("[checkpoint] Corrupted checkpoint: %s", e)
+        return None
+
+    # Version check
+    if data.get("version", 0) != CHECKPOINT_VERSION:
+        logger.warning("[checkpoint] Version mismatch (got %s, expected %s)",
+                       data.get("version"), CHECKPOINT_VERSION)
+        return None
+
+    steps = [
+        StepState(
+            name=s["name"],
+            status=s.get("status", "pending"),
+            started_at=s.get("started_at", 0),
+            completed_at=s.get("completed_at", 0),
+            message=s.get("message", ""),
+            details=s.get("details", {}),
+        )
+        for s in data.get("steps", [])
+    ]
+
+    return Checkpoint(
+        pipeline_name=data.get("pipeline_name", ""),
+        kb_name=data.get("kb_name", ""),
+        version=data.get("version", CHECKPOINT_VERSION),
+        created_at=data.get("created_at", 0),
+        updated_at=data.get("updated_at", 0),
+        steps=steps,
+        state_snapshot=data.get("state_snapshot", {}),
+    )
+
+
+def create_checkpoint(
+    pipeline_name: str, kb_name: str, step_names: list[str]
+) -> Checkpoint:
+    """Create a new checkpoint for a pipeline run."""
+    steps = [StepState(name=name, status="pending") for name in step_names]
+    return Checkpoint(pipeline_name=pipeline_name, kb_name=kb_name, steps=steps)
+
+
+def clear_checkpoint(knowledge_dir: Path) -> None:
+    """Remove checkpoint file after successful completion."""
+    cp_path = knowledge_dir / ".meta" / CHECKPOINT_FILENAME
+    if cp_path.exists():
+        cp_path.unlink()
+        logger.debug("[checkpoint] Cleared: %s", cp_path)
+
+
+# ---------------------------------------------------------------------------
+# Resume logic
+# ---------------------------------------------------------------------------
+
+def get_resume_point(knowledge_dir: Path) -> str | None:
+    """Get the step name to resume from, or None if no resume needed."""
+    cp = load_checkpoint(knowledge_dir)
+    if cp is None:
+        return None
+    return cp.next_step
+
+
+def should_skip_step(checkpoint: Checkpoint | None, step_name: str) -> bool:
+    """Check if a step should be skipped during resume."""
+    if checkpoint is None:
+        return False
+    return step_name in checkpoint.completed_steps
+
+
+# ---------------------------------------------------------------------------
+# Recovery (LLM failure handling) — see generate_recovery.py for full chain
+# ---------------------------------------------------------------------------
+
+def attempt_recovery(
+    task: SubagentTask,
+    strategy: LlmStrategy,
+    failure_reason: str,
+    max_retries: int = 2,
+) -> SubagentResult:
+    """Try retrying the task. For full recovery chain, use generate_recovery.py."""
+    for attempt in range(max_retries):
+        result = run_single(task, strategy)
+        if result.status == "done":
+            return result
+
+    return SubagentResult(
+        task_id=task.task_id, status="failed",
+        error=f"Recovery exhausted after {max_retries} retries: {failure_reason}",
+    )
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _sanitize_state(state: dict) -> dict:
+    """Remove non-serializable items from state for checkpoint storage."""
+    safe: dict[str, Any] = {}
+    for key, value in state.items():
+        if key.startswith("_"):
+            continue  # Skip private state (locks, etc.)
+        if isinstance(value, Path):
+            safe[key] = str(value)
+        elif isinstance(value, dict):
+            safe[key] = {
+                k: str(v) if isinstance(v, Path) else v
+                for k, v in value.items()
+                if not k.startswith("_")
+            }
+        elif isinstance(value, (str, int, float, bool, list)):
+            safe[key] = value
+    return safe

engine/pipeline/steps/__init__.py

@@ -0,0 +1,23 @@
+"""CLI pipeline steps — all step implementations.
+
+Steps are organized by pipeline phase:
+    lock.py      — AcquireLock / ReleaseLock
+    source.py    — FetchSource / BuildSource / ResolveDeps
+    skeleton.py  — ExtractSkeleton / ExtractFileList / Classify
+    generate.py  — GenerateDocs (main LLM generation step)
+    quality.py   — Validate / Sampling / Links / Duplicates
+    finalize.py  — Merge / Dedup / SharedDocs / Publish / RebuildIndex / Clean
+    sync.py      — Sync-specific steps
+    audit.py     — Audit-specific steps
+
+Import all step modules to trigger registration.
+"""
+
+# Import all step modules to register them
+from engine.pipeline.steps import lock  # noqa: F401
+from engine.pipeline.steps import source  # noqa: F401
+from engine.pipeline.steps import skeleton  # noqa: F401
+from engine.pipeline.steps import pregenerate  # noqa: F401
+from engine.pipeline.steps import generate  # noqa: F401
+from engine.pipeline.steps import quality  # noqa: F401
+from engine.pipeline.steps import finalize  # noqa: F401

engine/pipeline/steps/audit.py

@@ -0,0 +1,220 @@
+"""Audit pipeline steps — document quality audit via LLM.
+
+Steps:
+- FetchAndScaleStep: fetch source and assess scale for strategy selection
+- AuditDocsStep: LLM-powered document audit against source
+
+Requirements: Req 12, 18, 20, 27
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+from core.interfaces import Step, StepResult, PipelineContext, LlmRequest
+from engine.pipeline import register_step
+from engine.pipeline.steps.audit_helpers import (
+    compute_doc_size, get_audit_order, load_audit_progress,
+    update_audit_progress, build_audit_system_prompt, build_audit_user_prompt,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@register_step
+class FetchAndScaleStep(Step):
+    """Fetch source and assess scale for audit strategy selection."""
+
+    default_name = "fetch-and-scale"
+
+    def __init__(self):
+        super().__init__("fetch-and-scale")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        from core.git import ensure_repo
+        from core.preset import load_preset
+        from engine.pipeline.steps.source import _resolve_modules
+
+        modules = _resolve_modules(ctx)
+        strategies: dict[str, str] = {}
+
+        preset_name = ctx.kb_config.get("preset", "generic")
+        preset = load_preset(preset_name)
+        audit_cfg = preset.get("audit", {})
+        scale_cfg = audit_cfg.get("scale_thresholds", {})
+        single_agent_threshold = scale_cfg.get("single_agent_kb", 50) * 1024
+        skeleton_threshold = scale_cfg.get("skeleton_kb", 80) * 1024
+
+        for mod in modules:
+            name = mod["name"]
+            branch = mod.get("branch", "main")
+            try:
+                repo_path = ensure_repo(
+                    mod.get("url"), mod.get("local"),
+                    ctx.cache_dir, name, branch,
+                )
+                ctx.state.setdefault("module_repos", {})[name] = repo_path
+
+                module_dir = ctx.knowledge_dir / name
+                doc_size = compute_doc_size(module_dir)
+
+                if doc_size < single_agent_threshold:
+                    strategies[name] = "single-agent"
+                elif doc_size < skeleton_threshold:
+                    strategies[name] = "single-agent+skeleton"
+                else:
+                    strategies[name] = "multi-agent"
+            except Exception as e:
+                logger.warning("[fetch-and-scale] %s: %s", name, e)
+                continue
+
+        ctx.state["audit_strategies"] = strategies
+        return StepResult(status="ok", message=f"Scale assessed: {len(strategies)} modules")
+
+
+@register_step
+class AuditDocsStep(Step):
+    """LLM-powered document audit with structured findings.
+
+    Processes documents in dependency order (batch field from doc_types.yaml).
+    Supports progress persistence for resume and scope filtering.
+    """
+
+    default_name = "audit-docs"
+
+    def __init__(self):
+        super().__init__("audit-docs", checkpoint="audit-docs")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        from engine.strategies import create_strategy
+        from engine.pipeline.steps.generate import _make_config_obj
+        from core.prompt.response_parser import parse_audit_response
+        from core.preset import load_preset
+
+        config_obj = _make_config_obj(ctx)
+        if config_obj.agent_backend == "delegated":
+            return self._delegated_mode(ctx)
+
+        strategy = create_strategy(config_obj)
+        module_repos: dict[str, Path] = ctx.state.get("module_repos", {})
+
+        preset_name = ctx.kb_config.get("preset", "generic")
+        preset = load_preset(preset_name)
+        audit_cfg = preset.get("audit", {})
+        max_tokens = audit_cfg.get("max_tokens", 8192)
+        doc_truncate_chars = audit_cfg.get("doc_truncate_chars", 50000)
+
+        scope_doc_type = ctx.state.get("audit_scope")
+        scope_module = ctx.state.get("audit_module") or ctx.module
+        force = ctx.state.get("audit_force", False)
+
+        all_findings: dict[str, list] = {}
+        total_audited = 0
+
+        for name, repo_path in module_repos.items():
+            if scope_module and name != scope_module:
+                continue
+
+            module_dir = ctx.knowledge_dir / name
+            if not module_dir.is_dir():
+                continue
+
+            progress = load_audit_progress(module_dir) if not force else {}
+            audit_order = get_audit_order(ctx)
+
+            for doc_type in audit_order:
+                if scope_doc_type and doc_type != scope_doc_type:
+                    continue
+
+                doc_file = module_dir / f"{doc_type}.md"
+                if not doc_file.exists():
+                    continue
+
+                if progress.get(doc_type, {}).get("status") == "DONE":
+                    logger.info("[audit] %s/%s: already DONE, skipping", name, doc_type)
+                    continue
+
+                update_audit_progress(module_dir, doc_type, "IN_PROGRESS")
+
+                system_prompt = build_audit_system_prompt(doc_type)
+                user_prompt = build_audit_user_prompt(doc_file, doc_type, module_dir, doc_truncate_chars)
+
+                request = LlmRequest(system=system_prompt, user=user_prompt, max_tokens=max_tokens)
+                try:
+                    response = strategy.call(request)
+                except Exception as e:
+                    logger.error("[audit] LLM failed for %s/%s: %s", name, doc_type, e)
+                    update_audit_progress(module_dir, doc_type, "FAILED")
+                    total_audited += 1
+                    continue
+
+                if response.status != "done":
+                    logger.warning("[audit] LLM returned status '%s' for %s/%s: %s",
+                                   response.status, name, doc_type, response.error[:200] if response.error else "")
+                    update_audit_progress(module_dir, doc_type, "FAILED")
+                    total_audited += 1
+                    continue
+
+                result = parse_audit_response(response.content)
+                if not result.success:
+                    logger.warning("[audit] %s/%s: parse failed: %s", name, doc_type, result.error)
+                    update_audit_progress(module_dir, doc_type, "FAILED")
+                    continue
+
+                key = f"{name}/{doc_type}"
+                all_findings[key] = result.findings
+                update_audit_progress(module_dir, doc_type, "DONE", len(result.findings))
+                total_audited += 1
+
+        ctx.state["audit_findings"] = all_findings
+        total_issues = sum(
+            len([f for f in findings if f.status == "fail"])
+            for findings in all_findings.values()
+        )
+
+        if total_audited > 0 and not all_findings:
+            return StepResult(
+                status="ok",
+                message=f"Attempted {total_audited} docs, all LLM calls failed (check LLM connection)",
+                details={"total_audited": total_audited, "total_issues": 0, "all_failed": True},
+            )
+
+        return StepResult(
+            status="ok",
+            message=f"Audited {total_audited} docs, {total_issues} issues found",
+            details={"total_audited": total_audited, "total_issues": total_issues},
+        )
+
+    def _delegated_mode(self, ctx: PipelineContext) -> StepResult:
+        """Generate audit manifest for Agent dispatch."""
+        module_repos: dict[str, Path] = ctx.state.get("module_repos", {})
+        manifest = {"operation": "kb-audit", "modules": []}
+
+        for name, repo_path in module_repos.items():
+            module_dir = ctx.knowledge_dir / name
+            if not module_dir.is_dir():
+                continue
+
+            docs = []
+            for doc_type in get_audit_order(ctx):
+                doc_file = module_dir / f"{doc_type}.md"
+                if doc_file.exists():
+                    docs.append({
+                        "doc_type": doc_type,
+                        "doc_path": str(doc_file.relative_to(ctx.knowledge_dir)),
+                        "template": f"presets/java-spring/templates/subagent-audit-{doc_type}.md",
+                    })
+
+            if docs:
+                manifest["modules"].append({"name": name, "documents": docs})
+
+        manifest_path = ctx.knowledge_dir / ".audit-manifest.json"
+        manifest_path.write_text(
+            json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8"
+        )
+        return StepResult(
+            status="delegated",
+            message=f"Audit manifest written: {manifest_path}",
+        )

engine/pipeline/steps/audit_apply.py

@@ -0,0 +1,195 @@
+"""Audit fix and report steps — apply fixes and generate reports.
+
+Steps:
+- ApplyFixesStep: apply LLM-generated fixes to documents
+- GenerateReportStep: generate audit summary report
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+from core.interfaces import Step, StepResult, PipelineContext
+from engine.pipeline import register_step
+
+logger = logging.getLogger(__name__)
+
+
+@register_step
+class ApplyFixesStep(Step):
+    """Apply LLM-generated fixes to documents.
+
+    Fix scope determination:
+    - fix contains section heading -> Section_Updater.replace_section()
+    - fix is new section -> Section_Updater.append_section()
+    - fix is inline content -> direct line replacement
+    """
+
+    default_name = "apply-fixes"
+
+    def __init__(self):
+        super().__init__("apply-fixes")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        findings = ctx.state.get("audit_findings", {})
+        if not findings:
+            return StepResult(status="skipped", message="No findings to apply")
+
+        dry_run = ctx.state.get("dry_run", False)
+        applied = 0
+        skipped = 0
+        failed = 0
+        fix_log: list[dict] = []
+
+        for key, finding_list in findings.items():
+            parts = key.split("/", 1)
+            if len(parts) != 2:
+                continue
+            module_name, doc_type = parts
+            module_dir = ctx.knowledge_dir / module_name
+            doc_path = module_dir / f"{doc_type}.md"
+
+            if not doc_path.exists():
+                continue
+
+            for finding in finding_list:
+                if finding.status != "fail" or not finding.fix:
+                    skipped += 1
+                    continue
+
+                if dry_run:
+                    fix_log.append({
+                        "doc": doc_type, "dimension": finding.dimension,
+                        "action": "would_apply", "preview": finding.fix[:200],
+                    })
+                    applied += 1
+                    continue
+
+                success = _apply_single_fix(doc_path, finding)
+                if success:
+                    applied += 1
+                    fix_log.append({
+                        "doc": doc_type, "dimension": finding.dimension,
+                        "status": "applied",
+                    })
+                else:
+                    failed += 1
+                    fix_log.append({
+                        "doc": doc_type, "dimension": finding.dimension,
+                        "status": "failed",
+                    })
+
+        if fix_log:
+            for key in findings:
+                parts = key.split("/", 1)
+                if len(parts) == 2:
+                    module_dir = ctx.knowledge_dir / parts[0]
+                    if module_dir.is_dir():
+                        tracking_path = module_dir / ".audit-fixes.json"
+                        tracking_path.write_text(
+                            json.dumps(fix_log, ensure_ascii=False, indent=2),
+                            encoding="utf-8",
+                        )
+                    break
+
+        ctx.state["fixes_applied"] = applied
+        ctx.state["fixes_failed"] = failed
+
+        return StepResult(
+            status="ok",
+            message=f"Applied {applied} fixes, {skipped} skipped, {failed} failed",
+            details={"applied": applied, "skipped": skipped, "failed": failed},
+        )
+
+
+@register_step
+class GenerateReportStep(Step):
+    """Generate audit summary report."""
+
+    default_name = "generate-report"
+
+    def __init__(self):
+        super().__init__("generate-report")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        findings = ctx.state.get("audit_findings", {})
+        fixes_applied = ctx.state.get("fixes_applied", 0)
+        fixes_failed = ctx.state.get("fixes_failed", 0)
+        report_path = ctx.knowledge_dir / ".audit-report.md"
+
+        scope_doc = ctx.state.get("audit_scope", "all")
+        scope_mod = ctx.state.get("audit_module") or ctx.module or "all"
+
+        lines = ["# Audit Report", ""]
+        lines.append(f"Knowledge base: {ctx.kb_name}")
+        lines.append(f"Scope: {scope_doc}/{scope_mod}")
+        lines.append(f"Fixes: {fixes_applied} applied, {fixes_failed} failed")
+        lines.append("")
+
+        total_issues = 0
+        for key, finding_list in sorted(findings.items()):
+            fails = [f for f in finding_list if f.status == "fail"]
+            passes = [f for f in finding_list if f.status == "pass"]
+            total_issues += len(fails)
+
+            lines.append(f"## {key}")
+            lines.append(f"- Pass: {len(passes)}, Fail: {len(fails)}")
+            lines.append("")
+
+            if fails:
+                for f in fails[:10]:
+                    fix_indicator = " (fixed)" if f.fix else ""
+                    lines.append(f"  - [{f.dimension}] {f.detail}{fix_indicator}")
+                if len(fails) > 10:
+                    lines.append(f"  - ... and {len(fails) - 10} more")
+                lines.append("")
+
+        lines.append("---")
+        lines.append(f"Total: {total_issues} issues")
+
+        report_path.write_text("\n".join(lines), encoding="utf-8")
+
+        module_repos = ctx.state.get("module_repos", {})
+        for name in module_repos:
+            module_dir = ctx.knowledge_dir / name
+            progress_file = module_dir / ".audit-progress.json"
+            if progress_file.exists():
+                progress_file.unlink()
+
+        return StepResult(status="ok", message=f"Report: {report_path} ({total_issues} issues)")
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _apply_single_fix(doc_path: Path, finding) -> bool:
+    """Apply a single fix from an audit finding."""
+    from core.docs.section_updater import replace_section, append_section
+    from core.utils import validate_path_within_bounds
+
+    fix_content = finding.fix
+    if not fix_content:
+        return False
+
+    if not validate_path_within_bounds(doc_path, doc_path.parent):
+        logger.warning("[apply-fix] Path traversal blocked: %s", doc_path)
+        return False
+
+    lines = fix_content.strip().split("\n")
+    first_line = lines[0] if lines else ""
+
+    if first_line.startswith("## ") or first_line.startswith("### "):
+        heading = first_line
+        body = "\n".join(lines[1:]).strip()
+        success = replace_section(doc_path, heading, body)
+        if not success:
+            success = append_section(doc_path, heading, body)
+        return success
+    else:
+        logger.info("[apply-fix] Inline fix for %s/%s (deferred): %s",
+                    doc_path.name, finding.dimension, fix_content[:100])
+        return False