참고소스 수정본
This commit is contained in:
LASTA_DEV01\lasta
@@ -0,0 +1,410 @@
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from guardrails.utils.openai_utils.streaming_utils import (
|
||||
num_tokens_from_messages,
|
||||
num_tokens_from_string,
|
||||
)
|
||||
|
||||
|
||||
class TestNumTokensFromString:
|
||||
"""Test num_tokens_from_string function."""
|
||||
|
||||
def test_simple_text_gpt35(self):
|
||||
"""Test token counting for simple text with GPT-3.5."""
|
||||
text = "Hello, world!"
|
||||
result = num_tokens_from_string(text, "gpt-3.5-turbo")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_simple_text_gpt4(self):
|
||||
"""Test token counting for simple text with GPT-4."""
|
||||
text = "Hello, world!"
|
||||
result = num_tokens_from_string(text, "gpt-4")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_empty_string(self):
|
||||
"""Test token counting for empty string."""
|
||||
text = ""
|
||||
result = num_tokens_from_string(text, "gpt-3.5-turbo")
|
||||
assert result == 0
|
||||
|
||||
def test_longer_text(self):
|
||||
"""Test token counting for longer text."""
|
||||
text = """
|
||||
This is a longer text that contains multiple sentences.
|
||||
It should return a higher token count than a simple string.
|
||||
Let's add more content to make it even longer and see what happens.
|
||||
"""
|
||||
result = num_tokens_from_string(text, "gpt-3.5-turbo")
|
||||
assert isinstance(result, int)
|
||||
assert result > 10
|
||||
|
||||
def test_special_characters(self):
|
||||
"""Test token counting with special characters."""
|
||||
text = "Hello! @#$%^&*() 你好 🎉"
|
||||
result = num_tokens_from_string(text, "gpt-3.5-turbo")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_code_text(self):
|
||||
"""Test token counting for code."""
|
||||
text = """
|
||||
def hello_world():
|
||||
print("Hello, World!")
|
||||
return True
|
||||
"""
|
||||
result = num_tokens_from_string(text, "gpt-3.5-turbo")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_different_models_same_text(self):
|
||||
"""Test that different models can tokenize the same text."""
|
||||
text = "This is a test sentence."
|
||||
result_gpt35 = num_tokens_from_string(text, "gpt-3.5-turbo")
|
||||
result_gpt4 = num_tokens_from_string(text, "gpt-4")
|
||||
|
||||
# Both should return valid token counts
|
||||
assert isinstance(result_gpt35, int)
|
||||
assert isinstance(result_gpt4, int)
|
||||
assert result_gpt35 > 0
|
||||
assert result_gpt4 > 0
|
||||
|
||||
def test_unicode_text(self):
|
||||
"""Test token counting with Unicode characters."""
|
||||
text = "Hello 世界 مرحبا мир"
|
||||
result = num_tokens_from_string(text, "gpt-3.5-turbo")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_newlines_and_whitespace(self):
|
||||
"""Test token counting with various whitespace."""
|
||||
text = "Line 1\n\nLine 2\t\tTabbed\r\nCarriage return"
|
||||
result = num_tokens_from_string(text, "gpt-3.5-turbo")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
|
||||
class TestNumTokensFromMessages:
|
||||
"""Test num_tokens_from_messages function."""
|
||||
|
||||
def test_single_message_gpt35_turbo_0613(self):
|
||||
"""Test token counting for single message with gpt-3.5-turbo-0613."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
# Should be: 3 (tokens_per_message)
|
||||
# + tokens in "user"
|
||||
# + tokens in "Hello!"
|
||||
# + 3 (priming)
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_multiple_messages_gpt35_turbo_0613(self):
|
||||
"""Test token counting for multiple messages."""
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
assert isinstance(result, int)
|
||||
assert result > 10
|
||||
|
||||
def test_message_with_name_gpt35_turbo_0613(self):
|
||||
"""Test token counting for message with name field."""
|
||||
messages = [{"role": "user", "name": "Alice", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
# Should include tokens_per_name (1) for the name field
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_gpt4_0613(self):
|
||||
"""Test token counting with gpt-4-0613."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-4-0613")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_gpt4_32k_0613(self):
|
||||
"""Test token counting with gpt-4-32k-0613."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-4-32k-0613")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_gpt35_turbo_16k_0613(self):
|
||||
"""Test token counting with gpt-3.5-turbo-16k-0613."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-16k-0613")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_gpt35_turbo_0301(self):
|
||||
"""Test token counting with gpt-3.5-turbo-0301 (different token
|
||||
counting)."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0301")
|
||||
# tokens_per_message = 4, tokens_per_name = -1
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_gpt35_turbo_0301_with_name(self):
|
||||
"""Test token counting with gpt-3.5-turbo-0301 with name field."""
|
||||
messages = [{"role": "user", "name": "Alice", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0301")
|
||||
# tokens_per_name = -1, so having a name should reduce count
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
@patch("guardrails.utils.openai_utils.streaming_utils.logger")
|
||||
def test_generic_gpt35_turbo_fallback(self, mock_logger):
|
||||
"""Test that generic gpt-3.5-turbo falls back to gpt-3.5-turbo-0613."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo")
|
||||
|
||||
# Should log a warning about the fallback
|
||||
assert mock_logger.warning.called
|
||||
warning_message = mock_logger.warning.call_args[0][0]
|
||||
assert "gpt-3.5-turbo may update over time" in warning_message
|
||||
|
||||
# Should still return a valid token count
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
@patch("guardrails.utils.openai_utils.streaming_utils.logger")
|
||||
def test_generic_gpt4_fallback(self, mock_logger):
|
||||
"""Test that generic gpt-4 falls back to gpt-4-0613."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-4")
|
||||
|
||||
# Should log a warning about the fallback
|
||||
assert mock_logger.warning.called
|
||||
warning_message = mock_logger.warning.call_args[0][0]
|
||||
assert "gpt-4 may update over time" in warning_message
|
||||
|
||||
# Should still return a valid token count
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
@patch("guardrails.utils.openai_utils.streaming_utils.logger")
|
||||
def test_custom_gpt35_turbo_variant(self, mock_logger):
|
||||
"""Test custom gpt-3.5-turbo variant falls back to
|
||||
gpt-3.5-turbo-0613."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-custom")
|
||||
|
||||
# Should log a warning
|
||||
assert mock_logger.warning.called
|
||||
|
||||
# Should still return a valid token count
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
@patch("guardrails.utils.openai_utils.streaming_utils.logger")
|
||||
def test_custom_gpt4_variant(self, mock_logger):
|
||||
"""Test custom gpt-4 variant falls back to gpt-4-0613."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-4-turbo")
|
||||
|
||||
# Should log a warning
|
||||
assert mock_logger.warning.called
|
||||
|
||||
# Should still return a valid token count
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
@patch("guardrails.utils.openai_utils.streaming_utils.logger")
|
||||
def test_unknown_model_with_keyerror(self, mock_logger):
|
||||
"""Test unknown model that raises KeyError."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
|
||||
# Mock tiktoken to raise KeyError for unknown model
|
||||
with patch(
|
||||
"guardrails.utils.openai_utils.streaming_utils.tiktoken"
|
||||
) as mock_tiktoken:
|
||||
mock_encoding = Mock()
|
||||
mock_encoding.encode.return_value = [1, 2, 3]
|
||||
|
||||
def encoding_for_model_side_effect(model):
|
||||
if model == "unknown-model":
|
||||
raise KeyError("Model not found")
|
||||
return mock_encoding
|
||||
|
||||
mock_tiktoken.encoding_for_model.side_effect = (
|
||||
encoding_for_model_side_effect # noqa
|
||||
)
|
||||
mock_tiktoken.get_encoding.return_value = mock_encoding
|
||||
|
||||
# Should raise NotImplementedError for truly unknown model
|
||||
with pytest.raises(
|
||||
NotImplementedError,
|
||||
match="num_tokens_from_messages\\(\\) is not implemented for model unknown-model", # noqa
|
||||
):
|
||||
num_tokens_from_messages(messages, "unknown-model")
|
||||
|
||||
# Should have logged warning about model not found
|
||||
mock_logger.warning.assert_called_with(
|
||||
"model not found. Using cl100k_base encoding."
|
||||
)
|
||||
|
||||
def test_empty_messages_list(self):
|
||||
"""Test token counting with empty messages list."""
|
||||
messages = []
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
# Should only have the 3 tokens for priming
|
||||
assert result == 3
|
||||
|
||||
def test_message_with_empty_content(self):
|
||||
"""Test token counting for message with empty content."""
|
||||
messages = [{"role": "user", "content": ""}]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
# Should have tokens_per_message
|
||||
# + tokens for "user"
|
||||
# + 0 for content
|
||||
# + 3 for priming
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_complex_conversation(self):
|
||||
"""Test token counting for a complex conversation."""
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful coding assistant."},
|
||||
{"role": "user", "content": "How do I write a function in Python?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Here's how to write a function in Python:\n\ndef my_function():\n pass", # noqa
|
||||
},
|
||||
{"role": "user", "content": "Thanks!"},
|
||||
]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
assert isinstance(result, int)
|
||||
assert result > 20
|
||||
|
||||
def test_message_with_multiple_fields(self):
|
||||
"""Test token counting for message with multiple fields."""
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"name": "Alice",
|
||||
"content": "Hello, how are you?",
|
||||
}
|
||||
]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_consistent_token_count(self):
|
||||
"""Test that token count is consistent for the same input."""
|
||||
messages = [{"role": "user", "content": "Test message"}]
|
||||
result1 = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
result2 = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
assert result1 == result2
|
||||
|
||||
def test_gpt4_0314(self):
|
||||
"""Test token counting with gpt-4-0314."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-4-0314")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_gpt4_32k_0314(self):
|
||||
"""Test token counting with gpt-4-32k-0314."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
result = num_tokens_from_messages(messages, "gpt-4-32k-0314")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_messages_with_unicode(self):
|
||||
"""Test token counting for messages with Unicode content."""
|
||||
messages = [
|
||||
{"role": "user", "content": "Hello 世界"},
|
||||
{"role": "assistant", "content": "مرحبا مир"},
|
||||
]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_long_message_content(self):
|
||||
"""Test token counting for messages with long content."""
|
||||
long_content = " ".join(["word"] * 1000)
|
||||
messages = [{"role": "user", "content": long_content}]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
assert isinstance(result, int)
|
||||
assert result > 1000
|
||||
|
||||
def test_default_model_parameter(self):
|
||||
"""Test that default model parameter is gpt-3.5-turbo-0613."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
# Call without specifying model - should use default
|
||||
result = num_tokens_from_messages(messages)
|
||||
assert isinstance(result, int)
|
||||
assert result > 0
|
||||
|
||||
def test_priming_tokens_included(self):
|
||||
"""Test that the 3 priming tokens are always included."""
|
||||
# Even with empty messages, should have 3 tokens for priming
|
||||
messages = []
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
assert result == 3
|
||||
|
||||
@patch("guardrails.utils.openai_utils.streaming_utils.logger")
|
||||
def test_unsupported_model_raises_error(self, _mock_logger):
|
||||
"""Test that unsupported model raises NotImplementedError."""
|
||||
messages = [{"role": "user", "content": "Hello!"}]
|
||||
|
||||
with pytest.raises(NotImplementedError) as exc_info:
|
||||
num_tokens_from_messages(messages, "claude-2")
|
||||
|
||||
assert (
|
||||
"num_tokens_from_messages() is not implemented for model claude-2"
|
||||
in str( # noqa
|
||||
exc_info.value
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class TestIntegrationScenarios:
|
||||
"""Test integration scenarios combining both functions."""
|
||||
|
||||
def test_string_tokens_vs_message_tokens(self):
|
||||
"""Test that message tokens include overhead beyond just content
|
||||
tokens."""
|
||||
content = "Hello, world!"
|
||||
string_tokens = num_tokens_from_string(content, "gpt-3.5-turbo")
|
||||
message_tokens = num_tokens_from_messages(
|
||||
[{"role": "user", "content": content}], "gpt-3.5-turbo-0613"
|
||||
)
|
||||
|
||||
# Message tokens should be higher due to role tokens and overhead
|
||||
assert message_tokens > string_tokens
|
||||
|
||||
def test_multiple_messages_sum(self):
|
||||
"""Test that token counting works correctly for multiple messages."""
|
||||
messages = [
|
||||
{"role": "user", "content": "First message"},
|
||||
{"role": "assistant", "content": "Second message"},
|
||||
]
|
||||
result = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
|
||||
# Should be more than just the sum of content tokens
|
||||
first_tokens = num_tokens_from_string("First message", "gpt-3.5-turbo")
|
||||
second_tokens = num_tokens_from_string("Second message", "gpt-3.5-turbo")
|
||||
|
||||
# Total should include overhead for roles, formatting, and priming
|
||||
assert result > first_tokens + second_tokens
|
||||
|
||||
def test_consistency_across_model_versions(self):
|
||||
"""Test that specific model versions produce consistent results."""
|
||||
messages = [{"role": "user", "content": "Test"}]
|
||||
|
||||
# These models should use the same token counting (tokens_per_message=3)
|
||||
result_gpt35 = num_tokens_from_messages(messages, "gpt-3.5-turbo-0613")
|
||||
result_gpt4 = num_tokens_from_messages(messages, "gpt-4-0613")
|
||||
|
||||
# Should both return valid counts (may differ slightly due to encoding)
|
||||
assert isinstance(result_gpt35, int)
|
||||
assert isinstance(result_gpt4, int)
|
||||
assert result_gpt35 > 0
|
||||
assert result_gpt4 > 0
|
||||
@@ -0,0 +1,654 @@
|
||||
import os
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from guardrails.classes.llm.llm_response import LLMResponse
|
||||
from guardrails.utils.openai_utils.v1 import OpenAIClientV1
|
||||
|
||||
|
||||
class TestOpenAIClientV1Init:
|
||||
"""Test OpenAIClientV1 initialization."""
|
||||
|
||||
def test_init_with_api_key(self):
|
||||
"""Test initialization with provided API key."""
|
||||
client = OpenAIClientV1(api_key="test-key-123")
|
||||
assert client.api_key == "test-key-123"
|
||||
assert client.api_base is None
|
||||
|
||||
def test_init_with_api_base(self):
|
||||
"""Test initialization with custom base URL."""
|
||||
client = OpenAIClientV1(
|
||||
api_key="test-key-123", api_base="https://custom.api.com"
|
||||
)
|
||||
assert client.api_key == "test-key-123"
|
||||
assert client.api_base == "https://custom.api.com"
|
||||
|
||||
def test_init_with_env_var(self):
|
||||
"""Test initialization with API key from environment variable."""
|
||||
with patch.dict(os.environ, {"OPENAI_API_KEY": "env-key-456"}):
|
||||
client = OpenAIClientV1()
|
||||
assert client.api_key == "env-key-456"
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.openai.Client")
|
||||
def test_init_without_api_key(self, _mock_openai_client):
|
||||
"""Test initialization without API key when env var not set."""
|
||||
with patch.dict(os.environ, {}, clear=True):
|
||||
client = OpenAIClientV1()
|
||||
assert client.api_key is None
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.openai.Client")
|
||||
def test_client_creation(self, mock_openai_client):
|
||||
"""Test that OpenAI client is created with correct parameters."""
|
||||
OpenAIClientV1(api_key="test-key", api_base="https://custom.api.com")
|
||||
mock_openai_client.assert_called_once_with(
|
||||
api_key="test-key", base_url="https://custom.api.com"
|
||||
)
|
||||
|
||||
|
||||
class TestCreateEmbedding:
|
||||
"""Test create_embedding method."""
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.openai.Client")
|
||||
def test_create_embedding_success(self, mock_openai_client):
|
||||
"""Test successful embedding creation."""
|
||||
# Setup mock
|
||||
mock_embedding_response = Mock()
|
||||
mock_embedding_data_1 = Mock()
|
||||
mock_embedding_data_1.embedding = [0.1, 0.2, 0.3]
|
||||
mock_embedding_data_2 = Mock()
|
||||
mock_embedding_data_2.embedding = [0.4, 0.5, 0.6]
|
||||
mock_embedding_response.data = [mock_embedding_data_1, mock_embedding_data_2]
|
||||
|
||||
mock_client_instance = Mock()
|
||||
mock_client_instance.embeddings.create.return_value = mock_embedding_response
|
||||
mock_openai_client.return_value = mock_client_instance
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.create_embedding(
|
||||
model="text-embedding-ada-002", input=["text1", "text2"]
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert result == [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
|
||||
mock_client_instance.embeddings.create.assert_called_once_with(
|
||||
model="text-embedding-ada-002", input=["text1", "text2"]
|
||||
)
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.openai.Client")
|
||||
def test_create_embedding_single_input(self, mock_openai_client):
|
||||
"""Test embedding creation with single input."""
|
||||
# Setup mock
|
||||
mock_embedding_response = Mock()
|
||||
mock_embedding_data = Mock()
|
||||
mock_embedding_data.embedding = [0.1, 0.2, 0.3, 0.4, 0.5]
|
||||
mock_embedding_response.data = [mock_embedding_data]
|
||||
|
||||
mock_client_instance = Mock()
|
||||
mock_client_instance.embeddings.create.return_value = mock_embedding_response
|
||||
mock_openai_client.return_value = mock_client_instance
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.create_embedding(
|
||||
model="text-embedding-ada-002", input=["single text"]
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert result == [[0.1, 0.2, 0.3, 0.4, 0.5]]
|
||||
assert len(result) == 1
|
||||
|
||||
|
||||
class TestCreateCompletion:
|
||||
"""Test create_completion method."""
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_operation")
|
||||
@patch("guardrails.utils.openai_utils.v1.openai.Client")
|
||||
def test_create_completion_non_streaming(
|
||||
self, mock_openai_client, mock_trace_operation, mock_trace_llm_call
|
||||
):
|
||||
"""Test non-streaming completion."""
|
||||
# Setup mock response
|
||||
mock_completion_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_choice.text = "This is a completion response"
|
||||
mock_completion_response.choices = [mock_choice]
|
||||
mock_usage = Mock()
|
||||
mock_usage.prompt_tokens = 10
|
||||
mock_usage.completion_tokens = 20
|
||||
mock_usage.total_tokens = 30
|
||||
mock_completion_response.usage = mock_usage
|
||||
|
||||
mock_client_instance = Mock()
|
||||
mock_client_instance.completions.create.return_value = mock_completion_response
|
||||
mock_openai_client.return_value = mock_client_instance
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.create_completion(
|
||||
engine="gpt-3.5-turbo-instruct",
|
||||
prompt="Test prompt",
|
||||
temperature=0.7,
|
||||
max_tokens=100,
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert isinstance(result, LLMResponse)
|
||||
assert result.output == "This is a completion response"
|
||||
assert result.prompt_token_count == 10
|
||||
assert result.response_token_count == 20
|
||||
|
||||
# Verify API call
|
||||
mock_client_instance.completions.create.assert_called_once_with(
|
||||
model="gpt-3.5-turbo-instruct",
|
||||
prompt="Test prompt",
|
||||
temperature=0.7,
|
||||
max_tokens=100,
|
||||
)
|
||||
|
||||
# Verify tracing calls
|
||||
assert mock_trace_operation.call_count == 2 # input and output
|
||||
mock_trace_llm_call.assert_called()
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_operation")
|
||||
@patch("guardrails.utils.openai_utils.v1.openai.Client")
|
||||
def test_create_completion_streaming(
|
||||
self, mock_openai_client, mock_trace_operation, mock_trace_llm_call
|
||||
):
|
||||
"""Test streaming completion."""
|
||||
# Setup mock streaming response
|
||||
mock_stream = iter([{"choices": [{"text": "chunk1"}]}])
|
||||
|
||||
mock_client_instance = Mock()
|
||||
mock_client_instance.completions.create.return_value = mock_stream
|
||||
mock_openai_client.return_value = mock_client_instance
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.create_completion(
|
||||
engine="gpt-3.5-turbo-instruct", prompt="Test prompt", stream=True
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert isinstance(result, LLMResponse)
|
||||
assert result.output == ""
|
||||
assert result.stream_output is not None
|
||||
|
||||
|
||||
class TestConstructNonchatResponse:
|
||||
"""Test construct_nonchat_response method."""
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
def test_construct_nonchat_response_non_streaming(self, mock_trace_llm_call):
|
||||
"""Test non-streaming response construction."""
|
||||
# Setup mock response
|
||||
mock_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_choice.text = "Response text"
|
||||
mock_response.choices = [mock_choice]
|
||||
mock_usage = Mock()
|
||||
mock_usage.prompt_tokens = 15
|
||||
mock_usage.completion_tokens = 25
|
||||
mock_usage.total_tokens = 40
|
||||
mock_response.usage = mock_usage
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.construct_nonchat_response(
|
||||
stream=False, openai_response=mock_response
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert isinstance(result, LLMResponse)
|
||||
assert result.output == "Response text"
|
||||
assert result.prompt_token_count == 15
|
||||
assert result.response_token_count == 25
|
||||
assert result.stream_output is None
|
||||
|
||||
def test_construct_nonchat_response_streaming(self):
|
||||
"""Test streaming response construction."""
|
||||
# Setup mock streaming response
|
||||
mock_stream = iter([{"choices": [{"text": "chunk"}]}])
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.construct_nonchat_response(
|
||||
stream=True, openai_response=mock_stream
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert isinstance(result, LLMResponse)
|
||||
assert result.output == ""
|
||||
assert result.stream_output is not None
|
||||
|
||||
def test_construct_nonchat_response_no_choices_error(self):
|
||||
"""Test error when no choices in response."""
|
||||
# Setup mock response with no choices
|
||||
mock_response = Mock()
|
||||
mock_response.choices = []
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
|
||||
with pytest.raises(ValueError, match="No choices returned from OpenAI"):
|
||||
client.construct_nonchat_response(
|
||||
stream=False, openai_response=mock_response
|
||||
)
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
def test_construct_nonchat_response_no_usage_error(self, mock_trace_llm_call):
|
||||
"""Test error when no usage info in response."""
|
||||
# Setup mock response with no usage
|
||||
mock_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_choice.text = "Response text"
|
||||
mock_response.choices = [mock_choice]
|
||||
mock_response.usage = None
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
|
||||
with pytest.raises(ValueError, match="No token counts returned from OpenAI"):
|
||||
client.construct_nonchat_response(
|
||||
stream=False, openai_response=mock_response
|
||||
)
|
||||
|
||||
|
||||
class TestCreateChatCompletion:
|
||||
"""Test create_chat_completion method."""
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_operation")
|
||||
@patch("guardrails.utils.openai_utils.v1.openai.Client")
|
||||
def test_create_chat_completion_non_streaming(
|
||||
self, mock_openai_client, mock_trace_operation, mock_trace_llm_call
|
||||
):
|
||||
"""Test non-streaming chat completion."""
|
||||
# Setup mock response
|
||||
mock_chat_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_message = Mock()
|
||||
mock_message.content = "This is a chat response"
|
||||
mock_choice.message = mock_message
|
||||
mock_chat_response.choices = [mock_choice]
|
||||
mock_usage = Mock()
|
||||
mock_usage.prompt_tokens = 12
|
||||
mock_usage.completion_tokens = 18
|
||||
mock_usage.total_tokens = 30
|
||||
mock_chat_response.usage = mock_usage
|
||||
|
||||
mock_client_instance = Mock()
|
||||
mock_client_instance.chat.completions.create.return_value = mock_chat_response
|
||||
mock_openai_client.return_value = mock_client_instance
|
||||
|
||||
# Create client and call method
|
||||
messages = [{"role": "user", "content": "Hello"}]
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.create_chat_completion(
|
||||
model="gpt-4", messages=messages, temperature=0.5
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert isinstance(result, LLMResponse)
|
||||
assert result.output == "This is a chat response"
|
||||
assert result.prompt_token_count == 12
|
||||
assert result.response_token_count == 18
|
||||
|
||||
# Verify API call
|
||||
mock_client_instance.chat.completions.create.assert_called_once_with(
|
||||
model="gpt-4", messages=messages, temperature=0.5
|
||||
)
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_operation")
|
||||
@patch("guardrails.utils.openai_utils.v1.openai.Client")
|
||||
def test_create_chat_completion_with_tools(
|
||||
self, mock_openai_client, mock_trace_operation, mock_trace_llm_call
|
||||
):
|
||||
"""Test chat completion with function calling tools."""
|
||||
# Setup mock response
|
||||
mock_chat_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_message = Mock()
|
||||
mock_message.content = "Response with tools"
|
||||
mock_choice.message = mock_message
|
||||
mock_chat_response.choices = [mock_choice]
|
||||
mock_usage = Mock()
|
||||
mock_usage.prompt_tokens = 20
|
||||
mock_usage.completion_tokens = 30
|
||||
mock_usage.total_tokens = 50
|
||||
mock_chat_response.usage = mock_usage
|
||||
|
||||
mock_client_instance = Mock()
|
||||
mock_client_instance.chat.completions.create.return_value = mock_chat_response
|
||||
mock_openai_client.return_value = mock_client_instance
|
||||
|
||||
# Create client and call method
|
||||
messages = [{"role": "user", "content": "Call a function"}]
|
||||
tools = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather info",
|
||||
"parameters": {},
|
||||
},
|
||||
}
|
||||
]
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.create_chat_completion(
|
||||
model="gpt-4", messages=messages, tools=tools
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert isinstance(result, LLMResponse)
|
||||
assert result.output == "Response with tools"
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_operation")
|
||||
@patch("guardrails.utils.openai_utils.v1.openai.Client")
|
||||
def test_create_chat_completion_streaming(
|
||||
self, mock_openai_client, mock_trace_operation, mock_trace_llm_call
|
||||
):
|
||||
"""Test streaming chat completion."""
|
||||
# Setup mock streaming response
|
||||
mock_stream = iter([{"choices": [{"delta": {"content": "chunk"}}]}])
|
||||
|
||||
mock_client_instance = Mock()
|
||||
mock_client_instance.chat.completions.create.return_value = mock_stream
|
||||
mock_openai_client.return_value = mock_client_instance
|
||||
|
||||
# Create client and call method
|
||||
messages = [{"role": "user", "content": "Stream this"}]
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.create_chat_completion(
|
||||
model="gpt-4", messages=messages, stream=True
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert isinstance(result, LLMResponse)
|
||||
assert result.output == ""
|
||||
assert result.stream_output is not None
|
||||
|
||||
|
||||
class TestConstructChatResponse:
|
||||
"""Test construct_chat_response method."""
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
def test_construct_chat_response_with_content(self, mock_trace_llm_call):
|
||||
"""Test chat response construction with message content."""
|
||||
# Setup mock response
|
||||
mock_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_message = Mock()
|
||||
mock_message.content = "Chat response content"
|
||||
mock_choice.message = mock_message
|
||||
mock_response.choices = [mock_choice]
|
||||
mock_usage = Mock()
|
||||
mock_usage.prompt_tokens = 10
|
||||
mock_usage.completion_tokens = 15
|
||||
mock_usage.total_tokens = 25
|
||||
mock_response.usage = mock_usage
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.construct_chat_response(
|
||||
stream=False, openai_response=mock_response
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert isinstance(result, LLMResponse)
|
||||
assert result.output == "Chat response content"
|
||||
assert result.prompt_token_count == 10
|
||||
assert result.response_token_count == 15
|
||||
assert result.stream_output is None
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
def test_construct_chat_response_with_function_call(self, mock_trace_llm_call):
|
||||
"""Test chat response construction with function call."""
|
||||
# Setup mock response
|
||||
mock_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_message = Mock()
|
||||
mock_message.content = None
|
||||
mock_function_call = Mock()
|
||||
mock_function_call.arguments = '{"arg": "value"}'
|
||||
mock_message.function_call = mock_function_call
|
||||
mock_choice.message = mock_message
|
||||
mock_response.choices = [mock_choice]
|
||||
mock_usage = Mock()
|
||||
mock_usage.prompt_tokens = 20
|
||||
mock_usage.completion_tokens = 10
|
||||
mock_usage.total_tokens = 30
|
||||
mock_response.usage = mock_usage
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.construct_chat_response(
|
||||
stream=False, openai_response=mock_response
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert isinstance(result, LLMResponse)
|
||||
assert result.output == '{"arg": "value"}'
|
||||
assert result.prompt_token_count == 20
|
||||
assert result.response_token_count == 10
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
def test_construct_chat_response_with_tool_calls(self, _mock_trace_llm_call):
|
||||
"""Test chat response construction with tool calls."""
|
||||
|
||||
# Create a custom class for message with controlled attribute access
|
||||
class MockMessage:
|
||||
content = None
|
||||
|
||||
@property
|
||||
def function_call(self):
|
||||
raise AttributeError("no function_call")
|
||||
|
||||
@property
|
||||
def tool_calls(self):
|
||||
mock_tool_call = Mock()
|
||||
mock_function = Mock()
|
||||
mock_function.arguments = '{"tool_arg": "tool_value"}'
|
||||
mock_tool_call.function = mock_function
|
||||
return [mock_tool_call]
|
||||
|
||||
# Setup mock response
|
||||
mock_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_choice.message = MockMessage()
|
||||
mock_response.choices = [mock_choice]
|
||||
mock_usage = Mock()
|
||||
mock_usage.prompt_tokens = 25
|
||||
mock_usage.completion_tokens = 15
|
||||
mock_usage.total_tokens = 40
|
||||
mock_response.usage = mock_usage
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.construct_chat_response(
|
||||
stream=False, openai_response=mock_response
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert isinstance(result, LLMResponse)
|
||||
assert result.output == '{"tool_arg": "tool_value"}'
|
||||
assert result.prompt_token_count == 25
|
||||
assert result.response_token_count == 15
|
||||
|
||||
def test_construct_chat_response_streaming(self):
|
||||
"""Test streaming chat response construction."""
|
||||
# Setup mock streaming response
|
||||
mock_stream = iter([{"choices": [{"delta": {"content": "stream chunk"}}]}])
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
result = client.construct_chat_response(
|
||||
stream=True, openai_response=mock_stream
|
||||
)
|
||||
|
||||
# Assertions
|
||||
assert isinstance(result, LLMResponse)
|
||||
assert result.output == ""
|
||||
assert result.stream_output is not None
|
||||
|
||||
def test_construct_chat_response_no_choices_error(self):
|
||||
"""Test error when no choices in response."""
|
||||
# Setup mock response with no choices
|
||||
mock_response = Mock()
|
||||
mock_response.choices = []
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
|
||||
with pytest.raises(ValueError, match="No choices returned from OpenAI"):
|
||||
client.construct_chat_response(stream=False, openai_response=mock_response)
|
||||
|
||||
def test_construct_chat_response_no_message_error(self):
|
||||
"""Test error when no message in response."""
|
||||
# Setup mock response with no message
|
||||
mock_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_choice.message = None
|
||||
mock_response.choices = [mock_choice]
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
|
||||
with pytest.raises(ValueError, match="No message returned from OpenAI"):
|
||||
client.construct_chat_response(stream=False, openai_response=mock_response)
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
def test_construct_chat_response_no_usage_error(self, mock_trace_llm_call):
|
||||
"""Test error when no usage info in response."""
|
||||
# Setup mock response with no usage
|
||||
mock_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_message = Mock()
|
||||
mock_message.content = "Content"
|
||||
mock_choice.message = mock_message
|
||||
mock_response.choices = [mock_choice]
|
||||
mock_response.usage = None
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
|
||||
with pytest.raises(ValueError, match="No token counts returned from OpenAI"):
|
||||
client.construct_chat_response(stream=False, openai_response=mock_response)
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
def test_construct_chat_response_no_content_or_function_error(
|
||||
self, _mock_trace_llm_call
|
||||
):
|
||||
"""Test error when no content, function_call, or tool_calls in
|
||||
response."""
|
||||
|
||||
# Create a custom class for message with controlled attribute access
|
||||
class MockMessage:
|
||||
content = None
|
||||
|
||||
@property
|
||||
def function_call(self):
|
||||
raise AttributeError("no function_call")
|
||||
|
||||
@property
|
||||
def tool_calls(self):
|
||||
raise AttributeError("no tool_calls")
|
||||
|
||||
# Setup mock response with no content or function calls
|
||||
mock_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_choice.message = MockMessage()
|
||||
mock_response.choices = [mock_choice]
|
||||
mock_usage = Mock()
|
||||
mock_usage.prompt_tokens = 10
|
||||
mock_usage.completion_tokens = 5
|
||||
mock_usage.total_tokens = 15
|
||||
mock_response.usage = mock_usage
|
||||
|
||||
# Create client and call method
|
||||
client = OpenAIClientV1(api_key="test-key")
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="No message content or function call arguments returned from OpenAI",
|
||||
):
|
||||
client.construct_chat_response(stream=False, openai_response=mock_response)
|
||||
|
||||
|
||||
class TestIntegrationScenarios:
|
||||
"""Test integration scenarios combining multiple methods."""
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_operation")
|
||||
@patch("guardrails.utils.openai_utils.v1.openai.Client")
|
||||
def test_full_completion_workflow(
|
||||
self, mock_openai_client, mock_trace_operation, mock_trace_llm_call
|
||||
):
|
||||
"""Test complete workflow from client creation to completion."""
|
||||
# Setup mock
|
||||
mock_completion_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_choice.text = "Full workflow response"
|
||||
mock_completion_response.choices = [mock_choice]
|
||||
mock_usage = Mock()
|
||||
mock_usage.prompt_tokens = 5
|
||||
mock_usage.completion_tokens = 10
|
||||
mock_usage.total_tokens = 15
|
||||
mock_completion_response.usage = mock_usage
|
||||
|
||||
mock_client_instance = Mock()
|
||||
mock_client_instance.completions.create.return_value = mock_completion_response
|
||||
mock_openai_client.return_value = mock_client_instance
|
||||
|
||||
# Execute workflow
|
||||
client = OpenAIClientV1(api_key="workflow-test-key")
|
||||
result = client.create_completion(
|
||||
engine="gpt-3.5-turbo-instruct", prompt="Test workflow"
|
||||
)
|
||||
|
||||
# Verify end-to-end result
|
||||
assert result.output == "Full workflow response"
|
||||
assert result.prompt_token_count == 5
|
||||
assert result.response_token_count == 10
|
||||
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_llm_call")
|
||||
@patch("guardrails.utils.openai_utils.v1.trace_operation")
|
||||
@patch("guardrails.utils.openai_utils.v1.openai.Client")
|
||||
def test_full_chat_workflow(
|
||||
self, mock_openai_client, mock_trace_operation, mock_trace_llm_call
|
||||
):
|
||||
"""Test complete workflow for chat completion."""
|
||||
# Setup mock
|
||||
mock_chat_response = Mock()
|
||||
mock_choice = Mock()
|
||||
mock_message = Mock()
|
||||
mock_message.content = "Full chat workflow response"
|
||||
mock_choice.message = mock_message
|
||||
mock_chat_response.choices = [mock_choice]
|
||||
mock_usage = Mock()
|
||||
mock_usage.prompt_tokens = 8
|
||||
mock_usage.completion_tokens = 12
|
||||
mock_usage.total_tokens = 20
|
||||
mock_chat_response.usage = mock_usage
|
||||
|
||||
mock_client_instance = Mock()
|
||||
mock_client_instance.chat.completions.create.return_value = mock_chat_response
|
||||
mock_openai_client.return_value = mock_client_instance
|
||||
|
||||
# Execute workflow
|
||||
client = OpenAIClientV1(api_key="chat-workflow-key")
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant"},
|
||||
{"role": "user", "content": "Hello"},
|
||||
]
|
||||
result = client.create_chat_completion(model="gpt-4", messages=messages)
|
||||
|
||||
# Verify end-to-end result
|
||||
assert result.output == "Full chat workflow response"
|
||||
assert result.prompt_token_count == 8
|
||||
assert result.response_token_count == 12
|
||||
Reference in New Issue
Block a user