Update app.py

app.py

@@ -1,6 +1,6 @@
 # ============================================================
 #                    X-RUDRA CHAT
-#          Hugging Face Gradio Space Edition
+#          Dual‑Model + Web Research · Gradio Space
 # ============================================================
 
 from __future__ import annotations
@@ -21,24 +21,19 @@ from gradio_client import Client
 # ============================================================
 
 APP_NAME = "X-RUDRA"
-VERSION = "3.3.1"
+VERSION = "3.5.0"   # bumped version
 
-# M1 and M2 Spaces – change these to your own if needed
 M1_REPO = os.getenv("M1_REPO", "Shrijanagain/M1")
 M2_REPO = os.getenv("M2_REPO", "Shrijanagain/M2")
 PORT = int(os.getenv("PORT", "7860"))
 
-# ------------------------------------------------------------
-# HF_TOKEN is recommended to avoid rate limits
-# Set it as a Secret in your Space settings.
-# ------------------------------------------------------------
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN:
-    os.environ["HF_TOKEN"] = HF_TOKEN  # ensures the transformers lib uses it
+    os.environ["HF_TOKEN"] = HF_TOKEN
 
 
 # ============================================================
-# GRADIO CLIENTS FOR M1 / M2 (LAZY)
+# GRADIO CLIENTS FOR M1 / M2
 # ============================================================
 
 _M1_CLIENT = None
@@ -48,7 +43,6 @@ def get_m1_client():
     global _M1_CLIENT
     if _M1_CLIENT is None:
         try:
-            # Construct the public URL of the Space
             url = f"https://{M1_REPO.replace('/', '-')}.hf.space"
             _M1_CLIENT = Client(url)
         except Exception as e:
@@ -73,10 +67,14 @@ def get_m2_client():
 # ============================================================
 
 def is_casual_query(text: str) -> bool:
-    """Return True if the query is casual/timepass."""
     text = text.lower().strip()
-    if len(text.split()) <= 2:
+    words = text.split()
+    if len(words) <= 2:
         return True
+    question_words = {"what", "how", "why", "when", "where", "who", "which",
+                      "can", "could", "would", "will", "is", "are", "do", "does"}
+    if words[0] in question_words and len(words) >= 3:
+        return False
     casual_patterns = [
         "hey", "hi", "hello", "yo", "what's up", "how are you",
         "good morning", "good evening", "good night",
@@ -85,13 +83,13 @@ def is_casual_query(text: str) -> bool:
         "who are you", "what can you do", "help", "thanks"
     ]
     for pattern in casual_patterns:
-        if text.startswith(pattern) or pattern in text:
+        if text == pattern or text.startswith(pattern + " "):
             return True
     return False
 
 
 # ============================================================
-# LAZY ENGINE (for serious queries)
+# LAZY ENGINE (web search)
 # ============================================================
 
 _ENGINE = None
@@ -104,6 +102,129 @@ def get_engine():
     return _ENGINE
 
 
+# ============================================================
+# HELPERS – CALL MODELS
+# ============================================================
+
+def call_model(client, prompt, max_tokens=512, temperature=0.7):
+    if client is None:
+        return None
+    try:
+        result = client.predict(
+            prompt=prompt,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            api_name="/generate"
+        )
+        if result and isinstance(result, str) and result.strip():
+            return result.strip()
+    except Exception as e:
+        print(f"Model call failed: {e}")
+    return None
+
+
+# ============================================================
+# SYNTHESIS: M1 (draft) → M2 (refine)
+# ============================================================
+
+def get_combined_model_answer(question, sources, max_tokens=512, temperature=0.7):
+    # Build a concise summary of top sources (max 5)
+    top_sources = sources[:5] if sources else []
+    sources_text = ""
+    if top_sources:
+        for i, src in enumerate(top_sources, 1):
+            title = src.get("title", "Untitled")
+            snippet = src.get("snippet", src.get("description", ""))
+            sources_text += f"{i}. {title}: {snippet[:300]}\n"
+    else:
+        sources_text = "No specific information available."
+
+    # ----- 1. M1 generates draft answer -----
+    prompt_m1 = f"""Question: {question}
+
+Information:
+{sources_text}
+
+Based on the information above and your knowledge, provide a comprehensive, accurate, and well‑structured answer to the question. Be direct and natural – write as if you are an expert answering a user.
+
+Answer:"""
+
+    m1_client = get_m1_client()
+    m2_client = get_m2_client()
+
+    draft = call_model(m1_client, prompt_m1, max_tokens, temperature)
+    if not draft:
+        # Fallback: try M2 directly for draft
+        draft = call_model(m2_client, prompt_m1, max_tokens, temperature)
+
+    if not draft:
+        # Ultimate fallback – simple summarization
+        if sources:
+            parts = ["Based on available information:"]
+            for i, src in enumerate(sources[:5], 1):
+                title = src.get("title", "Untitled")
+                snippet = src.get("snippet", src.get("description", ""))
+                parts.append(f"{i}. {title}: {snippet[:200]}..." if snippet else f"{i}. {title}")
+            return "\n\n".join(parts)
+        else:
+            return "I couldn't find specific information on that topic. Could you rephrase?"
+
+    # ----- 2. M2 refines the draft -----
+    prompt_m2 = f"""Question: {question}
+
+Draft answer:
+{draft}
+
+Please refine and improve this answer to make it more comprehensive, accurate, and natural. Ensure it directly addresses the question. Provide only the final improved answer, without any extra commentary or meta‑references.
+
+Improved answer:"""
+
+    refined = call_model(m2_client, prompt_m2, max_tokens, temperature)
+    if refined:
+        return refined
+    else:
+        return draft
+
+
+# ============================================================
+# CASUAL REPLY (calls M1 or M2)
+# ============================================================
+
+def get_casual_model_response(query: str) -> str:
+    client = get_m1_client()
+    if client is not None:
+        try:
+            result = client.predict(
+                prompt=f"User: {query}\nAssistant:",
+                max_tokens=64,
+                temperature=0.7,
+                api_name="/generate"
+            )
+            if result and isinstance(result, str) and result.strip():
+                return result.strip()
+        except Exception as e:
+            print(f"M1 casual failed: {e}")
+
+    client = get_m2_client()
+    if client is not None:
+        try:
+            result = client.predict(
+                prompt=f"User: {query}\nAssistant:",
+                max_tokens=64,
+                temperature=0.7,
+                api_name="/generate"
+            )
+            if result and isinstance(result, str) and result.strip():
+                return result.strip()
+        except Exception as e:
+            print(f"M2 casual failed: {e}")
+
+    return (
+        f"👋 Hi there! I'm X‑RUDRA, your research assistant. "
+        f"How can I help you today? (Your message `{query}` was casual, so I kept it light.)"
+    )
+
+
 # ============================================================
 # FORMATTERS (Sources, Evidence, Verification)
 # ============================================================
@@ -174,14 +295,6 @@ def format_verification(contradictions):
     return "\n\n".join(output)
 
 
-def extract_answer(data):
-    for key in ("final_answer", "answer", "response", "final", "synthesis", "summary"):
-        val = data.get(key)
-        if isinstance(val, str) and val.strip():
-            return val.strip()
-    return None
-
-
 def build_activity(data, elapsed_ms):
     sources = data.get("sources", []) or data.get("results", [])
     claims = data.get("claims", [])
@@ -193,13 +306,13 @@ def build_activity(data, elapsed_ms):
 | Stage | Status |
 |---|---|
 | Task analysis | ✅ Complete |
-| M1 research | {"✅ Enabled" if data.get("m1") is not None else "⚙️ Pipeline"} |
-| M2 research | {"✅ Enabled" if data.get("m2") is not None else "⚙️ Pipeline"} |
+| M1 research | ✅ Generated draft |
+| M2 research | ✅ Refined answer |
 | Web discovery | ✅ Complete |
-| Evidence extraction | {"✅" if data.get("claims") else "⚙️"} |
+| Evidence extraction | {"✅" if claims else "⚙️"} |
 | Source verification | ✅ Complete |
 | Contradiction check | {"⚠️ Found" if contradictions else "✅ Clear"} |
-| Final synthesis | {"✅" if extract_answer(data) else "⚙️"} |
+| Final synthesis | ✅ Complete |
 
 **Sources:** `{len(sources)}`  
 **Claims:** `{len(claims)}`  
@@ -240,50 +353,6 @@ def safe_dict(value):
     return {"result": str(value)}
 
 
-# ============================================================
-# GET MODEL RESPONSE FOR CASUAL QUERIES
-# ============================================================
-
-def get_casual_model_response(query: str) -> str:
-    """Call M1 (or M2) to generate a friendly reply for casual queries."""
-    # Try M1 first
-    client = get_m1_client()
-    if client is not None:
-        try:
-            # Assuming the endpoint is /generate with inputs: prompt, max_tokens, temperature
-            result = client.predict(
-                prompt=f"User: {query}\nAssistant:",
-                max_tokens=64,
-                temperature=0.7,
-                api_name="/generate"
-            )
-            if result and isinstance(result, str) and result.strip():
-                return result.strip()
-        except Exception as e:
-            print(f"M1 casual call failed: {e}")
-
-    # Try M2
-    client = get_m2_client()
-    if client is not None:
-        try:
-            result = client.predict(
-                prompt=f"User: {query}\nAssistant:",
-                max_tokens=64,
-                temperature=0.7,
-                api_name="/generate"
-            )
-            if result and isinstance(result, str) and result.strip():
-                return result.strip()
-        except Exception as e:
-            print(f"M2 casual call failed: {e}")
-
-    # Ultimate fallback
-    return (
-        f"👋 Hi there! I'm X‑RUDRA, your research assistant. "
-        f"How can I help you today? (Your message `{query}` was casual, so I kept it light.)"
-    )
-
-
 # ============================================================
 # MAIN RESEARCH FUNCTION
 # ============================================================
@@ -294,9 +363,7 @@ async def do_research(question, max_results, max_rounds, use_models, freshness):
 
     question = str(question).strip()
 
-    # ------------------------------------------------------------
-    # CASUAL QUERY – get a model‑generated reply (no web search)
-    # ------------------------------------------------------------
+    # Casual query – skip heavy research
     if is_casual_query(question):
         answer = get_casual_model_response(question)
         history = [
@@ -305,11 +372,8 @@ async def do_research(question, max_results, max_rounds, use_models, freshness):
         ]
         return history, "⚡ Casual chat (model reply, no search).", "", "", ""
 
-    # ------------------------------------------------------------
-    # SERIOUS QUERY – run the engine (web + models)
-    # ------------------------------------------------------------
+    # Serious query – web search + two‑stage synthesis
     started = time.perf_counter()
-
     try:
         engine = get_engine()
         report = await engine.search(
@@ -322,12 +386,6 @@ async def do_research(question, max_results, max_rounds, use_models, freshness):
         data = safe_dict(report)
         elapsed_ms = int((time.perf_counter() - started) * 1000)
 
-        # Debug (optional)
-        print("\n" + "="*60)
-        print("RAW ENGINE DATA (first 3000 chars):")
-        print(json.dumps(data, indent=2, default=str)[:3000])
-        print("="*60 + "\n")
-
         # Convert 'results' to 'sources' if needed
         sources = data.get("sources", [])
         if not sources:
@@ -344,29 +402,18 @@ async def do_research(question, max_results, max_rounds, use_models, freshness):
                     })
             data["sources"] = sources
 
-        # Extract or synthesise answer
-        answer = extract_answer(data)
-        if answer is None:
-            if sources:
-                top = sources[:5]
-                parts = [f"Based on the top results:"]
-                for i, src in enumerate(top, 1):
-                    title = src.get("title", "Untitled")
-                    snippet = src.get("snippet", src.get("description", ""))
-                    parts.append(f"{i}. **{title}** – {snippet[:200]}..." if snippet else f"{i}. **{title}**")
-                answer = "\n\n".join(parts)
-            else:
-                answer = "No information found. Try rephrasing your question."
-
-        claims = data.get("claims", [])
+        # Generate final answer using M1 + M2
+        final_answer = get_combined_model_answer(question, sources)
+
+        # Build outputs for tabs
         sources_md = format_sources(sources)
-        evidence_md = format_evidence(claims)
+        evidence_md = format_evidence(data.get("claims", []))
         verification_md = format_verification(data.get("contradictions", []))
         activity_md = build_activity(data, elapsed_ms)
 
         history = [
             {"role": "user", "content": question},
-            {"role": "assistant", "content": answer}
+            {"role": "assistant", "content": final_answer}
         ]
 
         return history, activity_md, sources_md, evidence_md, verification_md
@@ -408,7 +455,7 @@ def health_check():
 
 
 # ============================================================
-# CSS
+# CSS AND UI
 # ============================================================
 
 CSS = """
@@ -423,15 +470,11 @@ footer { display: none !important; }
 """
 
 
-# ============================================================
-# GRADIO UI
-# ============================================================
-
 with gr.Blocks(title=APP_NAME) as demo:
     gr.HTML("""
         <div id="header">
             <div id="logo">⚡ X-RUDRA</div>
-            <div id="tagline">Dual-Model AI · Live Web Research · Evidence</div>
+            <div id="tagline">Dual‑Model AI · Live Web Research · Evidence</div>
         </div>
     """)