dify/api/core/data_loader/loader/notion.py

import json
import logging
from typing import Any, Dict, List, Optional

import requests
from flask import current_app
from langchain.document_loaders.base import BaseLoader
from langchain.schema import Document

from extensions.ext_database import db
from models.dataset import Document as DocumentModel
from models.source import DataSourceBinding

logger = logging.getLogger(__name__)

BLOCK_CHILD_URL_TMPL = "https://api.notion.com/v1/blocks/{block_id}/children"
DATABASE_URL_TMPL = "https://api.notion.com/v1/databases/{database_id}/query"
SEARCH_URL = "https://api.notion.com/v1/search"

RETRIEVE_PAGE_URL_TMPL = "https://api.notion.com/v1/pages/{page_id}"
RETRIEVE_DATABASE_URL_TMPL = "https://api.notion.com/v1/databases/{database_id}"
HEADING_TYPE = ['heading_1', 'heading_2', 'heading_3']


class NotionLoader(BaseLoader):
    def __init__(
            self,
            notion_access_token: str,
            notion_workspace_id: str,
            notion_obj_id: str,
            notion_page_type: str,
            document_model: Optional[DocumentModel] = None
    ):
        self._document_model = document_model
        self._notion_workspace_id = notion_workspace_id
        self._notion_obj_id = notion_obj_id
        self._notion_page_type = notion_page_type
        self._notion_access_token = notion_access_token

        if not self._notion_access_token:
            integration_token = current_app.config.get('NOTION_INTEGRATION_TOKEN')
            if integration_token is None:
                raise ValueError(
                    "Must specify `integration_token` or set environment "
                    "variable `NOTION_INTEGRATION_TOKEN`."
                )

            self._notion_access_token = integration_token

    @classmethod
    def from_document(cls, document_model: DocumentModel):
        data_source_info = document_model.data_source_info_dict
        if not data_source_info or 'notion_page_id' not in data_source_info \
                or 'notion_workspace_id' not in data_source_info:
            raise ValueError("no notion page found")

        notion_workspace_id = data_source_info['notion_workspace_id']
        notion_obj_id = data_source_info['notion_page_id']
        notion_page_type = data_source_info['type']
        notion_access_token = cls._get_access_token(document_model.tenant_id, notion_workspace_id)

        return cls(
            notion_access_token=notion_access_token,
            notion_workspace_id=notion_workspace_id,
            notion_obj_id=notion_obj_id,
            notion_page_type=notion_page_type,
            document_model=document_model
        )

    def load(self) -> List[Document]:
        self.update_last_edited_time(
            self._document_model
        )

        text_docs = self._load_data_as_documents(self._notion_obj_id, self._notion_page_type)

        return text_docs

    def _load_data_as_documents(
            self, notion_obj_id: str, notion_page_type: str
    ) -> List[Document]:
        docs = []
        if notion_page_type == 'database':
            # get all the pages in the database
            page_text_documents = self._get_notion_database_data(notion_obj_id)
            docs.extend(page_text_documents)
        elif notion_page_type == 'page':
            page_text_list = self._get_notion_block_data(notion_obj_id)
            for page_text in page_text_list:
                docs.append(Document(page_content=page_text))
        else:
            raise ValueError("notion page type not supported")

        return docs

    def _get_notion_database_data(
            self, database_id: str, query_dict: Dict[str, Any] = {}
    ) -> List[Document]:
        """Get all the pages from a Notion database."""
        res = requests.post(
            DATABASE_URL_TMPL.format(database_id=database_id),
            headers={
                "Authorization": "Bearer " + self._notion_access_token,
                "Content-Type": "application/json",
                "Notion-Version": "2022-06-28",
            },
            json=query_dict,
        )

        data = res.json()

        database_content_list = []
        if 'results' not in data or data["results"] is None:
            return []
        for result in data["results"]:
            properties = result['properties']
            data = {}
            for property_name, property_value in properties.items():
                type = property_value['type']
                if type == 'multi_select':
                    value = []
                    multi_select_list = property_value[type]
                    for multi_select in multi_select_list:
                        value.append(multi_select['name'])
                elif type == 'rich_text' or type == 'title':
                    if len(property_value[type]) > 0:
                        value = property_value[type][0]['plain_text']
                    else:
                        value = ''
                elif type == 'select' or type == 'status':
                    if property_value[type]:
                        value = property_value[type]['name']
                    else:
                        value = ''
                else:
                    value = property_value[type]
                data[property_name] = value
            row_dict = {k: v for k, v in data.items() if v}
            row_content = ''
            for key, value in row_dict.items():
                if isinstance(value, dict):
                    value_dict = {k: v for k, v in value.items() if v}
                    value_content = ''.join(f'{k}:{v} ' for k, v in value_dict.items())
                    row_content = row_content + f'{key}:{value_content}\n'
                else:
                    row_content = row_content + f'{key}:{value}\n'
            document = Document(page_content=row_content)
            database_content_list.append(document)

        return database_content_list

    def _get_notion_block_data(self, page_id: str) -> List[str]:
        result_lines_arr = []
        cur_block_id = page_id
        while True:
            block_url = BLOCK_CHILD_URL_TMPL.format(block_id=cur_block_id)
            query_dict: Dict[str, Any] = {}

            res = requests.request(
                "GET",
                block_url,
                headers={
                    "Authorization": "Bearer " + self._notion_access_token,
                    "Content-Type": "application/json",
                    "Notion-Version": "2022-06-28",
                },
                json=query_dict
            )
            data = res.json()
            # current block's heading
            heading = ''
            for result in data["results"]:
                result_type = result["type"]
                result_obj = result[result_type]
                cur_result_text_arr = []
                if result_type == 'table':
                    result_block_id = result["id"]
                    text = self._read_table_rows(result_block_id)
                    text += "\n\n"
                    result_lines_arr.append(text)
                else:
                    if "rich_text" in result_obj:
                        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"]
                                cur_result_text_arr.append(text)
                                if result_type in HEADING_TYPE:
                                    heading = text

                    result_block_id = result["id"]
                    has_children = result["has_children"]
                    block_type = result["type"]
                    if has_children and block_type != 'child_page':
                        children_text = self._read_block(
                            result_block_id, num_tabs=1
                        )
                        cur_result_text_arr.append(children_text)

                    cur_result_text = "\n".join(cur_result_text_arr)
                    cur_result_text += "\n\n"
                    if result_type in HEADING_TYPE:
                        result_lines_arr.append(cur_result_text)
                    else:
                        result_lines_arr.append(f'{heading}\n{cur_result_text}')

            if data["next_cursor"] is None:
                break
            else:
                cur_block_id = data["next_cursor"]
        return result_lines_arr

    def _read_block(self, block_id: str, num_tabs: int = 0) -> str:
        """Read a block."""
        result_lines_arr = []
        cur_block_id = block_id
        while True:
            block_url = BLOCK_CHILD_URL_TMPL.format(block_id=cur_block_id)
            query_dict: Dict[str, Any] = {}

            res = requests.request(
                "GET",
                block_url,
                headers={
                    "Authorization": "Bearer " + self._notion_access_token,
                    "Content-Type": "application/json",
                    "Notion-Version": "2022-06-28",
                },
                json=query_dict
            )
            data = res.json()
            if 'results' not in data or data["results"] is None:
                break
            heading = ''
            for result in data["results"]:
                result_type = result["type"]
                result_obj = result[result_type]
                cur_result_text_arr = []
                if result_type == 'table':
                    result_block_id = result["id"]
                    text = self._read_table_rows(result_block_id)
                    result_lines_arr.append(text)
                else:
                    if "rich_text" in result_obj:
                        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)
                                if result_type in HEADING_TYPE:
                                    heading = text
                    result_block_id = result["id"]
                    has_children = result["has_children"]
                    block_type = result["type"]
                    if has_children and block_type != 'child_page':
                        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)
                    if result_type in HEADING_TYPE:
                        result_lines_arr.append(cur_result_text)
                    else:
                        result_lines_arr.append(f'{heading}\n{cur_result_text}')

            if data["next_cursor"] is None:
                break
            else:
                cur_block_id = data["next_cursor"]

        result_lines = "\n".join(result_lines_arr)
        return result_lines

    def _read_table_rows(self, block_id: str) -> str:
        """Read table rows."""
        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={
                    "Authorization": "Bearer " + self._notion_access_token,
                    "Content-Type": "application/json",
                    "Notion-Version": "2022-06-28",
                },
                json=query_dict
            )
            data = res.json()
            # get table headers text
            table_header_cell_texts = []
            tabel_header_cells = data["results"][0]['table_row']['cells']
            for tabel_header_cell in tabel_header_cells:
                if tabel_header_cell:
                    for table_header_cell_text in tabel_header_cell:
                        text = table_header_cell_text["text"]["content"]
                        table_header_cell_texts.append(text)
            # get table columns text and format
            results = data["results"]
            for i in range(len(results) - 1):
                column_texts = []
                tabel_column_cells = data["results"][i + 1]['table_row']['cells']
                for j in range(len(tabel_column_cells)):
                    if tabel_column_cells[j]:
                        for table_column_cell_text in tabel_column_cells[j]:
                            column_text = table_column_cell_text["text"]["content"]
                            column_texts.append(f'{table_header_cell_texts[j]}:{column_text}')

                cur_result_text = "\n".join(column_texts)
                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 update_last_edited_time(self, document_model: DocumentModel):
        if not document_model:
            return

        last_edited_time = self.get_notion_last_edited_time()
        data_source_info = document_model.data_source_info_dict
        data_source_info['last_edited_time'] = last_edited_time
        update_params = {
            DocumentModel.data_source_info: json.dumps(data_source_info)
        }

        DocumentModel.query.filter_by(id=document_model.id).update(update_params)
        db.session.commit()

    def get_notion_last_edited_time(self) -> str:
        obj_id = self._notion_obj_id
        page_type = self._notion_page_type
        if page_type == 'database':
            retrieve_page_url = RETRIEVE_DATABASE_URL_TMPL.format(database_id=obj_id)
        else:
            retrieve_page_url = RETRIEVE_PAGE_URL_TMPL.format(page_id=obj_id)

        query_dict: Dict[str, Any] = {}

        res = requests.request(
            "GET",
            retrieve_page_url,
            headers={
                "Authorization": "Bearer " + self._notion_access_token,
                "Content-Type": "application/json",
                "Notion-Version": "2022-06-28",
            },
            json=query_dict
        )

        data = res.json()
        return data["last_edited_time"]

    @classmethod
    def _get_access_token(cls, tenant_id: str, notion_workspace_id: str) -> str:
        data_source_binding = DataSourceBinding.query.filter(
            db.and_(
                DataSourceBinding.tenant_id == tenant_id,
                DataSourceBinding.provider == 'notion',
                DataSourceBinding.disabled == False,
                DataSourceBinding.source_info['workspace_id'] == f'"{notion_workspace_id}"'
            )
        ).first()

        if not data_source_binding:
            raise Exception(f'No notion data source binding found for tenant {tenant_id} '
                            f'and notion workspace {notion_workspace_id}')

        return data_source_binding.access_token