Support Lemonade runtime context metadata

This commit is contained in:
Graham Neubig 2026-05-17 16:30:38 -04:00 • committed by openhands
parent 727a471ae9
commit 6f41a84c85
6 changed files with 273 additions and 11 deletions

View file

@ -868,6 +868,7 @@ openai_text_completion_compatible_providers: List = (
_openai_like_providers: List = [
"predibase",
"databricks",
"lemonade",
"watsonx",
] # private helper. similar to openai but require some custom auth / endpoint handling, so can't use the openai sdk
# well supported replicate llms

View file

@ -87,6 +87,7 @@ class ExceptionCheckers:
"is longer than the model's context length",
"input tokens exceed the configured limit",
"`inputs` tokens + `max_new_tokens` must be",
"exceeds the available context size", # llama.cpp/Lemonade
"exceeds the maximum number of tokens allowed", # Gemini
]
for substring in known_exception_substrings:
@ -891,12 +892,14 @@ def exception_type( # type: ignore # noqa: PLR0915
response=getattr(original_exception, "response", None),
litellm_debug_info=extra_information,
)
elif "model's maximum context limit" in error_str:
elif ExceptionCheckers.is_error_str_context_window_exceeded(error_str):
exception_mapping_worked = True
raise ContextWindowExceededError(
message=f"{custom_llm_provider.capitalize()}Exception: Context Window Error - {error_str}",
model=model,
llm_provider=custom_llm_provider,
response=getattr(original_exception, "response", None),
litellm_debug_info=extra_information,
)
elif "token_quota_reached" in error_str:
exception_mapping_worked = True

View file

@ -3,10 +3,12 @@ Translate from OpenAI's `/v1/chat/completions` to Lemonade's `/v1/chat/completio
"""
from typing import Any, List, Optional, Tuple, Union
from urllib.parse import quote
import httpx
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.secret_managers.main import get_secret_str
from litellm.types.llms.openai import (
@ -18,6 +20,8 @@ from ...openai_like.chat.transformation import OpenAILikeChatConfig
class LemonadeChatConfig(OpenAILikeChatConfig):
_DEFAULT_API_KEY = "lemonade"
repeat_penalty: Optional[float] = None
functions: Optional[list] = None
logit_bias: Optional[dict] = None
@ -68,7 +72,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
This method queries the Lemonade /models endpoint to retrieve the list of available models.
Args:
api_key: Optional API key (Lemonade doesn't require authentication)
api_key: Optional API key for authenticated Lemonade servers
api_base: Optional API base URL (defaults to LEMONADE_API_BASE env var or http://localhost:8000)
Returns:
@ -87,6 +91,7 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
try:
response = litellm.module_level_client.get(
url=f"{api_base}/models",
headers=self._get_auth_headers(api_key),
)
except Exception as e:
raise ValueError(
@ -101,6 +106,83 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
model_list = response.json().get("data", [])
return ["lemonade/" + model["id"] for model in model_list]
@staticmethod
def _get_positive_int(value: Any) -> Optional[int]:
if isinstance(value, bool):
return None
if isinstance(value, int) and value > 0:
return value
if isinstance(value, str):
try:
parsed = int(value)
except ValueError:
return None
if parsed > 0:
return parsed
return None
def _get_context_window(self, model_info: dict) -> Optional[int]:
recipe_options = model_info.get("recipe_options")
if not isinstance(recipe_options, dict):
recipe_options = {}
for value in (
recipe_options.get("ctx_size"),
model_info.get("max_input_tokens"),
model_info.get("context_window"),
model_info.get("max_context_window"),
):
parsed = self._get_positive_int(value)
if parsed is not None:
return parsed
return None
def _get_default_model_info(self, model: str) -> dict:
return {
"key": "lemonade/" + model,
"litellm_provider": "lemonade",
"mode": "chat",
"input_cost_per_token": 0.0,
"output_cost_per_token": 0.0,
"max_tokens": None,
"max_input_tokens": None,
"max_output_tokens": None,
}
def get_model_info(self, model: str, api_base: Optional[str] = None) -> Any:
if model.startswith("lemonade/"):
model = model.split("/", 1)[1]
api_base, api_key = self._get_openai_compatible_provider_info(
api_base=api_base, api_key=None
)
encoded_model = quote(model, safe="")
try:
response = litellm.module_level_client.get(
url=f"{api_base}/models/{encoded_model}",
headers=self._get_auth_headers(api_key),
)
response.raise_for_status()
model_info = response.json()
except Exception:
verbose_logger.debug("LemonadeError: Could not get model info.")
return self._get_default_model_info(model)
max_input_tokens = self._get_context_window(model_info)
max_output_tokens = self._get_positive_int(model_info.get("max_output_tokens"))
max_tokens = self._get_positive_int(model_info.get("max_tokens"))
model_info_response = self._get_default_model_info(model)
model_info_response.update(
{
"max_tokens": max_tokens or max_output_tokens,
"max_input_tokens": max_input_tokens,
"max_output_tokens": max_output_tokens,
}
)
return model_info_response
def _get_openai_compatible_provider_info(
self, api_base: Optional[str], api_key: Optional[str]
) -> Tuple[Optional[str], Optional[str]]:
@ -110,10 +192,19 @@ class LemonadeChatConfig(OpenAILikeChatConfig):
or get_secret_str("LEMONADE_API_BASE")
or "http://localhost:8000/api/v1"
) # type: ignore
# Lemonade doesn't check the key
key = "lemonade"
key = (
api_key
or litellm.lemonade_key
or get_secret_str("LEMONADE_API_KEY")
or self._DEFAULT_API_KEY
)
return api_base, key
def _get_auth_headers(self, api_key: Optional[str]) -> dict:
if api_key is None or api_key == self._DEFAULT_API_KEY:
return {}
return {"Authorization": f"Bearer {api_key}"}
def transform_response(
self,
model: str,

View file

@ -5727,6 +5727,10 @@ def _get_model_info_helper( # noqa: PLR0915
custom_llm_provider == "ollama" or custom_llm_provider == "ollama_chat"
) and not _is_potential_model_name_in_model_cost(potential_model_names):
return litellm.OllamaConfig().get_model_info(model, api_base=api_base)
elif custom_llm_provider == "lemonade":
return litellm.LemonadeChatConfig().get_model_info(
model=model, api_base=api_base
)
else:
"""
Check if: (in order of specificity)

View file

@ -14,6 +14,7 @@ from litellm.litellm_core_utils.exception_mapping_utils import (
exception_type,
extract_and_raise_litellm_exception,
)
from litellm.llms.openai.common_utils import OpenAIError
# Test cases for is_error_str_context_window_exceeded
# Tuple format: (error_message, expected_result)
@ -41,6 +42,10 @@ context_window_test_cases = [
"`inputs` tokens + `max_new_tokens` must be <= 4096",
True,
),
(
"request (67311 tokens) exceeds the available context size (65536 tokens), try increasing it",
True,
),
# Gemini 2.5/3 format
(
"The input token count exceeds the maximum number of tokens allowed 1048576.",
@ -182,7 +187,6 @@ class TestExceptionCheckers:
]
for error_str in positive_cases:
print("testing positive case=", error_str)
result = ExceptionCheckers.is_azure_content_policy_violation_error(
error_str
)
@ -255,6 +259,33 @@ def test_gemini_context_window_error_mapping(
)
def test_lemonade_context_window_error_mapping():
"""Lemonade's llama.cpp backend should map context overflows to LiteLLM's standard error."""
model = "lemonade/Qwen3.6-35B-A3B-GGUF"
error_message = (
'{"error":{"code":"context_length_exceeded","message":"request '
"(80010 tokens) exceeds the available context size (65536 tokens), "
'try increasing it","status_code":400,"type":"invalid_request_error"}}'
)
original_exception = OpenAIError(
status_code=400,
message=error_message,
headers={},
)
with pytest.raises(litellm.ContextWindowExceededError) as excinfo:
exception_type(
model=model,
original_exception=original_exception,
custom_llm_provider="lemonade",
)
assert excinfo.value.status_code == 400
assert excinfo.value.llm_provider == "lemonade"
assert excinfo.value.model == model
# Test cases for Vertex AI RateLimitError mapping
# As per https://github.com/BerriAI/litellm/issues/16189
vertex_rate_limit_test_cases = [

View file

@ -1,17 +1,14 @@
import json
import os
import sys
import pytest
sys.path.insert(
0, os.path.abspath("../../../../..")
) # Adds the parent directory to the system path
from unittest.mock import MagicMock, patch
import litellm
from litellm.llms.lemonade.chat.transformation import LemonadeChatConfig
from litellm.types.utils import ModelResponse
import httpx
def test_lemonade_config_initialization():
@ -28,8 +25,11 @@ def test_lemonade_config_initialization():
assert config.repeat_penalty == 1.1
def test_get_openai_compatible_provider_info():
def test_get_openai_compatible_provider_info(monkeypatch):
"""Test the provider info method returns correct API base and key"""
monkeypatch.delenv("LEMONADE_API_KEY", raising=False)
monkeypatch.setattr(litellm, "lemonade_key", None)
monkeypatch.setattr(litellm, "api_key", None)
config = LemonadeChatConfig()
api_base, key = config._get_openai_compatible_provider_info(
@ -40,8 +40,11 @@ def test_get_openai_compatible_provider_info():
assert key == "lemonade"
def test_get_openai_compatible_provider_info_with_custom_base():
def test_get_openai_compatible_provider_info_with_custom_base(monkeypatch):
"""Test the provider info method with custom API base"""
monkeypatch.delenv("LEMONADE_API_KEY", raising=False)
monkeypatch.setattr(litellm, "lemonade_key", None)
monkeypatch.setattr(litellm, "api_key", None)
config = LemonadeChatConfig()
custom_api_base = "https://custom.lemonade.ai/v1"
@ -53,6 +56,135 @@ def test_get_openai_compatible_provider_info_with_custom_base():
assert key == "lemonade"
def test_get_openai_compatible_provider_info_with_api_key_env(monkeypatch):
"""Test the provider info method reads Lemonade's API key from the environment."""
monkeypatch.setenv("LEMONADE_API_KEY", "test-key")
monkeypatch.setattr(litellm, "lemonade_key", None)
monkeypatch.setattr(litellm, "api_key", None)
config = LemonadeChatConfig()
api_base, key = config._get_openai_compatible_provider_info(
api_base=None, api_key=None
)
assert api_base == "http://localhost:8000/api/v1"
assert key == "test-key"
def test_get_openai_compatible_provider_info_ignores_global_api_key(monkeypatch):
"""Test that Lemonade discovery does not send unrelated global API keys."""
monkeypatch.delenv("LEMONADE_API_KEY", raising=False)
monkeypatch.setattr(litellm, "lemonade_key", None)
monkeypatch.setattr(litellm, "api_key", "global-openai-key")
config = LemonadeChatConfig()
api_base, key = config._get_openai_compatible_provider_info(
api_base="http://lemonade.test/v1", api_key=None
)
assert api_base == "http://lemonade.test/v1"
assert key == "lemonade"
assert config._get_auth_headers(key) == {}
def test_get_model_info_uses_loaded_context_size():
"""Test that Lemonade model info prefers the effective loaded ctx_size."""
config = LemonadeChatConfig()
response = MagicMock()
response.status_code = 200
response.json.return_value = {
"id": "Qwen3.6-35B-A3B-GGUF",
"recipe_options": {"ctx_size": 65536},
"max_context_window": 262144,
}
with patch.object(
litellm.module_level_client, "get", return_value=response
) as mock_get:
model_info = config.get_model_info(
model="lemonade/Qwen3.6-35B-A3B-GGUF",
api_base="http://lemonade.test/v1",
)
assert model_info["key"] == "lemonade/Qwen3.6-35B-A3B-GGUF"
assert model_info["litellm_provider"] == "lemonade"
assert model_info["max_input_tokens"] == 65536
assert "supports_function_calling" not in model_info
assert "supports_response_schema" not in model_info
assert "supports_tool_choice" not in model_info
assert mock_get.call_args.kwargs["headers"] == {}
def test_get_model_info_falls_back_when_server_unavailable():
"""Test that Lemonade metadata lookup failures return safe defaults."""
config = LemonadeChatConfig()
with patch.object(
litellm.module_level_client, "get", side_effect=Exception("boom")
):
model_info = config.get_model_info(
model="lemonade/Qwen3.6-35B-A3B-GGUF",
api_base="http://lemonade.test/v1",
)
assert model_info["key"] == "lemonade/Qwen3.6-35B-A3B-GGUF"
assert model_info["litellm_provider"] == "lemonade"
assert model_info["mode"] == "chat"
assert model_info["input_cost_per_token"] == 0.0
assert model_info["output_cost_per_token"] == 0.0
assert model_info["max_tokens"] is None
assert model_info["max_input_tokens"] is None
assert model_info["max_output_tokens"] is None
assert "supports_function_calling" not in model_info
assert "supports_response_schema" not in model_info
assert "supports_tool_choice" not in model_info
def test_get_model_info_sends_lemonade_api_key(monkeypatch):
"""Test that Lemonade model info uses auth for authenticated servers."""
monkeypatch.setenv("LEMONADE_API_KEY", "test-key")
monkeypatch.setattr(litellm, "lemonade_key", None)
monkeypatch.setattr(litellm, "api_key", None)
config = LemonadeChatConfig()
response = MagicMock()
response.status_code = 200
response.json.return_value = {
"id": "Qwen3.6-35B-A3B-GGUF",
"recipe_options": {"ctx_size": 65536},
}
with patch.object(
litellm.module_level_client, "get", return_value=response
) as mock_get:
config.get_model_info(
model="lemonade/Qwen3.6-35B-A3B-GGUF",
api_base="http://lemonade.test/v1",
)
assert mock_get.call_args.kwargs["headers"] == {"Authorization": "Bearer test-key"}
def test_litellm_get_model_info_uses_lemonade_api_base():
"""Test that LiteLLM model info is wired to Lemonade's model metadata API."""
response = MagicMock()
response.status_code = 200
response.json.return_value = {
"id": "Qwen3.6-35B-A3B-GGUF",
"max_input_tokens": 65536,
"max_context_window": 262144,
}
with patch.object(litellm.module_level_client, "get", return_value=response):
model_info = litellm.get_model_info(
model="lemonade/Qwen3.6-35B-A3B-GGUF",
api_base="http://lemonade.test/v1",
)
assert model_info["max_input_tokens"] == 65536
assert response.raise_for_status.called
assert response.json.called
def test_transform_response():
"""Test the response transformation adds lemonade prefix to model name"""
config = LemonadeChatConfig()