Support <think> tags (#117)

Author: LionsAd

src/api/models/bedrock.py

@@ -268,6 +268,7 @@ async def chat_stream(self, chat_request: ChatRequest) -> AsyncIterable[bytes]:
             response = await self._invoke_bedrock(chat_request, stream=True)
             message_id = self.generate_message_id()
             stream = response.get("stream")
+            self.think_emitted = False
             async for chunk in self._async_iterate(stream):
                 args = {"model_id": chat_request.model, "message_id": message_id, "chunk": chunk}
                 stream_response = self._create_response_stream(**args)
@@ -288,6 +289,7 @@ async def chat_stream(self, chat_request: ChatRequest) -> AsyncIterable[bytes]:
 
             # return an [DONE] message at the end.
             yield self.stream_response_to_bytes()
+            self.think_emitted = False  # Cleanup
         except Exception as e:
             logger.error("Stream error for model %s: %s", chat_request.model, str(e))
             error_event = Error(error=ErrorMessage(message=str(e)))
@@ -634,6 +636,9 @@ def _create_response(
                     logger.warning(
                         "Unknown tag in message content " + ",".join(c.keys())
                     )
+            if message.reasoning_content:
+                message.content = f"<think>{message.reasoning_content}</think>{message.content}"
+                message.reasoning_content = None
 
         response = ChatResponse(
             id=message_id,
@@ -702,11 +707,19 @@ def _create_response_stream(
                     content=delta["text"],
                 )
             elif "reasoningContent" in delta:
-                # ignore "signature" in the delta.
                 if "text" in delta["reasoningContent"]:
-                    message = ChatResponseMessage(
-                        reasoning_content=delta["reasoningContent"]["text"],
-                    )
+                    content = delta["reasoningContent"]["text"]
+                    if not self.think_emitted:
+                        # Port of "content_block_start" with "thinking"
+                        content = "<think>" + content
+                        self.think_emitted = True
+                    message = ChatResponseMessage(content=content)
+                elif "signature" in delta["reasoningContent"]:
+                    # Port of "signature_delta"
+                    if self.think_emitted:
+                        message = ChatResponseMessage(content="\n </think> \n\n")
+                    else:
+                        return None  # Ignore signature if no <think> started
             else:
                 # tool use
                 index = chunk["contentBlockDelta"]["contentBlockIndex"] - 1