Coverage for slack_bot / obsidian / dispatcher.py: 0%

1import os 

2from typing import Optional, List, Dict, Any 

3from slack_sdk import WebClient 

4 

5from health.utils.logging_config import setup_logger 

6from slack_bot.obsidian.indexer import ObsidianIndexer 

7from slack_bot.obsidian.generators import WritingAssistant, ReplyGenerator, DecisionSupport, SearchAnalyzer, DeAIReviser, ZhihuGenerator 

8from slack_bot.context.storage import ContextStorage as BaseStorage 

9 

10logger = setup_logger(__name__) 

11 

12class ObsidianContextStorage(BaseStorage): 

13 """Specific storage for Obsidian bot to avoid mixing with Health bot.""" 

14 def __init__(self, channel_id: str): 

15 super().__init__(channel_id) 

16 from health import config 

17 self.context_dir = config.DATA_DIR / "obsidian_context" 

18 self.context_dir.mkdir(parents=True, exist_ok=True) 

19 self.file_path = self.context_dir / f"{channel_id}.json" 

20 self.state_file = self.context_dir / f"{channel_id}_state.json" 

21 

22 def save_state(self, state: Dict[str, Any]): 

23 import json 

24 with open(self.state_file, "w") as f: 

25 json.dump(state, f) 

26 

27 def load_state(self) -> Dict[str, Any]: 

28 import json 

29 if not self.state_file.exists(): 

30 return {"mode": "none"} 

31 try: 

32 with open(self.state_file, "r") as f: 

33 return json.load(f) 

34 except: 

35 return {"mode": "none"} 

36 

37class ObsidianDispatcher: 

38 """Routes Obsidian Slack Bot messages.""" 

39 

40 MODE_WRITE = "write" 

41 MODE_REPLY = "reply" 

42 MODE_DECIDE = "decide" 

43 MODE_SEARCH = "search" 

44 MODE_DEAI = "deai" 

45 MODE_ZHIHU = "zhihu" 

46 

47 def __init__(self, vault_path: str): 

48 self.indexer = ObsidianIndexer(vault_path) 

49 self.indexer.scan_vault() 

50 

51 self.generators = { 

52 self.MODE_WRITE: WritingAssistant(self.indexer), 

53 self.MODE_REPLY: ReplyGenerator(self.indexer), 

54 self.MODE_DECIDE: DecisionSupport(self.indexer), 

55 self.MODE_SEARCH: SearchAnalyzer(self.indexer), 

56 self.MODE_DEAI: DeAIReviser(self.indexer), 

57 self.MODE_ZHIHU: ZhihuGenerator(self.indexer) 

58 } 

59 

60 token = os.environ.get("OBSIDIAN_SLACK_BOT_TOKEN") 

61 self.client = WebClient(token=token) 

62 logger.info("ObsidianDispatcher initialized") 

63 

64 def dispatch(self, text: str, channel_id: str, user_id: str, response_ts: Optional[str] = None): 

65 storage = ObsidianContextStorage(channel_id) 

66 state = storage.load_state() 

67 current_mode = state.get("mode", "none") 

68 

69 cmd = text.strip().lower() 

70 

71 # 1. Handle Commands 

72 if cmd.startswith("mode "): 

73 new_mode = cmd.split(" ", 1)[1].strip() 

74 if new_mode in self.generators: 

75 storage.clear() 

76 storage.save_state({"mode": new_mode}) 

77 self._reply(channel_id, f"✅ Switched to mode: *{new_mode}* (History cleared)", response_ts) 

78 return 

79 else: 

80 self._reply(channel_id, f"❌ Invalid mode. Available: {', '.join(self.generators.keys())}", response_ts) 

81 return 

82 

83 if cmd == "clear": 

84 storage.clear() 

85 self._reply(channel_id, "🧹 Context cleared.", response_ts) 

86 return 

87 

88 if cmd == "reload": 

89 self.indexer.scan_vault() 

90 self._reply(channel_id, "🔄 Vault reloaded.", response_ts) 

91 return 

92 

93 # 2. Handle Content based on mode 

94 if current_mode == "none": 

95 self._reply(channel_id, "⚠️ No mode selected. Use `mode write/reply/decide/search/deai/zhihu` to start.", response_ts) 

96 return 

97 

98 # Get history 

99 history = storage.get_context() 

100 generator = self.generators[current_mode] 

101 

102 try: 

103 logger.info(f"Generating Obsidian response in mode: {current_mode}") 

104 # The chat method handles first turn logic and returns updated history 

105 # But wait, our generators' chat method appends to history internally.  

106 # In our Slack setup, we want to append the result to ContextStorage. 

107 

108 # Note: generator.chat(text, history) returns (response, new_history_list) 

109 # new_history_list contains [...history, {user: rich_prompt}, {assistant: response}] 

110 response_text, updated_history = generator.chat(text, history) 

111 

112 # Since generator.chat already did the LLM call and returned history, 

113 # we just need to save the ONLY NEW messages to our storage. 

114 # However, during FIRST TURN, text is replaced by rich_prompt. 

115 

116 # For simplicity: completely overwrite storage with updated_history 

117 # ContextStorage doesn't support overwrite easily, but we can clear and re-add. 

118 # Better: Modify ObsidianContextStorage to support set_context. 

119 

120 # Actually, let's just use the updated_history and manually add to storage 

121 # Warning: first turn 'user' message in storage will be the RICH PROMPT 

122 # This is actually better for context continuity in Slack too. 

123 

124 storage.clear() 

125 for msg in updated_history: 

126 storage.add_message(msg["role"], msg["content"]) 

127 

128 self._reply(channel_id, response_text, response_ts) 

129 

130 except Exception as e: 

131 logger.error(f"Obsidian generation failed: {e}", exc_info=True) 

132 self._reply(channel_id, f"🐛 Error: {str(e)}", response_ts) 

133 

134 def _reply(self, channel_id: str, text: str, ts: Optional[str] = None): 

135 from slack_bot.utils.mrkdwn import SlackFormatter 

136 

137 # Slack API limits (official documentation): 

138 # - chat.postMessage: 40,000 chars 

139 # - chat.update: 4,000 chars 

140 # Slack markdown formatting can expand text by 50-100% (bold, links, code blocks, etc.) 

141 TRUNCATE_WARNING = "\n\n⚠️ (Response truncated, full text uploaded as file)" 

142 FILE_UPLOAD_HINT = "📄 Full response uploaded as file above ☝️" 

143 

144 # Calculate safe pre-formatting limit with aggressive overhead buffer 

145 if ts: 

146 # For updates: 4000 limit - 60% overhead for safety = ~1600 chars 

147 max_raw_length = 1500 

148 else: 

149 # For new messages: 40000 limit - 40% overhead = ~24000 chars 

150 max_raw_length = 24000 

151 

152 original_length = len(text) 

153 should_upload_file = original_length > max_raw_length 

154 

155 if should_upload_file: 

156 # Upload full text as file first 

157 self._upload_as_file(channel_id, text, ts) 

158 

159 # Then send a truncated preview 

160 truncate_at = text.rfind('\n\n', 0, 800) 

161 if truncate_at == -1 or truncate_at < 400: 

162 truncate_at = text.rfind('\n', 0, 800) 

163 if truncate_at == -1 or truncate_at < 400: 

164 truncate_at = 800 

165 

166 preview_text = text[:truncate_at] + f"\n\n...\n\n{FILE_UPLOAD_HINT}" 

167 formatted_text = SlackFormatter.convert(preview_text) 

168 logger.info(f"Uploaded full response ({original_length} chars) as file, sending preview ({len(preview_text)} chars)") 

169 else: 

170 # Text is short enough, just format and send 

171 formatted_text = SlackFormatter.convert(text) 

172 

173 # Final safety check after formatting 

174 final_limit = 3900 if ts else 39900 

175 if len(formatted_text) > final_limit: 

176 logger.warning(f"Formatted text ({len(formatted_text)} chars) exceeds limit after formatting. Uploading as file.") 

177 self._upload_as_file(channel_id, text, ts) 

178 # Send simple notification 

179 formatted_text = FILE_UPLOAD_HINT 

180 

181 # Send the message 

182 try: 

183 if ts: 

184 self.client.chat_update(channel=channel_id, ts=ts, text=formatted_text) 

185 else: 

186 self.client.chat_postMessage(channel=channel_id, text=formatted_text) 

187 except Exception as e: 

188 if "msg_too_long" in str(e): 

189 logger.error(f"Still got msg_too_long error even after truncation. Text length: {len(formatted_text)}") 

190 # Ultimate fallback: send error message only 

191 try: 

192 error_msg = "❌ Response formatting error. Please check the uploaded file above." 

193 if ts: 

194 self.client.chat_update(channel=channel_id, ts=ts, text=error_msg) 

195 else: 

196 self.client.chat_postMessage(channel=channel_id, text=error_msg) 

197 except: 

198 logger.error("Failed to send even error message", exc_info=True) 

199 else: 

200 raise 

201 

202 def _upload_as_file(self, channel_id: str, content: str, ts: Optional[str] = None): 

203 """Fallback: upload long responses as text files.""" 

204 try: 

205 import tempfile 

206 from datetime import datetime 

207 

208 filename = f"response_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt" 

209 

210 with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False, encoding='utf-8') as tmp: 

211 tmp.write(content) 

212 tmp_path = tmp.name 

213 

214 self.client.files_upload_v2( 

215 channel=channel_id, 

216 file=tmp_path, 

217 filename=filename, 

218 initial_comment="📄 Response too long for Slack message, uploaded as file:", 

219 thread_ts=ts 

220 ) 

221 

222 os.unlink(tmp_path) 

223 logger.info(f"Uploaded long response as file: {filename}") 

224 except Exception as e: 

225 logger.error(f"Failed to upload file: {e}", exc_info=True) 

226 # Ultimate fallback: send error message 

227 error_msg = "❌ Response too long to display. Please try a more specific query." 

228 if ts: 

229 self.client.chat_update(channel=channel_id, ts=ts, text=error_msg) 

230 else: 

231 self.client.chat_postMessage(channel=channel_id, text=error_msg) 

232