diff --git a/README.md b/README.md index 64eccba..021126a 100644 --- a/README.md +++ b/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 diff --git a/src/candy/curation/gemini.py b/src/candy/curation/gemini.py index c26a633..ec9203b 100644 --- a/src/candy/curation/gemini.py +++ b/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,7 +45,12 @@ 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( @@ -46,6 +58,7 @@ def __init__(self, api_key: str | None = None, model: str = "gemini-flash-latest "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 diff --git a/src/candy/pipeline.py b/src/candy/pipeline.py index 23ced03..8b6791f 100644 --- a/src/candy/pipeline.py +++ b/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: diff --git a/tests/test_curation.py b/tests/test_curation.py index b7c3018..9eb9941 100644 --- a/tests/test_curation.py +++ b/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 diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 6d2b42d..897c483 100644 --- a/tests/test_pipeline.py +++ b/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")