source-kb 0.2.2__py3-none-any.whl

engine/pipeline/steps/finalize.py

@@ -0,0 +1,160 @@
+"""Finalization steps — merge, dedup, publish, index, cleanup."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from core.interfaces import Step, StepResult, PipelineContext
+from engine.pipeline import register_step
+
+
+@register_step
+class MergeShardsStep(Step):
+    """Merge split document shards into single files."""
+
+    default_name = "merge-shards"
+
+    def __init__(self):
+        super().__init__("merge-shards")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        from core.skeleton.merge import find_shards, merge_shards, refine_merged_doc
+
+        merged_count = 0
+        for module_dir in _module_dirs(ctx):
+            # Find all shard prefixes from preset doc_type_mapping
+            from core.preset import load_preset, get_doc_type_mapping
+            preset_name = ctx.kb_config.get("preset", "generic")
+            preset = load_preset(preset_name)
+            dt_mapping = get_doc_type_mapping(preset)
+            known_types = {fn.removesuffix(".md") for fn in dt_mapping.values()} if dt_mapping else set()
+
+            for prefix in known_types:
+                shards = find_shards(module_dir, prefix)
+                if not shards:
+                    continue
+                merged_file = module_dir / f"{prefix}.md"
+                if merged_file.exists() and merged_file.stat().st_size > 100:
+                    continue
+
+                content = merge_shards(shards)
+                if content.strip():
+                    merged_file.write_text(content, encoding="utf-8")
+                    for s in shards:
+                        s.unlink(missing_ok=True)
+                    merged_count += 1
+
+                    # Post-merge refinement
+                    result = refine_merged_doc(merged_file)
+                    if result.changed:
+                        result.apply()
+
+        return StepResult(status="ok", message=f"Merged {merged_count} documents")
+
+
+@register_step
+class DeduplicateStep(Step):
+    """LLM-based deduplication of large documents."""
+
+    default_name = "deduplicate"
+
+    def __init__(self):
+        super().__init__("deduplicate")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        from engine.dedup import dedup_document
+        from engine.strategies import create_strategy
+        from engine.pipeline.steps.generate import _make_config_obj
+        from core.preset import load_preset
+
+        config_obj = _make_config_obj(ctx)
+        strategy = create_strategy(config_obj)
+
+        # Probe: delegated strategy means Agent handles dedup
+        from core.interfaces import LlmRequest
+        probe = strategy.call(LlmRequest(system="probe", user="probe"))
+        if probe.status == "delegated":
+            return StepResult(status="skipped", message="Delegated mode — dedup handled by Agent")
+
+        preset_name = ctx.kb_config.get("preset", "generic")
+        preset = load_preset(preset_name)
+        limits = preset.get("limits", {})
+        dedup_min_size = limits.get("dedup_min_doc_size_bytes", 15 * 1024)
+
+        total_deduped = 0
+        total_saved = 0.0
+
+        for module_dir in _module_dirs(ctx):
+            md_files = sorted(module_dir.glob("*.md"))
+            sibling_names = [f.name for f in md_files]
+
+            for md_file in md_files:
+                if md_file.stat().st_size < dedup_min_size:
+                    continue
+                changed, saved = dedup_document(md_file, sibling_names, strategy)
+                if changed:
+                    total_deduped += 1
+                    total_saved += saved
+
+        if total_deduped == 0:
+            return StepResult(status="ok", message="No redundancy found")
+        return StepResult(status="ok", message=f"Deduped {total_deduped} docs (saved {total_saved:.1f}KB)")
+
+
+@register_step
+class RebuildIndexStep(Step):
+    """Rebuild vector index for generated documents."""
+
+    default_name = "rebuild-index"
+
+    def __init__(self):
+        super().__init__("rebuild-index", checkpoint="cp9")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        # TODO: Import from core.rag when migrated
+        return StepResult(status="ok", message="Index rebuild (placeholder)")
+
+
+@register_step
+class CleanProgressStep(Step):
+    """Remove progress files after successful completion."""
+
+    default_name = "clean-progress"
+
+    def __init__(self):
+        super().__init__("clean-progress")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        from core.monitor.progress import cleanup_progress
+
+        cleaned = 0
+        for module_dir in _module_dirs(ctx):
+            cleaned += cleanup_progress(module_dir)
+        return StepResult(status="ok", message=f"Cleaned {cleaned} progress files")
+
+
+@register_step
+class SharedDocsStep(Step):
+    """Generate cross-module shared documents."""
+
+    default_name = "shared-docs"
+
+    def __init__(self):
+        super().__init__("shared-docs")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        from core.docs.shared import generate_shared_docs
+
+        generated = generate_shared_docs(ctx.knowledge_dir, ctx.config, ctx.kb_name)
+        if not generated:
+            return StepResult(status="skipped", message="No shared docs to generate")
+        return StepResult(status="ok", message=f"Generated {len(generated)} shared docs",
+                          details={"generated": generated})
+
+
+def _module_dirs(ctx: PipelineContext) -> list[Path]:
+    if ctx.module:
+        d = ctx.knowledge_dir / ctx.module
+        return [d] if d.is_dir() else []
+    return sorted(d for d in ctx.knowledge_dir.iterdir()
+                  if d.is_dir() and not d.name.startswith(".") and d.name != "_shared")

engine/pipeline/steps/generate.py

@@ -0,0 +1,169 @@
+"""Document generation step — the main LLM-powered step.
+
+Orchestrates prompt rendering, split planning, batch execution,
+and two-phase generation for large modules.
+
+Requirements: 5.1, 5.5
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from pathlib import Path
+
+from core.interfaces import Step, StepResult, PipelineContext
+from core.monitor.progress import write_progress
+from engine.pipeline import register_step
+from engine.runner import BatchAbortError
+
+logger = logging.getLogger(__name__)
+
+MIN_DOC_SIZE_BYTES = 500
+
+
+@register_step
+class GenerateDocsStep(Step):
+    """Generate knowledge base documents via LLM."""
+
+    default_name = "generate-docs"
+
+    def __init__(self):
+        super().__init__("generate-docs", checkpoint="cp4")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        import json
+        from core.preset import load_preset
+        from core.prompt.renderer import render_prompt
+        from engine.strategies import create_strategy, ConfigProxy
+        from engine.assembler import InlinePromptAssembler
+        from engine.runner import SubagentTask, run_batch, configure_garbage_patterns, CircuitBreakerConfig
+
+        config = ctx.config
+        preset_name = ctx.kb_config.get("preset", "generic")
+        preset = load_preset(preset_name)
+
+        # Configure runner from preset limits
+        limits = preset.get("limits", {})
+        configure_garbage_patterns(limits.get("garbage_patterns"))
+        breaker_config = CircuitBreakerConfig.from_config({"limits": limits})
+
+        strategy = create_strategy(ConfigProxy(config))
+        assembler = InlinePromptAssembler(preset=preset)
+
+        snippet_path = Path(ctx.project_root) / "skills" / "kb-init" / "templates" / "core" / "output-only.md"
+        execution_snippet = snippet_path.read_text(encoding="utf-8") if snippet_path.exists() else ""
+
+        from core.monitor.prompt_store import should_save_prompts
+        save_prompts = should_save_prompts(config)
+
+        module_repos: dict[str, Path] = ctx.state.get("module_repos", {})
+        generated: list[str] = []
+        max_concurrent = ctx.config.get("agent", {}).get("max_concurrent_subagents", 5)
+
+        for name in module_repos:
+            module_dir = ctx.knowledge_dir / name
+            tasks_file = module_dir / ".meta" / "dispatch-tasks.json"
+
+            if not tasks_file.exists():
+                logger.warning("No dispatch-tasks.json for %s, skipping", name)
+                continue
+
+            manifest = json.loads(tasks_file.read_text(encoding="utf-8"))
+            dispatch_tasks = manifest.get("tasks", [])
+
+            # Group by batch for sequential execution
+            batches: dict[str, list[dict]] = {}
+            for t in dispatch_tasks:
+                batches.setdefault(t["batch"], []).append(t)
+
+            for batch_id in sorted(batches.keys()):
+                batch_tasks: list[SubagentTask] = []
+
+                for t in batches[batch_id]:
+                    output_path = Path(t["output_file"])
+                    if output_path.exists() and output_path.stat().st_size > MIN_DOC_SIZE_BYTES:
+                        continue
+
+                    template_path = _find_template(ctx.project_root, preset_name, t["doc_type"], preset)
+                    if not template_path:
+                        continue
+
+                    extras = {}
+                    if t.get("shard_file_list"):
+                        extras["file_list_override"] = t["shard_file_list"]
+
+                    prompt = render_prompt(
+                        template_path=template_path,
+                        config=config,
+                        kb_name=ctx.kb_name,
+                        module_name=name,
+                        doc_type=t["doc_type"],
+                        assembler=assembler,
+                        extras=extras,
+                        execution_snippet=execution_snippet,
+                        preset=preset,
+                    )
+
+                    batch_tasks.append(SubagentTask(
+                        task_id=f"{name}__{t['id']}",
+                        prompt=prompt,
+                        output_path=output_path,
+                        doc_type=t["doc_type"],
+                    ))
+
+                if batch_tasks:
+                    _write_heartbeat(module_dir, [t["doc_type"] for t in batches[batch_id]])
+                    try:
+                        results = run_batch(batch_tasks, strategy, max_concurrent,
+                                           save_prompts=save_prompts, breaker_config=breaker_config)
+                    except BatchAbortError as e:
+                        logger.error("Batch %s aborted: %s", batch_id, e.reason)
+                        logger.error("Diagnosis: %s", e.diagnosis)
+                        return StepResult(
+                            status="failed",
+                            message=f"Generation aborted: {e.reason}\nDiagnosis: {e.diagnosis}",
+                            details={"generated": generated, "abort_reason": e.reason, "diagnosis": e.diagnosis},
+                        )
+                    for r in results:
+                        if r.status == "done" and r.output_path and r.output_path.exists():
+                            generated.append(f"{name}/{r.output_path.name}")
+                            write_progress(module_dir, r.task_id.split("__")[-1], "DONE")
+
+        if not generated:
+            return StepResult(status="skipped", message="No documents generated")
+        return StepResult(status="ok", message=f"Generated {len(generated)} documents",
+                          details={"generated": generated})
+
+
+def _make_config_obj(ctx: PipelineContext):
+    """Create a ConfigProxy from pipeline context. Used by multiple steps."""
+    from engine.strategies import ConfigProxy
+    return ConfigProxy(ctx.config)
+
+
+def _find_template(project_root: Path, preset_name: str, doc_type: str, preset: dict = None) -> Path | None:
+    """Find the template file for a doc type from preset config.
+
+    Uses preset search path: env > local > built-in.
+    """
+    from core.preset import find_preset_template, get_template_path
+
+    template_name = None
+    if preset:
+        template_name = get_template_path(preset, doc_type, preset_name)
+
+    if not template_name:
+        return None
+
+    return find_preset_template(preset_name, template_name)
+
+
+def _write_heartbeat(module_dir: Path, batch: list[str]) -> None:
+    """Write heartbeat file for current batch progress tracking."""
+    hb_path = module_dir / ".meta" / "heartbeat.txt"
+    hb_path.parent.mkdir(parents=True, exist_ok=True)
+    hb_path.write_text(
+        f"batch={','.join(batch)} ts={time.time():.0f}\n",
+        encoding="utf-8",
+    )

engine/pipeline/steps/generate_batch.py

@@ -0,0 +1,197 @@
+"""Batch execution module for document generation.
+
+Handles single-batch execution logic:
+- Shard generation via LlmStrategy
+- Two-phase coordination (call cli/two_phase.py)
+- Concurrency control
+
+Requirements: 5.2, 5.3, 5.4
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+from core.interfaces import LlmStrategy, PipelineContext
+from core.paths import resolve_file_list
+from core.monitor.progress import write_progress
+from engine.runner import SubagentTask, SubagentResult, run_batch
+from engine.two_phase import run_two_phase
+
+logger = logging.getLogger(__name__)
+
+MIN_DOC_SIZE_BYTES = 500
+
+
+def execute_batch(
+    ctx: PipelineContext,
+    module_name: str,
+    doc_types: list[str],
+    strategy: LlmStrategy,
+    max_concurrent: int = 5,
+) -> list[SubagentResult]:
+    """Execute a single batch of doc_types for a module.
+
+    For each doc_type, determines if splitting is needed and dispatches
+    either direct generation or two-phase split generation.
+
+    Args:
+        ctx: Pipeline context
+        module_name: Name of the module being processed
+        doc_types: List of doc types in this batch
+        strategy: LLM execution strategy
+        max_concurrent: Max concurrent tasks
+
+    Returns:
+        List of SubagentResults from all tasks in this batch
+    """
+    from core.preset import load_preset, get_doc_type_mapping
+    from core.skeleton.split import SplitConfig, compute_splits
+    from core.skeleton.file_list import load_skeleton
+    from core.prompt.renderer import render_prompt
+    from engine.assembler import InlinePromptAssembler
+    from engine.pipeline.steps.generate import _find_template, _make_config_obj
+
+    preset_name = ctx.kb_config.get("preset", "generic")
+    preset = load_preset(preset_name)
+    split_config = SplitConfig.from_preset(preset, mode="output-only")
+    assembler = InlinePromptAssembler()
+
+    module_dir = ctx.knowledge_dir / module_name
+    config_path = str(ctx.project_root / "kb-project.yaml")
+
+    tasks: list[SubagentTask] = []
+    split_tasks: list[SubagentTask] = []
+
+    for doc_type in doc_types:
+        doc_path = module_dir / f"{doc_type}.md"
+        if doc_path.exists() and doc_path.stat().st_size > MIN_DOC_SIZE_BYTES:
+            continue
+
+        if not resolve_file_list(module_dir, doc_type):
+            continue
+
+        template_path = _find_template(ctx.project_root, preset_name, doc_type, preset)
+        if not template_path:
+            continue
+
+        # Compute splits based on file list size
+        num_splits = _compute_doc_splits(module_dir, doc_type, split_config)
+
+        prompt = render_prompt(
+            template_path=template_path,
+            config=ctx.config,
+            kb_name=ctx.kb_name,
+            module_name=module_name,
+            doc_type=doc_type,
+            assembler=assembler,
+        )
+
+        if num_splits > 1:
+            # Create shard tasks for two-phase generation
+            shard_tasks = _create_shard_tasks(
+                ctx, module_dir, module_name, doc_type,
+                prompt, num_splits, template_path, assembler,
+            )
+            split_tasks.extend(shard_tasks)
+        else:
+            tasks.append(SubagentTask(
+                task_id=f"{module_name}__{doc_type}",
+                prompt=prompt,
+                output_path=doc_path,
+                doc_type=doc_type,
+            ))
+
+    results: list[SubagentResult] = []
+
+    # Execute direct tasks
+    if tasks:
+        logger.info("[batch] Running %d direct tasks for %s", len(tasks), module_name)
+        results.extend(run_batch(tasks, strategy, max_concurrent))
+
+    # Execute split tasks via two-phase
+    if split_tasks:
+        logger.info("[batch] Running %d split tasks (two-phase) for %s",
+                    len(split_tasks), module_name)
+        results.extend(run_two_phase(split_tasks, strategy, max_concurrent))
+
+    return results
+
+
+def _compute_doc_splits(
+    module_dir: Path, doc_type: str, split_config: "SplitConfig"
+) -> int:
+    """Compute number of splits needed for a doc type."""
+    from core.skeleton.file_list import load_skeleton
+
+    fl_path = resolve_file_list(module_dir, doc_type)
+    if not fl_path:
+        return 1
+
+    try:
+        lines = fl_path.read_text(encoding="utf-8").strip().splitlines()
+        file_count = len(lines)
+    except OSError:
+        return 1
+
+    # Simple heuristic: split if file count exceeds threshold
+    threshold = split_config.max_files_per_shard if split_config else 30
+    if file_count <= threshold:
+        return 1
+    return min(4, max(2, file_count // threshold + 1))
+
+
+def _create_shard_tasks(
+    ctx: PipelineContext,
+    module_dir: Path,
+    module_name: str,
+    doc_type: str,
+    base_prompt: str,
+    num_splits: int,
+    template_path: Path,
+    assembler,
+) -> list[SubagentTask]:
+    """Create shard tasks for split generation.
+
+    Divides the file list into shards and creates a task per shard.
+    """
+    from core.paths import shard_doc_path
+
+    fl_path = resolve_file_list(module_dir, doc_type)
+    if not fl_path:
+        return []
+
+    try:
+        all_files = fl_path.read_text(encoding="utf-8").strip().splitlines()
+    except OSError:
+        return []
+
+    if not all_files:
+        return []
+
+    # Divide files into shards
+    shard_size = max(1, len(all_files) // num_splits)
+    shards: list[list[str]] = []
+    for i in range(0, len(all_files), shard_size):
+        shards.append(all_files[i:i + shard_size])
+
+    # Merge last shard if too small
+    if len(shards) > 1 and len(shards[-1]) < shard_size // 2:
+        shards[-2].extend(shards[-1])
+        shards.pop()
+
+    tasks: list[SubagentTask] = []
+    for idx, shard_files in enumerate(shards, 1):
+        shard_suffix = f"\n\n## Files assigned to this shard\n\n" + "\n".join(f"- {f}" for f in shard_files)
+        output_path = shard_doc_path(module_dir, doc_type, idx)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+
+        tasks.append(SubagentTask(
+            task_id=f"{module_name}__{doc_type}__shard{idx:02d}",
+            prompt=base_prompt + shard_suffix,
+            output_path=output_path,
+            doc_type=doc_type,
+        ))
+
+    return tasks

engine/pipeline/steps/generate_recovery.py

@@ -0,0 +1,170 @@
+"""Recovery strategy chain for failed generation steps.
+
+Provides a chain-of-responsibility pattern for handling LLM failures:
+1. SplitRetryStrategy — re-split into smaller shards and retry
+2. ModelSwitchStrategy — try a different/larger model
+3. MarkFailedStrategy — give up and mark as failed
+
+Requirements: 5.2
+"""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+
+from core.interfaces import LlmStrategy, LlmRequest, LlmResponse
+from engine.runner import SubagentTask, SubagentResult, run_single
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class RecoveryContext:
+    """Context passed through the recovery chain."""
+    task: SubagentTask
+    strategy: LlmStrategy
+    failure_reason: str
+    attempt: int = 0
+    max_attempts: int = 3
+
+
+class RecoveryStrategy(ABC):
+    """Base class for recovery strategies."""
+
+    @abstractmethod
+    def can_handle(self, ctx: RecoveryContext) -> bool:
+        """Check if this strategy can handle the failure."""
+
+    @abstractmethod
+    def recover(self, ctx: RecoveryContext) -> SubagentResult | None:
+        """Attempt recovery. Returns result or None to pass to next strategy."""
+
+
+class SplitRetryStrategy(RecoveryStrategy):
+    """Re-split a failed shard into smaller pieces and retry."""
+
+    def can_handle(self, ctx: RecoveryContext) -> bool:
+        # Only applicable if the task has a large prompt (likely too much content)
+        return len(ctx.task.prompt) > 20000 and ctx.attempt < 2
+
+    def recover(self, ctx: RecoveryContext) -> SubagentResult | None:
+        logger.info("[recovery:split] %s: splitting prompt", ctx.task.task_id)
+
+        # Split the prompt roughly in half by file list
+        prompt = ctx.task.prompt
+        marker = "## Files assigned to this shard"
+        if marker not in prompt:
+            return None
+
+        parts = prompt.split(marker, 1)
+        if len(parts) < 2:
+            return None
+
+        file_section = parts[1]
+        files = [l.strip("- \n") for l in file_section.strip().splitlines() if l.strip().startswith("-")]
+        if len(files) <= 2:
+            return None
+
+        mid = len(files) // 2
+        first_half = files[:mid]
+        second_half = files[mid:]
+
+        results: list[SubagentResult] = []
+        for idx, chunk in enumerate([first_half, second_half], 1):
+            chunk_suffix = f"\n\n{marker}\n\n" + "\n".join(f"- {f}" for f in chunk)
+            sub_task = SubagentTask(
+                task_id=f"{ctx.task.task_id}__split{idx}",
+                prompt=parts[0] + chunk_suffix,
+                output_path=ctx.task.output_path.parent / f"{ctx.task.output_path.stem}-part{idx}.md",
+                doc_type=ctx.task.doc_type,
+            )
+            result = run_single(sub_task, ctx.strategy)
+            results.append(result)
+
+        # Merge successful results
+        merged_content = ""
+        for r in results:
+            if r.status == "done" and r.content:
+                merged_content += r.content + "\n\n"
+
+        if merged_content:
+            ctx.task.output_path.parent.mkdir(parents=True, exist_ok=True)
+            ctx.task.output_path.write_text(merged_content.strip(), encoding="utf-8")
+            return SubagentResult(
+                task_id=ctx.task.task_id, status="done",
+                output_path=ctx.task.output_path, content=merged_content,
+            )
+        return None
+
+
+class ModelSwitchStrategy(RecoveryStrategy):
+    """Switch to a different model and retry."""
+
+    def __init__(self, fallback_model: str | None = None):
+        self._fallback = fallback_model
+
+    def can_handle(self, ctx: RecoveryContext) -> bool:
+        return self._fallback is not None and ctx.attempt < 2
+
+    def recover(self, ctx: RecoveryContext) -> SubagentResult | None:
+        if not self._fallback:
+            return None
+        logger.info("[recovery:model-switch] %s: switching to %s",
+                    ctx.task.task_id, self._fallback)
+        task = SubagentTask(
+            task_id=ctx.task.task_id,
+            prompt=ctx.task.prompt,
+            output_path=ctx.task.output_path,
+            doc_type=ctx.task.doc_type,
+            model=self._fallback,
+        )
+        result = run_single(task, ctx.strategy)
+        return result if result.status == "done" else None
+
+
+class MarkFailedStrategy(RecoveryStrategy):
+    """Final fallback — mark the task as failed and continue."""
+
+    def can_handle(self, ctx: RecoveryContext) -> bool:
+        return True  # Always applicable as last resort
+
+    def recover(self, ctx: RecoveryContext) -> SubagentResult | None:
+        logger.warning("[recovery:mark-failed] %s: giving up after %d attempts",
+                       ctx.task.task_id, ctx.attempt)
+        return SubagentResult(
+            task_id=ctx.task.task_id, status="failed",
+            error=f"All recovery exhausted: {ctx.failure_reason}",
+        )
+
+
+class RecoveryChain:
+    """Chain of recovery strategies executed in order."""
+
+    def __init__(self, strategies: list[RecoveryStrategy] | None = None):
+        self._strategies = strategies or [
+            SplitRetryStrategy(),
+            ModelSwitchStrategy(),
+            MarkFailedStrategy(),
+        ]
+
+    def attempt(
+        self, task: SubagentTask, strategy: LlmStrategy, failure_reason: str
+    ) -> SubagentResult:
+        """Try recovery strategies in order until one succeeds."""
+        ctx = RecoveryContext(task=task, strategy=strategy, failure_reason=failure_reason)
+
+        for recovery in self._strategies:
+            ctx.attempt += 1
+            if not recovery.can_handle(ctx):
+                continue
+            result = recovery.recover(ctx)
+            if result and result.status == "done":
+                return result
+
+        return SubagentResult(
+            task_id=task.task_id, status="failed",
+            error=f"Recovery chain exhausted: {failure_reason}",
+        )