Update app.py

app.py

@@ -1,23 +1,8 @@
-"""
-Veda Programming Assistant (Gradio 6.x)
-- Hidden teacher fallback (OpenRouter) when student fails
-- IMPORTANT FIX: Always use teacher for CODE requests (bubble sort etc.)
-- Auto-training in background using teacher responses
-- Math solver for simple arithmetic
-- Compatible with Gradio messages format + multimodal inputs
-"""
-
+import gradio as gr
+import threading
+import time
 import os
 import json
-import time
-import threading
-import re
-import ast
-import operator as op
-
-import gradio as gr
-import tensorflow as tf
-
 from model import VedaProgrammingLLM
 from tokenizer import VedaTokenizer
 from database import db
@@ -25,634 +10,87 @@ from train import VedaTrainer
 from teacher import teacher
 from config import MODEL_DIR
 
-
-# -----------------------------
-# GLOBALS
-# -----------------------------
 model = None
 tokenizer = None
+current_id = -1
 
-# For building student prompt context
-conversation_history = []  # list of dicts: {"user": "...", "assistant": "..."}
-
-current_conv_id = -1
-
-# Teacher usage stats (not shown in chat)
-teacher_used_count = 0
-teacher_failed_count = 0
-student_used_count = 0
-
-# ---- IMPORTANT BEHAVIOR SWITCH ----
-# Forces teacher for code requests so user sees correct code now.
-FORCE_TEACHER_FOR_CODE_REQUESTS = True
-
-# Auto-training control
-AUTO_TRAIN_ENABLED = True
-AUTO_TRAIN_MIN_TEACHER_SAMPLES = 10     # retrain after this many new teacher samples
-AUTO_TRAIN_CHECK_EVERY_SEC = 120        # check every 2 minutes
-AUTO_TRAIN_EPOCHS = 3                   # keep small for Spaces CPU
-AUTO_TRAIN_COOLDOWN_SEC = 60 * 20       # 20 minutes between trainings
-
-_is_training = False
-_last_train_time = 0
-_train_lock = threading.Lock()
-
-
-# -----------------------------
-# GRADIO INPUT HELPERS
-# -----------------------------
-def extract_text(message):
-    """
-    Convert Gradio multimodal/messages -> plain string.
-    Handles:
-      - str
-      - dict {"text": "..."} or {"content": ...}
-      - list [{"type":"text","text":"..."}]
-    """
-    if message is None:
-        return ""
-    if isinstance(message, str):
-        return message
-
-    if isinstance(message, dict):
-        if "text" in message:
-            return str(message.get("text", ""))
-        if "content" in message:
-            return extract_text(message["content"])
-        return ""
-
-    if isinstance(message, list):
-        out = []
-        for part in message:
-            if isinstance(part, dict) and part.get("type") == "text":
-                out.append(str(part.get("text", "")))
-            elif isinstance(part, str):
-                out.append(part)
-        return "".join(out).strip()
-
-    return str(message)
-
-
-def ensure_messages_history(history):
-    """
-    Ensure history is messages-format list:
-      [{"role":"user","content":"..."}, {"role":"assistant","content":"..."}]
-    Convert tuple format if needed.
-    """
-    if history is None:
-        return []
-
-    # already messages format
-    if len(history) > 0 and isinstance(history[0], dict) and "role" in history[0] and "content" in history[0]:
-        fixed = []
-        for m in history:
-            fixed.append({"role": m["role"], "content": extract_text(m["content"])})
-        return fixed
-
-    # tuple format -> messages format
-    fixed = []
-    for pair in history:
-        if isinstance(pair, (list, tuple)) and len(pair) == 2:
-            fixed.append({"role": "user", "content": extract_text(pair[0])})
-            fixed.append({"role": "assistant", "content": extract_text(pair[1])})
-    return fixed
-
-
-# -----------------------------
-# SAFE MATH SOLVER
-# -----------------------------
-_ALLOWED_OPS = {
-    ast.Add: op.add,
-    ast.Sub: op.sub,
-    ast.Mult: op.mul,
-    ast.Div: op.truediv,
-    ast.Mod: op.mod,
-    ast.Pow: op.pow,
-    ast.USub: op.neg,
-    ast.UAdd: op.pos,
-}
-
-def safe_eval_math(expr: str):
-    node = ast.parse(expr, mode="eval").body
-
-    def _eval(n):
-        if isinstance(n, ast.Constant) and isinstance(n.value, (int, float)):
-            return n.value
-        if isinstance(n, ast.BinOp) and type(n.op) in _ALLOWED_OPS:
-            return _ALLOWED_OPS[type(n.op)](_eval(n.left), _eval(n.right))
-        if isinstance(n, ast.UnaryOp) and type(n.op) in _ALLOWED_OPS:
-            return _ALLOWED_OPS[type(n.op)](_eval(n.operand))
-        raise ValueError("Unsupported expression")
-
-    return _eval(node)
-
-def try_math_answer(user_text: str):
-    if not user_text:
-        return None
-    s = user_text.strip()
-    s = s.replace("=", "").replace("?", "").strip()
-    s = s.replace("^", "**")  # allow ^
-
-    if not re.fullmatch(r"[0-9\.\s\+\-\*\/\(\)%]+", s):
-        return None
-
-    try:
-        val = safe_eval_math(s)
-        if isinstance(val, float) and val.is_integer():
-            val = int(val)
-        return str(val)
-    except Exception:
-        return None
-
-
-# -----------------------------
-# QUALITY CHECK + TEACHER TRIGGER
-# -----------------------------
-def is_code_request(user_text: str) -> bool:
-    t = user_text.lower()
-    triggers = [
-        "write", "implement", "code", "function", "algorithm",
-        "bubble sort", "binary search", "merge sort", "quick sort", "quicksort",
-        "linked list", "stack", "queue", "class ", "def ",
-        "sort "
-    ]
-    return any(k in t for k in triggers)
-
-def looks_like_python_code(text: str) -> bool:
-    """
-    Stronger code detector.
-    Only returns True if we see real python structure.
-    """
-    if not text:
-        return False
-    t = text.strip()
-
-    if "```" in t:
-        return True
-
-    # must contain python keywords + structure
-    if "def " in t and ":" in t:
-        return True
-    if "class " in t and ":" in t:
-        return True
-
-    # allow indented blocks only if also includes python keywords
-    if "\n    " in t and ("for " in t or "while " in t or "if " in t or "return " in t):
-        return True
-
-    return False
-
-def is_gibberish(text: str) -> bool:
-    if not text:
-        return True
-    t = text.strip()
-
-    if len(t) < 25:
-        return True
-
-    # repeated greeting
-    if t.lower().count("hello how are you") >= 1:
-        return True
-
-    # lots of symbols vs letters
-    letters = sum(c.isalpha() for c in t)
-    special = sum(c in "[]{}()=<>|\\" for c in t)
-    if letters > 0 and (special / max(letters, 1)) > 0.35:
-        return True
-
-    # low unique word ratio
-    words = re.findall(r"[a-zA-Z_]+", t.lower())
-    if len(words) >= 20:
-        uniq_ratio = len(set(words)) / len(words)
-        if uniq_ratio < 0.35:
-            return True
-
-    # junk patterns
-    junk_patterns = [
-        r"return\s+if\s+is",
-        r"=\s*=\s*=",
-        r"def\s+def",
-        r"class\s+class",
-        r"return\s+return",
-        r"\[\s*\"?\s*\]",
-    ]
-    for p in junk_patterns:
-        if re.search(p, t):
-            return True
-
-    # “p y t h o n” style (too many single-letter tokens)
-    single_letter_words = re.findall(r"\b[a-zA-Z]\b", t)
-    word_count = len(re.findall(r"\b[a-zA-Z_]+\b", t))
-    if word_count > 0 and (len(single_letter_words) / word_count) > 0.4:
-        return True
-
-    return False
-
-def should_use_teacher(user_text: str, student_text: str) -> bool:
-    if not teacher.is_available():
-        return False
-
-    # IMPORTANT: Force teacher for code requests (until student becomes good)
-    if FORCE_TEACHER_FOR_CODE_REQUESTS and is_code_request(user_text):
-        # If student actually produced code, you could skip teacher,
-        # but early stage student is bad, so use teacher always.
-        return True
-
-    # fallback to teacher if gibberish or not code when code asked
-    if is_code_request(user_text) and not looks_like_python_code(student_text):
-        return True
-
-    if is_gibberish(student_text):
-        return True
-
-    return False
-
-
-# -----------------------------
-# MODEL LOAD
-# -----------------------------
-def initialize():
+# Initialize
+def init():
     global model, tokenizer
-
-    print("Initializing Veda Programming Assistant...")
-
-    config_path = os.path.join(MODEL_DIR, "config.json")
-    weights_path = os.path.join(MODEL_DIR, "weights.h5")
-    tok_path = os.path.join(MODEL_DIR, "tokenizer.json")
-
-    if os.path.exists(config_path) and os.path.exists(weights_path) and os.path.exists(tok_path):
-        print("Loading existing model...")
-
-        with open(config_path, "r") as f:
-            config = json.load(f)
-
+    if os.path.exists(os.path.join(MODEL_DIR, "weights.h5")):
+        with open(os.path.join(MODEL_DIR, "config.json")) as f: conf = json.load(f)
         tokenizer = VedaTokenizer()
-        tokenizer.load(tok_path)
-
-        model = VedaProgrammingLLM(
-            vocab_size=config["vocab_size"],
-            max_length=config["max_length"],
-            d_model=config["d_model"],
-            num_heads=config["num_heads"],
-            num_layers=config["num_layers"],
-            ff_dim=config["ff_dim"],
-        )
-
-        dummy = tf.zeros((1, config["max_length"]), dtype=tf.int32)
-        model(dummy)
-        model.load_weights(weights_path)
-
-        print("Model loaded.")
-    else:
-        print("No saved model found. Training initial model...")
-        trainer = VedaTrainer()
-        trainer.train(epochs=10)
-        model = trainer.model
-        tokenizer = trainer.tokenizer
-        print("Initial model trained.")
-
-
-def clean_response(text: str) -> str:
-    if not text:
-        return ""
-
-    text = text.replace("<CODE>", "\n```python\n")
-    text = text.replace("<ENDCODE>", "\n```\n")
-
-    for token in ["<PAD>", "<UNK>", "<START>", "<END>", "<USER>", "<ASSISTANT>"]:
-        text = text.replace(token, "")
-
-    lines = text.split("\n")
-    cleaned = []
-    empty = 0
-    for line in lines:
-        if line.strip() == "":
-            empty += 1
-            if empty <= 2:
-                cleaned.append(line)
-        else:
-            empty = 0
-            cleaned.append(line)
-    return "\n".join(cleaned).strip()
-
-
-# -----------------------------
-# STUDENT + TEACHER RESPONSE
-# -----------------------------
-def get_student_response(user_text: str, temperature: float, max_tokens: int) -> str:
-    global student_used_count
-    if model is None or tokenizer is None:
-        return ""
-
-    context = ""
-    for m in conversation_history[-3:]:
-        context += f"<USER> {m['user']}\n<ASSISTANT> {m['assistant']}\n"
-
-    prompt = context + f"<USER> {user_text}\n<ASSISTANT>"
-    tokens = tokenizer.encode(prompt)
-
-    if len(tokens) > model.max_length - max_tokens:
-        tokens = tokens[-(model.max_length - max_tokens):]
-
-    generated = model.generate(
-        tokens,
-        max_new_tokens=max_tokens,
-        temperature=temperature,
-        top_k=50,
-        top_p=0.9,
-        repetition_penalty=1.2,
-    )
-
-    out = tokenizer.decode(generated)
-
-    if "<ASSISTANT>" in out:
-        out = out.split("<ASSISTANT>")[-1].strip()
-    if "<USER>" in out:
-        out = out.split("<USER>")[0].strip()
-
-    student_used_count += 1
-    return clean_response(out)
-
-
-def get_teacher_response(user_text: str) -> str:
-    teacher_hist = []
-    for m in conversation_history[-4:]:
-        teacher_hist.append({"role": "user", "content": m["user"]})
-        teacher_hist.append({"role": "assistant", "content": m["assistant"]})
-
-    return teacher.ask(user_message=user_text, conversation_history=teacher_hist) or ""
-
-
-# -----------------------------
-# MAIN GENERATION (HIDDEN TEACHER)
-# -----------------------------
-def generate_response(user_input, temperature=0.7, max_tokens=200) -> str:
-    global current_conv_id, teacher_used_count, teacher_failed_count
-
-    user_text = extract_text(user_input).strip()
-    if not user_text:
-        return "Please type a message."
-
-    # math first
-    math_ans = try_math_answer(user_text)
-    if math_ans is not None:
-        conversation_history.append({"user": user_text, "assistant": math_ans})
-        current_conv_id = db.save_conversation(user_text, math_ans)
-        return math_ans
-
-    # student attempt
-    student = get_student_response(user_text, float(temperature), int(max_tokens))
-
-    if should_use_teacher(user_text, student):
-        teacher_resp = get_teacher_response(user_text)
-
-        if teacher_resp.strip():
-            teacher_used_count += 1
-
-            # Save distillation sample
-            try:
-                db.save_distillation_data(
-                    user_input=user_text,
-                    teacher_response=teacher_resp,
-                    student_response=student,
-                    quality_score=1.0,
-                )
-            except Exception as e:
-                print("Could not save distillation sample:", e)
-
-            final = teacher_resp
-        else:
-            teacher_failed_count += 1
-            final = student if student else "Please try again."
+        tokenizer.load(os.path.join(MODEL_DIR, "tokenizer.json"))
+        model = VedaProgrammingLLM(**conf)
+        model(tf.zeros((1, conf['max_length'])))
+        model.load_weights(os.path.join(MODEL_DIR, "weights.h5"))
     else:
-        final = student
-
-    final = clean_response(final)
-    if not final:
-        final = "Please try asking in a different way."
-
-    conversation_history.append({"user": user_text, "assistant": final})
-    current_conv_id = db.save_conversation(user_text, final)
-    return final
-
-
-# -----------------------------
-# AUTO TRAINING
-# -----------------------------
-def auto_train_loop():
-    global _is_training, _last_train_time, model, tokenizer
+        print("Training initial model...")
+        VedaTrainer().train(epochs=15)
+        init()
 
+# Auto-train loop
+def auto_train():
     while True:
-        time.sleep(AUTO_TRAIN_CHECK_EVERY_SEC)
-
-        if not AUTO_TRAIN_ENABLED:
-            continue
-
-        if time.time() - _last_train_time < AUTO_TRAIN_COOLDOWN_SEC:
-            continue
-
-        if _train_lock.locked():
-            continue
-
-        try:
-            unused = db.get_unused_distillation_data(limit=1000)
-        except Exception as e:
-            print("[AutoTrain] Could not read distillation data:", e)
-            continue
-
-        if len(unused) < AUTO_TRAIN_MIN_TEACHER_SAMPLES:
-            continue
-
-        with _train_lock:
-            _is_training = True
-            print(f"[AutoTrain] Training on {len(unused)} teacher samples...")
-
-            try:
-                distill_text = ""
-                ids = []
-                for row in unused:
-                    ids.append(row["id"])
-                    distill_text += f"<USER> {row['user_input']}\n<ASSISTANT> {row['teacher_response']}\n\n"
-
-                # include user-positive feedback too
-                extra = ""
-                try:
-                    good = db.get_good_conversations(limit=200)
-                    for conv in good:
-                        extra += f"<USER> {conv['user_input']}\n<ASSISTANT> {conv['assistant_response']}\n\n"
-                except Exception:
-                    pass
-
-                trainer = VedaTrainer()
-                hist = trainer.train(
-                    epochs=AUTO_TRAIN_EPOCHS,
-                    extra_data=extra,
-                    distillation_data=distill_text,
-                )
-
-                model = trainer.model
-                tokenizer = trainer.tokenizer
-
-                try:
-                    db.mark_distillation_used(ids)
-                except Exception as e:
-                    print("[AutoTrain] Could not mark distillation used:", e)
-
-                loss = float(hist.history["loss"][-1])
-                try:
-                    db.save_training_history(
-                        training_type="auto",
-                        samples_used=len(unused),
-                        epochs=AUTO_TRAIN_EPOCHS,
-                        final_loss=loss,
-                    )
-                except Exception:
-                    pass
-
-                _last_train_time = time.time()
-                print(f"[AutoTrain] Done. loss={loss:.4f}")
-
-            except Exception as e:
-                print("[AutoTrain] Training failed:", e)
-
-            _is_training = False
-
-
-# -----------------------------
-# GRADIO HANDLERS
-# -----------------------------
-def respond(message, history, temperature, max_tokens):
-    history = ensure_messages_history(history)
-    user_text = extract_text(message).strip()
-    if not user_text:
-        return "", history
-
-    bot_text = generate_response(user_text, temperature=float(temperature), max_tokens=int(max_tokens))
-
-    history.append({"role": "user", "content": user_text})
-    history.append({"role": "assistant", "content": bot_text})
-
+        time.sleep(300) # Check every 5 mins
+        data = db.get_unused_distillation()
+        if len(data) >= 5:
+            print("Auto-training on teacher data...")
+            text = "\n".join([f"<USER> {r[1]}\n<ASSISTANT> {r[2]}" for r in data])
+            VedaTrainer().train(epochs=5, extra_data=text)
+            db.mark_used([r[0] for r in data])
+            # Reload
+            init()
+
+threading.Thread(target=auto_train, daemon=True).start()
+
+def is_good(text):
+    if not text or len(text) < 10: return False
+    if "arr[" in text and "return" not in text: return False # Gibberish check
+    return True
+
+def respond(msg, history):
+    global current_id
+    if not msg.strip(): return "", history
+    
+    # 1. Try student
+    prompt = f"<USER> {msg}\n<ASSISTANT>"
+    toks = tokenizer.encode(prompt)
+    out = model.generate(toks, max_new_tokens=200)
+    resp = tokenizer.decode(out).split("<ASSISTANT>")[-1].split("<USER>")[0].strip()
+    
+    # Clean code tags
+    if "<CODE>" in resp:
+        resp = resp.replace("<CODE>", "```python\n").replace("</CODE>", "\n```")
+    elif "```" in resp and not ("def " in resp or "print" in resp):
+        # If model hallucinated code blocks around text
+        resp = resp.replace("```", "")
+
+    # 2. Check quality & fallback
+    if not is_good(resp) and teacher.is_available():
+        teacher_resp = teacher.ask(msg)
+        if teacher_resp:
+            resp = teacher_resp
+            db.save_distillation(msg, teacher_resp) # Save for learning
+    
+    current_id = db.save_conversation(msg, resp)
+    history.append({"role": "user", "content": msg})
+    history.append({"role": "assistant", "content": resp})
     return "", history
 
+def feedback(vote):
+    if current_id > 0: db.update_feedback(current_id, 1 if vote=="good" else -1)
 
-def feedback_good():
-    if current_conv_id > 0:
-        db.update_feedback(current_conv_id, 1)
-        return "Thanks!"
-    return "No message to rate yet."
-
-
-def feedback_bad():
-    if current_conv_id > 0:
-        db.update_feedback(current_conv_id, -1)
-        return "Thanks!"
-    return "No message to rate yet."
-
-
-def clear_chat():
-    global conversation_history
-    conversation_history = []
-    return [], ""
-
-
-def get_stats_md():
-    stats = db.get_stats()
-    teacher_ok = teacher.is_available()
-    return f"""
-## Statistics
-
-**Teacher available:** `{teacher_ok}`  
-**Teacher used (this runtime):** `{teacher_used_count}`  
-**Teacher failed (this runtime):** `{teacher_failed_count}`  
-**Student calls (this runtime):** `{student_used_count}`  
-**Auto-training enabled:** `{AUTO_TRAIN_ENABLED}`  
-**Currently training:** `{_is_training}`  
-
-### Conversations
-- Total: **{stats.get('total', 0)}**
-- Positive: **{stats.get('positive', 0)}**
-- Negative: **{stats.get('negative', 0)}**
-
-### Distillation (teacher lessons)
-- Total saved: **{stats.get('distillation_total', 0)}**
-- Pending for training: **{stats.get('distillation_unused', 0)}**
-"""
-
-
-# -----------------------------
-# STARTUP
-# -----------------------------
-print("=== Booting Veda Assistant ===")
-initialize()
-
-print("Teacher available:", teacher.is_available())
-if AUTO_TRAIN_ENABLED:
-    t = threading.Thread(target=auto_train_loop, daemon=True)
-    t.start()
-    print("Auto-training thread started.")
-print("=== Ready ===")
-
-
-# -----------------------------
 # UI
-# -----------------------------
-with gr.Blocks(title="Veda Programming Assistant") as demo:
-    gr.Markdown(
-        """
-# Veda Programming Assistant
-
-Ask programming questions, request code, or do math like `2+2=?` or `(10+5)/3`.
-"""
-    )
-
-    with gr.Tabs():
-        with gr.TabItem("Chat"):
-            chatbot = gr.Chatbot(label="Conversation", height=420, value=[])
-
-            with gr.Row():
-                msg = gr.Textbox(
-                    label="Message",
-                    placeholder="Example: Write bubble sort in python",
-                    lines=2,
-                    scale=4,
-                )
-                send = gr.Button("Send", variant="primary", scale=1)
-
-            with gr.Row():
-                temperature = gr.Slider(0.1, 1.5, 0.7, step=0.1, label="Temperature")
-                max_tokens = gr.Slider(50, 400, 200, step=50, label="Max tokens")
-
-            with gr.Row():
-                good = gr.Button("Helpful", variant="secondary")
-                bad = gr.Button("Not helpful", variant="secondary")
-                clear = gr.Button("Clear", variant="secondary")
-
-            status = gr.Textbox(label="", show_label=False, lines=1)
-
-            send.click(respond, inputs=[msg, chatbot, temperature, max_tokens], outputs=[msg, chatbot])
-            msg.submit(respond, inputs=[msg, chatbot, temperature, max_tokens], outputs=[msg, chatbot])
-
-            good.click(feedback_good, outputs=status)
-            bad.click(feedback_bad, outputs=status)
-            clear.click(clear_chat, outputs=[chatbot, status])
-
-            gr.Examples(
-                examples=[
-                    ["Write bubble sort in python"],
-                    ["Write binary search in python"],
-                    ["Explain recursion with example"],
-                    ["2+2=?"],
-                    ["(10+5)/3"],
-                    ["2^5"],
-                ],
-                inputs=msg,
-            )
-
-        with gr.TabItem("Statistics"):
-            stats_md = gr.Markdown()
-            refresh = gr.Button("Refresh")
-            refresh.click(get_stats_md, outputs=stats_md)
-            demo.load(get_stats_md, outputs=stats_md)
-
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+init()
+with gr.Blocks(title="Veda") as demo:
+    gr.Markdown("# 🕉️ Veda Assistant")
+    chat = gr.Chatbot(type="messages", height=400)
+    msg = gr.Textbox(label="Message")
+    with gr.Row():
+        gr.Button("👍").click(lambda: feedback("good"))
+        gr.Button("👎").click(lambda: feedback("bad"))
+    
+    msg.submit(respond, [msg, chat], [msg, chat])
+
+demo.launch(server_name="0.0.0.0", server_port=7860)