Recuperação Agêntica com LangChain e Amazon Bedrock Knowledge Bases

O problema com perguntas complexas em RAG

Imagine que um usuário pergunta ao assistente de suporte: “Compare o serviço de checkout e o de inventário em termos de escalonamento de plantão, metas de backup e restauração, e procedimento de rollback de deploy. Onde eles diferem?” Essa única pergunta, na prática, contém seis intenções distintas — dois serviços multiplicados por três dimensões.

O problema é que a busca por similaridade tradicional usa um único vetor para representar todas essas intenções ao mesmo tempo. O resultado é uma espécie de “média” das intenções, e os trechos recuperados, embora topicamente relevantes, cobrem apenas parte do que foi perguntado. A resposta chega concisa, sem erros aparentes, com scores razoáveis — mas incompleta.

É exatamente esse cenário que a recuperação agêntica no Amazon Bedrock Managed Knowledge Base foi projetada para resolver. Em vez de uma única busca, o serviço planeja a recuperação: decompõe a pergunta em sub-consultas, executa cada uma, avalia se as evidências coletadas são suficientes e, se não forem, busca novamente.

Visão geral da solução

O Amazon Bedrock Managed Knowledge Base é a capacidade de Geração Aumentada por Recuperação (RAG) totalmente gerenciada do Amazon Bedrock. Ela elimina a necessidade de gerenciar o próprio armazenamento vetorial, modelos de embedding e re-ranqueamento. Basta configurar uma fonte de dados — neste caso, um bucket do Serviço de Armazenamento Simples da Amazon (Amazon S3) — e o serviço cuida de chunking, embedding, armazenamento e recuperação.

O serviço expõe duas APIs principais:

  • Retrieve: executa uma única busca híbrida e retorna trechos com pontuação de relevância.
  • AgenticRetrieveStream: executa um loop de planejamento e transmite os passos de volta como eventos de rastreamento.

No pacote langchain-aws, a primeira API se comporta como um retriever LangChain padrão que pode ser encaixado em qualquer chain. A segunda é exposta como uma função standalone, já que o modelo de streaming não se encaixa na interface síncrona do BaseRetriever.

Imagem original — fonte: Aws

Pré-requisitos e permissões

Para seguir o walkthrough, são necessários: uma conta AWS com acesso ao Amazon Bedrock na região US East (N. Virginia) — us-east-1 —, Python 3.12 ou superior, um bucket S3 com documentos de amostra e os seguintes pacotes instalados:

langchain-aws>=1.6.3
langchain>=1.0
boto3>=1.43.32

A versão do Boto3 importa: a função agentic_retrieve_stream não existia antes da versão 1.43.32. Verifique a documentação da AWS para disponibilidade em outras regiões.

Duas identidades do Controle de Acesso e Identidade (IAM) estão envolvidas e vale separá-las deliberadamente. A service role é assumida pelo Knowledge Base para ler os documentos e chamar o modelo de embedding. A identidade que chama as APIs é a da sua aplicação. Nenhuma precisa das permissões da outra.

A service role precisa de uma política de confiança que permita ao Amazon Bedrock assumi-la, com as condições aws:SourceAccount e aws:SourceArn para evitar o problema de “confused deputy”:

{
  "Version": "2012-10-17",
  "Statement": [{
    "Effect": "Allow",
    "Principal": {"Service": "bedrock.amazonaws.com"},
    "Action": "sts:AssumeRole",
    "Condition": {
      "StringEquals": {"aws:SourceAccount": "111122223333"},
      "ArnLike": {
        "aws:SourceArn": "arn:aws:bedrock:us-east-1:111122223333:knowledge-base/*"
      }
    }
  }]
}

A identidade chamadora precisa de um conjunto diferente de permissões. Note que bedrock:AgenticRetrieveStream e bedrock:InvokeModelWithResponseStream não podem ser escopadas a um ARN de knowledge base específico, mas bedrock:Retrieve e bedrock:GetDocumentContent podem:

{
  "Version": "2012-10-17",
  "Statement": [
    {
      "Sid": "AgenticRetrievalAndPlannerModel",
      "Effect": "Allow",
      "Action": [
        "bedrock:AgenticRetrieveStream",
        "bedrock:InvokeModelWithResponseStream"
      ],
      "Resource": "*"
    },
    {
      "Sid": "RetrieveAndFullDocumentExpansion",
      "Effect": "Allow",
      "Action": ["bedrock:Retrieve", "bedrock:GetDocumentContent"],
      "Resource": "arn:aws:bedrock::111122223333:knowledge-base/"
    },
    {
      "Sid": "GenerateAnswersInTheChains",
      "Effect": "Allow",
      "Action": ["bedrock:InvokeModel", "bedrock:Converse", "bedrock:ConverseStream"],
      "Resource": "*"
    }
  ]
}

Um detalhe frequentemente negligenciado: bedrock:GetDocumentContent é chamado quando o planejador decide que um trecho não tem contexto suficiente e precisa buscar o documento inteiro — etapa chamada de FullDocumentExpansion. Uma política que tem apenas bedrock:Retrieve vai funcionar até o planejador tentar expandir um documento e falhar no meio da consulta.

Criando e populando o Knowledge Base

O Knowledge Base é criado com managedKnowledgeBaseConfiguration. Definir embeddingModelType como MANAGED usa o modelo de embedding gerenciado pelo serviço. Note a ausência de storageConfiguration na requisição — para um Knowledge Base gerenciado, o Amazon Bedrock é responsável pela camada de armazenamento:

import boto3
import os

REGION = os.environ["AWS_REGION"]
bedrock_agent = boto3.client("bedrock-agent", region_name=REGION)

response = bedrock_agent.create_knowledge_base(
    name=KB_NAME,
    roleArn=KB_ROLE_ARN,
    knowledgeBaseConfiguration={
        "type": "MANAGED",
        "managedKnowledgeBaseConfiguration": {
            "embeddingModelType": "MANAGED",
        },
    },
)
KB_ID = response["knowledgeBase"]["knowledgeBaseId"]

Após criar o Knowledge Base, o bucket S3 é anexado como fonte de dados e um job de ingestão é iniciado. A ingestão é assíncrona — o código abaixo faz polling até o job atingir um estado terminal:

import time

SUCCESS_STATES = frozenset({"COMPLETE"})
FAILURE_STATES = frozenset({"FAILED", "STOPPED"})

def wait_for_ingestion(kb_id, ds_id, job_id, timeout_s=1800):
    """Poll an ingestion job until it reaches a terminal state."""
    deadline = time.time() + timeout_s
    while time.time() < deadline:
        job = bedrock_agent.get_ingestion_job(
            knowledgeBaseId=kb_id,
            dataSourceId=ds_id,
            ingestionJobId=job_id,
        )["ingestionJob"]
        status = job["status"]
        if status in SUCCESS_STATES:
            return job
        if status in FAILURE_STATES:
            reasons = job.get("failureReasons") or ["no reason reported"]
            raise RuntimeError(f"Ingestion job {job_id} finished as {status}: " + "; ".join(reasons))
        time.sleep(15)
    raise TimeoutError(f"Ingestion job {job_id} did not finish in {timeout_s}s")

A configuração completa da fonte de dados e o tratamento de erros estão no repositório de exemplos.

Consulta com o retriever padrão do LangChain

O AmazonKnowledgeBasesRetriever encapsula a API Retrieve e se comporta como qualquer outro retriever LangChain. Para o Managed Knowledge Base, é necessário passar managedSearchConfiguration — e não vectorSearchConfiguration, que é o caminho para knowledge bases com armazenamento vetorial próprio (o que a maioria dos exemplos existentes ainda mostra):

from langchain_aws.retrievers import AmazonKnowledgeBasesRetriever

SIMPLE_QUERY = "What is the restore time objective for the checkout service?"

retriever = AmazonKnowledgeBasesRetriever(
    knowledge_base_id=KB_ID,
    region_name=REGION,
    retrieval_config={
        "managedSearchConfiguration": {
            "numberOfResults": 5,
        }
    },
)

docs = retriever.invoke(SIMPLE_QUERY)

Cada resultado retorna como um Document do LangChain. O score de relevância fica em metadata["score"] e os metadados do documento fonte ficam em metadata["source_metadata"]. Para perguntas com uma única intenção clara, essa é a abordagem correta: uma chamada, latência baixa, controle total sobre como a resposta é gerada.

O limite da busca simples em perguntas complexas

Ao passar a pergunta comparativa com seis intenções para o mesmo retriever, o resultado revela a limitação estrutural. Com cinco chunks retornados, a busca cobre apenas 4 das 6 sub-intenções. Com dez chunks, cobre todas as seis — mas com redundância visível: duas sub-intenções aparecem cobertas duas vezes e um chunk não cobre nenhuma. O problema não é de implementação; é estrutural. Um único vetor não consegue representar seis intenções, e não há etapa no processo que avalie se as evidências retornadas são suficientes para responder a pergunta.

Executando a recuperação agêntica

A recuperação agêntica não é um retriever LangChain, mas uma funcionalidade do Amazon Bedrock Managed Knowledge Base. O pacote langchain-aws a expõe como a função standalone agentic_retrieve:

from langchain_aws.retrievers.bedrock import agentic_retrieve

result = agentic_retrieve(
    knowledge_base_id=KB_ID,
    query=COMPLEX_QUERY,
    region_name=REGION,
    generate_response=True,
    number_of_results=10,
)

print(result["generatedResponse"]["answer"])

Com generate_response=True, o serviço retorna uma resposta fundamentada e citações junto com os chunks recuperados — sem necessidade de uma chamada de modelo separada. Internamente, o serviço planeja, recupera, avalia se as evidências são suficientes e itera se necessário.

Lendo os eventos de rastreamento do planejador

Para observar como o modelo decompõe a pergunta, é necessário chamar agentic_retrieve_stream diretamente no cliente bedrock-agent-runtime do Boto3. Esse é o único ponto do walkthrough em que se contorna o langchain-aws, pois o helper descarta os eventos de rastreamento e não expõe maxAgentIteration nem um modelo de planejador customizado:

runtime = boto3.client("bedrock-agent-runtime", region_name=REGION)

response = runtime.agentic_retrieve_stream(
    messages=[{"role": "user", "content": {"text": COMPLEX_QUERY}}],
    retrievers=[{
        "configuration": {
            "knowledgeBase": {
                "knowledgeBaseId": KB_ID,
                "retrievalOverrides": {"maxNumberOfResults": 10},
            }
        }
    }],
    agenticRetrieveConfiguration={
        "foundationModelType": "MANAGED",
        "rerankingModelType": "MANAGED",
        # 5 is the API default. Below 4 the planner stops decomposing entirely.
        "maxAgentIteration": 5,
    },
    generateResponse=False,
)

for event in response["stream"]:
    if "traceEvent" in event:
        attrs = event["traceEvent"]["attributes"]
        print(f"{attrs.get('step')}: {attrs.get('status')}")
        for action in attrs.get("actions", []) or []:
            if "retrieve" in action:
                query = action["retrieve"].get("inputQuery", {}).get("text", "")
                print(f"  sub-query: {query}")
    elif "result" in event:
        for chunk in event["result"].get("results", []):
            print(chunk.get("content", {}).get("text", "")[:120])

Definir generateResponse como False é recomendado quando o objetivo é apenas inspecionar o plano — evita uma chamada de modelo extra desnecessária.

Cada evento de rastreamento tem um campo step que indica onde o planejador está:

  • SpeculativeRetrieval: executa antes do primeiro plano para reduzir latência e não conta contra o orçamento de iterações.
  • Planning: o modelo lê a pergunta e os resultados anteriores e emite sub-consultas.
  • Retrieval: disparado uma vez por sub-consulta.
  • FullDocumentExpansion: aparece quando o modelo decide que um trecho não tem contexto suficiente e busca o documento inteiro.

O texto das sub-consultas — a parte mais valiosa para logging — fica em attributes.actions[].retrieve.inputQuery.text, não nos campos de nível superior do trace. Um handler que lê apenas step e status mostra que o planejamento aconteceu, mas não o que foi decidido.

Imagem original — fonte: Aws

Dois detalhes importantes antes de construir sobre essa API: a deduplicação se aplica apenas ao evento result — um chunk recuperado por três sub-consultas aparece uma vez no final, mas três vezes nos traces. E ao contrário da API Retrieve, os resultados da recuperação agêntica não têm um campo score tipado. Código que lê result["score"] após trocar de API não vai encontrar nada.

Sobre o parâmetro maxAgentIteration: ele aceita valores entre 2 e 10 e tem padrão 5. Com valores 2 ou 3, o planejador executa apenas um ciclo, não emite sub-consultas e retorna o que a etapa de SpeculativeRetrieval já encontrou — comportamento de busca simples com custo agêntico. A decomposição começa a partir de 4. O planejador frequentemente para antes de atingir o limite quando julga que as evidências são suficientes.

Construindo a RAG chain com LangChain

Para o retriever padrão, a composição usual com a LangChain Expression Language (LCEL) funciona diretamente:

from langchain_aws import ChatBedrockConverse
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

llm = ChatBedrockConverse(model=MODEL_ID, region_name=REGION)
prompt = ChatPromptTemplate.from_template(PROMPT_TEMPLATE)

chain = (
    {"context": retriever | format_docs, "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

Para encaixar a recuperação agêntica na mesma posição, é necessário envolvê-la em um RunnableLambda, já que ela é uma função e não um retriever:

from langchain_core.runnables import RunnableLambda

def agentic_context(question: str) -> str:
    result = agentic_retrieve(
        knowledge_base_id=KB_ID,
        query=question,
        region_name=REGION,
        number_of_results=10,
    )
    return "\n\n".join(
        item.get("content", {}).get("text", "")
        for item in result.get("results", [])
    )

agentic_chain = (
    {"context": RunnableLambda(agentic_context), "question": RunnablePassthrough()}
    | prompt
    | llm
    | StrOutputParser()
)

Note que generate_response está desativado aqui. Dentro de uma chain, o usual é querer controlar o próprio prompt e modelo — então os chunks são recuperados e a geração acontece downstream. Use a geração do próprio serviço quando quiser uma única chamada com menos código.

Quando usar cada abordagem

A recomendação é rotear com base no formato da consulta, em vez de escolher uma abordagem para tudo:

  • Use Retrieve para perguntas curtas e bem delimitadas. É mais barato, mais rápido, funciona com knowledge bases autogerenciados e retorna scores nos resultados. A maior parte do tráfego de produção tem esse formato.
  • Use AgenticRetrieveStream quando as perguntas forem multi-parte, comparativas ou exploratórias, ou quando as evidências estiverem espalhadas por mais de um knowledge base. A API suporta até cinco knowledge bases em uma única requisição, roteando sub-consultas com base em descrições em linguagem natural. Custa mais por chamada, faz várias invocações de modelo e tem latência mais alta.

Um classificador ou heurística sobre o formato da pergunta pode enviar a maior parte do tráfego pelo caminho barato e reservar o planejador para as consultas que realmente precisam dele. A AWS avaliou a recuperação agêntica no benchmark público MuSiQue (multi-hop) e observou melhora no recall em relação à busca simples, com os maiores ganhos nas perguntas mais difíceis. Perguntas de um único salto tiveram ganhos abaixo de cinco pontos — o que ilustra bem o trade-off: a decomposição ajuda quando há algo para decompor.

Em produção, o Amazon Bedrock Guardrails pode ser usado para aplicar políticas de conteúdo e verificações de fundamentação nas respostas geradas. Ambos os caminhos de recuperação suportam guardrails — mas a recuperação agêntica usa policyConfiguration.bedrockGuardrailConfiguration e suporta apenas o modo BLOCK. Se o modo MASK for necessário, isso é um motivo para permanecer na API Retrieve.

Limpeza dos recursos

Ao finalizar os experimentos, é importante remover os recursos criados para evitar cobranças contínuas de armazenamento:

bedrock_agent.delete_data_source(knowledgeBaseId=KB_ID, dataSourceId=DS_ID)
bedrock_agent.delete_knowledge_base(knowledgeBaseId=KB_ID)

O repositório de exemplos inclui um script de limpeza que também esvazia o bucket e remove a role.

Conclusão

A recuperação agêntica no Amazon Bedrock Managed Knowledge Base resolve um problema estrutural das aplicações RAG tradicionais: a incapacidade de um único vetor representar múltiplas intenções simultâneas. O loop de planejamento decompõe perguntas complexas, avalia as evidências coletadas e itera quando necessário — entregando cobertura superior em consultas multi-parte.

A integração com LangChain tem algumas arestas: a recuperação agêntica é uma função e não um retriever, então precisa de um RunnableLambda para entrar em uma chain. Os eventos de rastreamento que mostram o plano de consulta exigem uma chamada direta via Boto3. São trade-offs conhecidos que valem a pena para as consultas certas.

Para começar, consulte a documentação do Amazon Bedrock Knowledge Bases e o código de exemplo disponível no repositório de amostras. Consulte também a seção Knowledge Bases da página de preços do Amazon Bedrock antes de colocar a solução em produção.

Fonte

Agentic retrieval with LangChain and Amazon Bedrock Knowledge Bases (https://aws.amazon.com/blogs/machine-learning/agentic-retrieval-with-langchain-and-amazon-bedrock-knowledge-bases/)

Comments

Leave a Reply

Your email address will not be published. Required fields are marked *