mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
"""
|
|
Tests for Z.AI (Zhipu AI) provider - GLM models
|
|
"""
|
|
|
|
import math
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm import completion
|
|
from litellm.cost_calculator import cost_per_token
|
|
|
|
|
|
@pytest.fixture
|
|
def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
|
monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url=""))
|
|
|
|
|
|
@pytest.fixture
|
|
def zai_response():
|
|
"""Mock response from Z.AI API"""
|
|
return {
|
|
"id": "chatcmpl-zai-123",
|
|
"object": "chat.completion",
|
|
"created": 1677652288,
|
|
"model": "glm-4.6",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "Hello! How can I help you today?",
|
|
},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 15, "total_tokens": 25},
|
|
}
|
|
|
|
|
|
def test_get_llm_provider_zai():
|
|
"""Test that get_llm_provider correctly identifies zai provider"""
|
|
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
|
|
|
model, provider, api_key, api_base = get_llm_provider("zai/glm-4.6")
|
|
assert model == "glm-4.6"
|
|
assert provider == "zai"
|
|
assert api_base == "https://api.z.ai/api/paas/v4"
|
|
|
|
|
|
def test_zai_in_provider_lists():
|
|
"""Test that zai is registered in all necessary provider lists"""
|
|
assert "zai" in litellm.openai_compatible_providers
|
|
assert "zai" in litellm.provider_list
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_zai_completion_call(respx_mock, zai_response, monkeypatch):
|
|
"""Test completion call with zai provider using mocked response"""
|
|
monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
|
|
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
|
|
|
respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(
|
|
json=zai_response
|
|
)
|
|
|
|
response = await litellm.acompletion(
|
|
model="zai/glm-4.6",
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
max_tokens=20,
|
|
)
|
|
|
|
assert response.choices[0].message.content == "Hello! How can I help you today?"
|
|
assert response.usage.total_tokens == 25
|
|
|
|
assert len(respx_mock.calls) == 1
|
|
request = respx_mock.calls[0].request
|
|
assert request.method == "POST"
|
|
assert "api.z.ai" in str(request.url)
|
|
assert "Authorization" in request.headers
|
|
assert request.headers["Authorization"] == "Bearer test-api-key"
|
|
|
|
|
|
def test_zai_sync_completion(respx_mock, zai_response, monkeypatch):
|
|
"""Test synchronous completion call"""
|
|
monkeypatch.setenv("ZAI_API_KEY", "test-api-key")
|
|
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
|
|
|
respx_mock.post("https://api.z.ai/api/paas/v4/chat/completions").respond(
|
|
json=zai_response
|
|
)
|
|
|
|
response = completion(
|
|
model="zai/glm-4.6",
|
|
messages=[{"role": "user", "content": "Hello"}],
|
|
max_tokens=20,
|
|
)
|
|
|
|
assert response.choices[0].message.content == "Hello! How can I help you today?"
|
|
assert response.usage.total_tokens == 25
|