invoke_agent API on bedrock-agent-runtime inconsistently returns retrievedReferences

[tim-finnigan]

Original issue: boto/boto3#4197. (ref: P140606677)

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]

@wadebev11 Can you try updating your version of Boto3 to confirm this is still an issue in the latest version?

[wadebev11]

Yes, after updating to boto3==1.34.143, I'm still seeing the same behavior

[edu2105]

Hi @tim-finnigan is there any update regarding this issue?
I noticed that, in the agent playground from the AWS console, the metadata information from the knowledge base documents are now being displayed when accessing the trace event details, however, even using the latest boto3 version 1.35.2 the metadata object it is not present in the invoke_agent response.

[Wildstorm1]

Hi I am wondering if this has been fixed? I am noticing the same issue with Java SDK V2

[TailorDone]

@tim-finnigan Also looking for an update on this issue. It's a big problem for our customers as we can only provide a citation for them the first time and no subsequent times.