Add agent_with_memory.py

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)