Bedrock

AmazonKnowledgeBasesRetriever

Bases: BaseRetriever

Amazon Bedrock Knowledge Bases 检索。

请参阅 https://aws.amazon.com/bedrock/knowledge-bases 了解更多信息。

Parameters:

Name	Type	Description	Default
knowledge_base_id	str	知识库ID。	required
retrieval_config	Optional[Dict[str, Any]]	检索配置。	None
profile_name	Optional[str]	~/.aws/credentials 或 ~/.aws/config 文件中配置的配置文件名称，其中包含访问密钥或角色信息。如果未指定，则将使用默认凭据配置文件，或者如果在EC2实例上，则将使用IMDS中的凭据。	None
region_name	Optional[str]	AWS区域，例如 us-west-2。 回退到AWS_DEFAULT_REGION环境变量或~/.aws/config中指定的区域。	None
aws_access_key_id	Optional[str]	AWS访问密钥ID。	None
aws_secret_access_key	Optional[str]	AWS秘密访问密钥。	None
aws_session_token	Optional[str]	AWS临时会话令牌。	None

示例

.. code-block:: python

from llama_index.retrievers.bedrock import AmazonKnowledgeBasesRetriever

retriever = AmazonKnowledgeBasesRetriever(
    knowledge_base_id="<knowledge-base-id>",
    retrieval_config={
        "vectorSearchConfiguration": {
            "numberOfResults": 4,
            "overrideSearchType": "SEMANTIC",
            "filter": {
                "equals": {
                    "key": "tag",
                    "value": "space"
                }
            }
        }
    },
)

Source code in llama_index/retrievers/bedrock/base.py

class AmazonKnowledgeBasesRetriever(BaseRetriever):
    """`Amazon Bedrock Knowledge Bases` 检索。

    请参阅 https://aws.amazon.com/bedrock/knowledge-bases 了解更多信息。

    Args:
        knowledge_base_id: 知识库ID。
        retrieval_config: 检索配置。
        profile_name: ~/.aws/credentials 或 ~/.aws/config 文件中配置的配置文件名称，其中包含访问密钥或角色信息。如果未指定，则将使用默认凭据配置文件，或者如果在EC2实例上，则将使用IMDS中的凭据。
        region_name: AWS区域，例如 `us-west-2`。
            回退到AWS_DEFAULT_REGION环境变量或~/.aws/config中指定的区域。
        aws_access_key_id: AWS访问密钥ID。
        aws_secret_access_key: AWS秘密访问密钥。
        aws_session_token: AWS临时会话令牌。

    示例:
        .. code-block:: python

            from llama_index.retrievers.bedrock import AmazonKnowledgeBasesRetriever

            retriever = AmazonKnowledgeBasesRetriever(
                knowledge_base_id="<knowledge-base-id>",
                retrieval_config={
                    "vectorSearchConfiguration": {
                        "numberOfResults": 4,
                        "overrideSearchType": "SEMANTIC",
                        "filter": {
                            "equals": {
                                "key": "tag",
                                "value": "space"
                            }
                        }
                    }
                },
            )"""

    def __init__(
        self,
        knowledge_base_id: str,
        retrieval_config: Optional[Dict[str, Any]] = None,
        profile_name: Optional[str] = None,
        region_name: Optional[str] = None,
        aws_access_key_id: Optional[str] = None,
        aws_secret_access_key: Optional[str] = None,
        aws_session_token: Optional[str] = None,
        callback_manager: Optional[CallbackManager] = None,
    ):
        self._client = get_aws_service_client(
            service_name="bedrock-agent-runtime",
            profile_name=profile_name,
            region_name=region_name,
            aws_access_key_id=aws_access_key_id,
            aws_secret_access_key=aws_secret_access_key,
            aws_session_token=aws_session_token,
        )
        self.knowledge_base_id = knowledge_base_id
        self.retrieval_config = retrieval_config
        super().__init__(callback_manager)

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        query = query_bundle.query_str

        response = self._client.retrieve(
            retrievalQuery={"text": query.strip()},
            knowledgeBaseId=self.knowledge_base_id,
            retrievalConfiguration=self.retrieval_config,
        )
        results = response["retrievalResults"]
        node_with_score = []
        for result in results:
            metadata = {}
            if "location" in result:
                metadata["location"] = result["location"]
            if "metadata" in result:
                metadata["sourceMetadata"] = result["metadata"]
            node_with_score.append(
                NodeWithScore(
                    node=TextNode(
                        text=result["content"]["text"],
                        metadata=metadata,
                    ),
                    score=result["score"] if "score" in result else 0,
                )
            )
        return node_with_score