fix(API): fixed swagger docs error in nginx external port

api/apps/__init__.py

@@ -44,7 +44,8 @@
 Request.json = property(lambda self: self.get_json(force=True, silent=True))
 
 # Integrate APIFlask: Flask class -> APIFlask class.
-app = APIFlask(__name__, title=RAG_FLOW_SERVICE_NAME, version=API_VERSION, docs_path=f'/{API_VERSION}/docs')
+app = APIFlask(__name__, title=RAG_FLOW_SERVICE_NAME, version=API_VERSION, docs_path=f'/{API_VERSION}/docs',
+               spec_path=f'/{API_VERSION}/openapi.json')
 # Integrate APIFlask: Use apiflask.HTTPTokenAuth for the HTTP Bearer or API Keys authentication.
 http_token_auth = HTTPTokenAuth()
 

api/apps/apis/datasets.py

@@ -16,7 +16,8 @@
 
 from api.apps import http_token_auth
 from api.apps.services import dataset_service
-from api.utils.api_utils import server_error_response, http_basic_auth_required
+from api.settings import RetCode
+from api.utils.api_utils import server_error_response, http_basic_auth_required, get_json_result
 
 
 @manager.post('')
@@ -58,7 +59,7 @@ def get_dataset_by_id(kb_id):
 @manager.input(dataset_service.SearchDatasetReq, location='query')
 @manager.auth_required(http_token_auth)
 def get_dataset_by_name(query_data):
-    """Query Dataset(Knowledgebase) by Dataset(Knowledgebase) Name."""
+    """Query Dataset(Knowledgebase) by Name."""
     try:
         tenant_id = http_token_auth.current_user.id
         return dataset_service.get_dataset_by_name(tenant_id, query_data["name"])
@@ -94,3 +95,18 @@ def delete_dataset(kb_id):
         return dataset_service.delete_dataset(tenant_id, kb_id)
     except Exception as e:
         return server_error_response(e)
+
+
+@manager.post('/retrieval')
+@manager.input(dataset_service.RetrievalReq, location='json')
+@manager.auth_required(http_token_auth)
+def retrieval_in_dataset(json_data):
+    """Run document retrieval in one or more Datasets(Knowledgebase)."""
+    try:
+        tenant_id = http_token_auth.current_user.id
+        return dataset_service.retrieval_in_dataset(tenant_id, json_data)
+    except Exception as e:
+        if str(e).find("not_found") > 0:
+            return get_json_result(data=False, retmsg=f'No chunk found! Check the chunk status please!',
+                                   retcode=RetCode.DATA_ERROR)
+        return server_error_response(e)

api/apps/apis/documents.py

@@ -22,7 +22,7 @@
 @manager.input(document_service.ChangeDocumentParserReq, location='json')
 @manager.auth_required(http_token_auth)
 def change_document_parser(json_data):
-    """Change document file parser."""
+    """Change document file parsing method."""
     try:
         return document_service.change_document_parser(json_data)
     except Exception as e:

api/apps/services/dataset_service.py

@@ -16,17 +16,19 @@
 
 from apiflask import Schema, fields, validators
 
-from api.db import StatusEnum, FileSource, ParserType
+from api.db import StatusEnum, FileSource, ParserType, LLMType
 from api.db.db_models import File
 from api.db.services import duplicate_name
 from api.db.services.document_service import DocumentService
 from api.db.services.file2document_service import File2DocumentService
 from api.db.services.file_service import FileService
 from api.db.services.knowledgebase_service import KnowledgebaseService
-from api.db.services.user_service import TenantService
-from api.settings import RetCode
+from api.db.services.llm_service import TenantLLMService
+from api.db.services.user_service import TenantService, UserTenantService
+from api.settings import RetCode, retrievaler, kg_retrievaler
 from api.utils import get_uuid
 from api.utils.api_utils import get_json_result, get_data_error_result
+from rag.nlp import keyword_extraction
 
 
 class QueryDatasetReq(Schema):
@@ -48,14 +50,28 @@ class UpdateDatasetReq(Schema):
     kb_id = fields.String(required=True)
     name = fields.String(validate=validators.Length(min=1, max=128))
     description = fields.String(allow_none=True)
-    permission = fields.String(validate=validators.OneOf(['me', 'team']))
+    permission = fields.String(load_default="me", validate=validators.OneOf(['me', 'team']))
     embd_id = fields.String(validate=validators.Length(min=1, max=128))
     language = fields.String(validate=validators.OneOf(['Chinese', 'English']))
     parser_id = fields.String(validate=validators.OneOf([parser_type.value for parser_type in ParserType]))
     parser_config = fields.Dict()
     avatar = fields.String()
 
 
+class RetrievalReq(Schema):
+    kb_id = fields.String(required=True)
+    question = fields.String(required=True)
+    page = fields.Integer(load_default=1)
+    page_size = fields.Integer(load_default=30)
+    doc_ids = fields.List(fields.String())
+    similarity_threshold = fields.Float(load_default=0.0)
+    vector_similarity_weight = fields.Float(load_default=0.3)
+    top_k = fields.Integer(load_default=1024)
+    rerank_id = fields.String()
+    keyword = fields.Boolean(load_default=False)
+    highlight = fields.Boolean(load_default=False)
+
+
 def get_all_datasets(user_id, offset, count, orderby, desc):
     tenants = TenantService.get_joined_tenants_by_user_id(user_id)
     datasets = KnowledgebaseService.get_by_tenant_ids_by_offset(
@@ -159,3 +175,51 @@ def delete_dataset(tenant_id, kb_id):
         return get_data_error_result(
             retmsg="Database error (Knowledgebase removal)!")
     return get_json_result(data=True)
+
+
+def retrieval_in_dataset(tenant_id, json_data):
+    page = json_data["page"]
+    size = json_data["size"]
+    question = json_data["question"]
+    kb_id = json_data["kb_id"]
+    if isinstance(kb_id, str): kb_id = [kb_id]
+    doc_ids = json_data["doc_ids"]
+    similarity_threshold = json_data["similarity_threshold"]
+    vector_similarity_weight = json_data["vector_similarity_weight"]
+    top = json_data["top_k"]
+
+    tenants = UserTenantService.query(user_id=tenant_id)
+    for kid in kb_id:
+        for tenant in tenants:
+            if KnowledgebaseService.query(
+                    tenant_id=tenant.tenant_id, id=kid):
+                break
+        else:
+            return get_json_result(
+                data=False, retmsg=f'Only owner of knowledgebase authorized for this operation.',
+                retcode=RetCode.OPERATING_ERROR)
+
+    e, kb = KnowledgebaseService.get_by_id(kb_id[0])
+    if not e:
+        return get_data_error_result(retmsg="Knowledgebase not found!")
+
+    embd_mdl = TenantLLMService.model_instance(
+        kb.tenant_id, LLMType.EMBEDDING.value, llm_name=kb.embd_id)
+
+    rerank_mdl = None
+    if json_data["rerank_id"]:
+        rerank_mdl = TenantLLMService.model_instance(
+            kb.tenant_id, LLMType.RERANK.value, llm_name=json_data["rerank_id"])
+
+    if json_data["keyword"]:
+        chat_mdl = TenantLLMService.model_instance(kb.tenant_id, LLMType.CHAT)
+        question += keyword_extraction(chat_mdl, question)
+
+    retr = retrievaler if kb.parser_id != ParserType.KG else kg_retrievaler
+    ranks = retr.retrieval(
+        question, embd_mdl, kb.tenant_id, kb_id, page, size, similarity_threshold, vector_similarity_weight, top,
+        doc_ids, rerank_mdl=rerank_mdl, highlight=json_data["highlight"])
+    for c in ranks["chunks"]:
+        if "vector" in c:
+            del c["vector"]
+    return get_json_result(data=ranks)