Spaces:
Runtime error
Runtime error
File size: 7,888 Bytes
b699122 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 |
"""Slack reader."""
import logging
import os
import time
from datetime import datetime
from ssl import SSLContext
from typing import List, Optional
from gpt_index.readers.base import BaseReader
from gpt_index.readers.schema.base import Document
logger = logging.getLogger(__name__)
class SlackReader(BaseReader):
"""Slack reader.
Reads conversations from channels. If an earliest_date is provided, an
optional latest_date can also be provided. If no latest_date is provided,
we assume the latest date is the current timestamp.
Args:
slack_token (Optional[str]): Slack token. If not provided, we
assume the environment variable `SLACK_BOT_TOKEN` is set.
ssl (Optional[str]): Custom SSL context. If not provided, it is assumed
there is already an SSL context available.
earliest_date (Optional[datetime]): Earliest date from which
to read conversations. If not provided, we read all messages.
latest_date (Optional[datetime]): Latest date from which to
read conversations. If not provided, defaults to current timestamp
in combination with earliest_date.
"""
def __init__(
self,
slack_token: Optional[str] = None,
ssl: Optional[SSLContext] = None,
earliest_date: Optional[datetime] = None,
latest_date: Optional[datetime] = None,
) -> None:
"""Initialize with parameters."""
from slack_sdk import WebClient
if slack_token is None:
slack_token = os.environ["SLACK_BOT_TOKEN"]
if slack_token is None:
raise ValueError(
"Must specify `slack_token` or set environment "
"variable `SLACK_BOT_TOKEN`."
)
if ssl is None:
self.client = WebClient(token=slack_token)
else:
self.client = WebClient(token=slack_token, ssl=ssl)
if latest_date is not None and earliest_date is None:
raise ValueError(
"Must specify `earliest_date` if `latest_date` is specified."
)
if earliest_date is not None:
self.earliest_date_timestamp: Optional[float] = earliest_date.timestamp()
else:
self.earliest_date_timestamp = None
if latest_date is not None:
self.latest_date_timestamp = latest_date.timestamp()
else:
self.latest_date_timestamp = datetime.now().timestamp()
res = self.client.api_test()
if not res["ok"]:
raise ValueError(f"Error initializing Slack API: {res['error']}")
def _read_message(self, channel_id: str, message_ts: str) -> str:
from slack_sdk.errors import SlackApiError
"""Read a message."""
messages_text: List[str] = []
next_cursor = None
while True:
try:
# https://slack.com/api/conversations.replies
# List all replies to a message, including the message itself.
if self.earliest_date_timestamp is None:
result = self.client.conversations_replies(
channel=channel_id, ts=message_ts, cursor=next_cursor
)
else:
conversations_replies_kwargs = {
"channel": channel_id,
"ts": message_ts,
"cursor": next_cursor,
"latest": str(self.latest_date_timestamp),
}
if self.earliest_date_timestamp is not None:
conversations_replies_kwargs["oldest"] = str(
self.earliest_date_timestamp
)
result = self.client.conversations_replies(
**conversations_replies_kwargs # type: ignore
)
messages = result["messages"]
messages_text.extend(message["text"] for message in messages)
if not result["has_more"]:
break
next_cursor = result["response_metadata"]["next_cursor"]
except SlackApiError as e:
if e.response["error"] == "ratelimited":
logger.error(
"Rate limit error reached, sleeping for: {} seconds".format(
e.response.headers["retry-after"]
)
)
time.sleep(int(e.response.headers["retry-after"]))
else:
logger.error("Error parsing conversation replies: {}".format(e))
return "\n\n".join(messages_text)
def _read_channel(self, channel_id: str, reverse_chronological: bool) -> str:
from slack_sdk.errors import SlackApiError
"""Read a channel."""
result_messages: List[str] = []
next_cursor = None
while True:
try:
# Call the conversations.history method using the WebClient
# conversations.history returns the first 100 messages by default
# These results are paginated,
# see: https://api.slack.com/methods/conversations.history$pagination
conversations_history_kwargs = {
"channel": channel_id,
"cursor": next_cursor,
"latest": str(self.latest_date_timestamp),
}
if self.earliest_date_timestamp is not None:
conversations_history_kwargs["oldest"] = str(
self.earliest_date_timestamp
)
result = self.client.conversations_history(
**conversations_history_kwargs # type: ignore
)
conversation_history = result["messages"]
# Print results
logger.info(
"{} messages found in {}".format(
len(conversation_history), channel_id
)
)
result_messages.extend(
self._read_message(channel_id, message["ts"])
for message in conversation_history
)
if not result["has_more"]:
break
next_cursor = result["response_metadata"]["next_cursor"]
except SlackApiError as e:
if e.response["error"] == "ratelimited":
logger.error(
"Rate limit error reached, sleeping for: {} seconds".format(
e.response.headers["retry-after"]
)
)
time.sleep(int(e.response.headers["retry-after"]))
else:
logger.error("Error parsing conversation replies: {}".format(e))
return (
"\n\n".join(result_messages)
if reverse_chronological
else "\n\n".join(result_messages[::-1])
)
def load_data(
self, channel_ids: List[str], reverse_chronological: bool = True
) -> List[Document]:
"""Load data from the input directory.
Args:
channel_ids (List[str]): List of channel ids to read.
Returns:
List[Document]: List of documents.
"""
results = []
for channel_id in channel_ids:
channel_content = self._read_channel(
channel_id, reverse_chronological=reverse_chronological
)
results.append(
Document(channel_content, extra_info={"channel": channel_id})
)
return results
if __name__ == "__main__":
reader = SlackReader()
logger.info(reader.load_data(channel_ids=["C04DC2VUY3F"]))
|