# agent.py

"""
Простой пример агента на LangChain с поддержкой потокового вывода.
Используем только указанный стек:
- create_agent из langchain.agents
- @tool из langchain.tools
- rich.print как rprint
"""

from langchain.agents import create_agent
from langchain.tools import tool
from rich import print as rprint


# ------------------------------
# Определяем инструменты (tools)
# ------------------------------

@tool
def echo(text: str) -> str:
    """
    Простейший инструмент, который просто возвращает переданный текст.
    """
    return f"Echo: {text}"


@tool
def add(a: int, b: int) -> int:
    """
    Возвращает сумму двух целых чисел.
    """
    return a + b


# ------------------------------
# Создаём агента с инструментами
# ------------------------------

agent = create_agent(
    tools=[echo, add],
    # Можно задать любой шаблон диалога; здесь используем простой
    system_message="Ты помощник, который может отвечать и выполнять инструменты.",
)

# ------------------------------------
# Функция запуска агента в stream‑режиме
# ------------------------------------

def run_agent(prompt: str):
    """
    Запускает агента с заданным запросом и выводит результат токен за токеном.
    """
    # Стартуем потоковый вызов
    stream = agent.stream(
        {"messages": [{"role": "human", "content": prompt}]},
        stream_mode=["messages", "updates"],
    )

    # Перебираем события из потока и печатаем их в реальном времени
    for event in stream:
        if isinstance(event, dict):
            # В режиме 'updates' событие может быть словарём с информацией о шаге
            rprint(f"[bold cyan]Update:[/bold cyan] {event}")
        else:
            # В режиме 'messages' это строка‑токен
            rprint(event, end="", flush=True)

    rprint("\n[green]--- Конец ответа ---[/green]")


# ------------------------------
# Пример использования (можно удалить)
# ------------------------------

if __name__ == "__main__":
    sample_prompt = (
        "Привет! Скажи мне сумму 7 и 5, а потом эхо 'Hello World'."
    )
    run_agent(sample_prompt)