import os
from langchain_openai import ChatOpenAI
from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.messages import HumanMessage
from langgraph.prebuilt import create_react_agent

llm = ChatOpenAI(model="gpt-4o-mini")
search_tool = TavilySearchResults(max_results=3)
agent = create_react_agent(llm, tools=[search_tool])

def format_message(message) -> str:
    if message.content:
        return message.content
    return f"{message.tool_calls[0]['name']}({message.tool_calls[0]['args']})"

current_step = 1

def format_chunk_message(chunk):
    global current_step
    message, meta = chunk
    step_num = meta.get("langgraph_step", 0)
    if step_num != current_step:
        current_step = step_num
        print("\n --- --- --- \n")
    if message.content:
        print(message.content, end="", flush=False)

if __name__ == "__main__":
    user_query = input("Введите ваш вопрос: ")
    messages = [HumanMessage(content=user_query)]

    stream = agent.stream(
        {"messages": messages},
        stream_mode=["messages", "updates"]
    )

    for chunk in stream:
        chunk_type, chunk_data = chunk
        if chunk_type == "messages":
            format_chunk_message(chunk_data)
        elif chunk_type == "updates":
            if chunk_data.get("model", None):
                last_msg = chunk_data["model"]["messages"][-1]
                print(format_message(last_msg))

    print()