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.openllm.llm.llm import OpenLLMLargeLanguageModel10 11 12def test_validate_credentials_for_chat_model():13    model = OpenLLMLargeLanguageModel()14 15    with pytest.raises(CredentialsValidateFailedError):16        model.validate_credentials(17            model="NOT IMPORTANT",18            credentials={19                "server_url": "invalid_key",20            },21        )22 23    model.validate_credentials(24        model="NOT IMPORTANT",25        credentials={26            "server_url": os.environ.get("OPENLLM_SERVER_URL"),27        },28    )29 30 31def test_invoke_model():32    model = OpenLLMLargeLanguageModel()33 34    response = model.invoke(35        model="NOT IMPORTANT",36        credentials={37            "server_url": os.environ.get("OPENLLM_SERVER_URL"),38        },39        prompt_messages=[UserPromptMessage(content="Hello World!")],40        model_parameters={41            "temperature": 0.7,42            "top_p": 1.0,43            "top_k": 1,44        },45        stop=["you"],46        user="abc-123",47        stream=False,48    )49 50    assert isinstance(response, LLMResult)51    assert len(response.message.content) > 052    assert response.usage.total_tokens > 053 54 55def test_invoke_stream_model():56    model = OpenLLMLargeLanguageModel()57 58    response = model.invoke(59        model="NOT IMPORTANT",60        credentials={61            "server_url": os.environ.get("OPENLLM_SERVER_URL"),62        },63        prompt_messages=[UserPromptMessage(content="Hello World!")],64        model_parameters={65            "temperature": 0.7,66            "top_p": 1.0,67            "top_k": 1,68        },69        stop=["you"],70        stream=True,71        user="abc-123",72    )73 74    assert isinstance(response, Generator)75    for chunk in response:76        assert isinstance(chunk, LLMResultChunk)77        assert isinstance(chunk.delta, LLMResultChunkDelta)78        assert isinstance(chunk.delta.message, AssistantPromptMessage)79        assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True80 81 82def test_get_num_tokens():83    model = OpenLLMLargeLanguageModel()84 85    response = model.get_num_tokens(86        model="NOT IMPORTANT",87        credentials={88            "server_url": os.environ.get("OPENLLM_SERVER_URL"),89        },90        prompt_messages=[UserPromptMessage(content="Hello World!")],91        tools=[],92    )93 94    assert isinstance(response, int)95    assert response == 396