Update app.py

app.py

@@ -39,7 +39,7 @@ if not REPLICATE_API_TOKEN:
 
 # Replicate API配置
 REPLICATE_BASE_URL = "https://api.replicate.com/v1"
-DEFAULT_MODEL = "anthropic/claude-3.5-sonnet"  # 使用实际存在的模型
+DEFAULT_MODEL = "anthropic/claude-3.5-sonnet"
 
 # 全局异常处理器
 @app.exception_handler(Exception)
@@ -85,22 +85,21 @@ def transform_openai_to_replicate(openai_request: Dict[str, Any], model_override
         if prompt_parts and not prompt.endswith("\n\nAssistant:"):
             prompt += "\n\nAssistant:"
         
-        # 确定使用的模型 - 使用正确的 Replicate 模型名称
+        # 确定使用的模型
         model = model_override or openai_request.get("model", DEFAULT_MODEL)
         
-        # 正确的模型名称映射（基于搜索结果）
+        # 正确的模型名称映射
         model_mapping = {
-            "claude-4-sonnet": "anthropic/claude-4-sonnet",  # 最新的 Claude 4
-            "claude-3.5-sonnet": "anthropic/claude-3.5-sonnet",  # Claude 3.5 Sonnet
-            "claude-3-sonnet": "anthropic/claude-3-sonnet",  # Claude 3 Sonnet
-            "claude-3.5-haiku": "anthropic/claude-3.5-haiku",  # Claude 3.5 Haiku
-            "claude-3-haiku": "anthropic/claude-3-haiku",  # Claude 3 Haiku
+            "claude-4-sonnet": "anthropic/claude-4-sonnet",
+            "claude-3.5-sonnet": "anthropic/claude-3.5-sonnet",
+            "claude-3-sonnet": "anthropic/claude-3-sonnet",
+            "claude-3.5-haiku": "anthropic/claude-3.5-haiku",
+            "claude-3-haiku": "anthropic/claude-3-haiku",
         }
         
         if model in model_mapping:
             model = model_mapping[model]
         elif not model.startswith("anthropic/"):
-            # 默认使用 claude-3.5-sonnet
             model = "anthropic/claude-3.5-sonnet"
         
         replicate_request = {
@@ -130,12 +129,10 @@ async def create_replicate_prediction(session: aiohttp.ClientSession, model: str
         }
         
         logger.info(f"Creating prediction for model: {model}")
-        logger.info(f"Request URL: {url}")
         
         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}")
@@ -153,31 +150,6 @@ async def create_replicate_prediction(session: aiohttp.ClientSession, model: str
         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响应"""
-    try:
-        headers = {
-            "Accept": "text/event-stream",
-            "Cache-Control": "no-store"
-        }
-        
-        logger.info(f"Starting stream from: {stream_url}")
-        
-        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格式"""
     if not event_data.startswith("data: "):
@@ -290,32 +262,62 @@ async def chat_completions(request: Request):
     try:
         body = await request.json()
         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)
         
-        async with aiohttp.ClientSession() as session:
-            # 创建预测
-            prediction = await create_replicate_prediction(session, model, replicate_data)
-            prediction_id = prediction.get('id')
-            logger.info(f"Created prediction: {prediction_id}")
-            
-            if body.get("stream", False):
-                # 流式响应
-                stream_url = prediction.get("urls", {}).get("stream")
-                if not stream_url:
-                    raise HTTPException(status_code=500, detail="Stream URL not available")
-                
-                async def generate_stream():
+        if body.get("stream", False):
+            # 流式响应 - 修复会话管理问题
+            async def generate_stream():
+                async with aiohttp.ClientSession() as session:
                     try:
-                        async for event in stream_replicate_response(session, stream_url):
-                            openai_event = transform_replicate_to_openai_stream(event, model)
-                            if openai_event:
-                                yield openai_event
+                        # 创建预测
+                        prediction = await create_replicate_prediction(session, model, replicate_data)
+                        prediction_id = prediction.get('id')
+                        logger.info(f"Created prediction: {prediction_id}")
+                        
+                        # 获取流式URL
+                        stream_url = prediction.get("urls", {}).get("stream")
+                        if not stream_url:
+                            error_response = {
+                                "error": {
+                                    "message": "Stream URL not available",
+                                    "type": "stream_error"
+                                }
+                            }
+                            yield f"data: {json.dumps(error_response)}\n\n"
+                            return
+                        
+                        logger.info(f"Starting stream from: {stream_url}")
+                        
+                        # 流式读取响应
+                        headers = {
+                            "Accept": "text/event-stream",
+                            "Cache-Control": "no-store"
+                        }
+                        
+                        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}")
+                                error_response = {
+                                    "error": {
+                                        "message": f"Stream error: {error_text}",
+                                        "type": "stream_error"
+                                    }
+                                }
+                                yield f"data: {json.dumps(error_response)}\n\n"
+                                return
+                            
+                            async for line in response.content:
+                                line = line.decode('utf-8').strip()
+                                if line:
+                                    openai_event = transform_replicate_to_openai_stream(line, model)
+                                    if openai_event:
+                                        yield openai_event
+                                        
                     except Exception as e:
                         logger.error(f"Stream generation error: {e}")
-                        # 发送错误响应
                         error_response = {
                             "error": {
                                 "message": str(e),
@@ -323,23 +325,29 @@ 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": "*",
-                    }
-                )
             
-            else:
-                # 非流式响应 - 等待预测完成
+            return StreamingResponse(
+                generate_stream(),
+                media_type="text/event-stream",
+                headers={
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "Access-Control-Allow-Origin": "*",
+                }
+            )
+        
+        else:
+            # 非流式响应
+            async with aiohttp.ClientSession() as session:
+                # 创建预测
+                prediction = await create_replicate_prediction(session, model, replicate_data)
+                prediction_id = prediction.get('id')
+                logger.info(f"Created prediction: {prediction_id}")
+                
+                # 轮询等待结果
                 prediction_url = f"{REPLICATE_BASE_URL}/predictions/{prediction_id}"
                 headers = {"Authorization": f"Bearer {REPLICATE_API_TOKEN}"}
                 
-                # 轮询等待结果
                 max_attempts = 60  # 最多等待60秒
                 attempt = 0