Update agent.py

agent.py

@@ -1,55 +1,79 @@
-import os
+# agent.py
+
+import os, requests
 from langchain.tools import tool
 from langchain.agents import initialize_agent, AgentType
-# HF-hub integration lives in langchain_community now:
-from langchain_community.llms.huggingface_hub import HuggingFaceHub
+
 from langchain_community.document_loaders import WikipediaLoader
 
-# 1) Define Tools
+# 1) Define your tools
 
 @tool
 def calculator(expr: str) -> str:
-    """
-    Safely evaluates a math expression and returns the result.
-    """
+    """Safely evaluate a math expression."""
     try:
-        # sandbox a bit
-        result = eval(expr, {"__builtins__": {}})
-        return str(result)
+        return str(eval(expr, {"__builtins__": {}}))
     except Exception as e:
         return f"Error: {e}"
 
 @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()
+    """Fetch up to 2 Wikipedia pages for the query."""
+    docs = WikipediaLoader(query=query, load_max_docs=2).load()
     return "\n\n".join(d.page_content for d in docs)
 
-# 2) Build the Agent
+# 2) Build your Agent
 
 class BasicAgent:
     def __init__(self):
-        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},
+        token = os.environ.get("HF_TOKEN")
+        assert token, "HF_TOKEN secret is missing!"
+        # We call the free inference endpoint directly
+        self.api_url = "https://api-inference.huggingface.co/models/google/flan-t5-large"
+        self.headers = {"Authorization": f"Bearer {token}"}
+
+        # LangChain’s HF wrapper
+        from langchain.llms import HuggingFaceEndpoint
+        self.llm = HuggingFaceEndpoint(
+            endpoint_url=self.api_url,
+            headers=self.headers,
+            model_kwargs={"temperature": 0.0, "max_new_tokens": 200},
         )
-        # Register your tools
+
+        # Register tools and initialize a React agent
         self.agent = initialize_agent(
             [calculator, wiki_search],
             self.llm,
             agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-            verbose=False,
-            max_iterations=3,
-            early_stopping_method="generate",
+            verbose=True,            # see what it’s doing in the logs
+            max_iterations=5,        # let it call up to 5 tools
+            early_stopping_method="generate"
         )
 
     def __call__(self, question: str) -> str:
-        # Run the agent; it will decide when (and how) to call tools
-        return self.agent.run(question)
+        # (Optional) Inject 3 hard-coded examples to guide format
+        EXAMPLES = """
+Q: What is 2+2?
+A: 4
+
+Q: If a car goes 60 km/h for 2 hours, how far?
+A: 120
+
+Q: What is the capital of France?
+A: Paris
+"""
+        prompt = (
+            f"Answer the following question using the tools below. "
+            f"First think (internally), then output **only** the final answer—no chain-of-thought.\n\n"
+            f"Tools:\n"
+            f"  • calculator(expr: str) -> str\n"
+            f"  • wiki_search(query: str) -> str\n\n"
+            f"### Examples ###{EXAMPLES}\n"
+            f"### New Question ###\n{question}"
+        )
+
+        # Run the agent
+        raw = self.agent.run(prompt)
+
+        # Extract the last line as the answer
+        return raw.splitlines()[-1].strip()