ConversationEnding_v1.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conversation Ending\n",
    "### Goal\n",
    "\n",
    "Can I use a model to tell me when a conversation has ended?\n",
    "\n",
    "### Method\n",
    "\n",
    "Start gathering up non-threaded messages in a channel. Use a UDF to call a LLM to determine if the conversation is continuing or new.\n",
    "\n",
    "Use Kaskada to gather messages. Include up to 5 messages in the call to the LLM, independent of whether or not they are part of the conversation (as deemed by previous calls)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### The code:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Install the tools, initiate the things"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -q kaskada openai llama-cpp-python ipywidgets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai, getpass\n",
    "\n",
    "# Initialize OpenAI\n",
    "openai.api_key = getpass.getpass('OpenAI: API Key')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import kaskada as kd\n",
    "\n",
    "# Initialize Kaskada with a local execution context.\n",
    "kd.init_session()\n",
    "\n",
    "# set pandas to display all floats with 6 decimal places\n",
    "pd.options.display.float_format = '{:.6f}'.format"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pull in the user list, create a `format_user()` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "users_df = pd.read_json(\"slack-generation.users.json\")\n",
    "\n",
    "columns_to_keep = [\"id\", \"team_id\", \"name\", \"deleted\", \"real_name\", \"is_bot\", \"updated\"]\n",
    "\n",
    "users_df.drop(columns=users_df.columns.difference(columns_to_keep), inplace=True)\n",
    "\n",
    "users = {}\n",
    "for user in users_df.to_dict(orient='index').values():\n",
    "    users[user[\"id\"]] = user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_user(user_id):\n",
    "    return users[user_id] if user_id in users.keys() else None\n",
    "\n",
    "def format_user(user_id):\n",
    "    user = get_user(user_id)\n",
    "    return f\"{user['name']} ({user_id})\" if user else f\"({user_id})\"\n",
    "\n",
    "format_user(\"UBB9D2B01\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load the slack data, clean the message text, format message users"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load events from a Parquet file\n",
    "#\n",
    "# if you wan to load in your own slack data, change this to the path of your output file from 1.1 above\n",
    "# otherwise continue with `slack-generation.parquet`, which contains generated slack data for\n",
    "# example purposes. See the `slack-generation/notebook.ipynb` notebook for more info.\n",
    "input_file = \"slack-generation.parquet\"\n",
    "\n",
    "# Use the \"ts\" column as the time associated with each row,\n",
    "# and the \"channel\" column as the entity associated with each row.\n",
    "raw_msgs = await kd.sources.Parquet.create(\n",
    "    input_file,\n",
    "    time_column = \"ts\",\n",
    "    key_column = \"channel\",\n",
    "    time_unit = \"s\"\n",
    ")\n",
    "raw_msgs.preview(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@kd.udf(\"f<N: any>(x: N) -> string\")\n",
    "def format_users(batch: pd.Series):\n",
    "    # Apply to each row in the batch\n",
    "    return batch.map(format_user)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean Text\n",
    "import re\n",
    "\n",
    "def strip_code_blocks(line):\n",
    "    return re.sub(r\"```.*?```\", '', line)\n",
    "\n",
    "def user_repl(match_obj):\n",
    "    user_id = match_obj.group(1)\n",
    "    return format_user(user_id)\n",
    "\n",
    "def update_users(line):\n",
    "    return re.sub(r\"<@(.*?)>\", user_repl, line)\n",
    "\n",
    "def clean_message(text):\n",
    "        text = strip_code_blocks(update_users(text)).strip()\n",
    "        return None if text == \"\" else text\n",
    "\n",
    "@kd.udf(\"f<N: any>(x: N) -> string\")\n",
    "def clean_text(batch: pd.Series):\n",
    "    # Apply to each row in the batch\n",
    "    return batch.map(clean_message)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "formatted_msgs = raw_msgs.extend({\n",
    "    \"text\": raw_msgs.col(\"text\").pipe(clean_text),\n",
    "    \"user\": raw_msgs.col(\"user\").pipe(format_users)\n",
    "})\n",
    "formatted_msgs.preview(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Try to use a LLM to determine if a non-threaded message is part of the previous conversation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thread_ts = formatted_msgs.col(\"thread_ts\")\n",
    "\n",
    "# split messages into two subgroups: threads and non-threads\n",
    "non_threads = formatted_msgs.filter(thread_ts.is_null())\n",
    "\n",
    "# hack in a new, empty string column: `is_new`\n",
    "non_threads = non_threads.extend({\"is_new\": non_threads.col(\"text\").substring(0,0)})\n",
    "\n",
    "#filter to one channel for now\n",
    "non_threads = non_threads.filter(non_threads.col(\"channel\").eq(\"Project\"))\n",
    "\n",
    "@kd.udf(\"f<N: any>(x: N) -> string\")\n",
    "def format_message(batch: pd.Series):\n",
    "    def formatter(raw):\n",
    "        return f\"{raw['user']} --> {raw['text']}\" # --> {raw['reactions']}\"\n",
    "    return batch.map(formatter)\n",
    "\n",
    "# prefix message with user\n",
    "non_threads = non_threads.extend({\"text\": non_threads.select(\"user\", \"text\").pipe(format_message)})\n",
    "\n",
    "# collect previous messages into groups of potential conversations\n",
    "prev_messages = non_threads.lag(1)\n",
    "collect_window = kd.windows.Since(prev_messages.col(\"is_new\").eq(\"yes\"))\n",
    "\n",
    "non_threads = non_threads.extend({\"conversation\" : prev_messages.col(\"text\").collect(max=5, window=collect_window)})\n",
    "\n",
    "non_threads.preview()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Setup instructions and few-shot learning examples for the LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "system = \"\"\"\n",
    "You are a helpful assistant. You will be passed an existing conversation and a next\n",
    "line. Your job is to determine if the next line is part of the existing conversation\n",
    "or the start of a new conversation. You should respond `yes` if you think the line\n",
    "is part a new conversation or `no` otherwise. No explanation is needed.\n",
    "\n",
    "Lines of the conversation, and the next line, will be passed in plain text, where the\n",
    "user and their text is separated by an arrow like this: `-->`.\n",
    "\n",
    "The user field contains an username and an user_id in parenthesis, like this:\n",
    "`name (U1292934)` the username is lowercase and could match names in the conversation\n",
    "text in a case-insensitive way.\n",
    "\n",
    "Inside a conversation, `---` characters on their own line indicate that the next line will\n",
    "contain the text from the next user in the conversation. Conversations may contain no\n",
    "lines. When this is the case, the next line should always be a new conversation.\n",
    "\n",
    "The conversation will be prefixed by `Conversation:` on it's on line and the next line\n",
    "will be prefixed by `Next Line:` on its own line.\n",
    "\"\"\"\n",
    "\n",
    "user_empty_convo = \"\"\"\n",
    "Conversation:\n",
    "\n",
    "\n",
    "Next Line:\n",
    "userc (UFB3DA5BF) --> Risk mitigation is indeed essential, especially when we are relying\n",
    "on real-time data for our inventory tracking and resource allocation. Let's prioritize\n",
    "this aspect and design our system to be resilient.\n",
    "\"\"\"\n",
    "\n",
    "assistant_empty_convo = \"yes\"\n",
    "\n",
    "user_existing_convo = \"\"\"\n",
    "Conversation:\n",
    "userc (UFB3DA5BF) --> Risk mitigation is indeed essential, especially when we are relying\n",
    "on real-time data for our inventory tracking and resource allocation. Let's prioritize\n",
    "this aspect and design our system to be resilient.\n",
    "---\n",
    "userf (UEA27BBFF) --> Scaling our system effectively will be crucial for accommodating\n",
    "future growth. We should also keep an eye on performance metrics and fine-tune\n",
    "our resource allocation strategies as needed.\n",
    "\n",
    "\n",
    "Next Line:\n",
    "userb (UBB9D2B01) --> Sounds good, UserC. A short break will be refreshing. I'll be\n",
    "back in 15 minutes with more ideas for the next steps.\n",
    "\"\"\"\n",
    "\n",
    "assistant_existing_convo = \"no\"\n",
    "\n",
    "user_new_convo = \"\"\"\n",
    "Conversation:\n",
    "userc (UFB3DA5BF) --> Risk mitigation is indeed essential, especially when we are relying\n",
    "on real-time data for our inventory tracking and resource allocation. Let's prioritize\n",
    "this aspect and design our system to be resilient.\n",
    "---\n",
    "userb (UBB9D2B01) --> Sounds good, UserC. A short break will be refreshing. I'll be\n",
    "back in 15 minutes with more ideas for the next steps.\n",
    "\n",
    "\n",
    "Next Line:\n",
    "userb (UBB9D2B01) --> Good afternoon, everyone! The topic of multi-cloud strategies\n",
    "is intriguing. I'm excited to explore how it can help us achieve high availability\n",
    "for our inventory tracking tool.\n",
    "\"\"\"\n",
    "\n",
    "assistant_new_convo = 'yes'\n",
    "\n",
    "ai_prompt_messages = [\n",
    "        {\"role\": \"system\", \"content\": system},\n",
    "        {\"role\": \"user\", \"content\": user_empty_convo},\n",
    "        {\"role\": \"assistant\", \"content\": assistant_empty_convo},\n",
    "        {\"role\": \"user\", \"content\": user_existing_convo},\n",
    "        {\"role\": \"assistant\", \"content\": assistant_existing_convo},\n",
    "        {\"role\": \"user\", \"content\": user_new_convo},\n",
    "        {\"role\": \"assistant\", \"content\": assistant_new_convo},\n",
    "      ]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Declare a function that checks if a conversation is \"new\" by calling a LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def check_if_new_conversation(existing_conversation: [str], next_line: str) -> str:\n",
    "    user_text = \"Conversation:\\n\" + \"\\n---\\n\".join(existing_conversation) + \"\\n\\n\\nNext Line:\\n\" + next_line\n",
    "\n",
    "    prompt = ai_prompt_messages.copy()\n",
    "    prompt.append({\"role\": \"user\", \"content\": user_text})\n",
    "\n",
    "    attempts = 0\n",
    "    while True:\n",
    "      try:\n",
    "        attempts += 1\n",
    "        completion = openai.ChatCompletion.create(\n",
    "          # model choices: gpt-4, gpt-4-32k, gpt-3.5-turbo, gpt-3.5-turbo-16k\n",
    "          model=\"gpt-3.5-turbo\",\n",
    "          messages=prompt,\n",
    "          temperature=0\n",
    "        )\n",
    "        return completion.choices[0].message.content\n",
    "      except:\n",
    "        if attempts > 3:\n",
    "           return \"timeout\"\n",
    "        time.sleep(attempts * 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Call the function from a Kaskada UDF via the `pipe()` functionality, and output the results as jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@kd.udf(\"f<N: any>(x: N) -> string\")\n",
    "def determine_if_new(batch: pd.Series):\n",
    "    def is_it_new(raw):\n",
    "        return check_if_new_conversation(raw[\"conversation\"], raw[\"text\"])\n",
    "    return batch.map(is_it_new)\n",
    "\n",
    "non_threads = non_threads.extend({\"is_new\": non_threads.select(\"conversation\", \"text\").pipe(determine_if_new)})\n",
    "\n",
    "df = non_threads.to_pandas()\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.to_json(\"ConversationEnding_v1_results.jsonl\", orient=\"records\", lines=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Also test if we can use text in a string column to determine the breakpoint of a conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# can we use the text in a column to determine the breakpoint of a conversation? YES\n",
    "\n",
    "test = await kd.sources.JsonlFile.create(\n",
    "    \"ConversationEnding_v1_input.jsonl\",\n",
    "    time_column = \"ts\",\n",
    "    key_column = \"channel\",\n",
    "    time_unit = \"s\"\n",
    ")\n",
    "test.preview(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import timedelta\n",
    "\n",
    "ts = test.col(\"ts\")\n",
    "new_convo = test.col(\"new_convo\")\n",
    "\n",
    "is_new = new_convo.eq(\"yes\")\n",
    "\n",
    "# Eventually this will just be: `thread_ts = ts.first(window=kd.windows.Since(is_new, start=\"inclusive\"))`\n",
    "#\n",
    "# However, `Since()` is currently exclusive on the start of the window, inclusive on the end.\n",
    "# But we need inclusive on the start and exclusive on the end.\n",
    "#\n",
    "# The hack below does what we need until `Since()` provides additional options for inclusivity\n",
    "shifted_non_threads = test.shift_by(timedelta(microseconds=0.001))\n",
    "shifted_ts = test.lag(1).col(\"ts\").first(window=kd.windows.Since(is_new))\n",
    "thread_ts = ts.if_(is_new).else_(shifted_ts)\n",
    "\n",
    "# create threads_ts column for non-threaded messages\n",
    "non_threads_threads = test.extend({\"thread_ts\": thread_ts}).filter(ts.is_not_null().and_(thread_ts.is_not_null()))\n",
    "non_threads_threads.filter(non_threads_threads.col(\"channel\").eq(\"General\")).preview(20)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Yes, we can use text in a string column to determine the breakpoint of a conversation."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}