source-kb 0.2.2__py3-none-any.whl

engine/pipeline/steps/llm_plan_split.py

@@ -0,0 +1,253 @@
+"""LLM smart split step — fallback when community detection fails.
+
+When deterministic community detection cannot produce a good split plan,
+this module calls LLM to analyze the skeleton and suggest groupings.
+
+Requirements: 9.3
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import re
+from pathlib import Path
+
+from core.interfaces import LlmStrategy, LlmRequest
+
+logger = logging.getLogger(__name__)
+
+# Prompt template for LLM split planning
+SPLIT_PROMPT_TEMPLATE = """# Task: Analyze source code structure and split by business domain
+
+You are a senior system architect. Analyze the following source code skeleton, identify business domain boundaries, and group files by business domain.
+
+## Target Document Type
+{doc_type}
+
+## Source Code Skeleton Summary
+{skeleton_summary}
+
+## Split Requirements
+
+1. **Identify business domains**: Identify independent business domains based on class names, method names, and package structure
+2. **Granularity control**: Keep each shard within 2000-5000 lines of code
+3. **High cohesion**: Files belonging to the same business function should be in the same shard
+4. **Low coupling**: Different business domains should be separated as much as possible
+5. **Maximum {max_splits} shards**
+
+## Output Format
+
+Output JSON directly:
+{{
+  "analysis": "Brief rationale for domain partitioning",
+  "splits": [
+    {{
+      "name": "domain-name (English, kebab-case)",
+      "description": "Domain description",
+      "files": ["FileName1.java", "FileName2.java"],
+      "estimated_lines": 3000
+    }}
+  ]
+}}
+
+Ensure all files are assigned to a shard.
+"""
+
+def _get_focus_hint(doc_type: str, preset_name: str = "java-spring") -> str:
+    """Get focus hint from hooks, with empty fallback."""
+    try:
+        from core.preset_hooks import load_hooks
+        hooks = load_hooks(preset_name)
+        return hooks.get_focus_hint(doc_type)
+    except Exception:
+        return ""
+
+
+def plan_splits_via_llm(
+    module_dir: Path,
+    doc_type: str,
+    module_name: str,
+    strategy: LlmStrategy,
+    max_splits: int = 4,
+) -> dict:
+    """Use LLM to generate a split plan for a doc type.
+
+    Called as fallback when community detection fails or produces
+    poor results (e.g., all files in one community).
+
+    Args:
+        module_dir: Module knowledge directory
+        doc_type: Target document type
+        module_name: Module name
+        strategy: LLM strategy for making calls
+        max_splits: Maximum number of splits
+
+    Returns:
+        Split plan dict with keys: module, doc_type, analysis, splits, source
+    """
+    skeleton_summary = _load_skeleton_summary(module_dir)
+    if not skeleton_summary:
+        return {"error": "No skeleton data available", "splits": []}
+
+    focus = _get_focus_hint(doc_type) or "Group by file functionality"
+    prompt = SPLIT_PROMPT_TEMPLATE.format(
+        doc_type=f"{doc_type} ({focus})",
+        skeleton_summary=skeleton_summary,
+        max_splits=max_splits,
+    )
+
+    try:
+        resp = strategy.call(LlmRequest(
+            system="You are a code architecture analyst. Output only valid JSON.",
+            user=prompt,
+            max_tokens=2000,
+            temperature=0.0,
+        ))
+    except Exception as e:
+        logger.error("[llm-split] LLM call failed: %s", e)
+        return {"error": str(e), "splits": []}
+
+    if resp.status != "done":
+        return {"error": f"LLM status: {resp.status}", "splits": []}
+
+    result = _parse_response(resp.content)
+    if not result.get("splits"):
+        return {"error": "LLM returned no valid splits", "splits": []}
+
+    # Enforce max_splits limit
+    splits = result["splits"]
+    if len(splits) > max_splits:
+        splits = sorted(splits, key=lambda x: x.get("estimated_lines", 0), reverse=True)
+        splits = splits[:max_splits]
+
+    return {
+        "module": module_name,
+        "doc_type": doc_type,
+        "analysis": result.get("analysis", ""),
+        "recommended_agents": len(splits),
+        "splits": splits,
+        "source": "llm",
+    }
+
+
+def apply_split_plan(
+    module_dir: Path,
+    doc_type: str,
+    plan: dict,
+    source_cache: Path,
+) -> list[list[str]]:
+    """Convert a split plan into file groups for shard generation.
+
+    Args:
+        module_dir: Module knowledge directory
+        doc_type: Document type
+        plan: Split plan from plan_splits_via_llm
+        source_cache: Source cache directory for resolving paths
+
+    Returns:
+        List of file groups (each group is a list of file paths)
+    """
+    splits = plan.get("splits", [])
+    if not splits:
+        return []
+
+    # Build filename → full path mapping from file list
+    from core.paths import resolve_file_list
+    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 []
+
+    # Map basenames to full paths
+    basename_map: dict[str, list[str]] = {}
+    for fp in all_files:
+        basename = fp.rsplit("/", 1)[-1] if "/" in fp else fp
+        basename_map.setdefault(basename, []).append(fp)
+
+    groups: list[list[str]] = []
+    assigned: set[str] = set()
+
+    for split in splits:
+        group: list[str] = []
+        for filename in split.get("files", []):
+            # Try exact basename match
+            matches = basename_map.get(filename, [])
+            for m in matches:
+                if m not in assigned:
+                    group.append(m)
+                    assigned.add(m)
+        groups.append(group)
+
+    # Assign unmatched files to the largest group
+    unassigned = [f for f in all_files if f not in assigned]
+    if unassigned and groups:
+        largest = max(groups, key=len)
+        largest.extend(unassigned)
+
+    return [g for g in groups if g]
+
+
+def _load_skeleton_summary(module_dir: Path, max_entries: int = 80) -> str:
+    """Load skeleton summary for prompt context."""
+    summary_file = module_dir / ".meta" / "skeleton" / "summary.json"
+    if not summary_file.exists():
+        summary_file = module_dir / ".skeleton-summary.json"
+    if not summary_file.exists():
+        return ""
+
+    try:
+        data = json.loads(summary_file.read_text(encoding="utf-8"))
+    except (json.JSONDecodeError, OSError):
+        return ""
+
+    lines = []
+    for entry in data[:max_entries]:
+        file_path = entry.get("file", "")
+        classes = entry.get("classes", [])
+        methods = entry.get("methods", [])
+
+        class_names = [c.get("name", "") for c in classes]
+        method_names = [m.get("name", "") for m in methods[:5]]
+
+        parts = [file_path]
+        if class_names:
+            parts.append(f"  classes: {', '.join(class_names)}")
+        if method_names:
+            parts.append(f"  methods: {', '.join(method_names)}")
+        lines.append("\n".join(parts))
+
+    return "\n".join(lines)
+
+
+def _parse_response(text: str) -> dict:
+    """Parse LLM response, extracting JSON."""
+    text = text.strip()
+
+    # Remove markdown code blocks
+    if text.startswith("```"):
+        lines = text.split("\n")
+        if lines[0].startswith("```"):
+            lines = lines[1:]
+        if lines and lines[-1].strip() == "```":
+            lines = lines[:-1]
+        text = "\n".join(lines)
+
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        pass
+
+    # Try extracting JSON block
+    match = re.search(r'\{[\s\S]*\}', text)
+    if match:
+        try:
+            return json.loads(match.group())
+        except json.JSONDecodeError:
+            pass
+
+    return {}

engine/pipeline/steps/lock.py

@@ -0,0 +1,64 @@
+"""Lock management steps — acquire and release .kb-lock."""
+
+from __future__ import annotations
+
+from filelock import FileLock, Timeout
+
+from core.interfaces import Step, StepResult, PipelineContext
+from engine.pipeline import register_step
+
+
+@register_step
+class AcquireLockStep(Step):
+    """Acquire file lock on knowledge directory."""
+
+    default_name = "acquire-lock"
+
+    def __init__(self, timeout_minutes: int = 120):
+        super().__init__("acquire-lock")
+        self._timeout = timeout_minutes * 60
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        lock_path = ctx.knowledge_dir / ".kb-lock"
+        lock_path.parent.mkdir(parents=True, exist_ok=True)
+        lock = FileLock(str(lock_path), timeout=self._timeout)
+        try:
+            lock.acquire()
+        except Timeout:
+            return StepResult(status="failed", message=f"Lock timeout ({self._timeout // 60}min)")
+        ctx.state["_lock"] = lock
+        ctx.state["_lock_path"] = lock_path
+        return StepResult(status="ok", message=f"Lock acquired: {lock_path}")
+
+    def rollback(self, ctx: PipelineContext) -> None:
+        _release(ctx)
+
+
+@register_step
+class ReleaseLockStep(Step):
+    """Release file lock."""
+
+    default_name = "release-lock"
+
+    def __init__(self):
+        super().__init__("release-lock")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        return _release(ctx)
+
+
+def _release(ctx: PipelineContext) -> StepResult:
+    lock: FileLock | None = ctx.state.pop("_lock", None)
+    if lock is None:
+        return StepResult(status="skipped", message="No lock to release")
+    try:
+        lock.release()
+    except Exception:
+        pass
+    lock_path = ctx.state.pop("_lock_path", None)
+    if lock_path and lock_path.exists():
+        try:
+            lock_path.unlink()
+        except OSError:
+            pass
+    return StepResult(status="ok", message="Lock released")

engine/pipeline/steps/preflight.py

@@ -0,0 +1,237 @@
+"""Model capacity preflight check — validates prompt sizes against model context window.
+
+Three-phase check:
+1. Render actual prompts to get precise byte sizes
+2. LLM probe call (max_tokens=1) to verify connectivity + get token count
+3. Auto-adjust split thresholds if over limit, re-split and rewrite dispatch-tasks
+
+Runs before document generation in CLI pipeline.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from pathlib import Path
+
+from core.interfaces import Step, StepResult, PipelineContext, LlmRequest
+from engine.pipeline import register_step
+
+logger = logging.getLogger(__name__)
+
+_MODEL_CAPACITY: dict[str, int] = {
+    "claude-opus-4-6": 200000,
+    "claude-sonnet-4-6": 200000,
+    "claude-sonnet-4.6": 200000,
+    "claude-opus-4.6": 200000,
+    "claude-3-5-sonnet": 200000,
+    "claude-3-opus": 200000,
+    "claude-3-haiku": 200000,
+    "gpt-4o": 128000,
+    "gpt-4o-mini": 128000,
+    "gpt-4-turbo": 128000,
+    "gpt-4": 8192,
+    "gpt-3.5-turbo": 16384,
+    "deepseek-chat": 64000,
+    "deepseek-coder": 64000,
+    "qwen-turbo": 8000,
+    "qwen-plus": 32000,
+    "qwen-max": 32000,
+}
+
+BYTES_PER_TOKEN = 3
+INPUT_BUDGET_RATIO = 0.80
+
+
+def _get_model_capacity_tokens(model: str, config: dict) -> int:
+    """Get model max input tokens. Resolution: config override > known table > 128K default."""
+    explicit = config.get("agent", {}).get("model_max_input_tokens")
+    if explicit:
+        return int(explicit)
+
+    model_lower = model.lower().strip()
+    for known, capacity in _MODEL_CAPACITY.items():
+        if known in model_lower or model_lower in known:
+            return capacity
+
+    return 128000
+
+
+@register_step
+class PreflightCheckStep(Step):
+    """Validate model capacity, probe connectivity, auto-adjust if needed."""
+
+    default_name = "preflight-check"
+
+    def __init__(self):
+        super().__init__("preflight-check")
+
+    def run(self, ctx: PipelineContext) -> StepResult:
+        config = ctx.config
+        model = os.getenv("LLM_MODEL", "") or config.get("agent", {}).get("model", "unknown")
+        model_capacity = _get_model_capacity_tokens(model, config)
+        input_budget = int(model_capacity * INPUT_BUDGET_RATIO)
+
+        phase1 = self._phase1_render_prompts(ctx, config, model)
+        if phase1["error"]:
+            return StepResult(status="ok", message=f"Preflight: {phase1['error']}", details=phase1)
+
+        max_prompt_bytes = phase1["max_bytes"]
+        max_task_id = phase1["max_task_id"]
+        max_prompt = phase1["max_prompt"]
+        all_sizes = phase1["all_sizes"]
+
+        logger.info("Preflight phase 1: max prompt %dKB (%s), %d tasks total",
+                    max_prompt_bytes // 1024, max_task_id, len(all_sizes))
+
+        phase2 = self._phase2_probe(ctx, config, max_prompt, model)
+        actual_input_tokens = phase2.get("input_tokens")
+        probe_error = phase2.get("error", "")
+        context_exceeded = phase2.get("context_exceeded", False)
+
+        if probe_error and not context_exceeded:
+            msg = f"Preflight probe failed: {probe_error}"
+            logger.warning(msg)
+            return StepResult(status="failed", message=msg, details={
+                "model": model, "probe_error": probe_error,
+            })
+
+        if actual_input_tokens:
+            utilization = actual_input_tokens / input_budget
+        elif context_exceeded:
+            utilization = 1.5
+        else:
+            estimated_tokens = max_prompt_bytes // BYTES_PER_TOKEN
+            utilization = estimated_tokens / input_budget
+
+        details = {
+            "model": model,
+            "model_capacity_tokens": model_capacity,
+            "input_budget_tokens": input_budget,
+            "max_prompt_bytes": max_prompt_bytes,
+            "max_task_id": max_task_id,
+            "actual_input_tokens": actual_input_tokens,
+            "utilization": f"{utilization:.0%}",
+            "task_count": len(all_sizes),
+        }
+
+        if utilization > 0.85 or context_exceeded:
+            from engine.pipeline.steps.preflight_adjust import auto_adjust
+            adjust_result = auto_adjust(ctx, config, utilization, input_budget, all_sizes)
+            if adjust_result:
+                details["adjustment"] = adjust_result
+                msg = (
+                    f"Preflight: {model} capacity exceeded (utilization={utilization:.0%}), "
+                    f"auto-adjusted splits: {adjust_result['summary']}"
+                )
+                logger.info(msg)
+                return StepResult(status="ok", message=msg, details=details)
+            else:
+                msg = (
+                    f"Preflight: {model} capacity exceeded (utilization={utilization:.0%}), "
+                    f"auto-adjust failed — reduce max_lines/max_bytes or switch model"
+                )
+                return StepResult(status="failed", message=msg, details=details)
+
+        msg = f"Preflight OK: {model}, max prompt {max_prompt_bytes//1024}KB, utilization {utilization:.0%}"
+        return StepResult(status="ok", message=msg, details=details)
+
+    def _phase1_render_prompts(self, ctx: PipelineContext, config: dict, model: str) -> dict:
+        """Render all prompts and measure actual sizes."""
+        from core.preset import load_preset, get_template_path
+        from core.prompt.renderer import render_prompt
+        from engine.assembler import InlinePromptAssembler
+
+        preset_name = ctx.kb_config.get("preset", "generic")
+        preset = load_preset(preset_name)
+        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 ""
+
+        all_sizes: list[dict] = []
+        max_bytes = 0
+        max_task_id = ""
+        max_prompt = ""
+
+        for name in ctx.state.get("module_repos", {}):
+            module_dir = ctx.knowledge_dir / name
+            tasks_file = module_dir / ".meta" / "dispatch-tasks.json"
+            if not tasks_file.exists():
+                continue
+
+            manifest = json.loads(tasks_file.read_text(encoding="utf-8"))
+            for task in manifest.get("tasks", []):
+                tpl_name = get_template_path(preset, task["doc_type"], preset_name)
+                if not tpl_name:
+                    continue
+                tpl_path = ctx.project_root / "presets" / preset_name / "templates" / tpl_name
+                if not tpl_path.exists():
+                    continue
+
+                extras = {}
+                if task.get("shard_file_list"):
+                    extras["file_list_override"] = task["shard_file_list"]
+
+                try:
+                    prompt = render_prompt(
+                        template_path=tpl_path, config=config,
+                        kb_name=ctx.kb_name, module_name=name,
+                        doc_type=task["doc_type"], assembler=assembler,
+                        extras=extras, execution_snippet=execution_snippet,
+                        preset=preset,
+                    )
+                except Exception as e:
+                    logger.debug("Preflight render failed for %s: %s", task["id"], e)
+                    continue
+
+                size = len(prompt.encode("utf-8"))
+                all_sizes.append({"task_id": task["id"], "doc_type": task["doc_type"], "bytes": size, "module": name})
+
+                if size > max_bytes:
+                    max_bytes = size
+                    max_task_id = task["id"]
+                    max_prompt = prompt
+
+        if not all_sizes:
+            return {"error": "No renderable tasks found", "max_bytes": 0, "max_task_id": "",
+                    "max_prompt": "", "all_sizes": []}
+
+        return {"error": "", "max_bytes": max_bytes, "max_task_id": max_task_id,
+                "max_prompt": max_prompt, "all_sizes": all_sizes}
+
+    def _phase2_probe(self, ctx: PipelineContext, config: dict, prompt: str, model: str) -> dict:
+        """Send a probe request (max_tokens=1) to verify connectivity and get token count."""
+        from engine.strategies import create_strategy, ConfigProxy
+
+        if not prompt:
+            return {"error": "No prompt to probe"}
+
+        try:
+            strategy = create_strategy(ConfigProxy(config, timeout_override=30))
+        except Exception as e:
+            return {"error": f"Strategy init failed: {e}"}
+
+        try:
+            resp = strategy.call(LlmRequest(
+                system="Reply with a single dot.",
+                user=prompt,
+                model=None,
+                max_tokens=1,
+                temperature=0,
+            ))
+        except Exception as e:
+            error_str = str(e)
+            if "context_length" in error_str or "too long" in error_str.lower():
+                return {"context_exceeded": True, "error": error_str}
+            return {"error": error_str}
+
+        if resp.status == "failed":
+            error_str = resp.error
+            if "context_length" in error_str or "too long" in error_str.lower() or "maximum" in error_str.lower():
+                return {"context_exceeded": True, "error": error_str}
+            return {"error": error_str}
+
+        input_tokens = resp.usage.get("prompt_tokens", 0)
+        return {"input_tokens": input_tokens, "error": ""}

engine/pipeline/steps/preflight_adjust.py

@@ -0,0 +1,147 @@
+"""Preflight auto-adjust — re-split oversized doc types when model capacity is exceeded."""
+
+from __future__ import annotations
+
+import copy
+import json
+import logging
+from pathlib import Path
+
+from core.interfaces import PipelineContext
+
+logger = logging.getLogger(__name__)
+
+BYTES_PER_TOKEN = 3
+
+
+def auto_adjust(
+    ctx: PipelineContext,
+    config: dict,
+    utilization: float,
+    input_budget: int,
+    all_sizes: list[dict],
+) -> dict | None:
+    """Auto-adjust split thresholds for oversized doc types and re-split."""
+    from core.preset import load_preset
+    from core.skeleton.split import SplitConfig
+
+    preset_name = ctx.kb_config.get("preset", "generic")
+    preset = load_preset(preset_name)
+
+    budget_bytes = int(input_budget * BYTES_PER_TOKEN)
+    oversized = [s for s in all_sizes if s["bytes"] > budget_bytes * 0.85]
+
+    if not oversized:
+        oversized = sorted(all_sizes, key=lambda x: -x["bytes"])[:3]
+
+    target_ratio = 0.70 / utilization
+    target_ratio = max(target_ratio, 0.3)
+
+    current_config = SplitConfig.from_preset(preset, "output-only")
+    new_max_lines = max(2000, int(current_config.max_lines * target_ratio))
+    new_max_bytes = max(100 * 1024, int(current_config.max_bytes * target_ratio))
+
+    logger.info("Auto-adjust: ratio=%.2f, max_lines %d->%d, max_bytes %d->%d",
+                target_ratio, current_config.max_lines, new_max_lines,
+                current_config.max_bytes, new_max_bytes)
+
+    affected_doc_types = set(s["doc_type"] for s in oversized)
+    resplit_count = 0
+
+    for name in ctx.state.get("module_repos", {}):
+        module_dir = ctx.knowledge_dir / name
+        tasks_file = module_dir / ".meta" / "dispatch-tasks.json"
+        if not tasks_file.exists():
+            continue
+
+        try:
+            resplit_count += _resplit_module(
+                ctx, name, module_dir, preset, affected_doc_types,
+                new_max_lines, new_max_bytes,
+            )
+        except Exception as e:
+            logger.error("Re-split failed for %s: %s", name, e)
+            return None
+
+    if resplit_count == 0:
+        return None
+
+    return {
+        "summary": f"max_lines={new_max_lines}, max_bytes={new_max_bytes//1024}KB, re-split {resplit_count} doc types",
+        "new_max_lines": new_max_lines,
+        "new_max_bytes": new_max_bytes,
+        "affected_doc_types": list(affected_doc_types),
+        "target_ratio": round(target_ratio, 2),
+    }
+
+
+def _resplit_module(
+    ctx: PipelineContext,
+    module_name: str,
+    module_dir: Path,
+    preset: dict,
+    affected_doc_types: set[str],
+    new_max_lines: int,
+    new_max_bytes: int,
+) -> int:
+    """Re-run split planning for a module with adjusted thresholds."""
+    from core.skeleton.dispatch import compute_dispatch_plan
+    from core.skeleton.dispatch_render import plan_to_tasks, write_shard_files
+
+    adjusted_preset = copy.deepcopy(preset)
+    split_section = adjusted_preset.setdefault("split", {})
+    output_only = split_section.setdefault("output-only", {})
+    output_only["max_lines"] = new_max_lines
+    output_only["max_bytes"] = new_max_bytes
+
+    source = ctx.kb_config.get("source", {})
+    cache_dir = Path(source.get("cache_dir", "./.source-cache"))
+    base_dir = Path(ctx.config.get("_config_dir", ".")).resolve()
+    if not cache_dir.is_absolute():
+        cache_dir = (base_dir / cache_dir).resolve()
+
+    if source.get("structure") == "monorepo":
+        repo_name = source.get("repo_name", "repo")
+        for mod in source.get("modules", []):
+            if mod["name"] == module_name:
+                source_cache = cache_dir / repo_name / mod.get("path", module_name)
+                break
+        else:
+            source_cache = cache_dir / module_name
+    else:
+        source_cache = cache_dir / module_name
+
+    module_type = "service"
+    for repo in source.get("repos", []):
+        if repo["name"] == module_name:
+            module_type = repo.get("type", "service")
+            break
+
+    plan = compute_dispatch_plan(
+        preset=adjusted_preset,
+        module_dir=module_dir,
+        source_cache=source_cache,
+        mode="output-only",
+        module_name=module_name,
+        module_type=module_type,
+    )
+
+    if not plan or not plan.entries:
+        return 0
+
+    preset_name = adjusted_preset.get("name", "java-spring")
+    tasks_manifest = plan_to_tasks(
+        plan=plan, kb_name=ctx.kb_name, preset_name=preset_name,
+        preset=adjusted_preset, knowledge_dir=ctx.knowledge_dir, mode="output-only",
+    )
+    tasks_file = module_dir / ".meta" / "dispatch-tasks.json"
+    tasks_file.write_text(
+        json.dumps(tasks_manifest, ensure_ascii=False, indent=2),
+        encoding="utf-8",
+    )
+
+    write_shard_files(plan, module_dir)
+
+    resplit_count = sum(1 for e in plan.entries if e.doc_type in affected_doc_types)
+    logger.info("Re-split %s: %d doc types adjusted", module_name, resplit_count)
+    return resplit_count