From 0bae2ed8a0421d78504bb91c658e9b59f5faeaa5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Tue, 2 Jun 2026 04:40:45 +0000 Subject: [PATCH] Add agent_with_memory.py --- agent_with_memory.py | 167 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 167 insertions(+) create mode 100644 agent_with_memory.py diff --git a/agent_with_memory.py b/agent_with_memory.py new file mode 100644 index 0000000..656737f --- /dev/null +++ b/agent_with_memory.py @@ -0,0 +1,167 @@ +""" +Agent with memory and confirmation of tool calls. + +This script demonstrates how to create an agent using LangGraph with: +- Memory (MemorySaver) to keep conversation history. +- Interrupt-before to pause before calling a tool. +- Rich console for pretty printing. + +Run: + python agent_with_memory.py + +Make sure to install dependencies: + pip install -r requirements.txt +""" + +import json +from typing import Any, Dict, List, Optional + +# Rich console for pretty printing +from rich.console import Console + +# LangGraph components +from langgraph import AgentBuilder +from langgraph.checkpoint.memory import MemorySaver + +# LangChain components +from langchain_ollama import ChatOllama +from langchain_ollama import OllamaEmbeddings + +# Tool definition +from langchain.tools import BaseTool + +console = Console() + +# ----- Tool definition ----- +class GetPriceTool(BaseTool): + name: str = "get_price" + description: str = "Get the price of a product for a given city and date. Returns a string." + + def _run(self, city: str, date: str) -> str: + # Dummy implementation – in real life, query an API + return f"The price in {city} on {date} is $42." + + def _arun(self, city: str, date: str) -> str: # async version + return self._run(city, date) + +# ----- Agent creation ----- + +def create_agent( + model: Any, + tools: List[BaseTool], + system_prompt: str, + checkpointer: MemorySaver, + interrupt_before: Optional[List[str]] = None, +) -> Any: + """Build and return a LangGraph agent. + + Parameters + ---------- + model: The language model (e.g., ChatOllama). + tools: List of tools the agent can use. + system_prompt: System prompt for the agent. + checkpointer: MemorySaver instance for conversation memory. + interrupt_before: List of nodes to interrupt before (e.g., ['tools']). + """ + builder = AgentBuilder( + model=model, + tools=tools, + system_prompt=system_prompt, + checkpointer=checkpointer, + ) + if interrupt_before: + builder = builder.with_interrupt_before(interrupt_before) + return builder.build() + +# ----- Conversation loop ----- + +def ask_and_run( + agent: Any, + user_input: Optional[Dict[str, Any]], + config: Dict[str, Any], +) -> None: + """Handle streaming from the agent, including pauses for tool confirmation. + + Parameters + ---------- + agent: The LangGraph agent. + user_input: The user message dict or None to resume. + config: Configuration dict with thread_id. + """ + # Prepare the input for the agent + input_data = user_input if user_input is not None else None + + # Stream the agent's response + for chunk in agent.stream( + input_data, config=config, stream_mode=["messages", "updates"] + ): + chunk_type, chunk_data = chunk + + # Handle message chunks (token streaming) + if chunk_type == "messages": + # chunk_data is a list of messages; we print the last message content + if chunk_data: + last_msg = chunk_data[-1] + if last_msg.get("role") == "assistant": + console.print(f"[bold cyan]Agent:[/bold cyan] {last_msg.get("content", "")}") + + # Handle updates (tool calls, etc.) + if chunk_type == "updates": + # chunk_data contains the updated state; we can inspect tool calls + pass # For simplicity, we ignore updates here + + # Detect interrupt before tool call + if "__interrupt__" in chunk_data and agent.get_state(config).next == ("tools",): + # Agent is pausing before a tool call + state = agent.get_state(config) + # The last message should contain the tool call + last_msg = state.values["messages"][-1] + tool_call = last_msg["tool_calls"][0] + tool_name = tool_call["name"] + tool_args = tool_call["args"] + console.print(f"\n[bold yellow]Agent wants to call tool:[/bold yellow] {tool_name}({json.dumps(tool_args)})") + answer = input("Разрешить? (Y/n): ") + if answer.lower().strip() == "y": + console.print("[green]Tool call allowed. Resuming...[/green]") + # Recursively call ask_and_run to resume + ask_and_run(agent, None, config) + else: + console.print("[red]Tool call cancelled.[/red]") + break + +if __name__ == "__main__": + # Initialize the model and embeddings + llm = ChatOllama(model="llama2") + embeddings = OllamaEmbeddings(model="llama2") + + # Create memory saver + memory = MemorySaver() + + # Define system prompt + system_prompt = "You are a helpful assistant. Use the get_price tool to answer queries about prices." + + # Create the tool + get_price_tool = GetPriceTool() + + # Build the agent + agent = create_agent( + model=llm, + tools=[get_price_tool], + system_prompt=system_prompt, + checkpointer=memory, + interrupt_before=["tools"], + ) + + # Conversation loop + thread_id = "thread-1" + config = {"configurable": {"thread_id": thread_id}} + + console.print("[bold green]Agent ready. Type 'exit' to quit.[/bold green]") + while True: + user_input = input("\nВы: ") + if user_input.lower().strip() == "exit": + console.print("[bold magenta]Goodbye![/bold magenta]") + break + # Wrap user input into a message dict + message = {"messages": [{"role": "human", "content": user_input}]} + ask_and_run(agent, message, config) \ No newline at end of file