From 1b2c1c40ab0597b01b0f758cc59b50790e604664 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=90=D0=B4=D0=B5=D0=BB=D0=B8=D0=BD=D0=B0=20=D0=A1=D0=B0?= =?UTF-8?q?=D1=82=D1=82=D0=B0=D1=80=D0=BE=D0=B2=D0=B0?= Date: Thu, 28 May 2026 10:41:01 +0000 Subject: [PATCH] add agent_with_memory.py --- agent_with_memory.py | 173 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 173 insertions(+) create mode 100644 agent_with_memory.py diff --git a/agent_with_memory.py b/agent_with_memory.py new file mode 100644 index 0000000..39682c3 --- /dev/null +++ b/agent_with_memory.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +""" +Практическое задание №3 – Memory + Confirmation (Rich UI). +Запуск: + python agent_with_memory.py +""" + +from __future__ import annotations + +import json +import uuid +from pathlib import Path +from typing import Any, Dict, Iterable, List, Optional, Tuple + +# ────────────────────── 1. Библиотеки ─────────────────────── +from langchain_openai import ChatOpenAI # пример LLM‑модели +from langgraph.checkpoint.memory import MemorySaver +from langgraph.constants import START +from langgraph.graph import StateGraph +from langgraph.types import Command, interrupt +from rich.console import Console + +# ────────────────────── 2. Настройки и утилиты ─────────────── +console = Console() +OPENAI_API_KEY = "sk-..." # <-- ваш ключ OpenAI (или используйте переменную окружения) +llm = ChatOpenAI(temperature=0, model="gpt-4o-mini", openai_api_key=OPENAI_API_KEY) + +# Пример простого инструмента +def get_price(args: Dict[str, Any]) -> str: + """Возвращает цену (фиктивный ответ).""" + city = args.get("city") + date = args.get("date") + return f"Цена в {city} на {date}: 123₽" + +# ────────────────────── 3. Структура состояния ─────────────── +class State(dict): + """Состояние агента – список сообщений.""" + messages: List[Dict[str, Any]] + + +# ────────────────────── 4. Создание графа (агента) ──────── +def build_agent() -> Tuple[StateGraph, MemorySaver]: + # ├─ хранилище памяти + memory = MemorySaver() + + # └─ граф + builder = StateGraph(State) + + # Узел генерации ответа LLM + @builder.node() + def llm_node(state: State) -> dict: + user_msg = state["messages"][-1] + response = llm.invoke( + [ + {"role": "system", "content": SYSTEM_PROMPT}, + *state["messages"], + ] + ) + # Добавляем системный ответ + new_message = { + "role": "assistant", + "content": response.content, + "tool_calls": response.tool_calls or [], + } + return {"messages": state["messages"] + [new_message]} + + # Узел вызова инструмента + @builder.node() + def tool_node(state: State) -> dict: + # Последнее сообщение содержит tool_calls + last_msg = state["messages"][-1] + tool_call = last_msg["tool_calls"][0] + name, args = tool_call["name"], json.loads(tool_call["arguments"]) + if name == "get_price": + result = get_price(args) + else: + result = f"Unknown tool {name}" + # Добавляем результат как новое сообщение + return { + "messages": state["messages"] + + [ + { + "role": "tool", + "content": result, + "tool_call_id": tool_call["id"], + } + ] + } + + builder.add_node("llm", llm_node) + builder.add_node("tool", tool_node) + + # Переходы + builder.set_entry_point("llm") + builder.add_conditional_edges( + "llm", + lambda x: "tool" if any(m.get("tool_calls") for m in x["messages"]) else END, + {"tool": "tool", END: END}, + ) + builder.add_edge("tool", END) + + # Компилируем с чекпоинтером + graph = builder.compile(checkpointer=memory, interrupt_before=["tools"]) + return graph, memory + + +SYSTEM_PROMPT = """ +Ты — помощник. В процессе разговора можешь вызывать инструмент get_price(city, date). +Перед каждым вызовом инструмента агент должен остановиться и запросить подтверждение у пользователя. +""" + + +# ────────────────────── 5. Функция «разговаривать» ─────────── +def ask_and_run(user_input: Optional[Dict[str, Any]], config: dict): + """Обрабатывает поток от агента, ловит паузы и спрашивает подтверждение.""" + for chunk_type, chunk_data in agent.stream( + user_input, + config=config, + stream_mode=["messages", "updates"], + ): + # ├─ 1. Вывод токенов + if chunk_type == "messages": + console.print(chunk_data["content"], end="", style="green") + continue + + # ├─ 2. Информация о вызове инструмента (если есть) + if chunk_type == "updates" and "tool_calls" in chunk_data: + tool_call = chunk_data["tool_calls"][0] + name, args = tool_call["name"], json.loads(tool_call["arguments"]) + console.print(f"\n[bold cyan]Агент хочет вызвать инструмент:[/bold cyan]") + console.print(f"[yellow]{name}({args})[/yellow]\n") + + # ├─ 3. Пауза перед вызовом инструмента + if "__interrupt__" in chunk_data and state.next == ("tools",): + # Получаем состояние (поскольку после pause state ещё не обновился) + state = agent.get_state(config) + last_msg = state["messages"][-1] + tool_call = last_msg["tool_calls"][0] + name, args = tool_call["name"], json.loads(tool_call["arguments"]) + console.print(f"[bold red]Пауза: запрос на вызов инструмента {name}[/bold red]") + answer = console.input("[green]Разрешить? (Y/n): [/green]").strip().lower() + if answer in ("y", ""): + # Возобновляем выполнение с тем же config + ask_and_run(None, config) + else: + console.print("[red]Отменено[/red]") + break + + # └─ 4. Любой другой чанк – просто выводим + if chunk_type == "updates": + console.print(json.dumps(chunk_data, ensure_ascii=False), style="magenta") + + +# ────────────────────── 6. Основной цикл ───────────────────── +if __name__ == "__main__": + agent, memory = build_agent() + + thread_id = str(uuid.uuid4()) # можно задать свой id (например "chat-1") + config = {"configurable": {"thread_id": thread_id}} + + console.print("[bold blue]Добро пожаловать![/bold blue]") + while True: + user_msg = console.input("\n[bold]Вы:[/bold] ") + if user_msg.lower() in ("exit", "quit"): + break + + # Формируем сообщение от пользователя + payload = {"messages": [{"role": "human", "content": user_msg}]} + ask_and_run(payload, config) + + console.print("[bold blue]До свидания![/bold blue]")