hello-agents/Co-creation-projects/YYHDBL-HelloCodeAgentCli/memory/storage/neo4j_store.py

"""
Neo4j图数据库存储实现
"""

import logging
from typing import Dict, List, Optional, Any, Tuple
from datetime import datetime

try:
    from neo4j import GraphDatabase
    from neo4j.exceptions import ServiceUnavailable, AuthError
    NEO4J_AVAILABLE = True
except ImportError:
    NEO4J_AVAILABLE = False
    GraphDatabase = None

logger = logging.getLogger(__name__)

class Neo4jGraphStore:
    """Neo4j图数据库存储实现"""
    
    def __init__(
        self,
        uri: str = "bolt://localhost:7687",
        username: str = "neo4j", 
        password: str = "hello-agents-password",
        database: str = "neo4j",
        max_connection_lifetime: int = 3600,
        max_connection_pool_size: int = 50,
        connection_acquisition_timeout: int = 60,
        **kwargs
    ):
        """
        初始化Neo4j图存储 (支持云API)
        
        Args:
            uri: Neo4j连接URI (本地: bolt://localhost:7687, 云: neo4j+s://xxx.databases.neo4j.io)
            username: 用户名
            password: 密码
            database: 数据库名称
            max_connection_lifetime: 最大连接生命周期(秒)
            max_connection_pool_size: 最大连接池大小
            connection_acquisition_timeout: 连接获取超时(秒)
        """
        if not NEO4J_AVAILABLE:
            raise ImportError(
                "neo4j未安装。请运行: pip install neo4j>=5.0.0"
            )
        
        self.uri = uri
        self.username = username
        self.password = password
        self.database = database
        
        # 初始化驱动
        self.driver = None
        self._initialize_driver(
            max_connection_lifetime=max_connection_lifetime,
            max_connection_pool_size=max_connection_pool_size,
            connection_acquisition_timeout=connection_acquisition_timeout
        )
        
        # 创建索引
        self._create_indexes()
    
    def _initialize_driver(self, **config):
        """初始化Neo4j驱动"""
        try:
            self.driver = GraphDatabase.driver(
                self.uri,
                auth=(self.username, self.password),
                **config
            )
            
            # 验证连接
            self.driver.verify_connectivity()
            
            # 检查是否是云服务
            if "neo4j.io" in self.uri or "aura" in self.uri.lower():
                logger.info(f"✅ 成功连接到Neo4j云服务: {self.uri}")
            else:
                logger.info(f"✅ 成功连接到Neo4j服务: {self.uri}")
                
        except AuthError as e:
            logger.error(f"❌ Neo4j认证失败: {e}")
            logger.info("💡 请检查用户名和密码是否正确")
            raise
        except ServiceUnavailable as e:
            logger.error(f"❌ Neo4j服务不可用: {e}")
            if "localhost" in self.uri:
                logger.info("💡 本地连接失败，可以考虑使用Neo4j Aura云服务")
                logger.info("💡 或启动本地服务: docker run -p 7474:7474 -p 7687:7687 neo4j:5.14")
            else:
                logger.info("💡 请检查URL和网络连接")
            raise
        except Exception as e:
            logger.error(f"❌ Neo4j连接失败: {e}")
            raise
    
    def _create_indexes(self):
        """创建必要的索引以提高查询性能"""
        indexes = [
            # 实体索引
            "CREATE INDEX entity_id_index IF NOT EXISTS FOR (e:Entity) ON (e.id)",
            "CREATE INDEX entity_name_index IF NOT EXISTS FOR (e:Entity) ON (e.name)",
            "CREATE INDEX entity_type_index IF NOT EXISTS FOR (e:Entity) ON (e.type)",
            
            # 记忆索引
            "CREATE INDEX memory_id_index IF NOT EXISTS FOR (m:Memory) ON (m.id)",
            "CREATE INDEX memory_type_index IF NOT EXISTS FOR (m:Memory) ON (m.memory_type)",
            "CREATE INDEX memory_timestamp_index IF NOT EXISTS FOR (m:Memory) ON (m.timestamp)",
        ]
        
        with self.driver.session(database=self.database) as session:
            for index_query in indexes:
                try:
                    session.run(index_query)
                except Exception as e:
                    logger.debug(f"索引创建跳过 (可能已存在): {e}")
        
        logger.info("✅ Neo4j索引创建完成")
    
    def add_entity(self, entity_id: str, name: str, entity_type: str, properties: Dict[str, Any] = None) -> bool:
        """
        添加实体节点
        
        Args:
            entity_id: 实体ID
            name: 实体名称
            entity_type: 实体类型
            properties: 附加属性
        
        Returns:
            bool: 是否成功
        """
        try:
            props = properties or {}
            props.update({
                "id": entity_id,
                "name": name,
                "type": entity_type,
                "created_at": datetime.now().isoformat(),
                "updated_at": datetime.now().isoformat()
            })
            
            query = """
            MERGE (e:Entity {id: $entity_id})
            SET e += $properties
            RETURN e
            """
            
            with self.driver.session(database=self.database) as session:
                result = session.run(query, entity_id=entity_id, properties=props)
                record = result.single()
                
                if record:
                    logger.debug(f"✅ 添加实体: {name} ({entity_type})")
                    return True
                return False
                
        except Exception as e:
            logger.error(f"❌ 添加实体失败: {e}")
            return False
    
    def add_relationship(
        self, 
        from_entity_id: str, 
        to_entity_id: str, 
        relationship_type: str,
        properties: Dict[str, Any] = None
    ) -> bool:
        """
        添加实体间关系
        
        Args:
            from_entity_id: 源实体ID
            to_entity_id: 目标实体ID  
            relationship_type: 关系类型
            properties: 关系属性
        
        Returns:
            bool: 是否成功
        """
        try:
            props = properties or {}
            props.update({
                "type": relationship_type,
                "created_at": datetime.now().isoformat(),
                "updated_at": datetime.now().isoformat()
            })
            
            query = f"""
            MATCH (from:Entity {{id: $from_id}})
            MATCH (to:Entity {{id: $to_id}})
            MERGE (from)-[r:{relationship_type}]->(to)
            SET r += $properties
            RETURN r
            """
            
            with self.driver.session(database=self.database) as session:
                result = session.run(
                    query,
                    from_id=from_entity_id,
                    to_id=to_entity_id,
                    properties=props
                )
                record = result.single()
                
                if record:
                    logger.debug(f"✅ 添加关系: {from_entity_id} -{relationship_type}-> {to_entity_id}")
                    return True
                return False
                
        except Exception as e:
            logger.error(f"❌ 添加关系失败: {e}")
            return False
    
    def find_related_entities(
        self, 
        entity_id: str, 
        relationship_types: List[str] = None,
        max_depth: int = 2,
        limit: int = 50
    ) -> List[Dict[str, Any]]:
        """
        查找相关实体
        
        Args:
            entity_id: 起始实体ID
            relationship_types: 关系类型过滤
            max_depth: 最大搜索深度
            limit: 结果限制
        
        Returns:
            List[Dict]: 相关实体列表
        """
        try:
            # 构建关系类型过滤
            rel_filter = ""
            if relationship_types:
                rel_types = "|".join(relationship_types)
                rel_filter = f":{rel_types}"
            
            query = f"""
            MATCH path = (start:Entity {{id: $entity_id}})-[r{rel_filter}*1..{max_depth}]-(related:Entity)
            WHERE start.id <> related.id
            RETURN DISTINCT related, 
                   length(path) as distance,
                   [rel in relationships(path) | type(rel)] as relationship_path
            ORDER BY distance, related.name
            LIMIT $limit
            """
            
            with self.driver.session(database=self.database) as session:
                result = session.run(query, entity_id=entity_id, limit=limit)
                
                entities = []
                for record in result:
                    entity_data = dict(record["related"])
                    entity_data["distance"] = record["distance"]
                    entity_data["relationship_path"] = record["relationship_path"]
                    entities.append(entity_data)
                
                logger.debug(f"🔍 找到 {len(entities)} 个相关实体")
                return entities
                
        except Exception as e:
            logger.error(f"❌ 查找相关实体失败: {e}")
            return []
    
    def search_entities_by_name(self, name_pattern: str, entity_types: List[str] = None, limit: int = 20) -> List[Dict[str, Any]]:
        """
        按名称搜索实体
        
        Args:
            name_pattern: 名称模式 (支持部分匹配)
            entity_types: 实体类型过滤
            limit: 结果限制
        
        Returns:
            List[Dict]: 匹配的实体列表
        """
        try:
            # 构建类型过滤
            type_filter = ""
            params = {"pattern": f".*{name_pattern}.*", "limit": limit}
            
            if entity_types:
                type_filter = "AND e.type IN $types"
                params["types"] = entity_types
            
            query = f"""
            MATCH (e:Entity)
            WHERE e.name =~ $pattern {type_filter}
            RETURN e
            ORDER BY e.name
            LIMIT $limit
            """
            
            with self.driver.session(database=self.database) as session:
                result = session.run(query, **params)
                
                entities = []
                for record in result:
                    entity_data = dict(record["e"])
                    entities.append(entity_data)
                
                logger.debug(f"🔍 按名称搜索到 {len(entities)} 个实体")
                return entities
                
        except Exception as e:
            logger.error(f"❌ 按名称搜索实体失败: {e}")
            return []
    
    def get_entity_relationships(self, entity_id: str) -> List[Dict[str, Any]]:
        """
        获取实体的所有关系
        
        Args:
            entity_id: 实体ID
        
        Returns:
            List[Dict]: 关系列表
        """
        try:
            query = """
            MATCH (e:Entity {id: $entity_id})-[r]-(other:Entity)
            RETURN r, other, 
                   CASE WHEN startNode(r).id = $entity_id THEN 'outgoing' ELSE 'incoming' END as direction
            """
            
            with self.driver.session(database=self.database) as session:
                result = session.run(query, entity_id=entity_id)
                
                relationships = []
                for record in result:
                    rel_data = dict(record["r"])
                    other_data = dict(record["other"])
                    
                    relationship = {
                        "relationship": rel_data,
                        "other_entity": other_data,
                        "direction": record["direction"]
                    }
                    relationships.append(relationship)
                
                return relationships
                
        except Exception as e:
            logger.error(f"❌ 获取实体关系失败: {e}")
            return []
    
    def delete_entity(self, entity_id: str) -> bool:
        """
        删除实体及其所有关系
        
        Args:
            entity_id: 实体ID
        
        Returns:
            bool: 是否成功
        """
        try:
            query = """
            MATCH (e:Entity {id: $entity_id})
            DETACH DELETE e
            """
            
            with self.driver.session(database=self.database) as session:
                result = session.run(query, entity_id=entity_id)
                summary = result.consume()
                
                deleted_count = summary.counters.nodes_deleted
                logger.info(f"✅ 删除实体: {entity_id} (删除 {deleted_count} 个节点)")
                return deleted_count > 0
                
        except Exception as e:
            logger.error(f"❌ 删除实体失败: {e}")
            return False
    
    def clear_all(self) -> bool:
        """
        清空所有数据
        
        Returns:
            bool: 是否成功
        """
        try:
            query = "MATCH (n) DETACH DELETE n"
            
            with self.driver.session(database=self.database) as session:
                result = session.run(query)
                summary = result.consume()
                
                deleted_nodes = summary.counters.nodes_deleted
                deleted_relationships = summary.counters.relationships_deleted
                
                logger.info(f"✅ 清空Neo4j数据库: 删除 {deleted_nodes} 个节点, {deleted_relationships} 个关系")
                return True
                
        except Exception as e:
            logger.error(f"❌ 清空数据库失败: {e}")
            return False
    
    def get_stats(self) -> Dict[str, Any]:
        """
        获取图数据库统计信息
        
        Returns:
            Dict: 统计信息
        """
        try:
            queries = {
                "total_nodes": "MATCH (n) RETURN count(n) as count",
                "total_relationships": "MATCH ()-[r]->() RETURN count(r) as count",
                "entity_nodes": "MATCH (n:Entity) RETURN count(n) as count",
                "memory_nodes": "MATCH (n:Memory) RETURN count(n) as count",
            }
            
            stats = {}
            with self.driver.session(database=self.database) as session:
                for key, query in queries.items():
                    result = session.run(query)
                    record = result.single()
                    stats[key] = record["count"] if record else 0
            
            return stats
            
        except Exception as e:
            logger.error(f"❌ 获取统计信息失败: {e}")
            return {}
    
    def health_check(self) -> bool:
        """
        健康检查
        
        Returns:
            bool: 服务是否健康
        """
        try:
            with self.driver.session(database=self.database) as session:
                result = session.run("RETURN 1 as health")
                record = result.single()
                return record["health"] == 1
        except Exception as e:
            logger.error(f"❌ Neo4j健康检查失败: {e}")
            return False
    
    def __del__(self):
        """析构函数，清理资源"""
        if hasattr(self, 'driver') and self.driver:
            try:
                self.driver.close()
            except:
                pass