add b2

qwen_server/utils.py

@@ -10,6 +10,54 @@ from qwen_agent.log import logger
 from qwen_agent.utils.doc_parser import parse_doc, parse_html_bs
 from qwen_agent.utils.utils import print_traceback, save_text_to_file
 from qwen_server.schema import Record
+from b2sdk.v2 import B2Api
+from b2sdk.v2 import InMemoryAccountInfo
+import hashlib
+from datetime import datetime
+from io import BytesIO
+
+
+class B2Manager():
+    def __init__(self):
+        info = InMemoryAccountInfo()
+        b2_api = B2Api(info)
+        application_key_id = os.environ.get("b2_key_id")
+        application_key = os.environ.get("b2_key")
+        b2_api.authorize_account("production", application_key_id, application_key)
+        self.b2_bucket = b2_api.get_bucket_by_name(os.environ.get("b2_bucket_name"))
+        self.b2_api = b2_api
+        self.file_name = None
+
+    def gen_file_name(self, access_token, url):
+        url_md5 = hashlib.md5(b'%s' % url.encode(encoding='UTF-8')).hexdigest()
+        self.file_name = f"{access_token}/{url_md5}"
+
+    def get(self):
+        in_memory_file = BytesIO()
+        self.b2_bucket.download_file_by_name(self.file_name).save(in_memory_file)
+        # export_file = self.b2_bucket.download_file_by_name(self.file_name)
+        # export_file.save(in_memory_file)
+        in_memory_file.seek(0)
+        return in_memory_file.read()
+
+    def upsert(self, file_path, content):
+        with open(file_path, 'rb') as file_data:
+            # file_info = {'description': ''}
+            self.b2_bucket.upload_bytes(file_data.read(), self.file_name, file_infos=None)
+
+    def delete(self):
+        file_version_info = self.b2_bucket.get_file_info_by_name(self.file_name)
+        self.b2_bucket.hide_file(file_version_info.file_name)
+        # for version in self.b2_bucket.list_file_versions(self.file_name):
+        #     self.b2_bucket.delete_file_version(version.id_, version.file_name)
+
+    def list_files(self, access_token):
+        files = []
+        for file_version_info, folder_name in self.b2_bucket.ls(folder_to_list="access_token/", show_versions=False):
+            # The upload timestamp is in milliseconds, so we divide by 1000 to convert it to seconds
+            upload_timestamp = datetime.fromtimestamp(file_version_info.upload_timestamp / 1000.0)
+            files.append(f"File Name: {file_version_info.file_name}, \nUpload timestamp: {upload_timestamp}, \nMetadata: {file_version_info.file_info}")
+        return files
 
 
 def is_local_path(path):

requirements.txt

@@ -35,4 +35,5 @@ torch==2.0.1
 transformers_stream_generator
 einops
 accelerate
-itsdangerous
+itsdangerous
+b2sdk