Update app.py

app.py

@@ -54,7 +54,6 @@ def load_ai_model():
     global ai_generator
     print("🤖 Loading 500M Parameter AI Model (Qwen2.5-0.5B)... This may take a minute.")
     try:
-        # Qwen 0.5B ek chota aur bohot smart model hai jo CPU par achha chalta hai
         model_id = "Qwen/Qwen2.5-0.5B-Instruct"
         ai_generator = pipeline(
             "text-generation", 
@@ -69,28 +68,24 @@ def load_ai_model():
 
 # --- 3. ASYNC AI INFERENCE HELPERS ---
 def _generate_text(prompt: str, max_tokens: int):
-    # Chat template format for Instruct models
     messages = [{"role": "user", "content": prompt}]
     result = ai_generator(
         messages, 
         max_new_tokens=max_tokens, 
         do_sample=True, 
-        temperature=0.7,
+        temperature=0.6, # Thoda kam temperature taaki precise answer de
         top_p=0.9
     )
-    # Return extracted AI response
     return result[0]['generated_text'][-1]['content']
 
 async def run_ai_task(prompt: str, max_tokens: int = 150):
     if not ai_generator:
         return "AI model is not initialized."
         
-    # Check if AI is overloaded
     if ai_semaphore.locked():
         return "⚠️ AI Servers are currently at full capacity handling other users. Please try again in a few seconds."
     
     async with ai_semaphore:
-        # Run Heavy AI task in a separate thread so FastAPI doesn't block
         loop = asyncio.get_event_loop()
         response = await loop.run_in_executor(None, _generate_text, prompt, max_tokens)
         return response
@@ -109,11 +104,10 @@ async def log_query_encrypted(req_id: str, query: str):
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     load_dataset()
-    # Load AI in background so server starts immediately
     asyncio.get_event_loop().run_in_executor(None, load_ai_model)
     yield
 
-app = FastAPI(title="AI Enhanced Search Engine", lifespan=lifespan)
+app = FastAPI(title="Nebulix AI Search Engine", lifespan=lifespan)
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
 
 
@@ -132,7 +126,6 @@ def perform_core_search(query: str):
             
     scored_results.sort(key=lambda x: x[0], reverse=True)
     
-    # 🌟 SMART FILTER PREDICTION 🌟
     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."
@@ -164,12 +157,10 @@ async def standard_search(q: str = Query(...), background_tasks: BackgroundTasks
     background_tasks.add_task(log_query_encrypted, req_id, q)
     
     is_adult = any(word in q.lower() for word in ["porn", "sex", "xxx", "nude"])
-    
     results, prediction = perform_core_search(q)
     
-    # AI Summary Generation (Short summary based on top 2 results)
     ai_summary = "AI Summary not available."
-    if results and not ai_semaphore.locked(): # Only generate if not overloaded
+    if results and not ai_semaphore.locked():
         context = " ".join([r.get("title", "") + ". " + r.get("description", "") for r in results[:2]])
         prompt = f"Summarize the following search results about '{q}' in 2 short sentences. Results: {context}"
         ai_summary = await run_ai_task(prompt, max_tokens=60)
@@ -189,9 +180,39 @@ async def standard_search(q: str = Query(...), background_tasks: BackgroundTasks
         return JSONResponse(content=response, headers={"X-Safety-Message": "you are safe you can do anything"})
     return response
 
-# 🌟 NEW: AI RAG MODE ENDPOINT 🌟
+
+# 🌟 FAST AI MODE (Short Answer: 20-25 Tokens) 🌟
 @app.get("/aimode")
-async def ai_mode_search(q: str = Query(...), background_tasks: BackgroundTasks = BackgroundTasks()):
+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)
+    
+    results, prediction = perform_core_search(q)
+    
+    if not results:
+        return {"ai_response": "No relevant info found.", "results": []}
+        
+    context = " ".join([f"{r.get('title', '')} - {r.get('description', '')}" for r in results[:2]])
+        
+    prompt = f"""You are a smart Search Assistant. Answer the user's query based on the context in exactly 20 to 25 words. Be direct.
+User Query: {q}
+Context: {context}
+Short Answer:"""
+
+    # Max tokens limited to 35 to ensure the answer stays around 20-25 words
+    ai_reply = await run_ai_task(prompt, max_tokens=35)
+
+    return {
+        "request_id": req_id,
+        "ai_response": ai_reply,
+        "smart_prediction": prediction,
+        "results": results[:5] 
+    }
+
+
+# 🌟 NEW: SMART CHAT RAG MODE (Deep Answer: 250 Tokens + URLs) 🌟
+@app.get("/aimodechat")
+async def ai_mode_chat(q: str = Query(...), background_tasks: BackgroundTasks = BackgroundTasks()):
     req_id = str(uuid.uuid4())
     background_tasks.add_task(log_query_encrypted, req_id, q)
     
@@ -200,30 +221,43 @@ async def ai_mode_search(q: str = Query(...), background_tasks: BackgroundTasks
     if not results:
         return {
             "ai_response": f"I couldn't find any information about '{q}' in my database to give you an accurate answer.",
+            "sources": [],
             "results": []
         }
         
-    # Build RAG Context from Top 3 results
+    # Build Deep RAG Context & Extract Sources
     context = ""
-    for i, r in enumerate(results[:3]):
-        context += f"Source {i+1}: {r.get('title', '')} - {r.get('description', '')}\n"
+    sources_used = []
+    
+    for i, r in enumerate(results[:4]): # Top 4 results for deep context
+        url = r.get("url", "#")
+        title = r.get("title", "Unknown Title")
+        desc = r.get("description", "")
         
-    prompt = f"""You are a helpful Search Engine AI. Answer the user's query based ONLY on the provided search results. Be concise and polite.
+        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. 
+Your task is to provide a comprehensive, well-structured, and accurate answer to the user's query using ONLY the provided search results.
+If the search results do not contain the answer, politely say so.
+
 User Query: {q}
 
 Search Results Data:
 {context}
 
+Provide a detailed response (up to 250 tokens). Do not include the prompt in your output.
 Answer:"""
 
+    # Generating full 250 tokens deep answer
     ai_reply = await run_ai_task(prompt, max_tokens=250)
 
     return {
         "request_id": req_id,
         "ai_response": ai_reply,
         "smart_prediction": prediction,
-        "sources_used": len(results[:3]),
-        "results": results[:5] # Send top 5 results along with AI answer
+        "sources": sources_used, # Includes URLs explicitly
+        "results": results[:5] 
     }
 
 if __name__ == "__main__":