feat: openai_api_compatible support config stream_mode_delimiter (#2190)

Co-authored-by: wanggang <wanggy01@servyou.com.cn>
Co-authored-by: Chenhe Gu <guchenhe@gmail.com>
This commit is contained in:
geosmart
2024-01-26 00:31:59 +08:00
committed by GitHub
parent 5fc1bd026a
commit 21450b8a51
3 changed files with 87 additions and 26 deletions

View File

@@ -12,6 +12,7 @@ from core.model_runtime.model_providers.openai_api_compatible.llm.llm import OAI
Using Together.ai's OpenAI-compatible API as testing endpoint
"""
def test_validate_credentials():
model = OAIAPICompatLargeLanguageModel()
@@ -34,6 +35,7 @@ def test_validate_credentials():
}
)
def test_invoke_model():
model = OAIAPICompatLargeLanguageModel()
@@ -65,9 +67,47 @@ def test_invoke_model():
assert isinstance(response, LLMResult)
assert len(response.message.content) > 0
def test_invoke_stream_model():
model = OAIAPICompatLargeLanguageModel()
response = model.invoke(
model='mistralai/Mixtral-8x7B-Instruct-v0.1',
credentials={
'api_key': os.environ.get('TOGETHER_API_KEY'),
'endpoint_url': 'https://api.together.xyz/v1/',
'mode': 'chat',
'stream_mode_delimiter': '\\n\\n'
},
prompt_messages=[
SystemPromptMessage(
content='You are a helpful AI assistant.',
),
UserPromptMessage(
content='Who are you?'
)
],
model_parameters={
'temperature': 1.0,
'top_k': 2,
'top_p': 0.5,
},
stop=['How'],
stream=True,
user="abc-123"
)
assert isinstance(response, Generator)
for chunk in response:
assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage)
def test_invoke_stream_model_without_delimiter():
model = OAIAPICompatLargeLanguageModel()
response = model.invoke(
model='mistralai/Mixtral-8x7B-Instruct-v0.1',
credentials={
@@ -100,6 +140,7 @@ def test_invoke_stream_model():
assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage)
# using OpenAI's ChatGPT-3.5 as testing endpoint
def test_invoke_chat_model_with_tools():
model = OAIAPICompatLargeLanguageModel()
@@ -126,22 +167,22 @@ def test_invoke_chat_model_with_tools():
parameters={
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state e.g. San Francisco, CA"
},
"unit": {
"type": "string",
"enum": [
"celsius",
"fahrenheit"
]
}
"location": {
"type": "string",
"description": "The city and state e.g. San Francisco, CA"
},
"unit": {
"type": "string",
"enum": [
"celsius",
"fahrenheit"
]
}
},
"required": [
"location"
"location"
]
}
}
),
],
model_parameters={
@@ -156,6 +197,7 @@ def test_invoke_chat_model_with_tools():
assert isinstance(result.message, AssistantPromptMessage)
assert len(result.message.tool_calls) > 0
def test_get_num_tokens():
model = OAIAPICompatLargeLanguageModel()