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.anthropic.llm.llm import AnthropicLargeLanguageModel10from tests.integration_tests.model_runtime.__mock.anthropic import setup_anthropic_mock11 12 13@pytest.mark.parametrize("setup_anthropic_mock", [["none"]], indirect=True)14def test_validate_credentials(setup_anthropic_mock):15    model = AnthropicLargeLanguageModel()16 17    with pytest.raises(CredentialsValidateFailedError):18        model.validate_credentials(model="claude-instant-1.2", credentials={"anthropic_api_key": "invalid_key"})19 20    model.validate_credentials(21        model="claude-instant-1.2", credentials={"anthropic_api_key": os.environ.get("ANTHROPIC_API_KEY")}22    )23 24 25@pytest.mark.parametrize("setup_anthropic_mock", [["none"]], indirect=True)26def test_invoke_model(setup_anthropic_mock):27    model = AnthropicLargeLanguageModel()28 29    response = model.invoke(30        model="claude-instant-1.2",31        credentials={32            "anthropic_api_key": os.environ.get("ANTHROPIC_API_KEY"),33            "anthropic_api_url": os.environ.get("ANTHROPIC_API_URL"),34        },35        prompt_messages=[36            SystemPromptMessage(37                content="You are a helpful AI assistant.",38            ),39            UserPromptMessage(content="Hello World!"),40        ],41        model_parameters={"temperature": 0.0, "top_p": 1.0, "max_tokens": 10},42        stop=["How"],43        stream=False,44        user="abc-123",45    )46 47    assert isinstance(response, LLMResult)48    assert len(response.message.content) > 049 50 51@pytest.mark.parametrize("setup_anthropic_mock", [["none"]], indirect=True)52def test_invoke_stream_model(setup_anthropic_mock):53    model = AnthropicLargeLanguageModel()54 55    response = model.invoke(56        model="claude-instant-1.2",57        credentials={"anthropic_api_key": os.environ.get("ANTHROPIC_API_KEY")},58        prompt_messages=[59            SystemPromptMessage(60                content="You are a helpful AI assistant.",61            ),62            UserPromptMessage(content="Hello World!"),63        ],64        model_parameters={"temperature": 0.0, "max_tokens": 100},65        stream=True,66        user="abc-123",67    )68 69    assert isinstance(response, Generator)70 71    for chunk in response:72        assert isinstance(chunk, LLMResultChunk)73        assert isinstance(chunk.delta, LLMResultChunkDelta)74        assert isinstance(chunk.delta.message, AssistantPromptMessage)75        assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True76 77 78def test_get_num_tokens():79    model = AnthropicLargeLanguageModel()80 81    num_tokens = model.get_num_tokens(82        model="claude-instant-1.2",83        credentials={"anthropic_api_key": os.environ.get("ANTHROPIC_API_KEY")},84        prompt_messages=[85            SystemPromptMessage(86                content="You are a helpful AI assistant.",87            ),88            UserPromptMessage(content="Hello World!"),89        ],90    )91 92    assert num_tokens == 1893