Update agent.py

agent.py

@@ -8,26 +8,18 @@ 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 whichever you have access to
+# ——— 1) Gemini SDK client ———
+GENAI_CLIENT = genai.Client(api_key=os.environ.get("GEMINI_API_KEY", ""))
+GEMINI_MODEL = "gemini-1.5-pro"  # or your trial model
 
+# ——— 2) LLM wrapper with class‐level defaults ———
 class GeminiLLM(LLM):
     """
-    A LangChain-compatible wrapper around Google Gemini via google-genai.
+    A LangChain-compatible LLM wrapper around Google Gemini.
     """
-    # 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)
+    client: genai.Client = GENAI_CLIENT
+    model: str = GEMINI_MODEL
 
     @property
     def _llm_type(self) -> str:
@@ -42,14 +34,13 @@ class GeminiLLM(LLM):
         prompt: str,
         stop: Optional[List[str]] = None,
     ) -> str:
-        # Use the google-genai SDK
         response = self.client.generate_content(
             model=self.model,
             contents=[prompt]
         )
         return response.text
 
-# ——— 2) Tools definitions ———
+# ——— 3) Tools ———
 
 @tool
 def calculator(expr: str) -> str:
@@ -64,39 +55,36 @@ def calculator(expr: str) -> str:
 @tool
 def wiki_search(query: str) -> str:
     """
-    Fetches up to 2 Wikipedia pages matching the query and concatenates their text.
+    Loads up to 2 Wikipedia pages for the query and concatenates their text.
     """
     docs = WikipediaLoader(query=query, load_max_docs=2).load()
-    return "\n\n".join(doc.page_content for doc in docs)
+    return "\n\n".join(d.page_content for d in docs)
 
-# ——— 3) Agent construction ———
+# ——— 4) Agent ———
 
 class BasicAgent:
     """
-    A Zero-Shot React agent powered by Google Gemini,
-    with access to calculator and wiki_search tools.
+    Zero-Shot React agent powered by Gemini + two tools.
     """
     def __init__(self):
-        # Ensure your secret is set
-        assert "GEMINI_API_KEY" in os.environ, "❌ GEMINI_API_KEY not found in Secrets"
+        # Ensure your key is present
+        assert os.environ.get("GEMINI_API_KEY"), "❌ GEMINI_API_KEY not set in Secrets"
 
-        # Wrap Gemini as an LLM instance
-        gemini_llm = GeminiLLM(client=GENAI_CLIENT, model=GEMINI_MODEL)
+        # Instantiate our LLM (defaults are on the class)
+        gemini_llm = GeminiLLM()
 
-        # Initialize the agent with our two tools
+        # Build the agent
         self.agent = initialize_agent(
             tools=[calculator, wiki_search],
             llm=gemini_llm,
             agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
-            verbose=True,             # see tool calls in logs
+            verbose=True,
             max_iterations=5,
             early_stopping_method="generate",
         )
 
     def __call__(self, question: str) -> str:
-        """
-        Runs the agent on the question and returns only the final answer line.
-        """
+        # Build the prompt
         prompt = (
             "You have access to two tools:\n"
             "  • calculator(expr)\n"
@@ -104,8 +92,7 @@ class BasicAgent:
             "Think internally; output ONLY the final answer.\n\n"
             f"Question: {question}"
         )
-
-        raw_output = self.agent.run(prompt)
-        # 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()
+        # Run & extract
+        raw = self.agent.run(prompt)
+        lines = [l.strip() for l in raw.splitlines() if l.strip()]
+        return lines[-1] if lines else raw.strip()