Add task1_agent.py

This commit is contained in:
2026-05-07 14:45:16 +00:00
commit 97f4eebc7a
+131
View File
@@ -0,0 +1,131 @@
import random
from typing import Literal
from langchain.agents import create_agent
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from pydantic import SecretStr
# ---------------------------------------------------------------------------
# 1. Подключение к локальной модели (LM Studio)
# ---------------------------------------------------------------------------
llm = ChatOpenAI(
model="<название_модели>", # замените на имя модели, запущенной в LM Studio
base_url="http://localhost:1234/v1",
api_key=SecretStr("fake"),
temperature=0.7,
)
# ---------------------------------------------------------------------------
# 2. Инструмент с суб‑агентом: get_price
# ---------------------------------------------------------------------------
@tool
def get_price(product: str, city: str) -> str:
"""Возвращает таблицу с примерной ценой продукта в указанном городе.
Суб‑агент генерирует цену и форматирует её в виде Markdown‑таблицы.
"""
# ------------------- суб‑агент -------------------
# Простой суб‑агент, использующий тот же LLM, но с более специфическим
# промптом, чтобы сгенерировать реалистичную цену.
sub_prompt = (
f"Ты экономический аналитик. Дай примерную цену продукта '{product}' в городе '{city}'. "
"Ответ дай в виде строки цены (число) и названия магазина, без лишних символов. "
"Если нет данных – придумай разумную цену в диапазоне 30‑200 рублей."
)
sub_agent = create_agent(
model=llm,
tools=[], # суб‑агент не использует внешних инструментов
system_prompt=sub_prompt,
)
# Запрос к суб‑агенту
sub_answer = sub_agent.invoke({"messages": [{"role": "human", "content": "Цена?"}]})
# Получаем текст ответа
price_text = sub_answer["messages"][-1]["content"].strip()
# Пытаемся извлечь число из ответа; если не получится – генерируем случайно
try:
# ищем первое число в строке
import re
match = re.search(r"\d+", price_text)
price = int(match.group()) if match else random.randint(30, 200)
except Exception:
price = random.randint(30, 200)
# Выбираем магазин произвольно
stores = ["Магнит", "Пятёрочка", "Перекрёсток", "Ашан", "Дикси"]
store = random.choice(stores)
# Формируем таблицу Markdown
table = (
"| Продукт | Цена (руб.) | Магазин |\n"
"|---------|-------------|---------|\n"
f"| {product.capitalize()} | {price} | {store} |"
)
return table
# ---------------------------------------------------------------------------
# 3. Главный агент
# ---------------------------------------------------------------------------
main_agent = create_agent(
model=llm,
tools=[get_price],
system_prompt="Ты помощник по планированию покупок. Используй инструмент get_price, чтобы узнать цены, и подсчитай итоговую стоимость.",
)
# ---------------------------------------------------------------------------
# 4. Утилита для подсчёта итоговой стоимости из полученных таблиц
# ---------------------------------------------------------------------------
def parse_price(table: str) -> int:
"""Извлекает цену из строки таблицы Markdown."""
import re
match = re.search(r"\|\s*[^|]+\|\s*(\d+)\s*\|", table)
return int(match.group(1)) if match else 0
# ---------------------------------------------------------------------------
# 5. Основная функция: принимает запрос пользователя и выводит цепочку сообщений
# ---------------------------------------------------------------------------
def run_shopping_assistant(user_query: str):
"""Запускает главного агента с пользовательским запросом и выводит все шаги."""
# Запускаем агент в режиме stream, чтобы увидеть промежуточные вызовы
stream = main_agent.stream({"messages": [{"role": "human", "content": user_query}]},
stream_mode=["messages", "updates"])
step = 0
for chunk_type, chunk_data in stream:
if chunk_type == "messages":
message, meta = chunk_data
if meta.get("langgraph_step") != step:
step = meta.get("langgraph_step")
print("\n--- step", step, "---\n")
if message.get("content"):
print(message["content"], end="")
elif chunk_type == "updates":
# При завершении шага модели мы получаем полное сообщение
if "model" in chunk_data:
last_msg = chunk_data["model"]["messages"][-1]
if last_msg.get("content"):
print(last_msg["content"], end="")
elif last_msg.get("tool_calls"):
# Выводим вызов инструмента
tc = last_msg["tool_calls"][0]
print(f"{tc['name']}({tc['args']})")
print("\n\n--- Финальный вывод ---\n")
# После завершения получаем полную историю
final = main_agent.invoke({"messages": [{"role": "human", "content": user_query}]})
# Вычислим итоговую стоимость, если в ответе есть таблицы
total = 0
for msg in final["messages"]:
if msg.get("content") and "|" in msg["content"]:
total += parse_price(msg["content"])
if total:
print(f"Итоговая стоимость: {total} руб.")
# ---------------------------------------------------------------------------
# 6. Тестовый запуск (можно закомментировать при импортировании)
# ---------------------------------------------------------------------------
if __name__ == "__main__":
query = "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
run_shopping_assistant(query)