import json import traceback from dotenv import load_dotenv load_dotenv() import io from unittest.mock import AsyncMock, MagicMock, patch import pytest import litellm from litellm.llms.triton.embedding.transformation import TritonEmbeddingConfig from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE def test_split_embedding_by_shape_passes(): try: data = [ { "shape": [2, 3], "data": [1, 2, 3, 4, 5, 6], } ] split_output_data = TritonEmbeddingConfig.split_embedding_by_shape( data[0]["data"], data[0]["shape"] ) assert split_output_data == [[1, 2, 3], [4, 5, 6]] except Exception as e: pytest.fail(f"An exception occured: {e}") def test_split_embedding_by_shape_fails_with_shape_value_error(): data = [ { "shape": [2], "data": [1, 2, 3, 4, 5, 6], } ] with pytest.raises(ValueError, match='Shape must be of length'): TritonEmbeddingConfig.split_embedding_by_shape( data[0]["data"], data[0]["shape"] ) def test_triton_embedding_response_sets_usage_with_token_counter(): config = TritonEmbeddingConfig() mock_http_response = MagicMock() mock_http_response.status_code = 200 mock_http_response.json.return_value = { "model_name": "gte-base-en-v1", "outputs": [ { "name": "embedding", "shape": [1, 2], "data": [0.1, 0.2], } ], } model_response = litellm.EmbeddingResponse() request_data = { "inputs": [ { "name": "input_text", "shape": [1], "datatype": "BYTES", "data": ["hello from triton"], } ] } with patch( "litellm.llms.triton.embedding.transformation.token_counter", return_value=7, ): transformed = config.transform_embedding_response( model="triton/gte-base-en-v1", raw_response=mock_http_response, model_response=model_response, logging_obj=MagicMock(), request_data=request_data, ) assert transformed.usage is not None assert transformed.usage.prompt_tokens == 7 assert transformed.usage.completion_tokens == 0 assert transformed.usage.total_tokens == 7 def test_triton_embedding_response_sets_usage_with_word_count_fallback(): config = TritonEmbeddingConfig() mock_http_response = MagicMock() mock_http_response.status_code = 200 mock_http_response.json.return_value = { "model_name": "gte-base-en-v1", "outputs": [ { "name": "embedding", "shape": [1, 2], "data": [0.1, 0.2], } ], } model_response = litellm.EmbeddingResponse() request_data = { "inputs": [ { "name": "input_text", "shape": [1], "datatype": "BYTES", "data": ["hello from triton"], } ] } with patch( "litellm.llms.triton.embedding.transformation.token_counter", side_effect=Exception("tokenizer error"), ): transformed = config.transform_embedding_response( model="triton/gte-base-en-v1", raw_response=mock_http_response, model_response=model_response, logging_obj=MagicMock(), request_data=request_data, ) assert transformed.usage is not None assert transformed.usage.prompt_tokens == 3 assert transformed.usage.completion_tokens == 0 assert transformed.usage.total_tokens == 3 def test_triton_embedding_batch_usage_sums_per_input_token_counts(): """Batch inputs must not be joined before token counting (avoids extra newline tokens).""" config = TritonEmbeddingConfig() mock_http_response = MagicMock() mock_http_response.status_code = 200 mock_http_response.json.return_value = { "model_name": "gte-base-en-v1", "outputs": [ { "name": "embedding", "shape": [2, 2], "data": [0.1, 0.2, 0.3, 0.4], } ], } model_response = litellm.EmbeddingResponse() request_data = { "inputs": [ { "name": "input_text", "shape": [2], "datatype": "BYTES", "data": ["first input", "second input"], } ] } with patch( "litellm.llms.triton.embedding.transformation.token_counter", side_effect=[5, 7], ): transformed = config.transform_embedding_response( model="triton/gte-base-en-v1", raw_response=mock_http_response, model_response=model_response, logging_obj=MagicMock(), request_data=request_data, ) assert transformed.usage is not None assert transformed.usage.prompt_tokens == 12 assert transformed.usage.total_tokens == 12 @pytest.mark.parametrize("stream", [True, False]) def test_completion_triton_generate_api(stream): try: mock_response = MagicMock() if stream: def mock_iter_lines(): mock_output = "".join( [ 'data: {"model_name":"ensemble","model_version":"1","sequence_end":false,"sequence_id":0,"sequence_start":false,"text_output":"' + t + '"}\n\n' for t in ["I", " am", " an", " AI", " assistant"] ] ) for out in mock_output.split("\n"): yield out mock_response.iter_lines = mock_iter_lines else: def return_val(): return { "text_output": "I am an AI assistant", } mock_response.json = return_val mock_response.status_code = 200 with patch( "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", return_value=mock_response, ) as mock_post: response = litellm.completion( model="triton/llama-3-8b-instruct", messages=[{"role": "user", "content": "who are u?"}], max_tokens=10, timeout=5, api_base="http://localhost:8000/generate", stream=stream, ) # Verify the call was made mock_post.assert_called_once() # Get the arguments passed to the post request print("call args", mock_post.call_args) call_kwargs = mock_post.call_args.kwargs # Access kwargs directly # Verify URL if stream: assert call_kwargs["url"] == "http://localhost:8000/generate_stream" else: assert call_kwargs["url"] == "http://localhost:8000/generate" # Parse the request data from the JSON string request_data = json.loads(call_kwargs["data"]) # Verify request data assert request_data["text_input"] == "who are u?" assert request_data["parameters"]["max_tokens"] == 10 # Verify response if stream: tokens = ["I", " am", " an", " AI", " assistant", None] idx = 0 for chunk in response: assert chunk.choices[0].delta.content == tokens[idx] idx += 1 assert idx == len(tokens) else: assert response.choices[0].message.content == "I am an AI assistant" except Exception as e: print("exception", e) import traceback traceback.print_exc() pytest.fail(f"Error occurred: {e}") def test_completion_triton_infer_api(): litellm.set_verbose = True try: mock_response = MagicMock() def return_val(): return { "model_name": "basketgpt", "model_version": "2", "outputs": [ { "name": "text_output", "datatype": "BYTES", "shape": [1], "data": [ "0004900005024 0004900006774 0004900005024 0004900005027 0004900005026 0004900005025 0004900005027 0004900005024 0004900006774 0004900005027" ], }, { "name": "debug_probs", "datatype": "FP32", "shape": [0], "data": [], }, { "name": "debug_tokens", "datatype": "BYTES", "shape": [0], "data": [], }, ], } mock_response.json = return_val mock_response.status_code = 200 with patch( "litellm.llms.custom_httpx.http_handler.HTTPHandler.post", return_value=mock_response, ) as mock_post: response = litellm.completion( model="triton/llama-3-8b-instruct", messages=[ { "role": "user", "content": "0004900005025 0004900005026 0004900005027", } ], api_base="http://localhost:8000/infer", ) print("litellm response", response.model_dump_json(indent=4)) # Verify the call was made mock_post.assert_called_once() # Get the arguments passed to the post request call_kwargs = mock_post.call_args.kwargs # Verify URL assert call_kwargs["url"] == "http://localhost:8000/infer" # Parse the request data from the JSON string request_data = json.loads(call_kwargs["data"]) # Verify request matches expected Triton format assert request_data["inputs"][0]["name"] == "text_input" assert request_data["inputs"][0]["shape"] == [1] assert request_data["inputs"][0]["datatype"] == "BYTES" assert request_data["inputs"][0]["data"] == [ "0004900005025 0004900005026 0004900005027" ] assert request_data["inputs"][1]["shape"] == [1] assert request_data["inputs"][1]["datatype"] == "INT32" assert request_data["inputs"][1]["data"] == [20] # Verify response format matches expected completion format assert ( response.choices[0].message.content == "0004900005024 0004900006774 0004900005024 0004900005027 0004900005026 0004900005025 0004900005027 0004900005024 0004900006774 0004900005027" ) assert response.choices[0].finish_reason == "stop" assert response.choices[0].index == 0 assert response.object == "chat.completion" except Exception as e: print("exception", e) traceback.print_exc() pytest.fail(f"Error occurred: {e}") @pytest.mark.asyncio async def test_triton_embeddings(): try: litellm.set_verbose = True response = await litellm.aembedding( model="triton/my-triton-model", api_base=f"{FAKE_OPENAI_API_BASE}/triton/embeddings", input=["good morning from litellm"], ) print(f"response: {response}") # stubbed endpoint is setup to return this assert response.data[0]["embedding"] == [0.1, 0.2] except Exception as e: pytest.fail(f"Error occurred: {e}") def test_triton_generate_raw_request(): from litellm.utils import return_raw_request from litellm.types.utils import CallTypes try: kwargs = { "model": "triton/llama-3-8b-instruct", "messages": [{"role": "user", "content": "who are u?"}], "api_base": "http://localhost:8000/generate", } raw_request = return_raw_request(endpoint=CallTypes.completion, kwargs=kwargs) print("raw_request", raw_request) assert raw_request is not None assert "bad_words" not in json.dumps(raw_request["raw_request_body"]) assert "stop_words" not in json.dumps(raw_request["raw_request_body"]) except Exception as e: pytest.fail(f"Error occurred: {e}")