SolusOps/Study-with-ChampAI

1from __future__ import annotations2from dataclasses import dataclass, field3from typing import List, Dict4from services.model_router import ModelRouter5from services.json_parser import extract_json6from config.prompts import DOCUMENT_EXTRACT_SYSTEM, DOCUMENT_VISION_PROMPT7 8@dataclass9class DocumentConcepts:10    topics: List[str]11    definitions: List[Dict[str, str]]12    facts: List[str]13    formulae: List[str]14    ocr_text: str = ""  # raw OCR output (if from image)15 16class DocumentAgent:17    """18    Uses MiniCPM-V for all document understanding.19    - Text input: MiniCPM in text mode → concept extraction20    - Image input: MiniCPM in vision mode → OCR + concept extraction in one call21    Nemotron is never called here.22    """23    def __init__(self, router: ModelRouter): self._router = router24 25    def extract(self, raw_text: str) -> DocumentConcepts:26        """Text path: MiniCPM extracts structured concepts from text."""27        prompt = f"{DOCUMENT_EXTRACT_SYSTEM}\n\nExtract concepts from:\n\n{raw_text}"28        raw = self._router.understand(prompt=prompt)29        try:30            data = extract_json(raw)31        except ValueError as exc:32            raise ValueError(f"DocumentAgent: could not parse JSON. {exc}") from exc33        return DocumentConcepts(topics=data.get("topics",[]), definitions=data.get("definitions",[]),34                                facts=data.get("facts",[]), formulae=data.get("formulae",[]))35 36    def extract_from_image(self, image_b64: str) -> DocumentConcepts:37        """38        Image path: MiniCPM does OCR + concept extraction in a single vision call.39        Returns concepts AND the raw OCR text (stored in ocr_text for downstream use).40        """41        raw = self._router.understand(prompt=DOCUMENT_VISION_PROMPT, image_b64=image_b64)42        try:43            data = extract_json(raw)44        except ValueError as exc:45            raise ValueError(f"DocumentAgent: could not parse JSON from image. {exc}") from exc46        return DocumentConcepts(topics=data.get("topics",[]), definitions=data.get("definitions",[]),47                                facts=data.get("facts",[]), formulae=data.get("formulae",[]),48                                ocr_text=data.get("ocr_text",""))49