lambda_function.py

#!/usr/bin/env python3
"""
AWS Bedrock API Service Lambda Function - API Gateway Version with True SSE Streaming
支持 API Gateway HTTP API v2.0 的真正流式输出
"""

import json
import time
import uuid
import base64
from datetime import datetime
from typing import Dict, Any, Optional, Iterator

# 导入配置
from config import Config

# 导入服务模块
from services.apigw_authentication_handler import APIGatewayAuthenticationHandler
from services.request_parser import RequestParser
from services.image_processor import ImageProcessor
from services.bedrock_client import BedrockClient
from utils.logging_utils import setup_logging, StructuredLogger

# 设置结构化日志记录
setup_logging()
logger = StructuredLogger("bedrock-api-service-apigw")

# 初始化服务组件
auth_handler = APIGatewayAuthenticationHandler(
    secret_name=Config.API_KEYS_SECRET_NAME + "-apigw",
    region_name=Config.BEDROCK_REGION
)
request_parser = RequestParser()
image_processor = ImageProcessor(bucket_name=Config.S3_BUCKET + "-apigw")
bedrock_client = BedrockClient(region=Config.BEDROCK_REGION)


def create_sse_response(data: Dict[str, Any]) -> str:
    """
    创建 SSE 格式的响应
    
    Args:
        data: 响应数据
        
    Returns:
        SSE 格式的字符串
    """
    return f"data: {json.dumps(data, ensure_ascii=False)}\n\n"


def create_streaming_response_chunk(content: str, finish_reason: str, meta: Dict[str, Any], 
                                  created_time: int, model_name: str, image_id: str = "") -> Dict[str, Any]:
    """
    创建符合接口规范的流式响应块
    
    Args:
        content: 内容文本
        finish_reason: 完成原因 ("stream" 或 "stop")
        meta: 元数据
        created_time: 创建时间戳
        model_name: 模型名称
        image_id: 图像ID
        
    Returns:
        响应块字典
    """
    return {
        "id": image_id,
        "object": "understand_image",
        "created": created_time,
        "model": model_name,
        "choices": [{
            "index": 0,
            "delta": {
                "content": content
            },
            "finish_reason": finish_reason
        }],
        "meta": meta,
        "runtime_status": ""
    }


def create_non_streaming_response(content: str, meta: Dict[str, Any], 
                                created_time: int, model_name: str, 
                                runtime_status: str, image_id: str = "") -> Dict[str, Any]:
    """
    创建非流式响应
    
    Args:
        content: 完整内容
        meta: 元数据
        created_time: 创建时间戳
        model_name: 模型名称
        runtime_status: 运行时状态
        image_id: 图像ID
        
    Returns:
        完整响应字典
    """
    return {
        "id": image_id,
        "object": "understand_image",
        "created": created_time,
        "model": model_name,
        "choices": [{
            "index": 0,
            "delta": {
                "content": content
            },
            "finish_reason": "stop"
        }],
        "meta": meta,
        "runtime_status": runtime_status
    }


def create_json_response(status_code: int, data: Dict[str, Any], 
                        headers: Dict[str, str] = None) -> Dict[str, Any]:
    """
    创建标准 JSON 响应
    
    Args:
        status_code: HTTP 状态码
        data: 响应数据
        headers: 额外的响应头
        
    Returns:
        API Gateway 响应格式
    """
    default_headers = {
        'Content-Type': 'application/json',
        'Access-Control-Allow-Origin': '*',
        'Access-Control-Allow-Methods': 'POST, OPTIONS',
        'Access-Control-Allow-Headers': 'Content-Type, Authorization',
        'Cache-Control': 'no-store'
    }
    
    if headers:
        default_headers.update(headers)
    
    return {
        'statusCode': status_code,
        'headers': default_headers,
        'body': json.dumps(data, ensure_ascii=False),
        'isBase64Encoded': False
    }


def create_streaming_response(status_code: int, headers: Dict[str, str] = None) -> Dict[str, Any]:
    """
    创建流式响应头
    
    Args:
        status_code: HTTP 状态码
        headers: 额外的响应头
        
    Returns:
        API Gateway 流式响应格式
    """
    default_headers = {
        'Content-Type': 'text/event-stream',
        'Cache-Control': 'no-cache',
        'Connection': 'keep-alive',
        'Access-Control-Allow-Origin': '*',
        'Access-Control-Allow-Methods': 'POST, OPTIONS',
        'Access-Control-Allow-Headers': 'Content-Type, Authorization'
    }
    
    if headers:
        default_headers.update(headers)
    
    return {
        'statusCode': status_code,
        'headers': default_headers,
        'isBase64Encoded': False
    }


def create_cors_response() -> Dict[str, Any]:
    """创建 CORS 预检响应"""
    return {
        'statusCode': 200,
        'headers': {
            'Access-Control-Allow-Origin': '*',
            'Access-Control-Allow-Methods': 'POST, OPTIONS',
            'Access-Control-Allow-Headers': 'Content-Type, Authorization',
            'Access-Control-Max-Age': '86400'
        },
        'body': '',
        'isBase64Encoded': False
    }


def parse_apigw_event(event: Dict[str, Any]) -> Dict[str, Any]:
    """
    解析API Gateway v2.0事件格式
    
    Args:
        event: API Gateway v2.0事件
        
    Returns:
        解析后的事件数据
    """
    request_context = event.get('requestContext', {})
    http_context = request_context.get('http', {})
    
    return {
        'http_method': http_context.get('method', ''),
        'path': event.get('rawPath', ''),
        'headers': event.get('headers', {}),
        'body': event.get('body', ''),
        'is_base64_encoded': event.get('isBase64Encoded', False),
        'source_ip': http_context.get('sourceIp', 'unknown'),
        'user_agent': event.get('headers', {}).get('user-agent', 'unknown'),
        'query_parameters': event.get('queryStringParameters') or {}
    }


def lambda_handler(event: Dict[str, Any], context: Any) -> Any:
    """
    Lambda 处理函数，支持 API Gateway HTTP API v2.0 的流式输出
    
    Args:
        event: API Gateway v2.0 事件
        context: Lambda 上下文
        
    Returns:
        HTTP 响应或流式响应
    """
    request_id = str(uuid.uuid4())
    start_time = time.time()
    
    # 解析API Gateway事件
    parsed_event = parse_apigw_event(event)
    
    # 确定是否为流式请求
    is_streaming = False
    try:
        if parsed_event['body']:
            body_data = json.loads(parsed_event['body'])
            is_streaming = body_data.get('stream', False)
    except Exception:
        pass
    
    # 开始请求跟踪
    logger.start_request_tracking(
        request_id=request_id,
        is_streaming=is_streaming,
        model_name=Config.BEDROCK_MODEL_ID
    )
    
    try:
        # 记录请求信息
        logger.log_request(
            event=event,
            request_id=request_id
        )
        
        # 1. 处理CORS预检请求
        if parsed_event['http_method'] == 'OPTIONS':
            response = create_cors_response()
            logger.log_response(
                status_code=200,
                response_size=0,
                request_id=request_id
            )
            return response
        
        # 2. 处理健康检查端点
        if parsed_event['path'] == '/health' or parsed_event['path'].endswith('/health'):
            health_response = {
                'status': 'healthy',
                'timestamp': datetime.utcnow().isoformat() + 'Z',
                'service': 'bedrock-api-service-apigw',
                'version': '1.0.0',
                'streaming_support': True
            }
            return create_json_response(200, health_response)
        
        # 3. 验证HTTP方法
        if parsed_event['http_method'] != 'POST':
            error_response = {
                'error': {
                    'type': 'validation_error',
                    'message': 'Only POST method is supported',
                    'code': 400
                }
            }
            return create_json_response(400, error_response)
        
        # 4. 验证路径
        if not (parsed_event['path'] == '/mllm/completions' or 
                parsed_event['path'].endswith('/mllm/completions')):
            error_response = {
                'error': {
                    'type': 'validation_error',
                    'message': 'Invalid endpoint path',
                    'code': 404
                }
            }
            return create_json_response(404, error_response)
        
        # 5. 身份验证
        auth_start = time.time()
        auth_result = auth_handler.authenticate_request(parsed_event['headers'])
        auth_duration = (time.time() - auth_start) * 1000
        
        logger.log_performance_phase(
            phase="authentication",
            duration_ms=auth_duration,
            request_id=request_id,
            additional_info={
                "authenticated": auth_result.authenticated
            }
        )
        
        if not auth_result.authenticated:
            error_response = {
                'error': {
                    'type': 'authentication_error',
                    'message': 'Authentication failed',
                    'code': 401
                }
            }
            return create_json_response(401, error_response)
        
        # 6. 解析和验证请求
        parsing_start = time.time()
        request_body = parsed_event['body']
        
        if parsed_event['is_base64_encoded']:
            request_body = base64.b64decode(request_body).decode('utf-8')
        
        try:
            request_data = json.loads(request_body)
        except json.JSONDecodeError as e:
            error_response = {
                'error': {
                    'type': 'validation_error',
                    'message': f'Invalid JSON: {str(e)}',
                    'code': 400
                }
            }
            return create_json_response(400, error_response)
        
        # 解析请求
        request_event = {'body': json.dumps(request_data)}
        parsed_request = request_parser.parse_request(request_event)
        validation_result = request_parser.validate_request(parsed_request)
        
        parsing_duration = (time.time() - parsing_start) * 1000
        
        logger.log_performance_phase(
            phase="parsing",
            duration_ms=parsing_duration,
            request_id=request_id,
            additional_info={
                "model": parsed_request.model,
                "stream": parsed_request.stream,
                "message_count": len(parsed_request.messages)
            }
        )
        
        if not validation_result.is_valid:
            error_response = {
                'error': {
                    'type': 'validation_error',
                    'message': validation_result.error_message,
                    'code': 400
                }
            }
            return create_json_response(400, error_response)
        
        # 7. 处理图像（如果有）
        image_processing_start = time.time()
        processed_messages = []
        
        for message in parsed_request.messages:
            processed_content = []
            for content in message.content:
                if content.type == 'image':
                    try:
                        # 验证图像
                        content_dict = {'image_base64': content.image_base64}
                        processed_image = image_processor.process_image_content(content_dict)
                        processed_content.append(content)
                    except Exception as e:
                        logger.log_error(
                            error=e,
                            context={"content_type": content.type},
                            request_id=request_id,
                            error_phase="image_processing"
                        )
                        error_response = {
                            'error': {
                                'type': 'image_processing_error',
                                'message': f'图像处理失败: {str(e)}',
                                'code': 400
                            }
                        }
                        return create_json_response(400, error_response)
                else:
                    processed_content.append(content)
            
            # 创建新的Message对象
            from models.request_models import Message
            processed_message = Message(role=message.role, content=processed_content)
            processed_messages.append(processed_message)
        
        # 创建新的ParsedRequest对象
        from models.request_models import ParsedRequest
        parsed_request = ParsedRequest(
            model=parsed_request.model,
            messages=processed_messages,
            stream=parsed_request.stream,
            meta=parsed_request.meta,
            temperature=parsed_request.temperature,
            top_p=parsed_request.top_p,
            seed=parsed_request.seed,
            reset=parsed_request.reset
        )
        
        image_processing_duration = (time.time() - image_processing_start) * 1000
        
        logger.log_performance_phase(
            phase="image_processing",
            duration_ms=image_processing_duration,
            request_id=request_id
        )
        
        # 8. 格式化Bedrock请求
        try:
            bedrock_request = bedrock_client.format_bedrock_request(parsed_request)
        except Exception as e:
            logger.log_error(
                error=e,
                context={"model": parsed_request.model},
                request_id=request_id,
                error_phase="bedrock_formatting"
            )
            error_response = {
                'error': {
                    'type': 'internal_error',
                    'message': '请求格式化失败',
                    'code': 500
                }
            }
            return create_json_response(500, error_response)
        
        # 9. 路由到流式或非流式处理
        if parsed_request.stream:
            # 流式响应处理
            return handle_streaming_request(bedrock_request, parsed_request.meta, request_id, start_time)
        else:
            # 非流式响应处理
            return handle_non_streaming_request(bedrock_request, parsed_request.meta, request_id, start_time)
            
    except Exception as e:
        logger.log_error(
            error=e,
            context={
                "event_path": parsed_event.get('path'), 
                "event_method": parsed_event.get('http_method')
            },
            request_id=request_id,
            error_phase="lambda_handler"
        )
        
        # 返回通用错误响应
        error_response = {
            'error': {
                'type': 'internal_error',
                'message': '内部服务器错误，请稍后重试',
                'code': 500,
                'request_id': request_id,
                'timestamp': int(time.time()),
                'processing_time_ms': round((time.time() - start_time) * 1000, 2)
            }
        }
        
        return create_json_response(500, error_response)
    
    finally:
        # 完成请求跟踪
        logger.finish_request_tracking(request_id=request_id)


def handle_streaming_request(bedrock_request, meta: Dict[str, Any], request_id: str, start_time: float) -> Dict[str, Any]:
    """
    处理流式请求 - API Gateway HTTP API v2.0 支持
    
    Args:
        bedrock_request: Bedrock 请求对象
        meta: 元数据
        request_id: 请求ID
        start_time: 开始时间
        
    Returns:
        API Gateway 流式响应
    """
    try:
        created_time = int(time.time())
        model_name = Config.BEDROCK_MODEL_ID.split('.')[-1] if '.' in Config.BEDROCK_MODEL_ID else Config.BEDROCK_MODEL_ID
        image_id = meta.get('image_id', '') if meta else ''
        
        logger.log_structured("api_gateway_streaming_start", {
            "request_id": request_id,
            "model": model_name,
            "image_id": image_id
        })
        
        # 创建流式响应生成器
        def stream_generator():
            try:
                total_content = ""
                chunk_count = 0
                
                # 调用 Bedrock 流式 API
                for chunk in bedrock_client.invoke_model_stream(bedrock_request):
                    chunk_count += 1
                    
                    if chunk.type == 'content_block_delta' and chunk.content_block_delta:
                        delta = chunk.content_block_delta.get('delta', {})
                        
                        if 'text' in delta:
                            content_text = delta['text']
                            total_content += content_text
                            
                            # 创建流式响应块
                            response_chunk = create_streaming_response_chunk(
                                content=content_text,
                                finish_reason="stream",
                                meta=meta,
                                created_time=created_time,
                                model_name=model_name,
                                image_id=image_id
                            )
                            
                            # 发送 SSE 格式的响应
                            yield create_sse_response(response_chunk)
                
                # 发送最终响应块
                if total_content:
                    # 计算性能指标
                    end_time = time.time()
                    total_time_ms = (end_time - start_time) * 1000
                    estimated_tokens = len(total_content) // 4  # 粗略估算
                    tokens_per_second = estimated_tokens / (total_time_ms / 1000) if total_time_ms > 0 else 0
                    
                    runtime_status = json.dumps({
                        "total_time_ms": round(total_time_ms, 2),
                        "tokens_generated": estimated_tokens,
                        "tokens_per_second": round(tokens_per_second, 2),
                        "initialization_time_ms": 0,
                        "prompt_processing_time_ms": 0,
                        "token_generation_time_ms": round(total_time_ms, 2)
                    })
                    
                    # 创建最终响应块
                    final_chunk = create_streaming_response_chunk(
                        content="",  # 最终块通常不包含新内容
                        finish_reason="stop",
                        meta=meta,
                        created_time=created_time,
                        model_name=model_name,
                        image_id=image_id
                    )
                    final_chunk["runtime_status"] = runtime_status
                    
                    yield create_sse_response(final_chunk)
                
                # 发送流结束标记
                yield "data: [DONE]\n\n"
                
                logger.log_structured("api_gateway_streaming_completed", {
                    "request_id": request_id,
                    "total_chunks": chunk_count,
                    "total_content_length": len(total_content),
                    "processing_time_ms": round((time.time() - start_time) * 1000, 2)
                })
                
            except Exception as e:
                logger.log_error(
                    error=e,
                    context={"request_id": request_id},
                    error_phase="streaming_generator"
                )
                
                # 发送错误响应
                error_chunk = {
                    "error": {
                        "type": "streaming_error",
                        "message": "流式响应处理失败",
                        "code": 500,
                        "request_id": request_id
                    }
                }
                yield create_sse_response(error_chunk)
                yield "data: [DONE]\n\n"
        
        # 返回流式响应 - API Gateway 伪流式（收集所有块后返回）
        sse_response_parts = []
        total_content = ""
        chunk_count = 0
        
        try:
            # 调用 Bedrock 流式 API
            for chunk in bedrock_client.invoke_model_stream(bedrock_request):
                chunk_count += 1
                
                if chunk.type == 'content_block_delta' and chunk.content_block_delta:
                    delta = chunk.content_block_delta.get('delta', {})
                    
                    if 'text' in delta:
                        content_text = delta['text']
                        total_content += content_text
                        
                        # 创建流式响应块
                        response_chunk = create_streaming_response_chunk(
                            content=content_text,
                            finish_reason="stream",
                            meta=meta,
                            created_time=created_time,
                            model_name=model_name,
                            image_id=image_id
                        )
                        
                        # 添加到 SSE 响应中
                        sse_response_parts.append(create_sse_response(response_chunk))
            
            # 添加最终响应块
            if total_content:
                # 计算性能指标
                end_time = time.time()
                total_time_ms = (end_time - start_time) * 1000
                estimated_tokens = len(total_content) // 4  # 粗略估算
                tokens_per_second = estimated_tokens / (total_time_ms / 1000) if total_time_ms > 0 else 0
                
                runtime_status = json.dumps({
                    "total_time_ms": round(total_time_ms, 2),
                    "tokens_generated": estimated_tokens,
                    "tokens_per_second": round(tokens_per_second, 2),
                    "initialization_time_ms": 0,
                    "prompt_processing_time_ms": 0,
                    "token_generation_time_ms": round(total_time_ms, 2)
                })
                
                # 创建最终响应块
                final_chunk = create_streaming_response_chunk(
                    content="",  # 最终块通常不包含新内容
                    finish_reason="stop",
                    meta=meta,
                    created_time=created_time,
                    model_name=model_name,
                    image_id=image_id
                )
                final_chunk["runtime_status"] = runtime_status
                
                sse_response_parts.append(create_sse_response(final_chunk))
            
            # 添加流结束标记
            sse_response_parts.append("data: [DONE]\n\n")
            
            # 合并所有 SSE 响应部分
            complete_sse_response = "".join(sse_response_parts)
            
            logger.log_structured("api_gateway_pseudo_streaming_completed", {
                "request_id": request_id,
                "total_chunks": chunk_count,
                "total_content_length": len(total_content),
                "sse_response_length": len(complete_sse_response),
                "processing_time_ms": round((time.time() - start_time) * 1000, 2)
            })
            
            # 返回 SSE 格式的响应
            return {
                'statusCode': 200,
                'headers': {
                    'Content-Type': 'text/event-stream',
                    'Cache-Control': 'no-cache',
                    'Connection': 'keep-alive',
                    'Access-Control-Allow-Origin': '*',
                    'Access-Control-Allow-Methods': 'POST, OPTIONS',
                    'Access-Control-Allow-Headers': 'Content-Type, Authorization',
                    'X-Total-Chunks': str(chunk_count),
                    'X-Content-Length': str(len(total_content)),
                    'X-Processing-Time': str(round((time.time() - start_time) * 1000, 2))
                },
                'body': complete_sse_response,
                'isBase64Encoded': False
            }
            
        except Exception as e:
            logger.log_error(
                error=e,
                context={"request_id": request_id},
                error_phase="bedrock_streaming"
            )
            
            # 创建错误的 SSE 响应
            error_chunk = {
                "error": {
                    "type": "streaming_error",
                    "message": "流式响应处理失败",
                    "code": 500,
                    "request_id": request_id
                }
            }
            error_sse = create_sse_response(error_chunk) + "data: [DONE]\n\n"
            
            return {
                'statusCode': 500,
                'headers': {
                    'Content-Type': 'text/event-stream',
                    'Cache-Control': 'no-cache',
                    'Access-Control-Allow-Origin': '*'
                },
                'body': error_sse,
                'isBase64Encoded': False
            }
        
    except Exception as e:
        logger.log_error(
            error=e,
            context={"request_id": request_id},
            error_phase="handle_streaming_request"
        )
        
        error_response = {
            'error': {
                'type': 'streaming_error',
                'message': 'API Gateway流式响应处理失败',
                'code': 500,
                'request_id': request_id
            }
        }
        return create_json_response(500, error_response)


def handle_non_streaming_request(bedrock_request, meta: Dict[str, Any], request_id: str, start_time: float) -> Dict[str, Any]:
    """
    处理非流式请求
    
    Args:
        bedrock_request: Bedrock 请求对象
        meta: 元数据
        request_id: 请求ID
        start_time: 开始时间
        
    Returns:
        API Gateway JSON 响应
    """
    try:
        created_time = int(time.time())
        model_name = Config.BEDROCK_MODEL_ID.split('.')[-1] if '.' in Config.BEDROCK_MODEL_ID else Config.BEDROCK_MODEL_ID
        image_id = meta.get('image_id', '') if meta else ''
        
        logger.log_structured("non_streaming_request_start", {
            "request_id": request_id,
            "model": model_name
        })
        
        # 调用 Bedrock 非流式 API
        bedrock_response = bedrock_client.invoke_model(bedrock_request)
        
        # 提取内容 - 修复内容提取逻辑
        content = ""
        if bedrock_response.content:
            for content_block in bedrock_response.content:
                # 检查不同的内容结构
                if hasattr(content_block, 'text') and content_block.text:
                    content += content_block.text
                elif isinstance(content_block, dict):
                    if 'text' in content_block:
                        content += content_block['text']
                    elif 'content' in content_block:
                        content += content_block['content']
                elif hasattr(content_block, 'content') and content_block.content:
                    content += content_block.content
        
        # 如果仍然没有内容，尝试从原始响应中提取
        if not content and hasattr(bedrock_response, '__dict__'):
            logger.log_structured("bedrock_response_debug", {
                "request_id": request_id,
                "response_type": type(bedrock_response).__name__,
                "response_attributes": list(bedrock_response.__dict__.keys()) if hasattr(bedrock_response, '__dict__') else [],
                "content_type": type(bedrock_response.content).__name__ if hasattr(bedrock_response, 'content') else 'no_content'
            })
        
        # 计算性能指标
        end_time = time.time()
        total_time_ms = (end_time - start_time) * 1000
        estimated_tokens = len(content) // 4  # 粗略估算
        tokens_per_second = estimated_tokens / (total_time_ms / 1000) if total_time_ms > 0 else 0
        
        runtime_status = json.dumps({
            "total_time_ms": round(total_time_ms, 2),
            "tokens_generated": estimated_tokens,
            "tokens_per_second": round(tokens_per_second, 2),
            "initialization_time_ms": 0,
            "prompt_processing_time_ms": 0,
            "token_generation_time_ms": round(total_time_ms, 2)
        })
        
        # 创建响应
        response_data = create_non_streaming_response(
            content=content,
            meta=meta,
            created_time=created_time,
            model_name=model_name,
            runtime_status=runtime_status,
            image_id=image_id
        )
        
        logger.log_structured("non_streaming_request_completed", {
            "request_id": request_id,
            "content_length": len(content),
            "processing_time_ms": round(total_time_ms, 2)
        })
        
        return create_json_response(200, response_data, {
            'X-Response-Time': str(round(total_time_ms, 2)),
            'X-Response-Tokens': str(estimated_tokens)
        })
        
    except Exception as e:
        logger.log_error(
            error=e,
            context={"request_id": request_id},
            error_phase="handle_non_streaming_request"
        )
        
        error_response = {
            'error': {
                'type': 'internal_error',
                'message': '非流式响应处理失败',
                'code': 500,
                'request_id': request_id
            }
        }
        return create_json_response(500, error_response)