diff --git a/litellm/__init__.py b/litellm/__init__.py index 6e2a03b7c7c..bf0f9be634a 100644 --- a/litellm/__init__.py +++ b/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 ### diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 5cec763bb5d..fc6e6b1fcbf 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/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 diff --git a/litellm/passthrough/__init__.py b/litellm/passthrough/__init__.py index bfd13e7a74e..b1929bffdfb 100644 --- a/litellm/passthrough/__init__.py +++ b/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", ] diff --git a/litellm/passthrough/classify.py b/litellm/passthrough/classify.py new file mode 100644 index 00000000000..194cb4d6d1c --- /dev/null +++ b/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) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index cdeedd7b522..a3d238e78a3 100644 --- a/litellm/passthrough/main.py +++ b/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, diff --git a/litellm/types/classification.py b/litellm/types/classification.py new file mode 100644 index 00000000000..f93fb2e1de5 --- /dev/null +++ b/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 diff --git a/tests/test_litellm/passthrough/test_classify.py b/tests/test_litellm/passthrough/test_classify.py new file mode 100644 index 00000000000..eaebd9efd67 --- /dev/null +++ b/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"