Bedrock (knowledge bases) integration

os.environ["LANGSMITH_API_KEY"] = getpass.getpass("Enter your LangSmith API key: ")
os.environ["LANGSMITH_TRACING"] = "true"

pip install -qU "langchain-aws>=1.6.3"

from langchain_aws.retrievers import AmazonKnowledgeBasesRetriever

retriever = AmazonKnowledgeBasesRetriever(
    knowledge_base_id="PUIJP4EQUA",
    retrieval_config={"vectorSearchConfiguration": {"numberOfResults": 4}},
)

from langchain_aws.retrievers import AmazonKnowledgeBasesRetriever

retriever = AmazonKnowledgeBasesRetriever(
    knowledge_base_id="YOUR_MANAGED_KB_ID",
    retrieval_config={"managedSearchConfiguration": {"numberOfResults": 4}},
)

from langchain_aws.retrievers import agentic_retrieve

result = agentic_retrieve(
    knowledge_base_id="YOUR_MANAGED_KB_ID",
    query="What are the differences between S3 storage classes?",
    region_name="us-west-2",
)

for doc in result["results"]:
    print(doc["content"]["text"])

query = "What did the president say about Ketanji Brown?"

retriever.invoke(query)

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

llm = ChatBedrock(model_id="anthropic.claude-sonnet-4-20250514-v1:0")

prompt = ChatPromptTemplate.from_template(
    "Answer the question based on the context:\n\nContext: {context}\n\nQuestion: {question}"
)

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

chain.invoke("What are the key features?")

{
  "Effect": "Allow",
  "Action": [
    "bedrock:Retrieve",
    "bedrock:AgenticRetrieveStream"
  ],
  "Resource": "arn:aws:bedrock:<region>:<account-id>:knowledge-base/<kb-id>"
}