Kunal8954 commented on code in PR #72608:
URL: https://github.com/apache/airflow/pull/72608#discussion_r4089953733
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -750,10 +757,24 @@ 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 self.role_arn is None:
Review Comment:
Thanks for the review. The change is fully additive — existing self-managed
Dags keep producing an identical API call.
- `knowledge_base_configuration` is a new optional parameter appended after
all existing ones; `create_knowledge_base_kwargs`, `embedding_model_arn`,
`role_arn`, and `storage_config` keep their original positions and meaning
(they're now optional so managed Dags can omit what doesn't apply). Any
existing `__init__` call resolves exactly as before.
- In `execute()`, when `knowledge_base_configuration` is not provided the
operator builds the same `VECTOR` config from `embedding_model_arn` and passes
`storageConfiguration` as before, and `create_knowledge_base_kwargs` are still
forwarded unchanged. The only new behaviour is confined to the `MANAGED` path
(omit `storageConfiguration`, skip the vector-index retry loop).
- Added `test_existing_self_managed_behavior_preserved`, which asserts an
existing self-managed Dag (with `create_knowledge_base_kwargs` passthrough)
issues the byte-identical `CreateKnowledgeBase` call.
If there's still a scenario you consider backwards-incompatible, I'll fix it.
---
Drafted-by: big-pickle (no human review before posting)
--
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]