Update app.py

app.py

@@ -1,71 +1,74 @@
 import json
 from functools import partial
-
 import gradio as gr
-from transformers import pipeline
-
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 
-# -----------------------------
-# Load module repository
-# -----------------------------
+# -------------------------------------------------------------
+# Load Modules
+# -------------------------------------------------------------
 with open("modules.json", "r", encoding="utf-8") as f:
     MODULES = json.load(f)["modules"]
 
-GENERATORS = [m for m in MODULES if m["type"] == "generator"]
-CHECKERS = {m["id"]: m for m in MODULES if m["type"] == "checker"}
+GENERATORS = [m for m in MODULES if m.get("type") == "generator"]
+CHECKERS = {m["id"]: m for m in MODULES if m.get("type") == "checker"}
 GEN_BY_ID = {m["id"]: m for m in GENERATORS}
-
-# Friendly names for classifier
-MODULE_LABELS = {
-    "Analysis Note": "analysis_note_v1",
-    "Document Explainer": "document_explainer_v1",
-    "Strategy Memo": "strategy_memo_v1",
-    "Message / Post Reply": "message_reply_v1",
-    "Profile / Application": "profile_application_v1",
-    "System / Architecture Blueprint": "system_blueprint_v1",
-    "Modular Brainstorm": "modular_brainstorm_v1",
-}
-
-LABEL_LIST = list(MODULE_LABELS.keys())
-
-
-# -----------------------------
-# Models
-# -----------------------------
-# 1) Generator engine (same as before)
-llm = pipeline("text-generation", model="gpt2", max_new_tokens=512)
-
-# 2) Task classifier (zero-shot)
-#    You can swap this to another classification model later.
-classifier = pipeline(
-    "zero-shot-classification",
-    model="facebook/bart-large-mnli"
-)
-
-
-# -----------------------------
-# Low-level helpers
-# -----------------------------
+LABEL_TO_ID = {m["label"]: m["id"] for m in GENERATORS}
+LABEL_LIST = list(LABEL_TO_ID.keys())
+
+# -------------------------------------------------------------
+# Load Model (base LLM)
+# Swappable engine: GPT-2 / Llama / Mistral etc.
+# -------------------------------------------------------------
+tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
+model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")
+llm = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=300)
+
+# -------------------------------------------------------------
+# Automatic Router Components
+# -------------------------------------------------------------
+from router.rules import rule_router
+from router.zero_shot import classify_task  # zero-shot classifier
+
+# -------------------------------------------------------------
+# Domain Adapters (LoRA)
+# -------------------------------------------------------------
+from domain_heads.loader import load_adapter  # load domain-specific adapter
+
+# -------------------------------------------------------------
+# Reasoning Scaffolds
+# -------------------------------------------------------------
+from reasoning_scaffolds.cot import apply_cot
+from reasoning_scaffolds.critique_loop import critique_and_refine
+
+
+# -------------------------------------------------------------
+# Helper: LLM call
+# -------------------------------------------------------------
 def call_llm(prompt: str) -> str:
-    """Call the generator LLM and strip prompt prefix if echoed."""
-    out = llm(prompt, max_new_tokens=512, do_sample=False)[0]["generated_text"]
-    return out[len(prompt):].strip() if out.startswith(prompt) else out
+    out = llm(prompt, do_sample=False)[0]["generated_text"]
+    if out.startswith(prompt):
+        out = out[len(prompt):]
+    return out.strip()
 
 
-def generator_prompt(module_id: str, *inputs: str) -> str:
-    """Build a structured prompt for a generator module."""
+# -------------------------------------------------------------
+# Build prompts
+# -------------------------------------------------------------
+def build_generator_prompt(module_id: str, *inputs: str) -> str:
     m = GEN_BY_ID[module_id]
     keys = list(m["input_placeholders"].keys())
     vals = {k: inputs[i] if i < len(inputs) else "" for i, k in enumerate(keys)}
     secs = m["output_sections"]
 
     p = []
-    p.append(f"MODULE: {m['label']}")
-    p.append("INPUT:")
+    p.append("You are a structured reasoning module.")
+    p.append(f"MODULE: {m['label']} (id={module_id})")
+    p.append("")
+    p.append("INPUTS:")
     for k, v in vals.items():
         p.append(f"{k.upper()}: {v}")
     p.append("")
-    p.append("OUTPUT SECTIONS:")
+    p.append("You must respond using these sections:")
     for s in secs:
         p.append(f"- {s}")
     p.append("")
@@ -76,27 +79,28 @@ def generator_prompt(module_id: str, *inputs: str) -> str:
     return "\n".join(p)
 
 
-def checker_prompt(checker_id: str, *vals: str) -> str:
-    """Build a prompt for a checker module."""
+def build_checker_prompt(checker_id: str, *vals: str) -> str:
     c = CHECKERS[checker_id]
     secs = c["output_sections"]
 
     if len(vals) < 2:
-        orig = ""
+        original_task = ""
         draft = vals[0] if vals else ""
     else:
-        orig = "\n\n".join(vals[:-1])
+        original_task = "\n\n".join(vals[:-1])
         draft = vals[-1]
 
     p = []
-    p.append(f"CHECKER: {c['label']}")
+    p.append("You are a strict reviewer.")
+    p.append(f"CHECKER: {c['label']} (id={checker_id})")
+    p.append("")
     p.append("ORIGINAL TASK:")
-    p.append(orig)
+    p.append(original_task)
     p.append("")
-    p.append("DRAFT:")
+    p.append("DRAFT OUTPUT:")
     p.append(draft)
     p.append("")
-    p.append("RESPOND WITH:")
+    p.append("You must respond using:")
     for s in secs:
         p.append(f"- {s}")
     p.append("")
@@ -107,120 +111,99 @@ def checker_prompt(checker_id: str, *vals: str) -> str:
     return "\n".join(p)
 
 
-def run_generator(mid: str, *inputs: str) -> str:
-    return call_llm(generator_prompt(mid, *inputs))
+# -------------------------------------------------------------
+# Generator & Checker Execution
+# -------------------------------------------------------------
+def run_generator(module_id: str, *inputs: str) -> str:
+    m = GEN_BY_ID[module_id]
+
+    # Load domain adapter if defined
+    if m.get("domain"):
+        load_adapter(model, m["domain"])
 
+    # Build prompt
+    prompt = build_generator_prompt(module_id, *inputs)
 
-def run_checker(cid: str, *inputs: str) -> str:
-    return call_llm(checker_prompt(cid, *inputs))
+    # Apply reasoning scaffolds (CoT + critique loop)
+    prompt = apply_cot(prompt)
+    draft = call_llm(prompt)
+    final = critique_and_refine(draft)
 
+    return final
 
-# -----------------------------
-# Task classifier
-# -----------------------------
-def classify_task(task_text: str):
-    """
-    Classify a free-form task description into one of the Modular Intelligence modules.
 
-    Returns:
-        predicted_label: human-friendly label (e.g. "Strategy Memo")
-        module_id: internal id (e.g. "strategy_memo_v1")
-        scores_text: formatted confidence scores per label
-    """
-    if not task_text.strip():
-        return "No input", "", ""
+def run_checker(checker_id: str, *inputs: str) -> str:
+    prompt = build_checker_prompt(checker_id, *inputs)
+    prompt = apply_cot(prompt)
+    reviewed = call_llm(prompt)
+    return reviewed
 
-    res = classifier(
-        task_text,
-        candidate_labels=LABEL_LIST,
-        multi_label=False
-    )
-    # Zero-shot pipeline returns labels sorted by score descending
-    predicted_label = res["labels"][0]
-    module_id = MODULE_LABELS[predicted_label]
 
-    # Build a small score table
-    lines = []
-    for label, score in zip(res["labels"], res["scores"]):
-        lines.append(f"{label}: {score:.2f}")
-    scores_text = "\n".join(lines)
+# -------------------------------------------------------------
+# Hybrid Router (rules + zero-shot)
+# -------------------------------------------------------------
+def hybrid_route(task_text: str):
+    if not task_text or not task_text.strip():
+        return "No task provided", "", ""
 
-    return predicted_label, module_id, scores_text
+    # 1. Rule-based (deterministic)
+    rule_choice = rule_router(task_text)
+    if rule_choice:
+        return GEN_BY_ID[rule_choice]["label"], rule_choice, "Rule-based match"
 
+    # 2. Zero-shot fallback
+    predicted_label, module_id, scores = classify_task(task_text)
+    return predicted_label, module_id, scores
 
-# -----------------------------
+
+# -------------------------------------------------------------
 # UI
-# -----------------------------
+# -------------------------------------------------------------
 def build_ui():
     with gr.Blocks(title="Modular Intelligence") as demo:
-        gr.Markdown(
-            "# Modular Intelligence Demo\n"
-            "Choose a module directly, or use the **Auto-route** tab to classify your task."
-        )
-
-        # ---- Auto-route tab (task classifier) ----
-        with gr.Tab("Auto-route (Task Classifier)"):
-            gr.Markdown(
-                "Paste any task or question. The classifier will suggest "
-                "which Modular Intelligence module to use."
-            )
-            task_box = gr.Textbox(
-                label="Describe what you want to do",
-                placeholder="E.g. 'I want a structured memo on options for exiting a JV under time pressure'...",
-                lines=6
-            )
-            predicted_label_box = gr.Textbox(
-                label="Predicted module (human-readable)",
-                interactive=False
-            )
-            predicted_id_box = gr.Textbox(
-                label="Internal module id",
-                interactive=False
-            )
-            scores_box = gr.Textbox(
-                label="Classifier scores",
-                interactive=False,
-                lines=10
-            )
-
-            classify_button = gr.Button("Classify task")
-            classify_button.click(
-                fn=classify_task,
+        gr.Markdown("# Modular Intelligence\nUpgraded architecture with routing, adapters, and reasoning layers.")
+
+        # -------------------- Auto-Route Tab --------------------
+        with gr.Tab("Auto-Route"):
+            task_box = gr.Textbox(label="Describe your task", lines=6)
+            module_name = gr.Textbox(label="Suggested Module", interactive=False)
+            module_id = gr.Textbox(label="Module ID", interactive=False)
+            scores = gr.Textbox(label="Routing Details", lines=12, interactive=False)
+
+            classify_btn = gr.Button("Classify Task")
+            classify_btn.click(
+                fn=hybrid_route,
                 inputs=[task_box],
-                outputs=[predicted_label_box, predicted_id_box, scores_box]
-            )
-
-            gr.Markdown(
-                "You can then go to the corresponding module tab below and "
-                "fill in its inputs using this guidance."
+                outputs=[module_name, module_id, scores],
             )
 
-        # ---- One tab per generator module ----
+        # -------------------- Module Tabs ------------------------
         for m in GENERATORS:
             with gr.Tab(m["label"]):
-                gr.Markdown(m["description"])
+                gr.Markdown(f"**Module ID:** `{m['id']}` | **Domain:** `{m.get('domain', 'general')}`")
+
                 inputs = []
-                for k, ph in m["input_placeholders"].items():
-                    t = gr.Textbox(label=k, placeholder=ph, lines=4)
+                for key, placeholder in m["input_placeholders"].items():
+                    t = gr.Textbox(label=key, placeholder=placeholder, lines=4)
                     inputs.append(t)
 
-                out = gr.Textbox(label="Output", lines=16)
-                gr.Button("Run").click(
+                output_box = gr.Textbox(label="Module Output", lines=20)
+                gr.Button("Run Module").click(
                     fn=partial(run_generator, m["id"]),
                     inputs=inputs,
-                    outputs=out
+                    outputs=output_box,
                 )
 
-                if m.get("has_checker"):
-                    cid = m.get("checker_id")
-                    if cid in CHECKERS:
-                        gr.Markdown("### Checker")
-                        chk = gr.Textbox(label="Checker Output", lines=14)
-                        gr.Button("Check").click(
-                            fn=partial(run_checker, cid),
-                            inputs=inputs + [out],
-                            outputs=chk
-                        )
+                checker_id = m.get("checker_id")
+                if checker_id and checker_id in CHECKERS:
+                    checker_output = gr.Textbox(label="Checker Output", lines=15)
+                    gr.Button("Run Checker").click(
+                        fn=partial(run_checker, checker_id),
+                        inputs=inputs + [output_box],
+                        outputs=checker_output,
+                    )
+                else:
+                    gr.Markdown("_No checker available for this module._")
 
     return demo