Spaces:
Runtime error
Runtime error
| """Notion reader. Taken and modified from https://github.com/jerryjliu/gpt_index/blob/main/gpt_index/readers/notion.py""" | |
| import os | |
| from typing import Any, Dict, List, Optional | |
| import requests | |
| from gpt_index.readers.base import BaseReader | |
| from gpt_index.schema import Document | |
| BLOCK_CHILD_URL_TMPL = "https://api.notion.com/v1/blocks/{block_id}/children" | |
| SEARCH_URL = "https://api.notion.com/v1/search" | |
| class NotionPageReader(BaseReader): | |
| """Notion Page reader. | |
| Reads a set of Notion pages. | |
| """ | |
| def __init__(self, integration_token: Optional[str] = None) -> None: | |
| self.token = integration_token | |
| self.headers = { | |
| "Authorization": "Bearer " + self.token, | |
| "Content-Type": "application/json", | |
| "Notion-Version": "2022-06-28", | |
| } | |
| def _read_block(self, block_id: str, num_tabs: int = 0) -> str: | |
| """Read a block.""" | |
| done = False | |
| result_lines_arr = [] | |
| cur_block_id = block_id | |
| while not done: | |
| block_url = BLOCK_CHILD_URL_TMPL.format(block_id=cur_block_id) | |
| query_dict: Dict[str, Any] = {} | |
| res = requests.request( | |
| "GET", block_url, headers=self.headers, json=query_dict | |
| ) | |
| data = res.json() | |
| for result in data["results"]: | |
| result_type = result["type"] | |
| result_obj = result[result_type] | |
| # NOTE: Notion reader doesn't support all block objects atm, only | |
| # block objects with rich text. | |
| if "rich_text" not in result_obj: | |
| continue | |
| cur_result_text_arr = [] | |
| for rich_text in result_obj["rich_text"]: | |
| # skip if doesn't have text object | |
| if "text" in rich_text: | |
| text = rich_text["text"]["content"] | |
| prefix = "\t" * num_tabs | |
| cur_result_text_arr.append(prefix + text) | |
| result_block_id = result["id"] | |
| has_children = result["has_children"] | |
| if has_children: | |
| children_text = self._read_block( | |
| result_block_id, num_tabs=num_tabs + 1 | |
| ) | |
| cur_result_text_arr.append(children_text) | |
| cur_result_text = "\n".join(cur_result_text_arr) | |
| result_lines_arr.append(cur_result_text) | |
| if data["next_cursor"] is None: | |
| done = True | |
| break | |
| else: | |
| cur_block_id = data["next_cursor"] | |
| result_lines = "\n".join(result_lines_arr) | |
| return result_lines | |
| def read_page(self, page_id: str) -> str: | |
| """Read a page.""" | |
| return self._read_block(page_id) | |
| def search(self, query: str) -> List[str]: | |
| """Search Notion page given a text query.""" | |
| done = False | |
| next_cursor: Optional[str] = None | |
| results = [] | |
| while not done: | |
| query_dict = { | |
| "query": query, | |
| } | |
| if next_cursor is not None: | |
| query_dict["start_cursor"] = next_cursor | |
| res = requests.post(SEARCH_URL, headers=self.headers, json=query_dict) | |
| data = res.json() | |
| for result in data["results"]: | |
| result_data = { | |
| "page_id": result["id"], | |
| "created_by": result["created_by"]["id"], | |
| "last_edited_time": result["last_edited_time"], | |
| "source": result["url"] | |
| } | |
| results.append(result_data) | |
| if data["next_cursor"] is None: | |
| done = True | |
| break | |
| else: | |
| next_cursor = data["next_cursor"] | |
| return results | |
| def load_data(self, **load_kwargs: Any) -> List[Document]: | |
| """Load data from the input directory.""" | |
| docs = [] | |
| for metadata in load_kwargs["pages"]: | |
| page_text = self.read_page(metadata["page_id"]) | |
| docs.append(Document(page_text, extra_info=metadata)) | |
| return docs | |