Update agent.py

agent.py

@@ -1,71 +1,83 @@
 # agent.py
 
-import os, requests
-
+import os
+from langchain.tools import tool
+from langchain import HuggingFaceHub
+from langchain.agents import initialize_agent, AgentType
 from retrieval import retrieve_examples, wiki_context
 
-# Load your system prompt once
+# Load the system prompt
 SYSTEM_PROMPT = open("system_prompt.txt", "r").read().strip()
 
+# --- 1) Define your tools with @tool ---
+
+@tool
+def few_shot_retriever(query: str) -> str:
+    """
+    Returns a few-shot snippet (top-k examples) for this question.
+    """
+    return retrieve_examples(query, k=3)
+
+@tool
+def wiki_search(query: str) -> str:
+    """
+    Returns concatenated Wikipedia pages for factual lookup.
+    """
+    return wiki_context(query, max_docs=1)
+
+@tool
+def calculator(expression: str) -> str:
+    """
+    Safely evaluates a math expression and returns the result.
+    """
+    try:
+        # VERY simple eval: you can sandbox this further
+        result = eval(expression, {"__builtins__": {}})
+        return str(result)
+    except Exception as e:
+        return f"Error: {e}"
+
+# --- 2) Build your agent class ---
+
 class BasicAgent:
     def __init__(self):
-        print("🔧 Initializing BasicAgent…")
+        print("🔧 Initializing tool-using agent…")
+        # Read your HF_TOKEN from Secrets
         self.hf_token = os.environ.get("HF_TOKEN")
-        assert self.hf_token, "❌ HF_TOKEN not set in Secrets!"
-        self.model_id = "HuggingFaceH4/zephyr-7b-beta"
-        self.api_url = f"https://api-inference.huggingface.co/models/{self.model_id}"
-        self.headers = {
-            "Authorization": f"Bearer {self.hf_token}",
-            "Content-Type": "application/json"
-        }
+        assert self.hf_token, "HF_TOKEN environment variable not set!"
 
-    def should_use_wiki(self, q: str) -> bool:
-        kws = ["who is", "what is", "when was", "where is", "capital", "continent", "studio album"]
-        return any(kw in q.lower() for kw in kws)
-
-    def __call__(self, question: str) -> str:
-        print(f"📝 Question: {question}")
-
-        # 1) Retrieve few-shot examples
-        few_shot = retrieve_examples(question)
+        # 2a) LLM binding via HuggingFaceHub
+        self.llm = HuggingFaceHub(
+            repo_id="HuggingFaceH4/zephyr-7b-beta",
+            model_kwargs={"temperature": 0.2, "max_new_tokens": 250},
+            huggingfacehub_api_token=self.hf_token
+        )
 
-        # 2) Optionally retrieve wiki facts
-        wiki_snip = ""
-        if self.should_use_wiki(question):
-            try:
-                wiki_snip = wiki_context(question) + "\n\n"
-            except Exception as e:
-                print(f"⚠️ Wiki failed: {e}")
+        # 2b) Collect the decorated tools
+        self.tools = [few_shot_retriever, wiki_search, calculator]
 
-        # 3) Build the final prompt
-        prompt = (
-            f"[INST] {SYSTEM_PROMPT}\n\n"
-            f"{wiki_snip}"
-            f"### Few-Shot Examples ###\n{few_shot}\n\n"
-            f"### Question ###\n{question}\n\n"
-            "[/INST]"
+        # 2c) Initialize a React-style agent
+        self.agent = initialize_agent(
+            self.tools,
+            self.llm,
+            agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+            verbose=False,
+            max_iterations=3,
+            early_stopping_method="generate"
         )
 
-        # 4) Call the LLM
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "temperature": 0.2,
-                "max_new_tokens": 250,
-                "return_full_text": False
-            }
-        }
+    def __call__(self, question: str) -> str:
+        # Prepend your system_prompt for clean FINAL ANSWER formatting
+        prompt = f"{SYSTEM_PROMPT}\n\nQUESTION: {question}"
         try:
-            r = requests.post(self.api_url, headers=self.headers, json=payload, timeout=30)
-            r.raise_for_status()
-            out = r.json()
-            txt = out[0].get("generated_text", "").strip()
+            answer = self.agent.run(prompt)
         except Exception as e:
-            print(f"❌ LLM call failed: {e}")
+            print(f"❌ Agent.run() error: {e}")
             return f"Error: {e}"
 
-        # 5) Extract only the final answer
-        if "FINAL ANSWER:" in txt:
-            return txt.split("FINAL ANSWER:")[-1].strip()
-        # fallback: last line
-        return txt.splitlines()[-1].strip()
+        # The agent should already obey your "FINAL ANSWER:" rule,
+        # but just in case, strip out everything before it:
+        if "FINAL ANSWER:" in answer:
+            return answer.split("FINAL ANSWER:")[-1].strip()
+        # fallback to last line
+        return answer.splitlines()[-1].strip()