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    PromptMessageTool,10    SystemPromptMessage,11    UserPromptMessage,12)13from core.model_runtime.errors.validate import CredentialsValidateFailedError14from core.model_runtime.model_providers.oci.llm.llm import OCILargeLanguageModel15 16 17def test_validate_credentials():18    model = OCILargeLanguageModel()19 20    with pytest.raises(CredentialsValidateFailedError):21        model.validate_credentials(22            model="cohere.command-r-plus",23            credentials={"oci_config_content": "invalid_key", "oci_key_content": "invalid_key"},24        )25 26    model.validate_credentials(27        model="cohere.command-r-plus",28        credentials={29            "oci_config_content": os.environ.get("OCI_CONFIG_CONTENT"),30            "oci_key_content": os.environ.get("OCI_KEY_CONTENT"),31        },32    )33 34 35def test_invoke_model():36    model = OCILargeLanguageModel()37 38    response = model.invoke(39        model="cohere.command-r-plus",40        credentials={41            "oci_config_content": os.environ.get("OCI_CONFIG_CONTENT"),42            "oci_key_content": os.environ.get("OCI_KEY_CONTENT"),43        },44        prompt_messages=[UserPromptMessage(content="Hi")],45        model_parameters={"temperature": 0.5, "max_tokens": 10},46        stop=["How"],47        stream=False,48        user="abc-123",49    )50 51    assert isinstance(response, LLMResult)52    assert len(response.message.content) > 053 54 55def test_invoke_stream_model():56    model = OCILargeLanguageModel()57 58    response = model.invoke(59        model="meta.llama-3-70b-instruct",60        credentials={61            "oci_config_content": os.environ.get("OCI_CONFIG_CONTENT"),62            "oci_key_content": os.environ.get("OCI_KEY_CONTENT"),63        },64        prompt_messages=[UserPromptMessage(content="Hi")],65        model_parameters={"temperature": 0.5, "max_tokens": 100, "seed": 1234},66        stream=True,67        user="abc-123",68    )69 70    assert isinstance(response, Generator)71 72    for chunk in response:73        assert isinstance(chunk, LLMResultChunk)74        assert isinstance(chunk.delta, LLMResultChunkDelta)75        assert isinstance(chunk.delta.message, AssistantPromptMessage)76        assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True77 78 79def test_invoke_model_with_function():80    model = OCILargeLanguageModel()81 82    response = model.invoke(83        model="cohere.command-r-plus",84        credentials={85            "oci_config_content": os.environ.get("OCI_CONFIG_CONTENT"),86            "oci_key_content": os.environ.get("OCI_KEY_CONTENT"),87        },88        prompt_messages=[UserPromptMessage(content="Hi")],89        model_parameters={"temperature": 0.5, "max_tokens": 100, "seed": 1234},90        stream=False,91        user="abc-123",92        tools=[93            PromptMessageTool(94                name="get_current_weather",95                description="Get the current weather in a given location",96                parameters={97                    "type": "object",98                    "properties": {99                        "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},100                        "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},101                    },102                    "required": ["location"],103                },104            )105        ],106    )107 108    assert isinstance(response, LLMResult)109    assert len(response.message.content) > 0110 111 112def test_get_num_tokens():113    model = OCILargeLanguageModel()114 115    num_tokens = model.get_num_tokens(116        model="cohere.command-r-plus",117        credentials={118            "oci_config_content": os.environ.get("OCI_CONFIG_CONTENT"),119            "oci_key_content": os.environ.get("OCI_KEY_CONTENT"),120        },121        prompt_messages=[122            SystemPromptMessage(123                content="You are a helpful AI assistant.",124            ),125            UserPromptMessage(content="Hello World!"),126        ],127    )128 129    assert num_tokens == 18130