app.py

# ============================================================
#                    X-RUDRA CHAT
#          Dual‑Model + Web Research · Gradio Space
# ============================================================

# ──────────────────────────────────────────────────────────────
#  REASONING‑ENFORCED AGENT CONTRACT (AGENT.md)
#  See full policy in the multi‑line comment below.
# ──────────────────────────────────────────────────────────────

"""
# REASONING-ENFORCED AGENT
... (full policy – keep as before)
"""

from __future__ import annotations

import os
import json
import re
import time
import asyncio
import traceback

import gradio as gr
import spaces
from gradio_client import Client


# ============================================================
# CONFIG
# ============================================================

APP_NAME = "X-RUDRA"
VERSION = "3.8.3"   # bumped

M1_REPO = os.getenv("M1_REPO", "Shrijanagain/M1")
M2_REPO = os.getenv("M2_REPO", "Shrijanagain/M2")
PORT = int(os.getenv("PORT", "7860"))

HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN:
    os.environ["HF_TOKEN"] = HF_TOKEN


# ============================================================
# GRADIO CLIENTS FOR M1 / M2
# ============================================================

_M1_CLIENT = None
_M2_CLIENT = None

def get_m1_client():
    global _M1_CLIENT
    if _M1_CLIENT is None:
        try:
            url = f"https://{M1_REPO.replace('/', '-')}.hf.space"
            _M1_CLIENT = Client(url)
        except Exception as e:
            print(f"Could not connect to M1: {e}")
            _M1_CLIENT = None
    return _M1_CLIENT

def get_m2_client():
    global _M2_CLIENT
    if _M2_CLIENT is None:
        try:
            url = f"https://{M2_REPO.replace('/', '-')}.hf.space"
            _M2_CLIENT = Client(url)
        except Exception as e:
            print(f"Could not connect to M2: {e}")
            _M2_CLIENT = None
    return _M2_CLIENT


# ============================================================
# MODEL-BASED INTENT CLASSIFIER
# ============================================================

CLASSIFIER_SYSTEM_PROMPT = """
You are an intelligent assistant that classifies user messages into two categories:
- ACTION: The user asks for information that requires research, fact‑checking, retrieval of current data, or external knowledge. This includes questions about news, comparisons, statistics, history, technology, science, politics, etc.
- CASUAL: The user is just chatting, greeting, making small talk, joking, or asking a simple question that can be answered from general knowledge without a search.

Respond with ONLY ONE WORD: ACTION or CASUAL.
Do NOT add any extra text, punctuation, or explanation.
"""

def classify_intent(question: str) -> str:
    prompt = f"{CLASSIFIER_SYSTEM_PROMPT}\n\nUser message: \"{question}\"\n\nClassification:"
    
    # Try M1
    m1_client = get_m1_client()
    if m1_client is not None:
        try:
            result = m1_client.predict(
                prompt=prompt,
                max_tokens=64,
                temperature=0.1,          # FIXED: was 0.0 → now 0.1
                api_name="/generate"
            )
            if result:
                result = result.strip().upper()
                if "ACTION" in result:
                    print(f"[Classifier] M1 → ACTION")
                    return "ACTION"
                elif "CASUAL" in result:
                    print(f"[Classifier] M1 → CASUAL")
                    return "CASUAL"
                else:
                    print(f"[Classifier] M1 ambiguous: {result}")
        except Exception as e:
            print(f"[Classifier] M1 call failed: {e}")

    # Try M2
    m2_client = get_m2_client()
    if m2_client is not None:
        try:
            result = m2_client.predict(
                prompt=prompt,
                max_tokens=64,
                temperature=0.1,          # FIXED
                api_name="/generate"
            )
            if result:
                result = result.strip().upper()
                if "ACTION" in result:
                    print(f"[Classifier] M2 → ACTION")
                    return "ACTION"
                elif "CASUAL" in result:
                    print(f"[Classifier] M2 → CASUAL")
                    return "CASUAL"
                else:
                    print(f"[Classifier] M2 ambiguous: {result}")
        except Exception as e:
            print(f"[Classifier] M2 call failed: {e}")

    # Fallback: if both fail, use simple heuristic (no keyword list)
    # For very short messages (<=3 words), assume CASUAL to avoid unnecessary search.
    word_count = len(question.split())
    if word_count <= 3:
        print("[Classifier] Fallback → CASUAL (short message)")
        return "CASUAL"
    else:
        print("[Classifier] Fallback → ACTION (longer message)")
        return "ACTION"


# ============================================================
# LAZY ENGINE (web search)
# ============================================================

_ENGINE = None

def get_engine():
    global _ENGINE
    if _ENGINE is None:
        from web_search import XrudraWebSearch
        _ENGINE = XrudraWebSearch()
    return _ENGINE


# ============================================================
# HELPERS – CALL MODELS
# ============================================================

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,
            temperature=temperature,
            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


# ============================================================
# SYNTHESIS: M1 + M2 drafts → M2 merges (with higher token budget)
# ============================================================

def get_combined_model_answer(question, sources, max_tokens=512, temperature=0.7):
    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."

    base_prompt = 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()

    # ----- 1. Get independent drafts from M1 and M2 -----
    draft_m1 = call_model(m1_client, base_prompt, max_tokens, temperature)
    draft_m2 = call_model(m2_client, base_prompt, max_tokens, temperature)

    # Fallback if both fail
    if not draft_m1 and not draft_m2:
        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?", ""

    # If only one draft exists, use that as final
    if not 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

    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:
{draft_m1}

Draft from Model B:
{draft_m2}

Combine these two drafts into a single, comprehensive, accurate, and natural answer. Keep the best parts from each. Ensure the final answer directly addresses the question, is well‑structured, and reads as a single coherent response. Do NOT mention that you are combining drafts or that you used multiple models. Just provide the final answer.

Final answer:"""

    merge_max_tokens = max(1024, max_tokens * 2)
    merged = call_model(m2_client, merge_prompt, merge_max_tokens, temperature)

    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:
        print("Merge failed, falling back to draft_m1")
        merged = draft_m1

    # 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


# ============================================================
# 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=128,
                temperature=0.7,
                api_name="/generate"
            )
            if result and isinstance(result, str) and 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}")

    client = get_m2_client()
    if client is not None:
        try:
            result = client.predict(
                prompt=f"User: {query}\nAssistant:",
                max_tokens=128,
                temperature=0.7,
                api_name="/generate"
            )
            if result and isinstance(result, str) and 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}")

    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)
# ============================================================

def format_sources(sources):
    if not sources:
        return "## 📚 Sources\n\nNo sources were returned."
    output = ["## 📚 Sources", ""]
    for idx, src in enumerate(sources, 1):
        if not isinstance(src, dict):
            continue
        title = src.get("title", "Untitled")
        url = src.get("url", "")
        method = src.get("fetch_method", "web")
        score = src.get("source_score", src.get("score", "N/A"))
        snippet = src.get("snippet", src.get("description", ""))
        if url:
            output.append(f"### {idx}. [{title}]({url})")
        else:
            output.append(f"### {idx}. {title}")
        output.append(f"**Fetcher:** `{method}`")
        output.append(f"**Source score:** `{score}`")
        if snippet:
            output.append(f"\n> {snippet}")
        output.append("")
    return "\n".join(output)


def format_evidence(claims):
    if not claims:
        return "## 🧠 Evidence\n\nNo structured evidence was returned."
    output = ["## 🧠 Evidence", ""]
    for idx, claim in enumerate(claims, 1):
        if not isinstance(claim, dict):
            continue
        text = claim.get("claim", claim.get("text", ""))
        score = claim.get("support_score", claim.get("score", "N/A"))
        source = claim.get("source_url", claim.get("url", ""))
        output.append(f"### Evidence {idx}")
        output.append(str(text))
        output.append(f"**Support:** `{score}`")
        if source:
            output.append(f"**Source:** {source}")
        output.append("---")
    return "\n\n".join(output)


def format_verification(contradictions):
    if not contradictions:
        return "## ⚖️ Verification\n\n✅ No major contradictions detected."
    output = ["## ⚖️ Verification", "", "⚠️ Potential contradictions detected:", ""]
    for idx, item in enumerate(contradictions, 1):
        if not isinstance(item, dict):
            continue
        claim_a = item.get("claim_a", "")
        claim_b = item.get("claim_b", "")
        source_a = item.get("source_a", "")
        source_b = item.get("source_b", "")
        output.append(f"### Contradiction {idx}")
        output.append(f"**A:** {claim_a}")
        if source_a:
            output.append(f"Source A: `{source_a}`")
        output.append("")
        output.append(f"**B:** {claim_b}")
        if source_b:
            output.append(f"Source B: `{source_b}`")
        output.append("---")
    return "\n\n".join(output)


def build_activity(data, elapsed_ms):
    sources = data.get("sources", []) or data.get("results", [])
    claims = data.get("claims", [])
    contradictions = data.get("contradictions", [])
    rounds = data.get("rounds", data.get("research_rounds", "N/A"))
    return f"""
## ⚡ X-RUDRA Research

| Stage | Status |
|---|---|
| Task analysis | ✅ Complete |
| M1 research | ✅ Draft generated |
| M2 research | ✅ Draft generated + merged |
| Web discovery | ✅ Complete |
| Evidence extraction | {"✅" if claims else "⚙️"} |
| Source verification | ✅ Complete |
| Contradiction check | {"⚠️ Found" if contradictions else "✅ Clear"} |
| Final synthesis | ✅ Complete |

**Sources:** `{len(sources)}`  
**Claims:** `{len(claims)}`  
**Rounds:** `{rounds}`  
**Time:** `{elapsed_ms} ms`

### Engine

`M1` → `{M1_REPO}`

`M2` → `{M2_REPO}`

`Web` → `DuckDuckGo`

`Fetcher` → `Scrapling`

`Browser` → `Playwright`
"""


# ============================================================
# SAFE DICT HELPER
# ============================================================

def safe_dict(value):
    if isinstance(value, dict):
        return value
    if hasattr(value, "model_dump"):
        try:
            return value.model_dump()
        except Exception:
            pass
    if hasattr(value, "dict"):
        try:
            return value.dict()
        except Exception:
            pass
    return {"result": str(value)}


# ============================================================
# MAIN RESEARCH FUNCTION
# ============================================================

async def do_research(question, max_results, max_rounds, use_models, freshness):
    empty_history = []
    empty_activity = "⚪ Enter a question to start."
    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, empty_thinking

    question = str(question).strip()

    # ---- Step 1: Classify intent using M1 (or M2) ----
    intent = classify_intent(question)
    print(f"[Intent] {intent} for: {question}")

    # ---- Step 2: If casual, reply directly ----
    if intent == "CASUAL":
        answer = get_casual_model_response(question)
        history = [
            {"role": "user", "content": question},
            {"role": "assistant", "content": answer}
        ]
        return history, "⚡ Casual chat (model reply, no search).", "", "", "", ""

    # ---- Step 3: ACTION – run research pipeline ----
    started = time.perf_counter()
    try:
        engine = get_engine()
        report = await engine.search(
            question=question,
            max_results=int(max_results),
            max_rounds=int(max_rounds),
            use_models=bool(use_models),
            freshness_mode=str(freshness),
        )
        data = safe_dict(report)
        elapsed_ms = int((time.perf_counter() - started) * 1000)

        # Convert 'results' to 'sources' if needed
        sources = data.get("sources", [])
        if not sources:
            results = data.get("results", [])
            for res in results:
                if isinstance(res, dict):
                    sources.append({
                        "title": res.get("title", ""),
                        "url": res.get("url", ""),
                        "snippet": res.get("snippet", ""),
                        "fetch_method": "web",
                        "source_score": res.get("rank", "N/A"),
                        "description": res.get("snippet", ""),
                    })
            data["sources"] = sources

        # Generate final answer and thinking
        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, thinking_md

    except Exception as exc:
        error = f"❌ **X-RUDRA Error**\n\n`{type(exc).__name__}: {exc}`"
        print("\n" + "="*70)
        print("X-RUDRA ERROR")
        print(traceback.format_exc())
        print("="*70 + "\n")
        history = [
            {"role": "user", "content": question},
            {"role": "assistant", "content": error}
        ]
        return history, "❌ Research failed.", "", "", "", ""


# ============================================================
# GRADIO SYNC WRAPPER WITH @spaces.GPU
# ============================================================

@spaces.GPU
def run_research(question, max_results, max_rounds, use_models, freshness):
    return asyncio.run(do_research(question, max_results, max_rounds, use_models, freshness))


# ============================================================
# HEALTH CHECK
# ============================================================

def health_check():
    return f"""
## 🟢 X-RUDRA Online
**Version:** `{VERSION}`
**M1:** `{M1_REPO}`
**M2:** `{M2_REPO}`
**Engine:** Lazy initialized
"""


# ============================================================
# CSS
# ============================================================

CSS = """
body { background: #f7f7f8; }
.gradio-container { max-width: 1500px !important; }
#header { text-align: center; padding: 20px 0 10px 0; }
#logo { font-size: 38px; font-weight: 800; }
#tagline { opacity: 0.65; font-size: 15px; }
#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 – 6 outputs
# ============================================================

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>
    """)

    with gr.Row():
        with gr.Column(scale=7):
            chatbot = gr.Chatbot(label="X-RUDRA", height=600, elem_id="chat")
            with gr.Row():
                question = gr.Textbox(placeholder="Ask X-RUDRA anything...", lines=2, show_label=False, scale=8)
                send = gr.Button("➤", variant="primary", elem_id="send", scale=1)

        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
**M1** `{M1_REPO}`
**M2** `{M2_REPO}`
### Web Stack
`DuckDuckGo` · `Scrapling` · `Playwright`
""")

    with gr.Accordion("⚙️ Research Controls", open=False):
        with gr.Row():
            max_results = gr.Slider(1, 30, value=10, step=1, label="Max Sources")
            max_rounds = gr.Slider(1, 5, value=3, step=1, label="Research Rounds")
        with gr.Row():
            use_models = gr.Checkbox(value=True, label="Use M1 + M2")
            freshness = gr.Dropdown(["auto","latest","recent","current"], value="auto", label="Freshness")

    with gr.Tabs():
        with gr.Tab("📚 Sources"):
            sources = gr.Markdown("Sources will appear here.")
        with gr.Tab("🧠 Evidence"):
            evidence = gr.Markdown("Evidence will appear here.")
        with gr.Tab("⚖️ Verification"):
            verification = gr.Markdown("Verification will appear here.")

    with gr.Accordion("🩺 System Health", open=False):
        health_button = gr.Button("Check X-RUDRA")
        health_output = gr.Markdown()

    gr.Markdown("### Try X-RUDRA")
    gr.Examples(
        examples=[
            ["What are the latest UNESCO AI education initiatives?"],
            ["What are the latest developments in open source AI?"],
            ["Compare the latest major AI models."],
            ["Research India's current AI ecosystem."]
        ],
        inputs=question
    )

    inputs = [question, max_results, max_rounds, use_models, freshness]
    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)
    health_button.click(fn=health_check, inputs=[], outputs=[health_output])


# ============================================================
# START
# ============================================================

if __name__ == "__main__":
    print(f"Starting {APP_NAME} {VERSION}")
    print("M1:", M1_REPO)
    print("M2:", M2_REPO)
    print("Lazy engine initialization: ON")
    if HF_TOKEN:
        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)