Update app.py

app.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 
 import os
 import json
+import re
 import time
 import asyncio
 import traceback
@@ -21,7 +22,7 @@ from gradio_client import Client
 # ============================================================
 
 APP_NAME = "X-RUDRA"
-VERSION = "3.5.0"   # bumped version
+VERSION = "3.6.0"
 
 M1_REPO = os.getenv("M1_REPO", "Shrijanagain/M1")
 M2_REPO = os.getenv("M2_REPO", "Shrijanagain/M2")
@@ -124,11 +125,10 @@ def call_model(client, prompt, max_tokens=512, temperature=0.7):
 
 
 # ============================================================
-# SYNTHESIS: M1 (draft) → M2 (refine)
+# SYNTHESIS: M1 (draft) → M2 (refine) + extract thinking
 # ============================================================
 
 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:
@@ -154,20 +154,18 @@ Answer:"""
 
     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)
+            return "\n\n".join(parts), ""
         else:
-            return "I couldn't find specific information on that topic. Could you rephrase?"
+            return "I couldn't find specific information on that topic. Could you rephrase?", ""
 
     # ----- 2. M2 refines the draft -----
     prompt_m2 = f"""Question: {question}
@@ -180,10 +178,18 @@ Please refine and improve this answer to make it more comprehensive, accurate, a
 Improved answer:"""
 
     refined = call_model(m2_client, prompt_m2, max_tokens, temperature)
-    if refined:
-        return refined
-    else:
-        return draft
+    if not refined:
+        refined = draft
+
+    # Extract thinking from <think> tags (if any)
+    thinking_content = ""
+    clean_answer = refined
+    think_match = re.search(r"<think>(.*?)</think>", refined, re.DOTALL)
+    if think_match:
+        thinking_content = think_match.group(1).strip()
+        clean_answer = re.sub(r"<think>.*?</think>", "", refined, flags=re.DOTALL).strip()
+
+    return clean_answer, thinking_content
 
 
 # ============================================================
@@ -354,12 +360,12 @@ def safe_dict(value):
 
 
 # ============================================================
-# MAIN RESEARCH FUNCTION
+# MAIN RESEARCH FUNCTION (returns thinking as well)
 # ============================================================
 
 async def do_research(question, max_results, max_rounds, use_models, freshness):
     if not question or not str(question).strip():
-        return [], "⚪ Enter a question to start.", "", ""
+        return [], "⚪ Enter a question to start.", "", "", ""
 
     question = str(question).strip()
 
@@ -402,8 +408,8 @@ async def do_research(question, max_results, max_rounds, use_models, freshness):
                     })
             data["sources"] = sources
 
-        # Generate final answer using M1 + M2
-        final_answer = get_combined_model_answer(question, sources)
+        # Generate final answer using M1 + M2, also get thinking
+        final_answer, thinking_content = get_combined_model_answer(question, sources)
 
         # Build outputs for tabs
         sources_md = format_sources(sources)
@@ -411,12 +417,18 @@ async def do_research(question, max_results, max_rounds, use_models, freshness):
         verification_md = format_verification(data.get("contradictions", []))
         activity_md = build_activity(data, elapsed_ms)
 
+        # If there's thinking content, format it nicely
+        if thinking_content:
+            thinking_md = f"### 🧠 Thinking\n\n{thinking_content}"
+        else:
+            thinking_md = ""
+
         history = [
             {"role": "user", "content": question},
             {"role": "assistant", "content": final_answer}
         ]
 
-        return history, activity_md, sources_md, evidence_md, verification_md
+        return history, activity_md, sources_md, evidence_md, verification_md, thinking_md
 
     except Exception as exc:
         error = f"❌ **X-RUDRA Error**\n\n`{type(exc).__name__}: {exc}`"
@@ -428,7 +440,7 @@ async def do_research(question, max_results, max_rounds, use_models, freshness):
             {"role": "user", "content": question},
             {"role": "assistant", "content": error}
         ]
-        return history, "❌ Research failed.", "", "", ""
+        return history, "❌ Research failed.", "", "", "", ""
 
 
 # ============================================================
@@ -455,7 +467,7 @@ def health_check():
 
 
 # ============================================================
-# CSS AND UI
+# CSS – includes spinner animation for thinking
 # ============================================================
 
 CSS = """
@@ -467,10 +479,40 @@ body { background: #f7f7f8; }
 #chat { border-radius: 18px; }
 #send { min-height: 52px; font-size: 18px; font-weight: 700; }
 footer { display: none !important; }
+
+/* Thinking spinner animation */
+@keyframes think-pulse {
+    0% { opacity: 0.3; transform: scale(0.95); }
+    50% { opacity: 1; transform: scale(1.05); }
+    100% { opacity: 0.3; transform: scale(0.95); }
+}
+.thinking-spinner {
+    display: inline-block;
+    width: 12px;
+    height: 12px;
+    border-radius: 50%;
+    background: #6b7280;
+    margin-right: 8px;
+    animation: think-pulse 1.2s ease-in-out infinite;
+}
+.thinking-container {
+    background: #f3f4f6;
+    border-left: 4px solid #6366f1;
+    padding: 12px 16px;
+    border-radius: 8px;
+    margin: 12px 0;
+    font-family: monospace;
+    white-space: pre-wrap;
+    word-wrap: break-word;
+}
 """
 
 
-with gr.Blocks(title=APP_NAME) as demo:
+# ============================================================
+# GRADIO UI – added thinking output
+# ============================================================
+
+with gr.Blocks(title=APP_NAME, css=CSS) as demo:
     gr.HTML("""
         <div id="header">
             <div id="logo">⚡ X-RUDRA</div>
@@ -488,6 +530,8 @@ with gr.Blocks(title=APP_NAME) as demo:
         with gr.Column(scale=4):
             gr.Markdown("## 🔬 Live Research")
             activity = gr.Markdown("⚪ Waiting for your question.")
+            # Thinking output will appear here
+            thinking = gr.Markdown("", elem_id="thinking", visible=False)
             gr.Markdown("---")
             gr.Markdown(f"""
 ### Model Spaces
@@ -528,8 +572,9 @@ with gr.Blocks(title=APP_NAME) as demo:
         inputs=question
     )
 
+    # Update outputs: added thinking
     inputs = [question, max_results, max_rounds, use_models, freshness]
-    outputs = [chatbot, activity, sources, evidence, verification]
+    outputs = [chatbot, activity, sources, evidence, verification, thinking]
 
     send.click(fn=run_research, inputs=inputs, outputs=outputs)
     question.submit(fn=run_research, inputs=inputs, outputs=outputs)
@@ -549,4 +594,4 @@ if __name__ == "__main__":
         print("HF_TOKEN set – rate limits reduced.")
     else:
         print("HF_TOKEN not set – you may experience rate limits. Set it as a Secret in your Space.")
-    demo.launch(server_name="0.0.0.0", server_port=PORT, css=CSS, show_error=True)
+    demo.launch(server_name="0.0.0.0", server_port=PORT, show_error=True)