{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "47ed5db3-bda5-49e1-bf75-23e08c9a3af0",
   "metadata": {},
   "source": [
    "# How to pass private state\n",
    "\n",
    "Oftentimes, you may want nodes to be able to pass state to each other that should NOT be part of the main schema of the graph. This is often useful because there may be information that is not needed as input/output (and therefore doesn't really make sense to have in the main schema) but is ABSOLUTELY needed as part of the intermediate working logic.\n",
    "\n",
    "Let's take a look at an example below. In this example, we will create a RAG pipeline that:\n",
    "1. Takes in a user question\n",
    "2. Uses an LLM to generate a search query\n",
    "3. Retrieves documents for that generated query\n",
    "4. Generates a final answer based on those documents\n",
    "\n",
    "We will have a separate node for each step. We will only have the `question` and `answer` on the overall state. However, we will need separate states for the `search_query` and the `documents` - we will pass these as private state keys.\n",
    "\n",
    "Let's look at an example!\n",
    "\n",
    "## Setup\n",
    "\n",
    "First, let's install the required packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32d79ebd",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langgraph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e30836ce",
   "metadata": {},
   "source": [
    "<div class=\"admonition tip\">\n",
    "    <p class=\"admonition-title\">Set up <a href=\"https://smith.langchain.com\">LangSmith</a> for LangGraph development</p>\n",
    "    <p style=\"padding-top: 5px;\">\n",
    "        Sign up for LangSmith to quickly spot issues and improve the performance of your LangGraph projects. LangSmith lets you use trace data to debug, test, and monitor your LLM apps built with LangGraph — read more about how to get started <a href=\"https://docs.smith.langchain.com\">here</a>. \n",
    "    </p>\n",
    "</div>    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0028ced",
   "metadata": {},
   "source": [
    "## Define and use the graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3114c3ad-0ade-47ba-9488-53d6f7671578",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'question': 'foo', 'answer': 'fo\\n\\nfo\\n\\nfoo'}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langgraph.graph import StateGraph, START, END\n",
    "from typing import TypedDict\n",
    "\n",
    "\n",
    "# The overall state of the graph\n",
    "class OverallState(TypedDict):\n",
    "    question: str\n",
    "    answer: str\n",
    "\n",
    "\n",
    "# This is what the node that generates the query will return\n",
    "class QueryOutputState(TypedDict):\n",
    "    query: str\n",
    "\n",
    "\n",
    "# This is what the node that retrieves the documents will return\n",
    "class DocumentOutputState(TypedDict):\n",
    "    docs: list[str]\n",
    "\n",
    "\n",
    "# This is what the node that generates the final answer will take in\n",
    "class GenerateInputState(OverallState, DocumentOutputState):\n",
    "    pass\n",
    "\n",
    "\n",
    "# Node to generate query\n",
    "def generate_query(state: OverallState) -> QueryOutputState:\n",
    "    # Replace this with real logic\n",
    "    return {\"query\": state[\"question\"][:2]}\n",
    "\n",
    "\n",
    "# Node to retrieve documents\n",
    "def retrieve_documents(state: QueryOutputState) -> DocumentOutputState:\n",
    "    # Replace this with real logic\n",
    "    return {\"docs\": [state[\"query\"]] * 2}\n",
    "\n",
    "\n",
    "# Node to generate answer\n",
    "def generate(state: GenerateInputState) -> OverallState:\n",
    "    return {\"answer\": \"\\n\\n\".join(state[\"docs\"] + [state[\"question\"]])}\n",
    "\n",
    "\n",
    "graph = StateGraph(OverallState)\n",
    "graph.add_node(generate_query)\n",
    "graph.add_node(retrieve_documents)\n",
    "graph.add_node(generate)\n",
    "graph.add_edge(START, \"generate_query\")\n",
    "graph.add_edge(\"generate_query\", \"retrieve_documents\")\n",
    "graph.add_edge(\"retrieve_documents\", \"generate\")\n",
    "graph.add_edge(\"generate\", END)\n",
    "graph = graph.compile()\n",
    "\n",
    "graph.invoke({\"question\": \"foo\"})"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}