Update agent.py

agent.py

@@ -1,85 +1,55 @@
-# agent.py
-
 import os
 from langchain.tools import tool
-# new, works with langchain-community>=0.2.17
-from langchain_community.llms.huggingface_hub import HuggingFaceHub
-
 from langchain.agents import initialize_agent, AgentType
-from retrieval import retrieve_examples, wiki_context
-
-# Load the system prompt
-SYSTEM_PROMPT = open("system_prompt.txt", "r").read().strip()
-
-# --- 1) Define your tools with @tool ---
+# HF-hub integration lives in langchain_community now:
+from langchain_community.llms.huggingface_hub import HuggingFaceHub
+from langchain_community.document_loaders import WikipediaLoader
 
-@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)
+# 1) Define Tools
 
 @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:
+def calculator(expr: str) -> str:
     """
     Safely evaluates a math expression and returns the result.
     """
     try:
-        # VERY simple eval: you can sandbox this further
-        result = eval(expression, {"__builtins__": {}})
+        # sandbox a bit
+        result = eval(expr, {"__builtins__": {}})
         return str(result)
     except Exception as e:
         return f"Error: {e}"
 
-# --- 2) Build your agent class ---
+@tool
+def wiki_search(query: str) -> str:
+    """
+    Loads 1–2 Wikipedia pages for the query and concatenates their text.
+    """
+    loader = WikipediaLoader(query=query, load_max_docs=2)
+    docs = loader.load()
+    return "\n\n".join(d.page_content for d in docs)
+
+# 2) Build the Agent
 
 class BasicAgent:
     def __init__(self):
-        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 environment variable not set!"
-
-        # 2a) LLM binding via HuggingFaceHub
+        hf_token = os.environ.get("HF_TOKEN")
+        assert hf_token, "HF_TOKEN secret is missing!"
+        # Bind to a free-tier HF model
         self.llm = HuggingFaceHub(
             repo_id="HuggingFaceH4/zephyr-7b-beta",
+            huggingfacehub_api_token=hf_token,
             model_kwargs={"temperature": 0.2, "max_new_tokens": 250},
-            huggingfacehub_api_token=self.hf_token
         )
-
-        # 2b) Collect the decorated tools
-        self.tools = [few_shot_retriever, wiki_search, calculator]
-
-        # 2c) Initialize a React-style agent
+        # Register your tools
         self.agent = initialize_agent(
-            self.tools,
+            [calculator, wiki_search],
             self.llm,
             agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
             verbose=False,
             max_iterations=3,
-            early_stopping_method="generate"
+            early_stopping_method="generate",
         )
 
     def __call__(self, question: str) -> str:
-        # Prepend your system_prompt for clean FINAL ANSWER formatting
-        prompt = f"{SYSTEM_PROMPT}\n\nQUESTION: {question}"
-        try:
-            answer = self.agent.run(prompt)
-        except Exception as e:
-            print(f"❌ Agent.run() error: {e}")
-            return f"Error: {e}"
-
-        # 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()
+        # Run the agent; it will decide when (and how) to call tools
+        return self.agent.run(question)