From 23c1e753d8ff08fdf45da7f877610aea8581373f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=AD=D0=BC=D0=B8=D0=BB=D1=8C=20=D0=90=D0=BC=D0=B8=D1=80?= =?UTF-8?q?=D0=BE=D0=B2?= Date: Mon, 25 May 2026 21:32:07 +0000 Subject: [PATCH] add main.py --- main.py | 125 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 125 insertions(+) create mode 100644 main.py diff --git a/main.py b/main.py new file mode 100644 index 0000000..b06345f --- /dev/null +++ b/main.py @@ -0,0 +1,125 @@ +"""Иерархический AI-агент для планирования списка покупок (LangChain + LM Studio).""" +from __future__ import annotations + +import os +import re + +from langchain.agents import create_agent +from langchain.tools import tool +from langchain_openai import ChatOpenAI +from pydantic import SecretStr + +LM_STUDIO_BASE_URL = os.getenv("LM_STUDIO_BASE_URL", "http://localhost:1234/v1") +LM_STUDIO_MODEL = os.getenv("LM_STUDIO_MODEL", "local-model") + + +def build_llm() -> ChatOpenAI: + return ChatOpenAI( + model=LM_STUDIO_MODEL, + base_url=LM_STUDIO_BASE_URL, + api_key=SecretStr(os.getenv("OPENAI_API_KEY", "fake")), + temperature=0.7, + ) + + +def _extract_table(text: str) -> str: + """Оставляет markdown-таблицу из ответа субагента.""" + lines = [line for line in text.splitlines() if "|" in line] + if lines: + return "\n".join(lines) + return text.strip() + + +def _build_price_subagent(llm: ChatOpenAI): + return create_agent( + model=llm, + system_prompt=( + "Ты аналитик цен на продукты питания. " + "По названию продукта и городу оцени реалистичную цену в рублях, " + "опираясь на типичные российские розничные цены. " + "Ответь ТОЛЬКО одной строкой markdown-таблицы в формате:\n" + "| Продукт | Цена (руб.) | Магазин |" + ), + ) + + +def make_get_price_tool(llm: ChatOpenAI): + price_subagent = _build_price_subagent(llm) + + @tool + def get_price(product: str, city: str) -> str: + """Возвращает примерную цену продукта в указанном городе. + + Args: + product: название продукта (молоко, хлеб, яблоки и т.д.) + city: город покупателя + """ + prompt = ( + f"Город: {city}. Продукт: {product}. " + "Верни одну строку таблицы с реалистичной ценой и названием магазина." + ) + result = price_subagent.invoke( + {"messages": [{"role": "human", "content": prompt}]} + ) + last = result["messages"][-1] + content = getattr(last, "content", str(last)) + return _extract_table(content) + + return get_price + + +def format_message(message) -> str: + content = getattr(message, "content", None) + if content: + return str(content) + tool_calls = getattr(message, "tool_calls", None) or [] + if tool_calls: + call = tool_calls[0] + name = call.get("name") if isinstance(call, dict) else getattr(call, "name", "") + args = call.get("args") if isinstance(call, dict) else getattr(call, "args", {}) + return f"{name}({args})" + return str(message) + + +def _parse_price_row(row: str) -> int | None: + cells = [c.strip() for c in row.strip("|").split("|") if c.strip()] + if len(cells) < 2: + return None + price_raw = cells[1] + numbers = re.findall(r"\d+", price_raw.replace(" ", "")) + if not numbers: + return None + return int(numbers[0]) + + +def run_shopping_assistant() -> None: + llm = build_llm() + get_price = make_get_price_tool(llm) + + main_agent = create_agent( + model=llm, + tools=[get_price], + system_prompt="Ты помощник по планированию покупок", + ) + + question = ( + "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани." + ) + answer = main_agent.invoke( + {"messages": [{"role": "human", "content": question}]} + ) + + print("--- Цепочка сообщений ---") + for message in answer["messages"]: + formatted = format_message(message) + if formatted.strip(): + print("---") + print(formatted) + + print("\n--- Итог ---") + final = answer["messages"][-1] + print(getattr(final, "content", final)) + + +if __name__ == "__main__": + run_shopping_assistant()