mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
6a0d03914c
* test: drop the cwd-relative sys.path.insert calls from the test suite
TQ003 stands at 1,077 across 1,058 files, and 1,015 of them are the same shape:
sys.path.insert(0, os.path.abspath("../..")) and its deeper siblings. The
argument resolves against the working directory rather than the file, so from
the repo root, where every job runs pytest, it inserts the directory two levels
above the checkout. It has never pointed at litellm. The package is installed
into the environment anyway, which is what actually makes the import work, and
what the rule's message has said all along.
Removing them leaves 1,634 imports of sys and os with no remaining reference,
and those go too, except where another test module imports the name back out of
the file. The rest of TQ003 is 62 call sites that resolve against __file__ or a
variable, which are a different question and are left alone.
Collection is identical either way: 45,871 tests and the same 51 pre-existing
collection errors before and after, and ruff reports no new undefined name.
* test: drop the duplicate imports the sys.path sweep exposed to F811
* test(pre-call-utils): restore the os import the new bedrock tests need
879 lines
30 KiB
Python
879 lines
30 KiB
Python
import traceback
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
import io
|
|
|
|
import json
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm import RateLimitError, Timeout, completion, completion_cost, embedding
|
|
from unittest.mock import AsyncMock, patch
|
|
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
|
|
|
|
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}")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cohere_request_body_with_allowed_params():
|
|
"""
|
|
Test to validate that when allowed_openai_params is provided, the request body contains
|
|
the correct response_format and reasoning_effort values.
|
|
"""
|
|
# Define test parameters
|
|
test_response_format = {"type": "json"}
|
|
test_reasoning_effort = "low"
|
|
test_tools = [
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_current_time",
|
|
"description": "Get the current time in a given location.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"location": {
|
|
"type": "string",
|
|
"description": "The city name, e.g. San Francisco",
|
|
}
|
|
},
|
|
"required": ["location"],
|
|
},
|
|
},
|
|
}
|
|
]
|
|
|
|
# Create a mock response
|
|
mock_response = AsyncMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"text": "I am Command, a language model developed by Cohere.",
|
|
"generation_id": "mock-generation-id",
|
|
"finish_reason": "COMPLETE",
|
|
}
|
|
|
|
# Mock the AsyncHTTPHandler.post method at the module level
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
|
return_value=mock_response,
|
|
) as mock_post:
|
|
try:
|
|
await litellm.acompletion(
|
|
model="cohere/v1/command",
|
|
messages=[{"content": "what llm are you", "role": "user"}],
|
|
allowed_openai_params=["tools", "response_format", "reasoning_effort"],
|
|
response_format=test_response_format,
|
|
reasoning_effort=test_reasoning_effort,
|
|
tools=test_tools,
|
|
)
|
|
except Exception:
|
|
pass # We only care about the request body validation
|
|
|
|
# Verify the API call was made
|
|
mock_post.assert_called_once()
|
|
|
|
# Get and parse the request body
|
|
request_data = json.loads(mock_post.call_args.kwargs["data"])
|
|
print(f"request_data: {request_data}")
|
|
|
|
# Validate request contains our specified parameters
|
|
assert "allowed_openai_params" not in request_data
|
|
assert request_data["response_format"] == test_response_format
|
|
assert request_data["reasoning_effort"] == test_reasoning_effort
|
|
|
|
|
|
def test_cohere_embedding_outout_dimensions():
|
|
litellm._turn_on_debug()
|
|
response = embedding(
|
|
model="cohere/embed-v4.0", input="Hello, world!", dimensions=512
|
|
)
|
|
print(f"response: {response}\n")
|
|
assert len(response.data[0]["embedding"]) == 512
|
|
|
|
|
|
# Comprehensive Cohere Embed v4 tests
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_cohere_embed_v4_basic_text(sync_mode):
|
|
"""Test basic text embedding functionality with Cohere Embed v4."""
|
|
try:
|
|
data = {
|
|
"model": "cohere/embed-v4.0",
|
|
"input": ["Hello world!", "This is a test sentence."],
|
|
"input_type": "search_document",
|
|
}
|
|
|
|
if sync_mode:
|
|
response = embedding(**data)
|
|
else:
|
|
response = await litellm.aembedding(**data)
|
|
|
|
# Validate response structure
|
|
assert response.model is not None
|
|
assert len(response.data) == 2
|
|
assert response.data[0]["object"] == "embedding"
|
|
assert len(response.data[0]["embedding"]) > 0
|
|
assert response.usage.prompt_tokens > 0
|
|
assert isinstance(response.usage, litellm.Usage)
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_cohere_embed_v4_with_dimensions(sync_mode):
|
|
"""Test Cohere Embed v4 with specific dimension parameter."""
|
|
try:
|
|
data = {
|
|
"model": "cohere/embed-v4.0",
|
|
"input": ["Test with custom dimensions"],
|
|
"dimensions": 512,
|
|
"input_type": "search_query",
|
|
}
|
|
|
|
if sync_mode:
|
|
response = embedding(**data)
|
|
else:
|
|
response = await litellm.aembedding(**data)
|
|
|
|
# Validate dimension
|
|
assert len(response.data[0]["embedding"]) == 512
|
|
assert isinstance(response.usage, litellm.Usage)
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_cohere_embed_v4_image_embedding(sync_mode):
|
|
"""Test Cohere Embed v4 image embedding functionality (multimodal)."""
|
|
try:
|
|
import base64
|
|
|
|
# 1x1 pixel red PNG (base64 encoded)
|
|
test_image_data = b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\tpHYs\x00\x00\x0b\x13\x00\x00\x0b\x13\x01\x00\x9a\x9c\x18\x00\x00\x00\x0cIDATx\x9cc\xf8\x00\x00\x00\x01\x00\x01\x00\x00\x00\x00"
|
|
test_image_b64 = base64.b64encode(test_image_data).decode("utf-8")
|
|
|
|
data = {
|
|
"model": "cohere/embed-v4.0",
|
|
"input": [test_image_b64],
|
|
"input_type": "image",
|
|
}
|
|
|
|
if sync_mode:
|
|
response = embedding(**data)
|
|
else:
|
|
response = await litellm.aembedding(**data)
|
|
|
|
# Validate response structure for image embedding
|
|
assert response.model is not None
|
|
assert len(response.data) == 1
|
|
assert response.data[0]["object"] == "embedding"
|
|
assert len(response.data[0]["embedding"]) > 0
|
|
assert isinstance(response.usage, litellm.Usage)
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"input_type", ["search_document", "search_query", "classification", "clustering"]
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_cohere_embed_v4_input_types(input_type):
|
|
"""Test Cohere Embed v4 with different input types."""
|
|
try:
|
|
response = await litellm.aembedding(
|
|
model="cohere/embed-v4.0",
|
|
input=[f"Test text for {input_type}"],
|
|
input_type=input_type,
|
|
)
|
|
|
|
assert response.model is not None
|
|
assert len(response.data) == 1
|
|
assert response.data[0]["object"] == "embedding"
|
|
assert len(response.data[0]["embedding"]) > 0
|
|
assert isinstance(response.usage, litellm.Usage)
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
def test_cohere_embed_v4_encoding_format():
|
|
"""Test Cohere Embed v4 with different encoding formats."""
|
|
try:
|
|
response = embedding(
|
|
model="cohere/embed-v4.0",
|
|
input=["Test encoding format"],
|
|
encoding_format="float",
|
|
)
|
|
|
|
assert response.model is not None
|
|
assert len(response.data) == 1
|
|
assert response.data[0]["object"] == "embedding"
|
|
assert len(response.data[0]["embedding"]) > 0
|
|
# Validate that embeddings are floats
|
|
assert all(isinstance(x, float) for x in response.data[0]["embedding"])
|
|
assert isinstance(response.usage, litellm.Usage)
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
def test_cohere_embed_v4_error_handling():
|
|
"""Test error handling for Cohere Embed v4 with invalid inputs."""
|
|
try:
|
|
# Test with empty input - should raise an error
|
|
try:
|
|
response = embedding(model="cohere/embed-v4.0", input=[]) # Empty input
|
|
pytest.fail("Should have failed with empty input")
|
|
except Exception:
|
|
pass # Expected to fail
|
|
|
|
# Test with None input - should raise an error
|
|
try:
|
|
response = embedding(model="cohere/embed-v4.0", input=None)
|
|
pytest.fail("Should have failed with None input")
|
|
except Exception:
|
|
pass # Expected to fail
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error in error handling test: {e}")
|
|
|
|
|
|
@pytest.mark.parametrize("sync_mode", [True, False])
|
|
@pytest.mark.asyncio
|
|
async def test_cohere_embed_v4_multiple_texts(sync_mode):
|
|
"""Test Cohere Embed v4 with multiple text inputs."""
|
|
try:
|
|
texts = [
|
|
"The quick brown fox jumps over the lazy dog",
|
|
"Machine learning is transforming the world",
|
|
"Python is a versatile programming language",
|
|
"Natural language processing enables human-computer interaction",
|
|
]
|
|
|
|
data = {
|
|
"model": "cohere/embed-v4.0",
|
|
"input": texts,
|
|
"input_type": "search_document",
|
|
}
|
|
|
|
if sync_mode:
|
|
response = embedding(**data)
|
|
else:
|
|
response = await litellm.aembedding(**data)
|
|
|
|
# Validate response structure
|
|
assert response.model is not None
|
|
assert len(response.data) == len(texts)
|
|
|
|
for i, data_item in enumerate(response.data):
|
|
assert data_item["object"] == "embedding"
|
|
assert data_item["index"] == i
|
|
assert len(data_item["embedding"]) > 0
|
|
assert all(isinstance(x, float) for x in data_item["embedding"])
|
|
|
|
assert isinstance(response.usage, litellm.Usage)
|
|
assert response.usage.prompt_tokens > 0
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
def test_cohere_embed_v4_with_optional_params():
|
|
"""Test Cohere Embed v4 with various optional parameters."""
|
|
try:
|
|
response = embedding(
|
|
model="cohere/embed-v4.0",
|
|
input=["Test with optional parameters"],
|
|
input_type="search_query",
|
|
dimensions=256,
|
|
encoding_format="float",
|
|
)
|
|
|
|
# Validate response
|
|
assert response.model is not None
|
|
assert len(response.data) == 1
|
|
assert response.data[0]["object"] == "embedding"
|
|
assert len(response.data[0]["embedding"]) == 256 # Custom dimensions
|
|
assert all(isinstance(x, float) for x in response.data[0]["embedding"])
|
|
assert isinstance(response.usage, litellm.Usage)
|
|
|
|
except Exception as e:
|
|
pytest.fail(f"Error occurred: {e}")
|
|
|
|
|
|
# ==================== 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
|
|
async def test_cohere_documents_options_in_request_body():
|
|
"""
|
|
Test that documents parameters is properly included
|
|
in the request body after transformation (sent via extra_body).
|
|
"""
|
|
# Create a mock response
|
|
mock_response = AsyncMock()
|
|
mock_response.status_code = 200
|
|
mock_response.json.return_value = {
|
|
"text": "Test response with citations",
|
|
"generation_id": "mock-generation-id",
|
|
"finish_reason": "COMPLETE",
|
|
}
|
|
|
|
# Mock the AsyncHTTPHandler.post method
|
|
with patch(
|
|
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
|
|
return_value=mock_response,
|
|
) as mock_post:
|
|
try:
|
|
# Test documents and citation_options parameters
|
|
test_documents = [
|
|
{
|
|
"data": {
|
|
"title": "Test Document 1",
|
|
"snippet": "This is test content 1",
|
|
}
|
|
},
|
|
{
|
|
"data": {
|
|
"title": "Test Document 2",
|
|
"snippet": "This is test content 2",
|
|
}
|
|
},
|
|
]
|
|
await litellm.acompletion(
|
|
model="cohere_chat/command-a-03-2025",
|
|
messages=[{"role": "user", "content": "Test message"}],
|
|
documents=test_documents,
|
|
)
|
|
except Exception:
|
|
pass # We only care about the request body validation
|
|
|
|
# Verify the API call was made
|
|
mock_post.assert_called_once()
|
|
|
|
# Get and parse the request body
|
|
request_data = json.loads(mock_post.call_args.kwargs["data"])
|
|
print(f"Request body: {request_data}")
|
|
|
|
# Validate that documents and citation_options are in the request body
|
|
assert "documents" in request_data
|
|
assert request_data["documents"] == test_documents
|
|
|
|
|
|
@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")
|