notebook-ta 0.1.0__py3-none-any.whl

notebook_ta/__init__.py

@@ -0,0 +1,226 @@
+"""notebook-ta public API."""
+
+from __future__ import annotations
+
+import contextlib
+from pathlib import Path
+from typing import Any, cast
+
+import nest_asyncio
+
+from notebook_ta.config.loader import load_exercises, load_global
+from notebook_ta.config.models import ConfigurationError, GlobalConfig, LLMConfig
+from notebook_ta.exercise.definition import Exercise
+from notebook_ta.exercise.registry import ExerciseRegistry
+from notebook_ta.llm.base import LLMProvider, create_provider
+from notebook_ta.logging import get_logger, setup_logging
+from notebook_ta.notebook.extractor import detect_notebook_path, extract_statements
+from notebook_ta.notebook.magic import load_ipython_extension
+from notebook_ta.notebook.session import SessionState
+
+# Module-level singletons updated on each call to load()
+_registry: ExerciseRegistry = ExerciseRegistry()
+_llm_provider: LLMProvider | None = None
+_global_config: GlobalConfig | None = None
+
+_log = get_logger("init")
+
+
+def load(
+    global_config: str | Path,
+    exercises_config: str | Path,
+    *,
+    notebook_path: str | Path | None = None,
+    llm_overrides: dict[str, Any] | None = None,
+    debug: bool = False,
+) -> None:
+    """Load configuration files, register exercises, run auto-setup if needed,
+    and register the %%notebook_ta IPython magic.
+
+    Must be called from within a Jupyter notebook cell.
+
+    Args:
+        global_config: Path or URL to the global configuration TOML.
+        exercises_config: Path or URL to the exercises TOML.
+        notebook_path: Optional path to the ``.ipynb`` file.  Required only
+            when one or more exercises omit ``statement`` from the TOML and
+            the statement should instead be extracted from the notebook's
+            markdown cells (``<div id="<exercise_id>">…</div>`` pattern).  If
+            not provided the system will try to detect the notebook path
+            automatically; pass this argument explicitly when auto-detection
+            fails.
+        llm_overrides: Optional dict of LLM settings that override the values
+            from *global_config*. Valid keys mirror :class:`LLMConfig` fields
+            (e.g. ``model``, ``base_url``, ``provider``, ``timeout``).
+        debug: When ``True``, enable DEBUG-level logging to the terminal and
+               display the final LLM prompt in the notebook output as a
+               collapsible widget before each LLM call.  Defaults to ``False``.
+    """
+    global _registry, _llm_provider, _global_config
+
+    # Allow asyncio event loop nesting required in Jupyter environments
+    nest_asyncio.apply()
+
+    # Configure logging before anything else so all subsequent calls are captured.
+    setup_logging(debug=debug)
+
+    # 1. Load and validate configuration
+    _log.info("Loading notebook-ta configuration")
+    cfg = load_global(global_config)
+    exercise_configs = load_exercises(exercises_config)
+
+    # 1b. Apply programmatic LLM overrides (validated through Pydantic)
+    if llm_overrides:
+        try:
+            updated_llm = LLMConfig.model_validate({**cfg.llm.model_dump(), **llm_overrides})
+        except Exception as exc:
+            raise ConfigurationError(f"Invalid LLM override: {exc}") from exc
+        cfg = cfg.model_copy(update={"llm": updated_llm})
+
+    # 2. Auto-setup wizard
+    if cfg.llm.model == "auto":
+        _run_setup_wizard(cfg)
+
+    # 3. Create LLM provider
+    provider = create_provider(cfg.llm)
+    _log.info("LLM provider created: %r (model=%r)", cfg.llm.provider, cfg.llm.model)
+
+    # 4. Resolve missing statements from the notebook file
+    missing = [ex for ex in exercise_configs if ex.statement is None]
+    if missing:
+        nb_path = Path(notebook_path) if notebook_path is not None else _resolve_notebook_path()
+        _log.info("Extracting %d statement(s) from notebook: %s", len(missing), nb_path)
+        extracted = extract_statements(nb_path)
+        for ex_cfg in missing:
+            text = extracted.get(ex_cfg.id)
+            if text is None:
+                raise ConfigurationError(
+                    f"Exercise {ex_cfg.id!r} has no statement in the TOML and no "
+                    f"<div id={ex_cfg.id!r}> was found in {nb_path}."
+                )
+            ex_cfg.statement = text
+
+    # 5. Build and populate the registry
+    registry = ExerciseRegistry()
+    for ex_cfg in exercise_configs:
+        registry.register(Exercise(config=ex_cfg, global_config=cfg))
+
+    session = SessionState(hint_history_length=cfg.prompts.hint_history_length)
+    _log.info("Registry populated: %d exercise(s)", len(exercise_configs))
+
+    # 6. Register the IPython magic
+    try:
+        from IPython import get_ipython  # type: ignore[attr-defined]
+
+        ip = get_ipython()  # type: ignore[no-untyped-call]
+        if ip is not None:
+            load_ipython_extension(
+                ip, registry=registry, llm_provider=provider, session=session, debug=debug
+            )
+        else:
+            import warnings
+
+            warnings.warn(
+                "notebook_ta.load() was called outside of an IPython/Jupyter environment. "
+                "The %%notebook_ta magic will not be registered.",
+                stacklevel=2,
+            )
+    except ImportError:
+        import warnings
+
+        warnings.warn(
+            "IPython is not installed. The %%notebook_ta magic will not be registered.",
+            stacklevel=2,
+        )
+
+    # Update module singletons
+    _registry = registry
+    _llm_provider = provider
+    _global_config = cfg
+    _log.info("notebook-ta loaded successfully")
+
+    from IPython import display as ipydisplay
+
+    with contextlib.suppress(Exception):
+        cast(Any, ipydisplay.display)(
+            cast(Any, ipydisplay.Markdown)(
+                "✅ **notebook-ta loaded.**  "
+                f"Provider: `{cfg.llm.provider}` — Model: `{cfg.llm.model}`  \n"
+                f"{len(exercise_configs)} exercise(s) registered."
+            )
+        )
+
+
+def get_registry() -> ExerciseRegistry:
+    """Return the active ExerciseRegistry for introspection."""
+    return _registry
+
+
+# ---------------------------------------------------------------------------
+# Internal helpers
+# ---------------------------------------------------------------------------
+
+
+def _resolve_notebook_path() -> Path:
+    """Auto-detect the current notebook path or raise ConfigurationError."""
+    try:
+        from IPython import get_ipython as _gip  # type: ignore[attr-defined]
+
+        ip = _gip()  # type: ignore[no-untyped-call]
+    except ImportError:
+        ip = None
+
+    path = detect_notebook_path(ip)
+    if path is None:
+        raise ConfigurationError(
+            "Could not automatically determine the notebook path. "
+            "Pass notebook_path= to notebook_ta.load() explicitly, e.g.:\n"
+            "  notebook_ta.load(global_cfg, exercises_cfg, "
+            "notebook_path='my_notebook.ipynb')"
+        )
+    return path
+
+
+def _run_setup_wizard(cfg: GlobalConfig) -> None:
+    """Run hardware detection and update cfg.llm.model in place."""
+    from notebook_ta.setup_wizard.detector import detect_hardware, select_model
+
+    profile = detect_hardware()
+    model_spec = select_model(cfg.llm.available_models, profile)
+
+    try:
+        from IPython import display as ipydisplay
+
+        if model_spec is not None:
+            cfg.llm.model = model_spec.name
+            gpu_text = (
+                f", GPU: {profile.gpu_name} ({profile.vram_gb:.1f} GB VRAM)"
+                if profile.gpu_name
+                else ""
+            )
+            cast(Any, ipydisplay.display)(
+                cast(Any, ipydisplay.Markdown)(
+                    "🔍 **Hardware detected** — auto-selecting LLM model:\n\n"
+                    f"- RAM: {profile.ram_gb:.1f} GB"
+                    + gpu_text
+                    + f"\n- **Selected model:** `{model_spec.name}` — {model_spec.description}"
+                )
+            )
+        else:
+            gpu_text = (
+                f", GPU: {profile.gpu_name} ({profile.vram_gb:.1f} GB VRAM)"
+                if profile.gpu_name
+                else ""
+            )
+            cast(Any, ipydisplay.display)(
+                cast(Any, ipydisplay.Markdown)(
+                    "⚠️ **Hardware auto-detection:** No suitable model found for your hardware.\n\n"
+                    f"- RAM: {profile.ram_gb:.1f} GB"
+                    + gpu_text
+                    + "\n\nThe LLM provider will be marked as unavailable. "
+                    "Please configure a model manually in your `global_config.toml`."
+                )
+            )
+    except Exception:
+        if model_spec is not None:
+            cfg.llm.model = model_spec.name

notebook_ta/bench/__init__.py

@@ -0,0 +1 @@
+"""Prompt/model benchmarking tool for notebook-ta instructors."""

notebook_ta/bench/app.py

@@ -0,0 +1,37 @@
+"""NiceGUI application bootstrap for the prompt benchmarking tool."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from nicegui import ui
+
+from notebook_ta.bench.state import BenchAppState
+from notebook_ta.bench.storage import ProjectStore, get_last_project_path
+from notebook_ta.bench.ui import layout
+from notebook_ta.logging import get_logger
+
+_log = get_logger("bench.app")
+
+
+def main(project_path: str | None = None) -> None:
+    """Launch the benchmarking GUI and show its project welcome dialog.
+
+    ``project_path`` is offered as the recent project. Otherwise the last project
+    opened across previous runs is offered.
+    """
+    candidate = Path(project_path) if project_path else get_last_project_path()
+    recent_path = candidate if candidate is not None and candidate.exists() else None
+    state = BenchAppState(
+        ProjectStore(None), project_open=False, recent_project_path=recent_path
+    )
+
+    @ui.page("/")
+    def index() -> None:
+        layout.build(state)
+
+    ui.run(title="Notebook-TA Benchmarking", reload=False, show=True, port=0)
+
+
+if __name__ in {"__main__", "__mp_main__"}:
+    main()

notebook_ta/bench/catalog.py

@@ -0,0 +1,73 @@
+"""Style-preserving mutations for a local exercise TOML catalog."""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import tempfile
+from pathlib import Path
+
+import tomlkit
+from tomlkit.items import Table
+
+from notebook_ta.config.models import ConfigurationError, ExerciseConfig
+
+
+def set_exercise_name(path: str | Path, exercise_id: str, name: str) -> None:
+    """Update an exercise display name while preserving unrelated TOML formatting."""
+    catalog_path, document, exercises = _load_editable_catalog(path)
+    exercise = exercises.get(exercise_id)
+    if not isinstance(exercise, Table):
+        raise ConfigurationError(f"Exercise {exercise_id!r} does not exist in {catalog_path}.")
+    exercise["name"] = name
+    _write_document(catalog_path, document.as_string())
+
+
+def add_exercise(path: str | Path, exercise: ExerciseConfig) -> None:
+    """Append a new exercise to a local TOML catalog without disturbing existing entries."""
+    catalog_path, document, exercises = _load_editable_catalog(path)
+    if exercise.id in exercises:
+        raise ConfigurationError(f"Exercise {exercise.id!r} already exists in {catalog_path}.")
+
+    item = tomlkit.table()
+    if exercise.name:
+        item.add("name", exercise.name)
+    if exercise.statement:
+        item.add("statement", exercise.statement)
+    if exercise.additional_info:
+        item.add("additional_info", exercise.additional_info)
+    exercises.add(exercise.id, item)
+    _write_document(catalog_path, document.as_string())
+
+
+def _load_editable_catalog(path: str | Path) -> tuple[Path, tomlkit.TOMLDocument, Table]:
+    """Load a local catalog and return its mutable exercises table."""
+    source = str(path)
+    if source.startswith(("http://", "https://")):
+        raise ConfigurationError("Remote exercise catalogs are read-only.")
+    catalog_path = Path(path)
+    if not catalog_path.exists():
+        raise ConfigurationError(f"Exercise catalog not found: {catalog_path}")
+    try:
+        document = tomlkit.parse(catalog_path.read_text(encoding="utf-8"))
+    except (OSError, tomlkit.exceptions.ParseError) as exc:
+        raise ConfigurationError(f"Could not edit exercise catalog {catalog_path}: {exc}") from exc
+
+    exercises = document.get("exercises")
+    if not isinstance(exercises, Table):
+        exercises = tomlkit.table()
+        document["exercises"] = exercises
+    return catalog_path, document, exercises
+
+
+def _write_document(path: Path, content: str) -> None:
+    """Atomically replace a TOML catalog with updated content."""
+    fd, temporary_name = tempfile.mkstemp(dir=str(path.parent), prefix=".exercises-", suffix=".tmp")
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8", newline="") as temporary_file:
+            temporary_file.write(content)
+        os.replace(temporary_name, path)
+    except BaseException:
+        with contextlib.suppress(OSError):
+            os.remove(temporary_name)
+        raise

notebook_ta/bench/cli.py

@@ -0,0 +1,26 @@
+"""CLI command for launching the prompt benchmarking GUI."""
+
+from __future__ import annotations
+
+import click
+
+
+@click.group()
+def cli() -> None:
+    """notebook-ta command line interface."""
+
+
+@click.command("bench")
+@click.argument("project_file", required=False, type=click.Path(dir_okay=False))
+def bench(project_file: str | None) -> None:
+    """Launch the benchmarking GUI, optionally offering PROJECT_FILE on its welcome screen."""
+    try:
+        from notebook_ta.bench.app import main
+    except ImportError as exc:
+        raise click.ClickException(
+            "The benchmarking UI requires the 'bench' extra: pip install 'notebook-ta[bench]'"
+        ) from exc
+    main(project_file)
+
+
+cli.add_command(bench)

notebook_ta/bench/executor.py

@@ -0,0 +1,307 @@
+"""Sequential benchmark execution engine with per-job metrics capture.
+
+Per the architecture decision, LLM calls are processed strictly one at a time
+across the whole exercise x solution x model matrix -- this keeps hardware load
+predictable and keeps TTFT/throughput measurements free of interference from
+concurrent requests.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import contextlib
+import io
+import sys
+import time
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+
+from notebook_ta.bench.hashing import build_input_snapshot
+from notebook_ta.bench.models import (
+    BenchmarkRun,
+    ExecutionMetrics,
+    ExecutionRecord,
+    ModelUnderTest,
+    PromptVersion,
+    StudentSolution,
+    TestResultModel,
+    TokenUsage,
+)
+from notebook_ta.config.models import ExerciseConfig, GlobalConfig, PromptConfig
+from notebook_ta.exercise.definition import Exercise
+from notebook_ta.llm.base import LLMProvider, create_provider
+from notebook_ta.logging import get_logger
+from notebook_ta.testing.runner import TestResult, TestRunner
+
+_log = get_logger("bench.executor")
+
+
+@dataclass
+class BenchJob:
+    """One (exercise, solution, model, prompt version) unit of work for a benchmark run."""
+
+    exercise_config: ExerciseConfig
+    solution: StudentSolution
+    model: ModelUnderTest
+    prompt_version: PromptVersion
+    setup_code: str = ""
+
+
+# Called as on_progress(job, status, message, record). `record` is populated only
+# when the job has just finished (status in {"completed", "failed"}).
+ProgressCallback = Callable[[BenchJob, str, "str | None", "ExecutionRecord | None"], None]
+
+
+@contextmanager
+def extended_sys_path(dirs: list[str]) -> Iterator[None]:
+    """Temporarily prepend `dirs` to `sys.path` for external test module imports."""
+    added = [d for d in dirs if d not in sys.path]
+    for d in added:
+        sys.path.insert(0, d)
+    try:
+        yield
+    finally:
+        for d in added:
+            if d in sys.path:
+                sys.path.remove(d)
+
+
+def _build_global_config(prompt_version: PromptVersion, model: ModelUnderTest) -> GlobalConfig:
+    """Build a synthetic GlobalConfig so `Exercise.build_prompt()` can be reused unchanged."""
+    return GlobalConfig(
+        llm=model.llm_config,
+        prompts=PromptConfig(
+            on_success=prompt_version.on_success,
+            on_failure=prompt_version.on_failure,
+            on_no_llm="",
+        ),
+    )
+
+
+def build_jobs(
+    exercises: list[ExerciseConfig],
+    solutions_by_exercise: dict[str, list[StudentSolution]],
+    models: list[ModelUnderTest],
+    prompt_version: PromptVersion,
+    setup_code_by_exercise: dict[str, str] | None = None,
+) -> list[BenchJob]:
+    """Build the full model x exercise x solution job matrix for a benchmark run.
+
+    Jobs are grouped by model first (outer loop) so that `BenchExecutor.run()` -- which
+    processes the list strictly in order -- finishes every job for one model before moving
+    to the next. This avoids repeatedly loading/unloading models on the LLM backend (e.g.
+    Ollama swapping models in and out of GPU/RAM) that would happen with an exercise-major
+    ordering.
+    """
+    setup_code_by_exercise = setup_code_by_exercise or {}
+    jobs: list[BenchJob] = []
+    for model in models:
+        for exercise_config in exercises:
+            for solution in solutions_by_exercise.get(exercise_config.id, []):
+                jobs.append(
+                    BenchJob(
+                        exercise_config,
+                        solution,
+                        model,
+                        prompt_version,
+                        setup_code_by_exercise.get(exercise_config.id, ""),
+                    )
+                )
+    return jobs
+
+
+def run_setup_code(
+    setup_code: str, namespace: dict[str, object], test_names: list[str]
+) -> list[TestResult] | None:
+    """Run benchmark setup code and return failed test results if setup fails."""
+    if not setup_code:
+        return None
+    stdout_buf = io.StringIO()
+    try:
+        with contextlib.redirect_stdout(stdout_buf):
+            exec(setup_code, namespace)  # noqa: S102 -- isolated authoring-time sandbox
+    except Exception as exc:
+        captured = stdout_buf.getvalue()
+        message = str(exc)
+        if captured:
+            message = f"{message}\nOutput: {captured}"
+        return [
+            TestResult(name=name, passed=False, message=f"Setup code failed: {message}")
+            for name in test_names
+        ]
+    return None
+
+
+class BenchExecutor:
+    """Runs a benchmark job matrix sequentially, one LLM call at a time."""
+
+    def __init__(
+        self,
+        python_path_dirs: Callable[[], list[str]] | None = None,
+        unit_test_timeout: Callable[[], float] | None = None,
+    ) -> None:
+        """`python_path_dirs`, if given, is called fresh before every job so that changes
+        made in the Settings tab while a run is in progress take effect immediately."""
+        self._get_python_path_dirs = python_path_dirs or (lambda: [])
+        self._get_unit_test_timeout = unit_test_timeout or (lambda: 5.0)
+        self._provider_cache: dict[str, LLMProvider] = {}
+        self._cancel_requested = False
+        self._retry_event: asyncio.Event | None = None
+
+    def cancel(self) -> None:
+        """Request cancellation; takes effect before the next job starts (or on retry-wait)."""
+        self._cancel_requested = True
+        if self._retry_event is not None:
+            self._retry_event.set()
+
+    def retry(self) -> None:
+        """Resume a paused run (e.g. after the user restarts a disconnected LLM server)."""
+        if self._retry_event is not None:
+            self._retry_event.set()
+
+    def _get_provider(self, model: ModelUnderTest) -> LLMProvider:
+        """Return a cached LLMProvider instance for `model`, creating one if needed."""
+        if model.label not in self._provider_cache:
+            self._provider_cache[model.label] = create_provider(model.llm_config)
+        return self._provider_cache[model.label]
+
+    async def run(
+        self, jobs: list[BenchJob], run: BenchmarkRun, on_progress: ProgressCallback
+    ) -> None:
+        """Execute all jobs sequentially, invoking `on_progress` as progress is made."""
+        self._cancel_requested = False
+        for job in jobs:
+            if self._cancel_requested:
+                run.status = "cancelled"
+                return
+
+            provider = self._get_provider(job.model)
+            while not provider.is_available():
+                run.status = "paused"
+                on_progress(job, "paused", f"{job.model.label} is unreachable", None)
+                self._retry_event = asyncio.Event()
+                await self._retry_event.wait()
+                if self._cancel_requested:
+                    run.status = "cancelled"
+                    return
+            run.status = "running"
+
+            await self._execute_job(job, run, on_progress)
+
+        if run.status == "running":
+            run.status = "completed"
+
+    async def run_single(self, job: BenchJob, run: BenchmarkRun) -> ExecutionRecord:
+        """Execute a single job immediately, bypassing the queue (used for a Compare tab Re-run)."""
+        record_box: list[ExecutionRecord] = []
+
+        def _capture(
+            _job: BenchJob,
+            _status: str,
+            _message: str | None,
+            record: ExecutionRecord | None,
+        ) -> None:
+            if record is not None:
+                record_box.append(record)
+
+        await self._execute_job(job, run, _capture)
+        return record_box[0]
+
+    async def _execute_job(
+        self, job: BenchJob, run: BenchmarkRun, on_progress: ProgressCallback
+    ) -> None:
+        """Run unit tests + the LLM call for one job and report the resulting record."""
+        on_progress(job, "generating", None, None)
+        snapshot = build_input_snapshot(job.exercise_config, job.solution, job.setup_code)
+        global_config = _build_global_config(job.prompt_version, job.model)
+        global_config.unit_test_timeout = self._get_unit_test_timeout()
+        exercise = Exercise(job.exercise_config, global_config)
+
+        try:
+            namespace: dict[str, object] = {}
+            with extended_sys_path(self._get_python_path_dirs()):
+                exec(job.solution.code, namespace)  # noqa: S102 -- isolated authoring-time sandbox
+                test_names = [test_def.name for test_def in exercise.tests]
+                test_results = run_setup_code(job.setup_code, namespace, test_names)
+                if test_results is None:
+                    test_results = TestRunner().run(exercise, namespace)
+            prompt = exercise.build_prompt(job.solution.code, test_results, hint_history=None)
+
+            provider = self._get_provider(job.model)
+            output, metrics = await self._run_llm(provider, prompt)
+
+            record = ExecutionRecord(
+                run_id=run.id,
+                exercise_id=job.exercise_config.id,
+                solution_id=job.solution.id,
+                model_label=job.model.label,
+                prompt_version_id=job.prompt_version.id,
+                input_snapshot=snapshot,
+                full_prompt=prompt,
+                test_results=[
+                    TestResultModel(name=r.name, passed=r.passed, message=r.message)
+                    for r in test_results
+                ],
+                llm_output=output,
+                metrics=metrics,
+                status="completed",
+            )
+            on_progress(job, "completed", None, record)
+        except Exception as exc:
+            _log.warning(
+                "Job failed for exercise=%r solution=%r model=%r: %s",
+                job.exercise_config.id,
+                job.solution.id,
+                job.model.label,
+                exc,
+            )
+            record = ExecutionRecord(
+                run_id=run.id,
+                exercise_id=job.exercise_config.id,
+                solution_id=job.solution.id,
+                model_label=job.model.label,
+                prompt_version_id=job.prompt_version.id,
+                input_snapshot=snapshot,
+                status="failed",
+                error=str(exc),
+            )
+            on_progress(job, "failed", str(exc), record)
+
+    @staticmethod
+    async def _run_llm(provider: LLMProvider, prompt: str) -> tuple[str, ExecutionMetrics]:
+        """Stream `prompt` through `provider`, timing TTFT/total time and capturing token usage."""
+        start = time.monotonic()
+        first_token_time: float | None = None
+        chunks: list[str] = []
+        async for chunk in provider.stream(prompt):
+            if first_token_time is None:
+                first_token_time = time.monotonic()
+            chunks.append(chunk)
+        end = time.monotonic()
+        output = "".join(chunks)
+
+        usage = provider.get_last_usage()
+        if usage and usage.completion_tokens:
+            completion_tokens = usage.completion_tokens
+            prompt_tokens = usage.prompt_tokens
+            approximate = False
+        else:
+            completion_tokens = len(output.split())
+            prompt_tokens = usage.prompt_tokens if usage else None
+            approximate = True
+
+        total_time = end - start
+        ttft = (first_token_time - start) if first_token_time is not None else None
+        throughput = (completion_tokens / total_time) if total_time > 0 else None
+
+        return output, ExecutionMetrics(
+            time_to_first_token_s=ttft,
+            total_generation_time_s=total_time,
+            throughput_tokens_per_s=throughput,
+            token_usage=TokenUsage(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                approximate=approximate,
+            ),
+        )