Update app.py

app.py

@@ -10,51 +10,7 @@
 
 """
 # REASONING-ENFORCED AGENT
-
-## 1. CORE DIRECTIVE
-You are a reasoning-first agent.
-For every meaningful request, you MUST determine:
-1. What the user wants.
-2. The actual goal.
-3. Mandatory requirements.
-4. Constraints.
-5. Missing information.
-6. Ambiguity.
-7. Task decomposition.
-8. Dependencies.
-9. Execution order.
-10. Priorities.
-11. Evidence supporting key decisions.
-12. How the result will be verified.
-
-Do NOT jump from input to output without reasoning.
-
-## 2. MANDATORY REASONING GATE
-INPUT → UNDERSTAND → GOAL → REQUIREMENTS → CONSTRAINTS
-→ MISSING INFO → AMBIGUITY → DECOMPOSITION → DEPENDENCIES
-→ PRIORITY → PLAN → EXECUTE → VERIFY → FINALIZE
-
-## 3. NEVER SKIP CRITICAL ANALYSIS
-Goal identification, requirement extraction, constraint checking,
-missing‑information detection, dependency checking, completion verification.
-
-## 4. RESPONSE MODE
-DIRECT, EXPLANATORY, PLANNED, CLARIFICATION, DIAGNOSTIC, EXECUTION, REVISION.
-
-## 5. ANTI‑FABRICATION
-Never invent missing requirements, tool results, files, test results,
-execution, verification, sources, or completion.
-
-## 6. REASONING VISIBILITY
-The agent MUST reason internally but NEVER expose private chain‑of‑thought.
-Provide concise summaries only when useful.
-
-## 7. FINAL CONTRACT
-UNDERSTAND → PLAN → ACT → VERIFY → ADAPT → DELIVER
-Primary objective:
-"Produce the most correct, goal‑aligned, constraint‑compliant,
-evidence‑supported, verified result using the minimum effective
-reasoning and execution required."
+... (full policy – keep as before)
 """
 
 from __future__ import annotations
@@ -76,7 +32,7 @@ from gradio_client import Client
 # ============================================================
 
 APP_NAME = "X-RUDRA"
-VERSION = "3.7.1"   # bumped
+VERSION = "3.7.3"   # bumped
 
 M1_REPO = os.getenv("M1_REPO", "Shrijanagain/M1")
 M2_REPO = os.getenv("M2_REPO", "Shrijanagain/M2")
@@ -165,6 +121,7 @@ def call_model(client, prompt, max_tokens=512, temperature=0.7):
     if client is None:
         return None
     try:
+        print(f"Calling model with max_tokens={max_tokens}, prompt length={len(prompt)}")
         result = client.predict(
             prompt=prompt,
             max_tokens=max_tokens,
@@ -172,18 +129,21 @@ def call_model(client, prompt, max_tokens=512, temperature=0.7):
             api_name="/generate"
         )
         if result and isinstance(result, str) and result.strip():
+            print(f"Response length: {len(result)} chars")
             return result.strip()
+        else:
+            print("Empty response")
+            return None
     except Exception as e:
         print(f"Model call failed: {e}")
-    return None
+        return None
 
 
 # ============================================================
-# SYNTHESIS: M1 + M2 independent drafts → M2 merges them
+# SYNTHESIS: M1 + M2 drafts → M2 merges (with higher token budget)
 # ============================================================
 
 def get_combined_model_answer(question, sources, max_tokens=512, temperature=0.7):
-    # Build a summary of top sources (max 5)
     top_sources = sources[:5] if sources else []
     sources_text = ""
     if top_sources:
@@ -210,7 +170,7 @@ Answer:"""
     draft_m1 = call_model(m1_client, base_prompt, max_tokens, temperature)
     draft_m2 = call_model(m2_client, base_prompt, max_tokens, temperature)
 
-    # Fallback if one fails
+    # Fallback if both fail
     if not draft_m1 and not draft_m2:
         if sources:
             parts = ["Based on available information:"]
@@ -218,17 +178,30 @@ Answer:"""
                 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?", ""
 
-    # If only one draft exists, use that as final
+    # If only one draft exists, use that as final (extract thinking if any)
     if not draft_m1:
-        return draft_m2
-    if not draft_m2:
-        return draft_m1
+        thinking = ""
+        clean = draft_m2
+        think_match = re.search(r"<think>(.*?)</think>", draft_m2, re.DOTALL)
+        if think_match:
+            thinking = think_match.group(1).strip()
+            clean = re.sub(r"<think>.*?</think>", "", draft_m2, flags=re.DOTALL).strip()
+        return clean, thinking
 
-    # ----- 2. Merge both drafts using M2 -----
+    if not draft_m2:
+        thinking = ""
+        clean = draft_m1
+        think_match = re.search(r"<think>(.*?)</think>", draft_m1, re.DOTALL)
+        if think_match:
+            thinking = think_match.group(1).strip()
+            clean = re.sub(r"<think>.*?</think>", "", draft_m1, flags=re.DOTALL).strip()
+        return clean, thinking
+
+    # ----- 2. Merge both drafts using M2 with a larger token budget -----
     merge_prompt = f"""Question: {question}
 
 Draft from Model A:
@@ -241,14 +214,29 @@ Combine these two drafts into a single, comprehensive, accurate, and natural ans
 
 Final answer:"""
 
-    merged = call_model(m2_client, merge_prompt, max_tokens, temperature)
+    # Use a higher token limit for the merge step
+    merge_max_tokens = max(1024, max_tokens * 2)  # at least 1024, double the original
+    merged = call_model(m2_client, merge_prompt, merge_max_tokens, temperature)
+
+    # If merged is too short (incomplete), retry with even more tokens
+    if merged and len(merged) < 100:
+        print(f"Merged answer too short ({len(merged)} chars), retrying with 2048 tokens...")
+        merged = call_model(m2_client, merge_prompt, 2048, temperature)
+
     if not merged:
-        # If merging fails, fallback to draft_m1
+        # Fallback to draft_m1 if merge fails completely
+        print("Merge failed, falling back to draft_m1")
         merged = draft_m1
 
-    # Strip any remaining <think> tags from the final answer
-    clean_answer = re.sub(r"<think>.*?</think>", "", merged, flags=re.DOTALL).strip()
-    return clean_answer
+    # Extract thinking and clean tags
+    thinking_content = ""
+    clean_answer = merged
+    think_match = re.search(r"<think>(.*?)</think>", merged, re.DOTALL)
+    if think_match:
+        thinking_content = think_match.group(1).strip()
+        clean_answer = re.sub(r"<think>.*?</think>", "", merged, flags=re.DOTALL).strip()
+
+    return clean_answer, thinking_content
 
 
 # ============================================================
@@ -266,7 +254,8 @@ def get_casual_model_response(query: str) -> str:
                 api_name="/generate"
             )
             if result and isinstance(result, str) and result.strip():
-                return result.strip()
+                clean = re.sub(r"<think>.*?</think>", "", result, flags=re.DOTALL).strip()
+                return clean
         except Exception as e:
             print(f"M1 casual failed: {e}")
 
@@ -280,7 +269,8 @@ def get_casual_model_response(query: str) -> str:
                 api_name="/generate"
             )
             if result and isinstance(result, str) and result.strip():
-                return result.strip()
+                clean = re.sub(r"<think>.*?</think>", "", result, flags=re.DOTALL).strip()
+                return clean
         except Exception as e:
             print(f"M2 casual failed: {e}")
 
@@ -419,7 +409,7 @@ def safe_dict(value):
 
 
 # ============================================================
-# MAIN RESEARCH FUNCTION (returns 5 outputs – no thinking)
+# MAIN RESEARCH FUNCTION (returns 6 outputs)
 # ============================================================
 
 async def do_research(question, max_results, max_rounds, use_models, freshness):
@@ -428,9 +418,10 @@ async def do_research(question, max_results, max_rounds, use_models, freshness):
     empty_sources = ""
     empty_evidence = ""
     empty_verification = ""
+    empty_thinking = ""
 
     if not question or not str(question).strip():
-        return empty_history, empty_activity, empty_sources, empty_evidence, empty_verification
+        return empty_history, empty_activity, empty_sources, empty_evidence, empty_verification, empty_thinking
 
     question = str(question).strip()
 
@@ -441,7 +432,7 @@ async def do_research(question, max_results, max_rounds, use_models, freshness):
             {"role": "user", "content": question},
             {"role": "assistant", "content": answer}
         ]
-        return history, "⚡ Casual chat (model reply, no search).", "", "", ""
+        return history, "⚡ Casual chat (model reply, no search).", "", "", "", ""
 
     # ---- Serious query ----
     started = time.perf_counter()
@@ -473,20 +464,22 @@ async def do_research(question, max_results, max_rounds, use_models, freshness):
                     })
             data["sources"] = sources
 
-        # Generate final answer (no thinking)
-        final_answer = get_combined_model_answer(question, sources)
+        # Generate final answer and thinking content (tags already stripped)
+        final_answer, thinking_content = get_combined_model_answer(question, sources)
 
         sources_md = format_sources(sources)
         evidence_md = format_evidence(data.get("claims", []))
         verification_md = format_verification(data.get("contradictions", []))
         activity_md = build_activity(data, elapsed_ms)
 
+        thinking_md = f"### 🧠 Reasoning\n\n{thinking_content}" if thinking_content else ""
+
         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}`"
@@ -498,7 +491,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.", "", "", "", ""
 
 
 # ============================================================
@@ -525,7 +518,7 @@ def health_check():
 
 
 # ============================================================
-# CSS – no thinking spinner needed
+# CSS – includes spinner animation for thinking
 # ============================================================
 
 CSS = """
@@ -537,11 +530,36 @@ body { background: #f7f7f8; }
 #chat { border-radius: 18px; }
 #send { min-height: 52px; font-size: 18px; font-weight: 700; }
 footer { display: none !important; }
+
+@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;
+}
 """
 
 
 # ============================================================
-# GRADIO UI – 5 outputs (no thinking)
+# GRADIO UI – 6 outputs
 # ============================================================
 
 with gr.Blocks(title=APP_NAME) as demo:
@@ -562,6 +580,7 @@ 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 = gr.Markdown("", visible=True)
             gr.Markdown("---")
             gr.Markdown(f"""
 ### Model Spaces
@@ -603,7 +622,7 @@ with gr.Blocks(title=APP_NAME) as demo:
     )
 
     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)