jyapayne/dify
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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
commit b035c02f78
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
155 changed files with 4279 additions and 5925 deletions
Download patch file Download diff file Expand all files Collapse all files

1
api/pyproject.toml
View file

@ -76,7 +76,6 @@ exclude = [
"migrations/**/*", "migrations/**/*",
"services/**/*.py", "services/**/*.py",
"tasks/**/*.py", "tasks/**/*.py",
"tests/**/*.py",
] ]
[tool.pytest_env] [tool.pytest_env]

56
api/tests/integration_tests/model_runtime/__mock/anthropic.py
View file

@ -22,23 +22,20 @@ from anthropic.types import (
) )
from anthropic.types.message_delta_event import Delta from anthropic.types.message_delta_event import Delta
MOCK = os.getenv('MOCK_SWITCH', 'false') == 'true' MOCK = os.getenv("MOCK_SWITCH", "false") == "true"
class MockAnthropicClass: class MockAnthropicClass:
@staticmethod @staticmethod
def mocked_anthropic_chat_create_sync(model: str) -> Message: def mocked_anthropic_chat_create_sync(model: str) -> Message:
return Message( return Message(
id='msg-123', id="msg-123",
type='message', type="message",
role='assistant', role="assistant",
content=[ContentBlock(text='hello, I\'m a chatbot from anthropic', type='text')], content=[ContentBlock(text="hello, I'm a chatbot from anthropic", type="text")],
model=model, model=model,
stop_reason='stop_sequence', stop_reason="stop_sequence",
usage=Usage( usage=Usage(input_tokens=1, output_tokens=1),
input_tokens=1,
output_tokens=1
)
) )
@staticmethod @staticmethod
@ -46,52 +43,43 @@ class MockAnthropicClass:
full_response_text = "hello, I'm a chatbot from anthropic" full_response_text = "hello, I'm a chatbot from anthropic"
yield MessageStartEvent( yield MessageStartEvent(
type='message_start', type="message_start",
message=Message( message=Message(
id='msg-123', id="msg-123",
content=[], content=[],
role='assistant', role="assistant",
model=model, model=model,
stop_reason=None, stop_reason=None,
type='message', type="message",
usage=Usage( usage=Usage(input_tokens=1, output_tokens=1),
input_tokens=1, ),
output_tokens=1
)
)
) )
index = 0 index = 0
for i in range(0, len(full_response_text)): for i in range(0, len(full_response_text)):
yield ContentBlockDeltaEvent( yield ContentBlockDeltaEvent(
type='content_block_delta', type="content_block_delta", delta=TextDelta(text=full_response_text[i], type="text_delta"), index=index
delta=TextDelta(text=full_response_text[i], type='text_delta'),
index=index
) )
index += 1 index += 1
yield MessageDeltaEvent( yield MessageDeltaEvent(
type='message_delta', type="message_delta", delta=Delta(stop_reason="stop_sequence"), usage=MessageDeltaUsage(output_tokens=1)
delta=Delta(
stop_reason='stop_sequence'
),
usage=MessageDeltaUsage(
output_tokens=1
)
) )
yield MessageStopEvent(type='message_stop') yield MessageStopEvent(type="message_stop")
def mocked_anthropic(self: Messages, *, def mocked_anthropic(
self: Messages,
*,
max_tokens: int, max_tokens: int,
messages: Iterable[MessageParam], messages: Iterable[MessageParam],
model: str, model: str,
stream: Literal[True], stream: Literal[True],
**kwargs: Any **kwargs: Any,
) -> Union[Message, Stream[MessageStreamEvent]]: ) -> Union[Message, Stream[MessageStreamEvent]]:
if len(self._client.api_key) < 18: if len(self._client.api_key) < 18:
raise anthropic.AuthenticationError('Invalid API key') raise anthropic.AuthenticationError("Invalid API key")
if stream: if stream:
return MockAnthropicClass.mocked_anthropic_chat_create_stream(model=model) return MockAnthropicClass.mocked_anthropic_chat_create_stream(model=model)
@ -102,7 +90,7 @@ class MockAnthropicClass:
@pytest.fixture @pytest.fixture
def setup_anthropic_mock(request, monkeypatch: MonkeyPatch): def setup_anthropic_mock(request, monkeypatch: MonkeyPatch):
if MOCK: if MOCK:
monkeypatch.setattr(Messages, 'create', MockAnthropicClass.mocked_anthropic) monkeypatch.setattr(Messages, "create", MockAnthropicClass.mocked_anthropic)
yield yield

46
api/tests/integration_tests/model_runtime/__mock/google.py
View file

@ -12,63 +12,46 @@ from google.generativeai.client import _ClientManager, configure
from google.generativeai.types import GenerateContentResponse from google.generativeai.types import GenerateContentResponse
from google.generativeai.types.generation_types import BaseGenerateContentResponse from google.generativeai.types.generation_types import BaseGenerateContentResponse
current_api_key = '' current_api_key = ""
class MockGoogleResponseClass: class MockGoogleResponseClass:
_done = False _done = False
def __iter__(self): def __iter__(self):
full_response_text = 'it\'s google!' full_response_text = "it's google!"
for i in range(0, len(full_response_text) + 1, 1): for i in range(0, len(full_response_text) + 1, 1):
if i == len(full_response_text): if i == len(full_response_text):
self._done = True self._done = True
yield GenerateContentResponse( yield GenerateContentResponse(
done=True, done=True, iterator=None, result=glm.GenerateContentResponse({}), chunks=[]
iterator=None,
result=glm.GenerateContentResponse({
}),
chunks=[]
) )
else: else:
yield GenerateContentResponse( yield GenerateContentResponse(
done=False, done=False, iterator=None, result=glm.GenerateContentResponse({}), chunks=[]
iterator=None,
result=glm.GenerateContentResponse({
}),
chunks=[]
) )
class MockGoogleResponseCandidateClass: class MockGoogleResponseCandidateClass:
finish_reason = 'stop' finish_reason = "stop"
@property @property
def content(self) -> gag_content.Content: def content(self) -> gag_content.Content:
return gag_content.Content( return gag_content.Content(parts=[gag_content.Part(text="it's google!")])
parts=[
gag_content.Part(text='it\'s google!')
]
)
class MockGoogleClass: class MockGoogleClass:
@staticmethod @staticmethod
def generate_content_sync() -> GenerateContentResponse: def generate_content_sync() -> GenerateContentResponse:
return GenerateContentResponse( return GenerateContentResponse(done=True, iterator=None, result=glm.GenerateContentResponse({}), chunks=[])
done=True,
iterator=None,
result=glm.GenerateContentResponse({
}),
chunks=[]
)
@staticmethod @staticmethod
def generate_content_stream() -> Generator[GenerateContentResponse, None, None]: def generate_content_stream() -> Generator[GenerateContentResponse, None, None]:
return MockGoogleResponseClass() return MockGoogleResponseClass()
def generate_content(self: GenerativeModel, def generate_content(
self: GenerativeModel,
contents: content_types.ContentsType, contents: content_types.ContentsType,
*, *,
generation_config: generation_config_types.GenerationConfigType | None = None, generation_config: generation_config_types.GenerationConfigType | None = None,
@ -79,7 +62,7 @@ class MockGoogleClass:
global current_api_key global current_api_key
if len(current_api_key) < 16: if len(current_api_key) < 16:
raise Exception('Invalid API key') raise Exception("Invalid API key")
if stream: if stream:
return MockGoogleClass.generate_content_stream() return MockGoogleClass.generate_content_stream()
@ -88,7 +71,7 @@ class MockGoogleClass:
@property @property
def generative_response_text(self) -> str: def generative_response_text(self) -> str:
return 'it\'s google!' return "it's google!"
@property @property
def generative_response_candidates(self) -> list[MockGoogleResponseCandidateClass]: def generative_response_candidates(self) -> list[MockGoogleResponseCandidateClass]:
@ -122,6 +105,7 @@ class MockGoogleClass:
if not self.default_metadata: if not self.default_metadata:
return client return client
@pytest.fixture @pytest.fixture
def setup_google_mock(request, monkeypatch: MonkeyPatch): def setup_google_mock(request, monkeypatch: MonkeyPatch):
monkeypatch.setattr(BaseGenerateContentResponse, "text", MockGoogleClass.generative_response_text) monkeypatch.setattr(BaseGenerateContentResponse, "text", MockGoogleClass.generative_response_text)

3
api/tests/integration_tests/model_runtime/__mock/huggingface.py
View file

@ -6,7 +6,8 @@ from huggingface_hub import InferenceClient
from tests.integration_tests.model_runtime.__mock.huggingface_chat import MockHuggingfaceChatClass from tests.integration_tests.model_runtime.__mock.huggingface_chat import MockHuggingfaceChatClass
MOCK = os.getenv('MOCK_SWITCH', 'false').lower() == 'true' MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture @pytest.fixture
def setup_huggingface_mock(request, monkeypatch: MonkeyPatch): def setup_huggingface_mock(request, monkeypatch: MonkeyPatch):

23
api/tests/integration_tests/model_runtime/__mock/huggingface_chat.py
View file

@ -22,10 +22,8 @@ class MockHuggingfaceChatClass:
details=Details( details=Details(
finish_reason="length", finish_reason="length",
generated_tokens=6, generated_tokens=6,
tokens=[ tokens=[Token(id=0, text="You", logprob=0.0, special=False) for i in range(0, 6)],
Token(id=0, text="You", logprob=0.0, special=False) for i in range(0, 6) ),
]
)
) )
return response return response
@ -36,26 +34,23 @@ class MockHuggingfaceChatClass:
for i in range(0, len(full_text)): for i in range(0, len(full_text)):
response = TextGenerationStreamResponse( response = TextGenerationStreamResponse(
token = Token(id=i, text=full_text[i], logprob=0.0, special=False), token=Token(id=i, text=full_text[i], logprob=0.0, special=False),
) )
response.generated_text = full_text[i] response.generated_text = full_text[i]
response.details = StreamDetails(finish_reason='stop_sequence', generated_tokens=1) response.details = StreamDetails(finish_reason="stop_sequence", generated_tokens=1)
yield response yield response
def text_generation(self: InferenceClient, prompt: str, *, def text_generation(
stream: Literal[False] = ..., self: InferenceClient, prompt: str, *, stream: Literal[False] = ..., model: Optional[str] = None, **kwargs: Any
model: Optional[str] = None,
**kwargs: Any
) -> Union[TextGenerationResponse, Generator[TextGenerationStreamResponse, None, None]]: ) -> Union[TextGenerationResponse, Generator[TextGenerationStreamResponse, None, None]]:
# check if key is valid # check if key is valid
if not re.match(r'Bearer\shf\-[a-zA-Z0-9]{16,}', self.headers['authorization']): if not re.match(r"Bearer\shf\-[a-zA-Z0-9]{16,}", self.headers["authorization"]):
raise BadRequestError('Invalid API key') raise BadRequestError("Invalid API key")
if model is None: if model is None:
raise BadRequestError('Invalid model') raise BadRequestError("Invalid model")
if stream: if stream:
return MockHuggingfaceChatClass.generate_create_stream(model) return MockHuggingfaceChatClass.generate_create_stream(model)
return MockHuggingfaceChatClass.generate_create_sync(model) return MockHuggingfaceChatClass.generate_create_sync(model)

42
api/tests/integration_tests/model_runtime/__mock/huggingface_tei.py
View file

@ -5,10 +5,10 @@ class MockTEIClass:
@staticmethod @staticmethod
def get_tei_extra_parameter(server_url: str, model_name: str) -> TeiModelExtraParameter: def get_tei_extra_parameter(server_url: str, model_name: str) -> TeiModelExtraParameter:
# During mock, we don't have a real server to query, so we just return a dummy value # During mock, we don't have a real server to query, so we just return a dummy value
if 'rerank' in model_name: if "rerank" in model_name:
model_type = 'reranker' model_type = "reranker"
else: else:
model_type = 'embedding' model_type = "embedding"
return TeiModelExtraParameter(model_type=model_type, max_input_length=512, max_client_batch_size=1) return TeiModelExtraParameter(model_type=model_type, max_input_length=512, max_client_batch_size=1)
@ -17,16 +17,16 @@ class MockTEIClass:
# Use space as token separator, and split the text into tokens # Use space as token separator, and split the text into tokens
tokenized_texts = [] tokenized_texts = []
for text in texts: for text in texts:
tokens = text.split(' ') tokens = text.split(" ")
current_index = 0 current_index = 0
tokenized_text = [] tokenized_text = []
for idx, token in enumerate(tokens): for idx, token in enumerate(tokens):
s_token = { s_token = {
'id': idx, "id": idx,
'text': token, "text": token,
'special': False, "special": False,
'start': current_index, "start": current_index,
'stop': current_index + len(token), "stop": current_index + len(token),
} }
current_index += len(token) + 1 current_index += len(token) + 1
tokenized_text.append(s_token) tokenized_text.append(s_token)
@ -55,18 +55,18 @@ class MockTEIClass:
embedding = [0.1] * 768 embedding = [0.1] * 768
embeddings.append( embeddings.append(
{ {
'object': 'embedding', "object": "embedding",
'embedding': embedding, "embedding": embedding,
'index': idx, "index": idx,
} }
) )
return { return {
'object': 'list', "object": "list",
'data': embeddings, "data": embeddings,
'model': 'MODEL_NAME', "model": "MODEL_NAME",
'usage': { "usage": {
'prompt_tokens': sum(len(text.split(' ')) for text in texts), "prompt_tokens": sum(len(text.split(" ")) for text in texts),
'total_tokens': sum(len(text.split(' ')) for text in texts), "total_tokens": sum(len(text.split(" ")) for text in texts),
}, },
} }
@ -83,9 +83,9 @@ class MockTEIClass:
for idx, text in enumerate(texts): for idx, text in enumerate(texts):
reranked_docs.append( reranked_docs.append(
{ {
'index': idx, "index": idx,
'text': text, "text": text,
'score': 0.9, "score": 0.9,
} }
) )
# For mock, only return the first document # For mock, only return the first document

11
api/tests/integration_tests/model_runtime/__mock/openai.py
View file

@ -21,13 +21,17 @@ from tests.integration_tests.model_runtime.__mock.openai_remote import MockModel
from tests.integration_tests.model_runtime.__mock.openai_speech2text import MockSpeech2TextClass from tests.integration_tests.model_runtime.__mock.openai_speech2text import MockSpeech2TextClass
def mock_openai(monkeypatch: MonkeyPatch, methods: list[Literal["completion", "chat", "remote", "moderation", "speech2text", "text_embedding"]]) -> Callable[[], None]: def mock_openai(
monkeypatch: MonkeyPatch,
methods: list[Literal["completion", "chat", "remote", "moderation", "speech2text", "text_embedding"]],
) -> Callable[[], None]:
""" """
mock openai module mock openai module
:param monkeypatch: pytest monkeypatch fixture :param monkeypatch: pytest monkeypatch fixture
:return: unpatch function :return: unpatch function
""" """
def unpatch() -> None: def unpatch() -> None:
monkeypatch.undo() monkeypatch.undo()
@ -52,11 +56,12 @@ def mock_openai(monkeypatch: MonkeyPatch, methods: list[Literal["completion", "c
return unpatch return unpatch
MOCK = os.getenv('MOCK_SWITCH', 'false').lower() == 'true' MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture @pytest.fixture
def setup_openai_mock(request, monkeypatch): def setup_openai_mock(request, monkeypatch):
methods = request.param if hasattr(request, 'param') else [] methods = request.param if hasattr(request, "param") else []
if MOCK: if MOCK:
unpatch = mock_openai(monkeypatch, methods=methods) unpatch = mock_openai(monkeypatch, methods=methods)

158
api/tests/integration_tests/model_runtime/__mock/openai_chat.py
View file

@ -43,59 +43,61 @@ class MockChatClass:
if not functions or len(functions) == 0: if not functions or len(functions) == 0:
return None return None
function: completion_create_params.Function = functions[0] function: completion_create_params.Function = functions[0]
function_name = function['name'] function_name = function["name"]
function_description = function['description'] function_description = function["description"]
function_parameters = function['parameters'] function_parameters = function["parameters"]
function_parameters_type = function_parameters['type'] function_parameters_type = function_parameters["type"]
if function_parameters_type != 'object': if function_parameters_type != "object":
return None return None
function_parameters_properties = function_parameters['properties'] function_parameters_properties = function_parameters["properties"]
function_parameters_required = function_parameters['required'] function_parameters_required = function_parameters["required"]
parameters = {} parameters = {}
for parameter_name, parameter in function_parameters_properties.items(): for parameter_name, parameter in function_parameters_properties.items():
if parameter_name not in function_parameters_required: if parameter_name not in function_parameters_required:
continue continue
parameter_type = parameter['type'] parameter_type = parameter["type"]
if parameter_type == 'string': if parameter_type == "string":
if 'enum' in parameter: if "enum" in parameter:
if len(parameter['enum']) == 0: if len(parameter["enum"]) == 0:
continue continue
parameters[parameter_name] = parameter['enum'][0] parameters[parameter_name] = parameter["enum"][0]
else: else:
parameters[parameter_name] = 'kawaii' parameters[parameter_name] = "kawaii"
elif parameter_type == 'integer': elif parameter_type == "integer":
parameters[parameter_name] = 114514 parameters[parameter_name] = 114514
elif parameter_type == 'number': elif parameter_type == "number":
parameters[parameter_name] = 1919810.0 parameters[parameter_name] = 1919810.0
elif parameter_type == 'boolean': elif parameter_type == "boolean":
parameters[parameter_name] = True parameters[parameter_name] = True
return FunctionCall(name=function_name, arguments=dumps(parameters)) return FunctionCall(name=function_name, arguments=dumps(parameters))
@staticmethod @staticmethod
def generate_tool_calls(tools = NOT_GIVEN) -> Optional[list[ChatCompletionMessageToolCall]]: def generate_tool_calls(tools=NOT_GIVEN) -> Optional[list[ChatCompletionMessageToolCall]]:
list_tool_calls = [] list_tool_calls = []
if not tools or len(tools) == 0: if not tools or len(tools) == 0:
return None return None
tool = tools[0] tool = tools[0]
if 'type' in tools and tools['type'] != 'function': if "type" in tools and tools["type"] != "function":
return None return None
function = tool['function'] function = tool["function"]
function_call = MockChatClass.generate_function_call(functions=[function]) function_call = MockChatClass.generate_function_call(functions=[function])
if function_call is None: if function_call is None:
return None return None
list_tool_calls.append(ChatCompletionMessageToolCall( list_tool_calls.append(
id='sakurajima-mai', ChatCompletionMessageToolCall(
id="sakurajima-mai",
function=Function( function=Function(
name=function_call.name, name=function_call.name,
arguments=function_call.arguments, arguments=function_call.arguments,
), ),
type='function' type="function",
)) )
)
return list_tool_calls return list_tool_calls
@ -111,28 +113,25 @@ class MockChatClass:
tool_calls = MockChatClass.generate_tool_calls(tools=tools) tool_calls = MockChatClass.generate_tool_calls(tools=tools)
return _ChatCompletion( return _ChatCompletion(
id='cmpl-3QJQa5jXJ5Z5X', id="cmpl-3QJQa5jXJ5Z5X",
choices=[ choices=[
_ChatCompletionChoice( _ChatCompletionChoice(
finish_reason='content_filter', finish_reason="content_filter",
index=0, index=0,
message=ChatCompletionMessage( message=ChatCompletionMessage(
content='elaina', content="elaina", role="assistant", function_call=function_call, tool_calls=tool_calls
role='assistant', ),
function_call=function_call,
tool_calls=tool_calls
)
) )
], ],
created=int(time()), created=int(time()),
model=model, model=model,
object='chat.completion', object="chat.completion",
system_fingerprint='', system_fingerprint="",
usage=CompletionUsage( usage=CompletionUsage(
prompt_tokens=2, prompt_tokens=2,
completion_tokens=1, completion_tokens=1,
total_tokens=3, total_tokens=3,
) ),
) )
@staticmethod @staticmethod
@ -150,36 +149,40 @@ class MockChatClass:
for i in range(0, len(full_text) + 1): for i in range(0, len(full_text) + 1):
if i == len(full_text): if i == len(full_text):
yield ChatCompletionChunk( yield ChatCompletionChunk(
id='cmpl-3QJQa5jXJ5Z5X', id="cmpl-3QJQa5jXJ5Z5X",
choices=[ choices=[
Choice( Choice(
delta=ChoiceDelta( delta=ChoiceDelta(
content='', content="",
function_call=ChoiceDeltaFunctionCall( function_call=ChoiceDeltaFunctionCall(
name=function_call.name, name=function_call.name,
arguments=function_call.arguments, arguments=function_call.arguments,
) if function_call else None, )
role='assistant', if function_call
else None,
role="assistant",
tool_calls=[ tool_calls=[
ChoiceDeltaToolCall( ChoiceDeltaToolCall(
index=0, index=0,
id='misaka-mikoto', id="misaka-mikoto",
function=ChoiceDeltaToolCallFunction( function=ChoiceDeltaToolCallFunction(
name=tool_calls[0].function.name, name=tool_calls[0].function.name,
arguments=tool_calls[0].function.arguments, arguments=tool_calls[0].function.arguments,
), ),
type='function' type="function",
) )
] if tool_calls and len(tool_calls) > 0 else None ]
if tool_calls and len(tool_calls) > 0
else None,
), ),
finish_reason='function_call', finish_reason="function_call",
index=0, index=0,
) )
], ],
created=int(time()), created=int(time()),
model=model, model=model,
object='chat.completion.chunk', object="chat.completion.chunk",
system_fingerprint='', system_fingerprint="",
usage=CompletionUsage( usage=CompletionUsage(
prompt_tokens=2, prompt_tokens=2,
completion_tokens=17, completion_tokens=17,
@ -188,30 +191,45 @@ class MockChatClass:
) )
else: else:
yield ChatCompletionChunk( yield ChatCompletionChunk(
id='cmpl-3QJQa5jXJ5Z5X', id="cmpl-3QJQa5jXJ5Z5X",
choices=[ choices=[
Choice( Choice(
delta=ChoiceDelta( delta=ChoiceDelta(
content=full_text[i], content=full_text[i],
role='assistant', role="assistant",
), ),
finish_reason='content_filter', finish_reason="content_filter",
index=0, index=0,
) )
], ],
created=int(time()), created=int(time()),
model=model, model=model,
object='chat.completion.chunk', object="chat.completion.chunk",
system_fingerprint='', system_fingerprint="",
) )
def chat_create(self: Completions, *, def chat_create(
self: Completions,
*,
messages: list[ChatCompletionMessageParam], messages: list[ChatCompletionMessageParam],
model: Union[str,Literal[ model: Union[
"gpt-4-1106-preview", "gpt-4-vision-preview", "gpt-4", "gpt-4-0314", "gpt-4-0613", str,
"gpt-4-32k", "gpt-4-32k-0314", "gpt-4-32k-0613", Literal[
"gpt-3.5-turbo-1106", "gpt-3.5-turbo", "gpt-3.5-turbo-16k", "gpt-3.5-turbo-0301", "gpt-4-1106-preview",
"gpt-3.5-turbo-0613", "gpt-3.5-turbo-16k-0613"], "gpt-4-vision-preview",
"gpt-4",
"gpt-4-0314",
"gpt-4-0613",
"gpt-4-32k",
"gpt-4-32k-0314",
"gpt-4-32k-0613",
"gpt-3.5-turbo-1106",
"gpt-3.5-turbo",
"gpt-3.5-turbo-16k",
"gpt-3.5-turbo-0301",
"gpt-3.5-turbo-0613",
"gpt-3.5-turbo-16k-0613",
],
], ],
functions: list[completion_create_params.Function] | NotGiven = NOT_GIVEN, functions: list[completion_create_params.Function] | NotGiven = NOT_GIVEN,
response_format: completion_create_params.ResponseFormat | NotGiven = NOT_GIVEN, response_format: completion_create_params.ResponseFormat | NotGiven = NOT_GIVEN,
@ -220,23 +238,31 @@ class MockChatClass:
**kwargs: Any, **kwargs: Any,
): ):
openai_models = [ openai_models = [
"gpt-4-1106-preview", "gpt-4-vision-preview", "gpt-4", "gpt-4-0314", "gpt-4-0613", "gpt-4-1106-preview",
"gpt-4-32k", "gpt-4-32k-0314", "gpt-4-32k-0613", "gpt-4-vision-preview",
"gpt-3.5-turbo-1106", "gpt-3.5-turbo", "gpt-3.5-turbo-16k", "gpt-3.5-turbo-0301", "gpt-4",
"gpt-3.5-turbo-0613", "gpt-3.5-turbo-16k-0613", "gpt-4-0314",
"gpt-4-0613",
"gpt-4-32k",
"gpt-4-32k-0314",
"gpt-4-32k-0613",
"gpt-3.5-turbo-1106",
"gpt-3.5-turbo",
"gpt-3.5-turbo-16k",
"gpt-3.5-turbo-0301",
"gpt-3.5-turbo-0613",
"gpt-3.5-turbo-16k-0613",
] ]
azure_openai_models = [ azure_openai_models = ["gpt35", "gpt-4v", "gpt-35-turbo"]
"gpt35", "gpt-4v", "gpt-35-turbo" if not re.match(r"^(https?):\/\/[^\s\/$.?#].[^\s]*$", self._client.base_url.__str__()):
] raise InvokeAuthorizationError("Invalid base url")
if not re.match(r'^(https?):\/\/[^\s\/$.?#].[^\s]*$', self._client.base_url.__str__()):
raise InvokeAuthorizationError('Invalid base url')
if model in openai_models + azure_openai_models: if model in openai_models + azure_openai_models:
if not re.match(r'sk-[a-zA-Z0-9]{24,}$', self._client.api_key) and type(self._client) == OpenAI: if not re.match(r"sk-[a-zA-Z0-9]{24,}$", self._client.api_key) and type(self._client) == OpenAI:
# sometime, provider use OpenAI compatible API will not have api key or have different api key format # sometime, provider use OpenAI compatible API will not have api key or have different api key format
# so we only check if model is in openai_models # so we only check if model is in openai_models
raise InvokeAuthorizationError('Invalid api key') raise InvokeAuthorizationError("Invalid api key")
if len(self._client.api_key) < 18 and type(self._client) == AzureOpenAI: if len(self._client.api_key) < 18 and type(self._client) == AzureOpenAI:
raise InvokeAuthorizationError('Invalid api key') raise InvokeAuthorizationError("Invalid api key")
if stream: if stream:
return MockChatClass.mocked_openai_chat_create_stream(model=model, functions=functions, tools=tools) return MockChatClass.mocked_openai_chat_create_stream(model=model, functions=functions, tools=tools)

69
api/tests/integration_tests/model_runtime/__mock/openai_completion.py
View file

@ -17,9 +17,7 @@ from core.model_runtime.errors.invoke import InvokeAuthorizationError
class MockCompletionsClass: class MockCompletionsClass:
@staticmethod @staticmethod
def mocked_openai_completion_create_sync( def mocked_openai_completion_create_sync(model: str) -> CompletionMessage:
model: str
) -> CompletionMessage:
return CompletionMessage( return CompletionMessage(
id="cmpl-3QJQa5jXJ5Z5X", id="cmpl-3QJQa5jXJ5Z5X",
object="text_completion", object="text_completion",
@ -38,13 +36,11 @@ class MockCompletionsClass:
prompt_tokens=2, prompt_tokens=2,
completion_tokens=1, completion_tokens=1,
total_tokens=3, total_tokens=3,
) ),
) )
@staticmethod @staticmethod
def mocked_openai_completion_create_stream( def mocked_openai_completion_create_stream(model: str) -> Generator[CompletionMessage, None, None]:
model: str
) -> Generator[CompletionMessage, None, None]:
full_text = "Hello, world!\n\n```python\nprint('Hello, world!')\n```" full_text = "Hello, world!\n\n```python\nprint('Hello, world!')\n```"
for i in range(0, len(full_text) + 1): for i in range(0, len(full_text) + 1):
if i == len(full_text): if i == len(full_text):
@ -76,45 +72,58 @@ class MockCompletionsClass:
model=model, model=model,
system_fingerprint="", system_fingerprint="",
choices=[ choices=[
CompletionChoice( CompletionChoice(text=full_text[i], index=0, logprobs=None, finish_reason="content_filter")
text=full_text[i],
index=0,
logprobs=None,
finish_reason="content_filter"
)
], ],
) )
def completion_create(self: Completions, *, model: Union[ def completion_create(
str, Literal["babbage-002", "davinci-002", "gpt-3.5-turbo-instruct", self: Completions,
"text-davinci-003", "text-davinci-002", "text-davinci-001", *,
"code-davinci-002", "text-curie-001", "text-babbage-001", model: Union[
"text-ada-001"], str,
Literal[
"babbage-002",
"davinci-002",
"gpt-3.5-turbo-instruct",
"text-davinci-003",
"text-davinci-002",
"text-davinci-001",
"code-davinci-002",
"text-curie-001",
"text-babbage-001",
"text-ada-001",
],
], ],
prompt: Union[str, list[str], list[int], list[list[int]], None], prompt: Union[str, list[str], list[int], list[list[int]], None],
stream: Optional[Literal[False]] | NotGiven = NOT_GIVEN, stream: Optional[Literal[False]] | NotGiven = NOT_GIVEN,
**kwargs: Any **kwargs: Any,
): ):
openai_models = [ openai_models = [
"babbage-002", "davinci-002", "gpt-3.5-turbo-instruct", "text-davinci-003", "text-davinci-002", "text-davinci-001", "babbage-002",
"code-davinci-002", "text-curie-001", "text-babbage-001", "text-ada-001", "davinci-002",
] "gpt-3.5-turbo-instruct",
azure_openai_models = [ "text-davinci-003",
"gpt-35-turbo-instruct" "text-davinci-002",
"text-davinci-001",
"code-davinci-002",
"text-curie-001",
"text-babbage-001",
"text-ada-001",
] ]
azure_openai_models = ["gpt-35-turbo-instruct"]
if not re.match(r'^(https?):\/\/[^\s\/$.?#].[^\s]*$', self._client.base_url.__str__()): if not re.match(r"^(https?):\/\/[^\s\/$.?#].[^\s]*$", self._client.base_url.__str__()):
raise InvokeAuthorizationError('Invalid base url') raise InvokeAuthorizationError("Invalid base url")
if model in openai_models + azure_openai_models: if model in openai_models + azure_openai_models:
if not re.match(r'sk-[a-zA-Z0-9]{24,}$', self._client.api_key) and type(self._client) == OpenAI: if not re.match(r"sk-[a-zA-Z0-9]{24,}$", self._client.api_key) and type(self._client) == OpenAI:
# sometime, provider use OpenAI compatible API will not have api key or have different api key format # sometime, provider use OpenAI compatible API will not have api key or have different api key format
# so we only check if model is in openai_models # so we only check if model is in openai_models
raise InvokeAuthorizationError('Invalid api key') raise InvokeAuthorizationError("Invalid api key")
if len(self._client.api_key) < 18 and type(self._client) == AzureOpenAI: if len(self._client.api_key) < 18 and type(self._client) == AzureOpenAI:
raise InvokeAuthorizationError('Invalid api key') raise InvokeAuthorizationError("Invalid api key")
if not prompt: if not prompt:
raise BadRequestError('Invalid prompt') raise BadRequestError("Invalid prompt")
if stream: if stream:
return MockCompletionsClass.mocked_openai_completion_create_stream(model=model) return MockCompletionsClass.mocked_openai_completion_create_stream(model=model)

42
api/tests/integration_tests/model_runtime/__mock/openai_embeddings.py
View file

File diff suppressed because one or more lines are too long

92
api/tests/integration_tests/model_runtime/__mock/openai_moderation.py
View file

@ -10,58 +10,92 @@ from core.model_runtime.errors.invoke import InvokeAuthorizationError
class MockModerationClass: class MockModerationClass:
def moderation_create(self: Moderations,*, def moderation_create(
self: Moderations,
*,
input: Union[str, list[str]], input: Union[str, list[str]],
model: Union[str, Literal["text-moderation-latest", "text-moderation-stable"]] | NotGiven = NOT_GIVEN, model: Union[str, Literal["text-moderation-latest", "text-moderation-stable"]] | NotGiven = NOT_GIVEN,
**kwargs: Any **kwargs: Any,
) -> ModerationCreateResponse: ) -> ModerationCreateResponse:
if isinstance(input, str): if isinstance(input, str):
input = [input] input = [input]
if not re.match(r'^(https?):\/\/[^\s\/$.?#].[^\s]*$', self._client.base_url.__str__()): if not re.match(r"^(https?):\/\/[^\s\/$.?#].[^\s]*$", self._client.base_url.__str__()):
raise InvokeAuthorizationError('Invalid base url') raise InvokeAuthorizationError("Invalid base url")
if len(self._client.api_key) < 18: if len(self._client.api_key) < 18:
raise InvokeAuthorizationError('Invalid API key') raise InvokeAuthorizationError("Invalid API key")
for text in input: for text in input:
result = [] result = []
if 'kill' in text: if "kill" in text:
moderation_categories = { moderation_categories = {
'harassment': False, 'harassment/threatening': False, 'hate': False, 'hate/threatening': False, "harassment": False,
'self-harm': False, 'self-harm/instructions': False, 'self-harm/intent': False, 'sexual': False, "harassment/threatening": False,
'sexual/minors': False, 'violence': False, 'violence/graphic': False "hate": False,
"hate/threatening": False,
"self-harm": False,
"self-harm/instructions": False,
"self-harm/intent": False,
"sexual": False,
"sexual/minors": False,
"violence": False,
"violence/graphic": False,
} }
moderation_categories_scores = { moderation_categories_scores = {
'harassment': 1.0, 'harassment/threatening': 1.0, 'hate': 1.0, 'hate/threatening': 1.0, "harassment": 1.0,
'self-harm': 1.0, 'self-harm/instructions': 1.0, 'self-harm/intent': 1.0, 'sexual': 1.0, "harassment/threatening": 1.0,
'sexual/minors': 1.0, 'violence': 1.0, 'violence/graphic': 1.0 "hate": 1.0,
"hate/threatening": 1.0,
"self-harm": 1.0,
"self-harm/instructions": 1.0,
"self-harm/intent": 1.0,
"sexual": 1.0,
"sexual/minors": 1.0,
"violence": 1.0,
"violence/graphic": 1.0,
} }
result.append(Moderation( result.append(
Moderation(
flagged=True, flagged=True,
categories=Categories(**moderation_categories), categories=Categories(**moderation_categories),
category_scores=CategoryScores(**moderation_categories_scores) category_scores=CategoryScores(**moderation_categories_scores),
)) )
)
else: else:
moderation_categories = { moderation_categories = {
'harassment': False, 'harassment/threatening': False, 'hate': False, 'hate/threatening': False, "harassment": False,
'self-harm': False, 'self-harm/instructions': False, 'self-harm/intent': False, 'sexual': False, "harassment/threatening": False,
'sexual/minors': False, 'violence': False, 'violence/graphic': False "hate": False,
"hate/threatening": False,
"self-harm": False,
"self-harm/instructions": False,
"self-harm/intent": False,
"sexual": False,
"sexual/minors": False,
"violence": False,
"violence/graphic": False,
} }
moderation_categories_scores = { moderation_categories_scores = {
'harassment': 0.0, 'harassment/threatening': 0.0, 'hate': 0.0, 'hate/threatening': 0.0, "harassment": 0.0,
'self-harm': 0.0, 'self-harm/instructions': 0.0, 'self-harm/intent': 0.0, 'sexual': 0.0, "harassment/threatening": 0.0,
'sexual/minors': 0.0, 'violence': 0.0, 'violence/graphic': 0.0 "hate": 0.0,
"hate/threatening": 0.0,
"self-harm": 0.0,
"self-harm/instructions": 0.0,
"self-harm/intent": 0.0,
"sexual": 0.0,
"sexual/minors": 0.0,
"violence": 0.0,
"violence/graphic": 0.0,
} }
result.append(Moderation( result.append(
Moderation(
flagged=False, flagged=False,
categories=Categories(**moderation_categories), categories=Categories(**moderation_categories),
category_scores=CategoryScores(**moderation_categories_scores) category_scores=CategoryScores(**moderation_categories_scores),
))
return ModerationCreateResponse(
id='shiroii kuloko',
model=model,
results=result
) )
)
return ModerationCreateResponse(id="shiroii kuloko", model=model, results=result)

7
api/tests/integration_tests/model_runtime/__mock/openai_remote.py
View file

@ -8,15 +8,16 @@ class MockModelClass:
""" """
mock class for openai.models.Models mock class for openai.models.Models
""" """
def list( def list(
self, self,
**kwargs, **kwargs,
) -> list[Model]: ) -> list[Model]:
return [ return [
Model( Model(
id='ft:gpt-3.5-turbo-0613:personal::8GYJLPDQ', id="ft:gpt-3.5-turbo-0613:personal::8GYJLPDQ",
created=int(time()), created=int(time()),
object='model', object="model",
owned_by='organization:org-123', owned_by="organization:org-123",
) )
] ]

15
api/tests/integration_tests/model_runtime/__mock/openai_speech2text.py
View file

@ -9,7 +9,8 @@ from core.model_runtime.errors.invoke import InvokeAuthorizationError
class MockSpeech2TextClass: class MockSpeech2TextClass:
def speech2text_create(self: Transcriptions, def speech2text_create(
self: Transcriptions,
*, *,
file: FileTypes, file: FileTypes,
model: Union[str, Literal["whisper-1"]], model: Union[str, Literal["whisper-1"]],
@ -17,14 +18,12 @@ class MockSpeech2TextClass:
prompt: str | NotGiven = NOT_GIVEN, prompt: str | NotGiven = NOT_GIVEN,
response_format: Literal["json", "text", "srt", "verbose_json", "vtt"] | NotGiven = NOT_GIVEN, response_format: Literal["json", "text", "srt", "verbose_json", "vtt"] | NotGiven = NOT_GIVEN,
temperature: float | NotGiven = NOT_GIVEN, temperature: float | NotGiven = NOT_GIVEN,
**kwargs: Any **kwargs: Any,
) -> Transcription: ) -> Transcription:
if not re.match(r'^(https?):\/\/[^\s\/$.?#].[^\s]*$', self._client.base_url.__str__()): if not re.match(r"^(https?):\/\/[^\s\/$.?#].[^\s]*$", self._client.base_url.__str__()):
raise InvokeAuthorizationError('Invalid base url') raise InvokeAuthorizationError("Invalid base url")
if len(self._client.api_key) < 18: if len(self._client.api_key) < 18:
raise InvokeAuthorizationError('Invalid API key') raise InvokeAuthorizationError("Invalid API key")
return Transcription( return Transcription(text="1, 2, 3, 4, 5, 6, 7, 8, 9, 10")
text='1, 2, 3, 4, 5, 6, 7, 8, 9, 10'
)

115
api/tests/integration_tests/model_runtime/__mock/xinference.py
View file

@ -19,40 +19,43 @@ from xinference_client.types import Embedding, EmbeddingData, EmbeddingUsage
class MockXinferenceClass: class MockXinferenceClass:
def get_chat_model(self: Client, model_uid: str) -> Union[RESTfulChatglmCppChatModelHandle, RESTfulGenerateModelHandle, RESTfulChatModelHandle]: def get_chat_model(
if not re.match(r'https?:\/\/[^\s\/$.?#].[^\s]*$', self.base_url): self: Client, model_uid: str
raise RuntimeError('404 Not Found') ) -> Union[RESTfulChatglmCppChatModelHandle, RESTfulGenerateModelHandle, RESTfulChatModelHandle]:
if not re.match(r"https?:\/\/[^\s\/$.?#].[^\s]*$", self.base_url):
raise RuntimeError("404 Not Found")
if 'generate' == model_uid: if "generate" == model_uid:
return RESTfulGenerateModelHandle(model_uid, base_url=self.base_url, auth_headers={}) return RESTfulGenerateModelHandle(model_uid, base_url=self.base_url, auth_headers={})
if 'chat' == model_uid: if "chat" == model_uid:
return RESTfulChatModelHandle(model_uid, base_url=self.base_url, auth_headers={}) return RESTfulChatModelHandle(model_uid, base_url=self.base_url, auth_headers={})
if 'embedding' == model_uid: if "embedding" == model_uid:
return RESTfulEmbeddingModelHandle(model_uid, base_url=self.base_url, auth_headers={}) return RESTfulEmbeddingModelHandle(model_uid, base_url=self.base_url, auth_headers={})
if 'rerank' == model_uid: if "rerank" == model_uid:
return RESTfulRerankModelHandle(model_uid, base_url=self.base_url, auth_headers={}) return RESTfulRerankModelHandle(model_uid, base_url=self.base_url, auth_headers={})
raise RuntimeError('404 Not Found') raise RuntimeError("404 Not Found")
def get(self: Session, url: str, **kwargs): def get(self: Session, url: str, **kwargs):
response = Response() response = Response()
if 'v1/models/' in url: if "v1/models/" in url:
# get model uid # get model uid
model_uid = url.split('/')[-1] or '' model_uid = url.split("/")[-1] or ""
if not re.match(r'[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}', model_uid) and \ if not re.match(
model_uid not in ['generate', 'chat', 'embedding', 'rerank']: r"[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}", model_uid
) and model_uid not in ["generate", "chat", "embedding", "rerank"]:
response.status_code = 404 response.status_code = 404
response._content = b'{}' response._content = b"{}"
return response return response
# check if url is valid # check if url is valid
if not re.match(r'^(https?):\/\/[^\s\/$.?#].[^\s]*$', url): if not re.match(r"^(https?):\/\/[^\s\/$.?#].[^\s]*$", url):
response.status_code = 404 response.status_code = 404
response._content = b'{}' response._content = b"{}"
return response return response
if model_uid in ['generate', 'chat']: if model_uid in ["generate", "chat"]:
response.status_code = 200 response.status_code = 200
response._content = b'''{ response._content = b"""{
"model_type": "LLM", "model_type": "LLM",
"address": "127.0.0.1:43877", "address": "127.0.0.1:43877",
"accelerators": [ "accelerators": [
@ -75,12 +78,12 @@ class MockXinferenceClass:
"revision": null, "revision": null,
"context_length": 2048, "context_length": 2048,
"replica": 1 "replica": 1
}''' }"""
return response return response
elif model_uid == 'embedding': elif model_uid == "embedding":
response.status_code = 200 response.status_code = 200
response._content = b'''{ response._content = b"""{
"model_type": "embedding", "model_type": "embedding",
"address": "127.0.0.1:43877", "address": "127.0.0.1:43877",
"accelerators": [ "accelerators": [
@ -93,51 +96,48 @@ class MockXinferenceClass:
], ],
"revision": null, "revision": null,
"max_tokens": 512 "max_tokens": 512
}''' }"""
return response return response
elif 'v1/cluster/auth' in url: elif "v1/cluster/auth" in url:
response.status_code = 200 response.status_code = 200
response._content = b'''{ response._content = b"""{
"auth": true "auth": true
}''' }"""
return response return response
def _check_cluster_authenticated(self): def _check_cluster_authenticated(self):
self._cluster_authed = True self._cluster_authed = True
def rerank(self: RESTfulRerankModelHandle, documents: list[str], query: str, top_n: int, return_documents: bool) -> dict: def rerank(
self: RESTfulRerankModelHandle, documents: list[str], query: str, top_n: int, return_documents: bool
) -> dict:
# check if self._model_uid is a valid uuid # check if self._model_uid is a valid uuid
if not re.match(r'[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}', self._model_uid) and \ if (
self._model_uid != 'rerank': not re.match(r"[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}", self._model_uid)
raise RuntimeError('404 Not Found') and self._model_uid != "rerank"
):
raise RuntimeError("404 Not Found")
if not re.match(r'^(https?):\/\/[^\s\/$.?#].[^\s]*$', self._base_url): if not re.match(r"^(https?):\/\/[^\s\/$.?#].[^\s]*$", self._base_url):
raise RuntimeError('404 Not Found') raise RuntimeError("404 Not Found")
if top_n is None: if top_n is None:
top_n = 1 top_n = 1
return { return {
'results': [ "results": [
{ {"index": i, "document": doc, "relevance_score": 0.9} for i, doc in enumerate(documents[:top_n])
'index': i,
'document': doc,
'relevance_score': 0.9
}
for i, doc in enumerate(documents[:top_n])
] ]
} }
def create_embedding( def create_embedding(self: RESTfulGenerateModelHandle, input: Union[str, list[str]], **kwargs) -> dict:
self: RESTfulGenerateModelHandle,
input: Union[str, list[str]],
**kwargs
) -> dict:
# check if self._model_uid is a valid uuid # check if self._model_uid is a valid uuid
if not re.match(r'[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}', self._model_uid) and \ if (
self._model_uid != 'embedding': not re.match(r"[a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12}", self._model_uid)
raise RuntimeError('404 Not Found') and self._model_uid != "embedding"
):
raise RuntimeError("404 Not Found")
if isinstance(input, str): if isinstance(input, str):
input = [input] input = [input]
@ -147,31 +147,26 @@ class MockXinferenceClass:
object="list", object="list",
model=self._model_uid, model=self._model_uid,
data=[ data=[
EmbeddingData( EmbeddingData(index=i, object="embedding", embedding=[1919.810 for _ in range(768)])
index=i,
object="embedding",
embedding=[1919.810 for _ in range(768)]
)
for i in range(ipt_len) for i in range(ipt_len)
], ],
usage=EmbeddingUsage( usage=EmbeddingUsage(prompt_tokens=ipt_len, total_tokens=ipt_len),
prompt_tokens=ipt_len,
total_tokens=ipt_len
)
) )
return embedding return embedding
MOCK = os.getenv('MOCK_SWITCH', 'false').lower() == 'true'
MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture @pytest.fixture
def setup_xinference_mock(request, monkeypatch: MonkeyPatch): def setup_xinference_mock(request, monkeypatch: MonkeyPatch):
if MOCK: if MOCK:
monkeypatch.setattr(Client, 'get_model', MockXinferenceClass.get_chat_model) monkeypatch.setattr(Client, "get_model", MockXinferenceClass.get_chat_model)
monkeypatch.setattr(Client, '_check_cluster_authenticated', MockXinferenceClass._check_cluster_authenticated) monkeypatch.setattr(Client, "_check_cluster_authenticated", MockXinferenceClass._check_cluster_authenticated)
monkeypatch.setattr(Session, 'get', MockXinferenceClass.get) monkeypatch.setattr(Session, "get", MockXinferenceClass.get)
monkeypatch.setattr(RESTfulEmbeddingModelHandle, 'create_embedding', MockXinferenceClass.create_embedding) monkeypatch.setattr(RESTfulEmbeddingModelHandle, "create_embedding", MockXinferenceClass.create_embedding)
monkeypatch.setattr(RESTfulRerankModelHandle, 'rerank', MockXinferenceClass.rerank) monkeypatch.setattr(RESTfulRerankModelHandle, "rerank", MockXinferenceClass.rerank)
yield yield
if MOCK: if MOCK:

75
api/tests/integration_tests/model_runtime/anthropic/test_llm.py
View file

@ -10,79 +10,60 @@ from core.model_runtime.model_providers.anthropic.llm.llm import AnthropicLargeL
from tests.integration_tests.model_runtime.__mock.anthropic import setup_anthropic_mock from tests.integration_tests.model_runtime.__mock.anthropic import setup_anthropic_mock
@pytest.mark.parametrize('setup_anthropic_mock', [['none']], indirect=True) @pytest.mark.parametrize("setup_anthropic_mock", [["none"]], indirect=True)
def test_validate_credentials(setup_anthropic_mock): def test_validate_credentials(setup_anthropic_mock):
model = AnthropicLargeLanguageModel() model = AnthropicLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="claude-instant-1.2", credentials={"anthropic_api_key": "invalid_key"})
model='claude-instant-1.2',
credentials={
'anthropic_api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(
model='claude-instant-1.2', model="claude-instant-1.2", credentials={"anthropic_api_key": os.environ.get("ANTHROPIC_API_KEY")}
credentials={
'anthropic_api_key': os.environ.get('ANTHROPIC_API_KEY')
}
) )
@pytest.mark.parametrize('setup_anthropic_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_anthropic_mock", [["none"]], indirect=True)
def test_invoke_model(setup_anthropic_mock): def test_invoke_model(setup_anthropic_mock):
model = AnthropicLargeLanguageModel() model = AnthropicLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='claude-instant-1.2', model="claude-instant-1.2",
credentials={ credentials={
'anthropic_api_key': os.environ.get('ANTHROPIC_API_KEY'), "anthropic_api_key": os.environ.get("ANTHROPIC_API_KEY"),
'anthropic_api_url': os.environ.get('ANTHROPIC_API_URL') "anthropic_api_url": os.environ.get("ANTHROPIC_API_URL"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={"temperature": 0.0, "top_p": 1.0, "max_tokens": 10},
'temperature': 0.0, stop=["How"],
'top_p': 1.0,
'max_tokens': 10
},
stop=['How'],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
@pytest.mark.parametrize('setup_anthropic_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_anthropic_mock", [["none"]], indirect=True)
def test_invoke_stream_model(setup_anthropic_mock): def test_invoke_stream_model(setup_anthropic_mock):
model = AnthropicLargeLanguageModel() model = AnthropicLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='claude-instant-1.2', model="claude-instant-1.2",
credentials={ credentials={"anthropic_api_key": os.environ.get("ANTHROPIC_API_KEY")},
'anthropic_api_key': os.environ.get('ANTHROPIC_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={"temperature": 0.0, "max_tokens": 100},
'temperature': 0.0,
'max_tokens': 100
},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -98,18 +79,14 @@ def test_get_num_tokens():
model = AnthropicLargeLanguageModel() model = AnthropicLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='claude-instant-1.2', model="claude-instant-1.2",
credentials={ credentials={"anthropic_api_key": os.environ.get("ANTHROPIC_API_KEY")},
'anthropic_api_key': os.environ.get('ANTHROPIC_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert num_tokens == 18 assert num_tokens == 18

12
api/tests/integration_tests/model_runtime/anthropic/test_provider.py
View file

@ -7,17 +7,11 @@ from core.model_runtime.model_providers.anthropic.anthropic import AnthropicProv
from tests.integration_tests.model_runtime.__mock.anthropic import setup_anthropic_mock from tests.integration_tests.model_runtime.__mock.anthropic import setup_anthropic_mock
@pytest.mark.parametrize('setup_anthropic_mock', [['none']], indirect=True) @pytest.mark.parametrize("setup_anthropic_mock", [["none"]], indirect=True)
def test_validate_provider_credentials(setup_anthropic_mock): def test_validate_provider_credentials(setup_anthropic_mock):
provider = AnthropicProvider() provider = AnthropicProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"anthropic_api_key": os.environ.get("ANTHROPIC_API_KEY")})
credentials={
'anthropic_api_key': os.environ.get('ANTHROPIC_API_KEY')
}
)

261
api/tests/integration_tests/model_runtime/azure_openai/test_llm.py
View file

File diff suppressed because one or more lines are too long

49
api/tests/integration_tests/model_runtime/azure_openai/test_text_embedding.py
View file

@ -8,45 +8,43 @@ from core.model_runtime.model_providers.azure_openai.text_embedding.text_embeddi
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock
@pytest.mark.parametrize('setup_openai_mock', [['text_embedding']], indirect=True) @pytest.mark.parametrize("setup_openai_mock", [["text_embedding"]], indirect=True)
def test_validate_credentials(setup_openai_mock): def test_validate_credentials(setup_openai_mock):
model = AzureOpenAITextEmbeddingModel() model = AzureOpenAITextEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='embedding', model="embedding",
credentials={ credentials={
'openai_api_base': os.environ.get('AZURE_OPENAI_API_BASE'), "openai_api_base": os.environ.get("AZURE_OPENAI_API_BASE"),
'openai_api_key': 'invalid_key', "openai_api_key": "invalid_key",
'base_model_name': 'text-embedding-ada-002' "base_model_name": "text-embedding-ada-002",
} },
) )
model.validate_credentials( model.validate_credentials(
model='embedding', model="embedding",
credentials={ credentials={
'openai_api_base': os.environ.get('AZURE_OPENAI_API_BASE'), "openai_api_base": os.environ.get("AZURE_OPENAI_API_BASE"),
'openai_api_key': os.environ.get('AZURE_OPENAI_API_KEY'), "openai_api_key": os.environ.get("AZURE_OPENAI_API_KEY"),
'base_model_name': 'text-embedding-ada-002' "base_model_name": "text-embedding-ada-002",
} },
) )
@pytest.mark.parametrize('setup_openai_mock', [['text_embedding']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["text_embedding"]], indirect=True)
def test_invoke_model(setup_openai_mock): def test_invoke_model(setup_openai_mock):
model = AzureOpenAITextEmbeddingModel() model = AzureOpenAITextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='embedding', model="embedding",
credentials={ credentials={
'openai_api_base': os.environ.get('AZURE_OPENAI_API_BASE'), "openai_api_base": os.environ.get("AZURE_OPENAI_API_BASE"),
'openai_api_key': os.environ.get('AZURE_OPENAI_API_KEY'), "openai_api_key": os.environ.get("AZURE_OPENAI_API_KEY"),
'base_model_name': 'text-embedding-ada-002' "base_model_name": "text-embedding-ada-002",
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -58,14 +56,7 @@ def test_get_num_tokens():
model = AzureOpenAITextEmbeddingModel() model = AzureOpenAITextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='embedding', model="embedding", credentials={"base_model_name": "text-embedding-ada-002"}, texts=["hello", "world"]
credentials={
'base_model_name': 'text-embedding-ada-002'
},
texts=[
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2

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

12
api/tests/integration_tests/model_runtime/baichuan/test_provider.py
View file

@ -10,14 +10,6 @@ def test_validate_provider_credentials():
provider = BaichuanProvider() provider = BaichuanProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"api_key": "hahahaha"})
credentials={
'api_key': 'hahahaha'
}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"api_key": os.environ.get("BAICHUAN_API_KEY")})
credentials={
'api_key': os.environ.get('BAICHUAN_API_KEY')
}
)

40
api/tests/integration_tests/model_runtime/baichuan/test_text_embedding.py
View file

@ -11,18 +11,10 @@ def test_validate_credentials():
model = BaichuanTextEmbeddingModel() model = BaichuanTextEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="baichuan-text-embedding", credentials={"api_key": "invalid_key"})
model='baichuan-text-embedding',
credentials={
'api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(
model='baichuan-text-embedding', model="baichuan-text-embedding", credentials={"api_key": os.environ.get("BAICHUAN_API_KEY")}
credentials={
'api_key': os.environ.get('BAICHUAN_API_KEY')
}
) )
@ -30,44 +22,40 @@ def test_invoke_model():
model = BaichuanTextEmbeddingModel() model = BaichuanTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='baichuan-text-embedding', model="baichuan-text-embedding",
credentials={ credentials={
'api_key': os.environ.get('BAICHUAN_API_KEY'), "api_key": os.environ.get("BAICHUAN_API_KEY"),
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
assert len(result.embeddings) == 2 assert len(result.embeddings) == 2
assert result.usage.total_tokens == 6 assert result.usage.total_tokens == 6
def test_get_num_tokens(): def test_get_num_tokens():
model = BaichuanTextEmbeddingModel() model = BaichuanTextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='baichuan-text-embedding', model="baichuan-text-embedding",
credentials={ credentials={
'api_key': os.environ.get('BAICHUAN_API_KEY'), "api_key": os.environ.get("BAICHUAN_API_KEY"),
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2
def test_max_chunks(): def test_max_chunks():
model = BaichuanTextEmbeddingModel() model = BaichuanTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='baichuan-text-embedding', model="baichuan-text-embedding",
credentials={ credentials={
'api_key': os.environ.get('BAICHUAN_API_KEY'), "api_key": os.environ.get("BAICHUAN_API_KEY"),
}, },
texts=[ texts=[
"hello", "hello",
@ -92,7 +80,7 @@ def test_max_chunks():
"world", "world",
"hello", "hello",
"world", "world",
] ],
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)

66
api/tests/integration_tests/model_runtime/bedrock/test_llm.py
View file

@ -13,77 +13,63 @@ def test_validate_credentials():
model = BedrockLargeLanguageModel() model = BedrockLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="meta.llama2-13b-chat-v1", credentials={"anthropic_api_key": "invalid_key"})
model='meta.llama2-13b-chat-v1',
credentials={
'anthropic_api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(
model='meta.llama2-13b-chat-v1', model="meta.llama2-13b-chat-v1",
credentials={ credentials={
"aws_region": os.getenv("AWS_REGION"), "aws_region": os.getenv("AWS_REGION"),
"aws_access_key": os.getenv("AWS_ACCESS_KEY"), "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
"aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY") "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
} },
) )
def test_invoke_model(): def test_invoke_model():
model = BedrockLargeLanguageModel() model = BedrockLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='meta.llama2-13b-chat-v1', model="meta.llama2-13b-chat-v1",
credentials={ credentials={
"aws_region": os.getenv("AWS_REGION"), "aws_region": os.getenv("AWS_REGION"),
"aws_access_key": os.getenv("AWS_ACCESS_KEY"), "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
"aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY") "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={"temperature": 0.0, "top_p": 1.0, "max_tokens_to_sample": 10},
'temperature': 0.0, stop=["How"],
'top_p': 1.0,
'max_tokens_to_sample': 10
},
stop=['How'],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
def test_invoke_stream_model(): def test_invoke_stream_model():
model = BedrockLargeLanguageModel() model = BedrockLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='meta.llama2-13b-chat-v1', model="meta.llama2-13b-chat-v1",
credentials={ credentials={
"aws_region": os.getenv("AWS_REGION"), "aws_region": os.getenv("AWS_REGION"),
"aws_access_key": os.getenv("AWS_ACCESS_KEY"), "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
"aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY") "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={"temperature": 0.0, "max_tokens_to_sample": 100},
'temperature': 0.0,
'max_tokens_to_sample': 100
},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -100,20 +86,18 @@ def test_get_num_tokens():
model = BedrockLargeLanguageModel() model = BedrockLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='meta.llama2-13b-chat-v1', model="meta.llama2-13b-chat-v1",
credentials = { credentials={
"aws_region": os.getenv("AWS_REGION"), "aws_region": os.getenv("AWS_REGION"),
"aws_access_key": os.getenv("AWS_ACCESS_KEY"), "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
"aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY") "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
}, },
messages=[ messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert num_tokens == 18 assert num_tokens == 18

6
api/tests/integration_tests/model_runtime/bedrock/test_provider.py
View file

@ -10,14 +10,12 @@ def test_validate_provider_credentials():
provider = BedrockProvider() provider = BedrockProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(
credentials={ credentials={
"aws_region": os.getenv("AWS_REGION"), "aws_region": os.getenv("AWS_REGION"),
"aws_access_key": os.getenv("AWS_ACCESS_KEY"), "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
"aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY") "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
} }
) )

203
api/tests/integration_tests/model_runtime/chatglm/test_llm.py
View file

@ -23,79 +23,64 @@ def test_predefined_models():
assert len(model_schemas) >= 1 assert len(model_schemas) >= 1
assert isinstance(model_schemas[0], AIModelEntity) assert isinstance(model_schemas[0], AIModelEntity)
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_validate_credentials_for_chat_model(setup_openai_mock): def test_validate_credentials_for_chat_model(setup_openai_mock):
model = ChatGLMLargeLanguageModel() model = ChatGLMLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="chatglm2-6b", credentials={"api_base": "invalid_key"})
model='chatglm2-6b',
credentials={
'api_base': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(model="chatglm2-6b", credentials={"api_base": os.environ.get("CHATGLM_API_BASE")})
model='chatglm2-6b',
credentials={
'api_base': os.environ.get('CHATGLM_API_BASE')
}
)
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_invoke_model(setup_openai_mock): def test_invoke_model(setup_openai_mock):
model = ChatGLMLargeLanguageModel() model = ChatGLMLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='chatglm2-6b', model="chatglm2-6b",
credentials={ credentials={"api_base": os.environ.get("CHATGLM_API_BASE")},
'api_base': os.environ.get('CHATGLM_API_BASE')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_invoke_stream_model(setup_openai_mock): def test_invoke_stream_model(setup_openai_mock):
model = ChatGLMLargeLanguageModel() model = ChatGLMLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='chatglm2-6b', model="chatglm2-6b",
credentials={ credentials={"api_base": os.environ.get("CHATGLM_API_BASE")},
'api_base': os.environ.get('CHATGLM_API_BASE')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -105,52 +90,41 @@ def test_invoke_stream_model(setup_openai_mock):
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_invoke_stream_model_with_functions(setup_openai_mock): def test_invoke_stream_model_with_functions(setup_openai_mock):
model = ChatGLMLargeLanguageModel() model = ChatGLMLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='chatglm3-6b', model="chatglm3-6b",
credentials={ credentials={"api_base": os.environ.get("CHATGLM_API_BASE")},
'api_base': os.environ.get('CHATGLM_API_BASE')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='你是一个天气机器人，你不知道今天的天气怎么样，你需要通过调用一个函数来获取天气信息。' content="你是一个天气机器人，你不知道今天的天气怎么样，你需要通过调用一个函数来获取天气信息。"
), ),
UserPromptMessage( UserPromptMessage(content="波士顿天气如何？"),
content='波士顿天气如何？'
)
], ],
model_parameters={ model_parameters={
'temperature': 0, "temperature": 0,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
user='abc-123', user="abc-123",
stream=True, stream=True,
tools=[ tools=[
PromptMessageTool( PromptMessageTool(
name='get_current_weather', name="get_current_weather",
description='Get the current weather in a given location', description="Get the current weather in a given location",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {
"location": { "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"type": "string", "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
"description": "The city and state e.g. San Francisco, CA"
}, },
"unit": { "required": ["location"],
"type": "string",
"enum": ["celsius", "fahrenheit"]
}
}, },
"required": [
"location"
]
}
) )
] ],
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -170,120 +144,85 @@ def test_invoke_stream_model_with_functions(setup_openai_mock):
break break
assert call is not None assert call is not None
assert call.delta.message.tool_calls[0].function.name == 'get_current_weather' assert call.delta.message.tool_calls[0].function.name == "get_current_weather"
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_invoke_model_with_functions(setup_openai_mock): def test_invoke_model_with_functions(setup_openai_mock):
model = ChatGLMLargeLanguageModel() model = ChatGLMLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='chatglm3-6b', model="chatglm3-6b",
credentials={ credentials={"api_base": os.environ.get("CHATGLM_API_BASE")},
'api_base': os.environ.get('CHATGLM_API_BASE') prompt_messages=[UserPromptMessage(content="What is the weather like in San Francisco?")],
},
prompt_messages=[
UserPromptMessage(
content='What is the weather like in San Francisco?'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
user='abc-123', user="abc-123",
stream=False, stream=False,
tools=[ tools=[
PromptMessageTool( PromptMessageTool(
name='get_current_weather', name="get_current_weather",
description='Get the current weather in a given location', description="Get the current weather in a given location",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {
"location": { "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"type": "string", "unit": {"type": "string", "enum": ["c", "f"]},
"description": "The city and state e.g. San Francisco, CA"
}, },
"unit": { "required": ["location"],
"type": "string",
"enum": [
"c",
"f"
]
}
}, },
"required": [
"location"
]
}
) )
] ],
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
assert response.message.tool_calls[0].function.name == 'get_current_weather' assert response.message.tool_calls[0].function.name == "get_current_weather"
def test_get_num_tokens(): def test_get_num_tokens():
model = ChatGLMLargeLanguageModel() model = ChatGLMLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='chatglm2-6b', model="chatglm2-6b",
credentials={ credentials={"api_base": os.environ.get("CHATGLM_API_BASE")},
'api_base': os.environ.get('CHATGLM_API_BASE')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
tools=[ tools=[
PromptMessageTool( PromptMessageTool(
name='get_current_weather', name="get_current_weather",
description='Get the current weather in a given location', description="Get the current weather in a given location",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {
"location": { "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"type": "string", "unit": {"type": "string", "enum": ["c", "f"]},
"description": "The city and state e.g. San Francisco, CA"
}, },
"unit": { "required": ["location"],
"type": "string",
"enum": [
"c",
"f"
]
}
}, },
"required": [
"location"
]
}
) )
] ],
) )
assert isinstance(num_tokens, int) assert isinstance(num_tokens, int)
assert num_tokens == 77 assert num_tokens == 77
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='chatglm2-6b', model="chatglm2-6b",
credentials={ credentials={"api_base": os.environ.get("CHATGLM_API_BASE")},
'api_base': os.environ.get('CHATGLM_API_BASE')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
) )

14
api/tests/integration_tests/model_runtime/chatglm/test_provider.py
View file

@ -7,19 +7,11 @@ from core.model_runtime.model_providers.chatglm.chatglm import ChatGLMProvider
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True) @pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_validate_provider_credentials(setup_openai_mock): def test_validate_provider_credentials(setup_openai_mock):
provider = ChatGLMProvider() provider = ChatGLMProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"api_base": "hahahaha"})
credentials={
'api_base': 'hahahaha'
}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"api_base": os.environ.get("CHATGLM_API_BASE")})
credentials={
'api_base': os.environ.get('CHATGLM_API_BASE')
}
)

179
api/tests/integration_tests/model_runtime/cohere/test_llm.py
View file

@ -13,87 +13,49 @@ def test_validate_credentials_for_chat_model():
model = CohereLargeLanguageModel() model = CohereLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="command-light-chat", credentials={"api_key": "invalid_key"})
model='command-light-chat',
credentials={
'api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(model="command-light-chat", credentials={"api_key": os.environ.get("COHERE_API_KEY")})
model='command-light-chat',
credentials={
'api_key': os.environ.get('COHERE_API_KEY')
}
)
def test_validate_credentials_for_completion_model(): def test_validate_credentials_for_completion_model():
model = CohereLargeLanguageModel() model = CohereLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="command-light", credentials={"api_key": "invalid_key"})
model='command-light',
credentials={
'api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(model="command-light", credentials={"api_key": os.environ.get("COHERE_API_KEY")})
model='command-light',
credentials={
'api_key': os.environ.get('COHERE_API_KEY')
}
)
def test_invoke_completion_model(): def test_invoke_completion_model():
model = CohereLargeLanguageModel() model = CohereLargeLanguageModel()
credentials = { credentials = {"api_key": os.environ.get("COHERE_API_KEY")}
'api_key': os.environ.get('COHERE_API_KEY')
}
result = model.invoke( result = model.invoke(
model='command-light', model="command-light",
credentials=credentials, credentials=credentials,
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage( model_parameters={"temperature": 0.0, "max_tokens": 1},
content='Hello World!'
)
],
model_parameters={
'temperature': 0.0,
'max_tokens': 1
},
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(result, LLMResult) assert isinstance(result, LLMResult)
assert len(result.message.content) > 0 assert len(result.message.content) > 0
assert model._num_tokens_from_string('command-light', credentials, result.message.content) == 1 assert model._num_tokens_from_string("command-light", credentials, result.message.content) == 1
def test_invoke_stream_completion_model(): def test_invoke_stream_completion_model():
model = CohereLargeLanguageModel() model = CohereLargeLanguageModel()
result = model.invoke( result = model.invoke(
model='command-light', model="command-light",
credentials={ credentials={"api_key": os.environ.get("COHERE_API_KEY")},
'api_key': os.environ.get('COHERE_API_KEY') prompt_messages=[UserPromptMessage(content="Hello World!")],
}, model_parameters={"temperature": 0.0, "max_tokens": 100},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={
'temperature': 0.0,
'max_tokens': 100
},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(result, Generator) assert isinstance(result, Generator)
@ -109,28 +71,24 @@ def test_invoke_chat_model():
model = CohereLargeLanguageModel() model = CohereLargeLanguageModel()
result = model.invoke( result = model.invoke(
model='command-light-chat', model="command-light-chat",
credentials={ credentials={"api_key": os.environ.get("COHERE_API_KEY")},
'api_key': os.environ.get('COHERE_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={
'temperature': 0.0, "temperature": 0.0,
'p': 0.99, "p": 0.99,
'presence_penalty': 0.0, "presence_penalty": 0.0,
'frequency_penalty': 0.0, "frequency_penalty": 0.0,
'max_tokens': 10 "max_tokens": 10,
}, },
stop=['How'], stop=["How"],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(result, LLMResult) assert isinstance(result, LLMResult)
@ -141,24 +99,17 @@ def test_invoke_stream_chat_model():
model = CohereLargeLanguageModel() model = CohereLargeLanguageModel()
result = model.invoke( result = model.invoke(
model='command-light-chat', model="command-light-chat",
credentials={ credentials={"api_key": os.environ.get("COHERE_API_KEY")},
'api_key': os.environ.get('COHERE_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={"temperature": 0.0, "max_tokens": 100},
'temperature': 0.0,
'max_tokens': 100
},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(result, Generator) assert isinstance(result, Generator)
@ -177,32 +128,22 @@ def test_get_num_tokens():
model = CohereLargeLanguageModel() model = CohereLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='command-light', model="command-light",
credentials={ credentials={"api_key": os.environ.get("COHERE_API_KEY")},
'api_key': os.environ.get('COHERE_API_KEY') prompt_messages=[UserPromptMessage(content="Hello World!")],
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
]
) )
assert num_tokens == 3 assert num_tokens == 3
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='command-light-chat', model="command-light-chat",
credentials={ credentials={"api_key": os.environ.get("COHERE_API_KEY")},
'api_key': os.environ.get('COHERE_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert num_tokens == 15 assert num_tokens == 15
@ -213,25 +154,17 @@ def test_fine_tuned_model():
# test invoke # test invoke
result = model.invoke( result = model.invoke(
model='85ec47be-6139-4f75-a4be-0f0ec1ef115c-ft', model="85ec47be-6139-4f75-a4be-0f0ec1ef115c-ft",
credentials={ credentials={"api_key": os.environ.get("COHERE_API_KEY"), "mode": "completion"},
'api_key': os.environ.get('COHERE_API_KEY'),
'mode': 'completion'
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={"temperature": 0.0, "max_tokens": 100},
'temperature': 0.0,
'max_tokens': 100
},
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(result, LLMResult) assert isinstance(result, LLMResult)
@ -242,25 +175,17 @@ def test_fine_tuned_chat_model():
# test invoke # test invoke
result = model.invoke( result = model.invoke(
model='94f2d55a-4c79-4c00-bde4-23962e74b170-ft', model="94f2d55a-4c79-4c00-bde4-23962e74b170-ft",
credentials={ credentials={"api_key": os.environ.get("COHERE_API_KEY"), "mode": "chat"},
'api_key': os.environ.get('COHERE_API_KEY'),
'mode': 'chat'
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={"temperature": 0.0, "max_tokens": 100},
'temperature': 0.0,
'max_tokens': 100
},
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(result, LLMResult) assert isinstance(result, LLMResult)

10
api/tests/integration_tests/model_runtime/cohere/test_provider.py
View file

@ -10,12 +10,6 @@ def test_validate_provider_credentials():
provider = CohereProvider() provider = CohereProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"api_key": os.environ.get("COHERE_API_KEY")})
credentials={
'api_key': os.environ.get('COHERE_API_KEY')
}
)

24
api/tests/integration_tests/model_runtime/cohere/test_rerank.py
View file

@ -11,29 +11,17 @@ def test_validate_credentials():
model = CohereRerankModel() model = CohereRerankModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="rerank-english-v2.0", credentials={"api_key": "invalid_key"})
model='rerank-english-v2.0',
credentials={
'api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(model="rerank-english-v2.0", credentials={"api_key": os.environ.get("COHERE_API_KEY")})
model='rerank-english-v2.0',
credentials={
'api_key': os.environ.get('COHERE_API_KEY')
}
)
def test_invoke_model(): def test_invoke_model():
model = CohereRerankModel() model = CohereRerankModel()
result = model.invoke( result = model.invoke(
model='rerank-english-v2.0', model="rerank-english-v2.0",
credentials={ credentials={"api_key": os.environ.get("COHERE_API_KEY")},
'api_key': os.environ.get('COHERE_API_KEY')
},
query="What is the capital of the United States?", query="What is the capital of the United States?",
docs=[ docs=[
"Carson City is the capital city of the American state of Nevada. At the 2010 United States " "Carson City is the capital city of the American state of Nevada. At the 2010 United States "
@ -41,9 +29,9 @@ def test_invoke_model():
"Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) " "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) "
"is the capital of the United States. It is a federal district. The President of the USA and many major " "is the capital of the United States. It is a federal district. The President of the USA and many major "
"national government offices are in the territory. This makes it the political center of the United " "national government offices are in the territory. This makes it the political center of the United "
"States of America." "States of America.",
], ],
score_threshold=0.8 score_threshold=0.8,
) )
assert isinstance(result, RerankResult) assert isinstance(result, RerankResult)

38
api/tests/integration_tests/model_runtime/cohere/test_text_embedding.py
View file

@ -11,18 +11,10 @@ def test_validate_credentials():
model = CohereTextEmbeddingModel() model = CohereTextEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="embed-multilingual-v3.0", credentials={"api_key": "invalid_key"})
model='embed-multilingual-v3.0',
credentials={
'api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(
model='embed-multilingual-v3.0', model="embed-multilingual-v3.0", credentials={"api_key": os.environ.get("COHERE_API_KEY")}
credentials={
'api_key': os.environ.get('COHERE_API_KEY')
}
) )
@ -30,17 +22,10 @@ def test_invoke_model():
model = CohereTextEmbeddingModel() model = CohereTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='embed-multilingual-v3.0', model="embed-multilingual-v3.0",
credentials={ credentials={"api_key": os.environ.get("COHERE_API_KEY")},
'api_key': os.environ.get('COHERE_API_KEY') texts=["hello", "world", " ".join(["long_text"] * 100), " ".join(["another_long_text"] * 100)],
}, user="abc-123",
texts=[
"hello",
"world",
" ".join(["long_text"] * 100),
" ".join(["another_long_text"] * 100)
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -52,14 +37,9 @@ def test_get_num_tokens():
model = CohereTextEmbeddingModel() model = CohereTextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='embed-multilingual-v3.0', model="embed-multilingual-v3.0",
credentials={ credentials={"api_key": os.environ.get("COHERE_API_KEY")},
'api_key': os.environ.get('COHERE_API_KEY') texts=["hello", "world"],
},
texts=[
"hello",
"world"
]
) )
assert num_tokens == 3 assert num_tokens == 3

175
api/tests/integration_tests/model_runtime/google/test_llm.py
View file

File diff suppressed because one or more lines are too long

12
api/tests/integration_tests/model_runtime/google/test_provider.py
View file

@ -7,17 +7,11 @@ from core.model_runtime.model_providers.google.google import GoogleProvider
from tests.integration_tests.model_runtime.__mock.google import setup_google_mock from tests.integration_tests.model_runtime.__mock.google import setup_google_mock
@pytest.mark.parametrize('setup_google_mock', [['none']], indirect=True) @pytest.mark.parametrize("setup_google_mock", [["none"]], indirect=True)
def test_validate_provider_credentials(setup_google_mock): def test_validate_provider_credentials(setup_google_mock):
provider = GoogleProvider() provider = GoogleProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"google_api_key": os.environ.get("GOOGLE_API_KEY")})
credentials={
'google_api_key': os.environ.get('GOOGLE_API_KEY')
}
)

260
api/tests/integration_tests/model_runtime/huggingface_hub/test_llm.py
View file

@ -10,87 +10,75 @@ from core.model_runtime.model_providers.huggingface_hub.llm.llm import Huggingfa
from tests.integration_tests.model_runtime.__mock.huggingface import setup_huggingface_mock from tests.integration_tests.model_runtime.__mock.huggingface import setup_huggingface_mock
@pytest.mark.parametrize('setup_huggingface_mock', [['none']], indirect=True) @pytest.mark.parametrize("setup_huggingface_mock", [["none"]], indirect=True)
def test_hosted_inference_api_validate_credentials(setup_huggingface_mock): def test_hosted_inference_api_validate_credentials(setup_huggingface_mock):
model = HuggingfaceHubLargeLanguageModel() model = HuggingfaceHubLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='HuggingFaceH4/zephyr-7b-beta', model="HuggingFaceH4/zephyr-7b-beta",
credentials={ credentials={"huggingfacehub_api_type": "hosted_inference_api", "huggingfacehub_api_token": "invalid_key"},
'huggingfacehub_api_type': 'hosted_inference_api',
'huggingfacehub_api_token': 'invalid_key'
}
) )
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='fake-model', model="fake-model",
credentials={ credentials={"huggingfacehub_api_type": "hosted_inference_api", "huggingfacehub_api_token": "invalid_key"},
'huggingfacehub_api_type': 'hosted_inference_api',
'huggingfacehub_api_token': 'invalid_key'
}
) )
model.validate_credentials( model.validate_credentials(
model='HuggingFaceH4/zephyr-7b-beta', model="HuggingFaceH4/zephyr-7b-beta",
credentials={ credentials={
'huggingfacehub_api_type': 'hosted_inference_api', "huggingfacehub_api_type": "hosted_inference_api",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY') "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
} },
) )
@pytest.mark.parametrize('setup_huggingface_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_huggingface_mock", [["none"]], indirect=True)
def test_hosted_inference_api_invoke_model(setup_huggingface_mock): def test_hosted_inference_api_invoke_model(setup_huggingface_mock):
model = HuggingfaceHubLargeLanguageModel() model = HuggingfaceHubLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='HuggingFaceH4/zephyr-7b-beta', model="HuggingFaceH4/zephyr-7b-beta",
credentials={ credentials={
'huggingfacehub_api_type': 'hosted_inference_api', "huggingfacehub_api_type": "hosted_inference_api",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY') "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Who are you?")],
UserPromptMessage(
content='Who are you?'
)
],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
@pytest.mark.parametrize('setup_huggingface_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_huggingface_mock", [["none"]], indirect=True)
def test_hosted_inference_api_invoke_stream_model(setup_huggingface_mock): def test_hosted_inference_api_invoke_stream_model(setup_huggingface_mock):
model = HuggingfaceHubLargeLanguageModel() model = HuggingfaceHubLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='HuggingFaceH4/zephyr-7b-beta', model="HuggingFaceH4/zephyr-7b-beta",
credentials={ credentials={
'huggingfacehub_api_type': 'hosted_inference_api', "huggingfacehub_api_type": "hosted_inference_api",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY') "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Who are you?")],
UserPromptMessage(
content='Who are you?'
)
],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -101,86 +89,81 @@ def test_hosted_inference_api_invoke_stream_model(setup_huggingface_mock):
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
@pytest.mark.parametrize('setup_huggingface_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_huggingface_mock", [["none"]], indirect=True)
def test_inference_endpoints_text_generation_validate_credentials(setup_huggingface_mock): def test_inference_endpoints_text_generation_validate_credentials(setup_huggingface_mock):
model = HuggingfaceHubLargeLanguageModel() model = HuggingfaceHubLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='openchat/openchat_3.5', model="openchat/openchat_3.5",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': 'invalid_key', "huggingfacehub_api_token": "invalid_key",
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_TEXT_GEN_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_TEXT_GEN_ENDPOINT_URL"),
'task_type': 'text-generation' "task_type": "text-generation",
} },
) )
model.validate_credentials( model.validate_credentials(
model='openchat/openchat_3.5', model="openchat/openchat_3.5",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_TEXT_GEN_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_TEXT_GEN_ENDPOINT_URL"),
'task_type': 'text-generation' "task_type": "text-generation",
} },
) )
@pytest.mark.parametrize('setup_huggingface_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_huggingface_mock", [["none"]], indirect=True)
def test_inference_endpoints_text_generation_invoke_model(setup_huggingface_mock): def test_inference_endpoints_text_generation_invoke_model(setup_huggingface_mock):
model = HuggingfaceHubLargeLanguageModel() model = HuggingfaceHubLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='openchat/openchat_3.5', model="openchat/openchat_3.5",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_TEXT_GEN_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_TEXT_GEN_ENDPOINT_URL"),
'task_type': 'text-generation' "task_type": "text-generation",
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Who are you?")],
UserPromptMessage(
content='Who are you?'
)
],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
@pytest.mark.parametrize('setup_huggingface_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_huggingface_mock", [["none"]], indirect=True)
def test_inference_endpoints_text_generation_invoke_stream_model(setup_huggingface_mock): def test_inference_endpoints_text_generation_invoke_stream_model(setup_huggingface_mock):
model = HuggingfaceHubLargeLanguageModel() model = HuggingfaceHubLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='openchat/openchat_3.5', model="openchat/openchat_3.5",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_TEXT_GEN_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_TEXT_GEN_ENDPOINT_URL"),
'task_type': 'text-generation' "task_type": "text-generation",
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Who are you?")],
UserPromptMessage(
content='Who are you?'
)
],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -191,86 +174,81 @@ def test_inference_endpoints_text_generation_invoke_stream_model(setup_huggingfa
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
@pytest.mark.parametrize('setup_huggingface_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_huggingface_mock", [["none"]], indirect=True)
def test_inference_endpoints_text2text_generation_validate_credentials(setup_huggingface_mock): def test_inference_endpoints_text2text_generation_validate_credentials(setup_huggingface_mock):
model = HuggingfaceHubLargeLanguageModel() model = HuggingfaceHubLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='google/mt5-base', model="google/mt5-base",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': 'invalid_key', "huggingfacehub_api_token": "invalid_key",
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_TEXT2TEXT_GEN_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_TEXT2TEXT_GEN_ENDPOINT_URL"),
'task_type': 'text2text-generation' "task_type": "text2text-generation",
} },
) )
model.validate_credentials( model.validate_credentials(
model='google/mt5-base', model="google/mt5-base",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_TEXT2TEXT_GEN_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_TEXT2TEXT_GEN_ENDPOINT_URL"),
'task_type': 'text2text-generation' "task_type": "text2text-generation",
} },
) )
@pytest.mark.parametrize('setup_huggingface_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_huggingface_mock", [["none"]], indirect=True)
def test_inference_endpoints_text2text_generation_invoke_model(setup_huggingface_mock): def test_inference_endpoints_text2text_generation_invoke_model(setup_huggingface_mock):
model = HuggingfaceHubLargeLanguageModel() model = HuggingfaceHubLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='google/mt5-base', model="google/mt5-base",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_TEXT2TEXT_GEN_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_TEXT2TEXT_GEN_ENDPOINT_URL"),
'task_type': 'text2text-generation' "task_type": "text2text-generation",
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Who are you?")],
UserPromptMessage(
content='Who are you?'
)
],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
@pytest.mark.parametrize('setup_huggingface_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_huggingface_mock", [["none"]], indirect=True)
def test_inference_endpoints_text2text_generation_invoke_stream_model(setup_huggingface_mock): def test_inference_endpoints_text2text_generation_invoke_stream_model(setup_huggingface_mock):
model = HuggingfaceHubLargeLanguageModel() model = HuggingfaceHubLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='google/mt5-base', model="google/mt5-base",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_TEXT2TEXT_GEN_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_TEXT2TEXT_GEN_ENDPOINT_URL"),
'task_type': 'text2text-generation' "task_type": "text2text-generation",
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Who are you?")],
UserPromptMessage(
content='Who are you?'
)
],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -286,18 +264,14 @@ def test_get_num_tokens():
model = HuggingfaceHubLargeLanguageModel() model = HuggingfaceHubLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='google/mt5-base', model="google/mt5-base",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_TEXT2TEXT_GEN_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_TEXT2TEXT_GEN_ENDPOINT_URL"),
'task_type': 'text2text-generation' "task_type": "text2text-generation",
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage(
content='Hello World!'
)
]
) )
assert num_tokens == 7 assert num_tokens == 7

89
api/tests/integration_tests/model_runtime/huggingface_hub/test_text_embedding.py
View file

@ -14,19 +14,19 @@ def test_hosted_inference_api_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='facebook/bart-base', model="facebook/bart-base",
credentials={ credentials={
'huggingfacehub_api_type': 'hosted_inference_api', "huggingfacehub_api_type": "hosted_inference_api",
'huggingfacehub_api_token': 'invalid_key', "huggingfacehub_api_token": "invalid_key",
} },
) )
model.validate_credentials( model.validate_credentials(
model='facebook/bart-base', model="facebook/bart-base",
credentials={ credentials={
'huggingfacehub_api_type': 'hosted_inference_api', "huggingfacehub_api_type": "hosted_inference_api",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
} },
) )
@ -34,15 +34,12 @@ def test_hosted_inference_api_invoke_model():
model = HuggingfaceHubTextEmbeddingModel() model = HuggingfaceHubTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='facebook/bart-base', model="facebook/bart-base",
credentials={ credentials={
'huggingfacehub_api_type': 'hosted_inference_api', "huggingfacehub_api_type": "hosted_inference_api",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -55,25 +52,25 @@ def test_inference_endpoints_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='all-MiniLM-L6-v2', model="all-MiniLM-L6-v2",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': 'invalid_key', "huggingfacehub_api_token": "invalid_key",
'huggingface_namespace': 'Dify-AI', "huggingface_namespace": "Dify-AI",
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_EMBEDDINGS_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_EMBEDDINGS_ENDPOINT_URL"),
'task_type': 'feature-extraction' "task_type": "feature-extraction",
} },
) )
model.validate_credentials( model.validate_credentials(
model='all-MiniLM-L6-v2', model="all-MiniLM-L6-v2",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
'huggingface_namespace': 'Dify-AI', "huggingface_namespace": "Dify-AI",
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_EMBEDDINGS_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_EMBEDDINGS_ENDPOINT_URL"),
'task_type': 'feature-extraction' "task_type": "feature-extraction",
} },
) )
@ -81,18 +78,15 @@ def test_inference_endpoints_invoke_model():
model = HuggingfaceHubTextEmbeddingModel() model = HuggingfaceHubTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='all-MiniLM-L6-v2', model="all-MiniLM-L6-v2",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
'huggingface_namespace': 'Dify-AI', "huggingface_namespace": "Dify-AI",
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_EMBEDDINGS_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_EMBEDDINGS_ENDPOINT_URL"),
'task_type': 'feature-extraction' "task_type": "feature-extraction",
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -104,18 +98,15 @@ def test_get_num_tokens():
model = HuggingfaceHubTextEmbeddingModel() model = HuggingfaceHubTextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='all-MiniLM-L6-v2', model="all-MiniLM-L6-v2",
credentials={ credentials={
'huggingfacehub_api_type': 'inference_endpoints', "huggingfacehub_api_type": "inference_endpoints",
'huggingfacehub_api_token': os.environ.get('HUGGINGFACE_API_KEY'), "huggingfacehub_api_token": os.environ.get("HUGGINGFACE_API_KEY"),
'huggingface_namespace': 'Dify-AI', "huggingface_namespace": "Dify-AI",
'huggingfacehub_endpoint_url': os.environ.get('HUGGINGFACE_EMBEDDINGS_ENDPOINT_URL'), "huggingfacehub_endpoint_url": os.environ.get("HUGGINGFACE_EMBEDDINGS_ENDPOINT_URL"),
'task_type': 'feature-extraction' "task_type": "feature-extraction",
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2

38
api/tests/integration_tests/model_runtime/huggingface_tei/test_embeddings.py
View file

@ -10,61 +10,59 @@ from core.model_runtime.model_providers.huggingface_tei.text_embedding.text_embe
) )
from tests.integration_tests.model_runtime.__mock.huggingface_tei import MockTEIClass from tests.integration_tests.model_runtime.__mock.huggingface_tei import MockTEIClass
MOCK = os.getenv('MOCK_SWITCH', 'false').lower() == 'true' MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture @pytest.fixture
def setup_tei_mock(request, monkeypatch: pytest.MonkeyPatch): def setup_tei_mock(request, monkeypatch: pytest.MonkeyPatch):
if MOCK: if MOCK:
monkeypatch.setattr(TeiHelper, 'get_tei_extra_parameter', MockTEIClass.get_tei_extra_parameter) monkeypatch.setattr(TeiHelper, "get_tei_extra_parameter", MockTEIClass.get_tei_extra_parameter)
monkeypatch.setattr(TeiHelper, 'invoke_tokenize', MockTEIClass.invoke_tokenize) monkeypatch.setattr(TeiHelper, "invoke_tokenize", MockTEIClass.invoke_tokenize)
monkeypatch.setattr(TeiHelper, 'invoke_embeddings', MockTEIClass.invoke_embeddings) monkeypatch.setattr(TeiHelper, "invoke_embeddings", MockTEIClass.invoke_embeddings)
monkeypatch.setattr(TeiHelper, 'invoke_rerank', MockTEIClass.invoke_rerank) monkeypatch.setattr(TeiHelper, "invoke_rerank", MockTEIClass.invoke_rerank)
yield yield
if MOCK: if MOCK:
monkeypatch.undo() monkeypatch.undo()
@pytest.mark.parametrize('setup_tei_mock', [['none']], indirect=True) @pytest.mark.parametrize("setup_tei_mock", [["none"]], indirect=True)
def test_validate_credentials(setup_tei_mock): def test_validate_credentials(setup_tei_mock):
model = HuggingfaceTeiTextEmbeddingModel() model = HuggingfaceTeiTextEmbeddingModel()
# model name is only used in mock # model name is only used in mock
model_name = 'embedding' model_name = "embedding"
if MOCK: if MOCK:
# TEI Provider will check model type by API endpoint, at real server, the model type is correct. # TEI Provider will check model type by API endpoint, at real server, the model type is correct.
# So we dont need to check model type here. Only check in mock # So we dont need to check model type here. Only check in mock
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='reranker', model="reranker",
credentials={ credentials={
'server_url': os.environ.get('TEI_EMBEDDING_SERVER_URL', ""), "server_url": os.environ.get("TEI_EMBEDDING_SERVER_URL", ""),
} },
) )
model.validate_credentials( model.validate_credentials(
model=model_name, model=model_name,
credentials={ credentials={
'server_url': os.environ.get('TEI_EMBEDDING_SERVER_URL', ""), "server_url": os.environ.get("TEI_EMBEDDING_SERVER_URL", ""),
} },
) )
@pytest.mark.parametrize('setup_tei_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_tei_mock", [["none"]], indirect=True)
def test_invoke_model(setup_tei_mock): def test_invoke_model(setup_tei_mock):
model = HuggingfaceTeiTextEmbeddingModel() model = HuggingfaceTeiTextEmbeddingModel()
model_name = 'embedding' model_name = "embedding"
result = model.invoke( result = model.invoke(
model=model_name, model=model_name,
credentials={ credentials={
'server_url': os.environ.get('TEI_EMBEDDING_SERVER_URL', ""), "server_url": os.environ.get("TEI_EMBEDDING_SERVER_URL", ""),
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)

38
api/tests/integration_tests/model_runtime/huggingface_tei/test_rerank.py
View file

@ -11,63 +11,65 @@ from core.model_runtime.model_providers.huggingface_tei.rerank.rerank import (
from core.model_runtime.model_providers.huggingface_tei.text_embedding.text_embedding import TeiHelper from core.model_runtime.model_providers.huggingface_tei.text_embedding.text_embedding import TeiHelper
from tests.integration_tests.model_runtime.__mock.huggingface_tei import MockTEIClass from tests.integration_tests.model_runtime.__mock.huggingface_tei import MockTEIClass
MOCK = os.getenv('MOCK_SWITCH', 'false').lower() == 'true' MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"
@pytest.fixture @pytest.fixture
def setup_tei_mock(request, monkeypatch: pytest.MonkeyPatch): def setup_tei_mock(request, monkeypatch: pytest.MonkeyPatch):
if MOCK: if MOCK:
monkeypatch.setattr(TeiHelper, 'get_tei_extra_parameter', MockTEIClass.get_tei_extra_parameter) monkeypatch.setattr(TeiHelper, "get_tei_extra_parameter", MockTEIClass.get_tei_extra_parameter)
monkeypatch.setattr(TeiHelper, 'invoke_tokenize', MockTEIClass.invoke_tokenize) monkeypatch.setattr(TeiHelper, "invoke_tokenize", MockTEIClass.invoke_tokenize)
monkeypatch.setattr(TeiHelper, 'invoke_embeddings', MockTEIClass.invoke_embeddings) monkeypatch.setattr(TeiHelper, "invoke_embeddings", MockTEIClass.invoke_embeddings)
monkeypatch.setattr(TeiHelper, 'invoke_rerank', MockTEIClass.invoke_rerank) monkeypatch.setattr(TeiHelper, "invoke_rerank", MockTEIClass.invoke_rerank)
yield yield
if MOCK: if MOCK:
monkeypatch.undo() monkeypatch.undo()
@pytest.mark.parametrize('setup_tei_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_tei_mock", [["none"]], indirect=True)
def test_validate_credentials(setup_tei_mock): def test_validate_credentials(setup_tei_mock):
model = HuggingfaceTeiRerankModel() model = HuggingfaceTeiRerankModel()
# model name is only used in mock # model name is only used in mock
model_name = 'reranker' model_name = "reranker"
if MOCK: if MOCK:
# TEI Provider will check model type by API endpoint, at real server, the model type is correct. # TEI Provider will check model type by API endpoint, at real server, the model type is correct.
# So we dont need to check model type here. Only check in mock # So we dont need to check model type here. Only check in mock
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='embedding', model="embedding",
credentials={ credentials={
'server_url': os.environ.get('TEI_RERANK_SERVER_URL'), "server_url": os.environ.get("TEI_RERANK_SERVER_URL"),
} },
) )
model.validate_credentials( model.validate_credentials(
model=model_name, model=model_name,
credentials={ credentials={
'server_url': os.environ.get('TEI_RERANK_SERVER_URL'), "server_url": os.environ.get("TEI_RERANK_SERVER_URL"),
} },
) )
@pytest.mark.parametrize('setup_tei_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_tei_mock", [["none"]], indirect=True)
def test_invoke_model(setup_tei_mock): def test_invoke_model(setup_tei_mock):
model = HuggingfaceTeiRerankModel() model = HuggingfaceTeiRerankModel()
# model name is only used in mock # model name is only used in mock
model_name = 'reranker' model_name = "reranker"
result = model.invoke( result = model.invoke(
model=model_name, model=model_name,
credentials={ credentials={
'server_url': os.environ.get('TEI_RERANK_SERVER_URL'), "server_url": os.environ.get("TEI_RERANK_SERVER_URL"),
}, },
query="Who is Kasumi?", query="Who is Kasumi?",
docs=[ docs=[
"Kasumi is a girl's name of Japanese origin meaning \"mist\".", 'Kasumi is a girl\'s name of Japanese origin meaning "mist".',
"Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ", "Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ",
"and she leads a team named PopiParty." "and she leads a team named PopiParty.",
], ],
score_threshold=0.8 score_threshold=0.8,
) )
assert isinstance(result, RerankResult) assert isinstance(result, RerankResult)

69
api/tests/integration_tests/model_runtime/hunyuan/test_llm.py
View file

@ -14,19 +14,15 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='hunyuan-standard', model="hunyuan-standard", credentials={"secret_id": "invalid_key", "secret_key": "invalid_key"}
credentials={
'secret_id': 'invalid_key',
'secret_key': 'invalid_key'
}
) )
model.validate_credentials( model.validate_credentials(
model='hunyuan-standard', model="hunyuan-standard",
credentials={ credentials={
'secret_id': os.environ.get('HUNYUAN_SECRET_ID'), "secret_id": os.environ.get("HUNYUAN_SECRET_ID"),
'secret_key': os.environ.get('HUNYUAN_SECRET_KEY') "secret_key": os.environ.get("HUNYUAN_SECRET_KEY"),
} },
) )
@ -34,23 +30,16 @@ def test_invoke_model():
model = HunyuanLargeLanguageModel() model = HunyuanLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='hunyuan-standard', model="hunyuan-standard",
credentials={ credentials={
'secret_id': os.environ.get('HUNYUAN_SECRET_ID'), "secret_id": os.environ.get("HUNYUAN_SECRET_ID"),
'secret_key': os.environ.get('HUNYUAN_SECRET_KEY') "secret_key": os.environ.get("HUNYUAN_SECRET_KEY"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hi")],
UserPromptMessage( model_parameters={"temperature": 0.5, "max_tokens": 10},
content='Hi' stop=["How"],
)
],
model_parameters={
'temperature': 0.5,
'max_tokens': 10
},
stop=['How'],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -61,23 +50,15 @@ def test_invoke_stream_model():
model = HunyuanLargeLanguageModel() model = HunyuanLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='hunyuan-standard', model="hunyuan-standard",
credentials={ credentials={
'secret_id': os.environ.get('HUNYUAN_SECRET_ID'), "secret_id": os.environ.get("HUNYUAN_SECRET_ID"),
'secret_key': os.environ.get('HUNYUAN_SECRET_KEY') "secret_key": os.environ.get("HUNYUAN_SECRET_KEY"),
},
prompt_messages=[
UserPromptMessage(
content='Hi'
)
],
model_parameters={
'temperature': 0.5,
'max_tokens': 100,
'seed': 1234
}, },
prompt_messages=[UserPromptMessage(content="Hi")],
model_parameters={"temperature": 0.5, "max_tokens": 100, "seed": 1234},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -93,19 +74,17 @@ def test_get_num_tokens():
model = HunyuanLargeLanguageModel() model = HunyuanLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='hunyuan-standard', model="hunyuan-standard",
credentials={ credentials={
'secret_id': os.environ.get('HUNYUAN_SECRET_ID'), "secret_id": os.environ.get("HUNYUAN_SECRET_ID"),
'secret_key': os.environ.get('HUNYUAN_SECRET_KEY') "secret_key": os.environ.get("HUNYUAN_SECRET_KEY"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert num_tokens == 14 assert num_tokens == 14

11
api/tests/integration_tests/model_runtime/hunyuan/test_provider.py
View file

@ -10,16 +10,11 @@ def test_validate_provider_credentials():
provider = HunyuanProvider() provider = HunyuanProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"secret_id": "invalid_key", "secret_key": "invalid_key"})
credentials={
'secret_id': 'invalid_key',
'secret_key': 'invalid_key'
}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(
credentials={ credentials={
'secret_id': os.environ.get('HUNYUAN_SECRET_ID'), "secret_id": os.environ.get("HUNYUAN_SECRET_ID"),
'secret_key': os.environ.get('HUNYUAN_SECRET_KEY') "secret_key": os.environ.get("HUNYUAN_SECRET_KEY"),
} }
) )

48
api/tests/integration_tests/model_runtime/hunyuan/test_text_embedding.py
View file

@ -12,19 +12,15 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='hunyuan-embedding', model="hunyuan-embedding", credentials={"secret_id": "invalid_key", "secret_key": "invalid_key"}
credentials={
'secret_id': 'invalid_key',
'secret_key': 'invalid_key'
}
) )
model.validate_credentials( model.validate_credentials(
model='hunyuan-embedding', model="hunyuan-embedding",
credentials={ credentials={
'secret_id': os.environ.get('HUNYUAN_SECRET_ID'), "secret_id": os.environ.get("HUNYUAN_SECRET_ID"),
'secret_key': os.environ.get('HUNYUAN_SECRET_KEY') "secret_key": os.environ.get("HUNYUAN_SECRET_KEY"),
} },
) )
@ -32,47 +28,43 @@ def test_invoke_model():
model = HunyuanTextEmbeddingModel() model = HunyuanTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='hunyuan-embedding', model="hunyuan-embedding",
credentials={ credentials={
'secret_id': os.environ.get('HUNYUAN_SECRET_ID'), "secret_id": os.environ.get("HUNYUAN_SECRET_ID"),
'secret_key': os.environ.get('HUNYUAN_SECRET_KEY') "secret_key": os.environ.get("HUNYUAN_SECRET_KEY"),
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
assert len(result.embeddings) == 2 assert len(result.embeddings) == 2
assert result.usage.total_tokens == 6 assert result.usage.total_tokens == 6
def test_get_num_tokens(): def test_get_num_tokens():
model = HunyuanTextEmbeddingModel() model = HunyuanTextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='hunyuan-embedding', model="hunyuan-embedding",
credentials={ credentials={
'secret_id': os.environ.get('HUNYUAN_SECRET_ID'), "secret_id": os.environ.get("HUNYUAN_SECRET_ID"),
'secret_key': os.environ.get('HUNYUAN_SECRET_KEY') "secret_key": os.environ.get("HUNYUAN_SECRET_KEY"),
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2
def test_max_chunks(): def test_max_chunks():
model = HunyuanTextEmbeddingModel() model = HunyuanTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='hunyuan-embedding', model="hunyuan-embedding",
credentials={ credentials={
'secret_id': os.environ.get('HUNYUAN_SECRET_ID'), "secret_id": os.environ.get("HUNYUAN_SECRET_ID"),
'secret_key': os.environ.get('HUNYUAN_SECRET_KEY') "secret_key": os.environ.get("HUNYUAN_SECRET_KEY"),
}, },
texts=[ texts=[
"hello", "hello",
@ -97,7 +89,7 @@ def test_max_chunks():
"world", "world",
"hello", "hello",
"world", "world",
] ],
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)

12
api/tests/integration_tests/model_runtime/jina/test_provider.py
View file

@ -10,14 +10,6 @@ def test_validate_provider_credentials():
provider = JinaProvider() provider = JinaProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"api_key": "hahahaha"})
credentials={
'api_key': 'hahahaha'
}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"api_key": os.environ.get("JINA_API_KEY")})
credentials={
'api_key': os.environ.get('JINA_API_KEY')
}
)

32
api/tests/integration_tests/model_runtime/jina/test_text_embedding.py
View file

@ -11,18 +11,10 @@ def test_validate_credentials():
model = JinaTextEmbeddingModel() model = JinaTextEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="jina-embeddings-v2-base-en", credentials={"api_key": "invalid_key"})
model='jina-embeddings-v2-base-en',
credentials={
'api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(
model='jina-embeddings-v2-base-en', model="jina-embeddings-v2-base-en", credentials={"api_key": os.environ.get("JINA_API_KEY")}
credentials={
'api_key': os.environ.get('JINA_API_KEY')
}
) )
@ -30,15 +22,12 @@ def test_invoke_model():
model = JinaTextEmbeddingModel() model = JinaTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='jina-embeddings-v2-base-en', model="jina-embeddings-v2-base-en",
credentials={ credentials={
'api_key': os.environ.get('JINA_API_KEY'), "api_key": os.environ.get("JINA_API_KEY"),
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -50,14 +39,11 @@ def test_get_num_tokens():
model = JinaTextEmbeddingModel() model = JinaTextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='jina-embeddings-v2-base-en', model="jina-embeddings-v2-base-en",
credentials={ credentials={
'api_key': os.environ.get('JINA_API_KEY'), "api_key": os.environ.get("JINA_API_KEY"),
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 6 assert num_tokens == 6

4
api/tests/integration_tests/model_runtime/localai/test_embedding.py
View file

@ -1,4 +1,4 @@
""" """
LocalAI Embedding Interface is temporarily unavailable due to LocalAI Embedding Interface is temporarily unavailable due to
we could not find a way to test it for now. we could not find a way to test it for now.
""" """

152
api/tests/integration_tests/model_runtime/localai/test_llm.py
View file

@ -21,99 +21,78 @@ def test_validate_credentials_for_chat_model():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='chinese-llama-2-7b', model="chinese-llama-2-7b",
credentials={ credentials={
'server_url': 'hahahaha', "server_url": "hahahaha",
'completion_type': 'completion', "completion_type": "completion",
} },
) )
model.validate_credentials( model.validate_credentials(
model='chinese-llama-2-7b', model="chinese-llama-2-7b",
credentials={ credentials={
'server_url': os.environ.get('LOCALAI_SERVER_URL'), "server_url": os.environ.get("LOCALAI_SERVER_URL"),
'completion_type': 'completion', "completion_type": "completion",
} },
) )
def test_invoke_completion_model(): def test_invoke_completion_model():
model = LocalAILanguageModel() model = LocalAILanguageModel()
response = model.invoke( response = model.invoke(
model='chinese-llama-2-7b', model="chinese-llama-2-7b",
credentials={ credentials={
'server_url': os.environ.get('LOCALAI_SERVER_URL'), "server_url": os.environ.get("LOCALAI_SERVER_URL"),
'completion_type': 'completion', "completion_type": "completion",
},
prompt_messages=[
UserPromptMessage(
content='ping'
)
],
model_parameters={
'temperature': 0.7,
'top_p': 1.0,
'max_tokens': 10
}, },
prompt_messages=[UserPromptMessage(content="ping")],
model_parameters={"temperature": 0.7, "top_p": 1.0, "max_tokens": 10},
stop=[], stop=[],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
def test_invoke_chat_model(): def test_invoke_chat_model():
model = LocalAILanguageModel() model = LocalAILanguageModel()
response = model.invoke( response = model.invoke(
model='chinese-llama-2-7b', model="chinese-llama-2-7b",
credentials={ credentials={
'server_url': os.environ.get('LOCALAI_SERVER_URL'), "server_url": os.environ.get("LOCALAI_SERVER_URL"),
'completion_type': 'chat_completion', "completion_type": "chat_completion",
},
prompt_messages=[
UserPromptMessage(
content='ping'
)
],
model_parameters={
'temperature': 0.7,
'top_p': 1.0,
'max_tokens': 10
}, },
prompt_messages=[UserPromptMessage(content="ping")],
model_parameters={"temperature": 0.7, "top_p": 1.0, "max_tokens": 10},
stop=[], stop=[],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
def test_invoke_stream_completion_model(): def test_invoke_stream_completion_model():
model = LocalAILanguageModel() model = LocalAILanguageModel()
response = model.invoke( response = model.invoke(
model='chinese-llama-2-7b', model="chinese-llama-2-7b",
credentials={ credentials={
'server_url': os.environ.get('LOCALAI_SERVER_URL'), "server_url": os.environ.get("LOCALAI_SERVER_URL"),
'completion_type': 'completion', "completion_type": "completion",
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage( model_parameters={"temperature": 0.7, "top_p": 1.0, "max_tokens": 10},
content='Hello World!' stop=["you"],
)
],
model_parameters={
'temperature': 0.7,
'top_p': 1.0,
'max_tokens': 10
},
stop=['you'],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -123,28 +102,21 @@ def test_invoke_stream_completion_model():
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
def test_invoke_stream_chat_model(): def test_invoke_stream_chat_model():
model = LocalAILanguageModel() model = LocalAILanguageModel()
response = model.invoke( response = model.invoke(
model='chinese-llama-2-7b', model="chinese-llama-2-7b",
credentials={ credentials={
'server_url': os.environ.get('LOCALAI_SERVER_URL'), "server_url": os.environ.get("LOCALAI_SERVER_URL"),
'completion_type': 'chat_completion', "completion_type": "chat_completion",
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage( model_parameters={"temperature": 0.7, "top_p": 1.0, "max_tokens": 10},
content='Hello World!' stop=["you"],
)
],
model_parameters={
'temperature': 0.7,
'top_p': 1.0,
'max_tokens': 10
},
stop=['you'],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -154,64 +126,48 @@ def test_invoke_stream_chat_model():
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
def test_get_num_tokens(): def test_get_num_tokens():
model = LocalAILanguageModel() model = LocalAILanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='????', model="????",
credentials={ credentials={
'server_url': os.environ.get('LOCALAI_SERVER_URL'), "server_url": os.environ.get("LOCALAI_SERVER_URL"),
'completion_type': 'chat_completion', "completion_type": "chat_completion",
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
tools=[ tools=[
PromptMessageTool( PromptMessageTool(
name='get_current_weather', name="get_current_weather",
description='Get the current weather in a given location', description="Get the current weather in a given location",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {
"location": { "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"type": "string", "unit": {"type": "string", "enum": ["c", "f"]},
"description": "The city and state e.g. San Francisco, CA"
}, },
"unit": { "required": ["location"],
"type": "string",
"enum": [
"c",
"f"
]
}
}, },
"required": [
"location"
]
}
) )
] ],
) )
assert isinstance(num_tokens, int) assert isinstance(num_tokens, int)
assert num_tokens == 77 assert num_tokens == 77
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='????', model="????",
credentials={ credentials={
'server_url': os.environ.get('LOCALAI_SERVER_URL'), "server_url": os.environ.get("LOCALAI_SERVER_URL"),
'completion_type': 'chat_completion', "completion_type": "chat_completion",
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage(
content='Hello World!'
)
],
) )
assert isinstance(num_tokens, int) assert isinstance(num_tokens, int)

54
api/tests/integration_tests/model_runtime/localai/test_rerank.py
View file

@ -12,30 +12,29 @@ def test_validate_credentials_for_chat_model():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='bge-reranker-v2-m3', model="bge-reranker-v2-m3",
credentials={ credentials={
'server_url': 'hahahaha', "server_url": "hahahaha",
'completion_type': 'completion', "completion_type": "completion",
} },
) )
model.validate_credentials( model.validate_credentials(
model='bge-reranker-base', model="bge-reranker-base",
credentials={ credentials={
'server_url': os.environ.get('LOCALAI_SERVER_URL'), "server_url": os.environ.get("LOCALAI_SERVER_URL"),
'completion_type': 'completion', "completion_type": "completion",
} },
) )
def test_invoke_rerank_model(): def test_invoke_rerank_model():
model = LocalaiRerankModel() model = LocalaiRerankModel()
response = model.invoke( response = model.invoke(
model='bge-reranker-base', model="bge-reranker-base",
credentials={ credentials={"server_url": os.environ.get("LOCALAI_SERVER_URL")},
'server_url': os.environ.get('LOCALAI_SERVER_URL') query="Organic skincare products for sensitive skin",
},
query='Organic skincare products for sensitive skin',
docs=[ docs=[
"Eco-friendly kitchenware for modern homes", "Eco-friendly kitchenware for modern homes",
"Biodegradable cleaning supplies for eco-conscious consumers", "Biodegradable cleaning supplies for eco-conscious consumers",
@ -45,43 +44,38 @@ def test_invoke_rerank_model():
"Sustainable gardening tools and compost solutions", "Sustainable gardening tools and compost solutions",
"Sensitive skin-friendly facial cleansers and toners", "Sensitive skin-friendly facial cleansers and toners",
"Organic food wraps and storage solutions", "Organic food wraps and storage solutions",
"Yoga mats made from recycled materials" "Yoga mats made from recycled materials",
], ],
top_n=3, top_n=3,
score_threshold=0.75, score_threshold=0.75,
user="abc-123" user="abc-123",
) )
assert isinstance(response, RerankResult) assert isinstance(response, RerankResult)
assert len(response.docs) == 3 assert len(response.docs) == 3
def test__invoke(): def test__invoke():
model = LocalaiRerankModel() model = LocalaiRerankModel()
# Test case 1: Empty docs # Test case 1: Empty docs
result = model._invoke( result = model._invoke(
model='bge-reranker-base', model="bge-reranker-base",
credentials={ credentials={"server_url": "https://example.com", "api_key": "1234567890"},
'server_url': 'https://example.com', query="Organic skincare products for sensitive skin",
'api_key': '1234567890'
},
query='Organic skincare products for sensitive skin',
docs=[], docs=[],
top_n=3, top_n=3,
score_threshold=0.75, score_threshold=0.75,
user="abc-123" user="abc-123",
) )
assert isinstance(result, RerankResult) assert isinstance(result, RerankResult)
assert len(result.docs) == 0 assert len(result.docs) == 0
# Test case 2: Valid invocation # Test case 2: Valid invocation
result = model._invoke( result = model._invoke(
model='bge-reranker-base', model="bge-reranker-base",
credentials={ credentials={"server_url": "https://example.com", "api_key": "1234567890"},
'server_url': 'https://example.com', query="Organic skincare products for sensitive skin",
'api_key': '1234567890'
},
query='Organic skincare products for sensitive skin',
docs=[ docs=[
"Eco-friendly kitchenware for modern homes", "Eco-friendly kitchenware for modern homes",
"Biodegradable cleaning supplies for eco-conscious consumers", "Biodegradable cleaning supplies for eco-conscious consumers",
@ -91,11 +85,11 @@ def test__invoke():
"Sustainable gardening tools and compost solutions", "Sustainable gardening tools and compost solutions",
"Sensitive skin-friendly facial cleansers and toners", "Sensitive skin-friendly facial cleansers and toners",
"Organic food wraps and storage solutions", "Organic food wraps and storage solutions",
"Yoga mats made from recycled materials" "Yoga mats made from recycled materials",
], ],
top_n=3, top_n=3,
score_threshold=0.75, score_threshold=0.75,
user="abc-123" user="abc-123",
) )
assert isinstance(result, RerankResult) assert isinstance(result, RerankResult)
assert len(result.docs) == 3 assert len(result.docs) == 3

30
api/tests/integration_tests/model_runtime/localai/test_speech2text.py
View file

@ -10,19 +10,9 @@ def test_validate_credentials():
model = LocalAISpeech2text() model = LocalAISpeech2text()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="whisper-1", credentials={"server_url": "invalid_url"})
model='whisper-1',
credentials={
'server_url': 'invalid_url'
}
)
model.validate_credentials( model.validate_credentials(model="whisper-1", credentials={"server_url": os.environ.get("LOCALAI_SERVER_URL")})
model='whisper-1',
credentials={
'server_url': os.environ.get('LOCALAI_SERVER_URL')
}
)
def test_invoke_model(): def test_invoke_model():
@ -32,23 +22,21 @@ def test_invoke_model():
current_dir = os.path.dirname(os.path.abspath(__file__)) current_dir = os.path.dirname(os.path.abspath(__file__))
# Get assets directory # Get assets directory
assets_dir = os.path.join(os.path.dirname(current_dir), 'assets') assets_dir = os.path.join(os.path.dirname(current_dir), "assets")
# Construct the path to the audio file # Construct the path to the audio file
audio_file_path = os.path.join(assets_dir, 'audio.mp3') audio_file_path = os.path.join(assets_dir, "audio.mp3")
# Open the file and get the file object # Open the file and get the file object
with open(audio_file_path, 'rb') as audio_file: with open(audio_file_path, "rb") as audio_file:
file = audio_file file = audio_file
result = model.invoke( result = model.invoke(
model='whisper-1', model="whisper-1",
credentials={ credentials={"server_url": os.environ.get("LOCALAI_SERVER_URL")},
'server_url': os.environ.get('LOCALAI_SERVER_URL')
},
file=file, file=file,
user="abc-123" user="abc-123",
) )
assert isinstance(result, str) assert isinstance(result, str)
assert result == '1, 2, 3, 4, 5, 6, 7, 8, 9, 10' assert result == "1, 2, 3, 4, 5, 6, 7, 8, 9, 10"

41
api/tests/integration_tests/model_runtime/minimax/test_embedding.py
View file

@ -12,54 +12,47 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='embo-01', model="embo-01",
credentials={ credentials={"minimax_api_key": "invalid_key", "minimax_group_id": os.environ.get("MINIMAX_GROUP_ID")},
'minimax_api_key': 'invalid_key',
'minimax_group_id': os.environ.get('MINIMAX_GROUP_ID')
}
) )
model.validate_credentials( model.validate_credentials(
model='embo-01', model="embo-01",
credentials={ credentials={
'minimax_api_key': os.environ.get('MINIMAX_API_KEY'), "minimax_api_key": os.environ.get("MINIMAX_API_KEY"),
'minimax_group_id': os.environ.get('MINIMAX_GROUP_ID') "minimax_group_id": os.environ.get("MINIMAX_GROUP_ID"),
} },
) )
def test_invoke_model(): def test_invoke_model():
model = MinimaxTextEmbeddingModel() model = MinimaxTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='embo-01', model="embo-01",
credentials={ credentials={
'minimax_api_key': os.environ.get('MINIMAX_API_KEY'), "minimax_api_key": os.environ.get("MINIMAX_API_KEY"),
'minimax_group_id': os.environ.get('MINIMAX_GROUP_ID') "minimax_group_id": os.environ.get("MINIMAX_GROUP_ID"),
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
assert len(result.embeddings) == 2 assert len(result.embeddings) == 2
assert result.usage.total_tokens == 16 assert result.usage.total_tokens == 16
def test_get_num_tokens(): def test_get_num_tokens():
model = MinimaxTextEmbeddingModel() model = MinimaxTextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='embo-01', model="embo-01",
credentials={ credentials={
'minimax_api_key': os.environ.get('MINIMAX_API_KEY'), "minimax_api_key": os.environ.get("MINIMAX_API_KEY"),
'minimax_group_id': os.environ.get('MINIMAX_GROUP_ID') "minimax_group_id": os.environ.get("MINIMAX_GROUP_ID"),
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2

105
api/tests/integration_tests/model_runtime/minimax/test_llm.py
View file

@ -17,79 +17,70 @@ def test_predefined_models():
assert len(model_schemas) >= 1 assert len(model_schemas) >= 1
assert isinstance(model_schemas[0], AIModelEntity) assert isinstance(model_schemas[0], AIModelEntity)
def test_validate_credentials_for_chat_model(): def test_validate_credentials_for_chat_model():
sleep(3) sleep(3)
model = MinimaxLargeLanguageModel() model = MinimaxLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='abab5.5-chat', model="abab5.5-chat", credentials={"minimax_api_key": "invalid_key", "minimax_group_id": "invalid_key"}
credentials={
'minimax_api_key': 'invalid_key',
'minimax_group_id': 'invalid_key'
}
) )
model.validate_credentials( model.validate_credentials(
model='abab5.5-chat', model="abab5.5-chat",
credentials={ credentials={
'minimax_api_key': os.environ.get('MINIMAX_API_KEY'), "minimax_api_key": os.environ.get("MINIMAX_API_KEY"),
'minimax_group_id': os.environ.get('MINIMAX_GROUP_ID') "minimax_group_id": os.environ.get("MINIMAX_GROUP_ID"),
} },
) )
def test_invoke_model(): def test_invoke_model():
sleep(3) sleep(3)
model = MinimaxLargeLanguageModel() model = MinimaxLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='abab5-chat', model="abab5-chat",
credentials={ credentials={
'minimax_api_key': os.environ.get('MINIMAX_API_KEY'), "minimax_api_key": os.environ.get("MINIMAX_API_KEY"),
'minimax_group_id': os.environ.get('MINIMAX_GROUP_ID') "minimax_group_id": os.environ.get("MINIMAX_GROUP_ID"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
'top_k': 1, "top_k": 1,
}, },
stop=['you'], stop=["you"],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
def test_invoke_stream_model(): def test_invoke_stream_model():
sleep(3) sleep(3)
model = MinimaxLargeLanguageModel() model = MinimaxLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='abab5.5-chat', model="abab5.5-chat",
credentials={ credentials={
'minimax_api_key': os.environ.get('MINIMAX_API_KEY'), "minimax_api_key": os.environ.get("MINIMAX_API_KEY"),
'minimax_group_id': os.environ.get('MINIMAX_GROUP_ID') "minimax_group_id": os.environ.get("MINIMAX_GROUP_ID"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
'top_k': 1, "top_k": 1,
}, },
stop=['you'], stop=["you"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -99,34 +90,31 @@ def test_invoke_stream_model():
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
def test_invoke_with_search(): def test_invoke_with_search():
sleep(3) sleep(3)
model = MinimaxLargeLanguageModel() model = MinimaxLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='abab5.5-chat', model="abab5.5-chat",
credentials={ credentials={
'minimax_api_key': os.environ.get('MINIMAX_API_KEY'), "minimax_api_key": os.environ.get("MINIMAX_API_KEY"),
'minimax_group_id': os.environ.get('MINIMAX_GROUP_ID') "minimax_group_id": os.environ.get("MINIMAX_GROUP_ID"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="北京今天的天气怎么样")],
UserPromptMessage(
content='北京今天的天气怎么样'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
'top_k': 1, "top_k": 1,
'plugin_web_search': True, "plugin_web_search": True,
}, },
stop=['you'], stop=["you"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
total_message = '' total_message = ""
for chunk in response: for chunk in response:
assert isinstance(chunk, LLMResultChunk) assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta) assert isinstance(chunk.delta, LLMResultChunkDelta)
@ -134,24 +122,21 @@ def test_invoke_with_search():
total_message += chunk.delta.message.content total_message += chunk.delta.message.content
assert len(chunk.delta.message.content) > 0 if not chunk.delta.finish_reason else True assert len(chunk.delta.message.content) > 0 if not chunk.delta.finish_reason else True
assert '参考资料' in total_message assert "参考资料" in total_message
def test_get_num_tokens(): def test_get_num_tokens():
sleep(3) sleep(3)
model = MinimaxLargeLanguageModel() model = MinimaxLargeLanguageModel()
response = model.get_num_tokens( response = model.get_num_tokens(
model='abab5.5-chat', model="abab5.5-chat",
credentials={ credentials={
'minimax_api_key': os.environ.get('MINIMAX_API_KEY'), "minimax_api_key": os.environ.get("MINIMAX_API_KEY"),
'minimax_group_id': os.environ.get('MINIMAX_GROUP_ID') "minimax_group_id": os.environ.get("MINIMAX_GROUP_ID"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage( tools=[],
content='Hello World!'
)
],
tools=[]
) )
assert isinstance(response, int) assert isinstance(response, int)

8
api/tests/integration_tests/model_runtime/minimax/test_provider.py
View file

@ -12,14 +12,14 @@ def test_validate_provider_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(
credentials={ credentials={
'minimax_api_key': 'hahahaha', "minimax_api_key": "hahahaha",
'minimax_group_id': '123', "minimax_group_id": "123",
} }
) )
provider.validate_provider_credentials( provider.validate_provider_credentials(
credentials={ credentials={
'minimax_api_key': os.environ.get('MINIMAX_API_KEY'), "minimax_api_key": os.environ.get("MINIMAX_API_KEY"),
'minimax_group_id': os.environ.get('MINIMAX_GROUP_ID'), "minimax_group_id": os.environ.get("MINIMAX_GROUP_ID"),
} }
) )

70
api/tests/integration_tests/model_runtime/novita/test_llm.py
View file

@ -19,19 +19,12 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='meta-llama/llama-3-8b-instruct', model="meta-llama/llama-3-8b-instruct", credentials={"api_key": "invalid_key", "mode": "chat"}
credentials={
'api_key': 'invalid_key',
'mode': 'chat'
}
) )
model.validate_credentials( model.validate_credentials(
model='meta-llama/llama-3-8b-instruct', model="meta-llama/llama-3-8b-instruct",
credentials={ credentials={"api_key": os.environ.get("NOVITA_API_KEY"), "mode": "chat"},
'api_key': os.environ.get('NOVITA_API_KEY'),
'mode': 'chat'
}
) )
@ -39,27 +32,22 @@ def test_invoke_model():
model = NovitaLargeLanguageModel() model = NovitaLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='meta-llama/llama-3-8b-instruct', model="meta-llama/llama-3-8b-instruct",
credentials={ credentials={"api_key": os.environ.get("NOVITA_API_KEY"), "mode": "completion"},
'api_key': os.environ.get('NOVITA_API_KEY'),
'mode': 'completion'
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Who are you?"),
content='Who are you?'
)
], ],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_p': 0.5, "top_p": 0.5,
'max_tokens': 10, "max_tokens": 10,
}, },
stop=['How'], stop=["How"],
stream=False, stream=False,
user="novita" user="novita",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -70,27 +58,17 @@ def test_invoke_stream_model():
model = NovitaLargeLanguageModel() model = NovitaLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='meta-llama/llama-3-8b-instruct', model="meta-llama/llama-3-8b-instruct",
credentials={ credentials={"api_key": os.environ.get("NOVITA_API_KEY"), "mode": "chat"},
'api_key': os.environ.get('NOVITA_API_KEY'),
'mode': 'chat'
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Who are you?"),
content='Who are you?'
)
], ],
model_parameters={ model_parameters={"temperature": 1.0, "top_k": 2, "top_p": 0.5, "max_tokens": 100},
'temperature': 1.0,
'top_k': 2,
'top_p': 0.5,
'max_tokens': 100
},
stream=True, stream=True,
user="novita" user="novita",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -105,18 +83,16 @@ def test_get_num_tokens():
model = NovitaLargeLanguageModel() model = NovitaLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='meta-llama/llama-3-8b-instruct', model="meta-llama/llama-3-8b-instruct",
credentials={ credentials={
'api_key': os.environ.get('NOVITA_API_KEY'), "api_key": os.environ.get("NOVITA_API_KEY"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert isinstance(num_tokens, int) assert isinstance(num_tokens, int)

6
api/tests/integration_tests/model_runtime/novita/test_provider.py
View file

@ -10,12 +10,10 @@ def test_validate_provider_credentials():
provider = NovitaProvider() provider = NovitaProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(
credentials={ credentials={
'api_key': os.environ.get('NOVITA_API_KEY'), "api_key": os.environ.get("NOVITA_API_KEY"),
} }
) )

190
api/tests/integration_tests/model_runtime/ollama/test_llm.py
View file

File diff suppressed because one or more lines are too long

48
api/tests/integration_tests/model_runtime/ollama/test_text_embedding.py
View file

@ -12,21 +12,21 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='mistral:text', model="mistral:text",
credentials={ credentials={
'base_url': 'http://localhost:21434', "base_url": "http://localhost:21434",
'mode': 'chat', "mode": "chat",
'context_size': 4096, "context_size": 4096,
} },
) )
model.validate_credentials( model.validate_credentials(
model='mistral:text', model="mistral:text",
credentials={ credentials={
'base_url': os.environ.get('OLLAMA_BASE_URL'), "base_url": os.environ.get("OLLAMA_BASE_URL"),
'mode': 'chat', "mode": "chat",
'context_size': 4096, "context_size": 4096,
} },
) )
@ -34,17 +34,14 @@ def test_invoke_model():
model = OllamaEmbeddingModel() model = OllamaEmbeddingModel()
result = model.invoke( result = model.invoke(
model='mistral:text', model="mistral:text",
credentials={ credentials={
'base_url': os.environ.get('OLLAMA_BASE_URL'), "base_url": os.environ.get("OLLAMA_BASE_URL"),
'mode': 'chat', "mode": "chat",
'context_size': 4096, "context_size": 4096,
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -56,16 +53,13 @@ def test_get_num_tokens():
model = OllamaEmbeddingModel() model = OllamaEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='mistral:text', model="mistral:text",
credentials={ credentials={
'base_url': os.environ.get('OLLAMA_BASE_URL'), "base_url": os.environ.get("OLLAMA_BASE_URL"),
'mode': 'chat', "mode": "chat",
'context_size': 4096, "context_size": 4096,
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2

288
api/tests/integration_tests/model_runtime/openai/test_llm.py
View file

File diff suppressed because one or more lines are too long

33
api/tests/integration_tests/model_runtime/openai/test_moderation.py
View file

@ -7,48 +7,37 @@ from core.model_runtime.model_providers.openai.moderation.moderation import Open
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock
@pytest.mark.parametrize('setup_openai_mock', [['moderation']], indirect=True) @pytest.mark.parametrize("setup_openai_mock", [["moderation"]], indirect=True)
def test_validate_credentials(setup_openai_mock): def test_validate_credentials(setup_openai_mock):
model = OpenAIModerationModel() model = OpenAIModerationModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="text-moderation-stable", credentials={"openai_api_key": "invalid_key"})
model='text-moderation-stable',
credentials={
'openai_api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(
model='text-moderation-stable', model="text-moderation-stable", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}
credentials={
'openai_api_key': os.environ.get('OPENAI_API_KEY')
}
) )
@pytest.mark.parametrize('setup_openai_mock', [['moderation']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["moderation"]], indirect=True)
def test_invoke_model(setup_openai_mock): def test_invoke_model(setup_openai_mock):
model = OpenAIModerationModel() model = OpenAIModerationModel()
result = model.invoke( result = model.invoke(
model='text-moderation-stable', model="text-moderation-stable",
credentials={ credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
'openai_api_key': os.environ.get('OPENAI_API_KEY')
},
text="hello", text="hello",
user="abc-123" user="abc-123",
) )
assert isinstance(result, bool) assert isinstance(result, bool)
assert result is False assert result is False
result = model.invoke( result = model.invoke(
model='text-moderation-stable', model="text-moderation-stable",
credentials={ credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
'openai_api_key': os.environ.get('OPENAI_API_KEY')
},
text="i will kill you", text="i will kill you",
user="abc-123" user="abc-123",
) )
assert isinstance(result, bool) assert isinstance(result, bool)

12
api/tests/integration_tests/model_runtime/openai/test_provider.py
View file

@ -7,17 +7,11 @@ from core.model_runtime.model_providers.openai.openai import OpenAIProvider
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True) @pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_validate_provider_credentials(setup_openai_mock): def test_validate_provider_credentials(setup_openai_mock):
provider = OpenAIProvider() provider = OpenAIProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")})
credentials={
'openai_api_key': os.environ.get('OPENAI_API_KEY')
}
)

35
api/tests/integration_tests/model_runtime/openai/test_speech2text.py
View file

@ -7,26 +7,17 @@ from core.model_runtime.model_providers.openai.speech2text.speech2text import Op
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock
@pytest.mark.parametrize('setup_openai_mock', [['speech2text']], indirect=True) @pytest.mark.parametrize("setup_openai_mock", [["speech2text"]], indirect=True)
def test_validate_credentials(setup_openai_mock): def test_validate_credentials(setup_openai_mock):
model = OpenAISpeech2TextModel() model = OpenAISpeech2TextModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="whisper-1", credentials={"openai_api_key": "invalid_key"})
model='whisper-1',
credentials={
'openai_api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(model="whisper-1", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")})
model='whisper-1',
credentials={
'openai_api_key': os.environ.get('OPENAI_API_KEY')
}
)
@pytest.mark.parametrize('setup_openai_mock', [['speech2text']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["speech2text"]], indirect=True)
def test_invoke_model(setup_openai_mock): def test_invoke_model(setup_openai_mock):
model = OpenAISpeech2TextModel() model = OpenAISpeech2TextModel()
@ -34,23 +25,21 @@ def test_invoke_model(setup_openai_mock):
current_dir = os.path.dirname(os.path.abspath(__file__)) current_dir = os.path.dirname(os.path.abspath(__file__))
# Get assets directory # Get assets directory
assets_dir = os.path.join(os.path.dirname(current_dir), 'assets') assets_dir = os.path.join(os.path.dirname(current_dir), "assets")
# Construct the path to the audio file # Construct the path to the audio file
audio_file_path = os.path.join(assets_dir, 'audio.mp3') audio_file_path = os.path.join(assets_dir, "audio.mp3")
# Open the file and get the file object # Open the file and get the file object
with open(audio_file_path, 'rb') as audio_file: with open(audio_file_path, "rb") as audio_file:
file = audio_file file = audio_file
result = model.invoke( result = model.invoke(
model='whisper-1', model="whisper-1",
credentials={ credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")},
'openai_api_key': os.environ.get('OPENAI_API_KEY')
},
file=file, file=file,
user="abc-123" user="abc-123",
) )
assert isinstance(result, str) assert isinstance(result, str)
assert result == '1, 2, 3, 4, 5, 6, 7, 8, 9, 10' assert result == "1, 2, 3, 4, 5, 6, 7, 8, 9, 10"

45
api/tests/integration_tests/model_runtime/openai/test_text_embedding.py
View file

@ -8,42 +8,27 @@ from core.model_runtime.model_providers.openai.text_embedding.text_embedding imp
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock
@pytest.mark.parametrize('setup_openai_mock', [['text_embedding']], indirect=True) @pytest.mark.parametrize("setup_openai_mock", [["text_embedding"]], indirect=True)
def test_validate_credentials(setup_openai_mock): def test_validate_credentials(setup_openai_mock):
model = OpenAITextEmbeddingModel() model = OpenAITextEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="text-embedding-ada-002", credentials={"openai_api_key": "invalid_key"})
model='text-embedding-ada-002',
credentials={
'openai_api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(
model='text-embedding-ada-002', model="text-embedding-ada-002", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}
credentials={
'openai_api_key': os.environ.get('OPENAI_API_KEY')
}
) )
@pytest.mark.parametrize('setup_openai_mock', [['text_embedding']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["text_embedding"]], indirect=True)
def test_invoke_model(setup_openai_mock): def test_invoke_model(setup_openai_mock):
model = OpenAITextEmbeddingModel() model = OpenAITextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='text-embedding-ada-002', model="text-embedding-ada-002",
credentials={ credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY"), "openai_api_base": "https://api.openai.com"},
'openai_api_key': os.environ.get('OPENAI_API_KEY'), texts=["hello", "world", " ".join(["long_text"] * 100), " ".join(["another_long_text"] * 100)],
'openai_api_base': 'https://api.openai.com' user="abc-123",
},
texts=[
"hello",
"world",
" ".join(["long_text"] * 100),
" ".join(["another_long_text"] * 100)
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -55,15 +40,9 @@ def test_get_num_tokens():
model = OpenAITextEmbeddingModel() model = OpenAITextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='text-embedding-ada-002', model="text-embedding-ada-002",
credentials={ credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY"), "openai_api_base": "https://api.openai.com"},
'openai_api_key': os.environ.get('OPENAI_API_KEY'), texts=["hello", "world"],
'openai_api_base': 'https://api.openai.com'
},
texts=[
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2

147
api/tests/integration_tests/model_runtime/openai_api_compatible/test_llm.py
View file

@ -23,21 +23,17 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='mistralai/Mixtral-8x7B-Instruct-v0.1', model="mistralai/Mixtral-8x7B-Instruct-v0.1",
credentials={ credentials={"api_key": "invalid_key", "endpoint_url": "https://api.together.xyz/v1/", "mode": "chat"},
'api_key': 'invalid_key',
'endpoint_url': 'https://api.together.xyz/v1/',
'mode': 'chat'
}
) )
model.validate_credentials( model.validate_credentials(
model='mistralai/Mixtral-8x7B-Instruct-v0.1', model="mistralai/Mixtral-8x7B-Instruct-v0.1",
credentials={ credentials={
'api_key': os.environ.get('TOGETHER_API_KEY'), "api_key": os.environ.get("TOGETHER_API_KEY"),
'endpoint_url': 'https://api.together.xyz/v1/', "endpoint_url": "https://api.together.xyz/v1/",
'mode': 'chat' "mode": "chat",
} },
) )
@ -45,28 +41,26 @@ def test_invoke_model():
model = OAIAPICompatLargeLanguageModel() model = OAIAPICompatLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='mistralai/Mixtral-8x7B-Instruct-v0.1', model="mistralai/Mixtral-8x7B-Instruct-v0.1",
credentials={ credentials={
'api_key': os.environ.get('TOGETHER_API_KEY'), "api_key": os.environ.get("TOGETHER_API_KEY"),
'endpoint_url': 'https://api.together.xyz/v1/', "endpoint_url": "https://api.together.xyz/v1/",
'mode': 'completion' "mode": "completion",
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Who are you?"),
content='Who are you?'
)
], ],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -77,29 +71,27 @@ def test_invoke_stream_model():
model = OAIAPICompatLargeLanguageModel() model = OAIAPICompatLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='mistralai/Mixtral-8x7B-Instruct-v0.1', model="mistralai/Mixtral-8x7B-Instruct-v0.1",
credentials={ credentials={
'api_key': os.environ.get('TOGETHER_API_KEY'), "api_key": os.environ.get("TOGETHER_API_KEY"),
'endpoint_url': 'https://api.together.xyz/v1/', "endpoint_url": "https://api.together.xyz/v1/",
'mode': 'chat', "mode": "chat",
'stream_mode_delimiter': '\\n\\n' "stream_mode_delimiter": "\\n\\n",
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Who are you?"),
content='Who are you?'
)
], ],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -114,28 +106,26 @@ def test_invoke_stream_model_without_delimiter():
model = OAIAPICompatLargeLanguageModel() model = OAIAPICompatLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='mistralai/Mixtral-8x7B-Instruct-v0.1', model="mistralai/Mixtral-8x7B-Instruct-v0.1",
credentials={ credentials={
'api_key': os.environ.get('TOGETHER_API_KEY'), "api_key": os.environ.get("TOGETHER_API_KEY"),
'endpoint_url': 'https://api.together.xyz/v1/', "endpoint_url": "https://api.together.xyz/v1/",
'mode': 'chat' "mode": "chat",
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Who are you?"),
content='Who are you?'
)
], ],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -151,51 +141,37 @@ def test_invoke_chat_model_with_tools():
model = OAIAPICompatLargeLanguageModel() model = OAIAPICompatLargeLanguageModel()
result = model.invoke( result = model.invoke(
model='gpt-3.5-turbo', model="gpt-3.5-turbo",
credentials={ credentials={
'api_key': os.environ.get('OPENAI_API_KEY'), "api_key": os.environ.get("OPENAI_API_KEY"),
'endpoint_url': 'https://api.openai.com/v1/', "endpoint_url": "https://api.openai.com/v1/",
'mode': 'chat' "mode": "chat",
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(
content="what's the weather today in London?", content="what's the weather today in London?",
) ),
], ],
tools=[ tools=[
PromptMessageTool( PromptMessageTool(
name='get_weather', name="get_weather",
description='Determine weather in my location', description="Determine weather in my location",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {
"location": { "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"type": "string", "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
"description": "The city and state e.g. San Francisco, CA"
}, },
"unit": { "required": ["location"],
"type": "string",
"enum": [
"celsius",
"fahrenheit"
]
}
}, },
"required": [
"location"
]
}
), ),
], ],
model_parameters={ model_parameters={"temperature": 0.0, "max_tokens": 1024},
'temperature': 0.0,
'max_tokens': 1024
},
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(result, LLMResult) assert isinstance(result, LLMResult)
@ -207,19 +183,14 @@ def test_get_num_tokens():
model = OAIAPICompatLargeLanguageModel() model = OAIAPICompatLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='mistralai/Mixtral-8x7B-Instruct-v0.1', model="mistralai/Mixtral-8x7B-Instruct-v0.1",
credentials={ credentials={"api_key": os.environ.get("OPENAI_API_KEY"), "endpoint_url": "https://api.openai.com/v1/"},
'api_key': os.environ.get('OPENAI_API_KEY'),
'endpoint_url': 'https://api.openai.com/v1/'
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert isinstance(num_tokens, int) assert isinstance(num_tokens, int)

17
api/tests/integration_tests/model_runtime/openai_api_compatible/test_speech2text.py
View file

@ -14,18 +14,12 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model="whisper-1", model="whisper-1",
credentials={ credentials={"api_key": "invalid_key", "endpoint_url": "https://api.openai.com/v1/"},
"api_key": "invalid_key",
"endpoint_url": "https://api.openai.com/v1/"
},
) )
model.validate_credentials( model.validate_credentials(
model="whisper-1", model="whisper-1",
credentials={ credentials={"api_key": os.environ.get("OPENAI_API_KEY"), "endpoint_url": "https://api.openai.com/v1/"},
"api_key": os.environ.get("OPENAI_API_KEY"),
"endpoint_url": "https://api.openai.com/v1/"
},
) )
@ -47,13 +41,10 @@ def test_invoke_model():
result = model.invoke( result = model.invoke(
model="whisper-1", model="whisper-1",
credentials={ credentials={"api_key": os.environ.get("OPENAI_API_KEY"), "endpoint_url": "https://api.openai.com/v1/"},
"api_key": os.environ.get("OPENAI_API_KEY"),
"endpoint_url": "https://api.openai.com/v1/"
},
file=file, file=file,
user="abc-123", user="abc-123",
) )
assert isinstance(result, str) assert isinstance(result, str)
assert result == '1, 2, 3, 4, 5, 6, 7, 8, 9, 10' assert result == "1, 2, 3, 4, 5, 6, 7, 8, 9, 10"

50
api/tests/integration_tests/model_runtime/openai_api_compatible/test_text_embedding.py
View file

@ -12,27 +12,23 @@ from core.model_runtime.model_providers.openai_api_compatible.text_embedding.tex
Using OpenAI's API as testing endpoint Using OpenAI's API as testing endpoint
""" """
def test_validate_credentials(): def test_validate_credentials():
model = OAICompatEmbeddingModel() model = OAICompatEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='text-embedding-ada-002', model="text-embedding-ada-002",
credentials={ credentials={"api_key": "invalid_key", "endpoint_url": "https://api.openai.com/v1/", "context_size": 8184},
'api_key': 'invalid_key',
'endpoint_url': 'https://api.openai.com/v1/',
'context_size': 8184
}
) )
model.validate_credentials( model.validate_credentials(
model='text-embedding-ada-002', model="text-embedding-ada-002",
credentials={ credentials={
'api_key': os.environ.get('OPENAI_API_KEY'), "api_key": os.environ.get("OPENAI_API_KEY"),
'endpoint_url': 'https://api.openai.com/v1/', "endpoint_url": "https://api.openai.com/v1/",
'context_size': 8184 "context_size": 8184,
} },
) )
@ -40,19 +36,14 @@ def test_invoke_model():
model = OAICompatEmbeddingModel() model = OAICompatEmbeddingModel()
result = model.invoke( result = model.invoke(
model='text-embedding-ada-002', model="text-embedding-ada-002",
credentials={ credentials={
'api_key': os.environ.get('OPENAI_API_KEY'), "api_key": os.environ.get("OPENAI_API_KEY"),
'endpoint_url': 'https://api.openai.com/v1/', "endpoint_url": "https://api.openai.com/v1/",
'context_size': 8184 "context_size": 8184,
}, },
texts=[ texts=["hello", "world", " ".join(["long_text"] * 100), " ".join(["another_long_text"] * 100)],
"hello", user="abc-123",
"world",
" ".join(["long_text"] * 100),
" ".join(["another_long_text"] * 100)
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -64,16 +55,13 @@ def test_get_num_tokens():
model = OAICompatEmbeddingModel() model = OAICompatEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='text-embedding-ada-002', model="text-embedding-ada-002",
credentials={ credentials={
'api_key': os.environ.get('OPENAI_API_KEY'), "api_key": os.environ.get("OPENAI_API_KEY"),
'endpoint_url': 'https://api.openai.com/v1/embeddings', "endpoint_url": "https://api.openai.com/v1/embeddings",
'context_size': 8184 "context_size": 8184,
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2

33
api/tests/integration_tests/model_runtime/openllm/test_embedding.py
View file

@ -12,17 +12,17 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'server_url': 'ww' + os.environ.get('OPENLLM_SERVER_URL'), "server_url": "ww" + os.environ.get("OPENLLM_SERVER_URL"),
} },
) )
model.validate_credentials( model.validate_credentials(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'server_url': os.environ.get('OPENLLM_SERVER_URL'), "server_url": os.environ.get("OPENLLM_SERVER_URL"),
} },
) )
@ -30,33 +30,28 @@ def test_invoke_model():
model = OpenLLMTextEmbeddingModel() model = OpenLLMTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'server_url': os.environ.get('OPENLLM_SERVER_URL'), "server_url": os.environ.get("OPENLLM_SERVER_URL"),
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
assert len(result.embeddings) == 2 assert len(result.embeddings) == 2
assert result.usage.total_tokens > 0 assert result.usage.total_tokens > 0
def test_get_num_tokens(): def test_get_num_tokens():
model = OpenLLMTextEmbeddingModel() model = OpenLLMTextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'server_url': os.environ.get('OPENLLM_SERVER_URL'), "server_url": os.environ.get("OPENLLM_SERVER_URL"),
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2

67
api/tests/integration_tests/model_runtime/openllm/test_llm.py
View file

@ -14,67 +14,61 @@ def test_validate_credentials_for_chat_model():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'server_url': 'invalid_key', "server_url": "invalid_key",
} },
) )
model.validate_credentials( model.validate_credentials(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'server_url': os.environ.get('OPENLLM_SERVER_URL'), "server_url": os.environ.get("OPENLLM_SERVER_URL"),
} },
) )
def test_invoke_model(): def test_invoke_model():
model = OpenLLMLargeLanguageModel() model = OpenLLMLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'server_url': os.environ.get('OPENLLM_SERVER_URL'), "server_url": os.environ.get("OPENLLM_SERVER_URL"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
'top_k': 1, "top_k": 1,
}, },
stop=['you'], stop=["you"],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
def test_invoke_stream_model(): def test_invoke_stream_model():
model = OpenLLMLargeLanguageModel() model = OpenLLMLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'server_url': os.environ.get('OPENLLM_SERVER_URL'), "server_url": os.environ.get("OPENLLM_SERVER_URL"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
'top_k': 1, "top_k": 1,
}, },
stop=['you'], stop=["you"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -84,20 +78,17 @@ def test_invoke_stream_model():
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
def test_get_num_tokens(): def test_get_num_tokens():
model = OpenLLMLargeLanguageModel() model = OpenLLMLargeLanguageModel()
response = model.get_num_tokens( response = model.get_num_tokens(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'server_url': os.environ.get('OPENLLM_SERVER_URL'), "server_url": os.environ.get("OPENLLM_SERVER_URL"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage( tools=[],
content='Hello World!'
)
],
tools=[]
) )
assert isinstance(response, int) assert isinstance(response, int)

71
api/tests/integration_tests/model_runtime/openrouter/test_llm.py
View file

@ -19,19 +19,12 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='mistralai/mixtral-8x7b-instruct', model="mistralai/mixtral-8x7b-instruct", credentials={"api_key": "invalid_key", "mode": "chat"}
credentials={
'api_key': 'invalid_key',
'mode': 'chat'
}
) )
model.validate_credentials( model.validate_credentials(
model='mistralai/mixtral-8x7b-instruct', model="mistralai/mixtral-8x7b-instruct",
credentials={ credentials={"api_key": os.environ.get("TOGETHER_API_KEY"), "mode": "chat"},
'api_key': os.environ.get('TOGETHER_API_KEY'),
'mode': 'chat'
}
) )
@ -39,27 +32,22 @@ def test_invoke_model():
model = OpenRouterLargeLanguageModel() model = OpenRouterLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='mistralai/mixtral-8x7b-instruct', model="mistralai/mixtral-8x7b-instruct",
credentials={ credentials={"api_key": os.environ.get("TOGETHER_API_KEY"), "mode": "completion"},
'api_key': os.environ.get('TOGETHER_API_KEY'),
'mode': 'completion'
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Who are you?"),
content='Who are you?'
)
], ],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -70,27 +58,22 @@ def test_invoke_stream_model():
model = OpenRouterLargeLanguageModel() model = OpenRouterLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='mistralai/mixtral-8x7b-instruct', model="mistralai/mixtral-8x7b-instruct",
credentials={ credentials={"api_key": os.environ.get("TOGETHER_API_KEY"), "mode": "chat"},
'api_key': os.environ.get('TOGETHER_API_KEY'),
'mode': 'chat'
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Who are you?"),
content='Who are you?'
)
], ],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -105,18 +88,16 @@ def test_get_num_tokens():
model = OpenRouterLargeLanguageModel() model = OpenRouterLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='mistralai/mixtral-8x7b-instruct', model="mistralai/mixtral-8x7b-instruct",
credentials={ credentials={
'api_key': os.environ.get('TOGETHER_API_KEY'), "api_key": os.environ.get("TOGETHER_API_KEY"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert isinstance(num_tokens, int) assert isinstance(num_tokens, int)

74
api/tests/integration_tests/model_runtime/replicate/test_llm.py
View file

@ -14,19 +14,19 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='meta/llama-2-13b-chat', model="meta/llama-2-13b-chat",
credentials={ credentials={
'replicate_api_token': 'invalid_key', "replicate_api_token": "invalid_key",
'model_version': 'f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d' "model_version": "f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d",
} },
) )
model.validate_credentials( model.validate_credentials(
model='meta/llama-2-13b-chat', model="meta/llama-2-13b-chat",
credentials={ credentials={
'replicate_api_token': os.environ.get('REPLICATE_API_KEY'), "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),
'model_version': 'f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d' "model_version": "f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d",
} },
) )
@ -34,27 +34,25 @@ def test_invoke_model():
model = ReplicateLargeLanguageModel() model = ReplicateLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='meta/llama-2-13b-chat', model="meta/llama-2-13b-chat",
credentials={ credentials={
'replicate_api_token': os.environ.get('REPLICATE_API_KEY'), "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),
'model_version': 'f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d' "model_version": "f4e2de70d66816a838a89eeeb621910adffb0dd0baba3976c96980970978018d",
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Who are you?"),
content='Who are you?'
)
], ],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -65,27 +63,25 @@ def test_invoke_stream_model():
model = ReplicateLargeLanguageModel() model = ReplicateLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='mistralai/mixtral-8x7b-instruct-v0.1', model="mistralai/mixtral-8x7b-instruct-v0.1",
credentials={ credentials={
'replicate_api_token': os.environ.get('REPLICATE_API_KEY'), "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),
'model_version': '2b56576fcfbe32fa0526897d8385dd3fb3d36ba6fd0dbe033c72886b81ade93e' "model_version": "2b56576fcfbe32fa0526897d8385dd3fb3d36ba6fd0dbe033c72886b81ade93e",
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Who are you?"),
content='Who are you?'
)
], ],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -100,19 +96,17 @@ def test_get_num_tokens():
model = ReplicateLargeLanguageModel() model = ReplicateLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='', model="",
credentials={ credentials={
'replicate_api_token': os.environ.get('REPLICATE_API_KEY'), "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),
'model_version': '2b56576fcfbe32fa0526897d8385dd3fb3d36ba6fd0dbe033c72886b81ade93e' "model_version": "2b56576fcfbe32fa0526897d8385dd3fb3d36ba6fd0dbe033c72886b81ade93e",
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert num_tokens == 14 assert num_tokens == 14

95
api/tests/integration_tests/model_runtime/replicate/test_text_embedding.py
View file

@ -12,19 +12,19 @@ def test_validate_credentials_one():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='replicate/all-mpnet-base-v2', model="replicate/all-mpnet-base-v2",
credentials={ credentials={
'replicate_api_token': 'invalid_key', "replicate_api_token": "invalid_key",
'model_version': 'b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305' "model_version": "b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305",
} },
) )
model.validate_credentials( model.validate_credentials(
model='replicate/all-mpnet-base-v2', model="replicate/all-mpnet-base-v2",
credentials={ credentials={
'replicate_api_token': os.environ.get('REPLICATE_API_KEY'), "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),
'model_version': 'b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305' "model_version": "b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305",
} },
) )
@ -33,19 +33,19 @@ def test_validate_credentials_two():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='nateraw/bge-large-en-v1.5', model="nateraw/bge-large-en-v1.5",
credentials={ credentials={
'replicate_api_token': 'invalid_key', "replicate_api_token": "invalid_key",
'model_version': '9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1' "model_version": "9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1",
} },
) )
model.validate_credentials( model.validate_credentials(
model='nateraw/bge-large-en-v1.5', model="nateraw/bge-large-en-v1.5",
credentials={ credentials={
'replicate_api_token': os.environ.get('REPLICATE_API_KEY'), "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),
'model_version': '9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1' "model_version": "9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1",
} },
) )
@ -53,16 +53,13 @@ def test_invoke_model_one():
model = ReplicateEmbeddingModel() model = ReplicateEmbeddingModel()
result = model.invoke( result = model.invoke(
model='nateraw/bge-large-en-v1.5', model="nateraw/bge-large-en-v1.5",
credentials={ credentials={
'replicate_api_token': os.environ.get('REPLICATE_API_KEY'), "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),
'model_version': '9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1' "model_version": "9cf9f015a9cb9c61d1a2610659cdac4a4ca222f2d3707a68517b18c198a9add1",
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -74,16 +71,13 @@ def test_invoke_model_two():
model = ReplicateEmbeddingModel() model = ReplicateEmbeddingModel()
result = model.invoke( result = model.invoke(
model='andreasjansson/clip-features', model="andreasjansson/clip-features",
credentials={ credentials={
'replicate_api_token': os.environ.get('REPLICATE_API_KEY'), "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),
'model_version': '75b33f253f7714a281ad3e9b28f63e3232d583716ef6718f2e46641077ea040a' "model_version": "75b33f253f7714a281ad3e9b28f63e3232d583716ef6718f2e46641077ea040a",
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -95,16 +89,13 @@ def test_invoke_model_three():
model = ReplicateEmbeddingModel() model = ReplicateEmbeddingModel()
result = model.invoke( result = model.invoke(
model='replicate/all-mpnet-base-v2', model="replicate/all-mpnet-base-v2",
credentials={ credentials={
'replicate_api_token': os.environ.get('REPLICATE_API_KEY'), "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),
'model_version': 'b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305' "model_version": "b6b7585c9640cd7a9572c6e129c9549d79c9c31f0d3fdce7baac7c67ca38f305",
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -116,16 +107,13 @@ def test_invoke_model_four():
model = ReplicateEmbeddingModel() model = ReplicateEmbeddingModel()
result = model.invoke( result = model.invoke(
model='nateraw/jina-embeddings-v2-base-en', model="nateraw/jina-embeddings-v2-base-en",
credentials={ credentials={
'replicate_api_token': os.environ.get('REPLICATE_API_KEY'), "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),
'model_version': 'f8367a1c072ba2bc28af549d1faeacfe9b88b3f0e475add7a75091dac507f79e' "model_version": "f8367a1c072ba2bc28af549d1faeacfe9b88b3f0e475add7a75091dac507f79e",
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -137,15 +125,12 @@ def test_get_num_tokens():
model = ReplicateEmbeddingModel() model = ReplicateEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='nateraw/jina-embeddings-v2-base-en', model="nateraw/jina-embeddings-v2-base-en",
credentials={ credentials={
'replicate_api_token': os.environ.get('REPLICATE_API_KEY'), "replicate_api_token": os.environ.get("REPLICATE_API_KEY"),
'model_version': 'f8367a1c072ba2bc28af549d1faeacfe9b88b3f0e475add7a75091dac507f79e' "model_version": "f8367a1c072ba2bc28af549d1faeacfe9b88b3f0e475add7a75091dac507f79e",
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2

8
api/tests/integration_tests/model_runtime/sagemaker/test_provider.py
View file

@ -10,10 +10,6 @@ def test_validate_provider_credentials():
provider = SageMakerProvider() provider = SageMakerProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)

12
api/tests/integration_tests/model_runtime/sagemaker/test_rerank.py
View file

@ -12,11 +12,11 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='bge-m3-rerank-v2', model="bge-m3-rerank-v2",
credentials={ credentials={
"aws_region": os.getenv("AWS_REGION"), "aws_region": os.getenv("AWS_REGION"),
"aws_access_key": os.getenv("AWS_ACCESS_KEY"), "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
"aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY") "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
}, },
query="What is the capital of the United States?", query="What is the capital of the United States?",
docs=[ docs=[
@ -25,7 +25,7 @@ def test_validate_credentials():
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that " "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that "
"are a political division controlled by the United States. Its capital is Saipan.", "are a political division controlled by the United States. Its capital is Saipan.",
], ],
score_threshold=0.8 score_threshold=0.8,
) )
@ -33,11 +33,11 @@ def test_invoke_model():
model = SageMakerRerankModel() model = SageMakerRerankModel()
result = model.invoke( result = model.invoke(
model='bge-m3-rerank-v2', model="bge-m3-rerank-v2",
credentials={ credentials={
"aws_region": os.getenv("AWS_REGION"), "aws_region": os.getenv("AWS_REGION"),
"aws_access_key": os.getenv("AWS_ACCESS_KEY"), "aws_access_key": os.getenv("AWS_ACCESS_KEY"),
"aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY") "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),
}, },
query="What is the capital of the United States?", query="What is the capital of the United States?",
docs=[ docs=[
@ -46,7 +46,7 @@ def test_invoke_model():
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that " "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that "
"are a political division controlled by the United States. Its capital is Saipan.", "are a political division controlled by the United States. Its capital is Saipan.",
], ],
score_threshold=0.8 score_threshold=0.8,
) )
assert isinstance(result, RerankResult) assert isinstance(result, RerankResult)

32
api/tests/integration_tests/model_runtime/sagemaker/test_text_embedding.py
View file

@ -11,45 +11,23 @@ def test_validate_credentials():
model = SageMakerEmbeddingModel() model = SageMakerEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="bge-m3", credentials={})
model='bge-m3',
credentials={
}
)
model.validate_credentials( model.validate_credentials(model="bge-m3-embedding", credentials={})
model='bge-m3-embedding',
credentials={
}
)
def test_invoke_model(): def test_invoke_model():
model = SageMakerEmbeddingModel() model = SageMakerEmbeddingModel()
result = model.invoke( result = model.invoke(model="bge-m3-embedding", credentials={}, texts=["hello", "world"], user="abc-123")
model='bge-m3-embedding',
credentials={
},
texts=[
"hello",
"world"
],
user="abc-123"
)
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
assert len(result.embeddings) == 2 assert len(result.embeddings) == 2
def test_get_num_tokens(): def test_get_num_tokens():
model = SageMakerEmbeddingModel() model = SageMakerEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(model="bge-m3-embedding", credentials={}, texts=[])
model='bge-m3-embedding',
credentials={
},
texts=[
]
)
assert num_tokens == 0 assert num_tokens == 0

69
api/tests/integration_tests/model_runtime/siliconflow/test_llm.py
View file

@ -13,41 +13,22 @@ def test_validate_credentials():
model = SiliconflowLargeLanguageModel() model = SiliconflowLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="deepseek-ai/DeepSeek-V2-Chat", credentials={"api_key": "invalid_key"})
model='deepseek-ai/DeepSeek-V2-Chat',
credentials={
'api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(model="deepseek-ai/DeepSeek-V2-Chat", credentials={"api_key": os.environ.get("API_KEY")})
model='deepseek-ai/DeepSeek-V2-Chat',
credentials={
'api_key': os.environ.get('API_KEY')
}
)
def test_invoke_model(): def test_invoke_model():
model = SiliconflowLargeLanguageModel() model = SiliconflowLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='deepseek-ai/DeepSeek-V2-Chat', model="deepseek-ai/DeepSeek-V2-Chat",
credentials={ credentials={"api_key": os.environ.get("API_KEY")},
'api_key': os.environ.get('API_KEY') prompt_messages=[UserPromptMessage(content="Who are you?")],
}, model_parameters={"temperature": 0.5, "max_tokens": 10},
prompt_messages=[ stop=["How"],
UserPromptMessage(
content='Who are you?'
)
],
model_parameters={
'temperature': 0.5,
'max_tokens': 10
},
stop=['How'],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -58,22 +39,12 @@ def test_invoke_stream_model():
model = SiliconflowLargeLanguageModel() model = SiliconflowLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='deepseek-ai/DeepSeek-V2-Chat', model="deepseek-ai/DeepSeek-V2-Chat",
credentials={ credentials={"api_key": os.environ.get("API_KEY")},
'api_key': os.environ.get('API_KEY') prompt_messages=[UserPromptMessage(content="Hello World!")],
}, model_parameters={"temperature": 0.5, "max_tokens": 100, "seed": 1234},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={
'temperature': 0.5,
'max_tokens': 100,
'seed': 1234
},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -89,18 +60,14 @@ def test_get_num_tokens():
model = SiliconflowLargeLanguageModel() model = SiliconflowLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='deepseek-ai/DeepSeek-V2-Chat', model="deepseek-ai/DeepSeek-V2-Chat",
credentials={ credentials={"api_key": os.environ.get("API_KEY")},
'api_key': os.environ.get('API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert num_tokens == 12 assert num_tokens == 12

10
api/tests/integration_tests/model_runtime/siliconflow/test_provider.py
View file

@ -10,12 +10,6 @@ def test_validate_provider_credentials():
provider = SiliconflowProvider() provider = SiliconflowProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"api_key": os.environ.get("API_KEY")})
credentials={
'api_key': os.environ.get('API_KEY')
}
)

12
api/tests/integration_tests/model_runtime/siliconflow/test_rerank.py
View file

@ -13,9 +13,7 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model="BAAI/bge-reranker-v2-m3", model="BAAI/bge-reranker-v2-m3",
credentials={ credentials={"api_key": "invalid_key"},
"api_key": "invalid_key"
},
) )
model.validate_credentials( model.validate_credentials(
@ -30,17 +28,17 @@ def test_invoke_model():
model = SiliconflowRerankModel() model = SiliconflowRerankModel()
result = model.invoke( result = model.invoke(
model='BAAI/bge-reranker-v2-m3', model="BAAI/bge-reranker-v2-m3",
credentials={ credentials={
"api_key": os.environ.get("API_KEY"), "api_key": os.environ.get("API_KEY"),
}, },
query="Who is Kasumi?", query="Who is Kasumi?",
docs=[ docs=[
"Kasumi is a girl's name of Japanese origin meaning \"mist\".", 'Kasumi is a girl\'s name of Japanese origin meaning "mist".',
"Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ", "Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ",
"and she leads a team named PopiParty." "and she leads a team named PopiParty.",
], ],
score_threshold=0.8 score_threshold=0.8,
) )
assert isinstance(result, RerankResult) assert isinstance(result, RerankResult)

16
api/tests/integration_tests/model_runtime/siliconflow/test_speech2text.py
View file

@ -12,16 +12,12 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model="iic/SenseVoiceSmall", model="iic/SenseVoiceSmall",
credentials={ credentials={"api_key": "invalid_key"},
"api_key": "invalid_key"
},
) )
model.validate_credentials( model.validate_credentials(
model="iic/SenseVoiceSmall", model="iic/SenseVoiceSmall",
credentials={ credentials={"api_key": os.environ.get("API_KEY")},
"api_key": os.environ.get("API_KEY")
},
) )
@ -42,12 +38,8 @@ def test_invoke_model():
file = audio_file file = audio_file
result = model.invoke( result = model.invoke(
model="iic/SenseVoiceSmall", model="iic/SenseVoiceSmall", credentials={"api_key": os.environ.get("API_KEY")}, file=file
credentials={
"api_key": os.environ.get("API_KEY")
},
file=file
) )
assert isinstance(result, str) assert isinstance(result, str)
assert result == '1,2,3,4,5,6,7,8,9,10.' assert result == "1,2,3,4,5,6,7,8,9,10."

4
api/tests/integration_tests/model_runtime/siliconflow/test_text_embedding.py
View file

@ -15,9 +15,7 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model="BAAI/bge-large-zh-v1.5", model="BAAI/bge-large-zh-v1.5",
credentials={ credentials={"api_key": "invalid_key"},
"api_key": "invalid_key"
},
) )
model.validate_credentials( model.validate_credentials(

77
api/tests/integration_tests/model_runtime/spark/test_llm.py
View file

@ -13,20 +13,15 @@ def test_validate_credentials():
model = SparkLargeLanguageModel() model = SparkLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="spark-1.5", credentials={"app_id": "invalid_key"})
model='spark-1.5',
credentials={
'app_id': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(
model='spark-1.5', model="spark-1.5",
credentials={ credentials={
'app_id': os.environ.get('SPARK_APP_ID'), "app_id": os.environ.get("SPARK_APP_ID"),
'api_secret': os.environ.get('SPARK_API_SECRET'), "api_secret": os.environ.get("SPARK_API_SECRET"),
'api_key': os.environ.get('SPARK_API_KEY') "api_key": os.environ.get("SPARK_API_KEY"),
} },
) )
@ -34,24 +29,17 @@ def test_invoke_model():
model = SparkLargeLanguageModel() model = SparkLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='spark-1.5', model="spark-1.5",
credentials={ credentials={
'app_id': os.environ.get('SPARK_APP_ID'), "app_id": os.environ.get("SPARK_APP_ID"),
'api_secret': os.environ.get('SPARK_API_SECRET'), "api_secret": os.environ.get("SPARK_API_SECRET"),
'api_key': os.environ.get('SPARK_API_KEY') "api_key": os.environ.get("SPARK_API_KEY"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Who are you?")],
UserPromptMessage( model_parameters={"temperature": 0.5, "max_tokens": 10},
content='Who are you?' stop=["How"],
)
],
model_parameters={
'temperature': 0.5,
'max_tokens': 10
},
stop=['How'],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -62,23 +50,16 @@ def test_invoke_stream_model():
model = SparkLargeLanguageModel() model = SparkLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='spark-1.5', model="spark-1.5",
credentials={ credentials={
'app_id': os.environ.get('SPARK_APP_ID'), "app_id": os.environ.get("SPARK_APP_ID"),
'api_secret': os.environ.get('SPARK_API_SECRET'), "api_secret": os.environ.get("SPARK_API_SECRET"),
'api_key': os.environ.get('SPARK_API_KEY') "api_key": os.environ.get("SPARK_API_KEY"),
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={
'temperature': 0.5,
'max_tokens': 100
}, },
prompt_messages=[UserPromptMessage(content="Hello World!")],
model_parameters={"temperature": 0.5, "max_tokens": 100},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -94,20 +75,18 @@ def test_get_num_tokens():
model = SparkLargeLanguageModel() model = SparkLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='spark-1.5', model="spark-1.5",
credentials={ credentials={
'app_id': os.environ.get('SPARK_APP_ID'), "app_id": os.environ.get("SPARK_APP_ID"),
'api_secret': os.environ.get('SPARK_API_SECRET'), "api_secret": os.environ.get("SPARK_API_SECRET"),
'api_key': os.environ.get('SPARK_API_KEY') "api_key": os.environ.get("SPARK_API_KEY"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert num_tokens == 14 assert num_tokens == 14

10
api/tests/integration_tests/model_runtime/spark/test_provider.py
View file

@ -10,14 +10,12 @@ def test_validate_provider_credentials():
provider = SparkProvider() provider = SparkProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(
credentials={ credentials={
'app_id': os.environ.get('SPARK_APP_ID'), "app_id": os.environ.get("SPARK_APP_ID"),
'api_secret': os.environ.get('SPARK_API_SECRET'), "api_secret": os.environ.get("SPARK_API_SECRET"),
'api_key': os.environ.get('SPARK_API_KEY') "api_key": os.environ.get("SPARK_API_KEY"),
} }
) )

115
api/tests/integration_tests/model_runtime/stepfun/test_llm.py
View file

@ -21,40 +21,22 @@ def test_validate_credentials():
model = StepfunLargeLanguageModel() model = StepfunLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="step-1-8k", credentials={"api_key": "invalid_key"})
model='step-1-8k',
credentials={ model.validate_credentials(model="step-1-8k", credentials={"api_key": os.environ.get("STEPFUN_API_KEY")})
'api_key': 'invalid_key'
}
)
model.validate_credentials(
model='step-1-8k',
credentials={
'api_key': os.environ.get('STEPFUN_API_KEY')
}
)
def test_invoke_model(): def test_invoke_model():
model = StepfunLargeLanguageModel() model = StepfunLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='step-1-8k', model="step-1-8k",
credentials={ credentials={"api_key": os.environ.get("STEPFUN_API_KEY")},
'api_key': os.environ.get('STEPFUN_API_KEY') prompt_messages=[UserPromptMessage(content="Hello World!")],
}, model_parameters={"temperature": 0.9, "top_p": 0.7},
prompt_messages=[ stop=["Hi"],
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={
'temperature': 0.9,
'top_p': 0.7
},
stop=['Hi'],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -65,24 +47,17 @@ def test_invoke_stream_model():
model = StepfunLargeLanguageModel() model = StepfunLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='step-1-8k', model="step-1-8k",
credentials={ credentials={"api_key": os.environ.get("STEPFUN_API_KEY")},
'api_key': os.environ.get('STEPFUN_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={"temperature": 0.9, "top_p": 0.7},
'temperature': 0.9,
'top_p': 0.7
},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -98,10 +73,7 @@ def test_get_customizable_model_schema():
model = StepfunLargeLanguageModel() model = StepfunLargeLanguageModel()
schema = model.get_customizable_model_schema( schema = model.get_customizable_model_schema(
model='step-1-8k', model="step-1-8k", credentials={"api_key": os.environ.get("STEPFUN_API_KEY")}
credentials={
'api_key': os.environ.get('STEPFUN_API_KEY')
}
) )
assert isinstance(schema, AIModelEntity) assert isinstance(schema, AIModelEntity)
@ -110,65 +82,42 @@ def test_invoke_chat_model_with_tools():
model = StepfunLargeLanguageModel() model = StepfunLargeLanguageModel()
result = model.invoke( result = model.invoke(
model='step-1-8k', model="step-1-8k",
credentials={ credentials={"api_key": os.environ.get("STEPFUN_API_KEY")},
'api_key': os.environ.get('STEPFUN_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(
content="what's the weather today in Shanghai?", content="what's the weather today in Shanghai?",
) ),
], ],
model_parameters={ model_parameters={"temperature": 0.9, "max_tokens": 100},
'temperature': 0.9,
'max_tokens': 100
},
tools=[ tools=[
PromptMessageTool( PromptMessageTool(
name='get_weather', name="get_weather",
description='Determine weather in my location', description="Determine weather in my location",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {
"location": { "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"type": "string", "unit": {"type": "string", "enum": ["c", "f"]},
"description": "The city and state e.g. San Francisco, CA"
}, },
"unit": { "required": ["location"],
"type": "string",
"enum": [
"c",
"f"
]
}
}, },
"required": [
"location"
]
}
), ),
PromptMessageTool( PromptMessageTool(
name='get_stock_price', name="get_stock_price",
description='Get the current stock price', description="Get the current stock price",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {"symbol": {"type": "string", "description": "The stock symbol"}},
"symbol": { "required": ["symbol"],
"type": "string",
"description": "The stock symbol"
}
}, },
"required": [ ),
"symbol"
]
}
)
], ],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(result, LLMResult) assert isinstance(result, LLMResult)

29
api/tests/integration_tests/model_runtime/test_model_provider_factory.py
View file

@ -24,13 +24,8 @@ def test_get_models():
providers = factory.get_models( providers = factory.get_models(
model_type=ModelType.LLM, model_type=ModelType.LLM,
provider_configs=[ provider_configs=[
ProviderConfig( ProviderConfig(provider="openai", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")})
provider='openai', ],
credentials={
'openai_api_key': os.environ.get('OPENAI_API_KEY')
}
)
]
) )
logger.debug(providers) logger.debug(providers)
@ -44,29 +39,21 @@ def test_get_models():
assert provider_model.model_type == ModelType.LLM assert provider_model.model_type == ModelType.LLM
providers = factory.get_models( providers = factory.get_models(
provider='openai', provider="openai",
provider_configs=[ provider_configs=[
ProviderConfig( ProviderConfig(provider="openai", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")})
provider='openai', ],
credentials={
'openai_api_key': os.environ.get('OPENAI_API_KEY')
}
)
]
) )
assert len(providers) == 1 assert len(providers) == 1
assert isinstance(providers[0], SimpleProviderEntity) assert isinstance(providers[0], SimpleProviderEntity)
assert providers[0].provider == 'openai' assert providers[0].provider == "openai"
def test_provider_credentials_validate(): def test_provider_credentials_validate():
factory = ModelProviderFactory() factory = ModelProviderFactory()
factory.provider_credentials_validate( factory.provider_credentials_validate(
provider='openai', provider="openai", credentials={"openai_api_key": os.environ.get("OPENAI_API_KEY")}
credentials={
'openai_api_key': os.environ.get('OPENAI_API_KEY')
}
) )
@ -79,4 +66,4 @@ def test__get_model_provider_map():
logger.debug(model_provider.provider_instance) logger.debug(model_provider.provider_instance)
assert len(model_providers) >= 1 assert len(model_providers) >= 1
assert isinstance(model_providers['openai'], ModelProviderExtension) assert isinstance(model_providers["openai"], ModelProviderExtension)

74
api/tests/integration_tests/model_runtime/togetherai/test_llm.py
View file

@ -19,76 +19,61 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='mistralai/Mixtral-8x7B-Instruct-v0.1', model="mistralai/Mixtral-8x7B-Instruct-v0.1", credentials={"api_key": "invalid_key", "mode": "chat"}
credentials={
'api_key': 'invalid_key',
'mode': 'chat'
}
) )
model.validate_credentials( model.validate_credentials(
model='mistralai/Mixtral-8x7B-Instruct-v0.1', model="mistralai/Mixtral-8x7B-Instruct-v0.1",
credentials={ credentials={"api_key": os.environ.get("TOGETHER_API_KEY"), "mode": "chat"},
'api_key': os.environ.get('TOGETHER_API_KEY'),
'mode': 'chat'
}
) )
def test_invoke_model(): def test_invoke_model():
model = TogetherAILargeLanguageModel() model = TogetherAILargeLanguageModel()
response = model.invoke( response = model.invoke(
model='mistralai/Mixtral-8x7B-Instruct-v0.1', model="mistralai/Mixtral-8x7B-Instruct-v0.1",
credentials={ credentials={"api_key": os.environ.get("TOGETHER_API_KEY"), "mode": "completion"},
'api_key': os.environ.get('TOGETHER_API_KEY'),
'mode': 'completion'
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Who are you?"),
content='Who are you?'
)
], ],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
def test_invoke_stream_model(): def test_invoke_stream_model():
model = TogetherAILargeLanguageModel() model = TogetherAILargeLanguageModel()
response = model.invoke( response = model.invoke(
model='mistralai/Mixtral-8x7B-Instruct-v0.1', model="mistralai/Mixtral-8x7B-Instruct-v0.1",
credentials={ credentials={"api_key": os.environ.get("TOGETHER_API_KEY"), "mode": "chat"},
'api_key': os.environ.get('TOGETHER_API_KEY'),
'mode': 'chat'
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Who are you?"),
content='Who are you?'
)
], ],
model_parameters={ model_parameters={
'temperature': 1.0, "temperature": 1.0,
'top_k': 2, "top_k": 2,
'top_p': 0.5, "top_p": 0.5,
}, },
stop=['How'], stop=["How"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -98,22 +83,21 @@ def test_invoke_stream_model():
assert isinstance(chunk.delta, LLMResultChunkDelta) assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
def test_get_num_tokens(): def test_get_num_tokens():
model = TogetherAILargeLanguageModel() model = TogetherAILargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='mistralai/Mixtral-8x7B-Instruct-v0.1', model="mistralai/Mixtral-8x7B-Instruct-v0.1",
credentials={ credentials={
'api_key': os.environ.get('TOGETHER_API_KEY'), "api_key": os.environ.get("TOGETHER_API_KEY"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert isinstance(num_tokens, int) assert isinstance(num_tokens, int)

67
api/tests/integration_tests/model_runtime/tongyi/test_llm.py
View file

@ -13,18 +13,10 @@ def test_validate_credentials():
model = TongyiLargeLanguageModel() model = TongyiLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="qwen-turbo", credentials={"dashscope_api_key": "invalid_key"})
model='qwen-turbo',
credentials={
'dashscope_api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(
model='qwen-turbo', model="qwen-turbo", credentials={"dashscope_api_key": os.environ.get("TONGYI_DASHSCOPE_API_KEY")}
credentials={
'dashscope_api_key': os.environ.get('TONGYI_DASHSCOPE_API_KEY')
}
) )
@ -32,22 +24,13 @@ def test_invoke_model():
model = TongyiLargeLanguageModel() model = TongyiLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='qwen-turbo', model="qwen-turbo",
credentials={ credentials={"dashscope_api_key": os.environ.get("TONGYI_DASHSCOPE_API_KEY")},
'dashscope_api_key': os.environ.get('TONGYI_DASHSCOPE_API_KEY') prompt_messages=[UserPromptMessage(content="Who are you?")],
}, model_parameters={"temperature": 0.5, "max_tokens": 10},
prompt_messages=[ stop=["How"],
UserPromptMessage(
content='Who are you?'
)
],
model_parameters={
'temperature': 0.5,
'max_tokens': 10
},
stop=['How'],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -58,22 +41,12 @@ def test_invoke_stream_model():
model = TongyiLargeLanguageModel() model = TongyiLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='qwen-turbo', model="qwen-turbo",
credentials={ credentials={"dashscope_api_key": os.environ.get("TONGYI_DASHSCOPE_API_KEY")},
'dashscope_api_key': os.environ.get('TONGYI_DASHSCOPE_API_KEY') prompt_messages=[UserPromptMessage(content="Hello World!")],
}, model_parameters={"temperature": 0.5, "max_tokens": 100, "seed": 1234},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={
'temperature': 0.5,
'max_tokens': 100,
'seed': 1234
},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -89,18 +62,14 @@ def test_get_num_tokens():
model = TongyiLargeLanguageModel() model = TongyiLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='qwen-turbo', model="qwen-turbo",
credentials={ credentials={"dashscope_api_key": os.environ.get("TONGYI_DASHSCOPE_API_KEY")},
'dashscope_api_key': os.environ.get('TONGYI_DASHSCOPE_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert num_tokens == 12 assert num_tokens == 12

8
api/tests/integration_tests/model_runtime/tongyi/test_provider.py
View file

@ -10,12 +10,8 @@ def test_validate_provider_credentials():
provider = TongyiProvider() provider = TongyiProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(
credentials={ credentials={"dashscope_api_key": os.environ.get("TONGYI_DASHSCOPE_API_KEY")}
'dashscope_api_key': os.environ.get('TONGYI_DASHSCOPE_API_KEY')
}
) )

16
api/tests/integration_tests/model_runtime/tongyi/test_response_format.py
View file

@ -39,21 +39,17 @@ def invoke_model_with_json_response(model_name="qwen-max-0403"):
response = model.invoke( response = model.invoke(
model=model_name, model=model_name,
credentials={ credentials={"dashscope_api_key": os.environ.get("TONGYI_DASHSCOPE_API_KEY")},
'dashscope_api_key': os.environ.get('TONGYI_DASHSCOPE_API_KEY')
},
prompt_messages=[ prompt_messages=[
UserPromptMessage( UserPromptMessage(content='output json data with format `{"data": "test", "code": 200, "msg": "success"}')
content='output json data with format `{"data": "test", "code": 200, "msg": "success"}'
)
], ],
model_parameters={ model_parameters={
'temperature': 0.5, "temperature": 0.5,
'max_tokens': 50, "max_tokens": 50,
'response_format': 'JSON', "response_format": "JSON",
}, },
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
print("=====================================") print("=====================================")
print(response) print(response)

171
api/tests/integration_tests/model_runtime/upstage/test_llm.py
View file

@ -26,151 +26,113 @@ def test_predefined_models():
assert len(model_schemas) >= 1 assert len(model_schemas) >= 1
assert isinstance(model_schemas[0], AIModelEntity) assert isinstance(model_schemas[0], AIModelEntity)
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_validate_credentials_for_chat_model(setup_openai_mock): def test_validate_credentials_for_chat_model(setup_openai_mock):
model = UpstageLargeLanguageModel() model = UpstageLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
# model name to gpt-3.5-turbo because of mocking # model name to gpt-3.5-turbo because of mocking
model.validate_credentials( model.validate_credentials(model="gpt-3.5-turbo", credentials={"upstage_api_key": "invalid_key"})
model='gpt-3.5-turbo',
credentials={
'upstage_api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(
model='solar-1-mini-chat', model="solar-1-mini-chat", credentials={"upstage_api_key": os.environ.get("UPSTAGE_API_KEY")}
credentials={
'upstage_api_key': os.environ.get('UPSTAGE_API_KEY')
}
) )
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_invoke_chat_model(setup_openai_mock): def test_invoke_chat_model(setup_openai_mock):
model = UpstageLargeLanguageModel() model = UpstageLargeLanguageModel()
result = model.invoke( result = model.invoke(
model='solar-1-mini-chat', model="solar-1-mini-chat",
credentials={ credentials={"upstage_api_key": os.environ.get("UPSTAGE_API_KEY")},
'upstage_api_key': os.environ.get('UPSTAGE_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={
'temperature': 0.0, "temperature": 0.0,
'top_p': 1.0, "top_p": 1.0,
'presence_penalty': 0.0, "presence_penalty": 0.0,
'frequency_penalty': 0.0, "frequency_penalty": 0.0,
'max_tokens': 10 "max_tokens": 10,
}, },
stop=['How'], stop=["How"],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(result, LLMResult) assert isinstance(result, LLMResult)
assert len(result.message.content) > 0 assert len(result.message.content) > 0
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_invoke_chat_model_with_tools(setup_openai_mock): def test_invoke_chat_model_with_tools(setup_openai_mock):
model = UpstageLargeLanguageModel() model = UpstageLargeLanguageModel()
result = model.invoke( result = model.invoke(
model='solar-1-mini-chat', model="solar-1-mini-chat",
credentials={ credentials={"upstage_api_key": os.environ.get("UPSTAGE_API_KEY")},
'upstage_api_key': os.environ.get('UPSTAGE_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(
content="what's the weather today in London?", content="what's the weather today in London?",
) ),
], ],
model_parameters={ model_parameters={"temperature": 0.0, "max_tokens": 100},
'temperature': 0.0,
'max_tokens': 100
},
tools=[ tools=[
PromptMessageTool( PromptMessageTool(
name='get_weather', name="get_weather",
description='Determine weather in my location', description="Determine weather in my location",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {
"location": { "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"type": "string", "unit": {"type": "string", "enum": ["c", "f"]},
"description": "The city and state e.g. San Francisco, CA"
}, },
"unit": { "required": ["location"],
"type": "string",
"enum": [
"c",
"f"
]
}
}, },
"required": [
"location"
]
}
), ),
PromptMessageTool( PromptMessageTool(
name='get_stock_price', name="get_stock_price",
description='Get the current stock price', description="Get the current stock price",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {"symbol": {"type": "string", "description": "The stock symbol"}},
"symbol": { "required": ["symbol"],
"type": "string",
"description": "The stock symbol"
}
}, },
"required": [ ),
"symbol"
]
}
)
], ],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(result, LLMResult) assert isinstance(result, LLMResult)
assert isinstance(result.message, AssistantPromptMessage) assert isinstance(result.message, AssistantPromptMessage)
assert len(result.message.tool_calls) > 0 assert len(result.message.tool_calls) > 0
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_invoke_stream_chat_model(setup_openai_mock): def test_invoke_stream_chat_model(setup_openai_mock):
model = UpstageLargeLanguageModel() model = UpstageLargeLanguageModel()
result = model.invoke( result = model.invoke(
model='solar-1-mini-chat', model="solar-1-mini-chat",
credentials={ credentials={"upstage_api_key": os.environ.get("UPSTAGE_API_KEY")},
'upstage_api_key': os.environ.get('UPSTAGE_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={"temperature": 0.0, "max_tokens": 100},
'temperature': 0.0,
'max_tokens': 100
},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(result, Generator) assert isinstance(result, Generator)
@ -189,57 +151,36 @@ def test_get_num_tokens():
model = UpstageLargeLanguageModel() model = UpstageLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='solar-1-mini-chat', model="solar-1-mini-chat",
credentials={ credentials={"upstage_api_key": os.environ.get("UPSTAGE_API_KEY")},
'upstage_api_key': os.environ.get('UPSTAGE_API_KEY') prompt_messages=[UserPromptMessage(content="Hello World!")],
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
]
) )
assert num_tokens == 13 assert num_tokens == 13
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='solar-1-mini-chat', model="solar-1-mini-chat",
credentials={ credentials={"upstage_api_key": os.environ.get("UPSTAGE_API_KEY")},
'upstage_api_key': os.environ.get('UPSTAGE_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
tools=[ tools=[
PromptMessageTool( PromptMessageTool(
name='get_weather', name="get_weather",
description='Determine weather in my location', description="Determine weather in my location",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {
"location": { "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"type": "string", "unit": {"type": "string", "enum": ["c", "f"]},
"description": "The city and state e.g. San Francisco, CA"
}, },
"unit": { "required": ["location"],
"type": "string",
"enum": [
"c",
"f"
]
}
}, },
"required": [
"location"
]
}
), ),
] ],
) )
assert num_tokens == 106 assert num_tokens == 106

12
api/tests/integration_tests/model_runtime/upstage/test_provider.py
View file

@ -7,17 +7,11 @@ from core.model_runtime.model_providers.upstage.upstage import UpstageProvider
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock
@pytest.mark.parametrize('setup_openai_mock', [['chat']], indirect=True) @pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_validate_provider_credentials(setup_openai_mock): def test_validate_provider_credentials(setup_openai_mock):
provider = UpstageProvider() provider = UpstageProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"upstage_api_key": os.environ.get("UPSTAGE_API_KEY")})
credentials={
'upstage_api_key': os.environ.get('UPSTAGE_API_KEY')
}
)

37
api/tests/integration_tests/model_runtime/upstage/test_text_embedding.py
View file

@ -8,41 +8,31 @@ from core.model_runtime.model_providers.upstage.text_embedding.text_embedding im
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock
@pytest.mark.parametrize('setup_openai_mock', [['text_embedding']], indirect=True) @pytest.mark.parametrize("setup_openai_mock", [["text_embedding"]], indirect=True)
def test_validate_credentials(setup_openai_mock): def test_validate_credentials(setup_openai_mock):
model = UpstageTextEmbeddingModel() model = UpstageTextEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='solar-embedding-1-large-passage', model="solar-embedding-1-large-passage", credentials={"upstage_api_key": "invalid_key"}
credentials={
'upstage_api_key': 'invalid_key'
}
) )
model.validate_credentials( model.validate_credentials(
model='solar-embedding-1-large-passage', model="solar-embedding-1-large-passage", credentials={"upstage_api_key": os.environ.get("UPSTAGE_API_KEY")}
credentials={
'upstage_api_key': os.environ.get('UPSTAGE_API_KEY')
}
) )
@pytest.mark.parametrize('setup_openai_mock', [['text_embedding']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock", [["text_embedding"]], indirect=True)
def test_invoke_model(setup_openai_mock): def test_invoke_model(setup_openai_mock):
model = UpstageTextEmbeddingModel() model = UpstageTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='solar-embedding-1-large-passage', model="solar-embedding-1-large-passage",
credentials={ credentials={
'upstage_api_key': os.environ.get('UPSTAGE_API_KEY'), "upstage_api_key": os.environ.get("UPSTAGE_API_KEY"),
}, },
texts=[ texts=["hello", "world", " ".join(["long_text"] * 100), " ".join(["another_long_text"] * 100)],
"hello", user="abc-123",
"world",
" ".join(["long_text"] * 100),
" ".join(["another_long_text"] * 100)
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -54,14 +44,11 @@ def test_get_num_tokens():
model = UpstageTextEmbeddingModel() model = UpstageTextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='solar-embedding-1-large-passage', model="solar-embedding-1-large-passage",
credentials={ credentials={
'upstage_api_key': os.environ.get('UPSTAGE_API_KEY'), "upstage_api_key": os.environ.get("UPSTAGE_API_KEY"),
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 5 assert num_tokens == 5

70
api/tests/integration_tests/model_runtime/volcengine_maas/test_embedding.py
View file

@ -14,26 +14,26 @@ def test_validate_credentials():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'api_endpoint_host': 'maas-api.ml-platform-cn-beijing.volces.com', "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",
'volc_region': 'cn-beijing', "volc_region": "cn-beijing",
'volc_access_key_id': 'INVALID', "volc_access_key_id": "INVALID",
'volc_secret_access_key': 'INVALID', "volc_secret_access_key": "INVALID",
'endpoint_id': 'INVALID', "endpoint_id": "INVALID",
'base_model_name': 'Doubao-embedding', "base_model_name": "Doubao-embedding",
} },
) )
model.validate_credentials( model.validate_credentials(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'api_endpoint_host': 'maas-api.ml-platform-cn-beijing.volces.com', "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",
'volc_region': 'cn-beijing', "volc_region": "cn-beijing",
'volc_access_key_id': os.environ.get('VOLC_API_KEY'), "volc_access_key_id": os.environ.get("VOLC_API_KEY"),
'volc_secret_access_key': os.environ.get('VOLC_SECRET_KEY'), "volc_secret_access_key": os.environ.get("VOLC_SECRET_KEY"),
'endpoint_id': os.environ.get('VOLC_EMBEDDING_ENDPOINT_ID'), "endpoint_id": os.environ.get("VOLC_EMBEDDING_ENDPOINT_ID"),
'base_model_name': 'Doubao-embedding', "base_model_name": "Doubao-embedding",
}, },
) )
@ -42,20 +42,17 @@ def test_invoke_model():
model = VolcengineMaaSTextEmbeddingModel() model = VolcengineMaaSTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'api_endpoint_host': 'maas-api.ml-platform-cn-beijing.volces.com', "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",
'volc_region': 'cn-beijing', "volc_region": "cn-beijing",
'volc_access_key_id': os.environ.get('VOLC_API_KEY'), "volc_access_key_id": os.environ.get("VOLC_API_KEY"),
'volc_secret_access_key': os.environ.get('VOLC_SECRET_KEY'), "volc_secret_access_key": os.environ.get("VOLC_SECRET_KEY"),
'endpoint_id': os.environ.get('VOLC_EMBEDDING_ENDPOINT_ID'), "endpoint_id": os.environ.get("VOLC_EMBEDDING_ENDPOINT_ID"),
'base_model_name': 'Doubao-embedding', "base_model_name": "Doubao-embedding",
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -67,19 +64,16 @@ def test_get_num_tokens():
model = VolcengineMaaSTextEmbeddingModel() model = VolcengineMaaSTextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'api_endpoint_host': 'maas-api.ml-platform-cn-beijing.volces.com', "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",
'volc_region': 'cn-beijing', "volc_region": "cn-beijing",
'volc_access_key_id': os.environ.get('VOLC_API_KEY'), "volc_access_key_id": os.environ.get("VOLC_API_KEY"),
'volc_secret_access_key': os.environ.get('VOLC_SECRET_KEY'), "volc_secret_access_key": os.environ.get("VOLC_SECRET_KEY"),
'endpoint_id': os.environ.get('VOLC_EMBEDDING_ENDPOINT_ID'), "endpoint_id": os.environ.get("VOLC_EMBEDDING_ENDPOINT_ID"),
'base_model_name': 'Doubao-embedding', "base_model_name": "Doubao-embedding",
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2

113
api/tests/integration_tests/model_runtime/volcengine_maas/test_llm.py
View file

@ -14,25 +14,25 @@ def test_validate_credentials_for_chat_model():
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'api_endpoint_host': 'maas-api.ml-platform-cn-beijing.volces.com', "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",
'volc_region': 'cn-beijing', "volc_region": "cn-beijing",
'volc_access_key_id': 'INVALID', "volc_access_key_id": "INVALID",
'volc_secret_access_key': 'INVALID', "volc_secret_access_key": "INVALID",
'endpoint_id': 'INVALID', "endpoint_id": "INVALID",
} },
) )
model.validate_credentials( model.validate_credentials(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'api_endpoint_host': 'maas-api.ml-platform-cn-beijing.volces.com', "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",
'volc_region': 'cn-beijing', "volc_region": "cn-beijing",
'volc_access_key_id': os.environ.get('VOLC_API_KEY'), "volc_access_key_id": os.environ.get("VOLC_API_KEY"),
'volc_secret_access_key': os.environ.get('VOLC_SECRET_KEY'), "volc_secret_access_key": os.environ.get("VOLC_SECRET_KEY"),
'endpoint_id': os.environ.get('VOLC_MODEL_ENDPOINT_ID'), "endpoint_id": os.environ.get("VOLC_MODEL_ENDPOINT_ID"),
} },
) )
@ -40,28 +40,24 @@ def test_invoke_model():
model = VolcengineMaaSLargeLanguageModel() model = VolcengineMaaSLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'api_endpoint_host': 'maas-api.ml-platform-cn-beijing.volces.com', "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",
'volc_region': 'cn-beijing', "volc_region": "cn-beijing",
'volc_access_key_id': os.environ.get('VOLC_API_KEY'), "volc_access_key_id": os.environ.get("VOLC_API_KEY"),
'volc_secret_access_key': os.environ.get('VOLC_SECRET_KEY'), "volc_secret_access_key": os.environ.get("VOLC_SECRET_KEY"),
'endpoint_id': os.environ.get('VOLC_MODEL_ENDPOINT_ID'), "endpoint_id": os.environ.get("VOLC_MODEL_ENDPOINT_ID"),
'base_model_name': 'Skylark2-pro-4k', "base_model_name": "Skylark2-pro-4k",
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
'top_k': 1, "top_k": 1,
}, },
stop=['you'], stop=["you"],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -73,28 +69,24 @@ def test_invoke_stream_model():
model = VolcengineMaaSLargeLanguageModel() model = VolcengineMaaSLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'api_endpoint_host': 'maas-api.ml-platform-cn-beijing.volces.com', "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",
'volc_region': 'cn-beijing', "volc_region": "cn-beijing",
'volc_access_key_id': os.environ.get('VOLC_API_KEY'), "volc_access_key_id": os.environ.get("VOLC_API_KEY"),
'volc_secret_access_key': os.environ.get('VOLC_SECRET_KEY'), "volc_secret_access_key": os.environ.get("VOLC_SECRET_KEY"),
'endpoint_id': os.environ.get('VOLC_MODEL_ENDPOINT_ID'), "endpoint_id": os.environ.get("VOLC_MODEL_ENDPOINT_ID"),
'base_model_name': 'Skylark2-pro-4k', "base_model_name": "Skylark2-pro-4k",
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
'top_k': 1, "top_k": 1,
}, },
stop=['you'], stop=["you"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -102,29 +94,24 @@ def test_invoke_stream_model():
assert isinstance(chunk, LLMResultChunk) assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta) assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len( assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
def test_get_num_tokens(): def test_get_num_tokens():
model = VolcengineMaaSLargeLanguageModel() model = VolcengineMaaSLargeLanguageModel()
response = model.get_num_tokens( response = model.get_num_tokens(
model='NOT IMPORTANT', model="NOT IMPORTANT",
credentials={ credentials={
'api_endpoint_host': 'maas-api.ml-platform-cn-beijing.volces.com', "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",
'volc_region': 'cn-beijing', "volc_region": "cn-beijing",
'volc_access_key_id': os.environ.get('VOLC_API_KEY'), "volc_access_key_id": os.environ.get("VOLC_API_KEY"),
'volc_secret_access_key': os.environ.get('VOLC_SECRET_KEY'), "volc_secret_access_key": os.environ.get("VOLC_SECRET_KEY"),
'endpoint_id': os.environ.get('VOLC_MODEL_ENDPOINT_ID'), "endpoint_id": os.environ.get("VOLC_MODEL_ENDPOINT_ID"),
'base_model_name': 'Skylark2-pro-4k', "base_model_name": "Skylark2-pro-4k",
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="Hello World!")],
UserPromptMessage( tools=[],
content='Hello World!'
)
],
tools=[]
) )
assert isinstance(response, int) assert isinstance(response, int)

44
api/tests/integration_tests/model_runtime/wenxin/test_embedding.py
View file

@ -10,13 +10,10 @@ def test_invoke_embedding_v1():
model = WenxinTextEmbeddingModel() model = WenxinTextEmbeddingModel()
response = model.invoke( response = model.invoke(
model='embedding-v1', model="embedding-v1",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), texts=["hello", "你好", "xxxxx"],
'secret_key': os.environ.get('WENXIN_SECRET_KEY') user="abc-123",
},
texts=['hello', '你好', 'xxxxx'],
user="abc-123"
) )
assert isinstance(response, TextEmbeddingResult) assert isinstance(response, TextEmbeddingResult)
@ -29,13 +26,10 @@ def test_invoke_embedding_bge_large_en():
model = WenxinTextEmbeddingModel() model = WenxinTextEmbeddingModel()
response = model.invoke( response = model.invoke(
model='bge-large-en', model="bge-large-en",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), texts=["hello", "你好", "xxxxx"],
'secret_key': os.environ.get('WENXIN_SECRET_KEY') user="abc-123",
},
texts=['hello', '你好', 'xxxxx'],
user="abc-123"
) )
assert isinstance(response, TextEmbeddingResult) assert isinstance(response, TextEmbeddingResult)
@ -48,13 +42,10 @@ def test_invoke_embedding_bge_large_zh():
model = WenxinTextEmbeddingModel() model = WenxinTextEmbeddingModel()
response = model.invoke( response = model.invoke(
model='bge-large-zh', model="bge-large-zh",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), texts=["hello", "你好", "xxxxx"],
'secret_key': os.environ.get('WENXIN_SECRET_KEY') user="abc-123",
},
texts=['hello', '你好', 'xxxxx'],
user="abc-123"
) )
assert isinstance(response, TextEmbeddingResult) assert isinstance(response, TextEmbeddingResult)
@ -67,13 +58,10 @@ def test_invoke_embedding_tao_8k():
model = WenxinTextEmbeddingModel() model = WenxinTextEmbeddingModel()
response = model.invoke( response = model.invoke(
model='tao-8k', model="tao-8k",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), texts=["hello", "你好", "xxxxx"],
'secret_key': os.environ.get('WENXIN_SECRET_KEY') user="abc-123",
},
texts=['hello', '你好', 'xxxxx'],
user="abc-123"
) )
assert isinstance(response, TextEmbeddingResult) assert isinstance(response, TextEmbeddingResult)

193
api/tests/integration_tests/model_runtime/wenxin/test_llm.py
View file

@ -17,161 +17,125 @@ def test_predefined_models():
assert len(model_schemas) >= 1 assert len(model_schemas) >= 1
assert isinstance(model_schemas[0], AIModelEntity) assert isinstance(model_schemas[0], AIModelEntity)
def test_validate_credentials_for_chat_model(): def test_validate_credentials_for_chat_model():
sleep(3) sleep(3)
model = ErnieBotLargeLanguageModel() model = ErnieBotLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='ernie-bot', model="ernie-bot", credentials={"api_key": "invalid_key", "secret_key": "invalid_key"}
credentials={
'api_key': 'invalid_key',
'secret_key': 'invalid_key'
}
) )
model.validate_credentials( model.validate_credentials(
model='ernie-bot', model="ernie-bot",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'),
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
}
) )
def test_invoke_model_ernie_bot(): def test_invoke_model_ernie_bot():
sleep(3) sleep(3)
model = ErnieBotLargeLanguageModel() model = ErnieBotLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='ernie-bot', model="ernie-bot",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), prompt_messages=[UserPromptMessage(content="Hello World!")],
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
def test_invoke_model_ernie_bot_turbo(): def test_invoke_model_ernie_bot_turbo():
sleep(3) sleep(3)
model = ErnieBotLargeLanguageModel() model = ErnieBotLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='ernie-bot-turbo', model="ernie-bot-turbo",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), prompt_messages=[UserPromptMessage(content="Hello World!")],
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
def test_invoke_model_ernie_8k(): def test_invoke_model_ernie_8k():
sleep(3) sleep(3)
model = ErnieBotLargeLanguageModel() model = ErnieBotLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='ernie-bot-8k', model="ernie-bot-8k",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), prompt_messages=[UserPromptMessage(content="Hello World!")],
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
def test_invoke_model_ernie_bot_4(): def test_invoke_model_ernie_bot_4():
sleep(3) sleep(3)
model = ErnieBotLargeLanguageModel() model = ErnieBotLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='ernie-bot-4', model="ernie-bot-4",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), prompt_messages=[UserPromptMessage(content="Hello World!")],
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
def test_invoke_stream_model(): def test_invoke_stream_model():
sleep(3) sleep(3)
model = ErnieBotLargeLanguageModel() model = ErnieBotLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='ernie-3.5-8k', model="ernie-3.5-8k",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), prompt_messages=[UserPromptMessage(content="Hello World!")],
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -181,63 +145,48 @@ def test_invoke_stream_model():
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
def test_invoke_model_with_system(): def test_invoke_model_with_system():
sleep(3) sleep(3)
model = ErnieBotLargeLanguageModel() model = ErnieBotLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='ernie-bot', model="ernie-bot",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), prompt_messages=[SystemPromptMessage(content="你是Kasumi"), UserPromptMessage(content="你是谁？")],
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
},
prompt_messages=[
SystemPromptMessage(
content='你是Kasumi'
),
UserPromptMessage(
content='你是谁？'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert 'kasumi' in response.message.content.lower() assert "kasumi" in response.message.content.lower()
def test_invoke_with_search(): def test_invoke_with_search():
sleep(3) sleep(3)
model = ErnieBotLargeLanguageModel() model = ErnieBotLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='ernie-bot', model="ernie-bot",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), prompt_messages=[UserPromptMessage(content="北京今天的天气怎么样")],
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
},
prompt_messages=[
UserPromptMessage(
content='北京今天的天气怎么样'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
'disable_search': True, "disable_search": True,
}, },
stop=[], stop=[],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
total_message = '' total_message = ""
for chunk in response: for chunk in response:
assert isinstance(chunk, LLMResultChunk) assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta) assert isinstance(chunk.delta, LLMResultChunkDelta)
@ -247,24 +196,18 @@ def test_invoke_with_search():
assert len(chunk.delta.message.content) > 0 if not chunk.delta.finish_reason else True assert len(chunk.delta.message.content) > 0 if not chunk.delta.finish_reason else True
# there should be 对不起、我不能、不支持…… # there should be 对不起、我不能、不支持……
assert ('不' in total_message or '抱歉' in total_message or '无法' in total_message) assert "不" in total_message or "抱歉" in total_message or "无法" in total_message
def test_get_num_tokens(): def test_get_num_tokens():
sleep(3) sleep(3)
model = ErnieBotLargeLanguageModel() model = ErnieBotLargeLanguageModel()
response = model.get_num_tokens( response = model.get_num_tokens(
model='ernie-bot', model="ernie-bot",
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
'api_key': os.environ.get('WENXIN_API_KEY'), prompt_messages=[UserPromptMessage(content="Hello World!")],
'secret_key': os.environ.get('WENXIN_SECRET_KEY') tools=[],
},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
tools=[]
) )
assert isinstance(response, int) assert isinstance(response, int)

12
api/tests/integration_tests/model_runtime/wenxin/test_provider.py
View file

@ -10,16 +10,8 @@ def test_validate_provider_credentials():
provider = WenxinProvider() provider = WenxinProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"api_key": "hahahaha", "secret_key": "hahahaha"})
credentials={
'api_key': 'hahahaha',
'secret_key': 'hahahaha'
}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(
credentials={ credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")}
'api_key': os.environ.get('WENXIN_API_KEY'),
'secret_key': os.environ.get('WENXIN_SECRET_KEY')
}
) )

46
api/tests/integration_tests/model_runtime/xinference/test_embeddings.py
View file

@ -8,61 +8,57 @@ from core.model_runtime.model_providers.xinference.text_embedding.text_embedding
from tests.integration_tests.model_runtime.__mock.xinference import MOCK, setup_xinference_mock from tests.integration_tests.model_runtime.__mock.xinference import MOCK, setup_xinference_mock
@pytest.mark.parametrize('setup_xinference_mock', [['none']], indirect=True) @pytest.mark.parametrize("setup_xinference_mock", [["none"]], indirect=True)
def test_validate_credentials(setup_xinference_mock): def test_validate_credentials(setup_xinference_mock):
model = XinferenceTextEmbeddingModel() model = XinferenceTextEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='bge-base-en', model="bge-base-en",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': 'www ' + os.environ.get('XINFERENCE_EMBEDDINGS_MODEL_UID') "model_uid": "www " + os.environ.get("XINFERENCE_EMBEDDINGS_MODEL_UID"),
} },
) )
model.validate_credentials( model.validate_credentials(
model='bge-base-en', model="bge-base-en",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_EMBEDDINGS_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_EMBEDDINGS_MODEL_UID"),
} },
) )
@pytest.mark.parametrize('setup_xinference_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_xinference_mock", [["none"]], indirect=True)
def test_invoke_model(setup_xinference_mock): def test_invoke_model(setup_xinference_mock):
model = XinferenceTextEmbeddingModel() model = XinferenceTextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='bge-base-en', model="bge-base-en",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_EMBEDDINGS_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_EMBEDDINGS_MODEL_UID"),
}, },
texts=[ texts=["hello", "world"],
"hello", user="abc-123",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
assert len(result.embeddings) == 2 assert len(result.embeddings) == 2
assert result.usage.total_tokens > 0 assert result.usage.total_tokens > 0
def test_get_num_tokens(): def test_get_num_tokens():
model = XinferenceTextEmbeddingModel() model = XinferenceTextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='bge-base-en', model="bge-base-en",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_EMBEDDINGS_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_EMBEDDINGS_MODEL_UID"),
}, },
texts=[ texts=["hello", "world"],
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2

195
api/tests/integration_tests/model_runtime/xinference/test_llm.py
View file

@ -20,92 +20,84 @@ from tests.integration_tests.model_runtime.__mock.openai import setup_openai_moc
from tests.integration_tests.model_runtime.__mock.xinference import setup_xinference_mock from tests.integration_tests.model_runtime.__mock.xinference import setup_xinference_mock
@pytest.mark.parametrize('setup_openai_mock, setup_xinference_mock', [['chat', 'none']], indirect=True) @pytest.mark.parametrize("setup_openai_mock, setup_xinference_mock", [["chat", "none"]], indirect=True)
def test_validate_credentials_for_chat_model(setup_openai_mock, setup_xinference_mock): def test_validate_credentials_for_chat_model(setup_openai_mock, setup_xinference_mock):
model = XinferenceAILargeLanguageModel() model = XinferenceAILargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='ChatGLM3', model="ChatGLM3",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': 'www ' + os.environ.get('XINFERENCE_CHAT_MODEL_UID') "model_uid": "www " + os.environ.get("XINFERENCE_CHAT_MODEL_UID"),
} },
) )
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="aaaaa", credentials={"server_url": "", "model_uid": ""})
model='aaaaa',
credentials={
'server_url': '',
'model_uid': ''
}
)
model.validate_credentials( model.validate_credentials(
model='ChatGLM3', model="ChatGLM3",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_CHAT_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_CHAT_MODEL_UID"),
} },
) )
@pytest.mark.parametrize('setup_openai_mock, setup_xinference_mock', [['chat', 'none']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock, setup_xinference_mock", [["chat", "none"]], indirect=True)
def test_invoke_chat_model(setup_openai_mock, setup_xinference_mock): def test_invoke_chat_model(setup_openai_mock, setup_xinference_mock):
model = XinferenceAILargeLanguageModel() model = XinferenceAILargeLanguageModel()
response = model.invoke( response = model.invoke(
model='ChatGLM3', model="ChatGLM3",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_CHAT_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_CHAT_MODEL_UID"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
@pytest.mark.parametrize('setup_openai_mock, setup_xinference_mock', [['chat', 'none']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock, setup_xinference_mock", [["chat", "none"]], indirect=True)
def test_invoke_stream_chat_model(setup_openai_mock, setup_xinference_mock): def test_invoke_stream_chat_model(setup_openai_mock, setup_xinference_mock):
model = XinferenceAILargeLanguageModel() model = XinferenceAILargeLanguageModel()
response = model.invoke( response = model.invoke(
model='ChatGLM3', model="ChatGLM3",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_CHAT_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_CHAT_MODEL_UID"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -114,6 +106,8 @@ def test_invoke_stream_chat_model(setup_openai_mock, setup_xinference_mock):
assert isinstance(chunk.delta, LLMResultChunkDelta) assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
""" """
Funtion calling of xinference does not support stream mode currently Funtion calling of xinference does not support stream mode currently
""" """
@ -241,86 +235,75 @@ def test_invoke_stream_chat_model(setup_openai_mock, setup_xinference_mock):
# assert response.usage.total_tokens > 0 # assert response.usage.total_tokens > 0
# assert response.message.tool_calls[0].function.name == 'get_current_weather' # assert response.message.tool_calls[0].function.name == 'get_current_weather'
@pytest.mark.parametrize('setup_openai_mock, setup_xinference_mock', [['completion', 'none']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock, setup_xinference_mock", [["completion", "none"]], indirect=True)
def test_validate_credentials_for_generation_model(setup_openai_mock, setup_xinference_mock): def test_validate_credentials_for_generation_model(setup_openai_mock, setup_xinference_mock):
model = XinferenceAILargeLanguageModel() model = XinferenceAILargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='alapaca', model="alapaca",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': 'www ' + os.environ.get('XINFERENCE_GENERATION_MODEL_UID') "model_uid": "www " + os.environ.get("XINFERENCE_GENERATION_MODEL_UID"),
} },
) )
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="alapaca", credentials={"server_url": "", "model_uid": ""})
model='alapaca',
credentials={
'server_url': '',
'model_uid': ''
}
)
model.validate_credentials( model.validate_credentials(
model='alapaca', model="alapaca",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_GENERATION_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_GENERATION_MODEL_UID"),
} },
) )
@pytest.mark.parametrize('setup_openai_mock, setup_xinference_mock', [['completion', 'none']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock, setup_xinference_mock", [["completion", "none"]], indirect=True)
def test_invoke_generation_model(setup_openai_mock, setup_xinference_mock): def test_invoke_generation_model(setup_openai_mock, setup_xinference_mock):
model = XinferenceAILargeLanguageModel() model = XinferenceAILargeLanguageModel()
response = model.invoke( response = model.invoke(
model='alapaca', model="alapaca",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_GENERATION_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_GENERATION_MODEL_UID"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="the United States is")],
UserPromptMessage(
content='the United States is'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
user="abc-123", user="abc-123",
stream=False stream=False,
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
assert len(response.message.content) > 0 assert len(response.message.content) > 0
assert response.usage.total_tokens > 0 assert response.usage.total_tokens > 0
@pytest.mark.parametrize('setup_openai_mock, setup_xinference_mock', [['completion', 'none']], indirect=True)
@pytest.mark.parametrize("setup_openai_mock, setup_xinference_mock", [["completion", "none"]], indirect=True)
def test_invoke_stream_generation_model(setup_openai_mock, setup_xinference_mock): def test_invoke_stream_generation_model(setup_openai_mock, setup_xinference_mock):
model = XinferenceAILargeLanguageModel() model = XinferenceAILargeLanguageModel()
response = model.invoke( response = model.invoke(
model='alapaca', model="alapaca",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_GENERATION_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_GENERATION_MODEL_UID"),
}, },
prompt_messages=[ prompt_messages=[UserPromptMessage(content="the United States is")],
UserPromptMessage(
content='the United States is'
)
],
model_parameters={ model_parameters={
'temperature': 0.7, "temperature": 0.7,
'top_p': 1.0, "top_p": 1.0,
}, },
stop=['you'], stop=["you"],
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -330,66 +313,52 @@ def test_invoke_stream_generation_model(setup_openai_mock, setup_xinference_mock
assert isinstance(chunk.delta.message, AssistantPromptMessage) assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
def test_get_num_tokens(): def test_get_num_tokens():
model = XinferenceAILargeLanguageModel() model = XinferenceAILargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='ChatGLM3', model="ChatGLM3",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_GENERATION_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_GENERATION_MODEL_UID"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
tools=[ tools=[
PromptMessageTool( PromptMessageTool(
name='get_current_weather', name="get_current_weather",
description='Get the current weather in a given location', description="Get the current weather in a given location",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {
"location": { "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"type": "string", "unit": {"type": "string", "enum": ["c", "f"]},
"description": "The city and state e.g. San Francisco, CA"
}, },
"unit": { "required": ["location"],
"type": "string",
"enum": [
"c",
"f"
]
}
}, },
"required": [
"location"
]
}
) )
] ],
) )
assert isinstance(num_tokens, int) assert isinstance(num_tokens, int)
assert num_tokens == 77 assert num_tokens == 77
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='ChatGLM3', model="ChatGLM3",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_GENERATION_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_GENERATION_MODEL_UID"),
}, },
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!'
)
], ],
) )

32
api/tests/integration_tests/model_runtime/xinference/test_rerank.py
View file

@ -8,44 +8,42 @@ from core.model_runtime.model_providers.xinference.rerank.rerank import Xinferen
from tests.integration_tests.model_runtime.__mock.xinference import MOCK, setup_xinference_mock from tests.integration_tests.model_runtime.__mock.xinference import MOCK, setup_xinference_mock
@pytest.mark.parametrize('setup_xinference_mock', [['none']], indirect=True) @pytest.mark.parametrize("setup_xinference_mock", [["none"]], indirect=True)
def test_validate_credentials(setup_xinference_mock): def test_validate_credentials(setup_xinference_mock):
model = XinferenceRerankModel() model = XinferenceRerankModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(
model='bge-reranker-base', model="bge-reranker-base",
credentials={ credentials={"server_url": "awdawdaw", "model_uid": os.environ.get("XINFERENCE_RERANK_MODEL_UID")},
'server_url': 'awdawdaw',
'model_uid': os.environ.get('XINFERENCE_RERANK_MODEL_UID')
}
) )
model.validate_credentials( model.validate_credentials(
model='bge-reranker-base', model="bge-reranker-base",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_RERANK_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_RERANK_MODEL_UID"),
} },
) )
@pytest.mark.parametrize('setup_xinference_mock', [['none']], indirect=True)
@pytest.mark.parametrize("setup_xinference_mock", [["none"]], indirect=True)
def test_invoke_model(setup_xinference_mock): def test_invoke_model(setup_xinference_mock):
model = XinferenceRerankModel() model = XinferenceRerankModel()
result = model.invoke( result = model.invoke(
model='bge-reranker-base', model="bge-reranker-base",
credentials={ credentials={
'server_url': os.environ.get('XINFERENCE_SERVER_URL'), "server_url": os.environ.get("XINFERENCE_SERVER_URL"),
'model_uid': os.environ.get('XINFERENCE_RERANK_MODEL_UID') "model_uid": os.environ.get("XINFERENCE_RERANK_MODEL_UID"),
}, },
query="Who is Kasumi?", query="Who is Kasumi?",
docs=[ docs=[
"Kasumi is a girl's name of Japanese origin meaning \"mist\".", 'Kasumi is a girl\'s name of Japanese origin meaning "mist".',
"Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ", "Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ",
"and she leads a team named PopiParty." "and she leads a team named PopiParty.",
], ],
score_threshold=0.8 score_threshold=0.8,
) )
assert isinstance(result, RerankResult) assert isinstance(result, RerankResult)

69
api/tests/integration_tests/model_runtime/zhinao/test_llm.py
View file

@ -13,41 +13,22 @@ def test_validate_credentials():
model = ZhinaoLargeLanguageModel() model = ZhinaoLargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="360gpt2-pro", credentials={"api_key": "invalid_key"})
model='360gpt2-pro',
credentials={
'api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(model="360gpt2-pro", credentials={"api_key": os.environ.get("ZHINAO_API_KEY")})
model='360gpt2-pro',
credentials={
'api_key': os.environ.get('ZHINAO_API_KEY')
}
)
def test_invoke_model(): def test_invoke_model():
model = ZhinaoLargeLanguageModel() model = ZhinaoLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='360gpt2-pro', model="360gpt2-pro",
credentials={ credentials={"api_key": os.environ.get("ZHINAO_API_KEY")},
'api_key': os.environ.get('ZHINAO_API_KEY') prompt_messages=[UserPromptMessage(content="Who are you?")],
}, model_parameters={"temperature": 0.5, "max_tokens": 10},
prompt_messages=[ stop=["How"],
UserPromptMessage(
content='Who are you?'
)
],
model_parameters={
'temperature': 0.5,
'max_tokens': 10
},
stop=['How'],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -58,22 +39,12 @@ def test_invoke_stream_model():
model = ZhinaoLargeLanguageModel() model = ZhinaoLargeLanguageModel()
response = model.invoke( response = model.invoke(
model='360gpt2-pro', model="360gpt2-pro",
credentials={ credentials={"api_key": os.environ.get("ZHINAO_API_KEY")},
'api_key': os.environ.get('ZHINAO_API_KEY') prompt_messages=[UserPromptMessage(content="Hello World!")],
}, model_parameters={"temperature": 0.5, "max_tokens": 100, "seed": 1234},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={
'temperature': 0.5,
'max_tokens': 100,
'seed': 1234
},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -89,18 +60,14 @@ def test_get_num_tokens():
model = ZhinaoLargeLanguageModel() model = ZhinaoLargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='360gpt2-pro', model="360gpt2-pro",
credentials={ credentials={"api_key": os.environ.get("ZHINAO_API_KEY")},
'api_key': os.environ.get('ZHINAO_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert num_tokens == 21 assert num_tokens == 21

10
api/tests/integration_tests/model_runtime/zhinao/test_provider.py
View file

@ -10,12 +10,6 @@ def test_validate_provider_credentials():
provider = ZhinaoProvider() provider = ZhinaoProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"api_key": os.environ.get("ZHINAO_API_KEY")})
credentials={
'api_key': os.environ.get('ZHINAO_API_KEY')
}
)

104
api/tests/integration_tests/model_runtime/zhipuai/test_llm.py
View file

@ -18,41 +18,22 @@ def test_validate_credentials():
model = ZhipuAILargeLanguageModel() model = ZhipuAILargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="chatglm_turbo", credentials={"api_key": "invalid_key"})
model='chatglm_turbo',
credentials={
'api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(model="chatglm_turbo", credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")})
model='chatglm_turbo',
credentials={
'api_key': os.environ.get('ZHIPUAI_API_KEY')
}
)
def test_invoke_model(): def test_invoke_model():
model = ZhipuAILargeLanguageModel() model = ZhipuAILargeLanguageModel()
response = model.invoke( response = model.invoke(
model='chatglm_turbo', model="chatglm_turbo",
credentials={ credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")},
'api_key': os.environ.get('ZHIPUAI_API_KEY') prompt_messages=[UserPromptMessage(content="Who are you?")],
}, model_parameters={"temperature": 0.9, "top_p": 0.7},
prompt_messages=[ stop=["How"],
UserPromptMessage(
content='Who are you?'
)
],
model_parameters={
'temperature': 0.9,
'top_p': 0.7
},
stop=['How'],
stream=False, stream=False,
user="abc-123" user="abc-123",
) )
assert isinstance(response, LLMResult) assert isinstance(response, LLMResult)
@ -63,21 +44,12 @@ def test_invoke_stream_model():
model = ZhipuAILargeLanguageModel() model = ZhipuAILargeLanguageModel()
response = model.invoke( response = model.invoke(
model='chatglm_turbo', model="chatglm_turbo",
credentials={ credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")},
'api_key': os.environ.get('ZHIPUAI_API_KEY') prompt_messages=[UserPromptMessage(content="Hello World!")],
}, model_parameters={"temperature": 0.9, "top_p": 0.7},
prompt_messages=[
UserPromptMessage(
content='Hello World!'
)
],
model_parameters={
'temperature': 0.9,
'top_p': 0.7
},
stream=True, stream=True,
user="abc-123" user="abc-123",
) )
assert isinstance(response, Generator) assert isinstance(response, Generator)
@ -93,63 +65,45 @@ def test_get_num_tokens():
model = ZhipuAILargeLanguageModel() model = ZhipuAILargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='chatglm_turbo', model="chatglm_turbo",
credentials={ credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")},
'api_key': os.environ.get('ZHIPUAI_API_KEY')
},
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert num_tokens == 14 assert num_tokens == 14
def test_get_tools_num_tokens(): def test_get_tools_num_tokens():
model = ZhipuAILargeLanguageModel() model = ZhipuAILargeLanguageModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='tools', model="tools",
credentials={ credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")},
'api_key': os.environ.get('ZHIPUAI_API_KEY')
},
tools=[ tools=[
PromptMessageTool( PromptMessageTool(
name='get_current_weather', name="get_current_weather",
description='Get the current weather in a given location', description="Get the current weather in a given location",
parameters={ parameters={
"type": "object", "type": "object",
"properties": { "properties": {
"location": { "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"type": "string", "unit": {"type": "string", "enum": ["c", "f"]},
"description": "The city and state e.g. San Francisco, CA"
}, },
"unit": { "required": ["location"],
"type": "string",
"enum": [
"c",
"f"
]
}
}, },
"required": [
"location"
]
}
) )
], ],
prompt_messages=[ prompt_messages=[
SystemPromptMessage( SystemPromptMessage(
content='You are a helpful AI assistant.', content="You are a helpful AI assistant.",
), ),
UserPromptMessage( UserPromptMessage(content="Hello World!"),
content='Hello World!' ],
)
]
) )
assert num_tokens == 88 assert num_tokens == 88

10
api/tests/integration_tests/model_runtime/zhipuai/test_provider.py
View file

@ -10,12 +10,6 @@ def test_validate_provider_credentials():
provider = ZhipuaiProvider() provider = ZhipuaiProvider()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={})
credentials={}
)
provider.validate_provider_credentials( provider.validate_provider_credentials(credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")})
credentials={
'api_key': os.environ.get('ZHIPUAI_API_KEY')
}
)

36
api/tests/integration_tests/model_runtime/zhipuai/test_text_embedding.py
View file

@ -11,34 +11,19 @@ def test_validate_credentials():
model = ZhipuAITextEmbeddingModel() model = ZhipuAITextEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError): with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials( model.validate_credentials(model="text_embedding", credentials={"api_key": "invalid_key"})
model='text_embedding',
credentials={
'api_key': 'invalid_key'
}
)
model.validate_credentials( model.validate_credentials(model="text_embedding", credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")})
model='text_embedding',
credentials={
'api_key': os.environ.get('ZHIPUAI_API_KEY')
}
)
def test_invoke_model(): def test_invoke_model():
model = ZhipuAITextEmbeddingModel() model = ZhipuAITextEmbeddingModel()
result = model.invoke( result = model.invoke(
model='text_embedding', model="text_embedding",
credentials={ credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")},
'api_key': os.environ.get('ZHIPUAI_API_KEY') texts=["hello", "world"],
}, user="abc-123",
texts=[
"hello",
"world"
],
user="abc-123"
) )
assert isinstance(result, TextEmbeddingResult) assert isinstance(result, TextEmbeddingResult)
@ -50,14 +35,7 @@ def test_get_num_tokens():
model = ZhipuAITextEmbeddingModel() model = ZhipuAITextEmbeddingModel()
num_tokens = model.get_num_tokens( num_tokens = model.get_num_tokens(
model='text_embedding', model="text_embedding", credentials={"api_key": os.environ.get("ZHIPUAI_API_KEY")}, texts=["hello", "world"]
credentials={
'api_key': os.environ.get('ZHIPUAI_API_KEY')
},
texts=[
"hello",
"world"
]
) )
assert num_tokens == 2 assert num_tokens == 2

15
api/tests/integration_tests/tools/__mock/http.py
View file

@ -7,20 +7,17 @@ from _pytest.monkeypatch import MonkeyPatch
class MockedHttp: class MockedHttp:
def httpx_request(method: Literal['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'HEAD'], def httpx_request(
url: str, **kwargs) -> httpx.Response: method: Literal["GET", "POST", "PUT", "DELETE", "PATCH", "HEAD"], url: str, **kwargs
) -> httpx.Response:
""" """
Mocked httpx.request Mocked httpx.request
""" """
request = httpx.Request( request = httpx.Request(
method, method, url, params=kwargs.get("params"), headers=kwargs.get("headers"), cookies=kwargs.get("cookies")
url,
params=kwargs.get('params'),
headers=kwargs.get('headers'),
cookies=kwargs.get('cookies')
) )
data = kwargs.get('data', None) data = kwargs.get("data", None)
resp = json.dumps(data).encode('utf-8') if data else b'OK' resp = json.dumps(data).encode("utf-8") if data else b"OK"
response = httpx.Response( response = httpx.Response(
status_code=200, status_code=200,
request=request, request=request,

6
api/tests/integration_tests/tools/__mock_server/openapi_todo.py
View file

@ -10,6 +10,7 @@ todos_data = {
"user1": ["Go for a run", "Read a book"], "user1": ["Go for a run", "Read a book"],
} }
class TodosResource(Resource): class TodosResource(Resource):
def get(self, username): def get(self, username):
todos = todos_data.get(username, []) todos = todos_data.get(username, [])
@ -32,7 +33,8 @@ class TodosResource(Resource):
return {"error": "Invalid todo index"}, 400 return {"error": "Invalid todo index"}, 400
api.add_resource(TodosResource, '/todos/<string:username>')
if __name__ == '__main__': api.add_resource(TodosResource, "/todos/<string:username>")
if __name__ == "__main__":
app.run(port=5003, debug=True) app.run(port=5003, debug=True)

47
api/tests/integration_tests/tools/api_tool/test_api_tool.py
View file

@ -3,37 +3,40 @@ from core.tools.tool.tool import Tool
from tests.integration_tests.tools.__mock.http import setup_http_mock from tests.integration_tests.tools.__mock.http import setup_http_mock
tool_bundle = { tool_bundle = {
'server_url': 'http://www.example.com/{path_param}', "server_url": "http://www.example.com/{path_param}",
'method': 'post', "method": "post",
'author': '', "author": "",
'openapi': {'parameters': [{'in': 'path', 'name': 'path_param'}, "openapi": {
{'in': 'query', 'name': 'query_param'}, "parameters": [
{'in': 'cookie', 'name': 'cookie_param'}, {"in": "path", "name": "path_param"},
{'in': 'header', 'name': 'header_param'}, {"in": "query", "name": "query_param"},
{"in": "cookie", "name": "cookie_param"},
{"in": "header", "name": "header_param"},
], ],
'requestBody': { "requestBody": {
'content': {'application/json': {'schema': {'properties': {'body_param': {'type': 'string'}}}}}} "content": {"application/json": {"schema": {"properties": {"body_param": {"type": "string"}}}}}
}, },
'parameters': [] },
"parameters": [],
} }
parameters = { parameters = {
'path_param': 'p_param', "path_param": "p_param",
'query_param': 'q_param', "query_param": "q_param",
'cookie_param': 'c_param', "cookie_param": "c_param",
'header_param': 'h_param', "header_param": "h_param",
'body_param': 'b_param', "body_param": "b_param",
} }
def test_api_tool(setup_http_mock): def test_api_tool(setup_http_mock):
tool = ApiTool(api_bundle=tool_bundle, runtime=Tool.Runtime(credentials={'auth_type': 'none'})) tool = ApiTool(api_bundle=tool_bundle, runtime=Tool.Runtime(credentials={"auth_type": "none"}))
headers = tool.assembling_request(parameters) headers = tool.assembling_request(parameters)
response = tool.do_http_request(tool.api_bundle.server_url, tool.api_bundle.method, headers, parameters) response = tool.do_http_request(tool.api_bundle.server_url, tool.api_bundle.method, headers, parameters)
assert response.status_code == 200 assert response.status_code == 200
assert '/p_param' == response.request.url.path assert "/p_param" == response.request.url.path
assert b'query_param=q_param' == response.request.url.query assert b"query_param=q_param" == response.request.url.query
assert 'h_param' == response.request.headers.get('header_param') assert "h_param" == response.request.headers.get("header_param")
assert 'application/json' == response.request.headers.get('content-type') assert "application/json" == response.request.headers.get("content-type")
assert 'cookie_param=c_param' == response.request.headers.get('cookie') assert "cookie_param=c_param" == response.request.headers.get("cookie")
assert 'b_param' in response.content.decode() assert "b_param" in response.content.decode()

3
api/tests/integration_tests/tools/test_all_provider.py
View file

@ -7,7 +7,8 @@ provider_names = [provider.identity.name for provider in provider_generator]
ToolManager.clear_builtin_providers_cache() ToolManager.clear_builtin_providers_cache()
provider_generator = ToolManager.list_builtin_providers() provider_generator = ToolManager.list_builtin_providers()
@pytest.mark.parametrize('name', provider_names)
@pytest.mark.parametrize("name", provider_names)
def test_tool_providers(benchmark, name): def test_tool_providers(benchmark, name):
""" """
Test that all tool providers can be loaded Test that all tool providers can be loaded

Some files were not shown because too many files have changed in this diff Show more