{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1f53f753-12c6-4fac-b910-6e96677d8a49",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/meta-llama/llama-recipes/blob/main/recipes/3p_integrations/langchain/langgraph_rag_agent_local.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b9ab14a-fd80-4ca2-afc5-efe1c39532bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -U langchain_community tiktoken langchainhub chromadb langchain langgraph tavily-python sentence-transformers"
   ]
  },
  {
   "attachments": {
    "7b00797e-fb85-4474-9a9e-c505b61add81.png": {
     "image/png": "[encoded data, 243196 characters]"
    }
   },
   "cell_type": "markdown",
   "id": "0216de30-29cf-4464-9cc3-6e9a6d6c3e40",
   "metadata": {},
   "source": [
    "# Local LangGraph RAG agent with Llama 3\n",
    "\n",
    "Previously, we showed how to build simple agents with LangGraph and Llama 3.\n",
    "\n",
    "Now, we'll pick a more advanced use-case: advanced RAG, with the requirment that it runs locally.\n",
    "\n",
    "## Ideas\n",
    "\n",
    "We'll combine ideas from three RAG papers into a RAG agent:\n",
    "\n",
    "- **Routing:**  Adaptive RAG ([paper](https://arxiv.org/abs/2403.14403)). Route questions to different retrieval approaches\n",
    "- **Fallback:** Corrective RAG ([paper](https://arxiv.org/pdf/2401.15884.pdf)). Fallback to web search if docs are not relevant to query\n",
    "- **Self-correction:** Self-RAG ([paper](https://arxiv.org/abs/2310.11511)). Fix answers w/ hallucinations or don’t address question\n",
    "\n",
    "![langgraph_adaptive_rag.png](attachment:7b00797e-fb85-4474-9a9e-c505b61add81.png)\n",
    "\n",
    "Note that this will incorperate [a few general ideas for agents](https://www.deeplearning.ai/the-batch/how-agents-can-improve-llm-performance/):\n",
    "\n",
    "- **Reflection**: The self-correction mechanism is a form of reflection, where the LangGraph agent reflects on its retrieval and generations\n",
    "- **Planning**: The control flow laid out in the graph is a form of planning \n",
    "- **Tool use**: Specific nodes in the control flow (e.g., web search) will use tools\n",
    "\n",
    "## Local models\n",
    "\n",
    "#### Embedding\n",
    " \n",
    "[GPT4All Embeddings](https://blog.nomic.ai/posts/nomic-embed-text-v1):\n",
    "\n",
    "```\n",
    "pip install langchain-nomic\n",
    "```\n",
    "\n",
    "### LLM\n",
    "\n",
    "Use [Ollama](https://ollama.ai/) and [llama3](https://ollama.ai/library/llama3):\n",
    "\n",
    "```\n",
    "ollama pull llama3\n",
    "```\n",
    "\n",
    "Prompt - \n",
    "\n",
    "https://llama.meta.com/docs/model-cards-and-prompt-formats/meta-llama-3/\n",
    "\n",
    "### Tracing\n",
    "\n",
    "```\n",
    "### Tracing (optional)\n",
    "os.environ['LANGCHAIN_TRACING_V2'] = 'true'\n",
    "os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'\n",
    "os.environ['LANGCHAIN_API_KEY'] = 'LANGCHAIN_API_KEY'\n",
    "```\n",
    "\n",
    "### Search\n",
    "\n",
    "Uses [Tavily](https://tavily.com/#api)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d2242c5-0dbb-4f21-9771-92f6d679b1a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ['LANGCHAIN_TRACING_V2'] = 'true'\n",
    "os.environ['LANGCHAIN_ENDPOINT'] = 'https://api.smith.langchain.com'\n",
    "os.environ['LANGCHAIN_API_KEY'] = 'LANGCHAIN_API_KEY'\n",
    "\n",
    "os.environ['TAVILY_API_KEY'] = 'TAVILY_API_KEY'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2096d49c-d3dc-4329-ada7-aff56d210198",
   "metadata": {},
   "outputs": [],
   "source": [
    "### LLM\n",
    "\n",
    "local_llm = 'llama3'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "267c63e1-4c2f-439d-8d95-4c6aa01f41cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Index\n",
    "\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from langchain_community.document_loaders import WebBaseLoader\n",
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_community.embeddings import HuggingFaceEmbeddings\n",
    "\n",
    "urls = [\n",
    "    \"https://lilianweng.github.io/posts/2023-06-23-agent/\",\n",
    "    \"https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/\",\n",
    "    \"https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/\",\n",
    "]\n",
    "\n",
    "docs = [WebBaseLoader(url).load() for url in urls]\n",
    "docs_list = [item for sublist in docs for item in sublist]\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=250, chunk_overlap=0\n",
    ")\n",
    "doc_splits = text_splitter.split_documents(docs_list)\n",
    "\n",
    "# Add to vectorDB\n",
    "vectorstore = Chroma.from_documents(\n",
    "    documents=doc_splits,\n",
    "    collection_name=\"rag-chroma\",\n",
    "    embedding=HuggingFaceEmbeddings(),\n",
    ")\n",
    "retriever = vectorstore.as_retriever()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b008df98-8394-49da-8fb8-aefe2c90d03c",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Retrieval Grader \n",
    "\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain_community.chat_models import ChatOllama\n",
    "from langchain_core.output_parsers import JsonOutputParser\n",
    "\n",
    "# LLM\n",
    "llm = ChatOllama(model=local_llm, format=\"json\", temperature=0)\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    template=\"\"\"You are a grader assessing relevance \n",
    "    of a retrieved document to a user question. If the document contains keywords related to the user question, \n",
    "    grade it as relevant. It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \n",
    "    \n",
    "    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.\n",
    "    Provide the binary score as a JSON with a single key 'score' and no premable or explaination.\n",
    "     \n",
    "    Here is the retrieved document: \n",
    "    {document}\n",
    "    \n",
    "    Here is the user question: \n",
    "    {question}\n",
    "    \"\"\",\n",
    "    input_variables=[\"question\", \"document\"],\n",
    ")\n",
    "\n",
    "retrieval_grader = prompt | llm | JsonOutputParser()\n",
    "question = \"agent memory\"\n",
    "docs = retriever.invoke(question)\n",
    "doc_txt = docs[1].page_content\n",
    "print(retrieval_grader.invoke({\"question\": question, \"document\": doc_txt}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d531a81-6d4d-405e-975a-01ef1c9679fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Generate\n",
    "\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain import hub\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "# Prompt\n",
    "prompt = PromptTemplate(\n",
    "    template=\"\"\"You are an assistant for question-answering tasks. \n",
    "    Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. \n",
    "    Use three sentences maximum and keep the answer concise:\n",
    "    Question: {question} \n",
    "    Context: {context} \n",
    "    Answer: \n",
    "    \"\"\",\n",
    "    input_variables=[\"question\", \"document\"],\n",
    ")\n",
    "\n",
    "llm = ChatOllama(model=local_llm, temperature=0)\n",
    "\n",
    "# Post-processing\n",
    "def format_docs(docs):\n",
    "    return \"\\n\\n\".join(doc.page_content for doc in docs)\n",
    "\n",
    "# Chain\n",
    "rag_chain = prompt | llm | StrOutputParser()\n",
    "\n",
    "# Run\n",
    "question = \"agent memory\"\n",
    "docs = retriever.invoke(question)\n",
    "generation = rag_chain.invoke({\"context\": docs, \"question\": question})\n",
    "print(generation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0261a9a4-de13-4dd8-b082-95305a3e43ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Hallucination Grader \n",
    "\n",
    "# LLM\n",
    "llm = ChatOllama(model=local_llm, format=\"json\", temperature=0)\n",
    "\n",
    "# Prompt\n",
    "prompt = PromptTemplate(\n",
    "    template=\"\"\"You are a grader assessing whether \n",
    "    an answer is grounded in / supported by a set of facts. Give a binary score 'yes' or 'no' score to indicate \n",
    "    whether the answer is grounded in / supported by a set of facts. Provide the binary score as a JSON with a \n",
    "    single key 'score' and no preamble or explanation.\n",
    "    \n",
    "    Here are the facts:\n",
    "    {documents} \n",
    "\n",
    "    Here is the answer: \n",
    "    {generation}\n",
    "    \"\"\",\n",
    "    input_variables=[\"generation\", \"documents\"],\n",
    ")\n",
    "\n",
    "hallucination_grader = prompt | llm | JsonOutputParser()\n",
    "hallucination_grader.invoke({\"documents\": docs, \"generation\": generation})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df9f6944-4fee-4971-b3a7-2b81b44ed433",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Answer Grader \n",
    "\n",
    "# LLM\n",
    "llm = ChatOllama(model=local_llm, format=\"json\", temperature=0)\n",
    "\n",
    "# Prompt\n",
    "prompt = PromptTemplate(\n",
    "    template=\"\"\"You are a grader assessing whether an \n",
    "    answer is useful to resolve a question. Give a binary score 'yes' or 'no' to indicate whether the answer is \n",
    "    useful to resolve a question. Provide the binary score as a JSON with a single key 'score' and no preamble or explanation.\n",
    "     \n",
    "    Here is the answer:\n",
    "    {generation} \n",
    "\n",
    "    Here is the question: {question}\n",
    "    \"\"\",\n",
    "    input_variables=[\"generation\", \"question\"],\n",
    ")\n",
    "\n",
    "answer_grader = prompt | llm | JsonOutputParser()\n",
    "answer_grader.invoke({\"question\": question,\"generation\": generation})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9c910c1-738c-4bf7-bf9e-801862b227eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Router\n",
    "\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain_community.chat_models import ChatOllama\n",
    "from langchain_core.output_parsers import JsonOutputParser\n",
    "\n",
    "# LLM\n",
    "llm = ChatOllama(model=local_llm, format=\"json\", temperature=0)\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    template=\"\"\"You are an expert at routing a \n",
    "    user question to a vectorstore or web search. Use the vectorstore for questions on LLM  agents, \n",
    "    prompt engineering, and adversarial attacks. You do not need to be stringent with the keywords \n",
    "    in the question related to these topics. Otherwise, use web-search. Give a binary choice 'web_search' \n",
    "    or 'vectorstore' based on the question. Return the a JSON with a single key 'datasource' and \n",
    "    no premable or explaination. \n",
    "    \n",
    "    Question to route: \n",
    "    {question}\"\"\",\n",
    "    input_variables=[\"question\"],\n",
    ")\n",
    "\n",
    "question_router = prompt | llm | JsonOutputParser()\n",
    "question = \"llm agent memory\"\n",
    "docs = retriever.get_relevant_documents(question)\n",
    "doc_txt = docs[1].page_content\n",
    "print(question_router.invoke({\"question\": question}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "023ff2db-eb4e-4d44-904c-ea061abc16d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Search\n",
    "\n",
    "from langchain_community.tools.tavily_search import TavilySearchResults\n",
    "web_search_tool = TavilySearchResults(k=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccd59cdf-a04d-4b2e-b9cc-6a1b1e80a6c6",
   "metadata": {},
   "source": [
    "We'll implement these as a control flow in LangGraph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07fa3d08-6a86-4705-a28b-e2721070bc5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing_extensions import TypedDict\n",
    "from typing import List\n",
    "\n",
    "### State\n",
    "\n",
    "class GraphState(TypedDict):\n",
    "    \"\"\"\n",
    "    Represents the state of our graph.\n",
    "\n",
    "    Attributes:\n",
    "        question: question\n",
    "        generation: LLM generation\n",
    "        web_search: whether to add search\n",
    "        documents: list of documents \n",
    "    \"\"\"\n",
    "    question : str\n",
    "    generation : str\n",
    "    web_search : str\n",
    "    documents : List[str]\n",
    "\n",
    "from langchain.schema import Document\n",
    "\n",
    "### Nodes\n",
    "\n",
    "def retrieve(state):\n",
    "    \"\"\"\n",
    "    Retrieve documents from vectorstore\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): New key added to state, documents, that contains retrieved documents\n",
    "    \"\"\"\n",
    "    print(\"---RETRIEVE---\")\n",
    "    question = state[\"question\"]\n",
    "\n",
    "    # Retrieval\n",
    "    documents = retriever.invoke(question)\n",
    "    return {\"documents\": documents, \"question\": question}\n",
    "\n",
    "def generate(state):\n",
    "    \"\"\"\n",
    "    Generate answer using RAG on retrieved documents\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): New key added to state, generation, that contains LLM generation\n",
    "    \"\"\"\n",
    "    print(\"---GENERATE---\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "    \n",
    "    # RAG generation\n",
    "    generation = rag_chain.invoke({\"context\": documents, \"question\": question})\n",
    "    return {\"documents\": documents, \"question\": question, \"generation\": generation}\n",
    "\n",
    "def grade_documents(state):\n",
    "    \"\"\"\n",
    "    Determines whether the retrieved documents are relevant to the question\n",
    "    If any document is not relevant, we will set a flag to run web search\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): Filtered out irrelevant documents and updated web_search state\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---CHECK DOCUMENT RELEVANCE TO QUESTION---\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "    \n",
    "    # Score each doc\n",
    "    filtered_docs = []\n",
    "    web_search = \"No\"\n",
    "    for d in documents:\n",
    "        score = retrieval_grader.invoke({\"question\": question, \"document\": d.page_content})\n",
    "        grade = score['score']\n",
    "        # Document relevant\n",
    "        if grade.lower() == \"yes\":\n",
    "            print(\"---GRADE: DOCUMENT RELEVANT---\")\n",
    "            filtered_docs.append(d)\n",
    "        # Document not relevant\n",
    "        else:\n",
    "            print(\"---GRADE: DOCUMENT NOT RELEVANT---\")\n",
    "            # We do not include the document in filtered_docs\n",
    "            # We set a flag to indicate that we want to run web search\n",
    "            web_search = \"Yes\"\n",
    "            continue\n",
    "    return {\"documents\": filtered_docs, \"question\": question, \"web_search\": web_search}\n",
    "    \n",
    "def web_search(state):\n",
    "    \"\"\"\n",
    "    Web search based based on the question\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): Appended web results to documents\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---WEB SEARCH---\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "\n",
    "    # Web search\n",
    "    docs = web_search_tool.invoke({\"query\": question})\n",
    "    web_results = \"\\n\".join([d[\"content\"] for d in docs])\n",
    "    web_results = Document(page_content=web_results)\n",
    "    if documents is not None:\n",
    "        documents.append(web_results)\n",
    "    else:\n",
    "        documents = [web_results]\n",
    "    return {\"documents\": documents, \"question\": question}\n",
    "\n",
    "### Conditional edge\n",
    "\n",
    "def route_question(state):\n",
    "    \"\"\"\n",
    "    Route question to web search or RAG.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        str: Next node to call\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---ROUTE QUESTION---\")\n",
    "    question = state[\"question\"]\n",
    "    print(question)\n",
    "    source = question_router.invoke({\"question\": question})  \n",
    "    print(source)\n",
    "    print(source['datasource'])\n",
    "    if source['datasource'] == 'web_search':\n",
    "        print(\"---ROUTE QUESTION TO WEB SEARCH---\")\n",
    "        return \"websearch\"\n",
    "    elif source['datasource'] == 'vectorstore':\n",
    "        print(\"---ROUTE QUESTION TO RAG---\")\n",
    "        return \"vectorstore\"\n",
    "\n",
    "def decide_to_generate(state):\n",
    "    \"\"\"\n",
    "    Determines whether to generate an answer, or add web search\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        str: Binary decision for next node to call\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---ASSESS GRADED DOCUMENTS---\")\n",
    "    question = state[\"question\"]\n",
    "    web_search = state[\"web_search\"]\n",
    "    filtered_documents = state[\"documents\"]\n",
    "\n",
    "    if web_search == \"Yes\":\n",
    "        # All documents have been filtered check_relevance\n",
    "        # We will re-generate a new query\n",
    "        print(\"---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, INCLUDE WEB SEARCH---\")\n",
    "        return \"websearch\"\n",
    "    else:\n",
    "        # We have relevant documents, so generate answer\n",
    "        print(\"---DECISION: GENERATE---\")\n",
    "        return \"generate\"\n",
    "\n",
    "### Conditional edge\n",
    "\n",
    "def grade_generation_v_documents_and_question(state):\n",
    "    \"\"\"\n",
    "    Determines whether the generation is grounded in the document and answers question.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        str: Decision for next node to call\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---CHECK HALLUCINATIONS---\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "    generation = state[\"generation\"]\n",
    "\n",
    "    score = hallucination_grader.invoke({\"documents\": documents, \"generation\": generation})\n",
    "    grade = score['score']\n",
    "\n",
    "    # Check hallucination\n",
    "    if grade == \"yes\":\n",
    "        print(\"---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\")\n",
    "        # Check question-answering\n",
    "        print(\"---GRADE GENERATION vs QUESTION---\")\n",
    "        score = answer_grader.invoke({\"question\": question,\"generation\": generation})\n",
    "        grade = score['score']\n",
    "        if grade == \"yes\":\n",
    "            print(\"---DECISION: GENERATION ADDRESSES QUESTION---\")\n",
    "            return \"useful\"\n",
    "        else:\n",
    "            print(\"---DECISION: GENERATION DOES NOT ADDRESS QUESTION---\")\n",
    "            return \"not useful\"\n",
    "    else:\n",
    "        pprint(\"---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---\")\n",
    "        return \"not supported\"\n",
    "\n",
    "from langgraph.graph import END, StateGraph\n",
    "workflow = StateGraph(GraphState)\n",
    "\n",
    "# Define the nodes\n",
    "workflow.add_node(\"websearch\", web_search) # web search\n",
    "workflow.add_node(\"retrieve\", retrieve) # retrieve\n",
    "workflow.add_node(\"grade_documents\", grade_documents) # grade documents\n",
    "workflow.add_node(\"generate\", generate) # generatae"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73f21594-00d4-48a8-ae2e-4e55a010b540",
   "metadata": {},
   "source": [
    "### Graph Build"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9a4b9e4-3ba8-47d6-958c-e5a7112ac6f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build graph\n",
    "workflow.set_conditional_entry_point(\n",
    "    route_question,\n",
    "    {\n",
    "        \"websearch\": \"websearch\",\n",
    "        \"vectorstore\": \"retrieve\",\n",
    "    },\n",
    ")\n",
    "\n",
    "workflow.add_edge(\"retrieve\", \"grade_documents\")\n",
    "workflow.add_conditional_edges(\n",
    "    \"grade_documents\",\n",
    "    decide_to_generate,\n",
    "    {\n",
    "        \"websearch\": \"websearch\",\n",
    "        \"generate\": \"generate\",\n",
    "    },\n",
    ")\n",
    "workflow.add_edge(\"websearch\", \"generate\")\n",
    "workflow.add_conditional_edges(\n",
    "    \"generate\",\n",
    "    grade_generation_v_documents_and_question,\n",
    "    {\n",
    "        \"not supported\": \"generate\",\n",
    "        \"useful\": END,\n",
    "        \"not useful\": \"websearch\",\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13043b0f-17c7-49d3-9ea7-8f2c0f0c8691",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile\n",
    "app = workflow.compile()\n",
    "\n",
    "# Test\n",
    "from pprint import pprint\n",
    "inputs = {\"question\": \"What are the types of agent memory?\"}\n",
    "for output in app.stream(inputs):\n",
    "    for key, value in output.items():\n",
    "        pprint(f\"Finished running: {key}:\")\n",
    "pprint(value[\"generation\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d733ab80-7e7b-4b1a-9519-9242de647eda",
   "metadata": {},
   "source": [
    "Trace: \n",
    "\n",
    "https://smith.langchain.com/public/8d449b67-6bc4-4ecf-9153-759cd21df24f/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbfcec3e-a09a-40b4-9c15-fead97bf4e0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compile\n",
    "app = workflow.compile()\n",
    "\n",
    "# Test\n",
    "from pprint import pprint\n",
    "inputs = {\"question\": \"Who are the Bears expected to draft first in the NFL draft?\"}\n",
    "for output in app.stream(inputs):\n",
    "    for key, value in output.items():\n",
    "        pprint(f\"Finished running: {key}:\")\n",
    "pprint(value[\"generation\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9051bea-13fa-4eb3-b671-16f59755931d",
   "metadata": {},
   "source": [
    "Trace: \n",
    "\n",
    "https://smith.langchain.com/public/c785f9c0-f519-4a38-ad5a-febb59a2139c/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1059b3e-d197-47fc-b55f-82a3406016f3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
