Underground-Digital/Workflow-Engine

1import os2from collections.abc import Generator3 4import pytest5 6from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta7from core.model_runtime.entities.message_entities import AssistantPromptMessage, UserPromptMessage8from core.model_runtime.errors.validate import CredentialsValidateFailedError9from core.model_runtime.model_providers.volcengine_maas.llm.llm import VolcengineMaaSLargeLanguageModel10 11 12def test_validate_credentials_for_chat_model():13    model = VolcengineMaaSLargeLanguageModel()14 15    with pytest.raises(CredentialsValidateFailedError):16        model.validate_credentials(17            model="NOT IMPORTANT",18            credentials={19                "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",20                "volc_region": "cn-beijing",21                "volc_access_key_id": "INVALID",22                "volc_secret_access_key": "INVALID",23                "endpoint_id": "INVALID",24            },25        )26 27    model.validate_credentials(28        model="NOT IMPORTANT",29        credentials={30            "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",31            "volc_region": "cn-beijing",32            "volc_access_key_id": os.environ.get("VOLC_API_KEY"),33            "volc_secret_access_key": os.environ.get("VOLC_SECRET_KEY"),34            "endpoint_id": os.environ.get("VOLC_MODEL_ENDPOINT_ID"),35        },36    )37 38 39def test_invoke_model():40    model = VolcengineMaaSLargeLanguageModel()41 42    response = model.invoke(43        model="NOT IMPORTANT",44        credentials={45            "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",46            "volc_region": "cn-beijing",47            "volc_access_key_id": os.environ.get("VOLC_API_KEY"),48            "volc_secret_access_key": os.environ.get("VOLC_SECRET_KEY"),49            "endpoint_id": os.environ.get("VOLC_MODEL_ENDPOINT_ID"),50            "base_model_name": "Skylark2-pro-4k",51        },52        prompt_messages=[UserPromptMessage(content="Hello World!")],53        model_parameters={54            "temperature": 0.7,55            "top_p": 1.0,56            "top_k": 1,57        },58        stop=["you"],59        user="abc-123",60        stream=False,61    )62 63    assert isinstance(response, LLMResult)64    assert len(response.message.content) > 065    assert response.usage.total_tokens > 066 67 68def test_invoke_stream_model():69    model = VolcengineMaaSLargeLanguageModel()70 71    response = model.invoke(72        model="NOT IMPORTANT",73        credentials={74            "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",75            "volc_region": "cn-beijing",76            "volc_access_key_id": os.environ.get("VOLC_API_KEY"),77            "volc_secret_access_key": os.environ.get("VOLC_SECRET_KEY"),78            "endpoint_id": os.environ.get("VOLC_MODEL_ENDPOINT_ID"),79            "base_model_name": "Skylark2-pro-4k",80        },81        prompt_messages=[UserPromptMessage(content="Hello World!")],82        model_parameters={83            "temperature": 0.7,84            "top_p": 1.0,85            "top_k": 1,86        },87        stop=["you"],88        stream=True,89        user="abc-123",90    )91 92    assert isinstance(response, Generator)93    for chunk in response:94        assert isinstance(chunk, LLMResultChunk)95        assert isinstance(chunk.delta, LLMResultChunkDelta)96        assert isinstance(chunk.delta.message, AssistantPromptMessage)97        assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True98 99 100def test_get_num_tokens():101    model = VolcengineMaaSLargeLanguageModel()102 103    response = model.get_num_tokens(104        model="NOT IMPORTANT",105        credentials={106            "api_endpoint_host": "maas-api.ml-platform-cn-beijing.volces.com",107            "volc_region": "cn-beijing",108            "volc_access_key_id": os.environ.get("VOLC_API_KEY"),109            "volc_secret_access_key": os.environ.get("VOLC_SECRET_KEY"),110            "endpoint_id": os.environ.get("VOLC_MODEL_ENDPOINT_ID"),111            "base_model_name": "Skylark2-pro-4k",112        },113        prompt_messages=[UserPromptMessage(content="Hello World!")],114        tools=[],115    )116 117    assert isinstance(response, int)118    assert response == 6119