"""bioai.orchestrator -- end-to-end biopesticide design pipeline. Flow: 1. ``OllamaClient.parse_pest_report(user_text)`` -> pest dict 2. Load pest transcripts (synthetic by default) 3. Tile into 200-nt dsRNA precursors (50% overlap) 4. Dice each precursor into 21-nt siRNAs (Dicer-style) 5. ``CandidateRanker.rank(sirnas)`` -> top-N scored 6. ``OllamaClient.generate_safety_card(...)`` -> markdown card per top-5 7. ``OllamaClient.generate_regulatory_memo(...)`` -> EPA-style memo Returns a dict with ``pest_report``, ``candidates`` (top 10), ``safety_cards``, ``regulatory_memo``, ``total_cost_estimate``. CLI:: python -m bioai.orchestrator --user-text "Brown planthopper infestation in rice paddy in Tamil Nadu" If Ollama is not running (or the model isn't pulled), the LLM calls are skipped and the returned dict contains a ``degraded_mode`` flag and best-effort strings so the rest of the pipeline (ranking, safety cards from PINN + off-target index only) still produces useful output. """ from __future__ import annotations import argparse import sys from pathlib import Path from typing import Dict, List, Optional from .agent.ollama_client import OllamaClient, DEFAULT_MODEL from .inference.ranker import CandidateRanker from .models.sirna_cnn import resolve_device from .sequence_utils import ( SAFETY_SPECIES, dice_precursor, fasta_iter, normalize_pest_species, read_fasta, tile_sequence, ) # --------------------------------------------------------------------------- # # Defaults (resolved from bioai.paths so they're portable across machines) # --------------------------------------------------------------------------- # from bioai.paths import ( # noqa: E402 DEFAULT_PEST_FASTA, DEFAULT_SAFETY_FASTA, SIRNA_CHECKPOINT, PINN_CHECKPOINT, ) # Rough estimate of local compute cost for Llama 3.2 3B on a laptop GPU. # Local Ollama is effectively free (no API billing), but we still report an # opportunity-cost estimate based on cloud GPU-equivalent pricing for context. USD_PER_1K_INPUT_TOKENS = 0.0001 USD_PER_1K_OUTPUT_TOKENS = 0.0001 USD_PER_1K_INPUT_TOKENS = 0.0009 USD_PER_1K_OUTPUT_TOKENS = 0.0009 # --------------------------------------------------------------------------- # # Orchestrator # --------------------------------------------------------------------------- # class BiopesticideOrchestrator: """End-to-end pipeline. See module docstring for the flow diagram.""" def __init__( self, pest_fasta: Path = DEFAULT_PEST_FASTA, safety_fasta: Path = DEFAULT_SAFETY_FASTA, sirna_checkpoint: Path = SIRNA_CHECKPOINT, pinn_checkpoint: Path = PINN_CHECKPOINT, device: str = "auto", max_transcripts: int = 5, max_precursors_per_transcript: int = 4, max_sirnas_per_precursor: int = 9, # 200 / 21 ~ 9 ollama_client: Optional[OllamaClient] = None, ): self.pest_fasta = Path(pest_fasta) self.safety_fasta = Path(safety_fasta) self.sirna_checkpoint = Path(sirna_checkpoint) self.pinn_checkpoint = Path(pinn_checkpoint) self.device = resolve_device(device) self.max_transcripts = max_transcripts self.max_precursors_per_transcript = max_precursors_per_transcript self.max_sirnas_per_precursor = max_sirnas_per_precursor # Ranker (loads CNN + PINN weights, builds off-target index) safety_paths = {sp: self.safety_fasta for sp in SAFETY_SPECIES} self.ranker = CandidateRanker( safety_fasta_paths=safety_paths, sirna_checkpoint=self.sirna_checkpoint, pinn_checkpoint=self.pinn_checkpoint, device=str(self.device), ) # Ollama client (None in degraded mode if Ollama server is unreachable) self.llm = ollama_client self.degraded_mode = False if self.llm is None: try: self.llm = OllamaClient() print("[orchestrator] Ollama client initialised.") except RuntimeError as exc: self.degraded_mode = True print(f"[orchestrator] DEGRADED MODE -- {exc}. LLM calls will be skipped.") # Token accounting (rough) self._input_tokens = 0 self._output_tokens = 0 # ------------------------------------------------------------------ # def _estimate_tokens(self, text: str) -> int: # 1 token ~= 4 chars for English text (Llama tokenizer) return max(1, len(text) // 4) def _record_tokens(self, prompt: str, response: str) -> None: self._input_tokens += self._estimate_tokens(prompt) self._output_tokens += self._estimate_tokens(response) # ------------------------------------------------------------------ # def _load_pest_transcripts(self, pest_species: str | None = None) -> Dict[str, str]: if not self.pest_fasta.exists(): print(f"[orchestrator] pest FASTA missing: {self.pest_fasta}") return {} all_tx = read_fasta(self.pest_fasta) # If we know the target pest species, filter transcripts to that species. # Synthetic transcript headers look like "NILAPARVATA_LUGENS_FAKE_001". if pest_species: prefix = pest_species.upper() filtered = {k: v for k, v in all_tx.items() if k.startswith(prefix)} if filtered: items = list(filtered.items())[: self.max_transcripts] return dict(items) # Fall back to all transcripts if no match (defensive). print(f"[orchestrator] no transcripts matched species '{pest_species}'; using all") # Take a few transcripts to keep the demo fast. items = list(all_tx.items())[: self.max_transcripts] return dict(items) def _tile_and_dice(self, transcripts: Dict[str, str]) -> List[Dict]: """Return a list of ``{precursor, source, sirnas}`` dicts.""" out: List[Dict] = [] for gene_id, seq in transcripts.items(): windows = tile_sequence( seq, window=200, step=100, max_candidates=self.max_precursors_per_transcript, ) for start, end, precursor in windows: sirnas = dice_precursor(precursor, sirna_len=21, step=21)[: self.max_sirnas_per_precursor] out.append({ "gene_id": gene_id, "precursor_start": start, "precursor_end": end, "precursor_seq": precursor, "sirnas": sirnas, }) return out # ------------------------------------------------------------------ # def design(self, user_text: str, top_k: int = 10, pest_species_override: str = None) -> Dict: """Run the full pipeline. Returns a result dict (see module docstring). If ``pest_species_override`` is provided (e.g. "nilaparvata_lugens"), the LLM pest-report parsing step is skipped entirely and the species is used directly. This is the fast path used by the pest-card UI: ~2 seconds end-to-end vs ~15 seconds with LLM parsing. """ self._input_tokens = 0 self._output_tokens = 0 # Step 1: parse pest report (or use the override) if pest_species_override: # Fast path: skip LLM, build pest_report directly from the override normalized = normalize_pest_species(pest_species_override) # Infer crop from species crop_map = { "nilaparvata_lugens": "rice", "spodoptera_frugiperda": "maize", "schistocerca_gregaria": "wheat", "chilo_suppressalis": "rice", "myzus_persicae": "vegetables", "leptinotarsa_decemlineata": "potato", "bemisia_tabaci": "tomato", } pest_report = { "pest_species": normalized, "crop": crop_map.get(normalized, "unknown"), "severity": "moderate", "location": "unspecified", "notes": f"Direct selection (LLM parsing skipped for speed)", "_raw": "", "_degraded": False, } print(f"[orchestrator] pest_species_override provided: '{pest_species_override}' -> '{normalized}' (LLM parsing skipped)") elif self.degraded_mode: pest_report = self._degraded_pest_report(user_text) else: try: prompt = user_text pest_report = self.llm.parse_pest_report(user_text) self._record_tokens(prompt, pest_report.get("_raw", "")) except Exception as exc: print(f"[orchestrator] parse_pest_report failed ({exc!r}); degraded pest report") pest_report = self._degraded_pest_report(user_text) # Step 2: load transcripts (filter by pest species if known). # Normalize the species name first — the LLM may return a common name # like "Brown Planthopper" but the FASTA headers use the scientific # name "NILAPARVATA_LUGENS_FAKE_001". pest_species_raw = pest_report.get("pest_species") if pest_report else None pest_species = normalize_pest_species(pest_species_raw) if pest_species_raw else None if pest_species and pest_species != pest_species_raw: print(f"[orchestrator] normalized pest species: '{pest_species_raw}' -> '{pest_species}'") pest_report["pest_species"] = pest_species # update so downstream uses the normalized name transcripts = self._load_pest_transcripts(pest_species=pest_species) if not transcripts: return { "pest_report": pest_report, "candidates": [], "safety_cards": [], "regulatory_memo": "No pest transcripts available; cannot design candidates.", "total_cost_estimate": 0.0, "degraded_mode": self.degraded_mode, "error": "no pest transcripts", } # Step 3+4: tile + dice precursors = self._tile_and_dice(transcripts) all_sirnas: List[str] = [] for p in precursors: all_sirnas.extend(p["sirnas"]) # Deduplicate all_sirnas = list(dict.fromkeys(all_sirnas)) print(f"[orchestrator] {len(transcripts)} transcripts -> " f"{len(precursors)} precursors -> {len(all_sirnas)} unique siRNAs") # Step 5: rank ranked = self.ranker.rank_detailed(all_sirnas, top_k=top_k) for r in ranked: r["source_gene"] = next( (p["gene_id"] for p in precursors if r["sirna_seq"] in p["sirnas"]), "unknown", ) # Step 6: safety cards for top 5 top5 = ranked[:5] safety_cards: List[Dict] = [] for cand in top5: if self.degraded_mode: card = self._degraded_safety_card(cand) else: try: card = self.llm.generate_safety_card( sirna_seq=cand["sirna_seq"], offtarget_risks=cand["offtarget_per_species"], half_life_hours=cand["half_life_hours"], ) self._record_tokens(cand["sirna_seq"], card) except Exception as exc: print(f"[orchestrator] generate_safety_card failed ({exc!r}); degraded card") card = self._degraded_safety_card(cand) safety_cards.append({ "sirna_seq": cand["sirna_seq"], "card_markdown": card, }) # Step 7: regulatory memo if self.degraded_mode: memo = self._degraded_regulatory_memo(pest_report, ranked[:5]) else: try: pest_name = ( pest_report.get("pest_species") or "nilaparvata_lugens (brown planthopper)" ) memo = self.llm.generate_regulatory_memo(pest_name, ranked[:5]) # rough token accounting: prompt ~ all candidate strings prompt_blob = pest_name + "".join( c.get("sirna_seq", "") for c in ranked[:5] ) self._record_tokens(prompt_blob, memo) except Exception as exc: print(f"[orchestrator] generate_regulatory_memo failed ({exc!r}); degraded memo") memo = self._degraded_regulatory_memo(pest_report, ranked[:5]) cost = ( self._input_tokens * USD_PER_1K_INPUT_TOKENS + self._output_tokens * USD_PER_1K_OUTPUT_TOKENS ) / 1000.0 return { "pest_report": pest_report, "candidates": ranked, "safety_cards": safety_cards, "regulatory_memo": memo, "total_cost_estimate": cost, "degraded_mode": self.degraded_mode, "n_transcripts": len(transcripts), "n_precursors": len(precursors), "n_sirnas": len(all_sirnas), } # ------------------------------------------------------------------ # # Degraded-mode helpers (used when Ollama is not running) # ------------------------------------------------------------------ # @staticmethod def _degraded_pest_report(user_text: str) -> Dict: """Heuristic pest-species guesser for offline mode. We don't try to be smart -- just match a few common pest/crop keywords so the demo output looks plausible. The real pipeline uses the LLM. """ text = user_text.lower() pest = "nilaparvata_lugens" crop = "rice" if "aphid" in text: pest, crop = "myzus_persicae", "vegetables" if "fall armyworm" in text or "spodoptera" in text: pest, crop = "spodoptera_frugiperda", "maize" if "planthopper" in text or "lugens" in text: pest, crop = "nilaparvata_lugens", "rice" if "stem borer" in text or "chilo" in text: pest, crop = "chilo_suppressalis", "rice" if "locust" in text or "schistocerca" in text: pest, crop = "schistocerca_gregaria", "wheat" if "colorado potato" in text or "leptinotarsa" in text: pest, crop = "leptinotarsa_decemlineata", "potato" if "whitefly" in text or "bemisia" in text: pest, crop = "bemisia_tabaci", "tomato" return { "pest_species": pest, "crop": crop, "severity": "moderate", "location": "Tamil Nadu, India" if "tamil" in text else "unspecified", "notes": "DEGRADED MODE: parsed without LLM (Ollama server not running).", "_raw": "", "_degraded": True, } @staticmethod def _degraded_safety_card(cand: Dict) -> str: """Plain-markdown safety card built from the ranker outputs only.""" ot_lines = "\n".join( f" - {sp}: {risk:.3f}" for sp, risk in cand.get("offtarget_per_species", {}).items() ) or " - (no off-target hits)" risk_tier = ( "high" if cand.get("offtarget_max", 0) > 0.3 else "moderate" if cand.get("offtarget_max", 0) > 0.05 else "low" ) return ( f"# Safety Card (DEGRADED MODE)\n\n" f"## Sequence\n" f"`{cand['sirna_seq']}`\n\n" f"## Off-Target Profile\n" f"{ot_lines}\n\n" f"Max off-target risk: **{cand.get('offtarget_max', 0):.3f}**\n\n" f"## Environmental Fate\n" f"Predicted half-life: **{cand.get('half_life_hours', 0):.2f} hours**\n\n" f"## Overall Risk Tier\n" f"**{risk_tier}**\n\n" f"_Generated without LLM (Ollama server not running)._" ) @staticmethod def _degraded_regulatory_memo(pest_report: Dict, candidates: List[Dict]) -> str: pest = pest_report.get("pest_species", "unknown_pest") crop = pest_report.get("crop", "unknown_crop") lines = [ f"# Regulatory Memo (DEGRADED MODE)\n", f"## Pest & Crop\n", f"Target: **{pest}** on **{crop}**.\n", f"## Candidate Summary\n", ] for i, c in enumerate(candidates, 1): lines.append( f"{i}. `{c['sirna_seq']}` " f"efficacy={c.get('efficacy', 0):.3f} " f"offtarget_max={c.get('offtarget_max', 0):.3f} " f"half_life={c.get('half_life_hours', 0):.2f}h " f"score={c.get('final_score', 0):.3f}" ) lines.append("\n## Risk Assessment\n") any_high = any(c.get("offtarget_max", 0) > 0.3 for c in candidates) any_short = any(c.get("half_life_hours", 99) < 6 for c in candidates) if any_high: lines.append("- At least one candidate has HIGH off-target risk; flag for further screening.") if any_short: lines.append("- At least one candidate has a predicted half-life under 6 hours; field efficacy may be limited.") lines.append("\n## Recommendation\n") if any_high: lines.append("Recommend additional off-target screening before issuing an Experimental Use Permit.") else: lines.append("Candidates look suitable for an Experimental Use Permit application, pending wet-lab validation.") lines.append("\n_Generated without LLM (Ollama server not running)._") return "\n".join(lines) # --------------------------------------------------------------------------- # # CLI # --------------------------------------------------------------------------- # def main(argv: Optional[List[str]] = None) -> int: p = argparse.ArgumentParser(description="Run the end-to-end biopesticide design pipeline.") p.add_argument("--user-text", type=str, required=True, help="Free-text pest report (e.g. 'Brown planthopper infestation in rice paddy in Tamil Nadu').") p.add_argument("--device", type=str, default="auto", choices=["auto", "cpu", "cuda"]) p.add_argument("--pest-fasta", type=str, default=str(DEFAULT_PEST_FASTA)) p.add_argument("--safety-fasta", type=str, default=str(DEFAULT_SAFETY_FASTA)) p.add_argument("--sirna-checkpoint", type=str, default=str(SIRNA_CHECKPOINT)) p.add_argument("--pinn-checkpoint", type=str, default=str(PINN_CHECKPOINT)) p.add_argument("--top-k", type=int, default=10) p.add_argument("--max-transcripts", type=int, default=5) p.add_argument("--pest-species", type=str, default=None, help="skip LLM parsing and use this species directly (e.g. nilaparvata_lugens)") args = p.parse_args(argv) orch = BiopesticideOrchestrator( pest_fasta=Path(args.pest_fasta), safety_fasta=Path(args.safety_fasta), sirna_checkpoint=Path(args.sirna_checkpoint), pinn_checkpoint=Path(args.pinn_checkpoint), device=args.device, max_transcripts=args.max_transcripts, ) result = orch.design(args.user_text, top_k=args.top_k, pest_species_override=args.pest_species) # Pretty-print to stdout print("\n" + "=" * 78) print("DESIGN RESULT") print("=" * 78) print(f"Pest species : {result['pest_report'].get('pest_species')}") print(f"Crop : {result['pest_report'].get('crop')}") print(f"Severity : {result['pest_report'].get('severity')}") print(f"Location : {result['pest_report'].get('location')}") print(f"Transcripts : {result['n_transcripts']}") print(f"Precursors : {result['n_precursors']}") print(f"siRNAs : {result['n_sirnas']}") print(f"Degraded mode: {result['degraded_mode']}") print(f"Cost estimate: ${result['total_cost_estimate']:.4f}") print() print(f"Top {len(result['candidates'])} candidates:") for i, c in enumerate(result["candidates"], 1): print( f" {i:2d}. {c['sirna_seq']} eff={c['efficacy']:.3f} " f"ot_max={c['offtarget_max']:.3f} hl={c['half_life_hours']:.2f}h " f"score={c['final_score']:.3f}" ) print() print("Safety cards (top 5):") for sc in result["safety_cards"]: print(f"--- {sc['sirna_seq']} ---") print(sc["card_markdown"]) print() print("Regulatory memo:") print(result["regulatory_memo"]) return 0 if __name__ == "__main__": sys.exit(main())