Coverage for src/qdrant_loader_core/graph/extractor/confluence.py: 100%

1from __future__ import annotations 

2 

3from typing import TYPE_CHECKING 

4 

5if TYPE_CHECKING: 

6 from qdrant_loader.core.document import Document 

7 

8from qdrant_loader_core.graph.extractor.base_extractor import ( 

9 BaseEntityExtractor, 

10) 

11from qdrant_loader_core.graph.models import ( 

12 CoreEdgeType, 

13 CoreNodeLabel, 

14 GraphEdge, 

15 GraphNode, 

16 PersonInfo, 

17) 

18 

19 

20class ConfluenceEntityExtractor(BaseEntityExtractor): 

21 """ 

22 Confluence graph extractor. 

23 

24 Extracts: 

25 - Document node from page 

26 - Person node from author 

27 - Container node from space 

28 - Parent / child page hierarchy 

29 """ 

30 

31 source_type = "confluence" 

32 

33 # ------------------------------------------------------------------ 

34 # Project / Space 

35 # ------------------------------------------------------------------ 

36 

37 def _project( 

38 self, 

39 doc: Document, 

40 ) -> str | None: 

41 return doc.metadata.get("space_key") 

42 

43 def _build_document_node( 

44 self, 

45 doc: Document, 

46 project: str | None, 

47 ) -> GraphNode: 

48 """Build a document node for a Confluence page.""" 

49 

50 return GraphNode( 

51 id=doc.id, 

52 label=CoreNodeLabel.DOCUMENT.value, 

53 project=project, 

54 properties={ 

55 "title": doc.title, 

56 "url": doc.url, 

57 "created_at": doc.metadata.get("created_at"), 

58 "updated_at": doc.metadata.get("updated_at"), 

59 "source_type": self.source_type, 

60 }, 

61 ) 

62 

63 # ------------------------------------------------------------------ 

64 # People 

65 # ------------------------------------------------------------------ 

66 

67 def _extract_people( 

68 self, 

69 doc: Document, 

70 ) -> list[tuple[PersonInfo, str]]: 

71 author = doc.metadata.get("author") 

72 

73 if not author: 

74 return [] 

75 

76 return [ 

77 ( 

78 PersonInfo( 

79 id=str(author), 

80 display_name=str(author), 

81 ), 

82 "author", 

83 ) 

84 ] 

85 

86 # ------------------------------------------------------------------ 

87 # Space container 

88 # ------------------------------------------------------------------ 

89 

90 def _extract_container( 

91 self, 

92 doc: Document, 

93 ) -> GraphNode | None: 

94 space_key = doc.metadata.get("space_key") 

95 

96 if not space_key: 

97 return None 

98 

99 return GraphNode( 

100 id=f"space:{space_key}", 

101 label=CoreNodeLabel.CONTAINER.value, 

102 project=space_key, 

103 properties={ 

104 "kind": "confluence_space", 

105 "space_key": space_key, 

106 }, 

107 ) 

108 

109 # ------------------------------------------------------------------ 

110 # Labels 

111 # ------------------------------------------------------------------ 

112 

113 def _extract_labels( 

114 self, 

115 doc: Document, 

116 ) -> list[GraphNode]: 

117 project = self._project(doc) 

118 

119 labels = doc.metadata.get("labels", []) 

120 

121 return [ 

122 GraphNode( 

123 id=f"label:{label}", 

124 label=CoreNodeLabel.LABEL.value, 

125 project=project, 

126 properties={ 

127 "name": label, 

128 }, 

129 ) 

130 for label in labels 

131 ] 

132 

133 # ------------------------------------------------------------------ 

134 # Hierarchy 

135 # ------------------------------------------------------------------ 

136 

137 def _extract_source_specific( 

138 self, 

139 doc: Document, 

140 ) -> tuple[list[GraphNode], list[GraphEdge]]: 

141 project = self._project(doc) 

142 

143 edges: list[GraphEdge] = [] 

144 

145 parent_id = doc.metadata.get("parent_id") 

146 

147 if parent_id: 

148 edges.append( 

149 GraphEdge( 

150 source=doc.id, 

151 target=str(parent_id), 

152 edge_type=CoreEdgeType.PART_OF.value, 

153 project=project, 

154 properties={ 

155 "kind": "page_child", 

156 }, 

157 ) 

158 ) 

159 

160 for child in doc.metadata.get("children", []): 

161 child_id = child.get("id") 

162 

163 if not child_id: 

164 continue 

165 

166 edges.append( 

167 GraphEdge( 

168 source=doc.id, 

169 target=str(child_id), 

170 edge_type=CoreEdgeType.HAS_CHILD.value, 

171 project=project, 

172 properties={ 

173 "kind": "page_child", 

174 }, 

175 ) 

176 ) 

177 

178 return [], edges