Update reranker.py

reranker.py

@@ -1,4 +1,5 @@
 import os
+import re
 import google.generativeai as genai
 import json
 
@@ -6,29 +7,41 @@ genai.configure(api_key=os.environ["GEMINI_API_KEY"])
 model = genai.GenerativeModel("models/gemini-2.0-flash")
 
 def rerank(query, candidates):
-    prompt = f"""
+    prompt = f""" 
     Job description: "{query}"
-
-    Candidate SHL assessments:
-    {json.dumps(candidates, indent=2)}
-
+    
+    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": ..., 
+          "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)
+        return json.loads(response_text)
     except Exception as e:
-        return {"error": str(e), "raw_response": response.text}
+        return {"error": str(e), "raw_response": response_text}