Update app.py

app.py

@@ -7,14 +7,19 @@ from bs4 import BeautifulSoup
 
 app = FastAPI()
 
-MODEL_NAME = "microsoft/phi-2"
+MODEL_NAME = "microsoft/phi-1_5"
 
-print("Loading Phi-2...")
+print("Loading model...")
 
 torch.set_num_threads(2)
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True
+)
+
 model.to("cpu")
 
 print("Model loaded!")
@@ -22,51 +27,131 @@ print("Model loaded!")
 # -------- REQUEST SCHEMA --------
 class RequestData(BaseModel):
     prompt: str
-    use_search: bool = False
+    history: list = []
+    use_search: bool = True
 
 
-# -------- WEB SEARCH FUNCTION --------
-def search_web(query):
+# -------- TOOL 1: SEARCH --------
+def search_links(query):
     url = f"https://duckduckgo.com/html/?q={query}"
     headers = {"User-Agent": "Mozilla/5.0"}
 
-    response = requests.get(url, headers=headers)
-    soup = BeautifulSoup(response.text, "html.parser")
+    try:
+        res = requests.get(url, headers=headers, timeout=10)
+        soup = BeautifulSoup(res.text, "html.parser")
+
+        links = []
+        for a in soup.select("a.result__a"):
+            href = a.get("href")
+            if href:
+                links.append(href)
+
+        return links[:3]
+    except:
+        return []
+
+
+# -------- TOOL 2: OPEN PAGE --------
+def extract_page_text(url):
+    try:
+        res = requests.get(url, timeout=10, headers={"User-Agent": "Mozilla/5.0"})
+        soup = BeautifulSoup(res.text, "html.parser")
+
+        for tag in soup(["script", "style"]):
+            tag.decompose()
+
+        text = soup.get_text(separator=" ")
+        return text[:2000]
+
+    except:
+        return ""
 
-    results = []
-    for a in soup.select("a.result__a"):
-        results.append(a.get_text())
 
-    return " ".join(results[:5])
+# -------- TOOL 3: BROWSE --------
+def browse_web(query):
+    links = search_links(query)
 
+    contents = []
+    for link in links:
+        page = extract_page_text(link)
+        if page:
+            contents.append(page)
 
-# -------- GENERATE FUNCTION --------
+    return "\n\n".join(contents[:3])
+
+
+# -------- MEMORY BUILDER --------
+def build_prompt(prompt, history):
+    convo = ""
+
+    for user, bot in history:
+        convo += f"User: {user}\nAssistant: {bot}\n"
+
+    convo += f"User: {prompt}\nAssistant:"
+    return convo
+
+
+# -------- GENERATION --------
 def generate_text(prompt):
-    formatted = f"Instruct: {prompt}\nOutput:"
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
 
-    inputs = tokenizer(formatted, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model.generate(
+            inputs["input_ids"],
+            max_new_tokens=120,
+            temperature=0.7,
+            do_sample=True
+        )
 
-    outputs = model.generate(
-        inputs["input_ids"],
-        max_new_tokens=60,
-        temperature=0.7
-    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return result.split("Output:")[-1].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.
+
+User question:
+{prompt}
+
+Should you search the web? Answer YES or NO.
+"""
+
+    decision = generate_text(decision_prompt).lower()
+
+    if use_search and "yes" in decision:
+        web_data = browse_web(prompt)
+
+        final_prompt = f"""
+You are an AI assistant with access to web data.
+
+Conversation:
+{base_prompt}
+
+Web Data:
+{web_data}
+
+Answer clearly and accurately:
+"""
+    else:
+        final_prompt = base_prompt
+
+    return generate_text(final_prompt)
 
 
 # -------- API ENDPOINT --------
 @app.post("/generate")
 def generate(data: RequestData):
-    prompt = data.prompt
-
-    if data.use_search:
-        web_data = search_web(prompt)
-        prompt = f"{prompt}\n\nWeb Info: {web_data}"
 
-    response = generate_text(prompt)
+    response = agent(
+        prompt=data.prompt,
+        history=data.history,
+        use_search=data.use_search
+    )
 
-    return {
-        "response": response
-    }
+    return {"response": response}