gauravSsinha opened a new pull request, #67787:
URL: https://github.com/apache/airflow/pull/67787
## Description
Adds `BedrockRerankOperator` that uses the Bedrock Agent Runtime `Rerank`
API to score and reorder documents by relevance to a query.
## Motivation
Reranking is a critical step in production RAG pipelines — it improves
answer quality by reordering retrieved documents before passing them to a
generative model. Amazon Bedrock now supports reranking via the Cohere Rerank
model, but there's no Airflow operator for it.
## Changes
Added `BedrockRerankOperator` to
`providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py`:
- Uses `BedrockAgentRuntimeHook` (bedrock-agent-runtime client)
- Configurable `model_id` (defaults to `cohere.rerank-v3-5:0`)
- `number_of_results` parameter to limit returned documents
- All key parameters are template-able
- Returns ranked results with relevance scores
## Example Usage
```python
rerank = BedrockRerankOperator(
task_id='rerank_results',
query='What is serverless computing?',
documents=[{'textDocument': {'text': doc}} for doc in retrieved_docs],
model_id='cohere.rerank-v3-5:0',
number_of_results=5,
)
```
## Related
- Bedrock Rerank API:
https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_Rerank.html
- Existing operators: `BedrockRaGOperator`, `BedrockRetrieveOperator`
---
##### Was generative AI tooling used to co-author this PR?
- [x] Yes
Generated-by: Kiro (AI IDE)
--
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]
