clm-kernel 0.0.0__py3-none-any.whl

clm/layer3/rag/graph/store.py

@@ -0,0 +1,641 @@
+"""
+Graph Store
+
+SQLite-based knowledge graph storage with traversal capabilities.
+"""
+
+import json
+import logging
+import sqlite3
+from datetime import datetime, timezone
+from typing import Any
+
+from .extractor import Entity, ExtractionResult
+from .schema import GRAPH_SCHEMAS
+
+logger = logging.getLogger(__name__)
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Graph Store
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+class GraphStore:
+    """
+    SQLite-based storage for knowledge graph.
+
+    Features:
+    - Entity and relationship storage
+    - Graph traversal (BFS)
+    - Multi-hop pathfinding
+    - Community management
+
+    Usage:
+        store = GraphStore(db_path="graph.db")
+
+        # Add entities
+        entity_id = store.add_entity(entity, source_hash)
+
+        # Add relationships
+        store.add_relationship(source_id, target_id, "relates_to", source_hash)
+
+        # Find related entities
+        related = store.find_related("MCard", hops=2)
+    """
+
+    def __init__(self, db_path: str = ":memory:"):
+        """
+        Initialize graph store.
+
+        Args:
+            db_path: Path to SQLite database
+        """
+        self.db_path = db_path
+        self.conn = self._init_database()
+
+    def _init_database(self) -> sqlite3.Connection:
+        """Initialize SQLite connection and create tables."""
+        conn = sqlite3.connect(self.db_path, check_same_thread=False)
+        cursor = conn.cursor()
+
+        # Create all graph tables
+        for _name, schema in GRAPH_SCHEMAS.items():
+            cursor.execute(schema)
+
+        conn.commit()
+        logger.debug(f"Initialized graph store at {self.db_path}")
+        return conn
+
+    # ─────────────────────────────────────────────────────────────────────────
+    # Entity Operations
+    # ─────────────────────────────────────────────────────────────────────────
+
+    def add_entity(
+        self, entity: Entity, source_hash: str, embedding: bytes = None
+    ) -> int:
+        """
+        Add an entity to the graph.
+
+        Args:
+            entity: Entity to add
+            source_hash: Source MCard hash
+            embedding: Optional entity embedding
+
+        Returns:
+            Entity ID
+        """
+        now = datetime.now(timezone.utc).isoformat()
+        cursor = self.conn.cursor()
+
+        try:
+            cursor.execute(
+                """
+                INSERT INTO graph_entities (name, type, description, source_hash, embedding, created_at)
+                VALUES (?, ?, ?, ?, ?, ?)
+            """,
+                (
+                    entity.name,
+                    entity.type,
+                    entity.description,
+                    source_hash,
+                    embedding,
+                    now,
+                ),
+            )
+            self.conn.commit()
+            return cursor.lastrowid
+
+        except sqlite3.IntegrityError:
+            # Entity already exists, get its ID
+            cursor.execute(
+                """
+                SELECT id FROM graph_entities
+                WHERE name = ? AND type = ? AND source_hash = ?
+            """,
+                (entity.name, entity.type, source_hash),
+            )
+            row = cursor.fetchone()
+            return row[0] if row else -1
+
+    def get_entity_by_name(self, name: str) -> dict[str, Any] | None:
+        """Get entity by name (case-insensitive)."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            SELECT id, name, type, description, source_hash, created_at
+            FROM graph_entities
+            WHERE LOWER(name) = LOWER(?)
+            LIMIT 1
+        """,
+            (name,),
+        )
+
+        row = cursor.fetchone()
+        if row:
+            return {
+                "id": row[0],
+                "name": row[1],
+                "type": row[2],
+                "description": row[3],
+                "source_hash": row[4],
+                "created_at": row[5],
+            }
+        return None
+
+    def get_entity_by_id(self, entity_id: int) -> dict[str, Any] | None:
+        """Get entity by ID."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            SELECT id, name, type, description, source_hash, created_at
+            FROM graph_entities
+            WHERE id = ?
+        """,
+            (entity_id,),
+        )
+
+        row = cursor.fetchone()
+        if row:
+            return {
+                "id": row[0],
+                "name": row[1],
+                "type": row[2],
+                "description": row[3],
+                "source_hash": row[4],
+                "created_at": row[5],
+            }
+        return None
+
+    def search_entities(
+        self, query: str, type_filter: str = None, limit: int = 10
+    ) -> list[dict[str, Any]]:
+        """Search entities by name pattern."""
+        cursor = self.conn.cursor()
+
+        if type_filter:
+            cursor.execute(
+                """
+                SELECT id, name, type, description, source_hash
+                FROM graph_entities
+                WHERE name LIKE ? AND type = ?
+                LIMIT ?
+            """,
+                (f"%{query}%", type_filter, limit),
+            )
+        else:
+            cursor.execute(
+                """
+                SELECT id, name, type, description, source_hash
+                FROM graph_entities
+                WHERE name LIKE ?
+                LIMIT ?
+            """,
+                (f"%{query}%", limit),
+            )
+
+        return [
+            {
+                "id": r[0],
+                "name": r[1],
+                "type": r[2],
+                "description": r[3],
+                "source_hash": r[4],
+            }
+            for r in cursor.fetchall()
+        ]
+
+    def get_entities_by_source(self, source_hash: str) -> list[dict[str, Any]]:
+        """Get all entities from a source MCard."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            SELECT id, name, type, description
+            FROM graph_entities
+            WHERE source_hash = ?
+        """,
+            (source_hash,),
+        )
+
+        return [
+            {"id": r[0], "name": r[1], "type": r[2], "description": r[3]}
+            for r in cursor.fetchall()
+        ]
+
+    # ─────────────────────────────────────────────────────────────────────────
+    # Community Operations
+    # ─────────────────────────────────────────────────────────────────────────
+
+    def add_community(
+        self,
+        title: str,
+        summary: str,
+        member_ids: list[int],
+        level: int = 0,
+        parent_id: int | None = None,
+        embedding: bytes = None,
+    ) -> int:
+        """Add a community summary."""
+        now = datetime.now(timezone.utc).isoformat()
+        cursor = self.conn.cursor()
+
+        cursor.execute(
+            """
+            INSERT INTO graph_communities
+            (title, summary, member_entity_ids, level, parent_community_id, embedding, created_at)
+            VALUES (?, ?, ?, ?, ?, ?, ?)
+        """,
+            (title, summary, json.dumps(member_ids), level, parent_id, embedding, now),
+        )
+
+        self.conn.commit()
+        return cursor.lastrowid
+
+    def get_communities(self, level: int = 0) -> list[dict[str, Any]]:
+        """Get communities by level."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            SELECT id, title, summary, member_entity_ids, level
+            FROM graph_communities
+            WHERE level = ?
+        """,
+            (level,),
+        )
+
+        results = []
+        for row in cursor.fetchall():
+            results.append(
+                {
+                    "id": row[0],
+                    "title": row[1],
+                    "summary": row[2],
+                    "member_ids": json.loads(row[3]) if row[3] else [],
+                    "level": row[4],
+                }
+            )
+        return results
+
+    # ─────────────────────────────────────────────────────────────────────────
+    # Relationship Operations
+    # ─────────────────────────────────────────────────────────────────────────
+
+    def add_relationship(
+        self,
+        source_entity_id: int,
+        target_entity_id: int,
+        relationship: str,
+        source_hash: str,
+        description: str = "",
+        weight: float = 1.0,
+    ) -> int:
+        """
+        Add a relationship between entities.
+
+        Args:
+            source_entity_id: Source entity ID
+            target_entity_id: Target entity ID
+            relationship: Relationship type/verb
+            source_hash: Source MCard hash
+            description: Optional description
+            weight: Relationship weight
+
+        Returns:
+            Relationship ID
+        """
+        now = datetime.now(timezone.utc).isoformat()
+        cursor = self.conn.cursor()
+
+        try:
+            cursor.execute(
+                """
+                INSERT INTO graph_relationships
+                (source_entity_id, target_entity_id, relationship, description, weight, source_hash, created_at)
+                VALUES (?, ?, ?, ?, ?, ?, ?)
+            """,
+                (
+                    source_entity_id,
+                    target_entity_id,
+                    relationship,
+                    description,
+                    weight,
+                    source_hash,
+                    now,
+                ),
+            )
+            self.conn.commit()
+            return cursor.lastrowid
+
+        except sqlite3.IntegrityError:
+            # Relationship already exists
+            return -1
+
+    def get_relationships_from(self, entity_id: int) -> list[dict[str, Any]]:
+        """Get outgoing relationships from an entity."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            SELECT r.id, r.relationship, r.description, r.weight,
+                   e.id, e.name, e.type
+            FROM graph_relationships r
+            JOIN graph_entities e ON r.target_entity_id = e.id
+            WHERE r.source_entity_id = ?
+        """,
+            (entity_id,),
+        )
+
+        return [
+            {
+                "rel_id": r[0],
+                "relationship": r[1],
+                "description": r[2],
+                "weight": r[3],
+                "target_id": r[4],
+                "target_name": r[5],
+                "target_type": r[6],
+            }
+            for r in cursor.fetchall()
+        ]
+
+    def get_relationships_to(self, entity_id: int) -> list[dict[str, Any]]:
+        """Get incoming relationships to an entity."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            SELECT r.id, r.relationship, r.description, r.weight,
+                   e.id, e.name, e.type
+            FROM graph_relationships r
+            JOIN graph_entities e ON r.source_entity_id = e.id
+            WHERE r.target_entity_id = ?
+        """,
+            (entity_id,),
+        )
+
+        return [
+            {
+                "rel_id": r[0],
+                "relationship": r[1],
+                "description": r[2],
+                "weight": r[3],
+                "source_id": r[4],
+                "source_name": r[5],
+                "source_type": r[6],
+            }
+            for r in cursor.fetchall()
+        ]
+
+    # ─────────────────────────────────────────────────────────────────────────
+    # Graph Traversal
+    # ─────────────────────────────────────────────────────────────────────────
+
+    def find_related(
+        self, entity_name: str, hops: int = 2, direction: str = "both"
+    ) -> list[dict[str, Any]]:
+        """
+        Find entities within N hops of a given entity.
+
+        Args:
+            entity_name: Starting entity name
+            hops: Maximum traversal depth
+            direction: 'outgoing', 'incoming', or 'both'
+
+        Returns:
+            List of related entities with their paths
+        """
+        # Find starting entity
+        start = self.get_entity_by_name(entity_name)
+        if not start:
+            return []
+
+        visited: set[int] = {start["id"]}
+        results = []
+        frontier = [(start["id"], 0, [start["name"]])]  # (id, depth, path)
+
+        while frontier:
+            current_id, depth, path = frontier.pop(0)
+
+            if depth >= hops:
+                continue
+
+            # Get adjacent entities
+            neighbors = []
+
+            if direction in ("outgoing", "both"):
+                for rel in self.get_relationships_from(current_id):
+                    neighbors.append(
+                        (
+                            rel["target_id"],
+                            rel["target_name"],
+                            rel["relationship"],
+                            "outgoing",
+                        )
+                    )
+
+            if direction in ("incoming", "both"):
+                for rel in self.get_relationships_to(current_id):
+                    neighbors.append(
+                        (
+                            rel["source_id"],
+                            rel["source_name"],
+                            rel["relationship"],
+                            "incoming",
+                        )
+                    )
+
+            for neighbor_id, neighbor_name, rel, _dir in neighbors:
+                if neighbor_id not in visited:
+                    visited.add(neighbor_id)
+                    new_path = path + [f"--{rel}->", neighbor_name]
+
+                    entity = self.get_entity_by_id(neighbor_id)
+                    if entity:
+                        results.append(
+                            {
+                                "entity": entity,
+                                "depth": depth + 1,
+                                "path": new_path,
+                                "relationship": rel,
+                            }
+                        )
+
+                    frontier.append((neighbor_id, depth + 1, new_path))
+
+        return results
+
+    def find_path(
+        self, source_name: str, target_name: str, max_depth: int = 4
+    ) -> list[str] | None:
+        """
+        Find shortest path between two entities.
+
+        Args:
+            source_name: Starting entity name
+            target_name: Target entity name
+            max_depth: Maximum path length
+
+        Returns:
+            List of entity names in path, or None if no path
+        """
+        source = self.get_entity_by_name(source_name)
+        target = self.get_entity_by_name(target_name)
+
+        if not source or not target:
+            return None
+
+        if source["id"] == target["id"]:
+            return [source_name]
+
+        # BFS for shortest path
+        visited = {source["id"]}
+        queue = [(source["id"], [source["name"]])]
+
+        while queue:
+            current_id, path = queue.pop(0)
+
+            if len(path) > max_depth:
+                continue
+
+            for rel in self.get_relationships_from(current_id):
+                neighbor_id = rel["target_id"]
+
+                if neighbor_id == target["id"]:
+                    return path + [f"--{rel['relationship']}->", target_name]
+
+                if neighbor_id not in visited:
+                    visited.add(neighbor_id)
+                    queue.append(
+                        (
+                            neighbor_id,
+                            path + [f"--{rel['relationship']}->", rel["target_name"]],
+                        )
+                    )
+
+        return None
+
+    # ─────────────────────────────────────────────────────────────────────────
+    # Extraction Tracking
+    # ─────────────────────────────────────────────────────────────────────────
+
+    def mark_extracted(self, source_hash: str, entity_count: int, rel_count: int):
+        """Mark a source as having been extracted."""
+        now = datetime.now(timezone.utc).isoformat()
+        cursor = self.conn.cursor()
+
+        cursor.execute(
+            """
+            INSERT OR REPLACE INTO graph_extractions
+            (hash, entity_count, relationship_count, extracted_at)
+            VALUES (?, ?, ?, ?)
+        """,
+            (source_hash, entity_count, rel_count, now),
+        )
+
+        self.conn.commit()
+
+    def is_extracted(self, source_hash: str) -> bool:
+        """Check if a source has been extracted."""
+        cursor = self.conn.cursor()
+        cursor.execute("SELECT 1 FROM graph_extractions WHERE hash = ?", (source_hash,))
+        return cursor.fetchone() is not None
+
+    # ─────────────────────────────────────────────────────────────────────────
+    # Statistics
+    # ─────────────────────────────────────────────────────────────────────────
+
+    def count_entities(self) -> int:
+        """Get total entity count."""
+        cursor = self.conn.cursor()
+        cursor.execute("SELECT COUNT(*) FROM graph_entities")
+        return cursor.fetchone()[0]
+
+    def count_relationships(self) -> int:
+        """Get total relationship count."""
+        cursor = self.conn.cursor()
+        cursor.execute("SELECT COUNT(*) FROM graph_relationships")
+        return cursor.fetchone()[0]
+
+    def count_extractions(self) -> int:
+        """Get count of extracted sources."""
+        cursor = self.conn.cursor()
+        cursor.execute("SELECT COUNT(*) FROM graph_extractions")
+        return cursor.fetchone()[0]
+
+    def get_entity_types(self) -> dict[str, int]:
+        """Get counts by entity type."""
+        cursor = self.conn.cursor()
+        cursor.execute(
+            """
+            SELECT type, COUNT(*) FROM graph_entities
+            GROUP BY type
+            ORDER BY COUNT(*) DESC
+        """
+        )
+        return {row[0]: row[1] for row in cursor.fetchall()}
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get graph statistics."""
+        return {
+            "entity_count": self.count_entities(),
+            "relationship_count": self.count_relationships(),
+            "extraction_count": self.count_extractions(),
+            "entity_types": self.get_entity_types(),
+        }
+
+    def clear(self):
+        """Clear all graph data."""
+        cursor = self.conn.cursor()
+        cursor.execute("DELETE FROM graph_relationships")
+        cursor.execute("DELETE FROM graph_entities")
+        cursor.execute("DELETE FROM graph_communities")
+        cursor.execute("DELETE FROM graph_extractions")
+        self.conn.commit()
+        logger.info("Graph store cleared")
+
+
+# ─────────────────────────────────────────────────────────────────────────────
+# Batch Operations
+# ─────────────────────────────────────────────────────────────────────────────
+
+
+def store_extraction_result(
+    store: GraphStore, result: ExtractionResult, source_hash: str
+) -> tuple[int, int]:
+    """
+    Store an extraction result in the graph.
+
+    Args:
+        store: GraphStore instance
+        result: ExtractionResult from GraphExtractor
+        source_hash: Source MCard hash
+
+    Returns:
+        Tuple of (entity_count, relationship_count) stored
+    """
+    if not result.success:
+        return 0, 0
+
+    entity_map = {}  # name -> id
+
+    # Add entities
+    for entity in result.entities:
+        entity_id = store.add_entity(entity, source_hash)
+        if entity_id > 0:
+            entity_map[entity.name.lower()] = entity_id
+
+    # Add relationships
+    rel_count = 0
+    for rel in result.relationships:
+        source_id = entity_map.get(rel.source.lower())
+        target_id = entity_map.get(rel.target.lower())
+
+        if source_id and target_id:
+            rel_id = store.add_relationship(
+                source_id, target_id, rel.relationship, source_hash, rel.description
+            )
+            if rel_id > 0:
+                rel_count += 1
+
+    # Mark as extracted
+    store.mark_extracted(source_hash, len(entity_map), rel_count)
+
+    return len(entity_map), rel_count