Update app.py

app.py

@@ -2,15 +2,19 @@ import os
 import json
 import asyncio
 import aiohttp
+import traceback
 from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 import uvicorn
 from typing import Dict, Any, AsyncGenerator
 import logging
 
-# 配置日志
-logging.basicConfig(level=logging.INFO)
+# 配置更详细的日志
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger(__name__)
 
 app = FastAPI(
@@ -31,90 +35,145 @@ app.add_middleware(
 # 从环境变量获取配置
 REPLICATE_API_TOKEN = os.getenv("REPLICATE_API_TOKEN")
 if not REPLICATE_API_TOKEN:
-    logger.warning("REPLICATE_API_TOKEN not found in environment variables")
+    logger.error("REPLICATE_API_TOKEN not found in environment variables")
 
 # Replicate API配置
 REPLICATE_BASE_URL = "https://api.replicate.com/v1"
-DEFAULT_MODEL = "anthropic/claude-4-sonnet"
+DEFAULT_MODEL = "anthropic/claude-3-5-sonnet"
+
+# 全局异常处理器
+@app.exception_handler(Exception)
+async def global_exception_handler(request: Request, exc: Exception):
+    logger.error(f"Global exception: {str(exc)}")
+    logger.error(f"Traceback: {traceback.format_exc()}")
+    return JSONResponse(
+        status_code=500,
+        content={
+            "error": {
+                "message": f"Internal server error: {str(exc)}",
+                "type": "internal_error"
+            }
+        }
+    )
 
 def transform_openai_to_replicate(openai_request: Dict[str, Any], model_override: str = None) -> Dict[str, Any]:
     """将OpenAI格式的请求转换为Replicate格式"""
-    messages = openai_request.get("messages", [])
-    
-    # 提取system prompt
-    system_prompt = ""
-    user_messages = []
-    
-    for message in messages:
-        if message.get("role") == "system":
-            system_prompt = message.get("content", "")
-        elif message.get("role") in ["user", "assistant"]:
-            user_messages.append(message)
-    
-    # 构建prompt
-    prompt_parts = []
-    for msg in user_messages:
-        role = msg.get("role", "")
-        content = msg.get("content", "")
-        if role == "user":
-            prompt_parts.append(f"User: {content}")
-        elif role == "assistant":
-            prompt_parts.append(f"Assistant: {content}")
-    
-    prompt = "\n\n".join(prompt_parts)
-    if prompt_parts and not prompt.endswith("\n\nAssistant:"):
-        prompt += "\n\nAssistant:"
-    
-    # 确定使用的模型
-    model = model_override or openai_request.get("model", DEFAULT_MODEL)
-    if not model.startswith("anthropic/"):
-        model = f"anthropic/{model}" if "/" not in model else model
-    
-    replicate_request = {
-        "stream": openai_request.get("stream", False),
-        "input": {
-            "prompt": prompt,
-            "system_prompt": system_prompt or "You are a helpful assistant",
-            "max_tokens": openai_request.get("max_tokens", 1000),
-            "temperature": openai_request.get("temperature", 0.7)
+    try:
+        messages = openai_request.get("messages", [])
+        
+        # 提取system prompt
+        system_prompt = "You are a helpful assistant"
+        user_messages = []
+        
+        for message in messages:
+            if message.get("role") == "system":
+                system_prompt = message.get("content", "You are a helpful assistant")
+            elif message.get("role") in ["user", "assistant"]:
+                user_messages.append(message)
+        
+        # 构建prompt
+        prompt_parts = []
+        for msg in user_messages:
+            role = msg.get("role", "")
+            content = msg.get("content", "")
+            if role == "user":
+                prompt_parts.append(f"Human: {content}")
+            elif role == "assistant":
+                prompt_parts.append(f"Assistant: {content}")
+        
+        prompt = "\n\n".join(prompt_parts)
+        if prompt_parts and not prompt.endswith("\n\nAssistant:"):
+            prompt += "\n\nAssistant:"
+        
+        # 确定使用的模型
+        model = model_override or openai_request.get("model", DEFAULT_MODEL)
+        
+        # 模型名称映射
+        model_mapping = {
+            "claude-4-sonnet": "anthropic/claude-3-5-sonnet",
+            "claude-3-sonnet": "anthropic/claude-3-sonnet-20240229",
+            "claude-3-haiku": "anthropic/claude-3-haiku-20240307"
         }
-    }
-    
-    return replicate_request, model
+        
+        if model in model_mapping:
+            model = model_mapping[model]
+        elif not model.startswith("anthropic/"):
+            model = f"anthropic/{model}"
+        
+        replicate_request = {
+            "stream": openai_request.get("stream", False),
+            "input": {
+                "prompt": prompt,
+                "system_prompt": system_prompt,
+                "max_tokens": openai_request.get("max_tokens", 4000),
+                "temperature": openai_request.get("temperature", 0.7)
+            }
+        }
+        
+        logger.info(f"Transformed request for model: {model}")
+        return replicate_request, model
+        
+    except Exception as e:
+        logger.error(f"Error transforming request: {str(e)}")
+        raise HTTPException(status_code=400, detail=f"Request transformation error: {str(e)}")
 
 async def create_replicate_prediction(session: aiohttp.ClientSession, model: str, data: Dict[str, Any]) -> Dict[str, Any]:
     """创建Replicate预测"""
-    url = f"{REPLICATE_BASE_URL}/models/{model}/predictions"
-    headers = {
-        "Authorization": f"Bearer {REPLICATE_API_TOKEN}",
-        "Content-Type": "application/json"
-    }
-    
-    async with session.post(url, headers=headers, json=data) as response:
-        if response.status != 201:
-            error_text = await response.text()
-            logger.error(f"Replicate API error: {response.status} - {error_text}")
-            raise HTTPException(status_code=response.status, detail=f"Replicate API error: {error_text}")
+    try:
+        url = f"{REPLICATE_BASE_URL}/models/{model}/predictions"
+        headers = {
+            "Authorization": f"Bearer {REPLICATE_API_TOKEN}",
+            "Content-Type": "application/json"
+        }
+        
+        logger.info(f"Creating prediction for model: {model}")
+        logger.info(f"Request URL: {url}")
         
-        return await response.json()
+        async with session.post(url, headers=headers, json=data, timeout=30) as response:
+            response_text = await response.text()
+            logger.info(f"Replicate response status: {response.status}")
+            logger.info(f"Replicate response: {response_text}")
+            
+            if response.status != 201:
+                logger.error(f"Replicate API error: {response.status} - {response_text}")
+                raise HTTPException(
+                    status_code=response.status, 
+                    detail=f"Replicate API error: {response_text}"
+                )
+            
+            return json.loads(response_text)
+            
+    except asyncio.TimeoutError:
+        logger.error("Timeout creating Replicate prediction")
+        raise HTTPException(status_code=504, detail="Timeout creating prediction")
+    except Exception as e:
+        logger.error(f"Error creating prediction: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Prediction creation error: {str(e)}")
 
 async def stream_replicate_response(session: aiohttp.ClientSession, stream_url: str) -> AsyncGenerator[str, None]:
     """流式读取Replicate响应"""
-    headers = {
-        "Accept": "text/event-stream",
-        "Cache-Control": "no-store"
-    }
-    
-    async with session.get(stream_url, headers=headers) as response:
-        if response.status != 200:
-            error_text = await response.text()
-            logger.error(f"Stream error: {response.status} - {error_text}")
-            raise HTTPException(status_code=response.status, detail=f"Stream error: {error_text}")
+    try:
+        headers = {
+            "Accept": "text/event-stream",
+            "Cache-Control": "no-store"
+        }
+        
+        logger.info(f"Starting stream from: {stream_url}")
         
-        async for line in response.content:
-            line = line.decode('utf-8').strip()
-            if line:
-                yield line
+        async with session.get(stream_url, headers=headers, timeout=300) as response:
+            if response.status != 200:
+                error_text = await response.text()
+                logger.error(f"Stream error: {response.status} - {error_text}")
+                raise HTTPException(status_code=response.status, detail=f"Stream error: {error_text}")
+            
+            async for line in response.content:
+                line = line.decode('utf-8').strip()
+                if line:
+                    yield line
+                    
+    except Exception as e:
+        logger.error(f"Stream error: {str(e)}")
+        raise
 
 def transform_replicate_to_openai_stream(event_data: str, model: str) -> str:
     """将Replicate流式响应转换为OpenAI格式"""
@@ -158,8 +217,8 @@ def transform_replicate_to_openai_stream(event_data: str, model: str) -> str:
         
         return ""
     
-    except json.JSONDecodeError:
-        logger.warning(f"Failed to parse event data: {event_data}")
+    except json.JSONDecodeError as e:
+        logger.warning(f"Failed to parse event data: {event_data}, error: {e}")
         return ""
 
 @app.get("/")
@@ -168,7 +227,17 @@ async def root():
     return {
         "message": "Replicate API Proxy for LobeChat",
         "status": "running",
-        "replicate_token_configured": bool(REPLICATE_API_TOKEN)
+        "replicate_token_configured": bool(REPLICATE_API_TOKEN),
+        "version": "1.0.0"
+    }
+
+@app.get("/health")
+async def health():
+    """详细健康检查"""
+    return {
+        "status": "healthy",
+        "replicate_token": "configured" if REPLICATE_API_TOKEN else "missing",
+        "timestamp": asyncio.get_event_loop().time()
     }
 
 @app.get("/v1/models")
@@ -200,20 +269,22 @@ async def list_models():
 async def chat_completions(request: Request):
     """处理聊天完成请求（兼容OpenAI API）"""
     if not REPLICATE_API_TOKEN:
+        logger.error("REPLICATE_API_TOKEN not configured")
         raise HTTPException(status_code=500, detail="REPLICATE_API_TOKEN not configured")
     
     try:
         body = await request.json()
-        logger.info(f"Received request: {json.dumps(body, indent=2)}")
+        logger.info(f"Received chat completion request")
+        logger.info(f"Request body: {json.dumps(body, indent=2)}")
         
         # 转换请求格式
         replicate_data, model = transform_openai_to_replicate(body)
-        logger.info(f"Transformed to Replicate format: {json.dumps(replicate_data, indent=2)}")
         
         async with aiohttp.ClientSession() as session:
             # 创建预测
             prediction = await create_replicate_prediction(session, model, replicate_data)
-            logger.info(f"Created prediction: {prediction.get('id')}")
+            prediction_id = prediction.get('id')
+            logger.info(f"Created prediction: {prediction_id}")
             
             if body.get("stream", False):
                 # 流式响应
@@ -250,18 +321,26 @@ async def chat_completions(request: Request):
             
             else:
                 # 非流式响应 - 等待预测完成
-                prediction_url = f"{REPLICATE_BASE_URL}/predictions/{prediction['id']}"
+                prediction_url = f"{REPLICATE_BASE_URL}/predictions/{prediction_id}"
                 headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}"}
                 
                 # 轮询等待结果
-                while True:
+                max_attempts = 60  # 最多等待60秒
+                attempt = 0
+                
+                while attempt < max_attempts:
                     async with session.get(prediction_url, headers=headers) as response:
                         result = await response.json()
+                        status = result.get("status")
                         
-                        if result.get("status") == "succeeded":
-                            content = "".join(result.get("output", []))
+                        logger.info(f"Prediction {prediction_id} status: {status}")
+                        
+                        if status == "succeeded":
+                            output = result.get("output", [])
+                            content = "".join(output) if isinstance(output, list) else str(output)
+                            
                             openai_response = {
-                                "id": f"chatcmpl-{result['id']}",
+                                "id": f"chatcmpl-{prediction_id}",
                                 "object": "chat.completion",
                                 "created": int(asyncio.get_event_loop().time()),
                                 "model": model,
@@ -275,22 +354,34 @@ async def chat_completions(request: Request):
                                 }],
                                 "usage": {
                                     "prompt_tokens": 0,
-                                    "completion_tokens": 0,
-                                    "total_tokens": 0
+                                    "completion_tokens": len(content.split()),
+                                    "total_tokens": len(content.split())
                                 }
                             }
                             return openai_response
                         
-                        elif result.get("status") == "failed":
-                            raise HTTPException(status_code=500, detail=f"Prediction failed: {result.get('error')}")
+                        elif status == "failed":
+                            error_msg = result.get('error', 'Unknown error')
+                            logger.error(f"Prediction failed: {error_msg}")
+                            raise HTTPException(status_code=500, detail=f"Prediction failed: {error_msg}")
+                        
+                        elif status in ["canceled", "cancelled"]:
+                            raise HTTPException(status_code=500, detail="Prediction was canceled")
                         
                         # 等待一秒后重试
                         await asyncio.sleep(1)
+                        attempt += 1
+                
+                raise HTTPException(status_code=504, detail="Prediction timeout")
     
+    except HTTPException:
+        raise
     except Exception as e:
-        logger.error(f"Error processing request: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
+        logger.error(f"Unexpected error processing request: {str(e)}")
+        logger.error(f"Traceback: {traceback.format_exc()}")
+        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
 if __name__ == "__main__":
     port = int(os.getenv("PORT", 7860))
-    uvicorn.run(app, host="0.0.0.0", port=port)
+    logger.info(f"Starting server on port {port}")
+    uvicorn.run(app, host="0.0.0.0", port=port, log_level="info")