import os import time from dotenv import load_dotenv from langgraph.graph import StateGraph, END from langgraph.prebuilt import ToolNode, tools_condition from langchain_google_genai import ChatGoogleGenerativeAI from langchain_community.tools import DuckDuckGoSearchRun from langchain_community.document_loaders import WikipediaLoader, ArxivLoader from langchain_core.messages import SystemMessage, AIMessage, HumanMessage from langchain_core.tools import tool from tenacity import retry, stop_after_attempt, wait_exponential # Load environment variables load_dotenv() google_api_key = os.getenv("GOOGLE_API_KEY") or os.environ.get("GOOGLE_API_KEY") if not google_api_key: raise ValueError("Missing GOOGLE_API_KEY environment variable") # --- Math Tools --- @tool def multiply(a: int, b: int) -> int: """Multiply two integers.""" return a * b @tool def add(a: int, b: int) -> int: """Add two integers.""" return a + b @tool def subtract(a: int, b: int) -> int: """Subtract b from a.""" return a - b @tool def divide(a: int, b: int) -> float: """Divide a by b, error on zero.""" if b == 0: raise ValueError("Cannot divide by zero.") return a / b @tool def modulus(a: int, b: int) -> int: """Compute a mod b.""" return a % b # --- Browser Tools --- @tool def wiki_search(query: str) -> str: """Search Wikipedia and return up to 3 relevant documents.""" try: docs = WikipediaLoader(query=query, load_max_docs=3).load() if not docs: return "No Wikipedia results found." results = [] for doc in docs: title = doc.metadata.get('title', 'Unknown Title') content = doc.page_content[:2000] # Limit content length results.append(f"Title: {title}\nContent: {content}") return "\n\n---\n\n".join(results) except Exception as e: return f"Wikipedia search error: {str(e)}" @tool def arxiv_search(query: str) -> str: """Search Arxiv and return up to 3 relevant papers.""" try: docs = ArxivLoader(query=query, load_max_docs=3).load() if not docs: return "No arXiv papers found." results = [] for doc in docs: title = doc.metadata.get('Title', 'Unknown Title') authors = ", ".join(doc.metadata.get('Authors', [])) content = doc.page_content[:2000] # Limit content length results.append(f"Title: {title}\nAuthors: {authors}\nContent: {content}") return "\n\n---\n\n".join(results) except Exception as e: return f"arXiv search error: {str(e)}" @tool def web_search(query: str) -> str: """Search the web using DuckDuckGo and return top results.""" try: search = DuckDuckGoSearchRun() result = search.run(query) return f"Web search results for '{query}':\n{result[:2000]}" # Limit content length except Exception as e: return f"Web search error: {str(e)}" # --- Load system prompt --- with open("system_prompt.txt", "r", encoding="utf-8") as f: system_prompt = f.read() # --- System message --- sys_msg = SystemMessage(content=system_prompt) # --- Tool Setup --- tools = [ multiply, add, subtract, divide, modulus, wiki_search, arxiv_search, web_search, ] # --- Graph Builder --- def build_graph(): # Initialize model (Gemini 2.5 Flash) llm = ChatGoogleGenerativeAI( model="gemini-2.5-flash", temperature=0.3, google_api_key=google_api_key, max_retries=3 ) # Bind tools to LLM llm_with_tools = llm.bind_tools(tools) # Define state class AgentState: def __init__(self, messages): self.messages = messages # Node definitions with error handling def agent_node(state: AgentState): """Main agent node that processes messages with retry logic""" try: # Add rate limiting time.sleep(1) # 1 second delay between requests # Add retry logic for API quota issues @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10)) def invoke_llm_with_retry(): return llm_with_tools.invoke(state.messages) response = invoke_llm_with_retry() return AgentState(state.messages + [response]) except Exception as e: # Handle specific errors error_type = "UNKNOWN" if "429" in str(e): error_type = "QUOTA_EXCEEDED" elif "400" in str(e): error_type = "INVALID_REQUEST" error_msg = f"AGENT ERROR ({error_type}): {str(e)[:200]}" return AgentState(state.messages + [AIMessage(content=error_msg)]) # Tool node def tool_node(state: AgentState): """Execute tools based on agent's request""" last_message = state.messages[-1] tool_calls = last_message.additional_kwargs.get("tool_calls", []) tool_responses = [] for tool_call in tool_calls: tool_name = tool_call["function"]["name"] tool_args = tool_call["function"].get("arguments", {}) # Find the tool tool_func = next((t for t in tools if t.name == tool_name), None) if not tool_func: tool_responses.append(f"Tool {tool_name} not found") continue try: # Execute the tool if isinstance(tool_args, str): # Parse JSON if arguments are in string format import json tool_args = json.loads(tool_args) result = tool_func.invoke(tool_args) tool_responses.append(f"Tool {tool_name} result: {result}") except Exception as e: tool_responses.append(f"Tool {tool_name} error: {str(e)}") return AgentState(state.messages + [AIMessage(content="\n".join(tool_responses)]) # Custom condition function def should_continue(state: AgentState): last_message = state.messages[-1] # If there was an error, end if "AGENT ERROR" in last_message.content: return "end" # Check for tool calls if hasattr(last_message, "tool_calls") and last_message.tool_calls: return "tools" # Check for final answer if "FINAL ANSWER" in last_message.content: return "end" # Otherwise, continue to agent return "agent" # Build the graph workflow = StateGraph(AgentState) # Add nodes workflow.add_node("agent", agent_node) workflow.add_node("tools", tool_node) # Set entry point workflow.set_entry_point("agent") # Define edges workflow.add_conditional_edges( "agent", should_continue, { "agent": "agent", "tools": "tools", "end": END } ) workflow.add_conditional_edges( "tools", lambda state: "agent", { "agent": "agent" } ) return workflow.compile() # Initialize the agent graph agent_graph = build_graph()