Final_Assignment_Template / langgraph_agent.py
devasurya's picture
Update langgraph_agent.py
d360421 verified
raw
history blame
1.92 kB
import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.messages import AnyMessage, HumanMessage
from langchain_core.runnables import RunnableConfig
from langgraph.prebuilt import create_react_agent
from langgraph.prebuilt.chat_agent_executor import AgentState
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_experimental.utilities import PythonREPL
from langchain_community.tools import WikipediaQueryRun
from langchain_community.utilities import WikipediaAPIWrapper
load_dotenv()
BASE_URL = os.getenv("BASE_URL")
OKTA_ACCESS_TOKEN = os.getenv("OKTA_ACCESS_TOKEN")
SUBSCRIPTION_KEY = os.getenv("SUBSCRIPTION_KEY")
chat = ChatOpenAI(
api_key=OKTA_ACCESS_TOKEN,
base_url=BASE_URL,
model="GPT_4_1",
default_headers={"Subscription-Key": SUBSCRIPTION_KEY},
)
duckduckgo = DuckDuckGoSearchRun()
python_repl = PythonREPL()
wikipedia = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
def prompt(state: AgentState, config: RunnableConfig) -> list[AnyMessage]:
file_path = config["configurable"].get("file_path", "")
system_msg = (
f"You are an AI assistant evaluated on the GAIA benchmark. "
f"Answer questions using only verified information. "
f"Use the available tools to find answers when needed. "
f"If you do not have enough information, reply: 'I do not have enough information to answer the question.' "
f"Provide only the direct answer, with no extra explanation, formatting, or restating the question. "
f"For example, if asked 'What is the capital of France?', answer 'Paris'. "
)
if file_path:
system_msg += f"If a file is referenced, use the file path: {file_path}."
return [{"role": "system", "content": system_msg}] + state["messages"]
react_agent = create_react_agent(
model=chat,
tools=[duckduckgo, wikipedia],
prompt=prompt,
)