fix: tongyi json output (#5396)

This commit is contained in:
LXM
2024-06-22 12:25:23 +08:00
committed by GitHub
parent 3bbd75f1f2
commit e8ad0339a3
2 changed files with 96 additions and 18 deletions

View File

@@ -0,0 +1,84 @@
import json
import os
from collections.abc import Generator
from core.model_runtime.entities.llm_entities import LLMResultChunk, LLMResultChunkDelta
from core.model_runtime.entities.message_entities import AssistantPromptMessage, UserPromptMessage
from core.model_runtime.model_providers.tongyi.llm.llm import TongyiLargeLanguageModel
def test_invoke_model_with_json_response():
"""
Test the invocation of a model with JSON response.
"""
model_list = [
"qwen-max-0403",
"qwen-max-1201",
"qwen-max-longcontext",
"qwen-max",
"qwen-plus-chat",
"qwen-plus",
"qwen-turbo-chat",
"qwen-turbo",
]
for model_name in model_list:
print("testing model: ", model_name)
invoke_model_with_json_response(model_name)
def invoke_model_with_json_response(model_name="qwen-max-0403"):
"""
Method to invoke the model with JSON response format.
Args:
model_name (str): The name of the model to invoke. Defaults to "qwen-max-0403".
Returns:
None
"""
model = TongyiLargeLanguageModel()
response = model.invoke(
model=model_name,
credentials={
'dashscope_api_key': os.environ.get('TONGYI_DASHSCOPE_API_KEY')
},
prompt_messages=[
UserPromptMessage(
content='output json data with format `{"data": "test", "code": 200, "msg": "success"}'
)
],
model_parameters={
'temperature': 0.5,
'max_tokens': 50,
'response_format': 'JSON',
},
stream=True,
user="abc-123"
)
print("=====================================")
print(response)
assert isinstance(response, Generator)
output = ""
for chunk in response:
assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage)
output += chunk.delta.message.content
assert is_json(output)
def is_json(s):
"""
Check if a string is a valid JSON.
Args:
s (str): The string to check.
Returns:
bool: True if the string is a valid JSON, False otherwise.
"""
try:
json.loads(s)
except ValueError:
return False
return True