Stream-режим AI-агента: agent.py

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)