2024_02_13 - Run Any Open LLM with OpenAI API/02_LangGraph_agent_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LangGraph Example: LLM with Tools\n",
    "This Notebook tests out using different LLMs to power a custom LangGraph agent"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Notebook setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 161,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install --upgrade langchain langsmith langgraph langchain_openai langchain-together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 162,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "import os\n",
    "\n",
    "# Load .env file\n",
    "load_dotenv()\n",
    "\n",
    "# Set model variables\n",
    "OPENAI_BASE_URL = \"https://api.openai.com/v1\"\n",
    "OPENAI_API_KEY = os.getenv(\"OPENAI_API_KEY\")\n",
    "os.environ[\"OPENAI_ORGANIZATION\"] = os.getenv(\"OPENAI_ORGANIZATION\")\n",
    "\n",
    "OLLAMA_BASE_URL = \"http://localhost:11434/v1\"\n",
    "OLLAMA_API_KEY = \"ollama\"\n",
    "\n",
    "TOGETHER_BASE_URL = \"https://api.together.xyz/v1\"\n",
    "TOGETHER_API_KEY = os.getenv(\"TOGETHER_API_KEY\")\n",
    "\n",
    "# Initialize LangSmith\n",
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_ENDPOINT\"] = \"https://api.smith.langchain.com\"\n",
    "os.environ[\"LANGCHAIN_API_KEY\"] = os.getenv(\"LANGCHAIN_API_KEY\")\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = \"Demos\"\n",
    "\n",
    "# Initialize OpenWeatherMap\n",
    "OPENWEATHERMAP_API_KEY = os.getenv(\"OPENWEATHERMAP_API_KEY\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the Tools\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 163,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import things that are needed generically for tools\n",
    "from langchain.pydantic_v1 import BaseModel, Field\n",
    "from langchain.tools import StructuredTool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 164,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "\n",
    "class City(BaseModel):\n",
    "    city: str = Field(description=\"City\")\n",
    "    country: str = Field(description=\"Country code\")\n",
    "\n",
    "\n",
    "def get_current_weather(city: str, country: str) -> int:\n",
    "    response = requests.get(\n",
    "        f\"http://api.openweathermap.org/data/2.5/weather?q={city},{country}&appid={OPENWEATHERMAP_API_KEY}\"\n",
    "    )\n",
    "    data = response.json()\n",
    "    temp_kelvin = data[\"main\"][\"temp\"]\n",
    "    temp_fahrenheit = (temp_kelvin - 273.15) * 9 / 5 + 32\n",
    "    return int(temp_fahrenheit)\n",
    "\n",
    "\n",
    "weather = StructuredTool.from_function(\n",
    "    func=get_current_weather,\n",
    "    name=\"Get_Weather\",\n",
    "    description=\"Get the current temperature from a city, in Fahrenheit\",\n",
    "    args_schema=City,\n",
    "    return_direct=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 165,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DifferenceInput(BaseModel):\n",
    "    minuend: int = Field(\n",
    "        description=\"The number from which another number is to be subtracted\"\n",
    "    )\n",
    "    subtrahend: int = Field(description=\"The number to be subtracted\")\n",
    "\n",
    "\n",
    "def get_difference(minuend: int, subtrahend: int) -> int:\n",
    "    return minuend - subtrahend\n",
    "\n",
    "\n",
    "difference = StructuredTool.from_function(\n",
    "    func=get_difference,\n",
    "    name=\"Calculate_difference\",\n",
    "    description=\"Calculate the difference between two temperatures in Fahrenheit\",\n",
    "    args_schema=DifferenceInput,\n",
    "    return_direct=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 166,
   "metadata": {},
   "outputs": [],
   "source": [
    "tools = [weather, difference]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 167,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the tools to execute them from the graph\n",
    "from langgraph.prebuilt import ToolExecutor\n",
    "\n",
    "tool_executor = ToolExecutor(tools)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the agent\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 168,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the response schema for our agent\n",
    "from langchain_core.pydantic_v1 import BaseModel, Field\n",
    "\n",
    "\n",
    "class Response(BaseModel):\n",
    "    \"\"\"Format the final answer to the user's question\"\"\"\n",
    "\n",
    "    warmest_city: str = Field(\n",
    "        description=\"The warmest city and its current temperature\"\n",
    "    )\n",
    "    explanation: str = Field(\n",
    "        description=\"How much warmer it is in the warmest city than the other cities\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 169,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.utils.function_calling import (\n",
    "    convert_to_openai_function,\n",
    "    convert_to_openai_tool,\n",
    ")\n",
    "from langchain_together import Together\n",
    "\n",
    "from langchain_core.prompts import (\n",
    "    ChatPromptTemplate,\n",
    "    MessagesPlaceholder,\n",
    ")\n",
    "from langchain_core.runnables import RunnablePassthrough\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are a helpful assistant\",\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"messages\", optional=True),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# Create the LLM\n",
    "# OpenAI\n",
    "# llm = ChatOpenAI(\n",
    "#     openai_api_base=OPENAI_BASE_URL,\n",
    "#     api_key=OPENAI_API_KEY,\n",
    "#     model=\"gpt-3.5-turbo-1106\",\n",
    "# )\n",
    "\n",
    "# # Local (Ollama)\n",
    "# llm = ChatOpenAI(\n",
    "#     openai_api_base=OLLAMA_BASE_URL,\n",
    "#     api_key=OLLAMA_API_KEY,\n",
    "#     model=\"llama2\",\n",
    "# )\n",
    "\n",
    "# # Serverless (with Together.ai)\n",
    "llm = ChatOpenAI(\n",
    "    openai_api_base=TOGETHER_BASE_URL,\n",
    "    api_key=TOGETHER_API_KEY,\n",
    "    # model=\"mistralai/Mixtral-8x7B-Instruct-v0.1\",\n",
    "    model=\"mistralai/Mistral-7B-Instruct-v0.1\",\n",
    "    # model=\"togethercomputer/CodeLlama-34b-Instruct\",\n",
    ")\n",
    "\n",
    "tools_to_pass_in = [convert_to_openai_function(t) for t in tools]\n",
    "# tools_to_pass_in.append(convert_to_openai_function(Response))\n",
    "\n",
    "model = {\"messages\": RunnablePassthrough()} | prompt | llm.bind_tools(tools_to_pass_in)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the Agent State\n",
    "\n",
    "Everything from here onwards is new\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 170,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import TypedDict, Annotated, Sequence\n",
    "import operator\n",
    "from langchain_core.messages import BaseMessage\n",
    "\n",
    "\n",
    "class AgentState(TypedDict):\n",
    "    messages: Sequence[BaseMessage]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up the node actions\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 171,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.prebuilt import ToolInvocation\n",
    "import json\n",
    "from langchain_core.messages import ToolMessage\n",
    "\n",
    "\n",
    "# Define the function that determines whether to continue or not\n",
    "def should_continue(state):\n",
    "    last_message = state[\"messages\"][-1]\n",
    "    # If there are no tool calls, then we finish\n",
    "    if \"tool_calls\" not in last_message.additional_kwargs:\n",
    "        return \"end\"\n",
    "    # If there is a Response tool call, then we finish\n",
    "    elif any(\n",
    "        tool_call[\"function\"][\"name\"] == \"Response\"\n",
    "        for tool_call in last_message.additional_kwargs[\"tool_calls\"]\n",
    "    ):\n",
    "        return \"end\"\n",
    "    # Otherwise, we continue\n",
    "    else:\n",
    "        return \"continue\"\n",
    "\n",
    "\n",
    "# Define the function that calls the model\n",
    "def call_model(state):\n",
    "    messages = state[\"messages\"]\n",
    "    response = model.invoke(messages)\n",
    "    return {\"messages\": messages + [response]}\n",
    "\n",
    "\n",
    "# Define the function to execute tools\n",
    "def call_tool(state):\n",
    "    messages = state[\"messages\"]\n",
    "    # We know the last message involves at least one tool call\n",
    "    last_message = messages[-1]\n",
    "\n",
    "    # We loop through all tool calls and append the message to our message log\n",
    "    for tool_call in last_message.additional_kwargs[\"tool_calls\"]:\n",
    "        action = ToolInvocation(\n",
    "            tool=tool_call[\"function\"][\"name\"],\n",
    "            tool_input=json.loads(tool_call[\"function\"][\"arguments\"]),\n",
    "            id=tool_call[\"id\"],\n",
    "        )\n",
    "\n",
    "        # We call the tool_executor and get back a response\n",
    "        response = tool_executor.invoke(action)\n",
    "        # We use the response to create a FunctionMessage\n",
    "        function_message = ToolMessage(\n",
    "            content=str(response), name=action.tool, tool_call_id=tool_call[\"id\"]\n",
    "        )\n",
    "\n",
    "        # Add the function message to the list\n",
    "        messages.append(function_message)\n",
    "\n",
    "    # We return a list, because this will get added to the existing list\n",
    "\n",
    "    return {\"messages\": messages}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define the Graph\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 172,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.graph import StateGraph, END\n",
    "\n",
    "# Initialize a new graph\n",
    "graph = StateGraph(AgentState)\n",
    "\n",
    "# Define the two Nodes we will cycle between\n",
    "graph.add_node(\"agent\", call_model)\n",
    "graph.add_node(\"action\", call_tool)\n",
    "\n",
    "# Set the Starting Edge\n",
    "graph.set_entry_point(\"agent\")\n",
    "\n",
    "# Set our Contitional Edges\n",
    "graph.add_conditional_edges(\n",
    "    \"agent\",\n",
    "    should_continue,\n",
    "    {\n",
    "        \"continue\": \"action\",\n",
    "        \"end\": END,\n",
    "    },\n",
    ")\n",
    "\n",
    "# Set the Normal Edges\n",
    "graph.add_edge(\"action\", \"agent\")\n",
    "\n",
    "# Compile the workflow\n",
    "app = graph.compile()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run our graph\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 173,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Output from node 'agent':\n",
      "---\n",
      "{'messages': [HumanMessage(content='Where is it warmest: Austin, Texas; Tokyo; or Seattle? And by how much is it warmer than the other cities?'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_toanzknowvq8ymydvjto0w0r', 'function': {'arguments': '{\"city\":\"Austin, Texas\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_x14745jh6z0z271gdwyjv1z6', 'function': {'arguments': '{\"city\":\"Tokyo\",\"country\":\"Japan\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_2z3jz218jt9o4righx2q0pp5', 'function': {'arguments': '{\"city\":\"Seattle\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}]})]}\n",
      "\n",
      "---\n",
      "\n",
      "Output from node 'action':\n",
      "---\n",
      "{'messages': [HumanMessage(content='Where is it warmest: Austin, Texas; Tokyo; or Seattle? And by how much is it warmer than the other cities?'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_toanzknowvq8ymydvjto0w0r', 'function': {'arguments': '{\"city\":\"Austin, Texas\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_x14745jh6z0z271gdwyjv1z6', 'function': {'arguments': '{\"city\":\"Tokyo\",\"country\":\"Japan\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_2z3jz218jt9o4righx2q0pp5', 'function': {'arguments': '{\"city\":\"Seattle\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}]}), ToolMessage(content='61', tool_call_id='call_toanzknowvq8ymydvjto0w0r', name='Get_Weather'), ToolMessage(content='46', tool_call_id='call_x14745jh6z0z271gdwyjv1z6', name='Get_Weather'), ToolMessage(content='43', tool_call_id='call_2z3jz218jt9o4righx2q0pp5', name='Get_Weather')]}\n",
      "\n",
      "---\n",
      "\n",
      "Output from node 'agent':\n",
      "---\n",
      "{'messages': [HumanMessage(content='Where is it warmest: Austin, Texas; Tokyo; or Seattle? And by how much is it warmer than the other cities?'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_toanzknowvq8ymydvjto0w0r', 'function': {'arguments': '{\"city\":\"Austin, Texas\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_x14745jh6z0z271gdwyjv1z6', 'function': {'arguments': '{\"city\":\"Tokyo\",\"country\":\"Japan\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_2z3jz218jt9o4righx2q0pp5', 'function': {'arguments': '{\"city\":\"Seattle\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}]}), ToolMessage(content='61', tool_call_id='call_toanzknowvq8ymydvjto0w0r', name='Get_Weather'), ToolMessage(content='46', tool_call_id='call_x14745jh6z0z271gdwyjv1z6', name='Get_Weather'), ToolMessage(content='43', tool_call_id='call_2z3jz218jt9o4righx2q0pp5', name='Get_Weather'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_mntkddx3gbk9njo8dts8srpe', 'function': {'arguments': '{\"city\":\"London\",\"country\":\"UK\"}', 'name': 'Get_Weather'}, 'type': 'function'}]})]}\n",
      "\n",
      "---\n",
      "\n",
      "Output from node 'action':\n",
      "---\n",
      "{'messages': [HumanMessage(content='Where is it warmest: Austin, Texas; Tokyo; or Seattle? And by how much is it warmer than the other cities?'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_toanzknowvq8ymydvjto0w0r', 'function': {'arguments': '{\"city\":\"Austin, Texas\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_x14745jh6z0z271gdwyjv1z6', 'function': {'arguments': '{\"city\":\"Tokyo\",\"country\":\"Japan\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_2z3jz218jt9o4righx2q0pp5', 'function': {'arguments': '{\"city\":\"Seattle\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}]}), ToolMessage(content='61', tool_call_id='call_toanzknowvq8ymydvjto0w0r', name='Get_Weather'), ToolMessage(content='46', tool_call_id='call_x14745jh6z0z271gdwyjv1z6', name='Get_Weather'), ToolMessage(content='43', tool_call_id='call_2z3jz218jt9o4righx2q0pp5', name='Get_Weather'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_mntkddx3gbk9njo8dts8srpe', 'function': {'arguments': '{\"city\":\"London\",\"country\":\"UK\"}', 'name': 'Get_Weather'}, 'type': 'function'}]}), ToolMessage(content='49', tool_call_id='call_mntkddx3gbk9njo8dts8srpe', name='Get_Weather')]}\n",
      "\n",
      "---\n",
      "\n",
      "Output from node 'agent':\n",
      "---\n",
      "{'messages': [HumanMessage(content='Where is it warmest: Austin, Texas; Tokyo; or Seattle? And by how much is it warmer than the other cities?'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_toanzknowvq8ymydvjto0w0r', 'function': {'arguments': '{\"city\":\"Austin, Texas\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_x14745jh6z0z271gdwyjv1z6', 'function': {'arguments': '{\"city\":\"Tokyo\",\"country\":\"Japan\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_2z3jz218jt9o4righx2q0pp5', 'function': {'arguments': '{\"city\":\"Seattle\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}]}), ToolMessage(content='61', tool_call_id='call_toanzknowvq8ymydvjto0w0r', name='Get_Weather'), ToolMessage(content='46', tool_call_id='call_x14745jh6z0z271gdwyjv1z6', name='Get_Weather'), ToolMessage(content='43', tool_call_id='call_2z3jz218jt9o4righx2q0pp5', name='Get_Weather'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_mntkddx3gbk9njo8dts8srpe', 'function': {'arguments': '{\"city\":\"London\",\"country\":\"UK\"}', 'name': 'Get_Weather'}, 'type': 'function'}]}), ToolMessage(content='49', tool_call_id='call_mntkddx3gbk9njo8dts8srpe', name='Get_Weather'), AIMessage(content=' 48')]}\n",
      "\n",
      "---\n",
      "\n",
      "Output from node '__end__':\n",
      "---\n",
      "{'messages': [HumanMessage(content='Where is it warmest: Austin, Texas; Tokyo; or Seattle? And by how much is it warmer than the other cities?'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_toanzknowvq8ymydvjto0w0r', 'function': {'arguments': '{\"city\":\"Austin, Texas\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_x14745jh6z0z271gdwyjv1z6', 'function': {'arguments': '{\"city\":\"Tokyo\",\"country\":\"Japan\"}', 'name': 'Get_Weather'}, 'type': 'function'}, {'id': 'call_2z3jz218jt9o4righx2q0pp5', 'function': {'arguments': '{\"city\":\"Seattle\",\"country\":\"USA\"}', 'name': 'Get_Weather'}, 'type': 'function'}]}), ToolMessage(content='61', tool_call_id='call_toanzknowvq8ymydvjto0w0r', name='Get_Weather'), ToolMessage(content='46', tool_call_id='call_x14745jh6z0z271gdwyjv1z6', name='Get_Weather'), ToolMessage(content='43', tool_call_id='call_2z3jz218jt9o4righx2q0pp5', name='Get_Weather'), AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_mntkddx3gbk9njo8dts8srpe', 'function': {'arguments': '{\"city\":\"London\",\"country\":\"UK\"}', 'name': 'Get_Weather'}, 'type': 'function'}]}), ToolMessage(content='49', tool_call_id='call_mntkddx3gbk9njo8dts8srpe', name='Get_Weather'), AIMessage(content=' 48')]}\n",
      "\n",
      "---\n",
      "\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "inputs = {\n",
    "    \"messages\": [\n",
    "        HumanMessage(\n",
    "            content=\"Where is it warmest: Austin, Texas; Tokyo; or Seattle? And by how much is it warmer than the other cities?\"\n",
    "        )\n",
    "    ]\n",
    "}\n",
    "for output in app.with_config({\"run_name\": \"LLM with Tools\"}).stream(inputs):\n",
    "    # stream() yields dictionaries with output keyed by node name\n",
    "    for key, value in output.items():\n",
    "        print(f\"Output from node '{key}':\")\n",
    "        print(\"---\")\n",
    "        print(value)\n",
    "    print(\"\\n---\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}