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, SystemPromptMessage, UserPromptMessage8from core.model_runtime.errors.validate import CredentialsValidateFailedError9from core.model_runtime.model_providers.siliconflow.llm.llm import SiliconflowLargeLanguageModel10 11 12def test_validate_credentials():13    model = SiliconflowLargeLanguageModel()14 15    with pytest.raises(CredentialsValidateFailedError):16        model.validate_credentials(model="deepseek-ai/DeepSeek-V2-Chat", credentials={"api_key": "invalid_key"})17 18    model.validate_credentials(model="deepseek-ai/DeepSeek-V2-Chat", credentials={"api_key": os.environ.get("API_KEY")})19 20 21def test_invoke_model():22    model = SiliconflowLargeLanguageModel()23 24    response = model.invoke(25        model="deepseek-ai/DeepSeek-V2-Chat",26        credentials={"api_key": os.environ.get("API_KEY")},27        prompt_messages=[UserPromptMessage(content="Who are you?")],28        model_parameters={"temperature": 0.5, "max_tokens": 10},29        stop=["How"],30        stream=False,31        user="abc-123",32    )33 34    assert isinstance(response, LLMResult)35    assert len(response.message.content) > 036 37 38def test_invoke_stream_model():39    model = SiliconflowLargeLanguageModel()40 41    response = model.invoke(42        model="deepseek-ai/DeepSeek-V2-Chat",43        credentials={"api_key": os.environ.get("API_KEY")},44        prompt_messages=[UserPromptMessage(content="Hello World!")],45        model_parameters={"temperature": 0.5, "max_tokens": 100, "seed": 1234},46        stream=True,47        user="abc-123",48    )49 50    assert isinstance(response, Generator)51 52    for chunk in response:53        assert isinstance(chunk, LLMResultChunk)54        assert isinstance(chunk.delta, LLMResultChunkDelta)55        assert isinstance(chunk.delta.message, AssistantPromptMessage)56        assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True57 58 59def test_get_num_tokens():60    model = SiliconflowLargeLanguageModel()61 62    num_tokens = model.get_num_tokens(63        model="deepseek-ai/DeepSeek-V2-Chat",64        credentials={"api_key": os.environ.get("API_KEY")},65        prompt_messages=[66            SystemPromptMessage(67                content="You are a helpful AI assistant.",68            ),69            UserPromptMessage(content="Hello World!"),70        ],71    )72 73    assert num_tokens == 1274