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 (8    AssistantPromptMessage,9    SystemPromptMessage,10    UserPromptMessage,11)12from core.model_runtime.errors.validate import CredentialsValidateFailedError13from core.model_runtime.model_providers.openrouter.llm.llm import OpenRouterLargeLanguageModel14 15 16def test_validate_credentials():17    model = OpenRouterLargeLanguageModel()18 19    with pytest.raises(CredentialsValidateFailedError):20        model.validate_credentials(21            model="mistralai/mixtral-8x7b-instruct", credentials={"api_key": "invalid_key", "mode": "chat"}22        )23 24    model.validate_credentials(25        model="mistralai/mixtral-8x7b-instruct",26        credentials={"api_key": os.environ.get("TOGETHER_API_KEY"), "mode": "chat"},27    )28 29 30def test_invoke_model():31    model = OpenRouterLargeLanguageModel()32 33    response = model.invoke(34        model="mistralai/mixtral-8x7b-instruct",35        credentials={"api_key": os.environ.get("TOGETHER_API_KEY"), "mode": "completion"},36        prompt_messages=[37            SystemPromptMessage(38                content="You are a helpful AI assistant.",39            ),40            UserPromptMessage(content="Who are you?"),41        ],42        model_parameters={43            "temperature": 1.0,44            "top_k": 2,45            "top_p": 0.5,46        },47        stop=["How"],48        stream=False,49        user="abc-123",50    )51 52    assert isinstance(response, LLMResult)53    assert len(response.message.content) > 054 55 56def test_invoke_stream_model():57    model = OpenRouterLargeLanguageModel()58 59    response = model.invoke(60        model="mistralai/mixtral-8x7b-instruct",61        credentials={"api_key": os.environ.get("TOGETHER_API_KEY"), "mode": "chat"},62        prompt_messages=[63            SystemPromptMessage(64                content="You are a helpful AI assistant.",65            ),66            UserPromptMessage(content="Who are you?"),67        ],68        model_parameters={69            "temperature": 1.0,70            "top_k": 2,71            "top_p": 0.5,72        },73        stop=["How"],74        stream=True,75        user="abc-123",76    )77 78    assert isinstance(response, Generator)79 80    for chunk in response:81        assert isinstance(chunk, LLMResultChunk)82        assert isinstance(chunk.delta, LLMResultChunkDelta)83        assert isinstance(chunk.delta.message, AssistantPromptMessage)84 85 86def test_get_num_tokens():87    model = OpenRouterLargeLanguageModel()88 89    num_tokens = model.get_num_tokens(90        model="mistralai/mixtral-8x7b-instruct",91        credentials={92            "api_key": os.environ.get("TOGETHER_API_KEY"),93        },94        prompt_messages=[95            SystemPromptMessage(96                content="You are a helpful AI assistant.",97            ),98            UserPromptMessage(content="Hello World!"),99        ],100    )101 102    assert isinstance(num_tokens, int)103    assert num_tokens == 21104