Update app.py

app.py

@@ -3,19 +3,18 @@ import uuid
 import asyncio
 import os
 import uvicorn
-import aiohttp
-import io
-import zipfile
-from bs4 import BeautifulSoup
 from contextlib import asynccontextmanager
 from datetime import datetime
-from urllib.parse import urlparse, urljoin
-from fastapi import FastAPI, Query, BackgroundTasks
+from fastapi import FastAPI, Query, BackgroundTasks, Request
 from fastapi.responses import JSONResponse, FileResponse
 from fastapi.middleware.cors import CORSMiddleware
 from rapidfuzz import fuzz
 from cryptography.fernet import Fernet
 
+# AI Models Imports
+import torch
+from transformers import pipeline
+
 # --- 1. SETUP & CONFIGURATION ---
 SECRET_KEY = Fernet.generate_key()
 cipher_suite = Fernet(SECRET_KEY)
@@ -26,271 +25,206 @@ os.makedirs(LOG_DIR, exist_ok=True)
 
 raw_data = []
 search_corpus = []
-unique_domains = set()
-visited_urls = set()
-
-# Live Tracking Variables
-scrape_queue = asyncio.Queue(maxsize=50000)
-merge_lock = asyncio.Lock()
-url_buffer = []  # 1,000 URLs jama karne ke liye
-last_scraped_url = "None"
-daily_domains_scraped = set()
-current_date_str = datetime.now().strftime("%Y-%m-%d")
-
-# Thresholds
-NEXT_THRESHOLD = 30000
-DAILY_DOMAIN_LIMIT = 20000
 
+# AI Concurrency Limit (CPU par hang hone se bachane ke liye max 2 AI generation ek time pe)
+MAX_CONCURRENT_AI = 2
+ai_semaphore = asyncio.Semaphore(MAX_CONCURRENT_AI)
+ai_generator = None
 
-# --- 2. DATASET MANAGEMENT & BUFFER FLUSH ---
-def get_domain(url):
-    return urlparse(url).netloc.lower()
 
-def process_item_into_memory(item):
-    url = item.get("url", "")
-    if not url or url in visited_urls:
-        return 
-
-    visited_urls.add(url)
-    domain = get_domain(url)
-    unique_domains.add(domain)
-    raw_data.append(item)
-    
-    # Ye search ke liye fast dictionary hai
-    search_corpus.append({
-        "url": url.lower(),
-        "domain": domain,
-        "title": item.get("title", "").lower(),
-        "description": item.get("description", "").lower()
-    })
-
-async def flush_buffer_to_disk():
-    global url_buffer
-    if not url_buffer: return
-    
-    daily_file = f"date-{current_date_str}.jsonl"
-    async with merge_lock:
-        print(f"💾 Saving {len(url_buffer)} URLs to {daily_file}...")
-        with open(daily_file, "a", encoding="utf-8") as f:
-            for item in url_buffer:
-                f.write(json.dumps(item) + "\n")
-        url_buffer.clear()
-
-def load_all_data():
-    global NEXT_THRESHOLD
-    print("🔄 Loading datasets into Search Memory...")
-    
-    # Load dataset.jsonl
+# --- 2. LOAD DATASET & AI MODEL ---
+def load_dataset():
+    print("🔄 Loading Search Dataset...")
     if os.path.exists("dataset.jsonl"):
         with open("dataset.jsonl", "r", encoding="utf-8") as f:
             for line in f:
                 if line.strip():
-                    try: process_item_into_memory(json.loads(line))
+                    try:
+                        item = json.loads(line)
+                        raw_data.append(item)
+                        search_corpus.append({
+                            "url": item.get("url", "").lower(),
+                            "title": item.get("title", "").lower(),
+                            "description": item.get("description", "").lower()
+                        })
                     except: pass
+    print(f"✅ Loaded {len(raw_data)} records into memory.")
 
-    # Load daily date-*.jsonl
-    for filename in os.listdir(DATA_DIR):
-        if filename.startswith("date-") and filename.endswith(".jsonl"):
-            with open(filename, "r", encoding="utf-8") as f:
-                for line in f:
-                    if line.strip():
-                        try: process_item_into_memory(json.loads(line))
-                        except: pass
-
-    current_domains = len(unique_domains)
-    NEXT_THRESHOLD = ((current_domains // 30000) + 1) * 30000
-    print(f"✅ Search Engine Ready! URLs: {len(visited_urls)} | Domains: {current_domains}")
-
-
-async def check_and_merge_datasets():
-    global NEXT_THRESHOLD
-    if len(unique_domains) >= NEXT_THRESHOLD:
-        print(f"🌟 Merge Threshold ({NEXT_THRESHOLD}) Reached! Combining datasets...")
-        await flush_buffer_to_disk() # Pehle buffer khali karo
-        
-        async with merge_lock:
-            temp_file = "dataset.jsonl.tmp"
-            with open(temp_file, "w", encoding="utf-8") as f:
-                for item in raw_data:
-                    f.write(json.dumps(item) + "\n")
-            
-            os.replace(temp_file, "dataset.jsonl")
-            
-            # Delete old daily files
-            for filename in os.listdir(DATA_DIR):
-                if filename.startswith("date-") and filename.endswith(".jsonl"):
-                    os.remove(filename)
-            
-            NEXT_THRESHOLD += 30000
-            print("✅ Merge Completed & Old files deleted!")
-
-
-# --- 3. SMART SPIDER CRAWLER ---
-async def fetch_tranco_seeds(session):
-    print("🌍 Fetching Tranco Dataset for Seeding...")
+def load_ai_model():
+    global ai_generator
+    print("🤖 Loading 500M Parameter AI Model (Qwen2.5-0.5B)... This may take a minute.")
     try:
-        async with session.get("https://tranco-list.eu/top-1m.csv.zip", timeout=30) as resp:
-            if resp.status == 200:
-                with zipfile.ZipFile(io.BytesIO(await resp.read())) as z:
-                    with z.open(z.namelist()[0]) as f:
-                        lines = f.read().decode('utf-8').splitlines()
-                        for line in lines[:5000]:
-                            parts = line.split(",")
-                            if len(parts) >= 2:
-                                try: scrape_queue.put_nowait(f"https://{parts[1].strip()}/")
-                                except asyncio.QueueFull: break
+        # 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", 
+            model=model_id, 
+            device=-1, # -1 means CPU
+            torch_dtype=torch.float32
+        )
+        print("✅ AI Model Loaded Successfully!")
     except Exception as e:
-        print("⚠️ Tranco fetch failed, using fallbacks.")
-        for url in ["https://wikipedia.org/", "https://github.com/"]:
-            await scrape_queue.put(url)
-
-async def extract_page_data(session, url):
-    headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) SearchBot/1.0"}
-    try:
-        async with session.get(url, headers=headers, timeout=8) as response:
-            if response.status != 200: return None
-            
-            html = await response.text()
-            soup = BeautifulSoup(html, 'html.parser')
-            
-            title = soup.title.string.strip() if soup.title and soup.title.string else ""
-            desc_tag = soup.find("meta", {"name": "description"})
-            desc = desc_tag["content"].strip() if desc_tag else ""
-            
-            # Sirf tabhi save karenge jab kuch matlab ka content ho
-            if not title and not desc: return None
-            
-            fav_tag = soup.find("link", {"rel": lambda x: x and 'icon' in x.lower()})
-            favicon = urljoin(url, fav_tag["href"]) if fav_tag else ""
-            
-            images = [urljoin(url, img['src']) for img in soup.find_all('img', src=True)][:5]
-            
-            new_links = []
-            for a in soup.find_all('a', href=True):
-                href = urljoin(url, a['href'])
-                if href.startswith('http') and "tranco" not in href.lower():
-                    new_links.append(href)
-            
-            return {
-                "url": url, "title": title, "description": desc,
-                "favicon": favicon, "images": images, "videos": [], "docs": [],
-                "new_links": new_links[:20]
-            }
-    except: return None
-
-async def scraper_worker(session):
-    global last_scraped_url, current_date_str
+        print("❌ AI Model Loading Failed:", e)
+
+
+# --- 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,
+        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."
     
-    while True:
-        # Check for Date Rollover (Naya din shuru hua ya nahi)
-        today = datetime.now().strftime("%Y-%m-%d")
-        if today != current_date_str:
-            await flush_buffer_to_disk()
-            current_date_str = today
-            daily_domains_scraped.clear() # Reset daily limit
-            
-        # Daily Limit Check (20k Unique Domains)
-        if len(daily_domains_scraped) >= DAILY_DOMAIN_LIMIT:
-            print("🛑 Daily 20k limit reached! Sleeping for 1 hour...")
-            await asyncio.sleep(3600)
-            continue
-            
-        url = await scrape_queue.get()
-        if url in visited_urls:
-            scrape_queue.task_done()
-            continue
-            
-        data = await extract_page_data(session, url)
-        if data:
-            domain = get_domain(url)
-            daily_domains_scraped.add(domain)
-            last_scraped_url = url
-            
-            # Put new links in queue
-            for link in data.pop("new_links"):
-                if link not in visited_urls:
-                    try: scrape_queue.put_nowait(link)
-                    except: pass
-            
-            # Memory me turant update for Search
-            process_item_into_memory(data)
-            
-            # 1,000 URL Buffer System
-            url_buffer.append(data)
-            if len(url_buffer) >= 1000:
-                await flush_buffer_to_disk()
-                await check_and_merge_datasets()
-                
-        scrape_queue.task_done()
-        await asyncio.sleep(0.1)
+    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
 
-async def start_scraper():
-    async with aiohttp.ClientSession() as session:
-        if scrape_queue.empty(): await fetch_tranco_seeds(session)
-        workers = [asyncio.create_task(scraper_worker(session)) for _ in range(15)]
-        await asyncio.gather(*workers)
 
+# --- 4. LIFESPAN & LOGGING ---
+async def log_query_encrypted(req_id: str, query: str):
+    try:
+        encrypted_q = cipher_suite.encrypt(query.encode()).decode()
+        log_file = os.path.join(LOG_DIR, f"search_logs_{datetime.now().strftime('%Y-%m-%d')}.jsonl")
+        log_entry = {"id": req_id, "query": encrypted_q, "timestamp": datetime.now().isoformat()}
+        with open(log_file, "a", encoding="utf-8") as f:
+            f.write(json.dumps(log_entry) + "\n")
+    except: pass
 
-# --- 4. FASTAPI ENDPOINTS ---
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    load_all_data()
-    scraper_task = asyncio.create_task(start_scraper())
+    load_dataset()
+    # Load AI in background so server starts immediately
+    asyncio.get_event_loop().run_in_executor(None, load_ai_model)
     yield
-    await flush_buffer_to_disk() # Save everything before shutdown
-    scraper_task.cancel()
 
-app = FastAPI(title="Search Engine & Spider API", lifespan=lifespan)
+app = FastAPI(title="AI Enhanced Search Engine", lifespan=lifespan)
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
 
+
+# --- 5. SEARCH LOGIC & SMART PREDICTION ---
+def perform_core_search(query: str):
+    query_lower = query.lower()
+    scored_results = []
+    
+    for index, data in enumerate(search_corpus):
+        title_score = fuzz.WRatio(query_lower, data["title"])
+        desc_score = fuzz.WRatio(query_lower, data["description"])
+        total_score = max(title_score, desc_score)
+        
+        if total_score > 40:
+            scored_results.append((total_score, index))
+            
+    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."
+    elif best_score > 55:
+        prediction = "Medium confidence: We found related content, but it might not be exact."
+    else:
+        prediction = "Low confidence: Results are poor. Try changing your keywords."
+
+    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)
+        
+    return final_results, prediction
+
+
+# --- 6. ENDPOINTS ---
+
 @app.get("/")
 async def serve_ui():
-    if os.path.exists("index.html"):
-        return FileResponse("index.html")
-    return {"message": "API is Running! Upload index.html for UI."}
+    if os.path.exists("index.html"): return FileResponse("index.html")
+    return {"message": "API Running!"}
 
-# 🌟 NEW: LIVE STATUS TRACKER 🌟
-@app.get("/status")
-async def get_status():
-    return {
-        "status": "Running",
-        "total_urls_in_database": len(raw_data),
-        "total_unique_domains": len(unique_domains),
-        "daily_limit": f"{len(daily_domains_scraped)} / {DAILY_DOMAIN_LIMIT}",
-        "urls_in_queue_to_scrape": scrape_queue.qsize(),
-        "urls_in_save_buffer": f"{len(url_buffer)} / 1000",
-        "last_scraped_url": last_scraped_url
+@app.get("/search")
+async def standard_search(q: str = Query(...), background_tasks: BackgroundTasks = BackgroundTasks()):
+    req_id = str(uuid.uuid4())
+    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
+        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)
+    elif ai_semaphore.locked():
+        ai_summary = "AI is currently busy with other users. Showing fast results."
+
+    response = {
+        "request_id": req_id,
+        "smart_prediction": prediction,
+        "ai_summary": ai_summary,
+        "total_results": len(results),
+        "results": results
     }
 
-@app.get("/search")
-async def search(q: str = Query(...)):
-    query_lower = q.lower()
-    scored_results = []
+    if is_adult:
+        response["safety_message"] = "you are safe you can do anything"
+        return JSONResponse(content=response, headers={"X-Safety-Message": "you are safe you can do anything"})
+    return response
+
+# 🌟 NEW: AI RAG MODE ENDPOINT 🌟
+@app.get("/aimode")
+async def ai_mode_search(q: str = Query(...), background_tasks: BackgroundTasks = BackgroundTasks()):
+    req_id = str(uuid.uuid4())
+    background_tasks.add_task(log_query_encrypted, req_id, q)
     
-    async with merge_lock:
-        for index, data in enumerate(search_corpus):
-            # Partial match using WRatio for better searching
-            title_score = fuzz.WRatio(query_lower, data["title"])
-            desc_score = fuzz.WRatio(query_lower, data["description"])
-            
-            total_score = max(title_score, desc_score)
-            
-            if total_score > 50: # Standard Search Threshold
-                scored_results.append((total_score, index))
-        
-        scored_results.sort(key=lambda x: x[0], reverse=True)
-        top_30 = scored_results[:30]
+    results, prediction = perform_core_search(q)
+    
+    if not results:
+        return {
+            "ai_response": f"I couldn't find any information about '{q}' in my database to give you an accurate answer.",
+            "results": []
+        }
         
-        final_results = []
-        for score, index in top_30:
-            item = raw_data[index].copy()
-            item["match_score"] = round(score, 2)
-            final_results.append(item)
-            
-    if not final_results:
-        return {"total_results": 0, "results": [], "message": "No results found in DB yet. Crawler is working..."}
+    # Build RAG Context from Top 3 results
+    context = ""
+    for i, r in enumerate(results[:3]):
+        context += f"Source {i+1}: {r.get('title', '')} - {r.get('description', '')}\n"
         
-    return {"total_results": len(final_results), "results": final_results}
+    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.
+User Query: {q}
+
+Search Results Data:
+{context}
+
+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
+    }
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)