feat: add LiteLLM as a first-class chat model provider

openjudge/models/__init__.py

@@ -4,13 +4,15 @@
 """
 
 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
 from openjudge.models.qwen_vl_model import QwenVLModel
 
 __all__ = [
     "BaseChatModel",
+    "LiteLLMChatModel",
     "MiniMaxChatModel",
     "OpenAIChatModel",
     "QiniuChatModel",

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)

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"

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")