reranker.py

import os
import re
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": [...]
        }}
      ]
    }}
    
    Important instructions:
    1. Return ONLY valid JSON without any markdown code blocks or extra text
    2. The url field MUST be copied exactly as provided in the candidates, do not modify it
    3. Make sure to include the full URL value as provided in the input candidates
    4. Preserve all original data fields and their exact values (especially URLs)
    """
    
    response = model.generate_content(prompt)
    response_text = response.text
    
    # Try to extract JSON from possible markdown code blocks
    json_match = re.search(r'```(?:json)?\s*(.*?)```', response_text, re.DOTALL)
    if json_match:
        response_text = json_match.group(1).strip()
    
    try:
        return json.loads(response_text)
    except Exception as e:
        return {"error": str(e), "raw_response": response_text}