95 lines
3.8 KiB
Python
95 lines
3.8 KiB
Python
# solution.py
|
||
|
||
"""
|
||
Основной скрипт агента с потоковым выводом.
|
||
Используется стек: LangChain, OpenAI и rich.
|
||
"""
|
||
|
||
from langchain import hub
|
||
from langchain.agents import initialize_agent, AgentType
|
||
from langchain.tools import tool
|
||
from langchain_community.llms.openai import OpenAI
|
||
from rich import print as rprint
|
||
|
||
# ---------- 1. Подключаем LLM ----------
|
||
llm = OpenAI(temperature=0, model="gpt-4o-mini", streaming=True)
|
||
|
||
# ---------- 2. Определяем инструмент ----------
|
||
@tool
|
||
def get_price(product: str, city: str) -> str:
|
||
"""
|
||
Возвращает таблицу цен на указанный продукт в заданном городе.
|
||
"""
|
||
# В реальном приложении здесь будет запрос к базе данных или API.
|
||
# Для демонстрации возвращаем статический ответ.
|
||
if product.lower() == "молоко":
|
||
return "| Продукт | Цена (руб.) | Магазин |\n| Молоко | 89 | Магнит |"
|
||
elif product.lower() == "хлеб":
|
||
return "| Продукт | Цена (руб.) | Магазин |\n| Хлеб | 25 | Пятёрочка |"
|
||
else:
|
||
return f"Информация о {product} в городе {city} недоступна."
|
||
|
||
# ---------- 3. Создаём агента ----------
|
||
tools = [get_price]
|
||
agent = initialize_agent(
|
||
tools,
|
||
llm,
|
||
agent=AgentType.OPENAI_FUNCTIONS,
|
||
verbose=False, # отключаем стандартный вывод
|
||
)
|
||
|
||
# ---------- 4. Функции форматирования ----------
|
||
def format_message(message) -> str:
|
||
"""
|
||
Форматируем сообщение для вывода в консоль.
|
||
Если есть content – возвращаем его,
|
||
иначе формируем строку вызова инструмента.
|
||
"""
|
||
if message.content:
|
||
return message.content
|
||
# В случае tool_calls: берём первый вызов
|
||
call = message.tool_calls[0]
|
||
name = call["name"]
|
||
args = call.get("args", {})
|
||
return f"{name}({args})"
|
||
|
||
def format_chunk_message(chunk):
|
||
"""
|
||
Обрабатываем чанк типа 'messages'.
|
||
Выводим токен без перевода строки, добавляем разделитель при смене шага.
|
||
"""
|
||
message, meta = chunk
|
||
global step
|
||
|
||
if meta.get("langgraph_step") != step:
|
||
step = meta["langgraph_step"]
|
||
rprint("\n --- --- --- \n")
|
||
|
||
content = format_message(message)
|
||
# Выводим без перевода строки, чтобы токены «текли»
|
||
print(content, end="", flush=True)
|
||
|
||
# ---------- 5. Запускаем потоковый вывод ----------
|
||
if __name__ == "__main__":
|
||
# Пример запроса
|
||
user_prompt = "Сколько стоит молоко в Казани? А как насчёт хлеба?"
|
||
|
||
stream = agent.stream(
|
||
{"messages": [{"role": "human", "content": user_prompt}]},
|
||
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" and chunk_data.get("model"):
|
||
# После завершения шага выводим итоговое сообщение модели
|
||
last_msg = chunk_data["model"]["messages"][-1]
|
||
rprint("\n")
|
||
print(format_message(last_msg), flush=True)
|
||
|
||
# Завершающий перевод строки, чтобы консоль оставалась аккуратной
|
||
print() |