t6_mem0_v2/memory_cleanup.py

"""
Enhanced memory cleanup utilities for T6 Mem0 v2
Ensures synchronization between Supabase (vector) and Neo4j (graph) stores
"""

import logging
from typing import Optional
from neo4j import GraphDatabase
from mem0 import Memory
from config import settings

logger = logging.getLogger(__name__)


class MemoryCleanup:
    """Utilities for cleaning up memories across both vector and graph stores"""

    def __init__(self, memory: Memory):
        """
        Initialize cleanup utilities

        Args:
            memory: Mem0 Memory instance
        """
        self.memory = memory
        self.neo4j_driver = None

    def _get_neo4j_driver(self):
        """Get or create Neo4j driver"""
        if self.neo4j_driver is None:
            self.neo4j_driver = GraphDatabase.driver(
                settings.neo4j_uri,
                auth=(settings.neo4j_user, settings.neo4j_password)
            )
        return self.neo4j_driver

    def cleanup_neo4j_for_user(self, user_id: str) -> int:
        """
        Clean up Neo4j graph nodes for a specific user

        Args:
            user_id: User identifier

        Returns:
            Number of nodes deleted
        """
        try:
            driver = self._get_neo4j_driver()
            with driver.session() as session:
                # Delete all nodes with this user_id
                result = session.run(
                    "MATCH (n {user_id: $user_id}) DETACH DELETE n RETURN count(n) as deleted",
                    user_id=user_id
                )
                deleted = result.single()['deleted']
                logger.info(f"Deleted {deleted} Neo4j nodes for user_id={user_id}")
                return deleted
        except Exception as e:
            logger.error(f"Error cleaning up Neo4j for user {user_id}: {e}")
            raise

    def cleanup_neo4j_for_agent(self, agent_id: str) -> int:
        """
        Clean up Neo4j graph nodes for a specific agent

        Args:
            agent_id: Agent identifier

        Returns:
            Number of nodes deleted
        """
        try:
            driver = self._get_neo4j_driver()
            with driver.session() as session:
                # Delete all nodes with this agent_id
                result = session.run(
                    "MATCH (n {agent_id: $agent_id}) DETACH DELETE n RETURN count(n) as deleted",
                    agent_id=agent_id
                )
                deleted = result.single()['deleted']
                logger.info(f"Deleted {deleted} Neo4j nodes for agent_id={agent_id}")
                return deleted
        except Exception as e:
            logger.error(f"Error cleaning up Neo4j for agent {agent_id}: {e}")
            raise

    def cleanup_all_neo4j(self) -> dict:
        """
        Clean up ALL Neo4j graph data

        Returns:
            Dict with deleted counts
        """
        try:
            driver = self._get_neo4j_driver()
            with driver.session() as session:
                # Delete all relationships
                result = session.run("MATCH ()-[r]->() DELETE r RETURN count(r) as deleted")
                rels_deleted = result.single()['deleted']

                # Delete all nodes
                result = session.run("MATCH (n) DELETE n RETURN count(n) as deleted")
                nodes_deleted = result.single()['deleted']

                logger.info(f"Deleted {nodes_deleted} nodes and {rels_deleted} relationships from Neo4j")
                return {
                    "nodes_deleted": nodes_deleted,
                    "relationships_deleted": rels_deleted
                }
        except Exception as e:
            logger.error(f"Error cleaning up all Neo4j data: {e}")
            raise

    def delete_all_synchronized(
        self,
        user_id: Optional[str] = None,
        agent_id: Optional[str] = None,
        run_id: Optional[str] = None
    ) -> dict:
        """
        Delete all memories from BOTH Supabase and Neo4j

        This is the recommended method to ensure complete cleanup.

        Args:
            user_id: User identifier filter
            agent_id: Agent identifier filter
            run_id: Run identifier filter

        Returns:
            Dict with deletion statistics
        """
        logger.info(f"Synchronized delete_all: user_id={user_id}, agent_id={agent_id}, run_id={run_id}")

        # Step 1: Delete from vector store (Supabase) using mem0's method
        logger.info("Step 1: Deleting from vector store (Supabase)...")
        try:
            self.memory.delete_all(user_id=user_id, agent_id=agent_id, run_id=run_id)
            supabase_deleted = True
        except Exception as e:
            logger.error(f"Error deleting from Supabase: {e}")
            supabase_deleted = False

        # Step 2: Delete from graph store (Neo4j)
        logger.info("Step 2: Deleting from graph store (Neo4j)...")
        neo4j_deleted = 0
        try:
            if user_id:
                neo4j_deleted = self.cleanup_neo4j_for_user(user_id)
            elif agent_id:
                neo4j_deleted = self.cleanup_neo4j_for_agent(agent_id)
            else:
                # If no specific filter, clean up everything
                result = self.cleanup_all_neo4j()
                neo4j_deleted = result['nodes_deleted']
        except Exception as e:
            logger.error(f"Error deleting from Neo4j: {e}")

        result = {
            "supabase_success": supabase_deleted,
            "neo4j_nodes_deleted": neo4j_deleted,
            "synchronized": True
        }

        logger.info(f"Synchronized deletion complete: {result}")
        return result

    def close(self):
        """Close Neo4j driver connection"""
        if self.neo4j_driver:
            self.neo4j_driver.close()
            self.neo4j_driver = None

    def __del__(self):
        """Cleanup on deletion"""
        self.close()


# Convenience function for easy imports
def create_cleanup(memory: Memory) -> MemoryCleanup:
    """
    Create a MemoryCleanup instance

    Args:
        memory: Mem0 Memory instance

    Returns:
        MemoryCleanup instance
    """
    return MemoryCleanup(memory)