Update main.py

main.py

@@ -11,6 +11,7 @@ from sqlalchemy import create_engine, Column, Integer, String, Boolean, ForeignK
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy import LargeBinary
 from sqlalchemy.orm import sessionmaker, relationship
+import json
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
@@ -222,25 +223,35 @@ def get_streams():
     )
 
     prompt = """
-    You need to act as a recommendation engine.
-    List all 40+ streams/branches like Computer Science, Chemical Engineering, Aerospace, etc.
-    Note: Output should be a **valid JSON list**:
-    ["branch1", "branch2", "branch3", ...]
-    Return only the list, **do not return text or explanations**.
+    You are a recommendation engine. 
+    List at least 40 branches of study (e.g., Computer Science, Chemical Engineering, Aerospace). 
+    **Output should be a valid JSON array with double quotes, like this:**
+    ["Computer Science", "Chemical Engineering", "Aerospace", ...]
+    Do not add extra text, explanations, or newlines.
     """
-    
+
     formatted_prompt = format_prompt(prompt)
 
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=False, details=False, return_full_text=False)
 
     try:
-        # Ensure output is valid JSON
+        # Ensure the model's response is a valid JSON array
         cleaned_data = stream.strip()
-        if not cleaned_data.startswith("["):
-            raise ValueError("Invalid response format")
 
+        # Fix incomplete or malformed JSON
+        if not cleaned_data.startswith("[") or not cleaned_data.endswith("]"):
+            cleaned_data = cleaned_data.split("[", 1)[-1]  # Keep text after first [
+            cleaned_data = "[" + cleaned_data  # Add missing opening bracket
+            cleaned_data = cleaned_data.rsplit("]", 1)[0]  # Keep text before last ]
+            cleaned_data = cleaned_data + "]"  # Add missing closing bracket
+
+        # Parse JSON safely
         parsed_data = json.loads(cleaned_data)
-        return jsonify({"ans": parsed_data})  # Send as proper JSON
+
+        if not isinstance(parsed_data, list):  # Ensure it's a list
+            raise ValueError("Response is not a valid list")
+
+        return jsonify({"ans": parsed_data})  # Return clean JSON list
     except Exception as e:
         return jsonify({"error": "Invalid response format", "details": str(e)})