Update remote_uploader.py

remote_uploader.py

@@ -1,16 +1,31 @@
-import requests
-import argparse
+import asyncio
+import websockets
+import json
+import logging
+import sys
+import base64
 import os
-import time
+import argparse
+
+# 配置日志
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
-class FileUploader:
+class WebSocketFileUploader:
     def __init__(self, server_url, api_key, space_id):
-        self.server_url = server_url
+        # 将HTTP URL转换为WebSocket URL
+        if server_url.startswith('http://'):
+            self.ws_url = server_url.replace('http://', 'ws://') + '/ws/upload'
+        elif server_url.startswith('https://'):
+            self.ws_url = server_url.replace('https://', 'wss://') + '/ws/upload'
+        else:
+            self.ws_url = f"ws://{server_url}/ws/upload"
+            
         self.api_key = api_key
         self.space_id = space_id
         
-    def upload_file(self, file_path):
-        """上传单个文件"""
+    async def upload_file(self, file_path):
+        """使用WebSocket上传单个文件"""
         if not os.path.exists(file_path):
             print(f"文件不存在: {file_path}")
             return False
@@ -18,65 +33,57 @@ class FileUploader:
         filename = os.path.basename(file_path)
         print(f"正在上传文件: {filename}")
         
-        # 构建完整的 API 端点 URL
-        upload_url = f"{self.server_url.rstrip('/')}/api/remote_upload"
-        
-        # 准备请求头和数据
-        headers = {
-            'X-API-Key': self.api_key
-        }
-        data = {
-            'space_id': self.space_id
-        }
-        
-        # 读取文件内容并上传
         try:
+            # 读取文件内容并编码为base64
             with open(file_path, 'rb') as f:
-                files = {
-                    'file': (filename, f)
+                file_content = f.read()
+                file_b64 = base64.b64encode(file_content).decode('utf-8')
+            
+            # 连接到WebSocket服务器
+            async with websockets.connect(self.ws_url) as websocket:
+                # 发送认证和文件数据
+                upload_msg = {
+                    "type": "file_upload",
+                    "api_key": self.api_key,
+                    "space_id": self.space_id,
+                    "filename": filename,
+                    "content": file_b64
                 }
                 
-                # 发送 POST 请求
-                response = requests.post(upload_url, headers=headers, data=data, files=files)
+                await websocket.send(json.dumps(upload_msg))
+                logger.debug(f"发送文件上传请求: {filename}")
+                
+                # 等待服务器响应
+                response = await websocket.recv()
+                data = json.loads(response)
                 
-                # 打印服务器响应
-                print(f"服务器响应 ({response.status_code}):")
-                try:
-                    result = response.json()
-                    print(result)
-                    if response.status_code == 200:
-                        print(f"✅ 文件 '{filename}' 上传成功!")
-                        return True
-                    else:
-                        print(f"❌ 上传失败: {result.get('error', '未知错误')}")
-                        return False
-                except requests.exceptions.JSONDecodeError:
-                    print(response.text)
+                if data.get("type") == "success":
+                    print(f"✅ 文件 '{filename}' 上传成功!")
+                    return True
+                else:
+                    print(f"❌ 上传失败: {data.get('error', '未知错误')}")
                     return False
                     
-        except IOError as e:
-            print(f"读取文件时出错: {e}")
-            return False
-        except requests.exceptions.RequestException as e:
-            print(f"请求时出错: {e}")
+        except Exception as e:
+            print(f"上传文件 {filename} 时出错: {e}")
             return False
 
-def upload_directory_once(watch_dir, server_url, api_key, space_id):
+async def upload_directory_once(upload_dir, server_url, api_key, space_id):
     """一次性扫描并上传目录中的所有文件"""
-    if not os.path.exists(watch_dir):
-        print(f"目录不存在: {watch_dir}")
+    if not os.path.exists(upload_dir):
+        print(f"目录不存在: {upload_dir}")
         return
     
-    print(f"🔍 开始扫描目录: {watch_dir}")
+    print(f"🔍 开始扫描目录: {upload_dir}")
     print(f"📡 服务器地址: {server_url}")
     print(f"🔑 Space ID: {space_id}")
     print("-" * 50)
     
-    uploader = FileUploader(server_url, api_key, space_id)
+    uploader = WebSocketFileUploader(server_url, api_key, space_id)
     
     # 获取所有文件
     all_files = []
-    for root, dirs, files in os.walk(watch_dir):
+    for root, dirs, files in os.walk(upload_dir):
         for file in files:
             file_path = os.path.join(root, file)
             if os.path.isfile(file_path):
@@ -93,12 +100,12 @@ def upload_directory_once(watch_dir, server_url, api_key, space_id):
     
     for file_path in all_files:
         try:
-            if uploader.upload_file(file_path):
+            if await uploader.upload_file(file_path):
                 success_count += 1
             else:
                 failed_count += 1
             # 稍微延迟一下，避免服务器压力过大
-            time.sleep(0.5)
+            await asyncio.sleep(0.5)
         except Exception as e:
             print(f"上传文件 {file_path} 时发生异常: {e}")
             failed_count += 1
@@ -111,14 +118,17 @@ def upload_directory_once(watch_dir, server_url, api_key, space_id):
     
     return success_count, failed_count
 
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="一次性文件上传器 - 扫描并上传指定文件夹中的所有文件")
+async def main():
+    parser = argparse.ArgumentParser(description="WebSocket文件上传器 - 扫描并上传指定文件夹中的所有文件")
     parser.add_argument("api_key", help="您的 API 密钥")
     parser.add_argument("space_id", help="Space ID")
-    parser.add_argument("--server", default="http://127.0.0.1:5001", help="服务器的 URL 地址 (默认: http://127.0.0.1:5001)")
+    parser.add_argument("--server", default="https://gkbtyo-rqvays-5001.preview.cloudstudio.work", help="服务器的 URL 地址")
     parser.add_argument("--upload-dir", default="output", help="要上传的目录 (默认: output)")
     
     args = parser.parse_args()
     
-    # 开始一次性上传
-    upload_directory_once(args.upload_dir, args.server, args.api_key, args.space_id)
+    # 开始上传
+    await upload_directory_once(args.upload_dir, args.server, args.api_key, args.space_id)
+
+if __name__ == "__main__":
+    asyncio.run(main())