llama-index

1
"""Notion tool spec."""
2

3
from typing import Any, Dict, List, Optional, Type
4

5
import requests
6

7
from llama_index.legacy.bridge.pydantic import BaseModel
8
from llama_index.legacy.readers.notion import NotionPageReader
9
from llama_index.legacy.tools.tool_spec.base import SPEC_FUNCTION_TYPE, BaseToolSpec
10

11
SEARCH_URL = "https://api.notion.com/v1/search"
12

13

14
class NotionLoadDataSchema(BaseModel):
15
    """Notion load data schema."""
16

17
    page_ids: Optional[List[str]] = None
18
    database_id: Optional[str] = None
19

20

21
class NotionSearchDataSchema(BaseModel):
22
    """Notion search data schema."""
23

24
    query: str
25
    direction: Optional[str] = None
26
    timestamp: Optional[str] = None
27
    value: Optional[str] = None
28
    property: Optional[str] = None
29
    page_size: int = 100
30

31

32
class NotionToolSpec(BaseToolSpec):
33
    """Notion tool spec.
34

35
    Currently a simple wrapper around the data loader.
36
    TODO: add more methods to the Notion spec.
37

38
    """
39

40
    spec_functions = ["load_data", "search_data"]
41

42
    def __init__(self, integration_token: Optional[str] = None) -> None:
43
        """Initialize with parameters."""
44
        self.reader = NotionPageReader(integration_token=integration_token)
45

46
    def get_fn_schema_from_fn_name(
47
        self, fn_name: str, spec_functions: Optional[List[SPEC_FUNCTION_TYPE]] = None
48
    ) -> Optional[Type[BaseModel]]:
49
        """Return map from function name."""
50
        if fn_name == "load_data":
51
            return NotionLoadDataSchema
52
        elif fn_name == "search_data":
53
            return NotionSearchDataSchema
54
        else:
55
            raise ValueError(f"Invalid function name: {fn_name}")
56

57
    def load_data(
58
        self, page_ids: Optional[List[str]] = None, database_id: Optional[str] = None
59
    ) -> str:
60
        """Loads content from a set of page ids or a database id.
61

62
        Don't use this endpoint if you don't know the page ids or database id.
63

64
        """
65
        page_ids = page_ids or []
66
        docs = self.reader.load_data(page_ids=page_ids, database_id=database_id)
67
        return "\n".join([doc.get_content() for doc in docs])
68

69
    def search_data(
70
        self,
71
        query: str,
72
        direction: Optional[str] = None,
73
        timestamp: Optional[str] = None,
74
        value: Optional[str] = None,
75
        property: Optional[str] = None,
76
        page_size: int = 100,
77
    ) -> str:
78
        """Search a list of relevant pages.
79

80
        Contains metadata for each page (but not the page content).
81

82
        """
83
        payload: Dict[str, Any] = {
84
            "query": query,
85
            "page_size": page_size,
86
        }
87
        if direction is not None or timestamp is not None:
88
            payload["sort"] = {}
89
            if direction is not None:
90
                payload["sort"]["direction"] = direction
91
            if timestamp is not None:
92
                payload["sort"]["timestamp"] = timestamp
93

94
        if value is not None or property is not None:
95
            payload["filter"] = {}
96
            if value is not None:
97
                payload["filter"]["value"] = value
98
            if property is not None:
99
                payload["filter"]["property"] = property
100

101
        response = requests.post(SEARCH_URL, json=payload, headers=self.reader.headers)
102
        response_json = response.json()
103
        return response_json["results"]
104