diff --git a/solutions/699cc158d6d3a5544a3ed35b_Stream-режим_AI-агента/agent.py b/solutions/699cc158d6d3a5544a3ed35b_Stream-режим_AI-агента/agent.py new file mode 100644 index 0000000..7aa8940 --- /dev/null +++ b/solutions/699cc158d6d3a5544a3ed35b_Stream-режим_AI-агента/agent.py @@ -0,0 +1,125 @@ +import os +from typing import Dict, Any + +from langgraph.graph import StateGraph, START +from langgraph.prebuilt.tool_executor import ToolExecutorNode +from langgraph.checkpoint.memory import MemorySaver +from langchain_ollama import ChatOllama +from nomic_embed_text import NomicEmbedText +from langchain.tools import BaseTool +from langchain.schema import HumanMessage, AIMessage, SystemMessage + +# ---------- Настройки LLM и эмбеддингов ---------- +LLM_MODEL = os.getenv("OLLAMA_MODEL", "llama3") +EMBEDDING_MODEL = os.getenv("NOMIC_EMBEDDING_MODEL", "nomic-embed-text") + +llm = ChatOllama(model=LLM_MODEL, temperature=0.2) +embedding = NomicEmbedText() + +# ---------- Пример инструмента ---------- +class DummyTool(BaseTool): + name: str = "get_price" + description: str = ( + "Получает цену товара в указанном городе. " + "Аргументы: {'product': 'название', 'city': 'город'}." + ) + + def _run(self, product: str, city: str) -> str: + # В реальном коде здесь будет запрос к API + return f"Цена {product} в {city}: 100₽" + +tool = DummyTool() +tools = [tool] + +# ---------- Создание агента ---------- +def create_agent() -> StateGraph: + graph = StateGraph() + + def agent(state: Dict[str, Any]) -> Dict[str, Any]: + messages = state.get("messages", []) + # Добавляем системное сообщение с инструкцией + system_msg = SystemMessage( + content="Ты помощник. Используй инструменты при необходимости." + ) + all_messages = [system_msg] + messages + + # Запускаем LLM в режиме stream + response = llm.invoke(all_messages, stream=True) + return {"messages": [response]} + + graph.add_node("agent", agent) + + def tool_executor(state: Dict[str, Any]) -> Dict[str, Any]: + last_msg = state["messages"][-1] + if hasattr(last_msg, "tool_calls") and last_msg.tool_calls: + # Выполняем первый вызов инструмента + call = last_msg.tool_calls[0] + name = call["name"] + args = eval(call["args"]) + result = tools_dict[name].invoke(args) + # Создаём сообщение с результатом + tool_msg = AIMessage( + content=result, + tool_calls=[call], + ) + return {"messages": [tool_msg]} + return {} + + graph.add_node("tool_executor", ToolExecutorNode(tools)) + graph.set_entry_point("agent") + graph.add_edge(START, "agent") + graph.add_conditional_edges( + "agent", + lambda x: "messages" in x and hasattr(x["messages"][-1], "tool_calls"), + { + True: "tool_executor", + False: START, + }, + ) + graph.set_finish_node("agent") + + return graph.compile(checkpointer=MemorySaver()) + +tools_dict = {t.name: t for t in tools} +graph = create_agent() + +# ---------- Функции форматирования ---------- +def format_message(message) -> str: + if message.content: + return message.content + # Если нет content, выводим вызов инструмента + call = message.tool_calls[0] + return f"{call['name']}({call['args']})" + +step = 1 + +def format_chunk_message(chunk): + global step + message, meta = chunk + if meta.get("langgraph_step") != step: + step = meta["langgraph_step"] + print("\n --- --- --- \n") + if message.content: + print(message.content, end="", flush=True) + +# ---------- Запуск агента в режиме stream ---------- +def run_agent(user_input: str): + global step + step = 1 + # Инициализируем состояние с пользовательским сообщением + init_state = {"messages": [HumanMessage(content=user_input)]} + stream = graph.stream(init_state, stream_mode=["messages", "updates"]) + + for chunk_type, chunk_data in stream: + if chunk_type == "messages": + format_chunk_message(chunk_data) + elif chunk_type == "updates": + # При завершении шага выводим итоговое сообщение + model_info = chunk_data.get("model") + if model_info: + last_msg = model_info["messages"][-1] + print("\n" + format_message(last_msg) + "\n") + +if __name__ == "__main__": + user_query = input("Введите запрос: ") + run_agent(user_query) \ No newline at end of file