Update agent.py

agent.py

@@ -1,59 +1,69 @@
-from langgraph.graph import StateGraph
-from typing import TypedDict, List, Any
+from langgraph import create_agent
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_openai import ChatOpenAI
 from rich.console import Console
-
-# Define state schema
-class AgentState(TypedDict):
-    conversation_history: List[dict]
-    last_tool_call_confirmed: bool | None
+from langgraph.tools import tool
+import asyncio
 
 console = Console()
 
-# Memory node: just returns the existing history
-async def memory_node(state: AgentState) -> dict:
-    return {"memory": state.get("conversation_history", [])}
+@tool
+def get_price(params: dict) -> str:
+    """Return a mock price for a city and date."""
+    city = params.get("city", "unknown")
+    date = params.get("date", "unknown")
+    return f"Цена в {city} на {date}: 1000₽"
 
-# Confirm tool call node
-async def confirm_tool_node(state: AgentState, tool_call: Any) -> dict:
+llm = ChatOpenAI(
+    base_url="http://localhost:11434/v1",
+    api_key="ollama",
+    model="llama3",
+)
+
+memory = MemorySaver()
+
+agent = create_agent(
+    model=llm,
+    tools=[get_price],
+    system_prompt="You are a helpful assistant.",
+    checkpointer=memory,
+    interrupt_before=["tools"],
+)
+
+config = {"configurable": {"thread_id": "conversation-1"}}
+
+async def ask_and_run(user_input, config):
     try:
-        console.print(f"[bold cyan]Tool call:[/bold cyan] {tool_call}")
-        confirmation = input("Confirm execution? (y/n): ")
-        confirmed = confirmation.lower().startswith("y")
-        state["last_tool_call_confirmed"] = confirmed
+        async for chunk in agent.stream(user_input, config=config, stream_mode=["messages", "updates"]):
+            chunk_type, chunk_data = chunk
+            if chunk_type == "messages":
+                console.print(chunk_data, end="")
+            elif chunk_type == "updates":
+                if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",):
+                    state = agent.get_state(config)
+                    tool_calls = state.values.get("messages", [])[-1].tool_calls
+                    if tool_calls:
+                        last_call = tool_calls[0]
+                        console.print(f"\n\nTool call detected: {last_call['name']}{last_call['args']}")
+                        answer = input("Разрешить? (Y/n): ")
+                        if answer.lower().strip() == "y":
+                            await ask_and_run(None, config)
+                        else:
+                            console.print("Отменено")
+                            break
     except Exception as e:
-        console.print(f"[red]Error during confirmation:[/red] {e}")
-        state["last_tool_call_confirmed"] = False
-    return {"confirmed_tool_call": state["last_tool_call_confirmed"]}
+        console.print(f"[red]Ошибка во время выполнения:[/red] {e}")
+        raise
 
-# Agent node: placeholder for actual LangChain agent logic
-async def agent_node(state: AgentState, memory: Any, confirmed_tool_call: Any) -> dict:
-    # In a real implementation we would invoke the agent here.
-    # For demonstration, echo back the conversation history and confirmation flag.
-    response = {
-        "memory": memory,
-        "confirmed": confirmed_tool_call,
-        "message": "Agent executed with memory and confirmation."
-    }
-    return {"response": response}
-
-# Build graph following execution flow: MemoryNode -> AgentNode -> ConfirmToolNode -> AgentNode
-builder = StateGraph(AgentState)
-builder.add_node("MemoryNode", memory_node)
-builder.add_node("AgentNode", agent_node)
-builder.add_node("ConfirmToolNode", confirm_tool_node)
-
-# Define transitions according to plan execution flow
-builder.set_entry_point("MemoryNode")
-builder.add_edge("MemoryNode", "AgentNode")
-builder.add_edge("AgentNode", "ConfirmToolNode")
-builder.add_edge("ConfirmToolNode", "AgentNode")
-builder.set_finish_node("AgentNode")
-
-# Compile graph
-graph = builder.compile()
+def main():
+    console.print("Добро пожаловать! Введите 'exit' для выхода.")
+    while True:
+        user_input = input("\nВы: ")
+        if user_input.lower() == "exit":
+            break
+        msg = {"messages": [{"role": "human", "content": user_input}]}
+        asyncio.run(ask_and_run(msg, config))
+        console.print("\n---")
 
 if __name__ == "__main__":
-    # Example initial state
-    init_state: AgentState = {"conversation_history": [], "last_tool_call_confirmed": None}
-    result = graph.invoke(init_state)
-    console.print(result)
+    main()