add main.py

main.py

@@ -0,0 +1,133 @@
+"""Иерархический AI-агент со stream-выводом (LangChain + LM Studio)."""
+from __future__ import annotations
+
+import os
+
+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")
+
+STEP_SEPARATOR = "\n --- --- --- \n"
+
+
+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:
+    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 run_shopping_assistant_stream() -> None:
+    llm = build_llm()
+    get_price = make_get_price_tool(llm)
+
+    main_agent = create_agent(
+        model=llm,
+        tools=[get_price],
+        system_prompt="Ты помощник по планированию покупок",
+    )
+
+    question = (
+        "Помоги составить список покупок: молоко, хлеб, яблоки. Я нахожусь в Казани."
+    )
+
+    stream = main_agent.stream(
+        {"messages": [{"role": "human", "content": question}]},
+        stream_mode=["messages", "updates"],
+    )
+
+    step = 1
+
+    def format_chunk_message(chunk) -> None:
+        nonlocal step
+        message, meta = chunk
+        graph_step = meta.get("langgraph_step", step)
+        if graph_step != step:
+            step = graph_step
+            print(STEP_SEPARATOR, end="")
+        if message.content:
+            print(message.content, end="", flush=True)
+
+    for chunk in stream:
+        chunk_type, chunk_data = chunk
+
+        if chunk_type == "messages":
+            format_chunk_message(chunk_data)
+
+        if chunk_type == "updates":
+            model_update = chunk_data.get("model")
+            if model_update:
+                last_message = model_update["messages"][-1]
+                formatted = format_message(last_message)
+                if formatted.strip():
+                    print(formatted)
+
+    print()
+
+
+if __name__ == "__main__":
+    run_shopping_assistant_stream()