397 lines
17 KiB
Python
397 lines
17 KiB
Python
|
import json
|
||
|
|
from collections.abc import Generator
|
||
|
|
from typing import Any, Literal, cast
|
||
|
|
from uuid import UUID
|
||
|
|
|
||
|
|
from flask import request
|
||
|
|
from flask_restx import Resource, fields, marshal_with
|
||
|
|
from pydantic import BaseModel, Field
|
||
|
|
from sqlalchemy import select
|
||
|
|
from sqlalchemy.orm import sessionmaker
|
||
|
|
from werkzeug.exceptions import NotFound
|
||
|
|
|
||
|
|
from controllers.common.fields import SimpleResultResponse, TextContentResponse
|
||
|
|
from controllers.common.schema import get_or_create_model, register_response_schema_models, register_schema_model
|
||
|
|
from core.datasource.entities.datasource_entities import DatasourceProviderType, OnlineDocumentPagesMessage
|
||
|
|
from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin
|
||
|
|
from core.indexing_runner import IndexingRunner
|
||
|
|
from core.rag.extractor.entity.datasource_type import DatasourceType
|
||
|
|
from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo
|
||
|
|
from core.rag.extractor.notion_extractor import NotionExtractor
|
||
|
|
from extensions.ext_database import db
|
||
|
|
from fields.data_source_fields import (
|
||
|
|
integrate_fields,
|
||
|
|
integrate_icon_fields,
|
||
|
|
integrate_list_fields,
|
||
|
|
integrate_notion_info_list_fields,
|
||
|
|
integrate_page_fields,
|
||
|
|
integrate_workspace_fields,
|
||
|
|
)
|
||
|
|
from libs.datetime_utils import naive_utc_now
|
||
|
|
from libs.login import current_account_with_tenant, login_required
|
||
|
|
from models import DataSourceOauthBinding, Document
|
||
|
|
from services.dataset_service import DatasetService, DocumentService
|
||
|
|
from services.datasource_provider_service import DatasourceProviderService
|
||
|
|
from tasks.document_indexing_sync_task import document_indexing_sync_task
|
||
|
|
|
||
|
|
from .. import console_ns
|
||
|
|
from ..wraps import account_initialization_required, setup_required
|
||
|
|
|
||
|
|
|
||
|
|
class NotionEstimatePayload(BaseModel):
|
||
|
|
notion_info_list: list[dict[str, Any]]
|
||
|
|
process_rule: dict[str, Any]
|
||
|
|
doc_form: str = Field(default="text_model")
|
||
|
|
doc_language: str = Field(default="English")
|
||
|
|
|
||
|
|
|
||
|
|
class DataSourceNotionListQuery(BaseModel):
|
||
|
|
dataset_id: str | None = Field(default=None, description="Dataset ID")
|
||
|
|
credential_id: str = Field(..., description="Credential ID", min_length=1)
|
||
|
|
|
||
|
|
|
||
|
|
class DataSourceNotionPreviewQuery(BaseModel):
|
||
|
|
credential_id: str = Field(..., description="Credential ID", min_length=1)
|
||
|
|
|
||
|
|
|
||
|
|
register_schema_model(console_ns, NotionEstimatePayload)
|
||
|
|
register_response_schema_models(console_ns, SimpleResultResponse, TextContentResponse)
|
||
|
|
|
||
|
|
|
||
|
|
integrate_icon_model = get_or_create_model("DataSourceIntegrateIcon", integrate_icon_fields)
|
||
|
|
|
||
|
|
integrate_page_fields_copy = integrate_page_fields.copy()
|
||
|
|
integrate_page_fields_copy["page_icon"] = fields.Nested(integrate_icon_model, allow_null=True)
|
||
|
|
integrate_page_model = get_or_create_model("DataSourceIntegratePage", integrate_page_fields_copy)
|
||
|
|
|
||
|
|
integrate_workspace_fields_copy = integrate_workspace_fields.copy()
|
||
|
|
integrate_workspace_fields_copy["pages"] = fields.List(fields.Nested(integrate_page_model))
|
||
|
|
integrate_workspace_model = get_or_create_model("DataSourceIntegrateWorkspace", integrate_workspace_fields_copy)
|
||
|
|
|
||
|
|
integrate_fields_copy = integrate_fields.copy()
|
||
|
|
integrate_fields_copy["source_info"] = fields.Nested(integrate_workspace_model)
|
||
|
|
integrate_model = get_or_create_model("DataSourceIntegrate", integrate_fields_copy)
|
||
|
|
|
||
|
|
integrate_list_fields_copy = integrate_list_fields.copy()
|
||
|
|
integrate_list_fields_copy["data"] = fields.List(fields.Nested(integrate_model))
|
||
|
|
integrate_list_model = get_or_create_model("DataSourceIntegrateList", integrate_list_fields_copy)
|
||
|
|
|
||
|
|
notion_page_fields = {
|
||
|
|
"page_name": fields.String,
|
||
|
|
"page_id": fields.String,
|
||
|
|
"page_icon": fields.Nested(integrate_icon_model, allow_null=True),
|
||
|
|
"is_bound": fields.Boolean,
|
||
|
|
"parent_id": fields.String,
|
||
|
|
"type": fields.String,
|
||
|
|
}
|
||
|
|
notion_page_model = get_or_create_model("NotionIntegratePage", notion_page_fields)
|
||
|
|
|
||
|
|
notion_workspace_fields = {
|
||
|
|
"workspace_name": fields.String,
|
||
|
|
"workspace_id": fields.String,
|
||
|
|
"workspace_icon": fields.String,
|
||
|
|
"pages": fields.List(fields.Nested(notion_page_model)),
|
||
|
|
}
|
||
|
|
notion_workspace_model = get_or_create_model("NotionIntegrateWorkspace", notion_workspace_fields)
|
||
|
|
|
||
|
|
integrate_notion_info_list_fields_copy = integrate_notion_info_list_fields.copy()
|
||
|
|
integrate_notion_info_list_fields_copy["notion_info"] = fields.List(fields.Nested(notion_workspace_model))
|
||
|
|
integrate_notion_info_list_model = get_or_create_model(
|
||
|
|
"NotionIntegrateInfoList", integrate_notion_info_list_fields_copy
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@console_ns.route(
|
||
|
|
"/data-source/integrates",
|
||
|
|
"/data-source/integrates/<uuid:binding_id>/<string:action>",
|
||
|
|
)
|
||
|
|
class DataSourceApi(Resource):
|
||
|
|
@setup_required
|
||
|
|
@login_required
|
||
|
|
@account_initialization_required
|
||
|
|
@marshal_with(integrate_list_model)
|
||
|
|
def get(self):
|
||
|
|
_, current_tenant_id = current_account_with_tenant()
|
||
|
|
|
||
|
|
# get workspace data source integrates
|
||
|
|
data_source_integrates = db.session.scalars(
|
||
|
|
select(DataSourceOauthBinding).where(
|
||
|
|
DataSourceOauthBinding.tenant_id == current_tenant_id,
|
||
|
|
DataSourceOauthBinding.disabled == False,
|
||
|
|
)
|
||
|
|
).all()
|
||
|
|
|
||
|
|
base_url = request.url_root.rstrip("/")
|
||
|
|
data_source_oauth_base_path = "/console/api/oauth/data-source"
|
||
|
|
providers = ["notion"]
|
||
|
|
|
||
|
|
integrate_data = []
|
||
|
|
for provider in providers:
|
||
|
|
# existing_integrate = next((ai for ai in data_source_integrates if ai.provider == provider), None)
|
||
|
|
existing_integrates = filter(lambda item: item.provider == provider, data_source_integrates)
|
||
|
|
if existing_integrates:
|
||
|
|
for existing_integrate in list(existing_integrates):
|
||
|
|
integrate_data.append(
|
||
|
|
{
|
||
|
|
"id": existing_integrate.id,
|
||
|
|
"provider": provider,
|
||
|
|
"created_at": existing_integrate.created_at,
|
||
|
|
"is_bound": True,
|
||
|
|
"disabled": existing_integrate.disabled,
|
||
|
|
"source_info": existing_integrate.source_info,
|
||
|
|
"link": f"{base_url}{data_source_oauth_base_path}/{provider}",
|
||
|
|
}
|
||
|
|
)
|
||
|
|
else:
|
||
|
|
integrate_data.append(
|
||
|
|
{
|
||
|
|
"id": None,
|
||
|
|
"provider": provider,
|
||
|
|
"created_at": None,
|
||
|
|
"source_info": None,
|
||
|
|
"is_bound": False,
|
||
|
|
"disabled": None,
|
||
|
|
"link": f"{base_url}{data_source_oauth_base_path}/{provider}",
|
||
|
|
}
|
||
|
|
)
|
||
|
|
return {"data": integrate_data}, 200
|
||
|
|
|
||
|
|
@setup_required
|
||
|
|
@login_required
|
||
|
|
@account_initialization_required
|
||
|
|
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
|
||
|
|
def patch(self, binding_id, action: Literal["enable", "disable"]):
|
||
|
|
_, current_tenant_id = current_account_with_tenant()
|
||
|
|
binding_id = str(binding_id)
|
||
|
|
with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
|
||
|
|
data_source_binding = session.execute(
|
||
|
|
select(DataSourceOauthBinding).where(
|
||
|
|
DataSourceOauthBinding.id == binding_id, DataSourceOauthBinding.tenant_id == current_tenant_id
|
||
|
|
)
|
||
|
|
).scalar_one_or_none()
|
||
|
|
if data_source_binding is None:
|
||
|
|
raise NotFound("Data source binding not found.")
|
||
|
|
# enable binding
|
||
|
|
match action:
|
||
|
|
case "enable":
|
||
|
|
if data_source_binding.disabled:
|
||
|
|
data_source_binding.disabled = False
|
||
|
|
data_source_binding.updated_at = naive_utc_now()
|
||
|
|
db.session.add(data_source_binding)
|
||
|
|
db.session.commit()
|
||
|
|
else:
|
||
|
|
raise ValueError("Data source is not disabled.")
|
||
|
|
# disable binding
|
||
|
|
case "disable":
|
||
|
|
if not data_source_binding.disabled:
|
||
|
|
data_source_binding.disabled = True
|
||
|
|
data_source_binding.updated_at = naive_utc_now()
|
||
|
|
db.session.add(data_source_binding)
|
||
|
|
db.session.commit()
|
||
|
|
else:
|
||
|
|
raise ValueError("Data source is disabled.")
|
||
|
|
return {"result": "success"}, 200
|
||
|
|
|
||
|
|
|
||
|
|
@console_ns.route("/notion/pre-import/pages")
|
||
|
|
class DataSourceNotionListApi(Resource):
|
||
|
|
@setup_required
|
||
|
|
@login_required
|
||
|
|
@account_initialization_required
|
||
|
|
@marshal_with(integrate_notion_info_list_model)
|
||
|
|
def get(self):
|
||
|
|
current_user, current_tenant_id = current_account_with_tenant()
|
||
|
|
|
||
|
|
query = DataSourceNotionListQuery.model_validate(request.args.to_dict())
|
||
|
|
|
||
|
|
datasource_provider_service = DatasourceProviderService()
|
||
|
|
credential = datasource_provider_service.get_datasource_credentials(
|
||
|
|
tenant_id=current_tenant_id,
|
||
|
|
credential_id=query.credential_id,
|
||
|
|
provider="notion_datasource",
|
||
|
|
plugin_id="langgenius/notion_datasource",
|
||
|
|
)
|
||
|
|
if not credential:
|
||
|
|
raise NotFound("Credential not found.")
|
||
|
|
exist_page_ids = []
|
||
|
|
with sessionmaker(db.engine).begin() as session:
|
||
|
|
# import notion in the exist dataset
|
||
|
|
if query.dataset_id:
|
||
|
|
dataset = DatasetService.get_dataset(query.dataset_id)
|
||
|
|
if not dataset:
|
||
|
|
raise NotFound("Dataset not found.")
|
||
|
|
if dataset.data_source_type != "notion_import":
|
||
|
|
raise ValueError("Dataset is not notion type.")
|
||
|
|
|
||
|
|
documents = session.scalars(
|
||
|
|
select(Document).where(
|
||
|
|
Document.dataset_id == query.dataset_id,
|
||
|
|
Document.tenant_id == current_tenant_id,
|
||
|
|
Document.data_source_type == "notion_import",
|
||
|
|
Document.enabled.is_(True),
|
||
|
|
)
|
||
|
|
).all()
|
||
|
|
if documents:
|
||
|
|
for document in documents:
|
||
|
|
data_source_info = json.loads(document.data_source_info)
|
||
|
|
exist_page_ids.append(data_source_info["notion_page_id"])
|
||
|
|
# get all authorized pages
|
||
|
|
from core.datasource.datasource_manager import DatasourceManager
|
||
|
|
|
||
|
|
datasource_runtime = DatasourceManager.get_datasource_runtime(
|
||
|
|
provider_id="langgenius/notion_datasource/notion_datasource",
|
||
|
|
datasource_name="notion_datasource",
|
||
|
|
tenant_id=current_tenant_id,
|
||
|
|
datasource_type=DatasourceProviderType.ONLINE_DOCUMENT,
|
||
|
|
)
|
||
|
|
datasource_provider_service = DatasourceProviderService()
|
||
|
|
if credential:
|
||
|
|
datasource_runtime.runtime.credentials = credential
|
||
|
|
datasource_runtime = cast(OnlineDocumentDatasourcePlugin, datasource_runtime)
|
||
|
|
online_document_result: Generator[OnlineDocumentPagesMessage, None, None] = (
|
||
|
|
datasource_runtime.get_online_document_pages(
|
||
|
|
user_id=current_user.id,
|
||
|
|
datasource_parameters={},
|
||
|
|
provider_type=datasource_runtime.datasource_provider_type(),
|
||
|
|
)
|
||
|
|
)
|
||
|
|
try:
|
||
|
|
pages = []
|
||
|
|
workspace_info = {}
|
||
|
|
for message in online_document_result:
|
||
|
|
result = message.result
|
||
|
|
for info in result:
|
||
|
|
workspace_info = {
|
||
|
|
"workspace_id": info.workspace_id,
|
||
|
|
"workspace_name": info.workspace_name,
|
||
|
|
"workspace_icon": info.workspace_icon,
|
||
|
|
}
|
||
|
|
for page in info.pages:
|
||
|
|
page_info = {
|
||
|
|
"page_id": page.page_id,
|
||
|
|
"page_name": page.page_name,
|
||
|
|
"type": page.type,
|
||
|
|
"parent_id": page.parent_id,
|
||
|
|
"is_bound": page.page_id in exist_page_ids,
|
||
|
|
"page_icon": page.page_icon,
|
||
|
|
}
|
||
|
|
pages.append(page_info)
|
||
|
|
except Exception as e:
|
||
|
|
raise e
|
||
|
|
return {"notion_info": {**workspace_info, "pages": pages}}, 200
|
||
|
|
|
||
|
|
|
||
|
|
@console_ns.route(
|
||
|
|
"/notion/pages/<uuid:page_id>/<string:page_type>/preview",
|
||
|
|
"/datasets/notion-indexing-estimate",
|
||
|
|
)
|
||
|
|
class DataSourceNotionApi(Resource):
|
||
|
|
@setup_required
|
||
|
|
@login_required
|
||
|
|
@account_initialization_required
|
||
|
|
@console_ns.response(200, "Success", console_ns.models[TextContentResponse.__name__])
|
||
|
|
def get(self, page_id: UUID, page_type: str):
|
||
|
|
_, current_tenant_id = current_account_with_tenant()
|
||
|
|
|
||
|
|
query = DataSourceNotionPreviewQuery.model_validate(request.args.to_dict())
|
||
|
|
|
||
|
|
datasource_provider_service = DatasourceProviderService()
|
||
|
|
credential = datasource_provider_service.get_datasource_credentials(
|
||
|
|
tenant_id=current_tenant_id,
|
||
|
|
credential_id=query.credential_id,
|
||
|
|
provider="notion_datasource",
|
||
|
|
plugin_id="langgenius/notion_datasource",
|
||
|
|
)
|
||
|
|
|
||
|
|
page_id_str = str(page_id)
|
||
|
|
|
||
|
|
extractor = NotionExtractor(
|
||
|
|
notion_workspace_id="",
|
||
|
|
notion_obj_id=page_id_str,
|
||
|
|
notion_page_type=page_type,
|
||
|
|
notion_access_token=credential.get("integration_secret"),
|
||
|
|
tenant_id=current_tenant_id,
|
||
|
|
)
|
||
|
|
|
||
|
|
text_docs = extractor.extract()
|
||
|
|
return {"content": "\n".join([doc.page_content for doc in text_docs])}, 200
|
||
|
|
|
||
|
|
@setup_required
|
||
|
|
@login_required
|
||
|
|
@account_initialization_required
|
||
|
|
@console_ns.expect(console_ns.models[NotionEstimatePayload.__name__])
|
||
|
|
def post(self):
|
||
|
|
_, current_tenant_id = current_account_with_tenant()
|
||
|
|
|
||
|
|
payload = NotionEstimatePayload.model_validate(console_ns.payload or {})
|
||
|
|
args = payload.model_dump()
|
||
|
|
# validate args
|
||
|
|
DocumentService.estimate_args_validate(args)
|
||
|
|
notion_info_list = payload.notion_info_list
|
||
|
|
extract_settings = []
|
||
|
|
for notion_info in notion_info_list:
|
||
|
|
workspace_id = notion_info["workspace_id"]
|
||
|
|
credential_id = notion_info.get("credential_id")
|
||
|
|
for page in notion_info["pages"]:
|
||
|
|
extract_setting = ExtractSetting(
|
||
|
|
datasource_type=DatasourceType.NOTION,
|
||
|
|
notion_info=NotionInfo.model_validate(
|
||
|
|
{
|
||
|
|
"credential_id": credential_id,
|
||
|
|
"notion_workspace_id": workspace_id,
|
||
|
|
"notion_obj_id": page["page_id"],
|
||
|
|
"notion_page_type": page["type"],
|
||
|
|
"tenant_id": current_tenant_id,
|
||
|
|
}
|
||
|
|
),
|
||
|
|
document_model=args["doc_form"],
|
||
|
|
)
|
||
|
|
extract_settings.append(extract_setting)
|
||
|
|
indexing_runner = IndexingRunner()
|
||
|
|
response = indexing_runner.indexing_estimate(
|
||
|
|
current_tenant_id,
|
||
|
|
extract_settings,
|
||
|
|
args["process_rule"],
|
||
|
|
args["doc_form"],
|
||
|
|
args["doc_language"],
|
||
|
|
)
|
||
|
|
return response.model_dump(), 200
|
||
|
|
|
||
|
|
|
||
|
|
@console_ns.route("/datasets/<uuid:dataset_id>/notion/sync")
|
||
|
|
class DataSourceNotionDatasetSyncApi(Resource):
|
||
|
|
@setup_required
|
||
|
|
@login_required
|
||
|
|
@account_initialization_required
|
||
|
|
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
|
||
|
|
def get(self, dataset_id: UUID):
|
||
|
|
dataset_id_str = str(dataset_id)
|
||
|
|
dataset = DatasetService.get_dataset(dataset_id_str)
|
||
|
|
if dataset is None:
|
||
|
|
raise NotFound("Dataset not found.")
|
||
|
|
|
||
|
|
documents = DocumentService.get_document_by_dataset_id(dataset_id_str)
|
||
|
|
for document in documents:
|
||
|
|
document_indexing_sync_task.delay(dataset_id_str, document.id)
|
||
|
|
return {"result": "success"}, 200
|
||
|
|
|
||
|
|
|
||
|
|
@console_ns.route("/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/notion/sync")
|
||
|
|
class DataSourceNotionDocumentSyncApi(Resource):
|
||
|
|
@setup_required
|
||
|
|
@login_required
|
||
|
|
@account_initialization_required
|
||
|
|
@console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
|
||
|
|
def get(self, dataset_id: UUID, document_id: UUID):
|
||
|
|
dataset_id_str = str(dataset_id)
|
||
|
|
document_id_str = str(document_id)
|
||
|
|
dataset = DatasetService.get_dataset(dataset_id_str)
|
||
|
|
if dataset is None:
|
||
|
|
raise NotFound("Dataset not found.")
|
||
|
|
|
||
|
|
document = DocumentService.get_document(dataset_id_str, document_id_str)
|
||
|
|
if document is None:
|
||
|
|
raise NotFound("Document not found.")
|
||
|
|
document_indexing_sync_task.delay(dataset_id_str, document_id_str)
|
||
|
|
return {"result": "success"}, 200
|