gateway/ensemble.py

# -*- coding: utf-8 -*-
"""
Darwin-60B-DUO Ensemble V_1 — MAJ@N self-consistency + cross-verification.

For MCQ / short-answer queries:
  1) Each backend produces N samples at temperature τ (default 0.7)
  2) Each backend's answer = its own majority vote (RSA / self-consistency)
  3) If both majorities agree → return that answer
  4) If they disagree → each backend verifies the pair (cross-verification)
     and the gateway picks the tournament winner
  5) Tiebreaker on split verdicts: majority-vote-count confidence
"""
import asyncio
import re
from collections import Counter
from typing import Any, Dict, List, Optional, Tuple


_LETTERS = "ABCD"


def _extract_letter(text: str) -> str:
    """Extract A/B/C/D letter answer from a free-form response."""
    if not text:
        return ""
    # Strip CoT / thinking tags
    cleaned = re.sub(r"<\|START_THINKING\|>.*?<\|END_THINKING\|>", "", text, flags=re.S)
    cleaned = re.sub(r"<think>.*?</think>", "", cleaned, flags=re.S)
    for tag in ["<|END_THINKING|>", "</think>", "<|START_RESPONSE|>", "<|END_RESPONSE|>"]:
        if tag in cleaned:
            cleaned = cleaned.split(tag)[-1]
    # Common answer patterns
    patterns = [
        r"ANSWER:\s*\(?([A-D])\)?",
        r"\\boxed\{\s*\(?([A-D])\)?\s*\}",
        r"final answer\s*(?:is|:)?\s*\(?([A-D])\)?",
        r"answer\s+is\s*\(?([A-D])\)?",
        r"\(([A-D])\)\s*$",
    ]
    for p in patterns:
        m = re.search(p, cleaned, re.I | re.M)
        if m:
            return m.group(1).upper()
    # Fallback: last A-D token
    candidates = re.findall(r"\b([A-D])\b", cleaned)
    return candidates[-1].upper() if candidates else ""


def _majority(letters: List[str]) -> Tuple[Optional[str], Dict[str, int]]:
    valid = [l for l in letters if l in _LETTERS]
    if not valid:
        return None, {}
    counter = Counter(valid)
    top, _ = counter.most_common(1)[0]
    return top, dict(counter)


_VERIFY_TEMPLATE = (
    "You are a graduate-level expert verifier. Given the following multiple-"
    "choice question and two candidate letter answers, decide which is more "
    "likely correct.\n\n"
    "QUESTION:\n{question}\n\n"
    "CANDIDATE 1 says answer = {a1}\n"
    "CANDIDATE 2 says answer = {a2}\n\n"
    "Think briefly, then respond with exactly one line:\n"
    "VERDICT: 1   (if candidate 1's letter is correct)\n"
    "VERDICT: 2   (if candidate 2's letter is correct)"
)


def _parse_verdict(text: str) -> Optional[int]:
    m = re.search(r"VERDICT:\s*([12])", text)
    return int(m.group(1)) if m else None


def _last_user_text(messages: List[Dict[str, str]]) -> str:
    for m in reversed(messages):
        if m.get("role") == "user":
            return m.get("content", "")
    return ""


async def ensemble_v1(
    darwin,
    awaxis,
    messages: List[Dict[str, str]],
    temperature: float = 0.7,
    max_tokens: int = 4096,
    n_rsa: int = 8,
) -> str:
    """
    Run V_1 ensemble. Returns the final answer string formatted as
    "ANSWER: X" so downstream tooling can parse uniformly.
    """
    # --- Phase 1: parallel RSA (each backend N samples) ---
    d_task = darwin.chat(messages, temperature=temperature, max_tokens=max_tokens, n=n_rsa)
    a_task = awaxis.chat(messages, temperature=temperature, max_tokens=max_tokens, n=n_rsa)
    d_outs, a_outs = await asyncio.gather(d_task, a_task)

    d_letters = [_extract_letter(o) for o in d_outs]
    a_letters = [_extract_letter(o) for o in a_outs]
    d_maj, d_votes = _majority(d_letters)
    a_maj, a_votes = _majority(a_letters)

    # --- Phase 2: agreement check ---
    if d_maj is None and a_maj is None:
        return "ANSWER: (no valid answer extracted)"
    if d_maj is None:
        return f"ANSWER: {a_maj}"
    if a_maj is None:
        return f"ANSWER: {d_maj}"
    if d_maj == a_maj:
        return f"ANSWER: {d_maj}"

    # --- Phase 3: cross-verification on mismatch ---
    question = _last_user_text(messages)
    verify_prompt = _VERIFY_TEMPLATE.format(question=question, a1=d_maj, a2=a_maj)
    verify_msgs = [{"role": "user", "content": verify_prompt}]

    d_verify_task = darwin.chat(verify_msgs, temperature=0.0, max_tokens=2048, n=1)
    a_verify_task = awaxis.chat(verify_msgs, temperature=0.0, max_tokens=2048, n=1)
    d_verify_outs, a_verify_outs = await asyncio.gather(d_verify_task, a_verify_task)
    d_verdict = _parse_verdict(d_verify_outs[0])
    a_verdict = _parse_verdict(a_verify_outs[0])

    # --- Phase 4: combine verdicts ---
    if d_verdict == a_verdict and d_verdict is not None:
        return f"ANSWER: {d_maj if d_verdict == 1 else a_maj}"
    if d_verdict is None and a_verdict is None:
        # Fall back to confidence (higher own-vote count wins)
        d_conf = d_votes.get(d_maj, 0)
        a_conf = a_votes.get(a_maj, 0)
        return f"ANSWER: {d_maj if d_conf >= a_conf else a_maj}"
    if d_verdict is None:
        return f"ANSWER: {d_maj if a_verdict == 1 else a_maj}"
    if a_verdict is None:
        return f"ANSWER: {d_maj if d_verdict == 1 else a_maj}"
    # Split — confidence tiebreaker
    d_conf = d_votes.get(d_maj, 0)
    a_conf = a_votes.get(a_maj, 0)
    return f"ANSWER: {d_maj if d_conf >= a_conf else a_maj}"