Update app.py

app.py

@@ -4,42 +4,49 @@ import asyncio
 import os
 import uvicorn
 import re
+import math
+import time
 import multiprocessing
 import urllib.parse
-from collections import Counter
+import unicodedata
+import difflib
+from datetime import datetime
+from collections import Counter, defaultdict
 from contextlib import asynccontextmanager
+from typing import List, Dict, Any, Tuple, Set
+
 from fastapi import FastAPI, Query, Form
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import hf_hub_download
-from rank_bm25 import BM25Okapi
-import torch
 import aiohttp
 from bs4 import BeautifulSoup
+import torch
 
 # Hugging Face imports for Q4 Quantization
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
 
-# --- 1. SETUP & CPU ALLOCATION ---
+# --- CONFIGURATION & MULTIPROCESSING ---
 TOTAL_CORES = multiprocessing.cpu_count()
 AI_CORES = max(1, int(TOTAL_CORES * 0.7)) 
 torch.set_num_threads(AI_CORES)
 print(f"🖥️ Total Cores: {TOTAL_CORES} | Dedicated AI Cores: {AI_CORES}")
 
-raw_data = []
-suggestion_tuples = []
-bm25_index = None
 MAX_CONCURRENT_AI = max(1, AI_CORES // 2)
 ai_semaphore = asyncio.Semaphore(MAX_CONCURRENT_AI)
 ai_generator = None
-URL_STATS = {}
+HF_TOKEN = os.getenv("HF_TOKEN")
+
+# Cache Structures
 AI_RESPONSE_CACHE = {}
-MAX_CACHE_SIZE = 5000
+MAX_AI_CACHE_SIZE = 5000
 
-TOKEN_RE = re.compile(r'[a-z0-9]+')
-STOP_WORDS = {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "is", "are", "by", "this", "that", "it"}
-HF_TOKEN = os.getenv("HF_TOKEN")
+CONTENT_CACHE = {}
+MAX_CONTENT_CONCURRENCY = 10
+content_semaphore = asyncio.Semaphore(MAX_CONTENT_CONCURRENCY)
 
+# --- 1. DETERMINISTIC PREPROCESSING & CONSTANTS ---
+STOP_WORDS = {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "is", "are", "by", "this", "that", "it"}
 ADULT_KEYWORDS = {
     "porn", "sex", "nude", "nsfw", "xxx", "adult", "erotic", "onlyfans", 
     "mms", "chudai", "bhabhi", "devar", "savita bhabhi", "xvideos", "pornhub", 
@@ -49,29 +56,332 @@ ADULT_KEYWORDS = {
     "sex tape", "milf", "stepmom", "incest", "playboy", "scandal", "leaked"
 }
 
-def tokenize(text):
+EXPANSION_DICT = {
+    "usa": ["united", "states"],
+    "uk": ["united", "kingdom"],
+    "ai": ["artificial", "intelligence"],
+    "ml": ["machine", "learning"],
+    "nlp": ["natural", "language", "processing"],
+    "js": ["javascript"]
+}
+
+WEIGHTS = {
+    "W_BM25": 1.0, "W_TITLE": 4.0, "W_HEADINGS": 3.0, "W_DESCRIPTION": 2.0, "W_BODY": 1.0,
+    "W_URL": 1.5, "W_DOMAIN": 1.5, "W_PHRASE": 2.5, "W_EXACT": 3.0, "W_FRESHNESS": 1.0,
+    "W_QUALITY": 1.0, "W_POPULARITY": 0.5, "W_PROXIMITY": 1.5
+}
+
+MAX_RESULTS_PER_DOMAIN = 3
+STRICT_RESULT_LIMIT = 20
+
+def normalize_text(text: str) -> str:
+    if not text: return ""
+    text = unicodedata.normalize('NFKD', str(text)).lower()
+    text = re.sub(r'[\'’]', '', text) 
+    text = re.sub(r'[^a-z0-9\-]', ' ', text)
+    return " ".join(text.split())
+
+def tokenize(text: str, remove_stopwords: bool = True) -> List[str]:
+    tokens = normalize_text(text).split()
+    if remove_stopwords:
+        tokens = [t for t in tokens if t not in STOP_WORDS]
+    return tokens
+
+# --- 2. DETERMINISTIC QUERY INTENT CLASSIFIER ---
+def detect_query_intent(query_tokens: List[str]) -> Dict[str, Any]:
+    intents = []
+    
+    news_terms = {"latest", "news", "today", "breaking", "recent", "updates", "current"}
+    if any(term in query_tokens for term in news_terms): intents.append("NEWS")
+        
+    def_terms = {"what", "who", "definition", "meaning"}
+    if any(term in query_tokens for term in def_terms): intents.append("DEFINITION")
+        
+    how_terms = {"how", "tutorial", "guide", "steps", "learn", "documentation"}
+    if any(term in query_tokens for term in how_terms): intents.append("HOW_TO")
+        
+    comp_terms = {"vs", "versus", "compare", "comparison", "difference", "best"}
+    if any(term in query_tokens for term in comp_terms): intents.append("COMPARISON")
+        
+    local_terms = {"near", "restaurant", "restaurants", "location", "weather"}
+    if any(term in query_tokens for term in local_terms): intents.append("LOCAL")
+        
+    if len(query_tokens) <= 2 and not intents: intents.append("NAVIGATIONAL")
+        
+    return {"intent": intents[0] if intents else "GENERAL", "signals": intents}
+
+
+# --- 3. DEDICATED WEB CONTENT EXTRACTION ---
+async def extract_page_content(url: str) -> Dict[str, Any]:
+    """Downloads webpage safely, extracts pure text, and caches it."""
+    if url in CONTENT_CACHE:
+        return CONTENT_CACHE[url]
+        
+    async with content_semaphore:
+        try:
+            async with aiohttp.ClientSession() as session:
+                async with session.get(url, timeout=4.0, headers={"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) NebulixBot"}) as resp:
+                    if resp.status != 200:
+                        return {"url": url, "title": "", "content": "", "success": False, "error": f"HTTP {resp.status}"}
+                    
+                    html = await resp.text()
+                    soup = BeautifulSoup(html, "html.parser")
+                    
+                    # Clean irrelevant elements
+                    for tag in soup(['script', 'style', 'nav', 'footer', 'header', 'aside', 'noscript', 'form', 'button', 'iframe']):
+                        tag.decompose()
+                        
+                    blocks = []
+                    for tag in soup.find_all(['h1', 'h2', 'h3', 'p', 'article', 'main', 'li']):
+                        text = tag.get_text(separator=' ', strip=True)
+                        if len(text) > 30: # Skip very tiny menu stubs missed by decompose
+                            blocks.append(text)
+                            
+                    # Deduplicate contiguous identical text
+                    clean_blocks = []
+                    seen_blocks = set()
+                    for b in blocks:
+                        if b not in seen_blocks:
+                            seen_blocks.add(b)
+                            clean_blocks.append(b)
+                            
+                    content_str = "\n\n".join(clean_blocks)
+                    page_title = soup.title.string.strip() if soup.title and soup.title.string else ""
+                    
+                    result = {
+                        "url": url, 
+                        "title": page_title, 
+                        "content": content_str, 
+                        "success": True, 
+                        "timestamp": time.time()
+                    }
+                    CONTENT_CACHE[url] = result
+                    return result
+        except asyncio.TimeoutError:
+            return {"url": url, "title": "", "content": "", "success": False, "error": "Timeout"}
+        except Exception as e:
+            return {"url": url, "title": "", "content": "", "success": False, "error": str(e)}
+
+def rank_chunks_for_rag(text: str, query_tokens: List[str]) -> List[str]:
+    """Deterministically extracts the most relevant chunks from an article for context."""
     if not text: return []
-    return [t for t in TOKEN_RE.findall(str(text).lower()) if t not in STOP_WORDS]
+    chunks = text.split("\n\n")
+    q_set = set(query_tokens)
+    scored_chunks = []
+    
+    for c in chunks:
+        c_toks = set(tokenize(c))
+        score = len(q_set.intersection(c_toks))
+        if score > 0:
+            scored_chunks.append((score, c))
+            
+    if not scored_chunks and chunks:
+        scored_chunks = [(0, c) for c in chunks[:3]] # Fallback to top paragraphs
+        
+    scored_chunks.sort(reverse=True, key=lambda x: x[0])
+    return [c for score, c in scored_chunks[:3]] # Max 3 context chunks per source
+
+
+# --- 4. CORE DETERMINISTIC SEARCH INDEX ---
+class HybridSearchIndex:
+    def __init__(self):
+        self.docs = []
+        self.total_docs = 0
+        self.vocab = set()
+        self.df = defaultdict(int)
+        self.avg_doc_len = 0.0
+        self.inverted_index = defaultdict(set)
+        self.suggestions = []
+        self.k1 = 1.5
+        self.b = 0.75
+
+    def build(self, raw_data: List[dict]):
+        print("📥 Building Advanced Deterministic Search Index...")
+        total_len = 0
+        suggestion_counter = Counter()
+
+        for doc_id, doc in enumerate(raw_data):
+            self.docs.append(doc)
+            title_tokens = tokenize(doc.get("_fast_title", ""))
+            desc_tokens = tokenize(doc.get("_fast_desc", ""))
+            domain_tokens = tokenize(doc.get("_fast_domain", ""))
+            all_tokens = title_tokens + desc_tokens + domain_tokens
+            
+            doc["_token_len"] = len(all_tokens)
+            total_len += len(all_tokens)
+            
+            for term in set(all_tokens):
+                self.inverted_index[term].add(doc_id)
+                self.df[term] += 1
+                self.vocab.add(term)
+                
+            if not doc.get("_is_adult", False):
+                clean_title = [w for w in title_tokens if w not in ADULT_KEYWORDS and len(w) > 2]
+                for w in clean_title: suggestion_counter[w] += 1
+                for j in range(len(clean_title) - 1): suggestion_counter[f"{clean_title[j]} {clean_title[j+1]}"] += 2
+                d_clean = doc.get("_fast_domain", "").split('.')[0]
+                if len(d_clean) > 2 and d_clean not in ADULT_KEYWORDS: suggestion_counter[d_clean] += 5
+
+        self.total_docs = len(self.docs)
+        self.avg_doc_len = total_len / max(1, self.total_docs)
+        self.suggestions = suggestion_counter.most_common(40000)
+        print(f"🚀 Indexing Complete! {self.total_docs} records ready.")
+
+    def idf(self, term: str) -> float:
+        df = self.df.get(term, 0)
+        return math.log(1.0 + (self.total_docs - df + 0.5) / (df + 0.5))
+
+    def _get_candidates(self, query_tokens: List[str]) -> Set[int]:
+        candidates = set()
+        for term in query_tokens:
+            if term in self.inverted_index: candidates.update(self.inverted_index[term])
+        
+        if len(candidates) < 10:
+            for term in query_tokens:
+                if term not in self.inverted_index:
+                    matches = difflib.get_close_matches(term, self.vocab, n=2, cutoff=0.8)
+                    for m in matches: candidates.update(self.inverted_index[m])
+        return candidates
+
+    def calculate_features(self, doc: dict, query_tokens: List[str], query_str: str) -> Dict[str, float]:
+        title = doc.get("_fast_title", "")
+        desc = doc.get("_fast_desc", "")
+        domain = doc.get("_fast_domain", "")
+        url = doc.get("_fast_path", "")
+        
+        title_toks = tokenize(title, False)
+        desc_toks = tokenize(desc, False)
+        query_toks_full = tokenize(query_str, False)
+        
+        scores = {"bm25": 0.0, "title": 0.0, "phrase": 0.0, "exact": 0.0, "url": 0.0, "domain": 0.0, "quality": 0.0, "freshness": 0.0, "proximity": 0.0}
+
+        doc_len = doc.get("_token_len", 1)
+        for term in query_tokens:
+            tf = title_toks.count(term) * 2 + desc_toks.count(term)
+            if tf > 0:
+                tf_num = tf * (self.k1 + 1)
+                tf_den = tf + self.k1 * (1 - self.b + self.b * (doc_len / self.avg_doc_len))
+                scores["bm25"] += self.idf(term) * (tf_num / tf_den)
+                
+        q_norm = normalize_text(query_str)
+        if q_norm and q_norm in normalize_text(title): scores["exact"] = 1.0
+        elif q_norm and q_norm in normalize_text(desc): scores["exact"] = 0.5
+            
+        if len(query_toks_full) > 1:
+            phrase_overlap = 0
+            for i in range(len(query_toks_full) - 1):
+                bigram = f"{query_toks_full[i]} {query_toks_full[i+1]}"
+                if bigram in normalize_text(title): phrase_overlap += 1.0
+                elif bigram in normalize_text(desc): phrase_overlap += 0.5
+            scores["phrase"] = phrase_overlap / (len(query_toks_full) - 1)
+
+        title_overlap = sum(1 for t in query_tokens if t in title_toks)
+        scores["title"] = title_overlap / max(1, len(query_tokens))
+
+        q_joined = "".join(query_tokens)
+        domain_clean = domain.split('.')[0]
+        if q_joined == domain_clean or domain.startswith(q_joined + "."): scores["domain"] = 1.0
+        elif q_joined in domain: scores["domain"] = 0.5
+        if any(t in url for t in query_tokens): scores["url"] = 0.5
+
+        if len(query_tokens) > 1:
+            positions = [title_toks.index(t) for t in query_tokens if t in title_toks]
+            if len(positions) > 1:
+                scores["proximity"] = 1.0 / (1.0 + (max(positions) - min(positions)))
+
+        year_match = re.search(r'\b(202[0-6])\b', url + " " + desc)
+        if year_match:
+            scores["freshness"] = math.exp(-(datetime.now().year - int(year_match.group(1))) / 2.0)
+        else: scores["freshness"] = 0.3 
+
+        q_score = 0.0
+        if len(title_toks) >= 3: q_score += 0.2
+        if len(desc_toks) >= 10: q_score += 0.2
+        if doc.get("url", "").startswith("https"): q_score += 0.2
+        if len(title_toks) > 0 and (len(set(title_toks)) / len(title_toks)) < 0.5: q_score -= 0.3
+        scores["quality"] = max(0.0, min(1.0, q_score))
+
+        return scores
+
+    def search(self, query: str, debug: bool = False) -> Tuple[List[dict], Dict[str, Any]]:
+        query_norm = normalize_text(query)
+        query_tokens = tokenize(query_norm)
+        
+        expanded_tokens = list(query_tokens)
+        for t in query_tokens:
+            if t in EXPANSION_DICT: expanded_tokens.extend(EXPANSION_DICT[t])
+        query_tokens = list(dict.fromkeys(expanded_tokens)) 
+        
+        intent_info = detect_query_intent(query_tokens)
+        intent = intent_info["intent"]
+        user_wants_adult = any(aw in query_norm for aw in ADULT_KEYWORDS)
+        
+        candidates_ids = self._get_candidates(query_tokens)
+        scored_results = []
+        seen_domains = defaultdict(int)
+        seen_urls = set()
+
+        for doc_id in candidates_ids:
+            doc = self.docs[doc_id]
+            if not user_wants_adult and doc.get("_is_adult"): continue
+                
+            base_url = doc.get("url", "").split("?")[0].rstrip("/")
+            if base_url in seen_urls: continue
+                
+            features = self.calculate_features(doc, query_tokens, query_norm)
+            norm_bm25 = features["bm25"] / (features["bm25"] + 10.0)
+            
+            w_freshness = WEIGHTS["W_FRESHNESS"] + (3.0 if intent == "NEWS" else 0.0)
+            w_domain = WEIGHTS["W_DOMAIN"] + (4.0 if intent == "NAVIGATIONAL" else 0.0)
+            w_title = WEIGHTS["W_TITLE"] + (2.0 if intent == "DEFINITION" else 0.0)
+            w_phrase = WEIGHTS["W_PHRASE"] + (1.5 if intent in ("DEFINITION", "COMPARISON") else 0.0)
+            
+            final_score = (
+                WEIGHTS["W_BM25"] * norm_bm25 + w_title * features["title"] + w_phrase * features["phrase"] +
+                WEIGHTS["W_EXACT"] * features["exact"] + WEIGHTS["W_URL"] * features["url"] +
+                w_domain * features["domain"] + w_freshness * features["freshness"] +
+                WEIGHTS["W_QUALITY"] * features["quality"] + WEIGHTS["W_PROXIMITY"] * features["proximity"]
+            )
+            
+            if final_score > 0.05:
+                doc_out = {
+                    "title": doc.get("_fast_title", ""),
+                    "description": doc.get("_fast_desc", ""),
+                    "url": doc.get("url", ""),
+                    "domain": doc.get("_fast_domain", ""),
+                    "match_score": round(final_score, 4)
+                }
+                if debug:
+                    doc_out["_debug_features"] = features
+                    doc_out["_debug_intent"] = intent_info
+                
+                scored_results.append((final_score, doc_out, base_url, doc["_fast_domain"]))
+
+        scored_results.sort(key=lambda x: x[0], reverse=True)
+        
+        diversified_results = []
+        for score, doc_info, b_url, domain in scored_results:
+            if seen_domains[domain] >= MAX_RESULTS_PER_DOMAIN: continue
+            seen_domains[domain] += 1
+            seen_urls.add(b_url)
+            diversified_results.append(doc_info)
+            if len(diversified_results) >= STRICT_RESULT_LIMIT: break # STRICT LIMIT
+                
+        # Assign final rank
+        for idx, item in enumerate(diversified_results): item["rank"] = idx + 1
+
+        return diversified_results, intent_info
+
+# --- INITIALIZATION ---
+search_index = HybridSearchIndex()
 
 def load_and_index_dataset():
-    global bm25_index, raw_data, suggestion_tuples
-    print("📥 Loading Dataset & Building Fast Index...")
     dataset_path = "dataset.jsonl"
-    
     if HF_TOKEN:
-        try:
-            dataset_path = hf_hub_download(
-                repo_id="Nebulixlabs/Ultimate-Search-Engine-Dataset",
-                filename="dataset.jsonl",
-                repo_type="dataset",
-                token=HF_TOKEN 
-            )
-        except Exception:
-            pass
+        try: dataset_path = hf_hub_download(repo_id="Nebulixlabs/Ultimate-Search-Engine-Dataset", filename="dataset.jsonl", repo_type="dataset", token=HF_TOKEN)
+        except Exception: pass
 
-    tokenized_corpus = []
-    suggestion_counter = Counter()
-    
+    raw_data = []
     if os.path.exists(dataset_path):
         with open(dataset_path, "r", encoding="utf-8") as f:
             for i, line in enumerate(f):
@@ -82,40 +392,18 @@ def load_and_index_dataset():
                         url = doc.get("url", "")
                         title_str = doc.get("title", "")
                         desc_str = doc.get("description", "")
-                        
                         parsed_url = urllib.parse.urlparse(url)
-                        domain = parsed_url.netloc.lower().replace("www.", "")
-                        clean_domain_name = domain.split('.')[0] if '.' in domain else domain
-
-                        doc_text = f"{title_str} {desc_str} {url}".lower()
-                        is_adult = doc.get("tag", "general").lower() in ["adult", "nsfw"] or any(aw in doc_text for aw in ADULT_KEYWORDS)
-
-                        doc["_fast_domain"] = domain
+                        
+                        doc["_fast_domain"] = parsed_url.netloc.lower().replace("www.", "")
                         doc["_fast_path"] = parsed_url.path.strip("/")
-                        doc["_fast_title"] = title_str.lower()
-                        doc["_fast_desc"] = desc_str.lower()
-                        doc["_is_adult"] = is_adult
+                        doc["_fast_title"] = title_str
+                        doc["_fast_desc"] = desc_str
+                        doc["_is_adult"] = doc.get("tag", "general").lower() in ["adult", "nsfw"] or any(aw in (title_str+desc_str+url).lower() for aw in ADULT_KEYWORDS)
                         raw_data.append(doc)
-                        
-                        if not is_adult:
-                            if len(clean_domain_name) > 2 and clean_domain_name not in ADULT_KEYWORDS:
-                                suggestion_counter[clean_domain_name] += 5 
-                            
-                            clean_t_tokens = [w for w in tokenize(title_str) if w not in ADULT_KEYWORDS]
-                            for w in clean_t_tokens:
-                                if len(w) > 2: suggestion_counter[w] += 1
-                            for j in range(len(clean_t_tokens) - 1):
-                                suggestion_counter[f"{clean_t_tokens[j]} {clean_t_tokens[j+1]}"] += 2
-
-                        tokenized_corpus.append((tokenize(title_str) * 5) + (tokenize(desc_str) * 3) + tokenize(url))
                     except: pass
-                    
-        if tokenized_corpus:
-            suggestion_tuples = suggestion_counter.most_common(40000)
-            bm25_index = BM25Okapi(tokenized_corpus)
-            print(f"🚀 Indexing Complete! {len(raw_data)} records ready.")
+    search_index.build(raw_data)
 
-# --- QWEN 2.5 (0.5B) + Q4 OPTIMIZATION ---
+# --- QWEN AI GENERATION MODULE (FOR RAG ONLY) ---
 def load_ai_models():
     global ai_generator
     print("🤖 Loading Qwen/Qwen2.5-0.5B-Instruct...")
@@ -123,224 +411,137 @@ def load_ai_models():
         model_id = "Qwen/Qwen2.5-0.5B-Instruct"
         tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-        # Smart Hardware Check for 4-bit (Q4) Quantization
         if torch.cuda.is_available():
-            print("⚡ GPU detected! Loading in 4-bit (Q4) using BitsAndBytes for Extreme Speed...")
-            quant_config = BitsAndBytesConfig(
-                load_in_4bit=True,
-                bnb_4bit_compute_dtype=torch.float16,
-                bnb_4bit_quant_type="nf4",
-            )
-            model = AutoModelForCausalLM.from_pretrained(
-                model_id,
-                quantization_config=quant_config,
-                device_map="auto"
-            )
+            quant_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4")
+            model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quant_config, device_map="auto")
         else:
-            print("🖥️ CPU detected! Loading optimized model (PyTorch execution)...")
-            # For CPU, 4-bit BitsAndBytes isn't natively optimal, so we load in standard float32
-            # Qwen 0.5B is already tiny (~1GB) so it runs perfectly fast on CPU too.
-            model = AutoModelForCausalLM.from_pretrained(
-                model_id,
-                device_map="cpu",
-                torch_dtype=torch.float32
-            )
+            model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cpu", torch_dtype=torch.float32)
 
-        ai_generator = pipeline(
-            "text-generation", 
-            model=model, 
-            tokenizer=tokenizer
-        )
-        print("✅ Qwen2.5-0.5B-Instruct (Superfast Mode) Ready!")
+        ai_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        print("✅ Qwen AI Ready!")
     except Exception as e:
         print(f"❌ Model Loading Failed: {e}")
 
-# --- AI TEXT GENERATION FIX ---
-def _generate_text(messages, max_tokens: int):
-    tokenizer = ai_generator.tokenizer
-    
-    # 1. Use Native Chat Template for Qwen (Makes answers highly accurate & stops hallucinations)
-    try:
-        prompt = tokenizer.apply_chat_template(
-            messages, 
-            tokenize=False, 
-            add_generation_prompt=True
-        )
-    except Exception:
-        prompt = "\n".join([f"{m['role'].capitalize()}: {m['content']}" for m in messages]) + "\nAssistant: "
-
-    # 2. Greedy Decoding (do_sample=False) + Return Setup
-    result = ai_generator(
-        prompt, 
-        max_new_tokens=max_tokens, 
-        do_sample=False, 
-        truncation=True,
-        return_full_text=False, # ❌ FIX: No more repeating the query in answers!
-        pad_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0
-    )
-    
-    text_out = result[0]['generated_text'].strip()
-    return text_out
-
-async def run_ai_task(messages, max_tokens: int, query_key: str):
+async def run_ai_task(messages, max_tokens: int, query_key: str, skip_cache: bool = False):
     if not ai_generator: return "⚠️ AI Model is starting up."
-    if query_key in AI_RESPONSE_CACHE: return AI_RESPONSE_CACHE[query_key]
-
+    if not skip_cache and query_key in AI_RESPONSE_CACHE: return AI_RESPONSE_CACHE[query_key]
+    
     async with ai_semaphore:
-        response = await asyncio.to_thread(_generate_text, messages, max_tokens)
-        if len(AI_RESPONSE_CACHE) > MAX_CACHE_SIZE: AI_RESPONSE_CACHE.clear()
-        AI_RESPONSE_CACHE[query_key] = response
+        def _gen():
+            tokenizer = ai_generator.tokenizer
+            try: prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            except: prompt = "\n".join([f"{m['role'].capitalize()}: {m['content']}" for m in messages]) + "\nAssistant: "
+
+            result = ai_generator(prompt, max_new_tokens=max_tokens, do_sample=False, truncation=True, return_full_text=False)
+            return result[0]['generated_text'].strip()
+
+        response = await asyncio.to_thread(_gen)
+        if not skip_cache:
+            if len(AI_RESPONSE_CACHE) > MAX_AI_CACHE_SIZE: AI_RESPONSE_CACHE.clear()
+            AI_RESPONSE_CACHE[query_key] = response
         return response
 
-def chunk_text(text, chunk_size=80):
-    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):
-    try:
-        async with aiohttp.ClientSession() as session:
-            async with session.get(url, timeout=1.5) as resp:
-                if resp.status == 200:
-                    html = await resp.text()
-                    soup = BeautifulSoup(html, "html.parser")
-                    text = " ".join([tag.get_text(strip=True) for tag in soup.find_all(['p', 'h1', 'h2'])])
-                    if not text: return ""
-                    
-                    chunks = chunk_text(text)
-                    query_tokens = set(tokenize(query))
-                    
-                    best_chunk, max_score = "", 0
-                    for chunk in chunks:
-                        score = len(query_tokens.intersection(tokenize(chunk)))
-                        if score > max_score:
-                            max_score, best_chunk = score, chunk
-                        if max_score >= max(1, len(query_tokens) // 2): break
-                    return best_chunk
-    except:
-        return ""
-    return ""
-
-async def perform_core_search(raw_query: str):
-    if not raw_data or not bm25_index: return [], "Search index not ready.", 0
-
-    query = raw_query.strip().lower()
-    query_tokens = tokenize(query)
-    user_wants_adult = any(aw in query for aw in ADULT_KEYWORDS)
-    
-    scores = bm25_index.get_scores(query_tokens) if query_tokens else [1.0] * len(raw_data)
-    scored_results, seen_urls = [], set()
-    query_clean = query.replace(" ", "")
-
-    for index, score in enumerate(scores):
-        if score > 0.1 or not query_tokens:
-            doc = raw_data[index]
-            doc_url = doc["url"]
-            if doc_url in seen_urls: continue
-            if not user_wants_adult and doc["_is_adult"]: continue 
-                
-            final_score, domain = score, doc["_fast_domain"]
-            domain_clean = domain.replace(" ", "")
-            doc_title, doc_desc = doc["_fast_title"], doc["_fast_desc"]
-            
-            if query_clean == domain_clean or domain_clean.startswith(query_clean + "."):
-                final_score *= 50.0 if not doc["_fast_path"] else 10.0
-            elif query_clean in domain_clean:
-                final_score *= 25.0 if not doc["_fast_path"] else 5.0
-                    
-            if query in doc_title: final_score *= 5.0 
-            elif all(w in doc_title for w in query_tokens): final_score *= 3.0 
-            if query in doc_desc: final_score *= 2.0 
-                
-            stats = URL_STATS.get(doc_url, {"user_visits": 0, "ai_visits": 0, "is_trusted": False})
-            if stats.get("is_trusted", False): final_score *= 2.0  
-            final_score += (stats.get("user_visits", 0) * 0.05) + (stats.get("ai_visits", 0) * 0.1)
-            
-            seen_urls.add(doc_url)
-            doc_copy = {k: v for k, v in doc.items() if not k.startswith("_")}
-            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
-    prediction = "High confidence" if best_score > 15 else "Medium confidence" if best_score > 5 else "Low confidence"
-    
-    return [{"match_score": round(s, 2), **doc} for s, doc in scored_results[:100]], prediction, best_score
-
+# --- FASTAPI APP ENDPOINTS ---
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     asyncio.create_task(asyncio.to_thread(load_and_index_dataset))
     asyncio.create_task(asyncio.to_thread(load_ai_models))
     yield
 
-app = FastAPI(title="Nebulix Fast Qwen AI Search", lifespan=lifespan)
+app = FastAPI(title="Nebulix Fast Hybrid Search Engine", lifespan=lifespan)
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
 
 @app.get("/suggest")
 async def get_suggestions(q: str = Query(..., min_length=1)):
-    q = q.lower().strip()
-    matches = [word for word, count in suggestion_tuples if word.startswith(q)][:7]
+    q = normalize_text(q).strip()
+    matches = [word for word, count in search_index.suggestions if word.startswith(q)][:7]
     return {"query": q, "suggestions": matches}
 
-@app.get("/search")
-async def pure_search(q: str = Query(...)):
-    req_id = str(uuid.uuid4())
-    results, prediction, _ = await perform_core_search(q)
-    return {"request_id": req_id, "query": q, "results_count": len(results), "results": results}
+@app.get("/content")
+async def get_content_endpoint(url: str = Query(...)):
+    """Dedicated endpoint to just extract webpage content."""
+    res = await extract_page_content(url)
+    return res
 
-@app.get("/aimode")
-async def ai_mode_short(q: str = Query(...)):
+@app.get("/search")
+async def pure_search(q: str = Query(...), content: bool = Query(False), debug: bool = Query(False)):
     req_id = str(uuid.uuid4())
-    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[:2]])
-    messages = [
-        {"role": "system", "content": "You are a concise search AI. Provide a direct, 1-sentence answer based on the context. Highlight keywords in (parentheses)."},
-        {"role": "user", "content": f"Context: {context}\n\nQuestion: {q}"}
-    ]
-    ai_reply = await run_ai_task(messages, 50, query_key=f"short_{q}")
-    return {"request_id": req_id, "ai_response": ai_reply, "results": results}
+    results, intent_info = await asyncio.to_thread(search_index.search, q, debug)
+    
+    # Strictly limit to 20 results output
+    top_20 = results[:STRICT_RESULT_LIMIT]
+    
+    if content and top_20:
+        tasks = [extract_page_content(r["url"]) for r in top_20]
+        extracted_data = await asyncio.gather(*tasks)
+        for r, ext in zip(top_20, extracted_data):
+            r["content"] = ext.get("content", "") if ext.get("success") else ""
+            
+    return {"request_id": req_id, "query": q, "results_count": len(top_20), "results": top_20}
 
-async def process_rag_chat(q: str):
+async def process_aimode(q: str, include_content: bool, debug: bool):
     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", "news"])
-    results, _, _ = await perform_core_search(q)
     
-    rag_context = ""
+    # 1. Deterministic Search (Max 20 internally outputted)
+    results, intent_info = await asyncio.to_thread(search_index.search, q, debug)
+    top_20 = results[:STRICT_RESULT_LIMIT]
+    query_tokens = tokenize(q)
+    
+    # Bypass AI cache if it's a News/Latest query
+    is_news = "NEWS" in intent_info["signals"]
+
+    rag_context_parts = []
     used_sources = []
 
-    if needs_rag and results:
-        top_urls = results[:3]
-        tasks = [scrape_and_find_best_chunk(doc["url"], q) for doc in top_urls]
-        scraped_chunks = await asyncio.gather(*tasks)
+    if top_20:
+        # 2. Extract content ONLY for Top ranked sources
+        tasks = [extract_page_content(r["url"]) for r in top_20]
+        extracted_data = await asyncio.gather(*tasks)
         
-        context_parts = []
-        for i, chunk in enumerate(scraped_chunks):
-            doc = top_urls[i]
-            url = doc["url"]
-            if chunk:
-                context_parts.append(f"Source [{i+1}] ({url}): {chunk}")
-                used_sources.append(doc)
-            else:
-                context_parts.append(f"Source [{i+1}] ({url}): {doc.get('description', '')}")
-                used_sources.append(doc)
+        # 3. Deterministic Chunk Selection for Context
+        valid_sources_count = 0
+        for i, (r, ext) in enumerate(zip(top_20, extracted_data)):
+            if include_content:
+                r["content"] = ext.get("content", "") if ext.get("success") else ""
                 
-        rag_context = "\n".join(context_parts)
-
-    system_prompt = "You are a Smart AI Search Assistant. Provide a short, direct answer. ENCLOSE IMPORTANT KEYWORDS IN (PARENTHESES). If you use information from a Source, cite it."
+            if ext.get("success") and ext.get("content") and valid_sources_count < 5: # Build context from top 5 valid pages
+                best_chunks = rank_chunks_for_rag(ext["content"], query_tokens)
+                if best_chunks:
+                    rag_context_parts.append(f"[Source {valid_sources_count+1}]\nTitle: {r['title']}\nURL: {r['url']}\nRelevant Content: {' ... '.join(best_chunks)}")
+                    used_sources.append({"id": valid_sources_count+1, "title": r["title"], "url": r["url"]})
+                    valid_sources_count += 1
+                    
+    rag_context = "\n\n".join(rag_context_parts)
+    system_prompt = (
+        "You are an intelligent search assistant. Answer the user's query DIRECTLY based ONLY on the provided Live Context Data. "
+        "Do not invent facts. If the context contains the answer, summarize it clearly and cite the sources used (e.g., [Source 1]). "
+        "If the context does not contain the answer, state that you do not have enough information."
+    )
+    
     messages = [
         {"role": "system", "content": system_prompt},
-        {"role": "user", "content": f"Live Context Data:\n{rag_context}\n\nQuestion: {q}" if needs_rag else f"Question: {q}"}
+        {"role": "user", "content": f"Live Context Data:\n{rag_context}\n\nQuestion: {q}" if rag_context else f"Question: {q}"}
     ]
 
-    ai_reply = await run_ai_task(messages, 120, query_key=f"chat_{q}_{needs_rag}") 
-    return {"request_id": req_id, "ai_response": ai_reply, "rag_used": needs_rag, "sources": used_sources, "results": results[:15]}
+    # 4. Generate AI Answer (Only Step utilizing LLM)
+    ai_reply = await run_ai_task(messages, max_tokens=150, query_key=f"ai_{q}", skip_cache=is_news)
+    
+    return {
+        "request_id": req_id, 
+        "query": q, 
+        "ai_response": ai_reply, 
+        "results_count": len(top_20), 
+        "results": top_20,
+        "sources": used_sources
+    }
 
-@app.get("/aimodechat")
-async def ai_mode_chat_get(q: str = Query(...)): return await process_rag_chat(q)
+@app.get("/aimode")
+async def ai_mode_get(q: str = Query(...), content: bool = Query(False), debug: bool = Query(False)): 
+    return await process_aimode(q, content, debug)
 
-@app.post("/aimodechat")
-async def ai_mode_chat_post(q: str = Form(...)): return await process_rag_chat(q)
+@app.post("/aimode")
+async def ai_mode_post(q: str = Form(...), content: bool = Form(False), debug: bool = Form(False)): 
+    return await process_aimode(q, content, debug)
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)