lambda_function.py

import json
import logging
from typing import Dict, Any

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

import json
import boto3
import uuid
from datetime import datetime

def extract_text_from_response(response_body):
    """
    Try different ways to extract text from Bedrock response (for non-streaming).
    """
    possible_paths = [
        ['choices', 0, 'text'],
        ['choices', 0, 'message', 'content'],
        ['generation'],
        ['outputs', 0, 'text'], 
        ['text'],
        ['output'],
        ['result'],
        ['response']
    ]
    
    for path in possible_paths:
        try:
            value = response_body
            for key in path:
                if isinstance(key, int):
                    value = value[key]
                else:
                    value = value.get(key)
                if value is None:
                    break
            if value and isinstance(value, str):
                return value.strip()
        except (KeyError, IndexError, TypeError):
            continue
    return str(response_body)

def error_response(status_code, message):
    return {
        'statusCode': status_code,
        'headers': {'Content-Type': 'application/json'},
        'body': json.dumps({
            'error': {
                'message': message,
                'type': 'api_error'
            }
        })
    }

def lambda_handler(event, context):
    try:
        images_base_64 = []
        body = json.loads(event['body'])
        messages = body.get('messages', [])
        
        print(f"Received messages: {json.dumps(messages)}")
        if not messages:
            return error_response(400, "Messages array is required")

        # --- Build prompt ---
        prompt_text = "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
        
        for mi, message in enumerate(messages):
            role = message.get('role', '')
            content = message.get('content')

            if isinstance(content, str):
                try:
                    content = json.loads(content)
                except json.JSONDecodeError:
                    content = [{"type": "text", "text": content}]

            if not isinstance(content, list):
                content = [{"type": "text", "text": str(content)}]

            # --- Last user message special handling ---
            if mi == len(messages) - 1 and role == "user":
                last_image = None
                for idx, item in enumerate(content):
                    if item.get('type') == 'text':
                        prompt_text += f"<|im_start|>{role}\n{item.get('text')}"
                        if idx == len(content) - 2:
                            prompt_text += "<|vision_start|><|image_pad|><|vision_end|>"
                        prompt_text += "<|im_end|>\n"
                    elif item.get('type') == 'image_url':
                        last_image = item.get('image_url', {}).get('url')

                if last_image:
                    if last_image.startswith("data:"):
                        b64_data = last_image.split(",", 1)[1]
                        images_base_64 = [b64_data]
                    else:
                        images_base_64 = [last_image]

            # --- Earlier messages ---
            else:
                for item in content:
                    if item.get('type') == 'text':
                        prompt_text += f"<|im_start|>{role}\n{item.get('text')}<|im_end|>\n"

        prompt_text += "<|im_start|>assistant\n"
        print(f"Generated prompt: {prompt_text}")

        # --- Bedrock invocation ---
        bedrock_runtime = boto3.client('bedrock-runtime')
        
        request_payload = {
            "prompt": prompt_text,
            "max_tokens_to_generate": body.get('max_tokens', 200),
            "temperature": body.get('temperature', 0.7),
            "top_p": body.get('top_p', 0.9)
        }
        if images_base_64:
            request_payload["images"] = images_base_64

        request_body = json.dumps(request_payload)
        model_id = 'arn:aws:bedrock:eu-central-1:535772221613:imported-model/93b67ujo8pps'

        print("Streaming not requested. Invoking model normally.")
        response = bedrock_runtime.invoke_model(
            modelId=model_id,
            contentType='application/json',
            accept='application/json',
            body=request_body
        )

        body_content = response['body']
        if hasattr(body_content, 'read'):
            body_content = body_content.read()
        response_body = json.loads(body_content)
        print(f"Raw Bedrock response: {json.dumps(response_body)}")

        generated_text = extract_text_from_response(response_body)
        if not generated_text:
            generated_text = "Sorry, I couldn't generate a response."

        openai_response = {
            "id": f"chatcmpl-{uuid.uuid4().hex[:8]}",
            "object": "chat.completion",
            "created": int(datetime.utcnow().timestamp()),
            "model": body.get('model', 'custom-model'),
            "system_fingerprint": "fp",
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": generated_text
                },
                "finish_reason": "stop",
                "logprobs": None,
            }],
            "usage": {
                "prompt_tokens": len(prompt_text.split()),
                "completion_tokens": len(generated_text.split()),
                "total_tokens": len(prompt_text.split()) + len(generated_text.split())
            }
        }
        
        print(f"Final response: {json.dumps(openai_response)}")

        return {
            'statusCode': 200,
            'headers': {'Content-Type': 'application/json'},
            'body': json.dumps(openai_response)
        }

    except Exception as e:
        print(f"Error: {str(e)}")
        return error_response(500, f"Internal server error: {str(e)}")


# For local testing
if __name__ == "__main__":
    # Mock context for local testing
    class MockContext:
        def __init__(self):
            self.aws_request_id = "local-test-request-id"
            self.function_name = "test-function"
            self.function_version = "1"
            self.invoked_function_arn = "arn:aws:lambda:us-east-1:123456789012:function:test-function"
            self.memory_limit_in_mb = 128
            self.remaining_time_in_millis = 30000
    
    # Test the function locally
    test_event = {
        'httpMethod': 'GET',
        'path': '/health',
        'headers': {},
        'body': None
    }
    
    context = MockContext()
    result = lambda_handler(test_event, context)
    print("Test Result:")
    print(json.dumps(result, indent=2))