app/agents/d3/extractor.py

from __future__ import annotations
import re
from decimal import Decimal, InvalidOperation
from typing import Any, Iterator, List, Optional
from pydantic import BaseModel
from app.agents.cerebras_client import CerebrasClient
from app.agents.d3.registry import D3Template
from app.agents.senses_agent import VISION_MODEL_ID


class D3SourceValidationError(ValueError):
    pass


class D3TransformNarrative(BaseModel):
    answer_markdown: str


def _numeric_leaves(value: Any) -> Iterator[Decimal]:
    if isinstance(value, BaseModel):
        for field_value in value.__dict__.values():
            yield from _numeric_leaves(field_value)
    elif isinstance(value, dict):
        for field_value in value.values():
            yield from _numeric_leaves(field_value)
    elif isinstance(value, (list, tuple)):
        for field_value in value:
            yield from _numeric_leaves(field_value)
    elif isinstance(value, (int, float)) and not isinstance(value, bool):
        yield Decimal(str(value)).normalize()
    elif isinstance(value, str) and re.fullmatch(r"\s*[-+]?\d[\d,]*(?:\.\d+)?%?\s*", value):
        cleaned = value.strip().replace(",", "").rstrip("%")
        try:
            yield Decimal(cleaned).normalize()
        except InvalidOperation:
            return


def _source_numbers(text: str) -> set[Decimal]:
    found: set[Decimal] = set()
    for token in re.findall(r"[-+]?\d[\d,]*(?:\.\d+)?", text):
        try:
            found.add(Decimal(token.replace(",", "")).normalize())
        except InvalidOperation:
            continue
    return found


class D3DataExtractor:
    def __init__(self, client: Optional[CerebrasClient] = None) -> None:
        self._client = client or CerebrasClient()

    def fill(
        self,
        template: D3Template,
        concept: str,
        chunks: List[dict],
        familiarity: str,
        *,
        strict_source: bool = False,
        image_base64: str = "",
        selected_metric: str = "",
    ) -> BaseModel:
        chunk_text = "\n\n".join(f"[{c.get('source','?')}]: {c['text']}" for c in chunks)[:3000]
        if strict_source:
            extraction_rule = (
                "Transcribe only labels and values that are visibly present in the SELECTED SOURCE. "
                "Never invent, estimate, interpolate, calculate an unstated aggregate, or substitute example data. "
                "Keep model/method names exactly associated with their source row. If a metric is specified, use "
                "only that column. If the schema cannot be filled faithfully, fail instead of fabricating output."
            )
        else:
            extraction_rule = (
                "Prefer real numbers, labels, and relationships found verbatim in the SOURCE MATERIAL. If the source "
                "does not contain the data this chart needs, invent a plausible representative illustrative example."
            )
        messages = [
            {"role": "system", "content": (
                f"You provide data for a specific trusted visualization. {extraction_rule}\n\n"
                f"CHART: {template.title} — {template.data_requirements}\n"
                f"SELECTED METRIC: {selected_metric or 'not explicitly named'}\n\nSELECTED SOURCE:\n{chunk_text}"
            )},
            {"role": "user", "content": (
                ([{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
                  {"type": "text", "text": f"REQUEST: {concept}\nLEVEL: {familiarity}\nFill the schema from this selected image and any selected text."}]
                 if image_base64 else f"REQUEST: {concept}\nLEVEL: {familiarity}\nFill the schema from the selected text.")
            )},
        ]
        # Never declines -- the system prompt above asks the model to synthesize a
        # plausible illustrative example when the source lacks real data, rather
        # than leave the schema empty. D3Engine reports whether the result actually
        # came from the source via a separate grounding heuristic (see d3/engine.py).
        data = self._client.structured_complete(
            messages,
            template.schema,
            model=VISION_MODEL_ID if image_base64 else None,
            reasoning_effort="low" if strict_source else "medium",
        )
        if strict_source and not image_base64:
            allowed_numbers = _source_numbers(chunk_text)
            missing = sorted({value for value in _numeric_leaves(data) if value not in allowed_numbers})
            if missing:
                raise D3SourceValidationError(
                    "Extracted visualization contains numeric values absent from the selected source: "
                    + ", ".join(str(value) for value in missing[:6])
                )
        return data

    def explain(
        self,
        *,
        concept: str,
        template: D3Template,
        data: BaseModel,
        chunks: List[dict],
        image_base64: str = "",
    ) -> str:
        """Answer the analytical part of a transformation request from the same source."""
        chunk_text = "\n\n".join(str(c.get("text", "")) for c in chunks)[:5000]
        system = (
            "Give a concise research-assistant answer accompanying a source-selected visualization. "
            "Directly answer the user's analytical question in at most four short sentences. Use only "
            "facts and numeric comparisons visible in the selected source or the extracted chart data. "
            "For improvement questions, identify the compared baseline and where values increased; do "
            "not call a larger value an improvement when the metric direction is unknown. Do not discuss "
            "the routing process, template selection, or extraction. If the source cannot support part of "
            "the requested interpretation, state that limitation briefly."
        )
        text = (
            f"REQUEST:\n{concept}\n\nVISUAL FORM: {template.title}\n\n"
            f"EXTRACTED DATA:\n{data.model_dump_json()}\n\nSELECTED TEXT:\n{chunk_text}"
        )
        content: Any = text
        if image_base64:
            content = [
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
                {"type": "text", "text": text},
            ]
        result = self._client.structured_complete(
            [{"role": "system", "content": system}, {"role": "user", "content": content}],
            D3TransformNarrative,
            model=VISION_MODEL_ID if image_base64 else None,
            reasoning_effort="low",
        )
        return result.answer_markdown.strip()