youngkingdom/dify
1
0
Fork 0
mirror of https://github.com/langgenius/dify.git synced 2026-04-27 22:18:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'main' into fix/chore-fix

Browse Source
This commit is contained in:
Yeuoly
2024-12-24 21:28:56 +08:00
parent 39656f7f84 cdaef30cc9
commit 4199998c7e
734 changed files with 7911 additions and 5007 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
api/tests/integration_tests/model_runtime/gitee_ai/__init__.py
View File

132
api/tests/integration_tests/model_runtime/gitee_ai/test_llm.py
View File

@ -1,132 +0,0 @@
import os
from collections.abc import Generator
import pytest
from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta
from core.model_runtime.entities.message_entities import (
AssistantPromptMessage,
PromptMessageTool,
SystemPromptMessage,
UserPromptMessage,
)
from core.model_runtime.entities.model_entities import AIModelEntity
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gitee_ai.llm.llm import GiteeAILargeLanguageModel
def test_predefined_models():
model = GiteeAILargeLanguageModel()
model_schemas = model.predefined_models()
assert len(model_schemas) >= 1
assert isinstance(model_schemas[0], AIModelEntity)
def test_validate_credentials_for_chat_model():
model = GiteeAILargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError):
# model name to gpt-3.5-turbo because of mocking
model.validate_credentials(model="gpt-3.5-turbo", credentials={"api_key": "invalid_key"})
model.validate_credentials(
model="Qwen2-7B-Instruct",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
)
def test_invoke_chat_model():
model = GiteeAILargeLanguageModel()
result = model.invoke(
model="Qwen2-7B-Instruct",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
model_parameters={
"temperature": 0.0,
"top_p": 1.0,
"presence_penalty": 0.0,
"frequency_penalty": 0.0,
"max_tokens": 10,
"stream": False,
},
stop=["How"],
stream=False,
user="foo",
)
assert isinstance(result, LLMResult)
assert len(result.message.content) > 0
def test_invoke_stream_chat_model():
model = GiteeAILargeLanguageModel()
result = model.invoke(
model="Qwen2-7B-Instruct",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
model_parameters={"temperature": 0.0, "max_tokens": 100, "stream": False},
stream=True,
user="foo",
)
assert isinstance(result, Generator)
for chunk in result:
assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
if chunk.delta.finish_reason is not None:
assert chunk.delta.usage is not None
def test_get_num_tokens():
model = GiteeAILargeLanguageModel()
num_tokens = model.get_num_tokens(
model="Qwen2-7B-Instruct",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
prompt_messages=[UserPromptMessage(content="Hello World!")],
)
assert num_tokens == 10
num_tokens = model.get_num_tokens(
model="Qwen2-7B-Instruct",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
tools=[
PromptMessageTool(
name="get_weather",
description="Determine weather in my location",
parameters={
"type": "object",
"properties": {
"location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"unit": {"type": "string", "enum": ["c", "f"]},
},
"required": ["location"],
},
),
],
)
assert num_tokens == 77

15
api/tests/integration_tests/model_runtime/gitee_ai/test_provider.py
View File

@ -1,15 +0,0 @@
import os
import pytest
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gitee_ai.gitee_ai import GiteeAIProvider
def test_validate_provider_credentials():
provider = GiteeAIProvider()
with pytest.raises(CredentialsValidateFailedError):
provider.validate_provider_credentials(credentials={"api_key": "invalid_key"})
provider.validate_provider_credentials(credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")})

47
api/tests/integration_tests/model_runtime/gitee_ai/test_rerank.py
View File

@ -1,47 +0,0 @@
import os
import pytest
from core.model_runtime.entities.rerank_entities import RerankResult
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gitee_ai.rerank.rerank import GiteeAIRerankModel
def test_validate_credentials():
model = GiteeAIRerankModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model="bge-reranker-v2-m3",
credentials={"api_key": "invalid_key"},
)
model.validate_credentials(
model="bge-reranker-v2-m3",
credentials={
"api_key": os.environ.get("GITEE_AI_API_KEY"),
},
)
def test_invoke_model():
model = GiteeAIRerankModel()
result = model.invoke(
model="bge-reranker-v2-m3",
credentials={
"api_key": os.environ.get("GITEE_AI_API_KEY"),
},
query="What is the capital of the United States?",
docs=[
"Carson City is the capital city of the American state of Nevada. At the 2010 United States "
"Census, Carson City had a population of 55,274.",
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that "
"are a political division controlled by the United States. Its capital is Saipan.",
],
top_n=1,
score_threshold=0.01,
)
assert isinstance(result, RerankResult)
assert len(result.docs) == 1
assert result.docs[0].score >= 0.01

45
api/tests/integration_tests/model_runtime/gitee_ai/test_speech2text.py
View File

@ -1,45 +0,0 @@
import os
import pytest
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gitee_ai.speech2text.speech2text import GiteeAISpeech2TextModel
def test_validate_credentials():
model = GiteeAISpeech2TextModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model="whisper-base",
credentials={"api_key": "invalid_key"},
)
model.validate_credentials(
model="whisper-base",
credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")},
)
def test_invoke_model():
model = GiteeAISpeech2TextModel()
# Get the directory of the current file
current_dir = os.path.dirname(os.path.abspath(__file__))
# Get assets directory
assets_dir = os.path.join(os.path.dirname(current_dir), "assets")
# Construct the path to the audio file
audio_file_path = os.path.join(assets_dir, "audio.mp3")
# Open the file and get the file object
with open(audio_file_path, "rb") as audio_file:
file = audio_file
result = model.invoke(
model="whisper-base", credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")}, file=file
)
assert isinstance(result, str)
assert result == "1 2 3 4 5 6 7 8 9 10"

46
api/tests/integration_tests/model_runtime/gitee_ai/test_text_embedding.py
View File

@ -1,46 +0,0 @@
import os
import pytest
from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gitee_ai.text_embedding.text_embedding import GiteeAIEmbeddingModel
def test_validate_credentials():
model = GiteeAIEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(model="bge-large-zh-v1.5", credentials={"api_key": "invalid_key"})
model.validate_credentials(model="bge-large-zh-v1.5", credentials={"api_key": os.environ.get("GITEE_AI_API_KEY")})
def test_invoke_model():
model = GiteeAIEmbeddingModel()
result = model.invoke(
model="bge-large-zh-v1.5",
credentials={
"api_key": os.environ.get("GITEE_AI_API_KEY"),
},
texts=["hello", "world"],
user="user",
)
assert isinstance(result, TextEmbeddingResult)
assert len(result.embeddings) == 2
def test_get_num_tokens():
model = GiteeAIEmbeddingModel()
num_tokens = model.get_num_tokens(
model="bge-large-zh-v1.5",
credentials={
"api_key": os.environ.get("GITEE_AI_API_KEY"),
},
texts=["hello", "world"],
)
assert num_tokens == 2

23
api/tests/integration_tests/model_runtime/gitee_ai/test_tts.py
View File

@ -1,23 +0,0 @@
import os
from core.model_runtime.model_providers.gitee_ai.tts.tts import GiteeAIText2SpeechModel
def test_invoke_model():
model = GiteeAIText2SpeechModel()
result = model.invoke(
model="speecht5_tts",
tenant_id="test",
credentials={
"api_key": os.environ.get("GITEE_AI_API_KEY"),
},
content_text="Hello, world!",
voice="",
)
content = b""
for chunk in result:
content += chunk
assert content != b""

0
api/tests/integration_tests/model_runtime/gpustack/__init__.py
View File

49
api/tests/integration_tests/model_runtime/gpustack/test_embedding.py
View File

@ -1,49 +0,0 @@
import os
import pytest
from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gpustack.text_embedding.text_embedding import (
GPUStackTextEmbeddingModel,
)
def test_validate_credentials():
model = GPUStackTextEmbeddingModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model="bge-m3",
credentials={
"endpoint_url": "invalid_url",
"api_key": "invalid_api_key",
},
)
model.validate_credentials(
model="bge-m3",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
},
)
def test_invoke_model():
model = GPUStackTextEmbeddingModel()
result = model.invoke(
model="bge-m3",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
"context_size": 8192,
},
texts=["hello", "world"],
user="abc-123",
)
assert isinstance(result, TextEmbeddingResult)
assert len(result.embeddings) == 2
assert result.usage.total_tokens == 7

162
api/tests/integration_tests/model_runtime/gpustack/test_llm.py
View File

@ -1,162 +0,0 @@
import os
from collections.abc import Generator
import pytest
from core.model_runtime.entities.llm_entities import (
LLMResult,
LLMResultChunk,
LLMResultChunkDelta,
)
from core.model_runtime.entities.message_entities import (
AssistantPromptMessage,
PromptMessageTool,
SystemPromptMessage,
UserPromptMessage,
)
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gpustack.llm.llm import GPUStackLanguageModel
def test_validate_credentials_for_chat_model():
model = GPUStackLanguageModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model="llama-3.2-1b-instruct",
credentials={
"endpoint_url": "invalid_url",
"api_key": "invalid_api_key",
"mode": "chat",
},
)
model.validate_credentials(
model="llama-3.2-1b-instruct",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
"mode": "chat",
},
)
def test_invoke_completion_model():
model = GPUStackLanguageModel()
response = model.invoke(
model="llama-3.2-1b-instruct",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
"mode": "completion",
},
prompt_messages=[UserPromptMessage(content="ping")],
model_parameters={"temperature": 0.7, "top_p": 1.0, "max_tokens": 10},
stop=[],
user="abc-123",
stream=False,
)
assert isinstance(response, LLMResult)
assert len(response.message.content) > 0
assert response.usage.total_tokens > 0
def test_invoke_chat_model():
model = GPUStackLanguageModel()
response = model.invoke(
model="llama-3.2-1b-instruct",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
"mode": "chat",
},
prompt_messages=[UserPromptMessage(content="ping")],
model_parameters={"temperature": 0.7, "top_p": 1.0, "max_tokens": 10},
stop=[],
user="abc-123",
stream=False,
)
assert isinstance(response, LLMResult)
assert len(response.message.content) > 0
assert response.usage.total_tokens > 0
def test_invoke_stream_chat_model():
model = GPUStackLanguageModel()
response = model.invoke(
model="llama-3.2-1b-instruct",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
"mode": "chat",
},
prompt_messages=[UserPromptMessage(content="Hello World!")],
model_parameters={"temperature": 0.7, "top_p": 1.0, "max_tokens": 10},
stop=["you"],
stream=True,
user="abc-123",
)
assert isinstance(response, Generator)
for chunk in response:
assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
def test_get_num_tokens():
model = GPUStackLanguageModel()
num_tokens = model.get_num_tokens(
model="????",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
"mode": "chat",
},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
tools=[
PromptMessageTool(
name="get_current_weather",
description="Get the current weather in a given location",
parameters={
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and state e.g. San Francisco, CA",
},
"unit": {"type": "string", "enum": ["c", "f"]},
},
"required": ["location"],
},
)
],
)
assert isinstance(num_tokens, int)
assert num_tokens == 80
num_tokens = model.get_num_tokens(
model="????",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
"mode": "chat",
},
prompt_messages=[UserPromptMessage(content="Hello World!")],
)
assert isinstance(num_tokens, int)
assert num_tokens == 10

107
api/tests/integration_tests/model_runtime/gpustack/test_rerank.py
View File

@ -1,107 +0,0 @@
import os
import pytest
from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.gpustack.rerank.rerank import (
GPUStackRerankModel,
)
def test_validate_credentials_for_rerank_model():
model = GPUStackRerankModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model="bge-reranker-v2-m3",
credentials={
"endpoint_url": "invalid_url",
"api_key": "invalid_api_key",
},
)
model.validate_credentials(
model="bge-reranker-v2-m3",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
},
)
def test_invoke_rerank_model():
model = GPUStackRerankModel()
response = model.invoke(
model="bge-reranker-v2-m3",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
},
query="Organic skincare products for sensitive skin",
docs=[
"Eco-friendly kitchenware for modern homes",
"Biodegradable cleaning supplies for eco-conscious consumers",
"Organic cotton baby clothes for sensitive skin",
"Natural organic skincare range for sensitive skin",
"Tech gadgets for smart homes: 2024 edition",
"Sustainable gardening tools and compost solutions",
"Sensitive skin-friendly facial cleansers and toners",
"Organic food wraps and storage solutions",
"Yoga mats made from recycled materials",
],
top_n=3,
score_threshold=-0.75,
user="abc-123",
)
assert isinstance(response, RerankResult)
assert len(response.docs) == 3
def test__invoke():
model = GPUStackRerankModel()
# Test case 1: Empty docs
result = model._invoke(
model="bge-reranker-v2-m3",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
},
query="Organic skincare products for sensitive skin",
docs=[],
top_n=3,
score_threshold=0.75,
user="abc-123",
)
assert isinstance(result, RerankResult)
assert len(result.docs) == 0
# Test case 2: Expected docs
result = model._invoke(
model="bge-reranker-v2-m3",
credentials={
"endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),
"api_key": os.environ.get("GPUSTACK_API_KEY"),
},
query="Organic skincare products for sensitive skin",
docs=[
"Eco-friendly kitchenware for modern homes",
"Biodegradable cleaning supplies for eco-conscious consumers",
"Organic cotton baby clothes for sensitive skin",
"Natural organic skincare range for sensitive skin",
"Tech gadgets for smart homes: 2024 edition",
"Sustainable gardening tools and compost solutions",
"Sensitive skin-friendly facial cleansers and toners",
"Organic food wraps and storage solutions",
"Yoga mats made from recycled materials",
],
top_n=3,
score_threshold=-0.75,
user="abc-123",
)
assert isinstance(result, RerankResult)
assert len(result.docs) == 3
assert all(isinstance(doc, RerankDocument) for doc in result.docs)

0
api/tests/integration_tests/model_runtime/vessl_ai/__init__.py
View File

131
api/tests/integration_tests/model_runtime/vessl_ai/test_llm.py
View File

@ -1,131 +0,0 @@
import os
from collections.abc import Generator
import pytest
from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta
from core.model_runtime.entities.message_entities import (
AssistantPromptMessage,
SystemPromptMessage,
UserPromptMessage,
)
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.vessl_ai.llm.llm import VesslAILargeLanguageModel
def test_validate_credentials():
model = VesslAILargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": "invalid_key",
"endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
"mode": "chat",
},
)
with pytest.raises(CredentialsValidateFailedError):
model.validate_credentials(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": os.environ.get("VESSL_AI_API_KEY"),
"endpoint_url": "http://invalid_url",
"mode": "chat",
},
)
model.validate_credentials(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": os.environ.get("VESSL_AI_API_KEY"),
"endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
"mode": "chat",
},
)
def test_invoke_model():
model = VesslAILargeLanguageModel()
response = model.invoke(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": os.environ.get("VESSL_AI_API_KEY"),
"endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
"mode": "chat",
},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Who are you?"),
],
model_parameters={
"temperature": 1.0,
"top_k": 2,
"top_p": 0.5,
},
stop=["How"],
stream=False,
user="abc-123",
)
assert isinstance(response, LLMResult)
assert len(response.message.content) > 0
def test_invoke_stream_model():
model = VesslAILargeLanguageModel()
response = model.invoke(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": os.environ.get("VESSL_AI_API_KEY"),
"endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
"mode": "chat",
},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Who are you?"),
],
model_parameters={
"temperature": 1.0,
"top_k": 2,
"top_p": 0.5,
},
stop=["How"],
stream=True,
user="abc-123",
)
assert isinstance(response, Generator)
for chunk in response:
assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage)
def test_get_num_tokens():
model = VesslAILargeLanguageModel()
num_tokens = model.get_num_tokens(
model=os.environ.get("VESSL_AI_MODEL_NAME"),
credentials={
"api_key": os.environ.get("VESSL_AI_API_KEY"),
"endpoint_url": os.environ.get("VESSL_AI_ENDPOINT_URL"),
},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
)
assert isinstance(num_tokens, int)
assert num_tokens == 21

21
api/tests/integration_tests/model_runtime/wenxin/test_rerank.py
View File

@ -1,21 +0,0 @@
import os
from time import sleep
from core.model_runtime.entities.rerank_entities import RerankResult
from core.model_runtime.model_providers.wenxin.rerank.rerank import WenxinRerankModel
def test_invoke_bce_reranker_base_v1():
sleep(3)
model = WenxinRerankModel()
response = model.invoke(
model="bce-reranker-base_v1",
credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
query="What is Deep Learning?",
docs=["Deep Learning is ...", "My Book is ..."],
user="abc-123",
)
assert isinstance(response, RerankResult)
assert len(response.docs) == 2

0
api/tests/integration_tests/model_runtime/x/__init__.py
View File

204
api/tests/integration_tests/model_runtime/x/test_llm.py
View File

@ -1,204 +0,0 @@
import os
from collections.abc import Generator
import pytest
from core.model_runtime.entities.llm_entities import LLMResult, LLMResultChunk, LLMResultChunkDelta
from core.model_runtime.entities.message_entities import (
AssistantPromptMessage,
PromptMessageTool,
SystemPromptMessage,
UserPromptMessage,
)
from core.model_runtime.entities.model_entities import AIModelEntity
from core.model_runtime.errors.validate import CredentialsValidateFailedError
from core.model_runtime.model_providers.x.llm.llm import XAILargeLanguageModel
"""FOR MOCK FIXTURES, DO NOT REMOVE"""
from tests.integration_tests.model_runtime.__mock.openai import setup_openai_mock
def test_predefined_models():
model = XAILargeLanguageModel()
model_schemas = model.predefined_models()
assert len(model_schemas) >= 1
assert isinstance(model_schemas[0], AIModelEntity)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_validate_credentials_for_chat_model(setup_openai_mock):
model = XAILargeLanguageModel()
with pytest.raises(CredentialsValidateFailedError):
# model name to gpt-3.5-turbo because of mocking
model.validate_credentials(
model="gpt-3.5-turbo",
credentials={"api_key": "invalid_key", "endpoint_url": os.environ.get("XAI_API_BASE"), "mode": "chat"},
)
model.validate_credentials(
model="grok-beta",
credentials={
"api_key": os.environ.get("XAI_API_KEY"),
"endpoint_url": os.environ.get("XAI_API_BASE"),
"mode": "chat",
},
)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_invoke_chat_model(setup_openai_mock):
model = XAILargeLanguageModel()
result = model.invoke(
model="grok-beta",
credentials={
"api_key": os.environ.get("XAI_API_KEY"),
"endpoint_url": os.environ.get("XAI_API_BASE"),
"mode": "chat",
},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
model_parameters={
"temperature": 0.0,
"top_p": 1.0,
"presence_penalty": 0.0,
"frequency_penalty": 0.0,
"max_tokens": 10,
},
stop=["How"],
stream=False,
user="foo",
)
assert isinstance(result, LLMResult)
assert len(result.message.content) > 0
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_invoke_chat_model_with_tools(setup_openai_mock):
model = XAILargeLanguageModel()
result = model.invoke(
model="grok-beta",
credentials={
"api_key": os.environ.get("XAI_API_KEY"),
"endpoint_url": os.environ.get("XAI_API_BASE"),
"mode": "chat",
},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(
content="what's the weather today in London?",
),
],
model_parameters={"temperature": 0.0, "max_tokens": 100},
tools=[
PromptMessageTool(
name="get_weather",
description="Determine weather in my location",
parameters={
"type": "object",
"properties": {
"location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"unit": {"type": "string", "enum": ["c", "f"]},
},
"required": ["location"],
},
),
PromptMessageTool(
name="get_stock_price",
description="Get the current stock price",
parameters={
"type": "object",
"properties": {"symbol": {"type": "string", "description": "The stock symbol"}},
"required": ["symbol"],
},
),
],
stream=False,
user="foo",
)
assert isinstance(result, LLMResult)
assert isinstance(result.message, AssistantPromptMessage)
@pytest.mark.parametrize("setup_openai_mock", [["chat"]], indirect=True)
def test_invoke_stream_chat_model(setup_openai_mock):
model = XAILargeLanguageModel()
result = model.invoke(
model="grok-beta",
credentials={
"api_key": os.environ.get("XAI_API_KEY"),
"endpoint_url": os.environ.get("XAI_API_BASE"),
"mode": "chat",
},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
model_parameters={"temperature": 0.0, "max_tokens": 100},
stream=True,
user="foo",
)
assert isinstance(result, Generator)
for chunk in result:
assert isinstance(chunk, LLMResultChunk)
assert isinstance(chunk.delta, LLMResultChunkDelta)
assert isinstance(chunk.delta.message, AssistantPromptMessage)
assert len(chunk.delta.message.content) > 0 if chunk.delta.finish_reason is None else True
if chunk.delta.finish_reason is not None:
assert chunk.delta.usage is not None
assert chunk.delta.usage.completion_tokens > 0
def test_get_num_tokens():
model = XAILargeLanguageModel()
num_tokens = model.get_num_tokens(
model="grok-beta",
credentials={"api_key": os.environ.get("XAI_API_KEY"), "endpoint_url": os.environ.get("XAI_API_BASE")},
prompt_messages=[UserPromptMessage(content="Hello World!")],
)
assert num_tokens == 10
num_tokens = model.get_num_tokens(
model="grok-beta",
credentials={"api_key": os.environ.get("XAI_API_KEY"), "endpoint_url": os.environ.get("XAI_API_BASE")},
prompt_messages=[
SystemPromptMessage(
content="You are a helpful AI assistant.",
),
UserPromptMessage(content="Hello World!"),
],
tools=[
PromptMessageTool(
name="get_weather",
description="Determine weather in my location",
parameters={
"type": "object",
"properties": {
"location": {"type": "string", "description": "The city and state e.g. San Francisco, CA"},
"unit": {"type": "string", "enum": ["c", "f"]},
},
"required": ["location"],
},
),
],
)
assert num_tokens == 77