MoBnJlal's picture
Update Agent.py
0261e3a verified
import os
from langchain.tools import Tool
from langgraph.graph import StateGraph, END, MessagesState
from langchain.agents import AgentExecutor, initialize_agent
from langchain.memory import ConversationBufferMemory
from langchain_community.tools import DuckDuckGoSearchRun
#from langchain.schema import SystemMessage
#from langchain.schema.messages import HumanMessage
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
GAIA_SYSTEM_PROMPT = """
You are an advanced autonomous agent designed to pass the GAIA benchmark.
You must:
- Think step-by-step before answering.
- Use available tools (e.g., Wikipedia, arXiv) to gather accurate information.
- Cite scientific sources when applicable.
- Reason before deciding what tool to use.
- When unsure, search or say "I need more information."
Available tools:
- Wikipedia for factual and general knowledge.
- arXiv for up-to-date scientific papers.
- You may call tools multiple times before giving a final answer.
Think aloud and explain your process clearly.
"""
#google_api_key = os.environ.get("GOOGLE_API_KEY")
from langchain_google_genai import ChatGoogleGenerativeAI
llm = ChatGoogleGenerativeAI(
model="gemini-2.0-flash",
google_api_key="",
temperature=0.2,
max_output_tokens=1024
)
messages = [
SystemMessage(content=GAIA_SYSTEM_PROMPT) ,
HumanMessage(content="Get started, messages on the way.")
]
response = llm.invoke(messages)
def calculator_tool(input: str) -> str:
try:
return str(eval(input))
except Exception as e:
return f"Error: {e}"
calc_tool = Tool(name="Calculator", func=calculator_tool, description="Perform basic math.")
search_tool = DuckDuckGoSearchRun()
import wikipedia
def search_wikipedia(query: str) -> str:
try:
return wikipedia.summary(query, sentences=3)
except Exception as e:
return f"Wikipedia Error: {str(e)}"
wikipedia_tool = Tool(
name="Wikipedia Search",
func=search_wikipedia,
description="Use this to search and summarize articles from Wikipedia."
)
import arxiv
def search_arxiv(query: str) -> str:
try:
search = arxiv.Search(
query=query,
max_results=1,
sort_by=arxiv.SortCriterion.Relevance
)
for result in search.results():
return f"{result.title}\n\n{result.summary}\n\nPDF: {result.pdf_url}"
return "No relevant arXiv papers found."
except Exception as e:
return f"arXiv Error: {str(e)}"
arxiv_tool = Tool(
name="arXiv Search",
func=search_arxiv,
description="Search for scientific papers on arXiv.org"
)
class GaiaState:
input: str
intermediate_steps: list = []
output: str = ""
# Initialize tools and agent
#llm = ChatOpenAI(temperature=0,model_name="gpt-4", system_message=SystemMessage(content=GAIA_SYSTEM_PROMPT))
tools = [calc_tool, search_tool, arxiv_tool, wikipedia_tool] # Add more tools as needed
agent = initialize_agent(
tools=tools,
llm=llm,
agent_type="openai-functions",
verbose=True
)
def run_agent(state: MessagesState) -> MessagesState:
# Let the agent respond based on the full message history
result = agent.invoke(state["messages"])
# Add the AI's response to the message history
state["messages"].append(AIMessage(content=result.content))
return state
graph = StateGraph(MessagesState)
graph.add_node("run_agent", run_agent)
graph.set_entry_point("run_agent")
graph.add_edge("run_agent", END)
# Compile the graph
app = graph.compile()