alma-memory 0.4.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

alma/graph/backends/neo4j.py

@@ -0,0 +1,417 @@
+"""
+ALMA Graph Memory - Neo4j Backend.
+
+Neo4j implementation of the GraphBackend interface.
+"""
+
+import json
+import logging
+from datetime import datetime, timezone
+from typing import Any, Dict, List, Optional
+
+from alma.graph.base import GraphBackend
+from alma.graph.store import Entity, Relationship
+
+logger = logging.getLogger(__name__)
+
+
+class Neo4jBackend(GraphBackend):
+    """
+    Neo4j graph database backend.
+
+    Requires neo4j Python driver: pip install neo4j
+
+    Example usage:
+        backend = Neo4jBackend(
+            uri="bolt://localhost:7687",
+            username="neo4j",
+            password="password"
+        )
+        backend.add_entity(entity)
+        backend.close()
+    """
+
+    def __init__(
+        self,
+        uri: str,
+        username: str,
+        password: str,
+        database: str = "neo4j",
+    ):
+        """
+        Initialize Neo4j connection.
+
+        Args:
+            uri: Neo4j connection URI (bolt:// or neo4j+s://)
+            username: Database username
+            password: Database password
+            database: Database name (default: "neo4j")
+        """
+        self.uri = uri
+        self.username = username
+        self.password = password
+        self.database = database
+        self._driver = None
+
+    def _get_driver(self):
+        """Lazy initialization of Neo4j driver."""
+        if self._driver is None:
+            try:
+                from neo4j import GraphDatabase
+
+                self._driver = GraphDatabase.driver(
+                    self.uri,
+                    auth=(self.username, self.password),
+                )
+            except ImportError as err:
+                raise ImportError(
+                    "neo4j package required for Neo4j graph backend. "
+                    "Install with: pip install neo4j"
+                ) from err
+        return self._driver
+
+    def _run_query(self, query: str, parameters: Optional[Dict] = None) -> List[Dict]:
+        """Execute a Cypher query."""
+        driver = self._get_driver()
+        with driver.session(database=self.database) as session:
+            result = session.run(query, parameters or {})
+            return [dict(record) for record in result]
+
+    def add_entity(self, entity: Entity) -> str:
+        """Add or update an entity in Neo4j."""
+        # Extract project_id and agent from properties if present
+        properties = entity.properties.copy()
+        project_id = properties.pop("project_id", None)
+        agent = properties.pop("agent", None)
+
+        query = """
+        MERGE (e:Entity {id: $id})
+        SET e.name = $name,
+            e.entity_type = $entity_type,
+            e.properties = $properties,
+            e.created_at = $created_at
+        """
+        params = {
+            "id": entity.id,
+            "name": entity.name,
+            "entity_type": entity.entity_type,
+            "properties": json.dumps(properties),
+            "created_at": entity.created_at.isoformat(),
+        }
+
+        # Add optional fields if present
+        if project_id:
+            query += ", e.project_id = $project_id"
+            params["project_id"] = project_id
+        if agent:
+            query += ", e.agent = $agent"
+            params["agent"] = agent
+
+        query += " RETURN e.id as id"
+
+        result = self._run_query(query, params)
+        return result[0]["id"] if result else entity.id
+
+    def add_relationship(self, relationship: Relationship) -> str:
+        """Add or update a relationship in Neo4j."""
+        # Sanitize relationship type for Cypher (remove special characters)
+        rel_type = (
+            relationship.relation_type.replace("-", "_").replace(" ", "_").upper()
+        )
+
+        query = f"""
+        MATCH (source:Entity {{id: $source_id}})
+        MATCH (target:Entity {{id: $target_id}})
+        MERGE (source)-[r:{rel_type}]->(target)
+        SET r.id = $id,
+            r.properties = $properties,
+            r.confidence = $confidence,
+            r.created_at = $created_at
+        RETURN r.id as id
+        """
+        result = self._run_query(
+            query,
+            {
+                "id": relationship.id,
+                "source_id": relationship.source_id,
+                "target_id": relationship.target_id,
+                "properties": json.dumps(relationship.properties),
+                "confidence": relationship.confidence,
+                "created_at": relationship.created_at.isoformat(),
+            },
+        )
+        return result[0]["id"] if result else relationship.id
+
+    def get_entity(self, entity_id: str) -> Optional[Entity]:
+        """Get an entity by ID."""
+        query = """
+        MATCH (e:Entity {id: $id})
+        RETURN e.id as id, e.name as name, e.entity_type as entity_type,
+               e.properties as properties, e.created_at as created_at,
+               e.project_id as project_id, e.agent as agent
+        """
+        result = self._run_query(query, {"id": entity_id})
+        if not result:
+            return None
+
+        r = result[0]
+        properties = json.loads(r["properties"]) if r["properties"] else {}
+
+        # Add project_id and agent back to properties if present
+        if r.get("project_id"):
+            properties["project_id"] = r["project_id"]
+        if r.get("agent"):
+            properties["agent"] = r["agent"]
+
+        return Entity(
+            id=r["id"],
+            name=r["name"],
+            entity_type=r["entity_type"],
+            properties=properties,
+            created_at=(
+                datetime.fromisoformat(r["created_at"])
+                if r["created_at"]
+                else datetime.now(timezone.utc)
+            ),
+        )
+
+    def get_entities(
+        self,
+        entity_type: Optional[str] = None,
+        project_id: Optional[str] = None,
+        agent: Optional[str] = None,
+        limit: int = 100,
+    ) -> List[Entity]:
+        """Get entities with optional filtering."""
+        conditions = []
+        params: Dict[str, Any] = {"limit": limit}
+
+        if entity_type:
+            conditions.append("e.entity_type = $entity_type")
+            params["entity_type"] = entity_type
+        if project_id:
+            conditions.append("e.project_id = $project_id")
+            params["project_id"] = project_id
+        if agent:
+            conditions.append("e.agent = $agent")
+            params["agent"] = agent
+
+        where_clause = f"WHERE {' AND '.join(conditions)}" if conditions else ""
+
+        query = f"""
+        MATCH (e:Entity)
+        {where_clause}
+        RETURN e.id as id, e.name as name, e.entity_type as entity_type,
+               e.properties as properties, e.created_at as created_at,
+               e.project_id as project_id, e.agent as agent
+        LIMIT $limit
+        """
+
+        results = self._run_query(query, params)
+        entities = []
+        for r in results:
+            properties = json.loads(r["properties"]) if r["properties"] else {}
+            if r.get("project_id"):
+                properties["project_id"] = r["project_id"]
+            if r.get("agent"):
+                properties["agent"] = r["agent"]
+
+            entities.append(
+                Entity(
+                    id=r["id"],
+                    name=r["name"],
+                    entity_type=r["entity_type"],
+                    properties=properties,
+                    created_at=(
+                        datetime.fromisoformat(r["created_at"])
+                        if r["created_at"]
+                        else datetime.now(timezone.utc)
+                    ),
+                )
+            )
+        return entities
+
+    def get_relationships(self, entity_id: str) -> List[Relationship]:
+        """Get all relationships for an entity (both directions)."""
+        query = """
+        MATCH (e:Entity {id: $entity_id})-[r]-(other:Entity)
+        RETURN r.id as id,
+               CASE WHEN startNode(r).id = $entity_id THEN e.id ELSE other.id END as source_id,
+               CASE WHEN endNode(r).id = $entity_id THEN e.id ELSE other.id END as target_id,
+               type(r) as relation_type, r.properties as properties,
+               r.confidence as confidence, r.created_at as created_at
+        """
+
+        results = self._run_query(query, {"entity_id": entity_id})
+        relationships = []
+        for r in results:
+            rel_id = (
+                r["id"] or f"{r['source_id']}-{r['relation_type']}-{r['target_id']}"
+            )
+            relationships.append(
+                Relationship(
+                    id=rel_id,
+                    source_id=r["source_id"],
+                    target_id=r["target_id"],
+                    relation_type=r["relation_type"],
+                    properties=json.loads(r["properties"]) if r["properties"] else {},
+                    confidence=r["confidence"] or 1.0,
+                    created_at=(
+                        datetime.fromisoformat(r["created_at"])
+                        if r["created_at"]
+                        else datetime.now(timezone.utc)
+                    ),
+                )
+            )
+        return relationships
+
+    def search_entities(
+        self,
+        query: str,
+        embedding: Optional[List[float]] = None,
+        top_k: int = 10,
+    ) -> List[Entity]:
+        """
+        Search for entities by name.
+
+        Note: Vector similarity search requires Neo4j 5.x with vector index.
+        Falls back to text search if embedding is provided but vector index
+        is not available.
+        """
+        # For now, we do text-based search
+        # Vector search can be added when Neo4j vector indexes are set up
+        cypher = """
+        MATCH (e:Entity)
+        WHERE toLower(e.name) CONTAINS toLower($query)
+        RETURN e.id as id, e.name as name, e.entity_type as entity_type,
+               e.properties as properties, e.created_at as created_at,
+               e.project_id as project_id, e.agent as agent
+        LIMIT $limit
+        """
+
+        results = self._run_query(cypher, {"query": query, "limit": top_k})
+        entities = []
+        for r in results:
+            properties = json.loads(r["properties"]) if r["properties"] else {}
+            if r.get("project_id"):
+                properties["project_id"] = r["project_id"]
+            if r.get("agent"):
+                properties["agent"] = r["agent"]
+
+            entities.append(
+                Entity(
+                    id=r["id"],
+                    name=r["name"],
+                    entity_type=r["entity_type"],
+                    properties=properties,
+                    created_at=(
+                        datetime.fromisoformat(r["created_at"])
+                        if r["created_at"]
+                        else datetime.now(timezone.utc)
+                    ),
+                )
+            )
+        return entities
+
+    def delete_entity(self, entity_id: str) -> bool:
+        """Delete an entity and its relationships."""
+        query = """
+        MATCH (e:Entity {id: $id})
+        DETACH DELETE e
+        RETURN count(e) as deleted
+        """
+        result = self._run_query(query, {"id": entity_id})
+        return result[0]["deleted"] > 0 if result else False
+
+    def delete_relationship(self, relationship_id: str) -> bool:
+        """Delete a specific relationship by ID."""
+        query = """
+        MATCH ()-[r]-()
+        WHERE r.id = $id
+        DELETE r
+        RETURN count(r) as deleted
+        """
+        result = self._run_query(query, {"id": relationship_id})
+        return result[0]["deleted"] > 0 if result else False
+
+    def close(self) -> None:
+        """Close the Neo4j driver connection."""
+        if self._driver:
+            self._driver.close()
+            self._driver = None
+
+    # Additional methods for compatibility with existing GraphStore API
+
+    def find_entities(
+        self,
+        name: Optional[str] = None,
+        entity_type: Optional[str] = None,
+        limit: int = 10,
+    ) -> List[Entity]:
+        """
+        Find entities by name or type.
+
+        This method provides compatibility with the existing GraphStore API.
+        """
+        if name:
+            return self.search_entities(query=name, top_k=limit)
+
+        return self.get_entities(entity_type=entity_type, limit=limit)
+
+    def get_relationships_directional(
+        self,
+        entity_id: str,
+        direction: str = "both",
+        relation_type: Optional[str] = None,
+    ) -> List[Relationship]:
+        """
+        Get relationships for an entity with direction control.
+
+        This method provides compatibility with the existing GraphStore API.
+
+        Args:
+            entity_id: The entity ID.
+            direction: "outgoing", "incoming", or "both".
+            relation_type: Optional filter by relationship type.
+
+        Returns:
+            List of matching relationships.
+        """
+        if direction == "outgoing":
+            pattern = "(e)-[r]->(other)"
+        elif direction == "incoming":
+            pattern = "(e)<-[r]-(other)"
+        else:
+            pattern = "(e)-[r]-(other)"
+
+        type_filter = f":{relation_type}" if relation_type else ""
+
+        query = f"""
+        MATCH (e:Entity {{id: $entity_id}}){pattern.replace("[r]", f"[r{type_filter}]")}
+        RETURN r.id as id, e.id as source_id, other.id as target_id,
+               type(r) as relation_type, r.properties as properties,
+               r.confidence as confidence, r.created_at as created_at
+        """
+
+        results = self._run_query(query, {"entity_id": entity_id})
+        relationships = []
+        for r in results:
+            rel_id = (
+                r["id"] or f"{r['source_id']}-{r['relation_type']}-{r['target_id']}"
+            )
+            relationships.append(
+                Relationship(
+                    id=rel_id,
+                    source_id=r["source_id"],
+                    target_id=r["target_id"],
+                    relation_type=r["relation_type"],
+                    properties=json.loads(r["properties"]) if r["properties"] else {},
+                    confidence=r["confidence"] or 1.0,
+                    created_at=(
+                        datetime.fromisoformat(r["created_at"])
+                        if r["created_at"]
+                        else datetime.now(timezone.utc)
+                    ),
+                )
+            )
+        return relationships

alma/graph/base.py

@@ -0,0 +1,159 @@
+"""
+ALMA Graph Memory - Abstract Backend Interface.
+
+Defines the abstract base class for graph database backends,
+enabling pluggable graph storage implementations.
+"""
+
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+from alma.graph.store import Entity, Relationship
+
+
+class GraphBackend(ABC):
+    """
+    Abstract base class for graph database backends.
+
+    This interface defines the contract that all graph storage implementations
+    must follow, enabling ALMA to work with different graph databases like
+    Neo4j, Amazon Neptune, or in-memory stores.
+    """
+
+    @abstractmethod
+    def add_entity(self, entity: Entity) -> str:
+        """
+        Add or update an entity in the graph.
+
+        Args:
+            entity: The entity to add or update.
+
+        Returns:
+            The entity ID.
+        """
+        pass
+
+    @abstractmethod
+    def add_relationship(self, relationship: Relationship) -> str:
+        """
+        Add or update a relationship between entities.
+
+        Args:
+            relationship: The relationship to add or update.
+
+        Returns:
+            The relationship ID.
+        """
+        pass
+
+    @abstractmethod
+    def get_entity(self, entity_id: str) -> Optional[Entity]:
+        """
+        Get an entity by its ID.
+
+        Args:
+            entity_id: The unique identifier of the entity.
+
+        Returns:
+            The entity if found, None otherwise.
+        """
+        pass
+
+    @abstractmethod
+    def get_entities(
+        self,
+        entity_type: Optional[str] = None,
+        project_id: Optional[str] = None,
+        agent: Optional[str] = None,
+        limit: int = 100,
+    ) -> List[Entity]:
+        """
+        Get entities with optional filtering.
+
+        Args:
+            entity_type: Filter by entity type (person, organization, etc.).
+            project_id: Filter by project ID.
+            agent: Filter by agent name.
+            limit: Maximum number of entities to return.
+
+        Returns:
+            List of matching entities.
+        """
+        pass
+
+    @abstractmethod
+    def get_relationships(self, entity_id: str) -> List[Relationship]:
+        """
+        Get all relationships for an entity.
+
+        Args:
+            entity_id: The entity ID to get relationships for.
+
+        Returns:
+            List of relationships where the entity is source or target.
+        """
+        pass
+
+    @abstractmethod
+    def search_entities(
+        self,
+        query: str,
+        embedding: Optional[List[float]] = None,
+        top_k: int = 10,
+    ) -> List[Entity]:
+        """
+        Search for entities by name or using vector similarity.
+
+        Args:
+            query: Text query to search for in entity names.
+            embedding: Optional embedding vector for semantic search.
+            top_k: Maximum number of results to return.
+
+        Returns:
+            List of matching entities, ordered by relevance.
+        """
+        pass
+
+    @abstractmethod
+    def delete_entity(self, entity_id: str) -> bool:
+        """
+        Delete an entity and all its relationships.
+
+        Args:
+            entity_id: The entity ID to delete.
+
+        Returns:
+            True if the entity was deleted, False if not found.
+        """
+        pass
+
+    @abstractmethod
+    def delete_relationship(self, relationship_id: str) -> bool:
+        """
+        Delete a specific relationship.
+
+        Args:
+            relationship_id: The relationship ID to delete.
+
+        Returns:
+            True if the relationship was deleted, False if not found.
+        """
+        pass
+
+    @abstractmethod
+    def close(self) -> None:
+        """
+        Close the backend connection and release resources.
+
+        Should be called when the backend is no longer needed.
+        """
+        pass
+
+    def __enter__(self):
+        """Context manager entry."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit - ensures close is called."""
+        self.close()
+        return False