46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
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() |