Chat Models and Tool Calling.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build a Chatbot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Let's Cover Some Prerequisits to Build ChatBot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####  Chat Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Modern LLMs are typically accessed through a chat model interface that takes a list of messages as input and returns a message as output."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Tool Calling\n",
    "Chat models can call tools to perform tasks such as fetching data from a database, making API requests, or running custom code."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Structured outputs\n",
    "Chat models can be requested to respond in a particular format (e.g., JSON or matching a particular schema).\n",
    "\n",
    "#### Multimodality\n",
    "Large Language Models (LLMs) are not limited to processing text. They can also be used to process other types of data, such as images, audio, and video. This is known as multimodality.\n",
    "\n",
    "### Context window\n",
    "A chat model's context window refers to the maximum size of the input sequence the model can process at one time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prompt templates\n",
    "Prompt templates help to translate user input and parameters into instructions for a language model. This can be used to guide a model's response, helping it understand the context and generate relevant and coherent language-based output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "StringPromptValue(text='Tell me a joke about cats')"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "prompt_template = PromptTemplate.from_template(\"Tell me a joke about {topic}\")\n",
    "\n",
    "prompt_template.invoke({\"topic\": \"cats\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ChatPromptTemplates\n",
    "These prompt templates are used to format a list of messages. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatPromptValue(messages=[SystemMessage(content='Translate the following to urdu:', additional_kwargs={}, response_metadata={}), HumanMessage(content='Where are you going?', additional_kwargs={}, response_metadata={})])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "\n",
    "prompt_temp = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        ('system','Translate the following to urdu:'),\n",
    "        ('user','{text}'),\n",
    "    ]\n",
    ")\n",
    "\n",
    "prompt_temp.invoke({'text':\"Where are you going?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatPromptValue(messages=[SystemMessage(content='You are a helpful assistant', additional_kwargs={}, response_metadata={}), HumanMessage(content='hi!', additional_kwargs={}, response_metadata={})])"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "prompt_template = ChatPromptTemplate([\n",
    "    (\"system\", \"You are a helpful assistant\"),\n",
    "    MessagesPlaceholder(\"msgs\")\n",
    "])\n",
    "\n",
    "prompt_template.invoke({\"msgs\": [HumanMessage(content=\"hi!\")]})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build a Chatbot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ['GROQ_API_KEY']=getpass.getpass()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_groq import ChatGroq\n",
    "model = ChatGroq(model='llama3-8b-8192', api_key=os.environ[\"GROQ_API_KEY\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='کہاں جائے گا بھائی?\\n\\n(Kahaan jaaye ga bhai?)', additional_kwargs={}, response_metadata={'token_usage': {'completion_tokens': 26, 'prompt_tokens': 27, 'total_tokens': 53, 'completion_time': 0.021666667, 'prompt_time': 0.001156882, 'queue_time': 0.012016616, 'total_time': 0.022823549}, 'model_name': 'llama3-8b-8192', 'system_fingerprint': 'fp_179b0f92c9', 'finish_reason': 'stop', 'logprobs': None}, id='run-e4268066-3007-4e0e-b15c-53fdac32ea72-0', usage_metadata={'input_tokens': 27, 'output_tokens': 26, 'total_tokens': 53})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "\n",
    "message = [\n",
    "    SystemMessage(content='Translate the following into Urdu:'),\n",
    "    HumanMessage(content='Where are you going bro?'),\n",
    "]\n",
    "\n",
    "\n",
    "model.invoke(message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "parser = StrOutputParser()\n",
    "\n",
    "# response = model.invoke(message)\n",
    "# parser.invoke(response)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'کہاں جارہے بھائی؟ (Kahaan jaareh bhai?)'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = model | parser\n",
    "\n",
    "chain.invoke(message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"Hi Bob! It's nice to meet you. Is there something I can help you with or would you like to chat?\", additional_kwargs={}, response_metadata={'token_usage': {'completion_tokens': 26, 'prompt_tokens': 15, 'total_tokens': 41, 'completion_time': 0.021666667, 'prompt_time': 0.000142649, 'queue_time': 0.014490630000000001, 'total_time': 0.021809316}, 'model_name': 'llama3-8b-8192', 'system_fingerprint': 'fp_179b0f92c9', 'finish_reason': 'stop', 'logprobs': None}, id='run-b9da5b5d-fcf9-4ad3-a2c5-be3c30c2b049-0', usage_metadata={'input_tokens': 15, 'output_tokens': 26, 'total_tokens': 41})"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage\n",
    "\n",
    "model.invoke([HumanMessage(content=\"Hi! I'm Bob\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\"I apologize, but I'm a large language model, I don't have any information about your name. I'm a new conversation every time you interact with me, and I don't retain any personal information about individuals. If you'd like to share your name with me, I'd be happy to know it!\", additional_kwargs={}, response_metadata={'token_usage': {'completion_tokens': 64, 'prompt_tokens': 15, 'total_tokens': 79, 'completion_time': 0.053333333, 'prompt_time': 0.002308046, 'queue_time': 0.011451444, 'total_time': 0.055641379}, 'model_name': 'llama3-8b-8192', 'system_fingerprint': 'fp_6a6771ae9c', 'finish_reason': 'stop', 'logprobs': None}, id='run-b52b4066-d60d-4d41-850d-e1d4897348c8-0', usage_metadata={'input_tokens': 15, 'output_tokens': 64, 'total_tokens': 79})"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.invoke([HumanMessage(content=\"What's my name?\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='I remember! Your name is Salman!', additional_kwargs={}, response_metadata={'token_usage': {'completion_tokens': 9, 'prompt_tokens': 44, 'total_tokens': 53, 'completion_time': 0.0075, 'prompt_time': 0.00557453, 'queue_time': 0.00770922, 'total_time': 0.01307453}, 'model_name': 'llama3-8b-8192', 'system_fingerprint': 'fp_6a6771ae9c', 'finish_reason': 'stop', 'logprobs': None}, id='run-83c84f9c-2a99-4546-b0f5-70c1a415292e-0', usage_metadata={'input_tokens': 44, 'output_tokens': 9, 'total_tokens': 53})"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.messages import AIMessage\n",
    "\n",
    "model.invoke(\n",
    "    [\n",
    "        HumanMessage(content=\"Hello! I am Salman\"),\n",
    "        AIMessage(content=\"Hello Salman! Nice to meet you. How can i assist you today?\"),\n",
    "        HumanMessage(content=\"What's my name?\")\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Message persistence\n",
    "\n",
    "LangGraph implements a built-in persistence layer, making it ideal for chat applications that support multiple conversational turns.\n",
    "\n",
    "Wrapping our chat model in a minimal LangGraph application allows us to automatically persist the message history, simplifying the development of multi-turn applications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "from langgraph.graph import MessagesState, StateGraph, START\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow = StateGraph(state_schema=MessagesState)\n",
    "\n",
    "def call_model(state: MessagesState):\n",
    "    response = model.invoke(state['messages'])\n",
    "    return {'messages':response}\n",
    "\n",
    "workflow.add_edge(START, \"model\")\n",
    "workflow.add_node(\"model\", call_model)\n",
    "\n",
    "memory = MemorySaver()\n",
    "app = workflow.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A Quick Look At async "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task 1 started\n",
      "Task 1 finished\n",
      "Task 2 started\n",
      "Task 2 finished\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "def task1():\n",
    "    print(\"Task 1 started\")\n",
    "    time.sleep(2)  # Simulate long-running task\n",
    "    print(\"Task 1 finished\")\n",
    "\n",
    "def task2():\n",
    "    print(\"Task 2 started\")\n",
    "    time.sleep(3)  # Simulate long-running task\n",
    "    print(\"Task 2 finished\")\n",
    "\n",
    "task1()\n",
    "task2()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Task 1 started\n",
      "Task 2 started\n",
      "Task 1 finished\n",
      "Task 2 finished\n",
      "Total time: 3.000969171524048 seconds\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "import time\n",
    "async def task1():\n",
    "    print(\"Task 1 started\")\n",
    "    await asyncio.sleep(2)  # Simulate long-running task\n",
    "    print(\"Task 1 finished\")\n",
    "\n",
    "async def task2():\n",
    "    print(\"Task 2 started\")\n",
    "    await asyncio.sleep(3)  # Simulate long-running task\n",
    "    print(\"Task 2 finished\")\n",
    "\n",
    "async def main():\n",
    "    start_time = time.time()\n",
    "    await asyncio.gather(task1(), task2())\n",
    "    print(f\"Total time: {time.time() - start_time} seconds\")\n",
    "\n",
    "# asyncio.run(main())\n",
    "await main()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building Chat Model Continue..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Hi Salman! Nice to meet you!\n"
     ]
    }
   ],
   "source": [
    "config = {'configurable':{'thread_id':'abc123'}}\n",
    "\n",
    "\n",
    "query = \"Hi! I am Salman.\"\n",
    "input_messages = [HumanMessage(content=query)]\n",
    "output = app.invoke(\n",
    "    {'messages':input_messages}, config\n",
    ")\n",
    "\n",
    "output['messages'][-1].pretty_print() # output contains all messages in state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "According to our conversation, your name is Salman!\n"
     ]
    }
   ],
   "source": [
    "query = \"what is my name?\"\n",
    "\n",
    "input_messages = [HumanMessage(content=query)]\n",
    "\n",
    "output = app.invoke(\n",
    "    {'messages':input_messages}, config\n",
    ")\n",
    "\n",
    "output['messages'][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "I apologize, but I'm a large language model, I don't have the ability to know your name. I'm a new conversation each time you interact with me, and I don't retain any information about you or your identity. If you'd like to share your name with me, I'd be happy to chat with you and remember it for our conversation!\n"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"abc234\"}}\n",
    "\n",
    "input_messages = [HumanMessage(query)]\n",
    "output = app.invoke({\"messages\": input_messages}, config)\n",
    "output[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Here's one:\n",
      "\n",
      "Why couldn't the bicycle stand up by itself?\n",
      "\n",
      "(Wait for it...)\n",
      "\n",
      "Because it was two-tired!\n",
      "\n",
      "Hope that made you laugh!\n"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"abc234\"}}\n",
    "query = \"tell me a joke\"\n",
    "input_messages = [HumanMessage(query)]\n",
    "output = app.invoke({\"messages\": input_messages}, config)\n",
    "output[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, we can always go back to the original conversation (since we are persisting it in a database)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Your name is Salman!\n"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"abc123\"}}\n",
    "query = \"What is my name?\"\n",
    "input_messages = [HumanMessage(query)]\n",
    "output = app.invoke({\"messages\": input_messages}, config)\n",
    "output[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ChatPrompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        ('system', \"You're a comedian. You will be roasting everyone with jokes.\"),\n",
    "        MessagesPlaceholder(variable_name='messages')\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow = StateGraph(state_schema=MessagesState)\n",
    "\n",
    "def call_model(state: MessagesState):\n",
    "    chain = prompt | model\n",
    "    response = chain.invoke(state)\n",
    "    return {'messages':response}\n",
    "\n",
    "workflow.add_edge(START, \"model\")\n",
    "workflow.add_node(\"model\", call_model)\n",
    "\n",
    "memory = MemorySaver()\n",
    "app = workflow.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Great to meet you, Jim! I noticed you introduced yourself with a lot of enthusiasm... or was that just the excitement of knowing you're about to be roasted?\n"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"abc345\"}}\n",
    "query = \"Hi! I'm Jim.\"\n",
    "\n",
    "input_messages = [HumanMessage(query)]\n",
    "output = app.invoke({\"messages\": input_messages}, config)\n",
    "output[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Making our chat prompt more complex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are a helpful assistant. Answer all questions to the best of your ability in {language}.\",\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"messages\"),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Some Notes on typing and typing_extensions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "12\n",
      "{'age': 22} [295, 310] (1062, 1035) {'age', 'college_rn', 'uni_rn'}\n"
     ]
    }
   ],
   "source": [
    "from typing import Set, Dict, Tuple, List, Callable, Optional, Union, Literal, Any\n",
    "\n",
    "# Callable\n",
    "def execute_operation(operation: Callable[[int, int], int], a: int, b: int) -> int:\n",
    "    return operation(a, b)\n",
    "print(execute_operation(lambda x,y:x+y, 5,7))\n",
    "\n",
    "from typing import Dict, List, Tuple, Set\n",
    "\n",
    "def analyze_data(data: Dict[str, int], items: List[str], points: Tuple[int, int], flags: Optional[Set[str]]=None) -> None:\n",
    "    print(data, items, points, flags)\n",
    "\n",
    "analyze_data({'age':22}, [295, 310], (1062, 1035), {'age', 'college_rn', 'uni_rn'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'https://linkedin.com/in/msalmanai62/'"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing_extensions import Final\n",
    "\n",
    "API_URL: Final = \"https://linkedin.com/in/msalmanai62/\"\n",
    "API_URL\n",
    "# This value cannot be changed or overridden\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing_extensions import Annotated\n",
    "\n",
    "def process_data(data: Annotated[int, \"Must be positive\"]) -> int:\n",
    "    return data * 2\n",
    "\n",
    "process_data(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Sequence\n",
    "\n",
    "from langchain_core.messages import BaseMessage\n",
    "from langgraph.graph.message import add_messages\n",
    "from typing_extensions import Annotated, TypedDict\n",
    "\n",
    "\n",
    "class State(TypedDict):\n",
    "    messages: Annotated[Sequence[BaseMessage], add_messages]\n",
    "    language: str\n",
    "\n",
    "\n",
    "workflow = StateGraph(state_schema=State)\n",
    "\n",
    "\n",
    "def call_model(state: State):\n",
    "    chain = prompt | model\n",
    "    response = chain.invoke(state)\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "workflow.add_edge(START, \"model\")\n",
    "workflow.add_node(\"model\", call_model)\n",
    "\n",
    "memory = MemorySaver()\n",
    "app = workflow.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Asalamu alaikum Salman! Main Salman ki madad karnay ki taiyari mein hoon. Kya aap mera madad chah rahi hain?\n"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"abc456\"}}\n",
    "query = \"Hi! I'm salman.\"\n",
    "language = \"urdu\"\n",
    "\n",
    "input_messages = [HumanMessage(query)]\n",
    "output = app.invoke(\n",
    "    {\"messages\": input_messages, \"language\": language},\n",
    "    config,\n",
    ")\n",
    "output[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Salman yaar, aapki nam Salman hai!\n"
     ]
    }
   ],
   "source": [
    "query = \"What is my name?\"\n",
    "\n",
    "input_messages = [HumanMessage(query)]\n",
    "output = app.invoke(\n",
    "    {\"messages\": input_messages},\n",
    "    config,\n",
    ")\n",
    "output[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Managing Conversation History"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[SystemMessage(content=\"you're a good assistant\", additional_kwargs={}, response_metadata={}),\n",
       " HumanMessage(content=\"hi! I'm bob\", additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='hi!', additional_kwargs={}, response_metadata={}),\n",
       " HumanMessage(content='I like vanilla ice cream', additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='nice', additional_kwargs={}, response_metadata={}),\n",
       " HumanMessage(content='whats 2 + 2', additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='4', additional_kwargs={}, response_metadata={}),\n",
       " HumanMessage(content='thanks', additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='no problem!', additional_kwargs={}, response_metadata={}),\n",
       " HumanMessage(content='having fun?', additional_kwargs={}, response_metadata={}),\n",
       " AIMessage(content='yes!', additional_kwargs={}, response_metadata={})]"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.messages import SystemMessage, trim_messages\n",
    "from langchain_core.messages import AIMessage\n",
    "\n",
    "trimmer = trim_messages(\n",
    "    max_tokens=65,\n",
    "    strategy=\"last\",\n",
    "    token_counter=model,\n",
    "    include_system=True,\n",
    "    allow_partial=False,\n",
    "    start_on=\"human\",\n",
    ")\n",
    "\n",
    "messages = [\n",
    "    SystemMessage(content=\"you're a good assistant\"),\n",
    "    HumanMessage(content=\"hi! I'm bob\"),\n",
    "    AIMessage(content=\"hi!\"),\n",
    "    HumanMessage(content=\"I like vanilla ice cream\"),\n",
    "    AIMessage(content=\"nice\"),\n",
    "    HumanMessage(content=\"whats 2 + 2\"),\n",
    "    AIMessage(content=\"4\"),\n",
    "    HumanMessage(content=\"thanks\"),\n",
    "    AIMessage(content=\"no problem!\"),\n",
    "    HumanMessage(content=\"having fun?\"),\n",
    "    AIMessage(content=\"yes!\"),\n",
    "]\n",
    "\n",
    "trimmer.invoke(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow = StateGraph(state_schema=State)\n",
    "\n",
    "\n",
    "def call_model(state: State):\n",
    "    chain = prompt | model\n",
    "    trimmed_messages = trimmer.invoke(state[\"messages\"])\n",
    "    response = chain.invoke(\n",
    "        {\"messages\": trimmed_messages, \"language\": state[\"language\"]}\n",
    "    )\n",
    "    return {\"messages\": [response]}\n",
    "\n",
    "\n",
    "workflow.add_edge(START, \"model\")\n",
    "workflow.add_node(\"model\", call_model)\n",
    "\n",
    "memory = MemorySaver()\n",
    "app = workflow.compile(checkpointer=memory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "Your name is Bob!\n"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"abc567\"}}\n",
    "query = \"What is my name?\"\n",
    "language = \"English\"\n",
    "\n",
    "input_messages = messages + [HumanMessage(query)]\n",
    "output = app.invoke(\n",
    "    {\"messages\": input_messages, \"language\": language},\n",
    "    config,\n",
    ")\n",
    "output[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi Todd! Here's one for you:\n",
      "\n",
      "Why couldn't the bicycle stand up by itself?\n",
      "\n",
      "(Wait for it...)\n",
      "\n",
      "Because it was two-tired!\n",
      "\n",
      "Hope that made you laugh! Do you want to hear another one?"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"abc789\"}}\n",
    "query = \"Hi I'm Todd, please tell me a joke.\"\n",
    "language = \"English\"\n",
    "\n",
    "input_messages = [HumanMessage(query)]\n",
    "for chunk, metadata in app.stream(\n",
    "    {\"messages\": input_messages, \"language\": language},\n",
    "    config,\n",
    "    stream_mode=\"messages\",\n",
    "):\n",
    "    if isinstance(chunk, AIMessage):  # Filter to just model responses\n",
    "        print(chunk.content, end=\"\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "hf",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}