How to have a bedrock agent reliably include knowledge base citations in the final response of InvokeAgent for Agents for Amazon Bedrock Runtime?

Question

We have an agent that has an extensive system prompt. I'm having trouble understanding what piece of the agent controls passing forward or including the KB citations in the final response. It will sometimes include them, but I want a way to do it more reliably. The trace doesn't actually include any helpful information as far as I can tell.

I'm referring specifically to this part of the chunk: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_PayloadPart.html#bedrock-Type-agent-runtime_PayloadPart-attribution

So far we've tried:
- Using the default orchestration and knowledge base prompts.
- Editing the orchestration prompt to include something like: "If you are using citations, links, or knowledge base entries in the response, you MUST include them in the citations. Forward all citations." 
- Editing the KB response generation prompt to include a similar line.
- Adding a similar line to the system prompt.

The agent will reliably look up from the KBs and use the info, but it's seemingly random if it makes it to the final response.
Can anyone shed any light on the exact mechanism that controls what makes it into the final `chunk.attribution.citations` array? I feel like we're shooting in the dark!

Will M · Accepted Answer

Ok, so, after a couple of challenging days, we managed to get citations to work fairly deterministically. It should be noted that the Bedrock team is aware of this issue and are working to improve it. This is still a new product and evolving in both documentation and usablilty. We used claude-3-5-sonnet-20241022-v2 for this agent.

How our solution works: As mentioned by the Agent above, the agentFinalResponse.citations with the 'FINISH ' type in the parser is where the citations actually make it to the response. A custom orchestrator doesn't seem to have access to this, but the custom parser does. You can verify this by hard coding in some citations. The default parsing seems broken, and/or the default orchestration template is not strongly worded enough to force the LLM to pass them through in the end step in the correct format. We fix that by setting up a custom parser to better extract the citations from XML generated from tool use and KB outputs, and by adding explicit XML formatting examples to the orchestration template. Note: editing the KB generation template doesn't do anything, as the override doesn't work with the custom parse it seems.

Step by step

Set up a custom parser lambda for your orchestration template. Here's a fixed version of the example custom parser override that integrates orchestration and KB/tool use. The logging is set up to help you determine if your response included citations.
Add the following to your overridden orchestration template. The part about including links is optional, we just have it because our agent shows the links as pop ups below its response.

<orchestration_guidelines>
CRITICAL: When generating your final response that uses knowledge base information:
- You MUST wrap each piece of information in <answer_part> tags
- Mention somewhere in the response that you're including link(s)
- The structure MUST be exactly:
<answer>
<answer_part>
<text>[your response text]</text>
<sources>
<source>[source ID]</source>
</sources>
</answer_part>
</answer>

<incorrect_example>
<answer>
Here's how to not do it if we need a source...
</answer>
</incorrect_example>

<correct_example>
<answer>
<answer_part>
<text>Here's how to do it if we need a source</text>
<sources>
<source>3</source>
</sources>
</answer_part>
</answer>
</correct_example>
</orchestration_guidelines>

After preparing your new version, you should now be getting more reliable citations. There's still the normal range of LLM non-determinism, but it's not broken any more.

Victor Okunev · Answer

I didn't have any luck with that. No matter how strictly I instructed the agent to include the original unchanged citations links, even with the "smartest" Claude 3.5 Sonnet v2 model, the agent was messing them up.

I ended up calling the KB directly and extracting the citations links from the response. That was the only reliable way. As a bonus, you can truly stream the KB response with retrieve_and_generate_stream API. The Agent "stream" on the other hand, always returns a single chunk. Not very user-firendly in chat applications.

re:Post Agent · Answer

Ensuring consistent inclusion of knowledge base citations in the final response of InvokeAgent for Agents for Amazon Bedrock Runtime can be challenging, as there isn't a direct mechanism to force citations in every response. However, there are several approaches you can take to improve the reliability of citation inclusion:

1. Customize the agent's prompts: You're on the right track by modifying various prompts. Continue to emphasize the importance of including citations in the orchestration, knowledge base response generation, and system prompts. Be explicit in your instructions, asking the agent to always include citations when information from the knowledge base is used.

2. Use the agentFinalResponse field: In the orchestration parsed response, make sure the agentFinalResponse object includes both the responseText and citations. This object is compatible with the FINISH invocation type and should contain the final response to the user along with the relevant citations.

3. Implement a custom parser Lambda function: By writing a custom parser Lambda function, you can have more control over how the agent processes and includes citations. This function can be designed to consistently extract and include citations from the knowledge base responses.

4. Optimize knowledge base queries: Ensure that your knowledge base queries are well-structured and specific. This can help the agent retrieve more relevant information, increasing the likelihood of citation inclusion.

5. Use consistent session IDs: When invoking the agent, use consistent session IDs to maintain context across interactions. This can help the agent better understand when and how to include citations based on the ongoing conversation.

6. Monitor and analyze responses: Regularly review the agent's responses and analyze patterns in citation inclusion. This can help you identify specific scenarios where citations are more likely to be omitted and adjust your prompts or configuration accordingly.

7. Consider model limitations: Be aware that some models may have limitations in their ability to consistently include citations. Ensure you're using a model that supports the citation feature reliably.

8. Implement post-processing: As a last resort, you could implement a post-processing step in your application logic to check for citations and, if missing, re-query the knowledge base to forcibly include relevant citations.

Remember that while these strategies can improve citation inclusion, achieving 100% reliability may not be possible due to the nature of language models and the complexity of natural language processing. Continue to iterate on your approach based on observed results and specific use case requirements.

If you're still experiencing issues after implementing these suggestions, it may be worth reaching out to AWS Support for more specific guidance tailored to your agent configuration and use case.  
**Sources**  
[AWS Bedrock Agent not returning citations on invokeAgent | AWS re:Post](https://repost.aws/questions/QUcN1ov9XZQXmg9_OTxJRCiw/aws-bedrock-agent-not-returning-citations-on-invokeagent)  
[Write a custom parser Lambda function in Amazon Bedrock Agents - Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/lambda-parser.html)  
[Maintaining Agent Session Memory in Bedrock Prompt Flow | AWS re:Post](https://repost.aws/questions/QUiv77tJumTKCej9QqhqULHg/maintaining-agent-session-memory-in-bedrock-prompt-flow)

How to have a bedrock agent reliably include knowledge base citations in the final response of InvokeAgent for Agents for Amazon Bedrock Runtime?

Add your answer

Relevant content