Added streaming support for Jamba-Instruct

This commit is contained in:
JGalego 2024-08-17 01:56:15 +01:00
parent d65e5556dd
commit 58a5720f01
2 changed files with 7 additions and 2 deletions

View file

@ -114,6 +114,11 @@ class Ai21Provider(BaseBedrockProvider):
)
return body
def get_choice_text_from_stream(self, event) -> str:
rsp_dict = json.loads(event["chunk"]["bytes"])
completions = rsp_dict.get("choices", [{}])[0].get("delta", {}).get("content", "")
return completions
def _get_completion_from_dict(self, rsp_dict: dict) -> str:
if self.model_type == "j2":
# See https://docs.ai21.com/reference/j2-complete-ref

View file

@ -2,8 +2,6 @@ from metagpt.logs import logger
# max_tokens for each model
NOT_SUUPORT_STREAM_MODELS = {
# Jamba-Instruct
"ai21.jamba-instruct-v1:0": 256000,
# Jurassic-2 Mid-v1 and Ultra-v1
# + Legacy date: 2024-04-30 (us-west-2/Oregon)
# + EOL date: 2024-08-31 (us-west-2/Oregon)
@ -12,6 +10,8 @@ NOT_SUUPORT_STREAM_MODELS = {
}
SUPPORT_STREAM_MODELS = {
# Jamba-Instruct
"ai21.jamba-instruct-v1:0": 256000,
# Titan Text Large
"amazon.titan-tg1-large": 8000,
# Titan Text G1 - Express