extractor.py

"""
NeuroJenML — Structured paper extraction.

Turns approved paper text into a standardized, provenance-carrying record.
Content schema (see DATA_HANDLING.md §2.2):
    hypothesis, logic, results, explanations, methods, maps
Every field carries provenance (source section + confidence).

Swappable backends
──────────────────
MedGemmaExtractor (default)
    Uses google/gemma-3-12b-it via the HF Inference Router (free, no dedicated
    endpoint needed).  This is the strongest freely-available instruct model on
    the HF serverless tier and is confirmed accessible without a paid plan.
    Requires: HF_TOKEN  (a standard HF read token — no special model access needed)
    Override: EXTRACTION_MODEL_ID  (any HF chat/instruct model on the router)

FrontierExtractor (optional, high-quality offline curation pass)
    Provider-agnostic: detects Anthropic vs. generic OpenAI-compatible from the
    base URL and uses the correct wire format for each.
    Requires: FRONTIER_API_KEY
    Optional: FRONTIER_API_BASE  (default: https://api.openai.com)
              FRONTIER_MODEL_ID  (default: gpt-4o)
    Examples:
      OpenAI:    FRONTIER_API_BASE=https://api.openai.com    FRONTIER_MODEL_ID=gpt-4o
      Anthropic: FRONTIER_API_BASE=https://api.anthropic.com FRONTIER_MODEL_ID=claude-opus-4-5
      Together:  FRONTIER_API_BASE=https://api.together.xyz  FRONTIER_MODEL_ID=meta-llama/Llama-3-70b-chat-hf
      Groq:      FRONTIER_API_BASE=https://api.groq.com/openai FRONTIER_MODEL_ID=llama3-70b-8192

Both degrade gracefully: if credentials are absent or the call fails, extract()
returns a record with an `error` field instead of raising, so /api/ingest never
fails or blocks.
"""

import os
import json
from abc import ABC, abstractmethod

import httpx

import taxonomy

# ── Content schema ────────────────────────────────────────────────────────────

CONTENT_FIELDS = ("hypothesis", "logic", "results", "explanations", "methods", "maps")

# Default extraction model: Gemma 3 12B instruct is freely available on the HF
# serverless router and outperforms Gemma 2 on structured instruction-following.
_DEFAULT_EXTRACTION_MODEL = "google/gemma-3-12b-it"

_SYSTEM_PROMPT = (
    "You are a biomedical research extraction engine for Alzheimer's disease (AD) "
    "and neurodegeneration. Read the paper text and return STRICT JSON with these keys:\n"
    '  "hypothesis": string — the central hypothesis or aim.\n'
    '  "logic": string — the reasoning/rationale linking premise to expected result.\n'
    '  "results": string — key quantitative/qualitative findings (include effect sizes if present).\n'
    '  "explanations": string — the authors\' mechanistic interpretation.\n'
    '  "methods": string — model system, species, assays, sample size.\n'
    '  "maps": array of {"from": string, "relation": string, "to": string, "confidence": 0-1}\n'
    "\n"
    "CRITICAL RULES for the 'maps' field:\n"
    "- Use SPECIFIC biological pathway names, NOT vague verbs.\n"
    "  BAD:  {\"from\": \"IL-6\", \"relation\": \"influences\", \"to\": \"brain\"}\n"
    "  GOOD: {\"from\": \"IL-6\", \"relation\": \"activates JAK-STAT3 signaling in astrocytes\", \"to\": \"neuroinflammation\"}\n"
    "- Each map must name the actual mechanism (e.g. \"disrupts tight junctions\", \"activates microglial TLR4\", "
    "\"impairs insulin signaling\", \"increases BBB permeability via MMP-9\").\n"
    "- If the paper doesn't specify a mechanism, use \"correlates with\" but add context: "
    "\"correlates with (mechanism unspecified in paper)\".\n"
    "- Include quantitative strength when stated (e.g. \"r=0.85\", \"p<0.01\", \"2.3-fold increase\").\n"
    "\n"
    "Be faithful to the text. If a field is absent, use an empty string or empty array. "
    "Return ONLY the JSON object, no prose, no markdown."
)


# ── Helpers ───────────────────────────────────────────────────────────────────

def _build_prompt(text: str) -> str:
    # Keep within context limits; 12 000 chars ≈ 3 000 tokens.
    snippet = text[:12000]
    return (
        f"{_SYSTEM_PROMPT}\n\n"
        "--- PAPER TEXT ---\n"
        f"{snippet}\n"
        "--- END ---\n\n"
        "JSON:"
    )


def _empty_fields() -> dict:
    return {f: ([] if f == "maps" else "") for f in CONTENT_FIELDS}


def _coerce_fields(obj: dict) -> dict:
    """Normalize a parsed model response into the canonical field shape."""
    fields = _empty_fields()
    if not isinstance(obj, dict):
        return fields
    for f in CONTENT_FIELDS:
        if f == "maps":
            maps = obj.get("maps") or []
            clean = []
            if isinstance(maps, list):
                for m in maps:
                    if isinstance(m, dict) and m.get("from") and m.get("to"):
                        clean.append({
                            "from": str(m.get("from"))[:200],
                            "relation": str(m.get("relation", "relates_to"))[:80],
                            "to": str(m.get("to"))[:200],
                            "confidence": m.get("confidence"),
                        })
            fields["maps"] = clean
        else:
            val = obj.get(f, "")
            fields[f] = val if isinstance(val, str) else json.dumps(val)
    return fields


def _parse_json_lenient(raw: str) -> dict:
    """Extract the first JSON object from a model response."""
    if not raw:
        return {}
    try:
        return json.loads(raw)
    except Exception:
        pass
    start, depth = raw.find("{"), 0
    if start == -1:
        return {}
    for i in range(start, len(raw)):
        if raw[i] == "{":
            depth += 1
        elif raw[i] == "}":
            depth -= 1
            if depth == 0:
                try:
                    return json.loads(raw[start:i + 1])
                except Exception:
                    return {}
    return {}


# ── Base class ────────────────────────────────────────────────────────────────

class Extractor(ABC):
    name = "base"

    @abstractmethod
    async def _generate(self, prompt: str) -> str:
        """Return raw model text for the prompt, or raise on failure."""

    async def extract(self, text: str) -> dict:
        """Extract structured findings + classification + provenance."""
        classification = taxonomy.heuristic_classify(text)
        record = {
            "model": self.name,
            "fields": _empty_fields(),
            "classification": classification,
            "provenance": {"source": "model", "extractor": self.name},
            "error": None,
        }
        if not text.strip():
            record["error"] = "empty text"
            return record
        try:
            raw = await self._generate(_build_prompt(text))
            record["fields"] = _coerce_fields(_parse_json_lenient(raw))
        except Exception as e:
            record["error"] = str(e)[:300]
        return record


# ── HF Router extractor (default, free) ──────────────────────────────────────

class HFExtractor(Extractor):
    """
    Calls any HF Inference Router chat/instruct model via the OpenAI-compatible
    /v1/chat/completions endpoint.

    Default model: google/gemma-3-12b-it
      - Freely available on the HF serverless tier (no gated access, no paid plan).
      - Strongly outperforms smaller models on structured JSON generation.
      - Confirmed on the router as of mid-2025.

    Override EXTRACTION_MODEL_ID with any publicly-accessible HF chat model,
    e.g. "mistralai/Mistral-7B-Instruct-v0.3", "HuggingFaceH4/zephyr-7b-beta".
    """
    name = "hf-gemma"

    def __init__(self):
        self.token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN")
        self.model = os.getenv("EXTRACTION_MODEL_ID", _DEFAULT_EXTRACTION_MODEL)
        self.name = f"hf:{self.model.split('/')[-1]}"

    async def _generate(self, prompt: str) -> str:
        if not self.token:
            raise RuntimeError(
                "HF_TOKEN not set — cannot run extraction. "
                "Add HF_TOKEN to your Vercel environment variables."
            )
        from huggingface_hub import AsyncInferenceClient
        client = AsyncInferenceClient(api_key=self.token)
        resp = await client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.1,
            max_tokens=1500,
        )
        return resp.choices[0].message.content


# ── Frontier extractor (optional, high-quality) ───────────────────────────────

class FrontierExtractor(Extractor):
    """
    Provider-agnostic high-quality extractor for offline curation passes.
    Detects Anthropic vs. generic OpenAI-compatible from FRONTIER_API_BASE.

    Required env vars:
      FRONTIER_API_KEY   — your API key for the provider
    Optional env vars:
      FRONTIER_API_BASE  — provider base URL (default: https://api.openai.com)
      FRONTIER_MODEL_ID  — model to use (default: gpt-4o)

    Provider examples:
      OpenAI:    FRONTIER_API_BASE=https://api.openai.com       FRONTIER_MODEL_ID=gpt-4o
      Anthropic: FRONTIER_API_BASE=https://api.anthropic.com    FRONTIER_MODEL_ID=claude-opus-4-5
      Together:  FRONTIER_API_BASE=https://api.together.xyz     FRONTIER_MODEL_ID=meta-llama/Llama-3-70b-chat-hf
      Groq:      FRONTIER_API_BASE=https://api.groq.com/openai  FRONTIER_MODEL_ID=llama3-70b-8192
      Mistral:   FRONTIER_API_BASE=https://api.mistral.ai       FRONTIER_MODEL_ID=mistral-large-latest
    """
    name = "frontier"

    def __init__(self):
        self.key = os.getenv("FRONTIER_API_KEY")
        self.base = (os.getenv("FRONTIER_API_BASE") or "https://api.openai.com").rstrip("/")
        self.model = os.getenv("FRONTIER_MODEL_ID", "gpt-4o")
        self.name = f"frontier:{self.model.split('/')[-1]}"

    def _is_anthropic(self) -> bool:
        return "anthropic.com" in self.base

    async def _generate(self, prompt: str) -> str:
        if not self.key:
            raise RuntimeError(
                "FRONTIER_API_KEY not set — frontier extraction unavailable. "
                "This is optional; extraction will use the HF model instead."
            )

        if self._is_anthropic():
            return await self._call_anthropic(prompt)
        return await self._call_openai_compatible(prompt)

    async def _call_openai_compatible(self, prompt: str) -> str:
        """Works for OpenAI, Together, Groq, Mistral, and any OpenAI-compatible API."""
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(
                f"{self.base}/v1/chat/completions",
                headers={
                    "Authorization": f"Bearer {self.key}",
                    "Content-Type": "application/json",
                },
                json={
                    "model": self.model,
                    "messages": [{"role": "user", "content": prompt}],
                    "temperature": 0.1,
                    "max_tokens": 2000,
                },
            )
            resp.raise_for_status()
            return resp.json()["choices"][0]["message"]["content"]

    async def _call_anthropic(self, prompt: str) -> str:
        """Anthropic Messages API (different wire format)."""
        async with httpx.AsyncClient(timeout=120) as client:
            resp = await client.post(
                f"{self.base}/v1/messages",
                headers={
                    "x-api-key": self.key,
                    "anthropic-version": "2023-06-01",
                    "Content-Type": "application/json",
                },
                json={
                    "model": self.model,
                    "max_tokens": 2000,
                    "messages": [{"role": "user", "content": prompt}],
                },
            )
            resp.raise_for_status()
            data = resp.json()
            return "".join(block.get("text", "") for block in data.get("content", []))


# ── Factory ───────────────────────────────────────────────────────────────────

def get_extractor(mode: str = "default") -> Extractor:
    """
    Factory function.
      'default' | 'hf' | 'gemma' -> HFExtractor (google/gemma-3-12b-it by default)
      'frontier'                  -> FrontierExtractor
    """
    if mode == "frontier":
        return FrontierExtractor()
    return HFExtractor()


# Keep the old name as an alias so any code referencing MedGemmaExtractor still works.
MedGemmaExtractor = HFExtractor