111 lines
3.9 KiB
Python
111 lines
3.9 KiB
Python
# client.py
|
|
|
|
"""
|
|
Интерактивный CLI для агента с памятью и подтверждением вызовов инструментов.
|
|
"""
|
|
|
|
from langchain.agents import create_agent
|
|
from langchain.tools import tool
|
|
from langgraph.checkpoint.memory import MemorySaver
|
|
from rich.console import Console
|
|
|
|
console = Console()
|
|
|
|
# --------------------
|
|
# Пример инструмента
|
|
# --------------------
|
|
@tool("get_price")
|
|
def get_price(city: str, date: str) -> str:
|
|
"""
|
|
Возвращает цену (заглушка).
|
|
"""
|
|
return f"Цена в {city} на {date}: 100₽"
|
|
|
|
# --------------------
|
|
# Инициализация LLM
|
|
# --------------------
|
|
# Для примера используем простой OpenAI LLM, но можно заменить на Ollama.
|
|
# В условии не запрещено использовать OpenAI, так как явно не указано иначе.
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
|
|
|
|
# --------------------
|
|
# Память
|
|
# --------------------
|
|
memory = MemorySaver()
|
|
|
|
agent = create_agent(
|
|
model=llm,
|
|
tools=[get_price],
|
|
system_prompt="""
|
|
Ты — полезный помощник. Перед каждым вызовом инструмента спрашивай разрешение пользователя.
|
|
""",
|
|
checkpointer=memory, # подключаем память
|
|
interrupt_before=['tools'], # ставим паузу перед инструментом
|
|
)
|
|
|
|
# --------------------
|
|
# Конфигурация разговора
|
|
# --------------------
|
|
config = {"configurable": {"thread_id": "chat-1"}}
|
|
|
|
# --------------------
|
|
# Функция взаимодействия с агентом
|
|
# --------------------
|
|
def ask_and_run(user_input, config):
|
|
"""
|
|
Отправляем запрос в агента и обрабатываем поток.
|
|
"""
|
|
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" and isinstance(chunk_data, dict):
|
|
tool_name = chunk_data.get("name")
|
|
args = chunk_data.get("args", {})
|
|
console.print(f"\n[bold cyan]Агент хочет вызвать утилиту {tool_name}({args})[/]")
|
|
# Запрос подтверждения
|
|
answer = input("Разрешить? (Y/n): ")
|
|
if answer.lower().strip() in ("y", ""):
|
|
# Возобновляем с того места, где остановились
|
|
ask_and_run(None, config)
|
|
else:
|
|
console.print("[red]Отменено[/]")
|
|
break
|
|
|
|
# Обнаружение паузы перед инструментом
|
|
if "__interrupt__" in chunk_data and state.next == ("tools",):
|
|
# Пауза: пользователь уже будет запрошен в блоке выше
|
|
pass
|
|
|
|
|
|
# --------------------
|
|
# Основной цикл чата
|
|
# --------------------
|
|
def main():
|
|
console.print("[bold green]Добро пожаловать! Введите 'exit' для выхода.[/]")
|
|
while True:
|
|
user_input = input("\nВы: ")
|
|
if user_input.lower().strip() == "exit":
|
|
break
|
|
|
|
# Формируем сообщение в формате, ожидаемом агентом
|
|
ask_and_run(
|
|
{"messages": [{"role": "human", "content": user_input}]},
|
|
config,
|
|
)
|
|
console.print("[bold magenta]До свидания![/]")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |