genesis_v0.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import OpenAI\n",
    "import json\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv()\n",
    "\n",
    "\n",
    "OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')\n",
    "client = OpenAI(api_key=OPENAI_API_KEY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel\n",
    "\n",
    "class Item(BaseModel):\n",
    "    id: str # the id of the item\n",
    "    class Config:\n",
    "        extra = \"allow\"\n",
    "\n",
    "class Category(BaseModel):\n",
    "    name: str # snake_case name of the category\n",
    "    description: str # description of the category, what it represents\n",
    "\n",
    "class ClassifiedItem(BaseModel):\n",
    "    item_id: str # the id of the item\n",
    "    category_name: str # the name of the category\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Proof of Concept"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "raw_items = [\"apple\", \"banana\", \"carrot\", \"dog\", \"broccoli\"]\n",
    "items = [Item(id=item) for item in raw_items]\n",
    "categories = [\n",
    "  Category(name=\"fruit\", description=\"fruits\"),\n",
    "  Category(name=\"vegetable\", description=\"vegetables\"),\n",
    "  Category(name=\"animal\", description=\"animals\"),\n",
    "  Category(name=\"mineral\", description=\"minerals\")\n",
    "]\n",
    "\n",
    "prompt = f\"\"\"I will provide you with items and categories. You need to classify the items into the correct category.\n",
    "ITEMS:\n",
    "```\n",
    "{[i.dict() for i in items]}\n",
    "```\n",
    "CATEGORIES:\n",
    "```\n",
    "{[c.dict() for c in categories]}\n",
    "```\"\"\" \n",
    "res = client.chat.completions.create(\n",
    "\tmodel=\"gpt-3.5-turbo\",\n",
    "\tmessages=[\n",
    "\t\t{\"role\": \"user\", \"content\": prompt}\n",
    "\t],\n",
    "\ttools=[\n",
    "\t\t{\n",
    "\t\t\t\"type\": \"function\",\n",
    "\t\t\t\"function\": {\n",
    "\t\t\t\t\t\"name\": \"classifier\",\n",
    "          \"description\": \"Classify items into categories\",\n",
    "\t\t\t\t\t\"strict\": True,\n",
    "\t\t\t\t\t\"parameters\": {\n",
    "\t\t\t\t\t\t\t\"type\": \"object\",\n",
    "              \"properties\": {\n",
    "\t\t\t\t\t\t\t\t\"classified_items\": {\n",
    "\t\t\t\t\t\t\t\t\t\"type\": \"array\",\n",
    "\t\t\t\t\t\t\t\t\t\"items\": {\n",
    "\t\t\t\t\t\t\t\t\t\t\"type\": \"object\",\n",
    "\t\t\t\t\t\t\t\t\t\t\"properties\": {\n",
    "\t\t\t\t\t\t\t\t\t\t\t\"item_id\": {\n",
    "\t\t\t\t\t\t\t\t\t\t\t\t\"type\": \"string\",\n",
    "\t\t\t\t\t\t\t\t\t\t\t\t\"description\": \"The id of the item\"\n",
    "\t\t\t\t\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t\t\t\t\t\"category_name\": {\n",
    "\t\t\t\t\t\t\t\t\t\t\t\t\"type\": \"string\",\n",
    "\t\t\t\t\t\t\t\t\t\t\t\t\"description\": \"The name of the category\",\n",
    "\t\t\t\t\t\t\t\t\t\t\t\t\"enum\": [c.name for c in categories]\n",
    "\t\t\t\t\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t\t\t\t\"required\": [\"item_id\", \"category_name\"],\n",
    "                    \"additionalProperties\": False\n",
    "\t\t\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t\"required\": [\"classified_items\"],\n",
    "\t\t\t\t\t\t\t\"additionalProperties\": False,\n",
    "\t\t\t\t\t},\n",
    "\t\t\t},\n",
    "\t\t}  \n",
    "\t],\n",
    "\ttool_choice={\"type\": \"function\", \"function\": {\"name\": \"classifier\"}},\n",
    ")\n",
    "json_res = json.loads(res.choices[0].message.tool_calls[0].function.arguments)\n",
    "print(json.dumps(json_res, indent=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_categories = 7\n",
    "categorization_method = \"categorize in a way so that there is minumum potential overlap between categories\"\n",
    "raw_items = [\"apple\", \"banana\", \"carrot\", \"dog\", \"broccoli\"]\n",
    "items = [Item(id=item) for item in raw_items]\n",
    "prompt = f\"\"\"I will provide you with items. You need to create categories according to the following guideline\n",
    "```\n",
    "There must be at most {max_categories} categories.\n",
    "{categorization_method}\n",
    "```\n",
    "ITEMS:\n",
    "```\n",
    "{[i.dict() for i in items]}\n",
    "```\"\"\" \n",
    "res = client.chat.completions.create(\n",
    "\tmodel=\"gpt-3.5-turbo\",\n",
    "\tmessages=[\n",
    "\t\t{\"role\": \"user\", \"content\": prompt}\n",
    "\t],\n",
    "\ttools=[\n",
    "\t\t{\n",
    "\t\t\t\"type\": \"function\",\n",
    "\t\t\t\"function\": {\n",
    "\t\t\t\t\t\"name\": \"category_creator\",\n",
    "          \"description\": \"Create categories for items\",\n",
    "\t\t\t\t\t\"strict\": True,\n",
    "\t\t\t\t\t\"parameters\": {\n",
    "\t\t\t\t\t\t\t\"type\": \"object\",\n",
    "              \"properties\": {\n",
    "\t\t\t\t\t\t\t\t\"categories\": {\n",
    "\t\t\t\t\t\t\t\t\t\"type\": \"array\",\n",
    "\t\t\t\t\t\t\t\t\t\"items\": {\n",
    "\t\t\t\t\t\t\t\t\t\t\"type\": \"object\",\n",
    "\t\t\t\t\t\t\t\t\t\t\"properties\": {\n",
    "\t\t\t\t\t\t\t\t\t\t\t\"name\": {\n",
    "\t\t\t\t\t\t\t\t\t\t\t\t\"type\": \"string\",\n",
    "\t\t\t\t\t\t\t\t\t\t\t\t\"description\": \"Snake case name of the category\"\n",
    "\t\t\t\t\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t\t\t\t\t\"description\": {\n",
    "\t\t\t\t\t\t\t\t\t\t\t\t\"type\": \"string\",\n",
    "\t\t\t\t\t\t\t\t\t\t\t\t\"description\": \"The description of the category\"\n",
    "\t\t\t\t\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t\t\t\t\"required\": [\"name\", \"description\"],\n",
    "                    \"additionalProperties\": False\n",
    "\t\t\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t},\n",
    "\t\t\t\t\t\t\t\"required\": [\"categories\"],\n",
    "\t\t\t\t\t\t\t\"additionalProperties\": False,\n",
    "\t\t\t\t\t},\n",
    "\t\t\t},\n",
    "\t\t}  \n",
    "\t],\n",
    "\ttool_choice={\"type\": \"function\", \"function\": {\"name\": \"category_creator\"}},\n",
    ")\n",
    "json_res = json.loads(res.choices[0].message.tool_calls[0].function.arguments)\n",
    "print(json.dumps(json_res, indent=2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res.choices[0].message"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# System"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Optional\n",
    "\n",
    "class TreeNode:\n",
    "    def __init__(self, category: Category):\n",
    "        self.category = category\n",
    "        self.items: List[Item] = []\n",
    "        self.children: List[TreeNode] = []\n",
    "        self.parent: Optional[TreeNode] = None\n",
    "\n",
    "    def add_child(self, child: 'TreeNode'):\n",
    "        child.parent = self\n",
    "        self.children.append(child)\n",
    "\n",
    "    def add_item(self, item: Item):\n",
    "        self.items.append(item)\n",
    "\n",
    "    def remove_child(self, child: 'TreeNode'):\n",
    "        if child in self.children:\n",
    "            self.children.remove(child)\n",
    "            child.parent = None\n",
    "\n",
    "    def find_subcategory(self, category_name: str) -> Optional['TreeNode']:\n",
    "        if self.category.name == category_name:\n",
    "            return self\n",
    "        for child in self.children:\n",
    "            result = child.find_subcategory(category_name)\n",
    "            if result:\n",
    "                return result\n",
    "        return None\n",
    "\n",
    "    def __repr__(self):\n",
    "        return f\"TreeNode(category={self.category.name}, items_count={len(self.items)}, children_count={len(self.children)})\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TreeNode(category=phones, items_count=1, children_count=0)\n"
     ]
    }
   ],
   "source": [
    "# Create root category\n",
    "root_category = Category(name=\"root\", description=\"Root category\")\n",
    "root_node = TreeNode(root_category)\n",
    "\n",
    "# Add subcategories\n",
    "electronics = TreeNode(Category(name=\"electronics\", description=\"Electronic devices\"))\n",
    "root_node.add_child(electronics)\n",
    "\n",
    "phones = TreeNode(Category(name=\"phones\", description=\"Mobile phones\"))\n",
    "electronics.add_child(phones)\n",
    "\n",
    "# Add items\n",
    "phone_item = Item(id=\"1\", brand=\"Apple\", model=\"iPhone 12\")\n",
    "phones.add_item(phone_item)\n",
    "\n",
    "# Find a category\n",
    "found_category = root_node.find_subcategory(\"phones\")\n",
    "print(found_category)  # Should print the 'phones' TreeNode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import sleep\n",
    "\n",
    "class CreateCategoriesChat:\n",
    "    def __init__(self, client: OpenAI, tree_node: TreeNode, max_categories: int, categorization_method: str):\n",
    "        self.tree_node = tree_node\n",
    "        self.max_categories = max_categories\n",
    "        self.categorization_method = categorization_method\n",
    "        self.chat_history = []\n",
    "        self.client = client\n",
    "        self.initialize_chat()\n",
    "\n",
    "    def initialize_chat(self):\n",
    "        items = self.tree_node.items\n",
    "        prompt = f\"\"\"I will provide you with items. You need to create categories according to the following guideline\n",
    "```\n",
    "There must be at most {self.max_categories} categories.\n",
    "{self.categorization_method}\n",
    "```\n",
    "ITEMS:\n",
    "```\n",
    "{[item.dict() for item in items]}\n",
    "```\"\"\"\n",
    "        self.chat_history = [{\"role\": \"user\", \"content\": prompt}]\n",
    "\n",
    "    def generate_categories(self) -> List[Category]:\n",
    "        if self.chat_history[-1][\"role\"] != \"user\":\n",
    "            raise ValueError(\"The latest message must be a user message before generating categories.\")\n",
    "\n",
    "        response = self.client.chat.completions.create(\n",
    "            model=\"gpt-4o-mini\",\n",
    "            messages=self.chat_history,\n",
    "            tools=[\n",
    "                {\n",
    "                    \"type\": \"function\",\n",
    "                    \"function\": {\n",
    "                        \"name\": \"category_creator\",\n",
    "                        \"description\": \"Create categories for items\",\n",
    "                        \"parameters\": {\n",
    "                            \"type\": \"object\",\n",
    "                            \"properties\": {\n",
    "                                \"categories\": {\n",
    "                                    \"type\": \"array\",\n",
    "                                    \"items\": {\n",
    "                                        \"type\": \"object\",\n",
    "                                        \"properties\": {\n",
    "                                            \"name\": {\n",
    "                                                \"type\": \"string\",\n",
    "                                                \"description\": \"Snake case name of the category\"\n",
    "                                            },\n",
    "                                            \"description\": {\n",
    "                                                \"type\": \"string\",\n",
    "                                                \"description\": \"The description of the category\"\n",
    "                                            },\n",
    "                                        },\n",
    "                                        \"required\": [\"name\", \"description\"],\n",
    "                                        \"additionalProperties\": False\n",
    "                                    },\n",
    "                                },\n",
    "                            },\n",
    "                            \"required\": [\"categories\"],\n",
    "                            \"additionalProperties\": False,\n",
    "                        },\n",
    "                    },\n",
    "                }\n",
    "            ],\n",
    "            tool_choice={\"type\": \"function\", \"function\": {\"name\": \"category_creator\"}},\n",
    "        )\n",
    "\n",
    "        # Add assistant's response to chat history\n",
    "        self.chat_history.append(response.choices[0].message.model_dump())\n",
    "\n",
    "        # Add placeholder tool response\n",
    "        tool_call = response.choices[0].message.tool_calls[0]\n",
    "        self.chat_history.append({\n",
    "            \"role\": \"tool\",\n",
    "            \"content\": json.dumps({\"status\": \"ok\"}),\n",
    "            \"tool_call_id\": tool_call.id\n",
    "        })\n",
    "\n",
    "        # Parse and return categories\n",
    "        categories_data = json.loads(tool_call.function.arguments)\n",
    "        return [Category(**cat) for cat in categories_data[\"categories\"]]\n",
    "\n",
    "    def submit_feedback(self, feedback: str):\n",
    "        self.chat_history.append({\"role\": \"user\", \"content\": feedback})\n",
    "\n",
    "    def get_latest_categories(self) -> List[Category]:\n",
    "        for message in reversed(self.chat_history):\n",
    "            if message[\"role\"] == \"assistant\" and message.get(\"tool_calls\"):\n",
    "                tool_call = message[\"tool_calls\"][0]\n",
    "                categories_data = json.loads(tool_call[\"function\"][\"arguments\"])\n",
    "                return [Category(**cat) for cat in categories_data[\"categories\"]]\n",
    "        return []\n",
    "\n",
    "    def run_category_generation_flow(self, max_iterations: int = 3):\n",
    "        for _ in range(max_iterations):\n",
    "            categories = self.generate_categories()\n",
    "            print(f\"Generated categories for \\\"{self.tree_node.category.name}\\\":\", categories)\n",
    "            sleep(0.2)  # Optional delay between iterations\n",
    "            \n",
    "            feedback = input(\"Provide feedback (or type 'done' to finish): \")\n",
    "            if feedback.lower() == 'done':\n",
    "                break\n",
    "            \n",
    "            self.submit_feedback(feedback)\n",
    "\n",
    "        return self.get_latest_categories()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated categories for \"root\": [Category(name='fruits', description='Edible fruits including apples and bananas.'), Category(name='vegetables', description='Edible vegetables including carrots and broccoli.'), Category(name='pets', description='Domestic animals that can be kept as pets, such as dogs.')]\n",
      "Final categories: [Category(name='fruits', description='Edible fruits including apples and bananas.'), Category(name='vegetables', description='Edible vegetables including carrots and broccoli.'), Category(name='pets', description='Domestic animals that can be kept as pets, such as dogs.')]\n"
     ]
    }
   ],
   "source": [
    "raw_items = [\"apple\", \"banana\", \"carrot\", \"dog\", \"broccoli\"]\n",
    "items = [Item(id=item) for item in raw_items]\n",
    "tree_node = TreeNode(Category(name=\"root\", description=\"Root category\"))\n",
    "for item in items:\n",
    "    tree_node.add_item(item)\n",
    "  \n",
    "chat = CreateCategoriesChat(client, tree_node, max_categories=5, categorization_method=\"Minimize overlap between categories\")\n",
    "final_categories = chat.run_category_generation_flow()\n",
    "print(\"Final categories:\", final_categories)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Classifier:\n",
    "    def __init__(self, client: OpenAI, tree_node: TreeNode):\n",
    "        self.client = client\n",
    "        self.tree_node = tree_node\n",
    "\n",
    "    def classify_items(self) -> List[ClassifiedItem]:\n",
    "        categories = [child.category for child in self.tree_node.children]\n",
    "        items = self.tree_node.items\n",
    "\n",
    "        prompt = f\"\"\"I will provide you with items and categories. You need to classify the items into the correct category.\n",
    "ITEMS:\n",
    "```\n",
    "{[item.dict() for item in items]}\n",
    "```\n",
    "CATEGORIES:\n",
    "```\n",
    "{[category.dict() for category in categories]}\n",
    "```\"\"\"\n",
    "\n",
    "        response = self.client.chat.completions.create(\n",
    "            model=\"gpt-4o-mini\",\n",
    "            messages=[{\"role\": \"user\", \"content\": prompt}],\n",
    "            tools=[\n",
    "                {\n",
    "                    \"type\": \"function\",\n",
    "                    \"function\": {\n",
    "                        \"name\": \"classifier\",\n",
    "                        \"description\": \"Classify items into categories\",\n",
    "                        \"parameters\": {\n",
    "                            \"type\": \"object\",\n",
    "                            \"properties\": {\n",
    "                                \"classified_items\": {\n",
    "                                    \"type\": \"array\",\n",
    "                                    \"items\": {\n",
    "                                        \"type\": \"object\",\n",
    "                                        \"properties\": {\n",
    "                                            \"item_id\": {\n",
    "                                                \"type\": \"string\",\n",
    "                                                \"description\": \"The id of the item\"\n",
    "                                            },\n",
    "                                            \"category_name\": {\n",
    "                                                \"type\": \"string\",\n",
    "                                                \"description\": \"The name of the category\",\n",
    "                                                \"enum\": [category.name for category in categories]\n",
    "                                            },\n",
    "                                        },\n",
    "                                        \"required\": [\"item_id\", \"category_name\"],\n",
    "                                        \"additionalProperties\": False\n",
    "                                    },\n",
    "                                },\n",
    "                            },\n",
    "                            \"required\": [\"classified_items\"],\n",
    "                            \"additionalProperties\": False,\n",
    "                        },\n",
    "                    },\n",
    "                }\n",
    "            ],\n",
    "            tool_choice={\"type\": \"function\", \"function\": {\"name\": \"classifier\"}},\n",
    "        )\n",
    "\n",
    "        classified_items_data = json.loads(response.choices[0].message.tool_calls[0].function.arguments)\n",
    "        return [ClassifiedItem(**item) for item in classified_items_data[\"classified_items\"]]\n",
    "\n",
    "    def update_tree(self, classified_items: List[ClassifiedItem]):\n",
    "        for classified_item in classified_items:\n",
    "            item = next((item for item in self.tree_node.items if item.id == classified_item.item_id), None)\n",
    "            if item:\n",
    "                category_node = self.tree_node.find_subcategory(classified_item.category_name)\n",
    "                if category_node:\n",
    "                    category_node.add_item(item)\n",
    "                    # self.tree_node.items.remove(item) # Remove the item from the root node\n",
    "\n",
    "    def run_classification(self):\n",
    "        classified_items = self.classify_items()\n",
    "        self.update_tree(classified_items)\n",
    "        return classified_items"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_tree(node, level=0):\n",
    "    print(\"  \" * level + f\"{node.category.name}: {[item.id for item in node.items]}\")\n",
    "    for child in node.children:\n",
    "        print_tree(child, level + 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generated categories for \"root\": [Category(name='fruits', description='Sweet and fleshy products of a tree or other plant that contain seeds, often used as food.'), Category(name='vegetables', description='Edible plants or parts of plants, usually savory and used in meals.'), Category(name='animals', description='Living organisms that are usually characterized by voluntary movement, including mammals, birds, and reptiles.')]\n",
      "Final categories: [Category(name='fruits', description='Sweet and fleshy products of a tree or other plant that contain seeds, often used as food.'), Category(name='vegetables', description='Edible plants or parts of plants, usually savory and used in meals.'), Category(name='animals', description='Living organisms that are usually characterized by voluntary movement, including mammals, birds, and reptiles.')]\n",
      "\n",
      "Classified Items:\n",
      "Item: apple, Category: fruits\n",
      "Item: banana, Category: fruits\n",
      "Item: carrot, Category: vegetables\n",
      "Item: dog, Category: animals\n",
      "Item: broccoli, Category: vegetables\n",
      "\n",
      "Updated Tree Structure:\n",
      "root: ['apple', 'banana', 'carrot', 'dog', 'broccoli']\n",
      "  fruits: ['apple', 'banana']\n",
      "  vegetables: ['carrot', 'broccoli']\n",
      "  animals: ['dog']\n"
     ]
    }
   ],
   "source": [
    "raw_items = [\"apple\", \"banana\", \"carrot\", \"dog\", \"broccoli\"]\n",
    "items = [Item(id=item) for item in raw_items]\n",
    "tree_node = TreeNode(Category(name=\"root\", description=\"Root category\"))\n",
    "for item in items:\n",
    "    tree_node.add_item(item)\n",
    "\n",
    "chat = CreateCategoriesChat(client, tree_node, max_categories=5, categorization_method=\"Minimize overlap between categories\")\n",
    "final_categories = chat.run_category_generation_flow()\n",
    "print(\"Final categories:\", final_categories)\n",
    "\n",
    "# Add the generated categories as children to the root node\n",
    "for category in final_categories:\n",
    "    tree_node.add_child(TreeNode(category))\n",
    "\n",
    "# Run the classification\n",
    "classifier = Classifier(client, tree_node)\n",
    "classified_items = classifier.run_classification()\n",
    "\n",
    "print(\"\\nClassified Items:\")\n",
    "for item in classified_items:\n",
    "    print(f\"Item: {item.item_id}, Category: {item.category_name}\")\n",
    "\n",
    "print(\"\\nUpdated Tree Structure:\")\n",
    "\n",
    "print_tree(tree_node)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[DocumentField(id='user_id', data_type='int'),\n",
       " DocumentField(id='username', data_type='varchar'),\n",
       " DocumentField(id='email', data_type='varchar'),\n",
       " DocumentField(id='password_hash', data_type='varchar'),\n",
       " DocumentField(id='first_name', data_type='varchar'),\n",
       " DocumentField(id='last_name', data_type='varchar'),\n",
       " DocumentField(id='date_of_birth', data_type='date'),\n",
       " DocumentField(id='gender', data_type='char'),\n",
       " DocumentField(id='profile_picture_url', data_type='varchar'),\n",
       " DocumentField(id='bio', data_type='text'),\n",
       " DocumentField(id='location', data_type='varchar'),\n",
       " DocumentField(id='join_date', data_type='datetime'),\n",
       " DocumentField(id='last_login', data_type='datetime'),\n",
       " DocumentField(id='account_status', data_type='varchar'),\n",
       " DocumentField(id='preferred_language', data_type='varchar'),\n",
       " DocumentField(id='time_zone', data_type='varchar'),\n",
       " DocumentField(id='newsletter_subscription', data_type='bit'),\n",
       " DocumentField(id='two_factor_auth_enabled', data_type='bit'),\n",
       " DocumentField(id='privacy_settings', data_type='varchar'),\n",
       " DocumentField(id='total_posts', data_type='int'),\n",
       " DocumentField(id='total_comments', data_type='int'),\n",
       " DocumentField(id='reputation_score', data_type='int'),\n",
       " DocumentField(id='last_password_change', data_type='datetime'),\n",
       " DocumentField(id='account_type', data_type='varchar'),\n",
       " DocumentField(id='premium_member', data_type='bit'),\n",
       " DocumentField(id='referral_code', data_type='varchar'),\n",
       " DocumentField(id='social_media_links', data_type='varchar'),\n",
       " DocumentField(id='profile_views', data_type='int'),\n",
       " DocumentField(id='last_profile_update', data_type='datetime'),\n",
       " DocumentField(id='favorite_categories', data_type='varchar'),\n",
       " DocumentField(id='notification_preferences', data_type='varchar'),\n",
       " DocumentField(id='account_deletion_date', data_type='datetime'),\n",
       " DocumentField(id='ip_address', data_type='varchar'),\n",
       " DocumentField(id='browser_info', data_type='varchar'),\n",
       " DocumentField(id='device_type', data_type='varchar'),\n",
       " DocumentField(id='failed_login_attempts', data_type='tinyint'),\n",
       " DocumentField(id='recovery_email', data_type='varchar'),\n",
       " DocumentField(id='custom_theme', data_type='varchar'),\n",
       " DocumentField(id='badge_count', data_type='smallint'),\n",
       " DocumentField(id='last_search_query', data_type='varchar')]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import re\n",
    "from typing import List, Dict, Any\n",
    "from pydantic import BaseModel\n",
    "\n",
    "# Assuming we have the previously defined Item, Category, TreeNode, and Classifier classes\n",
    "\n",
    "class DocumentField(BaseModel):\n",
    "    id: str\n",
    "    data_type: str\n",
    "\n",
    "def parse_document(file_path: str) -> List[DocumentField]:\n",
    "    \"\"\"\n",
    "    Parse the document from txt format into a list of DocumentField objects.\n",
    "    \"\"\"\n",
    "    fields = []\n",
    "    with open(file_path, 'r') as file:\n",
    "        for line in file:\n",
    "            if line.startswith('#'):\n",
    "                continue\n",
    "            parts = line.strip().split('\\t')\n",
    "            if len(parts) >= 3:\n",
    "                field_name = parts[1].strip('[]').lower()\n",
    "                data_type = parts[2].strip('[]')\n",
    "                fields.append(DocumentField(id=field_name, data_type=data_type))\n",
    "    return fields\n",
    "\n",
    "parse_document(\"consumer_dict.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_classification_structure() -> TreeNode:\n",
    "    \"\"\"\n",
    "    Create the root class and its first level of classes.\n",
    "    \"\"\"\n",
    "    root = TreeNode(Category(name=\"rootCategory\", description=\"Root category for all data fields\"))\n",
    "    \n",
    "    categories = [\n",
    "        (\"personalInformation\", \"Includes core personal details, contact information, and demographic data.\"),\n",
    "        (\"geographicData\", \"Encompasses location-related information, including address details and census data.\"),\n",
    "        (\"financialInformation\", \"Covers all finance-related data, including credit, investments, and property information.\"),\n",
    "        (\"behavioralData\", \"Includes purchasing behavior, interests, hobbies, and lifestyle indicators.\"),\n",
    "        (\"vehicleAndInsuranceInformation\", \"Covers all data related to vehicle ownership and insurance.\"),\n",
    "        (\"socioCulturalAttributes\", \"Encompasses ethnicity, language, political affiliation, and donation behavior.\")\n",
    "    ]\n",
    "    \n",
    "    for name, description in categories:\n",
    "        root.add_child(TreeNode(Category(name=name, description=description)))\n",
    "    \n",
    "    return root\n",
    "\n",
    "def classify_fields(client: OpenAI, root: TreeNode, fields: List[DocumentField]) -> List[ClassifiedItem]:\n",
    "    \"\"\"\n",
    "    Classify the fields into the first level categories.\n",
    "    \"\"\"\n",
    "    items = [Item(id=field.id, data_type=field.data_type) for field in fields]\n",
    "    for item in items:\n",
    "        root.add_item(item)\n",
    "    \n",
    "    classifier = Classifier(client, root)\n",
    "    return classifier.run_classification()\n",
    "\n",
    "def generate_subcategories(client: OpenAI, root: TreeNode) -> None:\n",
    "    \"\"\"\n",
    "    Generate one level of subcategories under each of the main categories.\n",
    "    \"\"\"\n",
    "    for child in root.children:\n",
    "        chat = CreateCategoriesChat(client, child, max_categories=5, categorization_method=f\"Create subcategories for {child.category.name}\")\n",
    "        subcategories = chat.run_category_generation_flow()\n",
    "        for subcat in subcategories:\n",
    "            child.add_child(TreeNode(subcat))\n",
    "\n",
    "def classify_into_subcategories(client: OpenAI, root: TreeNode) -> List[ClassifiedItem]:\n",
    "    \"\"\"\n",
    "    Classify the items into the subcategories.\n",
    "    \"\"\"\n",
    "    all_classified_items = []\n",
    "    for main_category in root.children:\n",
    "        classifier = Classifier(client, main_category)\n",
    "        classified_items = classifier.run_classification()\n",
    "        all_classified_items.extend(classified_items)\n",
    "    return all_classified_items\n",
    "\n",
    "def main(client: OpenAI, file_path: str):\n",
    "    # Parse the document\n",
    "    fields = parse_document(file_path)\n",
    "\n",
    "    print(\"Parsed Fields:\", len(fields))\n",
    "    \n",
    "    # Create the classification structure\n",
    "    root = create_classification_structure()\n",
    "\n",
    "    print(\"\\nInitial Tree Structure:\")\n",
    "    print_tree(root)\n",
    "    \n",
    "    # Classify fields into main categories\n",
    "    classified_items = classify_fields(client, root, fields)\n",
    "    print(\"First-level classification:\")\n",
    "    for item in classified_items:\n",
    "        print(f\"Item: {item.item_id}, Category: {item.category_name}\")\n",
    "    \n",
    "    # Generate subcategories\n",
    "    generate_subcategories(client, root)\n",
    "    \n",
    "    # Classify into subcategories\n",
    "    subclassified_items = classify_into_subcategories(client, root)\n",
    "    print(\"\\nSecond-level classification:\")\n",
    "    for item in subclassified_items:\n",
    "        print(f\"Item: {item.item_id}, Subcategory: {item.category_name}\")\n",
    "    \n",
    "    # Print the final tree structure\n",
    "    print(\"\\nFinal Tree Structure:\")    \n",
    "    print_tree(root)\n",
    "    return root\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Parsed Fields: 40\n",
      "\n",
      "Initial Tree Structure:\n",
      "rootCategory: []\n",
      "  personalInformation: []\n",
      "  geographicData: []\n",
      "  financialInformation: []\n",
      "  behavioralData: []\n",
      "  vehicleAndInsuranceInformation: []\n",
      "  socioCulturalAttributes: []\n",
      "First-level classification:\n",
      "Item: user_id, Category: personalInformation\n",
      "Item: username, Category: personalInformation\n",
      "Item: email, Category: personalInformation\n",
      "Item: password_hash, Category: personalInformation\n",
      "Item: first_name, Category: personalInformation\n",
      "Item: last_name, Category: personalInformation\n",
      "Item: date_of_birth, Category: personalInformation\n",
      "Item: gender, Category: personalInformation\n",
      "Item: profile_picture_url, Category: personalInformation\n",
      "Item: bio, Category: personalInformation\n",
      "Item: location, Category: geographicData\n",
      "Item: join_date, Category: personalInformation\n",
      "Item: last_login, Category: personalInformation\n",
      "Item: account_status, Category: personalInformation\n",
      "Item: preferred_language, Category: socioCulturalAttributes\n",
      "Item: time_zone, Category: personalInformation\n",
      "Item: newsletter_subscription, Category: personalInformation\n",
      "Item: two_factor_auth_enabled, Category: personalInformation\n",
      "Item: privacy_settings, Category: personalInformation\n",
      "Item: total_posts, Category: behavioralData\n",
      "Item: total_comments, Category: behavioralData\n",
      "Item: reputation_score, Category: behavioralData\n",
      "Item: last_password_change, Category: personalInformation\n",
      "Item: account_type, Category: personalInformation\n",
      "Item: premium_member, Category: personalInformation\n",
      "Item: referral_code, Category: personalInformation\n",
      "Item: social_media_links, Category: personalInformation\n",
      "Item: profile_views, Category: behavioralData\n",
      "Item: last_profile_update, Category: personalInformation\n",
      "Item: favorite_categories, Category: behavioralData\n",
      "Item: notification_preferences, Category: personalInformation\n",
      "Item: account_deletion_date, Category: personalInformation\n",
      "Item: ip_address, Category: geographicData\n",
      "Item: browser_info, Category: personalInformation\n",
      "Item: device_type, Category: personalInformation\n",
      "Item: failed_login_attempts, Category: personalInformation\n",
      "Item: recovery_email, Category: personalInformation\n",
      "Item: custom_theme, Category: personalInformation\n",
      "Item: badge_count, Category: behavioralData\n",
      "Item: last_search_query, Category: behavioralData\n",
      "Generated categories for \"personalInformation\": [Category(name='personal_information', description=\"Information relating to the user's identity and profile.\"), Category(name='account_settings', description=\"Settings related to user's account management.\"), Category(name='security', description='Details about the security measures of the account.'), Category(name='preferences', description='User preferences regarding notifications and subscriptions.'), Category(name='activity_log', description='Logs and records of user activities within the platform.')]\n",
      "Generated categories for \"geographicData\": [Category(name='geographic_data', description='Data related to geographical locations and coordinates.'), Category(name='network_data', description='Information related to network connections and configurations.')]\n",
      "Generated categories for \"financialInformation\": [Category(name='financial_information', description='Details about financial metrics, reports, and statuses.'), Category(name='investment_items', description='Assets or items that can be invested in, such as stocks, bonds, and real estate.'), Category(name='personal_items', description='Items for personal use or belongings.'), Category(name='services', description='Various services offered or available for purchase.'), Category(name='promotional_materials', description='Materials used for marketing, advertising, and promotions.')]\n",
      "Generated categories for \"behavioralData\": [Category(name='user_statistics', description='Overall statistics related to user interactions such as posts and comments.'), Category(name='engagement_metrics', description='Metrics that indicate user engagement with the platform, such as profile views and badge counts.'), Category(name='behavioral_data', description='Data reflecting user behavior, including searches and preferences.'), Category(name='user_reputation', description='Metrics assessing user reputation within the platform.'), Category(name='preferences', description=\"Categories indicating user's favorite categories and personal interests.\")]\n",
      "Generated categories for \"vehicleAndInsuranceInformation\": [Category(name='vehicle_and_insurance_information', description='Details regarding vehicles and associated insurance.'), Category(name='personal_information', description='Personal data of individuals.'), Category(name='financial_information', description='Financial records and data.'), Category(name='health_information', description='Health records and medical data.'), Category(name='legal_documents', description='Documents related to legal matters and agreements.')]\n",
      "Generated categories for \"socioCulturalAttributes\": [Category(name='socio_cultural_attributes', description='Attributes that relate to the social and cultural background of individuals.'), Category(name='languages', description='Different languages that may be preferred by individuals.')]\n",
      "\n",
      "Second-level classification:\n",
      "Item: user_id, Subcategory: personal_information\n",
      "Item: username, Subcategory: personal_information\n",
      "Item: email, Subcategory: personal_information\n",
      "Item: password_hash, Subcategory: security\n",
      "Item: first_name, Subcategory: personal_information\n",
      "Item: last_name, Subcategory: personal_information\n",
      "Item: date_of_birth, Subcategory: personal_information\n",
      "Item: gender, Subcategory: personal_information\n",
      "Item: profile_picture_url, Subcategory: personal_information\n",
      "Item: bio, Subcategory: personal_information\n",
      "Item: join_date, Subcategory: activity_log\n",
      "Item: last_login, Subcategory: activity_log\n",
      "Item: account_status, Subcategory: account_settings\n",
      "Item: time_zone, Subcategory: account_settings\n",
      "Item: newsletter_subscription, Subcategory: preferences\n",
      "Item: two_factor_auth_enabled, Subcategory: security\n",
      "Item: privacy_settings, Subcategory: account_settings\n",
      "Item: last_password_change, Subcategory: security\n",
      "Item: account_type, Subcategory: account_settings\n",
      "Item: premium_member, Subcategory: account_settings\n",
      "Item: referral_code, Subcategory: account_settings\n",
      "Item: social_media_links, Subcategory: personal_information\n",
      "Item: last_profile_update, Subcategory: activity_log\n",
      "Item: notification_preferences, Subcategory: preferences\n",
      "Item: account_deletion_date, Subcategory: activity_log\n",
      "Item: browser_info, Subcategory: activity_log\n",
      "Item: device_type, Subcategory: activity_log\n",
      "Item: failed_login_attempts, Subcategory: security\n",
      "Item: recovery_email, Subcategory: security\n",
      "Item: custom_theme, Subcategory: preferences\n",
      "Item: location, Subcategory: geographic_data\n",
      "Item: ip_address, Subcategory: network_data\n",
      "Item: total_posts, Subcategory: user_statistics\n",
      "Item: total_comments, Subcategory: user_statistics\n",
      "Item: reputation_score, Subcategory: user_reputation\n",
      "Item: profile_views, Subcategory: engagement_metrics\n",
      "Item: favorite_categories, Subcategory: preferences\n",
      "Item: badge_count, Subcategory: engagement_metrics\n",
      "Item: last_search_query, Subcategory: behavioral_data\n",
      "Item: preferred_language, Subcategory: languages\n",
      "\n",
      "Final Tree Structure:\n",
      "rootCategory: ['user_id', 'username', 'email', 'password_hash', 'first_name', 'last_name', 'date_of_birth', 'gender', 'profile_picture_url', 'bio', 'location', 'join_date', 'last_login', 'account_status', 'preferred_language', 'time_zone', 'newsletter_subscription', 'two_factor_auth_enabled', 'privacy_settings', 'total_posts', 'total_comments', 'reputation_score', 'last_password_change', 'account_type', 'premium_member', 'referral_code', 'social_media_links', 'profile_views', 'last_profile_update', 'favorite_categories', 'notification_preferences', 'account_deletion_date', 'ip_address', 'browser_info', 'device_type', 'failed_login_attempts', 'recovery_email', 'custom_theme', 'badge_count', 'last_search_query']\n",
      "  personalInformation: ['user_id', 'username', 'email', 'password_hash', 'first_name', 'last_name', 'date_of_birth', 'gender', 'profile_picture_url', 'bio', 'join_date', 'last_login', 'account_status', 'time_zone', 'newsletter_subscription', 'two_factor_auth_enabled', 'privacy_settings', 'last_password_change', 'account_type', 'premium_member', 'referral_code', 'social_media_links', 'last_profile_update', 'notification_preferences', 'account_deletion_date', 'browser_info', 'device_type', 'failed_login_attempts', 'recovery_email', 'custom_theme']\n",
      "    personal_information: ['user_id', 'username', 'email', 'first_name', 'last_name', 'date_of_birth', 'gender', 'profile_picture_url', 'bio', 'social_media_links']\n",
      "    account_settings: ['account_status', 'time_zone', 'privacy_settings', 'account_type', 'premium_member', 'referral_code']\n",
      "    security: ['password_hash', 'two_factor_auth_enabled', 'last_password_change', 'failed_login_attempts', 'recovery_email']\n",
      "    preferences: ['newsletter_subscription', 'notification_preferences', 'custom_theme']\n",
      "    activity_log: ['join_date', 'last_login', 'last_profile_update', 'account_deletion_date', 'browser_info', 'device_type']\n",
      "  geographicData: ['location', 'ip_address']\n",
      "    geographic_data: ['location']\n",
      "    network_data: ['ip_address']\n",
      "  financialInformation: []\n",
      "    financial_information: []\n",
      "    investment_items: []\n",
      "    personal_items: []\n",
      "    services: []\n",
      "    promotional_materials: []\n",
      "  behavioralData: ['total_posts', 'total_comments', 'reputation_score', 'profile_views', 'favorite_categories', 'badge_count', 'last_search_query']\n",
      "    user_statistics: ['total_posts', 'total_comments']\n",
      "    engagement_metrics: ['profile_views', 'badge_count']\n",
      "    behavioral_data: ['last_search_query']\n",
      "    user_reputation: ['reputation_score']\n",
      "    preferences: ['favorite_categories']\n",
      "  vehicleAndInsuranceInformation: []\n",
      "    vehicle_and_insurance_information: []\n",
      "    personal_information: []\n",
      "    financial_information: []\n",
      "    health_information: []\n",
      "    legal_documents: []\n",
      "  socioCulturalAttributes: ['preferred_language']\n",
      "    socio_cultural_attributes: []\n",
      "    languages: ['preferred_language']\n"
     ]
    }
   ],
   "source": [
    "og_root = main(client, \"consumer_dict.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TreeNode(category=rootCategory, items_count=40, children_count=6)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "og_root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Item(id='user_id', data_type='int'),\n",
       " Item(id='username', data_type='varchar'),\n",
       " Item(id='email', data_type='varchar'),\n",
       " Item(id='password_hash', data_type='varchar'),\n",
       " Item(id='first_name', data_type='varchar'),\n",
       " Item(id='last_name', data_type='varchar'),\n",
       " Item(id='date_of_birth', data_type='date'),\n",
       " Item(id='gender', data_type='char'),\n",
       " Item(id='profile_picture_url', data_type='varchar'),\n",
       " Item(id='bio', data_type='text'),\n",
       " Item(id='join_date', data_type='datetime'),\n",
       " Item(id='last_login', data_type='datetime'),\n",
       " Item(id='account_status', data_type='varchar'),\n",
       " Item(id='time_zone', data_type='varchar'),\n",
       " Item(id='newsletter_subscription', data_type='bit'),\n",
       " Item(id='two_factor_auth_enabled', data_type='bit'),\n",
       " Item(id='privacy_settings', data_type='varchar'),\n",
       " Item(id='last_password_change', data_type='datetime'),\n",
       " Item(id='account_type', data_type='varchar'),\n",
       " Item(id='premium_member', data_type='bit'),\n",
       " Item(id='referral_code', data_type='varchar'),\n",
       " Item(id='social_media_links', data_type='varchar'),\n",
       " Item(id='last_profile_update', data_type='datetime'),\n",
       " Item(id='notification_preferences', data_type='varchar'),\n",
       " Item(id='account_deletion_date', data_type='datetime'),\n",
       " Item(id='browser_info', data_type='varchar'),\n",
       " Item(id='device_type', data_type='varchar'),\n",
       " Item(id='failed_login_attempts', data_type='tinyint'),\n",
       " Item(id='recovery_email', data_type='varchar'),\n",
       " Item(id='custom_theme', data_type='varchar')]"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pi_node = og_root.find_subcategory(\"personalInformation\")\n",
    "pi_node.items"
   ]
  }
 ],
 "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.9.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}