o-nikolas commented on code in PR #72597:
URL: https://github.com/apache/airflow/pull/72597#discussion_r3984153227
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -750,10 +777,49 @@ def execute_complete(self, context: Context, event:
dict[str, Any] | None = None
return validated_event["knowledge_base_id"]
def execute(self, context: Context) -> str:
- knowledge_base_config = {
- "type": "VECTOR",
- "vectorKnowledgeBaseConfiguration": {"embeddingModelArn":
self.embedding_model_arn},
+ if not self.role_arn:
+ raise ValueError("`role_arn` must be specified to create a
knowledge base.")
+
+ create_kwargs = self.create_knowledge_base_kwargs.copy()
+
+ if "knowledgeBaseConfiguration" in create_kwargs:
+ knowledge_base_config =
create_kwargs.pop("knowledgeBaseConfiguration")
+ elif "knowledge_base_configuration" in create_kwargs:
+ knowledge_base_config =
create_kwargs.pop("knowledge_base_configuration")
+ elif "knowledge_base_config" in create_kwargs:
+ knowledge_base_config = create_kwargs.pop("knowledge_base_config")
+ elif self.knowledge_base_configuration is not None:
+ knowledge_base_config = self.knowledge_base_configuration
+ elif self.embedding_model_arn is not None:
+ knowledge_base_config = {
+ "type": "VECTOR",
+ "vectorKnowledgeBaseConfiguration": {"embeddingModelArn":
self.embedding_model_arn},
+ }
+ else:
+ raise ValueError(
+ "Either `knowledge_base_configuration` or
`embedding_model_arn` must be provided."
+ )
+
+ if "storageConfiguration" in create_kwargs:
+ storage_config = create_kwargs.pop("storageConfiguration")
+ elif "storage_config" in create_kwargs:
Review Comment:
Same as above?
##########
providers/amazon/docs/operators/bedrock.rst:
##########
@@ -189,6 +189,23 @@
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]
+You can also create a fully managed knowledge base where Bedrock handles the
underlying vector store
+and infrastructure automatically:
+
+.. code-block:: python
Review Comment:
This should be added to the system test and included from there.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -750,10 +777,49 @@ def execute_complete(self, context: Context, event:
dict[str, Any] | None = None
return validated_event["knowledge_base_id"]
def execute(self, context: Context) -> str:
- knowledge_base_config = {
- "type": "VECTOR",
- "vectorKnowledgeBaseConfiguration": {"embeddingModelArn":
self.embedding_model_arn},
+ if not self.role_arn:
+ raise ValueError("`role_arn` must be specified to create a
knowledge base.")
+
+ create_kwargs = self.create_knowledge_base_kwargs.copy()
+
+ if "knowledgeBaseConfiguration" in create_kwargs:
+ knowledge_base_config =
create_kwargs.pop("knowledgeBaseConfiguration")
+ elif "knowledge_base_configuration" in create_kwargs:
+ knowledge_base_config =
create_kwargs.pop("knowledge_base_configuration")
+ elif "knowledge_base_config" in create_kwargs:
+ knowledge_base_config = create_kwargs.pop("knowledge_base_config")
Review Comment:
Why are we accepting all of these? We do not usually do this, instead
conform to the boto3 api snake case.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -712,9 +719,10 @@ class
BedrockCreateKnowledgeBaseOperator(AwsBaseOperator[BedrockAgentHook]):
def __init__(
self,
name: str,
- embedding_model_arn: str,
- role_arn: str,
- storage_config: dict[str, Any],
+ embedding_model_arn: str | None = None,
+ role_arn: str | None = None,
Review Comment:
Why is this being set to None? This is still required even by a Managed KB
typ, no?
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -725,12 +733,16 @@ def __init__(
deferrable: bool = conf.getboolean("operators", "default_deferrable",
fallback=False),
**kwargs,
):
+ if knowledge_base_configuration is None and "knowledge_base_config" in
kwargs:
Review Comment:
Was `knowledge_base_config` ever an accepted init arg? Were folks ever
passing this and it being pulled out of kwargs? I think this check, the test
and much of the other property/setter machinery added is not needed.
--
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]
