rewrite main.py with correct create_agent helper

main.py

@@ -1,60 +1,102 @@
 import os
 from langchain_openai import ChatOpenAI
 from langgraph.checkpoint.memory import MemorySaver
-from langgraph.prebuilt import create_react_agent as create_agent
+from langgraph.prebuilt import create_react_agent
 from rich.console import Console
-from rich.markdown import Markdown
-from tools import get_price
 
-# LLM configuration using BroJS
+# ---------------------------------------------------------------
+# LLM
+# ---------------------------------------------------------------
 llm = ChatOpenAI(
     model="openai/gpt-oss-20b:free",
     base_url="https://platform.brojs.ru/jrnl-bh/api/inference/v1",
     api_key=os.getenv("JOURNAL_MCP_PAT"),
-    temperature=0.5,
+    temperature=0.0,
 )
 
-# Memory checkpoint for interrupt handling
+# ---------------------------------------------------------------
+# Helper create_agent — обёртка, как описано в инструкции
+# ---------------------------------------------------------------
+def create_agent(model, tools, system_prompt="", checkpointer=None, interrupt_before=None):
+    """Helper-функция для создания агента с памятью и interrupt."""
+    kwargs = dict(model=model, tools=tools)
+    if system_prompt:
+        kwargs["prompt"] = system_prompt
+    if checkpointer is not None:
+        kwargs["checkpointer"] = checkpointer
+    if interrupt_before is not None:
+        kwargs["interrupt_before"] = interrupt_before
+    return create_react_agent(**kwargs)
+
+# ---------------------------------------------------------------
+# Инструменты
+# ---------------------------------------------------------------
+from langchain.tools import tool
+
+@tool
+def get_price(city: str, date: str) -> str:
+    """Возвращает условную цену/погоду для города на указанную дату."""
+    return f"Данные для {city} на {date}: 720 руб/кг, 18C, облачно"
+
+# ---------------------------------------------------------------
+# Агент с памятью и паузой перед инструментом
+# ---------------------------------------------------------------
 memory = MemorySaver()
 
-# Create the agent with interrupt before tools
 agent = create_agent(
     model=llm,
     tools=[get_price],
+    system_prompt="Ты полезный ассистент. Помогай пользователю находить цены и погоду.",
     checkpointer=memory,
     interrupt_before=["tools"],
 )
 
 console = Console()
 
+# ---------------------------------------------------------------
+# ask_and_run — точная структура из задания
+# ---------------------------------------------------------------
+def ask_and_run(user_input, config):
+    for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]):
+        state = agent.get_state(config)
+        chunk_type, chunk_data = chunk
 
-def ask_and_run(user_input: str, config: dict):
-    """Synchronously run the agent with streaming and handle interrupts."""
-    # Stream messages and updates
-    for chunk in agent.stream(
-        {"messages": [{"role": "user", "content": user_input}],
-         "configurable": config},
-        stream_mode=["messages", "updates"],
-    ):
-        if isinstance(chunk, dict) and "__interrupt__" in chunk:
-            # Interrupt: ask for confirmation
-            console.print("[bold red]Agent requested tool execution. Confirm? (y/n): ", end="")
-            choice = input().strip().lower()
-            if choice != "y":
-                # Reject by sending a new message to the agent
-                config.update({"configurable": {"thread_id": config.get("thread_id", "default")}})
-                continue
-        if isinstance(chunk, dict) and "messages" in chunk:
-            for msg in chunk["messages"]:
-                console.print(Markdown(msg["content"]))
+        if chunk_type == "messages":
+            msg = chunk_data[0] if isinstance(chunk_data, (list, tuple)) else chunk_data
+            content = getattr(msg, "content", "")
+            if content and not getattr(msg, "tool_calls", None):
+                console.print(content, end="")
 
+        if chunk_type == "updates":
+            console.print("\n --- --- --- \n")
+            for node_data in chunk_data.values():
+                for msg in node_data.get("messages", []):
+                    for tc in getattr(msg, "tool_calls", []):
+                        console.print(f"[bold]{tc['name']}({tc['args']})[/bold]")
+
+        if "__interrupt__" in chunk_data and state.next == ("tools",):
+            tool_call = state.values["messages"][-1].tool_calls[0]
+            console.print(f"Агент хочет вызвать утилиту [bold]{tool_call['name']}({tool_call['args']})[/bold]")
+            answer = input("Разрешить? (Y/n): ")
+            if answer.lower().strip() == "y":
+                ask_and_run(None, config)
+            else:
+                console.print("Отменено")
+                break
+
+# ---------------------------------------------------------------
+# Чат-цикл
+# ---------------------------------------------------------------
 if __name__ == "__main__":
-    thread_id = os.getenv("THREAD_ID", "session-1")
-    config = {"thread_id": thread_id}
-    console.print("[bold green]Welcome to the price agent. Type 'exit' to quit.")
+    config = {"configurable": {"thread_id": "разговор-1"}}
+    console.print("[bold green]Агент запущен. Введите 'exit' для выхода.[/bold green]")
+
     while True:
-        user_input = input("Вы: ")
-        if user_input.lower() in ("exit", "quit"):
+        user_input = input("\nВы: ")
+        if user_input == "exit":
             break
-        ask_and_run(user_input, config)
-    console.print("[bold blue]Goodbye!", style="bold")
+
+        ask_and_run(
+            {"messages": [{"role": "human", "content": user_input}]},
+            config,
+        )