diff --git a/openjudge/models/__init__.py b/openjudge/models/__init__.py index d2eb8f998..d87bbe911 100644 --- a/openjudge/models/__init__.py +++ b/openjudge/models/__init__.py @@ -4,6 +4,7 @@ """ from openjudge.models.base_chat_model import BaseChatModel +from openjudge.models.litellm_chat_model import LiteLLMChatModel from openjudge.models.minimax_chat_model import MiniMaxChatModel from openjudge.models.openai_chat_model import OpenAIChatModel from openjudge.models.qiniu_chat_model import QiniuChatModel @@ -11,6 +12,7 @@ __all__ = [ "BaseChatModel", + "LiteLLMChatModel", "MiniMaxChatModel", "OpenAIChatModel", "QiniuChatModel", diff --git a/openjudge/models/litellm_chat_model.py b/openjudge/models/litellm_chat_model.py new file mode 100644 index 000000000..50c30c053 --- /dev/null +++ b/openjudge/models/litellm_chat_model.py @@ -0,0 +1,160 @@ +# -*- coding: utf-8 -*- +"""LiteLLM chat model. + +LiteLLM exposes 100+ LLM providers behind a single OpenAI-compatible +interface, routed by the model-name prefix (e.g. ``anthropic/claude-sonnet-4-6``, +``gemini/gemini-2.5-flash``, ``bedrock/...``). Because LiteLLM returns +OpenAI-shaped responses, this model reuses :class:`OpenAIChatModel`'s response +parsing (non-streaming and streaming) and only swaps the transport from the +OpenAI SDK to ``litellm.acompletion``. + +Calling the SDK directly (rather than pointing ``OpenAIChatModel`` at a base +URL) lets LiteLLM use each provider's native authentication (Bedrock SigV4, +Vertex ADC, Azure AD) instead of only an OpenAI-style bearer token. +""" + +from typing import Any, AsyncGenerator, Callable, Dict, Literal, Type + +from loguru import logger +from pydantic import BaseModel + +from openjudge.models.openai_chat_model import OpenAIChatModel +from openjudge.models.schema.oai.message import ChatMessage +from openjudge.models.schema.oai.response import ChatResponse + + +class LiteLLMChatModel(OpenAIChatModel): + """Chat model backed by the LiteLLM SDK. + + Reuses :class:`OpenAIChatModel`'s OpenAI-shaped response handling and routes + generation through ``litellm.acompletion``, so a single class reaches any + LiteLLM-supported provider selected by the ``model`` prefix. + """ + + def __init__( + self, + model: str, + api_key: str | None = None, + base_url: str | None = None, + stream: bool = False, + reasoning_effort: Literal["low", "medium", "high"] | None = None, + drop_params: bool = True, + max_retries: int | None = None, + timeout: float | None = None, + **kwargs: Any, + ) -> None: + """Initialize the LiteLLM chat model. + + Args: + model: The LiteLLM model id, including its provider prefix where + required (e.g. ``gpt-4o``, ``anthropic/claude-sonnet-4-6``, + ``gemini/gemini-2.5-flash``). + api_key: Optional API key. When omitted, LiteLLM falls back to each + target provider's own environment variable (``OPENAI_API_KEY``, + ``ANTHROPIC_API_KEY``, ...). Set it (with ``base_url``) to target + a LiteLLM proxy. + base_url: Optional API base, forwarded to LiteLLM as ``api_base`` + (e.g. a LiteLLM proxy URL). + stream: Whether to stream the model output. + reasoning_effort: Reasoning effort for models that support it. + drop_params: Drop per-provider-unsupported params instead of raising, + so one call shape works across providers. Defaults to ``True``. + max_retries: Optional retry count, forwarded to LiteLLM as + ``num_retries``. + timeout: Optional request timeout in seconds. + kwargs: Extra generation kwargs forwarded to ``litellm.acompletion`` + (e.g. ``temperature``, ``top_p``). + """ + # Set the base attributes directly rather than calling + # OpenAIChatModel.__init__, which builds a persistent AsyncOpenAI + # client: LiteLLM has no persistent client and routes per call. + self.model = model + self.stream = stream + self.reasoning_effort = reasoning_effort + self.kwargs = kwargs or {} + self.drop_params = drop_params + self.api_key = api_key + self.base_url = base_url + self.max_retries = max_retries + self.timeout = timeout + + async def achat( + self, + messages: list[dict | ChatMessage], + tools: list[dict] | None = None, + tool_choice: Literal["auto", "none", "any", "required"] | str | None = None, + structured_model: Type[BaseModel] | None = None, + callback: Callable | None = None, + **kwargs: Any, + ) -> ChatResponse | AsyncGenerator[ChatResponse, None]: + """Get a response from LiteLLM for the given arguments. + + The parameters mirror :meth:`OpenAIChatModel.achat`; see that method for + details on ``tools``, ``tool_choice``, ``structured_model`` and + ``callback``. + + Returns: + Either a single :class:`ChatResponse` or, when ``stream`` is set, an + async generator of :class:`ChatResponse` chunks. + """ + import litellm + + if not isinstance(messages, list): + raise ValueError( + f"LiteLLM `messages` field expected type `list`, got `{type(messages)}` instead.", + ) + messages = [msg.to_dict() if isinstance(msg, ChatMessage) else msg for msg in messages] + + call_kwargs: Dict[str, Any] = { + "model": self.model, + "messages": messages, + "stream": self.stream, + # Drop params a given provider does not support instead of erroring, + # so the same request shape works across every backend. + "drop_params": self.drop_params, + **self.kwargs, + **kwargs, + } + if self.reasoning_effort and "reasoning_effort" not in call_kwargs: + call_kwargs["reasoning_effort"] = self.reasoning_effort + + # Forward credentials only when set, so LiteLLM otherwise falls back to + # each provider's own env var; setting them targets a LiteLLM proxy. + if self.api_key: + call_kwargs.setdefault("api_key", self.api_key) + if self.base_url: + call_kwargs.setdefault("api_base", self.base_url) + if self.max_retries is not None: + call_kwargs.setdefault("num_retries", self.max_retries) + if self.timeout is not None: + call_kwargs.setdefault("timeout", self.timeout) + + if structured_model: + if tools or tool_choice: + logger.warning( + "structured_model is provided. Both 'tools' and 'tool_choice' parameters will be " + "overridden and ignored. The model will only perform structured output generation.", + ) + call_kwargs.pop("tools", None) + call_kwargs.pop("tool_choice", None) + + # Some providers reject a full JSON schema; use a simple json_object. + if any(name in self.model.lower() for name in ("qwen", "gemini", "pai-judge")): + logger.info( + f"Model '{self.model}' detected: switching to 'json_object' response_format for compatibility" + ) + call_kwargs["response_format"] = {"type": "json_object"} + else: + call_kwargs["response_format"] = structured_model + else: + if tools: + call_kwargs["tools"] = tools + if tool_choice: + self._validate_tool_choice(tool_choice, tools) + call_kwargs["tool_choice"] = tool_choice + + response = await litellm.acompletion(**call_kwargs) + + if self.stream: + return self._handle_streaming_response(response, structured_model, callback) + return self._handle_non_streaming_response(response, structured_model, callback) diff --git a/pyproject.toml b/pyproject.toml index 9dc735276..af3353bf0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -80,6 +80,9 @@ verl = [ "transformers>=4.52.4,<5.0.0", "verl" ] +litellm = [ + "litellm>=1.89.0,<2.0.0" +] [project.urls] Homepage = "https://github.com/agentscope-ai/OpenJudge" diff --git a/tests/models/test_litellm_chat_model.py b/tests/models/test_litellm_chat_model.py new file mode 100644 index 000000000..36407df63 --- /dev/null +++ b/tests/models/test_litellm_chat_model.py @@ -0,0 +1,163 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +"""Unit tests for LiteLLMChatModel. + +LiteLLM is an optional gateway backend that reuses OpenAIChatModel's response +handling while routing generation through ``litellm.acompletion``. These tests +inject a fake ``litellm`` module so the suite runs without the real dependency. +""" + +import sys +import types + +import pytest + +from openjudge.models import LiteLLMChatModel +from openjudge.models.openai_chat_model import OpenAIChatModel + + +def _install_litellm_stub(content="LITELLM_OK"): + """Register a fake ``litellm`` module; return (module, calls list).""" + fake = types.ModuleType("litellm") + calls = [] + + class _Message: + def __init__(self): + self.role = "assistant" + self.content = content + + def model_dump(self): + return {"role": self.role, "content": self.content} + + class _Choice: + def __init__(self): + self.message = _Message() + + class _Response: + def __init__(self): + self.choices = [_Choice()] + + async def _acompletion(**kwargs): + calls.append(kwargs) + return _Response() + + fake.acompletion = _acompletion + sys.modules["litellm"] = fake + return fake, calls + + +@pytest.mark.unit +class TestLiteLLMChatModelInit: + """Verify constructor behavior.""" + + def test_inherits_from_openai_chat_model(self): + assert issubclass(LiteLLMChatModel, OpenAIChatModel) + + def test_exported_from_models_package(self): + from openjudge.models import LiteLLMChatModel as Imported + + assert Imported is LiteLLMChatModel + + def test_no_persistent_client_built(self): + # LiteLLM routes per call; unlike OpenAIChatModel it builds no client. + model = LiteLLMChatModel(model="gpt-4o") + assert not hasattr(model, "client") + + def test_defaults(self): + model = LiteLLMChatModel(model="anthropic/claude-sonnet-4-6") + assert model.model == "anthropic/claude-sonnet-4-6" + assert model.stream is False + assert model.drop_params is True + assert model.api_key is None + assert model.base_url is None + + def test_drop_params_opt_out(self): + model = LiteLLMChatModel(model="gpt-4o", drop_params=False) + assert model.drop_params is False + + +@pytest.mark.unit +class TestLiteLLMChatModelAchat: + """Verify achat dispatches to litellm.acompletion correctly.""" + + async def test_dispatch_sets_drop_params_and_omits_blank_creds(self): + _, calls = _install_litellm_stub() + model = LiteLLMChatModel(model="anthropic/claude-sonnet-4-6") + resp = await model.achat( + messages=[{"role": "user", "content": "hi"}], + temperature=0.2, + ) + assert resp.content == "LITELLM_OK" + call = calls[-1] + assert call["model"] == "anthropic/claude-sonnet-4-6" + assert call["drop_params"] is True + assert call["temperature"] == 0.2 + assert call["stream"] is False + # creds omitted when unset so litellm uses each provider's own env var + assert "api_key" not in call + assert "api_base" not in call + + async def test_forwards_proxy_creds_when_set(self): + _, calls = _install_litellm_stub() + model = LiteLLMChatModel( + model="gpt-4o", + api_key="sk-proxy", + base_url="http://localhost:4000/v1", + max_retries=3, + timeout=30.0, + ) + await model.achat(messages=[{"role": "user", "content": "hi"}]) + call = calls[-1] + assert call["api_key"] == "sk-proxy" + assert call["api_base"] == "http://localhost:4000/v1" + assert call["num_retries"] == 3 + assert call["timeout"] == 30.0 + + async def test_structured_model_sets_response_format_and_drops_tools(self): + from pydantic import BaseModel + + class Verdict(BaseModel): + score: int + + _, calls = _install_litellm_stub(content='{"score": 5}') + model = LiteLLMChatModel(model="gpt-4o") + await model.achat( + messages=[{"role": "user", "content": "rate this"}], + tools=[{"type": "function", "function": {"name": "x"}}], + tool_choice="auto", + structured_model=Verdict, + ) + call = calls[-1] + assert call["response_format"] is Verdict + assert "tools" not in call + assert "tool_choice" not in call + + async def test_gemini_structured_uses_json_object(self): + from pydantic import BaseModel + + class Verdict(BaseModel): + score: int + + _, calls = _install_litellm_stub(content='{"score": 5}') + model = LiteLLMChatModel(model="gemini/gemini-2.5-flash") + await model.achat( + messages=[{"role": "user", "content": "rate"}], + structured_model=Verdict, + ) + assert calls[-1]["response_format"] == {"type": "json_object"} + + async def test_invalid_tool_choice_raises(self): + _install_litellm_stub() + model = LiteLLMChatModel(model="gpt-4o") + with pytest.raises(ValueError): + await model.achat( + messages=[{"role": "user", "content": "hi"}], + tool_choice="no_such_function", + tools=[{"type": "function", "function": {"name": "real_fn"}}], + ) + + async def test_non_list_messages_raises(self): + _install_litellm_stub() + model = LiteLLMChatModel(model="gpt-4o") + with pytest.raises(ValueError): + await model.achat(messages="not a list")