AVB/bedrock-access-gateway
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Clean up code

Browse Source
This commit is contained in:
Aiden Dai
2024-04-03 11:20:31 +08:00
parent f1440602ce
commit 268c5ef0f1
2 changed files with 13 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
src/api/routers/model.py
View File

@@ -10,7 +10,6 @@ router = APIRouter()
router = APIRouter( router = APIRouter(
prefix="/models", prefix="/models",
tags=["items"],
dependencies=[Depends(api_key_auth)], dependencies=[Depends(api_key_auth)],
# responses={404: {"description": "Not found"}}, # responses={404: {"description": "Not found"}},
) )
@@ -18,12 +17,13 @@ router = APIRouter(
async def validate_model_id(model_id: str): async def validate_model_id(model_id: str):
if model_id not in (SUPPORTED_BEDROCK_MODELS | SUPPORTED_BEDROCK_EMBEDDING_MODELS).keys(): if model_id not in (SUPPORTED_BEDROCK_MODELS | SUPPORTED_BEDROCK_EMBEDDING_MODELS).keys():
raise HTTPException(status_code=400, detail="Unsupported Model Id") raise HTTPException(status_code=500, detail="Unsupported Model Id")
@router.get("/", response_model=Models) @router.get("/", response_model=Models)
async def list_models(): async def list_models():
model_list = [Model(id=model_id) for model_id in (SUPPORTED_BEDROCK_MODELS | SUPPORTED_BEDROCK_EMBEDDING_MODELS).keys()] model_list = [Model(id=model_id) for model_id in
(SUPPORTED_BEDROCK_MODELS | SUPPORTED_BEDROCK_EMBEDDING_MODELS).keys()]
return Models(data=model_list) return Models(data=model_list)
@@ -32,10 +32,10 @@ async def list_models():
response_model=Model, response_model=Model,
) )
async def get_model( async def get_model(
model_id: Annotated[ model_id: Annotated[
str, str,
Path(description="Model ID", example="anthropic.claude-3-sonnet-20240229-v1:0"), Path(description="Model ID", example="anthropic.claude-3-sonnet-20240229-v1:0"),
] ]
): ):
await validate_model_id(model_id) await validate_model_id(model_id)
return Model(id=model_id) return Model(id=model_id)

10
src/api/setting.py
View File

@@ -26,11 +26,13 @@ List of Amazon Bedrock models currently supported:
# Embeddings # Embeddings
- cohere.embed-multilingual-v3 - cohere.embed-multilingual-v3
- cohere.embed-english-v3 - cohere.embed-english-v3
- amazon.titan-embed-text-v1
- amazon.titan-embed-image-v1
""" """
DEBUG = os.environ.get("DEBUG", "false").lower() != "false" DEBUG = os.environ.get("DEBUG", "false").lower() != "false"
AWS_REGION = os.environ.get("AWS_REGION", "us-west-2") AWS_REGION = os.environ.get("AWS_REGION", "us-west-2")
DEFAULT_MODEL = os.environ.get("DEFAULT_MODEL", "anthropic.claude-3-sonnet-20240229-v1:0") DEFAULT_MODEL = os.environ.get(
DEFAULT_EMBEDDING_MODEL = os.environ.get("DEFAULT_EMBEDDING_MODEL", "cohere.embed-multilingual-v3") "DEFAULT_MODEL", "anthropic.claude-3-sonnet-20240229-v1:0"
)
DEFAULT_EMBEDDING_MODEL = os.environ.get(
"DEFAULT_EMBEDDING_MODEL", "cohere.embed-multilingual-v3"
)