87 lines
3.9 KiB
Python
87 lines
3.9 KiB
Python
import os
|
|
from typing import Any, Dict, Iterable
|
|
|
|
# Библиотеки для вывода и работы с агентом
|
|
from rich.console import Console
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from langchain_openai import ChatOpenAI
|
|
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
|
|
from langgraph.prebuilt import create_react_agent
|
|
|
|
# Инициализация консоли для красивого вывода
|
|
console = Console()
|
|
|
|
# Настройка модели LLM (потребуется переменная окружения OPENAI_API_KEY)
|
|
llm = ChatOpenAI(temperature=0.7)
|
|
|
|
# Простой инструмент – эхо‑ответ
|
|
from langchain_core.tools import tool
|
|
|
|
@tool
|
|
def echo_tool(text: str) -> str:
|
|
"""Возвращает введённый текст."""
|
|
return text
|
|
|
|
# Создаём агент с памятью и возможностью паузы перед вызовом инструмента
|
|
memory = MemorySaver()
|
|
agent = create_react_agent(
|
|
model=llm,
|
|
tools=[echo_tool],
|
|
prompt="You are a helpful assistant. Use the echo tool when needed.",
|
|
checkpointer=memory, # обязательно для памяти и паузы
|
|
interrupt_before=["tools"], # остановка перед каждым вызовом инструмента
|
|
)
|
|
|
|
# Конфигурация разговора (thread_id)
|
|
config = {"configurable": {"thread_id": "chat-1"}}
|
|
|
|
def ask_and_run(user_input: Dict[str, Any] | None, cfg: Dict[str, Any]) -> None:
|
|
"""
|
|
Отправляем запрос агенту и обрабатываем потоковые ответы.
|
|
Если агент останавливается перед вызовом инструмента – спрашиваем подтверждение.
|
|
"""
|
|
# Запускаем поток
|
|
for chunk in agent.stream(
|
|
user_input,
|
|
config=cfg,
|
|
stream_mode=["messages", "updates"],
|
|
):
|
|
state = agent.get_state(cfg)
|
|
chunk_type, chunk_data = chunk
|
|
|
|
if chunk_type == "messages":
|
|
# Печатаем токены ответа LLM
|
|
for msg in chunk_data:
|
|
if isinstance(msg, AIMessage):
|
|
console.print(msg.content, end="")
|
|
|
|
elif chunk_type == "updates":
|
|
# Выводим информацию о вызовах инструментов (если они произошли)
|
|
for upd in chunk_data:
|
|
console.print(f"[bold cyan]Tool call:{upd}")
|
|
|
|
# Обнаруживаем паузу перед инструментом
|
|
if "__interrupt__" in chunk_data and state.next == ("tools",):
|
|
# Получаем детали последнего вызова инструмента
|
|
tool_call = state.values["messages"][-1].tool_calls[0]
|
|
name, args = tool_call["name"], tool_call["args"]
|
|
console.print(f"\n\n[bold yellow]Пауза: агент хочет вызвать инструмент {name}({args})")
|
|
answer = input("Разрешить? (Y/n): ").strip().lower()
|
|
if answer in ("", "y", "yes"):
|
|
# Возобновляем работу с того места, где остановились
|
|
ask_and_run(None, cfg)
|
|
else:
|
|
console.print("[bold red]Отменено пользователем.")
|
|
break
|
|
|
|
# Основной цикл чата
|
|
if __name__ == "__main__":
|
|
console.print("[bold green]Запуск агента. Введите 'exit' для выхода.\n")
|
|
while True:
|
|
user_text = input("\nВы: ")
|
|
if user_text.lower() == "exit":
|
|
console.print("[bold magenta]До свидания!")
|
|
break
|
|
# Формируем сообщение пользователя
|
|
msg = {"messages": [{"role": "human", "content": user_text}]}
|
|
ask_and_run(msg, config) |