singularity_llm/splitbit/recursive_link.py

"""Recursive Link Memory — knowledge graph for conversation contexts.

Every conversation turn is stored as a linked context.
Related contexts are auto-linked (keyword overlap, semantic similarity).
During inference: inject linked context as compressed prefix.
Links decay over time, reinforced by co-access.
Persistent storage in SQLite.
"""

from __future__ import annotations

import hashlib
import logging
import os
import sqlite3
import time
from collections import defaultdict, deque
from dataclasses import dataclass, field
from typing import Any

logger = logging.getLogger(__name__)


@dataclass
class Context:
    """A conversation context stored in the recursive link graph."""
    id: str
    user_message: str
    assistant_response: str
    session_id: str = ""
    channel: str = "cli"
    timestamp: float = field(default_factory=time.time)
    access_count: int = 0
    keywords: list[str] = field(default_factory=list)


@dataclass
class Link:
    """A link between two contexts in the recursive link graph."""
    source_id: str
    target_id: str
    strength: float = 1.0
    created_at: float = field(default_factory=time.time)
    last_accessed: float = field(default_factory=time.time)
    access_count: int = 0


class RecursiveLinkGraph:
    """Recursive link graph for conversation memory.

    Features:
    - Context storage (every conversation turn)
    - Auto-linking based on keyword overlap
    - Link traversal (find related contexts)
    - Link decay (stale links fade over time)
    - Link reinforcement (co-access strengthens links)
    - Persistent storage in SQLite
    """

    DECAY_RATE = 0.001  # per second
    MIN_STRENGTH = 0.01
    MAX_LINKS_PER_CONTEXT = 50

    def __init__(self, db_path: str | None = None) -> None:
        self.db_path = db_path
        self._contexts: dict[str, Context] = {}
        self._links: dict[str, list[Link]] = defaultdict(list)
        self._keyword_index: dict[str, set[str]] = defaultdict(set)
        self._stats = {
            "contexts_stored": 0,
            "links_created": 0,
            "links_traversed": 0,
            "contexts_injected": 0,
        }

        if db_path:
            os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True)
            self._init_db()
            self._load_from_db()

    def _init_db(self) -> None:
        """Initialize SQLite database."""
        with sqlite3.connect(self.db_path) as conn:
            conn.executescript("""
                CREATE TABLE IF NOT EXISTS contexts (
                    id TEXT PRIMARY KEY,
                    user_message TEXT,
                    assistant_response TEXT,
                    session_id TEXT,
                    channel TEXT,
                    timestamp REAL,
                    access_count INTEGER,
                    keywords TEXT
                );
                CREATE TABLE IF NOT EXISTS links (
                    source_id TEXT,
                    target_id TEXT,
                    strength REAL,
                    created_at REAL,
                    last_accessed REAL,
                    access_count INTEGER,
                    PRIMARY KEY (source_id, target_id)
                );
                CREATE INDEX IF NOT EXISTS idx_keywords ON contexts(keywords);
            """)

    def _load_from_db(self) -> None:
        """Load contexts and links from SQLite."""
        with sqlite3.connect(self.db_path) as conn:
            for row in conn.execute("SELECT * FROM contexts"):
                ctx = Context(
                    id=row[0], user_message=row[1], assistant_response=row[2],
                    session_id=row[3], channel=row[4], timestamp=row[5],
                    access_count=row[6],
                    keywords=row[7].split(",") if row[7] else [],
                )
                self._contexts[ctx.id] = ctx
                for kw in ctx.keywords:
                    self._keyword_index[kw].add(ctx.id)

            for row in conn.execute("SELECT * FROM links"):
                link = Link(
                    source_id=row[0], target_id=row[1], strength=row[2],
                    created_at=row[3], last_accessed=row[4], access_count=row[5],
                )
                self._links[link.source_id].append(link)

        logger.info("Loaded %d contexts, %d links from DB", len(self._contexts), sum(len(v) for v in self._links.values()))

    def _save_context(self, ctx: Context) -> None:
        if not self.db_path:
            return
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                "INSERT OR REPLACE INTO contexts VALUES (?,?,?,?,?,?,?,?)",
                (ctx.id, ctx.user_message, ctx.assistant_response, ctx.session_id,
                 ctx.channel, ctx.timestamp, ctx.access_count, ",".join(ctx.keywords))
            )

    def _save_link(self, link: Link) -> None:
        if not self.db_path:
            return
        with sqlite3.connect(self.db_path) as conn:
            conn.execute(
                "INSERT OR REPLACE INTO links VALUES (?,?,?,?,?,?)",
                (link.source_id, link.target_id, link.strength,
                 link.created_at, link.last_accessed, link.access_count)
            )

    def _extract_keywords(self, text: str) -> list[str]:
        """Extract keywords from text (simple word frequency)."""
        words = text.lower().split()
        # Remove very common words
        stop = {"the", "a", "an", "is", "are", "was", "were", "be", "been",
                "have", "has", "had", "do", "does", "did", "will", "would",
                "could", "should", "may", "might", "can", "to", "of", "in",
                "on", "at", "by", "for", "with", "about", "as", "into", "like",
                "through", "after", "over", "between", "out", "against",
                "during", "without", "before", "under", "around", "among"}
        keywords = [w for w in words if len(w) > 2 and w not in stop]
        return list(set(keywords))[:20]

    def add_context(self, user_message: str, assistant_response: str,
                    session_id: str = "", channel: str = "cli") -> str:
        """Add a conversation context to the graph and auto-link it."""
        ctx_id = hashlib.sha256(
            f"{user_message}:{assistant_response}:{time.time()}".encode()
        ).hexdigest()[:16]

        keywords = self._extract_keywords(user_message + " " + assistant_response)
        ctx = Context(
            id=ctx_id, user_message=user_message, assistant_response=assistant_response,
            session_id=session_id, channel=channel, keywords=keywords,
        )

        self._contexts[ctx_id] = ctx
        self._stats["contexts_stored"] += 1

        # Index keywords
        for kw in keywords:
            self._keyword_index[kw].add(ctx_id)

        # Auto-link to related contexts
        self._auto_link(ctx_id, keywords)

        # Save to DB
        self._save_context(ctx)

        return ctx_id

    def _auto_link(self, ctx_id: str, keywords: list[str]) -> None:
        """Automatically create links to contexts with overlapping keywords."""
        related: dict[str, int] = defaultdict(int)

        for kw in keywords:
            for other_id in self._keyword_index.get(kw, set()):
                if other_id != ctx_id:
                    related[other_id] += 1

        # Sort by overlap count and create links
        for other_id, overlap in sorted(related.items(), key=lambda x: -x[1])[:self.MAX_LINKS_PER_CONTEXT]:
            strength = min(1.0, overlap / max(len(keywords), 1))
            link = Link(source_id=ctx_id, target_id=other_id, strength=strength)
            self._links[ctx_id].append(link)
            self._stats["links_created"] += 1
            self._save_link(link)

            # Bidirectional link
            rev_link = Link(source_id=other_id, target_id=ctx_id, strength=strength)
            self._links[other_id].append(rev_link)
            self._save_link(rev_link)

    def find_related(self, message: str, max_results: int = 3) -> list[Context]:
        """Find contexts related to the given message."""
        keywords = self._extract_keywords(message)
        related: dict[str, float] = defaultdict(float)

        for kw in keywords:
            for ctx_id in self._keyword_index.get(kw, set()):
                if ctx_id in self._contexts:
                    # Apply link decay
                    ctx = self._contexts[ctx_id]
                    age = time.time() - ctx.timestamp
                    decayed_strength = max(self.MIN_STRENGTH, 1.0 - age * self.DECAY_RATE)
                    related[ctx_id] += decayed_strength

        # Sort by relevance
        sorted_ids = sorted(related.items(), key=lambda x: -x[1])[:max_results]
        results = []
        for ctx_id, score in sorted_ids:
            ctx = self._contexts[ctx_id]
            ctx.access_count += 1
            self._stats["links_traversed"] += 1
            results.append(ctx)

        return results

    def get_injection_context(self, message: str) -> str:
        """Get related context text to inject as a prefix for inference."""
        related = self.find_related(message, max_results=3)
        if not related:
            return ""

        self._stats["contexts_injected"] += 1
        parts = []
        for ctx in related:
            parts.append(f"Previous: Q: {ctx.user_message[:100]} A: {ctx.assistant_response[:100]}")
        return " | ".join(parts)

    def decay_links(self) -> None:
        """Apply time-based decay to all links. Call periodically."""
        now = time.time()
        for ctx_id, links in self._links.items():
            for link in links:
                age = now - link.last_accessed
                link.strength = max(self.MIN_STRENGTH, link.strength - age * self.DECAY_RATE)

    def get_stats(self) -> dict[str, Any]:
        return {
            **self._stats,
            "total_contexts": len(self._contexts),
            "total_links": sum(len(v) for v in self._links.values()),
            "keyword_index_size": len(self._keyword_index),
        }