Update app.py

app.py

@@ -16,7 +16,7 @@ torch.set_num_threads(2)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    torch_dtype=torch.float32,
+    dtype=torch.float32,
     low_cpu_mem_usage=True
 )
 
@@ -24,6 +24,7 @@ model.to("cpu")
 
 print("Model loaded!")
 
+
 # -------- REQUEST SCHEMA --------
 class RequestData(BaseModel):
     prompt: str
@@ -31,6 +32,12 @@ class RequestData(BaseModel):
     use_search: bool = True
 
 
+# -------- ROOT ROUTE --------
+@app.get("/")
+def home():
+    return {"message": "API is running"}
+
+
 # -------- TOOL 1: SEARCH --------
 def search_links(query):
     url = f"https://duckduckgo.com/html/?q={query}"
@@ -80,18 +87,22 @@ def browse_web(query):
     return "\n\n".join(contents[:3])
 
 
-# -------- MEMORY BUILDER --------
+# -------- MEMORY BUILDER (FIXED) --------
 def build_prompt(prompt, history):
     convo = ""
 
-    for user, bot in history:
-        convo += f"User: {user}\nAssistant: {bot}\n"
+    for msg in history:
+        if isinstance(msg, dict):
+            if msg.get("role") == "user":
+                convo += f"User: {msg.get('content')}\n"
+            elif msg.get("role") == "assistant":
+                convo += f"Assistant: {msg.get('content')}\n"
 
     convo += f"User: {prompt}\nAssistant:"
     return convo
 
 
-# -------- GENERATION --------
+# -------- GENERATION (FIXED OUTPUT) --------
 def generate_text(prompt):
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
 
@@ -103,16 +114,17 @@ def generate_text(prompt):
             do_sample=True
         )
 
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Remove prompt from output
+    return full_text[len(prompt):].strip()
 
 
 # -------- AGENT LOOP --------
 def agent(prompt, history, use_search=True):
 
-    # Step 1: Build conversation
     base_prompt = build_prompt(prompt, history)
 
-    # Step 2: Decide if search is needed
     decision_prompt = f"""
 You are an AI agent.