Fixed <think> </think> tags for GPT-OSS in bedrock.py (#200)

Added handling for message and content block deltas, including safety checks for open thinking tags.

Results in working reasoning and makes GPT-OSS 80/120b usable in frontends that expect closing thinking tags.
This commit is contained in:
user-error1
2025-11-20 14:29:20 +08:00
committed by GitHub
parent 7e03ab062d
commit ce4cfabb21
+31 -7
View File
@@ -964,11 +964,13 @@ class BedrockModel(BaseChatModel):
finish_reason = None finish_reason = None
message = None message = None
usage = None usage = None
if "messageStart" in chunk: if "messageStart" in chunk:
message = ChatResponseMessage( message = ChatResponseMessage(
role=chunk["messageStart"]["role"], role=chunk["messageStart"]["role"],
content="", content="",
) )
if "contentBlockStart" in chunk: if "contentBlockStart" in chunk:
# tool call start # tool call start
delta = chunk["contentBlockStart"]["start"] delta = chunk["contentBlockStart"]["start"]
@@ -988,25 +990,30 @@ class BedrockModel(BaseChatModel):
) )
] ]
) )
if "contentBlockDelta" in chunk: if "contentBlockDelta" in chunk:
delta = chunk["contentBlockDelta"]["delta"] delta = chunk["contentBlockDelta"]["delta"]
if "text" in delta: if "text" in delta:
# stream content # Regular text content - close thinking tag if open
message = ChatResponseMessage( content = delta["text"]
content=delta["text"], if self.think_emitted:
) # Transition from reasoning to regular text
content = "</think>" + content
self.think_emitted = False
message = ChatResponseMessage(content=content)
elif "reasoningContent" in delta: elif "reasoningContent" in delta:
if "text" in delta["reasoningContent"]: if "text" in delta["reasoningContent"]:
content = delta["reasoningContent"]["text"] content = delta["reasoningContent"]["text"]
if not self.think_emitted: if not self.think_emitted:
# Port of "content_block_start" with "thinking" # Start of reasoning content
content = "<think>" + content content = "<think>" + content
self.think_emitted = True self.think_emitted = True
message = ChatResponseMessage(content=content) message = ChatResponseMessage(content=content)
elif "signature" in delta["reasoningContent"]: elif "signature" in delta["reasoningContent"]:
# Port of "signature_delta" # Port of "signature_delta" (for models that send it)
if self.think_emitted: if self.think_emitted:
message = ChatResponseMessage(content="\n </think> \n\n") message = ChatResponseMessage(content="</think>")
self.think_emitted = False
else: else:
return None # Ignore signature if no <think> started return None # Ignore signature if no <think> started
else: else:
@@ -1022,7 +1029,23 @@ class BedrockModel(BaseChatModel):
) )
] ]
) )
if "messageStop" in chunk: if "messageStop" in chunk:
# Safety check: Close any open thinking tags before message stops
if self.think_emitted:
self.think_emitted = False
return ChatStreamResponse(
id=message_id,
model=model_id,
choices=[
ChoiceDelta(
index=0,
delta=ChatResponseMessage(content="</think>"),
logprobs=None,
finish_reason=None,
)
],
)
message = ChatResponseMessage() message = ChatResponseMessage()
finish_reason = chunk["messageStop"]["stopReason"] finish_reason = chunk["messageStop"]["stopReason"]
@@ -1063,6 +1086,7 @@ class BedrockModel(BaseChatModel):
prompt_tokens_details=prompt_tokens_details, prompt_tokens_details=prompt_tokens_details,
), ),
) )
if message: if message:
return ChatStreamResponse( return ChatStreamResponse(
id=message_id, id=message_id,