Update app.py

app.py

@@ -28,16 +28,18 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndB
 
 # --- CONFIGURATION & MULTIPROCESSING ---
 TOTAL_CORES = multiprocessing.cpu_count()
-AI_CORES = max(1, int(TOTAL_CORES * 0.7)) 
+# STRICTLY 40% CPU FOR AI, 60% FOR SEARCH
+AI_CORES = max(1, int(TOTAL_CORES * 0.40)) 
 torch.set_num_threads(AI_CORES)
-print(f"🖥️ Total Cores: {TOTAL_CORES} | Dedicated AI Cores: {AI_CORES}")
+print(f"🖥️ Total Cores: {TOTAL_CORES} | Dedicated AI Cores (40%): {AI_CORES} | Dedicated Search Cores: {TOTAL_CORES - AI_CORES}")
 
-MAX_CONCURRENT_AI = max(1, AI_CORES // 2)
+# First come first serve configuration
+MAX_CONCURRENT_AI = max(1, AI_CORES)
 ai_semaphore = asyncio.Semaphore(MAX_CONCURRENT_AI)
 ai_generator = None
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Strict Cache Structures (Guarantees no re-runs for identical queries)
+# Strict Cache Structures
 AI_RESPONSE_CACHE = {}
 MAX_AI_CACHE_SIZE = 10000
 
@@ -68,7 +70,8 @@ EXPANSION_DICT = {
 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
+    "W_QUALITY": 1.0, "W_POPULARITY": 0.5, "W_PROXIMITY": 1.5, 
+    "W_ROOT_DOMAIN": 8.0 # Extreme priority for root domains (domain.com over domain.com/path)
 }
 
 MAX_RESULTS_PER_DOMAIN = 3
@@ -163,7 +166,7 @@ def rank_chunks_for_rag(text: str, query_tokens: List[str]) -> List[str]:
             
     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[:2]] # Strictly reduced for token limits
+    return [c for score, c in scored_chunks[:2]] 
 
 
 # --- 4. CORE DETERMINISTIC SEARCH INDEX ---
@@ -237,7 +240,7 @@ class HybridSearchIndex:
         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}
+        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}
 
         doc_len = doc.get("_token_len", 1)
         for term in query_tokens:
@@ -284,6 +287,10 @@ class HybridSearchIndex:
         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))
 
+        # ROOT DOMAIN DETECTION (Domain.com is always higher than Domain.com/path)
+        if url == "" or url == "/":
+            scores["root_domain"] = 1.0
+
         return scores
 
     def search(self, query: str, debug: bool = False) -> Tuple[List[dict], Dict[str, Any]]:
@@ -323,7 +330,8 @@ class HybridSearchIndex:
                 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_QUALITY"] * features["quality"] + WEIGHTS["W_PROXIMITY"] * features["proximity"] + 
+                WEIGHTS["W_ROOT_DOMAIN"] * features["root_domain"] # Hugely boosts root domains
             )
             
             if final_score > 0.05:
@@ -383,33 +391,43 @@ def load_and_index_dataset():
                     except: pass
     search_index.build(raw_data)
 
-# --- AI GENERATION MODULE (SMOLLM 135M) ---
+
+# --- AI GENERATION MODULE (~200-360M Params, 4-bit Quantized) ---
 def load_ai_models():
     global ai_generator
-    print("🤖 Loading HuggingFaceTB/SmolLM-135M-Instruct (Ultra Fast)...")
+    print("🤖 Loading HuggingFaceTB/SmolLM2-360M-Instruct (Ultra Fast 4-Bit)...")
     try:
-        model_id = "HuggingFaceTB/SmolLM-135M-Instruct"
+        # Chosen ~360M parameter model. Extremely fast and follows instructions aggressively.
+        model_id = "HuggingFaceTB/SmolLM2-360M-Instruct"
         tokenizer = AutoTokenizer.from_pretrained(model_id)
 
+        # Loading strictly in 4-bit config for maximum efficiency
+        quant_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16, bnb_4bit_quant_type="nf4")
+        
         if torch.cuda.is_available():
-            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:
             model = AutoModelForCausalLM.from_pretrained(model_id, device_map="cpu", torch_dtype=torch.float32)
 
         ai_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-        print("✅ SmolLM AI Ready!")
+        print("✅ Fast 4-Bit AI Ready!")
     except Exception as e:
         print(f"❌ Model Loading Failed: {e}")
 
 async def run_ai_task(messages, max_tokens: int, query_key: str):
     if not ai_generator: return "⚠️ AI Model is starting up."
     
-    # 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:
+    # --- FCFS / GRACEFUL DEGRADATION LOGIC ---
+    # Try to enter the AI queue. If CPU is busy (locked after 0.1s), instantly fallback to Search!
+    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)"
+
+    try:
         def _gen():
             tokenizer = ai_generator.tokenizer
             try: prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
@@ -420,11 +438,13 @@ async def run_ai_task(messages, max_tokens: int, query_key: str):
 
         response = await asyncio.to_thread(_gen)
         
-        # 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
+    finally:
+        ai_semaphore.release()
+
 
 # --- FASTAPI APP ENDPOINTS ---
 @asynccontextmanager
@@ -465,7 +485,7 @@ 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}"
     
-    # 1. IMMEDIATE CACHE CHECK (0s latency if previously asked)
+    # 1. IMMEDIATE CACHE CHECK 
     if cache_key in AI_RESPONSE_CACHE:
         results, _ = await asyncio.to_thread(search_index.search, q, debug)
         return {
@@ -478,7 +498,7 @@ async def process_aimode(q: str, include_content: bool, debug: bool, max_ai_toke
             "cached": True
         }
 
-    # 2. Deterministic Search 
+    # 2. Deterministic Search (Now prioritizing Domain.com natively via Deep Analytics)
     results, intent_info = await asyncio.to_thread(search_index.search, q, debug)
     top_20 = results[:STRICT_RESULT_LIMIT]
     query_tokens = tokenize(q)
@@ -487,7 +507,7 @@ async def process_aimode(q: str, include_content: bool, debug: bool, max_ai_toke
     used_sources = []
 
     if top_20:
-        tasks = [extract_page_content(r["url"]) for r in top_20[:3]] # Only fetch top 3 to keep it fast
+        tasks = [extract_page_content(r["url"]) for r in top_20[:3]] 
         extracted_data = await asyncio.gather(*tasks)
         
         valid_sources_count = 0
@@ -501,14 +521,15 @@ async def process_aimode(q: str, include_content: bool, debug: bool, max_ai_toke
                     used_sources.append({"id": valid_sources_count+1, "title": r["title"], "url": r["url"]})
                     valid_sources_count += 1
                     
-    # 3. STRICT 100 TOKEN INPUT LIMIT (User Query + RAG)
-    system_prompt = "Answer directly based on context. Cite sources."
+    # 3. STRICT TOKEN INSTRUCTIONS (1 Token = ~4 words mapping handled contextually)
+    max_words_instruction = max_ai_tokens * 4
+    system_prompt = f"Answer concisely using the given context. You MUST conclude your entire response exactly within {max_ai_tokens} tokens (which is roughly {max_words_instruction} words limit). Do not leave sentences incomplete. Summarize directly."
     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
+        if len(encoded_input) > 90:
             encoded_input = encoded_input[:90]
             combined_input = tokenizer.decode(encoded_input, skip_special_tokens=True)
 
@@ -517,7 +538,7 @@ async def process_aimode(q: str, include_content: bool, debug: bool, max_ai_toke
         {"role": "user", "content": combined_input}
     ]
 
-    # 4. Generate Answer (Restricted to 20 or 50 tokens based on GET/POST)
+    # 4. Generate Answer FCFS
     ai_reply = await run_ai_task(messages, max_tokens=max_ai_tokens, query_key=cache_key)
     
     return {
@@ -532,11 +553,11 @@ 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)): 
-    return await process_aimode(q, content, debug, max_ai_tokens=20) # Strictly 20 Tokens
+    return await process_aimode(q, content, debug, max_ai_tokens=20) # Strictly 20 Tokens Limit (~80 Words)
 
 @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, max_ai_tokens=50) # Strictly 50 Tokens
+    return await process_aimode(q, content, debug, max_ai_tokens=50) # Strictly 50 Tokens Limit (~200 Words)
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)