Update Agent.py

Agent.py

@@ -1,3 +1,4 @@
+import os
 from langchain.tools import Tool
 from langgraph.graph import StateGraph, END
 from langchain.agents import AgentExecutor, initialize_agent
@@ -5,7 +6,6 @@ from langchain.memory import ConversationBufferMemory
 from langchain_community.tools import DuckDuckGoSearchRun
 from langchain.schema import SystemMessage
 from langchain.schema.messages import HumanMessage
-from langchain_huggingface import ChatHuggingFace
 GAIA_SYSTEM_PROMPT = """
 You are an advanced autonomous agent designed to pass the GAIA benchmark.
 You must:
@@ -23,21 +23,14 @@ Available tools:
 
 Think aloud and explain your process clearly.
 """
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-import torch
-
-# Load the HF model + tokenizer
-model_id = "HuggingFaceH4/zephyr-7b-beta"
-
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
-
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=1024)
-
-# Create the LangChain chat wrapper
-llm = ChatHuggingFace(pipeline=pipe)
-
-
+google_api_key = os.environ.get("GOOGLE_API_KEY")
+from langchain_google_genai import ChatGoogleGenerativeAI
+llm = ChatGoogleGenerativeAI(
+    model="gemini-pro",
+    google_api_key=google_api_key,
+    temperature=0.2,
+    max_output_tokens=1024
+)
 messages = [
     SystemMessage(content=GAIA_SYSTEM_PROMPT)    
 ]