estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

chore(api/tests): apply ruff reformat #7590 (#7591)

Browse Source 
Co-authored-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
Bowen Liang
2024-08-23 23:52:25 +08:00
committed by GitHub
parent 2da63654e5
commit b035c02f78
155 changed files with 4279 additions and 5925 deletions
Download Patch File Download Diff File Expand all files Collapse all files

132
api/tests/integration_tests/model_runtime/baichuan/test_llm.py
View File

@@ -17,111 +17,99 @@ def test_predefined_models():
assert len(model_schemas) >= 1
assert isinstance(model_schemas[0], AIModelEntity)
def test_validate_credentials_for_chat_model():
sleep(3)
model = BaichuanLarguageModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model='baichuan2-turbo',
credentials={
'api_key': 'invalid_key',
'secret_key': 'invalid_key'
}
model="baichuan2-turbo", credentials={"api_key": "invalid_key", "secret_key": "invalid_key"}
)
model.validate_credentials(
model='baichuan2-turbo',
model="baichuan2-turbo",
credentials={
'api_key': os.environ.get('BAICHUAN_API_KEY'),
'secret_key': os.environ.get('BAICHUAN_SECRET_KEY')
}
"api_key": os.environ.get("BAICHUAN_API_KEY"),
"secret_key": os.environ.get("BAICHUAN_SECRET_KEY"),
},
)
def test_invoke_model():
sleep(3)
model = BaichuanLarguageModel()
response = model.invoke(
model='baichuan2-turbo',
model="baichuan2-turbo",
credentials={
'api_key': os.environ.get('BAICHUAN_API_KEY'),
'secret_key': os.environ.get('BAICHUAN_SECRET_KEY')
"api_key": os.environ.get("BAICHUAN_API_KEY"),
"secret_key": os.environ.get("BAICHUAN_SECRET_KEY"),
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
prompt_messages=[UserPromptMessage(content="Hello World!")],
model_parameters={
'temperature': 0.7,
'top_p': 1.0,
'top_k': 1,
"temperature": 0.7,
"top_p": 1.0,
"top_k": 1,
},
stop=['you'],
stop=["you"],
user="abc-123",
stream=False
stream=False,
)
assert isinstance(response, LLMResult)
assert len(response.message.content) > 0
assert response.usage.total_tokens > 0
def test_invoke_model_with_system_message():
sleep(3)
model = BaichuanLarguageModel()
response = model.invoke(
model='baichuan2-turbo',
model="baichuan2-turbo",
credentials={
'api_key': os.environ.get('BAICHUAN_API_KEY'),
'secret_key': os.environ.get('BAICHUAN_SECRET_KEY')
"api_key": os.environ.get("BAICHUAN_API_KEY"),
"secret_key": os.environ.get("BAICHUAN_SECRET_KEY"),
},
prompt_messages=[
SystemPromptMessage(
content='请记住你是Kasumi。'
),
UserPromptMessage(
content='现在告诉我你是谁？'
)
SystemPromptMessage(content="请记住你是Kasumi。"),
UserPromptMessage(content="现在告诉我你是谁？"),
],
model_parameters={
'temperature': 0.7,
'top_p': 1.0,
'top_k': 1,
"temperature": 0.7,
"top_p": 1.0,
"top_k": 1,
},
stop=['you'],
stop=["you"],
user="abc-123",
stream=False
stream=False,
)
assert isinstance(response, LLMResult)
assert len(response.message.content) > 0
assert response.usage.total_tokens > 0
def test_invoke_stream_model():
sleep(3)
model = BaichuanLarguageModel()
response = model.invoke(
model='baichuan2-turbo',
model="baichuan2-turbo",
credentials={
'api_key': os.environ.get('BAICHUAN_API_KEY'),
'secret_key': os.environ.get('BAICHUAN_SECRET_KEY')
"api_key": os.environ.get("BAICHUAN_API_KEY"),
"secret_key": os.environ.get("BAICHUAN_SECRET_KEY"),
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
prompt_messages=[UserPromptMessage(content="Hello World!")],
model_parameters={
'temperature': 0.7,
'top_p': 1.0,
'top_k': 1,
"temperature": 0.7,
"top_p": 1.0,
"top_k": 1,
},
stop=['you'],
stop=["you"],
stream=True,
user="abc-123"
user="abc-123",
)
assert isinstance(response, Generator)
@@ -131,34 +119,31 @@ def test_invoke_stream_model():
assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
def test_invoke_with_search():
sleep(3)
model = BaichuanLarguageModel()
response = model.invoke(
model='baichuan2-turbo',
model="baichuan2-turbo",
credentials={
'api_key': os.environ.get('BAICHUAN_API_KEY'),
'secret_key': os.environ.get('BAICHUAN_SECRET_KEY')
"api_key": os.environ.get("BAICHUAN_API_KEY"),
"secret_key": os.environ.get("BAICHUAN_SECRET_KEY"),
},
prompt_messages=[
UserPromptMessage(
content='北京今天的天气怎么样'
)
],
prompt_messages=[UserPromptMessage(content="北京今天的天气怎么样")],
model_parameters={
'temperature': 0.7,
'top_p': 1.0,
'top_k': 1,
'with_search_enhance': True,
"temperature": 0.7,
"top_p": 1.0,
"top_k": 1,
"with_search_enhance": True,
},
stop=['you'],
stop=["you"],
stream=True,
user="abc-123"
user="abc-123",
)
assert isinstance(response, Generator)
total_message = ''
total_message = ""
for chunk in response:
assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta)
@@ -166,25 +151,22 @@ def test_invoke_with_search():
assert len(chunk.delta.message.content) > 0 if not chunk.delta.finish_reason else True
total_message += chunk.delta.message.content
assert '不' not in total_message
assert "不" not in total_message
def test_get_num_tokens():
sleep(3)
model = BaichuanLarguageModel()
response = model.get_num_tokens(
model='baichuan2-turbo',
model="baichuan2-turbo",
credentials={
'api_key': os.environ.get('BAICHUAN_API_KEY'),
'secret_key': os.environ.get('BAICHUAN_SECRET_KEY')
"api_key": os.environ.get("BAICHUAN_API_KEY"),
"secret_key": os.environ.get("BAICHUAN_SECRET_KEY"),
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
tools=[]
prompt_messages=[UserPromptMessage(content="Hello World!")],
tools=[],
)
assert isinstance(response, int)
assert response == 9
assert response == 9