llama-index

1
"""Simple reader that reads tweets of a twitter handle."""
2

3
from typing import Any, List, Optional
4

5
from llama_index.legacy.readers.base import BasePydanticReader
6
from llama_index.legacy.schema import Document
7

8

9
class TwitterTweetReader(BasePydanticReader):
10
    """Twitter tweets reader.
11

12
    Read tweets of user twitter handle.
13

14
    Check 'https://developer.twitter.com/en/docs/twitter-api/\
15
        getting-started/getting-access-to-the-twitter-api' \
16
        on how to get access to twitter API.
17

18
    Args:
19
        bearer_token (str): bearer_token that you get from twitter API.
20
        num_tweets (Optional[int]): Number of tweets for each user twitter handle.\
21
            Default is 100 tweets.
22
    """
23

24
    is_remote: bool = True
25
    bearer_token: str
26
    num_tweets: Optional[int]
27

28
    def __init__(
29
        self,
30
        bearer_token: str,
31
        num_tweets: Optional[int] = 100,
32
    ) -> None:
33
        """Initialize with parameters."""
34
        super().__init__(
35
            num_tweets=num_tweets,
36
            bearer_token=bearer_token,
37
        )
38

39
    @classmethod
40
    def class_name(cls) -> str:
41
        return "TwitterTweetReader"
42

43
    def load_data(
44
        self,
45
        twitterhandles: List[str],
46
        num_tweets: Optional[int] = None,
47
        **load_kwargs: Any
48
    ) -> List[Document]:
49
        """Load tweets of twitter handles.
50

51
        Args:
52
            twitterhandles (List[str]): List of user twitter handles to read tweets.
53

54
        """
55
        try:
56
            import tweepy
57
        except ImportError:
58
            raise ImportError(
59
                "`tweepy` package not found, please run `pip install tweepy`"
60
            )
61

62
        client = tweepy.Client(bearer_token=self.bearer_token)
63
        results = []
64
        for username in twitterhandles:
65
            # tweets = api.user_timeline(screen_name=user, count=self.num_tweets)
66
            user = client.get_user(username=username)
67
            tweets = client.get_users_tweets(
68
                user.data.id, max_results=num_tweets or self.num_tweets
69
            )
70
            response = " "
71
            for tweet in tweets.data:
72
                response = response + tweet.text + "\n"
73
            results.append(Document(text=response))
74
        return results
75