llama-index

1
"""Psychic reader."""
2

3
import logging
4
import os
5
from typing import List, Optional
6

7
from llama_index.legacy.readers.base import BaseReader
8
from llama_index.legacy.schema import Document
9

10
logger = logging.getLogger(__name__)
11

12

13
class PsychicReader(BaseReader):
14
    """Psychic reader.
15

16
    Psychic is a platform that allows syncing data from many SaaS apps through one
17
        universal API.
18
    This reader connects to an instance of Psychic and reads data from it, given a
19
        connector ID, account ID, and API key.
20

21
    Learn more at docs.psychic.dev.
22

23
    Args:
24
        psychic_key (str): Secret key for Psychic.
25
            Get one at https://dashboard.psychic.dev/api-keys.
26

27
    """
28

29
    def __init__(self, psychic_key: Optional[str] = None) -> None:
30
        """Initialize with parameters."""
31
        try:
32
            from psychicapi import ConnectorId, Psychic
33
        except ImportError:
34
            raise ImportError(
35
                "`psychicapi` package not found, please run `pip install psychicapi`"
36
            )
37
        if psychic_key is None:
38
            psychic_key = os.environ["PSYCHIC_SECRET_KEY"]
39
            if psychic_key is None:
40
                raise ValueError(
41
                    "Must specify `psychic_key` or set environment "
42
                    "variable `PSYCHIC_SECRET_KEY`."
43
                )
44

45
        self.psychic = Psychic(secret_key=psychic_key)
46
        self.ConnectorId = ConnectorId
47

48
    def load_data(
49
        self, connector_id: Optional[str] = None, account_id: Optional[str] = None
50
    ) -> List[Document]:
51
        """Load data from a Psychic connection.
52

53
        Args:
54
            connector_id (str): The connector ID to connect to
55
            account_id (str): The account ID to connect to
56

57
        Returns:
58
            List[Document]: List of documents.
59

60
        """
61
        if not connector_id or not account_id:
62
            raise ValueError("Must specify both `connector_id` and `account_id`.")
63
        if connector_id not in self.ConnectorId.__members__:
64
            raise ValueError("Invalid connector ID.")
65

66
        # get all the documents in the database
67
        docs = []
68
        data = self.psychic.get_documents(self.ConnectorId[connector_id], account_id)
69
        for resource in data:
70
            text = resource.get("content")
71
            doc_id = resource.get("uri")
72
            docs.append(
73
                Document(
74
                    text=text,
75
                    id_=doc_id,
76
                    metadata={"connector_id": connector_id, "account_id": account_id},
77
                )
78
            )
79

80
        return docs
81

82

83
if __name__ == "__main__":
84
    reader = PsychicReader(psychic_key="public_key")
85
    logger.info(reader.load_data(connector_id="connector_id", account_id="account_id"))
86