feat: enable Colab Gemini instructions for galaxy and fit + per-assistant cap (#324)

.github/workflows/release.yml

@@ -818,23 +818,23 @@ jobs:
         python-version: [3.12]
         # `assistant`: the assistant repo whose colab/contexts.yaml stamps each
         # generated notebook's Colab Gemini instructions (PyAutoHands#324).
-        # Empty = no instructions yet: autogalaxy/autofit mappings land in
-        # phase 2, and a root without the file would fail generation loudly.
-        # Merge order: the assistant's colab/contexts.yaml must be on its main
-        # before this is released (PyAutoBrain -> autolens_assistant -> PyAutoHands).
+        # Empty = no instructions; a root without the file fails generation
+        # loudly. Merge order: each named assistant's colab/contexts.yaml must be
+        # on its main before this is released (autolens_assistant, then
+        # autofit_assistant + autogalaxy_assistant, then PyAutoHands).
         workspace:
           - repository: PyAutoLabs/autofit_workspace
             name: autofit
             package: PyAutoFit
             generate_notebooks: true
             bump_colab_urls: true
-            assistant: ""
+            assistant: autofit_assistant
           - repository: PyAutoLabs/autogalaxy_workspace
             name: autogalaxy
             package: PyAutoGalaxy
             generate_notebooks: true
             bump_colab_urls: true
-            assistant: ""
+            assistant: autogalaxy_assistant
           - repository: PyAutoLabs/autolens_workspace
             name: autolens
             package: PyAutoLens
@@ -846,7 +846,7 @@ jobs:
             package: PyAutoGalaxy
             generate_notebooks: true
             bump_colab_urls: true
-            assistant: ""
+            assistant: autogalaxy_assistant
           - repository: PyAutoLabs/HowToLens
             name: howtolens
             package: PyAutoLens
@@ -858,7 +858,7 @@ jobs:
             package: PyAutoFit
             generate_notebooks: true
             bump_colab_urls: true
-            assistant: ""
+            assistant: autofit_assistant
           - repository: PyAutoLabs/euclid_strong_lens_modeling_pipeline
             name: euclid_pipeline
             package: PyAutoLens

autohands/build_util.py

@@ -415,8 +415,17 @@ def _colab_gemini_config() -> dict:
 GEMINI_CONTEXT_FILE = Path("colab") / "contexts.yaml"
 
 # Human choice (#324): one instruction never exceeds this. Over-cap fails
-# generation naming the mapping entry; the text is never truncated.
+# generation naming the mapping entry; the text is never truncated. An
+# assistant may set its own `cap:` in colab/contexts.yaml (its skills may be
+# longer), up to GEMINI_CONTEXT_CAP_MAX — the size the human verified Colab
+# honours (52k tested, 2026-10-26).
 GEMINI_CONTEXT_CAP = 20_000
+GEMINI_CONTEXT_CAP_MAX = 50_000
+
+
+def gemini_context_cap(contexts: dict) -> int:
+    """The cap in force for a loaded mapping: its `cap:`, else the default."""
+    return contexts.get("cap", GEMINI_CONTEXT_CAP)
 
 # Fixed URL template for the per-project instruction id (uuid5 over it), so
 # regenerating a notebook rewrites the same `aiContexts` entry rather than
@@ -451,6 +460,12 @@ def load_gemini_contexts(assistant_root, enforce_never_map: bool = True) -> dict
             f"pass --assistant-root / PYAUTO_ASSISTANT_ROOT for this project)."
         )
     data = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
+    cap = data.get("cap", GEMINI_CONTEXT_CAP)
+    if isinstance(cap, bool) or not isinstance(cap, int) or not 0 < cap <= GEMINI_CONTEXT_CAP_MAX:
+        raise ValueError(
+            f"Gemini instructions: {path} sets cap {cap!r}; it must be a positive "
+            f"integer no larger than {GEMINI_CONTEXT_CAP_MAX}."
+        )
     for key in ("preface", "rules", "default", "contexts"):
         if key not in data:
             raise ValueError(f"Gemini instructions: {path} has no '{key}' entry.")
@@ -549,7 +564,8 @@ def gemini_context_from(assistant_root, notebook_name: str, contexts=None):
     Assemble the Gemini instruction for ``notebook_name`` ("<repo>/<relpath>").
 
     Returns ``(entry_label, skills, text)``. Raises if a mapped skill file is
-    missing or the text exceeds ``GEMINI_CONTEXT_CAP`` — never truncates.
+    missing or the text exceeds the mapping's cap (``gemini_context_cap``) —
+    never truncates.
     """
     assistant_root = Path(assistant_root)
     if contexts is None:
@@ -562,11 +578,12 @@ def gemini_context_from(assistant_root, notebook_name: str, contexts=None):
             f"Gemini instructions: mapping entry '{label}' names skill(s) {missing} "
             f"with no file under {assistant_root / 'skills'}."
         )
-    if len(text) > GEMINI_CONTEXT_CAP:
+    cap = gemini_context_cap(contexts)
+    if len(text) > cap:
         raise ValueError(
             f"Gemini instructions for {notebook_name}: mapping entry '{label}' "
             f"(skills {skills}) assembles to {len(text)} characters, over the "
-            f"{GEMINI_CONTEXT_CAP} cap. Map fewer or smaller skills to it in "
+            f"{cap} cap. Map fewer or smaller skills to it in "
             f"{assistant_root / GEMINI_CONTEXT_FILE}; the text is never truncated."
         )
     return label, skills, text
@@ -657,8 +674,8 @@ def resolve_assistant_root(project: str, explicit=None):
     An explicit or environment root is returned as given, and a missing
     ``colab/contexts.yaml`` under it then fails generation loudly. A
     workspace-resolved root without the file is skipped with a warning
-    instead: that is an assistant whose mapping has not landed yet (galaxy and
-    fit until phase 2 of #324), and the release must not break before it does.
+    instead: that is an assistant whose mapping has not landed yet (e.g. a newly
+    added project), and the release must not break before it does.
     """
     if project not in COLAB_ASSISTANTS:
         return None

autohands/gemini_context.py

@@ -81,7 +81,7 @@ def report(workspace, project, assistant_root, repo_name=None):
     assistant_root = Path(assistant_root).resolve()
     repo_name = _repo_name(project, repo_name)
     contexts = build_util.load_gemini_contexts(assistant_root, enforce_never_map=False)
-    cap = build_util.GEMINI_CONTEXT_CAP
+    cap = build_util.gemini_context_cap(contexts)
 
     rows = []
     for notebook in iter_notebooks(workspace):

bin/autohands

@@ -248,7 +248,8 @@ autohands gemini_context report --workspace W --project P --assistant-root A
 
 Colab Gemini instructions (PyAutoHands#324): each assistant's
 colab/contexts.yaml maps notebook-path globs to skills; the assembled text
-(preface + core rules + mapped skills, at most 20,000 characters) is written
+(preface + core rules + mapped skills, at most 20,000 characters unless the
+mapping sets its own `cap:`, up to 50,000) is written
 into metadata.colab.aiContexts and selected.
 
   apply    Stamp every notebook under <notebooks_dir>. `generate` already does

tests/test_gemini_instructions.py

@@ -328,3 +328,56 @@ def test_generate_fails_before_touching_notebooks(tmp_path, assistant, folder):
                             capture_output=True, text=True, check=True)
     assert status.stdout == ""
     assert list(workspace.glob("scripts/**/*.ipynb")) == []
+
+
+def _set_cap(assistant, cap_line):
+    yaml_path = assistant / "colab" / "contexts.yaml"
+    yaml_path.write_text(cap_line + "\n" + yaml_path.read_text())
+
+
+def _add_mid_skill(assistant):
+    # A 25k-character skill mapped by the `big` entry: over the default cap,
+    # within a 30k one.
+    (assistant / "skills" / "skill_big.md").write_text("y" * 25_000)
+
+
+def test_cap_defaults_when_absent(assistant):
+    contexts = build_util.load_gemini_contexts(assistant)
+    assert "cap" not in contexts
+    assert build_util.gemini_context_cap(contexts) == build_util.GEMINI_CONTEXT_CAP == 20_000
+
+
+def test_custom_cap_honoured(tmp_path, assistant):
+    _add_mid_skill(assistant)
+    name = f"{REPO}/notebooks/big/a.ipynb"
+    with pytest.raises(ValueError, match="over the 20000 cap"):
+        build_util.gemini_context_from(assistant, name)
+
+    _set_cap(assistant, "cap: 30000")
+    label, _, text = build_util.gemini_context_from(assistant, name)
+    assert 25_000 < len(text) <= 30_000
+    nb = make_notebook(tmp_path / "a.ipynb", [code_cell("x = 1")])
+    assert build_util.inject_gemini_instructions(nb, "autolens", assistant, "notebooks/big/a.ipynb") == label
+
+
+@pytest.mark.parametrize("cap", ["50001", "0", "-5", "lots", "true"])
+def test_cap_out_of_range_raises_at_load(assistant, cap):
+    _set_cap(assistant, f"cap: {cap}")
+    with pytest.raises(ValueError, match="must be a positive integer no larger than 50000"):
+        build_util.load_gemini_contexts(assistant)
+
+
+def test_cap_at_ceiling_loads(assistant):
+    _set_cap(assistant, f"cap: {build_util.GEMINI_CONTEXT_CAP_MAX}")
+    assert build_util.gemini_context_cap(build_util.load_gemini_contexts(assistant)) == 50_000
+
+
+def test_report_shows_cap_in_force(tmp_path, assistant):
+    _add_mid_skill(assistant)
+    _set_cap(assistant, "cap: 30000")
+    workspace = tmp_path / REPO
+    make_notebook(workspace / "notebooks" / "big" / "a.ipynb", [code_cell("x = 1")])
+    result = gemini_context.report(workspace, "autolens", assistant)
+    assert result["cap"] == 30_000
+    assert result["sizes"]["over_cap"] == []
+    assert "cap 30000" in gemini_context.format_report(result)