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

apply ruff linter

Browse Source
This commit is contained in:
Aiden Dai
2025-03-13 13:50:57 +08:00
parent 33e8fcfd3b
commit f21b9a2e84
10 changed files with 128 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.flake8
-19
View File
@@ -1,19 +0,0 @@
[flake8]
max-line-length = 120
ignore =
E203,W191,W503
exclude =
build
.git
__pycache__
.tox
venv
.venv
.venv-test
tmp*
deployment
cdk.out
node_modules
max-complexity = 10
require-code = True
.pre-commit-config.yaml
+8
View File
@@ -0,0 +1,8 @@
- repo: https://github.com/astral-sh/ruff-pre-commit
# Ruff version.
rev: v0.9.10
hooks:
# Run the linter.
- id: ruff
# Run the formatter.
- id: ruff-format
ruff.toml
+23
View File
@@ -0,0 +1,23 @@
line-length = 120
indent-width = 4
target-version = "py312"
exclude = [
".venv",
".vscode",
"test/*"
]
[lint]
select = ["E", "F"]
ignore = [
"E501",
"B008",
"C901",
"F401",
"W191",
]
[format]
# use double quotes for strings.
quote-style = "double"
src/api/auth.py
+5 -11
View File
@@ -16,9 +16,7 @@ if api_key_param:
# For backward compatibility. # For backward compatibility.
# Please now use secrets manager instead. # Please now use secrets manager instead.
ssm = boto3.client("ssm") ssm = boto3.client("ssm")
api_key = ssm.get_parameter(Name=api_key_param, WithDecryption=True)["Parameter"][ api_key = ssm.get_parameter(Name=api_key_param, WithDecryption=True)["Parameter"]["Value"]
"Value"
]
elif api_key_secret_arn: elif api_key_secret_arn:
sm = boto3.client("secretsmanager") sm = boto3.client("secretsmanager")
try: try:
@@ -26,11 +24,9 @@ elif api_key_secret_arn:
if "SecretString" in response: if "SecretString" in response:
secret = json.loads(response["SecretString"]) secret = json.loads(response["SecretString"])
api_key = secret["api_key"] api_key = secret["api_key"]
except ClientError as e: except ClientError:
raise RuntimeError( raise RuntimeError("Unable to retrieve API KEY, please ensure the secret ARN is correct")
"Unable to retrieve API KEY, please ensure the secret ARN is correct" except KeyError:
)
except KeyError as e:
raise RuntimeError('Please ensure the secret contains a "api_key" field') raise RuntimeError('Please ensure the secret contains a "api_key" field')
elif api_key_env: elif api_key_env:
api_key = api_key_env api_key = api_key_env
@@ -45,6 +41,4 @@ def api_key_auth(
credentials: Annotated[HTTPAuthorizationCredentials, Depends(security)], credentials: Annotated[HTTPAuthorizationCredentials, Depends(security)],
): ):
if credentials.credentials != api_key: if credentials.credentials != api_key:
raise HTTPException( raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API Key")
status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API Key"
)
src/api/models/base.py
+1 -3
View File
@@ -43,9 +43,7 @@ class BaseChatModel(ABC):
return "chatcmpl-" + str(uuid.uuid4())[:8] return "chatcmpl-" + str(uuid.uuid4())[:8]
@staticmethod @staticmethod
def stream_response_to_bytes( def stream_response_to_bytes(response: ChatStreamResponse | None = None) -> bytes:
response: ChatStreamResponse | None = None
) -> bytes:
if response: if response:
# to populate other fields when using exclude_unset=True # to populate other fields when using exclude_unset=True
response.system_fingerprint = "fp" response.system_fingerprint = "fp"
src/api/models/bedrock.py
+41 -88
View File
@@ -36,7 +36,6 @@ from api.schema import (
EmbeddingsResponse, EmbeddingsResponse,
EmbeddingsUsage, EmbeddingsUsage,
Embedding, Embedding,
) )
from api.setting import DEBUG, AWS_REGION, ENABLE_CROSS_REGION_INFERENCE, DEFAULT_MODEL from api.setting import DEBUG, AWS_REGION, ENABLE_CROSS_REGION_INFERENCE, DEFAULT_MODEL
@@ -50,15 +49,15 @@ bedrock_runtime = boto3.client(
config=config, config=config,
) )
bedrock_client = boto3.client( bedrock_client = boto3.client(
service_name='bedrock', service_name="bedrock",
region_name=AWS_REGION, region_name=AWS_REGION,
config=config, config=config,
) )
def get_inference_region_prefix(): def get_inference_region_prefix():
if AWS_REGION.startswith('ap-'): if AWS_REGION.startswith("ap-"):
return 'apac' return "apac"
return AWS_REGION[:2] return AWS_REGION[:2]
@@ -88,49 +87,38 @@ def list_bedrock_models() -> dict:
profile_list = [] profile_list = []
if ENABLE_CROSS_REGION_INFERENCE: if ENABLE_CROSS_REGION_INFERENCE:
# List system defined inference profile IDs # List system defined inference profile IDs
response = bedrock_client.list_inference_profiles( response = bedrock_client.list_inference_profiles(maxResults=1000, typeEquals="SYSTEM_DEFINED")
maxResults=1000, profile_list = [p["inferenceProfileId"] for p in response["inferenceProfileSummaries"]]
typeEquals='SYSTEM_DEFINED'
)
profile_list = [p['inferenceProfileId'] for p in response['inferenceProfileSummaries']]
# List foundation models, only cares about text outputs here. # List foundation models, only cares about text outputs here.
response = bedrock_client.list_foundation_models( response = bedrock_client.list_foundation_models(byOutputModality="TEXT")
byOutputModality='TEXT'
)
for model in response['modelSummaries']: for model in response["modelSummaries"]:
model_id = model.get('modelId', 'N/A') model_id = model.get("modelId", "N/A")
stream_supported = model.get('responseStreamingSupported', True) stream_supported = model.get("responseStreamingSupported", True)
status = model['modelLifecycle'].get('status', 'ACTIVE') status = model["modelLifecycle"].get("status", "ACTIVE")
# currently, use this to filter out rerank models and legacy models # currently, use this to filter out rerank models and legacy models
if not stream_supported or status not in ["ACTIVE", "LEGACY"]: if not stream_supported or status not in ["ACTIVE", "LEGACY"]:
continue continue
inference_types = model.get('inferenceTypesSupported', []) inference_types = model.get("inferenceTypesSupported", [])
input_modalities = model['inputModalities'] input_modalities = model["inputModalities"]
# Add on-demand model list # Add on-demand model list
if 'ON_DEMAND' in inference_types: if "ON_DEMAND" in inference_types:
model_list[model_id] = { model_list[model_id] = {"modalities": input_modalities}
'modalities': input_modalities
}
# Add cross-region inference model list. # Add cross-region inference model list.
profile_id = cr_inference_prefix + '.' + model_id profile_id = cr_inference_prefix + "." + model_id
if profile_id in profile_list: if profile_id in profile_list:
model_list[profile_id] = { model_list[profile_id] = {"modalities": input_modalities}
'modalities': input_modalities
}
except Exception as e: except Exception as e:
logger.error(f"Unable to list models: {str(e)}") logger.error(f"Unable to list models: {str(e)}")
if not model_list: if not model_list:
# In case stack not updated. # In case stack not updated.
model_list[DEFAULT_MODEL] = { model_list[DEFAULT_MODEL] = {"modalities": ["TEXT", "IMAGE"]}
'modalities': ["TEXT", "IMAGE"]
}
return model_list return model_list
@@ -140,7 +128,6 @@ bedrock_model_list = list_bedrock_models()
class BedrockModel(BaseChatModel): class BedrockModel(BaseChatModel):
def list_models(self) -> list[str]: def list_models(self) -> list[str]:
"""Always refresh the latest model list""" """Always refresh the latest model list"""
global bedrock_model_list global bedrock_model_list
@@ -224,10 +211,7 @@ class BedrockModel(BaseChatModel):
logger.info("Proxy response :" + stream_response.model_dump_json()) logger.info("Proxy response :" + stream_response.model_dump_json())
if stream_response.choices: if stream_response.choices:
yield self.stream_response_to_bytes(stream_response) yield self.stream_response_to_bytes(stream_response)
elif ( elif chat_request.stream_options and chat_request.stream_options.include_usage:
chat_request.stream_options
and chat_request.stream_options.include_usage
):
# An empty choices for Usage as per OpenAI doc below: # An empty choices for Usage as per OpenAI doc below:
# if you set stream_options: {"include_usage": true}. # if you set stream_options: {"include_usage": true}.
# an additional chunk will be streamed before the data: [DONE] message. # an additional chunk will be streamed before the data: [DONE] message.
@@ -277,9 +261,7 @@ class BedrockModel(BaseChatModel):
messages.append( messages.append(
{ {
"role": message.role, "role": message.role,
"content": self._parse_content_parts( "content": self._parse_content_parts(message, chat_request.model),
message, chat_request.model
),
} }
) )
elif isinstance(message, AssistantMessage): elif isinstance(message, AssistantMessage):
@@ -288,9 +270,7 @@ class BedrockModel(BaseChatModel):
messages.append( messages.append(
{ {
"role": message.role, "role": message.role,
"content": self._parse_content_parts( "content": self._parse_content_parts(message, chat_request.model),
message, chat_request.model
),
} }
) )
if message.tool_calls: if message.tool_calls:
@@ -305,7 +285,7 @@ class BedrockModel(BaseChatModel):
"toolUse": { "toolUse": {
"toolUseId": tool_call.id, "toolUseId": tool_call.id,
"name": tool_call.function.name, "name": tool_call.function.name,
"input": tool_input "input": tool_input,
} }
} }
], ],
@@ -364,16 +344,13 @@ class BedrockModel(BaseChatModel):
# Search through the list of messages and combine messages from the same role into one list # Search through the list of messages and combine messages from the same role into one list
for message in messages: for message in messages:
next_role = message['role'] next_role = message["role"]
next_content = message['content'] next_content = message["content"]
# If the next role is different from the previous message, add the previous role's messages to the list # If the next role is different from the previous message, add the previous role's messages to the list
if next_role != current_role: if next_role != current_role:
if current_content: if current_content:
reformatted_messages.append({ reformatted_messages.append({"role": current_role, "content": current_content})
"role": current_role,
"content": current_content
})
# Switch to the new role # Switch to the new role
current_role = next_role current_role = next_role
current_content = [] current_content = []
@@ -386,10 +363,7 @@ class BedrockModel(BaseChatModel):
# Add the last role's messages to the list # Add the last role's messages to the list
if current_content: if current_content:
reformatted_messages.append({ reformatted_messages.append({"role": current_role, "content": current_content})
"role": current_role,
"content": current_content
})
return reformatted_messages return reformatted_messages
@@ -426,25 +400,20 @@ class BedrockModel(BaseChatModel):
# From OpenAI api, the max_token is not supported in reasoning mode # From OpenAI api, the max_token is not supported in reasoning mode
# Use max_completion_tokens if provided. # Use max_completion_tokens if provided.
max_tokens = chat_request.max_completion_tokens if chat_request.max_completion_tokens else chat_request.max_tokens max_tokens = (
chat_request.max_completion_tokens if chat_request.max_completion_tokens else chat_request.max_tokens
)
budget_tokens = self._calc_budget_tokens(max_tokens, chat_request.reasoning_effort) budget_tokens = self._calc_budget_tokens(max_tokens, chat_request.reasoning_effort)
inference_config["maxTokens"] = max_tokens inference_config["maxTokens"] = max_tokens
# unset topP - Not supported # unset topP - Not supported
inference_config.pop("topP") inference_config.pop("topP")
args["additionalModelRequestFields"] = { args["additionalModelRequestFields"] = {
"reasoning_config": { "reasoning_config": {"type": "enabled", "budget_tokens": budget_tokens}
"type": "enabled",
"budget_tokens": budget_tokens
}
} }
# add tool config # add tool config
if chat_request.tools: if chat_request.tools:
args["toolConfig"] = { args["toolConfig"] = {"tools": [self._convert_tool_spec(t.function) for t in chat_request.tools]}
"tools": [
self._convert_tool_spec(t.function) for t in chat_request.tools
]
}
if chat_request.tool_choice and not chat_request.model.startswith("meta.llama3-1-"): if chat_request.tool_choice and not chat_request.model.startswith("meta.llama3-1-"):
if isinstance(chat_request.tool_choice, str): if isinstance(chat_request.tool_choice, str):
@@ -458,7 +427,8 @@ class BedrockModel(BaseChatModel):
# Specific tool to use # Specific tool to use
assert "function" in chat_request.tool_choice assert "function" in chat_request.tool_choice
args["toolConfig"]["toolChoice"] = { args["toolConfig"]["toolChoice"] = {
"tool": {"name": chat_request.tool_choice["function"].get("name", "")}} "tool": {"name": chat_request.tool_choice["function"].get("name", "")}
}
return args return args
def _create_response( def _create_response(
@@ -470,7 +440,6 @@ class BedrockModel(BaseChatModel):
input_tokens: int = 0, input_tokens: int = 0,
output_tokens: int = 0, output_tokens: int = 0,
) -> ChatResponse: ) -> ChatResponse:
message = ChatResponseMessage( message = ChatResponseMessage(
role="assistant", role="assistant",
) )
@@ -524,9 +493,7 @@ class BedrockModel(BaseChatModel):
response.created = int(time.time()) response.created = int(time.time())
return response return response
def _create_response_stream( def _create_response_stream(self, model_id: str, message_id: str, chunk: dict) -> ChatStreamResponse | None:
self, model_id: str, message_id: str, chunk: dict
) -> ChatStreamResponse | None:
"""Parsing the Bedrock stream response chunk. """Parsing the Bedrock stream response chunk.
Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html#message-inference-examples Ref: https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference.html#message-inference-examples
@@ -583,7 +550,7 @@ class BedrockModel(BaseChatModel):
index=index, index=index,
function=ResponseFunction( function=ResponseFunction(
arguments=delta["toolUse"]["input"], arguments=delta["toolUse"]["input"],
) ),
) )
] ]
) )
@@ -641,7 +608,6 @@ class BedrockModel(BaseChatModel):
response = requests.get(image_url) response = requests.get(image_url)
# Check if the request was successful # Check if the request was successful
if response.status_code == 200: if response.status_code == 200:
content_type = response.headers.get("Content-Type") content_type = response.headers.get("Content-Type")
if not content_type.startswith("image"): if not content_type.startswith("image"):
content_type = "image/jpeg" content_type = "image/jpeg"
@@ -649,9 +615,7 @@ class BedrockModel(BaseChatModel):
image_content = response.content image_content = response.content
return image_content, content_type return image_content, content_type
else: else:
raise HTTPException( raise HTTPException(status_code=500, detail="Unable to access the image url")
status_code=500, detail="Unable to access the image url"
)
def _parse_content_parts( def _parse_content_parts(
self, self,
@@ -695,7 +659,7 @@ class BedrockModel(BaseChatModel):
@staticmethod @staticmethod
def is_supported_modality(model_id: str, modality: str = "IMAGE") -> bool: def is_supported_modality(model_id: str, modality: str = "IMAGE") -> bool:
model = bedrock_model_list.get(model_id) model = bedrock_model_list.get(model_id)
modalities = model.get('modalities', []) modalities = model.get("modalities", [])
if modality in modalities: if modality in modalities:
return True return True
return False return False
@@ -740,7 +704,7 @@ class BedrockModel(BaseChatModel):
"max_tokens": "length", "max_tokens": "length",
"stop_sequence": "stop", "stop_sequence": "stop",
"complete": "stop", "complete": "stop",
"content_filtered": "content_filter" "content_filtered": "content_filter",
} }
return finish_reason_mapping.get(finish_reason.lower(), finish_reason.lower()) return finish_reason_mapping.get(finish_reason.lower(), finish_reason.lower())
return None return None
@@ -803,7 +767,6 @@ class BedrockEmbeddingsModel(BaseEmbeddingsModel, ABC):
class CohereEmbeddingsModel(BedrockEmbeddingsModel): class CohereEmbeddingsModel(BedrockEmbeddingsModel):
def _parse_args(self, embeddings_request: EmbeddingsRequest) -> dict: def _parse_args(self, embeddings_request: EmbeddingsRequest) -> dict:
texts = [] texts = []
if isinstance(embeddings_request.input, str): if isinstance(embeddings_request.input, str):
@@ -834,9 +797,7 @@ class CohereEmbeddingsModel(BedrockEmbeddingsModel):
return args return args
def embed(self, embeddings_request: EmbeddingsRequest) -> EmbeddingsResponse: def embed(self, embeddings_request: EmbeddingsRequest) -> EmbeddingsResponse:
response = self._invoke_model( response = self._invoke_model(args=self._parse_args(embeddings_request), model_id=embeddings_request.model)
args=self._parse_args(embeddings_request), model_id=embeddings_request.model
)
response_body = json.loads(response.get("body").read()) response_body = json.loads(response.get("body").read())
if DEBUG: if DEBUG:
logger.info("Bedrock response body: " + str(response_body)) logger.info("Bedrock response body: " + str(response_body))
@@ -849,19 +810,13 @@ class CohereEmbeddingsModel(BedrockEmbeddingsModel):
class TitanEmbeddingsModel(BedrockEmbeddingsModel): class TitanEmbeddingsModel(BedrockEmbeddingsModel):
def _parse_args(self, embeddings_request: EmbeddingsRequest) -> dict: def _parse_args(self, embeddings_request: EmbeddingsRequest) -> dict:
if isinstance(embeddings_request.input, str): if isinstance(embeddings_request.input, str):
input_text = embeddings_request.input input_text = embeddings_request.input
elif ( elif isinstance(embeddings_request.input, list) and len(embeddings_request.input) == 1:
isinstance(embeddings_request.input, list)
and len(embeddings_request.input) == 1
):
input_text = embeddings_request.input[0] input_text = embeddings_request.input[0]
else: else:
raise ValueError( raise ValueError("Amazon Titan Embeddings models support only single strings as input.")
"Amazon Titan Embeddings models support only single strings as input."
)
args = { args = {
"inputText": input_text, "inputText": input_text,
# Note: inputImage is not supported! # Note: inputImage is not supported!
@@ -875,9 +830,7 @@ class TitanEmbeddingsModel(BedrockEmbeddingsModel):
return args return args
def embed(self, embeddings_request: EmbeddingsRequest) -> EmbeddingsResponse: def embed(self, embeddings_request: EmbeddingsRequest) -> EmbeddingsResponse:
response = self._invoke_model( response = self._invoke_model(args=self._parse_args(embeddings_request), model_id=embeddings_request.model)
args=self._parse_args(embeddings_request), model_id=embeddings_request.model
)
response_body = json.loads(response.get("body").read()) response_body = json.loads(response.get("body").read())
if DEBUG: if DEBUG:
logger.info("Bedrock response body: " + str(response_body)) logger.info("Bedrock response body: " + str(response_body))
src/api/routers/chat.py
+2 -4
View File
@@ -30,7 +30,7 @@ async def chat_completions(
} }
], ],
), ),
] ],
): ):
if chat_request.model.lower().startswith("gpt-"): if chat_request.model.lower().startswith("gpt-"):
chat_request.model = DEFAULT_MODEL chat_request.model = DEFAULT_MODEL
@@ -39,7 +39,5 @@ async def chat_completions(
model = BedrockModel() model = BedrockModel()
model.validate(chat_request) model.validate(chat_request)
if chat_request.stream: if chat_request.stream:
return StreamingResponse( return StreamingResponse(content=model.chat_stream(chat_request), media_type="text/event-stream")
content=model.chat_stream(chat_request), media_type="text/event-stream"
)
return model.chat(chat_request) return model.chat(chat_request)
src/api/routers/embeddings.py
+2 -4
View File
@@ -21,13 +21,11 @@ async def embeddings(
examples=[ examples=[
{ {
"model": "cohere.embed-multilingual-v3", "model": "cohere.embed-multilingual-v3",
"input": [ "input": ["Your text string goes here"],
"Your text string goes here"
],
} }
], ],
), ),
] ],
): ):
if embeddings_request.model.lower().startswith("text-embedding-"): if embeddings_request.model.lower().startswith("text-embedding-"):
embeddings_request.model = DEFAULT_EMBEDDING_MODEL embeddings_request.model = DEFAULT_EMBEDDING_MODEL
src/api/routers/model.py
+2 -4
View File
@@ -22,9 +22,7 @@ async def validate_model_id(model_id: str):
@router.get("", response_model=Models) @router.get("", response_model=Models)
async def list_models(): async def list_models():
model_list = [ model_list = [Model(id=model_id) for model_id in chat_model.list_models()]
Model(id=model_id) for model_id in chat_model.list_models()
]
return Models(data=model_list) return Models(data=model_list)
@@ -36,7 +34,7 @@ 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)
src/api/setting.py
+2 -6
View File
@@ -13,10 +13,6 @@ Use OpenAI-Compatible RESTful APIs for Amazon Bedrock models.
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 = os.environ.get("DEFAULT_MODEL", "anthropic.claude-3-sonnet-20240229-v1:0")
"DEFAULT_MODEL", "anthropic.claude-3-sonnet-20240229-v1:0" DEFAULT_EMBEDDING_MODEL = os.environ.get("DEFAULT_EMBEDDING_MODEL", "cohere.embed-multilingual-v3")
)
DEFAULT_EMBEDDING_MODEL = os.environ.get(
"DEFAULT_EMBEDDING_MODEL", "cohere.embed-multilingual-v3"
)
ENABLE_CROSS_REGION_INFERENCE = os.environ.get("ENABLE_CROSS_REGION_INFERENCE", "true").lower() != "false" ENABLE_CROSS_REGION_INFERENCE = os.environ.get("ENABLE_CROSS_REGION_INFERENCE", "true").lower() != "false"