Update app.py

app.py

@@ -1,50 +1,132 @@
+import re
+import ast
+import operator
 import gradio as gr
-from huggingface_hub import InferenceClient
-from typing import List, Tuple
-import os
-
-# Initialize client with token (add HF_TOKEN in your Space secrets)
-client = InferenceClient(
-    model="HuggingFaceH4/zephyr-7b-beta",
-    token=os.getenv("HF_TOKEN")
-)
 
-def respond(
-    message: str,
-    history: List[Tuple[str, str]],
-    system_message: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-):
-    # Build conversation history
-    messages = [{"role": "system", "content": system_message}]
-    for user_msg, bot_reply in history:
-        if user_msg:
-            messages.append({"role": "user", "content": user_msg})
-        if bot_reply:
-            messages.append({"role": "assistant", "content": bot_reply})
-    messages.append({"role": "user", "content": message})
-
-    # Call Hugging Face Inference API
-    msg = client.chat_completion(
-        model="HuggingFaceH4/zephyr-7b-beta",
-        messages=messages,
-        max_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
+# LLM: small, CPU-friendly
+from transformers import pipeline
+
+LLM = None  # lazy-load to speed up app boot
+
+
+# ---------- Agent 1: Planner (decides the next action) ----------
+def planner_agent(user_msg: str) -> dict:
+    """
+    Very small heuristic planner:
+    - If it's a simple math question/expression => CALCULATE
+    - Else => ANSWER (LLM)
+    """
+    text = user_msg.strip().lower()
+
+    # If it looks like a calculation
+    math_like = bool(re.search(r"[0-9][0-9\.\s\+\-\*\/\%\^\(\)]*[0-9\)]", text))
+    asks_calc = any(k in text for k in ["calc", "calculate", "evaluate", "what is", "what's"])
+
+    if math_like or asks_calc:
+        # Avoid "story" or "explain" with numbers
+        if not any(k in text for k in ["story", "poem", "explain why", "compare"]):
+            return {"action": "CALCULATE", "reason": "Looks like a numeric expression or calculation request."}
+
+    return {"action": "ANSWER", "reason": "General question; better handled by the LLM."}
+
+
+# ---------- Tiny safe calculator (tool used by Solver) ----------
+# Safe AST-based evaluator supporting +,-,*,/,**,%,// and parentheses
+OPS = {
+    ast.Add: operator.add,
+    ast.Sub: operator.sub,
+    ast.Mult: operator.mul,
+    ast.Div: operator.truediv,
+    ast.Pow: operator.pow,
+    ast.Mod: operator.mod,
+    ast.FloorDiv: operator.floordiv,
+    ast.USub: operator.neg,
+    ast.UAdd: operator.pos,
+}
+
+def _eval(node):
+    if isinstance(node, ast.Num):  # py<3.8 fallback
+        return node.n
+    if isinstance(node, ast.Constant):
+        if isinstance(node.value, (int, float)):
+            return node.value
+        raise ValueError("Only numeric constants allowed.")
+    if isinstance(node, ast.BinOp):
+        left = _eval(node.left)
+        right = _eval(node.right)
+        op_type = type(node.op)
+        if op_type in OPS:
+            return OPS[op_type](left, right)
+        raise ValueError("Unsupported operator.")
+    if isinstance(node, ast.UnaryOp) and type(node.op) in OPS:
+        return OPS[type(node.op)](_eval(node.operand))
+    if isinstance(node, ast.Expression):
+        return _eval(node.body)
+    raise ValueError("Unsupported expression.")
+
+def safe_calculate(expr: str) -> str:
+    try:
+        tree = ast.parse(expr, mode="eval")
+        val = _eval(tree)
+        return str(val)
+    except Exception as e:
+        return f"Sorry, I couldn't calculate that: {e}"
+
+
+# ---------- Agent 2: Writer (LLM or Tool) ----------
+def get_llm():
+    global LLM
+    if LLM is None:
+        # flan-t5-small is ~80M params, OK on CPU Basic
+        LLM = pipeline("text2text-generation", model="google/flan-t5-small")
+    return LLM
+
+def writer_agent(user_msg: str, plan: dict) -> str:
+    if plan["action"] == "CALCULATE":
+        # Extract the most likely expression from the message
+        # Keep digits, ops, dots, spaces, and parentheses
+        expr = "".join(ch for ch in user_msg if ch in "0123456789.+-*/()%^ //")
+        # Clean up accidental double spaces
+        expr = re.sub(r"\s+", "", expr.replace("^", "**"))
+        if not expr:
+            # fallback to LLM if no expression found
+            llm = get_llm()
+            prompt = f"Answer briefly:\n\nQuestion: {user_msg}\nAnswer:"
+            out = llm(prompt, max_new_tokens=128, do_sample=True)[0]["generated_text"]
+            return out.strip()
+        return safe_calculate(expr)
+
+    # ANSWER with LLM
+    llm = get_llm()
+    prompt = (
+        "You are a concise, friendly assistant. "
+        "Answer clearly in 1-4 sentences.\n\n"
+        f"Question: {user_msg}\nAnswer:"
     )
-    return msg.choices[0].message.content
+    out = llm(prompt, max_new_tokens=192, do_sample=True, temperature=0.6, top_p=0.95)[0]["generated_text"]
+    return out.strip()
+
+
+# ---------- Gradio Chat glue ----------
+def chat_fn(message, history, show_agent_trace=False):
+    plan = planner_agent(message)
+    answer = writer_agent(message, plan)
+
+    if show_agent_trace:
+        trace = f"\n\n---\n*Agent trace:* action = **{plan['action']}**, reason = _{plan['reason']}_"
+        return answer + trace
+    return answer
+
 
-# Gradio Chat UI
 demo = gr.ChatInterface(
-    respond,
+    fn=chat_fn,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
+        gr.Checkbox(label="Show agent trace", value=False),
     ],
+    theme="soft",
+    css="""
+    .gradio-container {max-width: 760px !important}
+    """,
 )
 
 if __name__ == "__main__":