From 093c6fa586be04964820baaf1e3dca431f1fe823 Mon Sep 17 00:00:00 2001
From: yytdfc <fuchen@foxmail.com>
Date: Tue, 31 Dec 2024 11:15:24 +0800
Subject: [PATCH] add stop parameter (#86)

---
 src/api/models/bedrock.py | 6 ++++++
 src/api/schema.py         | 1 +
 2 files changed, 7 insertions(+)

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):