Merge pull request #24778 from Sameerlite/litellm_triton_embedding_usage_tracking

feat(triton): add embedding usage tracking fallback

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:

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: