Knowledge Base Retrieval docs improvement

[ramirolc02]

Describe the bug
Knowledge Base integration in RAG pipeline not working. AIConversation component only returns model chat() output, with no signs of embedding retrieval.

To Reproduce
Steps to reproduce the behavior:

1. using NextJS App router, configure Knowledge Base as described in https://docs.amplify.aws/react/ai/conversation/knowledge-base/.

const schema = a.schema({
  knowledgeBase: a
    .query()
    .arguments({ input: a.string() })
    .handler(
      a.handler.custom({
        dataSource: "KnowledgeBaseDataSource",
        entry: "./kbResolver.js",
      }),
    )
    .returns(a.string())
    .authorization((allow) => allow.authenticated()),

  chat: a.conversation({
    aiModel: a.ai.model("Claude 3 Sonnet"),
    systemPrompt: "You are a knowledgeable pharmaceutical assistant specialized in medication regulations.",
    tools: [
      a.ai.dataTool({
        name: 'searchDocumentation',
        description: 'Searches pharmaceutical documentation.',
        query: a.ref('knowledgeBase'),
      }),
    ]
  })
    .authorization((allow) => allow.owner()),

2. Generate simple UI, similar to: Build a Travel Planner with React Native, AWS Amplify, and Amazon B...
3. When the sendMessage function is called, the model response is streamed back without performing database retrieval as specified in the backend tool.

const [
        {
            data: { messages, conversation },
            isLoading,
        },
        sendMessage,
    ] = useAIConversation("chat", { id })

Response received:

{
    "data": {
        "listConversationMessageChats": {
            "items": [
                {
                    "id": "8b1f9709-3776-48f2-ab93-c40386d66e10",
                    "conversationId": "48c6b6f6-66fc-43ef-9f14-15596388bf7e",
                    "role": "user",
                    "content": [
                        {
                            "text": "Europe regulations for pharma?",
                            "document": null,
                            "image": null,
                            "toolResult": null,
                            "toolUse": null
                        }
                    ],
                    "aiContext": null,
                    "toolConfiguration": null,
                    "createdAt": "2025-03-28T07:48:39.535Z",
                    "updatedAt": "2025-03-28T07:48:39.535Z",
                    "owner": "720564b4-b0c1-70d8-960b-25becf38c7d5"
                }
            ],
            "nextToken": null,
            "__typename": "ModelConversationMessageChatConnection"
        }
    }
}

Getting the following log:

"The model returned the following errors: Your API request included an assistant message in the final position, which would pre-fill the assistant response. When using tools, pre-filling the assistant response is not supported."

Expected behavior
Expected the retrieval step to be invoked before model inference. However, only model inference step is performed, without any prior embedding similarity search.

Screenshots
If applicable, add screenshots to help explain your problem.

Desktop (please complete the following information):

• OS: [iOS]
• Browser [firefox]
• Version [e.g. 22]

Smartphone (please complete the following information):

• Device: [e.g. iPhone6]
• OS: [e.g. iOS8.1]
• Browser [e.g. stock browser, safari]
• Version [e.g. 22]

Additional context
Region: eu-west-1 region.
Discord chat: https://discord.com/channels/705853757799399426/1353857342868947054

[malikalimoekhamedov]

+1

[jjarvisp]

Hi @ramirolc02, thanks for opening this. We will investigate and follow up here shortly.

[lets-getitnow]

+1 confirmed from us-west-2 this is also a fundamental issue. Bigger than a bug, because it's a major piece of functionality that doesn't work without explanation or erroring.

[stocaaro]

Hello @ramirolc02, @malikalimoekhamedov and @lets-getitnow,

I worked through the related tutorial / guides this morning, along the way I ran into a number of problems, but didn't encounter the assistant message related error noted above.

At a high level:

• I mostly depended on the amplify docs notes, but they correspond in details with the blog linked above.
• Debugging this was a pain as logging in the console is off by default
• I had the best experience by getting one piece working at a time
• I appreciated the blog knowledge base input for its specificity and went further by adding more fake precision to the population numbers so that I would know if they came back that they came from the knowledge base

What caused the issues I encountered:

• When I built the knowledge base and pushed a document into it, I expected it to automatically be synced, which didn't happen by default.
• My AppSync resolver was initially failing because in kbResolver.js I hasn't filled in the knowledge base ID (I caught the KBID in the backend.ts file and missed this instance leading to errors.

The pattern that helped me was to start with the knowledge base and test every layer until my Application chats work, using the kb for reference:

• Make sure the knowledge base works: Test your knowledge base in the console to make sure it is working. This involves picking a model an using the chat interface to pose questions. Location (AWS Console > Bedrock > Builder tools > Knowledge Bases which you created following this guide)
• Make sure the AppSync resolver works: Test your AppSync Knowledge Base Resolver (Location: AWS Console > AWS AppSync > Select your API > Queries > Folder Icon > knowledgeBase)
  • Here my knowledge base resolver needed me to be authenticated, so I had to click on the key icon and login otherwise I just got 401 errors
  • To find your API, look at the endpoint in amplify_outputs.json and find the AppSync instance with this endpoint if there are multiple.
• Make sure my chat conversation uses the resolver: Here I went back to my application chat.
  • My initial knowledge base was too ambiguous and I wasn't sure it was getting called so I:
    • Used the city location base with additional population specificity so I would know the data when I saw it
    • Turned on AppSync logging (Location: AWS Console > AWS AppSync > Select your API > Settings > API configurations - Click "Enable logging")
    • Next, you will need to create a log group by clicking on the log group name added to settings, which will tell you it doesn't exist, then copy the name /aws/appsync/apis/[replacewithyourloggroupid], click "Log groups" and then "Create log group" to create a log group for your appsync logs
    • Then under Monitoring click "View logs in Cloudwatch" then "Start tailing" then add a highlight term for your tool a.query name, knowledgebase from the tutorial I followed.
    • Last, the chat agent only occasionally used the tools until I added something like the following to its system prompt "Do not make up any information. Use available tools to answer user questions, otherwise tell the user you do not know"

This is what it took for me to troubleshoot through the tutorial and get it working for me. Hopefully there's something in here that helps you unblock. If there is a specific issue you are encountering details help, but better would be to share an example app, including your knowledge base so that I can try to reproduce your exact setup on my account to give specific help for the issue your encountering.

Thanks,
Aaron