Добавлен client.py
This commit is contained in:
@@ -0,0 +1,120 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user