docwright 0.1.0__py3-none-any.whl

ai_docgen/engine.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+
+from ai_docgen.analyzer import DiffAnalyzer
+from ai_docgen.config import Config
+from ai_docgen.outputs.base import Output
+from ai_docgen.providers.base import LLMProvider
+from ai_docgen.registry import DocumentEntry, ProjectEntry, Registry
+from ai_docgen.renderer import DocumentRenderer, TemplateLoader
+
+SYSTEM_PROMPT = (
+    "You are a technical documentation writer. You update specific sections of documentation "
+    "based on code changes. Return ONLY the updated section content — no markdown fences, "
+    "no explanations, no surrounding text. Write in clear, concise English. "
+    "Be accurate and specific to the actual code."
+)
+
+
+class DocsEngine:
+    def __init__(self, repo_root: Path, provider: LLMProvider, output: Output) -> None:
+        self.repo_root = repo_root
+        self.provider = provider
+        self.output = output
+        self.config = Config.load(repo_root)
+        self.renderer = DocumentRenderer()
+        self.loader = TemplateLoader(
+            source=self.config.templates.source,
+            local_path=(
+                repo_root / self.config.templates.local_path
+                if self.config.templates.source == "local"
+                else None
+            ),
+        )
+
+    async def init(self) -> None:
+        changed_files: list[Path] = []
+        for doc_config in self.config.documents:
+            target = self.repo_root / doc_config.target
+            template_text = self.loader.load(doc_config.template)
+            document = target.read_text() if target.exists() else template_text
+            section_names = self.renderer.auto_section_names(template_text)
+            repo_context = self.gather_repo_context()
+            for section_name in section_names:
+                user_prompt = (
+                    f"Repository context:\n{repo_context}\n\n"
+                    f"Document type: {doc_config.type}\n"
+                    f"Update the '{section_name}' section with accurate, detailed information."
+                )
+                updated_content = await self.provider.complete(
+                    system=SYSTEM_PROMPT, user=user_prompt
+                )
+                document = self.renderer.patch_section(
+                    document, section_name, updated_content + "\n"
+                )
+            target.parent.mkdir(parents=True, exist_ok=True)
+            target.write_text(document)
+            changed_files.append(target)
+        if changed_files:
+            self.output.apply(changed_files, "docs: generate initial documentation")
+        Config.mark_initialized(self.repo_root)
+        self.register_in_registry()
+
+    async def run(self, diff_text: str) -> bool:
+        if not Config.is_initialized(self.repo_root):
+            await self.init()
+            return False
+        triggers = self.config.triggers
+        analyzer = DiffAnalyzer(
+            diff_text=diff_text,
+            trigger_paths=triggers.paths if triggers else [],
+            ignore_paths=triggers.ignore if triggers else [],
+        )
+        if not analyzer.has_relevant_changes():
+            return True
+        changed_files: list[Path] = []
+        diff_summary = analyzer.diff_summary()
+        for doc_config in self.config.documents:
+            target = self.repo_root / doc_config.target
+            if not target.exists():
+                continue
+            document = target.read_text()
+            section_names = self.renderer.auto_section_names(document)
+            updated = False
+            for section_name in section_names:
+                user_prompt = (
+                    f"Diff summary:\n{diff_summary}\n\n"
+                    f"Current document:\n{document}\n\n"
+                    f"Update the '{section_name}' section if the diff affects it. "
+                    f"If no update is needed, return the current section content unchanged."
+                )
+                updated_content = await self.provider.complete(
+                    system=SYSTEM_PROMPT, user=user_prompt
+                )
+                new_document = self.renderer.patch_section(
+                    document, section_name, updated_content + "\n"
+                )
+                if new_document != document:
+                    document = new_document
+                    updated = True
+            if updated:
+                target.write_text(document)
+                changed_files.append(target)
+        if changed_files:
+            self.output.apply(changed_files, "docs: update documentation")
+        return False
+
+    async def sync(self) -> None:
+        for doc_config in self.config.documents:
+            target = self.repo_root / doc_config.target
+            template_text = self.loader.load(doc_config.template)
+            document = target.read_text() if target.exists() else template_text
+            section_names = self.renderer.auto_section_names(template_text)
+            repo_context = self.gather_repo_context()
+            for section_name in section_names:
+                user_prompt = (
+                    f"Repository context:\n{repo_context}\n\n"
+                    f"Current document:\n{document}\n\n"
+                    f"Re-sync the '{section_name}' section to be accurate and up to date."
+                )
+                updated_content = await self.provider.complete(
+                    system=SYSTEM_PROMPT, user=user_prompt
+                )
+                document = self.renderer.patch_section(
+                    document, section_name, updated_content + "\n"
+                )
+            target.parent.mkdir(parents=True, exist_ok=True)
+            target.write_text(document)
+
+    def gather_repo_context(self) -> str:
+        lines: list[str] = [f"Repo root: {self.repo_root.name}"]
+        for candidate in ["pyproject.toml", "package.json", "composer.json", "go.mod"]:
+            path = self.repo_root / candidate
+            if path.exists():
+                lines.append(f"\n{candidate}:\n{path.read_text()[:2000]}")
+                break
+        src_dirs = ["app", "src", "lib"]
+        for src_dir in src_dirs:
+            full = self.repo_root / src_dir
+            if full.exists():
+                files = [str(p.relative_to(self.repo_root)) for p in full.rglob("*.py")][:20]
+                lines.append(f"\nSource files in {src_dir}/: {', '.join(files)}")
+                break
+        return "\n".join(lines)
+
+    def register_in_registry(self) -> None:
+        registry_path = self.repo_root / self.config.registry.path
+        registry = Registry(registry_path)
+        try:
+            remote = (
+                subprocess.check_output(["git", "remote", "get-url", "origin"], cwd=self.repo_root)
+                .decode()
+                .strip()
+            )
+        except subprocess.CalledProcessError:
+            remote = ""
+        registry.register(
+            ProjectEntry(
+                name=self.repo_root.name,
+                path=str(self.repo_root),
+                remote=remote,
+                documents=[DocumentEntry(target=d.target) for d in self.config.documents],
+            )
+        )

ai_docgen/outputs/base.py

@@ -0,0 +1,7 @@
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+
+class Output(ABC):
+    @abstractmethod
+    def apply(self, changed_files: list[Path], message: str) -> None: ...

ai_docgen/outputs/direct.py

@@ -0,0 +1,16 @@
+from pathlib import Path
+
+from git import Repo
+
+from ai_docgen.outputs.base import Output
+
+
+class DirectOutput(Output):
+    def __init__(self, repo_root: Path) -> None:
+        self.repo_root = repo_root
+
+    def apply(self, changed_files: list[Path], message: str) -> None:
+        repo = Repo(self.repo_root)
+        str_paths = [str(f.relative_to(self.repo_root)) for f in changed_files]
+        repo.index.add(str_paths)
+        repo.index.commit(message)

ai_docgen/outputs/factory.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+
+from ai_docgen.config import OutputConfig
+from ai_docgen.outputs.base import Output
+from ai_docgen.outputs.direct import DirectOutput
+from ai_docgen.outputs.pull_request import PullRequestOutput
+
+
+def build_output(config: OutputConfig, repo_root: Path) -> Output:
+    if config.mode == "direct":
+        return DirectOutput(repo_root=repo_root)
+    return PullRequestOutput(repo_root=repo_root, config=config)

ai_docgen/outputs/pull_request.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import hashlib
+import os
+import re
+import subprocess
+from pathlib import Path
+
+import httpx
+
+from ai_docgen.config import OutputConfig
+from ai_docgen.outputs.base import Output
+
+
+class PullRequestOutput(Output):
+    def __init__(self, repo_root: Path, config: OutputConfig) -> None:
+        self.repo_root = repo_root
+        self.config = config
+
+    def apply(self, changed_files: list[Path], message: str) -> None:
+        short_hash = hashlib.sha1(message.encode()).hexdigest()[:8]
+        branch = f"{self.config.branch_prefix}{short_hash}"
+        self.run(["git", "checkout", "-b", branch])
+        rel_paths = [str(f.relative_to(self.repo_root)) for f in changed_files]
+        self.run(["git", "add"] + rel_paths)
+        self.run(["git", "commit", "-m", message])
+        self.run(["git", "push", "origin", branch])
+        token = os.environ.get("GITHUB_TOKEN", "")
+        if token:
+            self.create_github_pr(branch, message, token)
+
+    def run(self, cmd: list[str]) -> None:
+        subprocess.run(cmd, cwd=self.repo_root, check=True)
+
+    def create_github_pr(self, branch: str, title: str, token: str) -> None:
+        remote_url = (
+            subprocess.check_output(["git", "remote", "get-url", "origin"], cwd=self.repo_root)
+            .decode()
+            .strip()
+        )
+        owner_repo = self.extract_owner_repo(remote_url)
+        if not owner_repo:
+            return
+        default_branch = (
+            subprocess.check_output(
+                ["git", "symbolic-ref", "refs/remotes/origin/HEAD"],
+                cwd=self.repo_root,
+            )
+            .decode()
+            .strip()
+            .replace("refs/remotes/origin/", "")
+        )
+        httpx.post(
+            f"https://api.github.com/repos/{owner_repo}/pulls",
+            headers={
+                "Authorization": f"Bearer {token}",
+                "Accept": "application/vnd.github+json",
+            },
+            json={
+                "title": self.config.pr_title,
+                "head": branch,
+                "base": default_branch,
+                "body": title,
+            },
+        )
+
+    def extract_owner_repo(self, remote_url: str) -> str | None:
+        match = re.search(r"[:/]([^/]+/[^/]+?)(?:\.git)?$", remote_url)
+        return match.group(1) if match else None

ai_docgen/providers/base.py

@@ -0,0 +1,6 @@
+from abc import ABC, abstractmethod
+
+
+class LLMProvider(ABC):
+    @abstractmethod
+    async def complete(self, system: str, user: str) -> str: ...

ai_docgen/providers/claude.py

@@ -0,0 +1,22 @@
+from anthropic import AsyncAnthropic
+from anthropic.types import TextBlock
+
+from ai_docgen.providers.base import LLMProvider
+
+
+class ClaudeProvider(LLMProvider):
+    def __init__(self, model: str, api_key: str) -> None:
+        self.model = model
+        self.client = AsyncAnthropic(api_key=api_key)
+
+    async def complete(self, system: str, user: str) -> str:
+        message = await self.client.messages.create(
+            model=self.model,
+            max_tokens=4096,
+            system=system,
+            messages=[{"role": "user", "content": user}],
+        )
+        block = message.content[0]
+        if not isinstance(block, TextBlock):
+            raise ValueError(f"Unexpected content block type: {type(block)}")
+        return block.text

ai_docgen/providers/factory.py

@@ -0,0 +1,26 @@
+import os
+
+from ai_docgen.config import ProviderConfig
+from ai_docgen.providers.base import LLMProvider
+from ai_docgen.providers.claude import ClaudeProvider
+from ai_docgen.providers.ollama import OllamaProvider
+from ai_docgen.providers.openai import OpenAIProvider
+
+
+def build_provider(config: ProviderConfig) -> LLMProvider:
+    if config.type == "ollama":
+        return OllamaProvider(
+            model=config.model,
+            base_url=config.base_url or "http://localhost:11434",
+        )
+    api_key = os.environ.get(config.api_key_env, "")
+    if not api_key:
+        raise OSError(
+            f"Environment variable '{config.api_key_env}' is not set. "
+            f"Required for provider '{config.type}'."
+        )
+    if config.type == "claude":
+        return ClaudeProvider(model=config.model, api_key=api_key)
+    if config.type == "openai":
+        return OpenAIProvider(model=config.model, api_key=api_key)
+    raise ValueError(f"Unknown provider type: {config.type}")

ai_docgen/providers/ollama.py

@@ -0,0 +1,25 @@
+import httpx
+
+from ai_docgen.providers.base import LLMProvider
+
+
+class OllamaProvider(LLMProvider):
+    def __init__(self, model: str, base_url: str = "http://localhost:11434") -> None:
+        self.model = model
+        self.base_url = base_url.rstrip("/")
+
+    async def complete(self, system: str, user: str) -> str:
+        async with httpx.AsyncClient(timeout=120) as client:
+            response = await client.post(
+                f"{self.base_url}/api/chat",
+                json={
+                    "model": self.model,
+                    "stream": False,
+                    "messages": [
+                        {"role": "system", "content": system},
+                        {"role": "user", "content": user},
+                    ],
+                },
+            )
+            response.raise_for_status()
+            return str(response.json()["message"]["content"])

ai_docgen/providers/openai.py

@@ -0,0 +1,20 @@
+from openai import AsyncOpenAI
+
+from ai_docgen.providers.base import LLMProvider
+
+
+class OpenAIProvider(LLMProvider):
+    def __init__(self, model: str, api_key: str) -> None:
+        self.model = model
+        self.client = AsyncOpenAI(api_key=api_key)
+
+    async def complete(self, system: str, user: str) -> str:
+        response = await self.client.chat.completions.create(
+            model=self.model,
+            messages=[
+                {"role": "system", "content": system},
+                {"role": "user", "content": user},
+            ],
+        )
+        content = response.choices[0].message.content
+        return content or ""

ai_docgen/registry.py

@@ -0,0 +1,125 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import date, timedelta
+from pathlib import Path
+
+import yaml
+
+STALE_DAYS = 30
+
+
+def _as_date(value: date | str | None, fallback: date) -> date:
+    if value is None:
+        return fallback
+    if isinstance(value, date):
+        return value
+    return date.fromisoformat(value)
+
+
+def _as_optional_date(value: date | str | None) -> date | None:
+    if value is None:
+        return None
+    if isinstance(value, date):
+        return value
+    return date.fromisoformat(value)
+
+
+@dataclass
+class DocumentEntry:
+    target: str
+    last_updated: date | None = None
+
+
+@dataclass
+class ProjectEntry:
+    name: str
+    path: str
+    remote: str
+    documents: list[DocumentEntry] = field(default_factory=list)
+    registered_at: date = field(default_factory=date.today)
+    last_updated: date = field(default_factory=date.today)
+    status: str = "synced"
+
+    def compute_status(self) -> str:
+        if (date.today() - self.last_updated) > timedelta(days=STALE_DAYS):
+            return "stale"
+        return "synced"
+
+
+class Registry:
+    def __init__(self, path: Path) -> None:
+        self.path = path
+
+    def all_projects(self) -> list[ProjectEntry]:
+        if not self.path.exists():
+            return []
+        data = yaml.safe_load(self.path.read_text()) or {}
+        projects = []
+        for p in data.get("projects", []):
+            docs = [
+                DocumentEntry(
+                    target=d["target"],
+                    last_updated=_as_optional_date(d.get("last_updated")),
+                )
+                for d in p.get("documents", [])
+            ]
+            projects.append(
+                ProjectEntry(
+                    name=p["name"],
+                    path=p["path"],
+                    remote=p.get("remote", ""),
+                    documents=docs,
+                    registered_at=_as_date(p.get("registered_at"), date.today()),
+                    last_updated=_as_date(p.get("last_updated"), date.today()),
+                    status=p.get("status", "synced"),
+                )
+            )
+        return projects
+
+    def register(self, entry: ProjectEntry) -> None:
+        projects = self.all_projects()
+        existing = next((p for p in projects if p.name == entry.name), None)
+        if existing:
+            projects.remove(existing)
+        entry.last_updated = date.today()
+        entry.status = entry.compute_status()
+        projects.append(entry)
+        self.write(projects)
+
+    def update_document_timestamp(self, project_name: str, target: str, updated: date) -> None:
+        projects = self.all_projects()
+        for project in projects:
+            if project.name == project_name:
+                for doc in project.documents:
+                    if doc.target == target:
+                        doc.last_updated = updated
+                project.last_updated = updated
+                project.status = project.compute_status()
+        self.write(projects)
+
+    def write(self, projects: list[ProjectEntry]) -> None:
+        self.path.parent.mkdir(parents=True, exist_ok=True)
+        data = {
+            "projects": [
+                {
+                    "name": p.name,
+                    "path": p.path,
+                    "remote": p.remote,
+                    "registered_at": p.registered_at.isoformat(),
+                    "last_updated": p.last_updated.isoformat(),
+                    "status": p.status,
+                    "documents": [
+                        {
+                            "target": d.target,
+                            "last_updated": (
+                                d.last_updated.isoformat() if d.last_updated else None
+                            ),
+                        }
+                        for d in p.documents
+                    ],
+                }
+                for p in projects
+            ]
+        }
+        self.path.write_text(yaml.dump(data, default_flow_style=False, allow_unicode=True))

ai_docgen/renderer.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+from jinja2 import Environment, select_autoescape
+
+
+@dataclass
+class Section:
+    name: str
+    is_auto: bool
+    content: str
+    start: int
+    end: int
+
+
+AUTO_OPEN = re.compile(r"<!-- AUTO:(\w+) -->")
+AUTO_CLOSE = re.compile(r"<!-- /AUTO:(\w+) -->")
+MANUAL_OPEN = re.compile(r"<!-- MANUAL -->")
+MANUAL_CLOSE = re.compile(r"<!-- /MANUAL -->")
+
+BUILTIN_TEMPLATES_DIR = Path(__file__).parent / "built_in_templates"
+
+
+class DocumentRenderer:
+    def parse_sections(self, text: str) -> list[Section]:
+        sections: list[Section] = []
+        lines = text.splitlines(keepends=True)
+        i = 0
+        while i < len(lines):
+            auto_match = AUTO_OPEN.match(lines[i].strip())
+            manual_match = MANUAL_OPEN.match(lines[i].strip())
+            if auto_match:
+                name = auto_match.group(1)
+                start = i
+                content_lines: list[str] = []
+                i += 1
+                while i < len(lines) and not AUTO_CLOSE.match(lines[i].strip()):
+                    content_lines.append(lines[i])
+                    i += 1
+                sections.append(
+                    Section(
+                        name=name,
+                        is_auto=True,
+                        content="".join(content_lines),
+                        start=start,
+                        end=i,
+                    )
+                )
+            elif manual_match:
+                start = i
+                content_lines = []
+                i += 1
+                while i < len(lines) and not MANUAL_CLOSE.match(lines[i].strip()):
+                    content_lines.append(lines[i])
+                    i += 1
+                sections.append(
+                    Section(
+                        name="MANUAL",
+                        is_auto=False,
+                        content="".join(content_lines),
+                        start=start,
+                        end=i,
+                    )
+                )
+            i += 1
+        return sections
+
+    def auto_section_names(self, text: str) -> list[str]:
+        return [s.name for s in self.parse_sections(text) if s.is_auto]
+
+    def patch_section(self, text: str, section_name: str, new_content: str) -> str:
+        open_marker = f"<!-- AUTO:{section_name} -->"
+        close_marker = f"<!-- /AUTO:{section_name} -->"
+        if open_marker not in text:
+            return text
+        before, rest = text.split(open_marker, 1)
+        _, after = rest.split(close_marker, 1)
+        return f"{before}{open_marker}\n{new_content}{close_marker}{after}"
+
+
+class TemplateLoader:
+    def __init__(self, source: str = "builtin", local_path: Path | None = None) -> None:
+        self.source = source
+        self.local_path = local_path
+
+    def load(self, template_name: str) -> str:
+        template_path = f"{template_name}.md.j2"
+        if self.source == "local" and self.local_path:
+            full_path = self.local_path / template_path
+            if full_path.exists():
+                return full_path.read_text()
+        builtin_path = BUILTIN_TEMPLATES_DIR / template_path
+        if not builtin_path.exists():
+            raise FileNotFoundError(f"Template not found: {template_name}")
+        return builtin_path.read_text()
+
+    def render(self, template_name: str, context: dict[str, str]) -> str:
+        template_content = self.load(template_name)
+        env = Environment(autoescape=select_autoescape([]))
+        tmpl = env.from_string(template_content)
+        return tmpl.render(**context)

ai_docgen/reporters/html.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+
+from ai_docgen.registry import Registry
+
+HTML_TEMPLATE = """\
+<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<title>ai-docgen — Project Status</title>
+<style>
+  body {{
+    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif;
+    max-width: 900px; margin: 40px auto; padding: 0 20px; color: #333;
+  }}
+  h1 {{ color: #1a1a2e; }}
+  table {{ width: 100%; border-collapse: collapse; margin-top: 20px; }}
+  th {{ background: #1a1a2e; color: white; padding: 10px 14px; text-align: left; }}
+  td {{ padding: 10px 14px; border-bottom: 1px solid #eee; }}
+  tr:hover td {{ background: #f9f9f9; }}
+  .synced {{ color: #22863a; font-weight: 600; }}
+  .stale {{ color: #cb2431; font-weight: 600; }}
+  footer {{ margin-top: 40px; color: #999; font-size: 0.85em; }}
+</style>
+</head>
+<body>
+<h1>ai-docgen — Project Status</h1>
+<table>
+  <thead><tr><th>Project</th><th>Status</th><th>Last Updated</th><th>Documents</th></tr></thead>
+  <tbody>
+{rows}
+  </tbody>
+</table>
+<footer>Generated by ai-docgen</footer>
+</body>
+</html>
+"""
+
+
+def render_html_report(registry_path: Path, output_path: Path) -> None:
+    registry = Registry(registry_path)
+    projects = registry.all_projects()
+    rows: list[str] = []
+    for project in projects:
+        doc_links = ", ".join(f'<a href="{d.target}">{d.target}</a>' for d in project.documents)
+        rows.append(
+            f"    <tr>"
+            f"<td>{project.name}</td>"
+            f'<td class="{project.status}">{project.status}</td>'
+            f"<td>{project.last_updated}</td>"
+            f"<td>{doc_links}</td>"
+            f"</tr>"
+        )
+    output_path.write_text(HTML_TEMPLATE.format(rows="\n".join(rows)))