Практическое задание №3: Память и подтверждение действий: client.py

2026-05-27 14:51:44 +00:00
parent 8b6ad69b70
commit 32632ede1b
1 changed files with 118 additions and 0 deletions
@@ -0,0 +1,118 @@
+# client.py
+
+"""
+Интерактивный клиент для агента, реализованного в agent.py.
+Поддерживает команды:
+    /add <text>   – добавить текст в базу Qdrant
+    /search <query> – выполнить поиск по базе и вывести результаты
+    /quit          – выйти из программы
+"""
+
+import sys
+from typing import Dict
+
+# rich для красивого вывода
+from rich.console import Console
+from rich.table import Table
+from rich.prompt import Prompt
+
+console = Console()
+
+# Импортируем агент, созданный в agent.py
+try:
+    from agent import create_agent, llm  # llm – объект модели Ollama
+except Exception as e:  # pragma: no cover
+    console.print(f"[red]Ошибка при импорте агента: {e}[/red]")
+    sys.exit(1)
+
+# Создаём память и агент с паузой перед инструментом
+from langgraph.checkpoint.memory import MemorySaver
+
+memory = MemorySaver()
+
+agent = create_agent(
+    model=llm,
+    tools=[  # инструменты, которые определены в agent.py
+        "search",
+        "add",
+    ],
+    system_prompt="You are a helpful assistant that can search and add documents.",
+    checkpointer=memory,
+    interrupt_before=["tools"],  # пауза перед вызовом любого инструмента
+)
+
+# Текущий thread_id (разговор)
+thread_id = "cli-thread"
+config = {"configurable": {"thread_id": thread_id}}
+
+
+def ask_and_run(user_input: Dict, config: Dict):
+    """
+    Отправляем запрос агенту и обрабатываем потоковые ответы.
+    При паузе перед инструментом спрашиваем подтверждение у пользователя.
+    """
+    for chunk in agent.stream(
+        user_input,
+        config=config,
+        stream_mode=["messages", "updates"],
+    ):
+        state = agent.get_state(config)
+        chunk_type, chunk_data = chunk
+
+        # Печатаем токены сообщения
+        if chunk_type == "messages":
+            console.print(chunk_data, end="")
+
+        # Печатаем вызовы инструментов
+        if chunk_type == "updates":
+            for update in chunk_data:
+                if isinstance(update, dict) and "tool_calls" in update.get("values", {}):
+                    tool_call = update["values"]["messages"][-1].tool_calls[0]
+                    console.print(f"\n[bold cyan]Агент хочет вызвать утилиту {tool_call['name']}({tool_call['args']})[/bold cyan]")
+                    # Спрашиваем подтверждение
+                    answer = Prompt.ask("Разрешить? (Y/n)", default="y")
+                    if answer.lower().strip() in ("y", "yes"):
+                        # Возобновляем работу с тем же состоянием
+                        ask_and_run(None, config)
+                    else:
+                        console.print("[red]Отменено[/red]")
+                        return
+
+        # Обрабатываем паузу перед инструментом
+        if "__interrupt__" in chunk_data and state.next == ("tools",):
+            # Пауза – уже обработана в блоке выше
+            continue
+
+
+def main():
+    console.print("[bold green]Привет! Введите команду (/add, /search, /quit).[/bold green]")
+    while True:
+        try:
+            user_input = Prompt.ask("\nВы")
+        except KeyboardInterrupt:  # pragma: no cover
+            console.print("\n[red]Выход...[/red]")
+            break
+
+        if not user_input.strip():
+            continue
+
+        if user_input.lower() == "/quit":
+            console.print("[green]До свидания![/green]")
+            break
+
+        if user_input.startswith("/add "):
+            text = user_input[len("/add ") :]
+            ask_and_run(
+                {"messages": [{"role": "human", "content": f"/add {text}"}]}, config
+            )
+        elif user_input.startswith("/search "):
+            query = user_input[len("/search ") :]
+            ask_and_run(
+                {"messages": [{"role": "human", "content": f"/search {query}"}]}, config
+            )
+        else:
+            console.print("[yellow]Неизвестная команда. Попробуйте /add, /search или /quit.[/yellow]")
+
+
+if __name__ == "__main__":
+    main()