|
@@ -0,0 +1,131 @@
|
|
|
+import os
|
|
|
+from collections.abc import Generator
|
|
|
+
|
|
|
+import pytest
|
|
|
+
|
|
|
+from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta
|
|
|
+from core.model_runtime.entities.message_entities import (
|
|
|
+ AssistantPromptMessage,
|
|
|
+ SystemPromptMessage,
|
|
|
+ UserPromptMessage,
|
|
|
+)
|
|
|
+from core.model_runtime.errors.validate import CredentialsValidateFailedError
|
|
|
+from core.model_runtime.model_providers.vessl_ai.llm.llm import VesslAILargeLanguageModel
|
|
|
+
|
|
|
+
|
|
|
+def test_validate_credentials():
|
|
|
+ model = VesslAILargeLanguageModel()
|
|
|
+
|
|
|
+ with pytest.raises(CredentialsValidateFailedError):
|
|
|
+ model.validate_credentials(
|
|
|
+ model=os.environ.get("VESSL_AI_MODEL_NAME"),
|
|
|
+ credentials={
|
|
|
+ "api_key": "invalid_key",
|
|
|
+ "endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
|
|
|
+ "mode": "chat",
|
|
|
+ },
|
|
|
+ )
|
|
|
+
|
|
|
+ with pytest.raises(CredentialsValidateFailedError):
|
|
|
+ model.validate_credentials(
|
|
|
+ model=os.environ.get("VESSL_AI_MODEL_NAME"),
|
|
|
+ credentials={
|
|
|
+ "api_key": os.environ.get("VESSL_AI_API_KEY"),
|
|
|
+ "endpoint_url": "http://invalid_url",
|
|
|
+ "mode": "chat",
|
|
|
+ },
|
|
|
+ )
|
|
|
+
|
|
|
+ model.validate_credentials(
|
|
|
+ model=os.environ.get("VESSL_AI_MODEL_NAME"),
|
|
|
+ credentials={
|
|
|
+ "api_key": os.environ.get("VESSL_AI_API_KEY"),
|
|
|
+ "endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
|
|
|
+ "mode": "chat",
|
|
|
+ },
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def test_invoke_model():
|
|
|
+ model = VesslAILargeLanguageModel()
|
|
|
+
|
|
|
+ response = model.invoke(
|
|
|
+ model=os.environ.get("VESSL_AI_MODEL_NAME"),
|
|
|
+ credentials={
|
|
|
+ "api_key": os.environ.get("VESSL_AI_API_KEY"),
|
|
|
+ "endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
|
|
|
+ "mode": "chat",
|
|
|
+ },
|
|
|
+ prompt_messages=[
|
|
|
+ SystemPromptMessage(
|
|
|
+ content="You are a helpful AI assistant.",
|
|
|
+ ),
|
|
|
+ UserPromptMessage(content="Who are you?"),
|
|
|
+ ],
|
|
|
+ model_parameters={
|
|
|
+ "temperature": 1.0,
|
|
|
+ "top_k": 2,
|
|
|
+ "top_p": 0.5,
|
|
|
+ },
|
|
|
+ stop=["How"],
|
|
|
+ stream=False,
|
|
|
+ user="abc-123",
|
|
|
+ )
|
|
|
+
|
|
|
+ assert isinstance(response, LLMResult)
|
|
|
+ assert len(response.message.content) > 0
|
|
|
+
|
|
|
+
|
|
|
+def test_invoke_stream_model():
|
|
|
+ model = VesslAILargeLanguageModel()
|
|
|
+
|
|
|
+ response = model.invoke(
|
|
|
+ model=os.environ.get("VESSL_AI_MODEL_NAME"),
|
|
|
+ credentials={
|
|
|
+ "api_key": os.environ.get("VESSL_AI_API_KEY"),
|
|
|
+ "endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
|
|
|
+ "mode": "chat",
|
|
|
+ },
|
|
|
+ prompt_messages=[
|
|
|
+ SystemPromptMessage(
|
|
|
+ content="You are a helpful AI assistant.",
|
|
|
+ ),
|
|
|
+ UserPromptMessage(content="Who are you?"),
|
|
|
+ ],
|
|
|
+ model_parameters={
|
|
|
+ "temperature": 1.0,
|
|
|
+ "top_k": 2,
|
|
|
+ "top_p": 0.5,
|
|
|
+ },
|
|
|
+ stop=["How"],
|
|
|
+ stream=True,
|
|
|
+ user="abc-123",
|
|
|
+ )
|
|
|
+
|
|
|
+ assert isinstance(response, Generator)
|
|
|
+
|
|
|
+ for chunk in response:
|
|
|
+ assert isinstance(chunk, LLMResultChunk)
|
|
|
+ assert isinstance(chunk.delta, LLMResultChunkDelta)
|
|
|
+ assert isinstance(chunk.delta.message, AssistantPromptMessage)
|
|
|
+
|
|
|
+
|
|
|
+def test_get_num_tokens():
|
|
|
+ model = VesslAILargeLanguageModel()
|
|
|
+
|
|
|
+ num_tokens = model.get_num_tokens(
|
|
|
+ model=os.environ.get("VESSL_AI_MODEL_NAME"),
|
|
|
+ credentials={
|
|
|
+ "api_key": os.environ.get("VESSL_AI_API_KEY"),
|
|
|
+ "endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
|
|
|
+ },
|
|
|
+ prompt_messages=[
|
|
|
+ SystemPromptMessage(
|
|
|
+ content="You are a helpful AI assistant.",
|
|
|
+ ),
|
|
|
+ UserPromptMessage(content="Hello World!"),
|
|
|
+ ],
|
|
|
+ )
|
|
|
+
|
|
|
+ assert isinstance(num_tokens, int)
|
|
|
+ assert num_tokens == 21
|