bump version to 0.3.26 (#1307)

Co-authored-by: jyong <jyong@dify.ai>

.devcontainer/Dockerfile

@@ -1,11 +1,8 @@
-FROM mcr.microsoft.com/devcontainers/anaconda:0-3
+FROM mcr.microsoft.com/devcontainers/python:3.10
+
+COPY . .
 
-# Copy environment.yml (if found) to a temp location so we update the environment. Also
-# copy "noop.txt" so the COPY instruction does not fail if no environment.yml exists.
-COPY environment.yml* .devcontainer/noop.txt /tmp/conda-tmp/
-RUN if [ -f "/tmp/conda-tmp/environment.yml" ]; then umask 0002 && /opt/conda/bin/conda env update -n base -f /tmp/conda-tmp/environment.yml; fi \
-    && rm -rf /tmp/conda-tmp
 
 # [Optional] Uncomment this section to install additional OS packages.
 # RUN apt-get update && export DEBIAN_FRONTEND=noninteractive \
-#     && apt-get -y install --no-install-recommends <your-package-list-here>
+#     && apt-get -y install --no-install-recommends <your-package-list-here>

.devcontainer/devcontainer.json

@@ -1,13 +1,12 @@
 // For format details, see https://aka.ms/devcontainer.json. For config options, see the
 // README at: https://github.com/devcontainers/templates/tree/main/src/anaconda
 {
-	"name": "Anaconda (Python 3)",
+	"name": "Python 3.10",
 	"build": { 
 		"context": "..",
 		"dockerfile": "Dockerfile"
 	},
 	"features": {
-		"ghcr.io/dhoeric/features/act:1": {},
 		"ghcr.io/devcontainers/features/node:1": {
 			"nodeGypDependencies": true,
 			"version": "lts"

.github/ISSUE_TEMPLATE/help_wanted.yml

@@ -0,0 +1,11 @@
+name: "🤝 Help Wanted"
+description: "Request help from the community"
+labels:
+- help-wanted
+body:
+- type: textarea
+  attributes:
+    label: Provide a description of the help you need
+    placeholder: Briefly describe what you need help with.
+  validations:
+    required: true

.gitignore

@@ -144,6 +144,7 @@ docker/volumes/app/storage/*
 docker/volumes/db/data/*
 docker/volumes/redis/data/*
 docker/volumes/weaviate/*
+docker/volumes/qdrant/*
 
 sdks/python-client/build
 sdks/python-client/dist

api/.env.example

@@ -50,25 +50,7 @@ S3_REGION=your-region
 WEB_API_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,*
 CONSOLE_CORS_ALLOW_ORIGINS=http://127.0.0.1:3000,*
 
-# Cookie configuration
-COOKIE_HTTPONLY=true
-COOKIE_SAMESITE=None
-COOKIE_SECURE=true
-
-# Session configuration
-SESSION_PERMANENT=true
-SESSION_USE_SIGNER=true
-
-## support redis, sqlalchemy
-SESSION_TYPE=redis
-
-# session redis configuration
-SESSION_REDIS_HOST=localhost
-SESSION_REDIS_PORT=6379
-SESSION_REDIS_PASSWORD=difyai123456
-SESSION_REDIS_DB=2
-
-# Vector database configuration, support: weaviate, qdrant
+# Vector database configuration, support: weaviate, qdrant, milvus
 VECTOR_STORE=weaviate
 
 # Weaviate configuration
@@ -77,9 +59,16 @@ WEAVIATE_API_KEY=WVF5YThaHlkYwhGUSmCRgsX3tD5ngdN8pkih
 WEAVIATE_GRPC_ENABLED=false
 WEAVIATE_BATCH_SIZE=100
 
-# Qdrant configuration, use `path:` prefix for local mode or `https://your-qdrant-cluster-url.qdrant.io` for remote mode
-QDRANT_URL=path:storage/qdrant
-QDRANT_API_KEY=your-qdrant-api-key
+# Qdrant configuration, use `http://localhost:6333` for local mode or `https://your-qdrant-cluster-url.qdrant.io` for remote mode
+QDRANT_URL=http://localhost:6333
+QDRANT_API_KEY=difyai123456
+
+# Milvus configuration
+MILVUS_HOST=127.0.0.1
+MILVUS_PORT=19530
+MILVUS_USER=root
+MILVUS_PASSWORD=Milvus
+MILVUS_SECURE=false
 
 # Mail configuration, support: resend
 MAIL_TYPE=

api/app.py

@@ -1,8 +1,7 @@
 # -*- coding:utf-8 -*-
 import os
-from datetime import datetime, timedelta
 
-from werkzeug.exceptions import Forbidden
+from werkzeug.exceptions import Unauthorized
 
 if not os.environ.get("DEBUG") or os.environ.get("DEBUG").lower() != 'true':
     from gevent import monkey
@@ -12,12 +11,11 @@ import logging
 import json
 import threading
 
-from flask import Flask, request, Response, session
-import flask_login
+from flask import Flask, request, Response
 from flask_cors import CORS
 
 from core.model_providers.providers import hosted
-from extensions import ext_session, ext_celery, ext_sentry, ext_redis, ext_login, ext_migrate, \
+from extensions import ext_celery, ext_sentry, ext_redis, ext_login, ext_migrate, \
     ext_database, ext_storage, ext_mail, ext_stripe
 from extensions.ext_database import db
 from extensions.ext_login import login_manager
@@ -27,12 +25,10 @@ from models import model, account, dataset, web, task, source, tool
 from events import event_handlers
 # DO NOT REMOVE ABOVE
 
-import core
 from config import Config, CloudEditionConfig
 from commands import register_commands
-from models.account import TenantAccountJoin, AccountStatus
-from models.model import Account, EndUser, App
-from services.account_service import TenantService
+from services.account_service import AccountService
+from libs.passport import PassportService
 
 import warnings
 warnings.simplefilter("ignore", ResourceWarning)
@@ -85,81 +81,33 @@ def initialize_extensions(app):
     ext_redis.init_app(app)
     ext_storage.init_app(app)
     ext_celery.init_app(app)
-    ext_session.init_app(app)
     ext_login.init_app(app)
     ext_mail.init_app(app)
     ext_sentry.init_app(app)
     ext_stripe.init_app(app)
 
 
-def _create_tenant_for_account(account):
-    tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
-
-    TenantService.create_tenant_member(tenant, account, role='owner')
-    account.current_tenant = tenant
-
-    return tenant
-
-
 # Flask-Login configuration
-@login_manager.user_loader
-def load_user(user_id):
-    """Load user based on the user_id."""
+@login_manager.request_loader
+def load_user_from_request(request_from_flask_login):
+    """Load user based on the request."""
     if request.blueprint == 'console':
         # Check if the user_id contains a dot, indicating the old format
-        if '.' in user_id:
-            tenant_id, account_id = user_id.split('.')
-        else:
-            account_id = user_id
+        auth_header = request.headers.get('Authorization', '')
+        if ' ' not in auth_header:
+            raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
+        auth_scheme, auth_token = auth_header.split(None, 1)
+        auth_scheme = auth_scheme.lower()
+        if auth_scheme != 'bearer':
+            raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
+        
+        decoded = PassportService().verify(auth_token)
+        user_id = decoded.get('user_id')
 
-        account = db.session.query(Account).filter(Account.id == account_id).first()
-
-        if account:
-            if account.status == AccountStatus.BANNED.value or account.status == AccountStatus.CLOSED.value:
-                raise Forbidden('Account is banned or closed.')
-
-            workspace_id = session.get('workspace_id')
-            if workspace_id:
-                tenant_account_join = db.session.query(TenantAccountJoin).filter(
-                    TenantAccountJoin.account_id == account.id,
-                    TenantAccountJoin.tenant_id == workspace_id
-                ).first()
-
-                if not tenant_account_join:
-                    tenant_account_join = db.session.query(TenantAccountJoin).filter(
-                        TenantAccountJoin.account_id == account.id).first()
-
-                    if tenant_account_join:
-                        account.current_tenant_id = tenant_account_join.tenant_id
-                    else:
-                        _create_tenant_for_account(account)
-                    session['workspace_id'] = account.current_tenant_id
-                else:
-                    account.current_tenant_id = workspace_id
-            else:
-                tenant_account_join = db.session.query(TenantAccountJoin).filter(
-                    TenantAccountJoin.account_id == account.id).first()
-                if tenant_account_join:
-                    account.current_tenant_id = tenant_account_join.tenant_id
-                else:
-                    _create_tenant_for_account(account)
-                session['workspace_id'] = account.current_tenant_id
-
-            current_time = datetime.utcnow()
-
-            # update last_active_at when last_active_at is more than 10 minutes ago
-            if current_time - account.last_active_at > timedelta(minutes=10):
-                account.last_active_at = current_time
-                db.session.commit()
-
-            # Log in the user with the updated user_id
-            flask_login.login_user(account, remember=True)
-
-        return account
+        return AccountService.load_user(user_id)
     else:
         return None
 
-
 @login_manager.unauthorized_handler
 def unauthorized_handler():
     """Handle unauthorized requests."""
@@ -216,6 +164,7 @@ if app.config['TESTING']:
 @app.after_request
 def after_request(response):
     """Add Version headers to the response."""
+    response.set_cookie('remember_token', '', expires=0)
     response.headers.add('X-Version', app.config['CURRENT_VERSION'])
     response.headers.add('X-Env', app.config['DEPLOY_ENV'])
     return response

api/commands.py

@@ -3,12 +3,13 @@ import json
 import math
 import random
 import string
+import threading
 import time
 import uuid
 
 import click
 from tqdm import tqdm
-from flask import current_app
+from flask import current_app, Flask
 from langchain.embeddings import OpenAIEmbeddings
 from werkzeug.exceptions import NotFound
 
@@ -456,92 +457,92 @@ def update_qdrant_indexes():
 @click.command('normalization-collections', help='restore all collections in one')
 def normalization_collections():
     click.echo(click.style('Start normalization collections.', fg='green'))
-    normalization_count = 0
-
+    normalization_count = []
     page = 1
     while True:
         try:
             datasets = db.session.query(Dataset).filter(Dataset.indexing_technique == 'high_quality') \
-                .order_by(Dataset.created_at.desc()).paginate(page=page, per_page=50)
+                .order_by(Dataset.created_at.desc()).paginate(page=page, per_page=100)
         except NotFound:
             break
-
+        datasets_result = datasets.items
         page += 1
-        for dataset in datasets:
-            if not dataset.collection_binding_id:
-                try:
-                    click.echo('restore dataset index: {}'.format(dataset.id))
-                    try:
-                        embedding_model = ModelFactory.get_embedding_model(
-                            tenant_id=dataset.tenant_id,
-                            model_provider_name=dataset.embedding_model_provider,
-                            model_name=dataset.embedding_model
-                        )
-                    except Exception:
-                        provider = Provider(
-                            id='provider_id',
-                            tenant_id=dataset.tenant_id,
-                            provider_name='openai',
-                            provider_type=ProviderType.CUSTOM.value,
-                            encrypted_config=json.dumps({'openai_api_key': 'TEST'}),
-                            is_valid=True,
-                        )
-                        model_provider = OpenAIProvider(provider=provider)
-                        embedding_model = OpenAIEmbedding(name="text-embedding-ada-002",
-                                                          model_provider=model_provider)
-                    embeddings = CacheEmbedding(embedding_model)
-                    dataset_collection_binding = db.session.query(DatasetCollectionBinding). \
-                        filter(DatasetCollectionBinding.provider_name == embedding_model.model_provider.provider_name,
-                               DatasetCollectionBinding.model_name == embedding_model.name). \
-                        order_by(DatasetCollectionBinding.created_at). \
-                        first()
+        for i in range(0, len(datasets_result), 5):
+            threads = []
+            sub_datasets = datasets_result[i:i + 5]
+            for dataset in sub_datasets:
+                document_format_thread = threading.Thread(target=deal_dataset_vector, kwargs={
+                    'flask_app': current_app._get_current_object(),
+                    'dataset': dataset,
+                    'normalization_count': normalization_count
+                })
+                threads.append(document_format_thread)
+                document_format_thread.start()
+            for thread in threads:
+                thread.join()
 
-                    if not dataset_collection_binding:
-                        dataset_collection_binding = DatasetCollectionBinding(
-                            provider_name=embedding_model.model_provider.provider_name,
-                            model_name=embedding_model.name,
-                            collection_name="Vector_index_" + str(uuid.uuid4()).replace("-", "_") + '_Node'
-                        )
-                        db.session.add(dataset_collection_binding)
-                        db.session.commit()
+    click.echo(click.style('Congratulations! restore {} dataset indexes.'.format(len(normalization_count)), fg='green'))
 
-                    from core.index.vector_index.qdrant_vector_index import QdrantVectorIndex, QdrantConfig
 
-                    index = QdrantVectorIndex(
-                        dataset=dataset,
-                        config=QdrantConfig(
-                            endpoint=current_app.config.get('QDRANT_URL'),
-                            api_key=current_app.config.get('QDRANT_API_KEY'),
-                            root_path=current_app.root_path
-                        ),
-                        embeddings=embeddings
-                    )
-                    if index:
-                        index.restore_dataset_in_one(dataset, dataset_collection_binding)
-                    else:
-                        click.echo('passed.')
+def deal_dataset_vector(flask_app: Flask, dataset: Dataset, normalization_count: list):
+    with flask_app.app_context():
+        try:
+            click.echo('restore dataset index: {}'.format(dataset.id))
+            try:
+                embedding_model = ModelFactory.get_embedding_model(
+                    tenant_id=dataset.tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
+            except Exception:
+                provider = Provider(
+                    id='provider_id',
+                    tenant_id=dataset.tenant_id,
+                    provider_name='openai',
+                    provider_type=ProviderType.CUSTOM.value,
+                    encrypted_config=json.dumps({'openai_api_key': 'TEST'}),
+                    is_valid=True,
+                )
+                model_provider = OpenAIProvider(provider=provider)
+                embedding_model = OpenAIEmbedding(name="text-embedding-ada-002",
+                                                  model_provider=model_provider)
+            embeddings = CacheEmbedding(embedding_model)
+            dataset_collection_binding = db.session.query(DatasetCollectionBinding). \
+                filter(DatasetCollectionBinding.provider_name == embedding_model.model_provider.provider_name,
+                       DatasetCollectionBinding.model_name == embedding_model.name). \
+                order_by(DatasetCollectionBinding.created_at). \
+                first()
 
-                    original_index = QdrantVectorIndex(
-                        dataset=dataset,
-                        config=QdrantConfig(
-                            endpoint=current_app.config.get('QDRANT_URL'),
-                            api_key=current_app.config.get('QDRANT_API_KEY'),
-                            root_path=current_app.root_path
-                        ),
-                        embeddings=embeddings
-                    )
-                    if original_index:
-                        original_index.delete_original_collection(dataset, dataset_collection_binding)
-                        normalization_count += 1
-                    else:
-                        click.echo('passed.')
-                except Exception as e:
-                    click.echo(
-                        click.style('Create dataset index error: {} {}'.format(e.__class__.__name__, str(e)),
-                                    fg='red'))
-                    continue
+            if not dataset_collection_binding:
+                dataset_collection_binding = DatasetCollectionBinding(
+                    provider_name=embedding_model.model_provider.provider_name,
+                    model_name=embedding_model.name,
+                    collection_name="Vector_index_" + str(uuid.uuid4()).replace("-", "_") + '_Node'
+                )
+                db.session.add(dataset_collection_binding)
+                db.session.commit()
 
-    click.echo(click.style('Congratulations! restore {} dataset indexes.'.format(normalization_count), fg='green'))
+            from core.index.vector_index.qdrant_vector_index import QdrantVectorIndex, QdrantConfig
+
+            index = QdrantVectorIndex(
+                dataset=dataset,
+                config=QdrantConfig(
+                    endpoint=current_app.config.get('QDRANT_URL'),
+                    api_key=current_app.config.get('QDRANT_API_KEY'),
+                    root_path=current_app.root_path
+                ),
+                embeddings=embeddings
+            )
+            if index:
+                # index.delete_by_group_id(dataset.id)
+                index.restore_dataset_in_one(dataset, dataset_collection_binding)
+            else:
+                click.echo('passed.')
+            normalization_count.append(1)
+        except Exception as e:
+            click.echo(
+                click.style('Create dataset index error: {} {}'.format(e.__class__.__name__, str(e)),
+                            fg='red'))
 
 
 @click.command('update_app_model_configs', help='Migrate data to support paragraph variable.')
@@ -646,6 +647,76 @@ def update_app_model_configs(batch_size):
 
             pbar.update(len(data_batch))
 
+@click.command('migrate_default_input_to_dataset_query_variable')
+@click.option("--batch-size", default=500, help="Number of records to migrate in each batch.")
+def migrate_default_input_to_dataset_query_variable(batch_size):
+
+    click.secho("Starting...", fg='green')
+
+    total_records = db.session.query(AppModelConfig) \
+        .join(App, App.app_model_config_id == AppModelConfig.id) \
+        .filter(App.mode == 'completion') \
+        .filter(AppModelConfig.dataset_query_variable == None) \
+        .count()
+    
+    if total_records == 0:
+        click.secho("No data to migrate.", fg='green')
+        return
+
+    num_batches = (total_records + batch_size - 1) // batch_size
+    
+    with tqdm(total=total_records, desc="Migrating Data") as pbar:
+        for i in range(num_batches):
+            offset = i * batch_size
+            limit = min(batch_size, total_records - offset)
+
+            click.secho(f"Fetching batch {i + 1}/{num_batches} from source database...", fg='green')
+
+            data_batch = db.session.query(AppModelConfig) \
+                .join(App, App.app_model_config_id == AppModelConfig.id) \
+                .filter(App.mode == 'completion') \
+                .filter(AppModelConfig.dataset_query_variable == None) \
+                .order_by(App.created_at) \
+                .offset(offset).limit(limit).all()
+
+            if not data_batch:
+                click.secho("No more data to migrate.", fg='green')
+                break
+
+            try:
+                click.secho(f"Migrating {len(data_batch)} records...", fg='green')
+                for data in data_batch:
+                    config = AppModelConfig.to_dict(data)
+
+                    tools = config["agent_mode"]["tools"]
+                    dataset_exists = "dataset" in str(tools)
+                    if not dataset_exists:
+                        continue
+
+                    user_input_form = config.get("user_input_form", [])
+                    for form in user_input_form:
+                        paragraph = form.get('paragraph')
+                        if paragraph \
+                            and paragraph.get('variable') == 'query':
+                                data.dataset_query_variable = 'query'
+                                break
+                        
+                        if paragraph \
+                            and paragraph.get('variable') == 'default_input':
+                                data.dataset_query_variable = 'default_input'
+                                break
+
+                db.session.commit()
+
+            except Exception as e:
+                click.secho(f"Error while migrating data: {e}, app_id: {data.app_id}, app_model_config_id: {data.id}",
+                            fg='red')
+                continue
+            
+            click.secho(f"Successfully migrated batch {i + 1}/{num_batches}.", fg='green')
+
+            pbar.update(len(data_batch))
+
 
 def register_commands(app):
     app.cli.add_command(reset_password)
@@ -659,3 +730,4 @@ def register_commands(app):
     app.cli.add_command(update_qdrant_indexes)
     app.cli.add_command(update_app_model_configs)
     app.cli.add_command(normalization_collections)
+    app.cli.add_command(migrate_default_input_to_dataset_query_variable)

api/config.py

@@ -10,9 +10,6 @@ from extensions.ext_redis import redis_client
 dotenv.load_dotenv()
 
 DEFAULTS = {
-    'COOKIE_HTTPONLY': 'True',
-    'COOKIE_SECURE': 'True',
-    'COOKIE_SAMESITE': 'None',
     'DB_USERNAME': 'postgres',
     'DB_PASSWORD': '',
     'DB_HOST': 'localhost',
@@ -22,10 +19,6 @@ DEFAULTS = {
     'REDIS_PORT': '6379',
     'REDIS_DB': '0',
     'REDIS_USE_SSL': 'False',
-    'SESSION_REDIS_HOST': 'localhost',
-    'SESSION_REDIS_PORT': '6379',
-    'SESSION_REDIS_DB': '2',
-    'SESSION_REDIS_USE_SSL': 'False',
     'OAUTH_REDIRECT_PATH': '/console/api/oauth/authorize',
     'OAUTH_REDIRECT_INDEX_PATH': '/',
     'CONSOLE_WEB_URL': 'https://cloud.dify.ai',
@@ -36,9 +29,6 @@ DEFAULTS = {
     'STORAGE_TYPE': 'local',
     'STORAGE_LOCAL_PATH': 'storage',
     'CHECK_UPDATE_URL': 'https://updates.dify.ai',
-    'SESSION_TYPE': 'sqlalchemy',
-    'SESSION_PERMANENT': 'True',
-    'SESSION_USE_SIGNER': 'True',
     'DEPLOY_ENV': 'PRODUCTION',
     'SQLALCHEMY_POOL_SIZE': 30,
     'SQLALCHEMY_POOL_RECYCLE': 3600,
@@ -102,7 +92,7 @@ class Config:
         self.CONSOLE_URL = get_env('CONSOLE_URL')
         self.API_URL = get_env('API_URL')
         self.APP_URL = get_env('APP_URL')
-        self.CURRENT_VERSION = "0.3.23"
+        self.CURRENT_VERSION = "0.3.26"
         self.COMMIT_SHA = get_env('COMMIT_SHA')
         self.EDITION = "SELF_HOSTED"
         self.DEPLOY_ENV = get_env('DEPLOY_ENV')
@@ -115,20 +105,6 @@ class Config:
         # Alternatively you can set it with `SECRET_KEY` environment variable.
         self.SECRET_KEY = get_env('SECRET_KEY')
 
-        # cookie settings
-        self.REMEMBER_COOKIE_HTTPONLY = get_bool_env('COOKIE_HTTPONLY')
-        self.SESSION_COOKIE_HTTPONLY = get_bool_env('COOKIE_HTTPONLY')
-        self.REMEMBER_COOKIE_SAMESITE = get_env('COOKIE_SAMESITE')
-        self.SESSION_COOKIE_SAMESITE = get_env('COOKIE_SAMESITE')
-        self.REMEMBER_COOKIE_SECURE = get_bool_env('COOKIE_SECURE')
-        self.SESSION_COOKIE_SECURE = get_bool_env('COOKIE_SECURE')
-        self.PERMANENT_SESSION_LIFETIME = timedelta(days=7)
-
-        # session settings, only support sqlalchemy, redis
-        self.SESSION_TYPE = get_env('SESSION_TYPE')
-        self.SESSION_PERMANENT = get_bool_env('SESSION_PERMANENT')
-        self.SESSION_USE_SIGNER = get_bool_env('SESSION_USE_SIGNER')
-
         # redis settings
         self.REDIS_HOST = get_env('REDIS_HOST')
         self.REDIS_PORT = get_env('REDIS_PORT')
@@ -137,14 +113,6 @@ class Config:
         self.REDIS_DB = get_env('REDIS_DB')
         self.REDIS_USE_SSL = get_bool_env('REDIS_USE_SSL')
 
-        # session redis settings
-        self.SESSION_REDIS_HOST = get_env('SESSION_REDIS_HOST')
-        self.SESSION_REDIS_PORT = get_env('SESSION_REDIS_PORT')
-        self.SESSION_REDIS_USERNAME = get_env('SESSION_REDIS_USERNAME')
-        self.SESSION_REDIS_PASSWORD = get_env('SESSION_REDIS_PASSWORD')
-        self.SESSION_REDIS_DB = get_env('SESSION_REDIS_DB')
-        self.SESSION_REDIS_USE_SSL = get_bool_env('SESSION_REDIS_USE_SSL')
-
         # storage settings
         self.STORAGE_TYPE = get_env('STORAGE_TYPE')
         self.STORAGE_LOCAL_PATH = get_env('STORAGE_LOCAL_PATH')
@@ -167,6 +135,14 @@ class Config:
         self.QDRANT_URL = get_env('QDRANT_URL')
         self.QDRANT_API_KEY = get_env('QDRANT_API_KEY')
 
+        # milvus setting
+        self.MILVUS_HOST = get_env('MILVUS_HOST')
+        self.MILVUS_PORT = get_env('MILVUS_PORT')
+        self.MILVUS_USER = get_env('MILVUS_USER')
+        self.MILVUS_PASSWORD = get_env('MILVUS_PASSWORD')
+        self.MILVUS_SECURE = get_env('MILVUS_SECURE')
+
+
         # cors settings
         self.CONSOLE_CORS_ALLOW_ORIGINS = get_cors_allow_origins(
             'CONSOLE_CORS_ALLOW_ORIGINS', self.CONSOLE_WEB_URL)

api/controllers/console/apikey.py

@@ -1,5 +1,5 @@
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 import flask_restful
 from flask_restful import Resource, fields, marshal_with
 from werkzeug.exceptions import Forbidden
@@ -81,6 +81,7 @@ class BaseApiKeyListResource(Resource):
         key = ApiToken.generate_api_key(self.token_prefix, 24)
         api_token = ApiToken()
         setattr(api_token, self.resource_id_field, resource_id)
+        api_token.tenant_id = current_user.current_tenant_id
         api_token.token = key
         api_token.type = self.resource_type
         db.session.add(api_token)

api/controllers/console/app/app.py

@@ -3,10 +3,9 @@ import json
 import logging
 from datetime import datetime
 
-import flask
 from flask_login import current_user
-from core.login.login import login_required
-from flask_restful import Resource, reqparse, fields, marshal_with, abort, inputs
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal_with, abort, inputs
 from werkzeug.exceptions import Forbidden
 
 from constants.model_template import model_templates, demo_model_templates
@@ -17,42 +16,13 @@ from controllers.console.wraps import account_initialization_required
 from core.model_providers.error import ProviderTokenNotInitError, LLMBadRequestError
 from core.model_providers.model_factory import ModelFactory
 from core.model_providers.model_provider_factory import ModelProviderFactory
-from core.model_providers.models.entity.model_params import ModelType
 from events.app_event import app_was_created, app_was_deleted
-from libs.helper import TimestampField
+from fields.app_fields import app_pagination_fields, app_detail_fields, template_list_fields, \
+    app_detail_fields_with_site
 from extensions.ext_database import db
 from models.model import App, AppModelConfig, Site
 from services.app_model_config_service import AppModelConfigService
 
-model_config_fields = {
-    'opening_statement': fields.String,
-    'suggested_questions': fields.Raw(attribute='suggested_questions_list'),
-    'suggested_questions_after_answer': fields.Raw(attribute='suggested_questions_after_answer_dict'),
-    'speech_to_text': fields.Raw(attribute='speech_to_text_dict'),
-    'retriever_resource': fields.Raw(attribute='retriever_resource_dict'),
-    'more_like_this': fields.Raw(attribute='more_like_this_dict'),
-    'sensitive_word_avoidance': fields.Raw(attribute='sensitive_word_avoidance_dict'),
-    'model': fields.Raw(attribute='model_dict'),
-    'user_input_form': fields.Raw(attribute='user_input_form_list'),
-    'pre_prompt': fields.String,
-    'agent_mode': fields.Raw(attribute='agent_mode_dict'),
-}
-
-app_detail_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'mode': fields.String,
-    'icon': fields.String,
-    'icon_background': fields.String,
-    'enable_site': fields.Boolean,
-    'enable_api': fields.Boolean,
-    'api_rpm': fields.Integer,
-    'api_rph': fields.Integer,
-    'is_demo': fields.Boolean,
-    'model_config': fields.Nested(model_config_fields, attribute='app_model_config'),
-    'created_at': TimestampField
-}
-
 
 def _get_app(app_id, tenant_id):
     app = db.session.query(App).filter(App.id == app_id, App.tenant_id == tenant_id).first()
@@ -62,35 +32,6 @@ def _get_app(app_id, tenant_id):
 
 
 class AppListApi(Resource):
-    prompt_config_fields = {
-        'prompt_template': fields.String,
-    }
-
-    model_config_partial_fields = {
-        'model': fields.Raw(attribute='model_dict'),
-        'pre_prompt': fields.String,
-    }
-
-    app_partial_fields = {
-        'id': fields.String,
-        'name': fields.String,
-        'mode': fields.String,
-        'icon': fields.String,
-        'icon_background': fields.String,
-        'enable_site': fields.Boolean,
-        'enable_api': fields.Boolean,
-        'is_demo': fields.Boolean,
-        'model_config': fields.Nested(model_config_partial_fields, attribute='app_model_config'),
-        'created_at': TimestampField
-    }
-
-    app_pagination_fields = {
-        'page': fields.Integer,
-        'limit': fields.Integer(attribute='per_page'),
-        'total': fields.Integer,
-        'has_more': fields.Boolean(attribute='has_next'),
-        'data': fields.List(fields.Nested(app_partial_fields), attribute='items')
-    }
 
     @setup_required
     @login_required
@@ -162,7 +103,8 @@ class AppListApi(Resource):
             model_configuration = AppModelConfigService.validate_configuration(
                 tenant_id=current_user.current_tenant_id,
                 account=current_user,
-                config=model_config_dict
+                config=model_config_dict,
+                mode=args['mode']
             )
 
             app = App(
@@ -236,18 +178,6 @@ class AppListApi(Resource):
 
 
 class AppTemplateApi(Resource):
-    template_fields = {
-        'name': fields.String,
-        'icon': fields.String,
-        'icon_background': fields.String,
-        'description': fields.String,
-        'mode': fields.String,
-        'model_config': fields.Nested(model_config_fields),
-    }
-
-    template_list_fields = {
-        'data': fields.List(fields.Nested(template_fields)),
-    }
 
     @setup_required
     @login_required
@@ -266,38 +196,6 @@ class AppTemplateApi(Resource):
 
 
 class AppApi(Resource):
-    site_fields = {
-        'access_token': fields.String(attribute='code'),
-        'code': fields.String,
-        'title': fields.String,
-        'icon': fields.String,
-        'icon_background': fields.String,
-        'description': fields.String,
-        'default_language': fields.String,
-        'customize_domain': fields.String,
-        'copyright': fields.String,
-        'privacy_policy': fields.String,
-        'customize_token_strategy': fields.String,
-        'prompt_public': fields.Boolean,
-        'app_base_url': fields.String,
-    }
-
-    app_detail_fields_with_site = {
-        'id': fields.String,
-        'name': fields.String,
-        'mode': fields.String,
-        'icon': fields.String,
-        'icon_background': fields.String,
-        'enable_site': fields.Boolean,
-        'enable_api': fields.Boolean,
-        'api_rpm': fields.Integer,
-        'api_rph': fields.Integer,
-        'is_demo': fields.Boolean,
-        'model_config': fields.Nested(model_config_fields, attribute='app_model_config'),
-        'site': fields.Nested(site_fields),
-        'api_base_url': fields.String,
-        'created_at': TimestampField
-    }
 
     @setup_required
     @login_required

api/controllers/console/app/audio.py

@@ -2,8 +2,8 @@
 import logging
 
 from flask import request
-from core.login.login import login_required
-from werkzeug.exceptions import InternalServerError, NotFound
+from libs.login import login_required
+from werkzeug.exceptions import InternalServerError
 
 import services
 from controllers.console import api

api/controllers/console/app/completion.py

@@ -5,7 +5,7 @@ from typing import Generator, Union
 
 import flask_login
 from flask import Response, stream_with_context
-from core.login.login import login_required
+from libs.login import login_required
 from werkzeug.exceptions import InternalServerError, NotFound
 
 import services

api/controllers/console/app/conversation.py

@@ -2,8 +2,8 @@ from datetime import datetime
 
 import pytz
 from flask_login import current_user
-from core.login.login import login_required
-from flask_restful import Resource, reqparse, fields, marshal_with
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal_with
 from flask_restful.inputs import int_range
 from sqlalchemy import or_, func
 from sqlalchemy.orm import joinedload
@@ -13,107 +13,14 @@ from controllers.console import api
 from controllers.console.app import _get_app
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
-from libs.helper import TimestampField, datetime_string, uuid_value
+from fields.conversation_fields import conversation_pagination_fields, conversation_detail_fields, \
+    conversation_message_detail_fields, conversation_with_summary_pagination_fields
+from libs.helper import datetime_string
 from extensions.ext_database import db
 from models.model import Message, MessageAnnotation, Conversation
 
-account_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'email': fields.String
-}
-
-feedback_fields = {
-    'rating': fields.String,
-    'content': fields.String,
-    'from_source': fields.String,
-    'from_end_user_id': fields.String,
-    'from_account': fields.Nested(account_fields, allow_null=True),
-}
-
-annotation_fields = {
-    'content': fields.String,
-    'account': fields.Nested(account_fields, allow_null=True),
-    'created_at': TimestampField
-}
-
-message_detail_fields = {
-    'id': fields.String,
-    'conversation_id': fields.String,
-    'inputs': fields.Raw,
-    'query': fields.String,
-    'message': fields.Raw,
-    'message_tokens': fields.Integer,
-    'answer': fields.String,
-    'answer_tokens': fields.Integer,
-    'provider_response_latency': fields.Float,
-    'from_source': fields.String,
-    'from_end_user_id': fields.String,
-    'from_account_id': fields.String,
-    'feedbacks': fields.List(fields.Nested(feedback_fields)),
-    'annotation': fields.Nested(annotation_fields, allow_null=True),
-    'created_at': TimestampField
-}
-
-feedback_stat_fields = {
-    'like': fields.Integer,
-    'dislike': fields.Integer
-}
-
-model_config_fields = {
-    'opening_statement': fields.String,
-    'suggested_questions': fields.Raw,
-    'model': fields.Raw,
-    'user_input_form': fields.Raw,
-    'pre_prompt': fields.String,
-    'agent_mode': fields.Raw,
-}
-
 
 class CompletionConversationApi(Resource):
-    class MessageTextField(fields.Raw):
-        def format(self, value):
-            return value[0]['text'] if value else ''
-
-    simple_configs_fields = {
-        'prompt_template': fields.String,
-    }
-
-    simple_model_config_fields = {
-        'model': fields.Raw(attribute='model_dict'),
-        'pre_prompt': fields.String,
-    }
-
-    simple_message_detail_fields = {
-        'inputs': fields.Raw,
-        'query': fields.String,
-        'message': MessageTextField,
-        'answer': fields.String,
-    }
-
-    conversation_fields = {
-        'id': fields.String,
-        'status': fields.String,
-        'from_source': fields.String,
-        'from_end_user_id': fields.String,
-        'from_end_user_session_id': fields.String(),
-        'from_account_id': fields.String,
-        'read_at': TimestampField,
-        'created_at': TimestampField,
-        'annotation': fields.Nested(annotation_fields, allow_null=True),
-        'model_config': fields.Nested(simple_model_config_fields),
-        'user_feedback_stats': fields.Nested(feedback_stat_fields),
-        'admin_feedback_stats': fields.Nested(feedback_stat_fields),
-        'message': fields.Nested(simple_message_detail_fields, attribute='first_message')
-    }
-
-    conversation_pagination_fields = {
-        'page': fields.Integer,
-        'limit': fields.Integer(attribute='per_page'),
-        'total': fields.Integer,
-        'has_more': fields.Boolean(attribute='has_next'),
-        'data': fields.List(fields.Nested(conversation_fields), attribute='items')
-    }
 
     @setup_required
     @login_required
@@ -191,21 +98,11 @@ class CompletionConversationApi(Resource):
 
 
 class CompletionConversationDetailApi(Resource):
-    conversation_detail_fields = {
-        'id': fields.String,
-        'status': fields.String,
-        'from_source': fields.String,
-        'from_end_user_id': fields.String,
-        'from_account_id': fields.String,
-        'created_at': TimestampField,
-        'model_config': fields.Nested(model_config_fields),
-        'message': fields.Nested(message_detail_fields, attribute='first_message'),
-    }
 
     @setup_required
     @login_required
     @account_initialization_required
-    @marshal_with(conversation_detail_fields)
+    @marshal_with(conversation_message_detail_fields)
     def get(self, app_id, conversation_id):
         app_id = str(app_id)
         conversation_id = str(conversation_id)
@@ -234,44 +131,11 @@ class CompletionConversationDetailApi(Resource):
 
 
 class ChatConversationApi(Resource):
-    simple_configs_fields = {
-        'prompt_template': fields.String,
-    }
-
-    simple_model_config_fields = {
-        'model': fields.Raw(attribute='model_dict'),
-        'pre_prompt': fields.String,
-    }
-
-    conversation_fields = {
-        'id': fields.String,
-        'status': fields.String,
-        'from_source': fields.String,
-        'from_end_user_id': fields.String,
-        'from_end_user_session_id': fields.String,
-        'from_account_id': fields.String,
-        'summary': fields.String(attribute='summary_or_query'),
-        'read_at': TimestampField,
-        'created_at': TimestampField,
-        'annotated': fields.Boolean,
-        'model_config': fields.Nested(simple_model_config_fields),
-        'message_count': fields.Integer,
-        'user_feedback_stats': fields.Nested(feedback_stat_fields),
-        'admin_feedback_stats': fields.Nested(feedback_stat_fields)
-    }
-
-    conversation_pagination_fields = {
-        'page': fields.Integer,
-        'limit': fields.Integer(attribute='per_page'),
-        'total': fields.Integer,
-        'has_more': fields.Boolean(attribute='has_next'),
-        'data': fields.List(fields.Nested(conversation_fields), attribute='items')
-    }
 
     @setup_required
     @login_required
     @account_initialization_required
-    @marshal_with(conversation_pagination_fields)
+    @marshal_with(conversation_with_summary_pagination_fields)
     def get(self, app_id):
         app_id = str(app_id)
 
@@ -356,19 +220,6 @@ class ChatConversationApi(Resource):
 
 
 class ChatConversationDetailApi(Resource):
-    conversation_detail_fields = {
-        'id': fields.String,
-        'status': fields.String,
-        'from_source': fields.String,
-        'from_end_user_id': fields.String,
-        'from_account_id': fields.String,
-        'created_at': TimestampField,
-        'annotated': fields.Boolean,
-        'model_config': fields.Nested(model_config_fields),
-        'message_count': fields.Integer,
-        'user_feedback_stats': fields.Nested(feedback_stat_fields),
-        'admin_feedback_stats': fields.Nested(feedback_stat_fields)
-    }
 
     @setup_required
     @login_required

api/controllers/console/app/generator.py

@@ -1,5 +1,5 @@
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource, reqparse
 
 from controllers.console import api

api/controllers/console/app/message.py

@@ -16,8 +16,9 @@ from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.model_providers.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
     ProviderTokenNotInitError, LLMAPIUnavailableError, QuotaExceededError, ModelCurrentlyNotSupportError
-from core.login.login import login_required
-from libs.helper import uuid_value, TimestampField
+from libs.login import login_required
+from fields.conversation_fields import message_detail_fields
+from libs.helper import uuid_value
 from libs.infinite_scroll_pagination import InfiniteScrollPagination
 from extensions.ext_database import db
 from models.model import MessageAnnotation, Conversation, Message, MessageFeedback
@@ -27,44 +28,6 @@ from services.errors.conversation import ConversationNotExistsError
 from services.errors.message import MessageNotExistsError
 from services.message_service import MessageService
 
-account_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'email': fields.String
-}
-
-feedback_fields = {
-    'rating': fields.String,
-    'content': fields.String,
-    'from_source': fields.String,
-    'from_end_user_id': fields.String,
-    'from_account': fields.Nested(account_fields, allow_null=True),
-}
-
-annotation_fields = {
-    'content': fields.String,
-    'account': fields.Nested(account_fields, allow_null=True),
-    'created_at': TimestampField
-}
-
-message_detail_fields = {
-    'id': fields.String,
-    'conversation_id': fields.String,
-    'inputs': fields.Raw,
-    'query': fields.String,
-    'message': fields.Raw,
-    'message_tokens': fields.Integer,
-    'answer': fields.String,
-    'answer_tokens': fields.Integer,
-    'provider_response_latency': fields.Float,
-    'from_source': fields.String,
-    'from_end_user_id': fields.String,
-    'from_account_id': fields.String,
-    'feedbacks': fields.List(fields.Nested(feedback_fields)),
-    'annotation': fields.Nested(annotation_fields, allow_null=True),
-    'created_at': TimestampField
-}
-
 
 class ChatMessageListApi(Resource):
     message_infinite_scroll_pagination_fields = {

api/controllers/console/app/model_config.py

@@ -1,5 +1,4 @@
 # -*- coding:utf-8 -*-
-import json
 
 from flask import request
 from flask_restful import Resource
@@ -9,7 +8,7 @@ from controllers.console import api
 from controllers.console.app import _get_app
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
-from core.login.login import login_required
+from libs.login import login_required
 from events.app_event import app_model_config_was_updated
 from extensions.ext_database import db
 from models.model import AppModelConfig
@@ -31,7 +30,8 @@ class ModelConfigResource(Resource):
         model_configuration = AppModelConfigService.validate_configuration(
             tenant_id=current_user.current_tenant_id,
             account=current_user,
-            config=request.json
+            config=request.json,
+            mode=app_model.mode
         )
 
         new_app_model_config = AppModelConfig(

api/controllers/console/app/site.py

@@ -1,33 +1,18 @@
 # -*- coding:utf-8 -*-
 from flask_login import current_user
-from core.login.login import login_required
-from flask_restful import Resource, reqparse, fields, marshal_with
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal_with
 from werkzeug.exceptions import NotFound, Forbidden
 
 from controllers.console import api
 from controllers.console.app import _get_app
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
+from fields.app_fields import app_site_fields
 from libs.helper import supported_language
 from extensions.ext_database import db
 from models.model import Site
 
-app_site_fields = {
-    'app_id': fields.String,
-    'access_token': fields.String(attribute='code'),
-    'code': fields.String,
-    'title': fields.String,
-    'icon': fields.String,
-    'icon_background': fields.String,
-    'description': fields.String,
-    'default_language': fields.String,
-    'customize_domain': fields.String,
-    'copyright': fields.String,
-    'privacy_policy': fields.String,
-    'customize_token_strategy': fields.String,
-    'prompt_public': fields.Boolean
-}
-
 
 def parse_app_site_args():
     parser = reqparse.RequestParser()

api/controllers/console/app/statistic.py

@@ -5,7 +5,7 @@ from datetime import datetime
 import pytz
 from flask import jsonify
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource, reqparse
 
 from controllers.console import api

api/controllers/console/auth/data_source_oauth.py

@@ -1,16 +1,13 @@
 import logging
-from datetime import datetime
-from typing import Optional
 
-import flask_login
 import requests
-from flask import request, redirect, current_app, session
+from flask import request, redirect, current_app
 from flask_login import current_user
 
 from flask_restful import Resource
 from werkzeug.exceptions import Forbidden
 
-from core.login.login import login_required
+from libs.login import login_required
 from libs.oauth_data_source import NotionOAuth
 from controllers.console import api
 from ..setup import setup_required
@@ -45,15 +42,34 @@ class OAuthDataSource(Resource):
         if current_app.config.get('NOTION_INTEGRATION_TYPE') == 'internal':
             internal_secret = current_app.config.get('NOTION_INTERNAL_SECRET')
             oauth_provider.save_internal_access_token(internal_secret)
-            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source=success')
+            return { 'data': '' }
         else:
             auth_url = oauth_provider.get_authorization_url()
-            return redirect(auth_url)
+            return { 'data': auth_url }, 200
 
 
 
 
 class OAuthDataSourceCallback(Resource):
+    def get(self, provider: str):
+        OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()
+        with current_app.app_context():
+            oauth_provider = OAUTH_DATASOURCE_PROVIDERS.get(provider)
+        if not oauth_provider:
+            return {'error': 'Invalid provider'}, 400
+        if 'code' in request.args:
+            code = request.args.get('code')
+
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?type=notion&code={code}')
+        elif 'error' in request.args:
+            error = request.args.get('error')
+
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?type=notion&error={error}')
+        else:
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?type=notion&error=Access denied')
+        
+
+class OAuthDataSourceBinding(Resource):
     def get(self, provider: str):
         OAUTH_DATASOURCE_PROVIDERS = get_oauth_providers()
         with current_app.app_context():
@@ -69,12 +85,7 @@ class OAuthDataSourceCallback(Resource):
                     f"An error occurred during the OAuthCallback process with {provider}: {e.response.text}")
                 return {'error': 'OAuth data source process failed'}, 400
 
-            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source=success')
-        elif 'error' in request.args:
-            error = request.args.get('error')
-            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source={error}')
-        else:
-            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_data_source=access_denied')
+            return {'result': 'success'}, 200
 
 
 class OAuthDataSourceSync(Resource):
@@ -101,4 +112,5 @@ class OAuthDataSourceSync(Resource):
 
 api.add_resource(OAuthDataSource, '/oauth/data-source/<string:provider>')
 api.add_resource(OAuthDataSourceCallback, '/oauth/data-source/callback/<string:provider>')
+api.add_resource(OAuthDataSourceBinding, '/oauth/data-source/binding/<string:provider>')
 api.add_resource(OAuthDataSourceSync, '/oauth/data-source/<string:provider>/<uuid:binding_id>/sync')

api/controllers/console/auth/login.py

@@ -6,7 +6,6 @@ from flask_restful import Resource, reqparse
 
 import services
 from controllers.console import api
-from controllers.console.error import AccountNotLinkTenantError
 from controllers.console.setup import setup_required
 from libs.helper import email
 from libs.password import valid_password
@@ -37,12 +36,12 @@ class LoginApi(Resource):
         except Exception:
             pass
 
-        flask_login.login_user(account, remember=args['remember_me'])
         AccountService.update_last_login(account, request)
 
         # todo: return the user info
+        token = AccountService.get_account_jwt_token(account)
 
-        return {'result': 'success'}
+        return {'result': 'success', 'data': token}
 
 
 class LogoutApi(Resource):

api/controllers/console/auth/oauth.py

@@ -2,9 +2,8 @@ import logging
 from datetime import datetime
 from typing import Optional
 
-import flask_login
 import requests
-from flask import request, redirect, current_app, session
+from flask import request, redirect, current_app
 from flask_restful import Resource
 
 from libs.oauth import OAuthUserInfo, GitHubOAuth, GoogleOAuth
@@ -75,12 +74,11 @@ class OAuthCallback(Resource):
             account.initialized_at = datetime.utcnow()
             db.session.commit()
 
-        # login user
-        session.clear()
-        flask_login.login_user(account, remember=True)
         AccountService.update_last_login(account, request)
 
-        return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?oauth_login=success')
+        token = AccountService.get_account_jwt_token(account)
+
+        return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}?console_token={token}')
 
 
 def _get_account_by_openid_or_email(provider: str, user_info: OAuthUserInfo) -> Optional[Account]:

api/controllers/console/datasets/data_source.py

@@ -2,10 +2,10 @@ import datetime
 import json
 
 from cachetools import TTLCache
-from flask import request, current_app
+from flask import request
 from flask_login import current_user
-from core.login.login import login_required
-from flask_restful import Resource, marshal_with, fields, reqparse, marshal
+from libs.login import login_required
+from flask_restful import Resource, marshal_with, reqparse
 from werkzeug.exceptions import NotFound
 
 from controllers.console import api
@@ -14,7 +14,7 @@ from controllers.console.wraps import account_initialization_required
 from core.data_loader.loader.notion import NotionLoader
 from core.indexing_runner import IndexingRunner
 from extensions.ext_database import db
-from libs.helper import TimestampField
+from fields.data_source_fields import integrate_notion_info_list_fields, integrate_list_fields
 from models.dataset import Document
 from models.source import DataSourceBinding
 from services.dataset_service import DatasetService, DocumentService
@@ -24,37 +24,6 @@ cache = TTLCache(maxsize=None, ttl=30)
 
 
 class DataSourceApi(Resource):
-    integrate_icon_fields = {
-        'type': fields.String,
-        'url': fields.String,
-        'emoji': fields.String
-    }
-    integrate_page_fields = {
-        'page_name': fields.String,
-        'page_id': fields.String,
-        'page_icon': fields.Nested(integrate_icon_fields, allow_null=True),
-        'parent_id': fields.String,
-        'type': fields.String
-    }
-    integrate_workspace_fields = {
-        'workspace_name': fields.String,
-        'workspace_id': fields.String,
-        'workspace_icon': fields.String,
-        'pages': fields.List(fields.Nested(integrate_page_fields)),
-        'total': fields.Integer
-    }
-    integrate_fields = {
-        'id': fields.String,
-        'provider': fields.String,
-        'created_at': TimestampField,
-        'is_bound': fields.Boolean,
-        'disabled': fields.Boolean,
-        'link': fields.String,
-        'source_info': fields.Nested(integrate_workspace_fields)
-    }
-    integrate_list_fields = {
-        'data': fields.List(fields.Nested(integrate_fields)),
-    }
 
     @setup_required
     @login_required
@@ -131,28 +100,6 @@ class DataSourceApi(Resource):
 
 
 class DataSourceNotionListApi(Resource):
-    integrate_icon_fields = {
-        'type': fields.String,
-        'url': fields.String,
-        'emoji': fields.String
-    }
-    integrate_page_fields = {
-        'page_name': fields.String,
-        'page_id': fields.String,
-        'page_icon': fields.Nested(integrate_icon_fields, allow_null=True),
-        'is_bound': fields.Boolean,
-        'parent_id': fields.String,
-        'type': fields.String
-    }
-    integrate_workspace_fields = {
-        'workspace_name': fields.String,
-        'workspace_id': fields.String,
-        'workspace_icon': fields.String,
-        'pages': fields.List(fields.Nested(integrate_page_fields))
-    }
-    integrate_notion_info_list_fields = {
-        'notion_info': fields.List(fields.Nested(integrate_workspace_fields)),
-    }
 
     @setup_required
     @login_required

api/controllers/console/datasets/datasets.py

@@ -1,8 +1,11 @@
 # -*- coding:utf-8 -*-
-from flask import request
+import flask_restful
+from flask import request, current_app
 from flask_login import current_user
-from core.login.login import login_required
-from flask_restful import Resource, reqparse, fields, marshal, marshal_with
+
+from controllers.console.apikey import api_key_list, api_key_fields
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal, marshal_with
 from werkzeug.exceptions import NotFound, Forbidden
 import services
 from controllers.console import api
@@ -12,45 +15,16 @@ from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.indexing_runner import IndexingRunner
 from core.model_providers.error import LLMBadRequestError, ProviderTokenNotInitError
-from core.model_providers.model_factory import ModelFactory
 from core.model_providers.models.entity.model_params import ModelType
-from libs.helper import TimestampField
+from fields.app_fields import related_app_list
+from fields.dataset_fields import dataset_detail_fields, dataset_query_detail_fields
+from fields.document_fields import document_status_fields
 from extensions.ext_database import db
 from models.dataset import DocumentSegment, Document
-from models.model import UploadFile
+from models.model import UploadFile, ApiToken
 from services.dataset_service import DatasetService, DocumentService
 from services.provider_service import ProviderService
 
-dataset_detail_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'description': fields.String,
-    'provider': fields.String,
-    'permission': fields.String,
-    'data_source_type': fields.String,
-    'indexing_technique': fields.String,
-    'app_count': fields.Integer,
-    'document_count': fields.Integer,
-    'word_count': fields.Integer,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-    'updated_by': fields.String,
-    'updated_at': TimestampField,
-    'embedding_model': fields.String,
-    'embedding_model_provider': fields.String,
-    'embedding_available': fields.Boolean
-}
-
-dataset_query_detail_fields = {
-    "id": fields.String,
-    "content": fields.String,
-    "source": fields.String,
-    "source_app_id": fields.String,
-    "created_by_role": fields.String,
-    "created_by": fields.String,
-    "created_at": TimestampField
-}
-
 
 def _validate_name(name):
     if not name or len(name) < 1 or len(name) > 40:
@@ -82,7 +56,8 @@ class DatasetListApi(Resource):
 
         # check embedding setting
         provider_service = ProviderService()
-        valid_model_list = provider_service.get_valid_model_list(current_user.current_tenant_id, ModelType.EMBEDDINGS.value)
+        valid_model_list = provider_service.get_valid_model_list(current_user.current_tenant_id,
+                                                                 ModelType.EMBEDDINGS.value)
         # if len(valid_model_list) == 0:
         #     raise ProviderNotInitializeError(
         #         f"No Embedding Model available. Please configure a valid provider "
@@ -157,7 +132,8 @@ class DatasetApi(Resource):
         # check embedding setting
         provider_service = ProviderService()
         # get valid model list
-        valid_model_list = provider_service.get_valid_model_list(current_user.current_tenant_id, ModelType.EMBEDDINGS.value)
+        valid_model_list = provider_service.get_valid_model_list(current_user.current_tenant_id,
+                                                                 ModelType.EMBEDDINGS.value)
         model_names = []
         for valid_model in valid_model_list:
             model_names.append(f"{valid_model['model_name']}:{valid_model['model_provider']['provider_name']}")
@@ -271,7 +247,8 @@ class DatasetIndexingEstimateApi(Resource):
         parser.add_argument('indexing_technique', type=str, required=True, nullable=True, location='json')
         parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
         parser.add_argument('dataset_id', type=str, required=False, nullable=False, location='json')
-        parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False, location='json')
+        parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False,
+                            location='json')
         args = parser.parse_args()
         # validate args
         DocumentService.estimate_args_validate(args)
@@ -320,18 +297,6 @@ class DatasetIndexingEstimateApi(Resource):
 
 
 class DatasetRelatedAppListApi(Resource):
-    app_detail_kernel_fields = {
-        'id': fields.String,
-        'name': fields.String,
-        'mode': fields.String,
-        'icon': fields.String,
-        'icon_background': fields.String,
-    }
-
-    related_app_list = {
-        'data': fields.List(fields.Nested(app_detail_kernel_fields)),
-        'total': fields.Integer,
-    }
 
     @setup_required
     @login_required
@@ -363,24 +328,6 @@ class DatasetRelatedAppListApi(Resource):
 
 
 class DatasetIndexingStatusApi(Resource):
-    document_status_fields = {
-        'id': fields.String,
-        'indexing_status': fields.String,
-        'processing_started_at': TimestampField,
-        'parsing_completed_at': TimestampField,
-        'cleaning_completed_at': TimestampField,
-        'splitting_completed_at': TimestampField,
-        'completed_at': TimestampField,
-        'paused_at': TimestampField,
-        'error': fields.String,
-        'stopped_at': TimestampField,
-        'completed_segments': fields.Integer,
-        'total_segments': fields.Integer,
-    }
-
-    document_status_fields_list = {
-        'data': fields.List(fields.Nested(document_status_fields))
-    }
 
     @setup_required
     @login_required
@@ -400,16 +347,101 @@ class DatasetIndexingStatusApi(Resource):
                                                           DocumentSegment.status != 're_segment').count()
             document.completed_segments = completed_segments
             document.total_segments = total_segments
-            documents_status.append(marshal(document, self.document_status_fields))
+            documents_status.append(marshal(document, document_status_fields))
         data = {
             'data': documents_status
         }
         return data
 
 
+class DatasetApiKeyApi(Resource):
+    max_keys = 10
+    token_prefix = 'dataset-'
+    resource_type = 'dataset'
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(api_key_list)
+    def get(self):
+        keys = db.session.query(ApiToken). \
+            filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id). \
+            all()
+        return {"items": keys}
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(api_key_fields)
+    def post(self):
+        # The role of the current user in the ta table must be admin or owner
+        if current_user.current_tenant.current_role not in ['admin', 'owner']:
+            raise Forbidden()
+
+        current_key_count = db.session.query(ApiToken). \
+            filter(ApiToken.type == self.resource_type, ApiToken.tenant_id == current_user.current_tenant_id). \
+            count()
+
+        if current_key_count >= self.max_keys:
+            flask_restful.abort(
+                400,
+                message=f"Cannot create more than {self.max_keys} API keys for this resource type.",
+                code='max_keys_exceeded'
+            )
+
+        key = ApiToken.generate_api_key(self.token_prefix, 24)
+        api_token = ApiToken()
+        api_token.tenant_id = current_user.current_tenant_id
+        api_token.token = key
+        api_token.type = self.resource_type
+        db.session.add(api_token)
+        db.session.commit()
+        return api_token, 200
+
+
+class DatasetApiDeleteApi(Resource):
+    resource_type = 'dataset'
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def delete(self, api_key_id):
+        api_key_id = str(api_key_id)
+
+        # The role of the current user in the ta table must be admin or owner
+        if current_user.current_tenant.current_role not in ['admin', 'owner']:
+            raise Forbidden()
+
+        key = db.session.query(ApiToken). \
+            filter(ApiToken.tenant_id == current_user.current_tenant_id, ApiToken.type == self.resource_type,
+                   ApiToken.id == api_key_id). \
+            first()
+
+        if key is None:
+            flask_restful.abort(404, message='API key not found')
+
+        db.session.query(ApiToken).filter(ApiToken.id == api_key_id).delete()
+        db.session.commit()
+
+        return {'result': 'success'}, 204
+
+
+class DatasetApiBaseUrlApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        return {
+            'api_base_url': (current_app.config['SERVICE_API_URL'] if current_app.config['SERVICE_API_URL']
+                             else request.host_url.rstrip('/')) + '/v1'
+        }
+
+
 api.add_resource(DatasetListApi, '/datasets')
 api.add_resource(DatasetApi, '/datasets/<uuid:dataset_id>')
 api.add_resource(DatasetQueryApi, '/datasets/<uuid:dataset_id>/queries')
 api.add_resource(DatasetIndexingEstimateApi, '/datasets/indexing-estimate')
 api.add_resource(DatasetRelatedAppListApi, '/datasets/<uuid:dataset_id>/related-apps')
 api.add_resource(DatasetIndexingStatusApi, '/datasets/<uuid:dataset_id>/indexing-status')
+api.add_resource(DatasetApiKeyApi, '/datasets/api-keys')
+api.add_resource(DatasetApiDeleteApi, '/datasets/api-keys/<uuid:api_key_id>')
+api.add_resource(DatasetApiBaseUrlApi, '/datasets/api-base-info')

api/controllers/console/datasets/datasets_document.py

@@ -1,11 +1,10 @@
 # -*- coding:utf-8 -*-
-import random
 from datetime import datetime
 from typing import List
 
 from flask import request, current_app
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource, fields, marshal, marshal_with, reqparse
 from sqlalchemy import desc, asc
 from werkzeug.exceptions import NotFound, Forbidden
@@ -23,7 +22,8 @@ from core.model_providers.error import ProviderTokenNotInitError, QuotaExceededE
     LLMBadRequestError
 from core.model_providers.model_factory import ModelFactory
 from extensions.ext_redis import redis_client
-from libs.helper import TimestampField
+from fields.document_fields import document_with_segments_fields, document_fields, \
+    dataset_and_document_fields, document_status_fields
 from extensions.ext_database import db
 from models.dataset import DatasetProcessRule, Dataset
 from models.dataset import Document, DocumentSegment
@@ -32,64 +32,6 @@ from services.dataset_service import DocumentService, DatasetService
 from tasks.add_document_to_index_task import add_document_to_index_task
 from tasks.remove_document_from_index_task import remove_document_from_index_task
 
-dataset_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'description': fields.String,
-    'permission': fields.String,
-    'data_source_type': fields.String,
-    'indexing_technique': fields.String,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-}
-
-document_fields = {
-    'id': fields.String,
-    'position': fields.Integer,
-    'data_source_type': fields.String,
-    'data_source_info': fields.Raw(attribute='data_source_info_dict'),
-    'dataset_process_rule_id': fields.String,
-    'name': fields.String,
-    'created_from': fields.String,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-    'tokens': fields.Integer,
-    'indexing_status': fields.String,
-    'error': fields.String,
-    'enabled': fields.Boolean,
-    'disabled_at': TimestampField,
-    'disabled_by': fields.String,
-    'archived': fields.Boolean,
-    'display_status': fields.String,
-    'word_count': fields.Integer,
-    'hit_count': fields.Integer,
-    'doc_form': fields.String,
-}
-
-document_with_segments_fields = {
-    'id': fields.String,
-    'position': fields.Integer,
-    'data_source_type': fields.String,
-    'data_source_info': fields.Raw(attribute='data_source_info_dict'),
-    'dataset_process_rule_id': fields.String,
-    'name': fields.String,
-    'created_from': fields.String,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-    'tokens': fields.Integer,
-    'indexing_status': fields.String,
-    'error': fields.String,
-    'enabled': fields.Boolean,
-    'disabled_at': TimestampField,
-    'disabled_by': fields.String,
-    'archived': fields.Boolean,
-    'display_status': fields.String,
-    'word_count': fields.Integer,
-    'hit_count': fields.Integer,
-    'completed_segments': fields.Integer,
-    'total_segments': fields.Integer
-}
-
 
 class DocumentResource(Resource):
     def get_document(self, dataset_id: str, document_id: str) -> Document:
@@ -303,11 +245,6 @@ class DatasetDocumentListApi(Resource):
 
 
 class DatasetInitApi(Resource):
-    dataset_and_document_fields = {
-        'dataset': fields.Nested(dataset_fields),
-        'documents': fields.List(fields.Nested(document_fields)),
-        'batch': fields.String
-    }
 
     @setup_required
     @login_required
@@ -504,24 +441,6 @@ class DocumentBatchIndexingEstimateApi(DocumentResource):
 
 
 class DocumentBatchIndexingStatusApi(DocumentResource):
-    document_status_fields = {
-        'id': fields.String,
-        'indexing_status': fields.String,
-        'processing_started_at': TimestampField,
-        'parsing_completed_at': TimestampField,
-        'cleaning_completed_at': TimestampField,
-        'splitting_completed_at': TimestampField,
-        'completed_at': TimestampField,
-        'paused_at': TimestampField,
-        'error': fields.String,
-        'stopped_at': TimestampField,
-        'completed_segments': fields.Integer,
-        'total_segments': fields.Integer,
-    }
-
-    document_status_fields_list = {
-        'data': fields.List(fields.Nested(document_status_fields))
-    }
 
     @setup_required
     @login_required
@@ -541,7 +460,7 @@ class DocumentBatchIndexingStatusApi(DocumentResource):
             document.total_segments = total_segments
             if document.is_paused:
                 document.indexing_status = 'paused'
-            documents_status.append(marshal(document, self.document_status_fields))
+            documents_status.append(marshal(document, document_status_fields))
         data = {
             'data': documents_status
         }
@@ -549,20 +468,6 @@ class DocumentBatchIndexingStatusApi(DocumentResource):
 
 
 class DocumentIndexingStatusApi(DocumentResource):
-    document_status_fields = {
-        'id': fields.String,
-        'indexing_status': fields.String,
-        'processing_started_at': TimestampField,
-        'parsing_completed_at': TimestampField,
-        'cleaning_completed_at': TimestampField,
-        'splitting_completed_at': TimestampField,
-        'completed_at': TimestampField,
-        'paused_at': TimestampField,
-        'error': fields.String,
-        'stopped_at': TimestampField,
-        'completed_segments': fields.Integer,
-        'total_segments': fields.Integer,
-    }
 
     @setup_required
     @login_required
@@ -586,7 +491,7 @@ class DocumentIndexingStatusApi(DocumentResource):
         document.total_segments = total_segments
         if document.is_paused:
             document.indexing_status = 'paused'
-        return marshal(document, self.document_status_fields)
+        return marshal(document, document_status_fields)
 
 
 class DocumentDetailApi(DocumentResource):

api/controllers/console/datasets/datasets_segments.py

@@ -3,7 +3,7 @@ import uuid
 from datetime import datetime
 from flask import request
 from flask_login import current_user
-from flask_restful import Resource, reqparse, fields, marshal
+from flask_restful import Resource, reqparse, marshal
 from werkzeug.exceptions import NotFound, Forbidden
 
 import services
@@ -14,48 +14,18 @@ from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.model_providers.error import LLMBadRequestError, ProviderTokenNotInitError
 from core.model_providers.model_factory import ModelFactory
-from core.login.login import login_required
+from libs.login import login_required
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
+from fields.segment_fields import segment_fields
 from models.dataset import DocumentSegment
 
-from libs.helper import TimestampField
 from services.dataset_service import DatasetService, DocumentService, SegmentService
 from tasks.enable_segment_to_index_task import enable_segment_to_index_task
 from tasks.disable_segment_from_index_task import disable_segment_from_index_task
 from tasks.batch_create_segment_to_index_task import batch_create_segment_to_index_task
 import pandas as pd
 
-segment_fields = {
-    'id': fields.String,
-    'position': fields.Integer,
-    'document_id': fields.String,
-    'content': fields.String,
-    'answer': fields.String,
-    'word_count': fields.Integer,
-    'tokens': fields.Integer,
-    'keywords': fields.List(fields.String),
-    'index_node_id': fields.String,
-    'index_node_hash': fields.String,
-    'hit_count': fields.Integer,
-    'enabled': fields.Boolean,
-    'disabled_at': TimestampField,
-    'disabled_by': fields.String,
-    'status': fields.String,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-    'indexing_at': TimestampField,
-    'completed_at': TimestampField,
-    'error': fields.String,
-    'stopped_at': TimestampField
-}
-
-segment_list_response = {
-    'data': fields.List(fields.Nested(segment_fields)),
-    'has_more': fields.Boolean,
-    'limit': fields.Integer
-}
-
 
 class DatasetDocumentSegmentListApi(Resource):
     @setup_required

api/controllers/console/datasets/file.py

@@ -1,28 +1,19 @@
-import datetime
-import hashlib
-import tempfile
-import chardet
-import time
-import uuid
-from pathlib import Path
-
 from cachetools import TTLCache
 from flask import request, current_app
-from flask_login import current_user
-from core.login.login import login_required
-from flask_restful import Resource, marshal_with, fields
-from werkzeug.exceptions import NotFound
+
+import services
+from libs.login import login_required
+from flask_restful import Resource, marshal_with
 
 from controllers.console import api
 from controllers.console.datasets.error import NoFileUploadedError, TooManyFilesError, FileTooLargeError, \
     UnsupportedFileTypeError
+
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
-from core.data_loader.file_extractor import FileExtractor
-from extensions.ext_storage import storage
-from libs.helper import TimestampField
-from extensions.ext_database import db
-from models.model import UploadFile
+from fields.file_fields import upload_config_fields, file_fields
+
+from services.file_service import FileService
 
 cache = TTLCache(maxsize=None, ttl=30)
 
@@ -31,10 +22,6 @@ PREVIEW_WORDS_LIMIT = 3000
 
 
 class FileApi(Resource):
-    upload_config_fields = {
-        'file_size_limit': fields.Integer,
-        'batch_count_limit': fields.Integer
-    }
 
     @setup_required
     @login_required
@@ -48,16 +35,6 @@ class FileApi(Resource):
             'batch_count_limit': batch_count_limit
         }, 200
 
-    file_fields = {
-        'id': fields.String,
-        'name': fields.String,
-        'size': fields.Integer,
-        'extension': fields.String,
-        'mime_type': fields.String,
-        'created_by': fields.String,
-        'created_at': TimestampField,
-    }
-
     @setup_required
     @login_required
     @account_initialization_required
@@ -73,45 +50,13 @@ class FileApi(Resource):
 
         if len(request.files) > 1:
             raise TooManyFilesError()
-
-        file_content = file.read()
-        file_size = len(file_content)
-
-        file_size_limit = current_app.config.get("UPLOAD_FILE_SIZE_LIMIT") * 1024 * 1024
-        if file_size > file_size_limit:
-            message = "({file_size} > {file_size_limit})"
-            raise FileTooLargeError(message)
-
-        extension = file.filename.split('.')[-1]
-        if extension.lower() not in ALLOWED_EXTENSIONS:
+        try:
+            upload_file = FileService.upload_file(file)
+        except services.errors.file.FileTooLargeError as file_too_large_error:
+            raise FileTooLargeError(file_too_large_error.description)
+        except services.errors.file.UnsupportedFileTypeError:
             raise UnsupportedFileTypeError()
 
-        # user uuid as file name
-        file_uuid = str(uuid.uuid4())
-        file_key = 'upload_files/' + current_user.current_tenant_id + '/' + file_uuid + '.' + extension
-
-        # save file to storage
-        storage.save(file_key, file_content)
-
-        # save file to db
-        config = current_app.config
-        upload_file = UploadFile(
-            tenant_id=current_user.current_tenant_id,
-            storage_type=config['STORAGE_TYPE'],
-            key=file_key,
-            name=file.filename,
-            size=file_size,
-            extension=extension,
-            mime_type=file.mimetype,
-            created_by=current_user.id,
-            created_at=datetime.datetime.utcnow(),
-            used=False,
-            hash=hashlib.sha3_256(file_content).hexdigest()
-        )
-
-        db.session.add(upload_file)
-        db.session.commit()
-
         return upload_file, 201
 
 
@@ -121,26 +66,7 @@ class FilePreviewApi(Resource):
     @account_initialization_required
     def get(self, file_id):
         file_id = str(file_id)
-
-        key = file_id + request.path
-        cached_response = cache.get(key)
-        if cached_response and time.time() - cached_response['timestamp'] < cache.ttl:
-            return cached_response['response']
-
-        upload_file = db.session.query(UploadFile) \
-            .filter(UploadFile.id == file_id) \
-            .first()
-
-        if not upload_file:
-            raise NotFound("File not found")
-
-        # extract text from file
-        extension = upload_file.extension
-        if extension.lower() not in ALLOWED_EXTENSIONS:
-            raise UnsupportedFileTypeError()
-
-        text = FileExtractor.load(upload_file, return_text=True)
-        text = text[0:PREVIEW_WORDS_LIMIT] if text else ''
+        text = FileService.get_file_preview(file_id)
         return {'content': text}
 
 

api/controllers/console/datasets/hit_testing.py

@@ -1,8 +1,8 @@
 import logging
 
 from flask_login import current_user
-from core.login.login import login_required
-from flask_restful import Resource, reqparse, marshal, fields
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal
 from werkzeug.exceptions import InternalServerError, NotFound, Forbidden
 
 import services
@@ -14,48 +14,10 @@ from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.model_providers.error import ProviderTokenNotInitError, QuotaExceededError, ModelCurrentlyNotSupportError, \
     LLMBadRequestError
-from libs.helper import TimestampField
+from fields.hit_testing_fields import hit_testing_record_fields
 from services.dataset_service import DatasetService
 from services.hit_testing_service import HitTestingService
 
-document_fields = {
-    'id': fields.String,
-    'data_source_type': fields.String,
-    'name': fields.String,
-    'doc_type': fields.String,
-}
-
-segment_fields = {
-    'id': fields.String,
-    'position': fields.Integer,
-    'document_id': fields.String,
-    'content': fields.String,
-    'answer': fields.String,
-    'word_count': fields.Integer,
-    'tokens': fields.Integer,
-    'keywords': fields.List(fields.String),
-    'index_node_id': fields.String,
-    'index_node_hash': fields.String,
-    'hit_count': fields.Integer,
-    'enabled': fields.Boolean,
-    'disabled_at': TimestampField,
-    'disabled_by': fields.String,
-    'status': fields.String,
-    'created_by': fields.String,
-    'created_at': TimestampField,
-    'indexing_at': TimestampField,
-    'completed_at': TimestampField,
-    'error': fields.String,
-    'stopped_at': TimestampField,
-    'document': fields.Nested(document_fields),
-}
-
-hit_testing_record_fields = {
-    'segment': fields.Nested(segment_fields),
-    'score': fields.Float,
-    'tsne_position': fields.Raw
-}
-
 
 class HitTestingApi(Resource):
 

api/controllers/console/explore/conversation.py

@@ -7,26 +7,12 @@ from werkzeug.exceptions import NotFound
 from controllers.console import api
 from controllers.console.explore.error import NotChatAppError
 from controllers.console.explore.wraps import InstalledAppResource
+from fields.conversation_fields import conversation_infinite_scroll_pagination_fields, simple_conversation_fields
 from libs.helper import TimestampField, uuid_value
 from services.conversation_service import ConversationService
 from services.errors.conversation import LastConversationNotExistsError, ConversationNotExistsError
 from services.web_conversation_service import WebConversationService
 
-conversation_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'inputs': fields.Raw,
-    'status': fields.String,
-    'introduction': fields.String,
-    'created_at': TimestampField
-}
-
-conversation_infinite_scroll_pagination_fields = {
-    'limit': fields.Integer,
-    'has_more': fields.Boolean,
-    'data': fields.List(fields.Nested(conversation_fields))
-}
-
 
 class ConversationListApi(InstalledAppResource):
 
@@ -76,7 +62,7 @@ class ConversationApi(InstalledAppResource):
 
 class ConversationRenameApi(InstalledAppResource):
 
-    @marshal_with(conversation_fields)
+    @marshal_with(simple_conversation_fields)
     def post(self, installed_app, c_id):
         app_model = installed_app.app
         if app_model.mode != 'chat':

api/controllers/console/explore/installed_app.py

@@ -2,8 +2,8 @@
 from datetime import datetime
 
 from flask_login import current_user
-from core.login.login import login_required
-from flask_restful import Resource, reqparse, fields, marshal_with, inputs
+from libs.login import login_required
+from flask_restful import Resource, reqparse, marshal_with, inputs
 from sqlalchemy import and_
 from werkzeug.exceptions import NotFound, Forbidden, BadRequest
 
@@ -11,32 +11,10 @@ from controllers.console import api
 from controllers.console.explore.wraps import InstalledAppResource
 from controllers.console.wraps import account_initialization_required
 from extensions.ext_database import db
-from libs.helper import TimestampField
+from fields.installed_app_fields import installed_app_list_fields
 from models.model import App, InstalledApp, RecommendedApp
 from services.account_service import TenantService
 
-app_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'mode': fields.String,
-    'icon': fields.String,
-    'icon_background': fields.String
-}
-
-installed_app_fields = {
-    'id': fields.String,
-    'app': fields.Nested(app_fields),
-    'app_owner_tenant_id': fields.String,
-    'is_pinned': fields.Boolean,
-    'last_used_at': TimestampField,
-    'editable': fields.Boolean,
-    'uninstallable': fields.Boolean,
-}
-
-installed_app_list_fields = {
-    'installed_apps': fields.List(fields.Nested(installed_app_fields))
-}
-
 
 class InstalledAppsListApi(Resource):
     @login_required

api/controllers/console/explore/message.py

@@ -17,6 +17,7 @@ from controllers.console.explore.error import NotCompletionAppError, AppSuggeste
 from controllers.console.explore.wraps import InstalledAppResource
 from core.model_providers.error import LLMRateLimitError, LLMBadRequestError, LLMAuthorizationError, LLMAPIConnectionError, \
     ProviderTokenNotInitError, LLMAPIUnavailableError, QuotaExceededError, ModelCurrentlyNotSupportError
+from fields.message_fields import message_infinite_scroll_pagination_fields
 from libs.helper import uuid_value, TimestampField
 from services.completion_service import CompletionService
 from services.errors.app import MoreLikeThisDisabledError
@@ -26,45 +27,6 @@ from services.message_service import MessageService
 
 
 class MessageListApi(InstalledAppResource):
-    feedback_fields = {
-        'rating': fields.String
-    }
-
-    retriever_resource_fields = {
-        'id': fields.String,
-        'message_id': fields.String,
-        'position': fields.Integer,
-        'dataset_id': fields.String,
-        'dataset_name': fields.String,
-        'document_id': fields.String,
-        'document_name': fields.String,
-        'data_source_type': fields.String,
-        'segment_id': fields.String,
-        'score': fields.Float,
-        'hit_count': fields.Integer,
-        'word_count': fields.Integer,
-        'segment_position': fields.Integer,
-        'index_node_hash': fields.String,
-        'content': fields.String,
-        'created_at': TimestampField
-    }
-
-    message_fields = {
-        'id': fields.String,
-        'conversation_id': fields.String,
-        'inputs': fields.Raw,
-        'query': fields.String,
-        'answer': fields.String,
-        'feedback': fields.Nested(feedback_fields, attribute='user_feedback', allow_null=True),
-        'retriever_resources': fields.List(fields.Nested(retriever_resource_fields)),
-        'created_at': TimestampField
-    }
-
-    message_infinite_scroll_pagination_fields = {
-        'limit': fields.Integer,
-        'has_more': fields.Boolean,
-        'data': fields.List(fields.Nested(message_fields))
-    }
 
     @marshal_with(message_infinite_scroll_pagination_fields)
     def get(self, installed_app):

api/controllers/console/explore/recommended_app.py

@@ -1,6 +1,6 @@
 # -*- coding:utf-8 -*-
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource, fields, marshal_with
 from sqlalchemy import and_
 

api/controllers/console/explore/wraps.py

@@ -1,5 +1,5 @@
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource
 from functools import wraps
 

api/controllers/console/setup.py

@@ -1,7 +1,6 @@
 # -*- coding:utf-8 -*-
 from functools import wraps
 
-import flask_login
 from flask import request, current_app
 from flask_restful import Resource, reqparse
 
@@ -58,9 +57,6 @@ class SetupApi(Resource):
         )
 
         setup()
-
-        # Login
-        flask_login.login_user(account)
         AccountService.update_last_login(account, request)
 
         return {'result': 'success'}, 201

api/controllers/console/universal_chat/conversation.py

@@ -6,31 +6,17 @@ from werkzeug.exceptions import NotFound
 
 from controllers.console import api
 from controllers.console.universal_chat.wraps import UniversalChatResource
+from fields.conversation_fields import conversation_with_model_config_infinite_scroll_pagination_fields, \
+    conversation_with_model_config_fields
 from libs.helper import TimestampField, uuid_value
 from services.conversation_service import ConversationService
 from services.errors.conversation import LastConversationNotExistsError, ConversationNotExistsError
 from services.web_conversation_service import WebConversationService
 
-conversation_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'inputs': fields.Raw,
-    'status': fields.String,
-    'introduction': fields.String,
-    'created_at': TimestampField,
-    'model_config': fields.Raw,
-}
-
-conversation_infinite_scroll_pagination_fields = {
-    'limit': fields.Integer,
-    'has_more': fields.Boolean,
-    'data': fields.List(fields.Nested(conversation_fields))
-}
-
 
 class UniversalChatConversationListApi(UniversalChatResource):
 
-    @marshal_with(conversation_infinite_scroll_pagination_fields)
+    @marshal_with(conversation_with_model_config_infinite_scroll_pagination_fields)
     def get(self, universal_app):
         app_model = universal_app
 
@@ -73,7 +59,7 @@ class UniversalChatConversationApi(UniversalChatResource):
 
 class UniversalChatConversationRenameApi(UniversalChatResource):
 
-    @marshal_with(conversation_fields)
+    @marshal_with(conversation_with_model_config_fields)
     def post(self, universal_app, c_id):
         app_model = universal_app
         conversation_id = str(c_id)

api/controllers/console/universal_chat/wraps.py

@@ -2,7 +2,7 @@ import json
 from functools import wraps
 
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required

api/controllers/console/workspace/account.py

@@ -4,7 +4,7 @@ from datetime import datetime
 import pytz
 from flask import current_app, request
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource, reqparse, fields, marshal_with
 
 from services.errors.account import CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError

api/controllers/console/workspace/members.py

@@ -1,7 +1,7 @@
 # -*- coding:utf-8 -*-
 from flask import current_app
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource, reqparse, marshal_with, abort, fields, marshal
 
 import services

api/controllers/console/workspace/model_providers.py

@@ -1,5 +1,5 @@
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource, reqparse
 from werkzeug.exceptions import Forbidden
 

api/controllers/console/workspace/models.py

@@ -1,5 +1,5 @@
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource, reqparse
 
 from controllers.console import api

api/controllers/console/workspace/providers.py

@@ -1,6 +1,6 @@
 # -*- coding:utf-8 -*-
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource, reqparse
 from werkzeug.exceptions import Forbidden
 

api/controllers/console/workspace/tool_providers.py

@@ -1,7 +1,7 @@
 import json
 
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource, abort, reqparse
 from werkzeug.exceptions import Forbidden
 

api/controllers/console/workspace/workspace.py

@@ -3,9 +3,8 @@ import logging
 
 from flask import request
 from flask_login import current_user
-from core.login.login import login_required
+from libs.login import login_required
 from flask_restful import Resource, fields, marshal_with, reqparse, marshal, inputs
-from flask_restful.inputs import int_range
 
 from controllers.console import api
 from controllers.console.admin import admin_required

api/controllers/service_api/__init__.py

@@ -9,4 +9,4 @@ api = ExternalApi(bp)
 
 from .app import completion, app, conversation, message, audio
 
-from .dataset import document
+from .dataset import document, segment, dataset

api/controllers/service_api/app/conversation.py

@@ -8,25 +8,11 @@ from controllers.service_api import api
 from controllers.service_api.app import create_or_update_end_user_for_user_id
 from controllers.service_api.app.error import NotChatAppError
 from controllers.service_api.wraps import AppApiResource
+from fields.conversation_fields import conversation_infinite_scroll_pagination_fields, simple_conversation_fields
 from libs.helper import TimestampField, uuid_value
 import services
 from services.conversation_service import ConversationService
 
-conversation_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'inputs': fields.Raw,
-    'status': fields.String,
-    'introduction': fields.String,
-    'created_at': TimestampField
-}
-
-conversation_infinite_scroll_pagination_fields = {
-    'limit': fields.Integer,
-    'has_more': fields.Boolean,
-    'data': fields.List(fields.Nested(conversation_fields))
-}
-
 
 class ConversationApi(AppApiResource):
 
@@ -50,7 +36,7 @@ class ConversationApi(AppApiResource):
             raise NotFound("Last Conversation Not Exists.")
 
 class ConversationDetailApi(AppApiResource):
-    @marshal_with(conversation_fields)
+    @marshal_with(simple_conversation_fields)
     def delete(self, app_model, end_user, c_id):
         if app_model.mode != 'chat':
             raise NotChatAppError()
@@ -70,7 +56,7 @@ class ConversationDetailApi(AppApiResource):
 
 class ConversationRenameApi(AppApiResource):
 
-    @marshal_with(conversation_fields)
+    @marshal_with(simple_conversation_fields)
     def post(self, app_model, end_user, c_id):
         if app_model.mode != 'chat':
             raise NotChatAppError()

api/controllers/service_api/dataset/dataset.py

@@ -0,0 +1,81 @@
+from flask import request
+from flask_restful import reqparse, marshal
+import services.dataset_service
+from controllers.service_api import api
+from controllers.service_api.dataset.error import DatasetNameDuplicateError
+from controllers.service_api.wraps import DatasetApiResource
+from libs.login import current_user
+from core.model_providers.models.entity.model_params import ModelType
+from fields.dataset_fields import dataset_detail_fields
+from services.dataset_service import DatasetService
+from services.provider_service import ProviderService
+
+
+def _validate_name(name):
+    if not name or len(name) < 1 or len(name) > 40:
+        raise ValueError('Name must be between 1 to 40 characters.')
+    return name
+
+
+class DatasetApi(DatasetApiResource):
+    """Resource for get datasets."""
+
+    def get(self, tenant_id):
+        page = request.args.get('page', default=1, type=int)
+        limit = request.args.get('limit', default=20, type=int)
+        provider = request.args.get('provider', default="vendor")
+        datasets, total = DatasetService.get_datasets(page, limit, provider,
+                                                      tenant_id, current_user)
+        # check embedding setting
+        provider_service = ProviderService()
+        valid_model_list = provider_service.get_valid_model_list(current_user.current_tenant_id,
+                                                                 ModelType.EMBEDDINGS.value)
+        model_names = []
+        for valid_model in valid_model_list:
+            model_names.append(f"{valid_model['model_name']}:{valid_model['model_provider']['provider_name']}")
+        data = marshal(datasets, dataset_detail_fields)
+        for item in data:
+            if item['indexing_technique'] == 'high_quality':
+                item_model = f"{item['embedding_model']}:{item['embedding_model_provider']}"
+                if item_model in model_names:
+                    item['embedding_available'] = True
+                else:
+                    item['embedding_available'] = False
+            else:
+                item['embedding_available'] = True
+        response = {
+            'data': data,
+            'has_more': len(datasets) == limit,
+            'limit': limit,
+            'total': total,
+            'page': page
+        }
+        return response, 200
+
+    """Resource for datasets."""
+
+    def post(self, tenant_id):
+        parser = reqparse.RequestParser()
+        parser.add_argument('name', nullable=False, required=True,
+                            help='type is required. Name must be between 1 to 40 characters.',
+                            type=_validate_name)
+        parser.add_argument('indexing_technique', type=str, location='json',
+                            choices=('high_quality', 'economy'),
+                            help='Invalid indexing technique.')
+        args = parser.parse_args()
+
+        try:
+            dataset = DatasetService.create_empty_dataset(
+                tenant_id=tenant_id,
+                name=args['name'],
+                indexing_technique=args['indexing_technique'],
+                account=current_user
+            )
+        except services.errors.dataset.DatasetNameDuplicateError:
+            raise DatasetNameDuplicateError()
+
+        return marshal(dataset, dataset_detail_fields), 200
+
+
+api.add_resource(DatasetApi, '/datasets')
+

api/controllers/service_api/dataset/document.py

@@ -1,114 +1,287 @@
-import datetime
-import uuid
+import json
 
-from flask import current_app
-from flask_restful import reqparse
+from flask import request
+from flask_restful import reqparse, marshal
+from sqlalchemy import desc
 from werkzeug.exceptions import NotFound
 
 import services.dataset_service
 from controllers.service_api import api
 from controllers.service_api.app.error import ProviderNotInitializeError
 from controllers.service_api.dataset.error import ArchivedDocumentImmutableError, DocumentIndexingError, \
-    DatasetNotInitedError
+    NoFileUploadedError, TooManyFilesError
 from controllers.service_api.wraps import DatasetApiResource
+from libs.login import current_user
 from core.model_providers.error import ProviderTokenNotInitError
 from extensions.ext_database import db
-from extensions.ext_storage import storage
-from models.model import UploadFile
+from fields.document_fields import document_fields, document_status_fields
+from models.dataset import Dataset, Document, DocumentSegment
 from services.dataset_service import DocumentService
+from services.file_service import FileService
 
 
-class DocumentListApi(DatasetApiResource):
+class DocumentAddByTextApi(DatasetApiResource):
     """Resource for documents."""
 
-    def post(self, dataset):
-        """Create document."""
+    def post(self, tenant_id, dataset_id):
+        """Create document by text."""
         parser = reqparse.RequestParser()
         parser.add_argument('name', type=str, required=True, nullable=False, location='json')
         parser.add_argument('text', type=str, required=True, nullable=False, location='json')
-        parser.add_argument('doc_type', type=str, location='json')
-        parser.add_argument('doc_metadata', type=dict, location='json')
+        parser.add_argument('process_rule', type=dict, required=False, nullable=True, location='json')
+        parser.add_argument('original_document_id', type=str, required=False, location='json')
+        parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
+        parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False,
+                            location='json')
+        parser.add_argument('indexing_technique', type=str, choices=Dataset.INDEXING_TECHNIQUE_LIST, nullable=False,
+                            location='json')
         args = parser.parse_args()
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
 
-        if not dataset.indexing_technique:
-            raise DatasetNotInitedError("Dataset indexing technique must be set.")
+        if not dataset:
+            raise ValueError('Dataset is not exist.')
 
-        doc_type = args.get('doc_type')
-        doc_metadata = args.get('doc_metadata')
+        if not dataset.indexing_technique and not args['indexing_technique']:
+            raise ValueError('indexing_technique is required.')
 
-        if doc_type and doc_type not in DocumentService.DOCUMENT_METADATA_SCHEMA:
-            raise ValueError('Invalid doc_type.')
-
-        # user uuid as file name
-        file_uuid = str(uuid.uuid4())
-        file_key = 'upload_files/' + dataset.tenant_id + '/' + file_uuid + '.txt'
-
-        # save file to storage
-        storage.save(file_key, args.get('text'))
-
-        # save file to db
-        config = current_app.config
-        upload_file = UploadFile(
-            tenant_id=dataset.tenant_id,
-            storage_type=config['STORAGE_TYPE'],
-            key=file_key,
-            name=args.get('name') + '.txt',
-            size=len(args.get('text')),
-            extension='txt',
-            mime_type='text/plain',
-            created_by=dataset.created_by,
-            created_at=datetime.datetime.utcnow(),
-            used=True,
-            used_by=dataset.created_by,
-            used_at=datetime.datetime.utcnow()
-        )
-
-        db.session.add(upload_file)
-        db.session.commit()
-
-        document_data = {
-            'data_source': {
-                'type': 'upload_file',
-                'info': [
-                    {
-                        'upload_file_id': upload_file.id
-                    }
-                ]
+        upload_file = FileService.upload_text(args.get('text'), args.get('name'))
+        data_source = {
+            'type': 'upload_file',
+            'info_list': {
+                'data_source_type': 'upload_file',
+                'file_info_list': {
+                    'file_ids': [upload_file.id]
+                }
             }
         }
+        args['data_source'] = data_source
+        # validate args
+        DocumentService.document_create_args_validate(args)
 
         try:
             documents, batch = DocumentService.save_document_with_dataset_id(
                 dataset=dataset,
-                document_data=document_data,
-                account=dataset.created_by_account,
-                dataset_process_rule=dataset.latest_process_rule,
+                document_data=args,
+                account=current_user,
+                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
                 created_from='api'
             )
         except ProviderTokenNotInitError as ex:
             raise ProviderNotInitializeError(ex.description)
         document = documents[0]
-        if doc_type and doc_metadata:
-            metadata_schema = DocumentService.DOCUMENT_METADATA_SCHEMA[doc_type]
 
-            document.doc_metadata = {}
-
-            for key, value_type in metadata_schema.items():
-                value = doc_metadata.get(key)
-                if value is not None and isinstance(value, value_type):
-                    document.doc_metadata[key] = value
-
-            document.doc_type = doc_type
-            document.updated_at = datetime.datetime.utcnow()
-            db.session.commit()
-
-        return {'id': document.id}
+        documents_and_batch_fields = {
+            'document': marshal(document, document_fields),
+            'batch': batch
+        }
+        return documents_and_batch_fields, 200
 
 
-class DocumentApi(DatasetApiResource):
-    def delete(self, dataset, document_id):
+class DocumentUpdateByTextApi(DatasetApiResource):
+    """Resource for update documents."""
+
+    def post(self, tenant_id, dataset_id, document_id):
+        """Update document by text."""
+        parser = reqparse.RequestParser()
+        parser.add_argument('name', type=str, required=False, nullable=True, location='json')
+        parser.add_argument('text', type=str, required=False, nullable=True, location='json')
+        parser.add_argument('process_rule', type=dict, required=False, nullable=True, location='json')
+        parser.add_argument('doc_form', type=str, default='text_model', required=False, nullable=False, location='json')
+        parser.add_argument('doc_language', type=str, default='English', required=False, nullable=False,
+                            location='json')
+        args = parser.parse_args()
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+
+        if not dataset:
+            raise ValueError('Dataset is not exist.')
+
+        if args['text']:
+            upload_file = FileService.upload_text(args.get('text'), args.get('name'))
+            data_source = {
+                'type': 'upload_file',
+                'info_list': {
+                    'data_source_type': 'upload_file',
+                    'file_info_list': {
+                        'file_ids': [upload_file.id]
+                    }
+                }
+            }
+            args['data_source'] = data_source
+        # validate args
+        args['original_document_id'] = str(document_id)
+        DocumentService.document_create_args_validate(args)
+
+        try:
+            documents, batch = DocumentService.save_document_with_dataset_id(
+                dataset=dataset,
+                document_data=args,
+                account=current_user,
+                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
+                created_from='api'
+            )
+        except ProviderTokenNotInitError as ex:
+            raise ProviderNotInitializeError(ex.description)
+        document = documents[0]
+
+        documents_and_batch_fields = {
+            'document': marshal(document, document_fields),
+            'batch': batch
+        }
+        return documents_and_batch_fields, 200
+
+
+class DocumentAddByFileApi(DatasetApiResource):
+    """Resource for documents."""
+    def post(self, tenant_id, dataset_id):
+        """Create document by upload file."""
+        args = {}
+        if 'data' in request.form:
+            args = json.loads(request.form['data'])
+        if 'doc_form' not in args:
+            args['doc_form'] = 'text_model'
+        if 'doc_language' not in args:
+            args['doc_language'] = 'English'
+        # get dataset info
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+
+        if not dataset:
+            raise ValueError('Dataset is not exist.')
+        if not dataset.indexing_technique and not args['indexing_technique']:
+            raise ValueError('indexing_technique is required.')
+
+        # save file info
+        file = request.files['file']
+        # check file
+        if 'file' not in request.files:
+            raise NoFileUploadedError()
+
+        if len(request.files) > 1:
+            raise TooManyFilesError()
+
+        upload_file = FileService.upload_file(file)
+        data_source = {
+            'type': 'upload_file',
+            'info_list': {
+                'file_info_list': {
+                    'file_ids': [upload_file.id]
+                }
+            }
+        }
+        args['data_source'] = data_source
+        # validate args
+        DocumentService.document_create_args_validate(args)
+
+        try:
+            documents, batch = DocumentService.save_document_with_dataset_id(
+                dataset=dataset,
+                document_data=args,
+                account=dataset.created_by_account,
+                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
+                created_from='api'
+            )
+        except ProviderTokenNotInitError as ex:
+            raise ProviderNotInitializeError(ex.description)
+        document = documents[0]
+        documents_and_batch_fields = {
+            'document': marshal(document, document_fields),
+            'batch': batch
+        }
+        return documents_and_batch_fields, 200
+
+
+class DocumentUpdateByFileApi(DatasetApiResource):
+    """Resource for update documents."""
+
+    def post(self, tenant_id, dataset_id, document_id):
+        """Update document by upload file."""
+        args = {}
+        if 'data' in request.form:
+            args = json.loads(request.form['data'])
+        if 'doc_form' not in args:
+            args['doc_form'] = 'text_model'
+        if 'doc_language' not in args:
+            args['doc_language'] = 'English'
+
+        # get dataset info
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+
+        if not dataset:
+            raise ValueError('Dataset is not exist.')
+        if 'file' in request.files:
+            # save file info
+            file = request.files['file']
+
+
+            if len(request.files) > 1:
+                raise TooManyFilesError()
+
+            upload_file = FileService.upload_file(file)
+            data_source = {
+                'type': 'upload_file',
+                'info_list': {
+                    'file_info_list': {
+                        'file_ids': [upload_file.id]
+                    }
+                }
+            }
+            args['data_source'] = data_source
+        # validate args
+        args['original_document_id'] = str(document_id)
+        DocumentService.document_create_args_validate(args)
+
+        try:
+            documents, batch = DocumentService.save_document_with_dataset_id(
+                dataset=dataset,
+                document_data=args,
+                account=dataset.created_by_account,
+                dataset_process_rule=dataset.latest_process_rule if 'process_rule' not in args else None,
+                created_from='api'
+            )
+        except ProviderTokenNotInitError as ex:
+            raise ProviderNotInitializeError(ex.description)
+        document = documents[0]
+        documents_and_batch_fields = {
+            'document': marshal(document, document_fields),
+            'batch': batch
+        }
+        return documents_and_batch_fields, 200
+
+
+class DocumentDeleteApi(DatasetApiResource):
+    def delete(self, tenant_id, dataset_id, document_id):
         """Delete document."""
         document_id = str(document_id)
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+
+        # get dataset info
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+
+        if not dataset:
+            raise ValueError('Dataset is not exist.')
 
         document = DocumentService.get_document(dataset.id, document_id)
 
@@ -126,8 +299,85 @@ class DocumentApi(DatasetApiResource):
         except services.errors.document.DocumentIndexingError:
             raise DocumentIndexingError('Cannot delete document during indexing.')
 
-        return {'result': 'success'}, 204
+        return {'result': 'success'}, 200
 
 
-api.add_resource(DocumentListApi, '/documents')
-api.add_resource(DocumentApi, '/documents/<uuid:document_id>')
+class DocumentListApi(DatasetApiResource):
+    def get(self, tenant_id, dataset_id):
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        page = request.args.get('page', default=1, type=int)
+        limit = request.args.get('limit', default=20, type=int)
+        search = request.args.get('keyword', default=None, type=str)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        if not dataset:
+            raise NotFound('Dataset not found.')
+
+        query = Document.query.filter_by(
+            dataset_id=str(dataset_id), tenant_id=tenant_id)
+
+        if search:
+            search = f'%{search}%'
+            query = query.filter(Document.name.like(search))
+
+        query = query.order_by(desc(Document.created_at))
+
+        paginated_documents = query.paginate(
+            page=page, per_page=limit, max_per_page=100, error_out=False)
+        documents = paginated_documents.items
+
+        response = {
+            'data': marshal(documents, document_fields),
+            'has_more': len(documents) == limit,
+            'limit': limit,
+            'total': paginated_documents.total,
+            'page': page
+        }
+
+        return response
+
+
+class DocumentIndexingStatusApi(DatasetApiResource):
+    def get(self, tenant_id, dataset_id, batch):
+        dataset_id = str(dataset_id)
+        batch = str(batch)
+        tenant_id = str(tenant_id)
+        # get dataset
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        if not dataset:
+            raise NotFound('Dataset not found.')
+        # get documents
+        documents = DocumentService.get_batch_documents(dataset_id, batch)
+        if not documents:
+            raise NotFound('Documents not found.')
+        documents_status = []
+        for document in documents:
+            completed_segments = DocumentSegment.query.filter(DocumentSegment.completed_at.isnot(None),
+                                                              DocumentSegment.document_id == str(document.id),
+                                                              DocumentSegment.status != 're_segment').count()
+            total_segments = DocumentSegment.query.filter(DocumentSegment.document_id == str(document.id),
+                                                          DocumentSegment.status != 're_segment').count()
+            document.completed_segments = completed_segments
+            document.total_segments = total_segments
+            if document.is_paused:
+                document.indexing_status = 'paused'
+            documents_status.append(marshal(document, document_status_fields))
+        data = {
+            'data': documents_status
+        }
+        return data
+
+
+api.add_resource(DocumentAddByTextApi, '/datasets/<uuid:dataset_id>/document/create_by_text')
+api.add_resource(DocumentAddByFileApi, '/datasets/<uuid:dataset_id>/document/create_by_file')
+api.add_resource(DocumentUpdateByTextApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update_by_text')
+api.add_resource(DocumentUpdateByFileApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/update_by_file')
+api.add_resource(DocumentDeleteApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>')
+api.add_resource(DocumentListApi, '/datasets/<uuid:dataset_id>/documents')
+api.add_resource(DocumentIndexingStatusApi, '/datasets/<uuid:dataset_id>/documents/<string:batch>/indexing-status')

api/controllers/service_api/dataset/error.py

@@ -1,20 +1,73 @@
-# -*- coding:utf-8 -*-
 from libs.exception import BaseHTTPException
 
 
+class NoFileUploadedError(BaseHTTPException):
+    error_code = 'no_file_uploaded'
+    description = "Please upload your file."
+    code = 400
+
+
+class TooManyFilesError(BaseHTTPException):
+    error_code = 'too_many_files'
+    description = "Only one file is allowed."
+    code = 400
+
+
+class FileTooLargeError(BaseHTTPException):
+    error_code = 'file_too_large'
+    description = "File size exceeded. {message}"
+    code = 413
+
+
+class UnsupportedFileTypeError(BaseHTTPException):
+    error_code = 'unsupported_file_type'
+    description = "File type not allowed."
+    code = 415
+
+
+class HighQualityDatasetOnlyError(BaseHTTPException):
+    error_code = 'high_quality_dataset_only'
+    description = "Current operation only supports 'high-quality' datasets."
+    code = 400
+
+
+class DatasetNotInitializedError(BaseHTTPException):
+    error_code = 'dataset_not_initialized'
+    description = "The dataset is still being initialized or indexing. Please wait a moment."
+    code = 400
+
+
 class ArchivedDocumentImmutableError(BaseHTTPException):
     error_code = 'archived_document_immutable'
-    description = "Cannot operate when document was archived."
+    description = "The archived document is not editable."
     code = 403
 
 
+class DatasetNameDuplicateError(BaseHTTPException):
+    error_code = 'dataset_name_duplicate'
+    description = "The dataset name already exists. Please modify your dataset name."
+    code = 409
+
+
+class InvalidActionError(BaseHTTPException):
+    error_code = 'invalid_action'
+    description = "Invalid action."
+    code = 400
+
+
+class DocumentAlreadyFinishedError(BaseHTTPException):
+    error_code = 'document_already_finished'
+    description = "The document has been processed. Please refresh the page or go to the document details."
+    code = 400
+
+
 class DocumentIndexingError(BaseHTTPException):
     error_code = 'document_indexing'
-    description = "Cannot operate document during indexing."
-    code = 403
+    description = "The document is being processed and cannot be edited."
+    code = 400
 
 
-class DatasetNotInitedError(BaseHTTPException):
-    error_code = 'dataset_not_inited'
-    description = "The dataset is still being initialized or indexing. Please wait a moment."
-    code = 403
+class InvalidMetadataError(BaseHTTPException):
+    error_code = 'invalid_metadata'
+    description = "The metadata content is incorrect. Please check and verify."
+    code = 400

api/controllers/service_api/dataset/segment.py

@@ -0,0 +1,201 @@
+from flask_login import current_user
+from flask_restful import reqparse, marshal
+from werkzeug.exceptions import NotFound
+from controllers.service_api import api
+from controllers.service_api.app.error import ProviderNotInitializeError
+from controllers.service_api.wraps import DatasetApiResource
+from core.model_providers.error import ProviderTokenNotInitError, LLMBadRequestError
+from core.model_providers.model_factory import ModelFactory
+from extensions.ext_database import db
+from fields.segment_fields import segment_fields
+from models.dataset import Dataset, DocumentSegment
+from services.dataset_service import DatasetService, DocumentService, SegmentService
+
+
+class SegmentApi(DatasetApiResource):
+    """Resource for segments."""
+    def post(self, tenant_id, dataset_id, document_id):
+        """Create single segment."""
+        # check dataset
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        if not dataset:
+            raise NotFound('Dataset not found.')
+        # check document
+        document_id = str(document_id)
+        document = DocumentService.get_document(dataset.id, document_id)
+        if not document:
+            raise NotFound('Document not found.')
+        # check embedding model setting
+        if dataset.indexing_technique == 'high_quality':
+            try:
+                ModelFactory.get_embedding_model(
+                    tenant_id=current_user.current_tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
+        # validate args
+        parser = reqparse.RequestParser()
+        parser.add_argument('segments', type=list, required=False, nullable=True, location='json')
+        args = parser.parse_args()
+        for args_item in args['segments']:
+            SegmentService.segment_create_args_validate(args_item, document)
+        segments = SegmentService.multi_create_segment(args['segments'], document, dataset)
+        return {
+            'data': marshal(segments, segment_fields),
+            'doc_form': document.doc_form
+        }, 200
+
+    def get(self, tenant_id, dataset_id, document_id):
+        """Create single segment."""
+        # check dataset
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        if not dataset:
+            raise NotFound('Dataset not found.')
+        # check document
+        document_id = str(document_id)
+        document = DocumentService.get_document(dataset.id, document_id)
+        if not document:
+            raise NotFound('Document not found.')
+        # check embedding model setting
+        if dataset.indexing_technique == 'high_quality':
+            try:
+                ModelFactory.get_embedding_model(
+                    tenant_id=current_user.current_tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('status', type=str,
+                            action='append', default=[], location='args')
+        parser.add_argument('keyword', type=str, default=None, location='args')
+        args = parser.parse_args()
+
+        status_list = args['status']
+        keyword = args['keyword']
+
+        query = DocumentSegment.query.filter(
+            DocumentSegment.document_id == str(document_id),
+            DocumentSegment.tenant_id == current_user.current_tenant_id
+        )
+
+        if status_list:
+            query = query.filter(DocumentSegment.status.in_(status_list))
+
+        if keyword:
+            query = query.where(DocumentSegment.content.ilike(f'%{keyword}%'))
+
+        total = query.count()
+        segments = query.order_by(DocumentSegment.position).all()
+        return {
+            'data': marshal(segments, segment_fields),
+            'doc_form': document.doc_form,
+            'total': total
+        }, 200
+
+
+class DatasetSegmentApi(DatasetApiResource):
+    def delete(self, tenant_id, dataset_id, document_id, segment_id):
+        # check dataset
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        if not dataset:
+            raise NotFound('Dataset not found.')
+        # check user's model setting
+        DatasetService.check_dataset_model_setting(dataset)
+        # check document
+        document_id = str(document_id)
+        document = DocumentService.get_document(dataset_id, document_id)
+        if not document:
+            raise NotFound('Document not found.')
+        # check segment
+        segment = DocumentSegment.query.filter(
+            DocumentSegment.id == str(segment_id),
+            DocumentSegment.tenant_id == current_user.current_tenant_id
+        ).first()
+        if not segment:
+            raise NotFound('Segment not found.')
+        SegmentService.delete_segment(segment, document, dataset)
+        return {'result': 'success'}, 200
+
+    def post(self, tenant_id, dataset_id, document_id, segment_id):
+        # check dataset
+        dataset_id = str(dataset_id)
+        tenant_id = str(tenant_id)
+        dataset = db.session.query(Dataset).filter(
+            Dataset.tenant_id == tenant_id,
+            Dataset.id == dataset_id
+        ).first()
+        if not dataset:
+            raise NotFound('Dataset not found.')
+        # check user's model setting
+        DatasetService.check_dataset_model_setting(dataset)
+        # check document
+        document_id = str(document_id)
+        document = DocumentService.get_document(dataset_id, document_id)
+        if not document:
+            raise NotFound('Document not found.')
+        if dataset.indexing_technique == 'high_quality':
+            # check embedding model setting
+            try:
+                ModelFactory.get_embedding_model(
+                    tenant_id=current_user.current_tenant_id,
+                    model_provider_name=dataset.embedding_model_provider,
+                    model_name=dataset.embedding_model
+                )
+            except LLMBadRequestError:
+                raise ProviderNotInitializeError(
+                    f"No Embedding Model available. Please configure a valid provider "
+                    f"in the Settings -> Model Provider.")
+            except ProviderTokenNotInitError as ex:
+                raise ProviderNotInitializeError(ex.description)
+            # check segment
+        segment_id = str(segment_id)
+        segment = DocumentSegment.query.filter(
+            DocumentSegment.id == str(segment_id),
+            DocumentSegment.tenant_id == current_user.current_tenant_id
+        ).first()
+        if not segment:
+            raise NotFound('Segment not found.')
+
+        # validate args
+        parser = reqparse.RequestParser()
+        parser.add_argument('segments', type=dict, required=False, nullable=True, location='json')
+        args = parser.parse_args()
+
+        SegmentService.segment_create_args_validate(args['segments'], document)
+        segment = SegmentService.update_segment(args['segments'], segment, document, dataset)
+        return {
+            'data': marshal(segment, segment_fields),
+            'doc_form': document.doc_form
+        }, 200
+
+
+api.add_resource(SegmentApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments')
+api.add_resource(DatasetSegmentApi, '/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/<uuid:segment_id>')

api/controllers/service_api/wraps.py

@@ -2,12 +2,14 @@
 from datetime import datetime
 from functools import wraps
 
-from flask import request
+from flask import request, current_app
+from flask_login import user_logged_in
 from flask_restful import Resource
 from werkzeug.exceptions import NotFound, Unauthorized
 
+from libs.login import _get_user
 from extensions.ext_database import db
-from models.dataset import Dataset
+from models.account import Tenant, TenantAccountJoin, Account
 from models.model import ApiToken, App
 
 
@@ -43,12 +45,24 @@ def validate_dataset_token(view=None):
         @wraps(view)
         def decorated(*args, **kwargs):
             api_token = validate_and_get_api_token('dataset')
-
-            dataset = db.session.query(Dataset).filter(Dataset.id == api_token.dataset_id).first()
-            if not dataset:
-                raise NotFound()
-
-            return view(dataset, *args, **kwargs)
+            tenant_account_join = db.session.query(Tenant, TenantAccountJoin) \
+                .filter(Tenant.id == api_token.tenant_id) \
+                .filter(TenantAccountJoin.tenant_id == Tenant.id) \
+                .filter(TenantAccountJoin.role == 'owner') \
+                .one_or_none()
+            if tenant_account_join:
+                tenant, ta = tenant_account_join
+                account = Account.query.filter_by(id=ta.account_id).first()
+                # Login admin
+                if account:
+                    account.current_tenant = tenant
+                    current_app.login_manager._update_request_context_with_user(account)
+                    user_logged_in.send(current_app._get_current_object(), user=_get_user())
+                else:
+                    raise Unauthorized("Tenant owner account is not exist.")
+            else:
+                raise Unauthorized("Tenant is not exist.")
+            return view(api_token.tenant_id, *args, **kwargs)
         return decorated
 
     if view:

api/controllers/web/conversation.py

@@ -6,26 +6,12 @@ from werkzeug.exceptions import NotFound
 from controllers.web import api
 from controllers.web.error import NotChatAppError
 from controllers.web.wraps import WebApiResource
+from fields.conversation_fields import conversation_infinite_scroll_pagination_fields, simple_conversation_fields
 from libs.helper import TimestampField, uuid_value
 from services.conversation_service import ConversationService
 from services.errors.conversation import LastConversationNotExistsError, ConversationNotExistsError
 from services.web_conversation_service import WebConversationService
 
-conversation_fields = {
-    'id': fields.String,
-    'name': fields.String,
-    'inputs': fields.Raw,
-    'status': fields.String,
-    'introduction': fields.String,
-    'created_at': TimestampField
-}
-
-conversation_infinite_scroll_pagination_fields = {
-    'limit': fields.Integer,
-    'has_more': fields.Boolean,
-    'data': fields.List(fields.Nested(conversation_fields))
-}
-
 
 class ConversationListApi(WebApiResource):
 
@@ -73,7 +59,7 @@ class ConversationApi(WebApiResource):
 
 class ConversationRenameApi(WebApiResource):
 
-    @marshal_with(conversation_fields)
+    @marshal_with(simple_conversation_fields)
     def post(self, app_model, end_user, c_id):
         if app_model.mode != 'chat':
             raise NotChatAppError()

api/core/completion.py

@@ -108,12 +108,14 @@ class Completion:
                 retriever_from=retriever_from
             )
 
+            query_for_agent = cls.get_query_for_agent(app, app_model_config, query, inputs)
+
             # run agent executor
             agent_execute_result = None
-            if agent_executor:
-                should_use_agent = agent_executor.should_use_agent(query)
+            if query_for_agent and agent_executor:
+                should_use_agent = agent_executor.should_use_agent(query_for_agent)
                 if should_use_agent:
-                    agent_execute_result = agent_executor.run(query)
+                    agent_execute_result = agent_executor.run(query_for_agent)
 
             # When no extra pre prompt is specified,
             # the output of the agent can be used directly as the main output content without calling LLM again
@@ -142,6 +144,13 @@ class Completion:
             logging.warning(f'ChunkedEncodingError: {e}')
             conversation_message_task.end()
             return
+        
+    @classmethod
+    def get_query_for_agent(cls, app: App, app_model_config: AppModelConfig, query: str, inputs: dict) -> str:
+        if app.mode != 'completion':
+            return query
+        
+        return inputs.get(app_model_config.dataset_query_variable, "")
 
     @classmethod
     def run_final_llm(cls, model_instance: BaseLLM, mode: str, app_model_config: AppModelConfig, query: str,

api/core/conversation_message_task.py

@@ -94,7 +94,7 @@ class ConversationMessageTask:
         if not self.conversation:
             self.is_new_conversation = True
             self.conversation = Conversation(
-                app_id=self.app_model_config.app_id,
+                app_id=self.app.id,
                 app_model_config_id=self.app_model_config.id,
                 model_provider=self.provider_name,
                 model_id=self.model_name,
@@ -112,10 +112,10 @@ class ConversationMessageTask:
             )
 
             db.session.add(self.conversation)
-            db.session.flush()
+            db.session.commit()
 
         self.message = Message(
-            app_id=self.app_model_config.app_id,
+            app_id=self.app.id,
             model_provider=self.provider_name,
             model_id=self.model_name,
             override_model_configs=json.dumps(override_model_configs) if override_model_configs else None,
@@ -140,7 +140,7 @@ class ConversationMessageTask:
         )
 
         db.session.add(self.message)
-        db.session.flush()
+        db.session.commit()
 
     def append_message_text(self, text: str):
         if text is not None:
@@ -191,12 +191,13 @@ class ConversationMessageTask:
         )
 
         db.session.add(message_chain)
-        db.session.flush()
+        db.session.commit()
 
         return message_chain
 
     def on_chain_end(self, message_chain: MessageChain, chain_result: ChainResult):
         message_chain.output = json.dumps(chain_result.completion)
+        db.session.commit()
 
         self._pub_handler.pub_chain(message_chain)
 
@@ -217,7 +218,7 @@ class ConversationMessageTask:
         )
 
         db.session.add(message_agent_thought)
-        db.session.flush()
+        db.session.commit()
 
         self._pub_handler.pub_agent_thought(message_agent_thought)
 
@@ -249,7 +250,7 @@ class ConversationMessageTask:
         message_agent_thought.tokens = agent_loop.prompt_tokens + agent_loop.completion_tokens
         message_agent_thought.total_price = loop_total_price
         message_agent_thought.currency = agent_model_instance.get_currency()
-        db.session.flush()
+        db.session.commit()
 
     def on_dataset_query_end(self, dataset_query_obj: DatasetQueryObj):
         dataset_query = DatasetQuery(
@@ -262,6 +263,7 @@ class ConversationMessageTask:
         )
 
         db.session.add(dataset_query)
+        db.session.commit()
 
     def on_dataset_query_finish(self, resource: List):
         if resource and len(resource) > 0:
@@ -285,7 +287,7 @@ class ConversationMessageTask:
                     created_by=self.user.id
                 )
                 db.session.add(dataset_retriever_resource)
-                db.session.flush()
+                db.session.commit()
             self.retriever_resource = resource
 
     def message_end(self):

api/core/data_loader/loader/notion.py

@@ -16,6 +16,7 @@ logger = logging.getLogger(__name__)
 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']

api/core/index/keyword_table_index/keyword_table_index.py

@@ -246,11 +246,28 @@ class KeywordTableIndex(BaseIndex):
         keyword_table = self._add_text_to_keyword_table(keyword_table, node_id, keywords)
         self._save_dataset_keyword_table(keyword_table)
 
+    def multi_create_segment_keywords(self, pre_segment_data_list: list):
+        keyword_table_handler = JiebaKeywordTableHandler()
+        keyword_table = self._get_dataset_keyword_table()
+        for pre_segment_data in pre_segment_data_list:
+            segment = pre_segment_data['segment']
+            if pre_segment_data['keywords']:
+                segment.keywords = pre_segment_data['keywords']
+                keyword_table = self._add_text_to_keyword_table(keyword_table, segment.index_node_id,
+                                                                pre_segment_data['keywords'])
+            else:
+                keywords = keyword_table_handler.extract_keywords(segment.content,
+                                                                  self._config.max_keywords_per_chunk)
+                segment.keywords = list(keywords)
+                keyword_table = self._add_text_to_keyword_table(keyword_table, segment.index_node_id, list(keywords))
+        self._save_dataset_keyword_table(keyword_table)
+
     def update_segment_keywords_index(self, node_id: str, keywords: List[str]):
         keyword_table = self._get_dataset_keyword_table()
         keyword_table = self._add_text_to_keyword_table(keyword_table, node_id, keywords)
         self._save_dataset_keyword_table(keyword_table)
 
+
 class KeywordTableRetriever(BaseRetriever, BaseModel):
     index: KeywordTableIndex
     search_kwargs: dict = Field(default_factory=dict)

api/core/index/vector_index/base.py

@@ -113,8 +113,10 @@ class BaseVectorIndex(BaseIndex):
     def delete_by_group_id(self, group_id: str) -> None:
         vector_store = self._get_vector_store()
         vector_store = cast(self._get_vector_store_class(), vector_store)
-
-        vector_store.delete()
+        if self.dataset.collection_binding_id:
+            vector_store.delete_by_group_id(group_id)
+        else:
+            vector_store.delete()
 
     def delete(self) -> None:
         vector_store = self._get_vector_store()
@@ -283,7 +285,7 @@ class BaseVectorIndex(BaseIndex):
 
         if documents:
             try:
-                self.create_with_collection_name(documents, dataset_collection_binding.collection_name)
+                self.add_texts(documents)
             except Exception as e:
                 raise e
 

api/core/index/vector_index/milvus.py

@@ -0,0 +1,860 @@
+"""Wrapper around the Milvus vector database."""
+from __future__ import annotations
+
+import logging
+from typing import Any, Iterable, List, Optional, Tuple, Union, Sequence
+from uuid import uuid4
+
+import numpy as np
+
+from langchain.docstore.document import Document
+from langchain.embeddings.base import Embeddings
+from langchain.vectorstores.base import VectorStore
+from langchain.vectorstores.utils import maximal_marginal_relevance
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MILVUS_CONNECTION = {
+    "host": "localhost",
+    "port": "19530",
+    "user": "",
+    "password": "",
+    "secure": False,
+}
+
+
+class Milvus(VectorStore):
+    """Initialize wrapper around the milvus vector database.
+
+    In order to use this you need to have `pymilvus` installed and a
+    running Milvus
+
+    See the following documentation for how to run a Milvus instance:
+    https://milvus.io/docs/install_standalone-docker.md
+
+    If looking for a hosted Milvus, take a look at this documentation:
+    https://zilliz.com/cloud and make use of the Zilliz vectorstore found in
+    this project,
+
+    IF USING L2/IP metric IT IS HIGHLY SUGGESTED TO NORMALIZE YOUR DATA.
+
+    Args:
+        embedding_function (Embeddings): Function used to embed the text.
+        collection_name (str): Which Milvus collection to use. Defaults to
+            "LangChainCollection".
+        connection_args (Optional[dict[str, any]]): The connection args used for
+            this class comes in the form of a dict.
+        consistency_level (str): The consistency level to use for a collection.
+            Defaults to "Session".
+        index_params (Optional[dict]): Which index params to use. Defaults to
+            HNSW/AUTOINDEX depending on service.
+        search_params (Optional[dict]): Which search params to use. Defaults to
+            default of index.
+        drop_old (Optional[bool]): Whether to drop the current collection. Defaults
+            to False.
+
+    The connection args used for this class comes in the form of a dict,
+    here are a few of the options:
+        address (str): The actual address of Milvus
+            instance. Example address: "localhost:19530"
+        uri (str): The uri of Milvus instance. Example uri:
+            "http://randomwebsite:19530",
+            "tcp:foobarsite:19530",
+            "https://ok.s3.south.com:19530".
+        host (str): The host of Milvus instance. Default at "localhost",
+            PyMilvus will fill in the default host if only port is provided.
+        port (str/int): The port of Milvus instance. Default at 19530, PyMilvus
+            will fill in the default port if only host is provided.
+        user (str): Use which user to connect to Milvus instance. If user and
+            password are provided, we will add related header in every RPC call.
+        password (str): Required when user is provided. The password
+            corresponding to the user.
+        secure (bool): Default is false. If set to true, tls will be enabled.
+        client_key_path (str): If use tls two-way authentication, need to
+            write the client.key path.
+        client_pem_path (str): If use tls two-way authentication, need to
+            write the client.pem path.
+        ca_pem_path (str): If use tls two-way authentication, need to write
+            the ca.pem path.
+        server_pem_path (str): If use tls one-way authentication, need to
+            write the server.pem path.
+        server_name (str): If use tls, need to write the common name.
+
+    Example:
+        .. code-block:: python
+
+        from langchain import Milvus
+        from langchain.embeddings import OpenAIEmbeddings
+
+        embedding = OpenAIEmbeddings()
+        # Connect to a milvus instance on localhost
+        milvus_store = Milvus(
+            embedding_function = Embeddings,
+            collection_name = "LangChainCollection",
+            drop_old = True,
+        )
+
+    Raises:
+        ValueError: If the pymilvus python package is not installed.
+    """
+
+    def __init__(
+        self,
+        embedding_function: Embeddings,
+        collection_name: str = "LangChainCollection",
+        connection_args: Optional[dict[str, Any]] = None,
+        consistency_level: str = "Session",
+        index_params: Optional[dict] = None,
+        search_params: Optional[dict] = None,
+        drop_old: Optional[bool] = False,
+    ):
+        """Initialize the Milvus vector store."""
+        try:
+            from pymilvus import Collection, utility
+        except ImportError:
+            raise ValueError(
+                "Could not import pymilvus python package. "
+                "Please install it with `pip install pymilvus`."
+            )
+
+        # Default search params when one is not provided.
+        self.default_search_params = {
+            "IVF_FLAT": {"metric_type": "L2", "params": {"nprobe": 10}},
+            "IVF_SQ8": {"metric_type": "L2", "params": {"nprobe": 10}},
+            "IVF_PQ": {"metric_type": "L2", "params": {"nprobe": 10}},
+            "HNSW": {"metric_type": "L2", "params": {"ef": 10}},
+            "RHNSW_FLAT": {"metric_type": "L2", "params": {"ef": 10}},
+            "RHNSW_SQ": {"metric_type": "L2", "params": {"ef": 10}},
+            "RHNSW_PQ": {"metric_type": "L2", "params": {"ef": 10}},
+            "IVF_HNSW": {"metric_type": "L2", "params": {"nprobe": 10, "ef": 10}},
+            "ANNOY": {"metric_type": "L2", "params": {"search_k": 10}},
+            "AUTOINDEX": {"metric_type": "L2", "params": {}},
+        }
+
+        self.embedding_func = embedding_function
+        self.collection_name = collection_name
+        self.index_params = index_params
+        self.search_params = search_params
+        self.consistency_level = consistency_level
+
+        # In order for a collection to be compatible, pk needs to be auto'id and int
+        self._primary_field = "id"
+        # In order for compatibility, the text field will need to be called "text"
+        self._text_field = "page_content"
+        # In order for compatibility, the vector field needs to be called "vector"
+        self._vector_field = "vectors"
+        # In order for compatibility, the metadata field will need to be called "metadata"
+        self._metadata_field = "metadata"
+        self.fields: list[str] = []
+        # Create the connection to the server
+        if connection_args is None:
+            connection_args = DEFAULT_MILVUS_CONNECTION
+        self.alias = self._create_connection_alias(connection_args)
+        self.col: Optional[Collection] = None
+
+        # Grab the existing collection if it exists
+        if utility.has_collection(self.collection_name, using=self.alias):
+            self.col = Collection(
+                self.collection_name,
+                using=self.alias,
+            )
+        # If need to drop old, drop it
+        if drop_old and isinstance(self.col, Collection):
+            self.col.drop()
+            self.col = None
+
+        # Initialize the vector store
+        self._init()
+
+    @property
+
+
+    def embeddings(self) -> Embeddings:
+        return self.embedding_func
+
+    def _create_connection_alias(self, connection_args: dict) -> str:
+        """Create the connection to the Milvus server."""
+        from pymilvus import MilvusException, connections
+
+        # Grab the connection arguments that are used for checking existing connection
+        host: str = connection_args.get("host", None)
+        port: Union[str, int] = connection_args.get("port", None)
+        address: str = connection_args.get("address", None)
+        uri: str = connection_args.get("uri", None)
+        user = connection_args.get("user", None)
+
+        # Order of use is host/port, uri, address
+        if host is not None and port is not None:
+            given_address = str(host) + ":" + str(port)
+        elif uri is not None:
+            given_address = uri.split("https://")[1]
+        elif address is not None:
+            given_address = address
+        else:
+            given_address = None
+            logger.debug("Missing standard address type for reuse atttempt")
+
+        # User defaults to empty string when getting connection info
+        if user is not None:
+            tmp_user = user
+        else:
+            tmp_user = ""
+
+        # If a valid address was given, then check if a connection exists
+        if given_address is not None:
+            for con in connections.list_connections():
+                addr = connections.get_connection_addr(con[0])
+                if (
+                    con[1]
+                    and ("address" in addr)
+                    and (addr["address"] == given_address)
+                    and ("user" in addr)
+                    and (addr["user"] == tmp_user)
+                ):
+                    logger.debug("Using previous connection: %s", con[0])
+                    return con[0]
+
+        # Generate a new connection if one doesn't exist
+        alias = uuid4().hex
+        try:
+            connections.connect(alias=alias, **connection_args)
+            logger.debug("Created new connection using: %s", alias)
+            return alias
+        except MilvusException as e:
+            logger.error("Failed to create new connection using: %s", alias)
+            raise e
+
+    def _init(
+        self, embeddings: Optional[list] = None, metadatas: Optional[list[dict]] = None
+    ) -> None:
+        if embeddings is not None:
+            self._create_collection(embeddings, metadatas)
+        self._extract_fields()
+        self._create_index()
+        self._create_search_params()
+        self._load()
+
+    def _create_collection(
+        self, embeddings: list, metadatas: Optional[list[dict]] = None
+    ) -> None:
+        from pymilvus import (
+            Collection,
+            CollectionSchema,
+            DataType,
+            FieldSchema,
+            MilvusException,
+        )
+        from pymilvus.orm.types import infer_dtype_bydata
+
+        # Determine embedding dim
+        dim = len(embeddings[0])
+        fields = []
+        # Determine metadata schema
+        # if metadatas:
+        #     # Create FieldSchema for each entry in metadata.
+        #     for key, value in metadatas[0].items():
+        #         # Infer the corresponding datatype of the metadata
+        #         dtype = infer_dtype_bydata(value)
+        #         # Datatype isn't compatible
+        #         if dtype == DataType.UNKNOWN or dtype == DataType.NONE:
+        #             logger.error(
+        #                 "Failure to create collection, unrecognized dtype for key: %s",
+        #                 key,
+        #             )
+        #             raise ValueError(f"Unrecognized datatype for {key}.")
+        #         # Dataype is a string/varchar equivalent
+        #         elif dtype == DataType.VARCHAR:
+        #             fields.append(FieldSchema(key, DataType.VARCHAR, max_length=65_535))
+        #         else:
+        #             fields.append(FieldSchema(key, dtype))
+        if metadatas:
+            fields.append(FieldSchema(self._metadata_field, DataType.JSON, max_length=65_535))
+
+        # Create the text field
+        fields.append(
+            FieldSchema(self._text_field, DataType.VARCHAR, max_length=65_535)
+        )
+        # Create the primary key field
+        fields.append(
+            FieldSchema(
+                self._primary_field, DataType.INT64, is_primary=True, auto_id=True
+            )
+        )
+        # Create the vector field, supports binary or float vectors
+        fields.append(
+            FieldSchema(self._vector_field, infer_dtype_bydata(embeddings[0]), dim=dim)
+        )
+
+        # Create the schema for the collection
+        schema = CollectionSchema(fields)
+
+        # Create the collection
+        try:
+            self.col = Collection(
+                name=self.collection_name,
+                schema=schema,
+                consistency_level=self.consistency_level,
+                using=self.alias,
+            )
+        except MilvusException as e:
+            logger.error(
+                "Failed to create collection: %s error: %s", self.collection_name, e
+            )
+            raise e
+
+    def _extract_fields(self) -> None:
+        """Grab the existing fields from the Collection"""
+        from pymilvus import Collection
+
+        if isinstance(self.col, Collection):
+            schema = self.col.schema
+            for x in schema.fields:
+                self.fields.append(x.name)
+            # Since primary field is auto-id, no need to track it
+            self.fields.remove(self._primary_field)
+
+    def _get_index(self) -> Optional[dict[str, Any]]:
+        """Return the vector index information if it exists"""
+        from pymilvus import Collection
+
+        if isinstance(self.col, Collection):
+            for x in self.col.indexes:
+                if x.field_name == self._vector_field:
+                    return x.to_dict()
+        return None
+
+    def _create_index(self) -> None:
+        """Create a index on the collection"""
+        from pymilvus import Collection, MilvusException
+
+        if isinstance(self.col, Collection) and self._get_index() is None:
+            try:
+                # If no index params, use a default HNSW based one
+                if self.index_params is None:
+                    self.index_params = {
+                        "metric_type": "IP",
+                        "index_type": "HNSW",
+                        "params": {"M": 8, "efConstruction": 64},
+                    }
+
+                try:
+                    self.col.create_index(
+                        self._vector_field,
+                        index_params=self.index_params,
+                        using=self.alias,
+                    )
+
+                # If default did not work, most likely on Zilliz Cloud
+                except MilvusException:
+                    # Use AUTOINDEX based index
+                    self.index_params = {
+                        "metric_type": "L2",
+                        "index_type": "AUTOINDEX",
+                        "params": {},
+                    }
+                    self.col.create_index(
+                        self._vector_field,
+                        index_params=self.index_params,
+                        using=self.alias,
+                    )
+                logger.debug(
+                    "Successfully created an index on collection: %s",
+                    self.collection_name,
+                )
+
+            except MilvusException as e:
+                logger.error(
+                    "Failed to create an index on collection: %s", self.collection_name
+                )
+                raise e
+
+    def _create_search_params(self) -> None:
+        """Generate search params based on the current index type"""
+        from pymilvus import Collection
+
+        if isinstance(self.col, Collection) and self.search_params is None:
+            index = self._get_index()
+            if index is not None:
+                index_type: str = index["index_param"]["index_type"]
+                metric_type: str = index["index_param"]["metric_type"]
+                self.search_params = self.default_search_params[index_type]
+                self.search_params["metric_type"] = metric_type
+
+    def _load(self) -> None:
+        """Load the collection if available."""
+        from pymilvus import Collection
+
+        if isinstance(self.col, Collection) and self._get_index() is not None:
+            self.col.load()
+
+    def add_texts(
+        self,
+        texts: Iterable[str],
+        metadatas: Optional[List[dict]] = None,
+        timeout: Optional[int] = None,
+        batch_size: int = 1000,
+        **kwargs: Any,
+    ) -> List[str]:
+        """Insert text data into Milvus.
+
+        Inserting data when the collection has not be made yet will result
+        in creating a new Collection. The data of the first entity decides
+        the schema of the new collection, the dim is extracted from the first
+        embedding and the columns are decided by the first metadata dict.
+        Metada keys will need to be present for all inserted values. At
+        the moment there is no None equivalent in Milvus.
+
+        Args:
+            texts (Iterable[str]): The texts to embed, it is assumed
+                that they all fit in memory.
+            metadatas (Optional[List[dict]]): Metadata dicts attached to each of
+                the texts. Defaults to None.
+            timeout (Optional[int]): Timeout for each batch insert. Defaults
+                to None.
+            batch_size (int, optional): Batch size to use for insertion.
+                Defaults to 1000.
+
+        Raises:
+            MilvusException: Failure to add texts
+
+        Returns:
+            List[str]: The resulting keys for each inserted element.
+        """
+        from pymilvus import Collection, MilvusException
+
+        texts = list(texts)
+
+        try:
+            embeddings = self.embedding_func.embed_documents(texts)
+        except NotImplementedError:
+            embeddings = [self.embedding_func.embed_query(x) for x in texts]
+
+        if len(embeddings) == 0:
+            logger.debug("Nothing to insert, skipping.")
+            return []
+
+        # If the collection hasn't been initialized yet, perform all steps to do so
+        if not isinstance(self.col, Collection):
+            self._init(embeddings, metadatas)
+
+        # Dict to hold all insert columns
+        insert_dict: dict[str, list] = {
+            self._text_field: texts,
+            self._vector_field: embeddings,
+        }
+
+        # Collect the metadata into the insert dict.
+        # if metadatas is not None:
+        #     for d in metadatas:
+        #         for key, value in d.items():
+        #             if key in self.fields:
+        #                 insert_dict.setdefault(key, []).append(value)
+        if metadatas is not None:
+            for d in metadatas:
+                insert_dict.setdefault(self._metadata_field, []).append(d)
+
+        # Total insert count
+        vectors: list = insert_dict[self._vector_field]
+        total_count = len(vectors)
+
+        pks: list[str] = []
+
+        assert isinstance(self.col, Collection)
+        for i in range(0, total_count, batch_size):
+            # Grab end index
+            end = min(i + batch_size, total_count)
+            # Convert dict to list of lists batch for insertion
+            insert_list = [insert_dict[x][i:end] for x in self.fields]
+            # Insert into the collection.
+            try:
+                res: Collection
+                res = self.col.insert(insert_list, timeout=timeout, **kwargs)
+                pks.extend(res.primary_keys)
+            except MilvusException as e:
+                logger.error(
+                    "Failed to insert batch starting at entity: %s/%s", i, total_count
+                )
+                raise e
+        return pks
+
+    def similarity_search(
+        self,
+        query: str,
+        k: int = 4,
+        param: Optional[dict] = None,
+        expr: Optional[str] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
+    ) -> List[Document]:
+        """Perform a similarity search against the query string.
+
+        Args:
+            query (str): The text to search.
+            k (int, optional): How many results to return. Defaults to 4.
+            param (dict, optional): The search params for the index type.
+                Defaults to None.
+            expr (str, optional): Filtering expression. Defaults to None.
+            timeout (int, optional): How long to wait before timeout error.
+                Defaults to None.
+            kwargs: Collection.search() keyword arguments.
+
+        Returns:
+            List[Document]: Document results for search.
+        """
+        if self.col is None:
+            logger.debug("No existing collection to search.")
+            return []
+        res = self.similarity_search_with_score(
+            query=query, k=k, param=param, expr=expr, timeout=timeout, **kwargs
+        )
+        return [doc for doc, _ in res]
+
+    def similarity_search_by_vector(
+        self,
+        embedding: List[float],
+        k: int = 4,
+        param: Optional[dict] = None,
+        expr: Optional[str] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
+    ) -> List[Document]:
+        """Perform a similarity search against the query string.
+
+        Args:
+            embedding (List[float]): The embedding vector to search.
+            k (int, optional): How many results to return. Defaults to 4.
+            param (dict, optional): The search params for the index type.
+                Defaults to None.
+            expr (str, optional): Filtering expression. Defaults to None.
+            timeout (int, optional): How long to wait before timeout error.
+                Defaults to None.
+            kwargs: Collection.search() keyword arguments.
+
+        Returns:
+            List[Document]: Document results for search.
+        """
+        if self.col is None:
+            logger.debug("No existing collection to search.")
+            return []
+        res = self.similarity_search_with_score_by_vector(
+            embedding=embedding, k=k, param=param, expr=expr, timeout=timeout, **kwargs
+        )
+        return [doc for doc, _ in res]
+
+    def similarity_search_with_score(
+        self,
+        query: str,
+        k: int = 4,
+        param: Optional[dict] = None,
+        expr: Optional[str] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
+    ) -> List[Tuple[Document, float]]:
+        """Perform a search on a query string and return results with score.
+
+        For more information about the search parameters, take a look at the pymilvus
+        documentation found here:
+        https://milvus.io/api-reference/pymilvus/v2.2.6/Collection/search().md
+
+        Args:
+            query (str): The text being searched.
+            k (int, optional): The amount of results to return. Defaults to 4.
+            param (dict): The search params for the specified index.
+                Defaults to None.
+            expr (str, optional): Filtering expression. Defaults to None.
+            timeout (int, optional): How long to wait before timeout error.
+                Defaults to None.
+            kwargs: Collection.search() keyword arguments.
+
+        Returns:
+            List[float], List[Tuple[Document, any, any]]:
+        """
+        if self.col is None:
+            logger.debug("No existing collection to search.")
+            return []
+
+        # Embed the query text.
+        embedding = self.embedding_func.embed_query(query)
+
+        res = self.similarity_search_with_score_by_vector(
+            embedding=embedding, k=k, param=param, expr=expr, timeout=timeout, **kwargs
+        )
+        return res
+
+    def _similarity_search_with_relevance_scores(
+        self,
+        query: str,
+        k: int = 4,
+        **kwargs: Any,
+    ) -> List[Tuple[Document, float]]:
+        """Return docs and relevance scores in the range [0, 1].
+
+        0 is dissimilar, 1 is most similar.
+
+        Args:
+            query: input text
+            k: Number of Documents to return. Defaults to 4.
+            **kwargs: kwargs to be passed to similarity search. Should include:
+                score_threshold: Optional, a floating point value between 0 to 1 to
+                    filter the resulting set of retrieved docs
+
+        Returns:
+            List of Tuples of (doc, similarity_score)
+        """
+        return self.similarity_search_with_score(query, k, **kwargs)
+
+    def similarity_search_with_score_by_vector(
+        self,
+        embedding: List[float],
+        k: int = 4,
+        param: Optional[dict] = None,
+        expr: Optional[str] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
+    ) -> List[Tuple[Document, float]]:
+        """Perform a search on a query string and return results with score.
+
+        For more information about the search parameters, take a look at the pymilvus
+        documentation found here:
+        https://milvus.io/api-reference/pymilvus/v2.2.6/Collection/search().md
+
+        Args:
+            embedding (List[float]): The embedding vector being searched.
+            k (int, optional): The amount of results to return. Defaults to 4.
+            param (dict): The search params for the specified index.
+                Defaults to None.
+            expr (str, optional): Filtering expression. Defaults to None.
+            timeout (int, optional): How long to wait before timeout error.
+                Defaults to None.
+            kwargs: Collection.search() keyword arguments.
+
+        Returns:
+            List[Tuple[Document, float]]: Result doc and score.
+        """
+        if self.col is None:
+            logger.debug("No existing collection to search.")
+            return []
+
+        if param is None:
+            param = self.search_params
+
+        # Determine result metadata fields.
+        output_fields = self.fields[:]
+        output_fields.remove(self._vector_field)
+
+        # Perform the search.
+        res = self.col.search(
+            data=[embedding],
+            anns_field=self._vector_field,
+            param=param,
+            limit=k,
+            expr=expr,
+            output_fields=output_fields,
+            timeout=timeout,
+            **kwargs,
+        )
+        # Organize results.
+        ret = []
+        for result in res[0]:
+            meta = {x: result.entity.get(x) for x in output_fields}
+            doc = Document(page_content=meta.pop(self._text_field), metadata=meta.get('metadata'))
+            pair = (doc, result.score)
+            ret.append(pair)
+
+        return ret
+
+    def max_marginal_relevance_search(
+        self,
+        query: str,
+        k: int = 4,
+        fetch_k: int = 20,
+        lambda_mult: float = 0.5,
+        param: Optional[dict] = None,
+        expr: Optional[str] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
+    ) -> List[Document]:
+        """Perform a search and return results that are reordered by MMR.
+
+        Args:
+            query (str): The text being searched.
+            k (int, optional): How many results to give. Defaults to 4.
+            fetch_k (int, optional): Total results to select k from.
+                Defaults to 20.
+            lambda_mult: Number between 0 and 1 that determines the degree
+                        of diversity among the results with 0 corresponding
+                        to maximum diversity and 1 to minimum diversity.
+                        Defaults to 0.5
+            param (dict, optional): The search params for the specified index.
+                Defaults to None.
+            expr (str, optional): Filtering expression. Defaults to None.
+            timeout (int, optional): How long to wait before timeout error.
+                Defaults to None.
+            kwargs: Collection.search() keyword arguments.
+
+
+        Returns:
+            List[Document]: Document results for search.
+        """
+        if self.col is None:
+            logger.debug("No existing collection to search.")
+            return []
+
+        embedding = self.embedding_func.embed_query(query)
+
+        return self.max_marginal_relevance_search_by_vector(
+            embedding=embedding,
+            k=k,
+            fetch_k=fetch_k,
+            lambda_mult=lambda_mult,
+            param=param,
+            expr=expr,
+            timeout=timeout,
+            **kwargs,
+        )
+
+    def max_marginal_relevance_search_by_vector(
+        self,
+        embedding: list[float],
+        k: int = 4,
+        fetch_k: int = 20,
+        lambda_mult: float = 0.5,
+        param: Optional[dict] = None,
+        expr: Optional[str] = None,
+        timeout: Optional[int] = None,
+        **kwargs: Any,
+    ) -> List[Document]:
+        """Perform a search and return results that are reordered by MMR.
+
+        Args:
+            embedding (str): The embedding vector being searched.
+            k (int, optional): How many results to give. Defaults to 4.
+            fetch_k (int, optional): Total results to select k from.
+                Defaults to 20.
+            lambda_mult: Number between 0 and 1 that determines the degree
+                        of diversity among the results with 0 corresponding
+                        to maximum diversity and 1 to minimum diversity.
+                        Defaults to 0.5
+            param (dict, optional): The search params for the specified index.
+                Defaults to None.
+            expr (str, optional): Filtering expression. Defaults to None.
+            timeout (int, optional): How long to wait before timeout error.
+                Defaults to None.
+            kwargs: Collection.search() keyword arguments.
+
+        Returns:
+            List[Document]: Document results for search.
+        """
+        if self.col is None:
+            logger.debug("No existing collection to search.")
+            return []
+
+        if param is None:
+            param = self.search_params
+
+        # Determine result metadata fields.
+        output_fields = self.fields[:]
+        output_fields.remove(self._vector_field)
+
+        # Perform the search.
+        res = self.col.search(
+            data=[embedding],
+            anns_field=self._vector_field,
+            param=param,
+            limit=fetch_k,
+            expr=expr,
+            output_fields=output_fields,
+            timeout=timeout,
+            **kwargs,
+        )
+        # Organize results.
+        ids = []
+        documents = []
+        scores = []
+        for result in res[0]:
+            meta = {x: result.entity.get(x) for x in output_fields}
+            doc = Document(page_content=meta.pop(self._text_field), metadata=meta)
+            documents.append(doc)
+            scores.append(result.score)
+            ids.append(result.id)
+
+        vectors = self.col.query(
+            expr=f"{self._primary_field} in {ids}",
+            output_fields=[self._primary_field, self._vector_field],
+            timeout=timeout,
+        )
+        # Reorganize the results from query to match search order.
+        vectors = {x[self._primary_field]: x[self._vector_field] for x in vectors}
+
+        ordered_result_embeddings = [vectors[x] for x in ids]
+
+        # Get the new order of results.
+        new_ordering = maximal_marginal_relevance(
+            np.array(embedding), ordered_result_embeddings, k=k, lambda_mult=lambda_mult
+        )
+
+        # Reorder the values and return.
+        ret = []
+        for x in new_ordering:
+            # Function can return -1 index
+            if x == -1:
+                break
+            else:
+                ret.append(documents[x])
+        return ret
+
+    @classmethod
+    def from_texts(
+        cls,
+        texts: List[str],
+        embedding: Embeddings,
+        metadatas: Optional[List[dict]] = None,
+        collection_name: str = "LangChainCollection",
+        connection_args: dict[str, Any] = DEFAULT_MILVUS_CONNECTION,
+        consistency_level: str = "Session",
+        index_params: Optional[dict] = None,
+        search_params: Optional[dict] = None,
+        drop_old: bool = False,
+        batch_size: int = 100,
+        ids: Optional[Sequence[str]] = None,
+        **kwargs: Any,
+    ) -> Milvus:
+        """Create a Milvus collection, indexes it with HNSW, and insert data.
+
+        Args:
+            texts (List[str]): Text data.
+            embedding (Embeddings): Embedding function.
+            metadatas (Optional[List[dict]]): Metadata for each text if it exists.
+                Defaults to None.
+            collection_name (str, optional): Collection name to use. Defaults to
+                "LangChainCollection".
+            connection_args (dict[str, Any], optional): Connection args to use. Defaults
+                to DEFAULT_MILVUS_CONNECTION.
+            consistency_level (str, optional): Which consistency level to use. Defaults
+                to "Session".
+            index_params (Optional[dict], optional): Which index_params to use. Defaults
+                to None.
+            search_params (Optional[dict], optional): Which search params to use.
+                Defaults to None.
+            drop_old (Optional[bool], optional): Whether to drop the collection with
+                that name if it exists. Defaults to False.
+            batch_size:
+                How many vectors upload per-request.
+                Default: 100
+            ids: Optional[Sequence[str]] = None,
+
+        Returns:
+            Milvus: Milvus Vector Store
+        """
+        vector_db = cls(
+            embedding_function=embedding,
+            collection_name=collection_name,
+            connection_args=connection_args,
+            consistency_level=consistency_level,
+            index_params=index_params,
+            search_params=search_params,
+            drop_old=drop_old,
+            **kwargs,
+        )
+        vector_db.add_texts(texts=texts, metadatas=metadatas, batch_size=batch_size)
+        return vector_db

api/core/index/vector_index/milvus_vector_index.py

@@ -9,30 +9,44 @@ from core.index.base import BaseIndex
 from core.index.vector_index.base import BaseVectorIndex
 from core.vector_store.milvus_vector_store import MilvusVectorStore
 from core.vector_store.weaviate_vector_store import WeaviateVectorStore
-from models.dataset import Dataset
+from extensions.ext_database import db
+from models.dataset import Dataset, DatasetCollectionBinding
 
 
 class MilvusConfig(BaseModel):
-    endpoint: str
+    host: str
+    port: int
     user: str
     password: str
+    secure: bool = False
     batch_size: int = 100
 
     @root_validator()
     def validate_config(cls, values: dict) -> dict:
-        if not values['endpoint']:
-            raise ValueError("config MILVUS_ENDPOINT is required")
+        if not values['host']:
+            raise ValueError("config MILVUS_HOST is required")
+        if not values['port']:
+            raise ValueError("config MILVUS_PORT is required")
         if not values['user']:
             raise ValueError("config MILVUS_USER is required")
         if not values['password']:
             raise ValueError("config MILVUS_PASSWORD is required")
         return values
 
+    def to_milvus_params(self):
+        return {
+            'host': self.host,
+            'port': self.port,
+            'user': self.user,
+            'password': self.password,
+            'secure': self.secure
+        }
+
 
 class MilvusVectorIndex(BaseVectorIndex):
     def __init__(self, dataset: Dataset, config: MilvusConfig, embeddings: Embeddings):
         super().__init__(dataset, embeddings)
-        self._client = self._init_client(config)
+        self._client_config = config
 
     def get_type(self) -> str:
         return 'milvus'
@@ -49,7 +63,6 @@ class MilvusVectorIndex(BaseVectorIndex):
         dataset_id = dataset.id
         return "Vector_index_" + dataset_id.replace("-", "_") + '_Node'
 
-
     def to_index_struct(self) -> dict:
         return {
             "type": self.get_type(),
@@ -58,26 +71,29 @@ class MilvusVectorIndex(BaseVectorIndex):
 
     def create(self, texts: list[Document], **kwargs) -> BaseIndex:
         uuids = self._get_uuids(texts)
-        self._vector_store = WeaviateVectorStore.from_documents(
+        index_params = {
+            'metric_type': 'IP',
+            'index_type': "HNSW",
+            'params':  {"M": 8, "efConstruction": 64}
+        }
+        self._vector_store = MilvusVectorStore.from_documents(
             texts,
             self._embeddings,
-            client=self._client,
-            index_name=self.get_index_name(self.dataset),
-            uuids=uuids,
-            by_text=False
+            collection_name=self.get_index_name(self.dataset),
+            connection_args=self._client_config.to_milvus_params(),
+            index_params=index_params
         )
 
         return self
 
     def create_with_collection_name(self, texts: list[Document], collection_name: str, **kwargs) -> BaseIndex:
         uuids = self._get_uuids(texts)
-        self._vector_store = WeaviateVectorStore.from_documents(
+        self._vector_store = MilvusVectorStore.from_documents(
             texts,
             self._embeddings,
-            client=self._client,
-            index_name=collection_name,
-            uuids=uuids,
-            by_text=False
+            collection_name=collection_name,
+            ids=uuids,
+            content_payload_key='page_content'
         )
 
         return self
@@ -86,42 +102,53 @@ class MilvusVectorIndex(BaseVectorIndex):
         """Only for created index."""
         if self._vector_store:
             return self._vector_store
-
         attributes = ['doc_id', 'dataset_id', 'document_id']
-        if self._is_origin():
-            attributes = ['doc_id']
 
-        return WeaviateVectorStore(
-            client=self._client,
-            index_name=self.get_index_name(self.dataset),
-            text_key='text',
-            embedding=self._embeddings,
-            attributes=attributes,
-            by_text=False
+        return MilvusVectorStore(
+            collection_name=self.get_index_name(self.dataset),
+            embedding_function=self._embeddings,
+            connection_args=self._client_config.to_milvus_params()
         )
 
     def _get_vector_store_class(self) -> type:
         return MilvusVectorStore
 
     def delete_by_document_id(self, document_id: str):
-        if self._is_origin():
-            self.recreate_dataset(self.dataset)
-            return
+
+        vector_store = self._get_vector_store()
+        vector_store = cast(self._get_vector_store_class(), vector_store)
+        ids = vector_store.get_ids_by_document_id(document_id)
+        if ids:
+            vector_store.del_texts({
+                'filter': f'id in {ids}'
+            })
+
+    def delete_by_ids(self, doc_ids: list[str]) -> None:
+
+        vector_store = self._get_vector_store()
+        vector_store = cast(self._get_vector_store_class(), vector_store)
+        ids = vector_store.get_ids_by_doc_ids(doc_ids)
+        vector_store.del_texts({
+            'filter': f' id in {ids}'
+        })
+
+    def delete_by_group_id(self, group_id: str) -> None:
 
         vector_store = self._get_vector_store()
         vector_store = cast(self._get_vector_store_class(), vector_store)
 
-        vector_store.del_texts({
-            "operator": "Equal",
-            "path": ["document_id"],
-            "valueText": document_id
-        })
+        vector_store.delete()
 
-    def _is_origin(self):
-        if self.dataset.index_struct_dict:
-            class_prefix: str = self.dataset.index_struct_dict['vector_store']['class_prefix']
-            if not class_prefix.endswith('_Node'):
-                # original class_prefix
-                return True
+    def delete(self) -> None:
+        vector_store = self._get_vector_store()
+        vector_store = cast(self._get_vector_store_class(), vector_store)
 
-        return False
+        from qdrant_client.http import models
+        vector_store.del_texts(models.Filter(
+            must=[
+                models.FieldCondition(
+                    key="group_id",
+                    match=models.MatchValue(value=self.dataset.id),
+                ),
+            ],
+        ))

api/core/index/vector_index/qdrant.py

@@ -1390,70 +1390,12 @@ class Qdrant(VectorStore):
             path=path,
             **kwargs,
         )
-        try:
-            # Skip any validation in case of forced collection recreate.
-            if force_recreate:
-                raise ValueError
-
-            # Get the vector configuration of the existing collection and vector, if it
-            # was specified. If the old configuration does not match the current one,
-            # an exception is being thrown.
-            collection_info = client.get_collection(collection_name=collection_name)
-            current_vector_config = collection_info.config.params.vectors
-            if isinstance(current_vector_config, dict) and vector_name is not None:
-                if vector_name not in current_vector_config:
-                    raise QdrantException(
-                        f"Existing Qdrant collection {collection_name} does not "
-                        f"contain vector named {vector_name}. Did you mean one of the "
-                        f"existing vectors: {', '.join(current_vector_config.keys())}? "
-                        f"If you want to recreate the collection, set `force_recreate` "
-                        f"parameter to `True`."
-                    )
-                current_vector_config = current_vector_config.get(
-                    vector_name
-                )  # type: ignore[assignment]
-            elif isinstance(current_vector_config, dict) and vector_name is None:
-                raise QdrantException(
-                    f"Existing Qdrant collection {collection_name} uses named vectors. "
-                    f"If you want to reuse it, please set `vector_name` to any of the "
-                    f"existing named vectors: "
-                    f"{', '.join(current_vector_config.keys())}."  # noqa
-                    f"If you want to recreate the collection, set `force_recreate` "
-                    f"parameter to `True`."
-                )
-            elif (
-                not isinstance(current_vector_config, dict) and vector_name is not None
-            ):
-                raise QdrantException(
-                    f"Existing Qdrant collection {collection_name} doesn't use named "
-                    f"vectors. If you want to reuse it, please set `vector_name` to "
-                    f"`None`. If you want to recreate the collection, set "
-                    f"`force_recreate` parameter to `True`."
-                )
-
-            # Check if the vector configuration has the same dimensionality.
-            if current_vector_config.size != vector_size:  # type: ignore[union-attr]
-                raise QdrantException(
-                    f"Existing Qdrant collection is configured for vectors with "
-                    f"{current_vector_config.size} "  # type: ignore[union-attr]
-                    f"dimensions. Selected embeddings are {vector_size}-dimensional. "
-                    f"If you want to recreate the collection, set `force_recreate` "
-                    f"parameter to `True`."
-                )
-
-            current_distance_func = (
-                current_vector_config.distance.name.upper()  # type: ignore[union-attr]
-            )
-            if current_distance_func != distance_func:
-                raise QdrantException(
-                    f"Existing Qdrant collection is configured for "
-                    f"{current_vector_config.distance} "  # type: ignore[union-attr]
-                    f"similarity. Please set `distance_func` parameter to "
-                    f"`{distance_func}` if you want to reuse it. If you want to "
-                    f"recreate the collection, set `force_recreate` parameter to "
-                    f"`True`."
-                )
-        except (UnexpectedResponse, RpcError, ValueError):
+        all_collection_name = []
+        collections_response = client.get_collections()
+        collection_list = collections_response.collections
+        for collection in collection_list:
+            all_collection_name.append(collection.name)
+        if collection_name not in all_collection_name:
             vectors_config = rest.VectorParams(
                 size=vector_size,
                 distance=rest.Distance[distance_func],
@@ -1481,6 +1423,67 @@ class Qdrant(VectorStore):
                 timeout=timeout,  # type: ignore[arg-type]
             )
             is_new_collection = True
+        if force_recreate:
+            raise ValueError
+
+        # Get the vector configuration of the existing collection and vector, if it
+        # was specified. If the old configuration does not match the current one,
+        # an exception is being thrown.
+        collection_info = client.get_collection(collection_name=collection_name)
+        current_vector_config = collection_info.config.params.vectors
+        if isinstance(current_vector_config, dict) and vector_name is not None:
+            if vector_name not in current_vector_config:
+                raise QdrantException(
+                    f"Existing Qdrant collection {collection_name} does not "
+                    f"contain vector named {vector_name}. Did you mean one of the "
+                    f"existing vectors: {', '.join(current_vector_config.keys())}? "
+                    f"If you want to recreate the collection, set `force_recreate` "
+                    f"parameter to `True`."
+                )
+            current_vector_config = current_vector_config.get(
+                vector_name
+            )  # type: ignore[assignment]
+        elif isinstance(current_vector_config, dict) and vector_name is None:
+            raise QdrantException(
+                f"Existing Qdrant collection {collection_name} uses named vectors. "
+                f"If you want to reuse it, please set `vector_name` to any of the "
+                f"existing named vectors: "
+                f"{', '.join(current_vector_config.keys())}."  # noqa
+                f"If you want to recreate the collection, set `force_recreate` "
+                f"parameter to `True`."
+            )
+        elif (
+                not isinstance(current_vector_config, dict) and vector_name is not None
+        ):
+            raise QdrantException(
+                f"Existing Qdrant collection {collection_name} doesn't use named "
+                f"vectors. If you want to reuse it, please set `vector_name` to "
+                f"`None`. If you want to recreate the collection, set "
+                f"`force_recreate` parameter to `True`."
+            )
+
+        # Check if the vector configuration has the same dimensionality.
+        if current_vector_config.size != vector_size:  # type: ignore[union-attr]
+            raise QdrantException(
+                f"Existing Qdrant collection is configured for vectors with "
+                f"{current_vector_config.size} "  # type: ignore[union-attr]
+                f"dimensions. Selected embeddings are {vector_size}-dimensional. "
+                f"If you want to recreate the collection, set `force_recreate` "
+                f"parameter to `True`."
+            )
+
+        current_distance_func = (
+            current_vector_config.distance.name.upper()  # type: ignore[union-attr]
+        )
+        if current_distance_func != distance_func:
+            raise QdrantException(
+                f"Existing Qdrant collection is configured for "
+                f"{current_vector_config.distance} "  # type: ignore[union-attr]
+                f"similarity. Please set `distance_func` parameter to "
+                f"`{distance_func}` if you want to reuse it. If you want to "
+                f"recreate the collection, set `force_recreate` parameter to "
+                f"`True`."
+            )
         qdrant = cls(
             client=client,
             collection_name=collection_name,

api/core/index/vector_index/qdrant_vector_index.py

@@ -169,6 +169,19 @@ class QdrantVectorIndex(BaseVectorIndex):
             ],
         ))
 
+    def delete(self) -> None:
+        vector_store = self._get_vector_store()
+        vector_store = cast(self._get_vector_store_class(), vector_store)
+
+        from qdrant_client.http import models
+        vector_store.del_texts(models.Filter(
+            must=[
+                models.FieldCondition(
+                    key="group_id",
+                    match=models.MatchValue(value=self.dataset.id),
+                ),
+            ],
+        ))
 
     def _is_origin(self):
         if self.dataset.index_struct_dict:

api/core/index/vector_index/vector_index.py

@@ -47,6 +47,20 @@ class VectorIndex:
                 ),
                 embeddings=embeddings
             )
+        elif vector_type == "milvus":
+            from core.index.vector_index.milvus_vector_index import MilvusVectorIndex, MilvusConfig
+
+            return MilvusVectorIndex(
+                dataset=dataset,
+                config=MilvusConfig(
+                    host=config.get('MILVUS_HOST'),
+                    port=config.get('MILVUS_PORT'),
+                    user=config.get('MILVUS_USER'),
+                    password=config.get('MILVUS_PASSWORD'),
+                    secure=config.get('MILVUS_SECURE'),
+                ),
+                embeddings=embeddings
+            )
         else:
             raise ValueError(f"Vector store {config.get('VECTOR_STORE')} is not supported.")
 

api/core/model_providers/model_provider_factory.py

@@ -51,6 +51,9 @@ class ModelProviderFactory:
         elif provider_name == 'chatglm':
             from core.model_providers.providers.chatglm_provider import ChatGLMProvider
             return ChatGLMProvider
+        elif provider_name == 'baichuan':
+            from core.model_providers.providers.baichuan_provider import BaichuanProvider
+            return BaichuanProvider
         elif provider_name == 'azure_openai':
             from core.model_providers.providers.azure_openai_provider import AzureOpenAIProvider
             return AzureOpenAIProvider

api/core/model_providers/models/embedding/huggingface_embedding.py

@@ -0,0 +1,22 @@
+from core.model_providers.error import LLMBadRequestError
+from core.model_providers.providers.base import BaseModelProvider
+from core.third_party.langchain.embeddings.huggingface_hub_embedding import HuggingfaceHubEmbeddings
+from core.model_providers.models.embedding.base import BaseEmbedding
+
+
+class HuggingfaceEmbedding(BaseEmbedding):
+    def __init__(self, model_provider: BaseModelProvider, name: str):
+        credentials = model_provider.get_model_credentials(
+            model_name=name,
+            model_type=self.type
+        )
+
+        client = HuggingfaceHubEmbeddings(
+            model=name,
+            **credentials
+        )
+
+        super().__init__(model_provider, client, name)
+
+    def handle_exceptions(self, ex: Exception) -> Exception:
+        return LLMBadRequestError(f"Huggingface embedding: {str(ex)}")

api/core/model_providers/models/embedding/openllm_embedding.py

@@ -0,0 +1,22 @@
+from core.third_party.langchain.embeddings.openllm_embedding import OpenLLMEmbeddings
+
+from core.model_providers.error import LLMBadRequestError
+from core.model_providers.providers.base import BaseModelProvider
+from core.model_providers.models.embedding.base import BaseEmbedding
+
+
+class OpenLLMEmbedding(BaseEmbedding):
+    def __init__(self, model_provider: BaseModelProvider, name: str):
+        credentials = model_provider.get_model_credentials(
+            model_name=name,
+            model_type=self.type
+        )
+
+        client = OpenLLMEmbeddings(
+            server_url=credentials['server_url']
+        )
+
+        super().__init__(model_provider, client, name)
+
+    def handle_exceptions(self, ex: Exception) -> Exception:
+        return LLMBadRequestError(f"OpenLLM embedding: {str(ex)}")

api/core/model_providers/models/embedding/xinference_embedding.py

@@ -1,5 +1,4 @@
 from core.third_party.langchain.embeddings.xinference_embedding import XinferenceEmbedding as XinferenceEmbeddings
-from replicate.exceptions import ModelError, ReplicateError
 
 from core.model_providers.error import LLMBadRequestError
 from core.model_providers.providers.base import BaseModelProvider
@@ -21,7 +20,4 @@ class XinferenceEmbedding(BaseEmbedding):
         super().__init__(model_provider, client, name)
 
     def handle_exceptions(self, ex: Exception) -> Exception:
-        if isinstance(ex, (ModelError, ReplicateError)):
-            return LLMBadRequestError(f"Xinference embedding: {str(ex)}")
-        else:
-            return ex
+        return LLMBadRequestError(f"Xinference embedding: {str(ex)}")

api/core/model_providers/models/llm/baichuan_model.py

@@ -0,0 +1,67 @@
+from typing import List, Optional, Any
+
+from langchain.callbacks.manager import Callbacks
+from langchain.schema import LLMResult
+
+from core.model_providers.error import LLMBadRequestError
+from core.model_providers.models.llm.base import BaseLLM
+from core.model_providers.models.entity.message import PromptMessage
+from core.model_providers.models.entity.model_params import ModelMode, ModelKwargs
+from core.third_party.langchain.llms.baichuan_llm import BaichuanChatLLM
+
+
+class BaichuanModel(BaseLLM):
+    model_mode: ModelMode = ModelMode.CHAT
+
+    def _init_client(self) -> Any:
+        provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, self.model_kwargs)
+        return BaichuanChatLLM(
+            streaming=self.streaming,
+            callbacks=self.callbacks,
+            **self.credentials,
+            **provider_model_kwargs
+        )
+
+    def _run(self, messages: List[PromptMessage],
+             stop: Optional[List[str]] = None,
+             callbacks: Callbacks = None,
+             **kwargs) -> LLMResult:
+        """
+        run predict by prompt messages and stop words.
+
+        :param messages:
+        :param stop:
+        :param callbacks:
+        :return:
+        """
+        prompts = self._get_prompt_from_messages(messages)
+        return self._client.generate([prompts], stop, callbacks)
+
+    def prompt_file_name(self, mode: str) -> str:
+        if mode == 'completion':
+            return 'baichuan_completion'
+        else:
+            return 'baichuan_chat'
+
+    def get_num_tokens(self, messages: List[PromptMessage]) -> int:
+        """
+        get num tokens of prompt messages.
+
+        :param messages:
+        :return:
+        """
+        prompts = self._get_prompt_from_messages(messages)
+        return max(self._client.get_num_tokens_from_messages(prompts), 0)
+
+    def _set_model_kwargs(self, model_kwargs: ModelKwargs):
+        provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, model_kwargs)
+        for k, v in provider_model_kwargs.items():
+            if hasattr(self.client, k):
+                setattr(self.client, k, v)
+
+    def handle_exceptions(self, ex: Exception) -> Exception:
+        return LLMBadRequestError(f"Baichuan: {str(ex)}")
+
+    @property
+    def support_streaming(self):
+        return True

api/core/model_providers/models/llm/base.py

@@ -1,6 +1,7 @@
 import json
 import os
 import re
+import time
 from abc import abstractmethod
 from typing import List, Optional, Any, Union, Tuple
 import decimal
@@ -20,6 +21,8 @@ from core.prompt.prompt_template import JinjaPromptTemplate
 from core.third_party.langchain.llms.fake import FakeLLM
 import logging
 
+from extensions.ext_database import db
+
 logger = logging.getLogger(__name__)
 
 

api/core/model_providers/models/llm/openai_model.py

@@ -5,6 +5,7 @@ from typing import List, Optional, Any
 import openai
 from langchain.callbacks.manager import Callbacks
 from langchain.schema import LLMResult
+from openai import api_requestor
 
 from core.model_providers.providers.base import BaseModelProvider
 from core.third_party.langchain.llms.chat_open_ai import EnhanceChatOpenAI

api/core/model_providers/models/llm/wenxin_model.py

@@ -18,6 +18,7 @@ class WenxinModel(BaseLLM):
         provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, self.model_kwargs)
         # TODO load price_config from configs(db)
         return Wenxin(
+            model=self.name,
             streaming=self.streaming,
             callbacks=self.callbacks,
             **self.credentials,

api/core/model_providers/providers/baichuan_provider.py

@@ -0,0 +1,167 @@
+import json
+from json import JSONDecodeError
+from typing import Type
+
+from langchain.schema import HumanMessage
+
+from core.helper import encrypter
+from core.model_providers.models.base import BaseProviderModel
+from core.model_providers.models.entity.model_params import ModelKwargsRules, KwargRule, ModelType
+from core.model_providers.models.llm.baichuan_model import BaichuanModel
+from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
+from core.third_party.langchain.llms.baichuan_llm import BaichuanChatLLM
+from models.provider import ProviderType
+
+
+class BaichuanProvider(BaseModelProvider):
+
+    @property
+    def provider_name(self):
+        """
+        Returns the name of a provider.
+        """
+        return 'baichuan'
+
+    def _get_fixed_model_list(self, model_type: ModelType) -> list[dict]:
+        if model_type == ModelType.TEXT_GENERATION:
+            return [
+                {
+                    'id': 'baichuan2-53b',
+                    'name': 'Baichuan2-53B',
+                }
+            ]
+        else:
+            return []
+
+    def get_model_class(self, model_type: ModelType) -> Type[BaseProviderModel]:
+        """
+        Returns the model class.
+
+        :param model_type:
+        :return:
+        """
+        if model_type == ModelType.TEXT_GENERATION:
+            model_class = BaichuanModel
+        else:
+            raise NotImplementedError
+
+        return model_class
+
+    def get_model_parameter_rules(self, model_name: str, model_type: ModelType) -> ModelKwargsRules:
+        """
+        get model parameter rules.
+
+        :param model_name:
+        :param model_type:
+        :return:
+        """
+        return ModelKwargsRules(
+            temperature=KwargRule[float](min=0, max=1, default=0.3, precision=2),
+            top_p=KwargRule[float](min=0, max=0.99, default=0.85, precision=2),
+            presence_penalty=KwargRule[float](enabled=False),
+            frequency_penalty=KwargRule[float](enabled=False),
+            max_tokens=KwargRule[int](enabled=False),
+        )
+
+    @classmethod
+    def is_provider_credentials_valid_or_raise(cls, credentials: dict):
+        """
+        Validates the given credentials.
+        """
+        if 'api_key' not in credentials:
+            raise CredentialsValidateFailedError('Baichuan api_key must be provided.')
+
+        if 'secret_key' not in credentials:
+            raise CredentialsValidateFailedError('Baichuan secret_key must be provided.')
+
+        try:
+            credential_kwargs = {
+                'api_key': credentials['api_key'],
+                'secret_key': credentials['secret_key'],
+            }
+
+            llm = BaichuanChatLLM(
+                temperature=0,
+                **credential_kwargs
+            )
+
+            llm([HumanMessage(content='ping')])
+        except Exception as ex:
+            raise CredentialsValidateFailedError(str(ex))
+
+    @classmethod
+    def encrypt_provider_credentials(cls, tenant_id: str, credentials: dict) -> dict:
+        credentials['api_key'] = encrypter.encrypt_token(tenant_id, credentials['api_key'])
+        credentials['secret_key'] = encrypter.encrypt_token(tenant_id, credentials['secret_key'])
+        return credentials
+
+    def get_provider_credentials(self, obfuscated: bool = False) -> dict:
+        if self.provider.provider_type == ProviderType.CUSTOM.value:
+            try:
+                credentials = json.loads(self.provider.encrypted_config)
+            except JSONDecodeError:
+                credentials = {
+                    'api_key': None,
+                    'secret_key': None,
+                }
+
+            if credentials['api_key']:
+                credentials['api_key'] = encrypter.decrypt_token(
+                    self.provider.tenant_id,
+                    credentials['api_key']
+                )
+
+                if obfuscated:
+                    credentials['api_key'] = encrypter.obfuscated_token(credentials['api_key'])
+
+            if credentials['secret_key']:
+                credentials['secret_key'] = encrypter.decrypt_token(
+                    self.provider.tenant_id,
+                    credentials['secret_key']
+                )
+
+                if obfuscated:
+                    credentials['secret_key'] = encrypter.obfuscated_token(credentials['secret_key'])
+
+            return credentials
+        else:
+            return {}
+
+    def should_deduct_quota(self):
+        return True
+
+    @classmethod
+    def is_model_credentials_valid_or_raise(cls, model_name: str, model_type: ModelType, credentials: dict):
+        """
+        check model credentials valid.
+
+        :param model_name:
+        :param model_type:
+        :param credentials:
+        """
+        return
+
+    @classmethod
+    def encrypt_model_credentials(cls, tenant_id: str, model_name: str, model_type: ModelType,
+                                  credentials: dict) -> dict:
+        """
+        encrypt model credentials for save.
+
+        :param tenant_id:
+        :param model_name:
+        :param model_type:
+        :param credentials:
+        :return:
+        """
+        return {}
+
+    def get_model_credentials(self, model_name: str, model_type: ModelType, obfuscated: bool = False) -> dict:
+        """
+        get credentials for llm use.
+
+        :param model_name:
+        :param model_type:
+        :param obfuscated:
+        :return:
+        """
+        return self.get_provider_credentials(obfuscated)

api/core/model_providers/providers/huggingface_hub_provider.py

@@ -1,5 +1,6 @@
 import json
 from typing import Type
+import requests
 
 from huggingface_hub import HfApi
 
@@ -10,8 +11,12 @@ from core.model_providers.providers.base import BaseModelProvider, CredentialsVa
 
 from core.model_providers.models.base import BaseProviderModel
 from core.third_party.langchain.llms.huggingface_endpoint_llm import HuggingFaceEndpointLLM
+from core.third_party.langchain.embeddings.huggingface_hub_embedding import HuggingfaceHubEmbeddings
+from core.model_providers.models.embedding.huggingface_embedding import HuggingfaceEmbedding
 from models.provider import ProviderType
 
+HUGGINGFACE_ENDPOINT_API = 'https://api.endpoints.huggingface.cloud/v2/endpoint/'
+
 
 class HuggingfaceHubProvider(BaseModelProvider):
     @property
@@ -33,6 +38,8 @@ class HuggingfaceHubProvider(BaseModelProvider):
         """
         if model_type == ModelType.TEXT_GENERATION:
             model_class = HuggingfaceHubModel
+        elif model_type == ModelType.EMBEDDINGS:
+            model_class = HuggingfaceEmbedding
         else:
             raise NotImplementedError
 
@@ -63,7 +70,7 @@ class HuggingfaceHubProvider(BaseModelProvider):
         :param model_type:
         :param credentials:
         """
-        if model_type != ModelType.TEXT_GENERATION:
+        if model_type not in [ModelType.TEXT_GENERATION, ModelType.EMBEDDINGS]:
             raise NotImplementedError
 
         if 'huggingfacehub_api_type' not in credentials \
@@ -88,19 +95,15 @@ class HuggingfaceHubProvider(BaseModelProvider):
             if 'task_type' not in credentials:
                 raise CredentialsValidateFailedError('Task Type must be provided.')
 
-            if credentials['task_type'] not in ("text2text-generation", "text-generation", "summarization"):
+            if credentials['task_type'] not in ("text2text-generation", "text-generation", 'feature-extraction'):
                 raise CredentialsValidateFailedError('Task Type must be one of text2text-generation, '
-                                                     'text-generation, summarization.')
+                                                     'text-generation, feature-extraction.')
 
             try:
-                llm = HuggingFaceEndpointLLM(
-                    endpoint_url=credentials['huggingfacehub_endpoint_url'],
-                    task=credentials['task_type'],
-                    model_kwargs={"temperature": 0.5, "max_new_tokens": 200},
-                    huggingfacehub_api_token=credentials['huggingfacehub_api_token']
-                )
-
-                llm("ping")
+                if credentials['task_type'] == 'feature-extraction':
+                    cls.check_embedding_valid(credentials, model_name)
+                else:
+                    cls.check_llm_valid(credentials)    
             except Exception as e:
                 raise CredentialsValidateFailedError(f"{e.__class__.__name__}:{str(e)}")
         else:
@@ -112,13 +115,64 @@ class HuggingfaceHubProvider(BaseModelProvider):
                 if 'inference' in model_info.cardData and not model_info.cardData['inference']:
                     raise ValueError(f'Inference API has been turned off for this model {model_name}.')
 
-                VALID_TASKS = ("text2text-generation", "text-generation", "summarization")
+                VALID_TASKS = ("text2text-generation", "text-generation", "feature-extraction")
                 if model_info.pipeline_tag not in VALID_TASKS:
                     raise ValueError(f"Model {model_name} is not a valid task, "
                                      f"must be one of {VALID_TASKS}.")
             except Exception as e:
                 raise CredentialsValidateFailedError(f"{e.__class__.__name__}:{str(e)}")
 
+    @classmethod
+    def check_llm_valid(cls, credentials: dict):
+        llm = HuggingFaceEndpointLLM(
+            endpoint_url=credentials['huggingfacehub_endpoint_url'],
+            task=credentials['task_type'],
+            model_kwargs={"temperature": 0.5, "max_new_tokens": 200},
+            huggingfacehub_api_token=credentials['huggingfacehub_api_token']
+        )
+
+        llm("ping")
+
+    @classmethod
+    def check_embedding_valid(cls, credentials: dict, model_name: str):
+
+        cls.check_endpoint_url_model_repository_name(credentials, model_name)
+        
+        embedding_model = HuggingfaceHubEmbeddings(
+            model=model_name,
+            **credentials
+        )
+
+        embedding_model.embed_query("ping")
+
+    @classmethod
+    def check_endpoint_url_model_repository_name(cls, credentials: dict, model_name: str):
+        try:
+            url = f'{HUGGINGFACE_ENDPOINT_API}{credentials["huggingface_namespace"]}'
+            headers = {
+                'Authorization': f'Bearer {credentials["huggingfacehub_api_token"]}',
+                'Content-Type': 'application/json'
+            }
+
+            response =requests.get(url=url, headers=headers)
+
+            if response.status_code != 200:
+                raise ValueError('User Name or Organization Name is invalid.')
+
+            model_repository_name = ''
+
+            for item in response.json().get("items", []):
+                if item.get("status", {}).get("url") == credentials['huggingfacehub_endpoint_url']:
+                    model_repository_name = item.get("model", {}).get("repository")
+                    break
+            
+            if model_repository_name != model_name:
+                raise ValueError(f'Model Name {model_name} is invalid. Please check it on the inference endpoints console.')
+
+        except Exception as e:
+            raise ValueError(str(e))
+        
+
     @classmethod
     def encrypt_model_credentials(cls, tenant_id: str, model_name: str, model_type: ModelType,
                                   credentials: dict) -> dict:

api/core/model_providers/providers/openllm_provider.py

@@ -2,11 +2,13 @@ import json
 from typing import Type
 
 from core.helper import encrypter
+from core.model_providers.models.embedding.openllm_embedding import OpenLLMEmbedding
 from core.model_providers.models.entity.model_params import KwargRule, ModelKwargsRules, ModelType
 from core.model_providers.models.llm.openllm_model import OpenLLMModel
 from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
 
 from core.model_providers.models.base import BaseProviderModel
+from core.third_party.langchain.embeddings.openllm_embedding import OpenLLMEmbeddings
 from core.third_party.langchain.llms.openllm import OpenLLM
 from models.provider import ProviderType
 
@@ -31,6 +33,8 @@ class OpenLLMProvider(BaseModelProvider):
         """
         if model_type == ModelType.TEXT_GENERATION:
             model_class = OpenLLMModel
+        elif model_type== ModelType.EMBEDDINGS:
+            model_class = OpenLLMEmbedding
         else:
             raise NotImplementedError
 
@@ -69,14 +73,21 @@ class OpenLLMProvider(BaseModelProvider):
                 'server_url': credentials['server_url']
             }
 
-            llm = OpenLLM(
-                llm_kwargs={
-                    'max_new_tokens': 10
-                },
-                **credential_kwargs
-            )
+            if model_type == ModelType.TEXT_GENERATION:
+                llm = OpenLLM(
+                    llm_kwargs={
+                        'max_new_tokens': 10
+                    },
+                    **credential_kwargs
+                )
 
-            llm("ping")
+                llm("ping")
+            elif model_type == ModelType.EMBEDDINGS:
+                embedding = OpenLLMEmbeddings(
+                    **credential_kwargs
+                )
+
+                embedding.embed_query("ping")
         except Exception as ex:
             raise CredentialsValidateFailedError(str(ex))
 

api/core/model_providers/providers/wenxin_provider.py

@@ -61,13 +61,18 @@ class WenxinProvider(BaseModelProvider):
         :param model_type:
         :return:
         """
+        model_max_tokens = {
+            'ernie-bot': 4800,
+            'ernie-bot-turbo': 11200,
+        }
+
         if model_name in ['ernie-bot', 'ernie-bot-turbo']:
             return ModelKwargsRules(
                 temperature=KwargRule[float](min=0.01, max=1, default=0.95, precision=2),
                 top_p=KwargRule[float](min=0.01, max=1, default=0.8, precision=2),
                 presence_penalty=KwargRule[float](enabled=False),
                 frequency_penalty=KwargRule[float](enabled=False),
-                max_tokens=KwargRule[int](enabled=False),
+                max_tokens=KwargRule[int](enabled=False, max=model_max_tokens.get(model_name)),
             )
         else:
             return ModelKwargsRules(

api/core/model_providers/rules/_providers.json

@@ -7,10 +7,11 @@
   "spark",
   "wenxin",
   "zhipuai",
+  "baichuan",
   "chatglm",
   "replicate",
   "huggingface_hub",
   "xinference",
   "openllm",
   "localai"
-]
+]

api/core/model_providers/rules/baichuan.json

@@ -0,0 +1,15 @@
+{
+    "support_provider_types": [
+        "custom"
+    ],
+    "system_config": null,
+    "model_flexibility": "fixed",
+    "price_config": {
+        "baichuan2-53b": {
+            "prompt": "0.01",
+            "completion": "0.01",
+            "unit": "0.001",
+            "currency": "RMB"
+        }
+    }
+}

api/core/third_party/langchain/embeddings/huggingface_hub_embedding.py

@@ -0,0 +1,74 @@
+from typing import Any, Dict, List, Optional
+import json
+import numpy as np
+
+from pydantic import BaseModel, Extra, root_validator
+
+from langchain.embeddings.base import Embeddings
+from langchain.utils import get_from_dict_or_env
+from huggingface_hub import InferenceClient
+
+HOSTED_INFERENCE_API = 'hosted_inference_api'
+INFERENCE_ENDPOINTS = 'inference_endpoints'
+
+
+class HuggingfaceHubEmbeddings(BaseModel, Embeddings):
+    client: Any
+    model: str
+
+    huggingface_namespace: Optional[str] = None
+    task_type: Optional[str] = None
+    huggingfacehub_api_type: Optional[str] = None
+    huggingfacehub_api_token: Optional[str] = None
+    huggingfacehub_endpoint_url: Optional[str] = None
+
+    class Config:
+        extra = Extra.forbid
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        values['huggingfacehub_api_token'] = get_from_dict_or_env(
+            values, "huggingfacehub_api_token", "HUGGINGFACEHUB_API_TOKEN"
+        )
+
+        values['client'] = InferenceClient(token=values['huggingfacehub_api_token'])
+
+        return values
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        model = ''
+
+        if self.huggingfacehub_api_type == HOSTED_INFERENCE_API:
+            model = self.model
+        else:
+            model = self.huggingfacehub_endpoint_url
+
+        output = self.client.post(
+            json={
+                "inputs": texts,
+                "options": {
+                    "wait_for_model": False,
+                    "use_cache": False
+                }
+            }, model=model)
+        
+        embeddings =  json.loads(output.decode())
+        return self.mean_pooling(embeddings)
+
+    def embed_query(self, text: str) -> List[float]:
+        return self.embed_documents([text])[0]
+    
+    # https://huggingface.co/docs/api-inference/detailed_parameters#feature-extraction-task
+    # Returned values are a list of floats, or a list of list of floats 
+    # (depending on if you sent a string or a list of string, 
+    # and if the automatic reduction, usually mean_pooling for instance was applied for you or not. 
+    # This should be explained on the model's README.)
+    def mean_pooling(self, embeddings: List) -> List[float]:
+        # If automatic reduction by giving model, no need to mean_pooling.
+        # For example one: List[List[float]]
+        if not isinstance(embeddings[0][0], list):
+            return embeddings
+
+        # For example two: List[List[List[float]]], need to mean_pooling.
+        sentence_embeddings = [np.mean(embedding[0], axis=0).tolist() for embedding in embeddings]
+        return sentence_embeddings

api/core/third_party/langchain/embeddings/openllm_embedding.py

@@ -0,0 +1,67 @@
+"""Wrapper around OpenLLM embedding models."""
+from typing import Any, List, Optional
+
+import requests
+from pydantic import BaseModel, Extra
+
+from langchain.embeddings.base import Embeddings
+
+
+class OpenLLMEmbeddings(BaseModel, Embeddings):
+    """Wrapper around OpenLLM embedding models.
+    """
+
+    client: Any  #: :meta private:
+
+    server_url: Optional[str] = None
+    """Optional server URL that currently runs a LLMServer with 'openllm start'."""
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        """Call out to OpenLLM's embedding endpoint.
+
+        Args:
+            texts: The list of texts to embed.
+
+        Returns:
+            List of embeddings, one for each text.
+        """
+        embeddings = []
+        for text in texts:
+            result = self.invoke_embedding(text=text)
+            embeddings.append(result)
+
+        return [list(map(float, e)) for e in embeddings]
+
+    def invoke_embedding(self, text):
+        params = [
+            text
+        ]
+
+        headers = {"Content-Type": "application/json"}
+        response = requests.post(
+            f'{self.server_url}/v1/embeddings',
+            headers=headers,
+            json=params
+        )
+
+        if not response.ok:
+            raise ValueError(f"OpenLLM HTTP {response.status_code} error: {response.text}")
+
+        json_response = response.json()
+        return json_response[0]["embeddings"][0]
+
+    def embed_query(self, text: str) -> List[float]:
+        """Call out to OpenLLM's embedding endpoint.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            Embeddings for the text.
+        """
+        return self.embed_documents([text])[0]

api/core/third_party/langchain/llms/baichuan_llm.py

@@ -0,0 +1,315 @@
+"""Wrapper around Baichuan APIs."""
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import time
+from typing import (
+    Any,
+    Dict,
+    List,
+    Optional, Iterator,
+)
+
+import requests
+from langchain.chat_models.base import BaseChatModel
+from langchain.schema import BaseMessage, ChatMessage, HumanMessage, AIMessage, SystemMessage
+from langchain.schema.messages import AIMessageChunk
+from langchain.schema.output import ChatResult, ChatGenerationChunk, ChatGeneration
+from pydantic import Extra, root_validator, BaseModel
+
+from langchain.callbacks.manager import (
+    CallbackManagerForLLMRun,
+)
+from langchain.utils import get_from_dict_or_env
+
+logger = logging.getLogger(__name__)
+
+
+class BaichuanModelAPI(BaseModel):
+    api_key: str
+    secret_key: str
+
+    base_url: str = "https://api.baichuan-ai.com/v1"
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    def do_request(self, model: str, messages: list[dict], parameters: dict, **kwargs: Any):
+        stream = 'stream' in kwargs and kwargs['stream']
+
+        url = self.base_url + ("/stream/chat" if stream else "/chat")
+
+        data = {
+            "model": model,
+            "messages": messages,
+            "parameters": parameters
+        }
+
+        json_data = json.dumps(data)
+        time_stamp = int(time.time())
+        signature = self._calculate_md5(self.secret_key + json_data + str(time_stamp))
+
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": "Bearer " + self.api_key,
+            "X-BC-Request-Id": "your requestId",
+            "X-BC-Timestamp": str(time_stamp),
+            "X-BC-Signature": signature,
+            "X-BC-Sign-Algo": "MD5",
+        }
+
+        response = requests.post(url, data=json_data, headers=headers, stream=stream, timeout=(5, 60))
+
+        if not response.ok:
+            raise ValueError(f"HTTP {response.status_code} error: {response.text}")
+
+        if not stream:
+            json_response = response.json()
+            if json_response['code'] != 0:
+                raise ValueError(
+                    f"API {json_response['code']}"
+                    f" error: {json_response['msg']}"
+                )
+            return json_response
+        else:
+            return response
+
+    def _calculate_md5(self, input_string):
+        md5 = hashlib.md5()
+        md5.update(input_string.encode('utf-8'))
+        encrypted = md5.hexdigest()
+        return encrypted
+
+
+class BaichuanChatLLM(BaseChatModel):
+    """Wrapper around Baichuan large language models.
+    To use, you should pass the api_key as a named parameter to the constructor.
+    Example:
+     .. code-block:: python
+         from core.third_party.langchain.llms.baichuan_llm import BaichuanChatLLM
+         model = BaichuanChatLLM(model="<model_name>", api_key="my-api-key", secret_key="my-secret-key")
+    """
+
+    @property
+    def lc_secrets(self) -> Dict[str, str]:
+        return {"api_key": "API_KEY", "secret_key": "SECRET_KEY"}
+
+    @property
+    def lc_serializable(self) -> bool:
+        return True
+
+    client: Any = None  #: :meta private:
+    model: str = "Baichuan2-53B"
+    """Model name to use."""
+    temperature: float = 0.3
+    """A non-negative float that tunes the degree of randomness in generation."""
+    top_p: float = 0.85
+    """Total probability mass of tokens to consider at each step."""
+    streaming: bool = False
+    """Whether to stream the response or return it all at once."""
+    api_key: Optional[str] = None
+    secret_key: Optional[str] = None
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        """Validate that api key and python package exists in environment."""
+        values["api_key"] = get_from_dict_or_env(
+            values, "api_key", "BAICHUAN_API_KEY"
+        )
+
+        values["secret_key"] = get_from_dict_or_env(
+            values, "secret_key", "BAICHUAN_SECRET_KEY"
+        )
+
+        values['client'] = BaichuanModelAPI(
+            api_key=values['api_key'],
+            secret_key=values['secret_key']
+        )
+        return values
+
+    @property
+    def _default_params(self) -> Dict[str, Any]:
+        """Get the default parameters for calling OpenAI API."""
+        return {
+            "model": self.model,
+            "parameters": {
+                "temperature": self.temperature,
+                "top_p": self.top_p
+            }
+        }
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        """Get the identifying parameters."""
+        return self._default_params
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "baichuan"
+
+    def _convert_message_to_dict(self, message: BaseMessage) -> dict:
+        if isinstance(message, ChatMessage):
+            message_dict = {"role": message.role, "content": message.content}
+        elif isinstance(message, HumanMessage):
+            message_dict = {"role": "user", "content": message.content}
+        elif isinstance(message, AIMessage):
+            message_dict = {"role": "assistant", "content": message.content}
+        elif isinstance(message, SystemMessage):
+            message_dict = {"role": "user", "content": message.content}
+        else:
+            raise ValueError(f"Got unknown type {message}")
+        return message_dict
+
+    def _convert_dict_to_message(self, _dict: Dict[str, Any]) -> BaseMessage:
+        role = _dict["role"]
+        if role == "user":
+            return HumanMessage(content=_dict["content"])
+        elif role == "assistant":
+            return AIMessage(content=_dict["content"])
+        elif role == "system":
+            return SystemMessage(content=_dict["content"])
+        else:
+            return ChatMessage(content=_dict["content"], role=role)
+
+    def _create_message_dicts(
+        self, messages: List[BaseMessage]
+    ) -> List[Dict[str, Any]]:
+        dict_messages = []
+        for m in messages:
+            message = self._convert_message_to_dict(m)
+            if dict_messages:
+                previous_message = dict_messages[-1]
+                if previous_message['role'] == message['role']:
+                    dict_messages[-1]['content'] += f"\n{message['content']}"
+                else:
+                    dict_messages.append(message)
+            else:
+                dict_messages.append(message)
+
+        return dict_messages
+
+    def _generate(
+            self,
+            messages: List[BaseMessage],
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            **kwargs: Any,
+    ) -> ChatResult:
+        if self.streaming:
+            generation: Optional[ChatGenerationChunk] = None
+            llm_output: Optional[Dict] = None
+            for chunk in self._stream(
+                    messages=messages, stop=stop, run_manager=run_manager, **kwargs
+            ):
+                if generation is None:
+                    generation = chunk
+                else:
+                    generation += chunk
+
+                if chunk.generation_info is not None \
+                        and 'token_usage' in chunk.generation_info:
+                    llm_output = {"token_usage": chunk.generation_info['token_usage'], "model_name": self.model}
+
+            assert generation is not None
+            return ChatResult(generations=[generation], llm_output=llm_output)
+        else:
+            message_dicts = self._create_message_dicts(messages)
+            params = self._default_params
+            params["messages"] = message_dicts
+            params.update(kwargs)
+            response = self.client.do_request(**params)
+            return self._create_chat_result(response)
+
+    def _stream(
+            self,
+            messages: List[BaseMessage],
+            stop: Optional[List[str]] = None,
+            run_manager: Optional[CallbackManagerForLLMRun] = None,
+            **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        message_dicts = self._create_message_dicts(messages)
+        params = self._default_params
+        params["messages"] = message_dicts
+        params.update(kwargs)
+
+        for event in self.client.do_request(stream=True, **params).iter_lines():
+            if event:
+                event = event.decode("utf-8")
+
+                meta = json.loads(event)
+
+                if meta['code'] != 0:
+                    raise ValueError(
+                        f"API {meta['code']}"
+                        f" error: {meta['msg']}"
+                    )
+
+                content = meta['data']['messages'][0]['content']
+
+                chunk_kwargs = {
+                    'message': AIMessageChunk(content=content),
+                }
+
+                if 'usage' in meta:
+                    token_usage = meta['usage']
+                    overall_token_usage = {
+                        'prompt_tokens': token_usage.get('prompt_tokens', 0),
+                        'completion_tokens': token_usage.get('answer_tokens', 0),
+                        'total_tokens': token_usage.get('total_tokens', 0)
+                    }
+                    chunk_kwargs['generation_info'] = {'token_usage': overall_token_usage}
+
+                yield ChatGenerationChunk(**chunk_kwargs)
+                if run_manager:
+                    run_manager.on_llm_new_token(content)
+
+    def _create_chat_result(self, response: Dict[str, Any]) -> ChatResult:
+        data = response["data"]
+        generations = []
+        for res in data["messages"]:
+            message = self._convert_dict_to_message(res)
+            gen = ChatGeneration(
+                message=message
+            )
+            generations.append(gen)
+        usage = response.get("usage")
+        token_usage = {
+            'prompt_tokens': usage.get('prompt_tokens', 0),
+            'completion_tokens': usage.get('answer_tokens', 0),
+            'total_tokens': usage.get('total_tokens', 0)
+        }
+        llm_output = {"token_usage": token_usage, "model_name": self.model}
+        return ChatResult(generations=generations, llm_output=llm_output)
+
+    def get_num_tokens_from_messages(self, messages: List[BaseMessage]) -> int:
+        """Get the number of tokens in the messages.
+
+        Useful for checking if an input will fit in a model's context window.
+
+        Args:
+            messages: The message inputs to tokenize.
+
+        Returns:
+            The sum of the number of tokens across the messages.
+        """
+        return sum([self.get_num_tokens(m.content) for m in messages])
+
+    def _combine_llm_outputs(self, llm_outputs: List[Optional[dict]]) -> dict:
+        token_usage: dict = {}
+        for output in llm_outputs:
+            if output is None:
+                # Happens in streaming
+                continue
+            token_usage = output["token_usage"]
+
+        return {"token_usage": token_usage, "model_name": self.model}

api/core/third_party/langchain/llms/huggingface_hub_llm.py

@@ -16,7 +16,7 @@ class HuggingFaceHubLLM(HuggingFaceHub):
     environment variable ``HUGGINGFACEHUB_API_TOKEN`` set with your API token, or pass
     it as a named parameter to the constructor.
 
-    Only supports `text-generation`, `text2text-generation` and `summarization` for now.
+    Only supports `text-generation`, `text2text-generation` for now.
 
     Example:
         .. code-block:: python

api/core/third_party/langchain/llms/openllm.py

@@ -66,6 +66,7 @@ class OpenLLM(LLM):
 
         json_response = response.json()
         completion = json_response["responses"][0]
+        completion = completion.lstrip(prompt)
 
         if stop is not None:
             completion = enforce_stop_tokens(completion, stop)

api/core/vector_store/milvus_vector_store.py

@@ -1,4 +1,4 @@
-from langchain.vectorstores import  Milvus
+from core.index.vector_index.milvus import Milvus
 
 
 class MilvusVectorStore(Milvus):
@@ -6,33 +6,41 @@ class MilvusVectorStore(Milvus):
         if not where_filter:
             raise ValueError('where_filter must not be empty')
 
-        self._client.batch.delete_objects(
-            class_name=self._index_name,
-            where=where_filter,
-            output='minimal'
-        )
+        self.col.delete(where_filter.get('filter'))
 
     def del_text(self, uuid: str) -> None:
-        self._client.data_object.delete(
-            uuid,
-            class_name=self._index_name
-        )
+        expr = f"id == {uuid}"
+        self.col.delete(expr)
 
     def text_exists(self, uuid: str) -> bool:
-        result = self._client.query.get(self._index_name).with_additional(["id"]).with_where({
-            "path": ["doc_id"],
-            "operator": "Equal",
-            "valueText": uuid,
-        }).with_limit(1).do()
+        result = self.col.query(
+            expr=f'metadata["doc_id"] == "{uuid}"',
+            output_fields=["id"]
+        )
 
-        if "errors" in result:
-            raise ValueError(f"Error during query: {result['errors']}")
+        return len(result) > 0
 
-        entries = result["data"]["Get"][self._index_name]
-        if len(entries) == 0:
-            return False
+    def get_ids_by_document_id(self, document_id: str):
+        result = self.col.query(
+            expr=f'metadata["document_id"] == "{document_id}"',
+            output_fields=["id"]
+        )
+        if result:
+            return [item["id"] for item in result]
+        else:
+            return None
 
-        return True
+    def get_ids_by_doc_ids(self, doc_ids: list):
+        result = self.col.query(
+            expr=f'metadata["doc_id"] in {doc_ids}',
+            output_fields=["id"]
+        )
+        if result:
+            return [item["id"] for item in result]
+        else:
+            return None
 
     def delete(self):
-        self._client.schema.delete_class(self._index_name)
+        from pymilvus import utility
+        utility.drop_collection(self.collection_name, None, self.alias)
+

api/events/event_handlers/clean_when_dataset_deleted.py

@@ -5,4 +5,5 @@ from tasks.clean_dataset_task import clean_dataset_task
 @dataset_was_deleted.connect
 def handle(sender, **kwargs):
     dataset = sender
-    clean_dataset_task.delay(dataset.id, dataset.tenant_id, dataset.indexing_technique, dataset.index_struct)
+    clean_dataset_task.delay(dataset.id, dataset.tenant_id, dataset.indexing_technique,
+                             dataset.index_struct, dataset.collection_binding_id)

api/extensions/ext_session.py

@@ -1,174 +0,0 @@
-import redis
-from redis.connection import SSLConnection, Connection
-from flask import request
-from flask_session import Session, SqlAlchemySessionInterface, RedisSessionInterface
-from flask_session.sessions import total_seconds
-from itsdangerous import want_bytes
-
-from extensions.ext_database import db
-
-sess = Session()
-
-
-def init_app(app):
-    sqlalchemy_session_interface = CustomSqlAlchemySessionInterface(
-        app,
-        db,
-        app.config.get('SESSION_SQLALCHEMY_TABLE', 'sessions'),
-        app.config.get('SESSION_KEY_PREFIX', 'session:'),
-        app.config.get('SESSION_USE_SIGNER', False),
-        app.config.get('SESSION_PERMANENT', True)
-    )
-
-    session_type = app.config.get('SESSION_TYPE')
-    if session_type == 'sqlalchemy':
-        app.session_interface = sqlalchemy_session_interface
-    elif session_type == 'redis':
-        connection_class = Connection
-        if app.config.get('SESSION_REDIS_USE_SSL', False):
-            connection_class = SSLConnection
-
-        sess_redis_client = redis.Redis()
-        sess_redis_client.connection_pool = redis.ConnectionPool(**{
-            'host': app.config.get('SESSION_REDIS_HOST', 'localhost'),
-            'port': app.config.get('SESSION_REDIS_PORT', 6379),
-            'username': app.config.get('SESSION_REDIS_USERNAME', None),
-            'password': app.config.get('SESSION_REDIS_PASSWORD', None),
-            'db': app.config.get('SESSION_REDIS_DB', 2),
-            'encoding': 'utf-8',
-            'encoding_errors': 'strict',
-            'decode_responses': False
-        }, connection_class=connection_class)
-
-        app.extensions['session_redis'] = sess_redis_client
-
-        app.session_interface = CustomRedisSessionInterface(
-            sess_redis_client,
-            app.config.get('SESSION_KEY_PREFIX', 'session:'),
-            app.config.get('SESSION_USE_SIGNER', False),
-            app.config.get('SESSION_PERMANENT', True)
-        )
-
-
-class CustomSqlAlchemySessionInterface(SqlAlchemySessionInterface):
-
-    def __init__(
-        self,
-        app,
-        db,
-        table,
-        key_prefix,
-        use_signer=False,
-        permanent=True,
-        sequence=None,
-        autodelete=False,
-    ):
-        if db is None:
-            from flask_sqlalchemy import SQLAlchemy
-
-            db = SQLAlchemy(app)
-        self.db = db
-        self.key_prefix = key_prefix
-        self.use_signer = use_signer
-        self.permanent = permanent
-        self.autodelete = autodelete
-        self.sequence = sequence
-        self.has_same_site_capability = hasattr(self, "get_cookie_samesite")
-
-        class Session(self.db.Model):
-            __tablename__ = table
-
-            if sequence:
-                id = self.db.Column(  # noqa: A003, VNE003, A001
-                    self.db.Integer, self.db.Sequence(sequence), primary_key=True
-                )
-            else:
-                id = self.db.Column(  # noqa: A003, VNE003, A001
-                    self.db.Integer, primary_key=True
-                )
-
-            session_id = self.db.Column(self.db.String(255), unique=True)
-            data = self.db.Column(self.db.LargeBinary)
-            expiry = self.db.Column(self.db.DateTime)
-
-            def __init__(self, session_id, data, expiry):
-                self.session_id = session_id
-                self.data = data
-                self.expiry = expiry
-
-            def __repr__(self):
-                return f"<Session data {self.data}>"
-
-        self.sql_session_model = Session
-
-    def save_session(self, *args, **kwargs):
-        if request.blueprint == 'service_api':
-            return
-        elif request.method == 'OPTIONS':
-            return
-        elif request.endpoint and request.endpoint == 'health':
-            return
-        return super().save_session(*args, **kwargs)
-
-
-class CustomRedisSessionInterface(RedisSessionInterface):
-
-    def save_session(self, app, session, response):
-        if request.blueprint == 'service_api':
-            return
-        elif request.method == 'OPTIONS':
-            return
-        elif request.endpoint and request.endpoint == 'health':
-            return
-
-        if not self.should_set_cookie(app, session):
-            return
-        domain = self.get_cookie_domain(app)
-        path = self.get_cookie_path(app)
-        if not session:
-            if session.modified:
-                self.redis.delete(self.key_prefix + session.sid)
-                response.delete_cookie(
-                    app.config["SESSION_COOKIE_NAME"], domain=domain, path=path
-                )
-            return
-
-        # Modification case.  There are upsides and downsides to
-        # emitting a set-cookie header each request.  The behavior
-        # is controlled by the :meth:`should_set_cookie` method
-        # which performs a quick check to figure out if the cookie
-        # should be set or not.  This is controlled by the
-        # SESSION_REFRESH_EACH_REQUEST config flag as well as
-        # the permanent flag on the session itself.
-        # if not self.should_set_cookie(app, session):
-        #    return
-        conditional_cookie_kwargs = {}
-        httponly = self.get_cookie_httponly(app)
-        secure = self.get_cookie_secure(app)
-        if self.has_same_site_capability:
-            conditional_cookie_kwargs["samesite"] = self.get_cookie_samesite(app)
-        expires = self.get_expiration_time(app, session)
-
-        if session.permanent:
-            value = self.serializer.dumps(dict(session))
-            if value is not None:
-                self.redis.setex(
-                    name=self.key_prefix + session.sid,
-                    value=value,
-                    time=total_seconds(app.permanent_session_lifetime),
-                )
-
-        if self.use_signer:
-            session_id = self._get_signer(app).sign(want_bytes(session.sid)).decode("utf-8")
-        else:
-            session_id = session.sid
-        response.set_cookie(
-            app.config["SESSION_COOKIE_NAME"],
-            session_id,
-            expires=expires,
-            httponly=httponly,
-            domain=domain,
-            path=path,
-            secure=secure,
-            **conditional_cookie_kwargs,
-        )

api/fields/app_fields.py

@@ -0,0 +1,138 @@
+from flask_restful import fields
+
+from libs.helper import TimestampField
+
+app_detail_kernel_fields = {
+    'id': fields.String,
+    'name': fields.String,
+    'mode': fields.String,
+    'icon': fields.String,
+    'icon_background': fields.String,
+}
+
+related_app_list = {
+    'data': fields.List(fields.Nested(app_detail_kernel_fields)),
+    'total': fields.Integer,
+}
+
+model_config_fields = {
+    'opening_statement': fields.String,
+    'suggested_questions': fields.Raw(attribute='suggested_questions_list'),
+    'suggested_questions_after_answer': fields.Raw(attribute='suggested_questions_after_answer_dict'),
+    'speech_to_text': fields.Raw(attribute='speech_to_text_dict'),
+    'retriever_resource': fields.Raw(attribute='retriever_resource_dict'),
+    'more_like_this': fields.Raw(attribute='more_like_this_dict'),
+    'sensitive_word_avoidance': fields.Raw(attribute='sensitive_word_avoidance_dict'),
+    'model': fields.Raw(attribute='model_dict'),
+    'user_input_form': fields.Raw(attribute='user_input_form_list'),
+    'dataset_query_variable': fields.String,
+    'pre_prompt': fields.String,
+    'agent_mode': fields.Raw(attribute='agent_mode_dict'),
+}
+
+app_detail_fields = {
+    'id': fields.String,
+    'name': fields.String,
+    'mode': fields.String,
+    'icon': fields.String,
+    'icon_background': fields.String,
+    'enable_site': fields.Boolean,
+    'enable_api': fields.Boolean,
+    'api_rpm': fields.Integer,
+    'api_rph': fields.Integer,
+    'is_demo': fields.Boolean,
+    'model_config': fields.Nested(model_config_fields, attribute='app_model_config'),
+    'created_at': TimestampField
+}
+
+prompt_config_fields = {
+    'prompt_template': fields.String,
+}
+
+model_config_partial_fields = {
+    'model': fields.Raw(attribute='model_dict'),
+    'pre_prompt': fields.String,
+}
+
+app_partial_fields = {
+    'id': fields.String,
+    'name': fields.String,
+    'mode': fields.String,
+    'icon': fields.String,
+    'icon_background': fields.String,
+    'enable_site': fields.Boolean,
+    'enable_api': fields.Boolean,
+    'is_demo': fields.Boolean,
+    'model_config': fields.Nested(model_config_partial_fields, attribute='app_model_config'),
+    'created_at': TimestampField
+}
+
+app_pagination_fields = {
+    'page': fields.Integer,
+    'limit': fields.Integer(attribute='per_page'),
+    'total': fields.Integer,
+    'has_more': fields.Boolean(attribute='has_next'),
+    'data': fields.List(fields.Nested(app_partial_fields), attribute='items')
+}
+
+template_fields = {
+    'name': fields.String,
+    'icon': fields.String,
+    'icon_background': fields.String,
+    'description': fields.String,
+    'mode': fields.String,
+    'model_config': fields.Nested(model_config_fields),
+}
+
+template_list_fields = {
+    'data': fields.List(fields.Nested(template_fields)),
+}
+
+site_fields = {
+    'access_token': fields.String(attribute='code'),
+    'code': fields.String,
+    'title': fields.String,
+    'icon': fields.String,
+    'icon_background': fields.String,
+    'description': fields.String,
+    'default_language': fields.String,
+    'customize_domain': fields.String,
+    'copyright': fields.String,
+    'privacy_policy': fields.String,
+    'customize_token_strategy': fields.String,
+    'prompt_public': fields.Boolean,
+    'app_base_url': fields.String,
+}
+
+app_detail_fields_with_site = {
+    'id': fields.String,
+    'name': fields.String,
+    'mode': fields.String,
+    'icon': fields.String,
+    'icon_background': fields.String,
+    'enable_site': fields.Boolean,
+    'enable_api': fields.Boolean,
+    'api_rpm': fields.Integer,
+    'api_rph': fields.Integer,
+    'is_demo': fields.Boolean,
+    'model_config': fields.Nested(model_config_fields, attribute='app_model_config'),
+    'site': fields.Nested(site_fields),
+    'api_base_url': fields.String,
+    'created_at': TimestampField
+}
+
+app_site_fields = {
+    'app_id': fields.String,
+    'access_token': fields.String(attribute='code'),
+    'code': fields.String,
+    'title': fields.String,
+    'icon': fields.String,
+    'icon_background': fields.String,
+    'description': fields.String,
+    'default_language': fields.String,
+    'customize_domain': fields.String,
+    'copyright': fields.String,
+    'privacy_policy': fields.String,
+    'customize_token_strategy': fields.String,
+    'prompt_public': fields.Boolean
+}

api/fields/conversation_fields.py

@@ -0,0 +1,182 @@
+from flask_restful import fields
+
+from libs.helper import TimestampField
+
+
+class MessageTextField(fields.Raw):
+    def format(self, value):
+        return value[0]['text'] if value else ''
+
+
+account_fields = {
+    'id': fields.String,
+    'name': fields.String,
+    'email': fields.String
+}
+
+feedback_fields = {
+    'rating': fields.String,
+    'content': fields.String,
+    'from_source': fields.String,
+    'from_end_user_id': fields.String,
+    'from_account': fields.Nested(account_fields, allow_null=True),
+}
+
+annotation_fields = {
+    'content': fields.String,
+    'account': fields.Nested(account_fields, allow_null=True),
+    'created_at': TimestampField
+}
+
+message_detail_fields = {
+    'id': fields.String,
+    'conversation_id': fields.String,
+    'inputs': fields.Raw,
+    'query': fields.String,
+    'message': fields.Raw,
+    'message_tokens': fields.Integer,
+    'answer': fields.String,
+    'answer_tokens': fields.Integer,
+    'provider_response_latency': fields.Float,
+    'from_source': fields.String,
+    'from_end_user_id': fields.String,
+    'from_account_id': fields.String,
+    'feedbacks': fields.List(fields.Nested(feedback_fields)),
+    'annotation': fields.Nested(annotation_fields, allow_null=True),
+    'created_at': TimestampField
+}
+
+feedback_stat_fields = {
+    'like': fields.Integer,
+    'dislike': fields.Integer
+}
+
+model_config_fields = {
+    'opening_statement': fields.String,
+    'suggested_questions': fields.Raw,
+    'model': fields.Raw,
+    'user_input_form': fields.Raw,
+    'pre_prompt': fields.String,
+    'agent_mode': fields.Raw,
+}
+
+simple_configs_fields = {
+    'prompt_template': fields.String,
+}
+
+simple_model_config_fields = {
+    'model': fields.Raw(attribute='model_dict'),
+    'pre_prompt': fields.String,
+}
+
+simple_message_detail_fields = {
+    'inputs': fields.Raw,
+    'query': fields.String,
+    'message': MessageTextField,
+    'answer': fields.String,
+}
+
+conversation_fields = {
+    'id': fields.String,
+    'status': fields.String,
+    'from_source': fields.String,
+    'from_end_user_id': fields.String,
+    'from_end_user_session_id': fields.String(),
+    'from_account_id': fields.String,
+    'read_at': TimestampField,
+    'created_at': TimestampField,
+    'annotation': fields.Nested(annotation_fields, allow_null=True),
+    'model_config': fields.Nested(simple_model_config_fields),
+    'user_feedback_stats': fields.Nested(feedback_stat_fields),
+    'admin_feedback_stats': fields.Nested(feedback_stat_fields),
+    'message': fields.Nested(simple_message_detail_fields, attribute='first_message')
+}
+
+conversation_pagination_fields = {
+    'page': fields.Integer,
+    'limit': fields.Integer(attribute='per_page'),
+    'total': fields.Integer,
+    'has_more': fields.Boolean(attribute='has_next'),
+    'data': fields.List(fields.Nested(conversation_fields), attribute='items')
+}
+
+conversation_message_detail_fields = {
+    'id': fields.String,
+    'status': fields.String,
+    'from_source': fields.String,
+    'from_end_user_id': fields.String,
+    'from_account_id': fields.String,
+    'created_at': TimestampField,
+    'model_config': fields.Nested(model_config_fields),
+    'message': fields.Nested(message_detail_fields, attribute='first_message'),
+}
+
+simple_model_config_fields = {
+    'model': fields.Raw(attribute='model_dict'),
+    'pre_prompt': fields.String,
+}
+
+conversation_with_summary_fields = {
+    'id': fields.String,
+    'status': fields.String,
+    'from_source': fields.String,
+    'from_end_user_id': fields.String,
+    'from_end_user_session_id': fields.String,
+    'from_account_id': fields.String,
+    'summary': fields.String(attribute='summary_or_query'),
+    'read_at': TimestampField,
+    'created_at': TimestampField,
+    'annotated': fields.Boolean,
+    'model_config': fields.Nested(simple_model_config_fields),
+    'message_count': fields.Integer,
+    'user_feedback_stats': fields.Nested(feedback_stat_fields),
+    'admin_feedback_stats': fields.Nested(feedback_stat_fields)
+}
+
+conversation_with_summary_pagination_fields = {
+    'page': fields.Integer,
+    'limit': fields.Integer(attribute='per_page'),
+    'total': fields.Integer,
+    'has_more': fields.Boolean(attribute='has_next'),
+    'data': fields.List(fields.Nested(conversation_with_summary_fields), attribute='items')
+}
+
+conversation_detail_fields = {
+    'id': fields.String,
+    'status': fields.String,
+    'from_source': fields.String,
+    'from_end_user_id': fields.String,
+    'from_account_id': fields.String,
+    'created_at': TimestampField,
+    'annotated': fields.Boolean,
+    'model_config': fields.Nested(model_config_fields),
+    'message_count': fields.Integer,
+    'user_feedback_stats': fields.Nested(feedback_stat_fields),
+    'admin_feedback_stats': fields.Nested(feedback_stat_fields)
+}
+
+simple_conversation_fields = {
+    'id': fields.String,
+    'name': fields.String,
+    'inputs': fields.Raw,
+    'status': fields.String,
+    'introduction': fields.String,
+    'created_at': TimestampField
+}
+
+conversation_infinite_scroll_pagination_fields = {
+    'limit': fields.Integer,
+    'has_more': fields.Boolean,
+    'data': fields.List(fields.Nested(simple_conversation_fields))
+}
+
+conversation_with_model_config_fields = {
+    **simple_conversation_fields,
+    'model_config': fields.Raw,
+}
+
+conversation_with_model_config_infinite_scroll_pagination_fields = {
+    'limit': fields.Integer,
+    'has_more': fields.Boolean,
+    'data': fields.List(fields.Nested(conversation_with_model_config_fields))
+}

api/fields/data_source_fields.py

@@ -0,0 +1,65 @@
+from flask_restful import fields
+
+from libs.helper import TimestampField
+
+integrate_icon_fields = {
+    'type': fields.String,
+    'url': fields.String,
+    'emoji': fields.String
+}
+
+integrate_page_fields = {
+    'page_name': fields.String,
+    'page_id': fields.String,
+    'page_icon': fields.Nested(integrate_icon_fields, allow_null=True),
+    'is_bound': fields.Boolean,
+    'parent_id': fields.String,
+    'type': fields.String
+}
+
+integrate_workspace_fields = {
+    'workspace_name': fields.String,
+    'workspace_id': fields.String,
+    'workspace_icon': fields.String,
+    'pages': fields.List(fields.Nested(integrate_page_fields))
+}
+
+integrate_notion_info_list_fields = {
+    'notion_info': fields.List(fields.Nested(integrate_workspace_fields)),
+}
+
+integrate_icon_fields = {
+    'type': fields.String,
+    'url': fields.String,
+    'emoji': fields.String
+}
+
+integrate_page_fields = {
+    'page_name': fields.String,
+    'page_id': fields.String,
+    'page_icon': fields.Nested(integrate_icon_fields, allow_null=True),
+    'parent_id': fields.String,
+    'type': fields.String
+}
+
+integrate_workspace_fields = {
+    'workspace_name': fields.String,
+    'workspace_id': fields.String,
+    'workspace_icon': fields.String,
+    'pages': fields.List(fields.Nested(integrate_page_fields)),
+    'total': fields.Integer
+}
+
+integrate_fields = {
+    'id': fields.String,
+    'provider': fields.String,
+    'created_at': TimestampField,
+    'is_bound': fields.Boolean,
+    'disabled': fields.Boolean,
+    'link': fields.String,
+    'source_info': fields.Nested(integrate_workspace_fields)
+}
+
+integrate_list_fields = {
+    'data': fields.List(fields.Nested(integrate_fields)),
+}

api/fields/dataset_fields.py

@@ -0,0 +1,43 @@
+from flask_restful import fields
+from libs.helper import TimestampField
+
+dataset_fields = {
+    'id': fields.String,
+    'name': fields.String,
+    'description': fields.String,
+    'permission': fields.String,
+    'data_source_type': fields.String,
+    'indexing_technique': fields.String,
+    'created_by': fields.String,
+    'created_at': TimestampField,
+}
+
+dataset_detail_fields = {
+    'id': fields.String,
+    'name': fields.String,
+    'description': fields.String,
+    'provider': fields.String,
+    'permission': fields.String,
+    'data_source_type': fields.String,
+    'indexing_technique': fields.String,
+    'app_count': fields.Integer,
+    'document_count': fields.Integer,
+    'word_count': fields.Integer,
+    'created_by': fields.String,
+    'created_at': TimestampField,
+    'updated_by': fields.String,
+    'updated_at': TimestampField,
+    'embedding_model': fields.String,
+    'embedding_model_provider': fields.String,
+    'embedding_available': fields.Boolean
+}
+
+dataset_query_detail_fields = {
+    "id": fields.String,
+    "content": fields.String,
+    "source": fields.String,
+    "source_app_id": fields.String,
+    "created_by_role": fields.String,
+    "created_by": fields.String,
+    "created_at": TimestampField
+}

api/fields/document_fields.py

@@ -0,0 +1,76 @@
+from flask_restful import fields
+
+from fields.dataset_fields import dataset_fields
+from libs.helper import TimestampField
+
+document_fields = {
+    'id': fields.String,
+    'position': fields.Integer,
+    'data_source_type': fields.String,
+    'data_source_info': fields.Raw(attribute='data_source_info_dict'),
+    'dataset_process_rule_id': fields.String,
+    'name': fields.String,
+    'created_from': fields.String,
+    'created_by': fields.String,
+    'created_at': TimestampField,
+    'tokens': fields.Integer,
+    'indexing_status': fields.String,
+    'error': fields.String,
+    'enabled': fields.Boolean,
+    'disabled_at': TimestampField,
+    'disabled_by': fields.String,
+    'archived': fields.Boolean,
+    'display_status': fields.String,
+    'word_count': fields.Integer,
+    'hit_count': fields.Integer,
+    'doc_form': fields.String,
+}
+
+document_with_segments_fields = {
+    'id': fields.String,
+    'position': fields.Integer,
+    'data_source_type': fields.String,
+    'data_source_info': fields.Raw(attribute='data_source_info_dict'),
+    'dataset_process_rule_id': fields.String,
+    'name': fields.String,
+    'created_from': fields.String,
+    'created_by': fields.String,
+    'created_at': TimestampField,
+    'tokens': fields.Integer,
+    'indexing_status': fields.String,
+    'error': fields.String,
+    'enabled': fields.Boolean,
+    'disabled_at': TimestampField,
+    'disabled_by': fields.String,
+    'archived': fields.Boolean,
+    'display_status': fields.String,
+    'word_count': fields.Integer,
+    'hit_count': fields.Integer,
+    'completed_segments': fields.Integer,
+    'total_segments': fields.Integer
+}
+
+dataset_and_document_fields = {
+    'dataset': fields.Nested(dataset_fields),
+    'documents': fields.List(fields.Nested(document_fields)),
+    'batch': fields.String
+}
+
+document_status_fields = {
+    'id': fields.String,
+    'indexing_status': fields.String,
+    'processing_started_at': TimestampField,
+    'parsing_completed_at': TimestampField,
+    'cleaning_completed_at': TimestampField,
+    'splitting_completed_at': TimestampField,
+    'completed_at': TimestampField,
+    'paused_at': TimestampField,
+    'error': fields.String,
+    'stopped_at': TimestampField,
+    'completed_segments': fields.Integer,
+    'total_segments': fields.Integer,
+}
+
+document_status_fields_list = {
+    'data': fields.List(fields.Nested(document_status_fields))
+}

api/fields/file_fields.py

@@ -0,0 +1,18 @@
+from flask_restful import fields
+
+from libs.helper import TimestampField
+
+upload_config_fields = {
+    'file_size_limit': fields.Integer,
+    'batch_count_limit': fields.Integer
+}
+
+file_fields = {
+    'id': fields.String,
+    'name': fields.String,
+    'size': fields.Integer,
+    'extension': fields.String,
+    'mime_type': fields.String,
+    'created_by': fields.String,
+    'created_at': TimestampField,
+}

api/fields/hit_testing_fields.py

@@ -0,0 +1,41 @@
+from flask_restful import fields
+
+from libs.helper import TimestampField
+
+document_fields = {
+    'id': fields.String,
+    'data_source_type': fields.String,
+    'name': fields.String,
+    'doc_type': fields.String,
+}
+
+segment_fields = {
+    'id': fields.String,
+    'position': fields.Integer,
+    'document_id': fields.String,
+    'content': fields.String,
+    'answer': fields.String,
+    'word_count': fields.Integer,
+    'tokens': fields.Integer,
+    'keywords': fields.List(fields.String),
+    'index_node_id': fields.String,
+    'index_node_hash': fields.String,
+    'hit_count': fields.Integer,
+    'enabled': fields.Boolean,
+    'disabled_at': TimestampField,
+    'disabled_by': fields.String,
+    'status': fields.String,
+    'created_by': fields.String,
+    'created_at': TimestampField,
+    'indexing_at': TimestampField,
+    'completed_at': TimestampField,
+    'error': fields.String,
+    'stopped_at': TimestampField,
+    'document': fields.Nested(document_fields),
+}
+
+hit_testing_record_fields = {
+    'segment': fields.Nested(segment_fields),
+    'score': fields.Float,
+    'tsne_position': fields.Raw
+}

api/fields/installed_app_fields.py

@@ -0,0 +1,25 @@
+from flask_restful import fields
+
+from libs.helper import TimestampField
+
+app_fields = {
+    'id': fields.String,
+    'name': fields.String,
+    'mode': fields.String,
+    'icon': fields.String,
+    'icon_background': fields.String
+}
+
+installed_app_fields = {
+    'id': fields.String,
+    'app': fields.Nested(app_fields),
+    'app_owner_tenant_id': fields.String,
+    'is_pinned': fields.Boolean,
+    'last_used_at': TimestampField,
+    'editable': fields.Boolean,
+    'uninstallable': fields.Boolean,
+}
+
+installed_app_list_fields = {
+    'installed_apps': fields.List(fields.Nested(installed_app_fields))
+}

api/fields/message_fields.py

@@ -0,0 +1,43 @@
+from flask_restful import fields
+
+from libs.helper import TimestampField
+
+feedback_fields = {
+    'rating': fields.String
+}
+
+retriever_resource_fields = {
+    'id': fields.String,
+    'message_id': fields.String,
+    'position': fields.Integer,
+    'dataset_id': fields.String,
+    'dataset_name': fields.String,
+    'document_id': fields.String,
+    'document_name': fields.String,
+    'data_source_type': fields.String,
+    'segment_id': fields.String,
+    'score': fields.Float,
+    'hit_count': fields.Integer,
+    'word_count': fields.Integer,
+    'segment_position': fields.Integer,
+    'index_node_hash': fields.String,
+    'content': fields.String,
+    'created_at': TimestampField
+}
+
+message_fields = {
+    'id': fields.String,
+    'conversation_id': fields.String,
+    'inputs': fields.Raw,
+    'query': fields.String,
+    'answer': fields.String,
+    'feedback': fields.Nested(feedback_fields, attribute='user_feedback', allow_null=True),
+    'retriever_resources': fields.List(fields.Nested(retriever_resource_fields)),
+    'created_at': TimestampField
+}
+
+message_infinite_scroll_pagination_fields = {
+    'limit': fields.Integer,
+    'has_more': fields.Boolean,
+    'data': fields.List(fields.Nested(message_fields))
+}

api/fields/segment_fields.py

@@ -0,0 +1,32 @@
+from flask_restful import fields
+from libs.helper import TimestampField
+
+segment_fields = {
+    'id': fields.String,
+    'position': fields.Integer,
+    'document_id': fields.String,
+    'content': fields.String,
+    'answer': fields.String,
+    'word_count': fields.Integer,
+    'tokens': fields.Integer,
+    'keywords': fields.List(fields.String),
+    'index_node_id': fields.String,
+    'index_node_hash': fields.String,
+    'hit_count': fields.Integer,
+    'enabled': fields.Boolean,
+    'disabled_at': TimestampField,
+    'disabled_by': fields.String,
+    'status': fields.String,
+    'created_by': fields.String,
+    'created_at': TimestampField,
+    'indexing_at': TimestampField,
+    'completed_at': TimestampField,
+    'error': fields.String,
+    'stopped_at': TimestampField
+}
+
+segment_list_response = {
+    'data': fields.List(fields.Nested(segment_fields)),
+    'has_more': fields.Boolean,
+    'limit': fields.Integer
+}

api/libs/login.py

@@ -1,11 +1,10 @@
 import os
 from functools import wraps
 
-import flask_login
 from flask import current_app
 from flask import g
 from flask import has_request_context
-from flask import request
+from flask import request, session
 from flask_login import user_logged_in
 from flask_login.config import EXEMPT_METHODS
 from werkzeug.exceptions import Unauthorized

api/migrations/versions/2e9819ca5b28_add_tenant_id_in_api_token.py

@@ -0,0 +1,36 @@
+"""add_tenant_id_in_api_token
+
+Revision ID: 2e9819ca5b28
+Revises: 6e2cfb077b04
+Create Date: 2023-09-22 15:41:01.243183
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = '2e9819ca5b28'
+down_revision = 'ab23c11305d4'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('api_tokens', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('tenant_id', postgresql.UUID(), nullable=True))
+        batch_op.create_index('api_token_tenant_idx', ['tenant_id', 'type'], unique=False)
+        batch_op.drop_column('dataset_id')
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('api_tokens', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('dataset_id', postgresql.UUID(), autoincrement=False, nullable=True))
+        batch_op.drop_index('api_token_tenant_idx')
+        batch_op.drop_column('tenant_id')
+
+    # ### end Alembic commands ###

api/migrations/versions/ab23c11305d4_add_dataset_query_variable_at_app_model_.py

@@ -0,0 +1,31 @@
+"""add dataset query variable at app model configs.
+
+Revision ID: ab23c11305d4
+Revises: 6e2cfb077b04
+Create Date: 2023-09-26 12:22:59.044088
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = 'ab23c11305d4'
+down_revision = '6e2cfb077b04'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('app_model_configs', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('dataset_query_variable', sa.String(length=255), nullable=True))
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('app_model_configs', schema=None) as batch_op:
+        batch_op.drop_column('dataset_query_variable')
+
+    # ### end Alembic commands ###

api/models/model.py

@@ -88,6 +88,7 @@ class AppModelConfig(db.Model):
     more_like_this = db.Column(db.Text)
     model = db.Column(db.Text)
     user_input_form = db.Column(db.Text)
+    dataset_query_variable = db.Column(db.String(255))
     pre_prompt = db.Column(db.Text)
     agent_mode = db.Column(db.Text)
     sensitive_word_avoidance = db.Column(db.Text)
@@ -152,6 +153,7 @@ class AppModelConfig(db.Model):
             "sensitive_word_avoidance": self.sensitive_word_avoidance_dict,
             "model": self.model_dict,
             "user_input_form": self.user_input_form_list,
+            "dataset_query_variable": self.dataset_query_variable,
             "pre_prompt": self.pre_prompt,
             "agent_mode": self.agent_mode_dict
         }
@@ -170,6 +172,7 @@ class AppModelConfig(db.Model):
             if model_config.get('sensitive_word_avoidance') else None
         self.model = json.dumps(model_config['model'])
         self.user_input_form = json.dumps(model_config['user_input_form'])
+        self.dataset_query_variable = model_config.get('dataset_query_variable')
         self.pre_prompt = model_config['pre_prompt']
         self.agent_mode = json.dumps(model_config['agent_mode'])
         self.retriever_resource = json.dumps(model_config['retriever_resource']) \
@@ -191,6 +194,7 @@ class AppModelConfig(db.Model):
             sensitive_word_avoidance=self.sensitive_word_avoidance,
             model=self.model,
             user_input_form=self.user_input_form,
+            dataset_query_variable=self.dataset_query_variable,
             pre_prompt=self.pre_prompt,
             agent_mode=self.agent_mode
         )
@@ -625,12 +629,13 @@ class ApiToken(db.Model):
     __table_args__ = (
         db.PrimaryKeyConstraint('id', name='api_token_pkey'),
         db.Index('api_token_app_id_type_idx', 'app_id', 'type'),
-        db.Index('api_token_token_idx', 'token', 'type')
+        db.Index('api_token_token_idx', 'token', 'type'),
+        db.Index('api_token_tenant_idx', 'tenant_id', 'type')
     )
 
     id = db.Column(UUID, server_default=db.text('uuid_generate_v4()'))
     app_id = db.Column(UUID, nullable=True)
-    dataset_id = db.Column(UUID, nullable=True)
+    tenant_id = db.Column(UUID, nullable=True)
     type = db.Column(db.String(16), nullable=False)
     token = db.Column(db.String(255), nullable=False)
     last_used_at = db.Column(db.DateTime, nullable=True)

api/requirements.txt

@@ -49,6 +49,8 @@ huggingface_hub~=0.16.4
 transformers~=4.31.0
 stripe~=5.5.0
 pandas==1.5.3
-xinference==0.4.2
+xinference==0.5.2
 safetensors==0.3.2
 zhipuai==1.0.7
+werkzeug==2.3.7
+pymilvus==2.3.0

api/services/account_service.py

@@ -4,11 +4,12 @@ import json
 import logging
 import secrets
 import uuid
-from datetime import datetime
+from datetime import datetime, timedelta
 from hashlib import sha256
 from typing import Optional
 
-from flask import session
+from werkzeug.exceptions import Forbidden, Unauthorized
+from flask import session, current_app
 from sqlalchemy import func
 
 from events.tenant_event import tenant_was_created
@@ -19,16 +20,82 @@ from services.errors.account import AccountLoginError, CurrentPasswordIncorrectE
 from libs.helper import get_remote_ip
 from libs.password import compare_password, hash_password
 from libs.rsa import generate_key_pair
+from libs.passport import PassportService
 from models.account import *
 from tasks.mail_invite_member_task import send_invite_member_mail_task
 
+def _create_tenant_for_account(account):
+    tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
+
+    TenantService.create_tenant_member(tenant, account, role='owner')
+    account.current_tenant = tenant
+
+    return tenant
+
 
 class AccountService:
 
     @staticmethod
-    def load_user(account_id: int) -> Account:
+    def load_user(user_id: str) -> Account:
         # todo: used by flask_login
-        pass
+        if '.' in user_id:
+            tenant_id, account_id = user_id.split('.')
+        else:
+            account_id = user_id
+
+        account = db.session.query(Account).filter(Account.id == account_id).first()
+
+        if account:
+            if account.status == AccountStatus.BANNED.value or account.status == AccountStatus.CLOSED.value:
+                raise Forbidden('Account is banned or closed.')
+
+            workspace_id = session.get('workspace_id')
+            if workspace_id:
+                tenant_account_join = db.session.query(TenantAccountJoin).filter(
+                    TenantAccountJoin.account_id == account.id,
+                    TenantAccountJoin.tenant_id == workspace_id
+                ).first()
+
+                if not tenant_account_join:
+                    tenant_account_join = db.session.query(TenantAccountJoin).filter(
+                        TenantAccountJoin.account_id == account.id).first()
+
+                    if tenant_account_join:
+                        account.current_tenant_id = tenant_account_join.tenant_id
+                    else:
+                        _create_tenant_for_account(account)
+                    session['workspace_id'] = account.current_tenant_id
+                else:
+                    account.current_tenant_id = workspace_id
+            else:
+                tenant_account_join = db.session.query(TenantAccountJoin).filter(
+                    TenantAccountJoin.account_id == account.id).first()
+                if tenant_account_join:
+                    account.current_tenant_id = tenant_account_join.tenant_id
+                else:
+                    _create_tenant_for_account(account)
+                session['workspace_id'] = account.current_tenant_id
+
+            current_time = datetime.utcnow()
+
+            # update last_active_at when last_active_at is more than 10 minutes ago
+            if current_time - account.last_active_at > timedelta(minutes=10):
+                account.last_active_at = current_time
+                db.session.commit()
+
+        return account
+    
+    @staticmethod
+    def get_account_jwt_token(account):
+        payload = {
+            "user_id": account.id,
+            "exp": datetime.utcnow() + timedelta(days=30),
+            "iss":  current_app.config['EDITION'],
+            "sub": 'Console API Passport',
+        }
+
+        token = PassportService().issue(payload)
+        return token
 
     @staticmethod
     def authenticate(email: str, password: str) -> Account:

api/services/app_model_config_service.py

@@ -81,7 +81,7 @@ class AppModelConfigService:
         return filtered_cp
 
     @staticmethod
-    def validate_configuration(tenant_id: str, account: Account, config: dict) -> dict:
+    def validate_configuration(tenant_id: str, account: Account, config: dict, mode: str) -> dict:
         # opening_statement
         if 'opening_statement' not in config or not config["opening_statement"]:
             config["opening_statement"] = ""
@@ -335,6 +335,9 @@ class AppModelConfigService:
 
                 if not AppModelConfigService.is_dataset_exists(account, tool_item["id"]):
                     raise ValueError("Dataset ID does not exist, please check your permission.")
+        
+        # dataset_query_variable
+        AppModelConfigService.is_dataset_query_variable_valid(config, mode)
 
         # Filter out extra parameters
         filtered_config = {
@@ -351,8 +354,25 @@ class AppModelConfigService:
                 "completion_params": config["model"]["completion_params"]
             },
             "user_input_form": config["user_input_form"],
+            "dataset_query_variable": config.get('dataset_query_variable'),
             "pre_prompt": config["pre_prompt"],
             "agent_mode": config["agent_mode"]
         }
 
         return filtered_config
+    
+    @staticmethod
+    def is_dataset_query_variable_valid(config: dict, mode: str) -> None:
+        # Only check when mode is completion
+        if mode != 'completion':
+            return
+        
+        agent_mode = config.get("agent_mode", {})
+        tools = agent_mode.get("tools", [])
+        dataset_exists = "dataset" in str(tools)
+        
+        dataset_query_variable = config.get("dataset_query_variable")
+
+        if dataset_exists and not dataset_query_variable:
+            raise ValueError("Dataset query variable is required when dataset is exist")
+