estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

feat: update notion extractor (#3898)

Browse Source 
Co-authored-by: duyalei <>
This commit is contained in:
yalei
2024-05-24 20:30:48 +08:00
committed by GitHub
parent f156014daa
commit 026175c8f7
3 changed files with 118 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
api/core/rag/extractor/notion_extractor.py
View File

@@ -19,8 +19,12 @@ 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']
# if user want split by headings, use the corresponding splitter
HEADING_SPLITTER = {
'heading_1': '# ',
'heading_2': '## ',
'heading_3': '### ',
}
class NotionExtractor(BaseExtractor):
@@ -73,8 +77,7 @@ class NotionExtractor(BaseExtractor):
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))
docs.append(Document(page_content='\n'.join(page_text_list)))
else:
raise ValueError("notion page type not supported")
@@ -96,7 +99,7 @@ class NotionExtractor(BaseExtractor):
data = res.json()
database_content_list = []
database_content = []
if 'results' not in data or data["results"] is None:
return []
for result in data["results"]:
@@ -131,10 +134,9 @@ class NotionExtractor(BaseExtractor):
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)
database_content.append(row_content)
return database_content_list
return [Document(page_content='\n'.join(database_content))]
def _get_notion_block_data(self, page_id: str) -> list[str]:
result_lines_arr = []
@@ -154,8 +156,6 @@ class NotionExtractor(BaseExtractor):
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]
@@ -172,8 +172,6 @@ class NotionExtractor(BaseExtractor):
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"]
@@ -185,11 +183,10 @@ class NotionExtractor(BaseExtractor):
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)
if result_type in HEADING_SPLITTER:
result_lines_arr.append(f"{HEADING_SPLITTER[result_type]}{cur_result_text}")
else:
result_lines_arr.append(f'{heading}\n{cur_result_text}')
result_lines_arr.append(cur_result_text + '\n\n')
if data["next_cursor"] is None:
break
@@ -218,7 +215,6 @@ class NotionExtractor(BaseExtractor):
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]
@@ -235,8 +231,6 @@ class NotionExtractor(BaseExtractor):
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"]
@@ -247,10 +241,10 @@ class NotionExtractor(BaseExtractor):
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)
if result_type in HEADING_SPLITTER:
result_lines_arr.append(f'{HEADING_SPLITTER[result_type]}{cur_result_text}')
else:
result_lines_arr.append(f'{heading}\n{cur_result_text}')
result_lines_arr.append(cur_result_text + '\n\n')
if data["next_cursor"] is None:
break