from unittest import mock import pytest from aiohttp import ClientTimeout from fastapi.encoders import jsonable_encoder from pydantic import ValidationError from mlflow.environment_variables import MLFLOW_GATEWAY_ROUTE_TIMEOUT_SECONDS from mlflow.gateway.config import EndpointConfig from mlflow.gateway.exceptions import AIGatewayException from mlflow.gateway.providers.cohere import CohereProvider from mlflow.gateway.schemas import chat, completions, embeddings from tests.gateway.tools import MockAsyncResponse, MockAsyncStreamingResponse def chat_config(): return { "name": "chat", "endpoint_type": "llm/v1/chat", "model": { "provider": "cohere", "name": "command", "config": { "cohere_api_key": "key", }, }, } def chat_response(): return { "response_id": "abc123", "text": "\n\nThis is a test!", "generation_id": "def456", "token_count": { "prompt_tokens": 13, "response_tokens": 7, "total_tokens": 20, "billed_tokens": 20, }, "meta": { "api_version": {"version": "1"}, "billed_units": {"input_tokens": 13, "output_tokens": 7}, }, "tool_inputs": None, } def chat_payload(stream: bool = False): payload = { "messages": [ {"role": "user", "content": "Message 1"}, {"role": "assistant", "content": "Message 2"}, {"role": "user", "content": "Message 3"}, ], "temperature": 0.5, } if stream: payload["stream"] = True return payload @pytest.mark.asyncio async def test_chat(): resp = chat_response() config = chat_config() with ( mock.patch("time.time", return_value=1677858242), mock.patch("aiohttp.ClientSession.post", return_value=MockAsyncResponse(resp)) as mock_post, ): provider = CohereProvider(EndpointConfig(**config)) payload = chat_payload() response = await provider.chat(chat.RequestPayload(**payload)) assert jsonable_encoder(response) == { "id": "abc123", "object": "chat.completion", "created": 1677858242, "model": "command", "provider": "cohere", "choices": [ { "message": { "role": "assistant", "content": "\n\nThis is a test!", "tool_calls": None, "refusal": None, }, "finish_reason": None, "index": 0, } ], "usage": { "prompt_tokens": 13, "completion_tokens": 7, "total_tokens": 20, }, } mock_post.assert_called_once_with( "https://api.cohere.ai/v1/chat", json={ "model": "command", "chat_history": [ {"role": "USER", "message": "Message 1"}, {"role": "CHATBOT", "message": "Message 2"}, ], "message": "Message 3", "temperature": 1.25, }, timeout=ClientTimeout(total=MLFLOW_GATEWAY_ROUTE_TIMEOUT_SECONDS.get()), ) @pytest.mark.asyncio async def test_chat_with_system_messages(): resp = chat_response() config = chat_config() with ( mock.patch("time.time", return_value=1677858242), mock.patch("aiohttp.ClientSession.post", return_value=MockAsyncResponse(resp)) as mock_post, ): provider = CohereProvider(EndpointConfig(**config)) payload = { "messages": [ {"role": "system", "content": "System Message 1"}, {"role": "user", "content": "Message 1"}, {"role": "assistant", "content": "Message 2"}, {"role": "system", "content": "System Message 2"}, {"role": "user", "content": "Message 3"}, ], "temperature": 0.5, } await provider.chat(chat.RequestPayload(**payload)) mock_post.assert_called_once_with( "https://api.cohere.ai/v1/chat", json={ "model": "command", "chat_history": [ {"role": "USER", "message": "Message 1"}, {"role": "CHATBOT", "message": "Message 2"}, ], "message": "Message 3", "preamble_override": "System Message 1\nSystem Message 2", "temperature": 1.25, }, timeout=ClientTimeout(total=MLFLOW_GATEWAY_ROUTE_TIMEOUT_SECONDS.get()), ) @pytest.mark.parametrize( "params", [{"n": 2}, {"stop": ["test"]}], ) @pytest.mark.asyncio async def test_chat_throws_if_parameter_not_permitted(params): config = chat_config() provider = CohereProvider(EndpointConfig(**config)) payload = chat_payload() payload.update(params) with pytest.raises(AIGatewayException, match=r".*") as e: await provider.chat(chat.RequestPayload(**payload)) assert e.value.status_code == 422 def chat_stream_response(): return [ # first chunk b'{"is_finished":false,"event_type":"stream-start","generation_id":"test-id2"}', # subsequent chunks b'{"text":" Hi","is_finished":false,"event_type":"text-generation"}', b'{"text":" there","is_finished":false,"event_type":"text-generation"}', # final chunk b'{"is_finished":true,"event_type":"stream-end","response":{"response_id":"test-id1",' b'"text":"How are you","generation_id":"test-id2","token_count":{"prompt_tokens":83,' b'"response_tokens":63,"total_tokens":146,"billed_tokens":128},"tool_inputs":null},' b'"finish_reason":"COMPLETE"}', ] @pytest.mark.asyncio async def test_chat_stream(): resp = chat_stream_response() config = chat_config() with ( mock.patch("time.time", return_value=1677858242), mock.patch( "aiohttp.ClientSession.post", return_value=MockAsyncStreamingResponse(resp) ) as mock_post, ): provider = CohereProvider(EndpointConfig(**config)) payload = chat_payload(stream=True) response = provider.chat_stream(chat.RequestPayload(**payload)) chunks = [jsonable_encoder(chunk) async for chunk in response] assert chunks == [ { "choices": [ { "delta": { "role": None, "content": " Hi", "tool_calls": None, }, "finish_reason": None, "index": 0, } ], "created": 1677858242, "id": None, "model": "command", "provider": "cohere", "object": "chat.completion.chunk", "usage": { "prompt_tokens": None, "completion_tokens": None, "total_tokens": None, }, }, { "choices": [ { "delta": { "role": None, "content": " there", "tool_calls": None, }, "finish_reason": None, "index": 0, } ], "created": 1677858242, "id": None, "model": "command", "provider": "cohere", "object": "chat.completion.chunk", "usage": { "prompt_tokens": None, "completion_tokens": None, "total_tokens": None, }, }, { "choices": [ { "delta": { "role": None, "content": None, "tool_calls": None, }, "finish_reason": "COMPLETE", "index": 0, } ], "created": 1677858242, "id": "test-id1", "model": "command", "provider": "cohere", "object": "chat.completion.chunk", "usage": { "prompt_tokens": 83, "completion_tokens": 63, "total_tokens": 146, }, }, ] mock_post.assert_called_once_with( "https://api.cohere.ai/v1/chat", json={ "model": "command", "chat_history": [ {"role": "USER", "message": "Message 1"}, {"role": "CHATBOT", "message": "Message 2"}, ], "message": "Message 3", "temperature": 1.25, "stream": True, }, timeout=ClientTimeout(total=MLFLOW_GATEWAY_ROUTE_TIMEOUT_SECONDS.get()), ) def completions_config(): return { "name": "completions", "endpoint_type": "llm/v1/completions", "model": { "provider": "cohere", "name": "command", "config": { "cohere_api_key": "key", }, }, } def completions_response(): return { "id": "string", "generations": [ { "id": "string", "text": "This is a test", } ], "prompt": "string", "headers": {"Content-Type": "application/json"}, } @pytest.mark.asyncio async def test_completions(): resp = completions_response() config = completions_config() with ( mock.patch("time.time", return_value=1677858242), mock.patch("aiohttp.ClientSession.post", return_value=MockAsyncResponse(resp)) as mock_post, ): provider = CohereProvider(EndpointConfig(**config)) payload = { "prompt": "This is a test", "n": 1, "stop": ["foobar"], } response = await provider.completions(completions.RequestPayload(**payload)) assert jsonable_encoder(response) == { "id": None, "object": "text_completion", "created": 1677858242, "model": "command", "choices": [ { "text": "This is a test", "index": 0, "finish_reason": None, } ], "usage": {"prompt_tokens": None, "completion_tokens": None, "total_tokens": None}, } mock_post.assert_called_once_with( "https://api.cohere.ai/v1/generate", json={ "prompt": "This is a test", "model": "command", "num_generations": 1, "stop_sequences": ["foobar"], }, timeout=ClientTimeout(total=MLFLOW_GATEWAY_ROUTE_TIMEOUT_SECONDS.get()), ) @pytest.mark.asyncio async def test_completions_temperature_is_scaled_correctly(): resp = completions_response() config = completions_config() with mock.patch( "aiohttp.ClientSession.post", return_value=MockAsyncResponse(resp) ) as mock_post: provider = CohereProvider(EndpointConfig(**config)) payload = { "prompt": "This is a test", "temperature": 0.5, } await provider.completions(completions.RequestPayload(**payload)) assert mock_post.call_args[1]["json"]["temperature"] == 0.5 * 2.5 def completions_stream_response(): return [ b'{"text":" Hi","is_finished":false,"event_type":"text-generation"}', b'{"text":" there","is_finished":false,"event_type":"text-generation"}', # final chunk b'{"is_finished":true,"event_type":"stream-end","finish_reason":"COMPLETE",' b'"response":{"id":"test-id1","generations":' b'[{"id":"test-id2","text":" Hi there","finish_reason":"COMPLETE"}],' b'"prompt":"This is a test"}}', ] @pytest.mark.asyncio async def test_completions_stream(): resp = completions_stream_response() config = completions_config() with ( mock.patch("time.time", return_value=1677858242), mock.patch( "aiohttp.ClientSession.post", return_value=MockAsyncStreamingResponse(resp) ) as mock_post, ): provider = CohereProvider(EndpointConfig(**config)) payload = { "prompt": "This is a test", "n": 1, "stream": True, } response = provider.completions_stream(completions.RequestPayload(**payload)) chunks = [jsonable_encoder(chunk) async for chunk in response] assert chunks == [ { "choices": [ { "text": " Hi", "finish_reason": None, "index": 0, } ], "created": 1677858242, "id": None, "model": "command", "object": "text_completion_chunk", "usage": { "prompt_tokens": None, "completion_tokens": None, "total_tokens": None, }, }, { "choices": [ { "text": " there", "finish_reason": None, "index": 0, } ], "created": 1677858242, "id": None, "model": "command", "object": "text_completion_chunk", "usage": { "prompt_tokens": None, "completion_tokens": None, "total_tokens": None, }, }, { "choices": [ { "text": None, "finish_reason": "COMPLETE", "index": 0, } ], "created": 1677858242, "id": "test-id1", "model": "command", "object": "text_completion_chunk", "usage": { "prompt_tokens": None, "completion_tokens": None, "total_tokens": None, }, }, ] mock_post.assert_called_once_with( "https://api.cohere.ai/v1/generate", json={ "prompt": "This is a test", "model": "command", "num_generations": 1, "stream": True, }, timeout=ClientTimeout(total=MLFLOW_GATEWAY_ROUTE_TIMEOUT_SECONDS.get()), ) def embeddings_config(): return { "name": "embeddings", "endpoint_type": "llm/v1/embeddings", "model": { "provider": "cohere", "name": "embed-english-light-v2.0", "config": { "cohere_api_key": "key", }, }, } def embeddings_response(): return { "id": "bc57846a-3e56-4327-8acc-588ca1a37b8a", "texts": ["hello world"], "embeddings": [ [ 3.25, 0.7685547, 2.65625, -0.30126953, -2.3554688, 1.2597656, ] ], "meta": [ { "api_version": [ { "version": "1", } ] }, ], "headers": {"Content-Type": "application/json"}, } def embeddings_batch_response(): return { "id": "bc57846a-3e56-4327-8acc-588ca1a37b8a", "texts": ["hello world"], "embeddings": [ [ 3.25, 0.7685547, 2.65625, -0.30126953, -2.3554688, 1.2597656, ], [ 7.25, 0.7685547, 4.65625, -0.30126953, -2.3554688, 8.2597656, ], ], "meta": [ { "api_version": [ { "version": "1", } ] }, ], "headers": {"Content-Type": "application/json"}, } @pytest.mark.asyncio async def test_embeddings(): resp = embeddings_response() config = embeddings_config() with mock.patch( "aiohttp.ClientSession.post", return_value=MockAsyncResponse(resp) ) as mock_post: provider = CohereProvider(EndpointConfig(**config)) payload = {"input": "This is a test"} response = await provider.embeddings(embeddings.RequestPayload(**payload)) assert jsonable_encoder(response) == { "object": "list", "data": [ { "object": "embedding", "embedding": [ 3.25, 0.7685547, 2.65625, -0.30126953, -2.3554688, 1.2597656, ], "index": 0, } ], "model": "embed-english-light-v2.0", "usage": {"prompt_tokens": None, "total_tokens": None}, } mock_post.assert_called_once() @pytest.mark.asyncio async def test_batch_embeddings(): resp = embeddings_batch_response() config = embeddings_config() with mock.patch( "aiohttp.ClientSession.post", return_value=MockAsyncResponse(resp) ) as mock_post: provider = CohereProvider(EndpointConfig(**config)) payload = {"input": ["This is a", "batch test"]} response = await provider.embeddings(embeddings.RequestPayload(**payload)) assert jsonable_encoder(response) == { "object": "list", "data": [ { "object": "embedding", "embedding": [ 3.25, 0.7685547, 2.65625, -0.30126953, -2.3554688, 1.2597656, ], "index": 0, }, { "object": "embedding", "embedding": [ 7.25, 0.7685547, 4.65625, -0.30126953, -2.3554688, 8.2597656, ], "index": 1, }, ], "model": "embed-english-light-v2.0", "usage": {"prompt_tokens": None, "total_tokens": None}, } mock_post.assert_called_once() @pytest.mark.asyncio async def test_param_model_is_not_permitted(): config = embeddings_config() provider = CohereProvider(EndpointConfig(**config)) payload = { "prompt": "This should fail", "max_tokens": 5000, "model": "something-else", } with pytest.raises(AIGatewayException, match=r".*") as e: await provider.completions(completions.RequestPayload(**payload)) assert "The parameter 'model' is not permitted" in e.value.detail assert e.value.status_code == 422 @pytest.mark.parametrize("prompt", [{"set1", "set2"}, ["list1"], [1], ["list1", "list2"], [1, 2]]) @pytest.mark.asyncio async def test_completions_throws_if_prompt_contains_non_string(prompt): config = completions_config() provider = CohereProvider(EndpointConfig(**config)) payload = {"prompt": prompt} with pytest.raises(ValidationError, match=r"prompt"): await provider.completions(completions.RequestPayload(**payload))