08_Adaptive_Rag_Agent.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5afcaed0-3d55-4e1f-95d3-c32c751c29d8",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# Adaptive RAG\n",
    "\n",
    "Adaptive RAG is a strategy for RAG that unites (1) [query analysis](https://blog.langchain.dev/query-construction/) with (2) [active / self-corrective RAG](https://blog.langchain.dev/agentic-rag-with-langgraph/).\n",
    "\n",
    "In the [paper](https://arxiv.org/abs/2403.14403), they report query analysis to route across:\n",
    "\n",
    "* No Retrieval\n",
    "* Single-shot RAG\n",
    "* Iterative RAG\n",
    "\n",
    "Let's build on this using LangGraph. \n",
    "\n",
    "In our implementation, we will route between:\n",
    "\n",
    "* Web search: for questions related to recent events\n",
    "* Self-corrective RAG: for questions related to our index\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a85501ca-eb89-4795-aeab-cdab050ead6b",
   "metadata": {},
   "source": [
    "# Enviorment "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "53d1a740-9fea-4a6e-8f95-fb9dbf1c80a1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: langchain_community in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (0.0.27)\n",
      "Collecting langchain_community\n",
      "  Downloading langchain_community-0.0.31-py3-none-any.whl.metadata (8.4 kB)\n",
      "Requirement already satisfied: tiktoken in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (0.5.2)\n",
      "Collecting tiktoken\n",
      "  Downloading tiktoken-0.6.0-cp311-cp311-macosx_11_0_arm64.whl.metadata (6.6 kB)\n",
      "Requirement already satisfied: langchain-openai in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (0.0.2.post1)\n",
      "Collecting langchain-openai\n",
      "  Downloading langchain_openai-0.1.1-py3-none-any.whl.metadata (2.5 kB)\n",
      "Requirement already satisfied: langchain-cohere in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (0.1.0)\n",
      "Requirement already satisfied: langchainhub in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (0.1.15)\n",
      "Requirement already satisfied: chromadb in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (0.4.24)\n",
      "Requirement already satisfied: langchain in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (0.1.11)\n",
      "Collecting langchain\n",
      "  Downloading langchain-0.1.14-py3-none-any.whl.metadata (13 kB)\n",
      "Requirement already satisfied: langgraph in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (0.0.30)\n",
      "Collecting langgraph\n",
      "  Downloading langgraph-0.0.31-py3-none-any.whl.metadata (44 kB)\n",
      "\u001b[2K     \u001b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.8/44.8 kB\u001b[0m \u001b[31m782.7 kB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m31m1.0 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: tavily-python in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (0.3.1)\n",
      "Collecting tavily-python\n",
      "  Downloading tavily_python-0.3.3-py3-none-any.whl.metadata (4.4 kB)\n",
      "Requirement already satisfied: PyYAML>=5.3 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain_community) (6.0.1)\n",
      "Requirement already satisfied: SQLAlchemy<3,>=1.4 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain_community) (2.0.28)\n",
      "Requirement already satisfied: aiohttp<4.0.0,>=3.8.3 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain_community) (3.9.3)\n",
      "Requirement already satisfied: dataclasses-json<0.7,>=0.5.7 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain_community) (0.6.4)\n",
      "Requirement already satisfied: langchain-core<0.2.0,>=0.1.37 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain_community) (0.1.38)\n",
      "Requirement already satisfied: langsmith<0.2.0,>=0.1.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain_community) (0.1.23)\n",
      "Requirement already satisfied: numpy<2,>=1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain_community) (1.26.4)\n",
      "Requirement already satisfied: requests<3,>=2 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain_community) (2.31.0)\n",
      "Requirement already satisfied: tenacity<9.0.0,>=8.1.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain_community) (8.2.3)\n",
      "Requirement already satisfied: regex>=2022.1.18 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from tiktoken) (2023.12.25)\n",
      "Requirement already satisfied: openai<2.0.0,>=1.10.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain-openai) (1.13.3)\n",
      "Requirement already satisfied: cohere<6.0.0,>=5.1.4 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain-cohere) (5.1.7)\n",
      "Requirement already satisfied: types-requests<3.0.0.0,>=2.31.0.2 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchainhub) (2.31.0.20240311)\n",
      "Requirement already satisfied: build>=1.0.3 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (1.2.1)\n",
      "Requirement already satisfied: pydantic>=1.9 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (2.6.4)\n",
      "Requirement already satisfied: chroma-hnswlib==0.7.3 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (0.7.3)\n",
      "Requirement already satisfied: fastapi>=0.95.2 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (0.109.0)\n",
      "Requirement already satisfied: uvicorn>=0.18.3 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (0.25.0)\n",
      "Requirement already satisfied: posthog>=2.4.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (3.5.0)\n",
      "Requirement already satisfied: typing-extensions>=4.5.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (4.10.0)\n",
      "Requirement already satisfied: pulsar-client>=3.1.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (3.4.0)\n",
      "Requirement already satisfied: onnxruntime>=1.14.1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (1.17.1)\n",
      "Requirement already satisfied: opentelemetry-api>=1.2.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (1.24.0)\n",
      "Requirement already satisfied: opentelemetry-exporter-otlp-proto-grpc>=1.2.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (1.24.0)\n",
      "Requirement already satisfied: opentelemetry-instrumentation-fastapi>=0.41b0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (0.45b0)\n",
      "Requirement already satisfied: opentelemetry-sdk>=1.2.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (1.24.0)\n",
      "Requirement already satisfied: tokenizers>=0.13.2 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (0.15.2)\n",
      "Requirement already satisfied: pypika>=0.48.9 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (0.48.9)\n",
      "Requirement already satisfied: tqdm>=4.65.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (4.66.2)\n",
      "Requirement already satisfied: overrides>=7.3.1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (7.7.0)\n",
      "Requirement already satisfied: importlib-resources in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (6.4.0)\n",
      "Requirement already satisfied: grpcio>=1.58.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (1.62.1)\n",
      "Requirement already satisfied: bcrypt>=4.0.1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (4.1.2)\n",
      "Requirement already satisfied: typer>=0.9.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (0.12.0)\n",
      "Requirement already satisfied: kubernetes>=28.1.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (29.0.0)\n",
      "Requirement already satisfied: mmh3>=4.0.1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (4.1.0)\n",
      "Requirement already satisfied: orjson>=3.9.12 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from chromadb) (3.9.15)\n",
      "Requirement already satisfied: jsonpatch<2.0,>=1.33 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain) (1.33)\n",
      "Requirement already satisfied: langchain-text-splitters<0.1,>=0.0.1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from langchain) (0.0.1)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain_community) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain_community) (23.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain_community) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain_community) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from aiohttp<4.0.0,>=3.8.3->langchain_community) (1.9.4)\n",
      "Requirement already satisfied: packaging>=19.1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from build>=1.0.3->chromadb) (23.2)\n",
      "Requirement already satisfied: pyproject_hooks in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from build>=1.0.3->chromadb) (1.0.0)\n",
      "Requirement already satisfied: fastavro<2.0.0,>=1.9.4 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from cohere<6.0.0,>=5.1.4->langchain-cohere) (1.9.4)\n",
      "Requirement already satisfied: httpx>=0.21.2 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from cohere<6.0.0,>=5.1.4->langchain-cohere) (0.26.0)\n",
      "Requirement already satisfied: marshmallow<4.0.0,>=3.18.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from dataclasses-json<0.7,>=0.5.7->langchain_community) (3.21.1)\n",
      "Requirement already satisfied: typing-inspect<1,>=0.4.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from dataclasses-json<0.7,>=0.5.7->langchain_community) (0.9.0)\n",
      "Requirement already satisfied: starlette<0.36.0,>=0.35.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from fastapi>=0.95.2->chromadb) (0.35.1)\n",
      "Requirement already satisfied: jsonpointer>=1.9 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from jsonpatch<2.0,>=1.33->langchain) (2.4)\n",
      "Requirement already satisfied: certifi>=14.05.14 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from kubernetes>=28.1.0->chromadb) (2024.2.2)\n",
      "Requirement already satisfied: six>=1.9.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from kubernetes>=28.1.0->chromadb) (1.16.0)\n",
      "Requirement already satisfied: python-dateutil>=2.5.3 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from kubernetes>=28.1.0->chromadb) (2.9.0.post0)\n",
      "Requirement already satisfied: google-auth>=1.0.1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from kubernetes>=28.1.0->chromadb) (2.29.0)\n",
      "Requirement already satisfied: websocket-client!=0.40.0,!=0.41.*,!=0.42.*,>=0.32.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from kubernetes>=28.1.0->chromadb) (1.7.0)\n",
      "Requirement already satisfied: requests-oauthlib in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from kubernetes>=28.1.0->chromadb) (2.0.0)\n",
      "Requirement already satisfied: oauthlib>=3.2.2 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from kubernetes>=28.1.0->chromadb) (3.2.2)\n",
      "Requirement already satisfied: urllib3>=1.24.2 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from kubernetes>=28.1.0->chromadb) (2.2.1)\n",
      "Requirement already satisfied: coloredlogs in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from onnxruntime>=1.14.1->chromadb) (15.0.1)\n",
      "Requirement already satisfied: flatbuffers in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from onnxruntime>=1.14.1->chromadb) (24.3.25)\n",
      "Requirement already satisfied: protobuf in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from onnxruntime>=1.14.1->chromadb) (4.25.3)\n",
      "Requirement already satisfied: sympy in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from onnxruntime>=1.14.1->chromadb) (1.12)\n",
      "Requirement already satisfied: anyio<5,>=3.5.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from openai<2.0.0,>=1.10.0->langchain-openai) (4.3.0)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from openai<2.0.0,>=1.10.0->langchain-openai) (1.9.0)\n",
      "Requirement already satisfied: sniffio in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from openai<2.0.0,>=1.10.0->langchain-openai) (1.3.1)\n",
      "Requirement already satisfied: deprecated>=1.2.6 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-api>=1.2.0->chromadb) (1.2.14)\n",
      "Requirement already satisfied: importlib-metadata<=7.0,>=6.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-api>=1.2.0->chromadb) (7.0.0)\n",
      "Requirement already satisfied: googleapis-common-protos~=1.52 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb) (1.63.0)\n",
      "Requirement already satisfied: opentelemetry-exporter-otlp-proto-common==1.24.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb) (1.24.0)\n",
      "Requirement already satisfied: opentelemetry-proto==1.24.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-exporter-otlp-proto-grpc>=1.2.0->chromadb) (1.24.0)\n",
      "Requirement already satisfied: opentelemetry-instrumentation-asgi==0.45b0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.45b0)\n",
      "Requirement already satisfied: opentelemetry-instrumentation==0.45b0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.45b0)\n",
      "Requirement already satisfied: opentelemetry-semantic-conventions==0.45b0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.45b0)\n",
      "Requirement already satisfied: opentelemetry-util-http==0.45b0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (0.45b0)\n",
      "Requirement already satisfied: setuptools>=16.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-instrumentation==0.45b0->opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (69.0.2)\n",
      "Requirement already satisfied: wrapt<2.0.0,>=1.0.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-instrumentation==0.45b0->opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (1.16.0)\n",
      "Requirement already satisfied: asgiref~=3.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from opentelemetry-instrumentation-asgi==0.45b0->opentelemetry-instrumentation-fastapi>=0.41b0->chromadb) (3.8.1)\n",
      "Requirement already satisfied: monotonic>=1.5 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from posthog>=2.4.0->chromadb) (1.6)\n",
      "Requirement already satisfied: backoff>=1.10.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from posthog>=2.4.0->chromadb) (2.2.1)\n",
      "Requirement already satisfied: annotated-types>=0.4.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from pydantic>=1.9->chromadb) (0.6.0)\n",
      "Requirement already satisfied: pydantic-core==2.16.3 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from pydantic>=1.9->chromadb) (2.16.3)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from requests<3,>=2->langchain_community) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from requests<3,>=2->langchain_community) (3.6)\n",
      "Requirement already satisfied: huggingface_hub<1.0,>=0.16.4 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from tokenizers>=0.13.2->chromadb) (0.22.1)\n",
      "Requirement already satisfied: typer-slim==0.12.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from typer-slim[standard]==0.12.0->typer>=0.9.0->chromadb) (0.12.0)\n",
      "Requirement already satisfied: typer-cli==0.12.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from typer>=0.9.0->chromadb) (0.12.0)\n",
      "Requirement already satisfied: click>=8.0.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from typer-slim==0.12.0->typer-slim[standard]==0.12.0->typer>=0.9.0->chromadb) (8.1.7)\n",
      "Requirement already satisfied: shellingham>=1.3.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from typer-slim[standard]==0.12.0->typer>=0.9.0->chromadb) (1.5.4)\n",
      "Requirement already satisfied: rich>=10.11.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from typer-slim[standard]==0.12.0->typer>=0.9.0->chromadb) (13.7.1)\n",
      "Requirement already satisfied: h11>=0.8 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from uvicorn>=0.18.3->uvicorn[standard]>=0.18.3->chromadb) (0.14.0)\n",
      "Requirement already satisfied: httptools>=0.5.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (0.6.1)\n",
      "Requirement already satisfied: python-dotenv>=0.13 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (1.0.1)\n",
      "Requirement already satisfied: uvloop!=0.15.0,!=0.15.1,>=0.14.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (0.19.0)\n",
      "Requirement already satisfied: watchfiles>=0.13 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (0.21.0)\n",
      "Requirement already satisfied: websockets>=10.4 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from uvicorn[standard]>=0.18.3->chromadb) (12.0)\n",
      "Requirement already satisfied: cachetools<6.0,>=2.0.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (5.3.3)\n",
      "Requirement already satisfied: pyasn1-modules>=0.2.1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (0.4.0)\n",
      "Requirement already satisfied: rsa<5,>=3.1.4 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (4.9)\n",
      "Requirement already satisfied: httpcore==1.* in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from httpx>=0.21.2->cohere<6.0.0,>=5.1.4->langchain-cohere) (1.0.4)\n",
      "Requirement already satisfied: filelock in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from huggingface_hub<1.0,>=0.16.4->tokenizers>=0.13.2->chromadb) (3.13.3)\n",
      "Requirement already satisfied: fsspec>=2023.5.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from huggingface_hub<1.0,>=0.16.4->tokenizers>=0.13.2->chromadb) (2024.3.1)\n",
      "Requirement already satisfied: zipp>=0.5 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from importlib-metadata<=7.0,>=6.0->opentelemetry-api>=1.2.0->chromadb) (3.18.1)\n",
      "Requirement already satisfied: mypy-extensions>=0.3.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from typing-inspect<1,>=0.4.0->dataclasses-json<0.7,>=0.5.7->langchain_community) (1.0.0)\n",
      "Requirement already satisfied: humanfriendly>=9.1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from coloredlogs->onnxruntime>=1.14.1->chromadb) (10.0)\n",
      "Requirement already satisfied: mpmath>=0.19 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from sympy->onnxruntime>=1.14.1->chromadb) (1.3.0)\n",
      "Requirement already satisfied: pyasn1<0.7.0,>=0.4.6 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from pyasn1-modules>=0.2.1->google-auth>=1.0.1->kubernetes>=28.1.0->chromadb) (0.6.0)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from rich>=10.11.0->typer-slim[standard]==0.12.0->typer>=0.9.0->chromadb) (3.0.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from rich>=10.11.0->typer-slim[standard]==0.12.0->typer>=0.9.0->chromadb) (2.17.2)\n",
      "Requirement already satisfied: mdurl~=0.1 in /Users/nuno/dev/langgraph/.venv/lib/python3.11/site-packages (from markdown-it-py>=2.2.0->rich>=10.11.0->typer-slim[standard]==0.12.0->typer>=0.9.0->chromadb) (0.1.2)\n",
      "Downloading langchain_community-0.0.31-py3-none-any.whl (1.9 MB)\n",
      "\u001b[2K   \u001b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.9/1.9 MB\u001b[0m \u001b[31m9.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m[31m10.3 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hDownloading tiktoken-0.6.0-cp311-cp311-macosx_11_0_arm64.whl (949 kB)\n",
      "\u001b[2K   \u001b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m949.8/949.8 kB\u001b[0m \u001b[31m15.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m31m20.3 MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hDownloading langchain_openai-0.1.1-py3-none-any.whl (32 kB)\n",
      "Downloading langchain-0.1.14-py3-none-any.whl (812 kB)\n",
      "\u001b[2K   \u001b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m812.8/812.8 kB\u001b[0m \u001b[31m14.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m MB/s\u001b[0m eta \u001b[36m0:00:01\u001b[0m\n",
      "\u001b[?25hDownloading langgraph-0.0.31-py3-none-any.whl (55 kB)\n",
      "\u001b[2K   \u001b[38;2;114;156;31m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m55.5/55.5 kB\u001b[0m \u001b[31m4.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading tavily_python-0.3.3-py3-none-any.whl (5.4 kB)\n",
      "Installing collected packages: tiktoken, tavily-python, langgraph, langchain-openai, langchain_community, langchain\n",
      "  Attempting uninstall: tiktoken\n",
      "    Found existing installation: tiktoken 0.5.2\n",
      "    Uninstalling tiktoken-0.5.2:\n",
      "      Successfully uninstalled tiktoken-0.5.2\n",
      "  Attempting uninstall: tavily-python\n",
      "    Found existing installation: tavily-python 0.3.1\n",
      "    Uninstalling tavily-python-0.3.1:\n",
      "      Successfully uninstalled tavily-python-0.3.1\n",
      "  Attempting uninstall: langgraph\n",
      "    Found existing installation: langgraph 0.0.30\n",
      "    Uninstalling langgraph-0.0.30:\n",
      "      Successfully uninstalled langgraph-0.0.30\n",
      "  Attempting uninstall: langchain-openai\n",
      "    Found existing installation: langchain-openai 0.0.2.post1\n",
      "    Uninstalling langchain-openai-0.0.2.post1:\n",
      "      Successfully uninstalled langchain-openai-0.0.2.post1\n",
      "  Attempting uninstall: langchain_community\n",
      "    Found existing installation: langchain-community 0.0.27\n",
      "    Uninstalling langchain-community-0.0.27:\n",
      "      Successfully uninstalled langchain-community-0.0.27\n",
      "  Attempting uninstall: langchain\n",
      "    Found existing installation: langchain 0.1.11\n",
      "    Uninstalling langchain-0.1.11:\n",
      "      Successfully uninstalled langchain-0.1.11\n",
      "Successfully installed langchain-0.1.14 langchain-openai-0.1.1 langchain_community-0.0.31 langgraph-0.0.31 tavily-python-0.3.3 tiktoken-0.6.0\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "! pip install -U langchain_community tiktoken langchain-openai langchain-cohere langchainhub chromadb langchain langgraph  tavily-python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "222f204d-956f-4128-b597-2c698120edda",
   "metadata": {},
   "outputs": [],
   "source": [
    "### LLMs\n",
    "import os\n",
    "os.environ['OPENAI_API_KEY'] = <your-api-key>\n",
    "os.environ['COHERE_API_KEY'] = <your-api-key>\n",
    "os.environ['TAVILY_API_KEY'] = <your-api-key>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "018c9e9f-8199-4f33-b5b4-7adfb66d6219",
   "metadata": {},
   "source": [
    "### Tracing\n",
    "\n",
    "* Optionally, use [LangSmith](https://docs.smith.langchain.com/) for tracing (shown at bottom) by setting: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08edba00-988a-478b-96fc-ae0199cbef49",
   "metadata": {},
   "outputs": [],
   "source": [
    "### 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'] = <your-api-key>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ac1c2cd-81fb-40eb-8ba1-e9197800cba6",
   "metadata": {},
   "source": [
    "## Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b224e5ba-50ca-495a-a7fa-0f75a080e03c",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Build 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_openai import OpenAIEmbeddings\n",
    "### from langchain_cohere import CohereEmbeddings\n",
    "\n",
    "# Set embeddings\n",
    "embd = OpenAIEmbeddings()\n",
    "\n",
    "# Docs to index\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",
    "# Load\n",
    "docs = [WebBaseLoader(url).load() for url in urls]\n",
    "docs_list = [item for sublist in docs for item in sublist]\n",
    "\n",
    "# Split\n",
    "text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(\n",
    "    chunk_size=500, chunk_overlap=0\n",
    ")\n",
    "doc_splits = text_splitter.split_documents(docs_list)\n",
    "\n",
    "# Add to vectorstore\n",
    "vectorstore = Chroma.from_documents(\n",
    "    documents=doc_splits,\n",
    "    collection_name=\"rag-chroma\",\n",
    "    embedding=embd,\n",
    ")\n",
    "retriever = vectorstore.as_retriever()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f52b427-750c-40f8-8893-e9caab3afd8d",
   "metadata": {},
   "source": [
    "## LLMs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4dec9d98-f3dc-4b7f-abc0-9d01c754f2be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "datasource='web_search'\n",
      "datasource='vectorstore'\n"
     ]
    }
   ],
   "source": [
    "### Router\n",
    "\n",
    "from typing import Literal\n",
    "\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# Data model\n",
    "class RouteQuery(BaseModel):\n",
    "    \"\"\"Route a user query to the most relevant datasource.\"\"\"\n",
    "\n",
    "    datasource: Literal[\"vectorstore\", \"web_search\"] = Field(\n",
    "        ...,\n",
    "        description=\"Given a user question choose to route it to web search or a vectorstore.\",\n",
    "    )\n",
    "\n",
    "# LLM with function call \n",
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
    "structured_llm_router = llm.with_structured_output(RouteQuery)\n",
    "\n",
    "# Prompt \n",
    "system = \"\"\"You are an expert at routing a user question to a vectorstore or web search.\n",
    "The vectorstore contains documents related to agents, prompt engineering, and adversarial attacks.\n",
    "Use the vectorstore for questions on these topics. Otherwise, use web-search.\"\"\"\n",
    "route_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "question_router = route_prompt | structured_llm_router\n",
    "print(question_router.invoke({\"question\": \"Who will the Bears draft first in the NFL draft?\"}))\n",
    "print(question_router.invoke({\"question\": \"What are the types of agent memory?\"}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "856801cb-f42a-44e7-956f-47845e3664ca",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "binary_score='no'\n"
     ]
    }
   ],
   "source": [
    "### Retrieval Grader \n",
    "\n",
    "# Data model\n",
    "class GradeDocuments(BaseModel):\n",
    "    \"\"\"Binary score for relevance check on retrieved documents.\"\"\"\n",
    "\n",
    "    binary_score: str = Field(description=\"Documents are relevant to the question, 'yes' or 'no'\")\n",
    "\n",
    "# LLM with function call \n",
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
    "structured_llm_grader = llm.with_structured_output(GradeDocuments)\n",
    "\n",
    "# Prompt \n",
    "system = \"\"\"You are a grader assessing relevance of a retrieved document to a user question. \\n \n",
    "    If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \\n\n",
    "    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",
    "grade_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"Retrieved document: \\n\\n {document} \\n\\n User question: {question}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "retrieval_grader = grade_prompt | structured_llm_grader\n",
    "question = \"agent memory\"\n",
    "docs = retriever.get_relevant_documents(question)\n",
    "doc_txt = docs[1].page_content\n",
    "print(retrieval_grader.invoke({\"question\": question, \"document\": doc_txt}))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2272333e-50b2-42ab-b472-e1055a3b94a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The design of generative agents combines LLM with memory, planning, and reflection mechanisms to enable agents to behave based on past experience and interact with other agents. Memory stream is a long-term memory module that records agents' experiences in natural language. The retrieval model surfaces context to inform the agent's behavior based on relevance, recency, and importance.\n"
     ]
    }
   ],
   "source": [
    "### Generate\n",
    "\n",
    "from langchain import hub\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "# Prompt\n",
    "prompt = hub.pull(\"rlm/rag-prompt\")\n",
    "\n",
    "# LLM\n",
    "llm = ChatOpenAI(model_name=\"gpt-3.5-turbo\", 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",
    "generation = rag_chain.invoke({\"context\": docs, \"question\": question})\n",
    "print(generation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f0c08d14-77a0-4eed-b882-2d636abb22a3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GradeHallucinations(binary_score='yes')"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### Hallucination Grader \n",
    "\n",
    "# Data model\n",
    "class GradeHallucinations(BaseModel):\n",
    "    \"\"\"Binary score for hallucination present in generation answer.\"\"\"\n",
    "\n",
    "    binary_score: str = Field(description=\"Answer is grounded in the facts, 'yes' or 'no'\")\n",
    "\n",
    "# LLM with function call \n",
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
    "structured_llm_grader = llm.with_structured_output(GradeHallucinations)\n",
    "\n",
    "# Prompt \n",
    "system = \"\"\"You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. \\n \n",
    "     Give a binary score 'yes' or 'no'. 'Yes' means that the answer is grounded in / supported by the set of facts.\"\"\"\n",
    "hallucination_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"Set of facts: \\n\\n {documents} \\n\\n LLM generation: {generation}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "hallucination_grader = hallucination_prompt | structured_llm_grader\n",
    "hallucination_grader.invoke({\"documents\": docs, \"generation\": generation})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ded99680-437a-4c9d-b860-619c88949d84",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "GradeAnswer(binary_score='yes')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### Answer Grader \n",
    "\n",
    "# Data model\n",
    "class GradeAnswer(BaseModel):\n",
    "    \"\"\"Binary score to assess answer addresses question.\"\"\"\n",
    "\n",
    "    binary_score: str = Field(description=\"Answer addresses the question, 'yes' or 'no'\")\n",
    "\n",
    "# LLM with function call \n",
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
    "structured_llm_grader = llm.with_structured_output(GradeAnswer)\n",
    "\n",
    "# Prompt \n",
    "system = \"\"\"You are a grader assessing whether an answer addresses / resolves a question \\n \n",
    "     Give a binary score 'yes' or 'no'. Yes' means that the answer resolves the question.\"\"\"\n",
    "answer_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"User question: \\n\\n {question} \\n\\n LLM generation: {generation}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "answer_grader = answer_prompt | structured_llm_grader\n",
    "answer_grader.invoke({\"question\": question,\"generation\": generation})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "9d75f1d7-a47a-4577-bb0d-84b504b0867e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"What is the role of memory in an agent's functioning?\""
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "### Question Re-writer\n",
    "\n",
    "# LLM \n",
    "llm = ChatOpenAI(model=\"gpt-3.5-turbo-0125\", temperature=0)\n",
    "\n",
    "# Prompt \n",
    "system = \"\"\"You a question re-writer that converts an input question to a better version that is optimized \\n \n",
    "     for vectorstore retrieval. Look at the input and try to reason about the underlying sematic intent / meaning.\"\"\"\n",
    "re_write_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", system),\n",
    "        (\"human\", \"Here is the initial question: \\n\\n {question} \\n Formulate an improved question.\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "question_rewriter = re_write_prompt | llm | StrOutputParser()\n",
    "question_rewriter.invoke({\"question\": question})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d07c0b31-b919-4498-869f-9673125c2473",
   "metadata": {},
   "source": [
    "## Web Search Tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "01d829bb-1074-4976-b650-ead41dcb9788",
   "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": "efbbff0e-8843-45bb-b2ff-137bef707ef4",
   "metadata": {},
   "source": [
    "# Graph \n",
    "\n",
    "Capture the flow in as a graph.\n",
    "\n",
    "## Graph state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e723fcdb-06e6-402d-912e-899795b78408",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing_extensions import TypedDict\n",
    "from typing import List\n",
    "\n",
    "class GraphState(TypedDict):\n",
    "    \"\"\"\n",
    "    Represents the state of our graph.\n",
    "\n",
    "    Attributes:\n",
    "        question: question\n",
    "        generation: LLM generation\n",
    "        documents: list of documents \n",
    "    \"\"\"\n",
    "    question : str\n",
    "    generation : str\n",
    "    documents : List[str]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e2d6c0d-42e8-4399-9751-e315be16607a",
   "metadata": {},
   "source": [
    "## Graph Flow "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b76b5ec3-0720-443d-85b1-c0e79659ca0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.schema import Document\n",
    "\n",
    "def retrieve(state):\n",
    "    \"\"\"\n",
    "    Retrieve documents\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\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",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): Updates documents key with only filtered relevant documents\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",
    "    for d in documents:\n",
    "        score = retrieval_grader.invoke({\"question\": question, \"document\": d.page_content})\n",
    "        grade = score.binary_score\n",
    "        if grade == \"yes\":\n",
    "            print(\"---GRADE: DOCUMENT RELEVANT---\")\n",
    "            filtered_docs.append(d)\n",
    "        else:\n",
    "            print(\"---GRADE: DOCUMENT NOT RELEVANT---\")\n",
    "            continue\n",
    "    return {\"documents\": filtered_docs, \"question\": question}\n",
    "\n",
    "def transform_query(state):\n",
    "    \"\"\"\n",
    "    Transform the query to produce a better question.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): Updates question key with a re-phrased question\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---TRANSFORM QUERY---\")\n",
    "    question = state[\"question\"]\n",
    "    documents = state[\"documents\"]\n",
    "\n",
    "    # Re-write question\n",
    "    better_question = question_rewriter.invoke({\"question\": question})\n",
    "    return {\"documents\": documents, \"question\": better_question}\n",
    "\n",
    "def web_search(state):\n",
    "    \"\"\"\n",
    "    Web search based on the re-phrased question.\n",
    "\n",
    "    Args:\n",
    "        state (dict): The current graph state\n",
    "\n",
    "    Returns:\n",
    "        state (dict): Updates documents key with appended web results\n",
    "    \"\"\"\n",
    "\n",
    "    print(\"---WEB SEARCH---\")\n",
    "    question = state[\"question\"]\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",
    "\n",
    "    return {\"documents\": web_results, \"question\": question}\n",
    "\n",
    "### Edges ###\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",
    "    source = question_router.invoke({\"question\": question})   \n",
    "    if source.datasource == 'web_search':\n",
    "        print(\"---ROUTE QUESTION TO WEB SEARCH---\")\n",
    "        return \"web_search\"\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 re-generate a question.\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",
    "    filtered_documents = state[\"documents\"]\n",
    "\n",
    "    if not filtered_documents:\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, TRANSFORM QUERY---\")\n",
    "        return \"transform_query\"\n",
    "    else:\n",
    "        # We have relevant documents, so generate answer\n",
    "        print(\"---DECISION: GENERATE---\")\n",
    "        return \"generate\"\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.binary_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.binary_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\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ab01f36-5628-49ab-bfd3-84bb6f1a1b0f",
   "metadata": {},
   "source": [
    "## Build Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "67854e07-9293-4c3c-bf9a-bc9a605570ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import END, StateGraph\n",
    "\n",
    "workflow = StateGraph(GraphState)\n",
    "\n",
    "# Define the nodes\n",
    "workflow.add_node(\"web_search\", 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\n",
    "workflow.add_node(\"transform_query\", transform_query) # transform_query\n",
    "\n",
    "# Build graph\n",
    "workflow.set_conditional_entry_point(\n",
    "    route_question,\n",
    "    {\n",
    "        \"web_search\": \"web_search\",\n",
    "        \"vectorstore\": \"retrieve\",\n",
    "    },\n",
    ")\n",
    "workflow.add_edge(\"web_search\", \"generate\")\n",
    "workflow.add_edge(\"retrieve\", \"grade_documents\")\n",
    "workflow.add_conditional_edges(\n",
    "    \"grade_documents\",\n",
    "    decide_to_generate,\n",
    "    {\n",
    "        \"transform_query\": \"transform_query\",\n",
    "        \"generate\": \"generate\",\n",
    "    },\n",
    ")\n",
    "workflow.add_edge(\"transform_query\", \"retrieve\")\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\": \"transform_query\",\n",
    "    },\n",
    ")\n",
    "\n",
    "# Compile\n",
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "29acc541-d726-4b75-84d1-a215845fe88a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---ROUTE QUESTION---\n",
      "---ROUTE QUESTION TO WEB SEARCH---\n",
      "---WEB SEARCH---\n",
      "\"Node 'web_search':\"\n",
      "'\\n---\\n'\n",
      "---GENERATE---\n",
      "---CHECK HALLUCINATIONS---\n",
      "---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\n",
      "---GRADE GENERATION vs QUESTION---\n",
      "---DECISION: GENERATION ADDRESSES QUESTION---\n",
      "\"Node 'generate':\"\n",
      "'\\n---\\n'\n",
      "('It is expected that the Chicago Bears could have the opportunity to draft '\n",
      " 'the first defensive player in the 2024 NFL draft. The Bears have the first '\n",
      " 'overall pick in the draft, giving them a prime position to select top '\n",
      " 'talent. The top wide receiver Marvin Harrison Jr. from Ohio State is also '\n",
      " 'mentioned as a potential pick for the Cardinals.')\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "# Run \n",
    "inputs = {\"question\": \"What player at the Bears expected to draft first in the 2024 NFL draft?\"}\n",
    "for output in app.stream(inputs):\n",
    "    for key, value in output.items():\n",
    "        # Node\n",
    "        pprint(f\"Node '{key}':\")\n",
    "        # Optional: print full state at each node\n",
    "        # pprint.pprint(value[\"keys\"], indent=2, width=80, depth=None)\n",
    "    pprint(\"\\n---\\n\")\n",
    "\n",
    "# Final generation\n",
    "pprint(value[\"generation\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11fddd00-58bf-4910-bf36-be9e5bfba778",
   "metadata": {},
   "source": [
    "Trace: \n",
    "\n",
    "https://smith.langchain.com/public/7e3aa7e5-c51f-45c2-bc66-b34f17ff2263/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "69a985dd-03c6-45af-a67b-b15746a2cb5f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "---ROUTE QUESTION---\n",
      "---ROUTE QUESTION TO RAG---\n",
      "---RETRIEVE---\n",
      "\"Node 'retrieve':\"\n",
      "'\\n---\\n'\n",
      "---CHECK DOCUMENT RELEVANCE TO QUESTION---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "---GRADE: DOCUMENT NOT RELEVANT---\n",
      "---GRADE: DOCUMENT RELEVANT---\n",
      "---ASSESS GRADED DOCUMENTS---\n",
      "---DECISION: GENERATE---\n",
      "\"Node 'grade_documents':\"\n",
      "'\\n---\\n'\n",
      "---GENERATE---\n",
      "---CHECK HALLUCINATIONS---\n",
      "---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---\n",
      "---GRADE GENERATION vs QUESTION---\n",
      "---DECISION: GENERATION ADDRESSES QUESTION---\n",
      "\"Node 'generate':\"\n",
      "'\\n---\\n'\n",
      "('The types of agent memory include Sensory Memory, Short-Term Memory (STM) or '\n",
      " 'Working Memory, and Long-Term Memory (LTM) with subtypes of Explicit / '\n",
      " 'declarative memory and Implicit / procedural memory. Sensory memory retains '\n",
      " 'sensory information briefly, STM stores information for cognitive tasks, and '\n",
      " 'LTM stores information for a long time with different types of memories.')\n"
     ]
    }
   ],
   "source": [
    "# Run\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",
    "        # Node\n",
    "        pprint(f\"Node '{key}':\")\n",
    "        # Optional: print full state at each node\n",
    "        # pprint.pprint(value[\"keys\"], indent=2, width=80, depth=None)\n",
    "    pprint(\"\\n---\\n\")\n",
    "\n",
    "# Final generation\n",
    "pprint(value [\"generation\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebf41097-fc4c-4072-95b3-e7e07731ada1",
   "metadata": {},
   "source": [
    "Trace: \n",
    "\n",
    "https://smith.langchain.com/public/fdf0a180-6d15-4d09-bb92-f84f2105ca51/r"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19ac1f6f-2d84-488f-8a0e-7ee2a46b0f71",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}