Обновлен main.py с использованием create_agent

main.py

@@ -1,101 +1,79 @@
+# main.py
+# Реализация агента с памятью и подтверждением действий
+
 import os
-import json
-import asyncio
 from rich.console import Console
-from langchain_openai import ChatOpenAI
-from langchain.tools import tool
-from deepagents import create_deep_agent
-from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
 from langgraph.checkpoint.memory import MemorySaver
+from langgraph import create_agent
+from langchain_ollama import ChatOllama
+from langchain.tools import Tool
 
-# --- LLM -------------------------------------------------------------------
+# Rich console for pretty output
-llm = ChatOpenAI(
-    model="openai/gpt-oss-20b:free",
-    base_url="https://openrouter.ai/api/v1",
-    api_key=os.getenv("OPENAI_API_KEY"),
-    temperature=0.0,
-)
-
-# --- Backend --------------------------------------------------------------
-backend = CompositeBackend([
-    LocalShellBackend(workspace_dir="./workspace"),
-    FilesystemBackend(),
-])
-
-# --- Tool ------------------------------------------------------------------
-@tool
-def get_weather(query: str) -> str:
-    """Return a fake weather report for a given city and date.
-    The argument is a JSON string with keys "city" and "date".
-    """
-    try:
-        data = json.loads(query)
-        city = data.get("city", "unknown")
-        date = data.get("date", "unknown")
-        return f"Weather in {city} on {date}: sunny, 25°C."
-    except Exception as e:
-        return f"Error parsing query: {e}"
-
-# --- Agent -----------------------------------------------------------------
-memory = MemorySaver()
-
-agent = create_deep_agent(
-    model=llm,
-    tools=[get_weather],
-    backend=backend,
-    system_prompt="You are a helpful weather assistant. Use the get_weather tool when asked about weather.",
-    checkpointer=memory,
-    interrupt_before=["tools"],  # pause before each tool call
-)
-
 console = Console()
 
-# ---------------------------------------------------------------------------
+# ---------- Tool definition ----------
-async def ask_and_run(user_input, config):
-    """Run the agent with optional user input.
-    If user_input is None, the agent resumes from the paused state.
-    """
-    # Stream the agent's response
-    async for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]):
-        state = agent.get_state(config)
-        chunk_type, chunk_data = chunk
 
+def get_price(city: str, date: str) -> str:
+    """Return a mock price for a given city and date."""
+    return f"Price in {city} on {date} is $100"
+
+price_tool = Tool(
+    name="get_price",
+    description="Retrieve the price for a specified city and date.",
+    func=get_price,
+)
+
+# ---------- LLM and Agent setup ----------
+llm = ChatOllama(model="llama3")
+memory = MemorySaver()
+agent = create_agent(
+    model=llm,
+    tools=[price_tool],
+    system_prompt="You are a helpful assistant that can call the get_price tool.",
+    checkpointer=memory,
+    interrupt_before=["tools"],  # pause before each tool invocation
+)
+
+# ---------- Helper function to handle streaming and confirmation ----------
+
+def ask_and_run(user_input, config):
+    for chunk_type, chunk_data in agent.stream(
+        user_input if user_input is not None else None,
+        config=config,
+        stream_mode=["messages", "updates"],
+    ):
         if chunk_type == "messages":
-            # chunk_data can be a dict with 'content'
+            console.print(chunk_data["content"], end="")
-            if isinstance(chunk_data, dict) and "content" in chunk_data:
+            console.flush()
-                console.print(chunk_data["content"], end="")
+            continue
-            else:
+        if chunk_type == "updates":
-                console.print(chunk_data)
+            console.print(f"\n[bold cyan]Update:[/bold cyan] {chunk_data}")
-        elif chunk_type == "updates":
+            continue
-            console.print("\n[bold]Tool call:[/bold]", chunk_data)
+        if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",):
-
+            state = agent.get_state(config)
-        # Detect pause before tool call
-        if "__interrupt__" in chunk_data and state.next == ("tools",):
-            # Retrieve the pending tool call
             last_msg = state.values["messages"][-1]
-            tool_call = None
+            tool_call = last_msg.tool_calls[0]
-            if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
+            tool_name = tool_call["name"]
-                tool_call = last_msg.tool_calls[0]
+            tool_args = tool_call["args"]
-            if tool_call:
+            console.print(f"\n\n[bold yellow]Agent wants to call tool:[/bold yellow] {tool_name}({tool_args})")
-                tool_name = tool_call["name"]
+            answer = input("Разрешить? (Y/n): ")
-                tool_args = tool_call["args"]
+            if answer.lower().strip() == "y":
-                console.print(f"\n[bold]Agent wants to call {tool_name}({tool_args})[/bold]")
+                ask_and_run(None, config)
-                answer = input("Разрешить? (Y/n): ")
+            else:
-                if answer.lower().strip() == "y":
+                console.print("[red]Отменено[/red]")
-                    await ask_and_run(None, config)  # resume
+            break
-                else:
-                    console.print("[red]Отменено[/red]")
-                    break
 
-# ---------------------------------------------------------------------------
+# ---------- Main chat loop ----------
-async def main():
+if __name__ == "__main__":
-    console.print("\n[bold green]Weather Agent ready! Type 'exit' to quit.[/bold green]")
+    thread_id = os.getenv("THREAD_ID", "chat-1")
-    config = {"configurable": {"thread_id": "session-1"}}
+    config = {"configurable": {"thread_id": thread_id}}
+    console.print("[bold green]Добро пожаловать в интерактивный агент![/bold green]")
+    console.print("Введите 'exit' для выхода.")
     while True:
         user_input = input("\nВы: ")
-        if user_input.lower() == "exit":
+        if user_input.lower().strip() == "exit":
+            console.print("[bold magenta]До свидания![/bold magenta]")
             break
-        await ask_and_run({"messages": [{"role": "human", "content": user_input}]}, config)
+        message = {"messages": [{"role": "human", "content": user_input}]}
-
+        ask_and_run(message, config)
-if __name__ == "__main__":
+        console.print("\n--- --- ---")
-    asyncio.run(main())