89 lines
3.9 KiB
Python
89 lines
3.9 KiB
Python
|
#
|
||
|
|
# Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
|
||
|
|
#
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
#
|
||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
#
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
#
|
||
|
|
from typing import Optional, List, Dict, Any
|
||
|
|
from pydantic import BaseModel, Field
|
||
|
|
|
||
|
|
|
||
|
|
class ListChunkRequest(BaseModel):
|
||
|
|
"""列出块请求模型"""
|
||
|
|
doc_id: str = Field(..., description="文档ID")
|
||
|
|
page: Optional[int] = Field(1, description="页码")
|
||
|
|
size: Optional[int] = Field(30, description="每页大小")
|
||
|
|
keywords: Optional[str] = Field("", description="关键词")
|
||
|
|
available_int: Optional[int] = Field(None, description="可用性状态")
|
||
|
|
|
||
|
|
|
||
|
|
class GetChunkRequest(BaseModel):
|
||
|
|
"""获取块请求模型"""
|
||
|
|
chunk_id: str = Field(..., description="块ID")
|
||
|
|
|
||
|
|
|
||
|
|
class SetChunkRequest(BaseModel):
|
||
|
|
"""设置块请求模型"""
|
||
|
|
doc_id: str = Field(..., description="文档ID")
|
||
|
|
chunk_id: str = Field(..., description="块ID")
|
||
|
|
content_with_weight: str = Field(..., description="带权重的内容")
|
||
|
|
important_kwd: Optional[List[str]] = Field(None, description="重要关键词")
|
||
|
|
question_kwd: Optional[List[str]] = Field(None, description="问题关键词")
|
||
|
|
tag_kwd: Optional[str] = Field(None, description="标签关键词")
|
||
|
|
tag_feas: Optional[Any] = Field(None, description="标签特征")
|
||
|
|
available_int: Optional[int] = Field(None, description="可用性状态")
|
||
|
|
|
||
|
|
|
||
|
|
class SwitchChunkRequest(BaseModel):
|
||
|
|
"""切换块状态请求模型"""
|
||
|
|
chunk_ids: List[str] = Field(..., description="块ID列表")
|
||
|
|
available_int: int = Field(..., description="可用性状态")
|
||
|
|
doc_id: str = Field(..., description="文档ID")
|
||
|
|
|
||
|
|
|
||
|
|
class RemoveChunkRequest(BaseModel):
|
||
|
|
"""删除块请求模型"""
|
||
|
|
chunk_ids: List[str] = Field(..., description="块ID列表")
|
||
|
|
doc_id: str = Field(..., description="文档ID")
|
||
|
|
|
||
|
|
|
||
|
|
class CreateChunkRequest(BaseModel):
|
||
|
|
"""创建块请求模型"""
|
||
|
|
doc_id: str = Field(..., description="文档ID")
|
||
|
|
content_with_weight: str = Field(..., description="带权重的内容")
|
||
|
|
important_kwd: Optional[List[str]] = Field([], description="重要关键词")
|
||
|
|
question_kwd: Optional[List[str]] = Field([], description="问题关键词")
|
||
|
|
tag_feas: Optional[Any] = Field(None, description="标签特征")
|
||
|
|
|
||
|
|
|
||
|
|
class RetrievalTestRequest(BaseModel):
|
||
|
|
"""检索测试请求模型"""
|
||
|
|
kb_id: List[str] = Field(..., description="知识库ID列表")
|
||
|
|
question: str = Field(..., description="问题")
|
||
|
|
page: Optional[int] = Field(1, description="页码")
|
||
|
|
size: Optional[int] = Field(30, description="每页大小")
|
||
|
|
doc_ids: Optional[List[str]] = Field([], description="文档ID列表")
|
||
|
|
use_kg: Optional[bool] = Field(False, description="是否使用知识图谱")
|
||
|
|
top_k: Optional[int] = Field(1024, description="返回数量")
|
||
|
|
cross_languages: Optional[List[str]] = Field([], description="跨语言列表")
|
||
|
|
search_id: Optional[str] = Field("", description="搜索ID")
|
||
|
|
rerank_id: Optional[str] = Field(None, description="重排序ID")
|
||
|
|
keyword: Optional[bool] = Field(False, description="是否使用关键词")
|
||
|
|
similarity_threshold: Optional[float] = Field(0.0, description="相似度阈值")
|
||
|
|
vector_similarity_weight: Optional[float] = Field(0.3, description="向量相似度权重")
|
||
|
|
highlight: Optional[bool] = Field(None, description="是否高亮")
|
||
|
|
|
||
|
|
|
||
|
|
class KnowledgeGraphRequest(BaseModel):
|
||
|
|
"""知识图谱请求模型"""
|
||
|
|
doc_id: str = Field(..., description="文档ID")
|