diff --git a/AGENTS.md b/AGENTS.md index 89908d7..b88d959 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -5,15 +5,15 @@ Stable goals for this repository: 1. **Embeddable generator** — **`pip install docgen`** (git URL or editable install from this repo), a consumer **bundle** (`docgen.yaml` + hints/narration), and shell/CI are enough to build and maintain narrated demos. **Do not vendor** this library into a product repo’s `src/`; pin via `requirements-docgen.txt` / `pipx` / `uv tool`. **No IDE assistant is required**; optional **`docgen wizard`** is a local web app only. -2. **Hybrid config and prose** — **`docgen.yaml`** should stay maintainable: deterministic merges (**`yaml-generate`**, gap checks) plus **optional OpenAI** where it adds value (narration hints, declarative scene YAML). Prefer **Git-reviewed** changes over opaque single-shot generation. -3. **Video stack** — Long-form demos pair **Markdown narration**, **OpenAI TTS**, **Whisper-style timestamps**, **Manim** visuals, **`compose`** (ffmpeg), **`concat`**, and **`validate`** (sync and narration lint). The CLI also supports **`pages`** for static preview sites. +2. **Hybrid config and prose** — **`docgen.yaml`** should stay maintainable: deterministic merges (**`yaml-generate`**, gap checks) plus **optional LLM** (OpenAI or Grok) where it adds value (narration hints, declarative scene YAML). Prefer **Git-reviewed** changes over opaque single-shot generation. +3. **Video stack** — Long-form demos pair **Markdown narration**, **TTS** (OpenAI or xAI), **Whisper-style timestamps**, **Manim** visuals, **`compose`** (ffmpeg), **`concat`**, and **`validate`** (sync and narration lint). The CLI also supports **`pages`** for static preview sites. 4. **Stable contracts** — CLI, exit codes, and reusable workflows should stay predictable for downstream repos and automation. 5. **Library, not app** — There is **no in-repo dogfood bundle**. Consumer projects (e.g. `course-builder`) are the integration test of record. The library must not import or special-case any consumer. 6. **Tool-only generation** — Narration, merged **`docgen.yaml`**, compiled **`scenes.py`**, TTS audio, composed media, and other **generated** artifacts must come from **docgen** (CLI/library) and **committed wrapper scripts** that call it — not from hand-edited outputs passed off as sources. In a consumer bundle, prefer **`hints/*.md`** + **`yaml-generate`** over ad-hoc YAML surgery. Cursor rules: **`.cursor/rules/docgen-tools-only.mdc`**, **`.cursor/rules/no-asset-edits.mdc`**. ## Protected assets in a consumer bundle (Cursor must not edit) -`docgen` + OpenAI are the **only** path that produces category **C** outputs — see **`.cursor/rules/no-asset-edits.mdc`**. Summary (paths relative to a consumer bundle, typically `docs/demos/`): +`docgen` (OpenAI or Grok) is the **only** path that produces category **C** outputs — see **`.cursor/rules/no-asset-edits.mdc`**. Summary (paths relative to a consumer bundle, typically `docs/demos/`): - **Outputs (do not hand-edit):** `/docgen.yaml` (as emitted by **`yaml-generate`**); `/narration/*.md` (except `README.md`); `/animations/scenes.py`, `timing.json`, `animations/specs/*.scene.yaml` (scene pipeline); `/audio/*.mp3`; `/images/*.png` (scene image assets from **`image-generate`**); `/recordings/**` where applicable. - **Inputs (maintainer-owned):** `/hints/**` with YAML front matter (`docgen.segment`, `docgen.wiring`); maintainer scripts under the bundle; `tests/**` fixtures inside this library; `/narration/README.md`. @@ -44,9 +44,9 @@ Commands registered on the **`docgen`** CLI include: - **`wizard`** — local web UI for narration/bootstrap workflows (focus files, **in-place narration revise**, per-segment **asset freshness** + **rebuild-from-here**, Vue **Benchmark** view, **Tool** tab to pip-upgrade docgen and pin `requirements-docgen.txt`). - **`gui`** — desktop window over the same Vue/Flask UI (`pip install 'docgen[gui]'` for pywebview). ``--smoke`` is a headless HTTP check. PyInstaller spec: ``packaging/docgen-gui.spec``. Frozen apps resolve templates/static/benchmark JSON via ``docgen.resources``. - **`freeze`** — ``docgen freeze`` builds the **`docgen-gui`** onedir (`pip install 'docgen[packaging]'`). Optional ``--smoke`` runs the binary headless. Do not run a full freeze in routine pytest; set ``DOCGEN_FREEZE_SMOKE=1`` for the optional test. -- **`tts`** — text-to-speech for segment files. -- **`timestamps`** — word/segment timing (`timing.json`). Default engine **`local`** aligns the known narration text against the mp3 offline (ffmpeg silencedetect, no API); **`--engine whisper`** keeps OpenAI whisper-1 transcription. Both emit the same Whisper-shaped blocks. -- **`image-generate`** — render scene-spec **image elements** (`image:` + `prompt:` boxes) via the OpenAI Images API into the bundle (also runs for missing assets inside `generate-all`). +- **`tts`** — text-to-speech for segment files (OpenAI or xAI `/v1/tts`). +- **`timestamps`** — word/segment timing (`timing.json`). Default engine **`local`** aligns the known narration text against the mp3 offline (ffmpeg silencedetect, no API); **`--engine whisper`** uses OpenAI whisper-1 or xAI `/v1/stt` when `ai.provider` is grok. Both emit the same Whisper-shaped blocks. +- **`image-generate`** — render scene-spec **image elements** (`image:` + `prompt:` boxes) via OpenAI Images or xAI Imagine into the bundle (also runs for missing assets inside `generate-all`). - **`manim`** — render Manim scenes declared in config. - **`compose`** — mux narration audio with visual sources via ffmpeg. - **`validate`** / **`validate --pre-push`** — drift, narration lint, Manim hints, **`timing_sync`**, **`story_end`** (last paced reveal vs audio end; hard fail), **`scene_assets`** (pre-render: stuck-board cadence, frame-budget overlaps, `MANIM_FONT` consistency, stale helpers / stale compiled class — hard fail; also a `generate-all` gate before Manim), **`av_sync`** (soft; prefers scene-spec labels as OCR anchors), **`subject_beat_coverage`** (declarative specs vs narration topic beats; hard fail when enabled), and related checks. @@ -93,7 +93,11 @@ Tests should cover **CLI-visible behavior** and contracts that adopters rely on: - **Virtualenv:** the project is installed editable into **`/workspace/.venv`** (created by the startup update script). Shells do **not** auto-activate it — run `. /workspace/.venv/bin/activate` (or prefix the venv path) before `docgen`, `pytest`, or `ruff`. The `docgen` console script lives at `/workspace/.venv/bin/docgen`. - **System deps are pre-baked in the VM snapshot** (not the update script): `ffmpeg` + `tesseract-ocr` (validation/compose/OCR), plus `build-essential`, `python3-dev`, `libcairo2-dev`, `libpango1.0-dev`, `pkg-config` (needed to build the `manim` extra's `manimpango`/`pycairo` wheels). If a fresh VM ever lacks these, reinstall via apt before `pip install`. - **Standard commands** are in `README.md` / `pyproject.toml` / `.github/workflows/ci.yml`: lint `ruff check src/ tests/`; tests `pytest tests/ -v --tb=short`; **required** `docgen benchmark` (CI job `benchmark`). The CI unit job also exports `PYTHONPATH=src` (not needed locally because of the editable install, but harmless). -- **OpenAI-gated vs offline commands:** `tts`, `timestamps --engine whisper`, `image-generate`, `narration-generate`, `scene-spec-generate`, and `yaml-generate --llm` call OpenAI and need `OPENAI_API_KEY` (integration tests auto-skip without it). Fully offline: `init`, `scene-compile`, `manim`, `compose`, `validate`, `lint`, `pages`, `concat`, `yaml-generate` (no `--llm`), `timestamps` (default `local` engine), and `benchmark`. +- **OpenAI / Grok vs offline commands:** `tts`, `timestamps --engine whisper`, `image-generate`, `narration-generate`, `scene-spec-generate`, and `yaml-generate --llm` call a provider and need `OPENAI_API_KEY` or (`DOCGEN_AI_PROVIDER=grok` + `XAI_API_KEY`). Integration tests auto-skip without credentials. Fully offline: `init`, `scene-compile`, `manim`, `compose`, `validate`, `lint`, `pages`, `concat`, `yaml-generate` (no `--llm`), `timestamps` (default `local` engine), and `benchmark`. +- **Generate against another repo (do not vendor):** this environment already has `docgen` on PATH. Point at a consumer checkout or clone URL — nothing is copied into that project's `src/`: + `docgen --repo /path/to/consumer generate-all` + `docgen --repo github.com/org/consumer init --defaults` + `--repo` finds `docs/demos/docgen.yaml`. Add the consumer as a Cloud **repository dependency** if you clone by GitHub URL (uses `GITHUB_TOKEN`). Set `XAI_API_KEY` or `OPENAI_API_KEY` as an environment secret for networked stages. - **`scene-compile` gotcha:** paced specs (`wait_word`) need a `timing.json` entry for that stem (`docgen timestamps` after TTS). Prefer `scene-compile --retime` after fresh timestamps; for a fully offline smoke render, author rows without wait indices only if you accept unpaced reveals. - **No in-repo dogfood bundle:** exercise the pipeline against a scratch bundle (`docgen init /tmp/ --defaults` in a throwaway git dir). Do not hand-edit consumer generated assets (see `.cursor/rules/no-asset-edits.mdc`). - **Wizard / desktop GUI:** `docgen wizard --port 8501` is the Flask UI (bundle optional for the Benchmark view). `docgen gui` / `docgen benchmark --gui` open the Vue benchmark view in a desktop window when `pywebview` is installed. `docgen gui --smoke` is the headless HTTP check (default pytest). `docgen freeze` builds the GUI onedir; do not run a full PyInstaller freeze in routine CI/pytest. Keep `packaging/docgen-gui.spec` and `docgen.gui.packaging` in sync. diff --git a/README.md b/README.md index 963d737..532342c 100644 --- a/README.md +++ b/README.md @@ -1,8 +1,8 @@ # docgen — documentation generator Reusable Python library and CLI for **narrated demo videos** built around **Manim**, -**OpenAI TTS**, and **ffmpeg** composition. Aimed at long-form, scripted explainers -that walk through how a system works. +**TTS** (OpenAI or xAI Grok), and **ffmpeg** composition. Aimed at long-form, +scripted explainers that walk through how a system works. ## Suite handbook (Courseforge) @@ -31,12 +31,12 @@ If you still need the legacy behaviour, pin a pre-removal commit ## What docgen does today - **TTS narration** — generate MP3 audio from Markdown scripts via OpenAI - `gpt-4o-mini-tts`. + `gpt-4o-mini-tts`, or xAI `/v1/tts` when `ai.provider` is `grok`. - **Word-level timestamps without Whisper** — the default `local` engine aligns the known narration text against the TTS mp3 offline (ffmpeg `silencedetect` - + proportional interpolation); no API call or transcription. OpenAI - `whisper-1` remains available via `timestamps.engine: whisper` / - `docgen timestamps --engine whisper`. Both engines write the same + + proportional interpolation); no API call or transcription. Network + transcription (`timestamps.engine: whisper`) uses OpenAI `whisper-1` or xAI + `/v1/stt` when the provider is Grok. Both engines write the same `timing.json` shape. - **Manim animations (default: declarative scene specs)** — primary visual surface. Prefer **`animations/specs/*.scene.yaml`** via **`docgen scene-spec-generate`** @@ -47,11 +47,12 @@ If you still need the legacy behaviour, pin a pre-removal commit map each paced label to a **`wait_word`** index. Hand-maintained custom Manim classes may still live in `animations/scenes.py` outside `BEGIN/END GENERATED SCENE` markers (use **`--skip-scene-retime`** to bypass the declarative stage). -- **OpenAI image assets in Manim scenes** — a scene-spec box may be an **image +- **Image assets in Manim scenes** — a scene-spec box may be an **image element** (`image: images/.png` + `prompt:`); `docgen image-generate` - renders the prompt via the OpenAI Images API (default `gpt-image-1`) and the - compiled scene shows it with the `_image` helper (`ImageMobject`, `Group` - rows). `generate-all` fills in missing assets automatically. + renders the prompt via OpenAI Images (default `gpt-image-1`) or xAI Imagine + (`grok-imagine-image-2.0` when `ai.provider` is `grok`). The compiled scene + shows it with the `_image` helper. `generate-all` fills in missing assets + automatically. - **ffmpeg composition** — combine narration audio and Manim video into final segments, with a freeze-tail guard. - **Validation** — A/V drift, freeze ratio, OCR error scan, layout, narration lint, @@ -65,8 +66,56 @@ If you still need the legacy behaviour, pin a pre-removal commit docs. **No IDE lock-in:** maintenance workflows are `docgen` CLI + YAML + shell/CI (and -OpenAI where a command calls the API). The wizard is a local Flask app, not a -plugin tied to one editor. +an LLM/TTS provider where a command calls the API). The wizard is a local Flask +app, not a plugin tied to one editor. + +### Run against another repo (do not vendor) + +Install docgen once (venv / pipx / this Cloud environment). Pass the **consumer** +checkout or clone URL; nothing from this library is copied into that project's +`src/`. + +```bash +# Local checkout of the product repo +docgen --repo /path/to/course-builder init --defaults +docgen --repo /path/to/course-builder yaml-generate +docgen --repo /path/to/course-builder generate-all + +# GitHub URL or org/repo shorthand (shallow clone into DOCGEN_REPO_CACHE) +docgen --repo github.com/acme/course-builder generate-all +# equivalent: DOCGEN_REPO=acme/course-builder docgen generate-all +``` + +`--repo` looks for `docs/demos/docgen.yaml` (then any other `docgen.yaml` under +the checkout). `repo_root` in that yaml still points at the consumer so +narration/scene prompts read *their* sources. + +### OpenAI or Grok (xAI) + +Default provider is **OpenAI** (`OPENAI_API_KEY`). To substitute **Grok** for +chat, TTS, Whisper, and image calls: + +```yaml +ai: + provider: grok # openai | grok +``` + +```bash +export DOCGEN_AI_PROVIDER=grok # overrides yaml +export XAI_API_KEY=xai-... +docgen --repo /path/to/consumer generate-all +``` + +Existing OpenAI model names in YAML are remapped at call time (`gpt-4o` → +`grok-4.6`, `gpt-image-1` → `grok-imagine-image-2.0`, TTS voice `coral` → +`eve`). Chat and images use the OpenAI SDK at `https://api.x.ai/v1`. TTS and +STT use xAI `POST /v1/tts` and `POST /v1/stt`. Keep `timestamps.engine: local` +unless you specifically want network STT. + +A Cursor Cloud environment that already has ffmpeg / tesseract / Manim build +deps can run the full pipeline: add `XAI_API_KEY` or `OPENAI_API_KEY` as an +environment secret (and grant the consumer as a repository dependency if you +clone by URL). ## Install (external tool — do not vendor into project `src/`) @@ -115,6 +164,8 @@ CI also installs `ffmpeg` and `tesseract` via apt for unit tests — see `.githu ## Quick start (in a consumer repo) +From **inside** the consumer bundle (library already on PATH): + ```bash cd your-project/docs/demos # bundle only — library is on PATH via pip pip install -r requirements-docgen.txt # after docgen init, or use the pip line above @@ -123,10 +174,18 @@ docgen generate-all # TTS → timestamps → scene retime → Manim → c docgen validate --pre-push ``` +From a machine / Cloud environment that has **docgen installed** but is not the +consumer repo: + +```bash +docgen --repo /path/to/your-project generate-all +``` + ## CLI commands | Command | Description | |---------|-------------| +| `docgen --repo PATH_OR_URL …` | Target a consumer checkout or clone URL (also `DOCGEN_REPO`). Finds `docs/demos/docgen.yaml`. Does **not** copy this library into the consumer `src/` | | `docgen --version` | Show installed version + recommended `pip install` line (external tool) | | `docgen init [TARGET_DIR] [--defaults] [--segments-file FILE]` | Scaffold a bundle: `docgen.yaml`, `requirements-docgen.txt`, wrapper scripts, directories | | `docgen wizard [--port 8501]` | Local web GUI: focus files, **revise narration**, asset freshness / rebuild-from-here, Vue **Benchmark** view, and a **Tool** tab to upgrade the installed `docgen` package (pip) + rewrite `requirements-docgen.txt` | @@ -169,14 +228,15 @@ your IDE or CI) is **not** replaced by the file. To make the file win, set environment, or **`DOCGEN_ENV_OVERRIDES=OPENAI_API_KEY,OTHER_KEY`** for specific keys only. -When `OPENAI_API_KEY` is present in both the shell and `env_file`, docgen prints a -one-line hint to stderr so a silent 401 from the wrong key is easier to diagnose. +When `OPENAI_API_KEY` or `XAI_API_KEY` is present in both the shell and `env_file`, +docgen prints a one-line hint to stderr so a silent 401 from the wrong key is +easier to diagnose. ### Narration from source (owner hints) Under `narration_from_source` in `docgen.yaml`, the **project owner** lists optional `hints` (strings) that steer the model (audience, terminology, what to -avoid). OpenAI generates the narration `.md` from your repo context +avoid). The chat model generates the narration `.md` from your repo context (`context.paths` / `context.globs`, relative to `repo_root`) plus those hints; the result is what `docgen tts` reads. See `docgen.narrate_from_source`. @@ -222,8 +282,12 @@ validation: prefer_scene_spec_labels: true # OCR anchors from paced box labels when specs exist visual_types: [manim] # only check types with on-screen text +```yaml +ai: + provider: openai # openai | grok (xAI). Override with DOCGEN_AI_PROVIDER. + timestamps: - engine: local # local (default, offline) or whisper (OpenAI whisper-1) + engine: local # local (default, offline) or whisper (OpenAI whisper-1 / xAI STT) silence_noise_db: -35.0 # ffmpeg silencedetect threshold for the local engine min_silence_sec: 0.3 diff --git a/src/docgen/ai_client.py b/src/docgen/ai_client.py new file mode 100644 index 0000000..106ca99 --- /dev/null +++ b/src/docgen/ai_client.py @@ -0,0 +1,509 @@ +"""OpenAI-compatible AI client with first-class xAI / Grok support. + +Chat completions and (most) image generation go through the official +``openai`` SDK. xAI speaks that protocol at ``https://api.x.ai/v1``, so +switching providers is ``base_url`` + key + model aliases. + +TTS and speech-to-text are **not** drop-in compatible: OpenAI uses +``/v1/audio/speech`` and ``whisper-1``; xAI uses ``POST /v1/tts`` and +``POST /v1/stt``. Those paths are adapted here so ``docgen tts`` / +``timestamps --engine whisper`` keep working. + +Resolution order for ``ai.provider``: + +1. ``DOCGEN_AI_PROVIDER`` (``openai`` or ``grok`` / ``xai``) +2. ``ai.provider`` in ``docgen.yaml`` +3. default ``openai`` (unchanged behaviour) + +Grok auth reads ``XAI_API_KEY``, then ``OPENAI_API_KEY``. OpenAI auth +reads ``OPENAI_API_KEY`` only (unless ``ai.api_key_env`` overrides). +""" + +from __future__ import annotations + +import json +import os +import random +import time +import urllib.error +import urllib.request +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from docgen.config import Config + +GROK_BASE_URL = "https://api.x.ai/v1" +DEFAULT_GROK_CHAT_MODEL = "grok-4.6" +DEFAULT_GROK_IMAGE_MODEL = "grok-imagine-image-2.0" +DEFAULT_GROK_TTS_VOICE = "eve" +DEFAULT_GROK_TTS_LANGUAGE = "en" +GROK_TTS_MAX_CHARS = 15_000 + +_GROK_PROVIDERS = frozenset({"grok", "xai", "x.ai"}) +_OPENAI_PROVIDERS = frozenset({"openai", "oai", ""}) + +# Existing yaml/init defaults keep OpenAI model names; remap at call time. +_GROK_CHAT_ALIASES = { + "gpt-4o": DEFAULT_GROK_CHAT_MODEL, + "gpt-4o-mini": DEFAULT_GROK_CHAT_MODEL, + "gpt-4.1": DEFAULT_GROK_CHAT_MODEL, + "gpt-4.1-mini": DEFAULT_GROK_CHAT_MODEL, + "gpt-4.1-nano": DEFAULT_GROK_CHAT_MODEL, + "gpt-4": DEFAULT_GROK_CHAT_MODEL, + "gpt-3.5-turbo": DEFAULT_GROK_CHAT_MODEL, +} + +_GROK_IMAGE_ALIASES = { + "gpt-image-1": DEFAULT_GROK_IMAGE_MODEL, + "gpt-image-1-mini": DEFAULT_GROK_IMAGE_MODEL, + "dall-e-3": DEFAULT_GROK_IMAGE_MODEL, + "dall-e-2": "grok-imagine-image", +} + +_GROK_VOICES = frozenset({"eve", "ara", "rex", "sal", "leo"}) +_OPENAI_TO_GROK_VOICE = { + "alloy": "sal", + "ash": "rex", + "ballad": "leo", + "coral": "eve", + "echo": "rex", + "fable": "leo", + "nova": "ara", + "onyx": "rex", + "sage": "ara", + "shimmer": "eve", + "verse": "sal", +} + +_API_KEY_ENVS = ("OPENAI_API_KEY", "XAI_API_KEY") +_MAX_HTTP_ATTEMPTS = 10 +_BASE_DELAY_SEC = 1.0 +_MAX_BACKOFF_SEC = 120.0 + + +@dataclass(frozen=True) +class AISettings: + provider: str # "openai" | "grok" + base_url: str | None + api_key: str | None + api_key_env: str + tts_language: str + + @property + def is_grok(self) -> bool: + return self.provider == "grok" + + def auth_help(self) -> str: + if self.is_grok: + return ( + f"Set {self.api_key_env} (xAI / Grok), or OPENAI_API_KEY, " + "and ai.provider: grok / DOCGEN_AI_PROVIDER=grok." + ) + return ( + f"Set {self.api_key_env}, or switch to Grok with ai.provider: grok " + "and XAI_API_KEY." + ) + + +def normalize_provider(raw: str | None) -> str: + value = (raw or "").strip().lower() + if value in _GROK_PROVIDERS: + return "grok" + if value in _OPENAI_PROVIDERS: + return "openai" + raise ValueError( + f"Unknown AI provider {raw!r}; use 'openai' or 'grok' " + "(aliases: xai, x.ai)." + ) + + +def resolve_ai_settings(cfg: "Config | None" = None) -> AISettings: + """Resolve provider, base URL, and API key from env + optional bundle config.""" + block: dict[str, Any] = {} + tts_language = DEFAULT_GROK_TTS_LANGUAGE + if cfg is not None: + raw_ai = cfg.raw.get("ai") if isinstance(cfg.raw, dict) else None + if isinstance(raw_ai, dict): + block = raw_ai + tts = cfg.raw.get("tts") if isinstance(cfg.raw, dict) else None + if isinstance(tts, dict): + lang = str(tts.get("language") or "").strip() + if lang: + tts_language = lang + + env_provider = (os.environ.get("DOCGEN_AI_PROVIDER") or "").strip() + yaml_provider = str(block.get("provider") or "").strip() + provider = normalize_provider(env_provider or yaml_provider or "openai") + + env_base = (os.environ.get("DOCGEN_AI_BASE_URL") or "").strip() + yaml_base = str(block.get("base_url") or "").strip() + if env_base: + base_url = env_base + elif yaml_base: + base_url = yaml_base + elif provider == "grok": + base_url = GROK_BASE_URL + else: + base_url = None + + env_key_name = (os.environ.get("DOCGEN_AI_API_KEY_ENV") or "").strip() + yaml_key_name = str(block.get("api_key_env") or "").strip() + if env_key_name: + api_key_env = env_key_name + elif yaml_key_name: + api_key_env = yaml_key_name + elif provider == "grok": + api_key_env = "XAI_API_KEY" + else: + api_key_env = "OPENAI_API_KEY" + + api_key = (os.environ.get(api_key_env) or "").strip() or None + if not api_key and provider == "grok": + api_key = (os.environ.get("OPENAI_API_KEY") or "").strip() or None + if api_key: + api_key_env = "OPENAI_API_KEY" + + return AISettings( + provider=provider, + base_url=base_url, + api_key=api_key, + api_key_env=api_key_env, + tts_language=tts_language, + ) + + +def openai_client(cfg: "Config | None" = None) -> Any: + """Return an ``openai.OpenAI`` client, optionally pointed at xAI.""" + import openai + + settings = resolve_ai_settings(cfg) + kwargs: dict[str, str] = {} + if settings.api_key: + kwargs["api_key"] = settings.api_key + if settings.base_url: + kwargs["base_url"] = settings.base_url + return openai.OpenAI(**kwargs) + + +def resolve_chat_model(model: str, settings: AISettings | None = None, *, cfg: "Config | None" = None) -> str: + chosen = (model or "").strip() + st = settings or resolve_ai_settings(cfg) + if not st.is_grok: + return chosen + if not chosen: + return DEFAULT_GROK_CHAT_MODEL + aliased = _GROK_CHAT_ALIASES.get(chosen) + if aliased: + return aliased + if chosen.lower().startswith("gpt-"): + return DEFAULT_GROK_CHAT_MODEL + return chosen + + +def resolve_image_model(model: str, settings: AISettings | None = None, *, cfg: "Config | None" = None) -> str: + chosen = (model or "").strip() + st = settings or resolve_ai_settings(cfg) + if not st.is_grok: + return chosen + if not chosen: + return DEFAULT_GROK_IMAGE_MODEL + aliased = _GROK_IMAGE_ALIASES.get(chosen) + if aliased: + return aliased + if chosen.lower().startswith(("gpt-image", "dall-e")): + return DEFAULT_GROK_IMAGE_MODEL + return chosen + + +def resolve_tts_voice(voice: str, settings: AISettings | None = None, *, cfg: "Config | None" = None) -> str: + chosen = (voice or "").strip() + st = settings or resolve_ai_settings(cfg) + if not st.is_grok: + return chosen + if chosen.lower() in _GROK_VOICES: + return chosen.lower() + return _OPENAI_TO_GROK_VOICE.get(chosen.lower(), DEFAULT_GROK_TTS_VOICE) + + +def chat_completion( + *, + system_prompt: str, + user_message: str, + model: str, + temperature: float, + cfg: "Config | None" = None, +) -> str: + """Chat completions via the OpenAI SDK (OpenAI or xAI Grok).""" + import openai + + settings = resolve_ai_settings(cfg) + client = openai_client(cfg) + resolved = resolve_chat_model(model, settings) + try: + response = client.chat.completions.create( + model=resolved, + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_message}, + ], + temperature=float(temperature), + ) + except openai.AuthenticationError as exc: + raise RuntimeError( + f"{_vendor(settings)} rejected {settings.api_key_env} (authentication failed): {exc}. " + f"{settings.auth_help()}" + ) from exc + except openai.PermissionDeniedError as exc: + raise RuntimeError( + f"{_vendor(settings)} permission denied for model {resolved!r}: {exc}." + ) from exc + except openai.APIConnectionError as exc: + raise RuntimeError( + f"{_vendor(settings)} connection error: {exc} — re-run when connectivity is restored." + ) from exc + return response.choices[0].message.content or "" + + +def synthesize_speech( + *, + text: str, + model: str, + voice: str, + instructions: str, + output_path: Path, + cfg: "Config | None" = None, +) -> None: + """Write MP3 bytes for ``text`` using OpenAI TTS or xAI ``/v1/tts``.""" + settings = resolve_ai_settings(cfg) + if settings.is_grok: + _grok_tts( + text=text, + voice=resolve_tts_voice(voice, settings), + language=settings.tts_language, + output_path=output_path, + settings=settings, + ) + return + + client = openai_client(cfg) + + def _call() -> None: + response = client.audio.speech.create( + model=model, + voice=voice, + input=text, + instructions=instructions, + ) + response.stream_to_file(str(output_path)) + + from docgen.openai_retry import call_with_rate_limit_retries + + call_with_rate_limit_retries(_call) + + +def transcribe_audio(audio_path: str | Path, *, cfg: "Config | None" = None) -> dict[str, Any]: + """Return Whisper-shaped ``{text, segments, words}`` from OpenAI or xAI STT.""" + settings = resolve_ai_settings(cfg) + if settings.is_grok: + return _grok_stt(Path(audio_path), settings) + + client = openai_client(cfg) + with open(audio_path, "rb") as f: + result = client.audio.transcriptions.create( + model="whisper-1", + file=f, + response_format="verbose_json", + timestamp_granularities=["word", "segment"], + ) + return { + "text": result.text, + "segments": [ + {"start": s.start, "end": s.end, "text": s.text} + for s in (result.segments or []) + ], + "words": [ + {"start": w.start, "end": w.end, "word": w.word} + for w in (result.words or []) + ], + } + + +def fetch_url_bytes(url: str, *, timeout: int = 120) -> bytes: + """GET ``url`` and return the response body (image CDN / signed URLs).""" + req = urllib.request.Request(url, method="GET") + with urllib.request.urlopen(req, timeout=timeout) as resp: + return resp.read() + + +def _vendor(settings: AISettings) -> str: + return "xAI" if settings.is_grok else "OpenAI" + + +def _grok_tts( + *, + text: str, + voice: str, + language: str, + output_path: Path, + settings: AISettings, +) -> None: + if not settings.api_key: + raise RuntimeError(f"xAI TTS needs an API key. {settings.auth_help()}") + if len(text) > GROK_TTS_MAX_CHARS: + raise RuntimeError( + f"xAI TTS accepts at most {GROK_TTS_MAX_CHARS} characters " + f"({len(text)} in this segment). Split the narration or shorten it." + ) + payload = { + "text": text, + "voice_id": voice, + "language": language or DEFAULT_GROK_TTS_LANGUAGE, + } + body = _http_json( + f"{settings.base_url.rstrip('/')}/tts", + payload, + settings=settings, + accept="audio/mpeg", + ) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_bytes(body) + + +def _grok_stt(audio_path: Path, settings: AISettings) -> dict[str, Any]: + if not settings.api_key: + raise RuntimeError(f"xAI STT needs an API key. {settings.auth_help()}") + data = _http_multipart( + f"{settings.base_url.rstrip('/')}/stt", + fields={"language": settings.tts_language or DEFAULT_GROK_TTS_LANGUAGE}, + filename=audio_path.name, + file_bytes=audio_path.read_bytes(), + content_type="audio/mpeg", + settings=settings, + ) + parsed = json.loads(data.decode("utf-8")) + text = str(parsed.get("text") or "") + raw_words = parsed.get("words") or [] + words: list[dict[str, Any]] = [] + for w in raw_words: + if not isinstance(w, dict): + continue + token = str(w.get("word") or w.get("text") or "").strip() + if not token: + continue + words.append( + { + "start": float(w.get("start") or 0.0), + "end": float(w.get("end") or 0.0), + "word": token, + } + ) + duration = float(parsed.get("duration") or (words[-1]["end"] if words else 0.0)) + segments = parsed.get("segments") + if not isinstance(segments, list) or not segments: + segments = ( + [{"start": words[0]["start"], "end": words[-1]["end"], "text": text}] + if words + else [{"start": 0.0, "end": duration, "text": text}] + ) + else: + segments = [ + { + "start": float(s.get("start") or 0.0), + "end": float(s.get("end") or 0.0), + "text": str(s.get("text") or ""), + } + for s in segments + if isinstance(s, dict) + ] + return {"text": text, "segments": segments, "words": words} + + +def _http_json( + url: str, + payload: dict[str, Any], + *, + settings: AISettings, + accept: str = "application/json", +) -> bytes: + data = json.dumps(payload).encode("utf-8") + headers = { + "Authorization": f"Bearer {settings.api_key}", + "Content-Type": "application/json", + "Accept": accept, + } + return _http_with_retries(url, data=data, headers=headers) + + +def _http_multipart( + url: str, + *, + fields: dict[str, str], + filename: str, + file_bytes: bytes, + content_type: str, + settings: AISettings, +) -> bytes: + boundary = "----docgen" + uuid.uuid4().hex + chunks: list[bytes] = [] + for key, value in fields.items(): + chunks.append(f"--{boundary}\r\n".encode()) + chunks.append(f'Content-Disposition: form-data; name="{key}"\r\n\r\n'.encode()) + chunks.append(str(value).encode("utf-8") + b"\r\n") + chunks.append(f"--{boundary}\r\n".encode()) + chunks.append( + f'Content-Disposition: form-data; name="file"; filename="{filename}"\r\n'.encode() + ) + chunks.append(f"Content-Type: {content_type}\r\n\r\n".encode()) + chunks.append(file_bytes) + chunks.append(f"\r\n--{boundary}--\r\n".encode()) + body = b"".join(chunks) + headers = { + "Authorization": f"Bearer {settings.api_key}", + "Content-Type": f"multipart/form-data; boundary={boundary}", + "Accept": "application/json", + } + return _http_with_retries(url, data=body, headers=headers) + + +def _http_with_retries(url: str, *, data: bytes, headers: dict[str, str]) -> bytes: + last_exc: BaseException | None = None + for attempt in range(_MAX_HTTP_ATTEMPTS): + req = urllib.request.Request(url, data=data, headers=headers, method="POST") + try: + with urllib.request.urlopen(req, timeout=120) as resp: + return resp.read() + except urllib.error.HTTPError as exc: + last_exc = exc + if exc.code == 429 and attempt < _MAX_HTTP_ATTEMPTS - 1: + retry_after = exc.headers.get("Retry-After") if exc.headers else None + delay = _retry_delay_sec(retry_after, attempt) + time.sleep(delay) + continue + detail = exc.read().decode("utf-8", errors="replace") if exc.fp else "" + raise RuntimeError( + f"xAI HTTP {exc.code} for {url}: {detail or exc.reason}" + ) from exc + except urllib.error.URLError as exc: + raise RuntimeError(f"xAI connection error for {url}: {exc}") from exc + raise RuntimeError(f"xAI request failed after retries: {last_exc}") + + +def _retry_delay_sec(retry_after: str | None, attempt: int) -> float: + delay: float | None = None + if retry_after: + try: + delay = float(retry_after) + except ValueError: + delay = None + if delay is None or delay <= 0: + delay = min(_MAX_BACKOFF_SEC, _BASE_DELAY_SEC * (2**attempt)) + jitter = random.uniform(0, max(0.05, 0.15 * delay)) + return delay + jitter + + +def conflicting_api_key_envs() -> tuple[str, ...]: + """Env names that warn when both the shell and ``env_file`` set them.""" + return _API_KEY_ENVS diff --git a/src/docgen/cli.py b/src/docgen/cli.py index 9704614..5591464 100644 --- a/src/docgen/cli.py +++ b/src/docgen/cli.py @@ -41,7 +41,7 @@ def _docgen_env_override_mode() -> str | set[str] | None: def _load_env(cfg: Config | None) -> None: - """Load .env file from config if specified, so OPENAI_API_KEY etc. are available. + """Load .env file from config if specified, so OPENAI_API_KEY / XAI_API_KEY etc. are available. By default **shell environment wins**: ``os.environ.setdefault`` does not replace keys already exported. Set ``DOCGEN_ENV_OVERRIDES=1`` so every key @@ -57,20 +57,19 @@ def _load_env(cfg: Config | None) -> None: os.environ[k] = v return override_keys = mode if isinstance(mode, set) else set() + from docgen.ai_client import conflicting_api_key_envs + + warn_keys = set(conflicting_api_key_envs()) for k, v in pairs: if k in override_keys: os.environ[k] = v continue - if ( - k == "OPENAI_API_KEY" - and v - and os.environ.get("OPENAI_API_KEY") - ): + if k in warn_keys and v and os.environ.get(k): click.echo( - "[docgen] OPENAI_API_KEY already set in the process environment; " + f"[docgen] {k} already set in the process environment; " "env_file value is ignored for this key (shell wins). " - "Unset OPENAI_API_KEY or set DOCGEN_ENV_OVERRIDES=1 to load all keys " - "from env_file, or DOCGEN_ENV_OVERRIDES=OPENAI_API_KEY to override just " + f"Unset {k} or set DOCGEN_ENV_OVERRIDES=1 to load all keys " + f"from env_file, or DOCGEN_ENV_OVERRIDES={k} to override just " "this key.", err=True, ) @@ -95,6 +94,18 @@ def _cli_version_string(ctx: click.Context, param: click.Parameter, value: bool) type=click.Path(exists=False), help="Path to docgen.yaml (parents of cwd are searched when omitted).", ) +@click.option( + "--repo", + "repo_spec", + default=None, + envvar="DOCGEN_REPO", + help=( + "Consumer git checkout or clone URL (also DOCGEN_REPO). docgen stays an " + "external tool — nothing is copied into that repo's src/. Looks for " + "docs/demos/docgen.yaml (or any docgen.yaml under the checkout). " + "GitHub org/repo shorthand is accepted. Clones into DOCGEN_REPO_CACHE." + ), +) @click.option( "--version", is_flag=True, @@ -104,20 +115,63 @@ def _cli_version_string(ctx: click.Context, param: click.Parameter, value: bool) help="Show installed docgen version and the recommended pip install line.", ) @click.pass_context -def main(ctx: click.Context, config_path: str | None) -> None: +def main( + ctx: click.Context, + config_path: str | None, + repo_spec: str | None, +) -> None: """docgen — demo generation pipeline (install as an external tool; keep only the bundle in-repo). - Environment: keys already set in the shell are not replaced by ``env_file`` - (see ``DOCGEN_ENV_OVERRIDES``). If no docgen.yaml is found, pass ``--config``. + Point at a consumer project with ``--repo PATH_OR_URL`` (or ``DOCGEN_REPO``); + do not vendor this library into that repo's ``src/``. Environment: keys + already set in the shell are not replaced by ``env_file`` (see + ``DOCGEN_ENV_OVERRIDES``). LLM / TTS / image calls use OpenAI by default; + set ``ai.provider: grok`` or ``DOCGEN_AI_PROVIDER=grok`` plus ``XAI_API_KEY`` + to use xAI. """ ctx.ensure_object(dict) + repo_root = None + if repo_spec: + from docgen.target_repo import TargetRepoError, resolve_repo + + try: + repo_root = resolve_repo(repo_spec) + except TargetRepoError as exc: + raise click.ClickException(str(exc)) from exc + click.echo(f"[docgen] target repo: {repo_root}", err=True) + ctx.obj["repo_root"] = repo_root + + cfg = None try: - cfg = Config.from_yaml(config_path) if config_path else Config.discover() + if config_path: + cfg_path = Path(config_path) + if not cfg_path.is_absolute() and repo_root is not None: + nested = repo_root / cfg_path + if nested.exists(): + cfg_path = nested + cfg = Config.from_yaml(cfg_path) + elif repo_root is not None: + from docgen.target_repo import find_bundle_yaml + + found = find_bundle_yaml(repo_root) + if found is not None: + click.echo(f"[docgen] bundle: {found}", err=True) + cfg = Config.from_yaml(found) + elif ctx.invoked_subcommand != "init": + click.echo( + f"[docgen] No docgen.yaml under {repo_root}; " + "run `docgen --repo … init --defaults` to scaffold " + "docs/demos (docgen is not copied into the consumer src/).", + err=True, + ) + else: + cfg = Config.discover() except FileNotFoundError: cfg = None click.echo( "[docgen] No docgen.yaml found in this directory tree; pass " - "`--config PATH/to/docgen.yaml` or `cd` to your demos bundle directory.", + "`--config PATH/to/docgen.yaml`, `--repo PATH_OR_URL`, or `cd` " + "to your demos bundle directory.", err=True, ) ctx.obj["config"] = cfg @@ -160,15 +214,20 @@ def init( from docgen.init import build_defaults_plan, generate_files, print_summary, run_wizard td = Path(target_dir).resolve() if target_dir else None + repo_root = ctx.obj.get("repo_root") if defaults: plan = build_defaults_plan( td, segments_file=segments_file.resolve() if segments_file else None, + repo_root=Path(repo_root) if repo_root else None, ) else: if segments_file is not None: raise click.ClickException("--segments-file requires --defaults.") - plan = run_wizard(target_dir=td) + plan = run_wizard( + target_dir=td, + repo_root=Path(repo_root) if repo_root else None, + ) created = generate_files(plan) print_summary(plan, created) @@ -318,7 +377,7 @@ def tts(ctx: click.Context, segment: str | None, dry_run: bool) -> None: help=( "Timing engine (default: timestamps.engine in docgen.yaml, local). " "local = offline narration-text alignment via ffmpeg silencedetect; " - "whisper = OpenAI whisper-1 transcription." + "whisper = OpenAI whisper-1 or xAI STT when ai.provider is grok." ), ) @click.pass_context @@ -500,10 +559,10 @@ def narration_generate( revise: bool, revision_notes: str, ) -> None: - """Generate or revise narration ``.md`` from repo sources + owner hints via OpenAI chat. + """Generate or revise narration ``.md`` from repo sources + owner hints via chat completions. Configure ``narration_from_source`` in docgen.yaml (context paths/globs, hints, model). - Requires ``OPENAI_API_KEY`` unless using a future offline stub. + Requires ``OPENAI_API_KEY``, or ``XAI_API_KEY`` with ``ai.provider: grok``. Use ``--segment `` to drive a single segment, or ``--all`` to iterate every id in ``segments.all`` (used by full-reset orchestration). diff --git a/src/docgen/config.py b/src/docgen/config.py index 3e29ee3..8b9235b 100644 --- a/src/docgen/config.py +++ b/src/docgen/config.py @@ -110,6 +110,21 @@ def pipeline_manim_scene_names(self) -> list[str]: def concat_map(self) -> dict[str, list[str]]: return self.raw.get("concat", {}) + # -- AI provider (OpenAI / Grok) ------------------------------------------ + + @property + def ai_config(self) -> dict[str, Any]: + """``ai.provider`` plus optional ``base_url`` / ``api_key_env``. + + ``provider`` is ``openai`` (default) or ``grok`` (xAI). Environment + ``DOCGEN_AI_PROVIDER`` overrides YAML. See :mod:`docgen.ai_client`. + """ + defaults: dict[str, Any] = {"provider": "openai"} + block = self.raw.get("ai") + if isinstance(block, dict): + defaults.update(block) + return defaults + # -- TTS ------------------------------------------------------------------- @property @@ -135,7 +150,7 @@ def timestamps_config(self) -> dict[str, Any]: ``engine: local`` (default) aligns the known narration text against the mp3 offline (ffmpeg silencedetect); ``engine: whisper`` transcribes via - OpenAI whisper-1 (network + API key). + OpenAI whisper-1 (network + API key) or xAI STT when ``ai.provider`` is grok. """ defaults: dict[str, Any] = { "engine": "local", @@ -149,7 +164,7 @@ def timestamps_config(self) -> dict[str, Any]: @property def image_generation_config(self) -> dict[str, Any]: - """Settings for ``docgen image-generate`` (OpenAI Images API). + """Settings for ``docgen image-generate`` (OpenAI Images or xAI Imagine). ``quality`` is passed through only when set (model-specific values, e.g. ``low`` / ``medium`` / ``high`` for gpt-image-1). diff --git a/src/docgen/image_generate.py b/src/docgen/image_generate.py index 1a536b4..4819692 100644 --- a/src/docgen/image_generate.py +++ b/src/docgen/image_generate.py @@ -1,4 +1,4 @@ -"""OpenAI image generation for scene-spec **image elements**. +"""Image generation for scene-spec **image elements** (OpenAI or xAI Imagine). A ``*.scene.yaml`` box may be an image element:: @@ -9,8 +9,8 @@ prompt: "Clean flat diagram of ..." # used by `docgen image-generate` label: architecture # optional Whisper timing anchor -``docgen image-generate`` scans specs, calls the OpenAI Images API for -elements whose asset is missing (or ``--force``), and writes PNG bytes to +``docgen image-generate`` scans specs, calls the Images API (OpenAI or xAI Imagine) +for elements whose asset is missing (or ``--force``), and writes PNG bytes to ``/``. ``docgen manim`` then loads the asset via the ``_image`` helper in ``scenes.py``. """ @@ -50,43 +50,66 @@ def generate_image_bytes( model: str, size: str, quality: str | None = None, + cfg: "Config | None" = None, ) -> bytes: - """Call the OpenAI Images API and return decoded PNG bytes.""" + """Call the Images API (OpenAI or xAI Imagine) and return decoded PNG bytes.""" import openai - client = openai.OpenAI() - kwargs: dict = {"model": model, "prompt": prompt, "size": size, "n": 1} - if quality: - kwargs["quality"] = quality - # dall-e models return URLs unless b64 is requested; gpt-image-1 is b64-only. - if model.startswith("dall-e"): - kwargs["response_format"] = "b64_json" + from docgen.ai_client import ( + fetch_url_bytes, + openai_client, + resolve_ai_settings, + resolve_image_model, + ) + + settings = resolve_ai_settings(cfg) + resolved = resolve_image_model(model, settings) + client = openai_client(cfg) + kwargs: dict = {"model": resolved, "prompt": prompt, "n": 1} + if not settings.is_grok: + kwargs["size"] = size + if quality: + kwargs["quality"] = quality + # dall-e models return URLs unless b64 is requested; gpt-image-1 is b64-only. + if resolved.startswith("dall-e"): + kwargs["response_format"] = "b64_json" try: response = call_with_rate_limit_retries(lambda: client.images.generate(**kwargs)) except openai.AuthenticationError as exc: raise ImageGenerationError( - f"OpenAI rejected OPENAI_API_KEY (authentication failed): {exc}. " - "Set a valid key or use --dry-run to inspect prompts only." + f"{'xAI' if settings.is_grok else 'OpenAI'} rejected {settings.api_key_env} " + f"(authentication failed): {exc}. {settings.auth_help()} " + "Or use --dry-run to inspect prompts only." ) from exc except openai.PermissionDeniedError as exc: raise ImageGenerationError( - f"OpenAI permission denied for image model {model!r}: {exc}. " - "Pick a model your account may use, or set image_generation.model in docgen.yaml." + f"{'xAI' if settings.is_grok else 'OpenAI'} permission denied for image model " + f"{resolved!r}: {exc}. Pick a model your account may use, or set " + "image_generation.model in docgen.yaml." ) from exc except openai.APIConnectionError as exc: raise ImageGenerationError( - f"OpenAI connection error: {exc} — re-run when connectivity is restored." + f"{'xAI' if settings.is_grok else 'OpenAI'} connection error: {exc} — " + "re-run when connectivity is restored." ) from exc data = response.data[0] if response.data else None b64 = getattr(data, "b64_json", None) if data is not None else None - if not b64: - raise ImageGenerationError( - f"OpenAI image response for model {model!r} had no b64_json payload; " - "cannot write the asset." - ) - return base64.b64decode(b64) + if b64: + return base64.b64decode(b64) + url = getattr(data, "url", None) if data is not None else None + if url: + try: + return fetch_url_bytes(str(url)) + except Exception as exc: + raise ImageGenerationError( + f"Image model {resolved!r} returned a URL but download failed: {exc}." + ) from exc + raise ImageGenerationError( + f"Image response for model {resolved!r} had neither b64_json nor url; " + "cannot write the asset." + ) def _resolve_asset_path(cfg: "Config", relpath: str) -> Path: @@ -144,7 +167,9 @@ def generate_images_for_spec( continue fn = image_fn or ( - lambda p: generate_image_bytes(prompt=p, model=model, size=size, quality=quality) + lambda p: generate_image_bytes( + prompt=p, model=model, size=size, quality=quality, cfg=cfg + ) ) data = fn(prompt) out.parent.mkdir(parents=True, exist_ok=True) diff --git a/src/docgen/init.py b/src/docgen/init.py index 91fb145..5373348 100644 --- a/src/docgen/init.py +++ b/src/docgen/init.py @@ -75,12 +75,21 @@ def build_defaults_plan( target_dir: Path | None, *, segments_file: Path | None = None, + repo_root: Path | None = None, ) -> InitPlan: """Non-interactive plan: git root, demo dir, segments from a segments file, - existing narration filenames, or a single starter (in that order).""" + existing narration filenames, or a single starter (in that order). + + ``repo_root`` (from ``docgen --repo``) is the consumer checkout. docgen is + not copied into that tree; only ``docs/demos`` (or ``target_dir``) is written. + """ plan = InitPlan() - git_root = detect_git_root(target_dir) - plan.repo_root = git_root.resolve() if git_root else (target_dir or Path.cwd()).resolve() + if repo_root is not None: + explicit = Path(repo_root).resolve() + plan.repo_root = detect_git_root(explicit) or explicit + else: + git_root = detect_git_root(target_dir) + plan.repo_root = git_root.resolve() if git_root else (target_dir or Path.cwd()).resolve() if target_dir is not None: plan.demo_dir = Path(target_dir).resolve() else: @@ -149,7 +158,10 @@ def scan_existing_assets(demo_dir: Path) -> dict[str, int]: return counts -def run_wizard(target_dir: Path | None = None) -> InitPlan: +def run_wizard( + target_dir: Path | None = None, + repo_root: Path | None = None, +) -> InitPlan: """Interactive wizard that collects project info and returns an InitPlan.""" plan = InitPlan() @@ -158,11 +170,14 @@ def run_wizard(target_dir: Path | None = None) -> InitPlan: click.secho(" " + "=" * 22, fg="cyan") click.echo() - # Detect git root - git_root = detect_git_root(target_dir) + # Detect git root (prefer --repo consumer checkout over cwd / this library). + git_root = detect_git_root(repo_root or target_dir) if git_root: click.echo(f" Git root: {git_root}") - plan.repo_root = git_root or (target_dir or Path.cwd()).resolve() + if repo_root is not None: + plan.repo_root = Path(repo_root).resolve() + else: + plan.repo_root = git_root or (target_dir or Path.cwd()).resolve() # Project name default_name = plan.repo_root.name @@ -175,6 +190,8 @@ def run_wizard(target_dir: Path | None = None) -> InitPlan: default_demo = str(target_dir.resolve()) elif git_root: default_demo = str(git_root / "docs" / "demos") + elif repo_root is not None: + default_demo = str(Path(repo_root).resolve() / "docs" / "demos") else: default_demo = str(Path.cwd() / "docs" / "demos") @@ -199,7 +216,7 @@ def run_wizard(target_dir: Path | None = None) -> InitPlan: click.echo(f" Found .env: {plan.env_file_rel}") else: env_input = click.prompt( - " Path to .env (for OPENAI_API_KEY, blank to skip)", + " Path to .env (for OPENAI_API_KEY / XAI_API_KEY, blank to skip)", default="", type=str, ) if env_input: @@ -315,6 +332,9 @@ def _write_config(plan: InitPlan) -> str: }, "segment_names": segment_names, "visual_map": {}, + "ai": { + "provider": "openai", # openai | grok (xAI; set XAI_API_KEY) + }, "compose": { "ffmpeg_timeout_sec": 300, }, @@ -488,6 +508,19 @@ def _write_bundle_readme(plan: InitPlan) -> str: Verify: ``docgen --version`` + ## API keys + + Default provider is OpenAI (``OPENAI_API_KEY``). To use Grok / xAI instead: + + ```bash + export DOCGEN_AI_PROVIDER=grok + export XAI_API_KEY=... + ``` + + or set ``ai.provider: grok`` in ``docgen.yaml``. Existing OpenAI model + names in YAML are remapped (chat to ``grok-4.6``, images to + ``grok-imagine-image-2.0``, TTS voice ``coral`` to ``eve``). + ## Run from this bundle ```bash diff --git a/src/docgen/manim_scene_support.py b/src/docgen/manim_scene_support.py index 7c87641..f48fc5f 100644 --- a/src/docgen/manim_scene_support.py +++ b/src/docgen/manim_scene_support.py @@ -864,36 +864,24 @@ def extract_reference_classes(scenes_py_text: str, *, max_bytes: int = 30_000) - # ── LLM call (OpenAI chat completions) ───────────────────────────────────── -def call_llm(*, system_prompt: str, user_message: str, model: str, temperature: float) -> str: - """Call OpenAI chat completions; convert auth/network errors into actionable RuntimeErrors.""" - import openai +def call_llm( + *, + system_prompt: str, + user_message: str, + model: str, + temperature: float, + cfg: "Config | None" = None, +) -> str: + """Call chat completions (OpenAI or Grok); convert auth/network errors into RuntimeErrors.""" + from docgen.ai_client import chat_completion - client = openai.OpenAI() - try: - response = client.chat.completions.create( - model=model, - messages=[ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": user_message}, - ], - temperature=float(temperature), - ) - except openai.AuthenticationError as exc: - raise RuntimeError( - f"OpenAI rejected OPENAI_API_KEY (authentication failed): {exc}. " - "Set a valid key or pass --dry-run to inspect the prompt only." - ) from exc - except openai.PermissionDeniedError as exc: - raise RuntimeError( - f"OpenAI permission denied for model {model!r}: {exc}. " - "Pick a model your account is allowed to use, or update YAML " - "manim_scene_generation.model." - ) from exc - except openai.APIConnectionError as exc: - raise RuntimeError( - f"OpenAI connection error: {exc} — re-run when connectivity is restored." - ) from exc - return response.choices[0].message.content or "" + return chat_completion( + system_prompt=system_prompt, + user_message=user_message, + model=model, + temperature=temperature, + cfg=cfg, + ) # ── Lint (compiled / hand-authored scene bodies) ─────────────────────────── diff --git a/src/docgen/narrate_from_source.py b/src/docgen/narrate_from_source.py index 7aeac74..c825c37 100644 --- a/src/docgen/narrate_from_source.py +++ b/src/docgen/narrate_from_source.py @@ -1,4 +1,4 @@ -"""Generate narration ``.md`` from repository sources using OpenAI chat completions. +"""Generate narration ``.md`` from repository sources using chat completions (OpenAI or Grok). The **project owner** defines optional **hints** (plain strings) in ``docgen.yaml`` under ``narration_from_source.hints`` and/or per-segment ``narration_from_source.segments..hints``. @@ -223,7 +223,7 @@ def generate_narration_markdown( revision_notes: str = "", mode: str = "generate", ) -> str: - """Call OpenAI and return markdown body (does not write files). + """Call the chat model and return markdown body (does not write files). Owner hints from YAML and ``extra_hints`` from the caller are sent as guidance only; the returned markdown is model-generated. @@ -272,6 +272,7 @@ def generate_narration_markdown( topic_label=topic, current_narration=current, mode=mode_norm, + cfg=cfg, ) diff --git a/src/docgen/scene_spec_generate.py b/src/docgen/scene_spec_generate.py index 0c9dde7..cdda0d4 100644 --- a/src/docgen/scene_spec_generate.py +++ b/src/docgen/scene_spec_generate.py @@ -181,7 +181,12 @@ def strip_yaml_fences(text: str) -> str: def _invoke_llm( - *, system_prompt: str, user_message: str, model: str, temperature: float + *, + system_prompt: str, + user_message: str, + model: str, + temperature: float, + cfg: "Config | None" = None, ) -> str: from docgen.manim_scene_support import call_llm @@ -191,6 +196,7 @@ def _invoke_llm( user_message=user_message, model=model, temperature=temperature, + cfg=cfg, ) ) @@ -563,7 +569,7 @@ def generate_scene_spec( if temperature_override is not None else float(settings.temperature or DEFAULT_SCENE_SPEC_TEMPERATURE) ) - invoke = llm or _invoke_llm + invoke = llm or (lambda **kw: _invoke_llm(cfg=cfg, **kw)) n_beats = len(cluster_subject_beats(narration_sentences(narration_text))) # Near-miss: allow a couple uncovered beats after retry, not a blind label quota. near_miss_slack = max(1, n_beats // 8) if n_beats else 0 @@ -588,10 +594,13 @@ def generate_scene_spec( temperature=min(0.9, temperature + 0.15 * attempt), ) except RuntimeError as exc: + from docgen.ai_client import resolve_ai_settings + + settings = resolve_ai_settings(cfg) raise SceneGenerationError( - f"OpenAI/chat call failed ({exc}). " - "Check OPENAI_API_KEY, set DOCGEN_ENV_OVERRIDES=1 to load the bundle env_file, " - "or use --dry-run to inspect the prompt only." + f"Chat call failed ({exc}). " + f"{settings.auth_help()} Set DOCGEN_ENV_OVERRIDES=1 to load the bundle " + "env_file, or use --dry-run to inspect the prompt only." ) from exc try: merged_spec = _parse_and_harden_llm_spec( diff --git a/src/docgen/target_repo.py b/src/docgen/target_repo.py new file mode 100644 index 0000000..b8b2bb7 --- /dev/null +++ b/src/docgen/target_repo.py @@ -0,0 +1,180 @@ +"""Resolve a consumer git checkout for docgen without vendoring this library. + +``docgen`` is an external pip tool. The *target* repository should keep only a +bundle (``docs/demos/docgen.yaml`` + hints/narration). This module locates that +bundle from a local path or clone URL so a Cloud / CI environment that has +docgen installed can generate against any consumer repo. +""" + +from __future__ import annotations + +import os +import re +import subprocess +from pathlib import Path + +_GIT_URL_PREFIX = re.compile(r"^(?:https?://|git@|ssh://|git://)", re.I) +_GITHUB_SHORTHAND = re.compile(r"^[A-Za-z0-9_.-]+/[A-Za-z0-9_.-]+$") +_SKIP_DIR_NAMES = frozenset( + { + ".git", + "node_modules", + ".venv", + "venv", + "__pycache__", + "archive", + ".tox", + ".mypy_cache", + "dist", + "build", + } +) + + +class TargetRepoError(RuntimeError): + """Raised when a ``--repo`` spec cannot be resolved or cloned.""" + + +def looks_like_git_url(spec: str) -> bool: + s = spec.strip() + if not s: + return False + if _GIT_URL_PREFIX.match(s) or s.endswith(".git"): + return True + if s.startswith("github.com/"): + return True + if _GITHUB_SHORTHAND.fullmatch(s) and "/" in s: + return True + return False + + +def normalize_git_url(spec: str) -> str: + """Turn ``org/repo`` / ``github.com/org/repo`` into an https clone URL.""" + s = spec.strip() + if s.startswith("github.com/"): + s = "https://" + s + elif _GITHUB_SHORTHAND.fullmatch(s) and not _GIT_URL_PREFIX.match(s): + s = f"https://github.com/{s}" + if s.startswith("https://github.com/") and not s.endswith(".git"): + s = s.rstrip("/") + ".git" + return s + + +def default_cache_dir() -> Path: + override = (os.environ.get("DOCGEN_REPO_CACHE") or "").strip() + if override: + return Path(override).expanduser() + xdg = (os.environ.get("XDG_CACHE_HOME") or "").strip() + if xdg: + return Path(xdg) / "docgen" / "repos" + return Path.home() / ".cache" / "docgen" / "repos" + + +def repo_cache_name(url: str) -> str: + name = url.rstrip("/").split("/")[-1] + if name.endswith(".git"): + name = name[: -len(".git")] + return name or "repo" + + +def find_bundle_yaml(repo_root: Path) -> Path | None: + """Locate ``docgen.yaml`` under a consumer checkout (children, not parents). + + :meth:`Config.discover` walks *up*. A ``--repo`` path is usually the git + root, so the canonical bundle ``docs/demos/docgen.yaml`` would be missed. + """ + root = repo_root.resolve() + preferred = ( + root / "docs" / "demos" / "docgen.yaml", + root / "docgen.yaml", + ) + for path in preferred: + if path.is_file(): + return path + found: list[Path] = [] + if not root.is_dir(): + return None + for path in root.rglob("docgen.yaml"): + try: + rel_parts = path.relative_to(root).parts + except ValueError: + continue + if any(part in _SKIP_DIR_NAMES for part in rel_parts[:-1]): + continue + found.append(path) + if not found: + return None + found.sort( + key=lambda p: ( + 0 if "demos" in p.parts else 1, + len(p.relative_to(root).parts), + str(p), + ) + ) + return found[0] + + +def resolve_repo( + spec: str, + *, + cache_dir: Path | None = None, + clone: bool = True, +) -> Path: + """Return a local checkout for ``spec`` (existing path or git URL).""" + raw = (spec or "").strip() + if not raw: + raise TargetRepoError("empty --repo spec") + local = Path(raw).expanduser() + if local.exists(): + return local.resolve() + if not looks_like_git_url(raw): + raise TargetRepoError( + f"repo path does not exist: {local}. Pass a local checkout or a " + "git URL / GitHub org/repo." + ) + if not clone: + raise TargetRepoError(f"repo is a git URL but clone is disabled: {raw}") + url = normalize_git_url(raw) + dest = (cache_dir or default_cache_dir()) / repo_cache_name(url) + if (dest / ".git").exists(): + return dest.resolve() + if dest.exists() and any(dest.iterdir()): + raise TargetRepoError( + f"clone destination {dest} exists and is not a git checkout; " + "set DOCGEN_REPO_CACHE or remove the directory." + ) + clone_git_repo(url, dest) + return dest.resolve() + + +def clone_git_repo(url: str, dest: Path) -> None: + """Shallow-clone ``url`` into ``dest`` (uses GITHUB_TOKEN / GH_TOKEN when set).""" + dest.parent.mkdir(parents=True, exist_ok=True) + cmd = ["git", "clone", "--depth", "1", url, str(dest)] + token = (os.environ.get("GITHUB_TOKEN") or os.environ.get("GH_TOKEN") or "").strip() + if token and "github.com" in url and url.startswith("https://"): + cmd = [ + "git", + "-c", + f"url.https://x-access-token:{token}@github.com/.insteadOf=https://github.com/", + "clone", + "--depth", + "1", + url, + str(dest), + ] + try: + subprocess.run( + cmd, + check=True, + capture_output=True, + text=True, + timeout=300, + ) + except FileNotFoundError as exc: + raise TargetRepoError("git is not installed; cannot clone --repo") from exc + except subprocess.CalledProcessError as exc: + err = (exc.stderr or exc.stdout or "").strip() + raise TargetRepoError(f"git clone failed for {url}: {err or exc}") from exc + except subprocess.TimeoutExpired as exc: + raise TargetRepoError(f"git clone timed out for {url}") from exc diff --git a/src/docgen/timestamps.py b/src/docgen/timestamps.py index 3b20585..e4c2088 100644 --- a/src/docgen/timestamps.py +++ b/src/docgen/timestamps.py @@ -6,8 +6,8 @@ * **local** (default) — offline alignment of the known narration text against the mp3 using ffmpeg ``silencedetect`` + proportional interpolation (:mod:`docgen.align`). No API calls; requires ``narration/.md``. -* **whisper** — OpenAI ``whisper-1`` transcription (legacy; requires - ``OPENAI_API_KEY`` and network). +* **whisper** — network transcription (legacy). OpenAI ``whisper-1``, or + xAI ``/v1/stt`` when ``ai.provider`` is ``grok``. Requires an API key. Select via ``timestamps.engine`` in docgen.yaml or ``docgen timestamps --engine``. """ @@ -31,29 +31,10 @@ def __init__(self, config: Config) -> None: # ── Whisper engine (network) ───────────────────────────────────── def extract(self, audio_path: str | Path) -> dict[str, Any]: - """Transcribe audio via OpenAI whisper-1 and return word-level timestamps.""" - import openai - - client = openai.OpenAI() - with open(audio_path, "rb") as f: - result = client.audio.transcriptions.create( - model="whisper-1", - file=f, - response_format="verbose_json", - timestamp_granularities=["word", "segment"], - ) + """Transcribe audio and return word-level timestamps (OpenAI Whisper or xAI STT).""" + from docgen.ai_client import transcribe_audio - return { - "text": result.text, - "segments": [ - {"start": s.start, "end": s.end, "text": s.text} - for s in (result.segments or []) - ], - "words": [ - {"start": w.start, "end": w.end, "word": w.word} - for w in (result.words or []) - ], - } + return transcribe_audio(audio_path, cfg=self.config) # ── Local engine (offline alignment) ───────────────────────────── diff --git a/src/docgen/tts.py b/src/docgen/tts.py index fd84bb5..66b686b 100644 --- a/src/docgen/tts.py +++ b/src/docgen/tts.py @@ -1,4 +1,4 @@ -"""TTS narration generator using OpenAI gpt-4o-mini-tts.""" +"""TTS narration generator (OpenAI gpt-4o-mini-tts or xAI / Grok ``/v1/tts``).""" from __future__ import annotations @@ -9,7 +9,6 @@ if TYPE_CHECKING: from docgen.config import Config -from docgen.openai_retry import call_with_rate_limit_retries def _probe_duration(path: Path) -> float | None: @@ -83,7 +82,7 @@ def _generate_one(self, seg_id: str, dry_run: bool) -> None: print(f" ... ({len(plain) - 500} more chars)") return - import openai + from docgen.ai_client import synthesize_speech audio_dir.mkdir(parents=True, exist_ok=True) @@ -94,18 +93,14 @@ def _generate_one(self, seg_id: str, dry_run: bool) -> None: print(f"[tts] Generating audio for {seg_id} ({len(plain)} chars) -> {out_path}") - client = openai.OpenAI() - - def _call() -> None: - response = client.audio.speech.create( - model=self.config.tts_model, - voice=self.config.tts_voice, - input=plain, - instructions=self.config.tts_instructions, - ) - response.stream_to_file(str(out_path)) - - call_with_rate_limit_retries(_call) + synthesize_speech( + text=plain, + model=self.config.tts_model, + voice=self.config.tts_voice, + instructions=self.config.tts_instructions, + output_path=out_path, + cfg=self.config, + ) print(f"[tts] Wrote {out_path}") new_duration = _probe_duration(out_path) diff --git a/src/docgen/wizard.py b/src/docgen/wizard.py index 4e509e9..93e1210 100644 --- a/src/docgen/wizard.py +++ b/src/docgen/wizard.py @@ -231,8 +231,9 @@ def generate_narration_via_llm( topic_label: str | None = None, current_narration: str = "", mode: str = "generate", + cfg: Any | None = None, ) -> str: - """Call OpenAI to generate or revise a narration draft. + """Call chat completions (OpenAI or Grok) to generate or revise a narration draft. ``guidance`` is **caller-supplied** (e.g. project-owner hints from ``docgen.yaml``), not text returned from a prior model call. ``topic_label`` is a human-facing focus @@ -243,8 +244,6 @@ def generate_narration_via_llm( both non-empty) edits the existing script in place: address feedback, keep structure/phrasing that still work, do not rewrite from scratch unless needed. """ - import openai - focus = (topic_label or "").strip() or _strip_segment_prefix(segment_name) notes = (revision_notes or "").strip() current = (current_narration or "").strip() @@ -322,16 +321,15 @@ def generate_narration_via_llm( ] sys_prompt = system_prompt - client = openai.OpenAI() - response = client.chat.completions.create( + from docgen.ai_client import chat_completion + + return chat_completion( + system_prompt=sys_prompt, + user_message="\n".join(user_parts), model=model, - messages=[ - {"role": "system", "content": sys_prompt}, - {"role": "user", "content": "\n".join(user_parts)}, - ], temperature=float(temperature), + cfg=cfg, ) - return response.choices[0].message.content or "" # --------------------------------------------------------------------------- @@ -578,6 +576,7 @@ def api_generate_narration(): topic_label=topic_label, current_narration=current_narration, mode=mode, + cfg=cfg, ) except ValueError as exc: return jsonify({"error": str(exc)}), 400 diff --git a/src/docgen/yaml_generate.py b/src/docgen/yaml_generate.py index 8f66512..8ca5cba 100644 --- a/src/docgen/yaml_generate.py +++ b/src/docgen/yaml_generate.py @@ -1,4 +1,4 @@ -"""Hybrid ``docgen yaml-generate``: structural defaults from disk + OpenAI prose blocks. +"""Hybrid ``docgen yaml-generate``: structural defaults from disk + LLM prose blocks. This does **not** replace an entire ``docgen.yaml`` in one shot. It: @@ -19,7 +19,7 @@ ``project-context.md``) merge ``env_file``, top-level ``narration_from_source`` / ``manim_scene_generation`` blocks, optional ``concat``, and ``discovery`` (see ``merge_hint_project``). -7. Optionally calls **OpenAI** to draft ``tts.instructions`` and ``wizard.system_prompt``. +7. Optionally calls **OpenAI or Grok** to draft ``tts.instructions`` and ``wizard.system_prompt``. Writing the file uses PyYAML: **YAML comments and key order in the original file are not preserved.** Prefer version control for review; keep hand-maintained prose in Git @@ -105,6 +105,16 @@ def merge_defaults( ex.append(ARCHIVE_EXCLUDE) changes.append(f"wizard.exclude_patterns: added {ARCHIVE_EXCLUDE!r}") + ai = raw.get("ai") + if ai is None: + raw["ai"] = {"provider": "openai"} + changes.append( + "ai: added provider openai (set grok / DOCGEN_AI_PROVIDER=grok to use xAI)" + ) + elif isinstance(ai, dict) and not str(ai.get("provider") or "").strip(): + ai["provider"] = "openai" + changes.append("ai.provider: defaulted to openai") + nf_existing = raw.get("narration_from_source") if nf_existing is None: ctx_paths: list[str] = [] @@ -750,8 +760,9 @@ def _llm_yaml_hints_json( existing_tts: str, existing_wizard: str, model: str, + cfg: "Config | None" = None, ) -> dict[str, str]: - import openai + from docgen.ai_client import chat_completion parts = [ f"Project: {project_label}", @@ -768,23 +779,13 @@ def _llm_yaml_hints_json( parts.append(f"FILE: {label}\n```\n{body}\n```") user = "\n".join(parts) - client = openai.OpenAI() - try: - resp = client.chat.completions.create( - model=model, - messages=[ - {"role": "system", "content": DEFAULT_SYSTEM_PROMPT}, - {"role": "user", "content": user}, - ], - temperature=0.35, - ) - except openai.AuthenticationError as exc: - raise RuntimeError( - f"OpenAI rejected OPENAI_API_KEY: {exc}. Set a valid key or omit --llm." - ) from exc - except openai.APIConnectionError as exc: - raise RuntimeError(f"OpenAI connection error: {exc}") from exc - text = (resp.choices[0].message.content or "").strip() + text = (chat_completion( + system_prompt=DEFAULT_SYSTEM_PROMPT, + user_message=user, + model=model, + temperature=0.35, + cfg=cfg, + ) or "").strip() # tolerate markdown code fence m = re.match(r"^```(?:json)?\s*\n(?P[\s\S]*?)\n```\s*$", text) if m: @@ -800,7 +801,7 @@ def _llm_yaml_hints_json( def generate_llm_hints(cfg: "Config", *, model: str | None = None) -> dict[str, str]: - """Call OpenAI; return ``tts_instructions`` and ``wizard_system_prompt``.""" + """Call chat completions; return ``tts_instructions`` and ``wizard_system_prompt``.""" raw = cfg.raw tts = raw.get("tts") if isinstance(raw.get("tts"), dict) else {} wiz = raw.get("wizard") if isinstance(raw.get("wizard"), dict) else {} @@ -818,6 +819,7 @@ def generate_llm_hints(cfg: "Config", *, model: str | None = None) -> dict[str, existing_tts=existing_tts, existing_wizard=existing_wizard, model=model or DEFAULT_LLM_MODEL, + cfg=cfg, ) diff --git a/tests/conftest.py b/tests/conftest.py index 5307c58..606005c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,6 +6,24 @@ from __future__ import annotations +import pytest + +_CLEAR_ENV = ( + "DOCGEN_REPO", + "DOCGEN_AI_PROVIDER", + "DOCGEN_AI_BASE_URL", + "DOCGEN_AI_API_KEY_ENV", + "DOCGEN_REPO_CACHE", +) + + +@pytest.fixture(autouse=True) +def _clear_docgen_override_env(monkeypatch: pytest.MonkeyPatch) -> None: + """Keep provider/repo env from leaking into CLI and AI-client tests.""" + for key in _CLEAR_ENV: + monkeypatch.delenv(key, raising=False) + + _FFMPEG_ONLY_VALIDATE_TESTS = frozenset( { "tests/test_validate.py::TestComposeGuard::test_compose_rejects_short_video", diff --git a/tests/test_ai_client.py b/tests/test_ai_client.py new file mode 100644 index 0000000..7c2083f --- /dev/null +++ b/tests/test_ai_client.py @@ -0,0 +1,192 @@ +"""Tests for :mod:`docgen.ai_client` (OpenAI vs Grok resolution, no live network).""" + +from __future__ import annotations + +import json +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +import yaml + +from docgen.ai_client import ( + DEFAULT_GROK_CHAT_MODEL, + DEFAULT_GROK_IMAGE_MODEL, + GROK_BASE_URL, + chat_completion, + openai_client, + resolve_ai_settings, + resolve_chat_model, + resolve_image_model, + resolve_tts_voice, + synthesize_speech, + transcribe_audio, +) +from docgen.config import Config + + +def _cfg(tmp_path: Path, raw: dict) -> Config: + p = tmp_path / "docgen.yaml" + p.write_text(yaml.dump(raw), encoding="utf-8") + return Config.from_yaml(p) + + +def test_default_provider_is_openai(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DOCGEN_AI_PROVIDER", raising=False) + monkeypatch.delenv("XAI_API_KEY", raising=False) + cfg = _cfg(tmp_path, {}) + st = resolve_ai_settings(cfg) + assert st.provider == "openai" + assert st.base_url is None + assert st.api_key_env == "OPENAI_API_KEY" + + +def test_yaml_provider_grok(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DOCGEN_AI_PROVIDER", raising=False) + monkeypatch.setenv("XAI_API_KEY", "xai-test") + cfg = _cfg(tmp_path, {"ai": {"provider": "grok"}}) + st = resolve_ai_settings(cfg) + assert st.provider == "grok" + assert st.base_url == GROK_BASE_URL + assert st.api_key == "xai-test" + assert st.api_key_env == "XAI_API_KEY" + + +def test_env_overrides_yaml_provider(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DOCGEN_AI_PROVIDER", "grok") + monkeypatch.setenv("XAI_API_KEY", "xai-test") + cfg = _cfg(tmp_path, {"ai": {"provider": "openai"}}) + assert resolve_ai_settings(cfg).provider == "grok" + + +def test_xai_alias_and_openai_model_remap( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DOCGEN_AI_PROVIDER", "xai") + cfg = _cfg(tmp_path, {}) + st = resolve_ai_settings(cfg) + assert st.is_grok + assert resolve_chat_model("gpt-4o", st) == DEFAULT_GROK_CHAT_MODEL + assert resolve_chat_model("gpt-4o-mini", st) == DEFAULT_GROK_CHAT_MODEL + assert resolve_chat_model("grok-4.6", st) == "grok-4.6" + assert resolve_image_model("gpt-image-1", st) == DEFAULT_GROK_IMAGE_MODEL + assert resolve_tts_voice("coral", st) == "eve" + assert resolve_tts_voice("eve", st) == "eve" + + +def test_openai_client_grok_passes_base_url( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DOCGEN_AI_PROVIDER", "grok") + monkeypatch.setenv("XAI_API_KEY", "xai-test") + cfg = _cfg(tmp_path, {}) + captured: dict = {} + + def _fake_openai(**kwargs): # noqa: ANN003 + captured.update(kwargs) + return MagicMock() + + with patch("openai.OpenAI", side_effect=_fake_openai): + openai_client(cfg) + assert captured["api_key"] == "xai-test" + assert captured["base_url"] == GROK_BASE_URL + + +def test_openai_client_default_no_kwargs(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("DOCGEN_AI_PROVIDER", raising=False) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.delenv("XAI_API_KEY", raising=False) + with patch("openai.OpenAI") as m: + m.return_value = MagicMock() + openai_client(None) + m.assert_called_once_with() + + +def test_chat_completion_uses_remapped_model( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("DOCGEN_AI_PROVIDER", "grok") + monkeypatch.setenv("XAI_API_KEY", "xai-test") + cfg = _cfg(tmp_path, {}) + captured: dict = {} + + class _Msg: + content = "ok" + + class _Choice: + message = _Msg() + + class _Resp: + choices = [_Choice()] + + fake = MagicMock() + fake.chat.completions.create.side_effect = lambda **kw: captured.update(kw) or _Resp() + + with patch("docgen.ai_client.openai_client", return_value=fake): + out = chat_completion( + system_prompt="sys", + user_message="user", + model="gpt-4o", + temperature=0.2, + cfg=cfg, + ) + assert out == "ok" + assert captured["model"] == DEFAULT_GROK_CHAT_MODEL + + +def test_grok_tts_posts_json(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DOCGEN_AI_PROVIDER", "grok") + monkeypatch.setenv("XAI_API_KEY", "xai-test") + cfg = _cfg(tmp_path, {"tts": {"voice": "coral", "language": "en"}}) + out = tmp_path / "n.mp3" + + def _http(url: str, *, data: bytes, headers: dict) -> bytes: + assert url.endswith("/tts") + payload = json.loads(data.decode()) + assert payload["voice_id"] == "eve" + assert payload["text"] == "Hello" + assert "Bearer xai-test" in headers["Authorization"] + return b"ID3fake" + + with patch("docgen.ai_client._http_with_retries", side_effect=_http): + synthesize_speech( + text="Hello", + model="gpt-4o-mini-tts", + voice="coral", + instructions="unused on grok", + output_path=out, + cfg=cfg, + ) + assert out.read_bytes() == b"ID3fake" + + +def test_grok_stt_maps_words(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("DOCGEN_AI_PROVIDER", "grok") + monkeypatch.setenv("XAI_API_KEY", "xai-test") + cfg = _cfg(tmp_path, {}) + mp3 = tmp_path / "n.mp3" + mp3.write_bytes(b"fake-mp3") + body = json.dumps( + { + "text": "Hello world", + "duration": 1.2, + "words": [ + {"text": "Hello", "start": 0.0, "end": 0.4}, + {"text": "world", "start": 0.4, "end": 1.0}, + ], + } + ).encode() + + with patch("docgen.ai_client._http_with_retries", return_value=body): + result = transcribe_audio(mp3, cfg=cfg) + assert result["text"] == "Hello world" + assert result["words"][0]["word"] == "Hello" + assert result["words"][1]["end"] == 1.0 + assert result["segments"][0]["text"] == "Hello world" + + +def test_unknown_provider_raises() -> None: + with pytest.raises(ValueError, match="Unknown AI provider"): + from docgen.ai_client import normalize_provider + + normalize_provider("ollama") diff --git a/tests/test_asset_graph.py b/tests/test_asset_graph.py index 97c4bea..2a77a12 100644 --- a/tests/test_asset_graph.py +++ b/tests/test_asset_graph.py @@ -221,7 +221,7 @@ class _FakeChat: class _FakeClient: chat = _FakeChat() - monkeypatch.setattr("openai.OpenAI", lambda: _FakeClient()) + monkeypatch.setattr("openai.OpenAI", lambda **kwargs: _FakeClient()) out = generate_narration_via_llm( source_texts=["## File: a.md\nsource"], diff --git a/tests/test_cli_env.py b/tests/test_cli_env.py index d266db6..538a27d 100644 --- a/tests/test_cli_env.py +++ b/tests/test_cli_env.py @@ -41,6 +41,19 @@ def test_load_env_warns_when_openai_in_shell_and_env_file(tmp_path, monkeypatch, assert os.environ["OPENAI_API_KEY"] == "sk-from-shell" +def test_load_env_warns_when_xai_in_shell_and_env_file(tmp_path, monkeypatch, capsys) -> None: + (tmp_path / ".env").write_text("XAI_API_KEY=xai-from-file\n", encoding="utf-8") + cfg = _minimal_cfg(tmp_path, env_file=".env") + monkeypatch.setenv("XAI_API_KEY", "xai-from-shell") + monkeypatch.delenv("DOCGEN_ENV_OVERRIDES", raising=False) + + cli._load_env(cfg) + + err = capsys.readouterr().err + assert "XAI_API_KEY already set" in err + assert os.environ["XAI_API_KEY"] == "xai-from-shell" + + def test_load_env_docgen_env_overrides_all(tmp_path, monkeypatch) -> None: (tmp_path / ".env").write_text("OPENAI_API_KEY=sk-from-file\nOTHER=x\n", encoding="utf-8") cfg = _minimal_cfg(tmp_path, env_file=".env") diff --git a/tests/test_config.py b/tests/test_config.py index a854637..0b27f9b 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -115,6 +115,16 @@ def test_effective_max_freeze_ratio_uses_global(tmp_path): assert c.effective_max_freeze_ratio(None) == 0.4 +def test_ai_config_defaults_and_override(tmp_path): + p = tmp_path / "docgen.yaml" + p.write_text("{}", encoding="utf-8") + c = Config.from_yaml(p) + assert c.ai_config["provider"] == "openai" + p.write_text("ai: {provider: grok}\n", encoding="utf-8") + c = Config.from_yaml(p) + assert c.ai_config["provider"] == "grok" + + def test_minimal_config(tmp_path): c = Config.minimal(tmp_path) assert c.base_dir == tmp_path.resolve() diff --git a/tests/test_init.py b/tests/test_init.py index 9d61bed..b8f80e5 100644 --- a/tests/test_init.py +++ b/tests/test_init.py @@ -105,6 +105,7 @@ def test_generate_files_minimal(tmp_path: Path) -> None: cfg = yaml.safe_load(cfg_text.split("\n\n", 1)[-1]) assert cfg["segments"]["all"] == ["01", "02"] assert cfg["segment_names"]["01"] == "01-intro" + assert cfg["ai"]["provider"] == "openai" assert "manim" not in cfg assert "vhs" not in cfg assert "test-project" in cfg["tts"]["instructions"] diff --git a/tests/test_target_repo.py b/tests/test_target_repo.py new file mode 100644 index 0000000..629f166 --- /dev/null +++ b/tests/test_target_repo.py @@ -0,0 +1,151 @@ +"""Tests for :mod:`docgen.target_repo` and ``docgen --repo``.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import patch + +import pytest +import yaml +from click.testing import CliRunner + +from docgen.cli import main +from docgen.init import build_defaults_plan +from docgen.target_repo import ( + TargetRepoError, + find_bundle_yaml, + looks_like_git_url, + normalize_git_url, + repo_cache_name, + resolve_repo, +) + + +def test_looks_like_git_url() -> None: + assert looks_like_git_url("https://github.com/acme/app.git") + assert looks_like_git_url("git@github.com:acme/app.git") + assert looks_like_git_url("github.com/acme/app") + assert looks_like_git_url("acme/app") + assert not looks_like_git_url("/tmp/local-checkout") + assert not looks_like_git_url("") + + +def test_normalize_git_url() -> None: + assert normalize_git_url("acme/app") == "https://github.com/acme/app.git" + assert normalize_git_url("github.com/acme/app") == "https://github.com/acme/app.git" + assert ( + normalize_git_url("https://github.com/acme/app") + == "https://github.com/acme/app.git" + ) + + +def test_repo_cache_name() -> None: + assert repo_cache_name("https://github.com/acme/app.git") == "app" + + +def test_resolve_repo_local_path(tmp_path: Path) -> None: + (tmp_path / "README.md").write_text("x", encoding="utf-8") + assert resolve_repo(str(tmp_path)) == tmp_path.resolve() + + +def test_resolve_repo_missing_local_raises(tmp_path: Path) -> None: + with pytest.raises(TargetRepoError, match="does not exist"): + resolve_repo(str(tmp_path / "nope")) + + +def test_resolve_repo_reuses_existing_clone(tmp_path: Path) -> None: + dest = tmp_path / "cache" / "app" + dest.mkdir(parents=True) + (dest / ".git").mkdir() + out = resolve_repo("https://github.com/acme/app.git", cache_dir=tmp_path / "cache") + assert out == dest.resolve() + + +def test_resolve_repo_clones_when_missing(tmp_path: Path) -> None: + cache = tmp_path / "cache" + + def _fake_clone(url: str, dest: Path) -> None: + dest.mkdir(parents=True) + (dest / ".git").mkdir() + (dest / "README.md").write_text(url, encoding="utf-8") + + with patch("docgen.target_repo.clone_git_repo", side_effect=_fake_clone): + out = resolve_repo("acme/app", cache_dir=cache) + assert out == (cache / "app").resolve() + assert (out / "README.md").read_text(encoding="utf-8").endswith("acme/app.git") + + +def test_find_bundle_yaml_prefers_docs_demos(tmp_path: Path) -> None: + (tmp_path / "docs" / "demos").mkdir(parents=True) + canonical = tmp_path / "docs" / "demos" / "docgen.yaml" + canonical.write_text("segments: {}\n", encoding="utf-8") + (tmp_path / "docgen.yaml").write_text("wrong: true\n", encoding="utf-8") + assert find_bundle_yaml(tmp_path) == canonical + + +def test_find_bundle_yaml_skips_venv(tmp_path: Path) -> None: + nested = tmp_path / ".venv" / "lib" / "docgen.yaml" + nested.parent.mkdir(parents=True) + nested.write_text("nope: true\n", encoding="utf-8") + assert find_bundle_yaml(tmp_path) is None + + +def test_build_defaults_plan_repo_root_does_not_use_cwd_library( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """``--repo`` must scaffold the *consumer*, not this library's checkout.""" + library = tmp_path / "documentation-generator" + consumer = tmp_path / "course-builder" + library.mkdir() + (library / ".git").mkdir() + consumer.mkdir() + (consumer / ".git").mkdir() + monkeypatch.chdir(library) + plan = build_defaults_plan(target_dir=None, repo_root=consumer) + assert plan.repo_root == consumer.resolve() + assert plan.demo_dir == (consumer / "docs" / "demos").resolve() + + +def _consumer_bundle(root: Path) -> Path: + (root / ".git").mkdir() + bundle = root / "docs" / "demos" + bundle.mkdir(parents=True) + (bundle / "narration").mkdir() + (bundle / "narration" / "01-intro.md").write_text("Hello there.\n", encoding="utf-8") + raw = { + "repo_root": "../..", + "dirs": {"narration": "narration"}, + "segments": {"default": ["01"], "all": ["01"]}, + "segment_names": {"01": "01-intro"}, + "visual_map": {"01": {"type": "still", "source": "01.mp4"}}, + } + (bundle / "docgen.yaml").write_text(yaml.dump(raw), encoding="utf-8") + return bundle + + +def test_cli_repo_runs_against_consumer_without_config_flag(tmp_path: Path) -> None: + consumer = tmp_path / "app" + consumer.mkdir() + _consumer_bundle(consumer) + runner = CliRunner() + result = runner.invoke(main, ["--repo", str(consumer), "lint", "--segment", "01"]) + assert result.exit_code == 0, result.output + assert "target repo" in result.output + assert "bundle:" in result.output + assert str(consumer / "src") not in result.output + + +def test_cli_repo_init_defaults_writes_bundle_only(tmp_path: Path) -> None: + consumer = tmp_path / "app" + consumer.mkdir() + (consumer / ".git").mkdir() + runner = CliRunner() + result = runner.invoke(main, ["--repo", str(consumer), "init", "--defaults"]) + assert result.exit_code == 0, result.output + yaml_path = consumer / "docs" / "demos" / "docgen.yaml" + assert yaml_path.is_file() + assert not (consumer / "src").exists() + cfg = yaml.safe_load(yaml_path.read_text(encoding="utf-8").split("\n\n", 1)[-1]) + assert cfg["repo_root"] in ("../..", "..\\..") + assert cfg["ai"]["provider"] == "openai" + assert (consumer / "docs" / "demos" / "requirements-docgen.txt").is_file() diff --git a/tests/test_yaml_generate.py b/tests/test_yaml_generate.py index 9ee17e1..dfeed28 100644 --- a/tests/test_yaml_generate.py +++ b/tests/test_yaml_generate.py @@ -74,6 +74,7 @@ def test_merge_defaults_adds_archive_exclude(tmp_path: Path) -> None: ch = merge_defaults(raw, cfg) assert any("archive" in c for c in ch) assert "**/archive/**" in raw["wizard"]["exclude_patterns"] + assert raw["ai"]["provider"] == "openai" def test_merge_defaults_idempotent_archive(tmp_path: Path) -> None: