图数据库数据准备模块 - 从Neo4j读取数据并转换为文档
| 28 | properties: Dict[str, Any] |
| 29 | |
| 30 | class GraphDataPreparationModule: |
| 31 | """图数据库数据准备模块 - 从Neo4j读取数据并转换为文档""" |
| 32 | |
| 33 | def __init__(self, uri: str, user: str, password: str, database: str = "neo4j"): |
| 34 | """ |
| 35 | 初始化图数据库连接 |
| 36 | |
| 37 | Args: |
| 38 | uri: Neo4j连接URI |
| 39 | user: 用户名 |
| 40 | password: 密码 |
| 41 | database: 数据库名称 |
| 42 | """ |
| 43 | self.uri = uri |
| 44 | self.user = user |
| 45 | self.password = password |
| 46 | self.database = database |
| 47 | self.driver = None |
| 48 | self.documents: List[Document] = [] |
| 49 | self.chunks: List[Document] = [] |
| 50 | self.diseases: List[GraphNode] = [] |
| 51 | self.symptoms: List[GraphNode] = [] |
| 52 | self.treatments: List[GraphNode] = [] |
| 53 | self.medications: List[GraphNode] = [] |
| 54 | self.risk_factors: List[GraphNode] = [] |
| 55 | self.care_tips: List[GraphNode] = [] |
| 56 | |
| 57 | self._connect() |
| 58 | |
| 59 | def _connect(self): |
| 60 | """建立Neo4j连接""" |
| 61 | try: |
| 62 | self.driver = GraphDatabase.driver( |
| 63 | self.uri, |
| 64 | auth=(self.user, self.password), |
| 65 | database=self.database |
| 66 | ) |
| 67 | logger.info(f"已连接到Neo4j数据库: {self.uri}") |
| 68 | |
| 69 | # 测试连接 |
| 70 | with self.driver.session() as session: |
| 71 | result = session.run("RETURN 1 as test") |
| 72 | test_result = result.single() |
| 73 | if test_result: |
| 74 | logger.info("Neo4j连接测试成功") |
| 75 | |
| 76 | except Exception as e: |
| 77 | logger.error(f"连接Neo4j失败: {e}") |
| 78 | raise |
| 79 | |
| 80 | def close(self): |
| 81 | """关闭数据库连接""" |
| 82 | if hasattr(self, 'driver') and self.driver: |
| 83 | self.driver.close() |
| 84 | logger.info("Neo4j连接已关闭") |
| 85 | |
| 86 | def load_graph_data(self) -> Dict[str, Any]: |
| 87 | """ |