dynamics_inference.py

"""
Prompt construction and inference backend integration for the DYNAMICS
Persona Explorer Space.

This module builds a personality-conditioned prompt from DYNAMICS dimensions,
sends it to the best available inference backend, and constructs a
deterministic reasoning trace from DYNAMICS activations.

Backend priority:
  1. Local model (Animus/Imprint) via LOCAL_INFERENCE_URL
  2. Gemini 2.5 Flash via GEMINI_API_KEY (rate-limited: 5 req/min per user)
  3. HF Inference API (public demo model, no key required)
"""

from __future__ import annotations

import os
import re
import time
import threading
from typing import Any

import requests


# ---------------------------------------------------------------------------
# Prompt template (from PREPRINT_OUTLINE.md Appendix A -- safe to publish)
# ---------------------------------------------------------------------------

_PROMPT_TEMPLATE = """\
You are a synthetic research participant. Your personality profile is as follows:

Discipline: {D:.2f} (scale 0=very low, 1=very high)
Yielding: {Y:.2f}
Novelty: {N:.2f}
Acuity: {A:.2f}
Mercuriality: {M:.2f}
Impulsivity: {I:.2f}
Candour: {C:.2f}
Sociability: {S:.2f}

Your current economic situation:
- Income: approximately {monthly_income} per month
- Current account balance: {current_balance}
- Financial anxiety: {financial_anxiety_label}

Your current emotional state:
- Feeling: {dominant_emotion}
- Mood valence: {valence_label}

You are presented with the following:
"{stimulus}"

Respond as this person would in first person. Be authentic to the personality profile.
Do not mention your DYNAMICS scores or the fact that you are synthetic.
Give a detailed, thoughtful response of 6-10 sentences. Include your reasoning, how you feel about it, what you would actually do, and why. Show how your personality shapes your reaction."""


def _valence_label(valence: float) -> str:
    """Convert numeric valence to a human-readable label."""
    if valence < -0.5:
        return "quite negative"
    if valence < -0.15:
        return "somewhat negative"
    if valence < 0.15:
        return "neutral"
    if valence < 0.5:
        return "somewhat positive"
    return "quite positive"


def build_prompt(
    dynamics: dict[str, float],
    income_band: str,
    balance: float,
    emotional_state: dict[str, Any],
    stimulus: str,
    monthly_income: float | None = None,
    financial_anxiety_label: str = "moderate",
) -> str:
    """Construct the inference prompt from persona parameters.

    Parameters
    ----------
    dynamics : dict
        Eight DYNAMICS dimensions (D, Y, N, A, M, I, C, S) as floats.
    income_band : str
        One of low, mid, high, wealthy.
    balance : float
        Current balance in GBP.
    emotional_state : dict
        Keys: valence (float), dominant_emotion (str).
    stimulus : str
        The stimulus text presented to the persona.
    monthly_income : float or None
        Monthly income in GBP.  Derived from band if not provided.
    financial_anxiety_label : str
        One of low, moderate, high.

    Returns
    -------
    str  The fully formatted prompt.
    """
    from dynamics_rules import default_income_for_band

    mi = monthly_income if monthly_income is not None else default_income_for_band(income_band)

    valence = emotional_state.get("valence", 0.0)
    dominant_emotion = emotional_state.get("dominant_emotion", "neutral")

    return _PROMPT_TEMPLATE.format(
        D=dynamics.get("D", 0.5),
        Y=dynamics.get("Y", 0.5),
        N=dynamics.get("N", 0.5),
        A=dynamics.get("A", 0.5),
        M=dynamics.get("M", 0.5),
        I=dynamics.get("I", 0.5),
        C=dynamics.get("C", 0.5),
        S=dynamics.get("S", 0.5),
        monthly_income=f"\u00a3{mi:,.2f}",
        current_balance=f"\u00a3{balance:,.2f}",
        financial_anxiety_label=financial_anxiety_label,
        dominant_emotion=dominant_emotion,
        valence_label=_valence_label(valence),
        stimulus=stimulus,
    )


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

# Primary: local model endpoint (Animus/Imprint or any OpenAI-compatible API)
_LOCAL_INFERENCE_URL = os.environ.get("LOCAL_INFERENCE_URL", "")

# Fallback: Gemini 2.5 Flash via google-generativeai
_GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "")
_GEMINI_MODEL = "gemini-2.5-flash"

# Third tier: HF Inference API (public, no key needed)
_HF_API_URL = "https://api-inference.huggingface.co/models/{model_id}"
_DEFAULT_HF_MODEL = "kronaxis/dynamics-llm-reasoning-demo"

# Gemini rate limit: 5 requests per minute per user
_GEMINI_RATE_LIMIT_PER_MINUTE = 5
_gemini_request_log: dict[str, list[float]] = {}
_gemini_rate_lock = threading.Lock()


def _gemini_rate_limited(session_id: str) -> bool:
    """Check whether the session has exceeded the Gemini per-minute rate limit."""
    now = time.time()
    with _gemini_rate_lock:
        log = _gemini_request_log.get(session_id, [])
        log = [t for t in log if now - t < 60.0]
        if len(log) >= _GEMINI_RATE_LIMIT_PER_MINUTE:
            _gemini_request_log[session_id] = log
            return True
        log.append(now)
        _gemini_request_log[session_id] = log
        return False


# ---------------------------------------------------------------------------
# Lazy Gemini client initialisation
# ---------------------------------------------------------------------------

_gemini_client = None
_gemini_init_attempted = False


def _get_gemini_client():
    """Lazily initialise the Gemini client (new google.genai SDK)."""
    global _gemini_client, _gemini_init_attempted
    if _gemini_init_attempted:
        return _gemini_client
    _gemini_init_attempted = True
    if not _GEMINI_API_KEY:
        return None
    try:
        from google import genai
        _gemini_client = genai.Client(api_key=_GEMINI_API_KEY)
        return _gemini_client
    except Exception:
        # Fallback to old library
        try:
            import google.generativeai as genai_old
            genai_old.configure(api_key=_GEMINI_API_KEY)
            _gemini_client = genai_old.GenerativeModel(_GEMINI_MODEL)
            return _gemini_client
        except Exception:
            return None


# ---------------------------------------------------------------------------
# Backend 1: Local model (Animus/Imprint -- OpenAI-compatible endpoint)
# ---------------------------------------------------------------------------

def _call_local(prompt: str) -> dict[str, str] | None:
    """Call the local inference endpoint. Returns None on failure."""
    if not _LOCAL_INFERENCE_URL:
        return None
    url = _LOCAL_INFERENCE_URL.rstrip("/")
    # Try OpenAI-compatible /v1/chat/completions first
    chat_url = f"{url}/v1/chat/completions"
    payload = {
        "model": "default",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 300,
        "temperature": 0.7,
    }
    try:
        resp = requests.post(chat_url, json=payload, timeout=30)
        resp.raise_for_status()
        data = resp.json()
        choices = data.get("choices", [])
        if choices:
            text = choices[0].get("message", {}).get("content", "").strip()
            if text:
                return {"raw_text": text, "reasoning_trace": ""}
        return None
    except requests.RequestException:
        pass

    # Fallback: plain text /generate endpoint
    try:
        resp = requests.post(
            f"{url}/generate",
            json={"prompt": prompt, "max_new_tokens": 300, "temperature": 0.7},
            timeout=30,
        )
        resp.raise_for_status()
        data = resp.json()
        text = data.get("text", data.get("generated_text", "")).strip()
        if text:
            return {"raw_text": text, "reasoning_trace": ""}
        return None
    except requests.RequestException:
        return None


# ---------------------------------------------------------------------------
# Backend 2: Gemini 2.5 Flash via google-generativeai
# ---------------------------------------------------------------------------

def _call_gemini(prompt: str) -> dict[str, str] | None:
    """Call Gemini 2.5 Flash. Returns None on failure."""
    client = _get_gemini_client()
    if client is None:
        return None
    try:
        # New google.genai SDK
        from google import genai
        from google.genai import types
        response = client.models.generate_content(
            model=_GEMINI_MODEL,
            contents=prompt,
            config=types.GenerateContentConfig(
                max_output_tokens=1024,
                temperature=0.7,
                thinking_config=types.ThinkingConfig(thinking_budget=0),
            ),
        )
        text = response.text.strip() if response and response.text else ""
        if text:
            return {"raw_text": text, "reasoning_trace": ""}
        return {"raw_text": "No response generated.", "reasoning_trace": ""}
    except ImportError:
        # Fallback: old google.generativeai SDK (no thinking_config support)
        try:
            response = client.generate_content(
                prompt,
                generation_config={
                    "max_output_tokens": 1024,
                    "temperature": 0.7,
                },
            )
            text = response.text.strip() if response and response.text else ""
            if text:
                return {"raw_text": text, "reasoning_trace": ""}
            return {"raw_text": "No response generated.", "reasoning_trace": ""}
        except Exception as exc:
            return {"raw_text": f"Gemini error: {exc}", "reasoning_trace": ""}
    except Exception as exc:
        return {"raw_text": f"Gemini error: {exc}", "reasoning_trace": ""}


# ---------------------------------------------------------------------------
# Backend 3: HF Inference API (public demo model)
# ---------------------------------------------------------------------------

def _call_hf(prompt: str, model_id: str, max_new_tokens: int, temperature: float) -> dict[str, str] | None:
    """Call HF Inference API. Returns None on failure."""
    api_token = os.environ.get("HF_TOKEN", "")
    headers: dict[str, str] = {}
    if api_token:
        headers["Authorization"] = f"Bearer {api_token}"

    payload = {
        "inputs": prompt,
        "parameters": {
            "max_new_tokens": max_new_tokens,
            "temperature": temperature,
            "return_full_text": False,
        },
    }

    url = _HF_API_URL.format(model_id=model_id)
    try:
        resp = requests.post(url, headers=headers, json=payload, timeout=60)
        resp.raise_for_status()
        data = resp.json()
    except requests.RequestException:
        return None

    if isinstance(data, list) and len(data) > 0:
        generated = data[0].get("generated_text", "")
    elif isinstance(data, dict):
        generated = data.get("generated_text", str(data))
    else:
        generated = str(data)

    parts = re.split(r"\n\s*Reasoning:\s*", generated, maxsplit=1)
    raw_text = parts[0].strip()
    reasoning = parts[1].strip() if len(parts) > 1 else ""
    if raw_text:
        return {"raw_text": raw_text, "reasoning_trace": reasoning}
    return None


# ---------------------------------------------------------------------------
# Availability check (used by UI to show status)
# ---------------------------------------------------------------------------

def get_backend_status() -> dict[str, bool]:
    """Return availability of each backend for display in the UI."""
    return {
        "local": bool(_LOCAL_INFERENCE_URL),
        "gemini": bool(_GEMINI_API_KEY),
        "hf": True,  # always available (public endpoint)
    }


def get_available_provider_label() -> str:
    """Return a human-readable label for the best available provider."""
    if _LOCAL_INFERENCE_URL:
        return "local model"
    if _GEMINI_API_KEY:
        return "Gemini 2.5 Flash (cloud)"
    return "HF Inference API (public demo)"


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

def call_inference(
    prompt: str,
    model_id: str = _DEFAULT_HF_MODEL,
    max_new_tokens: int = 300,
    temperature: float = 0.7,
    session_id: str = "default",
) -> dict[str, str]:
    """Send a prompt to the best available inference backend.

    Priority order:
      1. Local model (LOCAL_INFERENCE_URL) -- primary, no rate limit
      2. Gemini 2.5 Flash (GEMINI_API_KEY) -- fallback, 5 req/min per user
      3. HF Inference API -- public demo model, always available

    Returns a dict with keys: raw_text, reasoning_trace, provider.
    """
    # 1. Try local model first
    if _LOCAL_INFERENCE_URL:
        result = _call_local(prompt)
        if result is not None:
            result["provider"] = "local model"
            return result

    # 2. Try Gemini if key is set and rate limit not exceeded
    if _GEMINI_API_KEY:
        if _gemini_rate_limited(session_id):
            # Rate limited but Gemini is configured: skip to HF or return error
            pass
        else:
            result = _call_gemini(prompt)
            if result is not None:
                result["provider"] = "Gemini 2.5 Flash (cloud)"
                return result

    # 3. Try HF Inference API
    result = _call_hf(prompt, model_id, max_new_tokens, temperature)
    if result is not None:
        result["provider"] = "HF Inference API"
        return result

    # All backends failed
    return {
        "raw_text": "All inference backends are currently unavailable. Please try again later.",
        "reasoning_trace": "",
        "provider": "none",
    }


# ---------------------------------------------------------------------------
# Deterministic reasoning trace construction
# ---------------------------------------------------------------------------

_DIMENSION_LABELS = {
    "D": "Discipline",
    "Y": "Yielding",
    "N": "Novelty",
    "A": "Acuity",
    "M": "Mercuriality",
    "I": "Impulsivity",
    "C": "Candour",
    "S": "Sociability",
}

_HIGH_DRIVERS = {
    "D": "D_high_cost_benefit",
    "Y": "Y_high_compliance",
    "N": "N_high_novelty_seeking",
    "A": "A_high_analytical",
    "M": "M_high_anxiety",
    "I": "I_high_spontaneous",
    "C": "C_high_ethical",
    "S": "S_high_social_engagement",
}

_LOW_DRIVERS = {
    "D": "D_low_disorganised",
    "Y": "Y_low_confrontational",
    "N": "N_low_conservative",
    "A": "A_low_surface_reasoning",
    "M": "M_low_calm",
    "I": "I_low_measured",
    "C": "C_low_self_serving",
    "S": "S_low_reserved",
}


def build_reasoning_trace(
    dynamics: dict[str, float],
    response: dict[str, str],
    stimulus: str,
    financial_anxiety_label: str = "moderate",
    income_band: str = "mid",
    balance: float = 0.0,
    monthly_income: float = 0.0,
) -> dict[str, Any]:
    """Construct a deterministic reasoning trace from DYNAMICS activations.

    The trace identifies which dimensions most strongly drove the response
    and assigns causal labels consistent with published literature.

    Returns a dict matching the reasoning_trace schema from DATASET_SPEC.md.
    """
    # Identify the two strongest drivers (furthest from 0.5 midpoint).
    deviations = sorted(
        ((dim, abs(dynamics.get(dim, 0.5) - 0.5)) for dim in _DIMENSION_LABELS),
        key=lambda x: x[1],
        reverse=True,
    )

    top_drivers: list[str] = []
    driver_explanations: list[str] = []
    for dim, dev in deviations[:3]:
        val = dynamics.get(dim, 0.5)
        label = _DIMENSION_LABELS[dim]
        if val >= 0.5:
            top_drivers.append(_HIGH_DRIVERS[dim])
            driver_explanations.append(
                f"high {label} ({val:.2f}) promotes "
                + _HIGH_DRIVERS[dim].split("_", 1)[1].replace("_", " ")
            )
        else:
            top_drivers.append(_LOW_DRIVERS[dim])
            driver_explanations.append(
                f"low {label} ({val:.2f}) promotes "
                + _LOW_DRIVERS[dim].split("_", 1)[1].replace("_", " ")
            )

    # Economic driver.
    if balance > 0 and monthly_income > 0:
        ratio = balance / monthly_income
        if ratio < 0.3:
            econ_driver = "low_balance_relative_to_income"
        elif ratio > 1.5:
            econ_driver = "high_savings_buffer"
        else:
            econ_driver = "moderate_financial_position"
    else:
        econ_driver = f"income_band_{income_band}"

    # Narrative assembly.
    narrative = (
        f"The persona ({income_band} income, balance \u00a3{balance:,.2f}, "
        f"financial anxiety: {financial_anxiety_label}) responded to the stimulus. "
        + "; ".join(driver_explanations)
        + "."
    )

    # Confidence: higher when top dimensions are far from midpoint.
    top_devs = [d[1] for d in deviations[:3]]
    confidence = min(0.95, 0.5 + sum(top_devs) / 3.0)

    return {
        "narrative": narrative,
        "dynamics_drivers": top_drivers,
        "economic_driver": econ_driver,
        "memory_influence": "Not applicable (live demo, no memory history)",
        "confidence": round(confidence, 2),
    }