LQAgentPlatform/foundation/database/models/kg/neo4j_models.py

"""
Neo4j图数据库模型定义

提供Neo4j图数据库相关的数据结构定义
"""

from typing import Optional, Dict, Any, List
from dataclasses import dataclass
from datetime import datetime


@dataclass
class Neo4jNode:
    """Neo4j节点数据模型"""
    id: Optional[int] = None
    labels: List[str] = None
    properties: Optional[Dict[str, Any]] = None
    created_at: Optional[datetime] = None

    def __post_init__(self):
        if self.labels is None:
            self.labels = []
        if self.properties is None:
            self.properties = {}

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            'id': self.id,
            'labels': self.labels,
            'properties': self.properties,
            'created_at': self.created_at.isoformat() if self.created_at else None
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Neo4jNode':
        """从字典创建实例"""
        return cls(
            id=data.get('id'),
            labels=data.get('labels', []),
            properties=data.get('properties', {}),
            created_at=datetime.fromisoformat(data['created_at']) if data.get('created_at') else None
        )


@dataclass
class Neo4jRelationship:
    """Neo4j关系数据模型"""
    id: Optional[int] = None
    type: str = ""
    start_node_id: Optional[int] = None
    end_node_id: Optional[int] = None
    properties: Optional[Dict[str, Any]] = None
    created_at: Optional[datetime] = None

    def __post_init__(self):
        if self.properties is None:
            self.properties = {}

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            'id': self.id,
            'type': self.type,
            'start_node_id': self.start_node_id,
            'end_node_id': self.end_node_id,
            'properties': self.properties,
            'created_at': self.created_at.isoformat() if self.created_at else None
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Neo4jRelationship':
        """从字典创建实例"""
        return cls(
            id=data.get('id'),
            type=data.get('type', ''),
            start_node_id=data.get('start_node_id'),
            end_node_id=data.get('end_node_id'),
            properties=data.get('properties', {}),
            created_at=datetime.fromisoformat(data['created_at']) if data.get('created_at') else None
        )


@dataclass
class Neo4jGraph:
    """Neo4j图数据模型"""
    nodes: List[Neo4jNode] = None
    relationships: List[Neo4jRelationship] = None
    metadata: Optional[Dict[str, Any]] = None

    def __post_init__(self):
        if self.nodes is None:
            self.nodes = []
        if self.relationships is None:
            self.relationships = []
        if self.metadata is None:
            self.metadata = {}

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典"""
        return {
            'nodes': [node.to_dict() for node in self.nodes],
            'relationships': [rel.to_dict() for rel in self.relationships],
            'metadata': self.metadata
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Neo4jGraph':
        """从字典创建实例"""
        nodes_data = data.get('nodes', [])
        relationships_data = data.get('relationships', [])

        nodes = [Neo4jNode.from_dict(node_data) for node_data in nodes_data]
        relationships = [Neo4jRelationship.from_dict(rel_data) for rel_data in relationships_data]

        return cls(
            nodes=nodes,
            relationships=relationships,
            metadata=data.get('metadata', {})
        )


__all__ = [
    "Neo4jNode",
    "Neo4jRelationship",
    "Neo4jGraph"
]