src/knowledge/repository.py

import sqlite3
from typing import List, Optional, Tuple, Dict, Any

class KnowledgeRepository:
    def __init__(self, db_path: str = "data/knowledge/build/knowledge.db"):
        self.db_path = db_path
        self._conn: Optional[sqlite3.Connection] = None

    def connect(self):
        self._conn = sqlite3.connect(self.db_path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._enable_foreign_keys()

    def close(self):
        if self._conn:
            self._conn.close()
            self._conn = None

    def _enable_foreign_keys(self):
        cursor = self._conn.cursor()
        cursor.execute("PRAGMA foreign_keys = ON;")
        cursor.execute("PRAGMA synchronous = NORMAL;")
        cursor.execute("PRAGMA journal_mode = WAL;")

    def init_schema(self):
        if not self._conn:
            self.connect()
        
        schema = """
        CREATE TABLE IF NOT EXISTS franchises (
            franchise_id TEXT PRIMARY KEY,
            canonical_name TEXT NOT NULL
        );

        CREATE TABLE IF NOT EXISTS entities (
            entity_id TEXT PRIMARY KEY,
            canonical_name TEXT NOT NULL,
            parent_entity_id TEXT NULL,
            is_default_variant BOOLEAN DEFAULT FALSE,
            franchise_id TEXT,
            FOREIGN KEY(parent_entity_id) REFERENCES entities(entity_id),
            FOREIGN KEY(franchise_id) REFERENCES franchises(franchise_id)
        );

        CREATE TABLE IF NOT EXISTS trait_categories (
            category_id TEXT PRIMARY KEY,
            mutually_exclusive BOOLEAN DEFAULT FALSE
        );

        CREATE TABLE IF NOT EXISTS canonical_traits (
            trait_id TEXT PRIMARY KEY,
            category_id TEXT,
            trait_type TEXT CHECK(trait_type IN ('visual', 'semantic', 'style', 'signature_item')),
            FOREIGN KEY(category_id) REFERENCES trait_categories(category_id)
        );

        CREATE TABLE IF NOT EXISTS trait_hierarchy (
            parent_trait_id TEXT,
            child_trait_id TEXT,
            PRIMARY KEY (parent_trait_id, child_trait_id),
            FOREIGN KEY(parent_trait_id) REFERENCES canonical_traits(trait_id),
            FOREIGN KEY(child_trait_id) REFERENCES canonical_traits(trait_id)
        );

        CREATE TABLE IF NOT EXISTS trait_conflicts (
            trait_a TEXT,
            trait_b TEXT,
            PRIMARY KEY (trait_a, trait_b),
            FOREIGN KEY(trait_a) REFERENCES canonical_traits(trait_id),
            FOREIGN KEY(trait_b) REFERENCES canonical_traits(trait_id)
        );

        CREATE TABLE IF NOT EXISTS entity_visual_traits (
            entity_id TEXT,
            trait_id TEXT,
            confidence REAL,
            frequency INTEGER,
            source_count INTEGER,
            PRIMARY KEY(entity_id, trait_id),
            FOREIGN KEY(entity_id) REFERENCES entities(entity_id),
            FOREIGN KEY(trait_id) REFERENCES canonical_traits(trait_id)
        );

        CREATE TABLE IF NOT EXISTS entity_relationships (
            entity_id_a TEXT,
            entity_id_b TEXT,
            relationship_type TEXT,
            weight REAL,
            PRIMARY KEY(entity_id_a, entity_id_b, relationship_type),
            FOREIGN KEY(entity_id_a) REFERENCES entities(entity_id),
            FOREIGN KEY(entity_id_b) REFERENCES entities(entity_id)
        );

        CREATE TABLE IF NOT EXISTS trait_provenance (
            provenance_id TEXT PRIMARY KEY,
            entity_id TEXT,
            trait_id TEXT,
            source_uri TEXT,
            extraction_timestamp DATETIME,
            extractor_version TEXT,
            raw_payload_hash TEXT,
            trust_score REAL,
            FOREIGN KEY(entity_id, trait_id) REFERENCES entity_visual_traits(entity_id, trait_id)
        );

        CREATE TABLE IF NOT EXISTS knowledge_builds (
            build_id TEXT PRIMARY KEY,
            build_timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
            entity_count INTEGER,
            trait_count INTEGER
        );

        -- Indexes for performance
        CREATE INDEX IF NOT EXISTS idx_entities_canonical_name ON entities(canonical_name);
        CREATE INDEX IF NOT EXISTS idx_visual_traits_entity ON entity_visual_traits(entity_id);
        CREATE INDEX IF NOT EXISTS idx_relationships_entity_a ON entity_relationships(entity_id_a);
        """
        cursor = self._conn.cursor()
        cursor.executescript(schema)
        self._conn.commit()

    def get_entity_traits(self, entity_id: str) -> List[sqlite3.Row]:
        if not self._conn:
            self.connect()
        cursor = self._conn.cursor()
        cursor.execute("SELECT * FROM entity_visual_traits WHERE entity_id = ?", (entity_id,))
        return cursor.fetchall()
        
    def add_entity(self, entity_id: str, canonical_name: str, franchise_id: Optional[str] = None):
        if not self._conn:
            self.connect()
        cursor = self._conn.cursor()
        cursor.execute(
            "INSERT OR REPLACE INTO entities (entity_id, canonical_name, franchise_id) VALUES (?, ?, ?)",
            (entity_id, canonical_name, franchise_id)
        )
        self._conn.commit()