Update app.py

app.py

@@ -23,7 +23,7 @@ import aiohttp
 from bs4 import BeautifulSoup
 import torch
 
-# Hugging Face imports for Q4 Quantization
+# Hugging Face imports
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
 
 # --- CONFIGURATION & MULTIPROCESSING ---
@@ -37,9 +37,9 @@ ai_semaphore = asyncio.Semaphore(MAX_CONCURRENT_AI)
 ai_generator = None
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Cache Structures
+# Strict Cache Structures (Guarantees no re-runs for identical queries)
 AI_RESPONSE_CACHE = {}
-MAX_AI_CACHE_SIZE = 5000
+MAX_AI_CACHE_SIZE = 10000
 
 CONTENT_CACHE = {}
 MAX_CONTENT_CONCURRENCY = 10
@@ -113,31 +113,27 @@ 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]:
-    """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:
+                async with session.get(url, timeout=3.0, headers={"User-Agent": "Mozilla/5.0 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)
+                        if len(text) > 30: blocks.append(text)
                             
-                    # Deduplicate contiguous identical text
                     clean_blocks = []
                     seen_blocks = set()
                     for b in blocks:
@@ -148,22 +144,13 @@ async def extract_page_content(url: str) -> Dict[str, Any]:
                     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()
-                    }
+                    result = {"url": url, "title": page_title, "content": content_str, "success": True}
                     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 []
     chunks = text.split("\n\n")
     q_set = set(query_tokens)
@@ -172,14 +159,11 @@ def rank_chunks_for_rag(text: str, query_tokens: List[str]) -> List[str]:
     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 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
-        
+    if not scored_chunks and chunks: scored_chunks = [(0, c) for c in chunks[:3]]
     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
+    return [c for score, c in scored_chunks[:2]] # Strictly reduced for token limits
 
 
 # --- 4. CORE DETERMINISTIC SEARCH INDEX ---
@@ -290,8 +274,7 @@ class HybridSearchIndex:
                 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)
+        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
@@ -367,12 +350,9 @@ class HybridSearchIndex:
             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():
@@ -403,12 +383,12 @@ def load_and_index_dataset():
                     except: pass
     search_index.build(raw_data)
 
-# --- QWEN AI GENERATION MODULE (FOR RAG ONLY) ---
+# --- AI GENERATION MODULE (SMOLLM 135M) ---
 def load_ai_models():
     global ai_generator
-    print("🤖 Loading Qwen/Qwen2.5-0.5B-Instruct...")
+    print("🤖 Loading HuggingFaceTB/SmolLM-135M-Instruct (Ultra Fast)...")
     try:
-        model_id = "Qwen/Qwen2.5-0.5B-Instruct"
+        model_id = "HuggingFaceTB/SmolLM-135M-Instruct"
         tokenizer = AutoTokenizer.from_pretrained(model_id)
 
         if torch.cuda.is_available():
@@ -418,13 +398,16 @@ def load_ai_models():
             model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cpu", torch_dtype=torch.float32)
 
         ai_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-        print("✅ Qwen AI Ready!")
+        print("✅ SmolLM AI Ready!")
     except Exception as e:
         print(f"❌ Model Loading Failed: {e}")
 
-async def run_ai_task(messages, max_tokens: int, query_key: str, skip_cache: bool = False):
+async def run_ai_task(messages, max_tokens: int, query_key: str):
     if not ai_generator: return "⚠️ AI Model is starting up."
-    if not skip_cache and query_key in AI_RESPONSE_CACHE: return AI_RESPONSE_CACHE[query_key]
+    
+    # EXACT CACHE MATCH - Guarantees User 1 and User 2 get identical cached response
+    if query_key in AI_RESPONSE_CACHE: 
+        return AI_RESPONSE_CACHE[query_key]
     
     async with ai_semaphore:
         def _gen():
@@ -436,9 +419,11 @@ async def run_ai_task(messages, max_tokens: int, query_key: str, skip_cache: boo
             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
+        
+        # Unconditionally cache to save computing power
+        if len(AI_RESPONSE_CACHE) > MAX_AI_CACHE_SIZE: AI_RESPONSE_CACHE.clear()
+        AI_RESPONSE_CACHE[query_key] = response
+        
         return response
 
 # --- FASTAPI APP ENDPOINTS ---
@@ -459,7 +444,6 @@ async def get_suggestions(q: str = Query(..., min_length=1)):
 
 @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
 
@@ -467,8 +451,6 @@ async def get_content_endpoint(url: str = Query(...)):
 async def pure_search(q: str = Query(...), content: bool = Query(False), debug: bool = Query(False)):
     req_id = str(uuid.uuid4())
     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:
@@ -479,52 +461,64 @@ async def pure_search(q: str = Query(...), content: bool = Query(False), debug:
             
     return {"request_id": req_id, "query": q, "results_count": len(top_20), "results": top_20}
 
-async def process_aimode(q: str, include_content: bool, debug: bool):
+async def process_aimode(q: str, include_content: bool, debug: bool, max_ai_tokens: int):
     req_id = str(uuid.uuid4())
+    cache_key = f"ai_{normalize_text(q)}_{max_ai_tokens}"
     
-    # 1. Deterministic Search (Max 20 internally outputted)
+    # 1. IMMEDIATE CACHE CHECK (0s latency if previously asked)
+    if cache_key in AI_RESPONSE_CACHE:
+        results, _ = await asyncio.to_thread(search_index.search, q, debug)
+        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],
+            "sources": [],
+            "cached": True
+        }
+
+    # 2. Deterministic Search 
     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 top_20:
-        # 2. Extract content ONLY for Top ranked sources
-        tasks = [extract_page_content(r["url"]) for r in top_20]
+        tasks = [extract_page_content(r["url"]) for r in top_20[:3]] # Only fetch top 3 to keep it fast
         extracted_data = await asyncio.gather(*tasks)
         
-        # 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 ""
+        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 ""
                 
-            if ext.get("success") and ext.get("content") and valid_sources_count < 5: # Build context from top 5 valid pages
+            if ext.get("success") and ext.get("content") and valid_sources_count < 3:
                 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)}")
+                    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
                     
-    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."
-    )
+    # 3. STRICT 100 TOKEN INPUT LIMIT (User Query + RAG)
+    system_prompt = "Answer directly based on context. Cite sources."
+    combined_input = f"Context:\n{chr(10).join(rag_context_parts)}\n\nQuery: {q}"
     
+    if ai_generator:
+        tokenizer = ai_generator.tokenizer
+        encoded_input = tokenizer.encode(combined_input, add_special_tokens=False)
+        if len(encoded_input) > 90: # 90 + 10 for system prompt = 100 max tokens input
+            encoded_input = encoded_input[:90]
+            combined_input = tokenizer.decode(encoded_input, skip_special_tokens=True)
+
     messages = [
         {"role": "system", "content": system_prompt},
-        {"role": "user", "content": f"Live Context Data:\n{rag_context}\n\nQuestion: {q}" if rag_context else f"Question: {q}"}
+        {"role": "user", "content": combined_input}
     ]
 
-    # 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)
+    # 4. Generate Answer (Restricted to 20 or 50 tokens based on GET/POST)
+    ai_reply = await run_ai_task(messages, max_tokens=max_ai_tokens, query_key=cache_key)
     
     return {
         "request_id": req_id, 
@@ -532,16 +526,17 @@ async def process_aimode(q: str, include_content: bool, debug: bool):
         "ai_response": ai_reply, 
         "results_count": len(top_20), 
         "results": top_20,
-        "sources": used_sources
+        "sources": used_sources,
+        "cached": False
     }
 
 @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)
+    return await process_aimode(q, content, debug, max_ai_tokens=20) # Strictly 20 Tokens
 
 @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)
+    return await process_aimode(q, content, debug, max_ai_tokens=50) # Strictly 50 Tokens
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)