Coverage for src / local_deep_research / advanced_search_system / findings / topic.py: 89%

1from typing import List, Dict, Optional, Any 

2from dataclasses import dataclass, field 

3from datetime import datetime 

4 

5 

6@dataclass 

7class Topic: 

8 """ 

9 Represents a topic cluster of related sources with a lead text. 

10 

11 A topic groups multiple search results around a central theme, 

12 with one source selected as the "lead" that best represents the topic. 

13 """ 

14 

15 id: str 

16 title: str 

17 lead_source: Dict[str, Any] # The main source that represents this topic 

18 supporting_sources: List[Dict[str, Any]] = field(default_factory=list) 

19 rejected_sources: List[Dict[str, Any]] = field(default_factory=list) 

20 

21 # Metadata 

22 created_at: datetime = field(default_factory=datetime.now) 

23 

24 # Relationships to other topics 

25 related_topic_ids: List[str] = field(default_factory=list) 

26 parent_topic_id: Optional[str] = None 

27 child_topic_ids: List[str] = field(default_factory=list) 

28 

29 def add_supporting_source(self, source: Dict[str, Any]) -> None: 

30 """Add a source that supports this topic.""" 

31 if source not in self.supporting_sources: 

32 self.supporting_sources.append(source) 

33 

34 def reject_source(self, source: Dict[str, Any]) -> None: 

35 """Move a source to the rejected list.""" 

36 if source in self.supporting_sources: 

37 self.supporting_sources.remove(source) 

38 if source not in self.rejected_sources: 

39 self.rejected_sources.append(source) 

40 

41 def update_lead_source(self, new_lead: Dict[str, Any]) -> None: 

42 """ 

43 Change the lead source for this topic. 

44 The old lead becomes a supporting source. 

45 """ 

46 if self.lead_source: 46 ↛ 48line 46 didn't jump to line 48 because the condition on line 46 was always true

47 self.supporting_sources.append(self.lead_source) 

48 self.lead_source = new_lead 

49 # Remove new lead from supporting if it was there 

50 if new_lead in self.supporting_sources: 

51 self.supporting_sources.remove(new_lead) 

52 

53 def get_all_sources(self) -> List[Dict[str, Any]]: 

54 """Get all sources (lead + supporting) for this topic.""" 

55 return [self.lead_source] + self.supporting_sources 

56 

57 def to_dict(self) -> Dict[str, Any]: 

58 """Convert topic to dictionary for serialization.""" 

59 return { 

60 "id": self.id, 

61 "title": self.title, 

62 "lead_source": self.lead_source, 

63 "supporting_sources": self.supporting_sources, 

64 "rejected_sources": self.rejected_sources, 

65 "created_at": self.created_at.isoformat(), 

66 "related_topic_ids": self.related_topic_ids, 

67 "parent_topic_id": self.parent_topic_id, 

68 "child_topic_ids": self.child_topic_ids, 

69 } 

70 

71 

72@dataclass 

73class TopicGraph: 

74 """ 

75 Manages the collection of topics and their relationships. 

76 Provides methods for traversing and organizing topics. 

77 """ 

78 

79 topics: Dict[str, Topic] = field(default_factory=dict) 

80 

81 def add_topic(self, topic: Topic) -> None: 

82 """Add a topic to the graph.""" 

83 self.topics[topic.id] = topic 

84 

85 def get_topic(self, topic_id: str) -> Optional[Topic]: 

86 """Get a topic by ID.""" 

87 return self.topics.get(topic_id) 

88 

89 def link_topics(self, topic1_id: str, topic2_id: str) -> None: 

90 """Create a bidirectional relationship between two topics.""" 

91 topic1 = self.get_topic(topic1_id) 

92 topic2 = self.get_topic(topic2_id) 

93 

94 if topic1 and topic2: 

95 if topic2_id not in topic1.related_topic_ids: 

96 topic1.related_topic_ids.append(topic2_id) 

97 if topic1_id not in topic2.related_topic_ids: 

98 topic2.related_topic_ids.append(topic1_id) 

99 

100 def set_parent_child(self, parent_id: str, child_id: str) -> None: 

101 """Set a parent-child relationship between topics.""" 

102 parent = self.get_topic(parent_id) 

103 child = self.get_topic(child_id) 

104 

105 if parent and child: 105 ↛ exitline 105 didn't return from function 'set_parent_child' because the condition on line 105 was always true

106 if child_id not in parent.child_topic_ids: 106 ↛ 108line 106 didn't jump to line 108 because the condition on line 106 was always true

107 parent.child_topic_ids.append(child_id) 

108 child.parent_topic_id = parent_id 

109 

110 def get_root_topics(self) -> List[Topic]: 

111 """Get all topics that have no parent.""" 

112 return [ 

113 topic 

114 for topic in self.topics.values() 

115 if topic.parent_topic_id is None 

116 ] 

117 

118 def get_related_topics(self, topic_id: str) -> List[Topic]: 

119 """Get all topics related to a given topic.""" 

120 topic = self.get_topic(topic_id) 

121 if not topic: 

122 return [] 

123 

124 related = [] 

125 for related_id in topic.related_topic_ids: 

126 related_topic = self.get_topic(related_id) 

127 if related_topic: 127 ↛ 125line 127 didn't jump to line 125 because the condition on line 127 was always true

128 related.append(related_topic) 

129 return related 

130 

131 def merge_topics( 

132 self, topic1_id: str, topic2_id: str, new_title: str = None 

133 ) -> Optional[Topic]: 

134 """ 

135 Merge two topics into one, combining their sources. 

136 Returns the merged topic or None if merge failed. 

137 """ 

138 topic1 = self.get_topic(topic1_id) 

139 topic2 = self.get_topic(topic2_id) 

140 

141 if not topic1 or not topic2: 

142 return None 

143 

144 # Use topic1 as the base, update title if provided 

145 if new_title: 

146 topic1.title = new_title 

147 

148 # Add topic2's sources to topic1 

149 for source in topic2.supporting_sources: 

150 topic1.add_supporting_source(source) 

151 

152 for source in topic2.rejected_sources: 152 ↛ 153line 152 didn't jump to line 153 because the loop on line 152 never started

153 if source not in topic1.rejected_sources: 

154 topic1.rejected_sources.append(source) 

155 

156 # Update relationships 

157 for related_id in topic2.related_topic_ids: 

158 if ( 158 ↛ 157line 158 didn't jump to line 157 because the condition on line 158 was always true

159 related_id != topic1_id 

160 and related_id not in topic1.related_topic_ids 

161 ): 

162 topic1.related_topic_ids.append(related_id) 

163 

164 # Update child relationships 

165 for child_id in topic2.child_topic_ids: 165 ↛ 166line 165 didn't jump to line 166 because the loop on line 165 never started

166 self.set_parent_child(topic1_id, child_id) 

167 

168 # Remove topic2 from the graph 

169 del self.topics[topic2_id] 

170 

171 # Update any references to topic2 in other topics 

172 for topic in self.topics.values(): 

173 if topic2_id in topic.related_topic_ids: 

174 topic.related_topic_ids.remove(topic2_id) 

175 if topic1_id not in topic.related_topic_ids: 175 ↛ 178line 175 didn't jump to line 178 because the condition on line 175 was always true

176 topic.related_topic_ids.append(topic1_id) 

177 

178 if topic.parent_topic_id == topic2_id: 178 ↛ 179line 178 didn't jump to line 179 because the condition on line 178 was never true

179 topic.parent_topic_id = topic1_id 

180 

181 return topic1 

182 

183 def to_dict(self) -> Dict[str, Any]: 

184 """Convert the entire graph to a dictionary.""" 

185 return { 

186 topic_id: topic.to_dict() for topic_id, topic in self.topics.items() 

187 }