knowledge_graph.py

# core/knowledge_graph.py
import sqlite3
from datetime import datetime
import threading
import os
import hashlib

BASE_DIR = os.path.dirname(os.path.abspath(__file__))


class KnowledgeGraph:
    def __init__(self, db_path=None):
        self.db_path = db_path or "/data/invicta_data/knowledge.db"
        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
        self.lock = threading.Lock()
        self._init_tables()

    def _get_cursor(self):
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        conn.execute("PRAGMA foreign_keys = ON")
        conn.row_factory = sqlite3.Row
        return conn, conn.cursor()

    def _init_tables(self):
        with self.lock:
            conn, cur = self._get_cursor()
            try:
                cur.execute("""
                CREATE TABLE IF NOT EXISTS facts (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    topic TEXT NOT NULL,
                    fact TEXT NOT NULL,
                    fact_hash TEXT,
                    source_url TEXT DEFAULT '',
                    confidence REAL DEFAULT 1.0,
                    times_seen INTEGER DEFAULT 1,
                    learned_at TEXT DEFAULT (datetime('now'))
                );
            """)
                cur.execute("""
                CREATE TABLE IF NOT EXISTS entities (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT UNIQUE NOT NULL,
                    type TEXT,
                    description TEXT,
                    learned_at TEXT DEFAULT (datetime('now'))
                );
            """)
                cur.execute("""
                CREATE TABLE IF NOT EXISTS relationships (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    entity_a TEXT NOT NULL,
                    relation TEXT NOT NULL,
                    entity_b TEXT NOT NULL,
                    source_url TEXT DEFAULT '',
                    learned_at TEXT DEFAULT (datetime('now')),
                    UNIQUE(entity_a, relation, entity_b)
                );
            """)
                cur.execute("""
                CREATE TABLE IF NOT EXISTS topics_learned (
                    topic TEXT PRIMARY KEY,
                    times_queried INTEGER DEFAULT 1,
                    last_updated TEXT DEFAULT (datetime('now'))
                );
            """)
                cur.execute("CREATE INDEX IF NOT EXISTS idx_facts_topic ON facts(topic);")
                cur.execute("CREATE INDEX IF NOT EXISTS idx_facts_hash ON facts(fact_hash);")
                cur.execute("CREATE INDEX IF NOT EXISTS idx_facts_learned ON facts(learned_at);")
                cur.execute("CREATE INDEX IF NOT EXISTS idx_relationships_a ON relationships(entity_a);")
                cur.execute("CREATE INDEX IF NOT EXISTS idx_relationships_b ON relationships(entity_b);")
                conn.commit()
            except Exception as e:
                conn.rollback()
                if "already exists" not in str(e).lower() and "duplicate" not in str(e).lower():
                    raise
            finally:
                conn.close()

    def _fact_hash(self, fact):
        return hashlib.md5(fact.strip().lower()[:100].encode()).hexdigest()

    # ── Facts ────────────────────────────────────────────────────────────────
    def store_fact(self, topic, fact, source_url=""):
        fh = self._fact_hash(fact)
        with self.lock:
            conn, cur = self._get_cursor()
            cur.execute("SELECT id FROM facts WHERE fact_hash = ?", (fh,))
            existing = cur.fetchone()

            if existing:
                cur.execute(
                    "UPDATE facts SET times_seen = times_seen + 1, confidence = min(confidence + 0.1, 2.0) WHERE id = ?",
                    (existing["id"],)
                )
            else:
                cur.execute(
                    "INSERT INTO facts (topic, fact, fact_hash, source_url, learned_at) VALUES (?, ?, ?, ?, ?)",
                    (topic.lower(), fact, fh, source_url, datetime.now().isoformat(sep=' ', timespec='seconds'))
                )
            conn.commit()
            conn.close()

    def recall_facts(self, topic, limit=10):
        conn, cur = self._get_cursor()
        cur.execute(
            """SELECT fact, source_url FROM facts
               WHERE topic LIKE ?
               ORDER BY confidence DESC, learned_at DESC
               LIMIT ?""",
            (f"%{topic.lower()}%", limit)
        )
        rows = cur.fetchall()
        conn.close()
        return [(r["fact"], r["source_url"]) for r in rows]

    def knows_topic(self, topic):
        conn, cur = self._get_cursor()
        cur.execute(
            "SELECT COUNT(*) as c FROM facts WHERE topic LIKE ?",
            (f"%{topic.lower()}%",)
        )
        row = cur.fetchone()
        conn.close()
        return (row["c"] if row else 0) > 0

    def get_fact_count(self):
        conn, cur = self._get_cursor()
        cur.execute("SELECT COUNT(*) as c FROM facts")
        row = cur.fetchone()
        conn.close()
        return row["c"] if row else 0

    def get_recent_facts(self, limit=25):
        conn, cur = self._get_cursor()
        cur.execute(
            "SELECT id, topic, fact, source_url, confidence, learned_at FROM facts ORDER BY id DESC LIMIT ?",
            (limit,)
        )
        rows = cur.fetchall()
        conn.close()
        return [dict(r) for r in rows]


    # ── Relationships ─────────────────────────────────────────────────────────
    def store_relationship(self, entity_a, relation, entity_b, source=""):
        with self.lock:
            conn, cur = self._get_cursor()
            cur.execute(
                """INSERT INTO relationships (entity_a, relation, entity_b, source_url, learned_at) 
                   VALUES (?, ?, ?, ?, ?)
                   ON CONFLICT DO NOTHING""",
                (entity_a.lower(), relation.lower(), entity_b.lower(), source, datetime.now().isoformat(sep=' ', timespec='seconds'))
            )
            conn.commit()
            conn.close()

    def recall_relationships(self, entity, limit=10):
        conn, cur = self._get_cursor()
        cur.execute(
            """SELECT entity_a, relation, entity_b FROM relationships
               WHERE entity_a LIKE ? OR entity_b LIKE ?
               LIMIT ?""",
            (f"%{entity.lower()}%", f"%{entity.lower()}%", limit)
        )
        rows = cur.fetchall()
        conn.close()
        return [(r["entity_a"], r["relation"], r["entity_b"]) for r in rows]

    # ── Topics ────────────────────────────────────────────────────────────────
    def mark_topic_learned(self, topic):
        with self.lock:
            conn, cur = self._get_cursor()
            cur.execute("""
                INSERT INTO topics_learned (topic, times_queried, last_updated)
                VALUES (?, 1, ?)
                ON CONFLICT(topic) DO UPDATE SET
                    times_queried = topics_learned.times_queried + 1,
                    last_updated = excluded.last_updated
            """, (topic.lower(), datetime.now().isoformat(sep=' ', timespec='seconds')))
            conn.commit()
            conn.close()

    def get_top_topics(self, limit=10):
        conn, cur = self._get_cursor()
        cur.execute(
            "SELECT topic, times_queried FROM topics_learned ORDER BY times_queried DESC LIMIT ?",
            (limit,)
        )
        rows = cur.fetchall()
        conn.close()
        return [(r["topic"], r["times_queried"]) for r in rows]

    # ── Maintenance ───────────────────────────────────────────────────────────
    def prune_old_facts(self, max_facts=5000):
        conn, cur = self._get_cursor()
        cur.execute("SELECT COUNT(*) as c FROM facts")
        count = cur.fetchone()["c"]
        conn.close()
        if count > max_facts:
            with self.lock:
                conn, cur = self._get_cursor()
                cur.execute("""
                    DELETE FROM facts WHERE id IN (
                        SELECT id FROM facts
                        ORDER BY confidence ASC, learned_at ASC
                        LIMIT ?
                    )
                """, (count - max_facts,))
                conn.commit()
                conn.close()
                print(f"🧹 Pruned {count - max_facts} old facts from knowledge graph")