mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-14 04:09:09 +00:00
test: migrate bedrock and cohere translation tests
This commit is contained in:
yuneng
parent
500ac9b9f3
commit
99da2ee218
12 changed files with 404 additions and 2761 deletions
70
tests/e2e/llm_translation/test_bedrock_guardrail_e2e.py
Normal file
70
tests/e2e/llm_translation/test_bedrock_guardrail_e2e.py
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from e2e_config import unique_marker
|
||||
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
|
||||
from lifecycle import ResourceManager
|
||||
from models import ChatMessage, LiteLLMParamsBody
|
||||
from proxy_client import ProxyClient
|
||||
from sdk_clients import SdkClients
|
||||
|
||||
pytestmark = pytest.mark.e2e
|
||||
|
||||
BEDROCK_MODEL: Final = "bedrock/us.anthropic.claude-opus-5-5"
|
||||
|
||||
|
||||
class TestBedrockGuardrail:
|
||||
@pytest.mark.covers("guardrail.bedrock.pre_call.blocks", exercised_on=["chat_completions"])
|
||||
@meta(
|
||||
Subject(
|
||||
domain=Domain.GUARDRAILS,
|
||||
route=Route.CHAT_COMPLETIONS,
|
||||
providers=(Provider.BEDROCK,),
|
||||
models=(BEDROCK_MODEL,),
|
||||
mode=Mode.NONSTREAM,
|
||||
)
|
||||
)
|
||||
def test_guardrail_blocks_and_returns_trace(
|
||||
self,
|
||||
proxy: ProxyClient,
|
||||
resources: ResourceManager,
|
||||
sdk: SdkClients,
|
||||
) -> None:
|
||||
identifier: Final = os.environ.get("BEDROCK_GUARDRAIL_IDENTIFIER")
|
||||
version: Final = os.environ.get("BEDROCK_GUARDRAIL_VERSION")
|
||||
if not identifier or not version:
|
||||
pytest.skip("BEDROCK_GUARDRAIL_IDENTIFIER and BEDROCK_GUARDRAIL_VERSION are required")
|
||||
|
||||
model_name: Final = f"bedrock-guardrail-{unique_marker()}"
|
||||
model_id: Final = proxy.create_model(
|
||||
model_name,
|
||||
LiteLLMParamsBody(
|
||||
model=BEDROCK_MODEL,
|
||||
aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID",
|
||||
aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY",
|
||||
aws_region_name="os.environ/AWS_REGION",
|
||||
),
|
||||
)
|
||||
resources.defer(lambda: proxy.delete_model(model_id))
|
||||
|
||||
prompt: Final = f"where do i buy coffee from? {unique_marker()}"
|
||||
response: Final = sdk.openai(resources.key()).chat.completions.create(
|
||||
model=model_name,
|
||||
messages=[ChatMessage(role="user", content=prompt).model_dump()],
|
||||
max_tokens=64,
|
||||
extra_body={
|
||||
"guardrailConfig": {
|
||||
"guardrailIdentifier": identifier,
|
||||
"guardrailVersion": version,
|
||||
"trace": "enabled",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
content: Final = response.choices[0].message.content or ""
|
||||
assert "coffee guardrail applied" in content.lower(), f"guardrail did not block the prompt: {content!r}"
|
||||
trace: Final = getattr(response, "trace", None)
|
||||
assert isinstance(trace, dict) and trace, f"guardrail trace missing from response: {response!r}"
|
||||
|
|
@ -1,63 +0,0 @@
|
|||
"""
|
||||
Test Bedrock AgentCore integration
|
||||
"""
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
import litellm
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_13sf6-cALnp38iZD", # non-streaming invocation
|
||||
"bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC", # streaming invocation
|
||||
],
|
||||
)
|
||||
def test_bedrock_agentcore_basic(model):
|
||||
"""
|
||||
Test AgentCore invocation parameterized by model
|
||||
"""
|
||||
litellm.turn_on_debug()
|
||||
response = litellm.completion(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "user", "content": "Explain machine learning in simple terms"}
|
||||
],
|
||||
)
|
||||
print("response from agentcore=", response.model_dump_json(indent=4))
|
||||
# Assert that the message content has a response with some length
|
||||
assert response.choices[0].message.content
|
||||
assert len(response.choices[0].message.content) > 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
[
|
||||
"bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_13sf6-cALnp38iZD", # streaming invocation
|
||||
],
|
||||
)
|
||||
async def test_bedrock_agentcore_with_streaming(model):
|
||||
"""
|
||||
Test AgentCore with streaming
|
||||
"""
|
||||
print("running streming test for model=", model)
|
||||
# litellm.turn_on_debug()
|
||||
response = await litellm.acompletion(
|
||||
model="bedrock/agentcore/arn:aws:bedrock-agentcore:us-west-2:888602223428:runtime/hosted_agent_r9jvp-3ySZuRHjLC",
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Explain machine learning in simple terms",
|
||||
}
|
||||
],
|
||||
stream=True,
|
||||
)
|
||||
|
||||
async for chunk in response:
|
||||
print("chunk=", chunk)
|
||||
File diff suppressed because one or more lines are too long
|
|
@ -1,149 +0,0 @@
|
|||
import json
|
||||
import os
|
||||
from unittest.mock import Mock, patch
|
||||
import pytest
|
||||
import base64
|
||||
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}
|
||||
|
||||
cohere_embedding_response = {"embeddings": [[0.1, 0.2, 0.3]], "inputTextTokenCount": 10}
|
||||
|
||||
img_base_64 = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAGQAAABkBAMAAACCzIhnAAAAG1BMVEURAAD///+ln5/h39/Dv79qX18uHx+If39MPz9oMSdmAAAACXBIWXMAAA7EAAAOxAGVKw4bAAABB0lEQVRYhe2SzWrEIBCAh2A0jxEs4j6GLDS9hqWmV5Flt0cJS+lRwv742DXpEjY1kOZW6HwHFZnPmVEBEARBEARB/jd0KYA/bcUYbPrRLh6amXHJ/K+ypMoyUaGthILzw0l+xI0jsO7ZcmCcm4ILd+QuVYgpHOmDmz6jBeJImdcUCmeBqQpuqRIbVmQsLCrAalrGpfoEqEogqbLTWuXCPCo+Ki1XGqgQ+jVVuhB8bOaHkvmYuzm/b0KYLWwoK58oFqi6XfxQ4Uz7d6WeKpna6ytUs5e8betMcqAv5YPC5EZB2Lm9FIn0/VP6R58+/GEY1X1egVoZ/3bt/EqF6malgSAIgiDIH+QL41409QMY0LMAAAAASUVORK5CYII="
|
||||
|
||||
|
||||
|
||||
|
||||
def test_e2e_bedrock_embedding():
|
||||
"""
|
||||
Test text embedding with TwelveLabs Marengo.
|
||||
Validates that the transformation properly extracts embedding data from TwelveLabs response format.
|
||||
"""
|
||||
print("Testing text embedding...")
|
||||
original_region_name = os.environ.get("AWS_REGION_NAME")
|
||||
|
||||
os.environ["AWS_REGION_NAME"] = "us-east-1"
|
||||
|
||||
litellm.turn_on_debug()
|
||||
response = litellm.embedding(
|
||||
model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0",
|
||||
input=["Hello world from LiteLLM with TwelveLabs Marengo!"],
|
||||
)
|
||||
|
||||
# Validate response structure
|
||||
assert isinstance(
|
||||
response, litellm.EmbeddingResponse
|
||||
), "Response should be EmbeddingResponse type"
|
||||
assert hasattr(response, "data"), "Response should have 'data' attribute"
|
||||
assert len(response.data) > 0, "Response data should not be empty"
|
||||
|
||||
# Validate first embedding
|
||||
embedding_obj = response.data[0]
|
||||
assert hasattr(
|
||||
embedding_obj, "embedding"
|
||||
), "Embedding object should have 'embedding' attribute"
|
||||
assert isinstance(
|
||||
embedding_obj.embedding, list
|
||||
), "Embedding should be a list of floats"
|
||||
assert len(embedding_obj.embedding) > 0, "Embedding vector should not be empty"
|
||||
assert all(
|
||||
isinstance(x, (int, float)) for x in embedding_obj.embedding
|
||||
), "All embedding values should be numeric"
|
||||
|
||||
# Validate embedding properties
|
||||
assert embedding_obj.index == 0, "First embedding should have index 0"
|
||||
assert (
|
||||
embedding_obj.object == "embedding"
|
||||
), "Embedding object type should be 'embedding'"
|
||||
|
||||
# Validate usage information
|
||||
assert hasattr(response, "usage"), "Response should have usage information"
|
||||
assert response.usage is not None, "Usage should not be None"
|
||||
assert response.usage.total_tokens >= 0, "Total tokens should be non-negative"
|
||||
|
||||
print(
|
||||
f"Text embedding successful! Vector size: {len(embedding_obj.embedding)}, Response: {response}"
|
||||
)
|
||||
|
||||
# Restore original region name
|
||||
if original_region_name:
|
||||
os.environ["AWS_REGION_NAME"] = original_region_name
|
||||
|
||||
|
||||
def test_e2e_bedrock_embedding_image_twelvelabs_marengo():
|
||||
"""
|
||||
Test image embedding with TwelveLabs Marengo.
|
||||
Validates that the transformation properly extracts embedding data from TwelveLabs response format for images.
|
||||
"""
|
||||
print("Testing image embedding...")
|
||||
original_region_name = os.environ.get("AWS_REGION_NAME")
|
||||
os.environ["AWS_REGION_NAME"] = "us-east-1"
|
||||
litellm.turn_on_debug()
|
||||
|
||||
# Load duck.png and convert to base64
|
||||
duck_img_path = os.path.join(os.path.dirname(__file__), "duck.png")
|
||||
with open(duck_img_path, "rb") as img_file:
|
||||
duck_img_data = base64.b64encode(img_file.read()).decode("utf-8")
|
||||
duck_img_base64 = f"data:image/png;base64,{duck_img_data}"
|
||||
|
||||
response = litellm.embedding(
|
||||
model="bedrock/us.twelvelabs.marengo-embed-2-7-v1:0",
|
||||
input=[duck_img_base64],
|
||||
aws_region_name="us-east-1",
|
||||
input_type="image",
|
||||
)
|
||||
|
||||
# Validate response structure
|
||||
assert isinstance(
|
||||
response, litellm.EmbeddingResponse
|
||||
), "Response should be EmbeddingResponse type"
|
||||
assert hasattr(response, "data"), "Response should have 'data' attribute"
|
||||
assert len(response.data) > 0, "Response data should not be empty"
|
||||
|
||||
# Validate first embedding
|
||||
embedding_obj = response.data[0]
|
||||
assert hasattr(
|
||||
embedding_obj, "embedding"
|
||||
), "Embedding object should have 'embedding' attribute"
|
||||
assert isinstance(
|
||||
embedding_obj.embedding, list
|
||||
), "Embedding should be a list of floats"
|
||||
assert len(embedding_obj.embedding) > 0, "Embedding vector should not be empty"
|
||||
assert all(
|
||||
isinstance(x, (int, float)) for x in embedding_obj.embedding
|
||||
), "All embedding values should be numeric"
|
||||
|
||||
# Validate embedding properties
|
||||
assert embedding_obj.index == 0, "First embedding should have index 0"
|
||||
assert (
|
||||
embedding_obj.object == "embedding"
|
||||
), "Embedding object type should be 'embedding'"
|
||||
|
||||
# Validate usage information
|
||||
assert hasattr(response, "usage"), "Response should have usage information"
|
||||
assert response.usage is not None, "Usage should not be None"
|
||||
assert response.usage.total_tokens >= 0, "Total tokens should be non-negative"
|
||||
|
||||
# TwelveLabs Marengo should return 1024-dimensional embeddings
|
||||
expected_dimension = 1024
|
||||
assert (
|
||||
len(embedding_obj.embedding) == expected_dimension
|
||||
), f"TwelveLabs Marengo should return {expected_dimension}-dimensional embeddings, got {len(embedding_obj.embedding)}"
|
||||
|
||||
print(
|
||||
f"Image embedding successful! Vector size: {len(embedding_obj.embedding)}, Response: {response}"
|
||||
)
|
||||
|
||||
# Restore original region name
|
||||
if original_region_name:
|
||||
os.environ["AWS_REGION_NAME"] = original_region_name
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
titan_embedding_response = {"embedding": [0.1, 0.2, 0.3], "inputTextTokenCount": 10}
|
||||
|
|
@ -1,19 +0,0 @@
|
|||
from base_llm_unit_tests import BaseLLMChatTest
|
||||
|
||||
class TestBedrockGPTOSS(BaseLLMChatTest):
|
||||
|
||||
def get_base_completion_call_args(self) -> dict:
|
||||
return {
|
||||
"model": "bedrock/converse/openai.gpt-oss-20b-1:0",
|
||||
}
|
||||
|
||||
def test_function_calling_with_tool_response(self):
|
||||
"""Bedrock GPT-OSS intermittently emits truncated toolUse.input deltas on
|
||||
the live endpoint, which makes the inherited live integration test flaky.
|
||||
The accumulation side is covered deterministically by
|
||||
tests/unit/llms/bedrock/chat/test_invoke_handler.py::test_transform_tool_calls_index;
|
||||
the GPT-OSS-specific request-body transformation is covered by
|
||||
test_function_calling_request_body_gpt_oss below.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
|
@ -1,40 +0,0 @@
|
|||
from base_llm_unit_tests import BaseLLMChatTest
|
||||
import pytest
|
||||
import os
|
||||
|
||||
import litellm
|
||||
from litellm.types.llms.bedrock import BedrockInvokeNovaRequest
|
||||
|
||||
_LITELLM_LOGO_IMAGE_URL = (
|
||||
"https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0/"
|
||||
"ui/litellm-dashboard/public/assets/logos/litellm_logo.jpg"
|
||||
)
|
||||
_AWSMP_LOGO_IMAGE_URL = (
|
||||
"https://awsmp-logos.s3.amazonaws.com/seller-xw5kijmvmzasy/"
|
||||
"c233c9ade2ccb5491072ae232c814942.png"
|
||||
)
|
||||
|
||||
@pytest.mark.flaky(retries=3, delay=5)
|
||||
class TestBedrockInvokeClaudeJson(BaseLLMChatTest):
|
||||
def get_base_completion_call_args(self) -> dict:
|
||||
litellm.turn_on_debug()
|
||||
return {
|
||||
"model": "bedrock/invoke/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
}
|
||||
|
||||
test_pdf_handling = None
|
||||
|
||||
class TestBedrockInvokeNovaJson(BaseLLMChatTest):
|
||||
|
||||
def get_base_completion_call_args(self) -> dict:
|
||||
return {
|
||||
"model": "bedrock/invoke/us.amazon.nova-micro-v1:0",
|
||||
}
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def skip_non_json_tests(self, request):
|
||||
if not "json" in request.function.__name__.lower():
|
||||
pytest.skip(
|
||||
f"Skipping non-JSON test: {request.function.__name__} does not contain 'json'"
|
||||
)
|
||||
|
||||
|
|
@ -1,13 +0,0 @@
|
|||
from base_llm_unit_tests import BaseLLMChatTest
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
|
||||
class TestBedrockTestSuite(BaseLLMChatTest):
|
||||
test_function_calling_with_tool_response = None
|
||||
|
||||
def get_base_completion_call_args(self) -> dict:
|
||||
litellm.turn_on_debug()
|
||||
return {
|
||||
"model": "bedrock/converse/us.meta.llama3-3-70b-instruct-v1:0",
|
||||
}
|
||||
|
|
@ -1,30 +0,0 @@
|
|||
"""
|
||||
Tests for Bedrock Moonshot (Kimi K2) integration.
|
||||
|
||||
This test suite verifies:
|
||||
1. Basic completion functionality
|
||||
2. Streaming responses
|
||||
3. System message support
|
||||
4. Temperature parameter handling
|
||||
5. Reasoning content extraction from <reasoning> tags
|
||||
6. Tool calling support (including tool response handling)
|
||||
7. Parameter validation (e.g., stop sequences not supported)
|
||||
"""
|
||||
|
||||
from base_llm_unit_tests import BaseLLMChatTest
|
||||
|
||||
import litellm
|
||||
|
||||
class TestBedrockMoonshotInvoke(BaseLLMChatTest):
|
||||
"""
|
||||
Test suite for Bedrock Moonshot via invoke route.
|
||||
Inherits all standard LLM tests from BaseLLMChatTest.
|
||||
"""
|
||||
|
||||
test_streaming = None
|
||||
|
||||
def get_base_completion_call_args(self) -> dict:
|
||||
litellm.turn_on_debug()
|
||||
return {
|
||||
"model": "bedrock/invoke/moonshot.kimi-k2-thinking",
|
||||
}
|
||||
|
|
@ -1,20 +0,0 @@
|
|||
from base_llm_unit_tests import BaseLLMChatTest
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
|
||||
class TestBedrockNovaJson(BaseLLMChatTest):
|
||||
test_function_calling_with_tool_response = None
|
||||
|
||||
def get_base_completion_call_args(self) -> dict:
|
||||
litellm.turn_on_debug()
|
||||
return {
|
||||
"model": "bedrock/converse/us.amazon.nova-micro-v1:0",
|
||||
}
|
||||
|
||||
# @pytest.fixture(autouse=True)
|
||||
# def skip_non_json_tests(self, request):
|
||||
# if not "json" in request.function.__name__.lower():
|
||||
# pytest.skip(
|
||||
# f"Skipping non-JSON test: {request.function.__name__} does not contain 'json'"
|
||||
# )
|
||||
|
|
@ -1,753 +0,0 @@
|
|||
import traceback
|
||||
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
import io
|
||||
import json
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm import RateLimitError, Timeout, completion, completion_cost, embedding
|
||||
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.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(
|
||||
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
|
||||
@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")
|
||||
|
|
@ -1,7 +1,7 @@
|
|||
import json
|
||||
from collections.abc import Callable, Iterator
|
||||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -9,9 +9,10 @@ import respx
|
|||
|
||||
import litellm
|
||||
from litellm.caching.llm_caching_handler import LLMClientCache
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import cohere_messages_pt_v2
|
||||
from litellm.llms.cohere.chat.transformation import CohereChatConfig
|
||||
from litellm.llms.cohere.chat.v2_transformation import CohereV2ChatConfig
|
||||
from unittest.mock import AsyncMock, patch
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
|
||||
COHERE_V1_CHAT_URL: Final = "https://api.cohere.ai/v1/chat"
|
||||
COHERE_V2_CHAT_URL: Final = "https://api.cohere.com/v2/chat"
|
||||
|
|
@ -98,6 +99,250 @@ class TestCohereV2Transform:
|
|||
|
||||
assert result == {"temperature": 0.7, "max_tokens": 200}
|
||||
|
||||
|
||||
def test_cohere_v1_response_parses_citations_tool_calls_and_billed_units() -> None:
|
||||
config: Final = CohereChatConfig()
|
||||
model_response: Final = litellm.ModelResponse()
|
||||
raw_response: Final = httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"text": "Searching",
|
||||
"citations": [{"start": 0, "end": 9, "text": "Searching", "document_ids": ["doc-1"]}],
|
||||
"tool_calls": [{"name": "lookup", "generation_id": "call-1", "parameters": {"query": "weather"}}],
|
||||
"meta": {"billed_units": {"input_tokens": 7, "output_tokens": 3}},
|
||||
},
|
||||
)
|
||||
|
||||
response: Final = config.transform_response(
|
||||
model="command-r-plus",
|
||||
raw_response=raw_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.tool_calls[0].function.name == "lookup"
|
||||
assert response.choices[0].message.tool_calls[0].function.arguments == '{"query": "weather"}'
|
||||
assert getattr(response, "citations") == [{"start": 0, "end": 9, "text": "Searching", "document_ids": ["doc-1"]}]
|
||||
assert response.usage.prompt_tokens == 7
|
||||
assert response.usage.completion_tokens == 3
|
||||
|
||||
|
||||
def test_cohere_v2_response_parses_annotations_and_tool_calls() -> None:
|
||||
config: Final = CohereV2ChatConfig()
|
||||
model_response: Final = litellm.ModelResponse()
|
||||
raw_response: Final = httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "response-1",
|
||||
"finish_reason": "COMPLETE",
|
||||
"message": {
|
||||
"content": [{"type": "text", "text": "The weather is clear."}],
|
||||
"citations": [
|
||||
{
|
||||
"start": 0,
|
||||
"end": 21,
|
||||
"sources": [
|
||||
{
|
||||
"type": "document",
|
||||
"id": "source-1",
|
||||
"url": "https://example.com/weather",
|
||||
"document": {"title": "Weather report"},
|
||||
}
|
||||
],
|
||||
}
|
||||
],
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call-1",
|
||||
"type": "function",
|
||||
"function": {"name": "lookup", "arguments": '{"city":"Paris"}'},
|
||||
}
|
||||
],
|
||||
},
|
||||
"usage": {"tokens": {"input_tokens": 11, "output_tokens": 5}},
|
||||
},
|
||||
)
|
||||
|
||||
response: Final = config.transform_response(
|
||||
model="command-r-plus",
|
||||
raw_response=raw_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.tool_calls[0].function.name == "lookup"
|
||||
assert response.choices[0].message.tool_calls[0].function.arguments == '{"city":"Paris"}'
|
||||
assert response.choices[0].message.annotations[0]["url_citation"]["url"] == "https://example.com/weather"
|
||||
assert response.choices[0].message.annotations[0]["url_citation"]["title"] == "Weather report"
|
||||
assert response.usage.prompt_tokens == 11
|
||||
assert response.usage.completion_tokens == 5
|
||||
|
||||
|
||||
def test_cohere_v2_request_preserves_conversation_messages_and_parameters() -> None:
|
||||
config: Final = CohereV2ChatConfig()
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "What is the capital of France?"},
|
||||
{"role": "assistant", "content": "Paris."},
|
||||
{"role": "user", "content": "What language is spoken there?"},
|
||||
]
|
||||
optional_params: Final = config.map_openai_params(
|
||||
non_default_params={
|
||||
"max_completion_tokens": 128,
|
||||
"top_p": 0.7,
|
||||
"frequency_penalty": 0.3,
|
||||
"presence_penalty": 0.2,
|
||||
"stop": ["END"],
|
||||
},
|
||||
optional_params={},
|
||||
model="command-r-plus",
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
request: Final = config.transform_request(
|
||||
model="command-r-plus",
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request["messages"] == messages
|
||||
assert request["max_tokens"] == 128
|
||||
assert request["p"] == 0.7
|
||||
assert request["frequency_penalty"] == 0.3
|
||||
assert request["presence_penalty"] == 0.2
|
||||
assert request["stop_sequences"] == ["END"]
|
||||
|
||||
|
||||
@pytest.mark.respx(assert_all_called=True)
|
||||
def test_cohere_v2_completion_sends_mapped_request_and_parses_response(respx_mock: respx.MockRouter) -> None:
|
||||
route: Final = respx_mock.post(COHERE_V2_CHAT_URL).mock(
|
||||
return_value=httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"id": "response-1",
|
||||
"finish_reason": "COMPLETE",
|
||||
"message": {"content": [{"type": "text", "text": "Paris is the capital."}]},
|
||||
"usage": {"tokens": {"input_tokens": 8, "output_tokens": 4}},
|
||||
},
|
||||
)
|
||||
)
|
||||
response: Final = litellm.completion(
|
||||
model="cohere_chat/command-r-plus",
|
||||
messages=[{"role": "user", "content": "What is the capital of France?"}],
|
||||
api_key="test-api-key",
|
||||
client=HTTPHandler(),
|
||||
max_completion_tokens=64,
|
||||
top_p=0.7,
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "Paris is the capital."
|
||||
assert response.usage.prompt_tokens == 8
|
||||
assert response.usage.completion_tokens == 4
|
||||
assert route.called
|
||||
request_body: Final = json.loads(route.calls.last.request.content)
|
||||
assert request_body["messages"] == [{"role": "user", "content": "What is the capital of France?"}]
|
||||
assert request_body["model"] == "command-r-plus"
|
||||
assert request_body["max_tokens"] == 64
|
||||
assert request_body["p"] == 0.7
|
||||
assert "top_p" not in request_body
|
||||
|
||||
|
||||
def test_cohere_v1_request_separates_conversation_history_from_latest_message() -> None:
|
||||
messages: Final = [
|
||||
{"role": "user", "content": "What is 2 + 2?"},
|
||||
{"role": "assistant", "content": "4."},
|
||||
{"role": "user", "content": "And 3 + 3?"},
|
||||
]
|
||||
|
||||
latest_message, history = cohere_messages_pt_v2(
|
||||
messages=messages.copy(),
|
||||
model="command-r-plus",
|
||||
llm_provider="cohere_chat",
|
||||
)
|
||||
|
||||
assert latest_message == "And 3 + 3?"
|
||||
assert history == [
|
||||
{"role": "USER", "message": "What is 2 + 2?"},
|
||||
{"role": "CHATBOT", "message": "4.", "tool_calls": []},
|
||||
]
|
||||
|
||||
|
||||
def test_cohere_v2_stream_parser_emits_text_tool_calls_and_usage() -> None:
|
||||
config: Final = CohereV2ChatConfig()
|
||||
response_iterator: Final = config.get_model_response_iterator(
|
||||
iter(
|
||||
[
|
||||
json.dumps(
|
||||
{
|
||||
"type": "content-delta",
|
||||
"delta": {"message": {"content": {"text": "Searching"}}},
|
||||
}
|
||||
),
|
||||
json.dumps(
|
||||
{
|
||||
"type": "tool-call-delta",
|
||||
"delta": {
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call-1",
|
||||
"name": "lookup",
|
||||
"arguments": '{"city":"Paris"}',
|
||||
}
|
||||
]
|
||||
},
|
||||
}
|
||||
),
|
||||
json.dumps(
|
||||
{
|
||||
"event": "message-end",
|
||||
"data": {
|
||||
"delta": {
|
||||
"finish_reason": "COMPLETE",
|
||||
"usage": {"tokens": {"input_tokens": 8, "output_tokens": 3}},
|
||||
}
|
||||
},
|
||||
}
|
||||
),
|
||||
]
|
||||
),
|
||||
sync_stream=True,
|
||||
)
|
||||
|
||||
chunks: Final = list(response_iterator)
|
||||
|
||||
assert chunks[0]["text"] == "Searching"
|
||||
assert chunks[1]["tool_use"] is not None
|
||||
assert chunks[1]["tool_use"]["id"] == "call-1"
|
||||
assert chunks[1]["tool_use"]["function"]["name"] == "lookup"
|
||||
assert chunks[2]["is_finished"]
|
||||
assert chunks[2]["finish_reason"] == "COMPLETE"
|
||||
assert chunks[2]["usage"]["prompt_tokens"] == 8
|
||||
assert chunks[2]["usage"]["completion_tokens"] == 3
|
||||
|
||||
|
||||
def test_cohere_v2_error_class_preserves_status_and_message() -> None:
|
||||
config: Final = CohereV2ChatConfig()
|
||||
|
||||
error: Final = config.get_error_class(
|
||||
error_message="rate limit exceeded",
|
||||
status_code=429,
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert error.status_code == 429
|
||||
assert error.message == "rate limit exceeded"
|
||||
|
||||
def test_v2_map_max_completion_tokens_overrides_max_tokens(self):
|
||||
"""max_completion_tokens maps to cohere max_tokens and overrides max_tokens, matching v1"""
|
||||
result = self.config.map_openai_params(
|
||||
|
|
@ -202,7 +447,6 @@ async def test_cohere_request_body_with_allowed_params():
|
|||
|
||||
# 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
|
||||
|
|
@ -259,7 +503,6 @@ async def test_cohere_documents_options_in_request_body():
|
|||
|
||||
# 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
|
||||
|
|
|
|||
87
tests/unit/llms/cohere/embed/test_transformation.py
Normal file
87
tests/unit/llms/cohere/embed/test_transformation.py
Normal file
|
|
@ -0,0 +1,87 @@
|
|||
from typing import Final
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.cohere.embed.transformation import CohereEmbeddingConfig
|
||||
from litellm.types.utils import EmbeddingResponse
|
||||
|
||||
|
||||
def test_cohere_embed_v4_request_maps_input_options() -> None:
|
||||
config: Final = CohereEmbeddingConfig()
|
||||
|
||||
request: Final = config.transform_embedding_request(
|
||||
model="embed-v4.0",
|
||||
input=["first", "second"],
|
||||
optional_params={
|
||||
"input_type": "search_document",
|
||||
"embedding_types": ["float"],
|
||||
"output_dimension": 256,
|
||||
},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request == {
|
||||
"model": "embed-v4.0",
|
||||
"texts": ["first", "second"],
|
||||
"input_type": "search_document",
|
||||
"embedding_types": ["float"],
|
||||
"output_dimension": 256,
|
||||
}
|
||||
|
||||
|
||||
def test_cohere_embed_v4_request_routes_base64_images() -> None:
|
||||
config: Final = CohereEmbeddingConfig()
|
||||
|
||||
request: Final = config.transform_embedding_request(
|
||||
model="embed-v4.0",
|
||||
input=["data:image/png;base64,aGVsbG8="],
|
||||
optional_params={"input_type": "search_document"},
|
||||
headers={},
|
||||
)
|
||||
|
||||
assert request == {
|
||||
"model": "embed-v4.0",
|
||||
"images": ["data:image/png;base64,aGVsbG8="],
|
||||
"input_type": "search_document",
|
||||
}
|
||||
|
||||
|
||||
def test_cohere_embed_v4_response_parses_multiple_embeddings() -> None:
|
||||
config: Final = CohereEmbeddingConfig()
|
||||
logging_obj: Final = MagicMock()
|
||||
logging_obj.model_call_details = {"input": ["first", "second"]}
|
||||
model_response: Final = EmbeddingResponse()
|
||||
raw_response: Final = httpx.Response(
|
||||
200,
|
||||
json={
|
||||
"embeddings": {
|
||||
"float": [[0.1, 0.2], [0.3, 0.4]],
|
||||
"int8": [[1, 2], [3, 4]],
|
||||
},
|
||||
"meta": {"billed_units": {"input_tokens": 6}},
|
||||
},
|
||||
)
|
||||
|
||||
response: Final = config.transform_embedding_response(
|
||||
model="embed-v4.0",
|
||||
raw_response=raw_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
api_key="test-api-key",
|
||||
request_data={"texts": ["first", "second"]},
|
||||
optional_params={"embedding_types": ["float", "int8"]},
|
||||
litellm_params={},
|
||||
)
|
||||
|
||||
assert response.object == "list"
|
||||
assert response.model == "embed-v4.0"
|
||||
assert [item["embedding"] for item in response.data] == [
|
||||
[0.1, 0.2],
|
||||
[0.3, 0.4],
|
||||
[1, 2],
|
||||
[3, 4],
|
||||
]
|
||||
assert [item["index"] for item in response.data] == [0, 1, 0, 1]
|
||||
assert response.usage.prompt_tokens == 6
|
||||
assert response.usage.total_tokens == 6
|
||||
Loading…
Add table
Reference in a new issue