external knowledge api (#8913)

Co-authored-by: Yi <yxiaoisme@gmail.com>

api/services/dataset_service.py

@@ -32,6 +32,7 @@ from models.dataset import (
     DatasetQuery,
     Document,
     DocumentSegment,
+    ExternalKnowledgeBindings,
 )
 from models.model import UploadFile
 from models.source import DataSourceOauthBinding
@@ -39,6 +40,7 @@ from services.errors.account import NoPermissionError
 from services.errors.dataset import DatasetNameDuplicateError
 from services.errors.document import DocumentIndexingError
 from services.errors.file import FileNotExistsError
+from services.external_knowledge_service import ExternalDatasetService
 from services.feature_service import FeatureModel, FeatureService
 from services.tag_service import TagService
 from services.vector_service import VectorService
@@ -56,10 +58,8 @@ from tasks.sync_website_document_indexing_task import sync_website_document_inde
 
 class DatasetService:
     @staticmethod
-    def get_datasets(page, per_page, provider="vendor", tenant_id=None, user=None, search=None, tag_ids=None):
-        query = Dataset.query.filter(Dataset.provider == provider, Dataset.tenant_id == tenant_id).order_by(
-            Dataset.created_at.desc()
-        )
+    def get_datasets(page, per_page, tenant_id=None, user=None, search=None, tag_ids=None):
+        query = Dataset.query.filter(Dataset.tenant_id == tenant_id).order_by(Dataset.created_at.desc())
 
         if user:
             # get permitted dataset ids
@@ -137,7 +137,14 @@ class DatasetService:
 
     @staticmethod
     def create_empty_dataset(
-        tenant_id: str, name: str, indexing_technique: Optional[str], account: Account, permission: Optional[str] = None
+        tenant_id: str,
+        name: str,
+        indexing_technique: Optional[str],
+        account: Account,
+        permission: Optional[str] = None,
+        provider: str = "vendor",
+        external_knowledge_api_id: Optional[str] = None,
+        external_knowledge_id: Optional[str] = None,
     ):
         # check if dataset name already exists
         if Dataset.query.filter_by(name=name, tenant_id=tenant_id).first():
@@ -156,12 +163,28 @@ class DatasetService:
         dataset.embedding_model_provider = embedding_model.provider if embedding_model else None
         dataset.embedding_model = embedding_model.model if embedding_model else None
         dataset.permission = permission or DatasetPermissionEnum.ONLY_ME
+        dataset.provider = provider
         db.session.add(dataset)
+        db.session.flush()
+
+        if provider == "external" and external_knowledge_api_id:
+            external_knowledge_api = ExternalDatasetService.get_external_knowledge_api(external_knowledge_api_id)
+            if not external_knowledge_api:
+                raise ValueError("External API template not found.")
+            external_knowledge_binding = ExternalKnowledgeBindings(
+                tenant_id=tenant_id,
+                dataset_id=dataset.id,
+                external_knowledge_api_id=external_knowledge_api_id,
+                external_knowledge_id=external_knowledge_id,
+                created_by=account.id,
+            )
+            db.session.add(external_knowledge_binding)
+
         db.session.commit()
         return dataset
 
     @staticmethod
-    def get_dataset(dataset_id):
+    def get_dataset(dataset_id) -> Dataset:
         return Dataset.query.filter_by(id=dataset_id).first()
 
     @staticmethod
@@ -202,81 +225,103 @@ class DatasetService:
 
     @staticmethod
     def update_dataset(dataset_id, data, user):
-        data.pop("partial_member_list", None)
-        filtered_data = {k: v for k, v in data.items() if v is not None or k == "description"}
         dataset = DatasetService.get_dataset(dataset_id)
+
         DatasetService.check_dataset_permission(dataset, user)
-        action = None
-        if dataset.indexing_technique != data["indexing_technique"]:
-            # if update indexing_technique
-            if data["indexing_technique"] == "economy":
-                action = "remove"
-                filtered_data["embedding_model"] = None
-                filtered_data["embedding_model_provider"] = None
-                filtered_data["collection_binding_id"] = None
-            elif data["indexing_technique"] == "high_quality":
-                action = "add"
-                # get embedding model setting
-                try:
-                    model_manager = ModelManager()
-                    embedding_model = model_manager.get_model_instance(
-                        tenant_id=current_user.current_tenant_id,
-                        provider=data["embedding_model_provider"],
-                        model_type=ModelType.TEXT_EMBEDDING,
-                        model=data["embedding_model"],
-                    )
-                    filtered_data["embedding_model"] = embedding_model.model
-                    filtered_data["embedding_model_provider"] = embedding_model.provider
-                    dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
-                        embedding_model.provider, embedding_model.model
-                    )
-                    filtered_data["collection_binding_id"] = dataset_collection_binding.id
-                except LLMBadRequestError:
-                    raise ValueError(
-                        "No Embedding Model available. Please configure a valid provider "
-                        "in the Settings -> Model Provider."
-                    )
-                except ProviderTokenNotInitError as ex:
-                    raise ValueError(ex.description)
-        else:
+        if dataset.provider == "external":
+            dataset.retrieval_model = data.get("external_retrieval_model", None)
+            dataset.name = data.get("name", dataset.name)
+            dataset.description = data.get("description", "")
+            external_knowledge_id = data.get("external_knowledge_id", None)
+            db.session.add(dataset)
+            if not external_knowledge_id:
+                raise ValueError("External knowledge id is required.")
+            external_knowledge_api_id = data.get("external_knowledge_api_id", None)
+            if not external_knowledge_api_id:
+                raise ValueError("External knowledge api id is required.")
+            external_knowledge_binding = ExternalKnowledgeBindings.query.filter_by(dataset_id=dataset_id).first()
             if (
-                data["embedding_model_provider"] != dataset.embedding_model_provider
-                or data["embedding_model"] != dataset.embedding_model
+                external_knowledge_binding.external_knowledge_id != external_knowledge_id
+                or external_knowledge_binding.external_knowledge_api_id != external_knowledge_api_id
             ):
-                action = "update"
-                try:
-                    model_manager = ModelManager()
-                    embedding_model = model_manager.get_model_instance(
-                        tenant_id=current_user.current_tenant_id,
-                        provider=data["embedding_model_provider"],
-                        model_type=ModelType.TEXT_EMBEDDING,
-                        model=data["embedding_model"],
-                    )
-                    filtered_data["embedding_model"] = embedding_model.model
-                    filtered_data["embedding_model_provider"] = embedding_model.provider
-                    dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
-                        embedding_model.provider, embedding_model.model
-                    )
-                    filtered_data["collection_binding_id"] = dataset_collection_binding.id
-                except LLMBadRequestError:
-                    raise ValueError(
-                        "No Embedding Model available. Please configure a valid provider "
-                        "in the Settings -> Model Provider."
-                    )
-                except ProviderTokenNotInitError as ex:
-                    raise ValueError(ex.description)
+                external_knowledge_binding.external_knowledge_id = external_knowledge_id
+                external_knowledge_binding.external_knowledge_api_id = external_knowledge_api_id
+                db.session.add(external_knowledge_binding)
+            db.session.commit()
+        else:
+            data.pop("partial_member_list", None)
+            filtered_data = {k: v for k, v in data.items() if v is not None or k == "description"}
+            action = None
+            if dataset.indexing_technique != data["indexing_technique"]:
+                # if update indexing_technique
+                if data["indexing_technique"] == "economy":
+                    action = "remove"
+                    filtered_data["embedding_model"] = None
+                    filtered_data["embedding_model_provider"] = None
+                    filtered_data["collection_binding_id"] = None
+                elif data["indexing_technique"] == "high_quality":
+                    action = "add"
+                    # get embedding model setting
+                    try:
+                        model_manager = ModelManager()
+                        embedding_model = model_manager.get_model_instance(
+                            tenant_id=current_user.current_tenant_id,
+                            provider=data["embedding_model_provider"],
+                            model_type=ModelType.TEXT_EMBEDDING,
+                            model=data["embedding_model"],
+                        )
+                        filtered_data["embedding_model"] = embedding_model.model
+                        filtered_data["embedding_model_provider"] = embedding_model.provider
+                        dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
+                            embedding_model.provider, embedding_model.model
+                        )
+                        filtered_data["collection_binding_id"] = dataset_collection_binding.id
+                    except LLMBadRequestError:
+                        raise ValueError(
+                            "No Embedding Model available. Please configure a valid provider "
+                            "in the Settings -> Model Provider."
+                        )
+                    except ProviderTokenNotInitError as ex:
+                        raise ValueError(ex.description)
+            else:
+                if (
+                    data["embedding_model_provider"] != dataset.embedding_model_provider
+                    or data["embedding_model"] != dataset.embedding_model
+                ):
+                    action = "update"
+                    try:
+                        model_manager = ModelManager()
+                        embedding_model = model_manager.get_model_instance(
+                            tenant_id=current_user.current_tenant_id,
+                            provider=data["embedding_model_provider"],
+                            model_type=ModelType.TEXT_EMBEDDING,
+                            model=data["embedding_model"],
+                        )
+                        filtered_data["embedding_model"] = embedding_model.model
+                        filtered_data["embedding_model_provider"] = embedding_model.provider
+                        dataset_collection_binding = DatasetCollectionBindingService.get_dataset_collection_binding(
+                            embedding_model.provider, embedding_model.model
+                        )
+                        filtered_data["collection_binding_id"] = dataset_collection_binding.id
+                    except LLMBadRequestError:
+                        raise ValueError(
+                            "No Embedding Model available. Please configure a valid provider "
+                            "in the Settings -> Model Provider."
+                        )
+                    except ProviderTokenNotInitError as ex:
+                        raise ValueError(ex.description)
 
-        filtered_data["updated_by"] = user.id
-        filtered_data["updated_at"] = datetime.datetime.now()
+            filtered_data["updated_by"] = user.id
+            filtered_data["updated_at"] = datetime.datetime.now()
 
-        # update Retrieval model
-        filtered_data["retrieval_model"] = data["retrieval_model"]
+            # update Retrieval model
+            filtered_data["retrieval_model"] = data["retrieval_model"]
 
-        dataset.query.filter_by(id=dataset_id).update(filtered_data)
+            dataset.query.filter_by(id=dataset_id).update(filtered_data)
 
-        db.session.commit()
-        if action:
-            deal_dataset_vector_index_task.delay(dataset_id, action)
+            db.session.commit()
+            if action:
+                deal_dataset_vector_index_task.delay(dataset_id, action)
         return dataset
 
     @staticmethod