ai-doc-creator 2.2.0__py3-none-any.whl

ai_doc_creator/__init__.py

@@ -0,0 +1,3 @@
+"""AI Document Creator — AI-powered documentation generator (MCP server + CLI)."""
+
+__version__ = "2.2.0"

ai_doc_creator/cli.py

@@ -0,0 +1,100 @@
+# ai_doc_creator/cli.py
+import argparse
+import asyncio
+import logging
+import sys
+
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from .core.config import resolve_config  # noqa: E402
+from .core.backends import pick_backend, BackendError  # noqa: E402
+from .core.sources import GitSource, LocalSource  # noqa: E402
+from .core.file_traverser import FileTraverser  # noqa: E402
+from .core.graph import app  # noqa: E402
+from .core.doc_writer import DocumentationWriter  # noqa: E402
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+logger = logging.getLogger(__name__)
+
+
+async def run(source, output_dir, config):
+    repo_path = None
+    try:
+        repo_path = source.prepare()
+        files = list(FileTraverser(repo_path, max_file_size_kb=config.max_file_size_kb).traverse())
+        logger.info("Found %d files to process.", len(files))
+        if not files:
+            logger.warning("No files found to document. Exiting.")
+            return
+
+        backend = pick_backend(config, ctx=None)  # CLI has no host -> requires a provider key
+        final_state = await app.ainvoke(
+            {
+                "repo_path": repo_path,
+                "files": files,
+                "documents": {},
+                "index_content": "",
+                "backend": backend,
+                "max_concurrency": config.max_concurrency,
+                "profile": config.profile,
+                "diagrams": config.diagrams,
+            }
+        )
+
+        DocumentationWriter(output_dir).write_docs(
+            final_state["documents"], final_state["index_content"]
+        )
+        logger.info("Documentation generation complete!")
+    finally:
+        if source:
+            source.cleanup()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="AI Document Creator")
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("--repo", help="GitHub repository URL")
+    group.add_argument("--path", help="Local project directory")
+    parser.add_argument("--output", default="docs", help="Output directory")
+    parser.add_argument(
+        "--provider", default=None, help="LLM provider (anthropic/openai/azure/bedrock/ollama)"
+    )
+    parser.add_argument("--model", default=None, help="Model name override")
+    parser.add_argument(
+        "--profile",
+        default="readme",
+        choices=["readme", "api", "architecture", "tutorial"],
+        help="Documentation style (default: readme)",
+    )
+    parser.add_argument(
+        "--no-diagrams",
+        action="store_true",
+        help="Disable Mermaid architecture diagrams and flow charts",
+    )
+    args = parser.parse_args()
+
+    config = resolve_config(
+        provider=args.provider,
+        model=args.model,
+        profile=args.profile,
+        diagrams=not args.no_diagrams,
+    )
+    source = GitSource(args.repo) if args.repo else LocalSource(args.path)
+
+    try:
+        asyncio.run(run(source, args.output, config))
+    except BackendError as exc:
+        logger.error("%s", exc)
+        sys.exit(1)
+    except Exception as exc:
+        logger.error("An error occurred: %s", exc)
+        import traceback
+
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

ai_doc_creator/core/__init__.py

@@ -0,0 +1 @@
+"""Core pipeline: config, backends, sources, guards, graph, cache, writer."""

ai_doc_creator/core/backends.py

@@ -0,0 +1,126 @@
+# core/backends.py
+from __future__ import annotations
+
+import os
+from abc import ABC, abstractmethod
+
+from langchain.chat_models import init_chat_model  # imported at top so tests can monkeypatch it
+
+from .config import DocConfig
+
+
+class BackendError(RuntimeError):
+    """Raised when no usable LLM backend can be constructed."""
+
+
+def _content_to_text(content) -> str:
+    """Flatten a chat model response (str, or list of content blocks) to plain text."""
+    if isinstance(content, str):
+        return content
+    if isinstance(content, list):
+        parts: list[str] = []
+        for block in content:
+            if isinstance(block, str):
+                parts.append(block)
+            elif isinstance(block, dict) and block.get("type") == "text":
+                parts.append(block.get("text", ""))
+        return "".join(parts)
+    return str(content)
+
+
+class CompletionBackend(ABC):
+    """Generates text from a prompt. The only thing the pipeline knows about an LLM."""
+
+    @abstractmethod
+    async def complete(self, prompt: str) -> str: ...
+
+
+class FakeBackend(CompletionBackend):
+    """Deterministic backend for tests — never calls a real model."""
+
+    def __init__(self, response: str = "FAKE DOC"):
+        self.response = response
+        self.calls: list[str] = []
+
+    async def complete(self, prompt: str) -> str:
+        self.calls.append(prompt)
+        return self.response
+
+
+class ProviderBackend(CompletionBackend):
+    """Calls a real provider (anthropic/openai/azure/bedrock/ollama) via LangChain."""
+
+    def __init__(self, config: DocConfig):
+        if not config.has_provider:
+            raise BackendError("ProviderBackend requires a provider")
+        if config.provider == "azure":
+            self._model = self._build_azure(config)
+        else:
+            # Pass model + provider separately so model names containing ':'
+            # (e.g. bedrock 'amazon.nova-pro-v1:0') are never mis-parsed.
+            # An explicit per-request key (BYOK) wins over any env credential.
+            extra: dict = {}
+            if config.api_key:
+                extra["api_key"] = config.api_key
+            self._model = init_chat_model(
+                model=config.resolved_model,
+                model_provider=config.lc_provider,
+                temperature=0,
+                **extra,
+            )
+
+    @staticmethod
+    def _build_azure(config: DocConfig):
+        from langchain_openai import AzureChatOpenAI
+
+        from pydantic import SecretStr
+
+        raw_key = config.api_key or os.getenv("AZURE_OPENAI_API_KEY")
+        return AzureChatOpenAI(
+            azure_deployment=config.model or os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME", "gpt-4o"),
+            api_version=os.getenv("AZURE_OPENAI_API_VERSION", "2024-12-01-preview"),
+            azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
+            api_key=SecretStr(raw_key) if raw_key else None,
+            temperature=0,
+        )
+
+    async def complete(self, prompt: str) -> str:
+        result = await self._model.ainvoke(prompt)
+        return _content_to_text(getattr(result, "content", result))
+
+
+class SamplingBackend(CompletionBackend):
+    """Asks the MCP host's own model to generate — zero API cost to the operator."""
+
+    def __init__(self, ctx, max_tokens: int = 4096):
+        self._ctx = ctx
+        self._max_tokens = max_tokens
+
+    async def complete(self, prompt: str) -> str:
+        from mcp.types import SamplingMessage, TextContent
+
+        result = await self._ctx.session.create_message(
+            messages=[
+                SamplingMessage(role="user", content=TextContent(type="text", text=prompt))
+            ],
+            max_tokens=self._max_tokens,
+        )
+        content = result.content
+        return content.text if getattr(content, "type", None) == "text" else str(content)
+
+
+def pick_backend(config: DocConfig, ctx=None) -> CompletionBackend:
+    """Choose a backend: provider key → retry-wrapped ProviderBackend;
+    MCP host ctx → SamplingBackend; otherwise raise BackendError.
+    """
+    from .retry import with_retry
+
+    if config.has_provider:
+        return with_retry(ProviderBackend(config))
+    if ctx is not None:
+        return SamplingBackend(ctx)
+    raise BackendError(
+        "No LLM available. Set a provider key (ANTHROPIC_API_KEY / OPENAI_API_KEY / "
+        "AZURE_OPENAI_API_KEY / AWS credentials), pass provider='ollama' for a local model, "
+        "or run inside an MCP host that supports sampling."
+    )

ai_doc_creator/core/cache.py

@@ -0,0 +1,71 @@
+# core/cache.py
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+
+MANIFEST_FILENAME = ".ai-docs-manifest.json"
+
+
+def hash_file(repo_path: str, file_path: str) -> str:
+    """Return the SHA-256 hex digest of the content of *repo_path/file_path*."""
+    full = os.path.join(repo_path, file_path)
+    h = hashlib.sha256()
+    with open(full, "rb") as fh:
+        while chunk := fh.read(65536):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def compute_hashes(files: list[str], repo_path: str) -> dict[str, str]:
+    """Return ``{relative_path: sha256}`` for every path in *files*."""
+    return {fp: hash_file(repo_path, fp) for fp in files}
+
+
+def load_manifest(output_dir: str) -> dict[str, str] | None:
+    """Load the hash manifest from *output_dir*.
+
+    Returns ``None`` when no manifest exists or the file is corrupt.
+    """
+    path = os.path.join(output_dir, MANIFEST_FILENAME)
+    try:
+        with open(path, "r", encoding="utf-8") as fh:
+            return json.load(fh)
+    except (FileNotFoundError, json.JSONDecodeError):
+        return None
+
+
+def save_manifest(file_hashes: dict[str, str], output_dir: str) -> None:
+    """Persist *file_hashes* as ``{output_dir}/.ai-docs-manifest.json``."""
+    os.makedirs(output_dir, exist_ok=True)
+    path = os.path.join(output_dir, MANIFEST_FILENAME)
+    with open(path, "w", encoding="utf-8") as fh:
+        json.dump(file_hashes, fh, indent=2, sort_keys=True)
+
+
+def filter_changed(
+    files: list[str],
+    repo_path: str,
+    manifest: dict[str, str] | None,
+) -> tuple[list[str], list[str]]:
+    """Split *files* into ``(changed, unchanged)`` relative to *manifest*.
+
+    All files are returned as changed when *manifest* is ``None`` (first run).
+    Files that cannot be read are treated as changed.
+    """
+    if manifest is None:
+        return list(files), []
+    changed: list[str] = []
+    unchanged: list[str] = []
+    for fp in files:
+        try:
+            current_hash = hash_file(repo_path, fp)
+        except OSError:
+            changed.append(fp)
+            continue
+        if manifest.get(fp) == current_hash:
+            unchanged.append(fp)
+        else:
+            changed.append(fp)
+    return changed, unchanged

ai_doc_creator/core/config.py

@@ -0,0 +1,92 @@
+# core/config.py
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field, fields
+
+DEFAULT_MODELS = {
+    "anthropic": "claude-sonnet-4-6",
+    "openai": "gpt-4o",
+    "azure": None,
+    "bedrock": "amazon.nova-pro-v1:0",
+    "ollama": "llama3.1",
+}
+
+LC_PROVIDER = {
+    "anthropic": "anthropic",
+    "openai": "openai",
+    "azure": "azure_openai",
+    "bedrock": "bedrock_converse",
+    "ollama": "ollama",
+}
+
+_PROVIDER_ENV = {
+    "anthropic": "ANTHROPIC_API_KEY",
+    "openai": "OPENAI_API_KEY",
+    "azure": "AZURE_OPENAI_API_KEY",
+    "bedrock": "AWS_ACCESS_KEY_ID",
+}
+
+
+@dataclass
+class DocConfig:
+    provider: str | None = None
+    model: str | None = None
+    profile: str = "readme"
+    incremental: bool = True
+    diagrams: bool = True
+    max_file_size_kb: int = 100
+    max_concurrency: int = 8
+    pipeline_timeout_s: int = 300
+    # Per-request provider key (BYOK). repr=False so it can never leak via
+    # logging/str(config); it is handed to the provider client explicitly and
+    # never written to the environment.
+    api_key: str | None = field(default=None, repr=False)
+
+    @property
+    def has_provider(self) -> bool:
+        return self.provider is not None
+
+    @property
+    def model_id(self) -> str | None:
+        if not self.provider:
+            return None
+        model = self.model or DEFAULT_MODELS.get(self.provider)
+        return f"{self.provider}:{model}" if model else self.provider
+
+    @property
+    def lc_provider(self) -> str | None:
+        return LC_PROVIDER.get(self.provider, self.provider) if self.provider else None
+
+    @property
+    def resolved_model(self) -> str | None:
+        if not self.provider:
+            return None
+        return self.model or DEFAULT_MODELS.get(self.provider)
+
+
+def detect_provider() -> str | None:
+    """Return the first provider whose credentials are present in the environment."""
+    for provider, env_var in _PROVIDER_ENV.items():
+        if os.getenv(env_var):
+            return provider
+    return None
+
+
+def resolve_config(
+    provider: str | None = None,
+    model: str | None = None,
+    **overrides,
+) -> DocConfig:
+    """Build a DocConfig, auto-detecting the provider from env when not given."""
+    cfg = DocConfig(provider=provider or detect_provider(), model=model)
+    timeout_env = os.getenv("PIPELINE_TIMEOUT_S")
+    if timeout_env and timeout_env.isdigit():
+        cfg.pipeline_timeout_s = int(timeout_env)
+    valid_fields = {f.name for f in fields(DocConfig)}
+    for key, value in overrides.items():
+        if key not in valid_fields:
+            raise TypeError(f"resolve_config() got an unexpected keyword argument '{key}'")
+        if value is not None:
+            setattr(cfg, key, value)
+    return cfg

ai_doc_creator/core/diagrams.py

@@ -0,0 +1,212 @@
+# ai_doc_creator/core/diagrams.py
+"""Deterministic Mermaid diagram generation + validation of LLM-drawn blocks.
+
+The dependency and structure diagrams are produced by static analysis, never
+by a model, so they are always syntactically valid. Model-drawn diagrams pass
+through sanitize_mermaid_blocks(), which downgrades anything suspect to a
+plain fenced block rather than letting it break the rendered page.
+"""
+from __future__ import annotations
+
+import os
+import re
+
+_PY_IMPORT = re.compile(
+    r"^\s*(?:from\s+([.\w]+)\s+import|import\s+([.\w]+))", re.MULTILINE
+)
+_JS_IMPORT = re.compile(
+    r"""(?:import\s+(?:[\w{}*,\s]+\s+from\s+)?|require\(\s*|import\(\s*)['"]([^'"]+)['"]"""
+)
+_JS_EXTENSIONS = (".js", ".jsx", ".ts", ".tsx", ".mjs", ".cjs")
+_JS_RESOLVE_SUFFIXES = ("", ".js", ".jsx", ".ts", ".tsx", ".mjs", ".cjs",
+                        "/index.js", "/index.ts", "/index.jsx", "/index.tsx")
+
+# First token of a fenced block must be one of these for the block to count
+# as plausibly-valid Mermaid.
+_MERMAID_TYPES = (
+    "graph", "flowchart", "sequencediagram", "classdiagram", "statediagram",
+    "statediagram-v2", "erdiagram", "journey", "pie", "gantt", "mindmap",
+)
+_MERMAID_FENCE = re.compile(r"```mermaid[ \t]*\n(.*?)```", re.DOTALL | re.IGNORECASE)
+
+# Per-file read cap for import scanning: imports live at the top of files.
+_SCAN_BYTES = 65536
+
+
+def _max_nodes() -> int:
+    try:
+        return max(2, int(os.getenv("MAX_DIAGRAM_NODES", "40")))
+    except ValueError:
+        return 40
+
+
+def _label(text: str) -> str:
+    """A node label safe inside Mermaid's ["..."] syntax."""
+    return text.replace('"', "'").replace("\n", " ")
+
+
+def _read_head(path: str) -> str:
+    try:
+        with open(path, "r", encoding="utf-8", errors="ignore") as fh:
+            return fh.read(_SCAN_BYTES)
+    except OSError:
+        return ""
+
+
+def _py_module(file_path: str) -> str:
+    mod = file_path[:-3].replace("\\", "/").replace("/", ".")
+    return mod[: -len(".__init__")] if mod.endswith(".__init__") else mod
+
+
+def _resolve_py_import(name: str, importer: str, module_map: dict[str, str]) -> str | None:
+    """Map an import name (possibly relative) to an internal file, or None."""
+    if name.startswith("."):
+        level = len(name) - len(name.lstrip("."))
+        pkg_parts = os.path.dirname(importer).replace("\\", "/").split("/")
+        pkg_parts = [p for p in pkg_parts if p]
+        if level > 1:
+            pkg_parts = pkg_parts[: len(pkg_parts) - (level - 1)]
+        remainder = name.lstrip(".")
+        name = ".".join(pkg_parts + ([remainder] if remainder else []))
+    # Longest-prefix match: "a.b.c" may refer to module a/b.py's attribute c.
+    parts = name.split(".")
+    for end in range(len(parts), 0, -1):
+        candidate = ".".join(parts[:end])
+        if candidate in module_map:
+            return module_map[candidate]
+    return None
+
+
+def _resolve_js_import(spec: str, importer: str, file_set: set[str]) -> str | None:
+    if not spec.startswith("."):
+        return None  # external package
+    base = os.path.normpath(os.path.join(os.path.dirname(importer), spec))
+    base = base.replace("\\", "/")
+    for suffix in _JS_RESOLVE_SUFFIXES:
+        candidate = base + suffix
+        if candidate in file_set:
+            return candidate
+    return None
+
+
+def _collect_edges(files: list[str], repo_path: str) -> set[tuple[str, str]]:
+    file_set = {f.replace("\\", "/") for f in files}
+    module_map = {_py_module(f): f for f in file_set if f.endswith(".py")}
+    edges: set[tuple[str, str]] = set()
+    for f in sorted(file_set):
+        content = _read_head(os.path.join(repo_path, f))
+        if not content:
+            continue
+        if f.endswith(".py"):
+            for match in _PY_IMPORT.finditer(content):
+                name = match.group(1) or match.group(2)
+                target = _resolve_py_import(name, f, module_map)
+                if target and target != f:
+                    edges.add((f, target))
+        elif f.endswith(_JS_EXTENSIONS):
+            for match in _JS_IMPORT.finditer(content):
+                target = _resolve_js_import(match.group(1), f, file_set)
+                if target and target != f:
+                    edges.add((f, target))
+    return edges
+
+
+def build_dependency_mermaid(files: list[str], repo_path: str) -> str | None:
+    """Internal-module dependency graph as a fenced Mermaid block, or None
+    when there is nothing meaningful to draw."""
+    edges = _collect_edges(files, repo_path)
+    if len(edges) < 2:
+        return None
+
+    cap = _max_nodes()
+    # Keep the most-connected nodes so the capped diagram stays informative.
+    degree: dict[str, int] = {}
+    for a, b in edges:
+        degree[a] = degree.get(a, 0) + 1
+        degree[b] = degree.get(b, 0) + 1
+    kept = set(sorted(degree, key=lambda n: (-degree[n], n))[:cap])
+    kept_edges = sorted((a, b) for a, b in edges if a in kept and b in kept)
+    if len(kept_edges) < 2:
+        return None
+    dropped = len(degree) - len(kept)
+
+    ids = {n: f"n{i}" for i, n in enumerate(sorted(kept))}
+    lines = ["```mermaid", "graph LR"]
+    lines += [f'    {ids[n]}["{_label(n)}"]' for n in sorted(kept)]
+    lines += [f"    {ids[a]} --> {ids[b]}" for a, b in kept_edges]
+    lines.append("```")
+    if dropped > 0:
+        lines.append(f"\n*…plus {dropped} more module(s) not shown (MAX_DIAGRAM_NODES).*")
+    return "\n".join(lines)
+
+
+def build_structure_mermaid(files: list[str]) -> str | None:
+    """Directory-tree chart as a fenced Mermaid block, or None for <2 files."""
+    normalized = sorted({f.replace("\\", "/") for f in files})
+    if len(normalized) < 2:
+        return None
+
+    cap = _max_nodes()
+    ids: dict[str, str] = {".": "n0"}
+    labels: dict[str, str] = {".": "(root)"}
+    edges: list[tuple[str, str]] = []
+    truncated = False
+
+    def _ensure(node: str, label: str) -> bool:
+        nonlocal truncated
+        if node in ids:
+            return True
+        if len(ids) >= cap:
+            truncated = True
+            return False
+        ids[node] = f"n{len(ids)}"
+        labels[node] = label
+        return True
+
+    for f in normalized:
+        parts = f.split("/")
+        parent = "."
+        for depth, part in enumerate(parts):
+            node = "/".join(parts[: depth + 1])
+            if not _ensure(node, part):
+                break
+            if (parent, node) not in edges:
+                edges.append((parent, node))
+            parent = node
+
+    lines = ["```mermaid", "graph TD"]
+    lines += [f'    {ids[n]}["{_label(labels[n])}"]' for n in ids]
+    lines += [f"    {ids[a]} --> {ids[b]}" for a, b in edges if a in ids and b in ids]
+    lines.append("```")
+    if truncated:
+        lines.append("\n*…tree truncated (MAX_DIAGRAM_NODES).*")
+    return "\n".join(lines)
+
+
+def _is_plausible_mermaid(body: str) -> bool:
+    stripped = body.strip()
+    if not stripped:
+        return False
+    first_token = stripped.split()[0].lower()
+    if first_token not in _MERMAID_TYPES:
+        return False
+    for open_ch, close_ch in ("()", "[]", "{}"):
+        if stripped.count(open_ch) != stripped.count(close_ch):
+            return False
+    return True
+
+
+def sanitize_mermaid_blocks(markdown: str) -> str:
+    """Downgrade implausible ```mermaid fences to plain ```text fences.
+
+    Model-drawn diagrams are useful but not trusted: an invalid block would
+    render as a raw error box on GitHub. Content is preserved either way.
+    """
+
+    def _check(match: re.Match) -> str:
+        body = match.group(1)
+        if _is_plausible_mermaid(body):
+            return match.group(0)
+        return f"```text\n{body}```"
+
+    return _MERMAID_FENCE.sub(_check, markdown)

ai_doc_creator/core/doc_writer.py

@@ -0,0 +1,43 @@
+import os
+import logging
+from typing import Dict
+
+logger = logging.getLogger(__name__)
+
+
+class DocumentationWriter:
+    """Writes generated documentation to files."""
+
+    def __init__(self, output_dir: str):
+        self.output_dir = output_dir
+
+    def write_docs(self, documents: Dict[str, str], index_content: str):
+        """Writes the documentation files and the index file."""
+        try:
+            # Create output directory if it doesn't exist
+            os.makedirs(self.output_dir, exist_ok=True)
+
+            # Write individual file docs
+            for file_path, content in documents.items():
+                # We want to mirror the structure inside the output dir
+                # e.g., src/main.py -> docs/src/main.py.md
+
+                relative_path = file_path + ".md"
+                full_path = os.path.join(self.output_dir, relative_path)
+
+                # Ensure the directory for this file exists
+                os.makedirs(os.path.dirname(full_path), exist_ok=True)
+
+                with open(full_path, "w", encoding="utf-8") as f:
+                    f.write(content)
+                logger.info(f"Written documentation for {file_path} to {full_path}")
+
+            # Write index file
+            index_path = os.path.join(self.output_dir, "README.md")
+            with open(index_path, "w", encoding="utf-8") as f:
+                f.write(index_content)
+            logger.info(f"Written index to {index_path}")
+
+        except Exception as e:
+            logger.error(f"Failed to write documentation: {e}")
+            raise