Update app.py

app.py

@@ -39,6 +39,10 @@ REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
 if not REPLICATE_API_TOKEN:
     logger.error("REPLICATE_API_TOKEN not found in environment variables")
 
+# imgbb API 配置
+IMGBB_API_KEY = "78f0c4360135e80c46b24b44e1e20a20"
+IMGBB_API_URL = "https://api.imgbb.com/1/upload"
+
 # Replicate API配置
 REPLICATE_BASE_URL = "https://api.replicate.com/v1"
 DEFAULT_MODEL = "anthropic/claude-3.5-sonnet"
@@ -178,10 +182,10 @@ def decode_base64_file(data_url: str) -> tuple[str, str, str]:
         logger.error(f"Failed to parse data URL: {e}")
         return None, None, None
 
-async def upload_image_to_temp_service(session: aiohttp.ClientSession, base64_data: str) -> str:
+async def upload_image_to_imgbb(session: aiohttp.ClientSession, base64_data: str) -> str:
     """
-    将 base64 图片上传到临时图片托管服务
-    这里使用 imgbb 作为示例，你也可以使用其他服务
+    将 base64 图片上传到 imgbb
+    返回图片的 URL
     """
     try:
         # 从 base64 data URL 中提取纯 base64 数据
@@ -190,63 +194,91 @@ async def upload_image_to_temp_service(session: aiohttp.ClientSession, base64_da
         else:
             base64_content = base64_data
         
-        # 使用 imgbb API（需要免费注册获取 API key）
-        # 这里暂时返回原始 data URL，你需要根据实际情况实现图片上传
-        logger.warning("Image upload to external service not implemented, using workaround")
-        
-        # 临时解决方案：对于 Claude 4，我们需要找到另一种方式
-        # 可以考虑：
-        # 1. 使用临时文件服务（如 imgbb, imgur 等）
-        # 2. 使用自己的文件服务器
-        # 3. 修改为使用 claude-3.5-sonnet 作为替代
-        
-        return None  # 返回 None 表示上传失败
+        # 准备上传数据
+        data = {
+            'key': IMGBB_API_KEY,
+            'image': base64_content,
+            'expiration': 300  # 5分钟过期，避免永久占用存储
+        }
         
+        logger.info(f"Uploading image to imgbb, size: {len(base64_content)} chars")
+        
+        # 上传到 imgbb
+        async with session.post(IMGBB_API_URL, data=data, timeout=30) as response:
+            if response.status == 200:
+                result = await response.json()
+                if result.get('success'):
+                    image_url = result['data']['url']
+                    logger.info(f"Image uploaded successfully: {image_url}")
+                    return image_url
+                else:
+                    logger.error(f"imgbb upload failed: {result}")
+                    return None
+            else:
+                error_text = await response.text()
+                logger.error(f"imgbb upload error: {response.status} - {error_text}")
+                return None
+                
+    except asyncio.TimeoutError:
+        logger.error("Timeout uploading image to imgbb")
+        return None
     except Exception as e:
-        logger.error(f"Failed to upload image: {e}")
+        logger.error(f"Failed to upload image to imgbb: {e}")
         return None
 
-def format_image_for_model(base64_data: str, model_config: Dict[str, Any]) -> str:
+async def format_image_for_model(session: aiohttp.ClientSession, base64_data: str, model_config: Dict[str, Any]) -> str:
     """
     根据模型配置格式化图片数据
     """
     image_format = model_config.get("image_format", "data_url")
     
-    if image_format == "data_url":
-        # 检查 base64 数据是否已经包含 data URL 前缀
-        if base64_data.startswith("data:"):
-            return base64_data
-        
-        # 如果没有前缀，添加默认的 JPEG data URL 前缀
-        try:
-            # 解码 base64 数据的前几个字节来检测格式
-            decoded_bytes = base64.b64decode(base64_data[:100])
-            
-            if decoded_bytes.startswith(b'\xff\xd8\xff'):
-                # JPEG
-                return f"data:image/jpeg;base64,{base64_data}"
-            elif decoded_bytes.startswith(b'\x89PNG\r\n\x1a\n'):
-                # PNG
-                return f"data:image/png;base64,{base64_data}"
-            elif decoded_bytes.startswith(b'GIF87a') or decoded_bytes.startswith(b'GIF89a'):
-                # GIF
-                return f"data:image/gif;base64,{base64_data}"
-            elif decoded_bytes.startswith(b'RIFF') and b'WEBP' in decoded_bytes[:20]:
-                # WebP
-                return f"data:image/webp;base64,{base64_data}"
-            else:
-                # 默认使用 JPEG
-                return f"data:image/jpeg;base64,{base64_data}"
-        except Exception as e:
-            logger.warning(f"Failed to detect image format: {e}, using JPEG as default")
-            return f"data:image/jpeg;base64,{base64_data}"
+    if image_format == "url":
+        # 需要上传图片到 imgbb 并返回 URL
+        image_url = await upload_image_to_imgbb(session, base64_data)
+        if image_url:
+            return image_url
+        else:
+            logger.error("Failed to upload image, falling back to data URL")
+            # 上传失败时降级到 data URL 格式
+            return format_image_as_data_url(base64_data)
     
-    elif image_format == "url":
-        # 对于需要 URL 的模型，返回 None 表示需要上传
-        return None
+    elif image_format == "data_url":
+        return format_image_as_data_url(base64_data)
     
     return base64_data
 
+def format_image_as_data_url(base64_data: str) -> str:
+    """
+    将 base64 数据格式化为 data URL
+    """
+    # 检查 base64 数据是否已经包含 data URL 前缀
+    if base64_data.startswith("data:"):
+        return base64_data
+    
+    # 如果没有前缀，添加默认的 JPEG data URL 前缀
+    try:
+        # 解码 base64 数据的前几个字节来检测格式
+        decoded_bytes = base64.b64decode(base64_data[:100])
+        
+        if decoded_bytes.startswith(b'\xff\xd8\xff'):
+            # JPEG
+            return f"data:image/jpeg;base64,{base64_data}"
+        elif decoded_bytes.startswith(b'\x89PNG\r\n\x1a\n'):
+            # PNG
+            return f"data:image/png;base64,{base64_data}"
+        elif decoded_bytes.startswith(b'GIF87a') or decoded_bytes.startswith(b'GIF89a'):
+            # GIF
+            return f"data:image/gif;base64,{base64_data}"
+        elif decoded_bytes.startswith(b'RIFF') and b'WEBP' in decoded_bytes[:20]:
+            # WebP
+            return f"data:image/webp;base64,{base64_data}"
+        else:
+            # 默认使用 JPEG
+            return f"data:image/jpeg;base64,{base64_data}"
+    except Exception as e:
+        logger.warning(f"Failed to detect image format: {e}, using JPEG as default")
+        return f"data:image/jpeg;base64,{base64_data}"
+
 def extract_content_from_message(message: Dict[str, Any]) -> tuple[str, List[str], List[Dict[str, str]]]:
     """
     从消息中提取文本内容、图片和文件
@@ -347,7 +379,7 @@ def format_files_for_prompt(files: List[Dict[str, str]]) -> str:
     
     return "\n".join(file_sections)
 
-async def transform_openai_to_replicate(openai_request: Dict[str, Any], model_override: str = None) -> Dict[str, Any]:
+async def transform_openai_to_replicate(session: aiohttp.ClientSession, openai_request: Dict[str, Any], model_override: str = None) -> Dict[str, Any]:
     """将OpenAI格式的请求转换为Replicate格式"""
     try:
         messages = openai_request.get("messages", [])
@@ -416,14 +448,12 @@ async def transform_openai_to_replicate(openai_request: Dict[str, Any], model_ov
         # 处理图片格式
         formatted_image = None
         if has_images and primary_image:
-            if model_config.get("image_format") == "url":
-                # Claude 4 需要 URL 格式，暂时降级到 Claude 3.5
-                logger.warning(f"Model {model} requires URL format for images, falling back to claude-3.5-sonnet")
-                model = "anthropic/claude-3.5-sonnet"
-                model_config = MODEL_CONFIGS[model]
-                formatted_image = format_image_for_model(primary_image, model_config)
-            else:
-                formatted_image = format_image_for_model(primary_image, model_config)
+            logger.info(f"Processing image for model {model} with format {model_config.get('image_format')}")
+            formatted_image = await format_image_for_model(session, primary_image, model_config)
+            
+            if not formatted_image:
+                logger.error("Failed to format image for model")
+                raise HTTPException(status_code=500, detail="Failed to process image")
         
         # 构建 Replicate 格式的输入
         replicate_input = {}
@@ -459,7 +489,10 @@ async def transform_openai_to_replicate(openai_request: Dict[str, Any], model_ov
         # 处理图片
         if formatted_image:
             replicate_input["image"] = formatted_image
-            logger.info(f"Added image to request for model {model}: {formatted_image[:100]}...")
+            if formatted_image.startswith("http"):
+                logger.info(f"Added image URL to request for model {model}: {formatted_image}")
+            else:
+                logger.info(f"Added image data to request for model {model}: {formatted_image[:100]}...")
         
         # 只在有 system_prompt 时才添加
         if system_prompt:
@@ -530,7 +563,10 @@ async def create_replicate_prediction(session: aiohttp.ClientSession, model: str
         if "input" in log_data:
             if "image" in log_data["input"]:
                 image_data = log_data["input"]["image"]
-                log_data["input"]["image"] = f"[IMAGE_DATA_{len(image_data)}]"
+                if image_data.startswith("http"):
+                    log_data["input"]["image"] = f"[IMAGE_URL: {image_data}]"
+                else:
+                    log_data["input"]["image"] = f"[IMAGE_DATA_{len(image_data)}]"
             if "prompt" in log_data["input"] and len(log_data["input"]["prompt"]) > 1000:
                 log_data["input"]["prompt"] = log_data["input"]["prompt"][:1000] + "...[TRUNCATED]"
         logger.info(f"Request data: {json.dumps(log_data, indent=2)}")
@@ -615,13 +651,14 @@ async def root():
         "message": "Replicate API Proxy for LobeChat with Vision and File Support",
         "status": "running",
         "replicate_token_configured": bool(REPLICATE_API_TOKEN),
-        "version": "1.1.2",
+        "imgbb_token_configured": bool(IMGBB_API_KEY),
+        "version": "1.2.0",
         "supported_models": list(MODEL_CONFIGS.keys()),
         "vision_support": True,
         "file_support": True,
         "supported_text_files": list(SUPPORTED_TEXT_EXTENSIONS),
         "supported_image_files": list(SUPPORTED_IMAGE_EXTENSIONS),
-        "notes": "Claude 4 Sonnet image support temporarily falls back to Claude 3.5 Sonnet"
+        "claude4_vision_support": "Full support via imgbb image hosting"
     }
 
 @app.get("/health")
@@ -630,6 +667,7 @@ async def health():
     return {
         "status": "healthy",
         "replicate_token": "configured" if REPLICATE_API_TOKEN else "missing",
+        "imgbb_token": "configured" if IMGBB_API_KEY else "missing",
         "timestamp": asyncio.get_event_loop().time(),
         "model_configs": MODEL_CONFIGS,
         "supported_file_types": {
@@ -665,13 +703,13 @@ async def chat_completions(request: Request):
         logger.info(f"Client parameters: max_tokens={body.get('max_tokens', 'not set')}, temperature={body.get('temperature', 'not set')}")
         logger.info(f"Message count: {len(body.get('messages', []))}")
         
-        # 转换请求格式
-        replicate_data, model = await transform_openai_to_replicate(body)
-        
-        if body.get("stream", False):
-            # 流式响应
-            async def generate_stream():
-                async with aiohttp.ClientSession() as session:
+        async with aiohttp.ClientSession() as session:
+            # 转换请求格式
+            replicate_data, model = await transform_openai_to_replicate(session, body)
+            
+            if body.get("stream", False):
+                # 流式响应
+                async def generate_stream():
                     try:
                         # 创建预测
                         prediction = await create_replicate_prediction(session, model, replicate_data)
@@ -755,21 +793,20 @@ async def chat_completions(request: Request):
                             }
                         }
                         yield f"data: {json.dumps(error_response)}\n\n"
+                
+                return StreamingResponse(
+                    generate_stream(),
+                    media_type="text/event-stream",
+                    headers={
+                        "Cache-Control": "no-cache",
+                        "Connection": "keep-alive",
+                        "Access-Control-Allow-Origin": "*",
+                        "X-Accel-Buffering": "no",
+                    }
+                )
             
-            return StreamingResponse(
-                generate_stream(),
-                media_type="text/event-stream",
-                headers={
-                    "Cache-Control": "no-cache",
-                    "Connection": "keep-alive",
-                    "Access-Control-Allow-Origin": "*",
-                    "X-Accel-Buffering": "no",
-                }
-            )
-        
-        else:
-            # 非流式响应
-            async with aiohttp.ClientSession() as session:
+            else:
+                # 非流式响应
                 # 创建预测
                 prediction = await create_replicate_prediction(session, model, replicate_data)
                 prediction_id = prediction.get('id')