{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-29T20:09:11.440091Z",
     "iopub.status.busy": "2025-01-29T20:09:11.439766Z",
     "iopub.status.idle": "2025-01-29T20:09:11.751153Z",
     "shell.execute_reply": "2025-01-29T20:09:11.750263Z",
     "shell.execute_reply.started": "2025-01-29T20:09:11.440060Z"
    },
    "id": "xaiioUQni_ga",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "from modules.data_class import DataState\n",
    "from modules.tools import data_node\n",
    "from modules.nodes import chatbot_with_tools, human_node, maybe_exit_human_node, maybe_route_to_tools\n",
    "\n",
    "from langgraph.graph import StateGraph, START, END\n",
    "from langgraph.checkpoint.memory import MemorySaver\n",
    "\n",
    "from IPython.display import Image, display\n",
    "from pprint import pprint\n",
    "from typing import Literal\n",
    "\n",
    "from langgraph.prebuilt import ToolNode\n",
    "\n",
    "from collections.abc import Iterable\n",
    "from IPython.display import display, clear_output\n",
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-29T20:09:11.906458Z",
     "iopub.status.busy": "2025-01-29T20:09:11.905241Z",
     "iopub.status.idle": "2025-01-29T20:09:11.994921Z",
     "shell.execute_reply": "2025-01-29T20:09:11.993761Z",
     "shell.execute_reply.started": "2025-01-29T20:09:11.906419Z"
    },
    "id": "9rqkQzlZxrzp",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "graph_builder = StateGraph(DataState)\n",
    "memory = MemorySaver()\n",
    "\n",
    "# Nodes\n",
    "graph_builder.add_node(\"chatbot_healthassistant\", chatbot_with_tools)\n",
    "graph_builder.add_node(\"patient\", human_node)\n",
    "graph_builder.add_node(\"documenting\", data_node)\n",
    "\n",
    "# Chatbot -> {ordering, tools, human, END}\n",
    "graph_builder.add_conditional_edges(\"chatbot_healthassistant\", maybe_route_to_tools)\n",
    "# Human -> {chatbot, END}\n",
    "graph_builder.add_conditional_edges(\"patient\", maybe_exit_human_node)\n",
    "# TestCase_Paintrek\n",
    "# Tools (both kinds) always route back to chat afterwards.\n",
    "graph_builder.add_edge(\"documenting\", \"chatbot_healthassistant\")\n",
    "\n",
    "graph_builder.add_edge(START, \"chatbot_healthassistant\")\n",
    "graph_with_order_tools = graph_builder.compile(checkpointer=memory)\n",
    "\n",
    "# Image(graph_with_order_tools.get_graph().draw_mermaid_png())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-01-29T20:09:38.185616Z",
     "iopub.status.busy": "2025-01-29T20:09:38.185131Z",
     "iopub.status.idle": "2025-01-29T20:10:08.474591Z",
     "shell.execute_reply": "2025-01-29T20:10:08.472926Z",
     "shell.execute_reply.started": "2025-01-29T20:09:38.185577Z"
    },
    "id": "NCRSgaBUfIHF",
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Executing the chatbot graph...\n",
      "Model: Welcome to the Paintrek world. I am a health assistant, an interactive clinical recording system. I will ask you questions about your pain and related symptoms and record your responses.  I will then store this information securely. At any time, you can type `q` to quit.\n",
      "User: Let's start\n",
      "Model: Great! Let’s begin by gathering some basic information. Could you please provide your full name, date of birth, and contact number? This will help us proceed with the medical record setup.\n",
      "User: My name is Frank my DOB is 19860214, my contact number is 123456\n",
      "Model: Thank you for providing your information, Frank. Could you please confirm your full date of birth and provide any additional contact details such as an email address or another phone number? Also, could you tell me about your emergency contacts if applicable?\n",
      "If you're ready to proceed with the rest of the intake process now, let me know!\n",
      "User: my date of birth is Feb 14 1986, emergency contact is Zoe and her phone number is 568790\n",
      "Your input data:\n",
      "{'ID': {'name': '', 'DOB': datetime.date(1900, 1, 1), 'gender': '', 'contact': '', 'emergency_contact': ''}, 'symptom': {'main_symptom': '', 'symptom_length': ''}, 'pain': {'pain_location': '', 'pain_side': '', 'pain_intensity': 0, 'pain_description': '', 'start_time': datetime.date(1900, 1, 1), 'radiation': False, 'triggers': '', 'symptom': ''}, 'medical_hist': {'medical_condition': '', 'first_time': datetime.date(1900, 1, 1), 'surgery_history': '', 'medication': '', 'allergy': ''}, 'family_hist': {'family_history': ''}, 'social_hist': {'occupation': '', 'smoke': False, 'alcohol': False, 'drug': False, 'support_system': '', 'living_condition': ''}, 'review_system': {'weight_change': '', 'fever': False, 'chill': False, 'night_sweats': False, 'sleep': '', 'gastrointestinal': '', 'urinary': ''}, 'pain_manage': {'pain_medication': '', 'specialist': False, 'other_therapy': '', 'effectiveness': False}, 'functional': {'life_quality': '', 'limit_activity': '', 'mood': ''}, 'plan': {'goal': '', 'expectation': '', 'alternative_treatment': ''}}\n",
      "Model: Let's start by gathering your personal information and medical history. Could you please provide me with the following details:\n",
      "\n",
      "1. Your full name, date of birth, and contact information.\n",
      "2. The reason for your visit today (symptoms or concerns).\n",
      "3. Any current pain issues you are experiencing.\n",
      "\n",
      "Once I have this initial information, we can proceed to gather more detailed medical history and other relevant data. Please start with the first two points if possible. For the third point, could you describe any areas where you're feeling pain, how long it has been going on, and what makes it better or worse?\n",
      "User: q\n"
     ]
    }
   ],
   "source": [
    "# This is for the checkpointer to have a thread id to remember for each user, and I don't know if letter mix with number will work here \n",
    "# The default recursion limit for traversing nodes is 25 - setting it higher means you can try a more complex order with multiple steps and round-trips.\n",
    "config = {\"configurable\": {\"thread_id\": \"1\"}, \"recursion_limit\": 1000}\n",
    "\n",
    "# Uncomment this line to execute the graph:\n",
    "# Clear output before running new states\n",
    "clear_output(wait=True)\n",
    "\n",
    "# Ensure messages print immediately\n",
    "print(\"Executing the chatbot graph...\", flush=True)\n",
    "state = graph_with_order_tools.invoke({\"messages\": []}, config)\n",
    "# display(state)  # Ensures state is shown in Jupyter\n",
    "# sys.stdout.flush()\n",
    "\n",
    "# pprint(state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'ID': {'name': '',\n",
       "  'DOB': datetime.date(1900, 1, 1),\n",
       "  'gender': '',\n",
       "  'contact': '',\n",
       "  'emergency_contact': ''},\n",
       " 'symptom': {'main_symptom': '', 'symptom_length': ''},\n",
       " 'pain': {'pain_location': '',\n",
       "  'pain_side': '',\n",
       "  'pain_intensity': 0,\n",
       "  'pain_description': '',\n",
       "  'start_time': datetime.date(1900, 1, 1),\n",
       "  'radiation': False,\n",
       "  'triggers': '',\n",
       "  'symptom': ''},\n",
       " 'medical_hist': {'medical_condition': '',\n",
       "  'first_time': datetime.date(1900, 1, 1),\n",
       "  'surgery_history': '',\n",
       "  'medication': '',\n",
       "  'allergy': ''},\n",
       " 'family_hist': {'family_history': ''},\n",
       " 'social_hist': {'occupation': '',\n",
       "  'smoke': False,\n",
       "  'alcohol': False,\n",
       "  'drug': False,\n",
       "  'support_system': '',\n",
       "  'living_condition': ''},\n",
       " 'review_system': {'weight_change': '',\n",
       "  'fever': False,\n",
       "  'chill': False,\n",
       "  'night_sweats': False,\n",
       "  'sleep': '',\n",
       "  'gastrointestinal': '',\n",
       "  'urinary': ''},\n",
       " 'pain_manage': {'pain_medication': '',\n",
       "  'specialist': False,\n",
       "  'other_therapy': '',\n",
       "  'effectiveness': False},\n",
       " 'functional': {'life_quality': '', 'limit_activity': '', 'mood': ''},\n",
       " 'plan': {'goal': '', 'expectation': '', 'alternative_treatment': ''}}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "state[\"data\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "name": "day-3-building-an-agent-with-langgraph.ipynb",
   "toc_visible": true
  },
  "kaggle": {
   "accelerator": "none",
   "dataSources": [],
   "dockerImageVersionId": 30786,
   "isGpuEnabled": false,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "paintrek",
   "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}