120 lines
5.0 KiB
Python
120 lines
5.0 KiB
Python
# client.py
|
||
|
||
"""
|
||
Интерактивный CLI для LangChain‑агента с потоковым выводом.
|
||
Используем только стек из условия:
|
||
- langchain.agents.create_agent
|
||
- langchain.tools.tool
|
||
- langchain_ollama.ChatOllama
|
||
- rich.print (rprint)
|
||
"""
|
||
|
||
from langchain.agents import create_agent
|
||
from langchain.tools import tool
|
||
from langchain_ollama import ChatOllama
|
||
from rich import print as rprint
|
||
|
||
# ------------------------------------------------------------------
|
||
# 1. Определяем LLM‑обёртку – Ollama
|
||
# ------------------------------------------------------------------
|
||
llm = ChatOllama(model="llama3") # модель, которую можно заменить
|
||
|
||
# ------------------------------------------------------------------
|
||
# 2. Создаём инструмент get_price
|
||
# В реальном проекте он бы вызывал подагент или API,
|
||
# здесь просто возвращаем фиктивную таблицу.
|
||
# ------------------------------------------------------------------
|
||
@tool("get_price", "Получить цену товара в указанном городе")
|
||
def get_price(product: str, city: str) -> str:
|
||
"""
|
||
Возвращает таблицу с ценой продукта в заданном городе.
|
||
"""
|
||
# В реальном коде здесь могла бы быть логика обращения к API
|
||
return f"| Продукт | Цена (руб.) | Магазин |\n| {product} | 89 | Магнит |"
|
||
|
||
# ------------------------------------------------------------------
|
||
# 3. Создаём агента, который использует LLM и инструмент get_price
|
||
# ------------------------------------------------------------------
|
||
agent = create_agent(
|
||
llm=llm,
|
||
tools=[get_price],
|
||
verbose=False, # будем выводить вручную
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 4. Функции форматирования чанков
|
||
# ------------------------------------------------------------------
|
||
def format_message(message) -> str:
|
||
"""
|
||
Форматируем сообщение для вывода.
|
||
Если в сообщении есть content – возвращаем его,
|
||
иначе (tool call) формируем строку вызова инструмента.
|
||
"""
|
||
if message.content:
|
||
return message.content
|
||
# Предполагаем, что tool_calls содержит один элемент
|
||
tc = message.tool_calls[0]
|
||
name = tc["name"]
|
||
args = tc.get("args", {})
|
||
return f"{name}({args})"
|
||
|
||
def format_chunk_message(chunk):
|
||
"""
|
||
Выводим токен текста в режиме stream.
|
||
При смене шага выводим разделитель.
|
||
"""
|
||
global step
|
||
message, meta = chunk
|
||
if meta["langgraph_step"] != step:
|
||
step = meta["langgraph_step"]
|
||
rprint("\n--- --- ---\n")
|
||
if message.content:
|
||
# flush=False – чтобы не добавлять перевод строки после каждого токена
|
||
print(message.content, end="", flush=True)
|
||
|
||
# ------------------------------------------------------------------
|
||
# 5. Основной цикл CLI
|
||
# ------------------------------------------------------------------
|
||
def main():
|
||
"""
|
||
Запускаем интерактивный режим: пользователь вводит запрос,
|
||
агент отвечает в режиме stream.
|
||
"""
|
||
rprint("[bold cyan]LangChain Stream‑Agent[/bold cyan]")
|
||
rprint("Введите ваш вопрос (Ctrl+C для выхода):")
|
||
while True:
|
||
try:
|
||
user_input = input("\n> ")
|
||
except KeyboardInterrupt:
|
||
rprint("\n[red]Выход из программы.[/red]")
|
||
break
|
||
|
||
if not user_input.strip():
|
||
continue
|
||
|
||
# Запускаем потоковый вывод
|
||
stream = agent.stream(
|
||
{"messages": [{"role": "human", "content": user_input}]},
|
||
stream_mode=["messages", "updates"],
|
||
)
|
||
|
||
step = 1 # номер текущего шага агента
|
||
|
||
for chunk in stream:
|
||
chunk_type, chunk_data = chunk
|
||
|
||
if chunk_type == "messages":
|
||
format_chunk_message(chunk_data)
|
||
|
||
elif chunk_type == "updates":
|
||
# При завершении шага выводим итоговое сообщение
|
||
model_update = chunk_data.get("model")
|
||
if model_update:
|
||
last_msg = model_update["messages"][-1]
|
||
rprint("\n" + format_message(last_msg))
|
||
|
||
# После завершения потока добавляем перевод строки для следующего запроса
|
||
print()
|
||
|
||
if __name__ == "__main__":
|
||
main() |