'commit'

dsLightRag/Liblib/LiblibUtil.py

@@ -79,33 +79,76 @@ class LiblibUtil:
         payload = {"generateUuid": generate_uuid}
         return self.post_request(uri, payload)
 
-    def wait_for_generation_completion(self, generate_uuid, check_interval=5, max_wait_time=300):
-        """等待生图任务完成并获取最终结果"""
-        start_time = time.time()
-        
-        while time.time() - start_time < max_wait_time:
+    def wait_for_generation_completion(self, generate_uuid, interval=5):
+        """等待生成完成并返回最终状态"""
+        while True:
             status_data = self.get_generation_status(generate_uuid)
-            
-            if status_data:
-                status = status_data.get('status')
-                progress = status_data.get('percentCompleted', 0) * 100
-                print(f"生成进度：{progress:.2f}%")
-                
-                if status == 'completed':
-                    print("生图任务已完成")
-                    images = status_data.get('images', [])
-                    if images:
-                        print("生成的最终图片：")
-                        for img in images:
-                            print(f"- {img.get('imageUrl')}")
-                    return status_data
-                elif status == 'failed':
-                    print(f"生图任务失败: {status_data.get('errorMessage', '未知错误')}")
-                    return None
-                else:
-                    print(f"生图任务处理中，当前状态: {status}")
-            
-            time.sleep(check_interval)
+            if not status_data:
+                print("获取状态失败，重试中...")
+                time.sleep(interval)
+                continue
+
+            # 获取当前状态和进度
+            generate_status = status_data.get('generateStatus')
+            percent = status_data.get('percentCompleted', 0) * 100
+            status_text = {
+                0: "初始化",
+                1: "处理中",
+                2: "排队中",
+                3: "已取消",
+                4: "失败",
+                5: "完成"
+            }.get(generate_status, f"未知状态({generate_status})")
+
+            # 显示进度和状态
+            print(f"生成进度：{percent:.2f}% | 状态：{status_text}")
+
+            # 检查是否完成
+            if generate_status in [3, 4, 5]:  # 取消、失败或完成
+                return status_data
+
+            time.sleep(interval)
         
         print(f"生图任务超时（{max_wait_time}秒）")
-        return None
+        return None
+
+    @staticmethod
+    def download_and_upload_to_obs(image_url, generate_uuid):
+        """下载图片并上传到OBS"""
+        try:
+            # 1. 清理URL（去除可能的引号和空格）
+            clean_url = image_url.strip('` ')
+            print(f"清理后的图片URL: {clean_url}")
+
+            # 2. 下载图片到临时文件
+            response = requests.get(clean_url, stream=True)
+            response.raise_for_status()  # 检查HTTP错误
+
+            # 创建临时文件
+            with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
+                for chunk in response.iter_content(chunk_size=8192):
+                    temp_file.write(chunk)
+                temp_file_path = temp_file.name
+            print(f"图片已下载至临时文件: {temp_file_path}")
+
+            # 3. 上传到OBS
+            obs_uploader = ObsUploader()
+            obs_object_key = f"HuangHai/Liblib/{generate_uuid}.jpg"
+            success, result = obs_uploader.upload_file(
+                object_key=obs_object_key,
+                file_path=temp_file_path
+            )
+
+            # 4. 清理临时文件
+            os.unlink(temp_file_path)
+
+            if success:
+                print(f"✅ 图片成功上传至OBS")
+                return obs_object_key
+            else:
+                print(f"❌ OBS上传失败: {result}")
+                return None
+
+        except Exception as e:
+            print(f"处理图片时发生错误: {str(e)}")
+            return None