Update app.py

app.py

@@ -4,11 +4,10 @@ import asyncio
 import os
 import uvicorn
 import re
-import time
 import multiprocessing
 import urllib.parse
 from contextlib import asynccontextmanager
-from fastapi import FastAPI, Query, Request, Form
+from fastapi import FastAPI, Query, Form
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import hf_hub_download
@@ -16,7 +15,6 @@ from rank_bm25 import BM25Okapi
 import torch
 import aiohttp
 from bs4 import BeautifulSoup
-from langdetect import detect
 
 # DIRECT IMPORTS
 from transformers import pipeline
@@ -42,8 +40,6 @@ AI_RESPONSE_CACHE = {}
 MAX_CACHE_SIZE = 5000
 STOP_WORDS = {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "is", "are"}
 HF_TOKEN = os.getenv("HF_TOKEN")
-BANNED_IPS = set()
-IP_TRACKER = {}
 
 # 🔥 MASSIVE DESI & MODERN ADULT KEYWORDS LIST
 ADULT_KEYWORDS = {
@@ -55,30 +51,12 @@ ADULT_KEYWORDS = {
     "sex tape", "milf", "stepmom", "incest", "playboy", "scandal", "leaked"
 }
 
-# 🔥 MASSIVE LANGUAGE MAPPER (50+ Languages)
-LANG_MAP = {
-    "en": "English", "hi": "Hindi", "mr": "Marathi", "bn": "Bengali", "ta": "Tamil", 
-    "te": "Telugu", "gu": "Gujarati", "ur": "Urdu", "kn": "Kannada", "ml": "Malayalam", 
-    "pa": "Punjabi", "or": "Odia", "ne": "Nepali", "fr": "French", "es": "Spanish", 
-    "de": "German", "ru": "Russian", "ja": "Japanese", "zh-cn": "Chinese (Simplified)", 
-    "zh-tw": "Chinese (Traditional)", "ar": "Arabic", "pt": "Portuguese", "it": "Italian", 
-    "ko": "Korean", "id": "Indonesian", "tr": "Turkish", "nl": "Dutch", "vi": "Vietnamese"
-}
-
-# --- 2. ADVANCED NLP, INDEXING & LANGUAGE ---
+# --- 2. ADVANCED NLP & INDEXING ---
 def tokenize(text):
     if not text: return []
     tokens = re.findall(r'[a-z0-9]+', str(text).lower())
     return [t for t in tokens if t not in STOP_WORDS]
 
-def detect_language_name(text):
-    try:
-        if not text or len(text.strip()) < 10: return "Unknown"
-        code = detect(text)
-        return LANG_MAP.get(code, code.upper()) # Agar list me nahi mila toh code return karega
-    except:
-        return "Unknown"
-
 def load_and_index_dataset():
     global bm25_index, raw_data, suggestion_list
     print("📥 Loading Dataset & Building Deep Index...")
@@ -96,10 +74,14 @@ def load_and_index_dataset():
             print(f"⚠️ Hugging Face Download failed: {e}")
 
     tokenized_corpus = []
+    MAX_RECORDS_TO_LOAD = 150000 # Prevent OOM Crash
 
     if os.path.exists(dataset_path):
         with open(dataset_path, "r", encoding="utf-8") as f:
-            for line in f:
+            for i, line in enumerate(f):
+                if i >= MAX_RECORDS_TO_LOAD:
+                    break 
+                
                 if line.strip():
                     try:
                         doc = json.loads(line)
@@ -108,10 +90,6 @@ def load_and_index_dataset():
                         if url not in URL_STATS:
                             URL_STATS[url] = {"user_visits": 0, "ai_visits": 0, "is_trusted": False}
                             
-                        # Extract language and append to document
-                        combined_text = str(doc.get("title", "")) + " " + str(doc.get("description", ""))
-                        doc["language"] = detect_language_name(combined_text)
-                        
                         raw_data.append(doc)
                         
                         title_str = doc.get("title", "")
@@ -131,7 +109,7 @@ def load_and_index_dataset():
             print("⚡ Building High-Speed BM25 Okapi Index...")
             suggestion_list = sorted(list(set(suggestion_list)))
             bm25_index = BM25Okapi(tokenized_corpus)
-            print(f"🚀 Indexing Complete! {len(raw_data)} records ready.")
+            print(f"🚀 Indexing Complete! {len(raw_data)} records ready for search.")
     else:
         print("❌ Dataset file not found!")
 
@@ -139,29 +117,24 @@ def load_ai_models():
     global ai_generator
     print("🤖 Loading Modern Fast Text AI Model...")
     try:
-        # Chota, modern aur tezz model
         ai_generator = pipeline("text-generation", model="Qwen/Qwen2.5-0.5B-Instruct", device=-1, torch_dtype=torch.float32)
         print("✅ AI Model Loaded! System ready.")
     except Exception as e:
         print(f"❌ Model Loading Failed: {e}")
 
-# --- 3. SMART SCRIPT RAG: EARLY STOPPING CHUNKING ---
+# --- 3. SMART SCRIPT RAG ---
 def chunk_text(text, chunk_size=120):
     words = text.split()
     return [" ".join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]
 
 async def scrape_and_find_best_chunk(url, query):
-    """Script website visit karegi, sirf Headers & Paragraphs legi, aur best chunk lekar ruk jayegi"""
     try:
         async with aiohttp.ClientSession() as session:
-            # Sirf 3 second ka timeout taki endpoint na atke
             async with session.get(url, timeout=3) as resp:
                 if resp.status == 200:
                     html = await resp.text()
                     soup = BeautifulSoup(html, "html.parser")
-                    
-                    # Sirf Paragraphs aur Headers nikal rahe hain kachra avoid karne ke liye
-                    tags_to_extract = ['p', 'h1', 'h2', 'h3', 'h4', 'h5', 'h6']
+                    tags_to_extract = ['p', 'h1', 'h2', 'h3']
                     text = " ".join([tag.get_text(strip=True) for tag in soup.find_all(tags_to_extract)])
                     
                     if not text: return ""
@@ -171,24 +144,19 @@ async def scrape_and_find_best_chunk(url, query):
                     
                     best_chunk = ""
                     max_score = 0
-                    # Threshold: Agar query ke aade tokens mil gaye chunk mein toh wahi ruk jao
                     early_stop_threshold = max(1, len(query_tokens) // 2) 
 
                     for chunk in chunks:
                         chunk_tokens = set(tokenize(chunk))
                         score = len(query_tokens.intersection(chunk_tokens))
-                        
                         if score > max_score:
                             max_score = score
                             best_chunk = chunk
-                            
-                        # 🔥 EARLY STOPPING: Jo kaam ka chunk mila wahi liya, aage process nahi kiya
                         if max_score >= early_stop_threshold and max_score > 0:
                             break
-                            
                     return best_chunk
-    except Exception:
-        return "" # Scraping failed (timeout/block), ignore quietly
+    except:
+        return ""
     return ""
 
 # --- 4. CORE SEARCH ALGORITHM & RANKING ---
@@ -198,8 +166,6 @@ async def perform_core_search(raw_query: str):
 
     query = raw_query.strip().lower()
     query_tokens = tokenize(query)
-
-    # 🚨 SMART ADULT FILTER (Using 'in' operator to catch hidden adult queries)
     user_wants_adult = any(adult_word in query for adult_word in ADULT_KEYWORDS)
     
     scores = bm25_index.get_scores(query_tokens) if query_tokens else [1.0] * len(raw_data)
@@ -215,31 +181,33 @@ async def perform_core_search(raw_query: str):
             parsed_url = urllib.parse.urlparse(doc_url)
             domain = parsed_url.netloc.lower()
             
-            # Text combine to check adult words in document
             doc_text = (str(doc.get("title", "")) + " " + str(doc.get("description", "")) + " " + doc_url).lower()
             is_adult_doc = doc.get("tag", "general").lower() in ["adult", "nsfw"] or any(aw in doc_text for aw in ADULT_KEYWORDS)
             
-            # Agar user ki query normal hai, toh adult result filter ho jayega
             if not user_wants_adult and is_adult_doc:
                 continue 
                 
             final_score = score
             
-            # 🔥 DOMAIN RANKING BOOST
+            # 🔥 1. RELEVANCE BOOST: Domain aur Title Match
             if query.replace(" ", "") in domain.replace(" ", ""):
-                final_score *= 10.0 # Huge boost for direct domain match
+                final_score *= 10.0
             elif query in doc.get("title", "").lower():
                 final_score *= 3.0
                 
-            # 🔥 TRUST RANKING BOOST
-            stats = URL_STATS.get(doc_url, {})
+            # 🔥 2. CONTEXTUAL TRUST BOOST: Sirf relevant links ko trust ka fayda milega
+            stats = URL_STATS.get(doc_url, {"user_visits": 0, "ai_visits": 0, "is_trusted": False})
             if stats.get("is_trusted", False):
-                final_score *= 1.5 
-            
-            final_score += (stats.get("user_visits", 0) * 0.05)
+                final_score *= 2.0  # Relevant trusted domains ko double boost
+                
+            final_score += (stats.get("user_visits", 0) * 0.05) + (stats.get("ai_visits", 0) * 0.1)
             
             seen_urls.add(doc_url)
-            scored_results.append((final_score, doc))
+            
+            doc_copy = doc.copy()
+            doc_copy["trust_stats"] = stats
+            
+            scored_results.append((final_score, doc_copy))
 
     scored_results.sort(key=lambda x: x[0], reverse=True)
     best_score = scored_results[0][0] if scored_results else 0
@@ -274,21 +242,15 @@ async def lifespan(app: FastAPI):
 app = FastAPI(title="Nebulix Advanced Text AI & RAG Engine", lifespan=lifespan)
 
 app.add_middleware(
-    CORSMiddleware, 
-    allow_origins=["*"], 
-    allow_credentials=True, 
-    allow_methods=["*"], 
-    allow_headers=["*"]
+    CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"]
 )
 
-# 🔥 1. AUTO-SUGGESTIONS ENDPOINT
 @app.get("/suggest")
 async def get_suggestions(q: str = Query(..., min_length=1)):
     q = q.lower()
     matches = [s for s in suggestion_list if s.startswith(q)][:7]
     return {"query": q, "suggestions": matches}
 
-# 🔥 2. URL TRACKING & TRUST ENDPOINTS
 @app.post("/track-click")
 async def track_url_click(url: str = Query(...)):
     if url not in URL_STATS:
@@ -307,22 +269,38 @@ async def pure_search(q: str = Query(...)):
     results, prediction, _ = await perform_core_search(q)
     return {"request_id": req_id, "query": q, "results_count": len(results), "results": results}
 
-# 🔥 3. SMART RAG CHAT API (Only Text & Scraped Chunks)
-@app.post("/aimodechat")
-async def ai_mode_chat(q: str = Form(...)):
+@app.get("/aimode")
+async def ai_mode_short(q: str = Query(...)):
     req_id = str(uuid.uuid4())
-    
-    # AI Logic to decide RAG need
-    needs_rag = any(w in q.lower() for w in ["what", "who", "where", "how", "why", "tell me", "explain", "latest"])
+    results, _, _ = await perform_core_search(q)
+    if not results: 
+        return {"request_id": req_id, "ai_response": "I couldn't find any results.", "results": []}
+
+    context = " ".join([f"{r.get('title', '')} - {r.get('description', '')}" for r in results[:3]])
+    messages = [
+        {"role": "system", "content": "Provide a brief 2-3 sentence summary based only on the provided context."},
+        {"role": "user", "content": f"Query: {q}\nContext: {context}"}
+    ]
+
+    ai_reply = await run_ai_task(messages, 150, query_key=f"short_{q}")
+    return {"request_id": req_id, "ai_response": ai_reply, "results": results}
+
+
+async def process_rag_chat(q: str):
+    req_id = str(uuid.uuid4())
+    needs_rag = any(w in q.lower() for w in ["what", "who", "where", "how", "why", "tell me", "explain", "latest", "best"])
     
     results, _, _ = await perform_core_search(q)
     rag_context = ""
     used_sources = []
 
     if needs_rag and results:
+        # 🔥 SMART SCRAPING FIX: 
+        # Ab AI blindly 'sabse jyada visits' wali site pe nahi jayega. 
+        # perform_core_search ne already relevant results ko top pe rakha hai, 
+        # isliye hum sidha Top 10 Context-Relevant URLs ko hi scrape karenge.
         top_urls = results[:10]
         
-        # Script asynchronously scrapes multiple URLs at once
         tasks = [scrape_and_find_best_chunk(doc["url"], q) for doc in top_urls]
         scraped_chunks = await asyncio.gather(*tasks)
         
@@ -331,13 +309,14 @@ async def ai_mode_chat(q: str = Form(...)):
             doc = top_urls[i]
             url = doc["url"]
             
-            # Trust Tracking: Mark AI visits
+            # Trust Tracking: Mark AI visits for this relevant URL
             if url not in URL_STATS: URL_STATS[url] = {"user_visits": 0, "ai_visits": 0, "is_trusted": False}
             URL_STATS[url]["ai_visits"] += 1
             if URL_STATS[url]["ai_visits"] >= 100:
                 URL_STATS[url]["is_trusted"] = True
+                
+            doc["trust_stats"] = URL_STATS[url]
 
-            # Sirf kaam ka chunk (Early stopped chunk) lenge ya fall back pe description
             if chunk:
                 context_parts.append(f"Source [{i+1}] ({url}): {chunk}")
                 used_sources.append(doc)
@@ -347,7 +326,7 @@ async def ai_mode_chat(q: str = Form(...)):
                 
         rag_context = "\n".join(context_parts)
 
-    system_prompt = "You are a Smart AI Search Assistant. Use the provided search chunks to answer. Always mention the Source URL if you use information from it."
+    system_prompt = "You are a Smart AI Search Assistant. Answer accurately using the provided live context data. Always mention the Source URL if you use its information."
     if not needs_rag:
         system_prompt = "You are a helpful AI Assistant. Answer the user's query conversationally."
 
@@ -366,7 +345,13 @@ async def ai_mode_chat(q: str = Form(...)):
         "results": results[:15]
     }
 
+@app.get("/aimodechat")
+async def ai_mode_chat_get(q: str = Query(...)):
+    return await process_rag_chat(q)
+
+@app.post("/aimodechat")
+async def ai_mode_chat_post(q: str = Form(...)):
+    return await process_rag_chat(q)
+
 if __name__ == "__main__":
-    # app object ko direct pass karein (string "main:app" ki jagah).
-    # reload=True bhi hata diya hai kyunki cloud (Hugging Face) pe iski jarurat nahi hoti.
     uvicorn.run(app, host="0.0.0.0", port=7860)