Update agent.py

agent.py

@@ -1,7 +1,9 @@
 import os
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.prebuilt import ToolNode, tools_condition
-from langchain_openai import ChatOpenAI
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_groq import ChatGroq
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
 from langchain_core.messages import SystemMessage, HumanMessage
 from langchain_core.tools import tool
 from langchain_community.tools.tavily_search import TavilySearchResults
@@ -14,6 +16,9 @@ from PIL import Image
 import re
 import requests
 from io import BytesIO
+from dotenv import load_dotenv
+
+load_dotenv()
 
 # Load system prompt
 with open("system_prompt.txt", "r", encoding="utf-8") as f:
@@ -22,7 +27,7 @@ with open("system_prompt.txt", "r", encoding="utf-8") as f:
 # Tool: Wikipedia search
 @tool
 def wiki_search(query: str) -> str:
-    """Wikipedia search tool."""
+    """Search Wikipedia for a query and return content from up to 2 documents."""
     try:
         docs = WikipediaLoader(query=query, load_max_docs=2).load()
         return "\n\n---\n\n".join([doc.page_content for doc in docs])
@@ -32,7 +37,7 @@ def wiki_search(query: str) -> str:
 # Tool: Tavily web search
 @tool
 def web_search(query: str) -> str:
-    """Tavily web search tool."""
+    """Search the web using Tavily and return content from up to 3 results."""
     try:
         results = TavilySearchResults(max_results=3).invoke(query)
         if isinstance(results, list):
@@ -44,7 +49,7 @@ def web_search(query: str) -> str:
 # Tool: DuckDuckGo search
 @tool
 def duckduckgo_search(query: str) -> str:
-    """DuckDuckGo search tool."""
+    """Search using DuckDuckGo and return summaries from up to 3 results."""
     try:
         with DDGS() as ddgs:
             results = ddgs.text(query, max_results=3)
@@ -55,23 +60,21 @@ def duckduckgo_search(query: str) -> str:
 # Tool: YouTube transcript or duration extractor
 @tool
 def youtube_transcript(video_title_or_url: str) -> str:
-    """YouTube transcript or duration extractor tool."""
+    """Get duration of a YouTube video using its title or URL."""
     try:
         with DDGS() as ddgs:
             results = ddgs.videos(video_title_or_url, max_results=1)
             if not results:
                 return "No video found by that title."
             video = results[0]
-            video_url = video["url"]
-            duration = video.get("duration")
-            return f"Duration: {duration}"
+            return f"Duration: {video.get('duration')}"
     except Exception as e:
         return f"YouTube search failed: {e}"
 
-# Tool: Arxiv paper fetcher (parse arXiv.org abstract directly)
+# Tool: Arxiv paper fetcher
 @tool
 def arxiv_fetch(query_or_id: str) -> str:
-    """Arxiv paper fetcher tool."""
+    """Fetch metadata from arXiv either by ID or search query."""
     try:
         if re.match(r"\d{4}\.\d{5}(v\d+)?", query_or_id):
             abs_url = f"https://arxiv.org/abs/{query_or_id}"
@@ -88,7 +91,7 @@ def arxiv_fetch(query_or_id: str) -> str:
 
 @tool
 def math_solver(expression: str) -> str:
-    """Math solver tool."""
+    """Evaluate a math expression and return the result."""
     try:
         result = sympify(expression).evalf()
         return str(result)
@@ -97,12 +100,12 @@ def math_solver(expression: str) -> str:
 
 @tool
 def reverse_text(text: str) -> str:
-    """Text reversal tool."""
+    """Reverse the input string."""
     return text[::-1]
 
 @tool
 def image_info(url: str) -> str:
-    """Image dimension fetcher tool."""
+    """Fetch image size (width x height) from a given URL."""
     try:
         response = requests.get(url)
         img = Image.open(BytesIO(response.content))
@@ -122,12 +125,21 @@ tools = [
     image_info
 ]
 
-def build_graph():
-    llm = ChatOpenAI(
-        model="gpt-4o",
-        temperature=0,
-        api_key=os.getenv("OPENAI_API_KEY")
-    )
+def build_graph(provider: str = "groq"):
+    if provider == "google":
+        llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash", temperature=0)
+    elif provider == "groq":
+        llm = ChatGroq(model="llama3-70b-8192", temperature=0)
+    elif provider == "huggingface":
+        llm = ChatHuggingFace(
+            llm=HuggingFaceEndpoint(
+                url="https://api-inference.huggingface.co/models/tiiuae/falcon-7b-instruct",
+                temperature=0,
+            ),
+        )
+    else:
+        raise ValueError("Invalid provider. Choose 'google', 'groq', or 'huggingface'.")
+
     llm_with_tools = llm.bind_tools(tools)
 
     def system_node(state: MessagesState):
@@ -147,9 +159,9 @@ def build_graph():
     return builder.compile()
 
 class BasicAgent:
-    def __init__(self):
-        print("GAIA LangGraph Agent Initialized")
-        self.graph = build_graph()
+    def __init__(self, provider="groq"):
+        print(f"GAIA LangGraph Agent Initialized using {provider}")
+        self.graph = build_graph(provider)
 
     def __call__(self, question: str) -> str:
         try:
@@ -163,7 +175,7 @@ class BasicAgent:
             return f"FINAL ANSWER: error - {str(e)}"
 
 if __name__ == "__main__":
-    agent = BasicAgent()
+    agent = BasicAgent(provider="groq")
     questions = [
         "What is the zip code of the Eiffel Tower?",
         "What is the capital city of Australia?",