Coverage for src / local_deep_research / followup_research / service.py: 52%

1""" 

2Service layer for follow-up research functionality. 

3 

4This service handles the business logic for follow-up research, 

5including loading parent research context and orchestrating the search. 

6""" 

7 

8from typing import Dict, Any 

9from loguru import logger 

10 

11from ..database.models import ResearchHistory 

12from ..database.session_context import get_user_db_session 

13from ..web.services.research_sources_service import ResearchSourcesService 

14from .models import FollowUpRequest 

15 

16 

17class FollowUpResearchService: 

18 """Service for handling follow-up research operations.""" 

19 

20 def __init__(self, username: str = None): 

21 """ 

22 Initialize the follow-up research service. 

23 

24 Args: 

25 username: Username for database access 

26 """ 

27 self.username = username 

28 

29 def load_parent_research(self, parent_research_id: str) -> Dict[str, Any]: 

30 """ 

31 Load parent research data from the database. 

32 

33 Args: 

34 parent_research_id: ID of the parent research 

35 

36 Returns: 

37 Dictionary containing parent research data including: 

38 - report_content: The generated report 

39 - resources: List of research resources/links 

40 - query: Original research query 

41 - strategy: Strategy used 

42 """ 

43 try: 

44 with get_user_db_session(self.username) as session: 

45 # Load research history 

46 research = ( 

47 session.query(ResearchHistory) 

48 .filter_by(id=parent_research_id) 

49 .first() 

50 ) 

51 

52 if not research: 52 ↛ 53line 52 didn't jump to line 53 because the condition on line 52 was never true

53 logger.warning( 

54 f"Parent research not found: {parent_research_id}" 

55 ) 

56 return {} 

57 

58 logger.info( 

59 f"Found research: {research.id}, has meta: {research.research_meta is not None}" 

60 ) 

61 

62 # Use the ResearchSourcesService to get sources properly from database 

63 sources_service = ResearchSourcesService() 

64 resource_list = sources_service.get_research_sources( 

65 parent_research_id, username=self.username 

66 ) 

67 

68 logger.info( 

69 f"Found {len(resource_list)} sources from ResearchResource table" 

70 ) 

71 

72 # If no sources in database, try to get from research_meta as fallback 

73 if not resource_list and research.research_meta: 73 ↛ 74line 73 didn't jump to line 74 because the condition on line 73 was never true

74 logger.info( 

75 "No sources in database, checking research_meta" 

76 ) 

77 logger.info( 

78 f"Research meta keys: {list(research.research_meta.keys()) if isinstance(research.research_meta, dict) else 'Not a dict'}" 

79 ) 

80 

81 # Try different possible locations for sources in research_meta 

82 meta_sources = ( 

83 research.research_meta.get("all_links_of_system", []) 

84 or research.research_meta.get("sources", []) 

85 or research.research_meta.get("links", []) 

86 or [] 

87 ) 

88 

89 if meta_sources: 

90 logger.info( 

91 f"Found {len(meta_sources)} sources in research_meta, saving to database" 

92 ) 

93 # Save them to the database for future use 

94 saved = sources_service.save_research_sources( 

95 parent_research_id, 

96 meta_sources, 

97 username=self.username, 

98 ) 

99 logger.info(f"Saved {saved} sources to database") 

100 

101 # Now retrieve them properly formatted 

102 resource_list = sources_service.get_research_sources( 

103 parent_research_id, username=self.username 

104 ) 

105 

106 # Convert to dictionary format 

107 parent_data = { 

108 "research_id": research.id, 

109 "query": research.query, 

110 "report_content": research.report_content, 

111 "formatted_findings": research.research_meta.get( 

112 "formatted_findings", "" 

113 ) 

114 if research.research_meta 

115 else "", 

116 "strategy": research.research_meta.get("strategy_name", "") 

117 if research.research_meta 

118 else "", 

119 "resources": resource_list, 

120 "all_links_of_system": resource_list, 

121 } 

122 

123 logger.info( 

124 f"Loaded parent research {parent_research_id} with " 

125 f"{len(resource_list)} sources" 

126 ) 

127 

128 return parent_data 

129 

130 except Exception: 

131 logger.exception("Error loading parent research") 

132 return {} 

133 

134 def prepare_research_context( 

135 self, parent_research_id: str 

136 ) -> Dict[str, Any]: 

137 """ 

138 Prepare the research context for the contextual follow-up strategy. 

139 

140 Args: 

141 parent_research_id: ID of the parent research 

142 

143 Returns: 

144 Research context dictionary for the strategy 

145 """ 

146 parent_data = self.load_parent_research(parent_research_id) 

147 

148 if not parent_data: 148 ↛ 149line 148 didn't jump to line 149 because the condition on line 148 was never true

149 logger.warning("No parent data found, returning empty context") 

150 return {} 

151 

152 # Format context for the strategy 

153 research_context = { 

154 "parent_research_id": parent_research_id, 

155 "past_links": parent_data.get("all_links_of_system", []), 

156 "past_findings": parent_data.get("formatted_findings", ""), 

157 "report_content": parent_data.get("report_content", ""), 

158 "resources": parent_data.get("resources", []), 

159 "all_links_of_system": parent_data.get("all_links_of_system", []), 

160 "original_query": parent_data.get("query", ""), 

161 } 

162 

163 return research_context 

164 

165 def perform_followup(self, request: FollowUpRequest) -> Dict[str, Any]: 

166 """ 

167 Perform a follow-up research based on parent research. 

168 

169 This method prepares the context and parameters for the research system 

170 to use the contextual follow-up strategy. 

171 

172 Args: 

173 request: FollowUpRequest with question and parent research ID 

174 

175 Returns: 

176 Dictionary with research parameters for the research system 

177 """ 

178 # Prepare the research context from parent 

179 research_context = self.prepare_research_context( 

180 request.parent_research_id 

181 ) 

182 

183 if not research_context: 

184 logger.warning( 

185 f"Parent research not found: {request.parent_research_id}, using empty context" 

186 ) 

187 # Use empty context to allow follow-up without parent 

188 research_context = { 

189 "parent_research_id": request.parent_research_id, 

190 "past_links": [], 

191 "past_findings": "", 

192 "report_content": "", 

193 "resources": [], 

194 "all_links_of_system": [], 

195 "original_query": "", 

196 } 

197 

198 # Prepare parameters for the research system 

199 research_params = { 

200 "query": request.question, 

201 "strategy": "contextual-followup", 

202 "delegate_strategy": request.strategy, 

203 "max_iterations": request.max_iterations, 

204 "questions_per_iteration": request.questions_per_iteration, 

205 "research_context": research_context, 

206 "parent_research_id": request.parent_research_id, 

207 } 

208 

209 logger.info( 

210 f"Prepared follow-up research for question: '{request.question}' " 

211 f"based on parent: {request.parent_research_id}" 

212 ) 

213 

214 return research_params