reranker.py

"""
reranker.py
-----------
Stage 2 of the pipeline: use a free HuggingFace-hosted LLM to re-rank
the BM25 candidate set.

Model: mistralai/Mistral-7B-Instruct-v0.2

Why re-rank with an LLM instead of just using BM25?
- BM25 is purely lexical — it matches keywords but doesn't understand meaning.
  A query like "low-impact knee exercises" might not match "Single-Leg Box Squat"
  even though it's highly relevant.
- The LLM understands intent — rehab vs performance, equipment constraints,
  difficulty level — and can weigh candidates against each other holistically.
- Keeping retrieval and re-ranking separate means we only send ~15 candidates
  to the LLM rather than the full dataset, keeping latency low.

The prompt asks the model to return a JSON list of IDs in ranked order with a
short reason for each. We parse that and return the top 3–5 exercises.
"""

import json
import os
import re
from huggingface_hub import InferenceClient

MODEL = "mistralai/Mistral-7B-Instruct-v0.2"


def _get_client() -> InferenceClient:
    token = os.environ.get("HF_TOKEN")
    return InferenceClient(model=MODEL, token=token)


def rerank(query: str, candidates: list[dict], top_n: int = 5) -> list[dict]:
    """
    Sends the query + candidate exercises to a free HuggingFace-hosted LLM
    and returns the top_n exercises re-ranked by relevance.

    Returns a list of dicts, each with all original exercise fields plus
    a 'reason' key explaining why it was ranked highly.
    """
    if not candidates:
        return []

    # Format candidates for the prompt
    candidates_text = ""
    for i, ex in enumerate(candidates, 1):
        candidates_text += (
            f"{i}. ID: {ex['id']}\n"
            f"   Title: {ex['title']}\n"
            f"   Description: {ex['description']}\n"
            f"   Tags: {ex['tags']}\n"
            f"   Body Part: {ex['body_part']} | Difficulty: {ex['difficulty']}\n"
            f"   Equipment: {ex['equipment']} | Injury Focus: {ex['injury_focus']}\n"
            f"   Intensity: {ex['intensity']}\n\n"
        )

    prompt = f"""<s>[INST] You are a sports coaching assistant. A user has submitted this query:

"{query}"

Here are {len(candidates)} candidate exercises. Re-rank them from most to least relevant for the user's needs. Consider their goal, any constraints (equipment, injury, intensity), and appropriate difficulty.

{candidates_text}

Return ONLY a JSON array of the top {top_n} exercises in ranked order.
Each item must have:
- "id": the exercise ID string
- "reason": one sentence explaining why it fits the query

Example:
[
  {{"id": "EX_001", "reason": "Low-impact and targets knee stability, ideal for rehab."}},
  {{"id": "EX_005", "reason": "Isometric hold builds quad strength without stressing the knee."}}
]

Return only the JSON array, nothing else. [/INST]"""

    try:
        client = _get_client()
        response = client.text_generation(
            prompt,
            max_new_tokens=600,
            temperature=0.2,      # low temperature for consistent structured output
            do_sample=True,
        )
        raw = response.strip()

        # Extract JSON array from the response (model sometimes adds preamble)
        match = re.search(r'\[.*\]', raw, re.DOTALL)
        if match:
            raw = match.group(0)

        ranked_list = json.loads(raw)

    except Exception:
        # Graceful fallback: return BM25 order without reasons
        return [{**ex, "reason": ""} for ex in candidates[:top_n]]

    # Map IDs back to full exercise dicts and attach reason
    id_to_exercise = {ex["id"]: ex for ex in candidates}
    results = []
    for item in ranked_list[:top_n]:
        ex_id = item.get("id")
        if ex_id in id_to_exercise:
            exercise = dict(id_to_exercise[ex_id])
            exercise["reason"] = item.get("reason", "")
            results.append(exercise)

    # If parsing returned fewer than expected, pad with remaining BM25 candidates
    seen_ids = {ex["id"] for ex in results}
    for ex in candidates:
        if len(results) >= top_n:
            break
        if ex["id"] not in seen_ids:
            results.append({**ex, "reason": ""})

    return results