app/agents/modality_router.py

"""Visual Modality Router -> decides WHAT kind of visual fits an explicit
student visualize request.

This is now only ever reached from an explicit student action (a /visualize
command or a natural-language "show me a chart of X" request in Chat) --
Wiki's old click-to-offer path was removed entirely. Because the student has
already asked for a visual, "decline" is not offered as an outcome: the
router always picks exactly one of the 5 real modalities. It does NOT verify
the source has ENOUGH data for a specific chart, nor does it pick a concrete
engine/template -- that finer extract-and-verify-and-synthesize work belongs
to each of the 5 downstream engines individually (mirroring the
D3TemplateRouter.select() -> D3DataExtractor.fill() two-stage split in
app/agents/d3/).

Decision rule: classify primarily from what the student's own wording asks
for (an explicit "chart"/"graph" -> graph, "animate"/"simulate" -> 2d_anim,
"3D model" -> 3d, "formula"/"derive" -> formula, otherwise a conceptual point
-> 2d_text), using the source material only as secondary supporting context
-- never as a reason to refuse an outcome. Downstream engines independently
decide whether to ground the result in the source, the web, or a clearly
labeled illustrative synthesis; this router's only job is picking the shape.
"""
from __future__ import annotations

from typing import List, Literal

from pydantic import BaseModel

from app.agents.cerebras_client import CerebrasClient

_MAX_CHUNK_CHARS = 3000  # mirrors BrainAgent.extract_curriculum cap


class ModalityDecision(BaseModel):
    modality: Literal["formula", "graph", "2d_text", "3d", "2d_anim"]
    reasoning: str


_SYSTEM_PROMPT = """\
You are the Visualization Router for a student research assistant. The student
has explicitly asked for a visual -- your only job is picking which of 5 shapes
best fits their request. Declining is not an option; always pick exactly one.

Classify primarily from the STUDENT'S REQUEST wording itself, using the source
material only as secondary context (it may be sparse, unrelated, or absent --
that never changes which modality is the right shape for what they asked for):

- formula: the student is asking for an equation, derivation, or formula.
- graph: the student is asking for a chart, plot, graph, or a comparison/series
  of data (e.g. "bar chart of X", "plot Y over time").
- 2d_text: the student is asking for a definitional or conceptual explanation
  best served by prose plus at most one citation, not a diagram.
- 3d: the student is asking for a spatial/structural object -- a molecule, an
  anatomical structure, a 3D geometric shape or mechanism.
- 2d_anim: the student is asking for a 2D dynamic process, motion, or
  simulation -- a mechanism, a waveform, a state transition, a physical
  process with movement.

If the request itself doesn't name a type explicitly, infer the best fit from
what's actually being asked about (a described mechanism or motion -> 2d_anim,
a described structure -> 3d, a described relationship worth charting -> graph,
otherwise -> 2d_text).

Output a single JSON object matching the schema. One sentence for reasoning.\
"""


class VisualModalityRouter:
    def __init__(self) -> None:
        self._client = CerebrasClient()

    def classify(
        self,
        selection_text: str,
        card_markdown: str,
        chunks: List[dict],
        familiarity: str,
    ) -> ModalityDecision:
        chunk_text = "\n\n".join(c["text"] for c in chunks)
        if len(chunk_text) > _MAX_CHUNK_CHARS:
            chunk_text = chunk_text[:_MAX_CHUNK_CHARS]

        messages = [
            {"role": "system", "content": _SYSTEM_PROMPT},
            {
                "role": "user",
                "content": (
                    f"STUDENT REQUEST: {selection_text}\n"
                    f"STUDENT LEVEL: {familiarity}\n\n"
                    f"SOURCE MATERIAL (may be sparse or unrelated):\n{chunk_text}\n\n"
                    f"WIKI CARD SUMMARY:\n{card_markdown[:800]}"
                ),
            },
        ]
        return self._client.structured_complete(messages, ModalityDecision, reasoning_effort="medium")