Coverage for src/qdrant_loader_core/graph/extractor/base_extractor.py: 99%

1from __future__ import annotations 

2 

3from typing import TYPE_CHECKING 

4 

5if TYPE_CHECKING: 

6 from qdrant_loader.core.document import Document 

7 

8from abc import ABC, abstractmethod 

9from typing import ClassVar 

10 

11from ..models import ( 

12 CoreEdgeType, 

13 CoreNodeLabel, 

14 GraphEdge, 

15 GraphNode, 

16 PersonInfo, 

17 SubGraph, 

18) 

19 

20 

21# ------------------------ 

22# EntityExtractor Interface 

23# ------------------------ 

24class EntityExtractor(ABC): 

25 """ 

26 Interface for all entity extractors. 

27 """ 

28 

29 _registry: dict[str, type[EntityExtractor]] = {} 

30 

31 @abstractmethod 

32 async def extract(self, doc: Document) -> SubGraph: 

33 """ 

34 Map raw source data into a SubGraph 

35 """ 

36 raise NotImplementedError 

37 

38 # -------- Registry -------- 

39 @classmethod 

40 def register_extractor(cls, source_type: str, extractor_cls: type[EntityExtractor]): 

41 if source_type in cls._registry: 

42 raise ValueError( 

43 f"Extractor already registered for source_type={source_type}" 

44 ) 

45 cls._registry[source_type] = extractor_cls 

46 

47 @classmethod 

48 def for_source(cls, source_type: str) -> EntityExtractor: 

49 if source_type not in cls._registry: 

50 raise ValueError(f"No extractor registered for source_type={source_type}") 

51 return cls._registry[source_type]() # instantiate 

52 

53 

54class BaseEntityExtractor(EntityExtractor): 

55 """ 

56 Base extractor implementing shared graph extraction. 

57 

58 Common nodes: 

59 - Document 

60 - Person 

61 - Container 

62 - Label 

63 

64 Source-specific nodes are extracted via 

65 _extract_source_specific(). 

66 """ 

67 

68 source_type: ClassVar[str] 

69 

70 @classmethod 

71 def get_source_type(cls) -> str: 

72 return cls.source_type 

73 

74 async def extract(self, doc: Document) -> SubGraph: 

75 project = self._project(doc) 

76 

77 nodes: list[GraphNode] = [] 

78 edges: list[GraphEdge] = [] 

79 

80 # 1. Document 

81 doc_node = self._build_document_node(doc, project) 

82 nodes.append(doc_node) 

83 

84 # 2. People 

85 people = self._extract_people(doc) 

86 

87 # Deduplicate people by canonical person id 

88 seen_people: set[str] = set() 

89 

90 for person_info, role in people: 

91 canonical_id = person_info.id 

92 

93 if canonical_id in seen_people: 

94 continue 

95 

96 seen_people.add(canonical_id) 

97 

98 person_node = self._person_node(person_info, project) 

99 nodes.append(person_node) 

100 

101 edges.append( 

102 GraphEdge( 

103 source=( 

104 doc.metadata.get("key") 

105 if self.get_source_type() == "jira" 

106 else doc.id 

107 ), 

108 target=person_node.id, 

109 edge_type=CoreEdgeType.AUTHORED_BY.value, 

110 project=project, 

111 properties={"role": role}, 

112 ) 

113 ) 

114 

115 # 3. Container (project / space / repo) 

116 container = self._extract_container(doc) 

117 

118 if container: 

119 nodes.append(container) 

120 

121 edges.append( 

122 GraphEdge( 

123 source=( 

124 doc.metadata.get("key") 

125 if self.get_source_type() == "jira" 

126 else doc.id 

127 ), 

128 target=container.id, 

129 edge_type=CoreEdgeType.BELONGS_TO.value, 

130 project=project, 

131 ) 

132 ) 

133 

134 # 4. Labels 

135 labels = self._extract_labels(doc) 

136 

137 for label in labels: 

138 nodes.append(label) 

139 

140 edges.append( 

141 GraphEdge( 

142 source=( 

143 doc.metadata.get("key") 

144 if self.get_source_type() == "jira" 

145 else doc.id 

146 ), 

147 target=label.id, 

148 edge_type=CoreEdgeType.HAS_LABEL.value, 

149 project=project, 

150 ) 

151 ) 

152 

153 # 5. Cross references 

154 link_nodes, link_edges = self._extract_links(doc) 

155 

156 nodes.extend(link_nodes) 

157 edges.extend(link_edges) 

158 

159 # 6. Source specific 

160 custom_nodes, custom_edges = self._extract_source_specific(doc) 

161 

162 nodes.extend(custom_nodes) 

163 edges.extend(custom_edges) 

164 

165 return SubGraph( 

166 nodes=nodes, 

167 edges=edges, 

168 ) 

169 

170 def _person_node( 

171 self, 

172 person_info: PersonInfo, 

173 project: str | None, 

174 ) -> GraphNode: 

175 return GraphNode( 

176 id=person_info.id, 

177 label=CoreNodeLabel.PERSON.value, 

178 project=project, 

179 properties={ 

180 "display_name": person_info.display_name, 

181 }, 

182 ) 

183 

184 # ------------------------------------------------------------------ 

185 # Required hooks 

186 # ------------------------------------------------------------------ 

187 

188 @abstractmethod 

189 def _build_document_node( 

190 self, 

191 doc: Document, 

192 project: str | None, 

193 ) -> GraphNode: 

194 """ 

195 Build the main document node. 

196 """ 

197 

198 @abstractmethod 

199 def _project(self, doc: Document) -> str | None: 

200 """ 

201 Return project / space / repository identifier. 

202 """ 

203 

204 @abstractmethod 

205 def _extract_people( 

206 self, 

207 doc: Document, 

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

209 """ 

210 Return [(person_info, role), ...] 

211 """ 

212 

213 @abstractmethod 

214 def _extract_container( 

215 self, 

216 doc: Document, 

217 ) -> GraphNode | None: 

218 """ 

219 Return container node (project / space / repository) 

220 """ 

221 

222 # ------------------------------------------------------------------ 

223 # Optional hooks 

224 # ------------------------------------------------------------------ 

225 

226 def _extract_labels( 

227 self, 

228 doc: Document, 

229 ) -> list[GraphNode]: 

230 return [] 

231 

232 def _extract_links( 

233 self, 

234 doc: Document, 

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

236 return [], [] 

237 

238 def _extract_source_specific( 

239 self, 

240 doc: Document, 

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

242 return [], []