Module 5: AWS Services for AI (S3, Lambda, SageMaker Basics)

3. Lambda for AI Inference

Next →

Overview

In this capsule you'll build Lambda functions designed specifically for AI inference in a real AWS context. In Module 3 you learned Lambda fundamentals — handler, packaging, cold starts. In Module 4 you tested it on LocalStack. Now you go deeper into the patterns that make an AI Lambda run robustly in production: invoking multiple LLM providers, retry patterns with exponential backoff, structured output for downstream consumption, and error handling that distinguishes between transient and permanent errors.

Context: This capsule extends what you built in M3. You already know the base handler. Here you add the layers an AI service needs in production: resilience (retries), structure (output parsing), and flexibility (multi-provider). By the end, you'll have a production-ready Lambda handler you can invoke from API Gateway, from another Lambda, or as part of the S3 → Lambda flow from capsule 04.

Lambda for AI in a Real AWS Context

Difference from M3

In Module 3 you built a Lambda that invokes an LLM. It worked, but it was a first step. Now you add what production demands:

M3 (Lambda fundamentals):
├── Basic handler
├── One provider (OpenAI)
├── Simple error handling (generic try/except)
├── Fixed timeout
└── Output: flat JSON

M5 (this capsule):
├── Production-ready handler
├── Multi-provider (OpenAI + Anthropic)
├── Granular error handling (transient vs permanent)
├── Retry with exponential backoff
├── Structured output (parsing, validation)
├── Logging for debugging in CloudWatch
└── Output: structured JSON with metadata

What real AWS adds

On LocalStack your Lambda always has permissions and you don't pay. On real AWS:

  • Your Lambda needs an IAM role to access other services (S3, Secrets Manager). Without it, AccessDenied.
  • Every millisecond counts for cost. An unnecessary retry to an LLM doubles the cost of that invocation.
  • CloudWatch Logs is your debugging tool. You don't have access to the container. The logs are all you see.
  • Concurrency limits are real. If 1000 requests arrive simultaneously, AWS limits the concurrent executions.

Production-Ready Handler with Retry

Retry pattern for LLM calls

LLM calls fail. Rate limits, network timeouts, provider server errors. A production-ready handler needs to distinguish between errors worth retrying and permanent errors:

# handler.py — Lambda for AI inference with retry
import json
import logging
import os
import time
from openai import OpenAI, APITimeoutError, RateLimitError, APIConnectionError

logger = logging.getLogger()
logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))

MAX_RETRIES = int(os.environ.get("MAX_RETRIES", "3"))
BASE_DELAY = float(os.environ.get("RETRY_BASE_DELAY", "1.0"))
MODEL_NAME = os.environ.get("MODEL_NAME", "gpt-4o-mini")

client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY", ""),
    max_retries=0,  # We handle retries ourselves
)

RETRYABLE_ERRORS = (APITimeoutError, RateLimitError, APIConnectionError)

CORS_HEADERS = {
    "Content-Type": "application/json",
    "Access-Control-Allow-Origin": os.environ.get("ALLOWED_ORIGIN", "*"),
    "Access-Control-Allow-Methods": "POST, OPTIONS",
    "Access-Control-Allow-Headers": "Content-Type",
}


def invoke_llm_with_retry(messages: list, max_tokens: int, context) -> dict:
    """Invokes an LLM with retry and exponential backoff."""
    last_error = None

    for attempt in range(1, MAX_RETRIES + 1):
        remaining_ms = context.get_remaining_time_in_millis()
        if remaining_ms < 10000:
            raise TimeoutError(
                f"Insufficient Lambda time: {remaining_ms}ms remaining"
            )

        try:
            start = time.time()
            response = client.chat.completions.create(
                model=MODEL_NAME,
                messages=messages,
                max_tokens=max_tokens,
                timeout=(remaining_ms / 1000) - 5,
            )
            duration_ms = int((time.time() - start) * 1000)

            logger.info(json.dumps({
                "event": "llm_success",
                "attempt": attempt,
                "duration_ms": duration_ms,
                "tokens": response.usage.total_tokens,
            }))

            return {
                "content": response.choices[0].message.content,
                "model": response.model,
                "tokens_used": response.usage.total_tokens,
                "input_tokens": response.usage.prompt_tokens,
                "output_tokens": response.usage.completion_tokens,
                "attempt": attempt,
                "llm_duration_ms": duration_ms,
            }

        except RETRYABLE_ERRORS as e:
            last_error = e
            delay = BASE_DELAY * (2 ** (attempt - 1))
            logger.warning(json.dumps({
                "event": "llm_retry",
                "attempt": attempt,
                "error_type": type(e).__name__,
                "error": str(e),
                "next_delay_s": delay,
            }))
            if attempt < MAX_RETRIES:
                time.sleep(delay)

        except Exception as e:
            logger.error(json.dumps({
                "event": "llm_permanent_error",
                "attempt": attempt,
                "error_type": type(e).__name__,
                "error": str(e),
            }))
            raise

    raise last_error


def handler(event, context):
    """Lambda handler for AI inference."""
    start_time = time.time()

    method = event.get("requestContext", {}).get("http", {}).get("method", "")
    if method == "OPTIONS":
        return {"statusCode": 200, "headers": CORS_HEADERS, "body": ""}

    try:
        body = json.loads(event.get("body", "{}"))
    except (json.JSONDecodeError, TypeError):
        return _response(400, {"error": "Invalid JSON body"})

    prompt = body.get("prompt", "").strip()
    if not prompt:
        return _response(400, {"error": "prompt is required"})

    system_prompt = body.get(
        "system_prompt",
        "Respond clearly, concisely, and helpfully."
    )
    max_tokens = min(body.get("max_tokens", 500), 2000)

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt},
    ]

    try:
        result = invoke_llm_with_retry(messages, max_tokens, context)
    except TimeoutError as e:
        return _response(408, {"error": str(e)})
    except RETRYABLE_ERRORS as e:
        return _response(502, {
            "error": f"LLM unavailable after {MAX_RETRIES} retries: {str(e)}"
        })
    except Exception as e:
        return _response(502, {"error": f"LLM error: {str(e)}"})

    total_ms = int((time.time() - start_time) * 1000)

    return _response(200, {
        "answer": result["content"],
        "model": result["model"],
        "tokens_used": result["tokens_used"],
        "duration_ms": total_ms,
        "llm_duration_ms": result["llm_duration_ms"],
        "retries": result["attempt"] - 1,
    })


def _response(status_code: int, body: dict) -> dict:
    return {
        "statusCode": status_code,
        "headers": CORS_HEADERS,
        "body": json.dumps(body),
    }

Anatomy of the retry

Attempt 1: Call the LLM
    → Success → Return result
    → RateLimitError → Wait 1s → Attempt 2

Attempt 2: Call the LLM
    → Success → Return result
    → APITimeoutError → Wait 2s → Attempt 3

Attempt 3: Call the LLM
    → Success → Return result
    → Error → Return error to the client (retries exhausted)

Permanent errors (not retried):
├── AuthenticationError → Invalid API key
├── BadRequestError → Invalid prompt
└── PermissionDeniedError → No access to the model

Multi-Provider Handler

Handler that supports OpenAI and Anthropic

# multi_provider_handler.py
import json
import logging
import os
import time
from openai import OpenAI
from anthropic import Anthropic

logger = logging.getLogger()
logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))

openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", ""))
anthropic_client = Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY", ""))


def _invoke_openai(messages: list, max_tokens: int, timeout: float) -> dict:
    response = openai_client.chat.completions.create(
        model=os.environ.get("OPENAI_MODEL", "gpt-4o-mini"),
        messages=messages,
        max_tokens=max_tokens,
        timeout=timeout,
    )
    return {
        "content": response.choices[0].message.content,
        "provider": "openai",
        "model": response.model,
        "tokens_used": response.usage.total_tokens,
    }


def _invoke_anthropic(messages: list, max_tokens: int, timeout: float) -> dict:
    system_msg = ""
    user_messages = []
    for msg in messages:
        if msg["role"] == "system":
            system_msg = msg["content"]
        else:
            user_messages.append(msg)

    response = anthropic_client.messages.create(
        model=os.environ.get("ANTHROPIC_MODEL", "claude-3-haiku-20240307"),
        system=system_msg,
        messages=user_messages,
        max_tokens=max_tokens,
        timeout=timeout,
    )
    return {
        "content": response.content[0].text,
        "provider": "anthropic",
        "model": response.model,
        "tokens_used": response.usage.input_tokens + response.usage.output_tokens,
    }


PROVIDERS = {
    "openai": _invoke_openai,
    "anthropic": _invoke_anthropic,
}


def handler(event, context):
    try:
        body = json.loads(event.get("body", "{}"))
    except (json.JSONDecodeError, TypeError):
        return _response(400, {"error": "Invalid JSON body"})

    prompt = body.get("prompt", "").strip()
    provider = body.get("provider", "openai").lower()
    system_prompt = body.get("system_prompt", "Respond clearly and helpfully.")
    max_tokens = min(body.get("max_tokens", 500), 2000)

    if not prompt:
        return _response(400, {"error": "prompt is required"})

    if provider not in PROVIDERS:
        return _response(400, {
            "error": f"Provider '{provider}' not supported. Use: {list(PROVIDERS.keys())}"
        })

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": prompt},
    ]

    remaining_s = context.get_remaining_time_in_millis() / 1000
    timeout = remaining_s - 5

    start = time.time()
    try:
        result = PROVIDERS[provider](messages, max_tokens, timeout)
    except Exception as e:
        return _response(502, {"error": f"{provider} error: {str(e)}"})

    duration_ms = int((time.time() - start) * 1000)
    result["duration_ms"] = duration_ms

    return _response(200, result)


def _response(status_code: int, body: dict) -> dict:
    return {
        "statusCode": status_code,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps(body),
    }

Structured Output

Parsing structured LLM responses

When Lambda produces output that another system consumes (another Lambda, S3, a database), you need structured output — not free text:

import json
import os
from openai import OpenAI

client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", ""))


def invoke_with_structured_output(
    prompt: str,
    output_schema: dict,
    system_prompt: str = "Always respond in valid JSON.",
) -> dict:
    """Invokes an LLM and parses the response as structured JSON."""
    schema_instruction = (
        f"Respond ONLY with JSON that follows this schema:\n"
        f"{json.dumps(output_schema, indent=2)}\n"
        f"Do not include any additional text, only the JSON."
    )

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": f"{system_prompt}\n\n{schema_instruction}"},
            {"role": "user", "content": prompt},
        ],
        max_tokens=1000,
        response_format={"type": "json_object"},
    )

    raw_content = response.choices[0].message.content

    try:
        parsed = json.loads(raw_content)
    except json.JSONDecodeError:
        raise ValueError(f"LLM did not return valid JSON: {raw_content[:200]}")

    return {
        "data": parsed,
        "tokens_used": response.usage.total_tokens,
        "model": response.model,
    }


result = invoke_with_structured_output(
    prompt="Analyze the sentiment of: 'The product is excellent but shipping was slow'",
    output_schema={
        "sentiment": "positive | negative | mixed",
        "confidence": 0.95,
        "aspects": [
            {"aspect": "product", "sentiment": "positive"},
            {"aspect": "shipping", "sentiment": "negative"},
        ],
    },
)

print(json.dumps(result["data"], indent=2))

Lambda handler with structured output

def handler(event, context):
    """Lambda that returns structured output for downstream consumption."""
    try:
        body = json.loads(event.get("body", "{}"))
    except (json.JSONDecodeError, TypeError):
        return _response(400, {"error": "Invalid JSON"})

    prompt = body.get("prompt", "").strip()
    output_format = body.get("output_format", "text")

    if not prompt:
        return _response(400, {"error": "prompt is required"})

    if output_format == "json":
        try:
            result = invoke_with_structured_output(
                prompt=prompt,
                output_schema=body.get("schema", {}),
            )
            return _response(200, {
                "result": result["data"],
                "format": "json",
                "tokens_used": result["tokens_used"],
            })
        except ValueError as e:
            return _response(422, {"error": str(e)})

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=500,
    )

    return _response(200, {
        "result": response.choices[0].message.content,
        "format": "text",
        "tokens_used": response.usage.total_tokens,
    })

Granular Error Handling

Error classification

from openai import (
    APITimeoutError,
    RateLimitError,
    APIConnectionError,
    AuthenticationError,
    BadRequestError,
    PermissionDeniedError,
    InternalServerError,
)


def classify_error(error: Exception) -> dict:
    """Classifies an LLM error to decide how to handle it."""
    error_map = {
        APITimeoutError: {
            "category": "transient",
            "retry": True,
            "status_code": 504,
            "message": "LLM request timed out",
        },
        RateLimitError: {
            "category": "transient",
            "retry": True,
            "status_code": 429,
            "message": "LLM rate limit exceeded",
        },
        APIConnectionError: {
            "category": "transient",
            "retry": True,
            "status_code": 502,
            "message": "Cannot connect to LLM provider",
        },
        InternalServerError: {
            "category": "transient",
            "retry": True,
            "status_code": 502,
            "message": "LLM provider internal error",
        },
        AuthenticationError: {
            "category": "permanent",
            "retry": False,
            "status_code": 401,
            "message": "Invalid LLM API key",
        },
        BadRequestError: {
            "category": "permanent",
            "retry": False,
            "status_code": 400,
            "message": "Invalid request to LLM",
        },
        PermissionDeniedError: {
            "category": "permanent",
            "retry": False,
            "status_code": 403,
            "message": "Permission denied by LLM provider",
        },
    }

    for error_type, info in error_map.items():
        if isinstance(error, error_type):
            return {**info, "original_error": str(error)}

    return {
        "category": "unknown",
        "retry": False,
        "status_code": 500,
        "message": f"Unexpected error: {type(error).__name__}",
        "original_error": str(error),
    }

Use in the handler

def handler_with_error_classification(event, context):
    """Handler that uses error classification for appropriate responses."""
    try:
        body = json.loads(event.get("body", "{}"))
    except (json.JSONDecodeError, TypeError):
        return _response(400, {"error": "Invalid JSON"})

    prompt = body.get("prompt", "").strip()
    if not prompt:
        return _response(400, {"error": "prompt is required"})

    try:
        result = invoke_llm_with_retry(
            messages=[{"role": "user", "content": prompt}],
            max_tokens=500,
            context=context,
        )
        return _response(200, {"answer": result["content"]})

    except Exception as e:
        error_info = classify_error(e)
        logger.error(json.dumps({
            "event": "handler_error",
            "category": error_info["category"],
            "error_type": type(e).__name__,
            "retry_attempted": error_info["retry"],
            "message": error_info["message"],
        }))
        return _response(error_info["status_code"], {
            "error": error_info["message"],
            "category": error_info["category"],
        })

Structured Logging for CloudWatch

Logs as JSON for CloudWatch Insights

import json
import logging
import os
import time

logger = logging.getLogger()
logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))


def log_inference(
    request_id: str,
    prompt_length: int,
    result: dict,
    duration_ms: int,
    cold_start: bool,
):
    """Structured log of an inference for CloudWatch Insights."""
    logger.info(json.dumps({
        "event": "inference_complete",
        "request_id": request_id,
        "prompt_length": prompt_length,
        "model": result.get("model"),
        "tokens_used": result.get("tokens_used"),
        "input_tokens": result.get("input_tokens"),
        "output_tokens": result.get("output_tokens"),
        "duration_ms": duration_ms,
        "llm_duration_ms": result.get("llm_duration_ms"),
        "overhead_ms": duration_ms - result.get("llm_duration_ms", 0),
        "retries": result.get("attempt", 1) - 1,
        "cold_start": cold_start,
    }))

With structured JSON logs, you can run queries in CloudWatch Insights:

-- Average inference latency
fields @timestamp, duration_ms, tokens_used, model
| filter event = "inference_complete"
| stats avg(duration_ms) as avg_latency, avg(tokens_used) as avg_tokens by model

-- Errors by category
fields @timestamp, category, error_type, message
| filter event = "handler_error"
| stats count(*) as error_count by category, error_type
| sort error_count desc

-- Cold starts vs warm starts
fields @timestamp, cold_start, duration_ms
| filter event = "inference_complete"
| stats count(*) as invocations, avg(duration_ms) as avg_ms by cold_start

Programmatic Lambda Invocation

Invoke Lambda from another service with boto3

import boto3
import json
import os

lambda_client = boto3.client(
    "lambda",
    endpoint_url=os.environ.get("AWS_ENDPOINT_URL"),
)


def invoke_ai_lambda(
    function_name: str,
    prompt: str,
    system_prompt: str = "",
    max_tokens: int = 500,
    async_mode: bool = False,
) -> dict:
    """Invokes an AI inference Lambda programmatically."""
    payload = {
        "body": json.dumps({
            "prompt": prompt,
            "system_prompt": system_prompt,
            "max_tokens": max_tokens,
        })
    }

    invocation_type = "Event" if async_mode else "RequestResponse"

    response = lambda_client.invoke(
        FunctionName=function_name,
        InvocationType=invocation_type,
        Payload=json.dumps(payload),
    )

    if async_mode:
        return {"status": "accepted", "status_code": response["StatusCode"]}

    response_payload = json.loads(response["Payload"].read())
    body = json.loads(response_payload.get("body", "{}"))
    return body


result = invoke_ai_lambda(
    function_name="ai-inference-prod",
    prompt="Summarize this document in 3 key points.",
    system_prompt="You are an assistant specialized in summaries.",
)
print(result)

Batch processing with parallel invocations

import concurrent.futures
import boto3
import json
import os

lambda_client = boto3.client(
    "lambda",
    endpoint_url=os.environ.get("AWS_ENDPOINT_URL"),
)


def batch_inference(
    function_name: str,
    prompts: list[str],
    max_workers: int = 5,
) -> list[dict]:
    """Processes multiple prompts in parallel by invoking Lambda."""

    def _invoke_single(prompt: str) -> dict:
        payload = {"body": json.dumps({"prompt": prompt, "max_tokens": 300})}
        response = lambda_client.invoke(
            FunctionName=function_name,
            InvocationType="RequestResponse",
            Payload=json.dumps(payload),
        )
        result = json.loads(response["Payload"].read())
        return json.loads(result.get("body", "{}"))

    results = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_to_prompt = {
            executor.submit(_invoke_single, p): p for p in prompts
        }
        for future in concurrent.futures.as_completed(future_to_prompt):
            prompt = future_to_prompt[future]
            try:
                result = future.result()
                results.append({"prompt": prompt, "result": result})
            except Exception as e:
                results.append({"prompt": prompt, "error": str(e)})

    return results

Troubleshooting

Problem 1: "Task timed out" after retry

Retries consume Lambda time. If your Lambda has a 60s timeout and you do 3 retries of 10s each, the third retry can cause a timeout.

remaining_ms = context.get_remaining_time_in_millis()
if remaining_ms < 15000:
    logger.warning("Skipping retry: insufficient time remaining")
    raise last_error

Problem 2: OpenAI rate limit with concurrent invocations

Multiple Lambdas invoking OpenAI simultaneously can exceed rate limits.

# Use reserved concurrency in Lambda to limit parallel invocations
# In template.yaml:
# ReservedConcurrentExecutions: 10
#
# This limits this Lambda to 10 simultaneous executions

Problem 3: Structured output with invalid JSON

The LLM sometimes returns malformed JSON or JSON with extra text.

import re

def safe_parse_json(raw: str) -> dict:
    """Tries to parse JSON from the LLM output, cleaning it up if needed."""
    try:
        return json.loads(raw)
    except json.JSONDecodeError:
        pass

    json_match = re.search(r'\{[\s\S]*\}', raw)
    if json_match:
        try:
            return json.loads(json_match.group())
        except json.JSONDecodeError:
            pass

    raise ValueError(f"Cannot parse JSON from LLM output: {raw[:200]}")

Problem 4: High cold start with anthropic + openai SDKs

Packaging both SDKs increases the cold start.

openai only:        cold start ~1.5s
openai + anthropic: cold start ~2.5s
openai + langchain: cold start ~8s

If you only use one provider, don't package the other. If you need both, consider container deployment.

Practical Exercises

Exercise 1: Handler with provider fallback

Implement a handler that tries OpenAI first. If it fails (timeout, rate limit), it tries Anthropic as a fallback. Log which provider responded.

See solution 
import json
import logging
import os
import time
from openai import OpenAI, APITimeoutError, RateLimitError, APIConnectionError
from anthropic import Anthropic

logger = logging.getLogger()
logger.setLevel("INFO")

openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", ""))
anthropic_client = Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY", ""))

OPENAI_RETRYABLE = (APITimeoutError, RateLimitError, APIConnectionError)


def _try_openai(prompt: str, max_tokens: int, timeout: float) -> dict:
    response = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_tokens,
        timeout=timeout,
    )
    return {
        "content": response.choices[0].message.content,
        "provider": "openai",
        "model": response.model,
        "tokens_used": response.usage.total_tokens,
    }


def _try_anthropic(prompt: str, max_tokens: int, timeout: float) -> dict:
    response = anthropic_client.messages.create(
        model="claude-3-haiku-20240307",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=max_tokens,
        timeout=timeout,
    )
    return {
        "content": response.content[0].text,
        "provider": "anthropic",
        "model": response.model,
        "tokens_used": response.usage.input_tokens + response.usage.output_tokens,
    }


def handler(event, context):
    try:
        body = json.loads(event.get("body", "{}"))
    except (json.JSONDecodeError, TypeError):
        return _response(400, {"error": "Invalid JSON"})

    prompt = body.get("prompt", "").strip()
    if not prompt:
        return _response(400, {"error": "prompt is required"})

    max_tokens = min(body.get("max_tokens", 500), 2000)
    timeout = (context.get_remaining_time_in_millis() / 1000) - 5

    start = time.time()

    try:
        result = _try_openai(prompt, max_tokens, timeout / 2)
        logger.info(json.dumps({"event": "inference", "provider": "openai", "fallback": False}))
    except OPENAI_RETRYABLE as openai_err:
        logger.warning(json.dumps({
            "event": "openai_failed",
            "error": str(openai_err),
            "falling_back": "anthropic",
        }))
        try:
            result = _try_anthropic(prompt, max_tokens, timeout / 2)
            logger.info(json.dumps({"event": "inference", "provider": "anthropic", "fallback": True}))
        except Exception as anthropic_err:
            return _response(502, {
                "error": f"Both providers failed. OpenAI: {openai_err}. Anthropic: {anthropic_err}"
            })
    except Exception as e:
        return _response(502, {"error": str(e)})

    result["duration_ms"] = int((time.time() - start) * 1000)
    return _response(200, result)


def _response(status_code: int, body: dict) -> dict:
    return {
        "statusCode": status_code,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps(body),
    }

Exercise 2: Lambda with internal rate limiting

Implement a simple rate-limiting mechanism inside the handler using a global variable (taking advantage of warm starts). Limit to N invocations per minute. If exceeded, return 429.

See solution 
import json
import os
import time
from collections import deque
from openai import OpenAI

client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", ""))
MAX_PER_MINUTE = int(os.environ.get("MAX_REQUESTS_PER_MINUTE", "30"))

invocation_timestamps = deque()


def _check_rate_limit() -> bool:
    """Checks the rate limit using a 60s sliding window."""
    now = time.time()
    while invocation_timestamps and invocation_timestamps[0] < now - 60:
        invocation_timestamps.popleft()
    return len(invocation_timestamps) < MAX_PER_MINUTE


def handler(event, context):
    if not _check_rate_limit():
        return {
            "statusCode": 429,
            "headers": {"Content-Type": "application/json", "Retry-After": "60"},
            "body": json.dumps({
                "error": "Rate limit exceeded",
                "limit": f"{MAX_PER_MINUTE} requests/minute",
                "current": len(invocation_timestamps),
            }),
        }

    invocation_timestamps.append(time.time())

    try:
        body = json.loads(event.get("body", "{}"))
    except (json.JSONDecodeError, TypeError):
        return _response(400, {"error": "Invalid JSON"})

    prompt = body.get("prompt", "").strip()
    if not prompt:
        return _response(400, {"error": "prompt is required"})

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        max_tokens=500,
    )

    return _response(200, {
        "answer": response.choices[0].message.content,
        "rate_limit_remaining": MAX_PER_MINUTE - len(invocation_timestamps),
    })


def _response(status_code: int, body: dict) -> dict:
    return {
        "statusCode": status_code,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps(body),
    }

Exercise 3: Structured output with schema validation

Create a handler that asks the LLM to classify a text and validates that the response matches the expected schema. If it doesn't, retry once asking the LLM to correct it.

See solution 
import json
import os
from openai import OpenAI

client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY", ""))

CLASSIFICATION_SCHEMA = {
    "category": "one of: bug_report, feature_request, question, feedback",
    "priority": "one of: low, medium, high, critical",
    "summary": "string, max 100 chars",
    "confidence": "float between 0.0 and 1.0",
}

VALID_CATEGORIES = {"bug_report", "feature_request", "question", "feedback"}
VALID_PRIORITIES = {"low", "medium", "high", "critical"}


def validate_classification(data: dict) -> list[str]:
    errors = []
    if data.get("category") not in VALID_CATEGORIES:
        errors.append(f"Invalid category: {data.get('category')}")
    if data.get("priority") not in VALID_PRIORITIES:
        errors.append(f"Invalid priority: {data.get('priority')}")
    if not isinstance(data.get("summary"), str) or len(data.get("summary", "")) > 100:
        errors.append("summary must be a string <= 100 chars")
    conf = data.get("confidence")
    if not isinstance(conf, (int, float)) or not (0.0 <= conf <= 1.0):
        errors.append("confidence must be float between 0.0 and 1.0")
    return errors


def classify_with_validation(text: str, max_attempts: int = 2) -> dict:
    schema_str = json.dumps(CLASSIFICATION_SCHEMA, indent=2)

    for attempt in range(1, max_attempts + 1):
        prompt = f"Classify this text:\n\n{text}\n\nRespond ONLY with JSON following this schema:\n{schema_str}"

        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            max_tokens=200,
            response_format={"type": "json_object"},
        )

        raw = response.choices[0].message.content
        try:
            data = json.loads(raw)
        except json.JSONDecodeError:
            if attempt < max_attempts:
                continue
            raise ValueError(f"Invalid JSON after {max_attempts} attempts")

        errors = validate_classification(data)
        if not errors:
            return {"classification": data, "attempt": attempt, "valid": True}

        if attempt < max_attempts:
            prompt += f"\n\nYour previous response had errors: {errors}. Correct them."

    return {"classification": data, "attempt": attempt, "valid": False, "errors": errors}


def handler(event, context):
    try:
        body = json.loads(event.get("body", "{}"))
    except (json.JSONDecodeError, TypeError):
        return _response(400, {"error": "Invalid JSON"})

    text = body.get("text", "").strip()
    if not text:
        return _response(400, {"error": "text is required"})

    try:
        result = classify_with_validation(text)
    except Exception as e:
        return _response(502, {"error": str(e)})

    return _response(200, result)


def _response(status_code, body):
    return {
        "statusCode": status_code,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps(body),
    }

Exercise 4: Invoke a Lambda from another Lambda

Write two Lambda functions: an "orchestrator" that receives a long document, splits it into chunks, and invokes another "summarizer" Lambda for each chunk. The orchestrator aggregates the summaries.

See solution 
# orchestrator.py
import boto3
import json
import os
import concurrent.futures

lambda_client = boto3.client(
    "lambda",
    endpoint_url=os.environ.get("AWS_ENDPOINT_URL"),
)
SUMMARIZER_FUNCTION = os.environ.get("SUMMARIZER_FUNCTION", "ai-summarizer")


def chunk_text(text: str, max_chars: int = 2000) -> list[str]:
    words = text.split()
    chunks = []
    current = []
    current_len = 0

    for word in words:
        if current_len + len(word) + 1 > max_chars and current:
            chunks.append(" ".join(current))
            current = [word]
            current_len = len(word)
        else:
            current.append(word)
            current_len += len(word) + 1

    if current:
        chunks.append(" ".join(current))
    return chunks


def invoke_summarizer(chunk: str, chunk_index: int) -> dict:
    payload = {
        "body": json.dumps({
            "prompt": f"Summarize this fragment in 2-3 sentences:\n\n{chunk}",
            "max_tokens": 200,
        })
    }
    response = lambda_client.invoke(
        FunctionName=SUMMARIZER_FUNCTION,
        InvocationType="RequestResponse",
        Payload=json.dumps(payload),
    )
    result = json.loads(response["Payload"].read())
    body = json.loads(result.get("body", "{}"))
    return {"chunk_index": chunk_index, "summary": body.get("answer", "")}


def handler(event, context):
    try:
        body = json.loads(event.get("body", "{}"))
    except (json.JSONDecodeError, TypeError):
        return _response(400, {"error": "Invalid JSON"})

    document = body.get("document", "").strip()
    if not document:
        return _response(400, {"error": "document is required"})

    chunks = chunk_text(document)

    with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
        futures = {
            executor.submit(invoke_summarizer, chunk, i): i
            for i, chunk in enumerate(chunks)
        }
        summaries = []
        for future in concurrent.futures.as_completed(futures):
            summaries.append(future.result())

    summaries.sort(key=lambda x: x["chunk_index"])
    combined = "\n\n".join(s["summary"] for s in summaries)

    return _response(200, {
        "total_chunks": len(chunks),
        "summaries": summaries,
        "combined_summary": combined,
    })


def _response(status_code, body):
    return {
        "statusCode": status_code,
        "headers": {"Content-Type": "application/json"},
        "body": json.dumps(body),
    }

Summary

  • Retry with exponential backoff is essential for an AI Lambda. LLMs have rate limits and timeouts — retrying with increasing waits resolves most transient errors.
  • Classify errors into transient (retry) and permanent (fail immediately). Don't retry an AuthenticationError.
  • Multi-provider (OpenAI + Anthropic) gives you resilience. If one provider fails, the other responds.
  • Structured output turns the LLM's free text into validatable JSON. Use response_format={"type": "json_object"} and validate the schema.
  • Structured logging (JSON) is your debugger in production. CloudWatch Insights lets you query your logs.
  • Programmatic invocation with boto3 lets another Lambda or service invoke your inference function.

Additional Resources

  1. OpenAI Python SDK — Error Handling — OpenAI SDK errors
  2. Anthropic Python SDK — Official Anthropic SDK
  3. Lambda Invocation Types — Sync vs Async
  4. CloudWatch Logs Insights — Queries over logs
  5. Lambda Concurrency — Reserved and provisioned concurrency
  6. OpenAI Structured Outputs — JSON mode and structured outputs
  7. Retry Pattern — AWS Architecture — Retry pattern in AWS
  8. Lambda Best Practices — Official best practices
← Previous2. S3 for AI AssetsNext →4. S3 + Lambda Integration