scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 100

test(harness): move long-tail live chat tests to chat_live_longtail suite

Browse source 
One suite of env-gated cells per audit 8c:
- moved whole: test_mistral_api.py, test_openrouter.py, test_snowflake.py,
  test_a2a.py
- test_cohere.py: the 12 live v1/v2 chat/citations/tool tests; the 9 live
  embed v4 tests stay (non-chat, Sameer)
- test_groq.py: TestGroq; structured-output and chunk_parser goldens stay
- test_together_ai.py: TestTogetherAI; supported-params golden stays
- test_xai.py: TestXAIChat, TestXAIReasoningEffort,
  test_xai_message_name_filtering, test_xai_streaming_with_include_usage;
  5 param-mapping goldens stay
- test_deepseek_completion.py: TestDeepSeekChatCompletion (@skip, API
  hanging), test_completion_cost_deepseek; fill_reasoning_content goldens stay
- test_lambda_ai.py / test_v0.py: the env-gated live completion call each;
  config/registry goldens stay
- test_langgraph.py: 2 live local-server tests; config transforms stay
This commit is contained in:
mateo-berri 2026-06-11 19:14:51 +00:00
parent 854f370fe7
commit d5d69b1569
17 changed files with 949 additions and 857 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

518
tests/harness_suites/chat_live_longtail/test_cohere_chat_live.py Normal file
View file

@ -0,0 +1,518 @@
import os
import sys
from dotenv import load_dotenv
load_dotenv()
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import pytest
import litellm
from litellm import completion, embedding
litellm.num_retries = 3
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_chat_completion_cohere_citations(stream):
try:
litellm.set_verbose = True
messages = [
{
"role": "user",
"content": "Which penguins are the tallest?",
},
]
response = await litellm.acompletion(
model="cohere_chat/v1/command-r",
messages=messages,
documents=[
{"title": "Tall penguins", "text": "Emperor penguins are the tallest."},
{
"title": "Penguin habitats",
"text": "Emperor penguins only live in Antarctica.",
},
],
stream=stream,
)
if stream:
citations_chunk = False
async for chunk in response:
print("received chunk", chunk)
if "citations" in chunk:
citations_chunk = True
break
assert citations_chunk
else:
assert response.citations is not None
except litellm.ServiceUnavailableError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_cohere_command_r_plus_function_call():
litellm.set_verbose = True
tools = [
{
"type": "function",
"function": {
"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": ["celsius", "fahrenheit"]},
},
"required": ["location"],
},
},
}
]
messages = [
{
"role": "user",
"content": "What's the weather like in Boston today in Fahrenheit?",
}
]
try:
# test without max tokens
response = completion(
model="cohere_chat/v1/command-r-plus",
messages=messages,
tools=tools,
tool_choice="auto",
)
# Add any assertions, here to check response args
print(response)
assert isinstance(response.choices[0].message.tool_calls[0].function.name, str)
assert isinstance(
response.choices[0].message.tool_calls[0].function.arguments, str
)
except litellm.Timeout:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_cohere():
try:
# litellm.set_verbose=True
messages = [
{"role": "system", "content": "You're a good bot"},
{"role": "assistant", "content": [{"text": "2", "type": "text"}]},
{"role": "assistant", "content": [{"text": "3", "type": "text"}]},
{
"role": "user",
"content": "Hey",
},
]
response = completion(
model="cohere_chat/v1/command-r",
messages=messages,
)
print(response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False])
async def test_chat_completion_cohere(sync_mode):
try:
litellm.set_verbose = True
messages = [
{"role": "system", "content": "You're a good bot"},
{
"role": "user",
"content": "Hey",
},
]
if sync_mode is False:
response = await litellm.acompletion(
model="cohere_chat/v1/command-r",
messages=messages,
max_tokens=10,
)
else:
response = completion(
model="cohere_chat/v1/command-r",
messages=messages,
max_tokens=10,
)
print(response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [False])
async def test_chat_completion_cohere_stream(sync_mode):
try:
litellm.set_verbose = True
messages = [
{"role": "system", "content": "You're a good bot"},
{
"role": "user",
"content": "Hey",
},
]
if sync_mode is False:
response = await litellm.acompletion(
model="cohere_chat/v1/command-r",
messages=messages,
max_tokens=10,
stream=True,
)
print("async cohere stream response", response)
async for chunk in response:
print(chunk)
else:
response = completion(
model="cohere_chat/v1/command-r",
messages=messages,
max_tokens=10,
stream=True,
)
print(response)
for chunk in response:
print(chunk)
except litellm.APIConnectionError as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
async def test_cohere_v2_chat_completion(sync_mode):
"""Test basic Cohere v2 chat completion functionality."""
try:
litellm.set_verbose = True
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello, how are you?"},
]
if sync_mode:
response = completion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
max_tokens=50,
)
else:
response = await litellm.acompletion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
max_tokens=50,
)
# Validate response structure
assert response.choices is not None
assert len(response.choices) > 0
assert response.choices[0].message.content is not None
assert response.usage is not None
assert response.usage.total_tokens > 0
print(f"Cohere v2 response: {response}")
except litellm.ServiceUnavailableError:
pass # Skip if service is unavailable
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_cohere_v2_streaming(stream):
"""Test Cohere v2 streaming functionality."""
try:
litellm.set_verbose = True
messages = [{"role": "user", "content": "Tell me a short story about a robot."}]
response = await litellm.acompletion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
max_tokens=100,
stream=stream,
)
if stream:
# Test streaming response
chunks = []
async for chunk in response:
chunks.append(chunk)
if len(chunks) >= 3: # Test first few chunks
break
assert len(chunks) > 0
print(f"Received {len(chunks)} streaming chunks")
else:
# Test non-streaming response
assert response.choices is not None
assert len(response.choices) > 0
assert response.choices[0].message.content is not None
print(f"Non-streaming response: {response.choices[0].message.content}")
except litellm.ServiceUnavailableError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_cohere_v2_tool_calling():
"""Test Cohere v2 tool calling functionality."""
try:
litellm.set_verbose = True
tools = [
{
"type": "function",
"function": {
"name": "get_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": ["celsius", "fahrenheit"],
},
},
"required": ["location"],
},
},
}
]
messages = [{"role": "user", "content": "What's the weather like in New York?"}]
response = completion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
tools=tools,
tool_choice="auto",
max_tokens=100,
)
# Validate tool calling response
assert response.choices is not None
assert len(response.choices) > 0
message = response.choices[0].message
# Check if tool calls are present
if hasattr(message, "tool_calls") and message.tool_calls:
assert len(message.tool_calls) > 0
tool_call = message.tool_calls[0]
assert tool_call.function.name == "get_weather"
assert tool_call.function.arguments is not None
print(
f"Tool call: {tool_call.function.name} - {tool_call.function.arguments}"
)
else:
# If no tool calls, check that we got a regular response
assert message.content is not None
print(f"Regular response: {message.content}")
except litellm.ServiceUnavailableError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
async def test_cohere_v2_annotations(stream):
"""Test Cohere v2 annotations functionality (replaces citations)."""
try:
litellm.set_verbose = True
messages = [
{"role": "user", "content": "What are the benefits of renewable energy?"}
]
documents = [
{
"data": {
"title": "Renewable Energy Benefits Document",
"snippet": "Renewable energy sources like solar and wind power provide clean electricity while reducing greenhouse gas emissions and dependence on fossil fuels.",
}
},
{
"data": {
"title": "Environmental Impact Study",
"snippet": "Studies show that renewable energy significantly reduces carbon footprint and helps combat climate change.",
}
},
]
response = await litellm.acompletion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
documents=documents,
max_tokens=100,
stream=stream,
)
if stream:
# Test streaming with annotations
annotations_found = False
async for chunk in response:
# Check if chunk has a message with annotations
if (
hasattr(chunk, "choices")
and chunk.choices
and len(chunk.choices) > 0
and hasattr(chunk.choices[0], "message")
and hasattr(chunk.choices[0].message, "annotations")
and chunk.choices[0].message.annotations
):
annotations_found = True
print(
f"Streaming annotations: {chunk.choices[0].message.annotations}"
)
break
# Note: Annotations might not appear in every chunk during streaming
else:
# Test non-streaming with annotations
assert response.choices is not None
assert len(response.choices) > 0
# Check for annotations in message
message = response.choices[0].message
if hasattr(message, "annotations") and message.annotations:
assert len(message.annotations) > 0
print(f"Annotations found: {len(message.annotations)}")
# Validate annotation structure
for annotation in message.annotations:
assert (
annotation.get("type") == "url_citation"
), f"Expected type 'url_citation', got {annotation.get('type')}"
assert "url_citation" in annotation, "Missing url_citation field"
url_citation = annotation["url_citation"]
assert "start_index" in url_citation, "Missing start_index"
assert "end_index" in url_citation, "Missing end_index"
assert "title" in url_citation, "Missing title"
assert "url" in url_citation, "Missing url"
print(f"First annotation: {message.annotations[0]}")
else:
# Annotations might not always be present depending on the response
print("No annotations in this response")
# Ensure citations field is NOT present (removed backward compatibility)
assert not hasattr(
response, "citations"
), "Citations field should be removed - no backward compatibility"
except litellm.ServiceUnavailableError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_cohere_v2_parameter_mapping():
"""Test Cohere v2 parameter mapping and validation."""
try:
litellm.set_verbose = True
messages = [{"role": "user", "content": "Generate a creative story."}]
# Test various parameters that should be mapped correctly
response = completion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
temperature=0.7,
max_tokens=50,
top_p=0.9,
frequency_penalty=0.1,
presence_penalty=0.1,
stop=["END", "STOP"],
seed=42,
)
# Validate response
assert response.choices is not None
assert len(response.choices) > 0
assert response.choices[0].message.content is not None
assert response.usage is not None
print(f"Parameter mapping test response: {response.choices[0].message.content}")
except litellm.ServiceUnavailableError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_cohere_v2_error_handling():
"""Test Cohere v2 error handling with invalid parameters."""
try:
# Test with invalid model name
try:
response = completion(
model="cohere_chat/v2/invalid-model",
messages=[{"role": "user", "content": "Hello"}],
max_tokens=10,
)
# If we get here, the test should fail
pytest.fail("Should have failed with invalid model")
except Exception as e:
# Expected to fail with invalid model
print(f"Expected error with invalid model: {e}")
# Test with empty messages
try:
response = completion(
model="cohere_chat/v2/command-a-03-2025",
messages=[], # Empty messages
max_tokens=10,
)
pytest.fail("Should have failed with empty messages")
except Exception as e:
# Expected to fail with empty messages
print(f"Expected error with empty messages: {e}")
except Exception as e:
pytest.fail(f"Unexpected error in error handling test: {e}")
@pytest.mark.asyncio
async def test_cohere_v2_conversation_history():
"""Test Cohere v2 with conversation history."""
try:
litellm.set_verbose = True
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is 2+2?"},
{"role": "assistant", "content": "2+2 equals 4."},
{"role": "user", "content": "What about 3+3?"},
]
response = await litellm.acompletion(
model="cohere_chat/v2/command-a-03-2025", messages=messages, max_tokens=50
)
# Validate response with conversation history
assert response.choices is not None
assert len(response.choices) > 0
assert response.choices[0].message.content is not None
print(f"Conversation history response: {response.choices[0].message.content}")
except (
litellm.ServiceUnavailableError,
litellm.InternalServerError,
litellm.Timeout,
litellm.APIConnectionError,
):
pytest.skip("Cohere service unavailable")
except litellm.RateLimitError:
pytest.skip("Rate limit exceeded")

84
tests/harness_suites/chat_live_longtail/test_deepseek_live.py Normal file
View file

@ -0,0 +1,84 @@
from base_llm_unit_tests import BaseLLMChatTest
import pytest
import litellm
# Test implementations
@pytest.mark.skip(reason="Deepseek API is hanging")
class TestDeepSeekChatCompletion(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
return {
"model": "deepseek/deepseek-reasoner",
}
def test_completion_cost_deepseek():
litellm.set_verbose = True
model_name = "deepseek/deepseek-chat"
messages_1 = [
{
"role": "system",
"content": "You are a history expert. The user will provide a series of questions, and your answers should be concise and start with `Answer:`",
},
{
"role": "user",
"content": "In what year did Qin Shi Huang unify the six states?",
},
{"role": "assistant", "content": "Answer: 221 BC"},
{"role": "user", "content": "Who was the founder of the Han Dynasty?"},
{"role": "assistant", "content": "Answer: Liu Bang"},
{"role": "user", "content": "Who was the last emperor of the Tang Dynasty?"},
{"role": "assistant", "content": "Answer: Li Zhu"},
{
"role": "user",
"content": "Who was the founding emperor of the Ming Dynasty?",
},
{"role": "assistant", "content": "Answer: Zhu Yuanzhang"},
{
"role": "user",
"content": "Who was the founding emperor of the Qing Dynasty?",
},
]
message_2 = [
{
"role": "system",
"content": "You are a history expert. The user will provide a series of questions, and your answers should be concise and start with `Answer:`",
},
{
"role": "user",
"content": "In what year did Qin Shi Huang unify the six states?",
},
{"role": "assistant", "content": "Answer: 221 BC"},
{"role": "user", "content": "Who was the founder of the Han Dynasty?"},
{"role": "assistant", "content": "Answer: Liu Bang"},
{"role": "user", "content": "Who was the last emperor of the Tang Dynasty?"},
{"role": "assistant", "content": "Answer: Li Zhu"},
{
"role": "user",
"content": "Who was the founding emperor of the Ming Dynasty?",
},
{"role": "assistant", "content": "Answer: Zhu Yuanzhang"},
{"role": "user", "content": "When did the Shang Dynasty fall?"},
]
try:
response_1 = litellm.completion(model=model_name, messages=messages_1)
response_2 = litellm.completion(model=model_name, messages=message_2)
# Add any assertions here to check the response
print(response_2)
assert response_2.usage.prompt_cache_hit_tokens is not None
assert response_2.usage.prompt_cache_miss_tokens is not None
assert (
response_2.usage.prompt_tokens
== response_2.usage.prompt_cache_miss_tokens
+ response_2.usage.prompt_cache_hit_tokens
)
assert (
response_2.usage._cache_read_input_tokens
== response_2.usage.prompt_cache_hit_tokens
)
except litellm.APIError as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")

35
tests/harness_suites/chat_live_longtail/test_groq_live.py Normal file
View file

@ -0,0 +1,35 @@
import os
import sys
import pytest
# sys.path.insert(
# 0, os.path.abspath("../..")
# ) # noqa
# ) # Adds the parent directory to the system path
import litellm
from base_llm_unit_tests import BaseLLMChatTest
from litellm.llms.groq.chat.transformation import (
GroqChatConfig,
GroqChatCompletionStreamingHandler,
)
class TestGroq(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
return {
"model": "groq/llama-3.3-70b-versatile",
}
def test_tool_call_with_empty_enum_property(self):
pass
@pytest.mark.parametrize(
"model",
["groq/qwen/qwen3-32b", "groq/openai/gpt-oss-20b", "groq/openai/gpt-oss-120b"],
)
def test_reasoning_effort_in_supported_params(self, model):
"""Test that reasoning_effort is in the list of supported parameters for Groq"""
supported_params = GroqChatConfig().get_supported_openai_params(model=model)
assert "reasoning_effort" in supported_params

35
tests/harness_suites/chat_live_longtail/test_lambda_ai_live.py Normal file
View file

@ -0,0 +1,35 @@
"""
Tests for Lambda AI provider integration
"""
import os
from unittest import mock
import pytest
import litellm
from litellm import completion
from litellm.llms.lambda_ai.chat.transformation import LambdaAIChatConfig
@pytest.mark.asyncio
async def test_lambda_ai_completion_call():
"""Test completion call with Lambda AI provider (requires LAMBDA_API_KEY)"""
# Skip if no API key is available
if not os.getenv("LAMBDA_API_KEY"):
pytest.skip("LAMBDA_API_KEY not set")
try:
response = await litellm.acompletion(
model="lambda_ai/llama3.1-8b-instruct",
messages=[{"role": "user", "content": "Hello, this is a test"}],
max_tokens=10,
)
assert response.choices[0].message.content
assert response.model
assert response.usage
except Exception as e:
# If the API key is invalid or there's a network issue, that's okay
# The important thing is that the provider was recognized
if "lambda_ai" not in str(e) and "provider" not in str(e).lower():
# Re-raise if it's not a provider-related error
raise

88
tests/harness_suites/chat_live_longtail/test_langgraph_live.py Normal file
View file

@ -0,0 +1,88 @@
"""
Tests for LangGraph provider integration.
These tests require a LangGraph server running locally on port 2024.
To start a LangGraph server, follow the LangGraph documentation.
Example test server curl commands:
Streaming:
curl -s --request POST \
--url "http://localhost:2024/runs/stream" \
--header 'Content-Type: application/json' \
--data '{"assistant_id": "agent", "input": {"messages": [{"role": "human", "content": "What is 25 * 4?"}]}, "stream_mode": "messages-tuple"}'
Non-streaming:
curl -s --request POST \
--url "http://localhost:2024/runs/wait" \
--header 'Content-Type: application/json' \
--data '{"assistant_id": "agent", "input": {"messages": [{"role": "human", "content": "What is 25 * 4?"}]}}'
"""
import os
import sys
sys.path.insert(0, os.path.abspath("../.."))
import pytest
import litellm
@pytest.mark.asyncio
async def test_langgraph_acompletion_non_streaming():
"""
Test non-streaming acompletion call to LangGraph server.
Uses the /runs/wait endpoint for synchronous response.
"""
api_base = os.environ.get("LANGGRAPH_API_BASE", "http://localhost:2024")
try:
response = await litellm.acompletion(
model="langgraph/agent",
messages=[{"role": "user", "content": "What is 25 * 4?"}],
api_base=api_base,
stream=False,
)
assert response is not None
assert response.choices is not None
assert len(response.choices) > 0
assert response.choices[0].message is not None
assert response.choices[0].message.content is not None
assert len(response.choices[0].message.content) > 0
except Exception as e:
pytest.skip(f"LangGraph server not available: {e}")
@pytest.mark.asyncio
async def test_langgraph_acompletion_streaming():
"""
Test streaming acompletion call to LangGraph server.
Uses the /runs/stream endpoint with stream_mode="messages-tuple".
"""
api_base = os.environ.get("LANGGRAPH_API_BASE", "http://localhost:2024")
try:
response = await litellm.acompletion(
model="langgraph/agent",
messages=[{"role": "user", "content": "What is the weather in Tokyo?"}],
api_base=api_base,
stream=True,
)
full_content = ""
chunk_count = 0
async for chunk in response:
chunk_count += 1
if (
chunk.choices
and chunk.choices[0].delta
and chunk.choices[0].delta.content
):
full_content += chunk.choices[0].delta.content
assert chunk_count > 0, "Should receive at least one chunk"
except Exception as e:
pytest.skip(f"LangGraph server not available: {e}")

1
tests/harness_suites/chat_live_longtail/test_mistral_api.py
View file

@ -28,7 +28,6 @@ from httpx import Headers
from base_llm_unit_tests import BaseLLMChatTest
@pytest.mark.flaky(retries=3, delay=2)
class TestMistralCompletion(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
litellm.set_verbose = True

17
tests/harness_suites/chat_live_longtail/test_together_ai_live.py Normal file
View file

@ -0,0 +1,17 @@
"""
Test TogetherAI LLM
"""
import os
import sys
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from base_llm_unit_tests import BaseLLMChatTest
class TestTogetherAI(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
litellm.set_verbose = True
return {"model": "together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo"}

35
tests/harness_suites/chat_live_longtail/test_v0_live.py Normal file
View file

@ -0,0 +1,35 @@
"""
Tests for v0 provider integration
"""
import os
from unittest import mock
import pytest
import litellm
from litellm import completion
from litellm.llms.v0.chat.transformation import V0ChatConfig
@pytest.mark.asyncio
async def test_v0_completion_call():
"""Test completion call with v0 provider (requires V0_API_KEY)"""
# Skip if no API key is available
if not os.getenv("V0_API_KEY"):
pytest.skip("V0_API_KEY not set")
try:
response = await litellm.acompletion(
model="v0/gpt-4-turbo",
messages=[{"role": "user", "content": "Hello, this is a test"}],
max_tokens=10,
)
assert response.choices[0].message.content
assert response.model
assert response.usage
except Exception as e:
# If the API key is invalid or there's a network issue, that's okay
# The important thing is that the provider was recognized
if "v0" not in str(e) and "provider" not in str(e).lower():
# Re-raise if it's not a provider-related error
raise

133
tests/harness_suites/chat_live_longtail/test_xai_live.py Normal file
View file

@ -0,0 +1,133 @@
import os
import sys
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import pytest
import litellm
from litellm import completion
from unittest.mock import patch
from litellm.llms.xai.chat.transformation import XAIChatConfig, XAI_API_BASE
from base_llm_unit_tests import BaseReasoningLLMTests, BaseLLMChatTest
def test_xai_message_name_filtering():
messages = [
{
"role": "system",
"content": "*I press the green button*",
"name": "example_user",
},
{"role": "user", "content": "Hello", "name": "John"},
{"role": "assistant", "content": "Hello", "name": "Jane"},
]
response = completion(
model="xai/grok-3-mini-beta",
messages=messages,
)
assert response is not None
assert response.choices[0].message.content is not None
class TestXAIReasoningEffort(BaseReasoningLLMTests):
def get_base_completion_call_args(self):
return {
"model": "xai/grok-3-mini-beta",
"messages": [{"role": "user", "content": "Hello"}],
}
class TestXAIChat(BaseLLMChatTest):
def get_base_completion_call_args(self):
return {
"model": "xai/grok-3-mini-beta",
}
def test_web_search(self):
"""Web search is only supported for Grok 4 family models"""
from litellm.utils import supports_web_search
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm._turn_on_debug()
# Use grok-4-1-fast which supports web search
model = "xai/grok-4-1-fast"
if not supports_web_search(model, None):
pytest.skip("Model does not support web search")
response = completion(
model=model,
messages=[
{"role": "user", "content": "What's the weather like in Boston today?"}
],
web_search_options={},
max_tokens=100,
)
assert response is not None
def test_xai_streaming_with_include_usage():
"""
Test that xAI streaming correctly handles usage in the last chunk
when stream_options={"include_usage": True} is set.
xAI sends usage in a chunk with empty choices array, which should be
handled by XAIChatCompletionStreamingHandler.
"""
try:
response = completion(
model="xai/grok-4-1-fast-non-reasoning",
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Say hello in one word"},
],
stream=True,
stream_options={"include_usage": True},
max_tokens=10,
)
chunks = []
usage_chunk = None
for chunk in response:
chunks.append(chunk)
if hasattr(chunk, "usage") and chunk.usage is not None:
usage_chunk = chunk
# Verify we got chunks
assert len(chunks) > 0, "Should receive streaming chunks"
# Verify usage was included in one of the chunks
assert usage_chunk is not None, "Should receive usage in streaming chunks"
# Verify usage has expected fields
assert hasattr(
usage_chunk.usage, "prompt_tokens"
), "Usage should have prompt_tokens"
assert hasattr(
usage_chunk.usage, "completion_tokens"
), "Usage should have completion_tokens"
assert hasattr(
usage_chunk.usage, "total_tokens"
), "Usage should have total_tokens"
# Verify usage values are positive
assert usage_chunk.usage.prompt_tokens > 0, "prompt_tokens should be positive"
assert (
usage_chunk.usage.completion_tokens > 0
), "completion_tokens should be positive"
assert usage_chunk.usage.total_tokens > 0, "total_tokens should be positive"
print(f"✓ Successfully received usage in streaming chunk: {usage_chunk.usage}")
except Exception as e:
if "API key" in str(e) or "authentication" in str(e).lower():
pytest.skip(f"Skipping test due to API key issue: {str(e)}")
raise

511
tests/llm_translation/test_cohere.py
View file

@ -12,192 +12,11 @@ sys.path.insert(
import pytest
import litellm
from litellm import completion, embedding
from litellm import embedding
litellm.num_retries = 3
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.flaky(retries=3, delay=1)
@pytest.mark.asyncio
async def test_chat_completion_cohere_citations(stream):
try:
litellm.set_verbose = True
messages = [
{
"role": "user",
"content": "Which penguins are the tallest?",
},
]
response = await litellm.acompletion(
model="cohere_chat/v1/command-r",
messages=messages,
documents=[
{"title": "Tall penguins", "text": "Emperor penguins are the tallest."},
{
"title": "Penguin habitats",
"text": "Emperor penguins only live in Antarctica.",
},
],
stream=stream,
)
if stream:
citations_chunk = False
async for chunk in response:
print("received chunk", chunk)
if "citations" in chunk:
citations_chunk = True
break
assert citations_chunk
else:
assert response.citations is not None
except litellm.ServiceUnavailableError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_completion_cohere_command_r_plus_function_call():
litellm.set_verbose = True
tools = [
{
"type": "function",
"function": {
"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": ["celsius", "fahrenheit"]},
},
"required": ["location"],
},
},
}
]
messages = [
{
"role": "user",
"content": "What's the weather like in Boston today in Fahrenheit?",
}
]
try:
# test without max tokens
response = completion(
model="cohere_chat/v1/command-r-plus",
messages=messages,
tools=tools,
tool_choice="auto",
)
# Add any assertions, here to check response args
print(response)
assert isinstance(response.choices[0].message.tool_calls[0].function.name, str)
assert isinstance(
response.choices[0].message.tool_calls[0].function.arguments, str
)
except litellm.Timeout:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# @pytest.mark.skip(reason="flaky test, times out frequently")
@pytest.mark.flaky(retries=6, delay=1)
def test_completion_cohere():
try:
# litellm.set_verbose=True
messages = [
{"role": "system", "content": "You're a good bot"},
{"role": "assistant", "content": [{"text": "2", "type": "text"}]},
{"role": "assistant", "content": [{"text": "3", "type": "text"}]},
{
"role": "user",
"content": "Hey",
},
]
response = completion(
model="cohere_chat/v1/command-r",
messages=messages,
)
print(response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
# FYI - cohere_chat looks quite unstable, even when testing locally
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.flaky(retries=3, delay=1)
async def test_chat_completion_cohere(sync_mode):
try:
litellm.set_verbose = True
messages = [
{"role": "system", "content": "You're a good bot"},
{
"role": "user",
"content": "Hey",
},
]
if sync_mode is False:
response = await litellm.acompletion(
model="cohere_chat/v1/command-r",
messages=messages,
max_tokens=10,
)
else:
response = completion(
model="cohere_chat/v1/command-r",
messages=messages,
max_tokens=10,
)
print(response)
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.asyncio
@pytest.mark.parametrize("sync_mode", [False])
async def test_chat_completion_cohere_stream(sync_mode):
try:
litellm.set_verbose = True
messages = [
{"role": "system", "content": "You're a good bot"},
{
"role": "user",
"content": "Hey",
},
]
if sync_mode is False:
response = await litellm.acompletion(
model="cohere_chat/v1/command-r",
messages=messages,
max_tokens=10,
stream=True,
)
print("async cohere stream response", response)
async for chunk in response:
print(chunk)
else:
response = completion(
model="cohere_chat/v1/command-r",
messages=messages,
max_tokens=10,
stream=True,
)
print(response)
for chunk in response:
print(chunk)
except litellm.APIConnectionError as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_cohere_embedding_outout_dimensions():
litellm._turn_on_debug()
response = embedding(
@ -425,331 +244,3 @@ def test_cohere_embed_v4_with_optional_params():
# ==================== COHERE V2 API TESTS ====================
@pytest.mark.parametrize("sync_mode", [True, False])
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_cohere_v2_chat_completion(sync_mode):
"""Test basic Cohere v2 chat completion functionality."""
try:
litellm.set_verbose = True
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Hello, how are you?"},
]
if sync_mode:
response = completion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
max_tokens=50,
)
else:
response = await litellm.acompletion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
max_tokens=50,
)
# Validate response structure
assert response.choices is not None
assert len(response.choices) > 0
assert response.choices[0].message.content is not None
assert response.usage is not None
assert response.usage.total_tokens > 0
print(f"Cohere v2 response: {response}")
except litellm.ServiceUnavailableError:
pass # Skip if service is unavailable
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_cohere_v2_streaming(stream):
"""Test Cohere v2 streaming functionality."""
try:
litellm.set_verbose = True
messages = [{"role": "user", "content": "Tell me a short story about a robot."}]
response = await litellm.acompletion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
max_tokens=100,
stream=stream,
)
if stream:
# Test streaming response
chunks = []
async for chunk in response:
chunks.append(chunk)
if len(chunks) >= 3: # Test first few chunks
break
assert len(chunks) > 0
print(f"Received {len(chunks)} streaming chunks")
else:
# Test non-streaming response
assert response.choices is not None
assert len(response.choices) > 0
assert response.choices[0].message.content is not None
print(f"Non-streaming response: {response.choices[0].message.content}")
except litellm.ServiceUnavailableError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_cohere_v2_tool_calling():
"""Test Cohere v2 tool calling functionality."""
try:
litellm.set_verbose = True
tools = [
{
"type": "function",
"function": {
"name": "get_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": ["celsius", "fahrenheit"],
},
},
"required": ["location"],
},
},
}
]
messages = [{"role": "user", "content": "What's the weather like in New York?"}]
response = completion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
tools=tools,
tool_choice="auto",
max_tokens=100,
)
# Validate tool calling response
assert response.choices is not None
assert len(response.choices) > 0
message = response.choices[0].message
# Check if tool calls are present
if hasattr(message, "tool_calls") and message.tool_calls:
assert len(message.tool_calls) > 0
tool_call = message.tool_calls[0]
assert tool_call.function.name == "get_weather"
assert tool_call.function.arguments is not None
print(
f"Tool call: {tool_call.function.name} - {tool_call.function.arguments}"
)
else:
# If no tool calls, check that we got a regular response
assert message.content is not None
print(f"Regular response: {message.content}")
except litellm.ServiceUnavailableError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
@pytest.mark.parametrize("stream", [True, False])
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_cohere_v2_annotations(stream):
"""Test Cohere v2 annotations functionality (replaces citations)."""
try:
litellm.set_verbose = True
messages = [
{"role": "user", "content": "What are the benefits of renewable energy?"}
]
documents = [
{
"data": {
"title": "Renewable Energy Benefits Document",
"snippet": "Renewable energy sources like solar and wind power provide clean electricity while reducing greenhouse gas emissions and dependence on fossil fuels.",
}
},
{
"data": {
"title": "Environmental Impact Study",
"snippet": "Studies show that renewable energy significantly reduces carbon footprint and helps combat climate change.",
}
},
]
response = await litellm.acompletion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
documents=documents,
max_tokens=100,
stream=stream,
)
if stream:
# Test streaming with annotations
annotations_found = False
async for chunk in response:
# Check if chunk has a message with annotations
if (
hasattr(chunk, "choices")
and chunk.choices
and len(chunk.choices) > 0
and hasattr(chunk.choices[0], "message")
and hasattr(chunk.choices[0].message, "annotations")
and chunk.choices[0].message.annotations
):
annotations_found = True
print(
f"Streaming annotations: {chunk.choices[0].message.annotations}"
)
break
# Note: Annotations might not appear in every chunk during streaming
else:
# Test non-streaming with annotations
assert response.choices is not None
assert len(response.choices) > 0
# Check for annotations in message
message = response.choices[0].message
if hasattr(message, "annotations") and message.annotations:
assert len(message.annotations) > 0
print(f"Annotations found: {len(message.annotations)}")
# Validate annotation structure
for annotation in message.annotations:
assert (
annotation.get("type") == "url_citation"
), f"Expected type 'url_citation', got {annotation.get('type')}"
assert "url_citation" in annotation, "Missing url_citation field"
url_citation = annotation["url_citation"]
assert "start_index" in url_citation, "Missing start_index"
assert "end_index" in url_citation, "Missing end_index"
assert "title" in url_citation, "Missing title"
assert "url" in url_citation, "Missing url"
print(f"First annotation: {message.annotations[0]}")
else:
# Annotations might not always be present depending on the response
print("No annotations in this response")
# Ensure citations field is NOT present (removed backward compatibility)
assert not hasattr(
response, "citations"
), "Citations field should be removed - no backward compatibility"
except litellm.ServiceUnavailableError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_cohere_v2_parameter_mapping():
"""Test Cohere v2 parameter mapping and validation."""
try:
litellm.set_verbose = True
messages = [{"role": "user", "content": "Generate a creative story."}]
# Test various parameters that should be mapped correctly
response = completion(
model="cohere_chat/v2/command-a-03-2025",
messages=messages,
temperature=0.7,
max_tokens=50,
top_p=0.9,
frequency_penalty=0.1,
presence_penalty=0.1,
stop=["END", "STOP"],
seed=42,
)
# Validate response
assert response.choices is not None
assert len(response.choices) > 0
assert response.choices[0].message.content is not None
assert response.usage is not None
print(f"Parameter mapping test response: {response.choices[0].message.content}")
except litellm.ServiceUnavailableError:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_cohere_v2_error_handling():
"""Test Cohere v2 error handling with invalid parameters."""
try:
# Test with invalid model name
try:
response = completion(
model="cohere_chat/v2/invalid-model",
messages=[{"role": "user", "content": "Hello"}],
max_tokens=10,
)
# If we get here, the test should fail
pytest.fail("Should have failed with invalid model")
except Exception as e:
# Expected to fail with invalid model
print(f"Expected error with invalid model: {e}")
# Test with empty messages
try:
response = completion(
model="cohere_chat/v2/command-a-03-2025",
messages=[], # Empty messages
max_tokens=10,
)
pytest.fail("Should have failed with empty messages")
except Exception as e:
# Expected to fail with empty messages
print(f"Expected error with empty messages: {e}")
except Exception as e:
pytest.fail(f"Unexpected error in error handling test: {e}")
@pytest.mark.asyncio
@pytest.mark.flaky(retries=3, delay=1)
async def test_cohere_v2_conversation_history():
"""Test Cohere v2 with conversation history."""
try:
litellm.set_verbose = True
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "What is 2+2?"},
{"role": "assistant", "content": "2+2 equals 4."},
{"role": "user", "content": "What about 3+3?"},
]
response = await litellm.acompletion(
model="cohere_chat/v2/command-a-03-2025", messages=messages, max_tokens=50
)
# Validate response with conversation history
assert response.choices is not None
assert len(response.choices) > 0
assert response.choices[0].message.content is not None
print(f"Conversation history response: {response.choices[0].message.content}")
except (
litellm.ServiceUnavailableError,
litellm.InternalServerError,
litellm.Timeout,
litellm.APIConnectionError,
):
pytest.skip("Cohere service unavailable")
except litellm.RateLimitError:
pytest.skip("Rate limit exceeded")

82
tests/llm_translation/test_deepseek_completion.py
View file

@ -1,85 +1,3 @@
from base_llm_unit_tests import BaseLLMChatTest
import pytest
import litellm
# Test implementations
@pytest.mark.skip(reason="Deepseek API is hanging")
class TestDeepSeekChatCompletion(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
return {
"model": "deepseek/deepseek-reasoner",
}
def test_completion_cost_deepseek():
litellm.set_verbose = True
model_name = "deepseek/deepseek-chat"
messages_1 = [
{
"role": "system",
"content": "You are a history expert. The user will provide a series of questions, and your answers should be concise and start with `Answer:`",
},
{
"role": "user",
"content": "In what year did Qin Shi Huang unify the six states?",
},
{"role": "assistant", "content": "Answer: 221 BC"},
{"role": "user", "content": "Who was the founder of the Han Dynasty?"},
{"role": "assistant", "content": "Answer: Liu Bang"},
{"role": "user", "content": "Who was the last emperor of the Tang Dynasty?"},
{"role": "assistant", "content": "Answer: Li Zhu"},
{
"role": "user",
"content": "Who was the founding emperor of the Ming Dynasty?",
},
{"role": "assistant", "content": "Answer: Zhu Yuanzhang"},
{
"role": "user",
"content": "Who was the founding emperor of the Qing Dynasty?",
},
]
message_2 = [
{
"role": "system",
"content": "You are a history expert. The user will provide a series of questions, and your answers should be concise and start with `Answer:`",
},
{
"role": "user",
"content": "In what year did Qin Shi Huang unify the six states?",
},
{"role": "assistant", "content": "Answer: 221 BC"},
{"role": "user", "content": "Who was the founder of the Han Dynasty?"},
{"role": "assistant", "content": "Answer: Liu Bang"},
{"role": "user", "content": "Who was the last emperor of the Tang Dynasty?"},
{"role": "assistant", "content": "Answer: Li Zhu"},
{
"role": "user",
"content": "Who was the founding emperor of the Ming Dynasty?",
},
{"role": "assistant", "content": "Answer: Zhu Yuanzhang"},
{"role": "user", "content": "When did the Shang Dynasty fall?"},
]
try:
response_1 = litellm.completion(model=model_name, messages=messages_1)
response_2 = litellm.completion(model=model_name, messages=message_2)
# Add any assertions here to check the response
print(response_2)
assert response_2.usage.prompt_cache_hit_tokens is not None
assert response_2.usage.prompt_cache_miss_tokens is not None
assert (
response_2.usage.prompt_tokens
== response_2.usage.prompt_cache_miss_tokens
+ response_2.usage.prompt_cache_hit_tokens
)
assert (
response_2.usage._cache_read_input_tokens
== response_2.usage.prompt_cache_hit_tokens
)
except litellm.APIError as e:
pass
except Exception as e:
pytest.fail(f"Error occurred: {e}")
def test_deepseek_fill_reasoning_content_multiturn():

22
tests/llm_translation/test_groq.py
View file

@ -1,5 +1,3 @@
import os
import sys
import pytest
@ -10,32 +8,12 @@ import pytest
# ) # Adds the parent directory to the system path
import litellm
from base_llm_unit_tests import BaseLLMChatTest
from litellm.llms.groq.chat.transformation import (
GroqChatConfig,
GroqChatCompletionStreamingHandler,
)
class TestGroq(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
return {
"model": "groq/llama-3.3-70b-versatile",
}
def test_tool_call_with_empty_enum_property(self):
pass
@pytest.mark.parametrize(
"model",
["groq/qwen/qwen3-32b", "groq/openai/gpt-oss-20b", "groq/openai/gpt-oss-120b"],
)
def test_reasoning_effort_in_supported_params(self, model):
"""Test that reasoning_effort is in the list of supported parameters for Groq"""
supported_params = GroqChatConfig().get_supported_openai_params(model=model)
assert "reasoning_effort" in supported_params
class TestGroqStructuredOutputs:
"""
Tests for Groq structured outputs handling.

26
tests/llm_translation/test_lambda_ai.py
View file

@ -5,10 +5,8 @@ Tests for Lambda AI provider integration
import os
from unittest import mock
import pytest
import litellm
from litellm import completion
from litellm.llms.lambda_ai.chat.transformation import LambdaAIChatConfig
@ -79,30 +77,6 @@ def test_lambda_ai_in_provider_lists():
assert "https://api.lambda.ai/v1" in litellm.openai_compatible_endpoints
@pytest.mark.asyncio
async def test_lambda_ai_completion_call():
"""Test completion call with Lambda AI provider (requires LAMBDA_API_KEY)"""
# Skip if no API key is available
if not os.getenv("LAMBDA_API_KEY"):
pytest.skip("LAMBDA_API_KEY not set")
try:
response = await litellm.acompletion(
model="lambda_ai/llama3.1-8b-instruct",
messages=[{"role": "user", "content": "Hello, this is a test"}],
max_tokens=10,
)
assert response.choices[0].message.content
assert response.model
assert response.usage
except Exception as e:
# If the API key is invalid or there's a network issue, that's okay
# The important thing is that the provider was recognized
if "lambda_ai" not in str(e) and "provider" not in str(e).lower():
# Re-raise if it's not a provider-related error
raise
def test_lambda_ai_models_configuration():
"""Test that Lambda AI models are configured correctly"""
from litellm import get_model_info

63
tests/llm_translation/test_langgraph.py
View file

@ -23,70 +23,7 @@ import sys
sys.path.insert(0, os.path.abspath("../.."))
import pytest
import litellm
@pytest.mark.asyncio
async def test_langgraph_acompletion_non_streaming():
"""
Test non-streaming acompletion call to LangGraph server.
Uses the /runs/wait endpoint for synchronous response.
"""
api_base = os.environ.get("LANGGRAPH_API_BASE", "http://localhost:2024")
try:
response = await litellm.acompletion(
model="langgraph/agent",
messages=[{"role": "user", "content": "What is 25 * 4?"}],
api_base=api_base,
stream=False,
)
assert response is not None
assert response.choices is not None
assert len(response.choices) > 0
assert response.choices[0].message is not None
assert response.choices[0].message.content is not None
assert len(response.choices[0].message.content) > 0
except Exception as e:
pytest.skip(f"LangGraph server not available: {e}")
@pytest.mark.asyncio
async def test_langgraph_acompletion_streaming():
"""
Test streaming acompletion call to LangGraph server.
Uses the /runs/stream endpoint with stream_mode="messages-tuple".
"""
api_base = os.environ.get("LANGGRAPH_API_BASE", "http://localhost:2024")
try:
response = await litellm.acompletion(
model="langgraph/agent",
messages=[{"role": "user", "content": "What is the weather in Tokyo?"}],
api_base=api_base,
stream=True,
)
full_content = ""
chunk_count = 0
async for chunk in response:
chunk_count += 1
if (
chunk.choices
and chunk.choices[0].delta
and chunk.choices[0].delta.content
):
full_content += chunk.choices[0].delta.content
assert chunk_count > 0, "Should receive at least one chunk"
except Exception as e:
pytest.skip(f"LangGraph server not available: {e}")
def test_langgraph_config_get_complete_url():

10
tests/llm_translation/test_together_ai.py
View file

@ -1,8 +1,8 @@
"""
Test TogetherAI LLM
TogetherAI supported-params goldens. The live BaseLLMChatTest subclass moved
to tests/harness_suites/chat_live_longtail/.
"""
from base_llm_unit_tests import BaseLLMChatTest
import os
import sys
@ -14,11 +14,7 @@ import litellm
import pytest
class TestTogetherAI(BaseLLMChatTest):
def get_base_completion_call_args(self) -> dict:
litellm.set_verbose = True
return {"model": "together_ai/Qwen/Qwen2.5-7B-Instruct-Turbo"}
class TestTogetherAIParams:
@pytest.mark.parametrize(
"model, expected_bool",
[

26
tests/llm_translation/test_v0.py
View file

@ -5,10 +5,8 @@ Tests for v0 provider integration
import os
from unittest import mock
import pytest
import litellm
from litellm import completion
from litellm.llms.v0.chat.transformation import V0ChatConfig
@ -72,30 +70,6 @@ def test_v0_in_provider_lists():
assert "https://api.v0.dev/v1" in litellm.openai_compatible_endpoints
@pytest.mark.asyncio
async def test_v0_completion_call():
"""Test completion call with v0 provider (requires V0_API_KEY)"""
# Skip if no API key is available
if not os.getenv("V0_API_KEY"):
pytest.skip("V0_API_KEY not set")
try:
response = await litellm.acompletion(
model="v0/gpt-4-turbo",
messages=[{"role": "user", "content": "Hello, this is a test"}],
max_tokens=10,
)
assert response.choices[0].message.content
assert response.model
assert response.usage
except Exception as e:
# If the API key is invalid or there's a network issue, that's okay
# The important thing is that the provider was recognized
if "v0" not in str(e) and "provider" not in str(e).lower():
# Re-raise if it's not a provider-related error
raise
def test_v0_supported_params():
"""Test that v0 returns only the supported parameters"""
config = V0ChatConfig()

120
tests/llm_translation/test_xai.py
View file

@ -8,11 +8,8 @@ sys.path.insert(
import pytest
import litellm
from litellm import completion
from unittest.mock import patch
from litellm.llms.xai.chat.transformation import XAIChatConfig, XAI_API_BASE
from base_llm_unit_tests import BaseReasoningLLMTests, BaseLLMChatTest
def test_xai_chat_config_get_openai_compatible_provider_info():
@ -142,120 +139,3 @@ def test_xai_grok_4_frequency_penalty_not_supported(model):
assert "frequency_penalty" not in supported_params
def test_xai_message_name_filtering():
messages = [
{
"role": "system",
"content": "*I press the green button*",
"name": "example_user",
},
{"role": "user", "content": "Hello", "name": "John"},
{"role": "assistant", "content": "Hello", "name": "Jane"},
]
response = completion(
model="xai/grok-3-mini-beta",
messages=messages,
)
assert response is not None
assert response.choices[0].message.content is not None
class TestXAIReasoningEffort(BaseReasoningLLMTests):
def get_base_completion_call_args(self):
return {
"model": "xai/grok-3-mini-beta",
"messages": [{"role": "user", "content": "Hello"}],
}
class TestXAIChat(BaseLLMChatTest):
def get_base_completion_call_args(self):
return {
"model": "xai/grok-3-mini-beta",
}
def test_web_search(self):
"""Web search is only supported for Grok 4 family models"""
from litellm.utils import supports_web_search
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
litellm._turn_on_debug()
# Use grok-4-1-fast which supports web search
model = "xai/grok-4-1-fast"
if not supports_web_search(model, None):
pytest.skip("Model does not support web search")
response = completion(
model=model,
messages=[
{"role": "user", "content": "What's the weather like in Boston today?"}
],
web_search_options={},
max_tokens=100,
)
assert response is not None
def test_xai_streaming_with_include_usage():
"""
Test that xAI streaming correctly handles usage in the last chunk
when stream_options={"include_usage": True} is set.
xAI sends usage in a chunk with empty choices array, which should be
handled by XAIChatCompletionStreamingHandler.
"""
try:
response = completion(
model="xai/grok-4-1-fast-non-reasoning",
messages=[
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Say hello in one word"},
],
stream=True,
stream_options={"include_usage": True},
max_tokens=10,
)
chunks = []
usage_chunk = None
for chunk in response:
chunks.append(chunk)
if hasattr(chunk, "usage") and chunk.usage is not None:
usage_chunk = chunk
# Verify we got chunks
assert len(chunks) > 0, "Should receive streaming chunks"
# Verify usage was included in one of the chunks
assert usage_chunk is not None, "Should receive usage in streaming chunks"
# Verify usage has expected fields
assert hasattr(
usage_chunk.usage, "prompt_tokens"
), "Usage should have prompt_tokens"
assert hasattr(
usage_chunk.usage, "completion_tokens"
), "Usage should have completion_tokens"
assert hasattr(
usage_chunk.usage, "total_tokens"
), "Usage should have total_tokens"
# Verify usage values are positive
assert usage_chunk.usage.prompt_tokens > 0, "prompt_tokens should be positive"
assert (
usage_chunk.usage.completion_tokens > 0
), "completion_tokens should be positive"
assert usage_chunk.usage.total_tokens > 0, "total_tokens should be positive"
print(f"✓ Successfully received usage in streaming chunk: {usage_chunk.usage}")
except Exception as e:
if "API key" in str(e) or "authentication" in str(e).lower():
pytest.skip(f"Skipping test due to API key issue: {str(e)}")
raise