From 37f7d5732aac8770a52924fee3d71dab5ecbab12 Mon Sep 17 00:00:00 2001 From: jyong <718720800@qq.com> Date: Wed, 18 Sep 2024 15:29:30 +0800 Subject: [PATCH] external knowledge api --- api/controllers/console/datasets/external.py | 4 +- .../console/datasets/test_external.py | 49 ++++++++++++ ...18_0659-ec3df697ebbb_external_knowledge.py | 74 +++++++++++++++++++ api/models/dataset.py | 2 +- api/services/external_knowledge_service.py | 17 +++++ 5 files changed, 143 insertions(+), 3 deletions(-) create mode 100644 api/controllers/console/datasets/test_external.py create mode 100644 api/migrations/versions/2024_09_18_0659-ec3df697ebbb_external_knowledge.py diff --git a/api/controllers/console/datasets/external.py b/api/controllers/console/datasets/external.py index 92ce1ffd2e..397d602d28 100644 --- a/api/controllers/console/datasets/external.py +++ b/api/controllers/console/datasets/external.py @@ -68,7 +68,7 @@ class ExternalApiTemplateListApi(Resource): ) parser.add_argument( "settings", - type=list, + type=dict, location="json", nullable=False, required=True, @@ -126,7 +126,7 @@ class ExternalApiTemplateApi(Resource): ) parser.add_argument( "settings", - type=list, + type=dict, location="json", nullable=False, required=True, diff --git a/api/controllers/console/datasets/test_external.py b/api/controllers/console/datasets/test_external.py new file mode 100644 index 0000000000..7772597839 --- /dev/null +++ b/api/controllers/console/datasets/test_external.py @@ -0,0 +1,49 @@ +from flask import request +from flask_login import current_user +from flask_restful import Resource, marshal, reqparse +from werkzeug.exceptions import Forbidden, NotFound + +import services +from controllers.console import api +from controllers.console.app.error import ProviderNotInitializeError +from controllers.console.datasets.error import DatasetNameDuplicateError +from controllers.console.setup import setup_required +from controllers.console.wraps import account_initialization_required +from fields.dataset_fields import dataset_detail_fields +from libs.login import login_required +from services.external_knowledge_service import ExternalDatasetService + +class TestExternalApi(Resource): + @setup_required + @login_required + @account_initialization_required + def post(self): + parser = reqparse.RequestParser() + parser.add_argument( + "top_k", + nullable=False, + required=True, + type=int, + ) + parser.add_argument( + "score_threshold", + nullable=False, + required=True, + type=float, + ) + args = parser.parse_args() + result = ExternalDatasetService.test_external_knowledge_retrival( + args["top_k"], args["score_threshold"] + ) + response = { + "data": [item.to_dict() for item in api_templates], + "has_more": len(api_templates) == limit, + "limit": limit, + "total": total, + "page": page, + } + return response, 200 + + + +api.add_resource(TestExternalApi, "/dify/external-knowledge/retrival-documents") diff --git a/api/migrations/versions/2024_09_18_0659-ec3df697ebbb_external_knowledge.py b/api/migrations/versions/2024_09_18_0659-ec3df697ebbb_external_knowledge.py new file mode 100644 index 0000000000..a5b7d47e67 --- /dev/null +++ b/api/migrations/versions/2024_09_18_0659-ec3df697ebbb_external_knowledge.py @@ -0,0 +1,74 @@ +"""external_knowledge + +Revision ID: ec3df697ebbb +Revises: 675b5321501b +Create Date: 2024-09-18 06:59:54.048478 + +""" +from alembic import op +import models as models +import sqlalchemy as sa +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = 'ec3df697ebbb' +down_revision = '675b5321501b' +branch_labels = None +depends_on = None + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + op.create_table('external_api_templates', + sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False), + sa.Column('name', sa.String(length=255), nullable=False), + sa.Column('description', sa.String(length=255), nullable=False), + sa.Column('tenant_id', models.types.StringUUID(), nullable=False), + sa.Column('settings', sa.Text(), nullable=True), + sa.Column('created_by', models.types.StringUUID(), nullable=False), + sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP(0)'), nullable=False), + sa.Column('updated_by', models.types.StringUUID(), nullable=True), + sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP(0)'), nullable=False), + sa.PrimaryKeyConstraint('id', name='external_api_template_pkey') + ) + with op.batch_alter_table('external_api_templates', schema=None) as batch_op: + batch_op.create_index('external_api_templates_name_idx', ['name'], unique=False) + batch_op.create_index('external_api_templates_tenant_idx', ['tenant_id'], unique=False) + + op.create_table('external_knowledge_bindings', + sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False), + sa.Column('tenant_id', models.types.StringUUID(), nullable=False), + sa.Column('external_api_template_id', models.types.StringUUID(), nullable=False), + sa.Column('dataset_id', models.types.StringUUID(), nullable=False), + sa.Column('external_knowledge_id', sa.Text(), nullable=False), + sa.Column('created_by', models.types.StringUUID(), nullable=False), + sa.Column('created_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP(0)'), nullable=False), + sa.Column('updated_by', models.types.StringUUID(), nullable=True), + sa.Column('updated_at', sa.DateTime(), server_default=sa.text('CURRENT_TIMESTAMP(0)'), nullable=False), + sa.PrimaryKeyConstraint('id', name='external_knowledge_bindings_pkey') + ) + with op.batch_alter_table('external_knowledge_bindings', schema=None) as batch_op: + batch_op.create_index('external_knowledge_bindings_dataset_idx', ['dataset_id'], unique=False) + batch_op.create_index('external_knowledge_bindings_external_api_template_idx', ['external_api_template_id'], unique=False) + batch_op.create_index('external_knowledge_bindings_external_knowledge_idx', ['external_knowledge_id'], unique=False) + batch_op.create_index('external_knowledge_bindings_tenant_idx', ['tenant_id'], unique=False) + + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + + with op.batch_alter_table('external_knowledge_bindings', schema=None) as batch_op: + batch_op.drop_index('external_knowledge_bindings_tenant_idx') + batch_op.drop_index('external_knowledge_bindings_external_knowledge_idx') + batch_op.drop_index('external_knowledge_bindings_external_api_template_idx') + batch_op.drop_index('external_knowledge_bindings_dataset_idx') + + op.drop_table('external_knowledge_bindings') + with op.batch_alter_table('external_api_templates', schema=None) as batch_op: + batch_op.drop_index('external_api_templates_tenant_idx') + batch_op.drop_index('external_api_templates_name_idx') + + op.drop_table('external_api_templates') + # ### end Alembic commands ### diff --git a/api/models/dataset.py b/api/models/dataset.py index 4221aa4bb1..7fb217f7c3 100644 --- a/api/models/dataset.py +++ b/api/models/dataset.py @@ -719,7 +719,7 @@ class ExternalApiTemplates(db.Model): 'description': self.description, 'settings': self.settings_dict, 'created_by': self.created_by, - 'created_at': self.created_at, + 'created_at': self.created_at.isoformat(), } @property diff --git a/api/services/external_knowledge_service.py b/api/services/external_knowledge_service.py index 850e77b8ad..f189aeacb8 100644 --- a/api/services/external_knowledge_service.py +++ b/api/services/external_knowledge_service.py @@ -289,3 +289,20 @@ class ExternalDatasetService: "params": external_retrival_parameters, } response = ExternalDatasetService.process_external_api(ApiTemplateSetting(**api_template_setting), None) + + + @staticmethod + def test_external_knowledge_retrival( + top_k: int, score_threshold: float + ): + api_template_setting = { + "url": f"{settings.get('endpoint')}/dify/external-knowledge/retrival-documents", + "request_method": "post", + "headers": settings.get("headers"), + "params": { + "top_k": top_k, + "score_threshold": score_threshold, + }, + } + response = ExternalDatasetService.process_external_api(ApiTemplateSetting(**api_template_setting), None) + return response.json() \ No newline at end of file