invoke_agent API on bedrock-agent-runtime inconsistently returns retrievedReferences

[wadebev11]

Describe the bug

I have an Agent and a Knowledge Base set up in AWS Bedrock.

I'm trying to provide a citation link to the document in S3 that was used as part of RAG for a particular user query. The first response in the session has retrievedReferences that I can use to create this link. Any subsequent responses generated using RAG either have an empty array for retrievedReferences or don't even have an attribution field.

Expected Behavior

When RAG is used by the agent and it's pulling information from the Knowledge Base. The attribution field should be in the response and the retrievedReferences field should include the appropriate references rather than being an empty array.

Current Behavior

The first time the Agent uses the Knowledge Base, it returns the citations to the S3 files correctly:
agent_response_with_retrieved_references.json

Subsequent responses from the Agent with the same session id are missing attribution field or have empty arrays for retrievedReferences:
agent_response_missing_retrieved_references.json
agent_response_missing_attribution.json

Reproduction Steps

Call the invoke_agent API with multiple questions that will hit the Knowledge Base in the same session.

import uuid
import boto3
from typing import Dict, List

def call_invoke_agent_api(prompt: str, chat_session_id: str):

    bedrock_agent_runtime_client = boto3.client("bedrock-agent-runtime", region_name="us-east-1")
    agent_response = bedrock_agent_runtime_client.invoke_agent(
        agentId=***,
        agentAliasId=***,
        inputText=prompt,
        sessionId=chat_session_id,
        enableTrace=True,
    )

    text_completions: List[Dict] = []
    for event in agent_response["completion"]:
        chunk = event.get("chunk")
        if chunk is None:
            continue

        attribution = chunk.get("attribution")
        if attribution is not None:
            # the agent used RAG to generate a response
            citations = attribution.get("citations")
            for citation in citations:
                text = citation["generatedResponsePart"]["textResponsePart"]["text"]

                uri = filename = ""
                retrieved_references = citation["retrievedReferences"]
                for reference in retrieved_references:
                    uri = reference["location"]["s3Location"]["uri"]
                    filename = uri.split("/")[-1]

                text_completion = dict(text=text)

                if uri:
                    text_completion["uri"] = uri
                    text_completion["filename"] = filename

                text_completions.append(text_completion)

        else:
            # the agent didn't use the knowledge base to generate a response, just get it from bytes
            text_completions.append(dict(text=chunk["bytes"].decode()))

    return dict(text_completions=text_completions)

chat_session_id = uuid.uuid4().hex
call_invoke_agent_api(prompt_1, chat_session_id)
call_invoke_agent_api(prompt_2, chat_session_id)

Possible Solution

No response

Additional Information/Context

It seems to always provide the retrievedReferences correctly for the first prompt that causes the Agent to use the Knowledge Base. After that it will sometimes return valid references in retrievedReferences, sometimes empty arrays, and sometimes not even an attribution field in the response, even though the trace for invoke_agent call shows that the Agent is performing RAG and using the Knowledge Base.

In my Agent's Advanced prompts configuration, I have the Pre-processing template and Post-processing template turned off

SDK version used

boto3==1.34.61

Environment details (OS name and version, etc.)

AmazonLinux2

[tim-finnigan]

Thanks for reaching out. The invoke_agent command involves a call to the underlying Bedrock InvokeAgent API. This issue is currently being tracked here in our cross-SDK repository: aws/aws-sdk#746, where I saw you also commented. I can confirm that the Bedrock team is actively working on this issue and we should have more updates to share in the near future. Please refer to that issue for updates going forward.

[github-actions]

This issue is now closed. Comments on closed issues are hard for our team to see.
If you need more assistance, please open a new issue that references this one.

[Tutaki-dzhu]

We are experiencing the same issue with the .NET SDK as well as directly using the Bedrock Console. Don't think it is a boto3 issue. We've reached to AWS and they have escalated this to their internal team. Looks like an issue on the Bedrock side.