This is the lambda function I designed for my retrieval augmented generation system:

import os
import json
import boto3
import numpy as np
from typing import List, Dict, Any
from sagemaker.session import Session
from sagemaker.model import Model
from sagemaker.predictor import Predictor
from sagemaker.huggingface import HuggingFaceModel
from langchain.chains.question_answering import load_qa_chain
from langchain import PromptTemplate
from langchain.embeddings import SagemakerEndpointEmbeddings
from langchain.schema import Document
from langchain.vectorstores import Chroma
from langchain.llms.sagemaker_endpoint import LLMContentHandler, SagemakerEndpoint
from sagemaker.jumpstart.model import JumpStartModel
from langchain.llms.sagemaker_endpoint import LLMContentHandler, SagemakerEndpoint
from langchain.embeddings.sagemaker_endpoint import EmbeddingsContentHandler

from langchain.document_transformers import LongContextReorder
from langchain.chains import StuffDocumentsChain, LLMChain
from langchain.prompts import PromptTemplate

# Initialize SageMaker session and clients
sagemaker_session = Session()
aws_role = sagemaker_session.get_caller_identity_arn()
aws_region = boto3.Session().region_name
sagemaker_client = boto3.client("sagemaker")

# Define your model, embeddings, and other parameters as in the provided script
model_id = "meta-textgeneration-llama-2-7b-f"
model_version = "2.*"
my_model = JumpStartModel(model_id=model_id, model_version=model_version)

# Define other necessary classes and functions (e.g., embed_docs, ContentHandler)
# This code defines a function called embed_docs that takes a list of strings (docs) as input and returns a list of lists of floats (List[List[float]]). The function performs the following steps:

# - out = encoder.predict({"inputs": docs}): This line uses the predict method of the encoder object (an instance of the HuggingFaceModel class) to generate predictions for the given docs input. The input is provided as a dictionary where the key is “inputs” and the value is the docs list.
# - embeddings = np.mean(np.array(out), axis=1): This line calculates the mean of the embeddings in the out result obtained from the prediction. It uses numpy (np) to convert out into an array, and then applies np.mean along the second axis (axis=1) to calculate the mean for each embedding vector.
# - return embeddings.tolist(): This line converts the embeddings array into a list of lists by using the tolist() method. The resulting list is then returned as the output of the embed_docs function.

# Note that the encoder variable is assumed to be defined and refers to an instance of the HuggingFaceModel class, which is used for generating predictions.
def embed_docs(docs: List[str]) -> List[List[float]]:
    out = encoder.predict({"inputs": docs})
    embeddings = np.mean(np.array(out), axis=1)
    return embeddings.tolist()

# This code defines two classes: “SagemakerEndpointEmbeddingsJumpStart” and “ContentHandler”.

# The “SagemakerEndpointEmbeddingsJumpStart” class inherits from the 
# “SagemakerEndpointEmbeddings” class and adds a method called 
# “embed_documents”. This method takes a list of texts and a 
# chunk size as input and returns a list of embeddings, one for 
# # each text. The method loops over the texts and sends requests 
# to the “_embedding_func” method with chunked texts. The responses 
# are then extended to the results list.

# The “ContentHandler” class inherits from the “EmbeddingsContentHandler”
# class. It defines the content type and accepts properties for the handler. 
# It also provides methods to transform input and output. The “transform_input” 
# method takes a prompt string and a dictionary of model keyword arguments and 
# returns a JSON string encoding the input. The “transform_output” method takes 
# the output bytes and decodes it to a JSON string. It then calculates the mean 
# of the embeddings and converts it to a list.

# Finally, an instance of the “ContentHandler” class is created and assigned 
# to the “content_handler” variable. An instance of the 
# “SagemakerEndpointEmbeddingsJumpStart” class is also 
# created with specific parameters such as the endpoint name and content handler.


class SagemakerEndpointEmbeddingsJumpStart(SagemakerEndpointEmbeddings):
    def embed_documents(self, texts: List[str], chunk_size: int = 1) -> List[List[float]]:
        """Compute doc embeddings using a SageMaker Inference Endpoint.

        Args:
            texts: The list of texts to embed.
            chunk_size: The chunk size defines how many input texts will
                be grouped together as request. If None, will use the
                chunk size specified by the class.

        Returns:
            List of embeddings, one for each text.
        """
        results = []
        _chunk_size = len(texts) if chunk_size > len(texts) else chunk_size

        for i in range(0, len(texts), _chunk_size):
            response = self._embedding_func(texts[i : i + _chunk_size])
            results.extend(response)
        return results


class ContentHandlerEmbeddings(EmbeddingsContentHandler):
    content_type = "application/json"
    accepts = "application/json"
    
    def transform_input(self, prompt: str, model_kwargs={}) -> bytes:
        input_str = json.dumps({"inputs": prompt, **model_kwargs})
        return input_str.encode("utf-8")

    def transform_output(self, output: bytes) -> str:
        out = json.loads(output.read().decode("utf-8"))
        embeddings = np.mean(np.array(out), axis=1).tolist()
        # if len(embeddings) == 1:
        #     return [embeddings[0]]
        return embeddings


# This code imports the necessary modules and defines a 
# dictionary called “parameters” with several key-value 
# pairs for model parameters. It also defines another 
# dictionary called “endpoint_kwargs” with a custom attribute 
# “CustomAttributes” set to “accept_eula=true”.

# The class “ContentHandler” is defined and it inherits from 
# the “LLMContentHandler” class. It sets the “content_type” and 
# “accepts” properties for the handler. The class 
# provides two methods: “transform_input” and “transform_output”.
# The “transform_input” method takes a prompt string and a 
# dictionary of model keyword arguments. It prints the prompt, 
# sets a system prompt, and then encodes the input as a JSON string. 
# The “transform_output” method takes the output bytes, decodes it to a 
# JSON string, and retrieves the generated content from the response.

# An instance of the “ContentHandler” class is created and assigned to 
# the “content_handler” variable. Another instance of the “SagemakerEndpoint” 
# class is created with specific parameters such as the endpoint name, 
# region name, model keyword arguments, content handler, and endpoint 
# keyword arguments.

parameters = {
    "max_length": 4096,
    "num_return_sequences": 1,
    "top_k": 50,
    "top_p": 0.95,
    "do_sample": True,
    "temperature": 0.3
}

endpoint_kwargs = {
"CustomAttributes":"accept_eula=true"    
}

class ContentHandlerLLM(LLMContentHandler):
    content_type = "application/json"
    accepts = "application/json"

    def transform_input(self, prompt: str, model_kwargs: dict) -> bytes:
        system_prompt = """
        You are a helpful, respectful and honest assistant. Always answer as 
        helpfully as possible, while being safe. Your answers should not include
        any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content.
        Please, If a question does not make any sense, or is not factually coherent, explain 
        why instead of answering something not correct. If you don't know the answer 
        to a question, please don't share false information. Your goal is to provide answers relating to the 
        question.<>ensure that your responses are socially unbiased and positive in nature. Also, do NOT mention any 
        individual's names and personally identiafiable information such as date of birth, address, etc.
        """
        input_str = json.dumps({"inputs" : [[{"role" : "system",
        "content" : system_prompt},
        {"role" : "user", "content" : prompt}]],
        "parameters" : {**model_kwargs}})
        return input_str.encode('utf-8')
    
    def transform_output(self, output: bytes) -> str:
        response_json = json.loads(output.read().decode("utf-8"))
        return response_json[0]["generation"]["content"]

def load_index_from_s3(index_dir, embedding_function):
    bucket_name = 'hnb-ds-pr-ai-coe'
    folder_path_s3 = f'wiki_search/{index_dir}'

    # Initialize the S3 client
    s3_client = boto3.client('s3')

    # Check if the folder exists in S3
    try:
        s3_client.head_object(Bucket=bucket_name, Key=f"{folder_path_s3}/")
        folder_exists = True
    except Exception as e:
        folder_exists = False

    # If the folder exists or needs to be updated
    if not folder_exists or UPDATE:
        # Download the entire folder and its contents recursively from S3 to a local directory
        local_folder_path = "./chroma_db_wikisearch"
        os.makedirs(local_folder_path, exist_ok=True)

        # List all objects in the specified S3 folder
        objects = s3_client.list_objects(Bucket=bucket_name, Prefix=folder_path_s3)

        # Download each object (file or subfolder) from S3 to the local directory
        for obj in objects.get('Contents', []):
            s3_object_key = obj['Key']
            local_file_path = os.path.join(local_folder_path, os.path.relpath(s3_object_key, folder_path_s3))

            # Create local directories if they don't exist
            os.makedirs(os.path.dirname(local_file_path), exist_ok=True)

            # Download the object from S3
            s3_client.download_file(bucket_name, s3_object_key, local_file_path)

    # Load from disk
    chromadb = Chroma(persist_directory=local_folder_path, embedding_function=embedding_function)
    documents_with_metadata = chromadb.get(include=["embeddings","metadatas"])
    
    return chromadb, documents_with_metadata

# # Run a sample prompt and get the answer
# prompt = "Which instances can I use with Managed Spot Training in SageMaker?"
# print(f"Response:\n {sm_llm(prompt)}")

# AWS Lambda Handler Function
def lambda_handler(event, context):
    # Process input event
    query = event.get("query", "Default query if not provided")
    
    # Initialize your components (SagemakerEndpoint, Chroma, etc.) here
    content_handler_embeddings = ContentHandlerEmbeddings()
    content_handler_llm = ContentHandlerLLM()
    
    # Get embedding function from endpoint
    embedding_function = SagemakerEndpointEmbeddingsJumpStart(
        endpoint_name="minilm-embedding",
        region_name=aws_region,
        content_handler=content_handler_embeddings,
    )
    
    # Get LLM response generator from endpoint
    sm_llm = SagemakerEndpoint(
        endpoint_name="llama-2-rag",
        region_name=aws_region,
        model_kwargs=parameters,
        content_handler=content_handler_llm,
        endpoint_kwargs=endpoint_kwargs
    )
    
    index_dir = "chroma_db_wikisearch"
    load_from_s3 = True
    
    if load_from_s3:
        # load from S3
        chromadb, documents_with_metadata = load_index_from_s3(index_dir, embedding_function)
    
    else:
    
        # Load from disk
        chromadb = Chroma(persist_directory=index_dir, embedding_function=embedding_function)
        documents_with_metadata = chromadb.get(include=["embeddings","metadatas"])

    # Execute your logic
    retriever = chromadb.as_retriever(search_kwargs={"k": 3})
    context_set = retriever.get_relevant_documents(query)
    reordering = LongContextReorder()
    reordered_docs = reordering.transform_documents(context_set)
    
    document_prompt = PromptTemplate(
        input_variables=["page_content"], template="{page_content}"
    )
    document_variable_name = "context"
    stuff_prompt_override = """\n
\n\nGiven this text extracts
    :\n-----\n{context}\n-----\n
    Please answer the following question:\n{query}
\n
    .If the provided content lacks enough information, just say you don't know and say nothing more!"""
    prompt = PromptTemplate(
        template=stuff_prompt_override, input_variables=["context", "query"]
    )
    llm_chain = LLMChain(llm=sm_llm, prompt=prompt)
    chain = StuffDocumentsChain(
        llm_chain=llm_chain,
        document_prompt=document_prompt,
        document_variable_name=document_variable_name,
    )
    result = chain.run(input_documents=reordered_docs, query=query)

    # Return response
    return {
        "statusCode": 200,
        "body": json.dumps(result)
    }

def test_lambda_handler():
    test_event = {
    "query": "How should I handle sample bias in a machine learning model?"
    }
    
#     - "how to move multiple files from desktop to AWS?",
#     - "how to run pyspark in AWS?",
#     - "how to access snowflake in AWS?",
#     - "how to access sas server?",
#     - "How do I clone a git repo?",
#     - "What is the lowest cost instance type?",
#     - "How do I use the AWS CLI?",
#     - "How do I build a sagemaker pipeline?",
#     - "How do I read csv files from S3 into a sagemaker notebook?",
#     - "What is the data science repo naming convention?",
#     - "How do I create an Athena table in the data science account?",
#     - "Where is a fun place to eat?"
#     - "How do I access share drives?"
#     - "How do I request a work phone?"
#     - "How to make sure tagging script is complete?"
#     - "How do I build test in python?"
#     - "How can I run a processing step?"
#     - "How do get initial access to our git repos?"
#     - "What are the available instance sizes for AWS?"
#     - "Where a fun place to eat in columbus?"
    
    print('Testing Lambda ...')
    response = lambda_handler(test_event, None)
    # print(f"Response:\n{response['body'].strip()}")
    assert 'statusCode' in response
    assert 'body' in response
    return response
    

import time

def _post_process(text):
    return text.replace("\\n", "\n")

if __name__ == '__main__':
    start_time = time.time()
    response = test_lambda_handler()
    text = _post_process(response['body'])
    end_time = time.time()

    execution_time = end_time - start_time
    print(f"Execution time: {execution_time} seconds")

        
What types of unit test you can think to write for this so our lambda becomes better production ready? Nmae the tests and some explanation.