Feat/dataset notion import (#392)
Co-authored-by: StyleZhang <jasonapring2015@outlook.com> Co-authored-by: JzoNg <jzongcode@gmail.com>
This commit is contained in:
Jyong
GitHub
parent
f350948bde
commit
9253f72dea
96 changed files with 4479 additions and 367 deletions
367
api/core/data_source/notion.py
Normal file
367
api/core/data_source/notion.py
Normal file
|
|
@ -0,0 +1,367 @@
|
|||
"""Notion reader."""
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import requests # type: ignore
|
||||
|
||||
from llama_index.readers.base import BaseReader
|
||||
from llama_index.readers.schema.base import Document
|
||||
|
||||
INTEGRATION_TOKEN_NAME = "NOTION_INTEGRATION_TOKEN"
|
||||
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']
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# TODO: Notion DB reader coming soon!
|
||||
class NotionPageReader(BaseReader):
|
||||
"""Notion Page reader.
|
||||
|
||||
Reads a set of Notion pages.
|
||||
|
||||
Args:
|
||||
integration_token (str): Notion integration token.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self, integration_token: Optional[str] = None) -> None:
|
||||
"""Initialize with parameters."""
|
||||
if integration_token is None:
|
||||
integration_token = os.getenv(INTEGRATION_TOKEN_NAME)
|
||||
if integration_token is None:
|
||||
raise ValueError(
|
||||
"Must specify `integration_token` or set environment "
|
||||
"variable `NOTION_INTEGRATION_TOKEN`."
|
||||
)
|
||||
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()
|
||||
if 'results' not in data or data["results"] is None:
|
||||
done = True
|
||||
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"]
|
||||
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)
|
||||
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:
|
||||
done = True
|
||||
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=self.headers, 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 _read_parent_blocks(self, block_id: str, num_tabs: int = 0) -> List[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()
|
||||
# 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"]
|
||||
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)
|
||||
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:
|
||||
done = True
|
||||
break
|
||||
else:
|
||||
cur_block_id = data["next_cursor"]
|
||||
return result_lines_arr
|
||||
|
||||
def read_page(self, page_id: str) -> str:
|
||||
"""Read a page."""
|
||||
return self._read_block(page_id)
|
||||
|
||||
def read_page_as_documents(self, page_id: str) -> List[str]:
|
||||
"""Read a page as documents."""
|
||||
return self._read_parent_blocks(page_id)
|
||||
|
||||
def query_database_data(
|
||||
self, database_id: str, query_dict: Dict[str, Any] = {}
|
||||
) -> str:
|
||||
"""Get all the pages from a Notion database."""
|
||||
res = requests.post\
|
||||
(
|
||||
DATABASE_URL_TMPL.format(database_id=database_id),
|
||||
headers=self.headers,
|
||||
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
|
||||
database_content_list.append(json.dumps(data))
|
||||
|
||||
return "\n\n".join(database_content_list)
|
||||
|
||||
def query_database(
|
||||
self, database_id: str, query_dict: Dict[str, Any] = {}
|
||||
) -> List[str]:
|
||||
"""Get all the pages from a Notion database."""
|
||||
res = requests.post\
|
||||
(
|
||||
DATABASE_URL_TMPL.format(database_id=database_id),
|
||||
headers=self.headers,
|
||||
json=query_dict,
|
||||
)
|
||||
data = res.json()
|
||||
page_ids = []
|
||||
for result in data["results"]:
|
||||
page_id = result["id"]
|
||||
page_ids.append(page_id)
|
||||
|
||||
return page_ids
|
||||
|
||||
def search(self, query: str) -> List[str]:
|
||||
"""Search Notion page given a text query."""
|
||||
done = False
|
||||
next_cursor: Optional[str] = None
|
||||
page_ids = []
|
||||
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"]:
|
||||
page_id = result["id"]
|
||||
page_ids.append(page_id)
|
||||
|
||||
if data["next_cursor"] is None:
|
||||
done = True
|
||||
break
|
||||
else:
|
||||
next_cursor = data["next_cursor"]
|
||||
return page_ids
|
||||
|
||||
def load_data(
|
||||
self, page_ids: List[str] = [], database_id: Optional[str] = None
|
||||
) -> List[Document]:
|
||||
"""Load data from the input directory.
|
||||
|
||||
Args:
|
||||
page_ids (List[str]): List of page ids to load.
|
||||
|
||||
Returns:
|
||||
List[Document]: List of documents.
|
||||
|
||||
"""
|
||||
if not page_ids and not database_id:
|
||||
raise ValueError("Must specify either `page_ids` or `database_id`.")
|
||||
docs = []
|
||||
if database_id is not None:
|
||||
# get all the pages in the database
|
||||
page_ids = self.query_database(database_id)
|
||||
for page_id in page_ids:
|
||||
page_text = self.read_page(page_id)
|
||||
docs.append(Document(page_text))
|
||||
else:
|
||||
for page_id in page_ids:
|
||||
page_text = self.read_page(page_id)
|
||||
docs.append(Document(page_text))
|
||||
|
||||
return docs
|
||||
|
||||
def load_data_as_documents(
|
||||
self, page_ids: List[str] = [], database_id: Optional[str] = None
|
||||
) -> List[Document]:
|
||||
if not page_ids and not database_id:
|
||||
raise ValueError("Must specify either `page_ids` or `database_id`.")
|
||||
docs = []
|
||||
if database_id is not None:
|
||||
# get all the pages in the database
|
||||
page_text = self.query_database_data(database_id)
|
||||
docs.append(Document(page_text))
|
||||
else:
|
||||
for page_id in page_ids:
|
||||
page_text_list = self.read_page_as_documents(page_id)
|
||||
for page_text in page_text_list:
|
||||
docs.append(Document(page_text))
|
||||
|
||||
return docs
|
||||
|
||||
def get_page_last_edited_time(self, page_id: str) -> str:
|
||||
retrieve_page_url = RETRIEVE_PAGE_URL_TMPL.format(page_id=page_id)
|
||||
query_dict: Dict[str, Any] = {}
|
||||
|
||||
res = requests.request(
|
||||
"GET", retrieve_page_url, headers=self.headers, json=query_dict
|
||||
)
|
||||
data = res.json()
|
||||
return data["last_edited_time"]
|
||||
|
||||
def get_database_last_edited_time(self, database_id: str) -> str:
|
||||
retrieve_page_url = RETRIEVE_DATABASE_URL_TMPL.format(database_id=database_id)
|
||||
query_dict: Dict[str, Any] = {}
|
||||
|
||||
res = requests.request(
|
||||
"GET", retrieve_page_url, headers=self.headers, json=query_dict
|
||||
)
|
||||
data = res.json()
|
||||
return data["last_edited_time"]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
reader = NotionPageReader()
|
||||
logger.info(reader.search("What I"))
|
||||
|
|
@ -5,6 +5,8 @@ import tempfile
|
|||
import time
|
||||
from pathlib import Path
|
||||
from typing import Optional, List
|
||||
|
||||
from flask_login import current_user
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
|
||||
from llama_index import SimpleDirectoryReader
|
||||
|
|
@ -13,6 +15,8 @@ from llama_index.data_structs.node_v2 import DocumentRelationship
|
|||
from llama_index.node_parser import SimpleNodeParser, NodeParser
|
||||
from llama_index.readers.file.base import DEFAULT_FILE_EXTRACTOR
|
||||
from llama_index.readers.file.markdown_parser import MarkdownParser
|
||||
|
||||
from core.data_source.notion import NotionPageReader
|
||||
from core.index.readers.xlsx_parser import XLSXParser
|
||||
from core.docstore.dataset_docstore import DatesetDocumentStore
|
||||
from core.index.keyword_table_index import KeywordTableIndex
|
||||
|
|
@ -27,6 +31,7 @@ from extensions.ext_redis import redis_client
|
|||
from extensions.ext_storage import storage
|
||||
from models.dataset import Document, Dataset, DocumentSegment, DatasetProcessRule
|
||||
from models.model import UploadFile
|
||||
from models.source import DataSourceBinding
|
||||
|
||||
|
||||
class IndexingRunner:
|
||||
|
|
@ -35,42 +40,43 @@ class IndexingRunner:
|
|||
self.storage = storage
|
||||
self.embedding_model_name = embedding_model_name
|
||||
|
||||
def run(self, document: Document):
|
||||
def run(self, documents: List[Document]):
|
||||
"""Run the indexing process."""
|
||||
# get dataset
|
||||
dataset = Dataset.query.filter_by(
|
||||
id=document.dataset_id
|
||||
).first()
|
||||
for document in documents:
|
||||
# get dataset
|
||||
dataset = Dataset.query.filter_by(
|
||||
id=document.dataset_id
|
||||
).first()
|
||||
|
||||
if not dataset:
|
||||
raise ValueError("no dataset found")
|
||||
if not dataset:
|
||||
raise ValueError("no dataset found")
|
||||
|
||||
# load file
|
||||
text_docs = self._load_data(document)
|
||||
# load file
|
||||
text_docs = self._load_data(document)
|
||||
|
||||
# get the process rule
|
||||
processing_rule = db.session.query(DatasetProcessRule). \
|
||||
filter(DatasetProcessRule.id == document.dataset_process_rule_id). \
|
||||
first()
|
||||
# get the process rule
|
||||
processing_rule = db.session.query(DatasetProcessRule). \
|
||||
filter(DatasetProcessRule.id == document.dataset_process_rule_id). \
|
||||
first()
|
||||
|
||||
# get node parser for splitting
|
||||
node_parser = self._get_node_parser(processing_rule)
|
||||
# get node parser for splitting
|
||||
node_parser = self._get_node_parser(processing_rule)
|
||||
|
||||
# split to nodes
|
||||
nodes = self._step_split(
|
||||
text_docs=text_docs,
|
||||
node_parser=node_parser,
|
||||
dataset=dataset,
|
||||
document=document,
|
||||
processing_rule=processing_rule
|
||||
)
|
||||
# split to nodes
|
||||
nodes = self._step_split(
|
||||
text_docs=text_docs,
|
||||
node_parser=node_parser,
|
||||
dataset=dataset,
|
||||
document=document,
|
||||
processing_rule=processing_rule
|
||||
)
|
||||
|
||||
# build index
|
||||
self._build_index(
|
||||
dataset=dataset,
|
||||
document=document,
|
||||
nodes=nodes
|
||||
)
|
||||
# build index
|
||||
self._build_index(
|
||||
dataset=dataset,
|
||||
document=document,
|
||||
nodes=nodes
|
||||
)
|
||||
|
||||
def run_in_splitting_status(self, document: Document):
|
||||
"""Run the indexing process when the index_status is splitting."""
|
||||
|
|
@ -164,38 +170,98 @@ class IndexingRunner:
|
|||
nodes=nodes
|
||||
)
|
||||
|
||||
def indexing_estimate(self, file_detail: UploadFile, tmp_processing_rule: dict) -> dict:
|
||||
def file_indexing_estimate(self, file_details: List[UploadFile], tmp_processing_rule: dict) -> dict:
|
||||
"""
|
||||
Estimate the indexing for the document.
|
||||
"""
|
||||
# load data from file
|
||||
text_docs = self._load_data_from_file(file_detail)
|
||||
|
||||
processing_rule = DatasetProcessRule(
|
||||
mode=tmp_processing_rule["mode"],
|
||||
rules=json.dumps(tmp_processing_rule["rules"])
|
||||
)
|
||||
|
||||
# get node parser for splitting
|
||||
node_parser = self._get_node_parser(processing_rule)
|
||||
|
||||
# split to nodes
|
||||
nodes = self._split_to_nodes(
|
||||
text_docs=text_docs,
|
||||
node_parser=node_parser,
|
||||
processing_rule=processing_rule
|
||||
)
|
||||
|
||||
tokens = 0
|
||||
preview_texts = []
|
||||
for node in nodes:
|
||||
if len(preview_texts) < 5:
|
||||
preview_texts.append(node.get_text())
|
||||
total_segments = 0
|
||||
for file_detail in file_details:
|
||||
# load data from file
|
||||
text_docs = self._load_data_from_file(file_detail)
|
||||
|
||||
tokens += TokenCalculator.get_num_tokens(self.embedding_model_name, node.get_text())
|
||||
processing_rule = DatasetProcessRule(
|
||||
mode=tmp_processing_rule["mode"],
|
||||
rules=json.dumps(tmp_processing_rule["rules"])
|
||||
)
|
||||
|
||||
# get node parser for splitting
|
||||
node_parser = self._get_node_parser(processing_rule)
|
||||
|
||||
# split to nodes
|
||||
nodes = self._split_to_nodes(
|
||||
text_docs=text_docs,
|
||||
node_parser=node_parser,
|
||||
processing_rule=processing_rule
|
||||
)
|
||||
total_segments += len(nodes)
|
||||
for node in nodes:
|
||||
if len(preview_texts) < 5:
|
||||
preview_texts.append(node.get_text())
|
||||
|
||||
tokens += TokenCalculator.get_num_tokens(self.embedding_model_name, node.get_text())
|
||||
|
||||
return {
|
||||
"total_segments": len(nodes),
|
||||
"total_segments": total_segments,
|
||||
"tokens": tokens,
|
||||
"total_price": '{:f}'.format(TokenCalculator.get_token_price(self.embedding_model_name, tokens)),
|
||||
"currency": TokenCalculator.get_currency(self.embedding_model_name),
|
||||
"preview": preview_texts
|
||||
}
|
||||
|
||||
def notion_indexing_estimate(self, notion_info_list: list, tmp_processing_rule: dict) -> dict:
|
||||
"""
|
||||
Estimate the indexing for the document.
|
||||
"""
|
||||
# load data from notion
|
||||
tokens = 0
|
||||
preview_texts = []
|
||||
total_segments = 0
|
||||
for notion_info in notion_info_list:
|
||||
workspace_id = notion_info['workspace_id']
|
||||
data_source_binding = DataSourceBinding.query.filter(
|
||||
db.and_(
|
||||
DataSourceBinding.tenant_id == current_user.current_tenant_id,
|
||||
DataSourceBinding.provider == 'notion',
|
||||
DataSourceBinding.disabled == False,
|
||||
DataSourceBinding.source_info['workspace_id'] == f'"{workspace_id}"'
|
||||
)
|
||||
).first()
|
||||
if not data_source_binding:
|
||||
raise ValueError('Data source binding not found.')
|
||||
reader = NotionPageReader(integration_token=data_source_binding.access_token)
|
||||
for page in notion_info['pages']:
|
||||
if page['type'] == 'page':
|
||||
page_ids = [page['page_id']]
|
||||
documents = reader.load_data_as_documents(page_ids=page_ids)
|
||||
elif page['type'] == 'database':
|
||||
documents = reader.load_data_as_documents(database_id=page['page_id'])
|
||||
else:
|
||||
documents = []
|
||||
processing_rule = DatasetProcessRule(
|
||||
mode=tmp_processing_rule["mode"],
|
||||
rules=json.dumps(tmp_processing_rule["rules"])
|
||||
)
|
||||
|
||||
# get node parser for splitting
|
||||
node_parser = self._get_node_parser(processing_rule)
|
||||
|
||||
# split to nodes
|
||||
nodes = self._split_to_nodes(
|
||||
text_docs=documents,
|
||||
node_parser=node_parser,
|
||||
processing_rule=processing_rule
|
||||
)
|
||||
total_segments += len(nodes)
|
||||
for node in nodes:
|
||||
if len(preview_texts) < 5:
|
||||
preview_texts.append(node.get_text())
|
||||
|
||||
tokens += TokenCalculator.get_num_tokens(self.embedding_model_name, node.get_text())
|
||||
|
||||
return {
|
||||
"total_segments": total_segments,
|
||||
"tokens": tokens,
|
||||
"total_price": '{:f}'.format(TokenCalculator.get_token_price(self.embedding_model_name, tokens)),
|
||||
"currency": TokenCalculator.get_currency(self.embedding_model_name),
|
||||
|
|
@ -204,25 +270,50 @@ class IndexingRunner:
|
|||
|
||||
def _load_data(self, document: Document) -> List[Document]:
|
||||
# load file
|
||||
if document.data_source_type != "upload_file":
|
||||
if document.data_source_type not in ["upload_file", "notion_import"]:
|
||||
return []
|
||||
|
||||
data_source_info = document.data_source_info_dict
|
||||
if not data_source_info or 'upload_file_id' not in data_source_info:
|
||||
raise ValueError("no upload file found")
|
||||
text_docs = []
|
||||
if document.data_source_type == 'upload_file':
|
||||
if not data_source_info or 'upload_file_id' not in data_source_info:
|
||||
raise ValueError("no upload file found")
|
||||
|
||||
file_detail = db.session.query(UploadFile). \
|
||||
filter(UploadFile.id == data_source_info['upload_file_id']). \
|
||||
one_or_none()
|
||||
|
||||
text_docs = self._load_data_from_file(file_detail)
|
||||
file_detail = db.session.query(UploadFile). \
|
||||
filter(UploadFile.id == data_source_info['upload_file_id']). \
|
||||
one_or_none()
|
||||
|
||||
text_docs = self._load_data_from_file(file_detail)
|
||||
elif document.data_source_type == 'notion_import':
|
||||
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")
|
||||
workspace_id = data_source_info['notion_workspace_id']
|
||||
page_id = data_source_info['notion_page_id']
|
||||
page_type = data_source_info['type']
|
||||
data_source_binding = DataSourceBinding.query.filter(
|
||||
db.and_(
|
||||
DataSourceBinding.tenant_id == document.tenant_id,
|
||||
DataSourceBinding.provider == 'notion',
|
||||
DataSourceBinding.disabled == False,
|
||||
DataSourceBinding.source_info['workspace_id'] == f'"{workspace_id}"'
|
||||
)
|
||||
).first()
|
||||
if not data_source_binding:
|
||||
raise ValueError('Data source binding not found.')
|
||||
if page_type == 'page':
|
||||
# add page last_edited_time to data_source_info
|
||||
self._get_notion_page_last_edited_time(page_id, data_source_binding.access_token, document)
|
||||
text_docs = self._load_page_data_from_notion(page_id, data_source_binding.access_token)
|
||||
elif page_type == 'database':
|
||||
# add page last_edited_time to data_source_info
|
||||
self._get_notion_database_last_edited_time(page_id, data_source_binding.access_token, document)
|
||||
text_docs = self._load_database_data_from_notion(page_id, data_source_binding.access_token)
|
||||
# update document status to splitting
|
||||
self._update_document_index_status(
|
||||
document_id=document.id,
|
||||
after_indexing_status="splitting",
|
||||
extra_update_params={
|
||||
Document.file_id: file_detail.id,
|
||||
Document.word_count: sum([len(text_doc.text) for text_doc in text_docs]),
|
||||
Document.parsing_completed_at: datetime.datetime.utcnow()
|
||||
}
|
||||
|
|
@ -259,6 +350,41 @@ class IndexingRunner:
|
|||
|
||||
return text_docs
|
||||
|
||||
def _load_page_data_from_notion(self, page_id: str, access_token: str) -> List[Document]:
|
||||
page_ids = [page_id]
|
||||
reader = NotionPageReader(integration_token=access_token)
|
||||
text_docs = reader.load_data_as_documents(page_ids=page_ids)
|
||||
return text_docs
|
||||
|
||||
def _load_database_data_from_notion(self, database_id: str, access_token: str) -> List[Document]:
|
||||
reader = NotionPageReader(integration_token=access_token)
|
||||
text_docs = reader.load_data_as_documents(database_id=database_id)
|
||||
return text_docs
|
||||
|
||||
def _get_notion_page_last_edited_time(self, page_id: str, access_token: str, document: Document):
|
||||
reader = NotionPageReader(integration_token=access_token)
|
||||
last_edited_time = reader.get_page_last_edited_time(page_id)
|
||||
data_source_info = document.data_source_info_dict
|
||||
data_source_info['last_edited_time'] = last_edited_time
|
||||
update_params = {
|
||||
Document.data_source_info: json.dumps(data_source_info)
|
||||
}
|
||||
|
||||
Document.query.filter_by(id=document.id).update(update_params)
|
||||
db.session.commit()
|
||||
|
||||
def _get_notion_database_last_edited_time(self, page_id: str, access_token: str, document: Document):
|
||||
reader = NotionPageReader(integration_token=access_token)
|
||||
last_edited_time = reader.get_database_last_edited_time(page_id)
|
||||
data_source_info = document.data_source_info_dict
|
||||
data_source_info['last_edited_time'] = last_edited_time
|
||||
update_params = {
|
||||
Document.data_source_info: json.dumps(data_source_info)
|
||||
}
|
||||
|
||||
Document.query.filter_by(id=document.id).update(update_params)
|
||||
db.session.commit()
|
||||
|
||||
def _get_node_parser(self, processing_rule: DatasetProcessRule) -> NodeParser:
|
||||
"""
|
||||
Get the NodeParser object according to the processing rule.
|
||||
|
|
@ -308,7 +434,7 @@ class IndexingRunner:
|
|||
embedding_model_name=self.embedding_model_name,
|
||||
document_id=document.id
|
||||
)
|
||||
|
||||
# add document segments
|
||||
doc_store.add_documents(nodes)
|
||||
|
||||
# update document status to indexing
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue