From 5288ceaca5131d94ec4cb817dc95de99acf39353 Mon Sep 17 00:00:00 2001 From: Robert Keus Date: Mon, 27 Jul 2026 08:44:47 +0200 Subject: [PATCH] fix(greenpt): enable embeddings and rerank cost tracking --- litellm/llms/greenpt/rerank/transformation.py | 19 +++++++++ litellm/llms/openai_like/json_loader.py | 7 ++++ litellm/main.py | 1 + .../greenpt/rerank/test_transformation.py | 26 +++++++++++++ .../llms/openai_like/test_greenpt_provider.py | 39 +++++++++++++++++-- 5 files changed, 88 insertions(+), 4 deletions(-) rename tests/{litellm => test_litellm}/llms/greenpt/rerank/test_transformation.py (70%) rename tests/{litellm => test_litellm}/llms/openai_like/test_greenpt_provider.py (67%) diff --git a/litellm/llms/greenpt/rerank/transformation.py b/litellm/llms/greenpt/rerank/transformation.py index 54daa1fe4c7..b36c3029dae 100644 --- a/litellm/llms/greenpt/rerank/transformation.py +++ b/litellm/llms/greenpt/rerank/transformation.py @@ -1,5 +1,7 @@ from litellm.llms.hosted_vllm.rerank.transformation import HostedVLLMRerankConfig from litellm.secret_managers.main import get_secret_str +from litellm.types.rerank import RerankBilledUnits +from litellm.types.utils import ModelInfo class GreenPTRerankConfig(HostedVLLMRerankConfig): @@ -38,3 +40,20 @@ class GreenPTRerankConfig(HostedVLLMRerankConfig): ) request.setdefault("top_n", len(request["documents"])) return request + + def calculate_rerank_cost( + self, + model: str, + custom_llm_provider: str | None = None, + billed_units: RerankBilledUnits | None = None, + model_info: ModelInfo | None = None, + ) -> tuple[float, float]: + if model_info is None or billed_units is None: + return 0.0, 0.0 + input_cost_per_token = model_info.get("input_cost_per_token") + if input_cost_per_token is None: + return 0.0, 0.0 + total_tokens = billed_units.get("total_tokens") + if total_tokens is None: + return 0.0, 0.0 + return input_cost_per_token * total_tokens, 0.0 diff --git a/litellm/llms/openai_like/json_loader.py b/litellm/llms/openai_like/json_loader.py index 4640bb8a422..451e05364a1 100644 --- a/litellm/llms/openai_like/json_loader.py +++ b/litellm/llms/openai_like/json_loader.py @@ -73,6 +73,13 @@ class JSONProviderRegistry: return False return "/v1/responses" in provider.supported_endpoints + @classmethod + def supports_embeddings_api(cls, slug: str) -> bool: + provider = cls._providers.get(slug) + if provider is None: + return False + return "/v1/embeddings" in provider.supported_endpoints + @classmethod def list_providers(cls) -> list: """List all registered provider slugs""" diff --git a/litellm/main.py b/litellm/main.py index dc3ec469a1b..c08efdf76f9 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6213,6 +6213,7 @@ def embedding( custom_llm_provider == "openai_like" or custom_llm_provider == "llamafile" or custom_llm_provider == "lm_studio" + or JSONProviderRegistry.supports_embeddings_api(custom_llm_provider) ): api_base = api_base or litellm.api_base or get_secret_str("OPENAI_LIKE_API_BASE") diff --git a/tests/litellm/llms/greenpt/rerank/test_transformation.py b/tests/test_litellm/llms/greenpt/rerank/test_transformation.py similarity index 70% rename from tests/litellm/llms/greenpt/rerank/test_transformation.py rename to tests/test_litellm/llms/greenpt/rerank/test_transformation.py index cf12bfd80fd..9ae9df8edcf 100644 --- a/tests/litellm/llms/greenpt/rerank/test_transformation.py +++ b/tests/test_litellm/llms/greenpt/rerank/test_transformation.py @@ -64,3 +64,29 @@ def test_transform_rerank_request_preserves_top_n(): headers={}, ) assert request["top_n"] == 1 + + +def test_calculate_rerank_cost(): + prompt_cost, completion_cost = GreenPTRerankConfig().calculate_rerank_cost( + model="green-rerank", + billed_units={"total_tokens": 1000}, + model_info={"input_cost_per_token": 1.36524e-07}, + ) + assert prompt_cost == pytest.approx(0.000136524) + assert completion_cost == 0.0 + + +@pytest.mark.parametrize( + ("billed_units", "model_info"), + [ + (None, None), + ({}, {"input_cost_per_token": 1.36524e-07}), + ({"total_tokens": 1000}, {}), + ], +) +def test_calculate_rerank_cost_missing_usage_or_pricing(billed_units, model_info): + assert GreenPTRerankConfig().calculate_rerank_cost( + model="green-rerank", + billed_units=billed_units, + model_info=model_info, + ) == (0.0, 0.0) diff --git a/tests/litellm/llms/openai_like/test_greenpt_provider.py b/tests/test_litellm/llms/openai_like/test_greenpt_provider.py similarity index 67% rename from tests/litellm/llms/openai_like/test_greenpt_provider.py rename to tests/test_litellm/llms/openai_like/test_greenpt_provider.py index efe5f2808b9..bf78b5e94be 100644 --- a/tests/litellm/llms/openai_like/test_greenpt_provider.py +++ b/tests/test_litellm/llms/openai_like/test_greenpt_provider.py @@ -1,8 +1,12 @@ +import json import os import sys +from unittest.mock import Mock, patch sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../.."))) +import litellm +from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.llms.openai_like.dynamic_config import create_config_class from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.types.utils import LlmProviders @@ -39,7 +43,7 @@ def test_greenpt_maps_max_completion_tokens(): params = config.map_openai_params( non_default_params={"max_completion_tokens": 256}, optional_params={}, - model="greenpt/gemma4", + model="greenpt/glm-5.2", drop_params=False, ) assert params == {"max_tokens": 256} @@ -50,7 +54,7 @@ def test_greenpt_complete_url_appends_endpoint(): url = config.get_complete_url( api_base=GREENPT_BASE_URL, api_key="test-key", - model="greenpt/gemma4", + model="greenpt/glm-5.2", optional_params={}, litellm_params={}, stream=False, @@ -62,17 +66,44 @@ def test_greenpt_provider_resolution(): from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider model, provider, _, api_base = get_llm_provider( - model="greenpt/gemma4", + model="greenpt/glm-5.2", custom_llm_provider=None, api_base=None, api_key=None, ) - assert model == "gemma4" + assert model == "glm-5.2" assert provider == "greenpt" assert api_base == GREENPT_BASE_URL +def test_greenpt_embedding_dispatches(): + response = Mock() + response.raise_for_status = Mock() + response.json.return_value = { + "object": "list", + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "model": "green-embedding", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + } + client = HTTPHandler() + + with patch.object(HTTPHandler, "post", return_value=response) as mock_post: + result = litellm.embedding( + model="greenpt/green-embedding", + input=["renewable inference"], + api_key="test-key", + client=client, + ) + + assert result.data[0]["embedding"] == [0.1, 0.2] + assert mock_post.call_args.args[0] == f"{GREENPT_BASE_URL}/embeddings" + assert json.loads(mock_post.call_args.kwargs["data"]) == { + "model": "green-embedding", + "input": ["renewable inference"], + } + + def test_greenpt_rerank_config(): config = ProviderConfigManager.get_provider_rerank_config( model="green-rerank",