src/agents.py

"""MedGemma/TxGemma adapters plus auditable deterministic fallbacks."""
from __future__ import annotations

import json
import os
import re

from huggingface_hub import InferenceClient

WOUND_TARGETS = {
    "FGF2": ("down", ["fibroblast growth factor", "fgf2"]),
    "VEGFA": ("down", ["vascular endothelial growth factor", "vegfa", "angiogenesis"]),
    "IL8": ("up", ["interleukin-8", "il-8", "il8"]),
    "TNF": ("up", ["tumor necrosis factor", "tnf-alpha", "tnf"]),
    "MMP9": ("up", ["matrix metalloproteinase-9", "mmp-9", "mmp9"]),
    "TGFB1": ("context-dependent", ["transforming growth factor beta", "tgf-beta", "tgfb1"]),
    "HIF1A": ("down", ["hypoxia-inducible factor", "hif-1", "hif1a"]),
    "EGF": ("down", ["epidermal growth factor", " egf"]),
}


def _extract_json(text: str):
    match = re.search(r"\[.*\]", text, re.DOTALL)
    return json.loads(match.group(0)) if match else None


def medgemma_targets(disease: str, articles: list[dict], max_targets: int = 6, custom_prompt: str = "") -> tuple[list[dict], str]:
    endpoint = os.getenv("MEDGEMMA_ENDPOINT")
    token = os.getenv("HF_TOKEN")
    context = "\n\n".join(f"PMID {a['pmid']}: {a['title']}\n{a['abstract'][:1800]}" for a in articles[:12])
    if endpoint:
        prompt = f"""{custom_prompt or 'Extract disease-relevant molecular targets and their direction in disease.'}
Disease or phenotype: {disease}.
Return only a JSON array. Each object must contain target (official gene symbol), disease_direction (up/down/context-dependent/unknown), confidence (0-1), pmids (array), and evidence_sentence. Never invent a PMID or claim.
EVIDENCE:\n{context}"""
        try:
            client = InferenceClient(base_url=endpoint, token=token)
            result = client.chat.completions.create(messages=[{"role": "user", "content": prompt}], max_tokens=1200, temperature=0.1)
            parsed = _extract_json(result.choices[0].message.content)
            if parsed:
                return parsed[:max_targets], "MedGemma endpoint"
        except Exception:
            pass
    corpus = " ".join(f"{a['title']} {a['abstract']}" for a in articles).lower()
    ranked = []
    for symbol, (direction, aliases) in WOUND_TARGETS.items():
        count = sum(corpus.count(alias) for alias in aliases)
        if count:
            supporting = [a for a in articles if any(alias in (a["title"] + " " + a["abstract"]).lower() for alias in aliases)]
            ranked.append({"target": symbol, "disease_direction": direction, "confidence": min(0.82, 0.38 + 0.07 * count), "pmids": [a["pmid"] for a in supporting[:4]], "evidence_sentence": f"Mentioned in {len(supporting)} retrieved PubMed records; direction is a teaching prior requiring review."})
    ranked.sort(key=lambda x: x["confidence"], reverse=True)
    if not ranked:
        ranked = [{"target": "VEGFA", "disease_direction": "down", "confidence": 0.25, "pmids": [], "evidence_sentence": "Fallback hypothesis; insufficient directional evidence retrieved."}]
    return ranked[:max_targets], "transparent literature heuristic"


def txgemma_note(row: dict) -> tuple[str, str]:
    endpoint = os.getenv("TXGEMMA_ENDPOINT")
    token = os.getenv("HF_TOKEN")
    if endpoint:
        try:
            client = InferenceClient(base_url=endpoint, token=token)
            prompt = f"Assess this therapeutic hypothesis cautiously in <=45 words. Separate measured evidence from inference: {json.dumps(row, default=str)}"
            result = client.chat.completions.create(messages=[{"role": "user", "content": prompt}], max_tokens=100, temperature=0.1)
            return result.choices[0].message.content, "TxGemma endpoint"
        except Exception:
            pass
    return "Measured ChEMBL target activity supports molecular interaction; disease-specific direction, exposure, safety, and wound efficacy remain unproven.", "rules-based interpretation"