diff --git a/litellm/llms/triton/embedding/transformation.py b/litellm/llms/triton/embedding/transformation.py index 8ab0277e369..32861930914 100644 --- a/litellm/llms/triton/embedding/transformation.py +++ b/litellm/llms/triton/embedding/transformation.py @@ -8,7 +8,8 @@ from litellm.llms.base_llm.embedding.transformation import ( LiteLLMLoggingObj, ) from litellm.types.llms.openai import AllEmbeddingInputValues -from litellm.types.utils import EmbeddingResponse +from litellm.types.utils import EmbeddingResponse, Usage +from litellm.utils import token_counter from ..common_utils import TritonError @@ -103,8 +104,35 @@ class TritonEmbeddingConfig(BaseEmbeddingConfig): model_response.model = raw_response_json.get("model_name", "None") model_response.data = _embedding_output + model_response.usage = self._build_embedding_usage( + model=model, request_data=request_data + ) return model_response + def _build_embedding_usage(self, model: str, request_data: dict) -> Usage: + input_data = request_data.get("inputs", []) + input_text_values: List[str] = [] + for item in input_data: + if isinstance(item, dict) and item.get("name") == "input_text": + data_values = item.get("data", []) + if isinstance(data_values, list): + input_text_values = [str(value) for value in data_values] + break + + prompt_tokens = 0 + input_text = "\n".join(input_text_values) + if len(input_text) > 0: + try: + prompt_tokens = token_counter(model=model, text=input_text) + except Exception: + prompt_tokens = len(input_text.split()) + + return Usage( + prompt_tokens=prompt_tokens, + completion_tokens=0, + total_tokens=prompt_tokens, + ) + def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] ) -> BaseLLMException: diff --git a/tests/llm_translation/test_triton.py b/tests/llm_translation/test_triton.py index 8a3bbb4661c..08b15745bec 100644 --- a/tests/llm_translation/test_triton.py +++ b/tests/llm_translation/test_triton.py @@ -50,6 +50,94 @@ def test_split_embedding_by_shape_fails_with_shape_value_error(): ) +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 + + @pytest.mark.parametrize("stream", [True, False]) def test_completion_triton_generate_api(stream): try: