from typing import TypedDict, Annotated

from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint, HuggingFaceEmbeddings
from langchain_community.tools import DuckDuckGoSearchRun
from langgraph.graph import StateGraph, START
from langchain_core.messages import AnyMessage, SystemMessage
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langchain_openai import ChatOpenAI

from tools import multiply, divide, add, substract, wiki_search

available_tools = [multiply, divide, add, substract, DuckDuckGoSearchRun(), wiki_search]

with open("system_prompt.txt", "r") as file:
    sys_prompt = SystemMessage(content=file.read())

class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]
    
    
chat = ChatOpenAI(model="o4-mini",)
chat.bind_tools(available_tools)

def assistant(state: AgentState):
    return {
        "messages": [chat.invoke(state["messages"])]
    }

def init_step(state: AgentState):
    return {
        "messages": [sys_prompt]
    }

builder = StateGraph(AgentState)

builder.add_node("assistant", assistant)
builder.add_node("init", init_step)
builder.add_node("tools", ToolNode(available_tools))

builder.add_edge(START, "init")
builder.add_edge("init", "assistant")
builder.add_conditional_edges("assistant", tools_condition)

builder.add_edge("tools", "assistant")

gaia_agent = builder.compile()