96 lines
3.7 KiB
Python
96 lines
3.7 KiB
Python
"""
|
|
AI-агент для планирования списка покупок.
|
|
|
|
Иерархический агент на LangChain + локальная LLM (LM Studio).
|
|
- Главный агент принимает список продуктов
|
|
- Инструмент get_price использует субагент для получения цен
|
|
- Итоговая таблица с ценами и общей стоимостью
|
|
"""
|
|
|
|
import os
|
|
from langchain_core.tools import tool
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.prebuilt import create_react_agent
|
|
|
|
|
|
# Подключение к LLM через OpenAI-совместимый API (LM Studio)
|
|
llm = ChatOpenAI(
|
|
model=os.environ.get("LM_STUDIO_MODEL", "local-model"),
|
|
openai_api_base=os.environ.get("LM_STUDIO_API_BASE", "http://localhost:1234/v1"),
|
|
openai_api_key=os.environ.get("LM_STUDIO_API_KEY", "not-needed"),
|
|
temperature=0,
|
|
)
|
|
|
|
|
|
@tool
|
|
def get_price(product: str, city: str) -> str:
|
|
"""Получить цену продукта в указанном городе. Возвращает таблицу с ценами.
|
|
|
|
Args:
|
|
product: Название продукта
|
|
city: Город для поиска цен
|
|
"""
|
|
# Субагент для поиска цен
|
|
sub_agent = create_react_agent(
|
|
model=llm,
|
|
tools=[],
|
|
prompt=(
|
|
"Ты помощник по поиску цен на продукты. "
|
|
"Сгенерируй реалистичную цену для указанного продукта в указанном городе. "
|
|
"Верни результат строго в формате таблицы:\n"
|
|
"| Продукт | Цена | Магазин |\n"
|
|
"| {product} | {цена} руб | {магазин} |"
|
|
),
|
|
)
|
|
|
|
result = sub_agent.invoke({
|
|
"messages": [{"role": "user", "content": f"Найди цену для: {product} в городе {city}"}]
|
|
})
|
|
|
|
# Извлекаем ответ из последнего сообщения
|
|
last_message = result["messages"][-1]
|
|
return last_message.content if hasattr(last_message, "content") else str(last_message)
|
|
|
|
|
|
def main():
|
|
"""Главный агент — помощник по планированию покупок."""
|
|
|
|
# Главный агент с инструментом get_price
|
|
agent = create_react_agent(
|
|
model=llm,
|
|
tools=[get_price],
|
|
prompt="Ты помощник по планированию покупок. Помоги составить список покупок с ценами.",
|
|
)
|
|
|
|
# Запрос пользователя
|
|
user_request = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
|
|
|
|
print("=" * 60)
|
|
print(" Запрос:", user_request)
|
|
print("=" * 60)
|
|
|
|
# Запуск агента с выводом всех промежуточных сообщений
|
|
result = agent.invoke(
|
|
{"messages": [{"role": "user", "content": user_request}]},
|
|
)
|
|
|
|
# Вывод всех сообщений (промежуточные вызовы + финальный ответ)
|
|
print("\n--- Все сообщения ---\n")
|
|
for i, msg in enumerate(result["messages"]):
|
|
role = msg.__class__.__name__
|
|
content = msg.content if hasattr(msg, "content") else str(msg)
|
|
print(f"[{i}] {role}:")
|
|
print(f" {content}")
|
|
print()
|
|
|
|
# Финальный ответ
|
|
final = result["messages"][-1]
|
|
print("=" * 60)
|
|
print(" Финальный ответ:")
|
|
print("=" * 60)
|
|
print(final.content if hasattr(final, "content") else str(final))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|