diff --git a/solutions/699cc158d6d3a5544a3ed35b_Stream_режим_AI_агента/stream_agent.py b/solutions/699cc158d6d3a5544a3ed35b_Stream_режим_AI_агента/stream_agent.py index 433a501..f93b6b2 100644 --- a/solutions/699cc158d6d3a5544a3ed35b_Stream_режим_AI_агента/stream_agent.py +++ b/solutions/699cc158d6d3a5544a3ed35b_Stream_режим_AI_агента/stream_agent.py @@ -1,106 +1,46 @@ import os -from typing import Dict, Tuple - +from langchain_openai import ChatOpenAI from langchain_community.tools.tavily_search import TavilySearchResults -from langchain_core.messages import HumanMessage, AIMessage -from langgraph.graph import StateGraph -from langgraph.prebuilt import create_agent_executor -from langgraph.schema import MessagesState -from langgraph.utils import format_messages -from openai import OpenAI +from langchain_core.messages import HumanMessage +from langgraph.prebuilt import create_react_agent -# ------------------------------------------------------------------ -# 1. Подключаемся к LLM (OpenAI GPT‑4o-mini) -# ------------------------------------------------------------------ -client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) -llm = client.chat.completions.create - -# ------------------------------------------------------------------ -# 2. Определяем инструмент поиска -# ------------------------------------------------------------------ +llm = ChatOpenAI(model="gpt-4o-mini") search_tool = TavilySearchResults(max_results=3) +agent = create_react_agent(llm, tools=[search_tool]) -# ------------------------------------------------------------------ -# 3. Создаём агент (используем готовый LangGraph‑агент) -# ------------------------------------------------------------------ -def agent_executor(messages: list[HumanMessage]) -> AIMessage: - """Вызов агента с использованием LangGraph.""" - # Создаём простую схему, где агент может вызвать инструмент поиска - graph = StateGraph(MessagesState) - - @graph.node - def start(state: MessagesState) -> MessagesState: - return state - - @graph.node - def tool(state: MessagesState) -> MessagesState: - last_msg = state.messages[-1] - if isinstance(last_msg, AIMessage) and last_msg.tool_calls: - # вызываем инструмент - tool_name = last_msg.tool_calls[0]["name"] - args = last_msg.tool_calls[0]["args"] - result = search_tool.run(args) - new_message = AIMessage(content=result) - state.messages.append(new_message) - return state - - graph.add_edge("start", "tool") - graph.set_entry_point("start") - - # Запускаем граф - final_state = graph.invoke({"messages": messages}) - return final_state["messages"][-1] - -# ------------------------------------------------------------------ -# 4. Функции форматирования сообщений -# ------------------------------------------------------------------ -def format_message(message: AIMessage) -> str: - """Возвращает строку для печати из сообщения.""" +def format_message(message) -> str: if message.content: return message.content - # Если сообщение содержит вызов инструмента - tool_call = message.tool_calls[0] - name = tool_call["name"] - args = tool_call["args"] - return f"{name}({args})" + return f"{message.tool_calls[0]['name']}({message.tool_calls[0]['args']})" -def format_chunk_message(chunk: Tuple) -> None: - """Обрабатывает чанк типа 'messages'.""" +current_step = 1 + +def format_chunk_message(chunk): global current_step - message, meta = chunk # type: ignore[assignment] + 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=True) + print(message.content, end="", flush=False) -# ------------------------------------------------------------------ -# 5. Запускаем потоковый вывод -# ------------------------------------------------------------------ if __name__ == "__main__": - # Вводим запрос от пользователя user_query = input("Введите ваш вопрос: ") - - # Инициализируем состояние сообщений messages = [HumanMessage(content=user_query)] - # Создаём итератор потока (используем готовый LangGraph executor) - stream = create_agent_executor( - llm=llm, - tools=[search_tool], - agent_name="stream_agent", - stream_mode=["messages", "updates"], - ).stream({"messages": messages}) + stream = agent.stream( + {"messages": messages}, + stream_mode=["messages", "updates"] + ) - current_step = 1 - - for chunk_type, chunk_data in stream: + for chunk in stream: + chunk_type, chunk_data = chunk if chunk_type == "messages": format_chunk_message(chunk_data) elif chunk_type == "updates": - # Обрабатываем события обновления (например, завершение шага) - if isinstance(chunk_data, dict) and "model" in chunk_data: + if chunk_data.get("model", None): last_msg = chunk_data["model"]["messages"][-1] - print("\n" + format_message(last_msg)) - print() # завершаем вывод новой строкой \ No newline at end of file + print(format_message(last_msg)) + + print() \ No newline at end of file