from typing import Annotated, Sequence, TypedDict

from dotenv import load_dotenv

from langchain_core.messages import BaseMessage

from langchain_core.messages import ToolMessage

from langchain_core.messages import SystemMessage

from langchain_openai import ChatOpenAI

from langchain_core.tools import tool

from langgraph.graph.message import add_messages

from langgraph.graph import StateGraph, START, END

from langgraph.prebuilt import ToolNode

load_dotenv()

class AgentState(TypedDict):

messages: Annotated[Sequence[BaseMessage], add_messages]

@tool

def add(a: int, b: int) -> int:

"""Addition function to add two numbers."""

return a + b

@tool

def subtract(a: int, b: int) -> int:

"""Subtraction function to subtract two numbers."""

return a - b

@tool

def multiply(a: int, b:int) -> int:

"""This function is used to multiply two nubers.

If you want to calculate square of a number pass the same value in two arguements"""

return a * b

@tool

def division(a: int, b: int) -> int:

"""This function is used to calculate the division"""

return a / b

tools = [add, subtract, multiply, division]

model = ChatOpenAI(model="gpt-4o").bind_tools(tools)

def model_call(state: AgentState) -> AgentState:

system_message = SystemMessage(content="You are my AI assistant help me choose the proper tool to the best of your ability.")

result = model.invoke([system_message]+state["messages"])

return {"messages": [result]}

def should_continue(state: AgentState) -> bool:

messages = state["messages"]

last_message = messages[-1]

if not last_message.tool_calls:

return "end"

return "continue"

graph = StateGraph(AgentState)

graph.add_node("our_agent", model_call)

tool_node = ToolNode(tools=tools)

graph.add_node("tools", tool_node)

graph.add_edge(START, "our_agent")

graph.add_conditional_edges("our_agent",

should_continue,

{"continue": "tools",

"end": END

})

app = graph.compile()

def print_stream(stream):

for s in stream:

message = s["messages"][-1]

if(isinstance(message, tuple)):

print(message)

else:

message.pretty_print()

inputs = {"messages": [("user", "Hello How are you can you add two numbers 5, 6 and calculate square of the result")]}

print_stream(app.stream(inputs, stream_mode="values"))

inputs = {"messages": [("user", "Hello How are you can you add two numbers 0, 0 and calculate division of the result by same number")]}

print_stream(app.stream(inputs, stream_mode="values"))