llama-index

1
"""Slack tool spec."""
2

3
import logging
4
from datetime import datetime
5
from ssl import SSLContext
6
from typing import List, Optional
7

8
from llama_index.legacy.readers.slack import SlackReader
9
from llama_index.legacy.schema import Document
10
from llama_index.legacy.tools.tool_spec.base import BaseToolSpec
11

12
logger = logging.getLogger(__name__)
13

14

15
class SlackToolSpec(BaseToolSpec):
16
    """Slack tool spec."""
17

18
    spec_functions = ["load_data", "send_message", "fetch_channels"]
19

20
    def __init__(
21
        self,
22
        slack_token: Optional[str] = None,
23
        ssl: Optional[SSLContext] = None,
24
        earliest_date: Optional[datetime] = None,
25
        latest_date: Optional[datetime] = None,
26
    ) -> None:
27
        """Initialize with parameters."""
28
        self.reader = SlackReader(
29
            slack_token=slack_token,
30
            ssl=ssl,
31
            earliest_date=earliest_date,
32
            latest_date=latest_date,
33
        )
34

35
    def load_data(
36
        self,
37
        channel_ids: List[str],
38
        reverse_chronological: bool = True,
39
    ) -> List[Document]:
40
        """Load data from the input directory."""
41
        return self.reader.load_data(
42
            channel_ids=channel_ids,
43
            reverse_chronological=reverse_chronological,
44
        )
45

46
    def send_message(
47
        self,
48
        channel_id: str,
49
        message: str,
50
    ) -> None:
51
        """Send a message to a channel given the channel ID."""
52
        slack_client = self.reader.client
53
        try:
54
            msg_result = slack_client.chat_postMessage(
55
                channel=channel_id,
56
                text=message,
57
            )
58
            logger.info(msg_result)
59
        except Exception as e:
60
            logger.error(e)
61
            raise
62

63
    def fetch_channels(
64
        self,
65
    ) -> List[str]:
66
        """Fetch a list of relevant channels."""
67
        slack_client = self.reader.client
68
        try:
69
            msg_result = slack_client.conversations_list()
70
            logger.info(msg_result)
71
        except Exception as e:
72
            logger.error(e)
73
            raise
74

75
        return msg_result["channels"]
76