AWS Bedrock - Rerank API

AWS Bedrock - Rerank API

Cohere /rerank 형식으로 Bedrock의 Rerank API를 사용해요.

출처: 문서

본문

비용 추적 ✅ Bedrock Rerank API 호출에 대한 비용 추적이 지원돼요.

지원 파라미터 (Supported Parameters)

  • model - 기반 모델 ARN
  • query - rerank 대상 쿼리
  • documents - rerank할 문서 목록
  • top_n - 반환할 결과 수

사용법 (Usage)

SDK

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)

PROXY

1. config.yaml 설정:

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

2. Proxy 서버 시작:

litellm --config config.yaml
# RUNNING on http://0.0.0.0:4000

3. 테스트:

curl http://0.0.0.0:4000/rerank \
  -H "Authorization: Bearer ***" \
  -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
  }'

더 알아보기 (Learn more)