Update app.py

app.py

@@ -14,6 +14,7 @@ from cryptography.fernet import Fernet
 # AI Models Imports
 import torch
 from transformers import pipeline
+from sentence_transformers import SentenceTransformer, util
 
 # --- 1. SETUP & CONFIGURATION ---
 SECRET_KEY = Fernet.generate_key()
@@ -25,14 +26,17 @@ os.makedirs(LOG_DIR, exist_ok=True)
 
 raw_data = []
 search_corpus = []
+corpus_texts = [] # AI embedding ke liye
+corpus_embeddings = None 
 
 # AI Concurrency Limit
 MAX_CONCURRENT_AI = 2
 ai_semaphore = asyncio.Semaphore(MAX_CONCURRENT_AI)
 ai_generator = None
+embedding_model = None
 
 
-# --- 2. LOAD DATASET & AI MODEL ---
+# --- 2. LOAD DATASET & AI MODELS ---
 def load_dataset():
     print("🔄 Loading Search Dataset...")
     if os.path.exists("dataset.jsonl"):
@@ -42,49 +46,65 @@ def load_dataset():
                     try:
                         item = json.loads(line)
                         raw_data.append(item)
+                        title = item.get("title", "").lower()
+                        desc = item.get("description", "").lower()
                         search_corpus.append({
-                            "url": item.get("url", "").lower(),
-                            "title": item.get("title", "").lower(),
-                            "description": item.get("description", "").lower()
+                            "url": item.get("url", ""),
+                            "title": title,
+                            "description": desc
                         })
+                        corpus_texts.append(f"{title} {desc}")
                     except: pass
     print(f"✅ Loaded {len(raw_data)} records into memory.")
 
-def load_ai_model():
-    global ai_generator
-    # 🌟 NAYA MODEL: SmolLM2-360M (Chota hai par bohot smart hai)
-    print("🤖 Loading 360M Parameter AI Model (SmolLM2-360M-Instruct)...")
+def load_ai_models():
+    global ai_generator, embedding_model, corpus_embeddings
+    
+    # 🌟 1. AI SEMANTIC SEARCH MODEL (Accurate URLs ke liye)
+    print("🔍 Loading Semantic Search AI (all-MiniLM-L6-v2)...")
+    try:
+        embedding_model = SentenceTransformer('all-MiniLM-L6-v2', device='cpu')
+        if corpus_texts:
+            print("🧠 Encoding dataset into AI Brain... This may take a minute.")
+            corpus_embeddings = embedding_model.encode(corpus_texts, convert_to_tensor=True)
+        print("✅ Semantic Search AI Loaded!")
+    except Exception as e:
+        print("❌ Semantic Search AI Failed:", e)
+
+    # 🌟 2. BETTER RAG MODEL: Qwen 2.5 (Chota lekin bohot jyada smart)
+    print("🤖 Loading RAG AI Model (Qwen2.5-0.5B-Instruct)...")
     try:
-        model_id = "HuggingFaceTB/SmolLM2-360M-Instruct"
+        model_id = "Qwen/Qwen2.5-0.5B-Instruct"
         ai_generator = pipeline(
             "text-generation", 
             model=model_id, 
-            device=-1, # -1 means CPU
+            device=-1, # CPU
             torch_dtype=torch.float32
         )
-        print("✅ AI Model Loaded Successfully!")
+        print("✅ AI Generator Model Loaded Successfully!")
     except Exception as e:
-        print("❌ AI Model Loading Failed:", e)
+        print("❌ AI Generator Model Failed:", e)
 
 
 # --- 3. ASYNC AI INFERENCE HELPERS ---
 def _generate_text(prompt: str, max_tokens: int):
+    # Chat Template for Models
     messages = [{"role": "user", "content": prompt}]
     result = ai_generator(
         messages, 
         max_new_tokens=max_tokens, 
         do_sample=True, 
-        temperature=0.5, # Thoda kam temperature for precise answers
-        top_p=0.9
+        temperature=0.3, # Low temperature for accurate answers
+        top_p=0.85
     )
     return result[0]['generated_text'][-1]['content']
 
-async def run_ai_task(prompt: str, max_tokens: int = 150):
+async def run_ai_task(prompt: str, max_tokens: int):
     if not ai_generator:
         return "AI model is not initialized."
         
     if ai_semaphore.locked():
-        return "⚠️ AI Servers are currently at full capacity handling other users. Please try again in a few seconds."
+        return "⚠️ AI Servers are currently full. Please wait a second."
     
     async with ai_semaphore:
         loop = asyncio.get_event_loop()
@@ -105,43 +125,69 @@ async def log_query_encrypted(req_id: str, query: str):
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     load_dataset()
-    asyncio.get_event_loop().run_in_executor(None, load_ai_model)
+    # Loading models in background thread so server starts immediately
+    asyncio.get_event_loop().run_in_executor(None, load_ai_models)
     yield
 
 app = FastAPI(title="Nebulix AI Search Engine", lifespan=lifespan)
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
 
 
-# --- 5. SEARCH LOGIC & SMART PREDICTION ---
+# --- 5. HYBRID SEARCH LOGIC (Fuzzy + AI Semantic) ---
 def perform_core_search(query: str):
+    if not raw_data:
+        return [], "No dataset found.", 0
+
     query_lower = query.lower()
     scored_results = []
     
+    # 🧠 Calculate AI Semantic Scores
+    ai_scores = []
+    if embedding_model and corpus_embeddings is not None:
+        query_embedding = embedding_model.encode(query, convert_to_tensor=True)
+        # Cosine Similarity gives result between -1 to 1 (scaled to 0-100 below)
+        cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0].tolist()
+        ai_scores = cos_scores
+    else:
+        ai_scores = [0] * len(raw_data)
+
+    # 🔀 Combine AI Score + Fuzzy Score
     for index, data in enumerate(search_corpus):
+        # 1. Fuzzy Match (Keywords)
         title_score = fuzz.WRatio(query_lower, data["title"])
         desc_score = fuzz.WRatio(query_lower, data["description"])
-        total_score = max(title_score, desc_score)
+        fuzzy_max = max(title_score, desc_score)
+        
+        # 2. AI Score (Meaning)
+        ai_score = max(0, ai_scores[index] * 100) 
+        
+        # 3. Hybrid Score (60% AI + 40% Fuzzy)
+        total_score = (ai_score * 0.6) + (fuzzy_max * 0.4)
         
-        if total_score > 40:
-            scored_results.append((total_score, index))
+        scored_results.append((total_score, index))
             
+    # Sort from best to worst match
     scored_results.sort(key=lambda x: x[0], reverse=True)
     
     best_score = scored_results[0][0] if scored_results else 0
-    if best_score > 85:
-        prediction = "High confidence: We found exact matches for your query."
-    elif best_score > 55:
-        prediction = "Medium confidence: We found related content, but it might not be exact."
+    if best_score > 75:
+        prediction = "High confidence: We found strong semantic and exact matches."
+    elif best_score > 40:
+        prediction = "Medium confidence: We found related context for your query."
     else:
-        prediction = "Low confidence: Results are poor. Try changing your keywords."
+        prediction = "Low confidence: Forcing best available results."
 
+    # 📌 EXACTLY 30 RESULTS (Yahan har baar top 30 hi return honge)
     top_30 = scored_results[:30]
     final_results = []
+    
     for score, index in top_30:
         item = raw_data[index].copy()
         item["match_score"] = round(score, 2)
         final_results.append(item)
         
+    # Agar dataset mein 30 se kam data hai, toh jo hai utna return hoga.
+    # Agar 30+ hai, toh EXACTLY 30 return hoga har baar.
     return final_results, prediction, best_score
 
 
@@ -152,42 +198,43 @@ async def serve_ui():
     if os.path.exists("index.html"): return FileResponse("index.html")
     return {"message": "API Running!"}
 
-# 🌟 SMART AI OVERVIEW MODE (Max 100 Tokens + () Highlighting + Conditional Trigger) 🌟
+
+# 🌟 AI OVERVIEW MODE (40 Tokens = ~160 Words Logic) 🌟
 @app.get("/aimode")
 async def ai_mode_short(q: str = Query(...), background_tasks: BackgroundTasks = BackgroundTasks()):
     req_id = str(uuid.uuid4())
     background_tasks.add_task(log_query_encrypted, req_id, q)
     
+    # Har baar exactly 30 aayenge backend se
     results, prediction, best_score = perform_core_search(q)
     
-    # 🛑 CONDITIONAL AI TRIGGER: Agar query faltu hai ya result nahi mila, toh AI skip karo
-    if not results or best_score < 45 or len(q.strip()) < 3:
+    if not results or best_score < 30 or len(q.strip()) < 3:
         return {
             "request_id": req_id,
-            "ai_response": None, # Frontend isko dekh kar AI box hide kar dega
+            "ai_response": None, 
             "smart_prediction": prediction,
-            "results": results[:5]
+            "results": results # Exactly top 30 list
         }
         
-    context = " ".join([f"{r.get('title', '')} - {r.get('description', '')}" for r in results[:2]])
+    context = " ".join([f"{r.get('title', '')} - {r.get('description', '')}" for r in results[:3]])
         
     prompt = f"""You are a smart Search Assistant. Answer the user's query based ONLY on the context.
 Rules:
-1. Keep the answer under 100 words.
-2. You MUST enclose the most important keywords, names, or main concepts in parentheses, like (this).
+1. Provide your answer in approximately 150-160 words (Equivalent to what the user calls 40 extended tokens).
+2. You MUST enclose the most important keywords or concepts in parentheses, like (this).
 
 User Query: {q}
 Context: {context}
 Answer:"""
 
-    # Max tokens 120 (approx 90-100 words)
-    ai_reply = await run_ai_task(prompt, max_tokens=120)
+    # max_new_tokens set to 200 to easily accommodate the ~160 words constraint.
+    ai_reply = await run_ai_task(prompt, max_tokens=200)
 
     return {
         "request_id": req_id,
         "ai_response": ai_reply,
         "smart_prediction": prediction,
-        "results": results[:5] 
+        "results": results # Har baar exactly 30 pass honge 
     }
 
 
@@ -199,21 +246,21 @@ async def ai_mode_chat(q: str = Query(...), background_tasks: BackgroundTasks =
     
     results, prediction, best_score = perform_core_search(q)
     
-    # Agar search me kuch nahi mila (Faltu query), toh normal chat karo bina sources ke
-    if not results or best_score < 45:
+    # AI Chat Fallback if no logical context found
+    if not results or best_score < 30:
         prompt = f"""You are Nebulix AI. The user said: "{q}". 
-You do not have any web search results for this. Answer politely as a helpful AI assistant."""
+Answer politely as a helpful AI assistant. Answer in detailed format."""
         
-        ai_reply = await run_ai_task(prompt, max_tokens=150)
+        ai_reply = await run_ai_task(prompt, max_tokens=250)
         return {
             "request_id": req_id,
             "ai_response": ai_reply,
-            "smart_prediction": "No search results. General AI chat activated.",
+            "smart_prediction": "No reliable search results. General AI chat activated.",
             "sources": [],
-            "results": []
+            "results": results 
         }
         
-    # Build Deep RAG Context & Extract Sources
+    # Deep RAG Context 
     context = ""
     sources_used = []
     
@@ -221,29 +268,28 @@ You do not have any web search results for this. Answer politely as a helpful AI
         url = r.get("url", "#")
         title = r.get("title", "Unknown Title")
         desc = r.get("description", "")
-        
         context += f"Source {i+1} [{title}]: {desc}\n"
         sources_used.append({"id": i+1, "title": title, "url": url})
         
-    prompt = f"""You are Nebulix AI, a highly intelligent and helpful search assistant. 
-Provide a comprehensive and accurate answer to the user's query using ONLY the provided search results.
+    prompt = f"""You are Nebulix AI, a highly intelligent search assistant. 
+Provide a comprehensive answer using ONLY the provided search results.
 
 User Query: {q}
 
-Search Results Data:
+Context:
 {context}
 
 Provide a detailed response. Do not include the prompt in your output.
 Answer:"""
 
-    ai_reply = await run_ai_task(prompt, max_tokens=250)
+    ai_reply = await run_ai_task(prompt, max_tokens=300)
 
     return {
         "request_id": req_id,
         "ai_response": ai_reply,
         "smart_prediction": prediction,
         "sources": sources_used, 
-        "results": results[:5] 
+        "results": results # Har baar exactly 30 results!
     }
 
 if __name__ == "__main__":