Source code for langchain_community.retrievers.bedrock
from typing import Any, Dict, List, Optional
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.pydantic_v1 import BaseModel, root_validator
from langchain_core.retrievers import BaseRetriever
[docs]class VectorSearchConfig(BaseModel, extra="allow"):  # type: ignore[call-arg]
    numberOfResults: int = 4
[docs]class RetrievalConfig(BaseModel, extra="allow"):  # type: ignore[call-arg]
    vectorSearchConfiguration: VectorSearchConfig
[docs]class AmazonKnowledgeBasesRetriever(BaseRetriever):
    """A retriever class for `Amazon Bedrock Knowledge Bases`.
    See https://aws.amazon.com/bedrock/knowledge-bases for more info.
    Args:
        knowledge_base_id: Knowledge Base ID.
        region_name: The aws region e.g., `us-west-2`.
            Fallback to AWS_DEFAULT_REGION env variable or region specified in
            ~/.aws/config.
        credentials_profile_name: The name of the profile in the ~/.aws/credentials
            or ~/.aws/config files, which has either access keys or role information
            specified. If not specified, the default credential profile or, if on an
            EC2 instance, credentials from IMDS will be used.
        client: boto3 client for bedrock agent runtime.
        retrieval_config: Configuration for retrieval.
    Example:
        .. code-block:: python
            from langchain_community.retrievers import AmazonKnowledgeBasesRetriever
            retriever = AmazonKnowledgeBasesRetriever(
                knowledge_base_id="<knowledge-base-id>",
                retrieval_config={
                    "vectorSearchConfiguration": {