anshuman0123 opened a new pull request, #72597:
URL: https://github.com/apache/airflow/pull/72597
### Description
Amazon Bedrock Managed Knowledge Bases provide a fully managed RAG offering
where Bedrock manages the vector store, ingestion pipeline, and retrieval
infrastructure automatically without requiring users to configure and maintain
a self-managed vector database (e.g. OpenSearch Serverless, Pinecone, etc.).
Previously, `BedrockCreateKnowledgeBaseOperator` could not be used to create
managed knowledge bases because:
1. `embedding_model_arn` and `storage_config` were required `__init__`
arguments, neither of which applies to a fully managed knowledge base.
2. `storageConfiguration=self.storage_config` was always passed to the
underlying `create_knowledge_base` API call, which fails for managed knowledge
bases.
3. `knowledgeBaseConfiguration` was hardcoded to `type="VECTOR"` using
`embedding_model_arn`. Supplying custom configuration through
`create_knowledge_base_kwargs` raised `TypeError: got multiple values for
keyword argument`.
4. The vector indexing retry loop (`wait_for_indexing`) retries upon index
propagation errors which are only applicable to self-managed vector indices.
### Changes Made
- **Optional Parameters**: Made `embedding_model_arn` and `storage_config`
optional (`str | None = None` and `dict[str, Any] | None = None`).
- **Custom Configuration Support**: Added `knowledge_base_configuration`
parameter (and alias property `knowledge_base_config`) to support `MANAGED`
knowledge bases as well as custom configurations. Handled
`knowledgeBaseConfiguration` and `storageConfiguration` in
`create_knowledge_base_kwargs` gracefully without keyword conflicts.
- **Conditional `storageConfiguration`**: Only passes `storageConfiguration`
to the boto3 `create_knowledge_base` API when `storage_config` is provided.
- **Validation**: Added validation ensuring `role_arn` is specified, either
`knowledge_base_configuration` or `embedding_model_arn` is supplied, and
`storage_config` is provided when creating a self-managed `VECTOR` knowledge
base.
- **Index Retry Bypass**: Bypasses the vector index propagation retry loop
when the knowledge base type is `MANAGED`.
- **Documentation**: Updated Bedrock operator documentation in
`providers/amazon/docs/operators/bedrock.rst` with an explanation and code
example for creating a managed knowledge base.
- **Unit Tests**: Added comprehensive unit tests in
`TestBedrockCreateKnowledgeBaseOperator` covering:
- Creating a managed knowledge base with `knowledge_base_configuration`
- Creating a managed knowledge base via `create_knowledge_base_kwargs`
- Verifying managed knowledge bases bypass indexing retry loops
- Validation errors (missing `role_arn`, missing config/model, missing
`storage_config` for VECTOR KB)
- Property aliasing between `knowledge_base_config` and
`knowledge_base_configuration`
closes: #72592
---
##### Was generative AI tooling used to co-author this PR?
- [X] Yes (Antigravity)
Generated-by: Antigravity following [the
guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
--
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]
