how-tos/time-travel.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to view and update past graph state\n",
    "\n",
    "Once you start [checkpointing](./persistence.ipynb) your graphs, you can easily\n",
    "**get** or **update** the state of the agent at any point in time. This permits\n",
    "a few things:\n",
    "\n",
    "1. You can surface a state during an interrupt to a user to let them accept an\n",
    "   action.\n",
    "2. You can **rewind** the graph to reproduce or avoid issues.\n",
    "3. You can **modify** the state to embed your agent into a larger system, or to\n",
    "   let the user better control its actions.\n",
    "\n",
    "The key methods used for this functionality are:\n",
    "\n",
    "- [getState](https://bsorrentino.github.io/langgraph4j/apidocs/org/bsc/langgraph4j/CompiledGraph.html#getState-org.bsc.langgraph4j.RunnableConfig-):\n",
    "  fetch the values from the target config\n",
    "- [updateState](https://bsorrentino.github.io/langgraph4j/apidocs/org/bsc/langgraph4j/CompiledGraph.html#updateState-org.bsc.langgraph4j.RunnableConfig-java.util.Map-java.lang.String-):\n",
    "  apply the given values to the target state\n",
    "\n",
    "**Note:** this requires passing in a checkpointer.\n",
    "\n",
    "This works for [StateGraph](https://bsorrentino.github.io/langgraph4j/apidocs/org/bsc/langgraph4j/StateGraph.html)\n",
    "\n",
    "Below is an example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "String userHomeDir = System.getProperty(\"user.home\");\n",
    "String localRespoUrl = \"file://\" + userHomeDir + \"/.m2/repository/\";\n",
    "String langchain4jVersion = \"0.36.2\";\n",
    "String langgraph4jVersion = \"1.1-SNAPSHOT\";"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0mRepository \u001b[1m\u001b[32mlocal\u001b[0m url: \u001b[1m\u001b[32mfile:///Users/bsorrentino/.m2/repository/\u001b[0m added.\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "%dependency /add-repo local \\{localRespoUrl} release|never snapshot|always\n",
    "// %dependency /list-repos"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remove installed package from Jupiter cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash \n",
    "rm -rf \\{userHomeDir}/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/langgraph4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adding dependency \u001b[0m\u001b[1m\u001b[32morg.slf4j:slf4j-jdk14:2.0.9\n",
      "\u001b[0mAdding dependency \u001b[0m\u001b[1m\u001b[32morg.bsc.langgraph4j:langgraph4j-core-jdk8:1.0-SNAPSHOT\n",
      "\u001b[0mAdding dependency \u001b[0m\u001b[1m\u001b[32morg.bsc.langgraph4j:langgraph4j-langchain4j:1.0-SNAPSHOT\n",
      "\u001b[0mAdding dependency \u001b[0m\u001b[1m\u001b[32mdev.langchain4j:langchain4j:0.35.0\n",
      "\u001b[0mAdding dependency \u001b[0m\u001b[1m\u001b[32mdev.langchain4j:langchain4j-open-ai:0.35.0\n",
      "\u001b[0mSolving dependencies\n",
      "Resolved artifacts count: 26\n",
      "Add to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/slf4j/slf4j-jdk14/2.0.9/slf4j-jdk14-2.0.9.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/slf4j/slf4j-api/2.0.9/slf4j-api-2.0.9.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/langgraph4j/langgraph4j-core-jdk8/1.0-SNAPSHOT/langgraph4j-core-jdk8-1.0-SNAPSHOT.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/async/async-generator-jdk8/2.2.0/async-generator-jdk8-2.2.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/bsc/langgraph4j/langgraph4j-langchain4j/1.0-SNAPSHOT/langgraph4j-langchain4j-1.0-SNAPSHOT.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/dev/langchain4j/langchain4j/0.35.0/langchain4j-0.35.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/dev/langchain4j/langchain4j-core/0.35.0/langchain4j-core-0.35.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/google/code/gson/gson/2.10.1/gson-2.10.1.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/apache/opennlp/opennlp-tools/1.9.4/opennlp-tools-1.9.4.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/dev/langchain4j/langchain4j-open-ai/0.35.0/langchain4j-open-ai-0.35.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/dev/ai4j/openai4j/0.22.0/openai4j-0.22.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/squareup/retrofit2/retrofit/2.9.0/retrofit-2.9.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/squareup/retrofit2/converter-jackson/2.9.0/converter-jackson-2.9.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/fasterxml/jackson/core/jackson-databind/2.17.2/jackson-databind-2.17.2.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/fasterxml/jackson/core/jackson-annotations/2.17.2/jackson-annotations-2.17.2.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/fasterxml/jackson/core/jackson-core/2.17.2/jackson-core-2.17.2.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/squareup/okhttp3/okhttp/4.12.0/okhttp-4.12.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/squareup/okio/okio/3.6.0/okio-3.6.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/squareup/okio/okio-jvm/3.6.0/okio-jvm-3.6.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/squareup/okhttp3/okhttp-sse/4.12.0/okhttp-sse-4.12.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/jetbrains/kotlin/kotlin-stdlib-jdk8/1.9.10/kotlin-stdlib-jdk8-1.9.10.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/jetbrains/kotlin/kotlin-stdlib/1.9.10/kotlin-stdlib-1.9.10.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/jetbrains/kotlin/kotlin-stdlib-common/1.9.10/kotlin-stdlib-common-1.9.10.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/jetbrains/annotations/13.0/annotations-13.0.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/org/jetbrains/kotlin/kotlin-stdlib-jdk7/1.9.10/kotlin-stdlib-jdk7-1.9.10.jar\u001b[0m\n",
      "\u001b[0mAdd to classpath: \u001b[0m\u001b[32m/Users/bsorrentino/Library/Jupyter/kernels/rapaio-jupyter-kernel/mima_cache/com/knuddels/jtokkit/1.1.0/jtokkit-1.1.0.jar\u001b[0m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "%dependency /add org.slf4j:slf4j-jdk14:2.0.9\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-core-jdk8:\\{langgraph4jVersion}\n",
    "%dependency /add org.bsc.langgraph4j:langgraph4j-langchain4j:\\{langgraph4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j:\\{langchain4jVersion}\n",
    "%dependency /add dev.langchain4j:langchain4j-open-ai:\\{langchain4jVersion}\n",
    "\n",
    "%dependency /resolve"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "try( var file = new java.io.FileInputStream(\"./logging.properties\")) {\n",
    "    var lm = java.util.logging.LogManager.getLogManager();\n",
    "    lm.checkAccess(); \n",
    "    lm.readConfiguration( file );\n",
    "}\n",
    "\n",
    "var log = org.slf4j.LoggerFactory.getLogger(\"time-travel\");\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the state\n",
    "\n",
    "State is an (immutable) data class, inheriting from [AgentState], shared with all nodes in our graph. A state is basically a wrapper of a `Map<String,Object>` that provides some enhancers:\n",
    "\n",
    "1. Schema (optional), that is a `Map<String,Channel>` where each [Channel] describe behaviour of the related property\n",
    "1. `value()` accessors that inspect Map an return an Optional of value contained and cast to the required type\n",
    "\n",
    "[Channel]: https://bsorrentino.github.io/langgraph4j/apidocs/org/bsc/langgraph4j/state/Channel.html\n",
    "[AgentState]: https://bsorrentino.github.io/langgraph4j/apidocs/org/bsc/langgraph4j/state/AgentState.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import org.bsc.langgraph4j.state.AgentState;\n",
    "import org.bsc.langgraph4j.state.Channel;\n",
    "import org.bsc.langgraph4j.state.AppenderChannel;\n",
    "import dev.langchain4j.data.message.AiMessage;\n",
    "import dev.langchain4j.data.message.ChatMessage;\n",
    "\n",
    "public class MessageState extends AgentState {\n",
    "\n",
    "    static Map<String, Channel<?>> SCHEMA = Map.of(\n",
    "            \"messages\", AppenderChannel.<ChatMessage>of(ArrayList::new)\n",
    "    );\n",
    "\n",
    "    public MessageState(Map<String, Object> initData) {\n",
    "        super( initData  );\n",
    "    }\n",
    "\n",
    "    List<? extends ChatMessage> messages() {\n",
    "        return this.<List<? extends ChatMessage>>value( \"messages\" )\n",
    "                .orElseThrow( () -> new RuntimeException( \"messages not found\" ) );\n",
    "    }\n",
    "\n",
    "    // utility method to quick access to last message\n",
    "    Optional<? extends ChatMessage> lastMessage() {\n",
    "        List<? extends ChatMessage> messages = messages();\n",
    "        return ( messages.isEmpty() ) ? \n",
    "            Optional.empty() :\n",
    "            Optional.of(messages.get( messages.size() - 1 ));\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Register Serializers\n",
    "\n",
    "Every object that should be stored into State **MUST BE SERIALIZABLE**. If the object is not `Serializable` by default, Langgraph4j provides a way to build and associate a custom [Serializer] to it. \n",
    "\n",
    "In the example, we use  [`Serializer`] for [ToolExecutionRequest] and [ChatMesssager] provided by langgraph4j integration module with langchain4j .\n",
    "\n",
    "[Serializer]: https://bsorrentino.github.io/langgraph4j/apidocs/org/bsc/langgraph4j/serializer/Serializer.html\n",
    "[ChatMesssager]: https://docs.langchain4j.dev/apidocs/dev/langchain4j/data/message/ChatMesssager.html\n",
    "[ToolExecutionRequest]: https://docs.langchain4j.dev/apidocs/dev/langchain4j/data/message/ToolExecutionRequest.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "SerializerMapper: \n",
       "java.util.Map\n",
       "java.util.Collection\n",
       "dev.langchain4j.agent.tool.ToolExecutionRequest\n",
       "dev.langchain4j.data.message.ChatMessage"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import dev.langchain4j.data.message.AiMessage;\n",
    "import dev.langchain4j.data.message.SystemMessage;\n",
    "import dev.langchain4j.data.message.UserMessage;\n",
    "import dev.langchain4j.data.message.ToolExecutionResultMessage;\n",
    "import dev.langchain4j.agent.tool.ToolExecutionRequest;\n",
    "import org.bsc.langgraph4j.serializer.std.ObjectStreamStateSerializer;\n",
    "import org.bsc.langgraph4j.langchain4j.serializer.std.ChatMesssageSerializer;\n",
    "import org.bsc.langgraph4j.langchain4j.serializer.std.ToolExecutionRequestSerializer;\n",
    "import org.bsc.langgraph4j.langchain4j.serializer.std.UserMessageSerializer;\n",
    "\n",
    "var stateSerializer = new ObjectStreamStateSerializer<MessageState>( MessageState::new );\n",
    "stateSerializer.mapper()\n",
    "    // Setup custom serializer for Langchain4j ToolExecutionRequest\n",
    "    .register(ToolExecutionRequest.class, new ToolExecutionRequestSerializer() )\n",
    "    // Setup custom serializer for Langchain4j AiMessage\n",
    "    .register(ChatMessage.class, new ChatMesssageSerializer() )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the tools\n",
    "\n",
    "Using [langchain4j], We will first define the tools we want to use. For this simple example, we will\n",
    "use create a placeholder search engine. However, it is really easy to create\n",
    "your own tools - see documentation\n",
    "[here][tools] on how to do\n",
    "that.\n",
    "\n",
    "[langchain4j]: https://docs.langchain4j.dev\n",
    "[tools]: https://docs.langchain4j.dev/tutorials/tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dev.langchain4j.agent.tool.P;\n",
    "import dev.langchain4j.agent.tool.Tool;\n",
    "\n",
    "import java.util.Optional;\n",
    "\n",
    "import static java.lang.String.format;\n",
    "\n",
    "public class SearchTool {\n",
    "\n",
    "    @Tool(\"Use to surf the web, fetch current information, check the weather, and retrieve other information.\")\n",
    "    String execQuery(@P(\"The query to use in your search.\") String query) {\n",
    "\n",
    "        // This is a placeholder for the actual implementation\n",
    "        return \"Cold, with a low of 13 degrees\";\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up the model\n",
    "\n",
    "Now we will load the\n",
    "[chat model].\n",
    "\n",
    "1. It should work with messages. We will represent all agent state in the form of messages, so it needs to be able to work well with them.\n",
    "2. It should work with [tool calling],meaning it can return function arguments in its response.\n",
    "\n",
    "Note:\n",
    "   >\n",
    "   > These model requirements are not general requirements for using LangGraph4j - they are just requirements for this one example.\n",
    "   >\n",
    "\n",
    "[chat model]: https://docs.langchain4j.dev/tutorials/chat-and-language-models\n",
    "[tool calling]: https://docs.langchain4j.dev/tutorials/tools   \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dev.langchain4j.model.openai.OpenAiChatModel;\n",
    "import dev.langchain4j.agent.tool.ToolSpecification;\n",
    "import dev.langchain4j.agent.tool.ToolSpecifications;\n",
    "\n",
    "OpenAiChatModel llm = OpenAiChatModel.builder()\n",
    "    .apiKey( System.getenv(\"OPENAI_API_KEY\") )\n",
    "    .modelName( \"gpt-4o\" )\n",
    "    .logResponses(true)\n",
    "    .maxRetries(2)\n",
    "    .temperature(0.0)\n",
    "    .maxTokens(2000)\n",
    "    .build()  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test function calling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "execute: execQuery \n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Optional[ToolExecutionResultMessage { id = \"call_clfKNL7okcr8bHNwEJUYp7Jn\" toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import dev.langchain4j.agent.tool.ToolSpecification;\n",
    "import dev.langchain4j.agent.tool.ToolSpecifications;\n",
    "import dev.langchain4j.data.message.UserMessage;\n",
    "import dev.langchain4j.data.message.AiMessage;\n",
    "import dev.langchain4j.model.output.Response;\n",
    "import dev.langchain4j.model.openai.OpenAiChatModel;\n",
    "import dev.langchain4j.service.tool.DefaultToolExecutor;\n",
    "import org.bsc.langgraph4j.langchain4j.tool.ToolNode;\n",
    "\n",
    "var toolNode = ToolNode.of( new SearchTool() );\n",
    "\n",
    "var tools = toolNode.toolSpecifications();\n",
    "\n",
    "UserMessage userMessage = UserMessage.from(\"What will the weather be like in London tomorrow?\");\n",
    "\n",
    "Response<AiMessage> response = llm.generate(Collections.singletonList(userMessage), tools );\n",
    "\n",
    "AiMessage aiMessage = response.content();\n",
    "\n",
    "var result = toolNode.execute( aiMessage.toolExecutionRequests() );\n",
    "\n",
    "result;\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the graph\n",
    "\n",
    "We can now put it all together. We will run it first without a checkpointer:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "import static org.bsc.langgraph4j.StateGraph.START;\n",
    "import static org.bsc.langgraph4j.StateGraph.END;\n",
    "import static org.bsc.langgraph4j.action.AsyncEdgeAction.edge_async;\n",
    "import static org.bsc.langgraph4j.action.AsyncNodeAction.node_async;\n",
    "import org.bsc.langgraph4j.StateGraph;\n",
    "import org.bsc.langgraph4j.action.EdgeAction;\n",
    "import org.bsc.langgraph4j.action.NodeAction;\n",
    "import dev.langchain4j.data.message.AiMessage;\n",
    "import dev.langchain4j.data.message.ChatMessage;\n",
    "import dev.langchain4j.service.tool.DefaultToolExecutor;\n",
    "import org.bsc.langgraph4j.checkpoint.MemorySaver; \n",
    "import org.bsc.langgraph4j.CompileConfig; \n",
    "import java.util.stream.Collectors;\n",
    "// Route Message \n",
    "EdgeAction<MessageState> routeMessage = state -> {\n",
    "  \n",
    "  var lastMessage = state.lastMessage();\n",
    "  \n",
    "  if ( !lastMessage.isPresent()) return \"exit\";\n",
    "\n",
    "  if( lastMessage.get() instanceof AiMessage message  ) {\n",
    "\n",
    "    // If tools should be called\n",
    "    if ( message.hasToolExecutionRequests() ) return \"next\";\n",
    "    \n",
    "  }\n",
    "  \n",
    "  // If no tools are called, we can finish (respond to the user)\n",
    "  return \"exit\";\n",
    "};\n",
    "\n",
    "// Call Model\n",
    "NodeAction<MessageState> callModel = state -> {\n",
    "  var tools = ToolSpecifications.toolSpecificationsFrom( SearchTool.class );\n",
    "\n",
    "  var response = llm.generate( (List<ChatMessage>)state.messages(), tools );\n",
    "\n",
    "  return Map.of( \"messages\", response.content() );\n",
    "};\n",
    "\n",
    "// Invoke Tool \n",
    "NodeAction<MessageState> invokeTool = state -> {\n",
    "  var lastMessage = (AiMessage)state.lastMessage()\n",
    "                          .orElseThrow( () -> ( new IllegalStateException( \"last message not found!\")) );\n",
    "\n",
    "  var toolNode = ToolNode.of( new SearchTool() );\n",
    "  \n",
    "  var result = toolNode.execute( lastMessage.toolExecutionRequests() )\n",
    "                    .orElseThrow( () -> ( new IllegalStateException( \"tool execution failed!\")));\n",
    "\n",
    "  return Map.of( \"messages\", result );\n",
    "};\n",
    "\n",
    "// Define Graph\n",
    "var workflow = new StateGraph<MessageState> ( MessageState.SCHEMA, stateSerializer )\n",
    "  .addNode(\"agent\", node_async(callModel) )\n",
    "  .addNode(\"tools\", node_async(invokeTool) )\n",
    "  .addEdge(START, \"agent\")\n",
    "  .addConditionalEdges(\"agent\", edge_async(routeMessage), Map.of( \"next\", \"tools\", \"exit\", END ))\n",
    "  .addEdge(\"tools\", \"agent\");\n",
    "\n",
    "// Here we only save in-memory\n",
    "var memory = new MemorySaver();\n",
    "\n",
    "var compileConfig = CompileConfig.builder()\n",
    "                    .checkpointSaver(memory)\n",
    "                    .build();\n",
    "\n",
    "var graph = workflow.compile(compileConfig);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Interacting with the Agent\n",
    "\n",
    "We can now interact with the agent. Between interactions you can get and update state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "START \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__START__\n",
      "{messages=[UserMessage { name = null contents = [TextContent { text = \"Hi I'm Bartolo.\" }] }]}\n",
      "agent\n",
      "{messages=[UserMessage { name = null contents = [TextContent { text = \"Hi I'm Bartolo.\" }] }, AiMessage { text = \"Hello Bartolo! How can I assist you today?\" toolExecutionRequests = null }]}\n",
      "__END__\n",
      "{messages=[UserMessage { name = null contents = [TextContent { text = \"Hi I'm Bartolo.\" }] }, AiMessage { text = \"Hello Bartolo! How can I assist you today?\" toolExecutionRequests = null }]}\n"
     ]
    }
   ],
   "source": [
    "import org.bsc.langgraph4j.RunnableConfig;\n",
    "\n",
    "var runnableConfig =  RunnableConfig.builder()\n",
    "                .threadId(\"conversation-num-1\" )\n",
    "                .build();\n",
    "\n",
    "Map<String,Object> inputs = Map.of( \"messages\", UserMessage.from(\"Hi I'm Bartolo.\") );\n",
    "\n",
    "var result = graph.stream( inputs, runnableConfig );\n",
    "\n",
    "for( var r : result ) {\n",
    "  System.out.println( r.node() );\n",
    "  System.out.println( r.state() );\n",
    "  \n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here you can see the \"`agent`\" node ran, and then our edge returned `__END__` so the graph stopped execution there.\n",
    "\n",
    "Let's check the current graph state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "StateSnapshot{node=agent, state={messages=[UserMessage { name = null contents = [TextContent { text = \"Hi I'm Bartolo.\" }] }, AiMessage { text = \"Hello Bartolo! How can I assist you today?\" toolExecutionRequests = null }]}, config=RunnableConfig(threadId=conversation-num-1, checkPointId=a524b711-4588-4859-8c50-acff0e7c1ed3, nextNode=__END__, streamMode=VALUES)}\n"
     ]
    }
   ],
   "source": [
    "import org.bsc.langgraph4j.checkpoint.Checkpoint;\n",
    "\n",
    "var checkpoint = graph.getState(runnableConfig);\n",
    "\n",
    "System.out.println(checkpoint);\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The current state is the two messages we've seen above, 1. the Human Message we sent in, 2. the AIMessage we got back from the model.\n",
    "\n",
    "The next value is `__END__`  since the graph has terminated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "__END__"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "checkpoint.getNext()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's get it to execute a tool\n",
    "\n",
    "When we call the graph again, it will create a checkpoint after each internal execution step. Let's get it to run a tool, then look at the checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "START \n",
      "execute: execQuery \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AiMessage { text = \"The current weather in San Francisco is cold, with a low of 13 degrees Celsius. If you need more details or have any other questions, feel free to ask!\" toolExecutionRequests = null }\n"
     ]
    }
   ],
   "source": [
    "\n",
    "Map<String,Object> inputs = Map.of( \"messages\", UserMessage.from(\"What's the weather like in SF currently?\") );\n",
    "\n",
    "var state = graph.invoke( inputs, runnableConfig ).orElseThrow( () ->(new IllegalStateException()) ) ;\n",
    "\n",
    "System.out.println( state.lastMessage().orElse(null) );\n",
    "  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pause before tools\n",
    "\n",
    "If you notice below, we now will add interruptBefore=[\"action\"] - this means that before any actions are taken we pause. This is a great moment to allow the user to correct and update the state! This is very useful when you want to have a human-in-the-loop to validate (and potentially change) the action to take.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "START \n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "__START__\n",
      "{messages=[UserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }]}\n",
      "agent\n",
      "{messages=[UserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\", name = \"execQuery\", arguments = \"{\"query\":\"current weather in San Francisco\"}\" }] }]}\n"
     ]
    }
   ],
   "source": [
    "var memory = new MemorySaver();\n",
    "\n",
    "var compileConfig = CompileConfig.builder()\n",
    "                    .checkpointSaver(memory)\n",
    "                    .interruptBefore( \"tools\")\n",
    "                    .build();\n",
    "\n",
    "var graphWithInterrupt = workflow.compile(compileConfig);\n",
    "\n",
    "var runnableConfig =  RunnableConfig.builder()\n",
    "                .threadId(\"conversation-2\" )\n",
    "                .build();\n",
    "\n",
    "Map<String,Object> inputs = Map.of( \"messages\", UserMessage.from(\"What's the weather like in SF currently?\") );\n",
    "\n",
    "var result = graphWithInterrupt.stream( inputs, runnableConfig );\n",
    "\n",
    "for( var r : result ) {\n",
    "  System.out.println( r.node() );\n",
    "  System.out.println( r.state() );\n",
    "  \n",
    "}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get State\n",
    "\n",
    "You can fetch the latest graph checkpoint using `getState(config)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tools"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "var snapshot = graphWithInterrupt.getState(runnableConfig);\n",
    "snapshot.getNext();\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Resume\n",
    "\n",
    "You can resume by running the graph with a null input. The checkpoint is loaded, and with no new inputs, it will execute as if no interrupt had occurred."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "RESUME REQUEST \n",
      "RESUME FROM agent \n",
      "execute: execQuery \n",
      "RESULT:\n",
      "tools\n",
      "{messages=[UserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\", name = \"execQuery\", arguments = \"{\"query\":\"current weather in San Francisco\"}\" }] }, ToolExecutionResultMessage { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\" toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }]} \n",
      "RESULT:\n",
      "agent\n",
      "{messages=[UserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\", name = \"execQuery\", arguments = \"{\"query\":\"current weather in San Francisco\"}\" }] }, ToolExecutionResultMessage { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\" toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }, AiMessage { text = \"The current weather in San Francisco is cold, with a low of 13 degrees Celsius.\" toolExecutionRequests = null }]} \n",
      "RESULT:\n",
      "__END__\n",
      "{messages=[UserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\", name = \"execQuery\", arguments = \"{\"query\":\"current weather in San Francisco\"}\" }] }, ToolExecutionResultMessage { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\" toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }, AiMessage { text = \"The current weather in San Francisco is cold, with a low of 13 degrees Celsius.\" toolExecutionRequests = null }]} \n"
     ]
    }
   ],
   "source": [
    "var result = graphWithInterrupt.stream( null, snapshot.getConfig() );\n",
    "\n",
    "for( var r : result ) {\n",
    "  log.trace( \"RESULT:\\n{}\\n{}\", r.node(), r.state() );\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check full history\n",
    "\n",
    "Let's browse the history of this thread, from newest to oldest."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "---\n",
      "StateSnapshot{node=agent, state={messages=[UserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\", name = \"execQuery\", arguments = \"{\"query\":\"current weather in San Francisco\"}\" }] }, ToolExecutionResultMessage { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\" toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }, AiMessage { text = \"The current weather in San Francisco is cold, with a low of 13 degrees Celsius.\" toolExecutionRequests = null }]}, config=RunnableConfig(threadId=conversation-2, checkPointId=af521ddb-0f26-4c5b-8610-4d374f6818f5, nextNode=__END__, streamMode=VALUES)}\n",
      "--- \n",
      "\n",
      "---\n",
      "StateSnapshot{node=tools, state={messages=[UserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\", name = \"execQuery\", arguments = \"{\"query\":\"current weather in San Francisco\"}\" }] }, ToolExecutionResultMessage { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\" toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }]}, config=RunnableConfig(threadId=conversation-2, checkPointId=b1f96095-9651-4af5-b058-0089f47b17a5, nextNode=agent, streamMode=VALUES)}\n",
      "--- \n",
      "\n",
      "---\n",
      "StateSnapshot{node=agent, state={messages=[UserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\", name = \"execQuery\", arguments = \"{\"query\":\"current weather in San Francisco\"}\" }] }, ToolExecutionResultMessage { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\" toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }, AiMessage { text = \"The current weather in San Francisco is cold, with a low of 13 degrees Celsius.\" toolExecutionRequests = null }]}, config=RunnableConfig(threadId=conversation-2, checkPointId=22008a3a-fd9d-4899-92e2-7151875de7f0, nextNode=tools, streamMode=VALUES)}\n",
      "--- \n",
      "\n",
      "---\n",
      "StateSnapshot{node=__START__, state={messages=[UserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }]}, config=RunnableConfig(threadId=conversation-2, checkPointId=8abd26da-710d-4a98-97da-6f63ae4ac4d5, nextNode=agent, streamMode=VALUES)}\n",
      "--- \n"
     ]
    }
   ],
   "source": [
    "RunnableConfig toReplay = null;\n",
    "var states = graphWithInterrupt.getStateHistory(runnableConfig);\n",
    "for( var state: states ) {\n",
    "  \n",
    "  log.trace( \"\\n---\\n{}\\n---\",state);\n",
    "\n",
    "  if (state.getState().messages().size() == 3) {\n",
    "     toReplay = state.getConfig();\n",
    "  }\n",
    "}\n",
    "if (toReplay==null) {\n",
    "  throw new IllegalStateException(\"No state to replay\");\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Replay a past state\n",
    "\n",
    "To replay from this place we just need to pass its config back to the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "RESUME REQUEST \n",
      "RESUME FROM tools \n",
      "RESULT:\n",
      "agent\n",
      "{messages=[UserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\", name = \"execQuery\", arguments = \"{\"query\":\"current weather in San Francisco\"}\" }] }, ToolExecutionResultMessage { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\" toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }, AiMessage { text = \"The current weather in San Francisco is cold, with a low of 13 degrees Celsius.\" toolExecutionRequests = null }]}\n",
      "--- \n",
      "RESULT:\n",
      "__END__\n",
      "{messages=[UserMessage { name = null contents = [TextContent { text = \"What's the weather like in SF currently?\" }] }, AiMessage { text = null toolExecutionRequests = [ToolExecutionRequest { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\", name = \"execQuery\", arguments = \"{\"query\":\"current weather in San Francisco\"}\" }] }, ToolExecutionResultMessage { id = \"call_HmirEZJNwBCzKSQauWGHPCiD\" toolName = \"execQuery\" text = \"Cold, with a low of 13 degrees\" }, AiMessage { text = \"The current weather in San Francisco is cold, with a low of 13 degrees Celsius.\" toolExecutionRequests = null }]}\n",
      "--- \n"
     ]
    }
   ],
   "source": [
    "var results = graphWithInterrupt.stream(null, toReplay ); \n",
    "\n",
    "for( var r : results ) {\n",
    "  log.trace( \"RESULT:\\n{}\\n{}\\n---\", r.node(), r.state() );\n",
    "}\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Java (rjk 2.2.0)",
   "language": "java",
   "name": "rapaio-jupyter-kernel"
  },
  "language_info": {
   "codemirror_mode": "java",
   "file_extension": ".jshell",
   "mimetype": "text/x-java-source",
   "name": "java",
   "nbconvert_exporter": "script",
   "pygments_lexer": "java",
   "version": "17.0.2+8-86"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}