aiecs 1.5.1__py3-none-any.whl

aiecs/infrastructure/graph_storage/in_memory.py

@@ -0,0 +1,514 @@
+"""
+In-Memory Graph Store Implementation
+
+Implements Tier 1 of GraphStore interface using networkx.
+Tier 2 methods work automatically via default implementations.
+
+This is ideal for:
+- Development and testing
+- Small graphs (< 100K nodes)
+- Prototyping
+- Scenarios where persistence is not required
+"""
+
+from typing import List, Optional, Dict
+import networkx as nx
+from aiecs.domain.knowledge_graph.models.entity import Entity
+from aiecs.domain.knowledge_graph.models.relation import Relation
+from aiecs.infrastructure.graph_storage.base import GraphStore
+
+
+class InMemoryGraphStore(GraphStore):
+    """
+    In-Memory Graph Store using NetworkX
+
+    **Implementation Strategy**:
+    - Uses networkx.DiGraph for graph structure
+    - Stores Entity objects as node attributes
+    - Stores Relation objects as edge attributes
+    - Implements ONLY Tier 1 methods
+    - Tier 2 methods (traverse, find_paths, etc.) work automatically!
+
+    **Features**:
+    - Fast for small-medium graphs
+    - No external dependencies
+    - Full Python ecosystem integration
+    - Rich graph algorithms from networkx
+
+    **Limitations**:
+    - Not persistent (lost on restart)
+    - Limited by RAM
+    - No concurrent access control
+    - No vector search optimization
+
+    Example:
+        ```python
+        store = InMemoryGraphStore()
+        await store.initialize()
+
+        # Add entities
+        entity = Entity(id="person_1", entity_type="Person", properties={"name": "Alice"})
+        await store.add_entity(entity)
+
+        # Tier 2 methods work automatically
+        paths = await store.traverse("person_1", max_depth=3)
+        ```
+    """
+
+    def __init__(self):
+        """Initialize in-memory graph store"""
+        self.graph: Optional[nx.DiGraph] = None
+        self.entities: Dict[str, Entity] = {}
+        self.relations: Dict[str, Relation] = {}
+        self._initialized = False
+
+    # =========================================================================
+    # TIER 1 IMPLEMENTATION - Core CRUD Operations
+    # =========================================================================
+
+    async def initialize(self) -> None:
+        """Initialize the in-memory graph"""
+        if self._initialized:
+            return
+
+        self.graph = nx.DiGraph()
+        self.entities = {}
+        self.relations = {}
+        self._initialized = True
+
+    async def close(self) -> None:
+        """Close and cleanup (nothing to do for in-memory)"""
+        self.graph = None
+        self.entities = {}
+        self.relations = {}
+        self._initialized = False
+
+    async def add_entity(self, entity: Entity) -> None:
+        """
+        Add entity to graph
+
+        Args:
+            entity: Entity to add
+
+        Raises:
+            ValueError: If entity already exists
+            RuntimeError: If store not initialized
+        """
+        if not self._initialized:
+            raise RuntimeError("GraphStore not initialized. Call initialize() first.")
+
+        if entity.id in self.entities:
+            raise ValueError(f"Entity with ID '{entity.id}' already exists")
+
+        # Add to networkx graph
+        self.graph.add_node(entity.id, entity=entity)
+
+        # Add to entity index
+        self.entities[entity.id] = entity
+
+    async def get_entity(self, entity_id: str) -> Optional[Entity]:
+        """
+        Get entity by ID
+
+        Args:
+            entity_id: Entity ID
+
+        Returns:
+            Entity if found, None otherwise
+        """
+        if not self._initialized:
+            return None
+
+        return self.entities.get(entity_id)
+
+    async def add_relation(self, relation: Relation) -> None:
+        """
+        Add relation to graph
+
+        Args:
+            relation: Relation to add
+
+        Raises:
+            ValueError: If relation already exists or entities don't exist
+            RuntimeError: If store not initialized
+        """
+        if not self._initialized:
+            raise RuntimeError("GraphStore not initialized. Call initialize() first.")
+
+        if relation.id in self.relations:
+            raise ValueError(f"Relation with ID '{relation.id}' already exists")
+
+        # Validate entities exist
+        if relation.source_id not in self.entities:
+            raise ValueError(f"Source entity '{relation.source_id}' not found")
+        if relation.target_id not in self.entities:
+            raise ValueError(f"Target entity '{relation.target_id}' not found")
+
+        # Add to networkx graph
+        self.graph.add_edge(
+            relation.source_id,
+            relation.target_id,
+            key=relation.id,
+            relation=relation,
+        )
+
+        # Add to relation index
+        self.relations[relation.id] = relation
+
+    async def get_relation(self, relation_id: str) -> Optional[Relation]:
+        """
+        Get relation by ID
+
+        Args:
+            relation_id: Relation ID
+
+        Returns:
+            Relation if found, None otherwise
+        """
+        if not self._initialized:
+            return None
+
+        return self.relations.get(relation_id)
+
+    async def get_neighbors(
+        self,
+        entity_id: str,
+        relation_type: Optional[str] = None,
+        direction: str = "outgoing",
+    ) -> List[Entity]:
+        """
+        Get neighboring entities
+
+        Args:
+            entity_id: Entity ID to get neighbors for
+            relation_type: Optional filter by relation type
+            direction: "outgoing", "incoming", or "both"
+
+        Returns:
+            List of neighboring entities
+        """
+        if not self._initialized or entity_id not in self.graph:
+            return []
+
+        neighbors = []
+
+        # Get outgoing neighbors
+        if direction in ("outgoing", "both"):
+            for target_id in self.graph.successors(entity_id):
+                # Check relation type filter
+                if relation_type:
+                    edge_data = self.graph.get_edge_data(entity_id, target_id)
+                    if edge_data:
+                        relation = edge_data.get("relation")
+                        if relation and relation.relation_type == relation_type:
+                            if target_id in self.entities:
+                                neighbors.append(self.entities[target_id])
+                else:
+                    if target_id in self.entities:
+                        neighbors.append(self.entities[target_id])
+
+        # Get incoming neighbors
+        if direction in ("incoming", "both"):
+            for source_id in self.graph.predecessors(entity_id):
+                # Check relation type filter
+                if relation_type:
+                    edge_data = self.graph.get_edge_data(source_id, entity_id)
+                    if edge_data:
+                        relation = edge_data.get("relation")
+                        if relation and relation.relation_type == relation_type:
+                            if source_id in self.entities:
+                                neighbors.append(self.entities[source_id])
+                else:
+                    if source_id in self.entities:
+                        neighbors.append(self.entities[source_id])
+
+        return neighbors
+
+    async def get_all_entities(
+        self, entity_type: Optional[str] = None, limit: Optional[int] = None
+    ) -> List[Entity]:
+        """
+        Get all entities, optionally filtered by type
+
+        Args:
+            entity_type: Optional filter by entity type
+            limit: Optional limit on number of entities
+
+        Returns:
+            List of entities
+        """
+        if not self._initialized:
+            return []
+
+        entities = list(self.entities.values())
+
+        # Filter by entity type if specified
+        if entity_type:
+            entities = [e for e in entities if e.entity_type == entity_type]
+
+        # Apply limit if specified
+        if limit:
+            entities = entities[:limit]
+
+        return entities
+
+    # =========================================================================
+    # TIER 2 METHODS - Inherited from base class
+    # =========================================================================
+    # traverse(), find_paths(), subgraph_query(), execute_query()
+    # work automatically through default implementations!
+
+    # =========================================================================
+    # OPTIONAL: Override Tier 2 methods for optimization
+    # =========================================================================
+
+    async def vector_search(
+        self,
+        query_embedding: List[float],
+        entity_type: Optional[str] = None,
+        max_results: int = 10,
+        score_threshold: float = 0.0,
+    ) -> List[tuple]:
+        """
+        Optimized vector search for in-memory store
+
+        Performs brute-force cosine similarity over all entities with embeddings.
+
+        Args:
+            query_embedding: Query vector
+            entity_type: Optional filter by entity type
+            max_results: Maximum number of results
+            score_threshold: Minimum similarity score (0.0-1.0)
+
+        Returns:
+            List of (entity, similarity_score) tuples, sorted descending
+        """
+        if not self._initialized:
+            return []
+
+        if not query_embedding:
+            raise ValueError("Query embedding cannot be empty")
+
+        import numpy as np
+
+        query_vec = np.array(query_embedding, dtype=np.float32)
+        query_norm = np.linalg.norm(query_vec)
+
+        if query_norm == 0:
+            return []
+
+        scored_entities = []
+
+        for entity in self.entities.values():
+            # Filter by entity type if specified
+            if entity_type and entity.entity_type != entity_type:
+                continue
+
+            # Skip entities without embeddings
+            if not entity.embedding:
+                continue
+
+            # Compute cosine similarity
+            entity_vec = np.array(entity.embedding, dtype=np.float32)
+            entity_norm = np.linalg.norm(entity_vec)
+
+            if entity_norm == 0:
+                continue
+
+            # Cosine similarity
+            similarity = np.dot(query_vec, entity_vec) / (query_norm * entity_norm)
+            # Normalize to 0-1 range
+            similarity = (similarity + 1) / 2
+
+            # Filter by threshold
+            if similarity >= score_threshold:
+                scored_entities.append((entity, float(similarity)))
+
+        # Sort by score descending and return top results
+        scored_entities.sort(key=lambda x: x[1], reverse=True)
+        return scored_entities[:max_results]
+
+    async def text_search(
+        self,
+        query_text: str,
+        entity_type: Optional[str] = None,
+        max_results: int = 10,
+        score_threshold: float = 0.0,
+        method: str = "bm25",
+    ) -> List[tuple]:
+        """
+        Optimized text search for in-memory store
+
+        Performs text similarity search over entity properties using BM25, Jaccard,
+        cosine similarity, or Levenshtein distance.
+
+        Args:
+            query_text: Query text string
+            entity_type: Optional filter by entity type
+            max_results: Maximum number of results
+            score_threshold: Minimum similarity score (0.0-1.0)
+            method: Similarity method ("bm25", "jaccard", "cosine", "levenshtein")
+
+        Returns:
+            List of (entity, similarity_score) tuples, sorted descending
+        """
+        if not self._initialized:
+            return []
+
+        if not query_text:
+            return []
+
+        from aiecs.application.knowledge_graph.search.text_similarity import (
+            BM25Scorer,
+            jaccard_similarity_text,
+            cosine_similarity_text,
+            normalized_levenshtein_similarity,
+        )
+
+        # Get candidate entities
+        entities = list(self.entities.values())
+        if entity_type:
+            entities = [e for e in entities if e.entity_type == entity_type]
+
+        if not entities:
+            return []
+
+        scored_entities = []
+
+        # Extract text from entities (combine properties into searchable text)
+        entity_texts = []
+        for entity in entities:
+            # Combine all string properties into searchable text
+            text_parts = []
+            for key, value in entity.properties.items():
+                if isinstance(value, str):
+                    text_parts.append(value)
+                elif isinstance(value, (list, tuple)):
+                    text_parts.extend(str(v) for v in value if isinstance(v, str))
+            entity_text = " ".join(text_parts)
+            entity_texts.append((entity, entity_text))
+
+        if method == "bm25":
+            # Use BM25 scorer
+            corpus = [text for _, text in entity_texts]
+            scorer = BM25Scorer(corpus)
+            scores = scorer.score(query_text)
+
+            for (entity, _), score in zip(entity_texts, scores):
+                if score >= score_threshold:
+                    scored_entities.append((entity, float(score)))
+
+        elif method == "jaccard":
+            for entity, text in entity_texts:
+                score = jaccard_similarity_text(query_text, text)
+                if score >= score_threshold:
+                    scored_entities.append((entity, score))
+
+        elif method == "cosine":
+            for entity, text in entity_texts:
+                score = cosine_similarity_text(query_text, text)
+                if score >= score_threshold:
+                    scored_entities.append((entity, score))
+
+        elif method == "levenshtein":
+            for entity, text in entity_texts:
+                score = normalized_levenshtein_similarity(query_text, text)
+                if score >= score_threshold:
+                    scored_entities.append((entity, score))
+
+        else:
+            raise ValueError(
+                f"Unknown text search method: {method}. Use 'bm25', 'jaccard', 'cosine', or 'levenshtein'"
+            )
+
+        # Sort by score descending and return top results
+        scored_entities.sort(key=lambda x: x[1], reverse=True)
+        return scored_entities[:max_results]
+
+    async def find_paths(
+        self,
+        source_entity_id: str,
+        target_entity_id: str,
+        max_depth: int = 5,
+        max_paths: int = 10,
+    ) -> List:
+        """
+        Optimized path finding using networkx algorithms
+
+        Overrides default implementation to use networkx.all_simple_paths
+        for better performance.
+        """
+        from aiecs.domain.knowledge_graph.models.path import Path
+
+        if not self._initialized:
+            return []
+
+        if source_entity_id not in self.graph or target_entity_id not in self.graph:
+            return []
+
+        try:
+            # Use networkx's optimized path finding
+            paths = []
+            for node_path in nx.all_simple_paths(
+                self.graph,
+                source_entity_id,
+                target_entity_id,
+                cutoff=max_depth,
+            ):
+                # Convert node IDs to Entity and Relation objects
+                entities = [
+                    self.entities[node_id] for node_id in node_path if node_id in self.entities
+                ]
+
+                # Get relations between consecutive nodes
+                edges = []
+                for i in range(len(node_path) - 1):
+                    edge_data = self.graph.get_edge_data(node_path[i], node_path[i + 1])
+                    if edge_data and "relation" in edge_data:
+                        edges.append(edge_data["relation"])
+
+                if len(entities) == len(node_path):
+                    paths.append(Path(nodes=entities, edges=edges))
+
+                if len(paths) >= max_paths:
+                    break
+
+            return paths
+
+        except nx.NetworkXNoPath:
+            return []
+
+    # =========================================================================
+    # UTILITY METHODS
+    # =========================================================================
+
+    def get_stats(self) -> Dict[str, int]:
+        """
+        Get graph statistics
+
+        Returns:
+            Dictionary with node count, edge count, etc.
+        """
+        if not self._initialized:
+            return {"nodes": 0, "edges": 0, "entities": 0, "relations": 0}
+
+        return {
+            "nodes": self.graph.number_of_nodes(),
+            "edges": self.graph.number_of_edges(),
+            "entities": len(self.entities),
+            "relations": len(self.relations),
+        }
+
+    def clear(self) -> None:
+        """Clear all data from the graph"""
+        if self._initialized:
+            self.graph.clear()
+            self.entities.clear()
+            self.relations.clear()
+
+    def __str__(self) -> str:
+        stats = self.get_stats()
+        return f"InMemoryGraphStore(entities={stats['entities']}, relations={stats['relations']})"
+
+    def __repr__(self) -> str:
+        return self.__str__()