graph.py

"""LangGraph Agent"""

import os
from langchain_core.messages import HumanMessage, SystemMessage

from langchain_openai import ChatOpenAI
from langgraph.graph import START, MessagesState, StateGraph, END
from langgraph.prebuilt import ToolNode, tools_condition

from tools import TOOLS


system_msg = SystemMessage(content="""You are a helpful assistant tasked with answering questions using a set of tools. 

Now, I will ask you a question. Finish your answer with the following template: 
FINAL ANSWER: [YOUR FINAL ANSWER]. 
YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string.
Your answer should only start with "FINAL ANSWER: ", then follows with the answer. 
Think internally, but do not output any intermediate reasoning or explanation.  
Only output the final answer in the correct format and nothing else. """)


def build_graph():
    """Build the graph"""
    # Load environment variables from .env file

    llm = ChatOpenAI(
        model="gpt-4o",
        api_key=os.environ["OPENAI_API_KEY"],
        temperature=0,
    )
    # Bind tools to LLM
    llm_with_tools = llm.bind_tools(TOOLS)

    # Node
    def assistant(state: MessagesState):
        """Assistant node"""
        return {"messages": [llm_with_tools.invoke(state["messages"])]}
    

    builder = StateGraph(MessagesState)
    builder.add_node("assistant", assistant)
    builder.add_node("tools", ToolNode(TOOLS))
    builder.add_edge(START, "assistant")
    builder.add_conditional_edges(
        "assistant",
        tools_condition
    )
    builder.add_edge("tools", "assistant")
    builder.add_edge("assistant", END)


    # Compile graph
    return builder.compile()