Feat/firecrawl data source (#5232)

Co-authored-by: Nicolas <[email protected]>
Co-authored-by: chenhe <[email protected]>
Co-authored-by: takatost <[email protected]>

api/.env.example

@@ -215,4 +215,5 @@ WORKFLOW_MAX_EXECUTION_TIME=1200
 WORKFLOW_CALL_MAX_DEPTH=5
 
 # App configuration
-APP_MAX_EXECUTION_TIME=1200
+APP_MAX_EXECUTION_TIME=1200
+

api/controllers/console/__init__.py

@@ -29,13 +29,13 @@
 )
 
 # Import auth controllers
-from .auth import activate, data_source_oauth, login, oauth
+from .auth import activate, data_source_bearer_auth, data_source_oauth, login, oauth
 
 # Import billing controllers
 from .billing import billing
 
 # Import datasets controllers
-from .datasets import data_source, datasets, datasets_document, datasets_segments, file, hit_testing
+from .datasets import data_source, datasets, datasets_document, datasets_segments, file, hit_testing, website
 
 # Import explore controllers
 from .explore import (

api/controllers/console/auth/data_source_bearer_auth.py

@@ -0,0 +1,67 @@
+from flask_login import current_user
+from flask_restful import Resource, reqparse
+from werkzeug.exceptions import Forbidden
+
+from controllers.console import api
+from controllers.console.auth.error import ApiKeyAuthFailedError
+from libs.login import login_required
+from services.auth.api_key_auth_service import ApiKeyAuthService
+
+from ..setup import setup_required
+from ..wraps import account_initialization_required
+
+
+class ApiKeyAuthDataSource(Resource):
+ @setup_required
+ @login_required
+ @account_initialization_required
+ def get(self):
+ # The role of the current user in the table must be admin or owner
+ if not current_user.is_admin_or_owner:
+ raise Forbidden()
+ data_source_api_key_bindings = ApiKeyAuthService.get_provider_auth_list(current_user.current_tenant_id)
+ if data_source_api_key_bindings:
+ return {
+ 'settings': [data_source_api_key_binding.to_dict() for data_source_api_key_binding in
+ data_source_api_key_bindings]}
+ return {'settings': []}
+
+
+class ApiKeyAuthDataSourceBinding(Resource):
+ @setup_required
+ @login_required
+ @account_initialization_required
+ def post(self):
+ # The role of the current user in the table must be admin or owner
+ if not current_user.is_admin_or_owner:
+ raise Forbidden()
+ parser = reqparse.RequestParser()
+ parser.add_argument('category', type=str, required=True, nullable=False, location='json')
+ parser.add_argument('provider', type=str, required=True, nullable=False, location='json')
+ parser.add_argument('credentials', type=dict, required=True, nullable=False, location='json')
+ args = parser.parse_args()
+ ApiKeyAuthService.validate_api_key_auth_args(args)
+ try:
+ ApiKeyAuthService.create_provider_auth(current_user.current_tenant_id, args)
+ except Exception as e:
+ raise ApiKeyAuthFailedError(str(e))
+ return {'result': 'success'}, 200
+
+
+class ApiKeyAuthDataSourceBindingDelete(Resource):
+ @setup_required
+ @login_required
+ @account_initialization_required
+ def delete(self, binding_id):
+ # The role of the current user in the table must be admin or owner
+ if not current_user.is_admin_or_owner:
+ raise Forbidden()
+
+ ApiKeyAuthService.delete_provider_auth(current_user.current_tenant_id, binding_id)
+
+ return {'result': 'success'}, 200
+
+
+api.add_resource(ApiKeyAuthDataSource, '/api-key-auth/data-source')
+api.add_resource(ApiKeyAuthDataSourceBinding, '/api-key-auth/data-source/binding')
+api.add_resource(ApiKeyAuthDataSourceBindingDelete, '/api-key-auth/data-source/<uuid:binding_id>')

api/controllers/console/auth/error.py

@@ -0,0 +1,7 @@
+from libs.exception import BaseHTTPException
+
+
+class ApiKeyAuthFailedError(BaseHTTPException):
+ error_code = 'auth_failed'
+ description = "{message}"
+ code = 500

api/controllers/console/datasets/data_source.py

@@ -16,7 +16,7 @@
 from fields.data_source_fields import integrate_list_fields, integrate_notion_info_list_fields
 from libs.login import login_required
 from models.dataset import Document
-from models.source import DataSourceBinding
+from models.source import DataSourceOauthBinding
 from services.dataset_service import DatasetService, DocumentService
 from tasks.document_indexing_sync_task import document_indexing_sync_task
 
@@ -29,9 +29,9 @@ class DataSourceApi(Resource):
  @marshal_with(integrate_list_fields)
  def get(self):
  # get workspace data source integrates
- data_source_integrates = db.session.query(DataSourceBinding).filter(
- DataSourceBinding.tenant_id == current_user.current_tenant_id,
- DataSourceBinding.disabled == False
+ data_source_integrates = db.session.query(DataSourceOauthBinding).filter(
+ DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
+ DataSourceOauthBinding.disabled == False
  ).all()
 
  base_url = request.url_root.rstrip('/')
@@ -71,7 +71,7 @@ def get(self):
  def patch(self, binding_id, action):
  binding_id = str(binding_id)
  action = str(action)
- data_source_binding = DataSourceBinding.query.filter_by(
+ data_source_binding = DataSourceOauthBinding.query.filter_by(
  id=binding_id
  ).first()
  if data_source_binding is None:
@@ -124,7 +124,7 @@ def get(self):
  data_source_info = json.loads(document.data_source_info)
  exist_page_ids.append(data_source_info['notion_page_id'])
  # get all authorized pages
- data_source_bindings = DataSourceBinding.query.filter_by(
+ data_source_bindings = DataSourceOauthBinding.query.filter_by(
  tenant_id=current_user.current_tenant_id,
  provider='notion',
  disabled=False
@@ -163,12 +163,12 @@ class DataSourceNotionApi(Resource):
  def get(self, workspace_id, page_id, page_type):
  workspace_id = str(workspace_id)
  page_id = str(page_id)
- data_source_binding = DataSourceBinding.query.filter(
+ data_source_binding = DataSourceOauthBinding.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}"'
+ DataSourceOauthBinding.tenant_id == current_user.current_tenant_id,
+ DataSourceOauthBinding.provider == 'notion',
+ DataSourceOauthBinding.disabled == False,
+ DataSourceOauthBinding.source_info['workspace_id'] == f'"{workspace_id}"'
  )
  ).first()
  if not data_source_binding:

api/controllers/console/datasets/datasets.py

@@ -315,6 +315,22 @@ def post(self):
  document_model=args['doc_form']
  )
  extract_settings.append(extract_setting)
+ elif args['info_list']['data_source_type'] == 'website_crawl':
+ website_info_list = args['info_list']['website_info_list']
+ for url in website_info_list['urls']:
+ extract_setting = ExtractSetting(
+ datasource_type="website_crawl",
+ website_info={
+ "provider": website_info_list['provider'],
+ "job_id": website_info_list['job_id'],
+ "url": url,
+ "tenant_id": current_user.current_tenant_id,
+ "mode": 'crawl',
+ "only_main_content": website_info_list['only_main_content']
+ },
+ document_model=args['doc_form']
+ )
+ extract_settings.append(extract_setting)
  else:
  raise ValueError('Data source type not support')
  indexing_runner = IndexingRunner()
@@ -519,6 +535,7 @@ def get(self, vector_type):
  raise ValueError(f"Unsupported vector db type {vector_type}.")
 
 
+
 class DatasetErrorDocs(Resource):
  @setup_required
  @login_required

api/controllers/console/datasets/datasets_document.py

@@ -465,6 +465,20 @@ def get(self, dataset_id, batch):
  document_model=document.doc_form
  )
  extract_settings.append(extract_setting)
+ elif document.data_source_type == 'website_crawl':
+ extract_setting = ExtractSetting(
+ datasource_type="website_crawl",
+ website_info={
+ "provider": data_source_info['provider'],
+ "job_id": data_source_info['job_id'],
+ "url": data_source_info['url'],
+ "tenant_id": current_user.current_tenant_id,
+ "mode": data_source_info['mode'],
+ "only_main_content": data_source_info['only_main_content']
+ },
+ document_model=document.doc_form
+ )
+ extract_settings.append(extract_setting)
 
  else:
  raise ValueError('Data source type not support')
@@ -952,6 +966,33 @@ def post(self, dataset_id, document_id):
  return document
 
 
+class WebsiteDocumentSyncApi(DocumentResource):
+ @setup_required
+ @login_required
+ @account_initialization_required
+ def get(self, dataset_id, document_id):
+ """sync website document."""
+ dataset_id = str(dataset_id)
+ dataset = DatasetService.get_dataset(dataset_id)
+ if not dataset:
+ raise NotFound('Dataset not found.')
+ document_id = str(document_id)
+ document = DocumentService.get_document(dataset.id, document_id)
+ if not document:
+ raise NotFound('Document not found.')
+ if document.tenant_id != current_user.current_tenant_id:
+ raise Forbidden('No permission.')
+ if document.data_source_type != 'website_crawl':
+ raise ValueError('Document is not a website document.')
+ # 403 if document is archived
+ if DocumentService.check_archived(document):
+ raise ArchivedDocumentImmutableError()
+ # sync document
+ DocumentService.sync_website_document(dataset_id, document)
+
+ return {'result': 'success'}, 200
+
+
 api.add_resource(GetProcessRuleApi, '/datasets/process-rule')
 api.add_resource(DatasetDocumentListApi,
  '/datasets/<uuid:dataset_id>/documents')
@@ -980,3 +1021,5 @@ def post(self, dataset_id, document_id):
 api.add_resource(DocumentRetryApi, '/datasets/<uuid:dataset_id>/retry')
 api.add_resource(DocumentRenameApi,
  '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/rename')
+
+api.add_resource(WebsiteDocumentSyncApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/website-sync')

api/controllers/console/datasets/error.py

@@ -73,6 +73,12 @@ class InvalidMetadataError(BaseHTTPException):
  code = 400
 
 
+class WebsiteCrawlError(BaseHTTPException):
+ error_code = 'crawl_failed'
+ description = "{message}"
+ code = 500
+
+
 class DatasetInUseError(BaseHTTPException):
  error_code = 'dataset_in_use'
  description = "The dataset is being used by some apps. Please remove the dataset from the apps before deleting it."

api/controllers/console/datasets/website.py

@@ -0,0 +1,49 @@
+from flask_restful import Resource, reqparse
+
+from controllers.console import api
+from controllers.console.datasets.error import WebsiteCrawlError
+from controllers.console.setup import setup_required
+from controllers.console.wraps import account_initialization_required
+from libs.login import login_required
+from services.website_service import WebsiteService
+
+
+class WebsiteCrawlApi(Resource):
+
+ @setup_required
+ @login_required
+ @account_initialization_required
+ def post(self):
+ parser = reqparse.RequestParser()
+ parser.add_argument('provider', type=str, choices=['firecrawl'],
+ required=True, nullable=True, location='json')
+ parser.add_argument('url', type=str, required=True, nullable=True, location='json')
+ parser.add_argument('options', type=dict, required=True, nullable=True, location='json')
+ args = parser.parse_args()
+ WebsiteService.document_create_args_validate(args)
+ # crawl url
+ try:
+ result = WebsiteService.crawl_url(args)
+ except Exception as e:
+ raise WebsiteCrawlError(str(e))
+ return result, 200
+
+
+class WebsiteCrawlStatusApi(Resource):
+ @setup_required
+ @login_required
+ @account_initialization_required
+ def get(self, job_id: str):
+ parser = reqparse.RequestParser()
+ parser.add_argument('provider', type=str, choices=['firecrawl'], required=True, location='args')
+ args = parser.parse_args()
+ # get crawl status
+ try:
+ result = WebsiteService.get_crawl_status(job_id, args['provider'])
+ except Exception as e:
+ raise WebsiteCrawlError(str(e))
+ return result, 200
+
+
+api.add_resource(WebsiteCrawlApi, '/website/crawl')
+api.add_resource(WebsiteCrawlStatusApi, '/website/crawl/status/<string:job_id>')

api/core/indexing_runner.py

@@ -339,7 +339,7 @@ def indexing_estimate(self, tenant_id: str, extract_settings: list[ExtractSettin
  def _extract(self, index_processor: BaseIndexProcessor, dataset_document: DatasetDocument, process_rule: dict) \
  -> list[Document]:
  # load file
- if dataset_document.data_source_type not in ["upload_file", "notion_import"]:
+ if dataset_document.data_source_type not in ["upload_file", "notion_import", "website_crawl"]:
  return []
 
  data_source_info = dataset_document.data_source_info_dict
@@ -375,6 +375,23 @@ def _extract(self, index_processor: BaseIndexProcessor, dataset_document: Datase
  document_model=dataset_document.doc_form
  )
  text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule['mode'])
+ elif dataset_document.data_source_type == 'website_crawl':
+ if (not data_source_info or 'provider' not in data_source_info
+ or 'url' not in data_source_info or 'job_id' not in data_source_info):
+ raise ValueError("no website import info found")
+ extract_setting = ExtractSetting(
+ datasource_type="website_crawl",
+ website_info={
+ "provider": data_source_info['provider'],
+ "job_id": data_source_info['job_id'],
+ "tenant_id": dataset_document.tenant_id,
+ "url": data_source_info['url'],
+ "mode": data_source_info['mode'],
+ "only_main_content": data_source_info['only_main_content']
+ },
+ document_model=dataset_document.doc_form
+ )
+ text_docs = index_processor.extract(extract_setting, process_rule_mode=process_rule['mode'])
  # update document status to splitting
  self._update_document_index_status(
  document_id=dataset_document.id,

api/core/model_runtime/model_providers/togetherai/llm/llm.py

@@ -124,7 +124,7 @@ def get_customizable_model_schema(self, model: str, credentials: dict) -> AIMode
  default=float(credentials.get('presence_penalty', 0)),
  min=-2,
  max=2
- )
+ ),
  ],
  pricing=PriceConfig(
  input=Decimal(cred_with_endpoint.get('input_price', 0)),

api/core/rag/extractor/entity/datasource_type.py

@@ -4,3 +4,4 @@
 class DatasourceType(Enum):
  FILE = "upload_file"
  NOTION = "notion_import"
+ WEBSITE = "website_crawl"