o-nikolas commented on code in PR #74177:
URL: https://github.com/apache/airflow/pull/74177#discussion_r4350119136
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -761,25 +786,55 @@ def execute_complete(self, context: Context, event:
dict[str, Any] | None = None
self.log.info("Bedrock knowledge base creation job `%s` complete.",
self.name)
return validated_event["knowledge_base_id"]
- def execute(self, context: Context) -> str:
- knowledge_base_config = {
+ def _build_knowledge_base_config(self) -> dict[str, Any]:
+ """
+ Return the knowledgeBaseConfiguration to send to the API.
+
+ ``embedding_model_arn`` is kept for backward compatibility and only
builds a VECTOR
+ configuration when ``knowledge_base_config`` is not provided.
+ """
+ if self.knowledge_base_config is not None:
+ return self.knowledge_base_config
+ return {
"type": "VECTOR",
"vectorKnowledgeBaseConfiguration": {"embeddingModelArn":
self.embedding_model_arn},
}
+ def execute(self, context: Context) -> str:
+
+ create_kb_kwargs: dict[str, Any] = {
+ "name": self.name,
+ "roleArn": self.role_arn,
+ "knowledgeBaseConfiguration": self._build_knowledge_base_config(),
+ }
+
+ # Only self-managed vector knowledge bases use a storage configuration.
+ if self.storage_config:
Review Comment:
Nit: `if self.storage_config:` treats `{}` as "not passed". `is not None`
would match the constructor checks. Same for `bool(self.storage_config)`
on line 826.
##########
providers/amazon/docs/operators/bedrock.rst:
##########
@@ -229,6 +231,30 @@
https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base-supported.ht
:start-after: [START howto_operator_bedrock_create_knowledge_base]
:end-before: [END howto_operator_bedrock_create_knowledge_base]
+
+Other types of knowledge bases can be created by passing the
``knowledgeBaseConfiguration`` of
+the API through ``knowledge_base_config``. For the expected structure of each
type, see
+https://docs.aws.amazon.com/boto3/latest/reference/services/bedrock-agent/client/create_knowledge_base.html
+
+For example, a managed knowledge base, where Amazon Bedrock owns the vector
store, the ingestion
+pipeline and the retrieval infrastructure, needs neither
``embedding_model_arn`` nor
+``storage_config``:
+
+.. code-block:: python
Review Comment:
Could you add the managed knowledge base to a system test and include it
here with `exampleinclude`, like the vector example above? That keeps the
docs snippet tested.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
