lambda_function.py

import json
import base64
import boto3

KB_ID = 'A2Z6TQEJ4M'
MODEL_ARN = (
    'arn:aws:bedrock:ap-northeast-2::foundation-model/'
    'anthropic.claude-3-5-sonnet-20240620-v1:0'
)

bedrock_agent = boto3.client(
    'bedrock-agent-runtime',
    region_name='ap-northeast-2'
)

COMMON_HEADERS = {
    'Content-Type': 'application/json; charset=utf-8',
    'Access-Control-Allow-Origin': '*',
}


def _json_response(status, payload):
    return {
        'statusCode': status,
        'headers': COMMON_HEADERS,
        'body': json.dumps(payload, ensure_ascii=False),
    }


def _extract_question(event):
    body = event.get('body')

    if body:
        if event.get('isBase64Encoded'):
            body = base64.b64decode(body).decode('utf-8')

        try:
            data = json.loads(body)

            if isinstance(data, dict) and data.get('question'):
                return data['question']

        except (ValueError, json.JSONDecodeError):
            pass

    qs = event.get('queryStringParameters') or {}

    if qs.get('question'):
        return qs['question']

    return None


def _summarize_citations(raw_citations):
    citations = []

    for i, citation in enumerate(raw_citations or [], start=1):
        text_part = (
            citation
            .get('generatedResponsePart', {})
            .get('textResponsePart', {})
        )

        span = text_part.get('span') or {}
        references = []

        for ref in citation.get('retrievedReferences', []) or []:
            location = ref.get('location', {}) or {}
            s3_location = location.get('s3Location', {}) or {}
            content_text = (ref.get('content', {}) or {}).get('text', '')

            references.append({
                'source_uri': s3_location.get('uri') or location.get('type', ''),
                'snippet': (
                    content_text[:300] + '…'
                    if len(content_text) > 300
                    else content_text
                ),
                'metadata': ref.get('metadata', {}),
            })

        citations.append({
            'index': i,
            'answer_span': {
                'text': text_part.get('text', ''),
                'start': span.get('start'),
                'end': span.get('end'),
            },
            'references': references,
        })

    return citations


def lambda_handler(event, context):
    question = _extract_question(event)

    if not question:
        return _json_response(
            400,
            {
                'error': (
                    'question 이 필요합니다. '
                    'GET ?question=... 또는 '
                    'POST {"question": "..."}'
                )
            },
        )

    try:
        response = bedrock_agent.retrieve_and_generate(
            input={
                'text': question
            },
            retrieveAndGenerateConfiguration={
                'type': 'KNOWLEDGE_BASE',
                'knowledgeBaseConfiguration': {
                    'knowledgeBaseId': KB_ID,
                    'modelArn': MODEL_ARN,
                },
            },
        )

    except Exception as e:
        return _json_response(
            500,
            {
                'error': str(e)
            },
        )

    return _json_response(
        200,
        {
            'answer': response['output']['text'],
            'citations': _summarize_citations(response.get('citations')),
            'session_id': response.get('sessionId'),
        },
    )