stream test

This commit is contained in:
usamimeri_renko 2024-04-26 17:03:06 +08:00
parent 4c394a1cac
commit 8fafa2eb4e
4 changed files with 27 additions and 32 deletions

View file

@ -9,7 +9,7 @@ from metagpt.provider.bedrock.bedrock_provider import get_provider
from metagpt.provider.bedrock.utils import NOT_SUUPORT_STREAM_MODELS, get_max_tokens
try:
import boto3
from botocore.response import StreamingBody
from botocore.eventstream import EventStream
except ImportError:
raise ImportError(
"boto3 not found! please install it by `pip install boto3` ")
@ -70,7 +70,7 @@ class AmazonBedrockLLM(BaseLLM):
response_body = self._get_response_body(response)
return response_body
def invoke_model_with_response_stream(self, request_body) -> StreamingBody:
def invoke_model_with_response_stream(self, request_body) -> EventStream:
response = self.__client.invoke_model_with_response_stream(
modelId=self.config.model, body=request_body
)
@ -106,7 +106,6 @@ class AmazonBedrockLLM(BaseLLM):
messages, **self._generate_kwargs)
response = self.invoke_model_with_response_stream(request_body)
collected_content = []
for event in response["body"]:
chunk_text = self.__provider.get_choice_text_from_stream(event)

View file

@ -73,12 +73,7 @@ class AmazonProvider(BaseBedrockProvider):
return body
def _get_completion_from_dict(self, rsp_dict: dict) -> str:
return rsp_dict['results'][0]['outputText'].strip()
def get_choice_text_from_stream(self, event) -> str:
rsp_dict = json.loads(event["chunk"]["bytes"])
completions = rsp_dict["outputText"]
return completions
return rsp_dict['results'][0]['outputText']
PROVIDERS = {