diff --git a/src/api/models/bedrock.py b/src/api/models/bedrock.py index 16bfd33..4b58951 100644 --- a/src/api/models/bedrock.py +++ b/src/api/models/bedrock.py @@ -408,6 +408,12 @@ class BedrockModel(BaseChatModel): "topP": chat_request.top_p, } + if chat_request.stop is not None: + stop = chat_request.stop + if isinstance(stop, str): + stop = [stop] + inference_config["stopSequences"] = stop + args = { "modelId": chat_request.model, "messages": messages, diff --git a/src/api/schema.py b/src/api/schema.py index 15a25fb..b8ec75d 100644 --- a/src/api/schema.py +++ b/src/api/schema.py @@ -97,6 +97,7 @@ class ChatRequest(BaseModel): n: int | None = 1 # Not used tools: list[Tool] | None = None tool_choice: str | object = "auto" + stop: list[str] | str | None = None class Usage(BaseModel):