Add a two-stage curation fallback: alternate Gemini model, then manual

README.md

@@ -106,6 +106,9 @@ To skip this entirely, use Gemini to curate automatically instead:
 2. Get a free API key at [aistudio.google.com/app/api-keys](https://aistudio.google.com/app/api-keys)
 3. Run with `--curation-backend gemini --curation-api-key YOUR_KEY`, or set it once via `$env:GOOGLE_API_KEY="YOUR_KEY"` (PowerShell) / `export GOOGLE_API_KEY=YOUR_KEY` (bash) and just pass `--curation-backend gemini`
 
+**If this step appears to hang with nothing printing** (v3.0.2+): CANDy now logs before and after the Gemini request, and bounds each attempt to a 60-second timeout, so a multi-minute wait during retries is visible and finite rather than silent. On an older version, this step had no request timeout at all and could hang indefinitely on a stalled connection with zero output -- if you hit this, `Ctrl+C`, upgrade (`pip install --upgrade candy-cazyme`), and rerun.
+
+**If Gemini itself fails** (v3.0.3+, e.g. a `503 UNAVAILABLE` "high demand" error, a rate limit, or a model being deprecated): CANDy automatically retries once against a different model (`gemini-2.5-flash-lite` by default) before giving up on Gemini entirely. If that also fails, it automatically falls back to the interactive manual-curation prompt described above, rather than crashing and discarding the (often several-minutes-long) fetching/clustering/domain-detection work already done in that run.
 
 ### Python API
 

src/candy/curation/gemini.py

@@ -28,6 +28,13 @@
 )
 
 
+# Tried automatically if the primary model fails (server overload, model-specific
+# outage, deprecation, ...) -- a different, typically lower-demand model often
+# succeeds even when the primary one is returning 503s. Set fallback_model=None
+# to disable and fail on the first error instead.
+_DEFAULT_FALLBACK_MODEL = "gemini-2.5-flash-lite"
+
+
 class GeminiCurationBackend:
     """Automated curation via Google Gemini, ported from the notebook's default path.
 
@@ -38,14 +45,20 @@ class GeminiCurationBackend:
 
     name = "gemini"
 
-    def __init__(self, api_key: str | None = None, model: str = "gemini-flash-latest") -> None:
+    def __init__(
+        self,
+        api_key: str | None = None,
+        model: str = "gemini-flash-latest",
+        fallback_model: str | None = _DEFAULT_FALLBACK_MODEL,
+    ) -> None:
         self.api_key = api_key or os.environ.get("GOOGLE_API_KEY") or os.environ.get("GEMINI_API_KEY")
         if not self.api_key:
             raise ValueError(
                 "Gemini curation requires an API key. Pass api_key=, or set the "
                 "GOOGLE_API_KEY or GEMINI_API_KEY environment variable."
             )
         self.model = model
+        self.fallback_model = fallback_model
 
     def curate(self, domain_names: list[str], *, family: str | None = None) -> dict[str, list[str]]:
         try:
@@ -61,16 +74,31 @@ def curate(self, domain_names: list[str], *, family: str | None = None) -> dict[
             f"belong to enzymes from CAZy family {family}" if family else "are carbohydrate-active enzymes"
         )
         prompt = _PROMPT_TEMPLATE.format(domain_names=domain_names, family_clause=family_clause)
-
-        logger.info(
-            "Requesting domain-name curation from Gemini (model=%s, %d domain names). This can take "
-            "a while if the API is transiently retrying -- it isn't stuck if it takes a few minutes.",
-            self.model,
-            len(domain_names),
-        )
         client = genai.Client(api_key=self.api_key, http_options=types.HttpOptions(timeout=_REQUEST_TIMEOUT_MS))
-        response = client.models.generate_content(model=self.model, contents=prompt)
-        logger.info("Received curation response from Gemini.")
 
-        cleaned = response.text.replace("python", "").replace("```", "")
-        return ast.literal_eval(cleaned)
+        models_to_try = [self.model]
+        if self.fallback_model and self.fallback_model != self.model:
+            models_to_try.append(self.fallback_model)
+
+        last_error: Exception | None = None
+        for attempt_model in models_to_try:
+            logger.info(
+                "Requesting domain-name curation from Gemini (model=%s, %d domain names). This can "
+                "take a while if the API is transiently retrying -- it isn't stuck if it takes a few "
+                "minutes.",
+                attempt_model,
+                len(domain_names),
+            )
+            try:
+                response = client.models.generate_content(model=attempt_model, contents=prompt)
+                cleaned = response.text.replace("python", "").replace("```", "")
+                curated = ast.literal_eval(cleaned)
+            except Exception as exc:  # noqa: BLE001 -- deliberately broad: any failure means "try the next model"
+                logger.warning("Gemini curation failed with model=%s (%s).", attempt_model, exc)
+                last_error = exc
+                continue
+            logger.info("Received curation response from Gemini (model=%s).", attempt_model)
+            return curated
+
+        assert last_error is not None  # models_to_try is never empty
+        raise last_error

src/candy/pipeline.py

@@ -175,7 +175,18 @@ def run_pipeline(config: PipelineConfig) -> PipelineResult:
             curation_kwargs["model"] = config.curation.model
     curation_backend = get_curation_backend(config.curation.backend, **curation_kwargs)
     family_label = config.input.family if is_cazy_query else None
-    curated_domains = curation_backend.curate(raw_domain_names, family=family_label)
+    try:
+        curated_domains = curation_backend.curate(raw_domain_names, family=family_label)
+    except Exception as exc:  # noqa: BLE001 -- any curation-backend failure falls back to manual
+        if config.curation.backend == "manual":
+            raise
+        logger.warning(
+            "Automated curation backend '%s' failed (%s); falling back to manual curation so the "
+            "work already done in this run (fetching, clustering, domain detection) isn't lost.",
+            config.curation.backend,
+            exc,
+        )
+        curated_domains = get_curation_backend("manual").curate(raw_domain_names, family=family_label)
 
     # --- Stage 9: build the result database ---
     with open(sequences_fasta) as handle:

tests/test_curation.py

@@ -78,3 +78,56 @@ def test_gemini_curate_bounds_request_with_an_explicit_timeout_and_logs_progress
     messages = [record.message for record in caplog.records]
     assert any("Requesting domain-name curation from Gemini" in m for m in messages)
     assert any("Received curation response from Gemini" in m for m in messages)
+
+
+def test_gemini_curate_falls_back_to_secondary_model_on_failure(caplog):
+    # Regression test: a real run hit a 503 UNAVAILABLE ("high demand") from
+    # the primary model. A different model is often not overloaded at the
+    # same time, so try it automatically before giving up.
+    pytest.importorskip("google.genai")
+    from candy.curation.gemini import GeminiCurationBackend
+
+    fake_response = MagicMock()
+    fake_response.text = "{'Catalytic domain': ['A']}"
+    fake_client = MagicMock()
+    fake_client.models.generate_content.side_effect = [Exception("503 UNAVAILABLE"), fake_response]
+
+    with patch("google.genai.Client", return_value=fake_client), caplog.at_level("INFO"):
+        backend = GeminiCurationBackend(api_key="fake-key", model="primary-model", fallback_model="backup-model")
+        result = backend.curate(["A"])
+
+    assert result == {"Catalytic domain": ["A"]}
+    calls = fake_client.models.generate_content.call_args_list
+    assert calls[0].kwargs["model"] == "primary-model"
+    assert calls[1].kwargs["model"] == "backup-model"
+    assert any("Gemini curation failed with model=primary-model" in r.message for r in caplog.records)
+
+
+def test_gemini_curate_raises_last_error_when_all_models_fail():
+    pytest.importorskip("google.genai")
+    from candy.curation.gemini import GeminiCurationBackend
+
+    fake_client = MagicMock()
+    fake_client.models.generate_content.side_effect = [Exception("503 primary"), Exception("503 fallback")]
+
+    with patch("google.genai.Client", return_value=fake_client):
+        backend = GeminiCurationBackend(api_key="fake-key", model="primary-model", fallback_model="backup-model")
+        with pytest.raises(Exception, match="503 fallback"):
+            backend.curate(["A"])
+
+    assert fake_client.models.generate_content.call_count == 2
+
+
+def test_gemini_curate_fallback_disabled_when_fallback_model_is_none():
+    pytest.importorskip("google.genai")
+    from candy.curation.gemini import GeminiCurationBackend
+
+    fake_client = MagicMock()
+    fake_client.models.generate_content.side_effect = Exception("503 primary")
+
+    with patch("google.genai.Client", return_value=fake_client):
+        backend = GeminiCurationBackend(api_key="fake-key", model="primary-model", fallback_model=None)
+        with pytest.raises(Exception, match="503 primary"):
+            backend.curate(["A"])
+
+    assert fake_client.models.generate_content.call_count == 1

tests/test_pipeline.py

@@ -13,7 +13,7 @@
 
 import pandas as pd
 
-from candy.config import CAZyFamilyInput, ClusteringConfig, CustomFastaInput, PipelineConfig, Taxonomy
+from candy.config import CAZyFamilyInput, ClusteringConfig, CurationConfig, CustomFastaInput, PipelineConfig, Taxonomy
 from candy.pipeline import _warn_if_running_under_rosetta, run_pipeline
 
 
@@ -120,6 +120,66 @@ def test_run_pipeline_custom_fasta_mode_without_tree(tmp_path):
     assert result.sequence_count == 1
 
 
+class FailingCurationBackend:
+    name = "gemini"
+
+    def curate(self, domain_names, *, family=None):
+        raise RuntimeError("503 UNAVAILABLE: This model is currently experiencing high demand.")
+
+
+def test_run_pipeline_falls_back_to_manual_curation_when_automated_backend_fails(tmp_path, caplog):
+    # Regression test: a Gemini 503 used to crash the whole run, discarding
+    # all the upstream work (fetching, clustering, domain detection) over a
+    # transient failure. It should fall back to manual curation instead.
+    fasta_path = tmp_path / "input.fasta"
+    fasta_path.write_text(">Protein1\nMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLA\n")
+
+    config = PipelineConfig(
+        input=CustomFastaInput(fasta_path=fasta_path),
+        jobname="testjob_curation_fallback",
+        output_dir=tmp_path / "out",
+        build_tree=False,
+        curation=CurationConfig(backend="gemini", api_key="fake-key"),
+    )
+
+    def fake_get_curation_backend(name, **kwargs):
+        return StubCurationBackend() if name == "manual" else FailingCurationBackend()
+
+    with patch("candy.interpro._query_md5_batch", side_effect=lambda batch: _fake_match_xml(batch)), patch(
+        "candy.interpro._fetch_entry_name", return_value="Catalytic domain"
+    ), patch("candy.pipeline.get_curation_backend", side_effect=fake_get_curation_backend), caplog.at_level(
+        "WARNING"
+    ):
+        result = run_pipeline(config)
+
+    assert result.database_path.exists()
+    assert any("falling back to manual curation" in r.message for r in caplog.records)
+
+
+def test_run_pipeline_manual_curation_failure_is_not_caught(tmp_path):
+    # The manual backend has nowhere further to fall back to -- its own
+    # failures (e.g. EOFError on a non-interactive stdin) should propagate.
+    fasta_path = tmp_path / "input.fasta"
+    fasta_path.write_text(">Protein1\nMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLA\n")
+
+    config = PipelineConfig(
+        input=CustomFastaInput(fasta_path=fasta_path),
+        jobname="testjob_manual_failure",
+        output_dir=tmp_path / "out",
+        build_tree=False,
+        curation=CurationConfig(backend="manual"),
+    )
+
+    with patch("candy.interpro._query_md5_batch", side_effect=lambda batch: _fake_match_xml(batch)), patch(
+        "candy.interpro._fetch_entry_name", return_value="Catalytic domain"
+    ), patch("candy.pipeline.get_curation_backend", return_value=FailingCurationBackend()):
+        try:
+            run_pipeline(config)
+            assert False, "expected RuntimeError"
+        except RuntimeError as exc:
+            assert "503 UNAVAILABLE" in str(exc)
+
+
 def test_run_pipeline_custom_fasta_mode_with_tree(tmp_path):
     fasta_path = tmp_path / "input.fasta"
     fasta_path.write_text(">Protein1\nMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLAMKVLA\n")