feat: add HitTest API for knowledge base query testing and update SyncWeb API to use knowledge_id

liuruibin · liuruibin · commit b6f5a8378a13 · 2025-05-08T16:49:12.000+08:00
diff --git a/apps/knowledge/api/knowledge.py b/apps/knowledge/api/knowledge.py
@@ -5,7 +5,7 @@
 from common.result import ResultSerializer, DefaultResultSerializer
 from knowledge.serializers.common import GenerateRelatedSerializer
 from knowledge.serializers.knowledge import KnowledgeBaseCreateRequest, KnowledgeModelSerializer, KnowledgeEditRequest, \
-    KnowledgeWebCreateRequest
+    KnowledgeWebCreateRequest, HitTestSerializer
 
 
 class KnowledgeCreateResponse(ResultSerializer):
@@ -238,3 +238,9 @@ class GenerateRelatedAPI(SyncWebAPI):
     @staticmethod
     def get_request():
         return GenerateRelatedSerializer
+
+
+class HitTestAPI(SyncWebAPI):
+    @staticmethod
+    def get_request():
+        return HitTestSerializer
diff --git a/apps/knowledge/serializers/knowledge.py b/apps/knowledge/serializers/knowledge.py
@@ -14,6 +14,7 @@
 from django.utils.translation import gettext_lazy as _
 from rest_framework import serializers
 
+from common.config.embedding_config import VectorStore
 from common.db.search import native_search, get_dynamics_model, native_page_search
 from common.db.sql_execute import select_list
 from common.event import ListenerManagement
@@ -22,9 +23,9 @@
 from common.utils.fork import Fork, ChildLink
 from common.utils.split_model import get_split_model
 from knowledge.models import Knowledge, KnowledgeScope, KnowledgeType, Document, Paragraph, Problem, \
-    ProblemParagraphMapping, ApplicationKnowledgeMapping, TaskType, State
+    ProblemParagraphMapping, ApplicationKnowledgeMapping, TaskType, State, SearchMode
 from knowledge.serializers.common import ProblemParagraphManage, get_embedding_model_id_by_knowledge_id, MetaSerializer, \
-    GenerateRelatedSerializer
+    GenerateRelatedSerializer, get_embedding_model_by_knowledge_id, list_paragraph
 from knowledge.serializers.document import DocumentSerializers
 from knowledge.task.embedding import embedding_by_knowledge, delete_embedding_by_knowledge
 from knowledge.task.generate import generate_related_by_knowledge_id
@@ -79,6 +80,14 @@ def is_valid(self, *, knowledge: Knowledge = None):
             valid_class = knowledge_meta_valid_map.get(knowledge.type)
             valid_class(data=self.data.get('meta')).is_valid(raise_exception=True)
 
+class HitTestSerializer(serializers.Serializer):
+    query_text = serializers.CharField(required=True, label=_('query text'))
+    top_number = serializers.IntegerField(required=True, max_value=10000, min_value=1, label=_("top number"))
+    similarity = serializers.FloatField(required=True, max_value=2, min_value=0, label=_('similarity'))
+    search_mode = serializers.CharField(required=True, label=_('search mode'), validators=[
+        validators.RegexValidator(regex=re.compile("^embedding|keywords|blend$"),
+                                  message=_('The type only supports embedding|keywords|blend'), code=500)
+    ])
 
 class KnowledgeSerializer(serializers.Serializer):
     class Query(serializers.Serializer):
@@ -152,7 +161,7 @@ def generate_related(self, instance: Dict, with_valid=True):
             if with_valid:
                 self.is_valid(raise_exception=True)
                 GenerateRelatedSerializer(data=instance).is_valid(raise_exception=True)
-            knowledge_id = self.data.get('id')
+            knowledge_id = self.data.get('knowledge_id')
             model_id = instance.get("model_id")
             prompt = instance.get("prompt")
             state_list = instance.get('state_list')
@@ -382,15 +391,16 @@ def save_web(self, instance: Dict, with_valid=True):
             return {**KnowledgeModelSerializer(knowledge).data, 'document_list': []}
 
     class SyncWeb(serializers.Serializer):
-        id = serializers.CharField(required=True, label=_('knowledge id'))
+        workspace_id = serializers.CharField(required=True, label=_('workspace id'))
+        knowledge_id = serializers.CharField(required=True, label=_('knowledge id'))
         user_id = serializers.UUIDField(required=False, label=_('user id'))
         sync_type = serializers.CharField(required=True, label=_('sync type'), validators=[
             validators.RegexValidator(regex=re.compile("^replace|complete$"),
                                       message=_('The synchronization type only supports:replace|complete'), code=500)])
 
         def is_valid(self, *, raise_exception=False):
             super().is_valid(raise_exception=True)
-            first = QuerySet(Knowledge).filter(id=self.data.get("id")).first()
+            first = QuerySet(Knowledge).filter(id=self.data.get("knowledge_id")).first()
             if first is None:
                 raise AppApiException(300, _('id does not exist'))
             if first.type != KnowledgeType.WEB:
@@ -400,7 +410,7 @@ def sync(self, with_valid=True):
             if with_valid:
                 self.is_valid(raise_exception=True)
             sync_type = self.data.get('sync_type')
-            knowledge_id = self.data.get('id')
+            knowledge_id = self.data.get('knowledge_id')
             knowledge = QuerySet(Knowledge).get(id=knowledge_id)
             self.__getattribute__(sync_type + '_sync')(knowledge)
             return True
@@ -454,6 +464,52 @@ def complete_sync(self, knowledge):
             # 删除段落
             QuerySet(Paragraph).filter(knowledge=knowledge).delete()
             # 删除向量
-            delete_embedding_by_knowledge(self.data.get('id'))
+            delete_embedding_by_knowledge(self.data.get('knowledge_id'))
             # 同步
             self.replace_sync(knowledge)
+
+    class HitTest(serializers.Serializer):
+        workspace_id = serializers.CharField(required=True, label=_('workspace id'))
+        knowledge_id = serializers.UUIDField(required=True, label=_("id"))
+        user_id = serializers.UUIDField(required=False, label=_('user id'))
+        query_text = serializers.CharField(required=True, label=_('query text'))
+        top_number = serializers.IntegerField(required=True, max_value=10000, min_value=1, label=_("top number"))
+        similarity = serializers.FloatField(required=True, max_value=2, min_value=0, label=_('similarity'))
+        search_mode = serializers.CharField(required=True, label=_('search mode'), validators=[
+            validators.RegexValidator(regex=re.compile("^embedding|keywords|blend$"),
+                                      message=_('The type only supports embedding|keywords|blend'), code=500)
+        ])
+
+        def is_valid(self, *, raise_exception=True):
+            super().is_valid(raise_exception=True)
+            if not QuerySet(Knowledge).filter(id=self.data.get("knowledge_id")).exists():
+                raise AppApiException(300, _('id does not exist'))
+
+        def hit_test(self):
+            self.is_valid()
+            vector = VectorStore.get_embedding_vector()
+            exclude_document_id_list = [
+                str(
+                    document.id
+                ) for document in QuerySet(Document).filter(knowledge_id=self.data.get('knowledge_id'), is_active=False)
+            ]
+            model = get_embedding_model_by_knowledge_id(self.data.get('knowledge_id'))
+            # 向量库检索
+            hit_list = vector.hit_test(
+                self.data.get('query_text'),
+                [self.data.get('knowledge_id')],
+                exclude_document_id_list,
+                self.data.get('top_number'),
+                self.data.get('similarity'),
+                SearchMode(self.data.get('search_mode')),
+                model
+            )
+            hit_dict = reduce(lambda x, y: {**x, **y}, [{hit.get('paragraph_id'): hit} for hit in hit_list], {})
+            p_list = list_paragraph([h.get('paragraph_id') for h in hit_list])
+            return [
+                {
+                    **p,
+                    'similarity': hit_dict.get(p.get('id')).get('similarity'),
+                    'comprehensive_score': hit_dict.get(p.get('id')).get('comprehensive_score')
+                } for p in p_list
+            ]
diff --git a/apps/knowledge/urls.py b/apps/knowledge/urls.py
@@ -10,6 +10,7 @@
     path('workspace/<str:workspace_id>/knowledge/<str:knowledge_id>', views.KnowledgeView.Operate.as_view()),
     path('workspace/<str:workspace_id>/knowledge/<str:knowledge_id>/sync', views.KnowledgeView.SyncWeb.as_view()),
     path('workspace/<str:workspace_id>/knowledge/<str:knowledge_id>/generate_related', views.KnowledgeView.GenerateRelated.as_view()),
+    path('workspace/<str:workspace_id>/knowledge/<str:knowledge_id>/hit_test', views.KnowledgeView.HitTest.as_view()),
     path('workspace/<str:workspace_id>/knowledge/<str:knowledge_id>/document', views.DocumentView.as_view()),
     path('workspace/<str:workspace_id>/knowledge/<str:knowledge_id>/document/split', views.DocumentView.Split.as_view()),
     path('workspace/<str:workspace_id>/knowledge/<str:knowledge_id>/document/split_pattern', views.DocumentView.SplitPattern.as_view()),
diff --git a/apps/knowledge/views/knowledge.py b/apps/knowledge/views/knowledge.py
@@ -8,7 +8,7 @@
 from common.constants.permission_constants import PermissionConstants
 from common.result import result
 from knowledge.api.knowledge import KnowledgeBaseCreateAPI, KnowledgeWebCreateAPI, KnowledgeTreeReadAPI, \
-    KnowledgeEditAPI, KnowledgeReadAPI, KnowledgePageAPI, SyncWebAPI, GenerateRelatedAPI
+    KnowledgeEditAPI, KnowledgeReadAPI, KnowledgePageAPI, SyncWebAPI, GenerateRelatedAPI, HitTestAPI
 from knowledge.serializers.knowledge import KnowledgeSerializer
 
 
@@ -128,11 +128,38 @@ def put(self, request: Request, workspace_id: str, knowledge_id: str):
                 data={
                     'workspace_id': workspace_id,
                     'sync_type': request.query_params.get('sync_type'),
-                    'id': knowledge_id,
+                    'knowledge_id': knowledge_id,
                     'user_id': str(request.user.id)
                 }
             ).sync())
 
+    class HitTest(APIView):
+        authentication_classes = [TokenAuth]
+
+        @extend_schema(
+            methods=['PUT'],
+            summary=_('Hit test list'),
+            description=_('Hit test list'),
+            operation_id=_('Hit test list'),
+            parameters=HitTestAPI.get_parameters(),
+            request=HitTestAPI.get_request(),
+            responses=HitTestAPI.get_response(),
+            tags=[_('Knowledge Base')]
+        )
+        @has_permissions(PermissionConstants.KNOWLEDGE_EDIT.get_workspace_permission())
+        def put(self, request: Request, workspace_id: str, knowledge_id: str):
+            return result.success(KnowledgeSerializer.HitTest(
+                data={
+                    'workspace_id': workspace_id,
+                    'knowledge_id': knowledge_id,
+                    'user_id': request.user.id,
+                    "query_text": request.query_params.get("query_text"),
+                    "top_number": request.query_params.get("top_number"),
+                    'similarity': request.query_params.get('similarity'),
+                    'search_mode': request.query_params.get('search_mode')
+                }
+            ).hit_test())
+
     class GenerateRelated(APIView):
         authentication_classes = [TokenAuth]
 
