reranker.py

import os
import google.generativeai as genai
import json

genai.configure(api_key=os.environ["GEMINI_API_KEY"])
model = genai.GenerativeModel("models/gemini-2.0-flash")

def rerank(query, candidates):
    prompt = f"""
    Job description: "{query}"

    Candidate SHL assessments:
    {json.dumps(candidates, indent=2)}

    Rank the most relevant assessments and return a JSON list in this format:
    {{
      "recommended_assessments": [
        {{
          "url": ..., 
          "adaptive_support": ..., 
          "remote_support": ..., 
          "description": ..., 
          "duration": ..., 
          "test_type": [...]
        }}
      ]
    }}
    """

    response = model.generate_content(prompt)
    try:
        return json.loads(response.text)
    except Exception as e:
        return {"error": str(e), "raw_response": response.text}