study-buddy / app /agents /formula_engine.py
GitHub Actions
deploy d092bea3608b7a29952f16357fda39b7a29e399b
2e818da
Raw
History Blame Contribute Delete
3.57 kB
from __future__ import annotations
from typing import List, Optional
from pydantic import BaseModel, Field
from app.agents.cerebras_client import CerebrasClient
from app.schemas.graph import FormulaContent, FormulaStep, HTML5VisualPayload
_MAX_CHUNK_CHARS = 3000 # mirrors ModalityRouter._MAX_CHUNK_CHARS / BrainAgent.extract_curriculum cap
_DEFAULT_DECLINE_REASON = "This concept doesn't have an explicit formula in the source β€” better explored in chat."
class FormulaGrounding(BaseModel):
renderable: bool
anchor: str = ""
main_latex: str = ""
steps: List[FormulaStep] = Field(default_factory=list)
decline_reason: str = ""
class FormulaEngine:
"""Two-stage grounded formula extraction: extract (with a self-checked verbatim
anchor) then render into an HTML5VisualPayload. Mirrors the D3TemplateRouter/
D3DataExtractor split and the old TutorAgent visual-grounding pipeline."""
def __init__(self, client: Optional[CerebrasClient] = None) -> None:
self._client = client or CerebrasClient()
def generate(self, concept: str, chunks: List[dict], familiarity: str) -> HTML5VisualPayload:
from app.agents.tutor_agent import TutorAgent
grounding = self._extract(concept, chunks, familiarity)
if not grounding.renderable:
reason = grounding.decline_reason or _DEFAULT_DECLINE_REASON
return HTML5VisualPayload(
html_code=TutorAgent._decline_html(concept, reason),
animation_type="2d_text",
explanation=reason,
)
return HTML5VisualPayload(
html_code="",
animation_type="formula",
formula=FormulaContent(main_latex=grounding.main_latex, steps=grounding.steps),
explanation=(f"From the source: {grounding.anchor}" if grounding.anchor else ""),
source="paper" if grounding.anchor else "model_knowledge",
)
def _extract(self, concept: str, chunks: List[dict], familiarity: str) -> FormulaGrounding:
chunk_text = "\n\n".join(f"[Source: {c.get('source', '?')}]: {c['text']}" for c in chunks)[:_MAX_CHUNK_CHARS]
messages = [
{"role": "system", "content": (
"You extract a grounded formula/equation for a concept from the SOURCE MATERIAL only. "
"If the source contains an explicit formula/equation for this concept, copy a VERBATIM "
"30-80 character substring of the source's actual equation/formula line as `anchor` β€” "
"it must match the source text literally, not a paraphrase. Fill `main_latex` and `steps` "
"using ONLY numbers, coefficients, and notation actually present in the source β€” never "
"invent or estimate values. If no explicit formula/equation exists in the source, set "
"renderable=false and give a short, calm, student-facing `decline_reason` (e.g. "
f"\"'{concept}' doesn't have an explicit formula in the source β€” better explored in "
"chat.\").\n\n"
f"SOURCE MATERIAL:\n{chunk_text}"
)},
{"role": "user", "content": f"Concept: '{concept}' (level: {familiarity}). Extract the formula or decline."},
]
grounding = self._client.structured_complete(messages, FormulaGrounding, reasoning_effort="medium")
if grounding.anchor and grounding.anchor.strip() not in chunk_text:
grounding = grounding.model_copy(update={"anchor": ""})
return grounding