from litellm import rerank
import os 

os.environ["AWS_ACCESS_KEY_ID"]=""
os.environ["AWS_SECRET_ACCESS_KEY"]=""
os.environ["AWS_REGION_NAME"]=""

response = rerank(
 model="bedrock/arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0",# provide the model ARN - get this here https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock/client/list_foundation_models.html
 query="hello",
 documents=["hello","world"],
 top_n=2,
)

print(response)

model_list:
-model_name: bedrock-rerank
litellm_params:
model: bedrock/arn:aws:bedrock:us-west-2::foundation-model/amazon.rerank-v1:0
aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID
aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY
aws_region_name: os.environ/AWS_REGION_NAME

litellm --config config.yaml

# RUNNING on http://0.0.0.0:4000

curl http://0.0.0.0:4000/rerank \
 -H "Authorization: Bearer sk-1234" \
 -H "Content-Type: application/json" \
 -d '{
 "model": "bedrock-rerank",
 "query": "What is the capital of the United States?",
 "documents": [
 "Carson City is the capital city of the American state of Nevada.",
 "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean. Its capital is Saipan.",
 "Washington, D.C. is the capital of the United States.",
 "Capital punishment has existed in the United States since before it was a country."
 ],
 "top_n": 3


 }'