Agentic retrieval with LangChain and Amazon Bedrock Knowledge Bases

來源原標題

Agentic retrieval with LangChain and Amazon Bedrock Knowledge Bases

作者

Manideep Reddy Gillela

原始標籤

Advanced (300) · Amazon Bedrock Knowledge Bases · Technical How-to

原始連結

https://aws.amazon.com/blogs/machine-learning/agentic-retrieval-with-langchain-and-amazon-bedrock-knowledge-bases/

摘要原文

When a user asks the support assistant, a Retrieval Augmented Generation (RAG) application built with LangChain to compare two products across three dimensions, they’re effectively posing six questions simultaneously. Similarity search uses a single query vector to encapsulate all the intents. The retriever then generates the best approximation of the average of those intents. The resulting answer comes back concise. The search executes without errors. The relevance scores look reasonable. Yet the retrieved chunks, while topically relevant, only cover a fraction of what the question actually asked. In this post, we showcase a RAG application on Amazon Bedrock Managed Knowledge Base with LangChain . We run the same multi-part question through standard and agentic retrieval, and read the trace events to see the plan the model produced. We also cover what the two retrieval paths cost and when the cheaper one is the right choice. Agentic retrieval is available on Amazon Bedrock Managed Knowledge Base. Instead of one search, Amazon Bedrock Managed Knowledge Base plans the retrieval. It breaks the question into sub-queries, runs them, judges whether it has enough evidence, and searches again if it doesn’t. The langchain-aws package exposes both agentic and standard retrieval, so you can use either from a LangChain application. Amazon Bedrock Managed Knowledge Base , the fully managed RAG capability in Amazon Bedrock, removes the self-managed vector store, embeddings, and re-ranking models from the RAG architecture. You configure a data source, and Amazon Bedrock Managed Knowledge Bases handles chunking, embedding, storage, and retrieval. This walkthrough uses Amazon Simple Storage Service (Amazon S3) . Amazon Bedrock Managed Knowledge Bases provides two APIs. We briefly discuss those differences in this post. The Retrieve API runs one hybrid search and returns scored chunks. The AgenticRetrieveStream API runs a planning loop and streams the steps back to you as trace events. In the langchain-aws package, the first is a standard LangChain retriever you can drop into a chain. The second is a function retrieval directly from a knowledge base. The following diagram shows the solution architecture. The application queries Amazon Bedrock Knowledge Bases using either the Retrieve API (standard, single-shot) or the AgenticRetrieveStream API (multi-step planning loop). Both paths return document chunks from the knowledge base, which the application then uses to generate a grounded response. Figure 1: Solution architecture for querying Amazon Bedrock Knowledge Bases with the Retrieve and AgenticRetrieveStream APIs The following sections walk you through creating a knowledge base, querying it with both retrieval methods, and reading the trace events the agentic planner produces. Install the packages. The Boto3 version matters: agentic_retrieve_stream did not exist before 1.43.32. Two identities are involved and separating them is worth doing deliberately. The knowledge base assumes a service role to read your documents and call the embedding model. Your application uses an AWS Security Token Service (AWS STS) caller identity to query. Neither needs the other’s permissions. Amazon Bedrock creates the service role for you if you let it. To supply your own, give it a trust policy that lets Amazon Bedrock assume it. Scope it with aws:SourceAccount and aws:SourceArn so that another account can’t use it as a confused deputy: The service role also needs s3:ListBucket on your bucket and s3:GetObject on its contents, both conditioned on aws:ResourceAccount . Scope the knowledge-base/* wildcard character down to specific knowledge base IDs after you have created them. The AWS STS caller identity needs a different set. bedrock:AgenticRetrieveStream and bedrock:InvokeModelWithResponseStream can’t be scoped to a knowledge base Amazon Resource Name (ARN). bedrock:Retrieve and bedrock:GetDocumentContent can: bedrock:GetDocumentContent is often overlooked. Agentic retrieval calls it when a FullDocumentExpansion step decides a passage lacks the context to answer. A policy with only bedrock:Retrieve works until the planner reaches for a whole document and then fails partway through a query. To create and manage the knowledge base itself, the calling role additionally needs bedrock:CreateKnowledgeBase on * , and the GetKnowledgeBase , UpdateKnowledgeBase , DeleteKnowledgeBase , StartIngestionJob , GetIngestionJob , and ListIngestionJobs actions on knowledge-base/* . If you’re using guardrails, add bedrock:GetGuardrail and bedrock:ApplyGuardrail . Running this walkthrough might incur costs for document storage and ingestion in the knowledge base, retrieval calls, and foundation model (FM) inference. For more information about pricing, see the Knowledge Bases section of Amazon Bedrock pricing . Delete the resources when you complete this experiment. Create the knowledge base with a managedKnowledgeBaseConfiguration . Setting embeddingModelType to MANAGED uses the service-managed embedding model. There’s no storageConfiguration in that request. For a self-managed knowledge base you would pass one describing your vector store. Amazon Bedrock Managed Knowledge Base does not take one, which is the clearest signal in the API that Amazon Bedrock owns the storage layer. Attach the S3 bucket as a data source, then start an ingestion job. Ingestion is asynchronous, so poll until the job reaches a terminal state rather than sleeping for a fixed interval and hoping. The full data source configuration and error handling are in the sample repository . AmazonKnowledgeBasesRetriever wraps the Retrieve API and behaves like any other LangChain retriever. For Amazon Bedrock Managed Knowledge Bases, pass managedSearchConfiguration . This is the part that trips people up: vectorSearchConfiguration is the previous path for knowledge bases where you run your own vector store. It is what most existing examples show. Each result comes back as a LangChain Document . The relevance score is in metadata["score"] , and the source document’s own metadata is under metadata["source_metadata"] , renamed so it does not collide. If you want to drop low-confidence results, set min_score_confidence on the retriever instead of filtering afterward. For a question with one clear intent, this is the right tool. It is one call. The latency is the lowest of the two options, and you keep full control of how the answer gets generated. Most of the queries a production assistant sees are this shape, and reaching for a planning loop to answer them wastes money and time. Now give the same retriever a question with several parts: Five chunks come back, ranked by hybrid score against one embedding of that whole question. That question contains six intents: two services across three dimensions. Scoring the retrieved text for evidence of each one gives a concrete measure of what a single embedding recovers. At five results, one embedding standing for six intents misses two of them. At ten it covers all six, with visible waste: two sub-intents are covered twice and one chunk carries none. The retriever did its job. The limitation is structural: one vector cannot represent six intents, and there is no step in the process that asks whether the returned evidence is enough to answer the question. Agentic retrieval is not a LangChain retriever, but a feature of Amazon Bedrock Managed Knowledge Bases. The langchain-aws package exposes it as a standalone function, agentic_retrieve , because the underlying API streams its results and doesn’t fit the synchronous BaseRetriever interface. There’s no flag on AmazonKnowledgeBasesRetriever that switches it on. With generate_response=True , the service returns a grounded answer and citations alongside the retrieved chunks, so you get an answer without wiring up a separate model call. The function works only against Amazon Bedrock Managed Knowledge Base.