插入或更新边
(
self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
)
| 367 | # self._graph.add_node(node_id, **node_data) |
| 368 | |
| 369 | async def upsert_edge( |
| 370 | self, source_node_id: str, target_node_id: str, edge_data: dict[str, str] |
| 371 | ): |
| 372 | """插入或更新边""" |
| 373 | # print("go into upsert edge method") |
| 374 | source_name = source_node_id |
| 375 | target_name = target_node_id |
| 376 | weight = edge_data["weight"] |
| 377 | keywords = edge_data["keywords"] |
| 378 | description = edge_data["description"] |
| 379 | source_chunk_id = edge_data["source_id"] |
| 380 | logger.debug( |
| 381 | f"source_name:{source_name}, target_name:{target_name}, keywords: {keywords}" |
| 382 | ) |
| 383 | |
| 384 | content = keywords + source_name + target_name + description |
| 385 | contents = [content] |
| 386 | batches = [ |
| 387 | contents[i : i + self._max_batch_size] |
| 388 | for i in range(0, len(contents), self._max_batch_size) |
| 389 | ] |
| 390 | embeddings_list = await asyncio.gather( |
| 391 | *[self.embedding_func(batch) for batch in batches] |
| 392 | ) |
| 393 | embeddings = np.concatenate(embeddings_list) |
| 394 | content_vector = embeddings[0] |
| 395 | merge_sql = SQL_TEMPLATES["merge_edge"] |
| 396 | data = { |
| 397 | "workspace": self.db.workspace, |
| 398 | "source_name": source_name, |
| 399 | "target_name": target_name, |
| 400 | "weight": weight, |
| 401 | "keywords": keywords, |
| 402 | "description": description, |
| 403 | "source_chunk_id": source_chunk_id, |
| 404 | "content": content, |
| 405 | "content_vector": content_vector, |
| 406 | } |
| 407 | # print(merge_sql) |
| 408 | await self.db.execute(merge_sql, data) |
| 409 | # self._graph.add_edge(source_node_id, target_node_id, **edge_data) |
| 410 | |
| 411 | async def embed_nodes(self, algorithm: str) -> tuple[np.ndarray, list[str]]: |
| 412 | """为节点生成向量""" |