# client.py
"""
Интерактивный CLI для LangChain‑агента, работающего в режиме stream.
"""

from langchain.agents import create_agent          # обязательный стек
from langchain.tools import tool                   # обязательный стек
from rich import print as rprint                  # обязательный стек

# ------------------------------------------------------------------
# Определяем простейший инструмент (для демонстрации)
# ------------------------------------------------------------------
@tool("echo")
def echo_tool(input: str) -> str:
    """
    Возвращает введённую строку без изменений.
    """
    return input


# ------------------------------------------------------------------
# Создаём агента с одним инструментом
# ------------------------------------------------------------------
agent = create_agent(
    tools=[echo_tool],          # список доступных инструментов
    verbose=True,               # включаем вывод логов (необязательно)
)

# ------------------------------------------------------------------
# Основной цикл CLI
# ------------------------------------------------------------------
def main() -> None:
    """
    Запускает бесконечный цикл чтения пользовательского ввода и
    передаёт его агенту в режиме stream. Ответы печатаются токен за токеном.
    """
    rprint("[bold green]LangChain Stream‑Agent CLI[/bold green]")
    rprint("Введите сообщение (Ctrl+C для выхода):")

    try:
        while True:
            # Читаем ввод пользователя
            user_input = input("> ").strip()
            if not user_input:
                continue

            # Формируем запрос к агенту
            stream = agent.stream(
                {"messages": [{"role": "human", "content": user_input}]},
                stream_mode=["messages", "updates"],
            )

            # Перебираем события из потока и выводим их
            for event in stream:
                # Событие может быть либо токеном сообщения, либо обновлением состояния
                if isinstance(event, dict) and "type" in event:
                    # Обновление (например, вызов инструмента)
                    rprint(f"[yellow]Update: {event}[/yellow]")
                else:
                    # Токен текста
                    rprint(event, end="", flush=True)

            # После завершения потока выводим перевод строки
            rprint("\n")

    except KeyboardInterrupt:
        rprint("\n[bold red]Выход из программы.[/bold red]")


if __name__ == "__main__":
    main()