Update agent.py

agent.py

@@ -1,32 +1,57 @@
 # agent.py
 
 import os
+from typing import Optional, List, Mapping, Any
+
+from google import genai
 from langchain.tools import tool
 from langchain.agents import initialize_agent, AgentType
 from langchain_community.document_loaders import WikipediaLoader
+from langchain.llms.base import LLM
 
-# ——— Import & initialize Gemini client ———
-from google import genai
+# ——— 1) Gemini LLM wrapper ———
 
-# Make sure GEMINI_API_KEY is set in your Secrets
+# Initialize your Gemini client once
 GENAI_CLIENT = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
-GEMINI_MODEL = "gemini-1.5-pro"   # <-- Use your trial model here
+GEMINI_MODEL = "gemini-1.5-pro"   # or "gemini-1.0-turbo" per your trial
+
+class GeminiLLM(LLM):
+    """
+    A LangChain-compatible wrapper around Google Gemini via the google-genai SDK.
+    """
+    def __init__(self, client: genai.Client, model: str):
+        self.client = client
+        self.model = model
+
+    @property
+    def _llm_type(self) -> str:
+        return "gemini"
 
-def gemini_generate(prompt: str) -> str:
-    resp = GENAI_CLIENT.generate_content(
-        model=GEMINI_MODEL,
-        contents=[prompt]
-    )
-    return resp.text
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        return {"model": self.model}
 
-# ——— Tools ———
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+    ) -> str:
+        # Call the Gemini SDK
+        response = self.client.generate_content(
+            model=self.model,
+            contents=[prompt]
+        )
+        return response.text
+
+# ——— 2) Tools definitions ———
 
 @tool
 def calculator(expr: str) -> str:
     """
-    Safely evaluates a math expression.
+    Safely evaluates a math expression and returns the result.
     """
     try:
+        # simple sandbox
         return str(eval(expr, {"__builtins__": {}}))
     except Exception as e:
         return f"Error: {e}"
@@ -34,36 +59,51 @@ def calculator(expr: str) -> str:
 @tool
 def wiki_search(query: str) -> str:
     """
-    Fetches up to 2 Wikipedia pages for the query.
+    Fetches up to 2 Wikipedia pages matching the query and concatenates their text.
     """
     docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    return "\n\n".join(d.page_content for d in docs)
+    return "\n\n".join(doc.page_content for doc in docs)
 
-# ——— Agent ———
+# ——— 3) Agent construction ———
 
 class BasicAgent:
+    """
+    A Zero-Shot React agent powered by Google Gemini,
+    with access to calculator and wiki_search tools.
+    """
     def __init__(self):
-        # Zero-Shot React agent with our two tools
+        # Ensure your secret is set
+        assert "GEMINI_API_KEY" in os.environ, "❌ GEMINI_API_KEY not found in Secrets"
+
+        # Wrap Gemini as an LLM
+        gemini_llm = GeminiLLM(GENAI_CLIENT, GEMINI_MODEL)
+
+        # Initialize the agent with our tools
         self.agent = initialize_agent(
-            [calculator, wiki_search],
-            # Wrap the gemini_generate function as the LLM
-            lambda prompt: gemini_generate(prompt),
+            tools=[calculator, wiki_search],
+            llm=gemini_llm,
             agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-            verbose=True,
-            max_iterations=5,
-            early_stopping_method="generate"
+            verbose=True,             # turn on logs so you can debug
+            max_iterations=5,         # allow up to 5 tool calls
+            early_stopping_method="generate",
         )
 
     def __call__(self, question: str) -> str:
+        """
+        Runs the agent on the question and returns only the final answer line.
+        """
+        # System prompt + question
         prompt = (
             "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"
+            "Think internally; output ONLY the final answer (no chain-of-thought).\n\n"
             f"Question: {question}"
         )
 
-        raw = self.agent.run(prompt)
-        # Return only the final line
-        lines = [l for l in raw.splitlines() if l.strip()]
-        return lines[-1].strip()
+        # Run the agent
+        raw_output = self.agent.run(prompt)
+
+        # Extract the last non-empty line as the answer
+        lines = [line.strip() for line in raw_output.splitlines() if line.strip()]
+        return lines[-1] if lines else raw_output.strip()