Update app.py

app.py

@@ -18,7 +18,7 @@ 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 huggingface_hub import snapshot_download
 import aiohttp
 from bs4 import BeautifulSoup
 import torch
@@ -48,7 +48,6 @@ MAX_CONTENT_CONCURRENCY = 10
 content_semaphore = asyncio.Semaphore(MAX_CONTENT_CONCURRENCY)
 
 # --- 1. DETERMINISTIC PREPROCESSING & CONSTANTS ---
-# Expanded stop words to include question words, so they don't overpower the core query
 STOP_WORDS = {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "is", "are", "by", "this", "that", "it", "what", "who", "where", "when", "why", "how"}
 ADULT_KEYWORDS = {
     "porn", "sex", "nude", "nsfw", "xxx", "adult", "erotic", "onlyfans", 
@@ -117,9 +116,17 @@ def detect_query_intent(query_tokens: List[str]) -> Dict[str, Any]:
 
 # --- 3. DEDICATED WEB CONTENT EXTRACTION ---
 async def extract_page_content(url: str) -> Dict[str, Any]:
+    # 1. First check if content is natively loaded inside our dataset to save time
+    if url in search_index.url_to_doc:
+        doc = search_index.url_to_doc[url]
+        if doc.get("_fast_content"):
+            return {"url": url, "title": doc.get("_fast_title", ""), "content": doc.get("_fast_content", ""), "success": True}
+    
+    # 2. Check general content cache
     if url in CONTENT_CACHE:
         return CONTENT_CACHE[url]
         
+    # 3. Fallback to live extraction if absolutely not in DB
     async with content_semaphore:
         try:
             async with aiohttp.ClientSession() as session:
@@ -174,6 +181,7 @@ def rank_chunks_for_rag(text: str, query_tokens: List[str]) -> List[str]:
 class HybridSearchIndex:
     def __init__(self):
         self.docs = []
+        self.url_to_doc = {}
         self.total_docs = 0
         self.vocab = set()
         self.df = defaultdict(int)
@@ -184,16 +192,24 @@ class HybridSearchIndex:
         self.b = 0.75
 
     def build(self, raw_data: List[dict]):
-        print("📥 Building Advanced Deterministic Search Index...")
+        print(f"📥 Building Advanced Deterministic Search Index from {len(raw_data)} documents...")
         total_len = 0
         suggestion_counter = Counter()
 
         for doc_id, doc in enumerate(raw_data):
             self.docs.append(doc)
+            self.url_to_doc[doc.get("url", "")] = 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
+            headings_tokens = tokenize(doc.get("_fast_headings", ""))
+            snippet_tokens = tokenize(doc.get("_fast_snippet", ""))
+            
+            # Using partial content tokens to save memory while maximizing recall
+            content_tokens = tokenize(doc.get("_fast_content", ""))[:400] 
+            
+            all_tokens = title_tokens + desc_tokens + domain_tokens + headings_tokens + snippet_tokens + content_tokens
             
             doc["_token_len"] = len(all_tokens)
             total_len += len(all_tokens)
@@ -213,7 +229,7 @@ class HybridSearchIndex:
         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.")
+        print(f"🚀 Indexing Complete! {self.total_docs} records fully embedded.")
 
     def idf(self, term: str) -> float:
         df = self.df.get(term, 0)
@@ -236,35 +252,62 @@ class HybridSearchIndex:
         desc = doc.get("_fast_desc", "")
         domain = doc.get("_fast_domain", "")
         url = doc.get("_fast_path", "")
+        headings = doc.get("_fast_headings", "")
+        content = doc.get("_fast_content", "")
+        snippet = doc.get("_fast_snippet", "")
         
         title_toks = tokenize(title, False)
         desc_toks = tokenize(desc, False)
+        headings_toks = tokenize(headings, 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, "root_domain": 0.0}
+        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, "root_domain": 0.0, "headings": 0.0, "body": 0.0}
 
         doc_len = doc.get("_token_len", 1)
+        content_lower = content.lower()
+        snippet_lower = snippet.lower()
+        
         for term in query_tokens_sorted:
-            tf = title_toks.count(term) * 2 + desc_toks.count(term)
+            # TF spans across newly added dataset fields
+            tf = (title_toks.count(term) * 4 + 
+                  headings_toks.count(term) * 3 + 
+                  desc_toks.count(term) * 2 + 
+                  snippet_lower.count(term) * 2 + 
+                  content_lower.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
+        title_norm = normalize_text(title)
+        desc_norm = normalize_text(desc)
+        content_norm = normalize_text(content)
+        headings_norm = normalize_text(headings)
+        
+        if q_norm and q_norm in title_norm: scores["exact"] = 1.0
+        elif q_norm and q_norm in desc_norm: scores["exact"] = 0.5
+        elif q_norm and q_norm in content_norm: scores["exact"] = 0.2
             
         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
+                if bigram in title_norm: phrase_overlap += 1.0
+                elif bigram in headings_norm: phrase_overlap += 0.8
+                elif bigram in desc_norm: phrase_overlap += 0.5
+                elif bigram in content_norm: phrase_overlap += 0.3
             scores["phrase"] = phrase_overlap / (len(query_toks_full) - 1)
 
         title_overlap = sum(1 for t in query_tokens_sorted if t in title_toks)
         scores["title"] = title_overlap / max(1, len(query_tokens_sorted))
+        
+        headings_overlap = sum(1 for t in query_tokens_sorted if t in headings_toks)
+        scores["headings"] = headings_overlap / max(1, len(query_tokens_sorted))
+        
+        body_overlap = sum(1 for t in query_tokens_sorted if t in content_lower)
+        scores["body"] = body_overlap / max(1, len(query_tokens_sorted))
 
         q_joined = "".join(query_tokens_sorted)
         domain_clean = domain.split('.')[0]
@@ -277,7 +320,7 @@ class HybridSearchIndex:
             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)
+        year_match = re.search(r'\b(202[0-6])\b', url + " " + desc + " " + snippet)
         if year_match: scores["freshness"] = math.exp(-(datetime.now().year - int(year_match.group(1))) / 2.0)
         else: scores["freshness"] = 0.3 
 
@@ -295,22 +338,16 @@ class HybridSearchIndex:
 
     def search(self, query: str, debug: bool = False) -> Tuple[List[dict], Dict[str, Any]]:
         query_norm = normalize_text(query)
-        # Tokenize (Removes standard stop words)
         query_tokens = tokenize(query_norm)
-        
-        # We append normal words that user typed but were omitted, if they exist
         raw_tokens = normalize_text(query).split()
         
         expanded_tokens = list(query_tokens)
         for t in query_tokens:
             if t in EXPANSION_DICT: expanded_tokens.extend(EXPANSION_DICT[t])
         
-        # STRING SORTING ALGORITHM BY IMPORTANCE (Solves 'what is github' issue)
-        # We sort tokens so the RAREST words (Highest IDF) are verified first.
-        # Example: 'github' or '5000' will be at the very front of the array. 'what' or 'is' drops to the back.
         sorted_tokens = list(dict.fromkeys(expanded_tokens)) 
         if not sorted_tokens:
-            sorted_tokens = raw_tokens # Fallback if query was literally just "what is"
+            sorted_tokens = raw_tokens 
             
         sorted_tokens.sort(key=lambda t: self.idf(t), reverse=True)
         
@@ -339,36 +376,46 @@ class HybridSearchIndex:
             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"] + 
-                WEIGHTS["W_ROOT_DOMAIN"] * features["root_domain"] 
+                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"] + 
+                WEIGHTS["W_ROOT_DOMAIN"] * features["root_domain"] +
+                WEIGHTS["W_HEADINGS"] * features["headings"] +
+                WEIGHTS["W_BODY"] * features["body"]
             )
             
-            # --- CORE TERM PENALTY (CRITICAL FIX) ---
-            # If the #1 most important sorted word (e.g., 'github' or '5000') is completely missing from the website,
-            # we apply a massive 90% score drop. This kills random "DNS IP" matches instantly!
+            # Penalty for Missing Core Word
             if len(sorted_tokens) > 0:
                 top_term = sorted_tokens[0]
                 doc_title = tokenize(doc.get("_fast_title", ""), False)
                 doc_desc = tokenize(doc.get("_fast_desc", ""), False)
-                if top_term not in doc_title and top_term not in doc_desc and top_term not in doc.get("_fast_domain", ""):
+                doc_content = doc.get("_fast_content", "").lower()
+                doc_headings = tokenize(doc.get("_fast_headings", ""), False)
+                
+                if top_term not in doc_title and top_term not in doc_desc and top_term not in doc.get("_fast_domain", "") and top_term not in doc_content and top_term not in doc_headings:
                     final_score *= 0.1 # 90% Penalty
                     
             if len(sorted_tokens) > 1:
                 second_term = sorted_tokens[1]
-                if second_term not in doc_title and second_term not in doc_desc and second_term not in doc.get("_fast_domain", ""):
-                    final_score *= 0.4 # 60% Penalty for missing second most important word
-            # -----------------------------------------
-
+                if second_term not in doc_title and second_term not in doc_desc and second_term not in doc.get("_fast_domain", "") and second_term not in doc_content and second_term not in doc_headings:
+                    final_score *= 0.4 # 60% Penalty
+            
             if final_score > 0.05:
                 doc_out = {
+                    "doc_id": doc.get("doc_id", ""),
                     "title": doc.get("_fast_title", ""),
                     "description": doc.get("_fast_desc", ""),
+                    "snippet": doc.get("_fast_snippet", ""),
                     "url": doc.get("url", ""),
                     "domain": doc.get("_fast_domain", ""),
-                    "match_score": round(final_score, 4)
+                    "match_score": round(final_score, 4),
+                    "_internal_content": doc.get("_fast_content", "") # Passed internally for Fast Content injection
                 }
                 if debug:
                     doc_out["_debug_features"] = features
@@ -392,31 +439,66 @@ class HybridSearchIndex:
 search_index = HybridSearchIndex()
 
 def load_and_index_dataset():
-    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
-
+    print("📦 Automatically Downloading 'Nebulixlabs/web-search-dataset' ...")
+    dataset_path = "."
+    try:
+        dataset_path = snapshot_download(
+            repo_id="Nebulixlabs/web-search-dataset", 
+            repo_type="dataset", 
+            allow_patterns="data/*.jsonl",
+            token=HF_TOKEN
+        )
+    except Exception as e:
+        print("⚠️ Direct HF Hub fetch failed, attempting local 'data' folder mapping. Error:", e)
+        
     raw_data = []
-    if os.path.exists(dataset_path):
-        with open(dataset_path, "r", encoding="utf-8") as f:
-            for i, line in enumerate(f):
-                if i >= 150000: break 
-                if line.strip():
+    data_folder = os.path.join(dataset_path, "data")
+    if not os.path.exists(data_folder):
+        data_folder = "data" # Fallback local
+        
+    if os.path.exists(data_folder):
+        jsonl_files = [f for f in os.listdir(data_folder) if f.endswith(".jsonl")]
+        for file in jsonl_files:
+            file_path = os.path.join(data_folder, file)
+            with open(file_path, "r", encoding="utf-8") as f:
+                for line in f:
+                    if not line.strip(): continue
                     try:
                         doc = json.loads(line)
                         url = doc.get("url", "")
-                        title_str = doc.get("title", "")
-                        desc_str = doc.get("description", "")
+                        
+                        # Parsing New Nested Schema
+                        metadata = doc.get("metadata") or {}
+                        content_dict = doc.get("content") or {}
+                        
+                        title_str = metadata.get("title") or ""
+                        desc_str = metadata.get("description") or ""
+                        snippet_str = content_dict.get("snippet") or ""
+                        main_text = content_dict.get("main_text") or ""
+                        
+                        headings = content_dict.get("headings") or {}
+                        h1_str = " ".join(headings.get("h1") or [])
+                        h2_str = " ".join(headings.get("h2") or [])
+                        
                         parsed_url = urllib.parse.urlparse(url)
                         
-                        doc["_fast_domain"] = parsed_url.netloc.lower().replace("www.", "")
+                        doc["_fast_domain"] = doc.get("domain") or parsed_url.netloc.lower().replace("www.", "")
                         doc["_fast_path"] = parsed_url.path.strip("/")
                         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)
+                        doc["_fast_snippet"] = snippet_str
+                        doc["_fast_content"] = main_text
+                        doc["_fast_headings"] = h1_str + " " + h2_str
+                        
+                        combined_adult_check = f"{title_str} {desc_str} {url} {snippet_str} {main_text[:200]}".lower()
+                        doc["_is_adult"] = any(aw in combined_adult_check for aw in ADULT_KEYWORDS)
+                        
                         raw_data.append(doc)
-                    except: pass
+                        if len(raw_data) >= 150000: break
+                    except Exception: 
+                        pass
+            if len(raw_data) >= 150000: break
+            
     search_index.build(raw_data)
 
 
@@ -442,13 +524,9 @@ def load_ai_models():
 
 async def run_ai_task(messages, max_tokens: int, query_key: str):
     if not ai_generator: return "⚠️ AI Model is starting up."
+    if query_key in AI_RESPONSE_CACHE: return AI_RESPONSE_CACHE[query_key]
     
-    if query_key in AI_RESPONSE_CACHE: 
-        return AI_RESPONSE_CACHE[query_key]
-    
-    # --- FCFS / GRACEFUL DEGRADATION LOGIC ---
-    try:
-        await asyncio.wait_for(ai_semaphore.acquire(), timeout=0.1)
+    try: await asyncio.wait_for(ai_semaphore.acquire(), timeout=0.1)
     except asyncio.TimeoutError:
         return "⚠️ AI Servers are fully occupied right now. Displaying your ultra-fast search results below! (Please try AI Mode again in a few seconds)"
 
@@ -498,11 +576,11 @@ async def pure_search(q: str = Query(...), content: bool = Query(False), debug:
     results, intent_info = await asyncio.to_thread(search_index.search, q, debug)
     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 ""
+    for r in top_20:
+        if content:
+            r["content"] = r.get("_internal_content", "")
+        # Remove massive string payload to keep payload size lightweight unless user explicitly asked for it
+        r.pop("_internal_content", None)
             
     return {"request_id": req_id, "query": q, "results_count": len(top_20), "results": top_20}
 
@@ -510,39 +588,47 @@ async def process_aimode(q: str, include_content: bool, debug: bool, max_ai_toke
     req_id = str(uuid.uuid4())
     cache_key = f"ai_{normalize_text(q)}_{max_ai_tokens}"
     
+    results, intent_info = await asyncio.to_thread(search_index.search, q, debug)
+    top_20 = results[:STRICT_RESULT_LIMIT]
+    
     if cache_key in AI_RESPONSE_CACHE:
-        results, _ = await asyncio.to_thread(search_index.search, q, debug)
+        for r in top_20:
+            if include_content: r["content"] = r.get("_internal_content", "")
+            r.pop("_internal_content", None)
+            
         return {
             "request_id": req_id, 
             "query": q, 
             "ai_response": AI_RESPONSE_CACHE[cache_key], 
-            "results_count": len(results[:STRICT_RESULT_LIMIT]), 
-            "results": results[:STRICT_RESULT_LIMIT],
+            "results_count": len(top_20), 
+            "results": top_20,
             "sources": [],
             "cached": True
         }
 
-    results, intent_info = await asyncio.to_thread(search_index.search, q, debug)
-    top_20 = results[:STRICT_RESULT_LIMIT]
     query_tokens = tokenize(q)
-    
     rag_context_parts = []
     used_sources = []
 
     if top_20:
-        tasks = [extract_page_content(r["url"]) for r in top_20[:3]] 
-        extracted_data = await asyncio.gather(*tasks)
-        
         valid_sources_count = 0
-        for i, (r, ext) in enumerate(zip(top_20[:3], extracted_data)):
-            if include_content: r["content"] = ext.get("content", "") if ext.get("success") else ""
+        for r in top_20[:3]:
+            # Directly parse context from our ultra-fast pre-loaded dataset without web scraping!
+            doc_content = r.get("_internal_content", "")
+            
+            if include_content: 
+                r["content"] = doc_content
                 
-            if ext.get("success") and ext.get("content") and valid_sources_count < 3:
-                best_chunks = rank_chunks_for_rag(ext["content"], query_tokens)
+            if doc_content and valid_sources_count < 3:
+                best_chunks = rank_chunks_for_rag(doc_content, query_tokens)
                 if best_chunks:
                     rag_context_parts.append(f"[{valid_sources_count+1}]: {' ... '.join(best_chunks)}")
                     used_sources.append({"id": valid_sources_count+1, "title": r["title"], "url": r["url"]})
                     valid_sources_count += 1
+        
+        # Cleanup Internal Fields globally before returning JSON Payload
+        for r in top_20:
+            r.pop("_internal_content", None)
                     
     # STRICT COMPLETION PROMPT: Guarantees it finishes its sentence within the token limit gracefully.
     system_prompt = (
@@ -564,7 +650,6 @@ async def process_aimode(q: str, include_content: bool, debug: bool, max_ai_toke
         {"role": "user", "content": combined_input}
     ]
 
-    # Generating the answer
     ai_reply = await run_ai_task(messages, max_tokens=max_ai_tokens, query_key=cache_key)
     
     return {
@@ -579,12 +664,10 @@ async def process_aimode(q: str, include_content: bool, debug: bool, max_ai_toke
 
 @app.get("/aimode")
 async def ai_mode_get(q: str = Query(...), content: bool = Query(False), debug: bool = Query(False)): 
-    # STRICTLY forced to 20 tokens as requested
     return await process_aimode(q, content, debug, max_ai_tokens=20) 
 
 @app.post("/aimode")
 async def ai_mode_post(q: str = Form(...), content: bool = Form(False), debug: bool = Form(False)): 
-    # STRICTLY forced to 20 tokens as requested
     return await process_aimode(q, content, debug, max_ai_tokens=20)
 
 if __name__ == "__main__":