创建知识库、上传知识库增加文档id返回

userpj · Nov 4, 2024 · 0edf787 · 0edf787
1 parent 54b4b3b
commit 0edf787
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 9 deletions.
diff --git a/appbuilder/core/console/knowledge_base/data_class.py b/appbuilder/core/console/knowledge_base/data_class.py
@@ -220,6 +220,15 @@ class KnowledgeBaseCreateDocumentsRequest(BaseModel):
         None, description="文档处理选项"
     )
 
+class KnowledgeBaseCreateDocumentsResponse(BaseModel):
+    requestId: str = Field(..., description="请求ID")
+    documentIds: list[str] = Field(..., description="文档ID列表")
+
+
+class KnowledgeBaseUploadDocumentsResponse(BaseModel):
+    requestId: str = Field(..., description="请求ID")
+    documentId: str = Field(..., description="文档ID")
+
 
 class CreateChunkRequest(BaseModel):
     documentId: str = Field(..., description="文档ID")

diff --git a/appbuilder/core/console/knowledge_base/knowledge_base.py b/appbuilder/core/console/knowledge_base/knowledge_base.py
@@ -485,7 +485,7 @@ def create_documents(
         source: data_class.DocumentSource = None,
         processOption: data_class.DocumentProcessOption = None,
         client_token: str = None,
-    ):
+    ) -> data_class.KnowledgeBaseCreateDocumentsResponse:
         r"""
         创建文档
 
@@ -497,7 +497,9 @@ def create_documents(
             client_token (str, optional): 客户端令牌。默认为None，此时会自动生成一个随机UUID作为客户端令牌。
 
         Returns:
-            dict: 响应数据，包含请求ID: requestId
+            KnowledgeBaseCreateDocumentsResponse: 创建知识库文档的响应消息，返回一个KnowledgeBaseCreateDocumentsResponse对象，包含以下属性：
+            - requestId (str): 请求ID
+            - documentIds (list[str]): 文档ID列表
         """
         if self.knowledge_id == None and id == None:
             raise ValueError(
@@ -527,7 +529,8 @@ def create_documents(
         self.http_client.check_console_response(response)
         data = response.json()
 
-        return data
+        resp = data_class.KnowledgeBaseCreateDocumentsResponse(**data)
+        return resp
 
     def get_knowledge_base_list(
         self,
@@ -587,7 +590,7 @@ def upload_documents(
         id: Optional[str] = None,
         processOption: data_class.DocumentProcessOption = None,
         client_token: str = None,
-    ):
+    ) -> data_class.KnowledgeBaseUploadDocumentsResponse:
         r"""
         上传文档
 
@@ -599,7 +602,9 @@ def upload_documents(
             client_token (str, optional): 客户端令牌。默认为None，此时会自动生成一个随机UUID作为客户端令牌。
 
         Returns:
-            dict: 响应数据，包含请求ID: requestId
+            KnowledgeBaseUploadDocumentsResponse: 创建知识库文档的响应消息，返回一个KnowledgeBaseUploadDocumentsResponse对象，包含以下属性：
+            - requestId (str): 请求ID
+            - documentId (str): 文档ID
         """
         if not os.path.exists(file_path):
             raise FileNotFoundError("File {} does not exist".format(file_path))
@@ -635,8 +640,9 @@ def upload_documents(
             self.http_client.check_response_header(response)
             self.http_client.check_console_response(response)
             data = response.json()
+            resp = data_class.KnowledgeBaseUploadDocumentsResponse(**data)
 
-        return data
+        return resp
 
     def create_chunk(
         self,
@@ -656,7 +662,7 @@ def create_chunk(
             CreateChunkResponse: 创建文档块的相应消息, 包含以下属性：
             - id (str): 切片ID
         """
-        
+
         headers = self.http_client.auth_header_v2()
         headers["content-type"] = "application/json"
 

diff --git a/appbuilder/tests/test_knowledge_base.py b/appbuilder/tests/test_knowledge_base.py
@@ -72,7 +72,7 @@ def test_create_knowledge_base(self):
             print("create_knowledge_base函数运行失败{},将调用本地DATASET_ID".format(e))
             knowledge_base_id = os.getenv('DATASET_ID', 'UNKNOWN')
 
-        knowledge.create_documents(
+        create_documents_response = knowledge.create_documents(
             id=knowledge_base_id,
             contentFormat="rawText",
             source=appbuilder.DocumentSource(
@@ -97,8 +97,9 @@ def test_create_knowledge_base(self):
                 knowledgeAugmentation=appbuilder.DocumentChoices(choices=["faq"]),
             ),
         )
+        self.assertIsInstance(create_documents_response.documentIds, list)
 
-        knowledge.upload_documents(
+        upload_documents_response = knowledge.upload_documents(
             id=knowledge_base_id,
             content_format="rawText",
             file_path="./data/qa_appbuilder_client_demo.pdf",
@@ -119,6 +120,7 @@ def test_create_knowledge_base(self):
                 knowledgeAugmentation=appbuilder.DocumentChoices(choices=["faq"]),
             ),
         )
+        self.assertIsInstance(upload_documents_response.documentId, str)
 
         list_res = knowledge.get_documents_list(knowledge_base_id=knowledge_base_id)
         document_id = list_res.data[-1].id