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.entities.model_entities import AIModelEntity14from core.model_runtime.errors.validate import CredentialsValidateFailedError15from core.model_runtime.model_providers.x.llm.llm import XAILargeLanguageModel16 17"""FOR MOCK FIXTURES, DO NOT REMOVE"""18from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock19 20 21def test_predefined_models():22    model = XAILargeLanguageModel()23    model_schemas = model.predefined_models()24 25    assert len(model_schemas) >= 126    assert isinstance(model_schemas[0], AIModelEntity)27 28 29@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)30def test_validate_credentials_for_chat_model(setup_openai_mock):31    model = XAILargeLanguageModel()32 33    with pytest.raises(CredentialsValidateFailedError):34        # model name to gpt-3.5-turbo because of mocking35        model.validate_credentials(36            model="gpt-3.5-turbo",37            credentials={"api_key": "invalid_key", "endpoint_url": os.environ.get("XAI_API_BASE"), "mode": "chat"},38        )39 40    model.validate_credentials(41        model="grok-beta",42        credentials={43            "api_key": os.environ.get("XAI_API_KEY"),44            "endpoint_url": os.environ.get("XAI_API_BASE"),45            "mode": "chat",46        },47    )48 49 50@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)51def test_invoke_chat_model(setup_openai_mock):52    model = XAILargeLanguageModel()53 54    result = model.invoke(55        model="grok-beta",56        credentials={57            "api_key": os.environ.get("XAI_API_KEY"),58            "endpoint_url": os.environ.get("XAI_API_BASE"),59            "mode": "chat",60        },61        prompt_messages=[62            SystemPromptMessage(63                content="You are a helpful AI assistant.",64            ),65            UserPromptMessage(content="Hello World!"),66        ],67        model_parameters={68            "temperature": 0.0,69            "top_p": 1.0,70            "presence_penalty": 0.0,71            "frequency_penalty": 0.0,72            "max_tokens": 10,73        },74        stop=["How"],75        stream=False,76        user="foo",77    )78 79    assert isinstance(result, LLMResult)80    assert len(result.message.content) > 081 82 83@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)84def test_invoke_chat_model_with_tools(setup_openai_mock):85    model = XAILargeLanguageModel()86 87    result = model.invoke(88        model="grok-beta",89        credentials={90            "api_key": os.environ.get("XAI_API_KEY"),91            "endpoint_url": os.environ.get("XAI_API_BASE"),92            "mode": "chat",93        },94        prompt_messages=[95            SystemPromptMessage(96                content="You are a helpful AI assistant.",97            ),98            UserPromptMessage(99                content="what's the weather today in London?",100            ),101        ],102        model_parameters={"temperature": 0.0, "max_tokens": 100},103        tools=[104            PromptMessageTool(105                name="get_weather",106                description="Determine weather in my location",107                parameters={108                    "type": "object",109                    "properties": {110                        "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},111                        "unit": {"type": "string", "enum": ["c", "f"]},112                    },113                    "required": ["location"],114                },115            ),116            PromptMessageTool(117                name="get_stock_price",118                description="Get the current stock price",119                parameters={120                    "type": "object",121                    "properties": {"symbol": {"type": "string", "description": "The stock symbol"}},122                    "required": ["symbol"],123                },124            ),125        ],126        stream=False,127        user="foo",128    )129 130    assert isinstance(result, LLMResult)131    assert isinstance(result.message, AssistantPromptMessage)132 133 134@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)135def test_invoke_stream_chat_model(setup_openai_mock):136    model = XAILargeLanguageModel()137 138    result = model.invoke(139        model="grok-beta",140        credentials={141            "api_key": os.environ.get("XAI_API_KEY"),142            "endpoint_url": os.environ.get("XAI_API_BASE"),143            "mode": "chat",144        },145        prompt_messages=[146            SystemPromptMessage(147                content="You are a helpful AI assistant.",148            ),149            UserPromptMessage(content="Hello World!"),150        ],151        model_parameters={"temperature": 0.0, "max_tokens": 100},152        stream=True,153        user="foo",154    )155 156    assert isinstance(result, Generator)157 158    for chunk in result:159        assert isinstance(chunk, LLMResultChunk)160        assert isinstance(chunk.delta, LLMResultChunkDelta)161        assert isinstance(chunk.delta.message, AssistantPromptMessage)162        assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True163        if chunk.delta.finish_reason is not None:164            assert chunk.delta.usage is not None165            assert chunk.delta.usage.completion_tokens > 0166 167 168def test_get_num_tokens():169    model = XAILargeLanguageModel()170 171    num_tokens = model.get_num_tokens(172        model="grok-beta",173        credentials={"api_key": os.environ.get("XAI_API_KEY"), "endpoint_url": os.environ.get("XAI_API_BASE")},174        prompt_messages=[UserPromptMessage(content="Hello World!")],175    )176 177    assert num_tokens == 10178 179    num_tokens = model.get_num_tokens(180        model="grok-beta",181        credentials={"api_key": os.environ.get("XAI_API_KEY"), "endpoint_url": os.environ.get("XAI_API_BASE")},182        prompt_messages=[183            SystemPromptMessage(184                content="You are a helpful AI assistant.",185            ),186            UserPromptMessage(content="Hello World!"),187        ],188        tools=[189            PromptMessageTool(190                name="get_weather",191                description="Determine weather in my location",192                parameters={193                    "type": "object",194                    "properties": {195                        "location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},196                        "unit": {"type": "string", "enum": ["c", "f"]},197                    },198                    "required": ["location"],199                },200            ),201        ],202    )203 204    assert num_tokens == 77205