ai-pipeline-core 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

ai_pipeline_core/docs_generator/__init__.py

@@ -0,0 +1,54 @@
+"""AI-focused documentation generator.
+
+Generates dense, self-contained guides from source code and test suite
+for AI coding agents. Uses AST parsing, dependency resolution, and
+size management for guides with a 50KB warning threshold.
+"""
+
+from ai_pipeline_core.docs_generator.extractor import (
+    ClassInfo,
+    FunctionInfo,
+    MethodInfo,
+    ModuleInfo,
+    SymbolTable,
+    is_public_name,
+    parse_module,
+)
+from ai_pipeline_core.docs_generator.guide_builder import (
+    GuideData,
+    TestExample,
+    build_guide,
+    discover_tests,
+    select_examples,
+)
+from ai_pipeline_core.docs_generator.trimmer import manage_guide_size
+from ai_pipeline_core.docs_generator.validator import (
+    ValidationResult,
+    compute_source_hash,
+    validate_all,
+    validate_completeness,
+    validate_freshness,
+    validate_size,
+)
+
+__all__ = [
+    "ClassInfo",
+    "FunctionInfo",
+    "GuideData",
+    "MethodInfo",
+    "ModuleInfo",
+    "SymbolTable",
+    "TestExample",
+    "ValidationResult",
+    "build_guide",
+    "compute_source_hash",
+    "discover_tests",
+    "is_public_name",
+    "manage_guide_size",
+    "parse_module",
+    "select_examples",
+    "validate_all",
+    "validate_completeness",
+    "validate_freshness",
+    "validate_size",
+]

ai_pipeline_core/docs_generator/__main__.py

@@ -0,0 +1,5 @@
+"""Entry point for python -m ai_pipeline_core.docs_generator."""
+
+from ai_pipeline_core.docs_generator.cli import main
+
+raise SystemExit(main())

ai_pipeline_core/docs_generator/cli.py

@@ -0,0 +1,196 @@
+"""CLI for AI documentation generation and validation."""
+
+import argparse
+import sys
+from pathlib import Path
+
+from ai_pipeline_core.docs_generator.extractor import build_symbol_table
+from ai_pipeline_core.docs_generator.guide_builder import build_guide, render_guide
+from ai_pipeline_core.docs_generator.trimmer import manage_guide_size
+from ai_pipeline_core.docs_generator.validator import (
+    HASH_FILE,
+    compute_source_hash,
+    validate_all,
+)
+
+EXCLUDED_MODULES: frozenset[str] = frozenset({"docs_generator"})
+
+
+def _normalize_whitespace(content: str) -> str:
+    """Strip trailing whitespace from each line and ensure final newline."""
+    lines = [line.rstrip() for line in content.splitlines()]
+    return "\n".join(lines) + "\n"
+
+
+TEST_DIR_OVERRIDES: dict[str, str] = {}  # nosemgrep: no-mutable-module-globals
+
+
+def _discover_modules(source_dir: Path) -> list[str]:
+    """Discover all public module groupings from package structure."""
+    modules: set[str] = set()
+    for py_file in sorted(source_dir.rglob("*.py")):
+        if py_file.name.startswith("_") and py_file.name != "__init__.py":
+            continue
+        relative = py_file.relative_to(source_dir)
+        if len(relative.parts) > 1:
+            modules.add(relative.parts[0])
+        else:
+            modules.add(relative.stem)
+    return sorted(modules - EXCLUDED_MODULES)
+
+
+def main(argv: list[str] | None = None) -> int:
+    """Entry point for AI docs CLI with generate/check subcommands."""
+    parser = argparse.ArgumentParser(description="AI documentation generator")
+    parser.add_argument("--source-dir", type=Path, help="Source package directory")
+    parser.add_argument("--tests-dir", type=Path, help="Tests directory")
+    parser.add_argument("--output-dir", type=Path, help="Output .ai-docs directory")
+    subparsers = parser.add_subparsers(dest="command")
+    subparsers.add_parser("generate", help="Generate .ai-docs/ documentation")
+    subparsers.add_parser("check", help="Validate .ai-docs/ is up-to-date")
+
+    args = parser.parse_args(argv)
+    if not args.command:
+        parser.print_help()
+        return 1
+
+    source_dir, tests_dir, output_dir, repo_root = _resolve_paths(args)
+
+    if args.command == "generate":
+        return _run_generate(source_dir, tests_dir, output_dir, repo_root)
+    return _run_check(source_dir, tests_dir, output_dir)
+
+
+def _resolve_paths(args: argparse.Namespace) -> tuple[Path, Path, Path, Path]:
+    """Resolve source, tests, output directories and repo root from args or auto-detect."""
+    cli_file = Path(__file__).resolve()
+    repo_root = cli_file.parent.parent.parent
+    source_dir = args.source_dir or (repo_root / "ai_pipeline_core")
+    tests_dir = args.tests_dir or (repo_root / "tests")
+    output_dir = args.output_dir or (repo_root / ".ai-docs")
+    return source_dir, tests_dir, output_dir, repo_root
+
+
+def _run_generate(source_dir: Path, tests_dir: Path, output_dir: Path, repo_root: Path) -> int:
+    """Generate all module guides, INDEX.md, and .hash file."""
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Clean stale files
+    for existing in output_dir.glob("*.md"):
+        existing.unlink()
+    hash_file = output_dir / HASH_FILE
+    if hash_file.exists():
+        hash_file.unlink()
+
+    table = build_symbol_table(source_dir)
+    generated: list[tuple[str, int]] = []
+
+    for module_name in _discover_modules(source_dir):
+        data = build_guide(module_name, source_dir, tests_dir, table, TEST_DIR_OVERRIDES, repo_root)
+        if not data.classes and not data.functions:
+            print(f"  skip {module_name} (no public symbols)")
+            continue
+
+        content = render_guide(data)
+        content = manage_guide_size(data, content)
+        content = _normalize_whitespace(content)
+
+        guide_path = output_dir / f"{module_name}.md"
+        guide_path.write_text(content)
+        size = len(content.encode("utf-8"))
+        generated.append((module_name, size))
+        print(f"  wrote {module_name}.md ({size:,} bytes)")
+
+    # INDEX.md
+    index_content = _normalize_whitespace(_render_index(generated))
+    (output_dir / "INDEX.md").write_text(index_content)
+    print(f"  wrote INDEX.md ({len(index_content):,} bytes)")
+
+    # .hash
+    source_hash = compute_source_hash(source_dir, tests_dir)
+    (output_dir / HASH_FILE).write_text(source_hash + "\n")
+    print(f"  wrote {HASH_FILE}")
+
+    total = sum(size for _, size in generated)
+    print(f"\nGenerated {len(generated)} guides ({total:,} bytes total)")
+    return 0
+
+
+def _run_check(source_dir: Path, tests_dir: Path, output_dir: Path) -> int:
+    """Validate .ai-docs/ freshness, completeness, and size."""
+    if not output_dir.is_dir():
+        print("FAIL: .ai-docs/ directory does not exist. Run 'generate' first.", file=sys.stderr)
+        return 1
+
+    result = validate_all(output_dir, source_dir, tests_dir, excluded_modules=EXCLUDED_MODULES)
+
+    if not result.is_fresh:
+        print("FAIL: .ai-docs/ is stale (source hash mismatch)")
+    if result.missing_symbols:
+        print(f"FAIL: {len(result.missing_symbols)} public symbols missing from guides:")
+        for sym in result.missing_symbols:
+            print(f"  - {sym}")
+    if result.size_violations:
+        print(f"WARNING: {len(result.size_violations)} guides exceed size limit:")
+        for name, size in result.size_violations:
+            print(f"  - {name}: {size:,} bytes")
+
+    if result.is_valid:
+        print("OK: .ai-docs/ is up-to-date")
+        return 0
+    return 1
+
+
+def _render_index(generated: list[tuple[str, int]]) -> str:
+    """Render INDEX.md with reading order, task lookup, imports, and size table."""
+    lines: list[str] = [
+        "# AI Documentation Index",
+        "",
+        "Auto-generated guide index. Do not edit manually.",
+        "",
+        "## Reading Order",
+        "",
+    ]
+    for i, (name, _) in enumerate(generated, 1):
+        lines.append(f"{i}. [{name}]({name}.md)")
+
+    lines.extend([
+        "",
+        "## Task-Based Lookup",
+        "",
+        "| Task | Guide |",
+        "| ---- | ----- |",
+    ])
+    task_map = {
+        "Create/read documents": "documents",
+        "Store/retrieve documents": "document_store",
+        "Call LLMs": "llm",
+        "Deploy pipelines": "deployment",
+        "Load templates": "prompt_manager",
+        "Process images": "images",
+        "Define flows/tasks": "pipeline",
+        "Configure settings": "settings",
+        "Handle errors": "exceptions",
+        "Log messages": "logging",
+        "Debug & observe traces": "observability",
+        "Test pipelines": "testing",
+    }
+    guide_set = {name for name, _ in generated}
+    for task, guide in task_map.items():
+        if guide in guide_set:
+            lines.append(f"| {task} | [{guide}]({guide}.md) |")
+
+    lines.extend([
+        "",
+        "## Module Sizes",
+        "",
+        "| Module | Size |",
+        "| ------ | ---- |",
+    ])
+    for name, size in generated:
+        lines.append(f"| {name} | {size:,} bytes |")
+    total = sum(size for _, size in generated)
+    lines.append(f"| **Total** | **{total:,} bytes** |")
+    lines.append("")
+
+    return "\n".join(lines)

ai_pipeline_core/docs_generator/extractor.py

@@ -0,0 +1,324 @@
+"""AST-based symbol extraction from Python source files.
+
+Extracts class/function signatures, inheritance chains,
+and builds a symbol table for dependency resolution.
+"""
+
+import ast
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass(frozen=True)
+class MethodInfo:
+    """Extracted method/property metadata from a class body."""
+
+    name: str
+    signature: str
+    docstring: str
+    source: str
+    is_property: bool
+    is_classmethod: bool
+    is_abstract: bool
+    line_count: int
+    is_inherited: bool = False
+    inherited_from: str | None = None
+
+
+@dataclass(frozen=True)
+class ClassInfo:
+    """Extracted class metadata including methods, validators, and class variables."""
+
+    name: str
+    bases: tuple[str, ...]
+    docstring: str
+    is_public: bool
+    class_vars: tuple[tuple[str, str, str], ...]  # (name, type_annotation, default_value)
+    methods: tuple[MethodInfo, ...]
+    validators: tuple[MethodInfo, ...]
+    module_path: str
+    decorators: tuple[str, ...] = ()
+
+
+@dataclass(frozen=True)
+class FunctionInfo:
+    """Extracted module-level function metadata."""
+
+    name: str
+    signature: str
+    docstring: str
+    source: str
+    is_public: bool
+    is_async: bool
+    line_count: int
+    module_path: str
+
+
+@dataclass(frozen=True)
+class ModuleInfo:
+    """Parsed module containing its classes and functions."""
+
+    name: str
+    path: Path
+    docstring: str
+    is_public: bool
+    classes: tuple[ClassInfo, ...]
+    functions: tuple[FunctionInfo, ...]
+
+
+@dataclass
+class SymbolTable:
+    """Mutable during construction, used read-only after building.
+
+    Maps class and function names to their ClassInfo/FunctionInfo objects,
+    and provides class_to_module/function_to_module lookups for dependency resolution.
+    """
+
+    classes: dict[str, ClassInfo] = field(default_factory=dict)
+    functions: dict[str, FunctionInfo] = field(default_factory=dict)
+    class_to_module: dict[str, str] = field(default_factory=dict)
+    function_to_module: dict[str, str] = field(default_factory=dict)
+
+
+# Known external base classes that get stub representations
+EXTERNAL_STUBS: dict[str, str] = {
+    "BaseModel": "Pydantic base model. Fields are typed class attributes.",
+    "BaseSettings": "Pydantic settings model. Loads values from environment variables.",
+    "ABC": "Python abstract base class marker.",
+    "Generic": "Python generic base class for parameterized types.",
+    "list": "Python built-in list.",
+    "dict": "Python built-in dictionary.",
+    "StrEnum": "String enumeration base class.",
+}
+
+
+def is_public_name(name: str) -> bool:
+    """Determine if a symbol is public based on Python naming convention."""
+    if name.startswith("__") and name.endswith("__"):
+        return True
+    return not name.startswith("_")
+
+
+def parse_module(path: Path) -> ModuleInfo:
+    """Parse a single .py file and return all extracted symbols."""
+    source = path.read_text(encoding="utf-8")
+    source_lines = source.splitlines()
+    tree = ast.parse(source)
+
+    module_doc = ast.get_docstring(tree) or ""
+    module_path = _module_path(path)
+
+    classes: list[ClassInfo] = []
+    functions: list[FunctionInfo] = []
+
+    for node in tree.body:
+        if isinstance(node, ast.ClassDef):
+            classes.append(_extract_class(node, source_lines, module_path))
+        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            functions.append(_extract_function(node, source_lines, module_path))
+
+    module_public = any(c.is_public for c in classes) or any(f.is_public for f in functions)
+
+    return ModuleInfo(
+        name=path.stem,
+        path=path,
+        docstring=module_doc,
+        is_public=module_public,
+        classes=tuple(classes),
+        functions=tuple(functions),
+    )
+
+
+def build_symbol_table(source_dir: Path) -> SymbolTable:
+    """Parse all .py files under source_dir and build a unified symbol table."""
+    table = SymbolTable()
+
+    for py_file in sorted(source_dir.rglob("*.py")):
+        if py_file.name.startswith("_") and py_file.name != "__init__.py":
+            continue
+        module = parse_module(py_file)
+
+        relative = py_file.relative_to(source_dir)
+        if len(relative.parts) > 1:
+            package_name = relative.parts[0]
+        else:
+            package_name = relative.stem
+
+        for cls in module.classes:
+            table.classes[cls.name] = cls
+            table.class_to_module[cls.name] = package_name
+        for func in module.functions:
+            table.functions[func.name] = func
+            table.function_to_module[func.name] = package_name
+
+    return table
+
+
+def resolve_dependencies(
+    root_classes: list[str],
+    table: SymbolTable,
+) -> tuple[list[ClassInfo], set[str]]:
+    """Resolve transitive dependencies for a set of root classes.
+
+    Returns (resolved ClassInfo list in topological order, external base names).
+    """
+    resolved: list[ClassInfo] = []
+    external_bases: set[str] = set()
+    visited: set[str] = set()
+
+    def visit(name: str) -> None:
+        if name in visited:
+            return
+        visited.add(name)
+
+        if name in EXTERNAL_STUBS or name not in table.classes:
+            external_bases.add(name)
+            return
+
+        cls = table.classes[name]
+        for base in cls.bases:
+            visit(base.split("[")[0])
+
+        resolved.append(cls)
+
+    for root in root_classes:
+        visit(root)
+
+    return resolved, external_bases
+
+
+# ---------------------------------------------------------------------------
+# Private helpers
+# ---------------------------------------------------------------------------
+
+
+def _module_path(path: Path) -> str:
+    """Convert filesystem path to dotted module path.
+
+    e.g. ai_pipeline_core/documents/document.py -> ai_pipeline_core.documents.document
+    """
+    parts = list(path.with_suffix("").parts)
+    # Find the package root (ai_pipeline_core)
+    for i, part in enumerate(parts):
+        if part == "ai_pipeline_core":
+            return ".".join(parts[i:])
+    return ".".join(parts)
+
+
+def _decorator_name(decorator: ast.expr) -> str:
+    if isinstance(decorator, ast.Call):
+        return _decorator_name(decorator.func)
+    if isinstance(decorator, ast.Attribute):
+        return decorator.attr
+    if isinstance(decorator, ast.Name):
+        return decorator.id
+    return ""
+
+
+def _body_line_count(node: ast.FunctionDef | ast.AsyncFunctionDef) -> int:
+    if not node.body:
+        return 0
+    first = node.body[0]
+    body_nodes = node.body
+    is_docstring = isinstance(first, ast.Expr) and isinstance(first.value, ast.Constant) and isinstance(first.value.value, str)
+    if is_docstring:
+        body_nodes = node.body[1:]
+    if not body_nodes:
+        return 0
+    start = body_nodes[0].lineno
+    end = body_nodes[-1].end_lineno or body_nodes[-1].lineno
+    return end - start + 1
+
+
+def _extract_signature(node: ast.FunctionDef | ast.AsyncFunctionDef) -> str:
+    args_str = ast.unparse(node.args)
+    ret = f" -> {ast.unparse(node.returns)}" if node.returns else ""
+    return f"({args_str}){ret}"
+
+
+def _get_source(source_lines: list[str], node: ast.AST) -> str:
+    decoratable = (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)
+    if isinstance(node, decoratable) and node.decorator_list:
+        start = node.decorator_list[0].lineno - 1
+    else:
+        start: int = node.lineno - 1  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType]
+    end: int = node.end_lineno or node.lineno  # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType, reportUnknownVariableType]
+    return "\n".join(source_lines[start:end])
+
+
+def _is_validator(node: ast.FunctionDef | ast.AsyncFunctionDef) -> bool:
+    validator_names = ("field_validator", "model_validator")
+    return any(_decorator_name(d) in validator_names for d in node.decorator_list)
+
+
+def _extract_method(
+    node: ast.FunctionDef | ast.AsyncFunctionDef,
+    source_lines: list[str],
+) -> MethodInfo:
+    decorator_names = {_decorator_name(d) for d in node.decorator_list}
+    return MethodInfo(
+        name=node.name,
+        signature=_extract_signature(node),
+        docstring=ast.get_docstring(node) or "",
+        source=_get_source(source_lines, node),
+        is_property="property" in decorator_names,
+        is_classmethod="classmethod" in decorator_names,
+        is_abstract="abstractmethod" in decorator_names,
+        line_count=_body_line_count(node),
+    )
+
+
+def _extract_class(node: ast.ClassDef, source_lines: list[str], module_path: str) -> ClassInfo:
+    docstring = ast.get_docstring(node) or ""
+    bases = [ast.unparse(base) for base in node.bases]
+
+    methods: list[MethodInfo] = []
+    validators: list[MethodInfo] = []
+    class_vars: list[tuple[str, str, str]] = []
+
+    for item in node.body:
+        if isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            method = _extract_method(item, source_lines)
+            methods.append(method)
+            if _is_validator(item):
+                validators.append(method)
+        elif isinstance(item, ast.AnnAssign) and item.target and isinstance(item.target, ast.Name):
+            name = item.target.id
+            if is_public_name(name):
+                type_ann = ast.unparse(item.annotation) if item.annotation else ""
+                default = ast.unparse(item.value) if item.value else ""
+                class_vars.append((name, type_ann, default))
+        elif isinstance(item, ast.Assign) and len(item.targets) == 1 and isinstance(item.targets[0], ast.Name):
+            name = item.targets[0].id
+            if is_public_name(name):
+                default = ast.unparse(item.value)
+                class_vars.append((name, "", default))
+
+    return ClassInfo(
+        name=node.name,
+        bases=tuple(bases),
+        docstring=docstring,
+        is_public=is_public_name(node.name),
+        class_vars=tuple(class_vars),
+        methods=tuple(methods),
+        validators=tuple(validators),
+        module_path=module_path,
+        decorators=tuple(ast.unparse(d) for d in node.decorator_list),
+    )
+
+
+def _extract_function(
+    node: ast.FunctionDef | ast.AsyncFunctionDef,
+    source_lines: list[str],
+    module_path: str,
+) -> FunctionInfo:
+    return FunctionInfo(
+        name=node.name,
+        signature=_extract_signature(node),
+        docstring=ast.get_docstring(node) or "",
+        source=_get_source(source_lines, node),
+        is_public=is_public_name(node.name),
+        is_async=isinstance(node, ast.AsyncFunctionDef),
+        line_count=_body_line_count(node),
+        module_path=module_path,
+    )