feat: add vLLM classify SDK endpoint

litellm/__init__.py

@@ -1420,7 +1420,19 @@ global_disable_no_log_param: bool = False
 from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key
 
 ### PASSTHROUGH ###
-from .passthrough import allm_passthrough_route, llm_passthrough_route
+from .passthrough import (
+    aclassify,
+    allm_passthrough_route,
+    classify,
+    llm_passthrough_route,
+)
+from .types.classification import (
+    ClassificationData,
+    ClassificationInput,
+    ClassificationRequest,
+    ClassificationResponse,
+    ClassificationUsage,
+)
 from .google_genai import agenerate_content
 
 ### GLOBAL CONFIG ###

litellm/llms/custom_httpx/http_handler.py

@@ -515,14 +515,12 @@ class AsyncHTTPHandler:
         client_alias: Optional[str] = None,  # name for client in logs
         ssl_verify: Optional[VerifyTypes] = None,
         shared_session: Optional["ClientSession"] = None,
+        client: Optional[httpx.AsyncClient] = None,
     ):
         self.timeout = timeout
         self.event_hooks = event_hooks
-        self.client = self.create_client(
-            timeout=timeout,
-            event_hooks=event_hooks,
-            ssl_verify=ssl_verify,
-            shared_session=shared_session,
+        self.client = client or self.create_client(
+            timeout=timeout, event_hooks=event_hooks, ssl_verify=ssl_verify, shared_session=shared_session
         )
         self.client_alias = client_alias
 

litellm/passthrough/__init__.py

@@ -1,8 +1,11 @@
+from .classify import aclassify, classify
 from .main import allm_passthrough_route, llm_passthrough_route
 from .utils import BasePassthroughUtils
 
 __all__ = [
+    "aclassify",
     "allm_passthrough_route",
+    "classify",
     "llm_passthrough_route",
     "BasePassthroughUtils",
 ]

litellm/passthrough/classify.py

@@ -0,0 +1,132 @@
+from collections.abc import Mapping
+from typing import Literal
+
+import httpx
+
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
+from litellm.passthrough.main import (
+    aclassification_passthrough_request,
+    classification_passthrough_request,
+)
+from litellm.types.classification import (
+    ClassificationInput,
+    ClassificationRequest,
+    ClassificationResponse,
+)
+
+
+def _classification_request_body(
+    *,
+    model: str,
+    input: ClassificationInput,
+    user: str | None,
+    truncate_prompt_tokens: int | None,
+    truncation_side: Literal["left", "right"] | None,
+    priority: int,
+    add_special_tokens: bool,
+    request_id: str | None,
+    use_activation: bool | None,
+    extra_body: Mapping[str, object] | None,
+) -> dict[str, object]:
+    request = ClassificationRequest(
+        model=model,
+        input=input,
+        user=user,
+        truncate_prompt_tokens=truncate_prompt_tokens,
+        truncation_side=truncation_side,
+        priority=priority,
+        add_special_tokens=add_special_tokens,
+        request_id=request_id,
+        use_activation=use_activation,
+    ).model_dump(exclude_none=True)
+    return {**request, **dict(extra_body or {})}
+
+
+def _classification_response(response: httpx.Response) -> ClassificationResponse:
+    return ClassificationResponse.model_validate(response.json())
+
+
+def classify(
+    *,
+    model: str,
+    input: ClassificationInput,
+    custom_llm_provider: str | None = None,
+    api_base: str | None = None,
+    api_key: str | None = None,
+    user: str | None = None,
+    truncate_prompt_tokens: int | None = None,
+    truncation_side: Literal["left", "right"] | None = None,
+    priority: int = 0,
+    add_special_tokens: bool = True,
+    request_id: str | None = None,
+    use_activation: bool | None = None,
+    extra_body: Mapping[str, object] | None = None,
+    request_headers: dict[str, str] | None = None,
+    client: HTTPHandler | None = None,
+    timeout: float | httpx.Timeout | None = None,
+) -> ClassificationResponse:
+    response = classification_passthrough_request(
+        model=model,
+        custom_llm_provider=custom_llm_provider,
+        api_base=api_base,
+        api_key=api_key,
+        request_headers=request_headers,
+        json=_classification_request_body(
+            model=model,
+            input=input,
+            user=user,
+            truncate_prompt_tokens=truncate_prompt_tokens,
+            truncation_side=truncation_side,
+            priority=priority,
+            add_special_tokens=add_special_tokens,
+            request_id=request_id,
+            use_activation=use_activation,
+            extra_body=extra_body,
+        ),
+        client=client,
+        timeout=timeout,
+    )
+    return _classification_response(response)
+
+
+async def aclassify(
+    *,
+    model: str,
+    input: ClassificationInput,
+    custom_llm_provider: str | None = None,
+    api_base: str | None = None,
+    api_key: str | None = None,
+    user: str | None = None,
+    truncate_prompt_tokens: int | None = None,
+    truncation_side: Literal["left", "right"] | None = None,
+    priority: int = 0,
+    add_special_tokens: bool = True,
+    request_id: str | None = None,
+    use_activation: bool | None = None,
+    extra_body: Mapping[str, object] | None = None,
+    request_headers: dict[str, str] | None = None,
+    client: AsyncHTTPHandler | None = None,
+    timeout: float | httpx.Timeout | None = None,
+) -> ClassificationResponse:
+    response = await aclassification_passthrough_request(
+        model=model,
+        custom_llm_provider=custom_llm_provider,
+        api_base=api_base,
+        api_key=api_key,
+        request_headers=request_headers,
+        json=_classification_request_body(
+            model=model,
+            input=input,
+            user=user,
+            truncate_prompt_tokens=truncate_prompt_tokens,
+            truncation_side=truncation_side,
+            priority=priority,
+            add_special_tokens=add_special_tokens,
+            request_id=request_id,
+            use_activation=use_activation,
+            extra_body=extra_body,
+        ),
+        client=client,
+        timeout=timeout,
+    )
+    return _classification_response(response)

litellm/passthrough/main.py

@@ -361,6 +361,60 @@ def llm_passthrough_route(
         )
 
 
+def classification_passthrough_request(
+    *,
+    model: str,
+    custom_llm_provider: str | None,
+    api_base: str | None,
+    api_key: str | None,
+    request_headers: dict[str, str] | None,
+    json: dict[str, object],
+    client: HTTPHandler | None,
+    timeout: float | httpx.Timeout | None,
+) -> httpx.Response:
+    response = llm_passthrough_route(
+        method="POST",
+        endpoint="classify",
+        model=model,
+        custom_llm_provider=custom_llm_provider,
+        api_base=api_base,
+        api_key=api_key,
+        request_headers=request_headers,
+        json=json,
+        client=client,
+        timeout=timeout,
+    )
+    assert isinstance(response, httpx.Response)
+    return response
+
+
+async def aclassification_passthrough_request(
+    *,
+    model: str,
+    custom_llm_provider: str | None,
+    api_base: str | None,
+    api_key: str | None,
+    request_headers: dict[str, str] | None,
+    json: dict[str, object],
+    client: AsyncHTTPHandler | None,
+    timeout: float | httpx.Timeout | None,
+) -> httpx.Response:
+    response = await allm_passthrough_route(
+        method="POST",
+        endpoint="classify",
+        model=model,
+        custom_llm_provider=custom_llm_provider,
+        api_base=api_base,
+        api_key=api_key,
+        request_headers=request_headers,
+        json=json,
+        client=client,
+        timeout=timeout,
+    )
+    assert isinstance(response, httpx.Response)
+    return response
+
+
 async def _async_passthrough_request(
     client: Union[HTTPHandler, AsyncHTTPHandler],
     request: httpx.Request,

litellm/types/classification.py

@@ -0,0 +1,46 @@
+from typing import Literal, TypeAlias
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+ClassificationInput: TypeAlias = str | list[str] | list[int] | list[list[int]]
+
+
+class ClassificationRequest(BaseModel):
+    model: str
+    input: ClassificationInput
+    user: str | None = None
+    truncate_prompt_tokens: int | None = Field(default=None, ge=-1)
+    truncation_side: Literal["left", "right"] | None = None
+    priority: int = 0
+    add_special_tokens: bool = True
+    request_id: str | None = None
+    use_activation: bool | None = None
+
+
+class ClassificationData(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    index: int
+    label: str | None
+    probs: list[float]
+    num_classes: int
+
+
+class ClassificationUsage(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    prompt_tokens: int = 0
+    total_tokens: int = 0
+    completion_tokens: int | None = 0
+
+
+class ClassificationResponse(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+    id: str
+    object: str = "list"
+    created: int
+    model: str
+    data: list[ClassificationData]
+    usage: ClassificationUsage

tests/test_litellm/passthrough/test_classify.py

@@ -0,0 +1,84 @@
+import json
+
+import httpx
+import pytest
+
+import litellm
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
+from litellm.types.classification import ClassificationResponse
+
+
+CLASSIFICATION_RESPONSE = {
+    "id": "classify-test",
+    "object": "list",
+    "created": 1,
+    "model": "classifier",
+    "data": [
+        {
+            "index": 0,
+            "label": "positive",
+            "probs": [0.1, 0.9],
+            "num_classes": 2,
+        }
+    ],
+    "usage": {
+        "prompt_tokens": 4,
+        "total_tokens": 4,
+    },
+}
+
+
+def test_classify_calls_vllm_classify_and_parses_response() -> None:
+    def transport(request: httpx.Request) -> httpx.Response:
+        assert request.url == httpx.URL("http://localhost:8000/classify")
+        assert request.method == "POST"
+        assert json.loads(request.read()) == {
+            "model": "classifier",
+            "input": "LiteLLM is useful",
+            "truncate_prompt_tokens": 128,
+            "priority": 0,
+            "add_special_tokens": True,
+            "use_activation": False,
+            "cache_salt": "request-salt",
+        }
+        return httpx.Response(status_code=200, json=CLASSIFICATION_RESPONSE)
+
+    response = litellm.classify(
+        model="hosted_vllm/classifier",
+        input="LiteLLM is useful",
+        api_base="http://localhost:8000",
+        truncate_prompt_tokens=128,
+        use_activation=False,
+        extra_body={"cache_salt": "request-salt"},
+        client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(transport))),
+    )
+
+    assert isinstance(response, ClassificationResponse)
+    assert response.data[0].label == "positive"
+    assert response.data[0].probs == [0.1, 0.9]
+    assert response.usage.total_tokens == 4
+
+
+@pytest.mark.asyncio
+async def test_aclassify_accepts_token_ids() -> None:
+    def transport(request: httpx.Request) -> httpx.Response:
+        assert request.url == httpx.URL("http://localhost:8000/classify")
+        assert json.loads(request.read()) == {
+            "model": "classifier",
+            "input": [[101, 2023, 102]],
+            "priority": 0,
+            "add_special_tokens": False,
+        }
+        return httpx.Response(status_code=200, json=CLASSIFICATION_RESPONSE)
+
+    async with httpx.AsyncClient(transport=httpx.MockTransport(transport)) as async_client:
+        response = await litellm.aclassify(
+            model="vllm/classifier",
+            input=[[101, 2023, 102]],
+            api_base="http://localhost:8000",
+            add_special_tokens=False,
+            client=AsyncHTTPHandler(client=async_client),
+        )
+
+    assert isinstance(response, ClassificationResponse)
+    assert response.id == "classify-test"