Spaces:
Running
Running
| import os | |
| import logging | |
| from app.config import settings | |
| from app.ai.prompts import get_prompt | |
| logger = logging.getLogger(__name__) | |
| class LLMClient: | |
| def __init__(self): | |
| self.api_key = settings.GROQ_API_KEY | |
| self.fallback_key = settings.GOOGLE_API_KEY | |
| self.model = settings.DEFAULT_MODEL | |
| self.fallback_model = "gemini/gemini-2.0-flash" | |
| self.pro_model = settings.PRO_MODEL | |
| def has_api_key(self) -> bool: | |
| return bool(self.api_key) or bool(self.fallback_key) | |
| def get_providers(self) -> list[dict]: | |
| providers = [] | |
| if self.api_key: | |
| providers.append({"model": self.model, "api_key": self.api_key, "name": "groq"}) | |
| if self.fallback_key: | |
| providers.append({"model": self.fallback_model, "api_key": self.fallback_key, "name": "gemini"}) | |
| return providers | |
| def get_all_candidates(self) -> list[dict]: | |
| """Primary providers plus any explicitly-configured backups, in preference order. | |
| Used by the streaming interpreter so a rate-limited provider can fall | |
| back to another model without the user seeing a raw litellm error. | |
| Only models actually configured in the environment are added — never | |
| guessed names, and never the built-in PRO_MODEL default (an Anthropic | |
| model) which would be routed through the Google API key. | |
| """ | |
| candidates = self.get_providers() | |
| seen = {(c["model"], c["api_key"]) for c in candidates} | |
| backups = [] | |
| pro = os.getenv("PRO_MODEL", "").strip() | |
| if self.fallback_key and pro: | |
| if (pro, self.fallback_key) not in seen: | |
| backups.append({"model": pro, "api_key": self.fallback_key, "name": "gemini-fallback"}) | |
| seen.add((pro, self.fallback_key)) | |
| return candidates + backups | |
| def build_prompt(self, pipeline_type: str, context: dict) -> str: | |
| template = get_prompt(pipeline_type) | |
| blast = context.get("blast", {}) | |
| top = blast.get("top_hit", {}) | |
| uniprot = context.get("uniprot", {}) or {} | |
| af = context.get("alphafold", {}) or {} | |
| values = { | |
| "blast_count": blast.get("count", 0), | |
| "top_hit_accession": top.get("accession", "N/A"), | |
| "top_hit_description": top.get("description", "N/A"), | |
| "top_hit_evalue": top.get("evalue", "N/A"), | |
| "top_hit_identity_pct": top.get("identity_pct", "N/A"), | |
| "top_hit_bit_score": top.get("bit_score", "N/A"), | |
| "uniprot_name": uniprot.get("full_name", "N/A"), | |
| "uniprot_organism": uniprot.get("organism", "N/A"), | |
| "uniprot_genes": ", ".join(uniprot.get("gene_names", []) or []) or "N/A", | |
| "uniprot_functions": "; ".join(uniprot.get("functions", []) or []) or "N/A", | |
| "uniprot_locations": "; ".join(uniprot.get("subcellular_locations", []) or []) or "N/A", | |
| "uniprot_keywords": ", ".join(uniprot.get("keywords", []) or []) or "N/A", | |
| "uniprot_go_terms": ", ".join(uniprot.get("go_terms", []) or []) or "N/A", | |
| "uniprot_features": "; ".join( | |
| f"{f.get('type', '')}: {f.get('description', '')}" for f in (uniprot.get("features", []) or []) | |
| ) or "N/A", | |
| "alphafold_available": "Yes" if af.get("structure_available") else "No", | |
| "alphafold_confidence": af.get("confidence", "N/A"), | |
| } | |
| # Live API text (BLAST descriptions, UniProt functions…) can contain | |
| # { or } characters, which str.format() would try to interpret as | |
| # placeholders and crash the prompt build. Escape them in every value. | |
| escaped = {k: _format_safe(v) for k, v in values.items()} | |
| try: | |
| return template.format(**escaped) | |
| except Exception as e: # last-resort: never let prompt build break the pipeline | |
| logger.warning("Prompt template %s failed to fill (%s); using fallback prompt", pipeline_type, e) | |
| return ( | |
| f"Analyze this protein. BLAST top hit: {values['top_hit_description']} " | |
| f"(E-value {values['top_hit_evalue']}, identity {values['top_hit_identity_pct']}%). " | |
| f"UniProt: {values['uniprot_name']} from {values['uniprot_organism']}. " | |
| f"Functions: {values['uniprot_functions']}." | |
| ) | |
| def _format_safe(value) -> str: | |
| """Escape braces so API-returned text cannot corrupt a str.format() template.""" | |
| return str(value).replace("{", "{{").replace("}", "}}") | |
| llm_client = LLMClient() | |