Update agent.py

agent.py

@@ -8,20 +8,26 @@ 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
+from pydantic import Field
 
 # ——— 1) Gemini LLM wrapper ———
 
 # Initialize your Gemini client once
 GENAI_CLIENT = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
-GEMINI_MODEL = "gemini-1.5-pro"   # or "gemini-1.0-turbo" per your trial
+GEMINI_MODEL = "gemini-1.5-pro"  # or whichever you have access to
 
 class GeminiLLM(LLM):
     """
-    A LangChain-compatible wrapper around Google Gemini via the google-genai SDK.
+    A LangChain-compatible wrapper around Google Gemini via google-genai.
     """
-    def __init__(self, client: genai.Client, model: str):
-        self.client = client
-        self.model = model
+    # Declare as Pydantic fields so they're accepted
+    client: genai.Client = Field(...)
+    model: str = Field(...)
+
+    def __init__(self, client: genai.Client, model: str, **kwargs):
+        super().__init__(**kwargs)
+        object.__setattr__(self, "client", client)
+        object.__setattr__(self, "model", model)
 
     @property
     def _llm_type(self) -> str:
@@ -36,7 +42,7 @@ class GeminiLLM(LLM):
         prompt: str,
         stop: Optional[List[str]] = None,
     ) -> str:
-        # Call the Gemini SDK
+        # Use the google-genai SDK
         response = self.client.generate_content(
             model=self.model,
             contents=[prompt]
@@ -51,7 +57,6 @@ def calculator(expr: str) -> str:
     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}"
@@ -75,16 +80,16 @@ class BasicAgent:
         # 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)
+        # Wrap Gemini as an LLM instance
+        gemini_llm = GeminiLLM(client=GENAI_CLIENT, model=GEMINI_MODEL)
 
-        # Initialize the agent with our tools
+        # Initialize the agent with our two tools
         self.agent = initialize_agent(
             tools=[calculator, wiki_search],
             llm=gemini_llm,
             agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-            verbose=True,             # turn on logs so you can debug
-            max_iterations=5,         # allow up to 5 tool calls
+            verbose=True,             # see tool calls in logs
+            max_iterations=5,
             early_stopping_method="generate",
         )
 
@@ -92,18 +97,15 @@ class BasicAgent:
         """
         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"
-            "Think internally; output ONLY the final answer (no chain-of-thought).\n\n"
+            "Think internally; output ONLY the final answer.\n\n"
             f"Question: {question}"
         )
 
-        # Run the agent
         raw_output = self.agent.run(prompt)
-
-        # Extract the last non-empty line as the answer
+        # Extract the last non-empty line
         lines = [line.strip() for line in raw_output.splitlines() if line.strip()]
         return lines[-1] if lines else raw_output.strip()