Add task1_agent.py
This commit is contained in:
+131
@@ -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)
|
||||
Reference in New Issue
Block a user