agent-brain-rag 1.2.0__py3-none-any.whl → 3.0.0__py3-none-any.whl

agent_brain_server/storage/graph_store.py

@@ -0,0 +1,519 @@
+"""Graph store manager for GraphRAG feature (Feature 113).
+
+Provides abstraction over graph storage backends:
+- SimplePropertyGraphStore: In-memory graph with JSON persistence (default)
+- Kuzu: High-performance embedded graph database (optional)
+
+All graph operations are no-ops when ENABLE_GRAPH_INDEX is False.
+"""
+
+import json
+import logging
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Optional
+
+from agent_brain_server.config import settings
+
+logger = logging.getLogger(__name__)
+
+
+class GraphStoreManager:
+    """Manages graph storage backends for GraphRAG.
+
+    Supports SimplePropertyGraphStore (default) and Kuzu (optional).
+    Implements singleton pattern for consistent graph access.
+
+    Attributes:
+        persist_dir: Directory for graph persistence.
+        store_type: Backend type - "simple" or "kuzu".
+    """
+
+    _instance: Optional["GraphStoreManager"] = None
+
+    def __init__(self, persist_dir: Path, store_type: str = "simple") -> None:
+        """Initialize graph store manager.
+
+        Args:
+            persist_dir: Directory for graph persistence.
+            store_type: Backend type - "simple" or "kuzu".
+        """
+        self.persist_dir = persist_dir
+        self.store_type = store_type
+        self._graph_store: Optional[Any] = None
+        self._initialized = False
+        self._entity_count = 0
+        self._relationship_count = 0
+        self._last_updated: Optional[datetime] = None
+
+    @classmethod
+    def get_instance(
+        cls,
+        persist_dir: Optional[Path] = None,
+        store_type: Optional[str] = None,
+    ) -> "GraphStoreManager":
+        """Get or create singleton instance.
+
+        Args:
+            persist_dir: Directory for graph persistence.
+            store_type: Backend type - "simple" or "kuzu".
+
+        Returns:
+            The singleton GraphStoreManager instance.
+        """
+        if cls._instance is None:
+            if persist_dir is None:
+                persist_dir = Path(settings.GRAPH_INDEX_PATH)
+            if store_type is None:
+                store_type = settings.GRAPH_STORE_TYPE
+            cls._instance = cls(persist_dir, store_type)
+        return cls._instance
+
+    @classmethod
+    def reset_instance(cls) -> None:
+        """Reset the singleton instance. Used for testing."""
+        cls._instance = None
+
+    def initialize(self) -> None:
+        """Initialize the graph store based on store_type.
+
+        For "simple": Uses SimplePropertyGraphStore with JSON persistence.
+        For "kuzu": Attempts to use Kuzu, falls back to simple with warning.
+
+        This is a no-op when ENABLE_GRAPH_INDEX is False.
+        """
+        if not settings.ENABLE_GRAPH_INDEX:
+            logger.debug("Graph indexing disabled, skipping initialization")
+            return
+
+        if self._initialized:
+            logger.debug("Graph store already initialized")
+            return
+
+        # Ensure persistence directory exists
+        self.persist_dir.mkdir(parents=True, exist_ok=True)
+
+        if self.store_type == "kuzu":
+            self._initialize_kuzu_store()
+        else:
+            self._initialize_simple_store()
+
+        # Try to load existing graph data
+        self.load()
+
+        self._initialized = True
+        logger.info(
+            f"Graph store initialized: type={self.store_type}, "
+            f"entities={self._entity_count}, relationships={self._relationship_count}"
+        )
+
+    def _initialize_simple_store(self) -> None:
+        """Initialize SimplePropertyGraphStore backend."""
+        try:
+            from llama_index.core.graph_stores import SimplePropertyGraphStore
+
+            self._graph_store = SimplePropertyGraphStore()
+            logger.debug("Initialized SimplePropertyGraphStore")
+        except ImportError as e:
+            logger.warning(f"Failed to import SimplePropertyGraphStore: {e}")
+            # Create a minimal fallback store
+            self._graph_store = _MinimalGraphStore()
+            logger.debug("Using minimal fallback graph store")
+
+    def _initialize_kuzu_store(self) -> None:
+        """Initialize Kuzu graph store with fallback to simple."""
+        try:
+            import kuzu  # noqa: F401 - just check import
+            from llama_index.graph_stores.kuzu import KuzuPropertyGraphStore
+
+            kuzu_db_path = self.persist_dir / "kuzu_db"
+            kuzu_db_path.mkdir(parents=True, exist_ok=True)
+
+            self._graph_store = KuzuPropertyGraphStore(
+                database_path=str(kuzu_db_path),
+            )
+            logger.debug(f"Initialized KuzuPropertyGraphStore at {kuzu_db_path}")
+        except ImportError as e:
+            logger.warning(
+                f"Kuzu not available ({e}), falling back to SimplePropertyGraphStore. "
+                "Install with: pip install llama-index-graph-stores-kuzu"
+            )
+            self.store_type = "simple"
+            self._initialize_simple_store()
+
+    def persist(self) -> None:
+        """Persist graph to disk.
+
+        For SimplePropertyGraphStore, serializes to JSON.
+        For Kuzu, data is automatically persisted.
+
+        This is a no-op when ENABLE_GRAPH_INDEX is False or not initialized.
+        """
+        if not settings.ENABLE_GRAPH_INDEX:
+            return
+
+        if not self._initialized or self._graph_store is None:
+            logger.debug("Graph store not initialized, skipping persist")
+            return
+
+        if self.store_type == "simple":
+            self._persist_simple_store()
+
+        self._last_updated = datetime.now(timezone.utc)
+        logger.debug(
+            f"Graph persisted: entities={self._entity_count}, "
+            f"relationships={self._relationship_count}"
+        )
+
+    def _persist_simple_store(self) -> None:
+        """Persist SimplePropertyGraphStore to JSON."""
+        persist_path = self.persist_dir / "graph_store.json"
+        llamaindex_persist_path = self.persist_dir / "graph_store_llamaindex.json"
+
+        try:
+            # Try LlamaIndex native persistence first
+            graph_store = self._graph_store
+            if graph_store is not None and hasattr(graph_store, "persist"):
+                graph_store.persist(str(llamaindex_persist_path))
+                logger.debug(
+                    f"Graph persisted via LlamaIndex to {llamaindex_persist_path}"
+                )
+            elif graph_store is not None and hasattr(graph_store, "_data"):
+                # Minimal store fallback - use our own format
+                data = getattr(graph_store, "_data", {})
+                with open(persist_path, "w") as f:
+                    json.dump(data, f, indent=2, default=str)
+                logger.debug(f"Graph persisted to {persist_path}")
+
+            # Always persist metadata separately
+            metadata = {
+                "entity_count": self._entity_count,
+                "relationship_count": self._relationship_count,
+                "last_updated": (
+                    self._last_updated.isoformat() if self._last_updated else None
+                ),
+                "store_type": self.store_type,
+            }
+            metadata_path = self.persist_dir / "graph_metadata.json"
+            with open(metadata_path, "w") as f:
+                json.dump(metadata, f, indent=2)
+
+        except (OSError, TypeError) as e:
+            logger.error(f"Failed to persist graph store: {e}")
+
+    def load(self) -> bool:
+        """Load graph from disk.
+
+        For SimplePropertyGraphStore, loads from JSON.
+        For Kuzu, data is automatically loaded.
+
+        Returns:
+            True if loaded successfully, False otherwise.
+        """
+        if not settings.ENABLE_GRAPH_INDEX:
+            return False
+
+        if self._graph_store is None:
+            return False
+
+        if self.store_type == "simple":
+            return self._load_simple_store()
+
+        # Kuzu loads automatically, just update counts
+        self._update_counts()
+        return True
+
+    def _load_simple_store(self) -> bool:
+        """Load SimplePropertyGraphStore from persisted data."""
+        llamaindex_persist_path = self.persist_dir / "graph_store_llamaindex.json"
+        persist_path = self.persist_dir / "graph_store.json"
+        metadata_path = self.persist_dir / "graph_metadata.json"
+
+        # Load metadata if available
+        if metadata_path.exists():
+            try:
+                with open(metadata_path) as f:
+                    metadata = json.load(f)
+                self._entity_count = metadata.get("entity_count", 0)
+                self._relationship_count = metadata.get("relationship_count", 0)
+                last_updated_str = metadata.get("last_updated")
+                if last_updated_str:
+                    self._last_updated = datetime.fromisoformat(last_updated_str)
+            except (OSError, json.JSONDecodeError) as e:
+                logger.warning(f"Failed to load graph metadata: {e}")
+
+        # Try LlamaIndex native load first
+        if llamaindex_persist_path.exists():
+            try:
+                from llama_index.core.graph_stores import SimplePropertyGraphStore
+
+                self._graph_store = SimplePropertyGraphStore.from_persist_path(
+                    str(llamaindex_persist_path)
+                )
+                logger.debug(
+                    f"Graph loaded from {llamaindex_persist_path}: "
+                    f"entities={self._entity_count}, "
+                    f"relationships={self._relationship_count}"
+                )
+                return True
+            except Exception as e:
+                logger.warning(f"Failed to load via LlamaIndex: {e}")
+
+        # Fall back to minimal store format
+        if persist_path.exists():
+            try:
+                with open(persist_path) as f:
+                    data = json.load(f)
+
+                # Restore minimal store data
+                graph_store = self._graph_store
+                if graph_store is not None and hasattr(graph_store, "_data"):
+                    graph_store._data = data
+                    if "entities" in data:
+                        graph_store._entities = data.get("entities", {})
+                    if "relationships" in data:
+                        graph_store._relationships = data.get("relationships", [])
+
+                logger.debug(
+                    f"Graph loaded from {persist_path}: "
+                    f"entities={self._entity_count}, "
+                    f"relationships={self._relationship_count}"
+                )
+                return True
+            except (OSError, json.JSONDecodeError) as e:
+                logger.error(f"Failed to load graph store: {e}")
+                return False
+
+        logger.debug("No graph data found to load")
+        return False
+
+    def _update_counts(self) -> None:
+        """Update entity and relationship counts from graph store."""
+        if self._graph_store is None:
+            return
+
+        try:
+            # Try to get counts from graph store
+            if hasattr(self._graph_store, "get_triplets"):
+                triplets = self._graph_store.get_triplets()
+                entities: set[str] = set()
+                for triplet in triplets:
+                    if hasattr(triplet, "subject"):
+                        entities.add(triplet.subject)
+                    if hasattr(triplet, "object"):
+                        entities.add(triplet.object)
+                self._entity_count = len(entities)
+                self._relationship_count = len(triplets)
+            elif hasattr(self._graph_store, "_entities"):
+                self._entity_count = len(self._graph_store._entities)
+                self._relationship_count = len(
+                    getattr(self._graph_store, "_relationships", [])
+                )
+        except Exception as e:
+            logger.warning(f"Failed to update graph counts: {e}")
+
+    def add_triplet(
+        self,
+        subject: str,
+        predicate: str,
+        obj: str,
+        subject_type: Optional[str] = None,
+        object_type: Optional[str] = None,
+        source_chunk_id: Optional[str] = None,
+    ) -> bool:
+        """Add a triplet to the graph.
+
+        Args:
+            subject: Subject entity.
+            predicate: Relationship type.
+            obj: Object entity.
+            subject_type: Optional type for subject.
+            object_type: Optional type for object.
+            source_chunk_id: Optional source chunk ID.
+
+        Returns:
+            True if added successfully, False otherwise.
+        """
+        if not settings.ENABLE_GRAPH_INDEX:
+            return False
+
+        if not self._initialized or self._graph_store is None:
+            return False
+
+        try:
+            if hasattr(self._graph_store, "upsert_triplet"):
+                self._graph_store.upsert_triplet(
+                    subject=subject,
+                    predicate=predicate,
+                    object_=obj,
+                )
+            elif hasattr(self._graph_store, "add_triplet"):
+                self._graph_store.add_triplet(subject, predicate, obj)
+            elif hasattr(self._graph_store, "_add_triplet"):
+                # Minimal store fallback
+                self._graph_store._add_triplet(
+                    subject, predicate, obj, subject_type, object_type, source_chunk_id
+                )
+
+            # Update counts
+            self._entity_count = max(
+                self._entity_count,
+                self._entity_count + 1,  # Approximate
+            )
+            self._relationship_count += 1
+            self._last_updated = datetime.now(timezone.utc)
+
+            return True
+        except Exception as e:
+            logger.error(f"Failed to add triplet: {e}")
+            return False
+
+    def clear(self) -> None:
+        """Clear all graph data.
+
+        This is a no-op when ENABLE_GRAPH_INDEX is False.
+        """
+        if not settings.ENABLE_GRAPH_INDEX:
+            return
+
+        if self._graph_store is not None:
+            if hasattr(self._graph_store, "clear"):
+                self._graph_store.clear()
+            elif hasattr(self._graph_store, "_data"):
+                self._graph_store._data = {}
+
+        self._entity_count = 0
+        self._relationship_count = 0
+        self._last_updated = None
+
+        # Remove persisted data
+        persist_path = self.persist_dir / "graph_store.json"
+        if persist_path.exists():
+            persist_path.unlink()
+
+        logger.info("Graph store cleared")
+
+    @property
+    def is_initialized(self) -> bool:
+        """Check if the graph store is initialized."""
+        return self._initialized
+
+    @property
+    def entity_count(self) -> int:
+        """Return number of entities in graph."""
+        return self._entity_count
+
+    @property
+    def relationship_count(self) -> int:
+        """Return number of relationships in graph."""
+        return self._relationship_count
+
+    @property
+    def last_updated(self) -> Optional[datetime]:
+        """Return timestamp of last update."""
+        return self._last_updated
+
+    @property
+    def graph_store(self) -> Optional[Any]:
+        """Return the underlying graph store instance."""
+        return self._graph_store
+
+
+class _MinimalGraphStore:
+    """Minimal fallback graph store when LlamaIndex is not available.
+
+    Provides basic in-memory graph storage with JSON serialization.
+    """
+
+    def __init__(self) -> None:
+        """Initialize minimal graph store."""
+        self._data: dict[str, Any] = {
+            "entities": {},
+            "relationships": [],
+        }
+        self._entities: dict[str, dict[str, Any]] = {}
+        self._relationships: list[dict[str, Any]] = []
+
+    def _add_triplet(
+        self,
+        subject: str,
+        predicate: str,
+        obj: str,
+        subject_type: Optional[str] = None,
+        object_type: Optional[str] = None,
+        source_chunk_id: Optional[str] = None,
+    ) -> None:
+        """Add a triplet to the minimal store."""
+        # Add entities
+        if subject not in self._entities:
+            self._entities[subject] = {"name": subject, "type": subject_type}
+        if obj not in self._entities:
+            self._entities[obj] = {"name": obj, "type": object_type}
+
+        # Add relationship
+        self._relationships.append(
+            {
+                "subject": subject,
+                "predicate": predicate,
+                "object": obj,
+                "source_chunk_id": source_chunk_id,
+            }
+        )
+
+        # Update data dict
+        self._data["entities"] = self._entities
+        self._data["relationships"] = self._relationships
+
+    def clear(self) -> None:
+        """Clear all data."""
+        self._data = {"entities": {}, "relationships": []}
+        self._entities = {}
+        self._relationships = []
+
+
+# Module-level singleton access
+_graph_store_manager: Optional[GraphStoreManager] = None
+
+
+def get_graph_store_manager(
+    persist_dir: Optional[Path] = None,
+    store_type: Optional[str] = None,
+) -> GraphStoreManager:
+    """Get the global graph store manager instance.
+
+    Args:
+        persist_dir: Directory for graph persistence.
+        store_type: Backend type - "simple" or "kuzu".
+
+    Returns:
+        The singleton GraphStoreManager instance.
+    """
+    global _graph_store_manager
+    if _graph_store_manager is None:
+        _graph_store_manager = GraphStoreManager.get_instance(persist_dir, store_type)
+    return _graph_store_manager
+
+
+def initialize_graph_store(
+    persist_dir: Optional[Path] = None,
+    store_type: Optional[str] = None,
+) -> GraphStoreManager:
+    """Initialize and return the global graph store manager.
+
+    Args:
+        persist_dir: Directory for graph persistence.
+        store_type: Backend type - "simple" or "kuzu".
+
+    Returns:
+        The initialized GraphStoreManager instance.
+    """
+    manager = get_graph_store_manager(persist_dir, store_type)
+    manager.initialize()
+    return manager
+
+
+def reset_graph_store_manager() -> None:
+    """Reset the global graph store manager. Used for testing."""
+    global _graph_store_manager
+    _graph_store_manager = None
+    GraphStoreManager.reset_instance()

agent_brain_server/storage/vector_store.py

@@ -260,6 +260,41 @@ class VectorStoreManager:
                 return 0
             return self._collection.count()
 
+    async def get_by_id(self, chunk_id: str) -> Optional[dict[str, Any]]:
+        """
+        Get a document by its chunk ID.
+
+        Args:
+            chunk_id: The unique identifier of the chunk.
+
+        Returns:
+            Dictionary with 'text' and 'metadata' keys, or None if not found.
+        """
+        if not self.is_initialized:
+            return None
+
+        async with self._lock:
+            assert self._collection is not None
+            try:
+                results = self._collection.get(
+                    ids=[chunk_id],
+                    include=["documents", "metadatas"],  # type: ignore[list-item]
+                )
+
+                if results["ids"] and results["ids"]:
+                    documents = results.get("documents", [[]])
+                    metadatas = results.get("metadatas", [[]])
+                    text = documents[0] if documents else ""
+                    metadata = metadatas[0] if metadatas else {}
+                    return {
+                        "text": text,
+                        "metadata": metadata if metadata else {},
+                    }
+            except Exception as e:
+                logger.warning(f"Failed to get document by ID {chunk_id}: {e}")
+
+            return None
+
     async def delete_collection(self) -> None:
         """
         Delete the entire collection.

agent_brain_server/storage_paths.py

@@ -5,13 +5,14 @@ from pathlib import Path
 
 logger = logging.getLogger(__name__)
 
-STATE_DIR_NAME = ".claude/doc-serve"
+STATE_DIR_NAME = ".claude/agent-brain"
 
 SUBDIRECTORIES = [
     "data",
     "data/chroma_db",
     "data/bm25_index",
     "data/llamaindex",
+    "data/graph_index",
     "logs",
 ]
 
@@ -19,7 +20,7 @@ SUBDIRECTORIES = [
 def resolve_state_dir(project_root: Path) -> Path:
     """Resolve the state directory for a project.
 
-    Returns <project_root>/.claude/doc-serve/
+    Returns <project_root>/.claude/agent-brain/
 
     Args:
         project_root: Resolved project root path.
@@ -48,6 +49,7 @@ def resolve_storage_paths(state_dir: Path) -> dict[str, Path]:
         "chroma_db": state_dir / "data" / "chroma_db",
         "bm25_index": state_dir / "data" / "bm25_index",
         "llamaindex": state_dir / "data" / "llamaindex",
+        "graph_index": state_dir / "data" / "graph_index",
         "logs": state_dir / "logs",
     }
 
@@ -67,6 +69,6 @@ def resolve_shared_project_dir(project_id: str) -> Path:
     Returns:
         Path to shared project data directory.
     """
-    shared_dir = Path.home() / ".doc-serve" / "projects" / project_id / "data"
+    shared_dir = Path.home() / ".agent-brain" / "projects" / project_id / "data"
     shared_dir.mkdir(parents=True, exist_ok=True)
     return shared_dir

agent_brain_rag-1.2.0.dist-info/RECORD

@@ -1,31 +0,0 @@
-agent_brain_server/__init__.py,sha256=Mv_H50b2L5Gz9er1hpu_YZB3CCazw-I7QJgg9ZUgjgM,95
-agent_brain_server/api/__init__.py,sha256=nvTvO_ahHAAsRDlV3dL_JlNruSdan4kav_P5sT_1PFk,93
-agent_brain_server/api/main.py,sha256=lz0AngQ57zpZhscxE2QuHwt_4d6yclofo62Tdv3CPhQ,11752
-agent_brain_server/api/routers/__init__.py,sha256=Z3PUKDwxeI8T88xMQeTcQ8nq1bQnu0nYjRnqNIsDUyY,254
-agent_brain_server/api/routers/health.py,sha256=KmXkyoJDVLYakEmcdfjdqklxhRbdrhMGhnMsfeNgwzM,3491
-agent_brain_server/api/routers/index.py,sha256=jTfeY8EIbM1tksJY6jNDw6qkmKBMaAi1ufDIdrg80WQ,6756
-agent_brain_server/api/routers/query.py,sha256=TAbpuCybjacUo_-7Zm3Jt8EW8EMDTYJDLv6aGtjaUzQ,2793
-agent_brain_server/config/__init__.py,sha256=zzDErZGUBwUm5Fk43OHJN7eWpeIw_1kWdnhsN6QQqSc,84
-agent_brain_server/config/settings.py,sha256=ePpGSJ2VVov_LOtQtaOedi_HM-Yf8YcpJQzSyLHfyZQ,2629
-agent_brain_server/indexing/__init__.py,sha256=6ISRNpfpokbI1x9bo8_SIRs9xx3x3YlYgeWYWHQj0oI,595
-agent_brain_server/indexing/bm25_index.py,sha256=CIeGZIa8PwvOm_Fnz3EMgETvM2Mu6GzANWeUims0P0Q,5268
-agent_brain_server/indexing/chunking.py,sha256=kW-ZOdwOS-PfXuPAYb8qL4-7yAY-rpY9a1el2YCN9CA,29848
-agent_brain_server/indexing/document_loader.py,sha256=D6U3mDl_2jFfT_JxD-5yudtqlJW8Fh0zj5lhz0V0tPs,16090
-agent_brain_server/indexing/embedding.py,sha256=JqiE9AOSooMSsALYtx7F1KFvjYVWwJt2dW-TpNmgGBw,9055
-agent_brain_server/locking.py,sha256=yRpswpBem4964gTh1VH4VUPT-vBIRnNEVgWdorJA4Hg,3365
-agent_brain_server/models/__init__.py,sha256=JLVsiVet-JDVC9s6ZefBcTXWZT5TVXrbI0E2B3ejpWM,478
-agent_brain_server/models/health.py,sha256=54S5_uS7Ca24heQcMmRZn-fTwNNVpMDuDLE1EQwqfHE,3661
-agent_brain_server/models/index.py,sha256=pjDv7phLS6dpiHLlEtcAuXQN_aHIfA_4lMkAZ-NkXZQ,5400
-agent_brain_server/models/query.py,sha256=phl-bJFxhu83k1iP4Cx3GH5qJUQ12Y1Ah_bAwXOJVZM,6054
-agent_brain_server/project_root.py,sha256=HIY5NMRDYWYIT7K6B_UMOGo1zXn9zwAdGI6ApViKI_8,2194
-agent_brain_server/runtime.py,sha256=bcchfDBt_tn2_r-lFhxqKp0RncKrz152D-oM5RcKGrU,3076
-agent_brain_server/services/__init__.py,sha256=E4VPN9Rqa2mxGQQEQn-5IYj63LSPTrA8aIx8ENO5xcc,296
-agent_brain_server/services/indexing_service.py,sha256=ichWj0z-PaIWLMmSp85mkE80J7JjpzmRx6-jT3aA-BU,18563
-agent_brain_server/services/query_service.py,sha256=RwU6eBxK0AhA0w4ZQzNxYUz_NgMGHW5yTMmKY0KTtgw,14881
-agent_brain_server/storage/__init__.py,sha256=T4N0DhDT3av7zn3Ra6uLhMsjL9N5wakqC4cl-QXWvmM,222
-agent_brain_server/storage/vector_store.py,sha256=gVaOJ9B9eRWYkteBoinb6vQHsKuRQEm8-vTo-nXuDBU,10883
-agent_brain_server/storage_paths.py,sha256=aHSsTwELwIk1VOLO0L6O-RsbRJ1U20jvJ-EQsFKG1a8,1761
-agent_brain_rag-1.2.0.dist-info/METADATA,sha256=LilsuwE92JbI_FJh9cYXhqhe2IdRnAfm7ijdQKEwi0s,5626
-agent_brain_rag-1.2.0.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-agent_brain_rag-1.2.0.dist-info/entry_points.txt,sha256=r6kfBauVV9OtMXQ7trM9ofbW5N_kdzqTouoJ9zc31-I,122
-agent_brain_rag-1.2.0.dist-info/RECORD,,