add main.py

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()