Added LiteLLM to the stack
This commit is contained in:
@@ -0,0 +1,318 @@
|
||||
"""
|
||||
Unit tests for CometAPI Chat Configuration
|
||||
|
||||
Tests the CometAPIChatConfig class methods using mocks
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.llms.cometapi.chat.transformation import (
|
||||
CometAPIChatCompletionStreamingHandler,
|
||||
CometAPIConfig,
|
||||
)
|
||||
from litellm.llms.cometapi.common_utils import CometAPIException
|
||||
|
||||
|
||||
class TestCometAPIChatCompletionStreamingHandler:
|
||||
def test_chunk_parser_successful(self):
|
||||
handler = CometAPIChatCompletionStreamingHandler(
|
||||
streaming_response=None, sync_stream=True
|
||||
)
|
||||
|
||||
# Test input chunk
|
||||
chunk = {
|
||||
"id": "test_id",
|
||||
"created": 1234567890,
|
||||
"model": "gpt-3.5-turbo",
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
||||
"choices": [
|
||||
{"delta": {"content": "test content", "reasoning": "test reasoning"}}
|
||||
],
|
||||
}
|
||||
|
||||
# Parse chunk
|
||||
result = handler.chunk_parser(chunk)
|
||||
|
||||
# Verify response
|
||||
assert result.id == "test_id"
|
||||
assert result.object == "chat.completion.chunk"
|
||||
assert result.created == 1234567890
|
||||
assert result.model == "gpt-3.5-turbo"
|
||||
assert result.usage.prompt_tokens == chunk["usage"]["prompt_tokens"]
|
||||
assert result.usage.completion_tokens == chunk["usage"]["completion_tokens"]
|
||||
assert result.usage.total_tokens == chunk["usage"]["total_tokens"]
|
||||
assert len(result.choices) == 1
|
||||
assert result.choices[0]["delta"]["reasoning_content"] == "test reasoning"
|
||||
|
||||
def test_chunk_parser_error_response(self):
|
||||
handler = CometAPIChatCompletionStreamingHandler(
|
||||
streaming_response=None, sync_stream=True
|
||||
)
|
||||
|
||||
# Test error chunk
|
||||
error_chunk = {
|
||||
"error": {
|
||||
"message": "test error",
|
||||
"code": 400,
|
||||
}
|
||||
}
|
||||
|
||||
# Verify error handling
|
||||
with pytest.raises(CometAPIException) as exc_info:
|
||||
handler.chunk_parser(error_chunk)
|
||||
|
||||
assert "CometAPI Error: test error" in str(exc_info.value)
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
def test_chunk_parser_key_error(self):
|
||||
handler = CometAPIChatCompletionStreamingHandler(
|
||||
streaming_response=None, sync_stream=True
|
||||
)
|
||||
|
||||
# Test invalid chunk missing required fields
|
||||
invalid_chunk = {"incomplete": "data"}
|
||||
|
||||
# Verify KeyError handling
|
||||
with pytest.raises(CometAPIException) as exc_info:
|
||||
handler.chunk_parser(invalid_chunk)
|
||||
|
||||
assert "KeyError" in str(exc_info.value)
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
|
||||
class TestCometAPIConfig:
|
||||
def test_transform_request_basic(self):
|
||||
"""Test basic request transformation"""
|
||||
config = CometAPIConfig()
|
||||
|
||||
transformed_request = config.transform_request(
|
||||
model="cometapi/gpt-3.5-turbo",
|
||||
messages=[
|
||||
{"role": "user", "content": "Hello, world!"}
|
||||
],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert transformed_request["model"] == "cometapi/gpt-3.5-turbo"
|
||||
assert transformed_request["messages"] == [
|
||||
{"role": "user", "content": "Hello, world!"}
|
||||
]
|
||||
|
||||
def test_transform_request_with_extra_body(self):
|
||||
"""Test request transformation with extra_body parameters"""
|
||||
config = CometAPIConfig()
|
||||
|
||||
transformed_request = config.transform_request(
|
||||
model="cometapi/gpt-4",
|
||||
messages=[{"role": "user", "content": "Hello, world!"}],
|
||||
optional_params={"extra_body": {"custom_param": "custom_value"}},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
# Validate that extra_body parameters are merged into the request
|
||||
assert transformed_request["custom_param"] == "custom_value"
|
||||
assert transformed_request["messages"] == [
|
||||
{"role": "user", "content": "Hello, world!"}
|
||||
]
|
||||
|
||||
def test_cache_control_flag_removal(self):
|
||||
"""Test cache control flag removal from messages"""
|
||||
config = CometAPIConfig()
|
||||
|
||||
transformed_request = config.transform_request(
|
||||
model="cometapi/gpt-3.5-turbo",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Hello, world!",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
# CometAPI should remove cache_control flags by default
|
||||
assert transformed_request["messages"][0].get("cache_control") is None
|
||||
|
||||
def test_map_openai_params(self):
|
||||
"""Test OpenAI parameter mapping"""
|
||||
config = CometAPIConfig()
|
||||
|
||||
non_default_params = {
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 100,
|
||||
"top_p": 0.9,
|
||||
}
|
||||
|
||||
mapped_params = config.map_openai_params(
|
||||
non_default_params=non_default_params,
|
||||
optional_params={},
|
||||
model="cometapi/gpt-3.5-turbo",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
assert mapped_params["temperature"] == 0.7
|
||||
assert mapped_params["max_tokens"] == 100
|
||||
assert mapped_params["top_p"] == 0.9
|
||||
|
||||
def test_get_error_class(self):
|
||||
"""Test error class creation"""
|
||||
config = CometAPIConfig()
|
||||
|
||||
error = config.get_error_class(
|
||||
error_message="Test error",
|
||||
status_code=400,
|
||||
headers={"Content-Type": "application/json"}
|
||||
)
|
||||
|
||||
assert isinstance(error, CometAPIException)
|
||||
assert error.message == "Test error"
|
||||
assert error.status_code == 400
|
||||
|
||||
|
||||
# Integration test example (requires real API key)
|
||||
@pytest.mark.skip(reason="Skipping integration test")
|
||||
def test_cometapi_integration():
|
||||
"""
|
||||
Integration test - requires real API key
|
||||
Run with: pytest -k test_cometapi_integration -s
|
||||
"""
|
||||
import os
|
||||
from litellm import completion
|
||||
|
||||
# Try to get API key from multiple environment variables
|
||||
api_key = (
|
||||
os.getenv("COMETAPI_API_KEY")
|
||||
or os.getenv("COMETAPI_KEY")
|
||||
or os.getenv("COMET_API_KEY")
|
||||
)
|
||||
|
||||
if not api_key:
|
||||
pytest.skip("COMETAPI_API_KEY not set - skipping integration test")
|
||||
|
||||
response = completion(
|
||||
model="cometapi/gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "Say hello in one word"}],
|
||||
api_key=api_key,
|
||||
max_tokens=10,
|
||||
temperature=0.7
|
||||
)
|
||||
|
||||
# Verify response structure
|
||||
assert response.choices[0].message.content
|
||||
assert len(response.choices[0].message.content.strip()) > 0
|
||||
assert response.model
|
||||
assert response.usage
|
||||
assert response.usage.total_tokens > 0
|
||||
|
||||
|
||||
def test_cometapi_streaming_integration():
|
||||
"""
|
||||
Integration test for streaming - requires real API key
|
||||
Run with: pytest -k test_cometapi_streaming_integration -s
|
||||
"""
|
||||
import os
|
||||
from litellm import completion
|
||||
|
||||
# Try to get API key from multiple environment variables
|
||||
api_key = (
|
||||
os.getenv("COMETAPI_API_KEY")
|
||||
or os.getenv("COMETAPI_KEY")
|
||||
or os.getenv("COMET_API_KEY")
|
||||
)
|
||||
|
||||
if not api_key:
|
||||
pytest.skip("COMETAPI_API_KEY not set - skipping streaming integration test")
|
||||
|
||||
try:
|
||||
print(f"🔍 Testing streaming with API key: {api_key[:6]}...{api_key[-4:]} (length: {len(api_key)})")
|
||||
print(f"🔍 API base URL: {os.getenv('COMETAPI_API_BASE', 'default')}")
|
||||
|
||||
# test streaming API call
|
||||
response = completion(
|
||||
model="cometapi/gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "Count from 1 to 5"}],
|
||||
api_key=api_key,
|
||||
max_tokens=50,
|
||||
stream=True
|
||||
)
|
||||
|
||||
# collect streaming response
|
||||
chunks = []
|
||||
content_parts = []
|
||||
|
||||
for chunk in response:
|
||||
chunks.append(chunk)
|
||||
if chunk.choices[0].delta.content:
|
||||
content_parts.append(chunk.choices[0].delta.content)
|
||||
|
||||
# Verify we received at least one chunk and content
|
||||
assert len(chunks) > 0, "Should receive at least one chunk"
|
||||
assert len(content_parts) > 0, "Should receive content in chunks"
|
||||
|
||||
full_content = "".join(content_parts)
|
||||
assert len(full_content.strip()) > 0, "Should have non-empty content"
|
||||
|
||||
print(f"✅ Received {len(chunks)} chunks")
|
||||
print(f"✅ Full content: {full_content}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Streaming integration test error details:")
|
||||
print(f" Error type: {type(e).__name__}")
|
||||
print(f" Error message: {str(e)}")
|
||||
if hasattr(e, 'status_code'):
|
||||
print(f" Status code: {e.status_code}")
|
||||
if hasattr(e, 'response'):
|
||||
print(f" Response: {e.response}")
|
||||
|
||||
# Re-raise with more context for pytest
|
||||
pytest.fail(f"Streaming integration test failed: {type(e).__name__}: {str(e)}")
|
||||
def test_cometapi_with_custom_base_url():
|
||||
"""
|
||||
Test CometAPI with custom base URL
|
||||
"""
|
||||
import os
|
||||
from litellm import completion
|
||||
|
||||
api_key = (
|
||||
os.getenv("COMETAPI_API_KEY")
|
||||
or os.getenv("COMETAPI_KEY")
|
||||
or os.getenv("COMET_API_KEY")
|
||||
)
|
||||
|
||||
custom_base_url = os.getenv("COMETAPI_API_BASE", "https://api.cometapi.com/v1")
|
||||
|
||||
if not api_key:
|
||||
pytest.skip("COMETAPI_API_KEY not set - skipping custom base URL test")
|
||||
|
||||
try:
|
||||
response = completion(
|
||||
model="cometapi/gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
api_key=api_key,
|
||||
api_base=custom_base_url,
|
||||
max_tokens=5
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content
|
||||
print(f"✅ Custom base URL test passed: {response.choices[0].message.content}")
|
||||
|
||||
except Exception as e:
|
||||
pytest.fail(f"Custom base URL test failed: {str(e)}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Quick test runner
|
||||
pytest.main([__file__, "-v"])
|
||||
Reference in New Issue
Block a user