graphoai_Db_V3.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "from typing import Dict, Any, List, Tuple\n",
    "import spacy\n",
    "from spacy.matcher import Matcher, PhraseMatcher\n",
    "\n",
    "class AdvancedSchemaToGraphConverter:\n",
    "    def __init__(self):\n",
    "        self.graph = nx.Graph()\n",
    "        self.nlp = spacy.load(\"en_core_web_sm\")\n",
    "        self.matcher = Matcher(self.nlp.vocab)\n",
    "        self.phrase_matcher = PhraseMatcher(self.nlp.vocab)\n",
    "        \n",
    "        # Define patterns for common SQL operations\n",
    "        self.matcher.add(\"AGGREGATION\", [[{\"LOWER\": {\"IN\": [\"average\", \"avg\", \"sum\", \"count\", \"max\", \"min\"]}}]])\n",
    "        self.matcher.add(\"GROUPING\", [[{\"LOWER\": \"group\"}, {\"LOWER\": \"by\"}]])\n",
    "        self.matcher.add(\"ORDERING\", [[{\"LOWER\": {\"IN\": [\"order\", \"sort\"]}}, {\"LOWER\": \"by\"}]])\n",
    "        self.matcher.add(\"LIMIT\", [[{\"LOWER\": {\"IN\": [\"top\", \"bottom\"]}}, {\"POS\": \"NUM\"}]])\n",
    "        self.matcher.add(\"TIME_RANGE\", [[{\"LOWER\": {\"IN\": [\"last\", \"past\"]}}, {\"POS\": \"NUM\"}, {\"LOWER\": {\"IN\": [\"day\", \"week\", \"month\", \"year\"]}}]])\n",
    "        \n",
    "        # Domain-specific enrichments\n",
    "        self.matcher.add(\"CUSTOMER_RETENTION\", [[{\"LOWER\": \"customer\"}, {\"LOWER\": \"retention\"}]])\n",
    "        self.matcher.add(\"POPULAR_PRODUCTS\", [[{\"LOWER\": \"popular\"}, {\"LOWER\": \"products\"}]])\n",
    "        self.matcher.add(\"BOUGHT_TOGETHER\", [[{\"LOWER\": \"bought\"}, {\"LOWER\": \"together\"}]])\n",
    "\n",
    "    def load_schema(self, schema_json: str):\n",
    "        schema = json.loads(schema_json)\n",
    "        self._process_schema(schema)\n",
    "        self._update_phrase_matcher()\n",
    "\n",
    "    def _process_schema(self, schema: Dict[str, Any]):\n",
    "        for table_name, table_info in schema.items():\n",
    "            self._add_table_node(table_name, table_info)\n",
    "            self._process_columns(table_name, table_info['columns'])\n",
    "            self._process_relationships(table_name, table_info.get('relationships', []))\n",
    "\n",
    "    def _add_table_node(self, table_name: str, table_info: Dict[str, Any]):\n",
    "        # Adding description and enriched metadata for the table\n",
    "        self.graph.add_node(table_name, type='table', \n",
    "                            description=table_info.get('description', ''),\n",
    "                            domain_keywords=self._generate_domain_keywords(table_name),\n",
    "                            temporal_columns=self._detect_temporal_columns(table_info['columns']))\n",
    "\n",
    "    def _process_columns(self, table_name: str, columns: Dict[str, Any]):\n",
    "        for column_name, column_info in columns.items():\n",
    "            column_node = f\"{table_name}.{column_name}\"\n",
    "            # Enriching the column with data type, constraints, and statistical metadata\n",
    "            self.graph.add_node(column_node, type='column', \n",
    "                                data_type=column_info['data_type'],\n",
    "                                constraints=column_info.get('constraints', []),\n",
    "                                common_values=self._get_common_values(column_name),\n",
    "                                functional_dependency=self._detect_functional_dependency(table_name, column_name))\n",
    "            self.graph.add_edge(table_name, column_node, type='has_column')\n",
    "\n",
    "    def _process_relationships(self, table_name: str, relationships: list):\n",
    "        for relationship in relationships:\n",
    "            related_table = relationship['related_table']\n",
    "            # Adding relationship type and cardinality for richer context\n",
    "            self.graph.add_edge(table_name, related_table, type='related_to',\n",
    "                                relationship_type=relationship['type'],\n",
    "                                cardinality=self._infer_cardinality(relationship['type']))\n",
    "\n",
    "    def _update_phrase_matcher(self):\n",
    "        patterns = []\n",
    "        for node in self.graph.nodes():\n",
    "            if self.graph.nodes[node]['type'] in ['table', 'column']:\n",
    "                patterns.append(self.nlp(node.lower()))\n",
    "        self.phrase_matcher.add(\"SCHEMA_ELEMENTS\", patterns)\n",
    "\n",
    "    def advanced_question_processing(self, question: str) -> Dict[str, Any]:\n",
    "        doc = self.nlp(question.lower())\n",
    "        \n",
    "        # Extract entities and operations\n",
    "        entities = [ent.text for ent in doc.ents]\n",
    "        operations = []\n",
    "        for match_id, start, end in self.matcher(doc):\n",
    "            operations.append(doc[start:end].text)\n",
    "        \n",
    "        # Get relevant tables and columns\n",
    "        matches = self.phrase_matcher(doc)\n",
    "        relevant_elements = [doc[start:end].text for _, start, end in matches]\n",
    "        \n",
    "        tables, columns = self._classify_relevant_elements(relevant_elements)\n",
    "        \n",
    "        # Expand tables and columns based on the question context\n",
    "        expanded_tables, expanded_columns = self._expand_relevant_elements(tables, columns, doc)\n",
    "        \n",
    "        # Identify potential joins and related tables\n",
    "        potential_joins = self.find_potential_joins(expanded_tables)\n",
    "        \n",
    "        return {\n",
    "            \"entities\": entities,\n",
    "            \"operations\": operations,\n",
    "            \"tables\": expanded_tables,\n",
    "            \"columns\": expanded_columns,\n",
    "            \"potential_joins\": potential_joins\n",
    "        }\n",
    "\n",
    "    def _classify_relevant_elements(self, elements: List[str]) -> Tuple[List[str], List[str]]:\n",
    "        tables = []\n",
    "        columns = []\n",
    "        for element in elements:\n",
    "            if '.' in element:\n",
    "                columns.append(element)\n",
    "            else:\n",
    "                tables.append(element)\n",
    "        return tables, columns\n",
    "\n",
    "    def _expand_relevant_elements(self, tables: List[str], columns: List[str], doc) -> Tuple[List[str], List[str]]:\n",
    "        expanded_tables = set(tables)\n",
    "        expanded_columns = set(columns)\n",
    "        \n",
    "        # Expand based on relationships\n",
    "        for table in tables:\n",
    "            expanded_tables.update(self.get_related_tables(table))\n",
    "        \n",
    "        # Expand based on common query patterns\n",
    "        if any(token.text in ['total', 'sum', 'amount'] for token in doc):\n",
    "            expanded_columns.add('orders.total_amount')\n",
    "            expanded_tables.add('orders')\n",
    "        \n",
    "        if any(token.text in ['customer', 'user'] for token in doc):\n",
    "            expanded_tables.add('users')\n",
    "        \n",
    "        if any(token.text in ['product', 'item'] for token in doc):\n",
    "            expanded_tables.add('products')\n",
    "            expanded_tables.add('order_items')\n",
    "        \n",
    "        if any(token.text in ['category'] for token in doc):\n",
    "            expanded_columns.add('products.category')\n",
    "        \n",
    "        if any(token.text in ['bought', 'purchased', 'together'] for token in doc):\n",
    "            expanded_tables.update(['orders', 'order_items', 'products'])\n",
    "        \n",
    "        # Expand columns for all tables\n",
    "        for table in expanded_tables:\n",
    "            expanded_columns.update(self.get_table_columns(table))\n",
    "        \n",
    "        return list(expanded_tables), list(expanded_columns)\n",
    "\n",
    "    def get_related_tables(self, table: str) -> List[str]:\n",
    "        if table not in self.graph:\n",
    "            return []\n",
    "        return [node for node in nx.neighbors(self.graph, table) \n",
    "                if self.graph.nodes[node]['type'] == 'table']\n",
    "\n",
    "    def get_table_columns(self, table: str) -> List[str]:\n",
    "        if table not in self.graph:\n",
    "            return []\n",
    "        return [node for node in nx.neighbors(self.graph, table) \n",
    "                if self.graph.nodes[node]['type'] == 'column']\n",
    "\n",
    "    def find_potential_joins(self, tables: List[str]) -> List[tuple]:\n",
    "        joins = []\n",
    "        for i in range(len(tables)):\n",
    "            for j in range(i+1, len(tables)):\n",
    "                if tables[i] in self.graph and tables[j] in self.graph:\n",
    "                    try:\n",
    "                        path = nx.shortest_path(self.graph, tables[i], tables[j])\n",
    "                        joins.append((tables[i], tables[j], path))\n",
    "                    except nx.NetworkXNoPath:\n",
    "                        continue\n",
    "        return joins\n",
    "\n",
    "    # Helper methods to enrich schema with richer metadata\n",
    "    def _generate_domain_keywords(self, table_name: str) -> List[str]:\n",
    "        # Generating domain-specific keywords\n",
    "        keywords = {\n",
    "            'users': ['customer', 'client', 'user'],\n",
    "            'orders': ['purchase', 'transaction', 'sale'],\n",
    "            'products': ['item', 'goods', 'merchandise'],\n",
    "        }\n",
    "        return keywords.get(table_name, [])\n",
    "\n",
    "    def _detect_temporal_columns(self, columns: Dict[str, Any]) -> List[str]:\n",
    "        # Detect columns related to time/date\n",
    "        temporal_columns = [col for col, info in columns.items() if 'date' in col or 'time' in info['data_type']]\n",
    "        return temporal_columns\n",
    "\n",
    "    def _get_common_values(self, column_name: str) -> List[Any]:\n",
    "        # Simulate common values for demonstration\n",
    "        common_values = {\n",
    "            'users.age': [25, 30, 35, 40],\n",
    "            'products.category': ['Electronics', 'Clothing', 'Home'],\n",
    "            'orders.total_amount': [100.0, 250.5, 500.0],\n",
    "        }\n",
    "        return common_values.get(column_name, [])\n",
    "\n",
    "    def _detect_functional_dependency(self, table_name: str, column_name: str) -> str:\n",
    "        # Simulate functional dependency for example\n",
    "        if table_name == 'orders' and column_name == 'total_amount':\n",
    "            return 'total_amount = quantity * price'\n",
    "        return ''\n",
    "\n",
    "    def _infer_cardinality(self, relationship_type: str) -> str:\n",
    "        if relationship_type == 'one-to-many':\n",
    "            return '1-to-N'\n",
    "        if relationship_type == 'many-to-one':\n",
    "            return 'N-to-1'\n",
    "        return 'N-to-N'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Example usage\n",
    "converter = AdvancedSchemaToGraphConverter()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema_json = \"\"\"\n",
    "{\n",
    "  \"users\": {\n",
    "    \"description\": \"Table containing user information\",\n",
    "    \"columns\": {\n",
    "      \"user_id\": {\"data_type\": \"int\", \"constraints\": [\"primary key\"]},\n",
    "      \"name\": {\"data_type\": \"varchar\"},\n",
    "      \"age\": {\"data_type\": \"int\"},\n",
    "      \"registration_date\": {\"data_type\": \"date\"}\n",
    "    }\n",
    "  },\n",
    "  \"orders\": {\n",
    "    \"description\": \"Table containing order information\",\n",
    "    \"columns\": {\n",
    "      \"order_id\": {\"data_type\": \"int\", \"constraints\": [\"primary key\"]},\n",
    "      \"user_id\": {\"data_type\": \"int\", \"constraints\": [\"foreign key\"]},\n",
    "      \"total_amount\": {\"data_type\": \"float\"},\n",
    "      \"order_date\": {\"data_type\": \"date\"}\n",
    "    },\n",
    "    \"relationships\": [\n",
    "      {\"related_table\": \"users\", \"type\": \"many-to-one\"}\n",
    "    ]\n",
    "  },\n",
    "  \"products\": {\n",
    "    \"description\": \"Table containing product information\",\n",
    "    \"columns\": {\n",
    "      \"product_id\": {\"data_type\": \"int\", \"constraints\": [\"primary key\"]},\n",
    "      \"name\": {\"data_type\": \"varchar\"},\n",
    "      \"category\": {\"data_type\": \"varchar\"}\n",
    "    }\n",
    "  },\n",
    "  \"order_items\": {\n",
    "    \"description\": \"Table containing ordered products\",\n",
    "    \"columns\": {\n",
    "      \"order_id\": {\"data_type\": \"int\", \"constraints\": [\"foreign key\"]},\n",
    "      \"product_id\": {\"data_type\": \"int\", \"constraints\": [\"foreign key\"]},\n",
    "      \"quantity\": {\"data_type\": \"int\"},\n",
    "      \"price\": {\"data_type\": \"float\"}\n",
    "    },\n",
    "    \"relationships\": [\n",
    "      {\"related_table\": \"orders\", \"type\": \"many-to-one\"},\n",
    "      {\"related_table\": \"products\", \"type\": \"many-to-one\"}\n",
    "    ]\n",
    "  }\n",
    "}\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'entities': [], 'operations': [], 'tables': ['products', 'order_items', 'orders'], 'columns': ['orders.total_amount', 'products.category'], 'potential_joins': []}\n"
     ]
    }
   ],
   "source": [
    "question = \"What is the month-over-month growth rate in total sales for each product category?\"\n",
    "query_context = converter.advanced_question_processing(question)\n",
    "print(query_context)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}