diff --git a/client.py b/client.py new file mode 100644 index 0000000..d44ab44 --- /dev/null +++ b/client.py @@ -0,0 +1,111 @@ +# 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() \ No newline at end of file