Update agent.py

agent.py

@@ -1,32 +1,30 @@
 # agent.py
 
 import os
-import requests
 from langchain.tools import tool
 from langchain.agents import initialize_agent, AgentType
 from langchain_community.document_loaders import WikipediaLoader
 
-# ——— 1) Gemini Client Setup ———
+# ——— Import & initialize Gemini client ———
 from google import genai
 
-# Initialize once at import time
+# Make sure GEMINI_API_KEY is set in your Secrets
 GENAI_CLIENT = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
-GEMINI_MODEL = "gemini-1.5-pro"  # or "gemini-1.0", "gemini-2.0-flash", etc.
+GEMINI_MODEL = "gemini-1.5-pro"   # <-- Use your trial model here
 
 def gemini_generate(prompt: str) -> str:
-    """Call Google Gemini via the GenAI SDK."""
-    response = GENAI_CLIENT.generate_content(
+    resp = GENAI_CLIENT.generate_content(
         model=GEMINI_MODEL,
         contents=[prompt]
     )
-    return response.text
+    return resp.text
 
-# ——— 2) Tools ———
+# ——— Tools ———
 
 @tool
 def calculator(expr: str) -> str:
-     """
-    Safely evaluates a math expression and returns the result.
+    """
+    Safely evaluates a math expression.
     """
     try:
         return str(eval(expr, {"__builtins__": {}}))
@@ -35,20 +33,20 @@ def calculator(expr: str) -> str:
 
 @tool
 def wiki_search(query: str) -> str:
-     """
-    Fetches up to 2 Wikipedia pages for the query and concatenates their text.
+    """
+    Fetches 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)
 
-# ——— 3) Agent Definition ———
+# ——— Agent ———
 
 class BasicAgent:
     def __init__(self):
-        # We’re not using Hugging Face anymore—Gemini handles LLM calls
+        # Zero-Shot React agent with our two tools
         self.agent = initialize_agent(
             [calculator, wiki_search],
-            # Wrap our gemini_generate as an LLM
+            # Wrap the gemini_generate function as the LLM
             lambda prompt: gemini_generate(prompt),
             agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
             verbose=True,
@@ -57,14 +55,15 @@ class BasicAgent:
         )
 
     def __call__(self, question: str) -> str:
-        # Prepend your toy examples or system prompt if you like
         prompt = (
-            "You have two tools:\n"
+            "You have access to two tools:\n"
             "  • calculator(expr)\n"
             "  • wiki_search(query)\n"
             "Use them internally, then OUTPUT ONLY the final answer.\n\n"
             f"Question: {question}"
         )
-        result = self.agent.run(prompt)
-        # Strip off anything but the last line
-        return result.splitlines()[-1].strip()
+
+        raw = self.agent.run(prompt)
+        # Return only the final line
+        lines = [l for l in raw.splitlines() if l.strip()]
+        return lines[-1].strip()