example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to connect a chat bot to your memory service"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import dotenv\n",
    "\n",
    "dotenv.load_dotenv(\".env\", override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langgraph_sdk import get_client\n",
    "\n",
    "# Update to your URL. Copy this from page of ryour LangGraph Deployment\n",
    "deployment_url = \"\"\n",
    "\n",
    "client = get_client(url=deployment_url)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example Chat Bot\n",
    "\n",
    "The bot fetches user memories my semantic similarity, templates them, then responds!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import uuid\n",
    "from datetime import datetime, timezone\n",
    "from typing import List, Optional\n",
    "\n",
    "import langsmith\n",
    "from langchain.chat_models import init_chat_model\n",
    "from langchain_core.messages import AnyMessage\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain_core.runnables import RunnableConfig\n",
    "from langgraph.checkpoint import MemorySaver\n",
    "from langgraph.graph import START, StateGraph, add_messages\n",
    "from langgraph_sdk import get_client\n",
    "from pydantic.v1 import BaseModel, Field\n",
    "from typing_extensions import Annotated, TypedDict\n",
    "\n",
    "from memory_service import (\n",
    "    _constants as constants,\n",
    "    _settings as settings,\n",
    "    _utils as utils,\n",
    ")\n",
    "\n",
    "\n",
    "class ChatState(TypedDict):\n",
    "    \"\"\"The state of the chatbot.\"\"\"\n",
    "\n",
    "    messages: Annotated[List[AnyMessage], add_messages]\n",
    "    user_memories: List[dict]\n",
    "\n",
    "\n",
    "class ChatConfigurable(TypedDict):\n",
    "    \"\"\"The configurable fields for the chatbot.\"\"\"\n",
    "\n",
    "    user_id: str\n",
    "    thread_id: str\n",
    "    memory_service_url: str = \"\"\n",
    "    model: str\n",
    "    delay: Optional[float]\n",
    "\n",
    "\n",
    "def _ensure_configurable(config: RunnableConfig) -> ChatConfigurable:\n",
    "    \"\"\"Ensure the configuration is valid.\"\"\"\n",
    "    return ChatConfigurable(\n",
    "        user_id=config[\"configurable\"][\"user_id\"],\n",
    "        thread_id=config[\"configurable\"][\"thread_id\"],\n",
    "        mem_assistant_id=config[\"configurable\"][\"mem_assistant_id\"],\n",
    "        memory_service_url=config[\"configurable\"].get(\n",
    "            \"memory_service_url\", os.environ.get(\"MEMORY_SERVICE_URL\", \"\")\n",
    "        ),\n",
    "        model=config[\"configurable\"].get(\n",
    "            \"model\", \"accounts/fireworks/models/firefunction-v2\"\n",
    "        ),\n",
    "        delay=config[\"configurable\"].get(\"delay\", 60),\n",
    "    )\n",
    "\n",
    "\n",
    "PROMPT = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You are a helpful and friendly chatbot. Get to know the user!\"\n",
    "            \" Ask questions! Be spontaneous!\"\n",
    "            \"{user_info}\\n\\nSystem Time: {time}\",\n",
    "        ),\n",
    "        (\"placeholder\", \"{messages}\"),\n",
    "    ]\n",
    ").partial(\n",
    "    time=lambda: datetime.now(timezone.utc).strftime(\"%Y-%m-%d %H:%M:%S\"),\n",
    ")\n",
    "\n",
    "\n",
    "@langsmith.traceable\n",
    "def format_query(messages: List[AnyMessage]) -> str:\n",
    "    \"\"\"Format the query for the user's memories.\"\"\"\n",
    "    # This is quite naive :)\n",
    "    return \" \".join([str(m.content) for m in messages if m.type == \"human\"][-5:])\n",
    "\n",
    "\n",
    "async def query_memories(state: ChatState, config: RunnableConfig) -> ChatState:\n",
    "    \"\"\"Query the user's memories.\"\"\"\n",
    "    configurable: ChatConfigurable = config[\"configurable\"]\n",
    "    user_id = configurable[\"user_id\"]\n",
    "    index = utils.get_index()\n",
    "    embeddings = utils.get_embeddings()\n",
    "\n",
    "    query = format_query(state[\"messages\"])\n",
    "    vec = await embeddings.aembed_query(query)\n",
    "    # You can also filter by memory type, etc. here.\n",
    "    with langsmith.trace(\n",
    "        \"pinecone_query\", inputs={\"query\": query, \"user_id\": user_id}\n",
    "    ) as rt:\n",
    "        response = index.query(\n",
    "            vector=vec,\n",
    "            filter={\"user_id\": {\"$eq\": str(user_id)}},\n",
    "            include_metadata=True,\n",
    "            top_k=10,\n",
    "            namespace=settings.SETTINGS.pinecone_namespace,\n",
    "        )\n",
    "        rt.outputs[\"response\"] = response\n",
    "    memories = []\n",
    "    if matches := response.get(\"matches\"):\n",
    "        memories = [m[\"metadata\"][constants.PAYLOAD_KEY] for m in matches]\n",
    "    return {\n",
    "        \"user_memories\": memories,\n",
    "    }\n",
    "\n",
    "\n",
    "@langsmith.traceable\n",
    "def format_memories(memories: List[dict]) -> str:\n",
    "    \"\"\"Format the user's memories.\"\"\"\n",
    "    if not memories:\n",
    "        return \"\"\n",
    "    # Note Bene: You can format better than this....\n",
    "    memories = \"\\n\".join(str(m) for m in memories)\n",
    "    return f\"\"\"\n",
    "\n",
    "## Memories\n",
    "\n",
    "You have noted the following memorable events from previous interactions with the user.\n",
    "<memories>\n",
    "{memories}\n",
    "</memories>\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "async def bot(state: ChatState, config: RunnableConfig) -> ChatState:\n",
    "    \"\"\"Prompt the bot to resopnd to the user, incorporating memories (if provided).\"\"\"\n",
    "    configurable = _ensure_configurable(config)\n",
    "    model = init_chat_model(configurable[\"model\"])\n",
    "    chain = PROMPT | model\n",
    "    memories = format_memories(state[\"user_memories\"])\n",
    "    m = await chain.ainvoke(\n",
    "        {\n",
    "            \"messages\": state[\"messages\"],\n",
    "            \"user_info\": memories,\n",
    "        },\n",
    "        config,\n",
    "    )\n",
    "\n",
    "    return {\n",
    "        \"messages\": [m],\n",
    "    }\n",
    "\n",
    "\n",
    "class MemorableEvent(BaseModel):\n",
    "    \"\"\"A memorable event.\"\"\"\n",
    "\n",
    "    description: str\n",
    "    participants: List[str] = Field(\n",
    "        description=\"Names of participants in the event and their relationship to the user.\"\n",
    "    )\n",
    "\n",
    "\n",
    "async def post_messages(state: ChatState, config: RunnableConfig) -> ChatState:\n",
    "    \"\"\"Query the user's memories.\"\"\"\n",
    "    configurable = _ensure_configurable(config)\n",
    "    langgraph_client = get_client(url=configurable[\"memory_service_url\"])\n",
    "    thread_id = config[\"configurable\"][\"thread_id\"]\n",
    "    # Hash \"memory_{thread_id}\" to get a new uuid5 for the memory id\n",
    "    memory_thread_id = uuid.uuid5(uuid.NAMESPACE_URL, f\"memory_{thread_id}\")\n",
    "    try:\n",
    "        await langgraph_client.threads.get(thread_id=memory_thread_id)\n",
    "    except Exception:\n",
    "        await langgraph_client.threads.create(thread_id=memory_thread_id)\n",
    "\n",
    "    await langgraph_client.runs.create(\n",
    "        memory_thread_id,\n",
    "        assistant_id=configurable[\"mem_assistant_id\"],\n",
    "        input={\n",
    "            \"messages\": state[\"messages\"],  # the service dedupes messages\n",
    "        },\n",
    "        config={\n",
    "            \"configurable\": {\n",
    "                \"user_id\": configurable[\"user_id\"],\n",
    "            },\n",
    "        },\n",
    "        multitask_strategy=\"rollback\",\n",
    "    )\n",
    "\n",
    "    return {\n",
    "        \"messages\": [],\n",
    "    }\n",
    "\n",
    "\n",
    "builder = StateGraph(ChatState, ChatConfigurable)\n",
    "builder.add_node(query_memories)\n",
    "builder.add_node(bot)\n",
    "builder.add_node(post_messages)\n",
    "builder.add_edge(START, \"query_memories\")\n",
    "builder.add_edge(\"query_memories\", \"bot\")\n",
    "builder.add_edge(\"bot\", \"post_messages\")\n",
    "\n",
    "chat_graph = builder.compile(checkpointer=MemorySaver())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "mem_assistant = await client.assistants.create(\n",
    "    graph_id=\"memory\",\n",
    "    config={\n",
    "        \"configurable\": {\n",
    "            \"delay\": 4,  # seconds wait before considering a thread as \"completed\"\n",
    "            \"schemas\": {\n",
    "                \"MemorableEvent\": {\n",
    "                    \"system_prompt\": \"Extract any memorable events from the user's\"\n",
    "                    \" messages that you would like to remember.\",\n",
    "                    \"update_mode\": \"insert\",\n",
    "                    \"function\": MemorableEvent.schema(),\n",
    "                },\n",
    "            },\n",
    "        }\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# mem_assistant = (await client.assistants.search(graph_id=\"memory\"))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uuid\n",
    "\n",
    "user_id = str(uuid.uuid4())  # more permanent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'thread_id': '3ff82998-b622-421f-8c8c-4b14d10c17b1',\n",
       " 'created_at': '2024-06-28T00:42:23.884229+00:00',\n",
       " 'updated_at': '2024-06-28T00:42:23.884229+00:00',\n",
       " 'metadata': {},\n",
       " 'status': 'idle'}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "thread_id = str(uuid.uuid4())  # can adjust\n",
    "await client.threads.create(thread_id=thread_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Chat:\n",
    "    def __init__(self, user_id: str, thread_id: str):\n",
    "        self.thread_id = thread_id\n",
    "        self.user_id = user_id\n",
    "\n",
    "    async def __call__(self, query: str) -> str:\n",
    "        chunks = chat_graph.astream_events(\n",
    "            input={\n",
    "                \"messages\": [(\"user\", query)],\n",
    "            },\n",
    "            config={\n",
    "                \"configurable\": {\n",
    "                    \"user_id\": self.user_id,\n",
    "                    \"thread_id\": self.thread_id,\n",
    "                    \"memory_service_url\": deployment_url,\n",
    "                    \"mem_assistant_id\": mem_assistant[\"assistant_id\"],\n",
    "                    \"delay\": 4,\n",
    "                }\n",
    "            },\n",
    "            version=\"v2\",\n",
    "        )\n",
    "        res = \"\"\n",
    "        async for event in chunks:\n",
    "            if event.get(\"event\") == \"on_chat_model_stream\":\n",
    "                tok = event[\"data\"][\"chunk\"].content\n",
    "                print(tok, end=\"\")\n",
    "                res += tok\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat = Chat(user_id, thread_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hi! It's nice to meet you. What brings you here today?"
     ]
    }
   ],
   "source": [
    "_ = await chat(\"Hi there\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "That's so sweet of you! I'm sure Steve will appreciate the effort you're putting into making him feel special. What's the theme of the party going to be? Has Steve mentioned anything he's been into lately that you could incorporate into the celebration?"
     ]
    }
   ],
   "source": [
    "_ = await chat(\n",
    "    \"I've been planning a surprise party for my friend steve. \"\n",
    "    \"He has been having a rough month and I want it to be special.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "That's a great idea! Crocheting is a unique interest, and incorporating it into the party could make it really special and personalized to Steve. You could decorate with crocheted items, have a \"crochet station\" where guests can make their own simple projects, or even have a crochet-themed cake. What do you think Steve's favorite colors or yarn types are? That could help you get started with planning."
     ]
    }
   ],
   "source": [
    "_ = await chat(\n",
    "    \"Steve really likes crocheting. Maybe I can do something with that? Or is that dumb... \"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Whoa, that's cool! Capoeira is such a dynamic and energetic activity. You could definitely incorporate elements of it into the party to make it more exciting. Maybe you could hire a capoeira instructor to lead a short workshop or demo, or even have a mini \"roda\" (that's the circle where capoeiristas play, right?) set up for guests to try out some moves. What do you think Steve would think of that?"
     ]
    }
   ],
   "source": [
    "_ = await chat(\"He's also into capoeira...\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "That's a great connection to have! It's always helpful to get recommendations from people who have experience with a particular activity or business. You could reach out to the studio and ask if they know of any instructors who might be available to lead a workshop or demo at the party. They might even have some suggestions for how to incorporate capoeira into the celebration in a way that would be fun and engaging for Steve and the other guests. Do you think you'll reach out to them today, or wait until later in the planning process?"
     ]
    }
   ],
   "source": [
    "_ = await chat(\n",
    "    \"Oh that's a cool idea. One time i took classes from this studio nearby. Wonder if they have any recs. \"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I'm doing great, thanks for asking! I'm just happy to be chatting with you and helping with your party planning. It's always exciting to see people come together to celebrate special occasions. But enough about me - let's get back to Steve's party! What do you think about serving some Brazilian-inspired food and drinks to tie in with the capoeira theme?"
     ]
    }
   ],
   "source": [
    "_ = await chat(\"Idk. Anyways - how are you doing?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Nice to meet you, Ken! I'm glad we could chat about Steve's party and get some ideas going. If you need any more help or just want to bounce some ideas off me, feel free to reach out anytime. Good luck with the planning, and I hope Steve has an amazing time!"
     ]
    }
   ],
   "source": [
    "_ = await chat(\"My name is Ken btw\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convo 2\n",
    "\n",
    "Our memory is configured only to consider a thread \"ready to process\" if has been inactive for a minute.\n",
    "We'll wait for things to populate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "\n",
    "await asyncio.sleep(60)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "thread_id_2 = uuid.uuid4()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat2 = Chat(user_id, thread_id_2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I remember you! We were planning a surprise party for Steve, and Ken was also involved. How's everything going? Did the party turn out well?"
     ]
    }
   ],
   "source": [
    "_ = await chat2(\"Remember me?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I remember because I have a special memory book where I keep track of all the fun conversations and events we've shared together! It's like a digital scrapbook, and it helps me remember important details about our chats."
     ]
    }
   ],
   "source": [
    "_ = await chat2(\"wdy remember??\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "That's great to hear! I'm glad to know that the planning is going smoothly. Are there any new developments or updates that you'd like to share about the party? Maybe I can even offer some suggestions or ideas to make it an even more special celebration for Steve!"
     ]
    }
   ],
   "source": [
    "_ = await chat2(\"Oh planning is going alright!\")"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}