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
+42 -89
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,
} }
} }
], ],
@@ -335,7 +315,7 @@ class BedrockModel(BaseChatModel):
return self._reframe_multi_payloard(messages) return self._reframe_multi_payloard(messages)
def _reframe_multi_payloard(self, messages: list) -> list: def _reframe_multi_payloard(self, messages: list) -> list:
""" Receive messages and reformat them to comply with the Claude format """Receive messages and reformat them to comply with the Claude format
With OpenAI format requests, it's not a problem to repeatedly receive messages from the same role, but With OpenAI format requests, it's not a problem to repeatedly receive messages from the same role, but
with Claude format requests, you cannot repeatedly receive messages from the same role. with Claude format requests, you cannot repeatedly receive messages from the same role.
@@ -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"