80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
# 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) |