notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "afd1e805-38c4-4f90-8d8a-53689dc86bbd",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6881e5e1-39f1-4f06-b650-adaf3937564b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b946a030",
   "metadata": {},
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'retry'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[4], line 21\u001b[0m\n\u001b[1;32m     18\u001b[0m \u001b[38;5;66;03m# Third party imports\u001b[39;00m\n\u001b[1;32m     20\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01msentence_transformers\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m SentenceTransformer, util\n\u001b[0;32m---> 21\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mretry\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m retry\n\u001b[1;32m     23\u001b[0m \u001b[38;5;66;03m# Greyparrot imports\u001b[39;00m\n\u001b[1;32m     24\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mgreyparrot\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mllm\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mprompting\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m prompts \u001b[38;5;28;01mas\u001b[39;00m prompts_utils\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'retry'"
     ]
    }
   ],
   "source": [
    "# Standard library imports\n",
    "from pydantic import BaseModel\n",
    "import threading\n",
    "import warnings\n",
    "import random\n",
    "import string\n",
    "import queue\n",
    "import json\n",
    "import csv\n",
    "import os\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "JWT_SECRET_API = !echo $(aws --profile \"chat-prod_ro\" secretsmanager get-secret-value --secret-id \"arn:aws:secretsmanager:us-west-2:242659714806:secret:shared/cresta-server-jwt_secret-VDn5My\" --query SecretString --output text) # type: ignore\n",
    "os.environ[\"JWT_SECRET_API\"] = json.loads(JWT_SECRET_API[0])[\"jwt-secret\"]\n",
    "os.environ[\"CONFIG_SERVICE_ADDR\"] = \"auth.chat-prod.internal.cresta.ai:443\"\n",
    "os.environ[\"CONFIG_USE_SECURE_CHANNEL\"] = \"true\"\n",
    "\n",
    "# Third party imports\n",
    "\n",
    "from sentence_transformers import SentenceTransformer, util\n",
    "from retry import retry\n",
    "\n",
    "# Greyparrot imports\n",
    "from greyparrot.llm.prompting import prompts as prompts_utils\n",
    "from greyparrot.conversations.db import ConversationsDBConn\n",
    "from greyparrot.multi_tenancy.v3_config import V3Config\n",
    "from greyparrot.chats_common import PartialChat\n",
    "from greyparrot.dataset_common import Dataset\n",
    "from greyparrot.common import get_logger\n",
    "\n",
    "logger = get_logger(__name__)\n",
    "\n",
    "# Local imports\n",
    "from llm_proxy import get_open_ai_client, LLMProxyProperties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7e4bd5d0-d369-4e36-92d6-a8d30198ecca",
   "metadata": {},
   "outputs": [],
   "source": [
    "embedder = SentenceTransformer(\"all-mpnet-base-v2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "be1a39f6-f947-40e6-8cc3-c129fe8610cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "customer_id = \"brinks\"\n",
    "profile_id = \"care-voice\"\n",
    "usecase_id = \"care-voice\"\n",
    "language_code = \"en-US\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0a5b2751-25c6-404f-9454-3747754271ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_chats_with_ids(chat_ids: list[str]):\n",
    "    customer_name = V3Config.short_name_from_ids(customer_id, profile_id)\n",
    "    conv_db_conn = ConversationsDBConn.from_customer_name(customer_name)\n",
    "    chats = conv_db_conn.get_detailed_chats(\n",
    "        customer_id=customer_id,\n",
    "        profile_id=profile_id,\n",
    "        usecase_id=usecase_id,\n",
    "        language_code=language_code,\n",
    "        conversation_ids=chat_ids,\n",
    "        is_dev_user=False,\n",
    "    )\n",
    "    return chats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "afcb1094-7197-466c-8a03-d3725ccaf380",
   "metadata": {},
   "outputs": [],
   "source": [
    "LLM_ENGINE = \"gpt-4-0125-preview\"\n",
    "CONCURRENCY = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "43492383-6f20-492d-9258-db38bb67267a",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_proxy_properties = LLMProxyProperties(\n",
    "    project_id=\"virtual-agent-simulation\",\n",
    "    customer_id=customer_id,\n",
    "    profile_id=profile_id,\n",
    "    usecase_id=\"\",\n",
    ")\n",
    "open_ai_client = get_open_ai_client(\n",
    "    llm_proxy_properties=llm_proxy_properties,\n",
    "    provider=\"openai\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "620967f2-bf35-4305-9889-f08021619e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "@retry(tries=3, delay=60, backoff=2, logger=logger)\n",
    "def chat_completion(**kwargs):\n",
    "    return open_ai_client.beta.chat.completions.parse(**kwargs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2d1cb11f-7c7c-447f-a925-0766a9d88134",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_jsonl_file(data, output_file):\n",
    "    with open(output_file, \"w\") as f:\n",
    "        for item in data:\n",
    "            f.write(json.dumps(item) + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6b8f602c-878c-478b-865a-0358d5ef7906",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_workflow(w: dict):    \n",
    "    for w_k, w_v in w.items():        \n",
    "        if isinstance(w_v, list):\n",
    "            val = \"\\n* \".join(w_v)\n",
    "            print(f\"{w_k.capitalize()}:\\n* {val}\")\n",
    "        else:\n",
    "            print(f\"{w_k.capitalize()}: {w_v}\")\n",
    "    print(\"-\" * 50)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc5c40f7-e0e3-4eb8-818e-6a43d2d1dd02",
   "metadata": {},
   "source": [
    "# Discovery prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e404cd18-a964-45c5-8977-64053d53bfc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO remove this after fixing speaker_role flips\n",
    "flips = {\"agent\": \"visitor\", \"visitor\": \"agent\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "d52de3a4-12a9-46a4-a930-5f14eb66db2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def chat_to_prompt_text(chat: PartialChat, speakers_flipped: bool = False):\n",
    "    return \"\\n\".join([\n",
    "        f\"{string.capwords(prompts_utils.speaker_role_str_for_prompts(flips[m.speaker_role.value] if speakers_flipped else m.speaker_role.value))}: {m.text}\"\n",
    "        for m in chat.messages\n",
    "    ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "099b0d7c-a431-4e1e-803a-2f69d8a1df70",
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM_PROMPT_AGENT_WORKFLOW_DISCOVERY = \"\"\"### Context and data description\n",
    "You are a conversation analyst working for a Call Center.\n",
    "\n",
    "You will be given 1 conversation at a time. Each conversation is between a Call Center Agent and a Customer. Your primary goal is to extract workflows of steps which the Agent takes in **the given conversation** to help resolve the Customer's needs related to a procut issue.\n",
    "\n",
    "The primary use case of these workflows is to create a troubleshooting template to address similar customer needs in the future.\n",
    "\n",
    "Each workflow should be a list of steps which the Agent needs to take.\n",
    "\n",
    "For each workflow, return the product, issue, and a list of steps which the Agent needs to take.\n",
    "\n",
    "Make sure the product is specific and not general.\n",
    "Make sure the issue is specific and not general.\n",
    "Make sure the steps are detailed.\n",
    "\n",
    "**Important**: There could be more than 1 workflow in a single conversation. There could also be no workflows in a single conversation. The workflows will be used to create troubleshooting guides to address similar customer needs in the future.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7522dc7d-8ba0-40ce-b7be-512664145158",
   "metadata": {},
   "outputs": [],
   "source": [
    "SYSTEM_PROMPT_VISITOR_WORKFLOW_DISCOVERY = \"\"\"### Context and data description\n",
    "You are a conversation analyst working for a Call Center.\n",
    "\n",
    "You will be given 1 conversation at a time. Each conversation is between a Call Center Agent and a Customer. Your primary goal is to extract flows of steps which Customer takes in **the given conversation**.\n",
    "\n",
    "The primary use case of this flow is to create a template to simulate similar customer scenarios.\n",
    "\n",
    "The flow should be a list of steps which the Customer takes.\n",
    "\n",
    "The schema of each flow should be as follows:\n",
    "- **title**: title of the flow\n",
    "- **steps**: a list of steps which Customer needs to follow\n",
    "\n",
    "**Important**: There could be more than 1 flow in a single conversation. The flows will be used create templates to simulate similar customer scenarios.\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "2f80ef46-1308-42f7-9ad7-4a5ee415c1ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "FLOW_PROMPTS = {\n",
    "    \"agent\": SYSTEM_PROMPT_AGENT_WORKFLOW_DISCOVERY,\n",
    "    \"visitor\": SYSTEM_PROMPT_VISITOR_WORKFLOW_DISCOVERY\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b37c0eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Flow(BaseModel):\n",
    "    product: str\n",
    "    issue: str\n",
    "    steps: list[str]\n",
    "\n",
    "class Flows(BaseModel):\n",
    "    flows: list[Flow]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "64e4a2f3-1e34-43ee-94e2-65015660af2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def discover_flow_in_chat(chat: PartialChat,\n",
    "                          speaker_role: str,\n",
    "                          speakers_flipped: bool = False,\n",
    "                          llm_engine: str = LLM_ENGINE):\n",
    "    logger.info(f\"Discovering {speaker_role} flow in chat {chat.chat_name}\")\n",
    "    messages = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": FLOW_PROMPTS[speaker_role]\n",
    "        },\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": chat_to_prompt_text(chat, speakers_flipped)\n",
    "        },\n",
    "    ]\n",
    "\n",
    "    # Note: temperature=0.1 to allow for some exploration\n",
    "    response = chat_completion(model=llm_engine,\n",
    "                               messages=messages,\n",
    "                               temperature=0.1,\n",
    "                               response_format=Flows)\n",
    "    workflow = response.choices[0].message.parsed\n",
    "\n",
    "    return workflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "034f33d5-e6fe-49f4-918b-6a88819af747",
   "metadata": {},
   "outputs": [],
   "source": [
    "def extract_flows_from_chats(chats: list[PartialChat], speaker_role: str,\n",
    "                                concurrency: int = 10):\n",
    "    lock = threading.Lock()\n",
    "    indexes = queue.Queue()\n",
    "\n",
    "    workflows = {}\n",
    "    for idx in range(len(chats)):\n",
    "        indexes.put(idx)\n",
    "\n",
    "    def workflow_labeler_worker():\n",
    "        while True:\n",
    "            try:\n",
    "                idx = indexes.get(block=False)\n",
    "            except queue.Empty:\n",
    "                return\n",
    "            chat = chats[idx]\n",
    "            try:\n",
    "                extracted_workflows = discover_flow_in_chat(chat, speaker_role)\n",
    "                with lock:\n",
    "                    workflows[str(chat)] = extracted_workflows\n",
    "                    if len(workflows) % 10 == 0:\n",
    "                        print(f\"Workflows from {len(workflows)} chats extracted!\")\n",
    "            except Exception as e:\n",
    "                logger.warning(e, str(chat))\n",
    "            indexes.task_done()\n",
    "\n",
    "    logger.info(\n",
    "        f\"Starting processing {len(chats)} chats with {concurrency} workers\")\n",
    "    workers = [\n",
    "        threading.Thread(target=workflow_labeler_worker)\n",
    "        for _ in range(concurrency)\n",
    "    ]\n",
    "    for worker in workers:\n",
    "        worker.start()\n",
    "    for worker in workers:\n",
    "        worker.join()\n",
    "    logger.info(f\"Finished processing all {len(chats)} chats\")\n",
    "\n",
    "    return workflows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aefb919-3d2e-4cea-a6f2-f593a3e251ff",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "test_chat = get_chats_with_ids([\"0843c54c-6487-45ce-946a-cc6257484f54\"])[0]\n",
    "str(test_chat), test_chat.messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56a6064d-1c44-4811-9245-d5e83d838178",
   "metadata": {},
   "outputs": [],
   "source": [
    "workflows = discover_flow_in_chat(test_chat, \"agent\", speakers_flipped=True)\n",
    "for idx, w in enumerate(workflows):\n",
    "    print(\"-\" * 50)\n",
    "    print(f\"[Agent Workflow#{idx + 1}]\")\n",
    "    for k, v in w.items():\n",
    "        if isinstance(v, list):\n",
    "            val = \"\\n* \".join(v)\n",
    "            print(f\"{k.capitalize()}:\\n* {val}\")\n",
    "        else:\n",
    "            print(f\"{k.capitalize()}: {v}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c84379c-64fc-465a-978d-310a51919ac6",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(workflows)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5593f36c-c4c0-4fac-9e09-b64e0f8d65d2",
   "metadata": {},
   "source": [
    "# Relevant chats (from KA-QE trainset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e663f09-a7df-4377-a5a7-12c80666817f",
   "metadata": {},
   "outputs": [],
   "source": [
    "qe_dataset = Dataset.pull_from_repo(\n",
    "    \"brinks-care-voice/hf:082720241521675544.train\")\n",
    "len(qe_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "086dad2f-dcb5-4bab-a4d1-2ef87db5ef4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "qe_dataset[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb47768b-bd30-4000-8dc2-73dc170bd110",
   "metadata": {},
   "outputs": [],
   "source": [
    "chat_ids = list(set([pc.chat_name for pc, _ in qe_dataset]))\n",
    "len(chat_ids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "add32431-3497-4257-9ef6-d676c5111ff8",
   "metadata": {},
   "outputs": [],
   "source": [
    "queries_to_chat = {query: pc.chat_name for pc, query in qe_dataset}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e6fd7c3-9046-4f1b-9740-914226e99a52",
   "metadata": {},
   "outputs": [],
   "source": [
    "queries = list(queries_to_chat.keys())\n",
    "len(queries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "998b3aa4-4022-4ecf-b5b5-92c88c4b390f",
   "metadata": {},
   "outputs": [],
   "source": [
    "source_embeddings = embedder.encode(queries,\n",
    "                                    convert_to_tensor=True,\n",
    "                                    show_progress_bar=False)\n",
    "len(source_embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd12c873-3521-4d59-862b-d5273663c7b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"Brinks Evaluation v.20240806_ - Response Evaluation (internal).csv\") as f:\n",
    "    reader = csv.DictReader(f)\n",
    "    data = list(reader)\n",
    "    customer_queries = [entry[\"Question\"] for entry in data]\n",
    "len(customer_queries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "def021c3-8d6e-4b10-b5f5-4373e28165a3",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "matched_queries = []\n",
    "for q in customer_queries:\n",
    "    target_embeddings = embedder.encode(q,\n",
    "                                        convert_to_tensor=True,\n",
    "                                        show_progress_bar=False)\n",
    "    scores = util.cos_sim(target_embeddings,\n",
    "                          source_embeddings).cpu().tolist()[0]\n",
    "    matches = [(idx, score) for idx, score in enumerate(scores) if score > 0.7]\n",
    "    if matches:\n",
    "        matched_queries.extend([queries[i] for i, s in matches])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32559261-df19-43bd-9ffd-b4d251757485",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(matched_queries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0bc5aefd-615b-44b3-97db-b7f926d245ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "matched_chats = [queries_to_chat[q] for q in matched_queries]\n",
    "len(matched_chats)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffe05dfb-a510-4065-ad25-c7ca82c29bce",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "num_chats = 1000\n",
    "sampled_chat_ids = list(set(random.sample(chat_ids, num_chats) + matched_chats))\n",
    "print(len(sampled_chat_ids))\n",
    "sampled_chats = get_chats_with_ids(sampled_chat_ids)\n",
    "len(sampled_chats)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a04fb3d-0c9e-4a97-b074-c967ef49f05b",
   "metadata": {},
   "source": [
    "# Agent Workflow Discovery"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bc0f6fa-e10a-41cf-bffc-04e6138493bf",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "workflows = extract_flows_from_chats(sampled_chats, \"agent\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85294de9-348e-483f-afa0-7be035c0c4ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(workflows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da7374d8-8423-41f9-9437-beed54b2ab50",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "for pc, ws in list(workflows.items())[:5]:\n",
    "    print(f\"\\n\\n<{pc}>\")\n",
    "    for idx, w in enumerate(ws):\n",
    "        print(\"-\" * 50)\n",
    "        print(f\"[Agent Workflow#{idx + 1}]\")\n",
    "        for k, v in w.items():\n",
    "            if isinstance(v, list):\n",
    "                val = \"\\n* \".join(v)\n",
    "                print(f\"{k.capitalize()}:\\n* {val}\")\n",
    "            else:\n",
    "                print(f\"{k.capitalize()}: {v}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "468aeb75",
   "metadata": {},
   "outputs": [],
   "source": [
    "## start here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "bbed0a88-bb89-472d-9858-63897b9c7f2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "title_to_workflow, title_to_pc_id = {}, {}\n",
    "for pc_id, ws in workflows.items():\n",
    "    for idx, w in enumerate(ws):\n",
    "        title_to_workflow[w[\"title\"]] = w\n",
    "        title_to_pc_id[w[\"title\"]] = pc_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6d1f8e5-b425-421b-9694-4e144792901e",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(title_to_workflow), len(title_to_pc_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "6999ea39-e3f3-472f-8e90-9ff8063bb491",
   "metadata": {},
   "outputs": [],
   "source": [
    "def group_semantically_similar_workflows(workflows: list[str], embedder: SentenceTransformer,\n",
    "                                           semantic_threshold: int = 0.8):\n",
    "    workflow_embeddings = embedder.encode(workflows, convert_to_tensor=True)\n",
    "    workflow_groups, solo_workflows = [], []\n",
    "    for idx, workflow in enumerate(workflows):\n",
    "        existing_groups = [group for group in workflow_groups if idx in group]\n",
    "        if existing_groups:\n",
    "            assert len(\n",
    "                existing_groups) == 1, \"A workflow should only be in 1 group\"\n",
    "            continue\n",
    "\n",
    "        print(f\"Finding similar Workflows for Workflow#{idx}\")\n",
    "\n",
    "        scores = util.cos_sim(workflow_embeddings[idx:idx + 1],\n",
    "                              workflow_embeddings).cpu().tolist()[0]\n",
    "        matches = [\n",
    "            i for i, score in enumerate(scores) if score > semantic_threshold\n",
    "        ]\n",
    "        filtered_matches = [i for i in matches if i != idx]\n",
    "        if filtered_matches:\n",
    "            outstanding_groups = []\n",
    "            new_group = filtered_matches + [idx]\n",
    "            for group in workflow_groups:\n",
    "                if set(filtered_matches).intersection(set(group)):\n",
    "                    new_group.extend(group)\n",
    "                else:\n",
    "                    outstanding_groups.append(group)\n",
    "            workflow_groups = outstanding_groups + [list(set(new_group))]\n",
    "        else:\n",
    "            solo_workflows.append(idx)\n",
    "\n",
    "    print(f\"Found {len(solo_workflows)} # of Solo Workflows..\")\n",
    "    print(f\"Found {len(workflow_groups)} # Groups of Workflows..\")\n",
    "\n",
    "    return [workflows[idx] for idx in solo_workflows\n",
    "           ], [[workflows[idx] for idx in group] for group in workflow_groups]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "811eb663-b988-4de3-90c7-3db995d55780",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "solo_workflows, grouped_workflows = group_semantically_similar_workflows(\n",
    "    list(title_to_workflow.keys()), embedder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dce35b42-acdc-4e70-ab1e-a17fae8c8ad5",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "for each in grouped_workflows:\n",
    "    print(\"-\" * 50)\n",
    "    print(len(each))\n",
    "    for e in each:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a86c07bf-691a-4a02-80cd-508fbb1f9514",
   "metadata": {},
   "outputs": [],
   "source": [
    "unique_workflow_titles = solo_workflows + [random.choice(w) for w in grouped_workflows]\n",
    "unique_workflows = [title_to_workflow[w] for w in unique_workflow_titles]\n",
    "len(unique_workflows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7de0bb98-0575-4611-aed7-608e3376bdc0",
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_unique_workflows = [w for w in unique_workflows if len(w[\"steps\"]) > 3] # skip too small workflows\n",
    "len(filtered_unique_workflows)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "900474f3-89c9-4d3e-b0d8-8059e16f7288",
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_unique_workflows[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20532109-1ab2-41ad-904f-f7b61d24b669",
   "metadata": {},
   "outputs": [],
   "source": [
    "for w in filtered_unique_workflows[:10]:\n",
    "    print(w[\"title\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "315b9b83-f1b9-4585-a2aa-6ea8a705d08f",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "for w in filtered_unique_workflows:\n",
    "    print(\"-\" * 100)\n",
    "    print(\"\\n\\n\")\n",
    "    for k, v in w.items():\n",
    "        if isinstance(v, list):\n",
    "            val = \"\\n* \".join(v)\n",
    "            print(f\"{k.capitalize()}:\\n* {val}\")\n",
    "        else:\n",
    "            print(f\"{k.capitalize()}: {v}\")    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a04ebdb-26a9-4453-9cc2-c3f41c2cb794",
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "[w[\"title\"] for w in filtered_unique_workflows]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8406c7fd-2906-4e2f-afb3-afe0760f1810",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(filtered_unique_workflows)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.13.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}