rnsr 0.1.0__py3-none-any.whl

rnsr/indexing/kv_store.py

@@ -0,0 +1,313 @@
+"""
+KV Store - SQLite-Backed Key-Value Storage for Full Text
+
+The Skeleton Index pattern requires separating:
+- Summaries (stored in vector index for retrieval)
+- Full Text (stored externally in this KV Store)
+
+This prevents full text from polluting the LLM context until
+explicitly requested during synthesis.
+
+Usage:
+    kv = SQLiteKVStore("./data/document_kv.db")
+    kv.put("node_123", "Full text content here...")
+    content = kv.get("node_123")
+"""
+
+from __future__ import annotations
+
+import hashlib
+import sqlite3
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Iterator, Union
+
+import structlog
+
+from rnsr.exceptions import IndexingError
+
+logger = structlog.get_logger(__name__)
+
+
+class SQLiteKVStore:
+    """
+    SQLite-backed key-value store for document content.
+    
+    Stores full text content separately from the vector index,
+    allowing the skeleton index to contain only summaries.
+    
+    Attributes:
+        db_path: Path to the SQLite database file.
+    """
+    
+    def __init__(self, db_path: Path | str):
+        """
+        Initialize the KV store.
+        
+        Args:
+            db_path: Path to the SQLite database file.
+                     Will be created if it doesn't exist.
+        """
+        self.db_path = Path(db_path)
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+        
+        self._init_db()
+        
+        logger.info("kv_store_initialized", db_path=str(self.db_path))
+    
+    def _init_db(self) -> None:
+        """Create the database schema if it doesn't exist."""
+        with self._connect() as conn:
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS documents (
+                    node_id TEXT PRIMARY KEY,
+                    content TEXT NOT NULL,
+                    content_hash TEXT NOT NULL,
+                    char_count INTEGER NOT NULL,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                )
+            """)
+            
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_content_hash 
+                ON documents(content_hash)
+            """)
+            
+            conn.commit()
+    
+    @contextmanager
+    def _connect(self) -> Iterator[sqlite3.Connection]:
+        """Context manager for database connections."""
+        conn = sqlite3.connect(self.db_path)
+        conn.row_factory = sqlite3.Row
+        try:
+            yield conn
+        finally:
+            conn.close()
+    
+    def put(self, node_id: str, content: str) -> str:
+        """
+        Store content for a node.
+        
+        Args:
+            node_id: Unique identifier for the node.
+            content: Full text content to store.
+            
+        Returns:
+            SHA256 hash of the content.
+        """
+        content_hash = hashlib.sha256(content.encode()).hexdigest()[:16]
+        char_count = len(content)
+        
+        with self._connect() as conn:
+            conn.execute(
+                """
+                INSERT INTO documents (node_id, content, content_hash, char_count)
+                VALUES (?, ?, ?, ?)
+                ON CONFLICT(node_id) DO UPDATE SET
+                    content = excluded.content,
+                    content_hash = excluded.content_hash,
+                    char_count = excluded.char_count,
+                    updated_at = CURRENT_TIMESTAMP
+                """,
+                (node_id, content, content_hash, char_count),
+            )
+            conn.commit()
+        
+        logger.debug(
+            "kv_put",
+            node_id=node_id,
+            char_count=char_count,
+            hash=content_hash,
+        )
+        
+        return content_hash
+    
+    def get(self, node_id: str) -> str | None:
+        """
+        Retrieve content for a node.
+        
+        Args:
+            node_id: Unique identifier for the node.
+            
+        Returns:
+            Full text content, or None if not found.
+        """
+        with self._connect() as conn:
+            cursor = conn.execute(
+                "SELECT content FROM documents WHERE node_id = ?",
+                (node_id,),
+            )
+            row = cursor.fetchone()
+        
+        if row is None:
+            logger.debug("kv_miss", node_id=node_id)
+            return None
+        
+        logger.debug("kv_hit", node_id=node_id)
+        return row["content"]
+    
+    def get_batch(self, node_ids: list[str]) -> dict[str, str | None]:
+        """
+        Retrieve content for multiple nodes.
+        
+        Args:
+            node_ids: List of node identifiers.
+            
+        Returns:
+            Dictionary mapping node_id to content (or None if not found).
+        """
+        result: dict[str, str | None] = {nid: None for nid in node_ids}
+        
+        if not node_ids:
+            return result
+        
+        placeholders = ",".join("?" * len(node_ids))
+        
+        with self._connect() as conn:
+            cursor = conn.execute(
+                f"SELECT node_id, content FROM documents WHERE node_id IN ({placeholders})",
+                node_ids,
+            )
+            for row in cursor:
+                result[row["node_id"]] = row["content"]
+        
+        found = sum(1 for v in result.values() if v is not None)
+        logger.debug("kv_batch_get", requested=len(node_ids), found=found)
+        
+        return result
+    
+    def delete(self, node_id: str) -> bool:
+        """
+        Delete content for a node.
+        
+        Args:
+            node_id: Unique identifier for the node.
+            
+        Returns:
+            True if deleted, False if not found.
+        """
+        with self._connect() as conn:
+            cursor = conn.execute(
+                "DELETE FROM documents WHERE node_id = ?",
+                (node_id,),
+            )
+            conn.commit()
+            deleted = cursor.rowcount > 0
+        
+        logger.debug("kv_delete", node_id=node_id, deleted=deleted)
+        return deleted
+    
+    def exists(self, node_id: str) -> bool:
+        """Check if a node exists in the store."""
+        with self._connect() as conn:
+            cursor = conn.execute(
+                "SELECT 1 FROM documents WHERE node_id = ? LIMIT 1",
+                (node_id,),
+            )
+            return cursor.fetchone() is not None
+    
+    def count(self) -> int:
+        """Get the total number of stored documents."""
+        with self._connect() as conn:
+            cursor = conn.execute("SELECT COUNT(*) FROM documents")
+            return cursor.fetchone()[0]
+    
+    def get_metadata(self, node_id: str) -> dict | None:
+        """
+        Get metadata about a stored document.
+        
+        Args:
+            node_id: Unique identifier for the node.
+            
+        Returns:
+            Dictionary with hash, char_count, timestamps, or None.
+        """
+        with self._connect() as conn:
+            cursor = conn.execute(
+                """
+                SELECT content_hash, char_count, created_at, updated_at 
+                FROM documents WHERE node_id = ?
+                """,
+                (node_id,),
+            )
+            row = cursor.fetchone()
+        
+        if row is None:
+            return None
+        
+        return {
+            "content_hash": row["content_hash"],
+            "char_count": row["char_count"],
+            "created_at": row["created_at"],
+            "updated_at": row["updated_at"],
+        }
+    
+    def clear(self) -> int:
+        """
+        Delete all documents from the store.
+        
+        Returns:
+            Number of documents deleted.
+        """
+        with self._connect() as conn:
+            cursor = conn.execute("DELETE FROM documents")
+            count = cursor.rowcount
+            conn.commit()
+        
+        logger.warning("kv_store_cleared", count=count)
+        return count
+
+
+class InMemoryKVStore:
+    """
+    In-memory key-value store for testing and ephemeral usage.
+    
+    API-compatible with SQLiteKVStore.
+    """
+    
+    def __init__(self):
+        self._store: dict[str, str] = {}
+        self._metadata: dict[str, dict] = {}
+    
+    def put(self, node_id: str, content: str) -> str:
+        content_hash = hashlib.sha256(content.encode()).hexdigest()[:16]
+        self._store[node_id] = content
+        self._metadata[node_id] = {
+            "content_hash": content_hash,
+            "char_count": len(content),
+        }
+        return content_hash
+    
+    def get(self, node_id: str) -> str | None:
+        return self._store.get(node_id)
+    
+    def get_batch(self, node_ids: list[str]) -> dict[str, str | None]:
+        return {nid: self._store.get(nid) for nid in node_ids}
+    
+    def delete(self, node_id: str) -> bool:
+        if node_id in self._store:
+            del self._store[node_id]
+            del self._metadata[node_id]
+            return True
+        return False
+    
+    def exists(self, node_id: str) -> bool:
+        return node_id in self._store
+    
+    def count(self) -> int:
+        return len(self._store)
+    
+    def get_metadata(self, node_id: str) -> dict | None:
+        return self._metadata.get(node_id)
+    
+    def clear(self) -> int:
+        count = len(self._store)
+        self._store.clear()
+        self._metadata.clear()
+        return count
+
+
+# Type alias for either store
+KVStore = Union[SQLiteKVStore, InMemoryKVStore]

rnsr/indexing/persistence.py

@@ -0,0 +1,323 @@
+"""
+Persistence Module - Save and Load RNSR Indexes
+
+Provides functionality to persist and restore:
+- Skeleton Index (SkeletonNode structures)
+- KV Store (already SQLite-backed, but needs export/import)
+- Document metadata
+
+Usage:
+    from rnsr.indexing import save_index, load_index
+    
+    # Save after indexing
+    skeleton, kv_store = build_skeleton_index(tree)
+    save_index(skeleton, kv_store, "./my_document_index/")
+    
+    # Load later (no re-processing needed!)
+    skeleton, kv_store = load_index("./my_document_index/")
+    answer = run_navigator("question", skeleton, kv_store)
+"""
+
+from __future__ import annotations
+
+import json
+import shutil
+from dataclasses import asdict
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import structlog
+
+from rnsr.exceptions import IndexingError
+from rnsr.indexing.kv_store import InMemoryKVStore, KVStore, SQLiteKVStore
+from rnsr.models import SkeletonNode
+
+logger = structlog.get_logger(__name__)
+
+# Version for format compatibility
+INDEX_FORMAT_VERSION = "1.0"
+
+
+def save_index(
+    skeleton: dict[str, SkeletonNode],
+    kv_store: KVStore,
+    index_dir: str | Path,
+    metadata: dict[str, Any] | None = None,
+) -> Path:
+    """
+    Save a skeleton index and KV store to disk.
+    
+    Creates a directory structure:
+        index_dir/
+            manifest.json      # Version, metadata, timestamps
+            skeleton.json      # SkeletonNode structures
+            content.db         # SQLite KV store (copied or created)
+    
+    Args:
+        skeleton: Dictionary of node_id -> SkeletonNode
+        kv_store: KV store containing full text
+        index_dir: Directory to save the index
+        metadata: Optional metadata (title, source, etc.)
+    
+    Returns:
+        Path to the index directory
+        
+    Example:
+        skeleton, kv = build_skeleton_index(tree)
+        save_index(skeleton, kv, "./indexes/contract_2024/")
+    """
+    index_path = Path(index_dir)
+    index_path.mkdir(parents=True, exist_ok=True)
+    
+    # Build manifest
+    manifest = {
+        "version": INDEX_FORMAT_VERSION,
+        "created_at": datetime.now().isoformat(),
+        "node_count": len(skeleton),
+        "metadata": metadata or {},
+    }
+    
+    # Find root node for extra info
+    root = next((n for n in skeleton.values() if n.level == 0), None)
+    if root:
+        manifest["root_id"] = root.node_id
+        manifest["root_header"] = root.header
+    
+    # Save manifest
+    manifest_path = index_path / "manifest.json"
+    with open(manifest_path, "w") as f:
+        json.dump(manifest, f, indent=2)
+    
+    # Save skeleton nodes
+    skeleton_path = index_path / "skeleton.json"
+    skeleton_data = {
+        node_id: _skeleton_node_to_dict(node)
+        for node_id, node in skeleton.items()
+    }
+    with open(skeleton_path, "w") as f:
+        json.dump(skeleton_data, f, indent=2)
+    
+    # Handle KV store
+    content_path = index_path / "content.db"
+    
+    if isinstance(kv_store, SQLiteKVStore):
+        # Copy the existing database
+        if kv_store.db_path != content_path:
+            shutil.copy2(kv_store.db_path, content_path)
+    elif isinstance(kv_store, InMemoryKVStore):
+        # Export in-memory store to SQLite
+        sqlite_kv = SQLiteKVStore(content_path)
+        for node_id in skeleton.keys():
+            content = kv_store.get(node_id)
+            if content:
+                sqlite_kv.put(node_id, content)
+    else:
+        raise IndexingError(f"Unsupported KV store type: {type(kv_store)}")
+    
+    logger.info(
+        "index_saved",
+        path=str(index_path),
+        nodes=len(skeleton),
+    )
+    
+    return index_path
+
+
+def load_index(
+    index_dir: str | Path,
+    in_memory: bool = False,
+) -> tuple[dict[str, SkeletonNode], KVStore]:
+    """
+    Load a skeleton index and KV store from disk.
+    
+    Args:
+        index_dir: Directory containing the saved index
+        in_memory: If True, load KV store into memory (faster but uses more RAM)
+    
+    Returns:
+        Tuple of (skeleton dict, kv_store)
+        
+    Example:
+        skeleton, kv = load_index("./indexes/contract_2024/")
+        answer = run_navigator("What are the payment terms?", skeleton, kv)
+    """
+    index_path = Path(index_dir)
+    
+    if not index_path.exists():
+        raise IndexingError(f"Index directory not found: {index_path}")
+    
+    # Load and validate manifest
+    manifest_path = index_path / "manifest.json"
+    if not manifest_path.exists():
+        raise IndexingError(f"Manifest not found: {manifest_path}")
+    
+    with open(manifest_path) as f:
+        manifest = json.load(f)
+    
+    version = manifest.get("version", "unknown")
+    if version != INDEX_FORMAT_VERSION:
+        logger.warning(
+            "index_version_mismatch",
+            expected=INDEX_FORMAT_VERSION,
+            found=version,
+        )
+    
+    # Load skeleton nodes
+    skeleton_path = index_path / "skeleton.json"
+    if not skeleton_path.exists():
+        raise IndexingError(f"Skeleton index not found: {skeleton_path}")
+    
+    with open(skeleton_path) as f:
+        skeleton_data = json.load(f)
+    
+    skeleton: dict[str, SkeletonNode] = {
+        node_id: _dict_to_skeleton_node(data)
+        for node_id, data in skeleton_data.items()
+    }
+    
+    # Load KV store
+    content_path = index_path / "content.db"
+    if not content_path.exists():
+        raise IndexingError(f"Content database not found: {content_path}")
+    
+    if in_memory:
+        # Load into memory for faster access
+        sqlite_kv = SQLiteKVStore(content_path)
+        kv_store = InMemoryKVStore()
+        for node_id in skeleton.keys():
+            content = sqlite_kv.get(node_id)
+            if content:
+                kv_store.put(node_id, content)
+    else:
+        # Use SQLite directly
+        kv_store = SQLiteKVStore(content_path)
+    
+    logger.info(
+        "index_loaded",
+        path=str(index_path),
+        nodes=len(skeleton),
+        version=version,
+    )
+    
+    return skeleton, kv_store
+
+
+def get_index_info(index_dir: str | Path) -> dict[str, Any]:
+    """
+    Get information about a saved index without loading it.
+    
+    Args:
+        index_dir: Directory containing the saved index
+        
+    Returns:
+        Dictionary with index metadata
+        
+    Example:
+        info = get_index_info("./indexes/contract_2024/")
+        print(f"Index has {info['node_count']} nodes")
+    """
+    index_path = Path(index_dir)
+    manifest_path = index_path / "manifest.json"
+    
+    if not manifest_path.exists():
+        raise IndexingError(f"Manifest not found: {manifest_path}")
+    
+    with open(manifest_path) as f:
+        manifest = json.load(f)
+    
+    # Add file size info
+    content_path = index_path / "content.db"
+    if content_path.exists():
+        manifest["content_size_bytes"] = content_path.stat().st_size
+    
+    skeleton_path = index_path / "skeleton.json"
+    if skeleton_path.exists():
+        manifest["skeleton_size_bytes"] = skeleton_path.stat().st_size
+    
+    return manifest
+
+
+def delete_index(index_dir: str | Path) -> bool:
+    """
+    Delete a saved index.
+    
+    Args:
+        index_dir: Directory containing the saved index
+        
+    Returns:
+        True if deleted, False if not found
+    """
+    index_path = Path(index_dir)
+    
+    if not index_path.exists():
+        return False
+    
+    shutil.rmtree(index_path)
+    logger.info("index_deleted", path=str(index_path))
+    return True
+
+
+def list_indexes(base_dir: str | Path) -> list[dict[str, Any]]:
+    """
+    List all indexes in a directory.
+    
+    Args:
+        base_dir: Directory to search for indexes
+        
+    Returns:
+        List of index info dictionaries
+        
+    Example:
+        indexes = list_indexes("./indexes/")
+        for idx in indexes:
+            print(f"{idx['path']}: {idx['node_count']} nodes")
+    """
+    base_path = Path(base_dir)
+    indexes = []
+    
+    if not base_path.exists():
+        return indexes
+    
+    for item in base_path.iterdir():
+        if item.is_dir() and (item / "manifest.json").exists():
+            try:
+                info = get_index_info(item)
+                info["path"] = str(item)
+                indexes.append(info)
+            except Exception as e:
+                logger.warning("failed_to_read_index", path=str(item), error=str(e))
+    
+    return indexes
+
+
+# =============================================================================
+# Serialization Helpers
+# =============================================================================
+
+def _skeleton_node_to_dict(node: SkeletonNode) -> dict[str, Any]:
+    """Convert SkeletonNode to JSON-serializable dict."""
+    return {
+        "node_id": node.node_id,
+        "parent_id": node.parent_id,
+        "level": node.level,
+        "header": node.header,
+        "summary": node.summary,
+        "child_ids": node.child_ids,
+        "page_num": node.page_num,
+        "metadata": node.metadata,
+    }
+
+
+def _dict_to_skeleton_node(data: dict[str, Any]) -> SkeletonNode:
+    """Convert dict back to SkeletonNode."""
+    return SkeletonNode(
+        node_id=data["node_id"],
+        parent_id=data.get("parent_id"),
+        level=data["level"],
+        header=data.get("header", ""),
+        summary=data.get("summary", ""),
+        child_ids=data.get("child_ids", []),
+        page_num=data.get("page_num"),
+        metadata=data.get("metadata", {}),
+    )