[Fix] modify sagemaker llm (#12274)

api/core/model_runtime/model_providers/sagemaker/llm/llm.py

@@ -1,6 +1,5 @@
 import json
 import logging
-import re
 from collections.abc import Generator, Iterator
 from typing import Any, Optional, Union, cast
 
@@ -132,58 +131,115 @@ class SageMakerLargeLanguageModel(LargeLanguageModel):
         """
         handle stream chat generate response
         """
+
+        class ChunkProcessor:
+            def __init__(self):
+                self.buffer = bytearray()
+
+            def try_decode_chunk(self, chunk: bytes) -> list[dict]:
+                """尝试从chunk中解码出完整的JSON对象"""
+                self.buffer.extend(chunk)
+                results = []
+
+                while True:
+                    try:
+                        start = self.buffer.find(b"{")
+                        if start == -1:
+                            self.buffer.clear()
+                            break
+
+                        bracket_count = 0
+                        end = start
+
+                        for i in range(start, len(self.buffer)):
+                            if self.buffer[i] == ord("{"):
+                                bracket_count += 1
+                            elif self.buffer[i] == ord("}"):
+                                bracket_count -= 1
+                                if bracket_count == 0:
+                                    end = i + 1
+                                    break
+
+                        if bracket_count != 0:
+                            # JSON不完整，等待更多数据
+                            if start > 0:
+                                self.buffer = self.buffer[start:]
+                            break
+
+                        json_bytes = self.buffer[start:end]
+                        try:
+                            data = json.loads(json_bytes)
+                            results.append(data)
+                            self.buffer = self.buffer[end:]
+                        except json.JSONDecodeError:
+                            self.buffer = self.buffer[start + 1 :]
+
+                    except Exception as e:
+                        logger.debug(f"Warning: Error processing chunk ({str(e)})")
+                        if start > 0:
+                            self.buffer = self.buffer[start:]
+                        break
+
+                return results
+
         full_response = ""
-        buffer = ""
+        processor = ChunkProcessor()
-        for chunk_bytes in resp:
+
-            buffer += chunk_bytes.decode("utf-8")
+        try:
-            last_idx = 0
+            for chunk in resp:
-            for match in re.finditer(r"^data:\s*(.+?)(\n\n)", buffer):
+                json_objects = processor.try_decode_chunk(chunk)
-                try:
+
-                    data = json.loads(match.group(1).strip())
+                for data in json_objects:
-                    last_idx = match.span()[1]
+                    if data.get("choices"):
-
+                        choice = data["choices"][0]
-                    if "content" in data["choices"][0]["delta"]:
+
-                        chunk_content = data["choices"][0]["delta"]["content"]
+                        if "delta" in choice and "content" in choice["delta"]:
-                        assistant_prompt_message = AssistantPromptMessage(content=chunk_content, tool_calls=[])
+                            chunk_content = choice["delta"]["content"]
-
+                            assistant_prompt_message = AssistantPromptMessage(content=chunk_content, tool_calls=[])
-                        if data["choices"][0]["finish_reason"] is not None:
+
-                            temp_assistant_prompt_message = AssistantPromptMessage(content=full_response, tool_calls=[])
+                            if choice.get("finish_reason") is not None:
-                            prompt_tokens = self._num_tokens_from_messages(messages=prompt_messages, tools=tools)
+                                temp_assistant_prompt_message = AssistantPromptMessage(
-                            completion_tokens = self._num_tokens_from_messages(
+                                    content=full_response, tool_calls=[]
-                                messages=[temp_assistant_prompt_message], tools=[]
+                                )
-                            )
+
-                            usage = self._calc_response_usage(
+                                prompt_tokens = self._num_tokens_from_messages(messages=prompt_messages, tools=tools)
-                                model=model,
+                                completion_tokens = self._num_tokens_from_messages(
-                                credentials=credentials,
+                                    messages=[temp_assistant_prompt_message], tools=[]
-                                prompt_tokens=prompt_tokens,
+                                )
-                                completion_tokens=completion_tokens,
+
-                            )
+                                usage = self._calc_response_usage(
-
+                                    model=model,
-                            yield LLMResultChunk(
+                                    credentials=credentials,
-                                model=model,
+                                    prompt_tokens=prompt_tokens,
-                                prompt_messages=prompt_messages,
+                                    completion_tokens=completion_tokens,
-                                system_fingerprint=None,
+                                )
-                                delta=LLMResultChunkDelta(
+
-                                    index=0,
+                                yield LLMResultChunk(
-                                    message=assistant_prompt_message,
+                                    model=model,
-                                    finish_reason=data["choices"][0]["finish_reason"],
+                                    prompt_messages=prompt_messages,
-                                    usage=usage,
+                                    system_fingerprint=None,
-                                ),
+                                    delta=LLMResultChunkDelta(
-                            )
+                                        index=0,
-                        else:
+                                        message=assistant_prompt_message,
-                            yield LLMResultChunk(
+                                        finish_reason=choice["finish_reason"],
-                                model=model,
+                                        usage=usage,
-                                prompt_messages=prompt_messages,
+                                    ),
-                                system_fingerprint=None,
+                                )
-                                delta=LLMResultChunkDelta(index=0, message=assistant_prompt_message),
+                            else:
-                            )
+                                yield LLMResultChunk(
-
+                                    model=model,
-                            full_response += chunk_content
+                                    prompt_messages=prompt_messages,
-                except (json.JSONDecodeError, KeyError, IndexError) as e:
+                                    system_fingerprint=None,
-                    logger.info("json parse exception, content: {}".format(match.group(1).strip()))
+                                    delta=LLMResultChunkDelta(index=0, message=assistant_prompt_message),
-                    pass
+                                )
-
+
-            buffer = buffer[last_idx:]
+                                full_response += chunk_content
+
+        except Exception as e:
+            raise
+
+        if not full_response:
+            logger.warning("No content received from stream response")
 
     def _invoke(
         self,