add stop parameter (#86)

This commit is contained in:
yytdfc
2024-12-31 11:15:24 +08:00
committed by GitHub
parent b2c187c716
commit 093c6fa586
2 changed files with 7 additions and 0 deletions

View File

@@ -408,6 +408,12 @@ class BedrockModel(BaseChatModel):
"topP": chat_request.top_p, "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 = { args = {
"modelId": chat_request.model, "modelId": chat_request.model,
"messages": messages, "messages": messages,

View File

@@ -97,6 +97,7 @@ class ChatRequest(BaseModel):
n: int | None = 1 # Not used n: int | None = 1 # Not used
tools: list[Tool] | None = None tools: list[Tool] | None = None
tool_choice: str | object = "auto" tool_choice: str | object = "auto"
stop: list[str] | str | None = None
class Usage(BaseModel): class Usage(BaseModel):