app.py

"""
Vitalis LOREIN MCP Server — Interactive Demo UI.
Pick any HuggingFace model, run it through the LOREIN cognitive pipeline.
Shows the difference: model alone vs model + MIRROR + Truth + Governor.
"""

import gradio as gr
import json
import os
import time

try:
    from huggingface_hub import InferenceClient
    HAS_HF = True
except ImportError:
    HAS_HF = False

os.environ["HF_HUB_DISABLE_SYMLINKS"] = "1"

from pipeline.lorein import LOREIN
from router import route as _route_task, classify_task, TASK_ROUTES
lorein = LOREIN()
client = None

DEFAULT_MODEL = "microsoft/Phi-3-mini-4k-instruct"
TEST_MODELS = [
    "microsoft/Phi-3-mini-4k-instruct",
    "HuggingFaceH4/zephyr-7b-beta",
    "mistralai/Mistral-7B-Instruct-v0.3",
    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    "google/gemma-2b-it",
]

def set_model(model_id):
    global client
    if not HAS_HF:
        return "⚠️ huggingface_hub not installed on this server."
    if not model_id or not model_id.strip():
        return "Please enter a model ID."
    try:
        token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN")
        client = InferenceClient(model=model_id.strip(), token=token)
        return f"✅ Connected to **{model_id.strip()}**"
    except Exception as e:
        return f"❌ Failed: {e}"

def run_model_direct(message, history):
    if not client:
        return None, "No model selected. Pick a model first."
    try:
        messages = [{"role": "user", "content": message}]
        if history:
            for h in history:
                messages.insert(0, {"role": "user", "content": h[0]})
                messages.insert(0, {"role": "assistant", "content": h[1]})
        response = client.chat_completion(messages=messages, max_tokens=256, temperature=0.7)
        return response.choices[0].message.content, None
    except Exception as e:
        return None, f"Error: {e}"

def run_lorein_pipeline(message, history):
    if not client:
        return None, None, "No model selected."

    try:
        # Step 1: MIRROR thinking
        narrative, threads = lorein.mirror.think(message)
        mirror_out = f"**Goals Check:** {'✅ Safe' if threads[0]['safe'] else '⚠️ Issues'}\n"
        mirror_out += f"**Reasoning:** {'Ambiguous query' if threads[1]['ambiguity'] else str(len(threads[1].get('hypotheses',[]))) + ' hypotheses'}\n"
        mirror_out += f"**Memory:** {threads[2].get('episodic_count',0)} episodes, {len(threads[2].get('semantic_retrievals',[]))} semantic hits\n"
        mirror_out += f"**Narrative:** {narrative[:200]}..."

        # Step 2: Generate with MIRROR context
        enriched = f"[Context]\n{narrative}\n\n[Query]\n{message}"
        messages = [{"role": "user", "content": enriched}]
        response = client.chat_completion(messages=messages, max_tokens=256, temperature=0.7)
        model_response = response.choices[0].message.content

        # Step 3: Truth check
        truth = lorein.truth_check(model_response)
        truth_out = f"**Truthful:** {truth['truthful']} | **Confidence:** {truth['confidence']}\n"
        if truth.get('issues'):
            truth_out += f"**Issues:** {', '.join(truth['issues'])}"
        else:
            truth_out += "**Issues:** None detected"

        # Step 4: Governor evaluation
        gov = lorein.governor.evaluate(model_response)
        gov_out = f"**Allowed:** {'✅ PASS' if gov['allowed'] else '❌ BLOCKED'}\n"
        if gov.get('violations'):
            gov_out += f"**Violations:** {', '.join(gov['violations'])}"
        else:
            gov_out += "**Violations:** None"

        # Step 5: Journal log
        journal_size = lorein.journal.size()
        chain_ok = lorein.journal.verify_chain()
        state = lorein.inference.state()

        pipeline_out = (
            f"### 🧠 MIRROR Thinking\n{mirror_out}\n\n"
            f"### 📝 Generated Response\n{model_response[:300]}...\n\n"
            f"### ✅ Truth Check\n{truth_out}\n\n"
            f"### 🔒 Governor\n{gov_out}\n\n"
            f"### 📊 System\n- **Journal entries:** {journal_size}\n"
            f"- **Chain integrity:** {'✅ Intact' if chain_ok else '❌ Broken'}\n"
            f"- **Confidence:** {state['confidence']}\n"
            f"- **Explore urgency:** {state['exploration_urgency']}"
        )
        return model_response, pipeline_out, None
    except Exception as e:
        return None, None, f"Error: {e}"

with gr.Blocks(
    title="LOREIN — Test Any HF Model",
    theme=gr.themes.Soft(primary_hue="purple", secondary_hue="violet"),
    css="""
    .header { text-align: center; padding: 24px; background: linear-gradient(135deg, #1a0a3e, #4c1d95); border-radius: 12px; margin-bottom: 16px; }
    .header h1 { color: #a78bfa; margin: 0; }
    .header p { color: #c4b5fd; margin: 4px 0 0; font-size: 0.9em; }
    .pipeline-box { background: #1e1e2e; border-left: 4px solid #7c3aed; border-radius: 0 8px 8px 0; padding: 12px; margin: 8px 0; }
    .step-label { font-weight: bold; color: #a78bfa; }
    """
) as demo:
    gr.HTML("""
    <div class="header">
        <h1>LOREIN — Cognitive Pipeline Demo</h1>
        <p>Pick any HuggingFace model. See the difference: raw vs LOREIN-enhanced.</p>
    </div>
    """)

    gr.Markdown("## 1. Pick a Model")
    with gr.Row():
        model_input = gr.Dropdown(
            choices=TEST_MODELS,
            label="Model ID",
            value=DEFAULT_MODEL,
            allow_custom_value=True,
            info="Type any HuggingFace model ID (e.g., microsoft/Phi-3-mini-4k-instruct)",
            scale=3,
        )
        connect_btn = gr.Button("Connect Model", variant="primary", scale=1, size="lg")

    model_status = gr.Markdown("Enter a model ID and click Connect.")

    connect_btn.click(set_model, inputs=[model_input], outputs=model_status)

    gr.Markdown("## 2. Test It")

    with gr.Tabs():
        with gr.TabItem("🔄 LOREIN Pipeline (Plan → Generate → Verify)"):
            lorein_msg = gr.Textbox(
                label="Ask something",
                placeholder="Type a question to see the full LOREIN pipeline in action...",
                lines=2,
            )
            with gr.Row():
                lorein_submit = gr.Button("Run Pipeline", variant="primary", size="lg", scale=1)
                clear_btn = gr.Button("Clear", size="sm", scale=1)

            with gr.Row():
                with gr.Column(scale=2):
                    lorein_response = gr.Textbox(label="Final Response", lines=6, max_lines=12)
                with gr.Column(scale=2):
                    lorein_pipeline = gr.Markdown("Pipeline output will appear here...", label="Pipeline Steps")

            lorein_msg.submit(run_lorein_pipeline, inputs=[lorein_msg, gr.State([])], outputs=[lorein_response, lorein_pipeline, gr.Markdown(visible=False)])
            lorein_submit.click(run_lorein_pipeline, inputs=[lorein_msg, gr.State([])], outputs=[lorein_response, lorein_pipeline, gr.Markdown(visible=False)])

        with gr.TabItem("🧭 Model Router"):
            gr.Markdown("### Route any task to the best model")
            gr.Markdown("Describe what you need — the router classifies your task and recommends the optimal model.")
            router_msg = gr.Textbox(
                label="Describe your task",
                placeholder="e.g., Write a Python function to sort a list, or Explain quantum computing...",
                lines=2,
            )
            router_out = gr.Markdown("Enter a task description to see routing.")

            def run_router(msg):
                if not msg:
                    return "Enter a task description."
                routing = _route_task(msg)
                cats = {k: v for k, v in TASK_ROUTES.items()}
                out = f"**Classified as:** `{routing['task']}`\n\n"
                out += f"**Recommended model:** `{routing['model']}`\n\n"
                out += f"**Task description:** {routing['task_description']}\n\n"
                out += "**Available task types:**\n"
                for task, info in cats.items():
                    out += f"- `{task}`: {info['description']}\n"
                return out

            router_msg.change(run_router, inputs=[router_msg], outputs=router_out)

    gr.Markdown("---")
    gr.Markdown("""
    ### How the Pipeline Works

    1. **MIRROR Thinking** — Decomposes your query into Goals/Reasoning/Memory threads
    2. **Model Generates** — The model answers with the MIRROR context as background
    3. **Truth Check** — Scans for hedging, contradictions, uncertainty
    4. **Governor** — Safety screen (mutual advancement principle)
    5. **Memory + Journal** — Logs everything to the hash-chained ledger
    """)

if __name__ == "__main__":
    demo.launch()