04_skills_docgen.py

#!/usr/bin/env python3
"""Document Generation using GitHub Copilot SDK Skills (SKILL.md).

See the tutorial for learning goals, prerequisites, and usage:
    docs/copilot_sdk_tutorial/tutorials/04_skills.md     (English)
    docs/copilot_sdk_tutorial/tutorials/04_skills.ja.md  (日本語)
"""

import argparse
import asyncio
import sys
from pathlib import Path

from _telemetry import add_telemetry_arguments, apply_telemetry_arguments, make_client
from copilot.generated.rpc import PermissionDecisionApproveOnce
from copilot.generated.session_events import (
    SessionEventType,
    PermissionRequest,
)
from copilot.session import (
    PermissionRequestResult,
    SystemMessageReplaceConfig,
)

# ---------------------------------------------------------------------------
# Sample Python code that needs docstrings generated
# ---------------------------------------------------------------------------

SAMPLE_CODE = """\
def calculate_discount(price: float, discount_pct: float) -> float:
    if discount_pct < 0 or discount_pct > 100:
        raise ValueError("discount_pct must be between 0 and 100")
    return price * (1 - discount_pct / 100)


def batch_process(items: list[str], handler) -> list[str]:
    results = []
    for item in items:
        try:
            results.append(handler(item))
        except Exception as exc:
            results.append(f"ERROR: {exc}")
    return results
"""


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Document Generation using GitHub Copilot SDK Skills",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        "--skills-dir",
        "-s",
        default=str(Path(__file__).parent / "skills"),
        help="Path to the skills directory containing SKILL.md files",
    )
    parser.add_argument(
        "--cli-url",
        "-c",
        default=None,
        help=(
            "Optional Copilot CLI server URL (e.g. localhost:3000). "
            "When omitted, the SDK launches the copilot CLI over stdio."
        ),
    )
    add_telemetry_arguments(parser)
    return parser.parse_args()


async def run(cli_url: str | None, skills_dir: str) -> None:
    skills_path = Path(skills_dir)
    if not skills_path.exists():
        print(
            f"[Warning] Skills directory not found: {skills_dir}. "
            "Running without skills.",
            file=sys.stderr,
        )
        resolved_skills_dir: str | None = None
    else:
        resolved_skills_dir = str(skills_path.resolve())
        print(f"[Info] Loading skills from: {resolved_skills_dir}")

    def approve_all(
        request: PermissionRequest,
        context: dict,
    ) -> PermissionRequestResult:
        return PermissionDecisionApproveOnce()

    client = make_client(cli_url)
    await client.start()

    session_config: dict = {
        "on_permission_request": approve_all,
        "tools": [],
        "streaming": True,
        "system_message": SystemMessageReplaceConfig(
            mode="replace",
            content=(
                "You are a Python documentation specialist. "
                "Generate clear, complete Google-style docstrings for all functions "
                "in the provided code. Return only the updated code with docstrings added."
            ),
        ),
    }
    if resolved_skills_dir:
        session_config["skill_directories"] = [resolved_skills_dir]

    session = await client.create_session(**session_config)

    print("=== Generating Documentation ===\n")

    def on_event(event) -> None:  # noqa: ANN001
        if event.type == SessionEventType.ASSISTANT_MESSAGE_DELTA:
            print(event.data.delta_content, end="", flush=True)
        elif event.type == SessionEventType.TOOL_EXECUTION_START:
            print(f"\n[Skill] Running: {event.data.tool_name}", file=sys.stderr)
        elif event.type == SessionEventType.SESSION_ERROR:
            print(f"\n[Error] {event.data.message}", file=sys.stderr)

    session.on(on_event)

    prompt = (
        f"Please add Google-style docstrings to all functions in the following code:\n\n"
        f"```python\n{SAMPLE_CODE}\n```"
    )
    await session.send_and_wait(prompt, timeout=300)
    print("\n\n=== Done ===")


def main() -> None:
    args = parse_args()
    apply_telemetry_arguments(args)
    try:
        asyncio.run(run(args.cli_url, args.skills_dir))
    except KeyboardInterrupt:
        print("\nBye!")


if __name__ == "__main__":
    main()