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.bedrock.llm.llm import BedrockLargeLanguageModel10 11 12def test_validate_credentials():13    model = BedrockLargeLanguageModel()14 15    with pytest.raises(CredentialsValidateFailedError):16        model.validate_credentials(model="meta.llama2-13b-chat-v1", credentials={"anthropic_api_key": "invalid_key"})17 18    model.validate_credentials(19        model="meta.llama2-13b-chat-v1",20        credentials={21            "aws_region": os.getenv("AWS_REGION"),22            "aws_access_key": os.getenv("AWS_ACCESS_KEY"),23            "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),24        },25    )26 27 28def test_invoke_model():29    model = BedrockLargeLanguageModel()30 31    response = model.invoke(32        model="meta.llama2-13b-chat-v1",33        credentials={34            "aws_region": os.getenv("AWS_REGION"),35            "aws_access_key": os.getenv("AWS_ACCESS_KEY"),36            "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),37        },38        prompt_messages=[39            SystemPromptMessage(40                content="You are a helpful AI assistant.",41            ),42            UserPromptMessage(content="Hello World!"),43        ],44        model_parameters={"temperature": 0.0, "top_p": 1.0, "max_tokens_to_sample": 10},45        stop=["How"],46        stream=False,47        user="abc-123",48    )49 50    assert isinstance(response, LLMResult)51    assert len(response.message.content) > 052 53 54def test_invoke_stream_model():55    model = BedrockLargeLanguageModel()56 57    response = model.invoke(58        model="meta.llama2-13b-chat-v1",59        credentials={60            "aws_region": os.getenv("AWS_REGION"),61            "aws_access_key": os.getenv("AWS_ACCESS_KEY"),62            "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),63        },64        prompt_messages=[65            SystemPromptMessage(66                content="You are a helpful AI assistant.",67            ),68            UserPromptMessage(content="Hello World!"),69        ],70        model_parameters={"temperature": 0.0, "max_tokens_to_sample": 100},71        stream=True,72        user="abc-123",73    )74 75    assert isinstance(response, Generator)76 77    for chunk in response:78        print(chunk)79        assert isinstance(chunk, LLMResultChunk)80        assert isinstance(chunk.delta, LLMResultChunkDelta)81        assert isinstance(chunk.delta.message, AssistantPromptMessage)82        assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True83 84 85def test_get_num_tokens():86    model = BedrockLargeLanguageModel()87 88    num_tokens = model.get_num_tokens(89        model="meta.llama2-13b-chat-v1",90        credentials={91            "aws_region": os.getenv("AWS_REGION"),92            "aws_access_key": os.getenv("AWS_ACCESS_KEY"),93            "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),94        },95        messages=[96            SystemPromptMessage(97                content="You are a helpful AI assistant.",98            ),99            UserPromptMessage(content="Hello World!"),100        ],101    )102 103    assert num_tokens == 18104