session = boto3.Session(region_name=region)
credentials = session.get_credentials()

awsauth = AWS4Auth(
    credentials.access_key,
    credentials.secret_key,
    region,
    'es',  
    session_token=credentials.token
)

os_client = OpenSearch(
    hosts=[{
        'host': opensearch_url.replace("https://", ""), 
        'port': 443
    }],
    http_compress=True,
    http_auth=awsauth,
    use_ssl=True,
    verify_certs=True,
    ssl_assert_hostname=False,
    ssl_show_warn=False,
    connection_class=RequestsHttpConnection
)

from langchain_community.vectorstores.opensearch_vector_search import OpenSearchVectorSearch
vectorstore = OpenSearchVectorSearch(
    index_name=index_name,  
    is_aoss = False,
    embedding_function=bedrock_embeddings,
    opensearch_url=opensearch_url,
    http_auth=awsauth,
    connection_class=RequestsHttpConnection
)

def get_contextual_text(whole_text, splitted_text, llm): 
    contextual_template = (
        "<document>"
        "{WHOLE_DOCUMENT}"
        "</document>"
        "Here is the chunk we want to situate within the whole document."
        "<chunk>"
        "{CHUNK_CONTENT}"
        "</chunk>"
        "Please give a short succinct context to situate this chunk within the overall document for the purposes of improving search retrieval of the chunk."
        "Answer only with the succinct context and nothing else in English."
        "Put it in <result> tags."
    )              
    contextual_prompt = ChatPromptTemplate([
        ('human', contextual_template)
    ])

    contextual_text = ""        
    contexual_chain = contextual_prompt | llm            
    response = contexual_chain.invoke(
        {
            "WHOLE_DOCUMENT": whole_text,
            "CHUNK_CONTENT": splitted_text
        }
    )    
    output = response.content
    return output[output.find('<result>')+8:output.find('</result>')]

def store_image_for_opensearch(key):
    image_obj = s3_client.get_object(Bucket=s3_bucket, Key=key)
    image_content = image_obj['Body'].read()
    img = Image.open(BytesIO(image_content))
                        
    width, height = img.size 
    print(f"width: {width}, height: {height}, size: {width*height}")
            
    isResized = False
    while(width*height > 5242880):
        width = int(width/2)
        height = int(height/2)
        isResized = True
        print(f"width: {width}, height: {height}, size: {width*height}")
           
    buffer = BytesIO()
    img.save(buffer, format="PNG")
    img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
                                                            
    llm = get_model()
    text = extract_text(llm, img_base64)
    extracted_text = text[text.find('<result>')+8:text.find('</result>')] 
    
    contextual_text = object_meta["contextual_text"]
    summary = summary_image(llm, img_base64, contextual_text)
    image_summary = summary[summary.find('<result>')+8:summary.find('</result>')]    
    contents = f"[이미지 요약]\n{image_summary}\n\n[추출된 텍스트]\n{extracted_text}"
    
    page = object_meta["page"]
    docs = []
    docs.append(
        Document(
            page_content=contents,
            metadata={
                'name': key,
                'page': page,
                'url': path+parse.quote(key)
            }
        )
    )
    return add_to_opensearch(docs)                                                                                                      

def summary_image(llm, img_base64, contextual_text):  
    query = "이미지가 의미하는 내용을 풀어서 자세히 알려주세요. markdown 포맷으로 답변을 작성합니다."
    if contextual_text:
        query += "\n아래 <reference>는 이미지와 관련된 내용입니다. 이미지 분석시 참고하세요. \n<reference>\n"+contextual_text+"\n</reference>"    
    messages = [
        HumanMessage(
            content=[
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{img_base64}", 
                    },
                },
                {
                    "type": "text", "text": query
                },
            ]
        )
    ]    
    result = llm.invoke(messages)
    extracted_text = result.content        
    return extracted_text

parent_splitter = RecursiveCharacterTextSplitter(
    chunk_size=2000,
    chunk_overlap=100,
    separators=["\n\n", "\n", ".", " ", ""],
    length_function = len,
)
child_splitter = RecursiveCharacterTextSplitter(
    chunk_size=400,
    chunk_overlap=50,
    length_function = len,
)

splitted_docs = parent_splitter.split_documents(docs)
parent_docs, contexualized_chunks = get_contextual_docs_using_parallel_processing(docs[-1], splitted_docs)

for i, doc in enumerate(parent_docs):
    doc.metadata["doc_level"] = "parent"
        
parent_doc_ids = vectorstore.add_documents(parent_docs, bulk_size = 10000)
ids = parent_doc_ids

for i, doc in enumerate(splitted_docs):
    _id = parent_doc_ids[i]
    child_docs = child_splitter.split_documents([doc])
    for _doc in child_docs:
        _doc.metadata["parent_doc_id"] = _id
        _doc.metadata["doc_level"] = "child"

    contexualized_child_docs = []
    for _doc in child_docs:
        contexualized_child_docs.append(
            Document(
                page_content=contexualized_chunks[i]+"\n\n"+_doc.page_content,
                metadata=_doc.metadata
            )
        )
    child_docs = contexualized_child_docs

    child_doc_ids = vectorstore.add_documents(child_docs, bulk_size = 10000)        
    ids += child_doc_ids           

{
  "mcpServers": {
    "awslabs.bedrock-kb-retrieval-mcp-server": {
      "command": "uvx",
      "args": ["awslabs.bedrock-kb-retrieval-mcp-server@latest"],
      "env": {
        "AWS_PROFILE": "your-profile-name",
        "AWS_REGION": "us-east-1",
        "FASTMCP_LOG_LEVEL": "ERROR",
        "KB_INCLUSION_TAG_KEY": "optional-tag-key-to-filter-kbs",
        "BEDROCK_KB_RERANKING_ENABLED": "false"
      },
      "disabled": false,
      "autoApprove": []
    }
  }
}

{
    "mcpServers": {
        "knowledge_base_custom": {
            "command": "python",
            "args": [
                "application/mcp_server_knowledge_base.py"
            ],
            "env": {
                "KB_INCLUSION_TAG_KEY": "mcp-rag"
            }
        }
    }
}

@mcp.tool()
def knowledge_base_search(keyword: str) -> list:
    """
    Search the knowledge base with the given keyword.
    keyword: the keyword to search
    return: the result of search
    """

    return rag.retrieve_knowledge_base(keyword)

def retrieve_knowledge_base(query):
    lambda_client = boto3.client(
        service_name='lambda',
        region_name=bedrock_region
    )
    functionName = f"knowledge-base-for-{projectName}"

    mcp_env = utils.load_mcp_env()
    grading_mode = mcp_env['grading_mode']
    multi_region = mcp_env['multi_region']

    payload = {
        'function': 'search_rag',
        'knowledge_base_name': knowledge_base_name,
        'keyword': query,
        'top_k': numberOfDocs,
        'grading': grading_mode,
        'model_name': model_name,
        'multi_region': multi_region
    }

    output = lambda_client.invoke(
        FunctionName=functionName,
        Payload=json.dumps(payload),
    )
    payload = json.load(output['Payload'])        
    return payload['response']

{
    "mcpServers": {
        "opensearch-mcp-server": {
            "command": "uvx",
            "args": [
                "opensearch-mcp-server-py"
            ],
            "env": {
                "OPENSEARCH_URL": managed_opensearch_url,
                "AWS_REGION": aws_region,
                "AWS_ACCESS_KEY_ID": credentials.access_key,
                "AWS_SECRET_ACCESS_KEY": credentials.secret_key
            }
        }
    }
}

@mcp.tool()
def opensearch_search(keyword: str) -> list:
    """
    Search the knowledge base with the given keyword.
    keyword: the keyword to search
    return: the result of search
    """

    return rag.retrieve_opensearch(keyword)

def retrieve_opensearch(query):
    lambda_client = boto3.client(
        service_name='lambda',
        region_name=bedrock_region
    )
    functionName = f"opensearch-for-{projectName}"

    mcp_env = utils.load_mcp_env()
    grading_mode = mcp_env['grading_mode']
    multi_region = mcp_env['multi_region']

    payload = {
        'function': 'search_rag',
        'keyword': query,
        'top_k': numberOfDocs,
        'grading': grading_mode,
        'model_name': model_name,
        'multi_region': multi_region
    }
    output = lambda_client.invoke(
        FunctionName=functionName,
        Payload=json.dumps(payload),
    )
    payload = json.load(output['Payload'])        
    return payload['response']

cd gateway/opensearch-retriever && cp config.json.sample config.json

{
  "projectName": "mcp-rag",
  "opensearch_url": "https://search-mcp-rag-mxtkul2z3qv5iiqprb7q3jx4wy.us-west-2.es.amazonaws.com",
  "sharing_url": "https://d20lfnyi6fvd87.cloudfront.net"
}

python create_gateway_role.py

python create_gateway_tool.py

test_mcp_remote.py

git clone https://github.com/kyopark2014/mcp-rag/

cd mcp-rag && pip install -r requirements.txt

cd cdk-mcp-rag/ && cdk deploy --all

{
    "projectName":"mcp-rag",
    "accountId":"862926741992",
    "region":"us-west-2",
    "roleKnowledgeBase":"arn:aws:iam::862926741992:role/role-knowledge-base-for-mcp-rag-us-west-2",
    "collectionArn":"arn:aws:aoss:us-west-2:862926741992:collection/8krsnuq4it9gpl70i3u6",
    "serverless_opensearch_url":"https://8krsnuq4it9gpl70i3u6.us-west-2.aoss.amazonaws.com",
    "managed_opensearch_url":"https://search-mcp-rag-mxtkul3z3qv5iiqprb7q3jx4wy.us-west-2.es.amazonaws.com",
    "knowledge_base_role":"arn:aws:iam::862926741992:role/role-knowledge-base-for-mcp-rag-us-west-2",
    "s3_bucket":"storage-for-mcp-rag-862926741992-us-west-2",
    "s3_arn":"arn:aws:s3:::storage-for-mcp-rag-862926741992-us-west-2",
    "sharing_url":"https://d3mo4kqj5cjiuy.cloudfront.net"
 }

streamlit run application/app.py