|
|
|
|
@ -1,5 +1,6 @@
|
|
|
|
|
import json
|
|
|
|
|
import logging
|
|
|
|
|
import re
|
|
|
|
|
from collections.abc import Generator, Iterator
|
|
|
|
|
from typing import Any, Optional, Union, cast
|
|
|
|
|
|
|
|
|
|
@ -131,115 +132,58 @@ 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 = ""
|
|
|
|
|
processor = ChunkProcessor()
|
|
|
|
|
|
|
|
|
|
try:
|
|
|
|
|
for chunk in resp:
|
|
|
|
|
json_objects = processor.try_decode_chunk(chunk)
|
|
|
|
|
|
|
|
|
|
for data in json_objects:
|
|
|
|
|
if data.get("choices"):
|
|
|
|
|
choice = data["choices"][0]
|
|
|
|
|
|
|
|
|
|
if "delta" in choice and "content" in choice["delta"]:
|
|
|
|
|
chunk_content = choice["delta"]["content"]
|
|
|
|
|
assistant_prompt_message = AssistantPromptMessage(content=chunk_content, tool_calls=[])
|
|
|
|
|
|
|
|
|
|
if choice.get("finish_reason") is not None:
|
|
|
|
|
temp_assistant_prompt_message = AssistantPromptMessage(
|
|
|
|
|
content=full_response, tool_calls=[]
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
prompt_tokens = self._num_tokens_from_messages(messages=prompt_messages, tools=tools)
|
|
|
|
|
completion_tokens = self._num_tokens_from_messages(
|
|
|
|
|
messages=[temp_assistant_prompt_message], tools=[]
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
usage = self._calc_response_usage(
|
|
|
|
|
model=model,
|
|
|
|
|
credentials=credentials,
|
|
|
|
|
prompt_tokens=prompt_tokens,
|
|
|
|
|
completion_tokens=completion_tokens,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
yield LLMResultChunk(
|
|
|
|
|
model=model,
|
|
|
|
|
prompt_messages=prompt_messages,
|
|
|
|
|
system_fingerprint=None,
|
|
|
|
|
delta=LLMResultChunkDelta(
|
|
|
|
|
index=0,
|
|
|
|
|
message=assistant_prompt_message,
|
|
|
|
|
finish_reason=choice["finish_reason"],
|
|
|
|
|
usage=usage,
|
|
|
|
|
),
|
|
|
|
|
)
|
|
|
|
|
else:
|
|
|
|
|
yield LLMResultChunk(
|
|
|
|
|
model=model,
|
|
|
|
|
prompt_messages=prompt_messages,
|
|
|
|
|
system_fingerprint=None,
|
|
|
|
|
delta=LLMResultChunkDelta(index=0, message=assistant_prompt_message),
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
full_response += chunk_content
|
|
|
|
|
|
|
|
|
|
except Exception as e:
|
|
|
|
|
raise
|
|
|
|
|
|
|
|
|
|
if not full_response:
|
|
|
|
|
logger.warning("No content received from stream response")
|
|
|
|
|
buffer = ""
|
|
|
|
|
for chunk_bytes in resp:
|
|
|
|
|
buffer += chunk_bytes.decode("utf-8")
|
|
|
|
|
last_idx = 0
|
|
|
|
|
for match in re.finditer(r"^data:\s*(.+?)(\n\n)", buffer):
|
|
|
|
|
try:
|
|
|
|
|
data = json.loads(match.group(1).strip())
|
|
|
|
|
last_idx = match.span()[1]
|
|
|
|
|
|
|
|
|
|
if "content" in data["choices"][0]["delta"]:
|
|
|
|
|
chunk_content = data["choices"][0]["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=[])
|
|
|
|
|
prompt_tokens = self._num_tokens_from_messages(messages=prompt_messages, tools=tools)
|
|
|
|
|
completion_tokens = self._num_tokens_from_messages(
|
|
|
|
|
messages=[temp_assistant_prompt_message], tools=[]
|
|
|
|
|
)
|
|
|
|
|
usage = self._calc_response_usage(
|
|
|
|
|
model=model,
|
|
|
|
|
credentials=credentials,
|
|
|
|
|
prompt_tokens=prompt_tokens,
|
|
|
|
|
completion_tokens=completion_tokens,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
yield LLMResultChunk(
|
|
|
|
|
model=model,
|
|
|
|
|
prompt_messages=prompt_messages,
|
|
|
|
|
system_fingerprint=None,
|
|
|
|
|
delta=LLMResultChunkDelta(
|
|
|
|
|
index=0,
|
|
|
|
|
message=assistant_prompt_message,
|
|
|
|
|
finish_reason=data["choices"][0]["finish_reason"],
|
|
|
|
|
usage=usage,
|
|
|
|
|
),
|
|
|
|
|
)
|
|
|
|
|
else:
|
|
|
|
|
yield LLMResultChunk(
|
|
|
|
|
model=model,
|
|
|
|
|
prompt_messages=prompt_messages,
|
|
|
|
|
system_fingerprint=None,
|
|
|
|
|
delta=LLMResultChunkDelta(index=0, message=assistant_prompt_message),
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
full_response += chunk_content
|
|
|
|
|
except (json.JSONDecodeError, KeyError, IndexError) as e:
|
|
|
|
|
logger.info("json parse exception, content: {}".format(match.group(1).strip()))
|
|
|
|
|
pass
|
|
|
|
|
|
|
|
|
|
buffer = buffer[last_idx:]
|
|
|
|
|
|
|
|
|
|
def _invoke(
|
|
|
|
|
self,
|
|
|
|
|
|
Warren Chen
1 year ago
committed by
GitHub