{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6e6a0a39-9a4c-47ae-a238-1a3a847eea5b",
   "metadata": {},
   "source": [
    "# How to add runtime configuration to your graph\n",
    "\n",
    "Sometimes you want to be able to configure your agent when calling it. \n",
    "Examples of this include configuring which LLM to use.\n",
    "Below we walk through an example of doing so.\n",
    "\n",
    "## Setup\n",
    "\n",
    "First, let's install the required packages and set our API keys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03df6e04",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install -U langgraph langchain_anthropic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a00c45e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "\n",
    "def _set_env(var: str):\n",
    "    if not os.environ.get(var):\n",
    "        os.environ[var] = getpass.getpass(f\"{var}: \")\n",
    "\n",
    "\n",
    "_set_env(\"ANTHROPIC_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55e8be3b",
   "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": "df1ff9cf-f8d2-4109-adf9-2adec83f5a95",
   "metadata": {},
   "source": [
    "## Define graph\n",
    "\n",
    "First, let's create a very simple graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "816523d0-0b59-47cf-9f4c-4838024efe22",
   "metadata": {},
   "outputs": [],
   "source": [
    "import operator\n",
    "from typing import Annotated, Sequence, TypedDict\n",
    "\n",
    "from langchain_anthropic import ChatAnthropic\n",
    "from langchain_core.messages import BaseMessage, HumanMessage\n",
    "\n",
    "from langgraph.graph import END, StateGraph, START\n",
    "\n",
    "model = ChatAnthropic(model_name=\"claude-2.1\")\n",
    "\n",
    "\n",
    "class AgentState(TypedDict):\n",
    "    messages: Annotated[Sequence[BaseMessage], operator.add]\n",
    "\n",
    "\n",
    "def _call_model(state):\n",
    "    response = model.invoke(state[\"messages\"])\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "# Define a new graph\n",
    "workflow = StateGraph(AgentState)\n",
    "workflow.add_node(\"model\", _call_model)\n",
    "workflow.add_edge(START, \"model\")\n",
    "workflow.add_edge(\"model\", END)\n",
    "\n",
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "070f11a6-2441-4db5-9df6-e318f110e281",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='hi'),\n",
       "  AIMessage(content='Hello!', response_metadata={'id': 'msg_012SakNGNitBcKJgc9yZ1Asv', 'model': 'claude-2.1', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 10, 'output_tokens': 6}}, id='run-9e375cd7-ae84-4db2-981c-c7e18ecabddf-0', usage_metadata={'input_tokens': 10, 'output_tokens': 6, 'total_tokens': 16})]}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "app.invoke({\"messages\": [HumanMessage(content=\"hi\")]})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69a1dd47-c5b3-4e04-af56-45682f74d61f",
   "metadata": {},
   "source": [
    "## Configure the graph\n",
    "\n",
    "Great! Now let's suppose that we want to extend this example so the user is able to choose from multiple llms.\n",
    "We can easily do that by passing in a config. Any configuration information needs to be passed inside `configurable` key as shown below.\n",
    "This config is meant to contain things are not part of the input (and therefore that we don't want to track as part of the state)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c01f1e7c-8e8b-4e26-98f7-56ac225077b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from typing import Optional\n",
    "from langchain_core.runnables.config import RunnableConfig\n",
    "\n",
    "openai_model = ChatOpenAI()\n",
    "\n",
    "models = {\n",
    "    \"anthropic\": model,\n",
    "    \"openai\": openai_model,\n",
    "}\n",
    "\n",
    "def _call_model(state: AgentState, config: RunnableConfig):\n",
    "    # Access the config through the configurable key\n",
    "    model_name = config[\"configurable\"].get(\"model\", \"anthropic\")\n",
    "    model = models[model_name]\n",
    "    response = model.invoke(state[\"messages\"])\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "# Define a new graph\n",
    "workflow = StateGraph(AgentState)\n",
    "workflow.add_node(\"model\", _call_model)\n",
    "workflow.add_edge(START, \"model\")\n",
    "workflow.add_edge(\"model\", END)\n",
    "\n",
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7741b75c-55ba-4c78-bbb1-5dc20a210f11",
   "metadata": {},
   "source": [
    "If we call it with no configuration, it will use the default as we defined it (Anthropic)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ef50f048-fc43-40c0-b713-346408fcf052",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='hi'),\n",
       "  AIMessage(content='Hello!', response_metadata={'id': 'msg_0133PAX5DyoUYL1gZiGR8NXs', 'model': 'claude-2.1', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 10, 'output_tokens': 6}}, id='run-03e8bd8b-fa09-4258-920d-8f53a7b91fcc-0', usage_metadata={'input_tokens': 10, 'output_tokens': 6, 'total_tokens': 16})]}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "app.invoke({\"messages\": [HumanMessage(content=\"hi\")]})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6896b32-9b25-4342-bfd0-29a3d329a06a",
   "metadata": {},
   "source": [
    "We can also call it with a config to get it to use a different model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f2f7c74b-9fb0-41c6-9728-dcf9d8a3c397",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='hi'),\n",
       "  AIMessage(content='Hello! How can I assist you today?', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 8, 'total_tokens': 17}, 'model_name': 'gpt-3.5-turbo-0125', 'system_fingerprint': None, 'finish_reason': 'stop', 'logprobs': None}, id='run-6d0c7c25-03de-49d6-b3be-ff0858d17122-0', usage_metadata={'input_tokens': 8, 'output_tokens': 9, 'total_tokens': 17})]}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config = {\"configurable\": {\"model\": \"openai\"}}\n",
    "app.invoke({\"messages\": [HumanMessage(content=\"hi\")]}, config=config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b4c7eaf1-4ee0-42b3-971d-273a108f205f",
   "metadata": {},
   "source": [
    "We can also adapt our graph to take in more configuration! Like a system message for example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f0393a43-9fbe-4056-972f-3e91ea329041",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.messages import SystemMessage\n",
    "\n",
    "# We can define a config schema to specify the configuration options for the graph\n",
    "# A config schema is useful for indicating which fields are available in the configurable dict inside the config\n",
    "class ConfigSchema(TypedDict):\n",
    "    model: Optional[str]\n",
    "    system_message: Optional[str]\n",
    "\n",
    "def _call_model(state: AgentState, config: RunnableConfig):\n",
    "    # Access the config through the configurable key\n",
    "    model_name = config[\"configurable\"].get(\"model\", \"anthropic\")\n",
    "    model = models[model_name]\n",
    "    messages = state[\"messages\"]\n",
    "    if \"system_message\" in config[\"configurable\"]:\n",
    "        messages = [\n",
    "            SystemMessage(content=config[\"configurable\"][\"system_message\"])\n",
    "        ] + messages\n",
    "    response = model.invoke(messages)\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "# Define a new graph - note that we pass in the configuration schema here, but it is not necessary\n",
    "workflow = StateGraph(AgentState, ConfigSchema)\n",
    "workflow.add_node(\"model\", _call_model)\n",
    "workflow.add_edge(START, \"model\")\n",
    "workflow.add_edge(\"model\", END)\n",
    "\n",
    "app = workflow.compile()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "718685f7-4cdd-4181-9fc8-e7762d584727",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='hi'),\n",
       "  AIMessage(content='Hello!', response_metadata={'id': 'msg_01TVJvxCXsCT9JVe7A4iUUi9', 'model': 'claude-2.1', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 10, 'output_tokens': 6}}, id='run-627eb685-c4d7-481d-9095-c0a1822e8c10-0', usage_metadata={'input_tokens': 10, 'output_tokens': 6, 'total_tokens': 16})]}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "app.invoke({\"messages\": [HumanMessage(content=\"hi\")]})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e043a719-f197-46ef-9d45-84740a39aeb0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [HumanMessage(content='hi'),\n",
       "  AIMessage(content='Ciao!', response_metadata={'id': 'msg_01CpBD1cMCYvvPX2cogUawJj', 'model': 'claude-2.1', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 14, 'output_tokens': 7}}, id='run-6ef2fea6-9bfa-4266-bd05-263160a1db7b-0', usage_metadata={'input_tokens': 14, 'output_tokens': 7, 'total_tokens': 21})]}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config = {\"configurable\": {\"system_message\": \"respond in italian\"}}\n",
    "app.invoke({\"messages\": [HumanMessage(content=\"hi\")]}, config=config)"
   ]
  }
 ],
 "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
}