source-kb 0.2.2__py3-none-any.whl

engine/two_phase.py

@@ -0,0 +1,143 @@
+"""Two-phase generation — CLI-only strategy for split documents.
+
+Phase 1: Each shard generates a structural outline (2-4KB)
+Phase 2: All outlines merged → each shard generates full content with global context
+
+This eliminates information isolation between shards in output-only mode
+(where source is inlined and shards can't read each other's files).
+
+Agent mode doesn't need this — sub-agents can read any file directly.
+"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+from core.interfaces import LlmStrategy
+from engine.runner import SubagentTask, SubagentResult, run_batch
+
+logger = logging.getLogger(__name__)
+
+OUTLINE_SUFFIX = """
+
+## ⚠️ Current Task: Generate document outline only (Phase 1/2)
+
+Please **output only the document outline**, not the full content.
+
+### Outline Format Requirements
+
+```markdown
+## [Section Title]
+
+**Summary**: 1-2 sentences describing the business process covered by this section
+
+**Methods covered**:
+- ClassName.method1 — one-sentence responsibility
+- ClassName.method2 — one-sentence responsibility
+
+**Key business rules** (numbered list, one sentence each):
+1. ...
+2. ...
+
+---
+```
+
+### Rules
+1. Each `##` section corresponds to one business process
+2. Only write title + summary + method list + key rules
+3. Keep the size within **2-4KB**
+"""
+
+FULL_PHASE_PREFIX = """
+
+## Global Document Outline (section plans for all shards)
+
+Below is the combined document outline generated by all shards in Phase 1. You only need to expand **the sections you are responsible for**.
+
+{combined_outline}
+
+---
+
+## Sections You Are Responsible For
+
+Based on your file list, expand the sections in the outline above that belong to you.
+"""
+
+
+def run_two_phase(
+    tasks: list[SubagentTask],
+    strategy: LlmStrategy,
+    max_concurrent: int = 5,
+) -> list[SubagentResult]:
+    """Execute two-phase generation for split documents.
+
+    Args:
+        tasks: Original full-generation tasks (one per shard)
+        strategy: LLM execution strategy
+        max_concurrent: Concurrency limit
+
+    Returns:
+        SubagentResult list from phase 2 (full content)
+    """
+    if not tasks:
+        return []
+
+    # Phase 1: Generate outlines
+    logger.info("[two-phase] Phase 1: generating %d outlines", len(tasks))
+    outline_tasks = [
+        SubagentTask(
+            task_id=f"{t.task_id}__outline",
+            prompt=t.prompt + OUTLINE_SUFFIX,
+            output_path=t.output_path.parent / ".meta" / "outlines" / f"{t.task_id}.md",
+            doc_type=t.doc_type,
+            timeout=300,
+        )
+        for t in tasks
+    ]
+
+    outline_results = run_batch(outline_tasks, strategy, max_concurrent)
+
+    # Collect outlines
+    outlines: list[tuple[str, str]] = []
+    for result in outline_results:
+        if result.status == "done" and result.content:
+            name = result.task_id.replace("__outline", "")
+            outlines.append((name, result.content))
+
+    if not outlines:
+        logger.warning("[two-phase] Phase 1 failed, falling back to direct generation")
+        return run_batch(tasks, strategy, max_concurrent)
+
+    # Merge outlines
+    combined = "\n\n---\n\n".join(f"### Shard: {name}\n\n{content}" for name, content in outlines)
+    logger.info("[two-phase] Phase 1 done: %d outlines, %d chars", len(outlines), len(combined))
+
+    # Phase 2: Full generation with global outline context
+    logger.info("[two-phase] Phase 2: generating full content")
+    prefix = FULL_PHASE_PREFIX.format(combined_outline=combined)
+
+    full_tasks = [
+        SubagentTask(
+            task_id=t.task_id,
+            prompt=_inject_outline(t.prompt, prefix),
+            output_path=t.output_path,
+            doc_type=t.doc_type,
+            timeout=t.timeout,
+        )
+        for t in tasks
+    ]
+
+    results = run_batch(full_tasks, strategy, max_concurrent)
+    logger.info("[two-phase] Phase 2 done: %d/%d succeeded",
+                sum(1 for r in results if r.status == "done"), len(results))
+    return results
+
+
+def _inject_outline(prompt: str, prefix: str) -> str:
+    """Inject outline context into the original prompt."""
+    marker = "## Rules you must follow"
+    if marker in prompt:
+        pos = prompt.index(marker)
+        return prompt[:pos] + prefix + "\n" + prompt[pos:]
+    return prompt + prefix

mcp_server/__init__.py

@@ -0,0 +1,73 @@
+"""source-kb MCP Server — code knowledge base documentation generation toolkit.
+
+Provides AI agents with project discovery, workflow orchestration, skeleton extraction,
+document generation, and coverage validation capabilities via the MCP protocol.
+
+After installation:
+    uvx source-kb-mcp
+
+For local development:
+    python -m mcp_server
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+PROJECT_ROOT = Path(__file__).resolve().parent.parent
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from mcp.server.fastmcp import FastMCP
+
+mcp = FastMCP(
+    "source-kb",
+    instructions=(
+        "A toolkit for automatically generating structured knowledge base documentation from source code. "
+        "Start by calling discover() to understand project status, then use get_workflow() to get workflow guidance."
+    ),
+)
+
+
+def project_root() -> Path:
+    """Get the user's project root directory."""
+    return Path(os.environ.get("SOURCE_KB_ROOT", os.getcwd()))
+
+
+def find_config(root: Path | None = None) -> Path | None:
+    """Find the kb-project.yaml configuration file."""
+    root = root or project_root()
+    for name in ("kb-project.yaml", "kb-project.yml"):
+        p = root / name
+        if p.exists():
+            return p
+    return None
+
+
+# Register all tools
+from mcp_server.tools.discovery import register as _reg_discovery  # noqa: E402
+from mcp_server.tools.workflow import register as _reg_workflow  # noqa: E402
+from mcp_server.tools.source import register as _reg_source  # noqa: E402
+from mcp_server.tools.planning import register as _reg_planning  # noqa: E402
+from mcp_server.tools.validation import register as _reg_validation  # noqa: E402
+from mcp_server.tools.generation import register as _reg_generation  # noqa: E402
+from mcp_server.tools.config import register as _reg_config  # noqa: E402
+
+_reg_discovery(mcp)
+_reg_workflow(mcp)
+_reg_source(mcp)
+_reg_planning(mcp)
+_reg_validation(mcp)
+_reg_generation(mcp)
+_reg_config(mcp)
+
+
+def main():
+    """MCP Server entry point."""
+    mcp.run()
+
+
+if __name__ == "__main__":
+    main()

mcp_server/__main__.py

@@ -0,0 +1,5 @@
+"""__main__ support for python -m mcp_server."""
+
+from mcp_server import main
+
+main()

mcp_server/tools/__init__.py

@@ -0,0 +1 @@
+"""MCP Server tool modules."""

mcp_server/tools/config.py

@@ -0,0 +1,63 @@
+"""Config tools — preset listing and configuration queries."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from mcp.server.fastmcp import FastMCP
+
+
+def register(mcp: FastMCP) -> None:
+    """Register config tools on the MCP server."""
+
+    @mcp.tool()
+    def list_presets() -> str:
+        """List available language presets and their supported document types.
+
+        Returns the name, description, and supported doc_types list for each preset.
+        """
+        from mcp_server import PROJECT_ROOT
+
+        presets_dir = PROJECT_ROOT / "presets"
+        result = {}
+
+        if presets_dir.is_dir():
+            import yaml
+            for preset_dir in sorted(presets_dir.iterdir()):
+                if not preset_dir.is_dir():
+                    continue
+                preset_yaml = preset_dir / "preset.yaml"
+                if not preset_yaml.exists():
+                    continue
+
+                cfg = yaml.safe_load(preset_yaml.read_text(encoding="utf-8"))
+
+                # Merge doc_types.yaml if exists
+                doc_types_yaml = preset_dir / "doc_types.yaml"
+                doc_types = {}
+                if doc_types_yaml.exists():
+                    raw = yaml.safe_load(doc_types_yaml.read_text(encoding="utf-8")) or {}
+                    doc_types = raw.get("doc_types", raw)
+                elif "doc_types" in cfg:
+                    doc_types = cfg["doc_types"]
+
+                doc_type_list = []
+                for dt_key, dt_cfg in doc_types.items():
+                    if isinstance(dt_cfg, dict):
+                        doc_type_list.append({
+                            "name": dt_key,
+                            "filename": dt_cfg.get("filename", f"{dt_key}.md"),
+                            "batch": dt_cfg.get("batch", 99),
+                            "conditional": dt_cfg.get("conditional", False),
+                            "global_view": dt_cfg.get("global_view", False),
+                        })
+
+                result[preset_dir.name] = {
+                    "name": cfg.get("name", preset_dir.name),
+                    "description": cfg.get("description", ""),
+                    "doc_types": doc_type_list,
+                    "doc_type_count": len(doc_type_list),
+                }
+
+        return json.dumps({"status": "ok", "presets": result}, ensure_ascii=False, indent=2)

mcp_server/tools/discovery.py

@@ -0,0 +1,276 @@
+"""Discovery tools — project detection and initialization."""
+
+from __future__ import annotations
+
+import json
+import re
+from pathlib import Path
+from typing import Any
+
+from mcp.server.fastmcp import FastMCP
+
+
+def register(mcp: FastMCP) -> None:
+    """Register discovery tools on the MCP server."""
+
+    @mcp.tool()
+    def discover() -> str:
+        """Discover project status and available workflows. This is the entry-point tool that agents should call first.
+
+        Scans the current directory, detects whether a kb-project.yaml configuration exists,
+        checks the state of any generated knowledge bases, and returns available workflows
+        with suggested next actions.
+        """
+        from mcp_server import project_root, find_config
+        from mcp_server.workflow_loader import list_workflows
+
+        root = project_root()
+        config_path = find_config(root)
+
+        result: dict[str, Any] = {
+            "project_root": str(root),
+            "project_detected": config_path is not None,
+            "config_path": str(config_path) if config_path else None,
+            "knowledge_bases": [],
+            "available_workflows": list_workflows(),
+            "suggested_action": "",
+        }
+
+        if config_path:
+            import yaml
+            try:
+                cfg = yaml.safe_load(config_path.read_text(encoding="utf-8"))
+                base_dir = config_path.parent
+                for kb_name, kb_cfg in cfg.get("knowledge_bases", {}).items():
+                    kb_dir = Path(kb_cfg.get("knowledge_dir", ""))
+                    if not kb_dir.is_absolute():
+                        kb_dir = (base_dir / kb_dir).resolve()
+
+                    source = kb_cfg.get("source", {})
+                    if source.get("structure") == "monorepo":
+                        modules = [m["name"] for m in source.get("modules", [])]
+                    else:
+                        modules = [r["name"] for r in source.get("repos", [])]
+
+                    has_docs = kb_dir.is_dir() and any(kb_dir.rglob("*.md"))
+
+                    result["knowledge_bases"].append({
+                        "name": kb_name,
+                        "preset": kb_cfg.get("preset", "generic"),
+                        "modules": modules,
+                        "has_docs": has_docs,
+                        "knowledge_dir": str(kb_dir),
+                    })
+
+                if result["knowledge_bases"]:
+                    has_any_docs = any(kb["has_docs"] for kb in result["knowledge_bases"])
+                    if has_any_docs:
+                        result["suggested_action"] = (
+                            "Knowledge base exists. Use get_workflow('kb-sync') for incremental updates, "
+                            "or get_workflow('kb-audit') to audit quality."
+                        )
+                    else:
+                        result["suggested_action"] = (
+                            "Configuration is ready but no documents exist yet. Call get_workflow('kb-init') to start initialization."
+                        )
+                else:
+                    result["suggested_action"] = "Configuration file is empty. Call init_project() to generate configuration."
+            except Exception as e:
+                result["error"] = f"Configuration parsing failed: {e}"
+                result["suggested_action"] = "Configuration file parsing failed. Please check the YAML format."
+        else:
+            result["suggested_action"] = (
+                "No kb-project.yaml detected. Call detect_project() to scan project structure, "
+                "then call init_project() to generate configuration."
+            )
+
+        return json.dumps(result, ensure_ascii=False, indent=2)
+
+    @mcp.tool()
+    def detect_project(path: str = ".") -> str:
+        """Detect project structure and recommend kb-project.yaml configuration.
+
+        Scans for pom.xml, build.gradle, .git, etc. to automatically infer project type,
+        module list, and recommended preset. Returns a configuration suggestion ready to be written.
+
+        Args:
+            path: Project root directory path (defaults to current directory)
+        """
+        from mcp_server import project_root
+
+        root = Path(path).resolve() if path != "." else project_root()
+
+        result: dict[str, Any] = {
+            "project_root": str(root),
+            "preset": "generic",
+            "structure": "unknown",
+            "modules": [],
+            "git_url": None,
+            "branch": "main",
+        }
+
+        # Detect Java/Spring
+        pom = root / "pom.xml"
+        if pom.exists():
+            result["preset"] = "java-spring"
+            content = pom.read_text(encoding="utf-8", errors="replace")
+            if "<modules>" in content:
+                result["structure"] = "monorepo"
+                modules = re.findall(r"<module>([^<]+)</module>", content)
+                for m in modules:
+                    mod_type = "service"
+                    if "api" in m.lower() or "contract" in m.lower():
+                        mod_type = "api-contract"
+                    elif "common" in m.lower() or "base" in m.lower() or "lib" in m.lower():
+                        mod_type = "library"
+                    result["modules"].append({"name": m, "path": m, "type": mod_type})
+            else:
+                result["structure"] = "multi-repo"
+                result["modules"].append({"name": root.name, "type": "service"})
+
+        # Detect git remote
+        git_config = root / ".git" / "config"
+        if git_config.exists():
+            content = git_config.read_text(encoding="utf-8", errors="replace")
+            urls = re.findall(r"url\s*=\s*(.+)", content)
+            if urls:
+                result["git_url"] = urls[0].strip()
+            # Detect current branch
+            head = root / ".git" / "HEAD"
+            if head.exists():
+                head_content = head.read_text(encoding="utf-8").strip()
+                if head_content.startswith("ref: refs/heads/"):
+                    result["branch"] = head_content.replace("ref: refs/heads/", "")
+
+        # Generate suggested config
+        project_name = root.name.lower().replace(" ", "-")
+        if result["structure"] == "monorepo" and result["modules"]:
+            source_block = {
+                "structure": "monorepo",
+                "url": result["git_url"] or "https://your-git-server/repo.git",
+                "repo_name": project_name,
+                "branch": result["branch"],
+                "cache_dir": "./.source-cache",
+                "modules": result["modules"],
+            }
+        else:
+            repos = []
+            for m in result["modules"]:
+                repos.append({
+                    "name": m["name"],
+                    "url": result["git_url"] or "https://your-git-server/repo.git",
+                    "branch": result["branch"],
+                    "type": m.get("type", "service"),
+                })
+            source_block = {
+                "structure": "multi-repo",
+                "cache_dir": "./.source-cache",
+                "repos": repos,
+            }
+
+        import yaml
+        suggested = {
+            "version": 1,
+            "knowledge_bases": {
+                project_name: {
+                    "name": project_name,
+                    "collection": f"{project_name.replace('-', '_')}_knowledge",
+                    "knowledge_dir": f"./knowledge/{project_name}",
+                    "preset": result["preset"],
+                    "source": source_block,
+                }
+            },
+            "embedding": {"backend": "chromadb-default"},
+            "agent": {"model": "delegated"},
+        }
+        result["suggested_config"] = yaml.dump(
+            suggested, allow_unicode=True, default_flow_style=False, sort_keys=False
+        )
+
+        return json.dumps(result, ensure_ascii=False, indent=2)
+
+    @mcp.tool()
+    def init_project(
+        project_name: str,
+        preset: str = "generic",
+        structure: str = "multi-repo",
+        knowledge_dir: str = "",
+        repos: str = "[]",
+        modules: str = "[]",
+        git_url: str = "",
+        branch: str = "main",
+    ) -> str:
+        """Generate a kb-project.yaml configuration file.
+
+        Args:
+            project_name: Project/knowledge base name (lowercase with hyphens)
+            preset: Language preset (generic | java-spring)
+            structure: Repository structure (multi-repo | monorepo)
+            knowledge_dir: Knowledge base output directory (defaults to ./knowledge/{project_name})
+            repos: JSON list of repositories for multi-repo [{name, url, branch, type}]
+            modules: JSON list of modules for monorepo [{name, path, type}]
+            git_url: Repository URL for monorepo
+            branch: Branch name for monorepo
+        """
+        import yaml
+        from mcp_server import project_root
+
+        root = project_root()
+        if not knowledge_dir:
+            knowledge_dir = f"./knowledge/{project_name}"
+
+        if structure == "monorepo":
+            try:
+                modules_list = json.loads(modules)
+            except (json.JSONDecodeError, TypeError):
+                modules_list = []
+            source_block: dict[str, Any] = {
+                "structure": "monorepo",
+                "url": git_url or "https://your-git-server/repo.git",
+                "repo_name": project_name,
+                "branch": branch,
+                "cache_dir": "./.source-cache",
+                "modules": modules_list if modules_list else [
+                    {"name": "example-service", "path": "services/example", "type": "service"}
+                ],
+            }
+        else:
+            try:
+                repos_list = json.loads(repos)
+            except (json.JSONDecodeError, TypeError):
+                repos_list = []
+            source_block = {
+                "structure": "multi-repo",
+                "cache_dir": "./.source-cache",
+                "repos": repos_list if repos_list else [
+                    {"name": "example-service", "url": "https://git.example.com/repo.git",
+                     "branch": "main", "type": "service"}
+                ],
+            }
+
+        config = {
+            "version": 1,
+            "knowledge_bases": {
+                project_name: {
+                    "name": project_name,
+                    "collection": f"{project_name.replace('-', '_')}_knowledge",
+                    "knowledge_dir": knowledge_dir,
+                    "preset": preset,
+                    "source": source_block,
+                }
+            },
+            "embedding": {"backend": "chromadb-default"},
+            "agent": {"model": "delegated"},
+        }
+
+        output_path = root / "kb-project.yaml"
+        output_path.write_text(
+            yaml.dump(config, allow_unicode=True, default_flow_style=False, sort_keys=False),
+            encoding="utf-8",
+        )
+
+        return json.dumps({
+            "status": "ok",
+            "config_path": str(output_path),
+            "message": f"Created {output_path}. Please review the repository URLs and other settings, then call get_workflow('kb-init') to start initialization.",
+        }, ensure_ascii=False, indent=2)