feat(segment): 新增文档片段管理功能

新增文档片段的创建、更新、删除、启用/禁用等功能，包括相关handler、service、schema的添加和修改。同时优化了分页查询逻辑和向量数据库的批量处理性能。

Author: v833

internal/handler/__init__.py

@@ -4,6 +4,7 @@
 from .upload_file_handler import UploadFileHandler
 from .dataset_handler import DatasetHandler
 from .document_handler import DocumentHandler
+from .segment_handler import SegmentHandler
 
 __all__ = [
     "AppHandler",
@@ -12,4 +13,5 @@
     "UploadFileHandler",
     "DatasetHandler",
     "DocumentHandler",
+    "SegmentHandler",
 ]

internal/handler/dataset_handler.py

@@ -34,6 +34,14 @@ def hit(self, dataset_id: UUID):
         """根据传递的知识库id+检索参数执行召回测试"""
         pass
 
+        q = "alembic==1.15.2"
+        retriever = self.vector_database_service.vector_store.as_retriever(
+            search_type="mmr",
+            search_kwargs={"k": 10},
+        )
+        documents = retriever.invoke(q)
+        return success_json({"documents": [doc.page_content for doc in documents]})
+
     def embeddings_query(self):
         # upload_file = self.db.session.query(UploadFile).get(
         #     "94acfe76-bbad-4d4c-9751-bfcd36bca124"

internal/handler/segment_handler.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+from uuid import UUID
+from injector import inject
+
+from internal.schema.segment_schema import (
+    GetSegmentsWithPageReq,
+    GetSegmentsWithPageResp,
+    GetSegmentResp,
+    UpdateSegmentEnabledReq,
+)
+from internal.service.segment_service import SegmentService
+from pkg.response.response import validate_error_json, success_json, success_message
+from pkg.paginator import PageModel
+
+
+@inject
+@dataclass
+class SegmentHandler:
+
+    segment_service: SegmentService
+
+    def get_segments_with_page(self, dataset_id: UUID, document_id: UUID):
+
+        req = GetSegmentsWithPageReq()
+
+        if not req.validate():
+            return validate_error_json(req.errors)
+
+        segments, paginator = self.segment_service.get_segments_with_page(
+            dataset_id, document_id, req
+        )
+
+        resp = GetSegmentsWithPageResp(many=True)
+
+        return success_json(PageModel(list=resp.dump(segments), paginator=paginator))
+
+    def create_segment(self, dataset_id: UUID, document_id: UUID):
+        pass
+
+    def get_segment(self, dataset_id: UUID, document_id: UUID, segment_id: UUID):
+        """获取指定的文档片段信息详情"""
+        segment = self.segment_service.get_segment(dataset_id, document_id, segment_id)
+        resp = GetSegmentResp()
+        return success_json(resp.dump(segment))
+
+    def update_segment_enabled(
+        self, dataset_id: UUID, document_id: UUID, segment_id: UUID
+    ):
+        """根据传递的信息更新指定的文档片段启用状态"""
+        # 1.提取请求并校验
+        req = UpdateSegmentEnabledReq()
+        if not req.validate():
+            return validate_error_json(req.errors)
+
+        # 2.调用服务更新文档片段的启用状态
+        self.segment_service.update_segment_enabled(
+            dataset_id, document_id, segment_id, req.enabled.data
+        )
+
+        return success_message("修改片段状态成功")

internal/model/dataset.py

@@ -14,7 +14,7 @@
 
 from internal.extension.database_extension import db
 from .app import AppDatasetJoin
-from internal.model.upload_file import UploadFile
+from .upload_file import UploadFile
 
 
 class Dataset(db.Model):
@@ -122,25 +122,27 @@ class Document(db.Model):
     )
 
     @property
-    def upload_file(self):
-        """只读属性，获取文档关联的上传文件"""
+    def upload_file(self) -> "UploadFile":
         return (
             db.session.query(UploadFile)
-            .filter(UploadFile.id == self.upload_file_id)
+            .filter(
+                UploadFile.id == self.upload_file_id,
+            )
             .one_or_none()
         )
 
     @property
-    def process_rule(self):
-        """只读属性，获取文档关联的处理规则"""
+    def process_rule(self) -> "ProcessRule":
         return (
             db.session.query(ProcessRule)
-            .filter(ProcessRule.id == self.process_rule_id)
+            .filter(
+                ProcessRule.id == self.process_rule_id,
+            )
             .one_or_none()
         )
 
     @property
-    def segment_count(self):
+    def segment_count(self) -> int:
         return (
             db.session.query(func.count(Segment.id))
             .filter(
@@ -150,7 +152,7 @@ def segment_count(self):
         )
 
     @property
-    def hit_count(self):
+    def hit_count(self) -> int:
         return (
             db.session.query(func.coalesce(func.sum(Segment.hit_count), 0))
             .filter(
@@ -200,6 +202,10 @@ class Segment(db.Model):
         DateTime, nullable=False, server_default=text("CURRENT_TIMESTAMP(0)")
     )
 
+    @property
+    def document(self) -> "Document":
+        return db.session.query(Document).get(self.document_id)
+
 
 class KeywordTable(db.Model):
     """关键词表模型"""

internal/router/router.py

@@ -6,6 +6,7 @@
 from internal.handler import AppHandler, BuiltinToolHandler, ApiToolHandler
 from internal.handler.dataset_handler import DatasetHandler
 from internal.handler.document_handler import DocumentHandler
+from internal.handler.segment_handler import SegmentHandler
 from internal.handler.upload_file_handler import UploadFileHandler
 
 
@@ -20,6 +21,7 @@ class Router:
     upload_file_handler: UploadFileHandler
     dataset_handler: DatasetHandler
     document_handler: DocumentHandler
+    segment_handler: SegmentHandler
 
     def register_router(self, app: Flask):
         """注册路由"""
@@ -182,5 +184,25 @@ def register_router(self, app: Flask):
             view_func=self.document_handler.delete_document,
         )
 
+        # 文档片段模块
+        bp.add_url_rule(
+            "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments",
+            view_func=self.segment_handler.get_segments_with_page,
+        )
+        bp.add_url_rule(
+            "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments",
+            methods=["POST"],
+            view_func=self.segment_handler.create_segment,
+        )
+        bp.add_url_rule(
+            "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/<uuid:segment_id>",
+            view_func=self.segment_handler.get_segment,
+        )
+        bp.add_url_rule(
+            "/datasets/<uuid:dataset_id>/documents/<uuid:document_id>/segments/<uuid:segment_id>/enabled",
+            methods=["POST"],
+            view_func=self.segment_handler.update_segment_enabled,
+        )
+
         # 在应用上去注册蓝图
         app.register_blueprint(bp)

internal/schema/__init__.py

@@ -9,6 +9,10 @@
     GetDatasetsWithPageReq,
     GetDatasetsWithPageResp,
 )
+from .segment_schema import (
+    GetSegmentsWithPageReq,
+    GetSegmentsWithPageResp,
+)
 
 __all__ = [
     "CompletionReq",
@@ -23,4 +27,6 @@
     "GetDatasetResp",
     "GetDatasetsWithPageReq",
     "GetDatasetsWithPageResp",
+    "GetSegmentsWithPageReq",
+    "GetSegmentsWithPageResp",
 ]

internal/schema/segment_schema.py

@@ -0,0 +1,162 @@
+from flask_wtf import FlaskForm
+from marshmallow import Schema, fields, pre_dump
+from wtforms import StringField, BooleanField
+from wtforms.validators import Optional, ValidationError, DataRequired
+
+from internal.lib.helper import datetime_to_timestamp
+from internal.model import Segment
+from pkg.paginator import PaginatorReq
+from .schema import ListField
+
+
+class GetSegmentsWithPageReq(PaginatorReq):
+    """获取文档片段列表请求"""
+
+    search_word = StringField("search_word", default="", validators=[Optional()])
+
+
+class GetSegmentsWithPageResp(Schema):
+    """获取文档片段列表响应结构"""
+
+    id = fields.UUID(dump_default="")
+    document_id = fields.UUID(dump_default="")
+    dataset_id = fields.UUID(dump_default="")
+    position = fields.Integer(dump_default=0)
+    content = fields.String(dump_default="")
+    keywords = fields.List(fields.String, dump_default=[])
+    character_count = fields.Integer(dump_default=0)
+    token_count = fields.Integer(dump_default=0)
+    hit_count = fields.Integer(dump_default=0)
+    enabled = fields.Boolean(dump_default=False)
+    disabled_at = fields.Integer(dump_default=0)
+    status = fields.String(dump_default="")
+    error = fields.String(dump_default="")
+    updated_at = fields.Integer(dump_default=0)
+    created_at = fields.Integer(dump_default=0)
+
+    @pre_dump
+    def process_data(self, data: Segment, **kwargs):
+        return {
+            "id": data.id,
+            "document_id": data.document_id,
+            "dataset_id": data.dataset_id,
+            "position": data.position,
+            "content": data.content,
+            "keywords": data.keywords,
+            "character_count": data.character_count,
+            "token_count": data.token_count,
+            "hit_count": data.hit_count,
+            "enabled": data.enabled,
+            "disabled_at": datetime_to_timestamp(data.disabled_at),
+            "status": data.status,
+            "error": data.error,
+            "updated_at": datetime_to_timestamp(data.updated_at),
+            "created_at": datetime_to_timestamp(data.created_at),
+        }
+
+
+class GetSegmentResp(Schema):
+    """获取文档详情响应结构"""
+
+    id = fields.UUID(dump_default="")
+    document_id = fields.UUID(dump_default="")
+    dataset_id = fields.UUID(dump_default="")
+    position = fields.Integer(dump_default=0)
+    content = fields.String(dump_default="")
+    keywords = fields.List(fields.String, dump_default=[])
+    character_count = fields.Integer(dump_default=0)
+    token_count = fields.Integer(dump_default=0)
+    hit_count = fields.Integer(dump_default=0)
+    hash = fields.String(dump_default="")
+    enabled = fields.Boolean(dump_default=False)
+    disabled_at = fields.Integer(dump_default=0)
+    status = fields.String(dump_default="")
+    error = fields.String(dump_default="")
+    updated_at = fields.Integer(dump_default=0)
+    created_at = fields.Integer(dump_default=0)
+
+    @pre_dump
+    def process_data(self, data: Segment, **kwargs):
+        return {
+            "id": data.id,
+            "document_id": data.document_id,
+            "dataset_id": data.dataset_id,
+            "position": data.position,
+            "content": data.content,
+            "keywords": data.keywords,
+            "character_count": data.character_count,
+            "token_count": data.token_count,
+            "hit_count": data.hit_count,
+            "hash": data.hash,
+            "enabled": data.enabled,
+            "disabled_at": datetime_to_timestamp(data.disabled_at),
+            "status": data.status,
+            "error": data.error,
+            "updated_at": datetime_to_timestamp(data.updated_at),
+            "created_at": datetime_to_timestamp(data.created_at),
+        }
+
+
+class UpdateSegmentEnabledReq(FlaskForm):
+    """更新文档片段启用状态请求"""
+
+    enabled = BooleanField("enabled")
+
+    def validate_enabled(self, field: BooleanField) -> None:
+        """校验文档启用状态enabled"""
+        if not isinstance(field.data, bool):
+            raise ValidationError("enabled状态不能为空且必须为布尔值")
+
+
+class CreateSegmentReq(FlaskForm):
+    """创建文档片段请求结构"""
+
+    content = StringField("content", validators=[DataRequired("片段内容不能为空")])
+    keywords = ListField("keywords")
+
+    def validate_keywords(self, field: ListField):
+        """校验关键词列表，涵盖长度不能为空，默认为值为空列表"""
+        # 1.校验数据类型+非空
+        if field.data is None:
+            field.data = []
+        if not isinstance(field.data, list):
+            raise ValidationError("关键词列表格式必须是数组")
+
+        # 2.校验数据的长度，最长不能超过10个关键词
+        if len(field.data) > 10:
+            raise ValidationError("关键词长度范围数量在1-10")
+
+        # 3.循环校验关键词信息，关键词必须是字符串
+        for keyword in field.data:
+            if not isinstance(keyword, str):
+                raise ValidationError("关键词必须是字符串")
+
+        # 4.删除重复数据并更新
+        field.data = list(dict.fromkeys(field.data))
+
+
+class UpdateSegmentReq(FlaskForm):
+    """更新文档片段请求"""
+
+    content = StringField("content", validators=[DataRequired("片段内容不能为空")])
+    keywords = ListField("keywords")
+
+    def validate_keywords(self, field: ListField):
+        """校验关键词列表，涵盖长度不能为空，默认为值为空列表"""
+        # 1.校验数据类型+非空
+        if field.data is None:
+            field.data = []
+        if not isinstance(field.data, list):
+            raise ValidationError("关键词列表格式必须是数组")
+
+        # 2.校验数据的长度，最长不能超过10个关键词
+        if len(field.data) > 10:
+            raise ValidationError("关键词长度范围数量在1-10")
+
+        # 3.循环校验关键词信息，关键词必须是字符串
+        for keyword in field.data:
+            if not isinstance(keyword, str):
+                raise ValidationError("关键词必须是字符串")
+
+        # 4.删除重复数据并更新
+        field.data = list(dict.fromkeys(field.data))

internal/service/__init__.py

@@ -12,6 +12,7 @@
 from .indexing_service import IndexingService
 from .keyword_table_service import KeywordTableService
 from .process_rule_service import ProcessRuleService
+from .segment_service import SegmentService
 
 __all__ = [
     "BuiltinToolService",
@@ -28,4 +29,5 @@
     "IndexingService",
     "KeywordTableService",
     "ProcessRuleService",
+    "SegmentService",
 ]