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

1from __future__ import annotations 

2 

3from typing import TYPE_CHECKING 

4from urllib.parse import urlparse 

5 

6if TYPE_CHECKING: 

7 from qdrant_loader.core.document import Document 

8 

9from qdrant_loader_core.graph.extractor.base_extractor import ( 

10 BaseEntityExtractor, 

11) 

12from qdrant_loader_core.graph.models import ( 

13 CoreEdgeType, 

14 CoreNodeLabel, 

15 GraphEdge, 

16 GraphNode, 

17) 

18 

19 

20class PublicDocsEntityExtractor(BaseEntityExtractor): 

21 """ 

22 Public documentation graph extractor. 

23 

24 Extracts: 

25 - Document node from page 

26 - Container node from website/domain 

27 - LINKS_TO edges between pages 

28 - Attachment nodes 

29 """ 

30 

31 source_type = "publicdocs" 

32 

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

34 # Project 

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

36 

37 def _project( 

38 self, 

39 doc: Document, 

40 ) -> str | None: 

41 url = getattr(doc, "url", None) or doc.metadata.get("url") 

42 

43 if not url: 

44 return None 

45 

46 parsed = urlparse(url) 

47 

48 return parsed.netloc 

49 

50 # ------------------------------------------------------------------ 

51 # Build Document Node 

52 # ------------------------------------------------------------------ 

53 def _build_document_node( 

54 self, 

55 doc: Document, 

56 project: str | None, 

57 ) -> GraphNode: 

58 return GraphNode( 

59 id=doc.id, 

60 label=CoreNodeLabel.DOCUMENT.value, 

61 project=project, 

62 properties={ 

63 "title": doc.title, 

64 "source_type": self.source_type, 

65 }, 

66 ) 

67 

68 # ------------------------------------------------------------------ 

69 # People 

70 # ------------------------------------------------------------------ 

71 

72 def _extract_people( 

73 self, 

74 doc: Document, 

75 ) -> list: 

76 return [] 

77 

78 # ------------------------------------------------------------------ 

79 # Container 

80 # ------------------------------------------------------------------ 

81 

82 def _extract_container( 

83 self, 

84 doc: Document, 

85 ) -> GraphNode | None: 

86 url = getattr(doc, "url", None) or doc.metadata.get("url") 

87 

88 if not url: 

89 return None 

90 

91 parsed = urlparse(url) 

92 

93 domain = parsed.netloc 

94 

95 if not domain: 

96 return None 

97 

98 return GraphNode( 

99 id=f"site:{domain}", 

100 label=CoreNodeLabel.CONTAINER.value, 

101 project=domain, 

102 properties={ 

103 "kind": "website", 

104 "domain": domain, 

105 }, 

106 ) 

107 

108 def _extract_labels( 

109 self, 

110 doc: Document, 

111 ) -> list[GraphNode]: 

112 project = self._project(doc) 

113 

114 return [ 

115 GraphNode( 

116 id=f"label:{tag}", 

117 label=CoreNodeLabel.LABEL.value, 

118 project=project, 

119 properties={"name": tag}, 

120 ) 

121 for tag in doc.metadata.get("tags", []) 

122 ] 

123 

124 # ------------------------------------------------------------------ 

125 # Links + Attachments 

126 # ------------------------------------------------------------------ 

127 

128 def _extract_source_specific( 

129 self, 

130 doc: Document, 

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

132 project = self._project(doc) 

133 

134 nodes: list[GraphNode] = [] 

135 edges: list[GraphEdge] = [] 

136 

137 # -------------------------------------------------------------- 

138 # Internal links 

139 # -------------------------------------------------------------- 

140 

141 for link in doc.metadata.get("links", []): 

142 edges.append( 

143 GraphEdge( 

144 source=doc.id, 

145 target=link, 

146 edge_type=CoreEdgeType.LINKS_TO.value, 

147 project=project, 

148 ) 

149 ) 

150 

151 # -------------------------------------------------------------- 

152 # Attachments 

153 # -------------------------------------------------------------- 

154 

155 for attachment in doc.metadata.get("attachments", []): 

156 attachment_id = attachment.get("id") 

157 

158 if not attachment_id: 

159 continue 

160 

161 nodes.append( 

162 GraphNode( 

163 id=f"attachment:{attachment_id}", 

164 label=CoreNodeLabel.ATTACHMENT.value, 

165 project=project, 

166 properties={ 

167 "filename": attachment.get("filename"), 

168 "mime_type": attachment.get("mime_type"), 

169 "download_url": attachment.get("download_url"), 

170 }, 

171 ) 

172 ) 

173 

174 edges.append( 

175 GraphEdge( 

176 source=doc.id, 

177 target=f"attachment:{attachment_id}", 

178 edge_type=CoreEdgeType.HAS_ATTACHMENT.value, 

179 project=project, 

180 ) 

181 ) 

182 

183 return nodes, edges