API: Download doc api (#1354)

### What problem does this PR solve?

Adds download_document api

### Type of change


- [x] New Feature (non-breaking change which adds functionality)
This commit is contained in:
cecilia-uu
2024-07-04 16:33:55 +08:00
committed by GitHub
parent daaf6aed50
commit 306108fe0e
4 changed files with 137 additions and 17 deletions

View File

@ -12,12 +12,12 @@
# 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.
import json
import os
import requests
from api.db.services.document_service import DocumentService
from api.settings import RetCode
@ -126,7 +126,22 @@ class RAGFlow:
return response.json()
# ----------------------------download a file-----------------------------------------------------
def download_file(self, dataset_id, document_id):
endpoint = f"{self.dataset_url}/{dataset_id}/documents/{document_id}"
res = requests.get(endpoint, headers=self.authorization_header)
content = res.content # binary data
# decode the binary data
try:
decoded_content = content.decode("utf-8")
json_data = json.loads(decoded_content)
return json_data # message
except json.JSONDecodeError: # binary data
_, document = DocumentService.get_by_id(document_id)
file_path = os.path.join(os.getcwd(), document.name)
with open(file_path, "wb") as file:
file.write(content)
return {"code": RetCode.SUCCESS, "data": content}
# ----------------------------start parsing-----------------------------------------------------
# ----------------------------stop parsing-----------------------------------------------------
@ -144,3 +159,4 @@ class RAGFlow:
# ----------------------------get a specific chunk-----------------------------------------------------
# ----------------------------retrieval test-----------------------------------------------------

View File

@ -1,4 +1,4 @@
API_KEY = 'ImFhMmJhZmUwMmQxNzExZWZhZDdmMzA0M2Q3ZWU1MzdlIg.ZnDsIQ.u-0-_qCRU6a4WICxyAPsjaafyOo'
API_KEY = 'IjJkOGQ4ZDE2MzkyMjExZWZhYTk0MzA0M2Q3ZWU1MzdlIg.ZoUfug.RmqcYyCrlAnLtkzk6bYXiXN3eEY'
HOST_ADDRESS = 'http://127.0.0.1:9380'

View File

@ -3,7 +3,6 @@ from test_sdkbase import TestSdk
from ragflow import RAGFlow
import pytest
from common import API_KEY, HOST_ADDRESS
from api.contants import NAME_LENGTH_LIMIT
class TestFile(TestSdk):
@ -625,8 +624,76 @@ class TestFile(TestSdk):
update_res = ragflow.update_file(created_res_id, doc_id, **params)
assert (update_res["code"] == RetCode.DATA_ERROR and
update_res["message"] == "Illegal value ? for 'template_type' field.")
# ----------------------------download a file-----------------------------------------------------
def test_download_nonexistent_document(self):
"""
Test downloading a document which does not exist.
"""
# create a dataset
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
created_res = ragflow.create_dataset("test_download_nonexistent_document")
created_res_id = created_res["data"]["dataset_id"]
res = ragflow.download_file(created_res_id, "imagination")
assert res["code"] == RetCode.ARGUMENT_ERROR and res["message"] == f"This document 'imagination' cannot be found!"
def test_download_document_in_nonexistent_dataset(self):
"""
Test downloading a document whose dataset is nonexistent.
"""
# create a dataset
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
created_res = ragflow.create_dataset("test_download_nonexistent_document")
created_res_id = created_res["data"]["dataset_id"]
# upload files
file_paths = ["test_data/test.txt"]
uploading_res = ragflow.upload_local_file(created_res_id, file_paths)
# get the doc_id
data = uploading_res["data"][0]
doc_id = data["id"]
# download file
res = ragflow.download_file("imagination", doc_id)
assert res["code"] == RetCode.DATA_ERROR and res["message"] == f"This dataset 'imagination' cannot be found!"
def test_download_document_with_success(self):
"""
Test the downloading of a document with success.
"""
# create a dataset
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
created_res = ragflow.create_dataset("test_download_nonexistent_document")
created_res_id = created_res["data"]["dataset_id"]
# upload files
file_paths = ["test_data/test.txt"]
uploading_res = ragflow.upload_local_file(created_res_id, file_paths)
# get the doc_id
data = uploading_res["data"][0]
doc_id = data["id"]
# download file
with open("test_data/test.txt", "rb") as file:
binary_data = file.read()
res = ragflow.download_file(created_res_id, doc_id)
assert res["code"] == RetCode.SUCCESS and res["data"] == binary_data
def test_download_an_empty_document(self):
"""
Test the downloading of an empty document.
"""
# create a dataset
ragflow = RAGFlow(API_KEY, HOST_ADDRESS)
created_res = ragflow.create_dataset("test_download_nonexistent_document")
created_res_id = created_res["data"]["dataset_id"]
# upload files
file_paths = ["test_data/empty.txt"]
uploading_res = ragflow.upload_local_file(created_res_id, file_paths)
# get the doc_id
data = uploading_res["data"][0]
doc_id = data["id"]
# download file
res = ragflow.download_file(created_res_id, doc_id)
assert res["code"] == RetCode.DATA_ERROR and res["message"] == "This file is empty."
# ----------------------------start parsing-----------------------------------------------------
# ----------------------------stop parsing-----------------------------------------------------