Update agent.py

agent.py

@@ -1,16 +1,30 @@
 # agent.py
 
-import os, requests
+import os
+import requests
 from langchain.tools import tool
 from langchain.agents import initialize_agent, AgentType
-
 from langchain_community.document_loaders import WikipediaLoader
 
-# 1) Define your tools
+# ——— 1) Gemini Client Setup ———
+from google import genai
+
+# Initialize once at import time
+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.
+
+def gemini_generate(prompt: str) -> str:
+    """Call Google Gemini via the GenAI SDK."""
+    response = GENAI_CLIENT.generate_content(
+        model=GEMINI_MODEL,
+        contents=[prompt]
+    )
+    return response.text
+
+# ——— 2) Tools ———
 
 @tool
 def calculator(expr: str) -> str:
-    """Safely evaluate a math expression."""
     try:
         return str(eval(expr, {"__builtins__": {}}))
     except Exception as e:
@@ -18,62 +32,33 @@ def calculator(expr: str) -> str:
 
 @tool
 def wiki_search(query: str) -> str:
-    """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 your Agent
+# ——— 3) Agent Definition ———
 
 class BasicAgent:
     def __init__(self):
-        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 tools and initialize a React agent
+        # We’re not using Hugging Face anymore—Gemini handles LLM calls
         self.agent = initialize_agent(
             [calculator, wiki_search],
-            self.llm,
+            # Wrap our gemini_generate as an LLM
+            lambda prompt: gemini_generate(prompt),
             agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-            verbose=True,            # see what it’s doing in the logs
-            max_iterations=5,        # let it call up to 5 tools
+            verbose=True,
+            max_iterations=5,
             early_stopping_method="generate"
         )
 
     def __call__(self, question: str) -> str:
-        # (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
-"""
+        # Prepend your toy examples or system prompt if you like
         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}"
+            "You have two tools:\n"
+            "  • calculator(expr)\n"
+            "  • wiki_search(query)\n"
+            "Use them internally, then OUTPUT ONLY the final answer.\n\n"
+            f"Question: {question}"
         )
-
-        # Run the agent
-        raw = self.agent.run(prompt)
-
-        # Extract the last line as the answer
-        return raw.splitlines()[-1].strip()
+        result = self.agent.run(prompt)
+        # Strip off anything but the last line
+        return result.splitlines()[-1].strip()