Coverage for slack_bot / main.py: 0%

1import os 

2import sys 

3import click 

4from slack_bolt import App 

5from slack_bolt.adapter.socket_mode import SocketModeHandler 

6from dotenv import load_dotenv 

7 

8# Add project root to path 

9sys.path.append(os.getcwd()) 

10 

11from health.utils.logging_config import setup_logger 

12from health.utils.env_loader import load_env_with_extras 

13from slack_bot.dispatcher import MessageDispatcher 

14from slack_bot.context.storage import ContextStorage 

15 

16logger = setup_logger(__name__) 

17 

18# Load environment variables 

19load_env_with_extras() 

20 

21 

22@click.command() 

23@click.option("--model", default=None, help="Gemini model to use (overrides .env)") 

24def main(model: str): 

25 """Butler Slack Bot - Gemini Edition.""" 

26 

27 # Override model in environment only if specified via CLI 

28 if model: 

29 os.environ["GEMINI_MODEL"] = model 

30 logger.info(f"Using model from CLI: {model}") 

31 

32 # Get the actual model being used (from CLI or .env) 

33 actual_model = os.environ.get("GEMINI_MODEL", "gemini-2.0-flash-exp") 

34 base_url = os.environ.get("GEMINI_BASE_URL", "Not set (will use direct API)") 

35 config_path = os.environ.get("GEMINI_CONFIG_PATH", ".env") 

36 

37 logger.info(f"Starting Butler Bot with configuration:") 

38 logger.info(f" Config file: {config_path}") 

39 logger.info(f" Model: {actual_model}") 

40 logger.info(f" Base URL: {base_url}") 

41 logger.info(f" Using proxy: {'Yes' if base_url and 'http' in base_url else 'No (Direct Google API)'}") 

42 

43 # Initialize Slack App 

44 bot_token = os.environ.get("SLACK_BOT_TOKEN") 

45 app_token = os.environ.get("SLACK_APP_TOKEN") 

46 

47 if not bot_token or not app_token: 

48 logger.error("Missing SLACK_BOT_TOKEN or SLACK_APP_TOKEN in .env") 

49 sys.exit(1) 

50 

51 app = App(token=bot_token) 

52 

53 # Tool mode: 'none' (no tools), 'light' (5), 'standard' (10), 'full' (15), 'all' (21) 

54 # Use 'none' if your proxy doesn't support function calling 

55 tool_mode = os.environ.get("TOOL_MODE", "light") 

56 dispatcher = MessageDispatcher(tool_mode=tool_mode) 

57 logger.info(f"Dispatcher initialized in '{tool_mode}' mode") 

58 

59 # Message deduplication: Track processed message IDs 

60 processed_messages = set() 

61 MAX_TRACKED_MESSAGES = 1000 

62 

63 # Command Handler: Clear Context 

64 @app.message(r"^/(clear|reset)$") 

65 def handle_clear_context(message, say): 

66 """Clear conversation context for the current channel.""" 

67 try: 

68 channel = message["channel"] 

69 storage = ContextStorage(channel) 

70 

71 # Get context size before clearing 

72 old_context = storage.get_context() 

73 old_count = len(old_context) 

74 

75 storage.clear() 

76 

77 logger.info(f"✅ Context cleared for channel {channel} (removed {old_count} messages)") 

78 say(f"✅ 对话历史已清空,我们重新开始!\n\n_Context cleared ({old_count} messages removed). Starting fresh conversation._") 

79 except Exception as e: 

80 logger.error(f"Error clearing context: {e}", exc_info=True) 

81 say(f"❌ 清空对话历史失败: {e}") 

82 

83 # Command Handler: Show Context Stats 

84 @app.message(r"^/stats$") 

85 def handle_context_stats(message, say): 

86 """Show conversation context statistics.""" 

87 try: 

88 channel = message["channel"] 

89 storage = ContextStorage(channel) 

90 context = storage.get_context() 

91 

92 # Calculate token estimate (rough: 1 token ≈ 3 chars for Chinese/English mix) 

93 total_chars = sum(len(msg.get("content", "")) for msg in context) 

94 estimated_tokens = total_chars // 3 

95 

96 stats_msg = ( 

97 f"📊 **对话上下文统计 (Context Stats)**\n\n" 

98 f"• 消息数量 (Messages): {len(context)}\n" 

99 f"• 总字符数 (Total chars): {total_chars:,}\n" 

100 f"• 预估 Token (Est. tokens): ~{estimated_tokens:,}\n\n" 

101 f"_使用 `/clear` 可以清空对话历史_" 

102 ) 

103 

104 say(stats_msg) 

105 logger.info(f"Context stats for {channel}: {len(context)} messages, ~{estimated_tokens} tokens") 

106 except Exception as e: 

107 logger.error(f"Error getting context stats: {e}", exc_info=True) 

108 say(f"❌ 获取统计信息失败: {e}") 

109 

110 # Message Handler 

111 @app.event("message") 

112 def handle_message(message, say): 

113 

114 try: 

115 # Skip command messages (they're handled by command handlers) 

116 text = message.get("text", "") 

117 if text.startswith("/clear") or text.startswith("/reset") or text.startswith("/stats"): 

118 logger.debug("Skipping command message in general handler") 

119 return 

120 

121 # Deduplicate using client_msg_id or event_ts 

122 msg_id = message.get("client_msg_id") or message.get("ts") 

123 

124 if msg_id in processed_messages: 

125 logger.debug(f"Skipping duplicate message: {msg_id}") 

126 return 

127 

128 # Add to processed set 

129 processed_messages.add(msg_id) 

130 

131 # Limit memory usage by keeping only recent message IDs 

132 if len(processed_messages) > MAX_TRACKED_MESSAGES: 

133 # Remove oldest half 

134 to_remove = list(processed_messages)[:MAX_TRACKED_MESSAGES // 2] 

135 for old_id in to_remove: 

136 processed_messages.discard(old_id) 

137 

138 text = message.get("text", "") 

139 channel = message["channel"] 

140 user = message["user"] 

141 files = message.get("files", []) 

142 

143 # Generate Request ID 

144 import uuid 

145 import time 

146 request_id = str(uuid.uuid4())[:8] 

147 start_time = time.time() 

148 

149 logger.info(f"[{request_id}] 📥 Received message from {user} in {channel}: {text} (Files: {len(files)})") 

150 

151 # Send temporary "Processing..." message 

152 ack_response = say(f"⏳ Butler is thinking...") 

153 ack_ts = ack_response["ts"] 

154 

155 # Dispatch to Gemini with the timestamp to update 

156 dispatcher.dispatch(text, channel, user, response_ts=ack_ts, request_id=request_id, files=files) 

157 

158 elapsed = time.time() - start_time 

159 logger.info(f"[{request_id}] ✅ Request completed in {elapsed:.2f}s") 

160 

161 except Exception as e: 

162 logger.error(f"Error handling message: {e}", exc_info=True) 

163 try: 

164 say(f"🐛 Butler encountered an error: {e}") 

165 except Exception: 

166 pass 

167 

168 # Start Socket Mode 

169 try: 

170 handler = SocketModeHandler(app, app_token) 

171 logger.info("⚡️ Butler (Gemini) is connected to Slack!") 

172 handler.start() 

173 except Exception as e: 

174 logger.error(f"Failed to start SocketModeHandler: {e}") 

175 sys.exit(1) 

176 

177 

178if __name__ == "__main__": 

179 main()