Upload api.py

api.py

@@ -0,0 +1,90 @@
+import json
+import pickle
+import numpy as np
+import faiss
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from sentence_transformers import SentenceTransformer
+import os
+import requests
+from bs4 import BeautifulSoup
+
+app = FastAPI(title="SHL Assessment Recommender")
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"]
+)
+
+print("Loading models...")
+embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
+faiss_index = faiss.read_index("vector_index.faiss")
+
+with open("assessment_data.pkl", "rb") as f:
+    stored_assessments = pickle.load(f)
+
+print("Models loaded!")
+
+class QueryInput(BaseModel):
+    query: str
+
+def find_similar(query, top_k=10):
+    vec = embedding_model.encode([query]).astype("float32")
+    _, indices = faiss_index.search(vec, top_k)
+    results = []
+    for idx in indices[0]:
+        if idx < len(stored_assessments):
+            results.append(stored_assessments[idx])
+    return results
+
+def balance(candidates):
+    p_types = ["P", "B"]
+    k_types = ["K", "A"]
+    p_items = [i for i in candidates if any(t in p_types for t in i.get("test_type", []))]
+    k_items = [i for i in candidates if any(t in k_types for t in i.get("test_type", []))]
+    o_items = [i for i in candidates if i not in p_items and i not in k_items]
+    result = []
+    result.extend(p_items[:3])
+    result.extend(k_items[:3])
+    result.extend(o_items[:4])
+    remaining = [i for i in candidates if i not in result]
+    while len(result) < 10 and remaining:
+        result.append(remaining.pop(0))
+    return result[:10]
+
+@app.get("/health")
+def health():
+    return {"status": "healthy"}
+
+@app.post("/recommend")
+def recommend(user_input: QueryInput):
+    query = user_input.query
+    if query.startswith("http"):
+        try:
+            r = requests.get(query, timeout=10)
+            soup = BeautifulSoup(r.content, "html.parser")
+            query = soup.get_text(separator=" ", strip=True)[:3000]
+        except:
+            pass
+    candidates = find_similar(query, top_k=20)
+    final = balance(candidates)
+    response = []
+    for item in final:
+        response.append({
+            "url": item.get("url", ""),
+            "name": item.get("name", ""),
+            "adaptive_support": item.get("adaptive_support", "No"),
+            "description": item.get("description", ""),
+            "duration": item.get("duration", 0),
+            "remote_support": item.get("remote_support", "No"),
+            "test_type": item.get("test_type", [])
+        })
+    return {"recommended_assessments": response}
+
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.environ.get("PORT", 8000))
+    uvicorn.run(app, host="0.0.0.0", port=port)