add main.py

main.py

@@ -0,0 +1,180 @@
+"""Агент с памятью разговора и подтверждением вызовов инструментов."""
+from __future__ import annotations
+
+import os
+
+from langchain.agents import create_agent
+from langchain.tools import tool
+from langchain_openai import ChatOpenAI
+from langgraph.checkpoint.memory import MemorySaver
+from pydantic import SecretStr
+from rich.console import Console
+
+console = Console()
+
+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"
+THREAD_ID = "разговор-1"
+
+
+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]
+    return "\n".join(lines) if lines else text.strip()
+
+
+def _build_price_subagent(llm: ChatOpenAI):
+    return create_agent(
+        model=llm,
+        system_prompt=(
+            "Ты аналитик цен. По продукту, городу и при необходимости дате "
+            "оцени реалистичную цену в рублях. "
+            "Ответь одной строкой markdown-таблицы: | Продукт | Цена (руб.) | Магазин |"
+        ),
+    )
+
+
+def make_get_price_tool(llm: ChatOpenAI):
+    price_subagent = _build_price_subagent(llm)
+
+    @tool
+    def get_price(product: str, city: str, date: str = "") -> str:
+        """Возвращает примерную цену продукта в городе (опционально с учётом даты/периода).
+
+        Args:
+            product: название продукта или категория (молоко, хлеб, погода как запрос цен и т.д.)
+            city: город
+            date: период (сегодня, завтра, через неделю) — необязательно
+        """
+        when = f", период: {date}" if date else ""
+        prompt = (
+            f"Город: {city}. Продукт/запрос: {product}{when}. "
+            "Верни одну строку таблицы с ценой и магазином."
+        )
+        result = price_subagent.invoke(
+            {"messages": [{"role": "human", "content": prompt}]}
+        )
+        last = result["messages"][-1]
+        return _extract_table(getattr(last, "content", str(last)))
+
+    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 _tool_call_repr(state) -> str:
+    messages = state.values.get("messages", [])
+    if not messages:
+        return "unknown_tool()"
+    last = messages[-1]
+    tool_calls = getattr(last, "tool_calls", None) or []
+    if not tool_calls:
+        return "unknown_tool()"
+    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})'
+
+
+def build_agent(llm: ChatOpenAI):
+    memory = MemorySaver()
+    get_price = make_get_price_tool(llm)
+    return create_agent(
+        model=llm,
+        tools=[get_price],
+        system_prompt="Ты помощник по планированию покупок",
+        checkpointer=memory,
+        interrupt_before=["tools"],
+    )
+
+
+def run_chat() -> None:
+    llm = build_llm()
+    agent = build_agent(llm)
+    config = {"configurable": {"thread_id": THREAD_ID}}
+    step_holder = {"step": 1}
+
+    def format_chunk_message(chunk) -> None:
+        message, meta = chunk
+        graph_step = meta.get("langgraph_step", step_holder["step"])
+        if graph_step != step_holder["step"]:
+            step_holder["step"] = graph_step
+            console.print(STEP_SEPARATOR, end="")
+        if message.content:
+            console.print(message.content, end="")
+
+    def ask_and_run(user_input, cfg) -> None:
+        stream = agent.stream(
+            user_input,
+            config=cfg,
+            stream_mode=["messages", "updates"],
+        )
+
+        for chunk in stream:
+            state = agent.get_state(cfg)
+            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():
+                        console.print(formatted)
+
+                if "__interrupt__" in chunk_data and state.next == ("tools",):
+                    tool_repr = _tool_call_repr(state)
+                    console.print(STEP_SEPARATOR, end="")
+                    console.print(tool_repr)
+                    console.print(
+                        f"Агент хочет вызвать утилиту {tool_repr}"
+                    )
+                    answer = input("Разрешить? (Y/n): ")
+                    if answer.lower().strip() == "y":
+                        ask_and_run(None, cfg)
+                    else:
+                        console.print("Отменено")
+                    return
+
+        console.print()
+
+    console.print(
+        "Чат с агентом (память + подтверждение инструментов). "
+        "Введите exit для выхода."
+    )
+
+    while True:
+        user_input = input("\nВы: ")
+        if user_input.strip().lower() == "exit":
+            break
+        ask_and_run(
+            {"messages": [{"role": "human", "content": user_input}]},
+            config,
+        )
+
+
+if __name__ == "__main__":
+    run_chat()