Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import logging | |
| from app.agents.cerebras_client import CerebrasClient | |
| from app.schemas.visual_lesson import NucleicVisualizationDraft, PlanCritique, VisualizationRequest | |
| from app.services.nucleic_resolver import ResolvedNucleic | |
| logger = logging.getLogger(__name__) | |
| class NucleicVisualAgent: | |
| def __init__(self, client: CerebrasClient | None = None) -> None: | |
| self.client = client or CerebrasClient() | |
| def fallback(request: VisualizationRequest, resolved: ResolvedNucleic) -> NucleicVisualizationDraft: | |
| text = f"{request.prompt} {request.selection_text}".lower() | |
| if resolved.structure: | |
| primary = "structure" | |
| elif resolved.intent.mode == "comparison": | |
| primary = "comparison" | |
| elif resolved.nucleobase_molecule and any( | |
| term in text for term in ("3d", "molecule", "element", "atom", "ball-and-stick", "model") | |
| ): | |
| primary = "molecule_3d" | |
| elif resolved.intent.focus_base: | |
| primary = "chemistry" | |
| elif any(term in text for term in ("chemistry", "chemical", "nucleotide", "atom", "bond")): | |
| primary = "chemistry" | |
| else: | |
| primary = "helix_3d" | |
| captions = { | |
| "structure": "Official coordinates; rotate, zoom, and focus a verified polymer chain.", | |
| "comparison": "Geometry and chemistry are aligned so structural differences remain visually comparable.", | |
| "chemistry": "Select a base to inspect its bond pattern and connection to sugar and phosphate.", | |
| "helix_3d": "Backbones wind in opposite directions around paired bases; select a base for its chemistry.", | |
| } | |
| if resolved.intent.focus_base and resolved.nucleobase_molecule: | |
| molecule = resolved.nucleobase_molecule | |
| caption = ( | |
| f"{molecule.name.title()} · {molecule.molecular_formula} · PubChem CID {molecule.pubchem_cid}. " | |
| + ("Select an atom to inspect its element and bonds." if primary == "molecule_3d" else "The bond graph shows the isolated nucleobase, without sugar or phosphate.") | |
| ) | |
| else: | |
| caption = captions[primary] | |
| return NucleicVisualizationDraft( | |
| primary_view=primary, | |
| visible_feature_ids=[feature.feature_id for feature in resolved.features[:8]], | |
| caption=caption, | |
| ) | |
| def _plan(self, request: VisualizationRequest, resolved: ResolvedNucleic, feedback: list[str] | None = None) -> NucleicVisualizationDraft: | |
| catalog = "\n".join(f"- {item.feature_id}: {item.label}" for item in resolved.features) | |
| context = "\n".join( | |
| f"- {source.origin}: {source.excerpt[:280]}" | |
| for source in resolved.sources | |
| if source.origin in {"selection", "project", "web"} and source.excerpt | |
| )[:2400] | |
| messages = [ | |
| { | |
| "role": "system", | |
| "content": ( | |
| "Choose a visual composition for a nucleic-acid visualization. The visual must answer the prompt; it is not a lesson. " | |
| "Choose only from the supplied feature IDs and one primary view. Use helix_3d for overall form, chemistry for bonds or nucleotides, " | |
| "molecule_3d only when an official nucleobase conformer exists, comparison only for an actual comparison request, and structure only " | |
| "when official macromolecular coordinates exist. The caption must be one short " | |
| "sentence under 180 characters. Do not create sequences, facts, identifiers, coordinates, labels, commands, or teaching steps." | |
| ), | |
| }, | |
| { | |
| "role": "user", | |
| "content": ( | |
| f"PROMPT:\n{request.prompt}\nMODE: {resolved.intent.mode}\nMOLECULE: {resolved.intent.molecule}\n" | |
| f"HAS OFFICIAL STRUCTURE: {bool(resolved.structure)}\nHAS OFFICIAL NUCLEOBASE CONFORMER: " | |
| f"{bool(resolved.nucleobase_molecule)}\nFEATURES:\n{catalog}\nVISUAL CONTEXT:\n{context or '- none'}" | |
| + ("\nREVISION:\n" + "\n".join(feedback) if feedback else "") | |
| ), | |
| }, | |
| ] | |
| return self.client.structured_complete( | |
| messages, | |
| NucleicVisualizationDraft, | |
| reasoning_effort="medium", | |
| timeout=8.0, | |
| trace_label=f"nucleic-plan:{request.request_id}", | |
| ) | |
| def _critic(self, request: VisualizationRequest, resolved: ResolvedNucleic, draft: NucleicVisualizationDraft) -> PlanCritique: | |
| messages = [ | |
| { | |
| "role": "system", | |
| "content": ( | |
| "Review a visual-first nucleic-acid composition. Reject it if it does not answer the prompt, selects unavailable features, " | |
| "uses structure or molecule_3d without the corresponding verified coordinates, uses comparison for a non-comparison request, " | |
| "or lets prose dominate the visual." | |
| ), | |
| }, | |
| {"role": "user", "content": f"PROMPT: {request.prompt}\nPLAN: {draft.model_dump_json()}\nAVAILABLE: {[item.feature_id for item in resolved.features]}"}, | |
| ] | |
| return self.client.structured_complete( | |
| messages, | |
| PlanCritique, | |
| reasoning_effort="low", | |
| timeout=8.0, | |
| trace_label=f"nucleic-critic:{request.request_id}", | |
| ) | |
| def plan_with_review(self, request: VisualizationRequest, resolved: ResolvedNucleic) -> tuple[NucleicVisualizationDraft, list[str]]: | |
| warnings: list[str] = [] | |
| logger.info("Nucleic visual planner start request=%s", request.request_id) | |
| try: | |
| draft = self._plan(request, resolved) | |
| except Exception as exc: | |
| logger.warning("Nucleic visual planner fallback request=%s error=%s", request.request_id, exc) | |
| return self.fallback(request, resolved), ["Visual composition planning was unavailable; a deterministic prompt-matched composition was used."] | |
| try: | |
| logger.info("Nucleic visual critic start request=%s", request.request_id) | |
| critique = self._critic(request, resolved, draft) | |
| if not critique.approved: | |
| logger.info("Nucleic visual revision start request=%s", request.request_id) | |
| draft = self._plan(request, resolved, critique.feedback) | |
| except Exception as exc: | |
| logger.warning("Nucleic visual critic fallback request=%s error=%s", request.request_id, exc) | |
| warnings.append("Visual composition review was unavailable; the schema-valid composition was used.") | |
| available = {feature.feature_id for feature in resolved.features} | |
| if any(feature_id not in available for feature_id in draft.visible_feature_ids): | |
| draft = self.fallback(request, resolved) | |
| warnings.append("The proposed callouts referenced unavailable features and were replaced deterministically.") | |
| if resolved.features and not draft.visible_feature_ids: | |
| draft = self.fallback(request, resolved) | |
| if draft.primary_view == "structure" and not resolved.structure: | |
| draft = self.fallback(request, resolved) | |
| if draft.primary_view == "molecule_3d" and not resolved.nucleobase_molecule: | |
| draft = self.fallback(request, resolved) | |
| if draft.primary_view == "comparison" and resolved.intent.mode != "comparison": | |
| draft = self.fallback(request, resolved) | |
| required_view = self.fallback(request, resolved).primary_view | |
| if draft.primary_view != required_view: | |
| draft = draft.model_copy(update={"primary_view": required_view}) | |
| logger.info("Nucleic visual planner done request=%s view=%s", request.request_id, draft.primary_view) | |
| return draft, warnings | |