Merge pull request #6 from DataLabTechTV/dev

feat: econ_graph graph analytics

Author: DataLabTechTV

.env.example

@@ -36,8 +36,8 @@ ANALYTICS_MART_DB=marts/analytics.sqlite
 # KùzuDB configurations
 # =====================
 
-MUSIC_TASTE_GRAPH_DB=graphs/music_taste.kz
-ECON_COMP_GRAPH_DB=graphs/econ_comp.kz
+MUSIC_TASTE_GRAPH_DB=graphs/music_taste.kuzu
+ECON_COMP_GRAPH_DB=graphs/econ_comp.kuzu
 
 # Ollama configurations
 # =====================

dlctl/cli.py

@@ -18,6 +18,7 @@
 from shared.storage import Storage, StoragePrefix
 
 LOG_FILE = Path(__file__).resolve().parents[1] / "logs/datalab.log"
+LOG_FILE_RELATIVE = os.path.relpath(LOG_FILE.resolve(), start=Path.cwd())
 
 
 @click.group(
@@ -34,7 +35,7 @@
     "logfile_enabled",
     is_flag=True,
     default=True,
-    help=f"Disable file logging ({LOG_FILE.relative_to(Path.cwd())})",
+    help=f"Disable file logging ({LOG_FILE_RELATIVE})",
 )
 @click.option(
     "--version",

graph/analytics.py

@@ -0,0 +1,71 @@
+from pathlib import Path
+
+import kuzu
+from loguru import logger as log
+
+from shared.lakehouse import Lakehouse
+from shared.settings import LOCAL_DIR, env
+
+
+class GraphAnalytics:
+    def __init__(self, schema: str):
+        dbname = env.str(f"{schema.upper()}_GRAPH_DB")
+        db_path = Path(LOCAL_DIR) / dbname
+
+        if not db_path.exists():
+            raise FileNotFoundError(f"db not found: {db_path}")
+
+        db = kuzu.Database(db_path)
+        self.conn = kuzu.Connection(db)
+
+        self.lh = Lakehouse()
+
+    def compute_con_scores(
+        self,
+        node_label: str,
+        rel_label: str,
+        column_name: str = "con_score",
+    ):
+        log.info(
+            "Computing CON scores for {} nodes via {} rels, storing to {} property",
+            node_label,
+            rel_label,
+            column_name,
+        )
+
+        log.debug("Adding {} to {}, if not exists", column_name, node_label)
+
+        self.conn.execute(
+            f"""
+            ALTER TABLE {node_label}
+            ADD IF NOT EXISTS {column_name} DOUBLE
+            """
+        )
+
+        log.debug("Resetting {} on {}", column_name, node_label)
+
+        self.conn.execute(
+            f"""
+            MATCH (c:{node_label})
+            SET c.`{column_name}` = 0.0
+            """
+        )
+
+        log.debug("Computing CON scores")
+
+        self.conn.execute(
+            f"""
+            MATCH (a:{node_label})-[ac:{rel_label}]->(c:{node_label})
+            MATCH (b:{node_label})-[bc:{rel_label}]->(c:{node_label})
+            WHERE a <> b
+            WITH a, b,
+                CASE
+                    WHEN ac.esi < bc.esi
+                    THEN ac.esi
+                    ELSE bc.esi
+                END AS min_esi
+            WITH a, b, sum(min_esi) AS con_pair
+            WITH a, sum(con_pair) AS con_score
+            SET a.`{column_name}` = con_score
+            """
+        )

graph/cli.py

@@ -4,6 +4,7 @@
 import click
 from loguru import logger as log
 
+from graph.analytics import GraphAnalytics
 from graph.embedding import NodeEmbedding, NodeEmbeddingAlgo
 from graph.ops import KuzuOps
 from graph.rag import ContextAssemblerException, GraphRAG, GraphRetrievalException
@@ -104,6 +105,15 @@ def embeddings(schema: str, dimension: int, batch_size: int, epochs: int, algo:
         log.exception(e)
 
 
+@compute.command(help="Compute common out-neighbors (CON) score")
+@click.argument("schema", type=click.STRING)
+@click.argument("node_label", type=click.STRING)
+@click.argument("rel_label", type=click.STRING)
+def con_score(schema: str, node_label: str, rel_label: str):
+    ga = GraphAnalytics(schema)
+    ga.compute_con_scores(node_label, rel_label)
+
+
 @graph.command(help="Reindex embedding property")
 @click.argument("schema", type=click.STRING)
 def reindex(schema: str):

graph/ops.py

@@ -1,7 +1,7 @@
-import os
 import shutil
 import tempfile
 from enum import Enum
+from pathlib import Path
 from string import Template
 from typing import Any, Optional
 
@@ -23,15 +23,15 @@ class KuzuTableType(Enum):
 class KuzuOps:
     def __init__(self, schema: str, overwrite: bool = False):
         dbname = env.str(f"{schema.upper()}_GRAPH_DB")
-        db_path = os.path.join(LOCAL_DIR, dbname)
+        db_path = Path(LOCAL_DIR) / dbname
 
-        if os.path.exists(db_path):
+        if db_path.exists():
             if overwrite:
                 log.warning(f"Overwriting database: {db_path}")
-                if os.path.isdir(db_path):
+                if db_path.is_dir():
                     shutil.rmtree(db_path)
-                elif os.path.isfile(db_path):
-                    os.unlink(db_path)
+                elif db_path.is_file():
+                    db_path.unlink()
 
         db = kuzu.Database(db_path)
         self.conn = kuzu.Connection(db)