diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 89c72acc06d..17f22030a0f 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -152,6 +152,7 @@ LLM_CONFIG_NAMES: Final = ( "InfinityRerankConfig", "JinaAIRerankConfig", "DeepinfraRerankConfig", + "GPUStackRerankConfig", "HostedVLLMRerankConfig", "NvidiaNimRerankConfig", "NvidiaNimRankingConfig", @@ -280,6 +281,7 @@ LLM_CONFIG_NAMES: Final = ( "AzureSpeechAudioTranscriptionConfig", "HostedVLLMChatConfig", "HostedVLLMEmbeddingConfig", + "GPUStackEmbeddingConfig", # Alias for backwards compatibility "VolcEngineConfig", # Alias for VolcEngineChatConfig "LlamafileChatConfig", @@ -673,6 +675,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.deepinfra.rerank.transformation", "DeepinfraRerankConfig", ), + "GPUStackRerankConfig": ( + ".llms.gpustack.rerank.transformation", + "GPUStackRerankConfig", + ), "HostedVLLMRerankConfig": ( ".llms.hosted_vllm.rerank.transformation", "HostedVLLMRerankConfig", @@ -1084,6 +1090,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.hosted_vllm.embedding.transformation", "HostedVLLMEmbeddingConfig", ), + "GPUStackEmbeddingConfig": ( + ".llms.gpustack.embedding.transformation", + "GPUStackEmbeddingConfig", + ), # Alias for backwards compatibility "VolcEngineConfig": ( ".llms.volcengine.chat.transformation", diff --git a/litellm/llms/gpustack/common_utils.py b/litellm/llms/gpustack/common_utils.py new file mode 100644 index 00000000000..f93c7977216 --- /dev/null +++ b/litellm/llms/gpustack/common_utils.py @@ -0,0 +1,61 @@ +from typing import Final +from urllib.parse import urlsplit, urlunsplit + +from litellm.secret_managers.main import get_secret_str + + +def get_gpustack_api_base(api_base: str | None) -> str: + resolved_api_base: Final = api_base or get_secret_str("GPUSTACK_API_BASE") + if resolved_api_base is None: + raise ValueError("api_base is required for GPUStack. Set it in the call or via GPUSTACK_API_BASE.") + return resolved_api_base + + +def get_gpustack_endpoint(api_base: str | None, endpoint: str) -> str: + parsed_api_base: Final = urlsplit(get_gpustack_api_base(api_base)) + normalized_path: Final = parsed_api_base.path.rstrip("/") + normalized_endpoint: Final = endpoint.strip("/") + endpoint_path: str + if normalized_path.endswith(f"/{normalized_endpoint}"): + endpoint_path = normalized_path + elif normalized_path.endswith("/v1"): + endpoint_path = f"{normalized_path}/{normalized_endpoint}" + else: + endpoint_path = f"{normalized_path}/v1/{normalized_endpoint}" + return urlunsplit(parsed_api_base._replace(path=endpoint_path)) + + +def get_gpustack_api_key(api_key: str | None) -> str | None: + return api_key or get_secret_str("GPUSTACK_API_KEY") + + +def get_gpustack_headers( + headers: dict[str, object], + api_key: str | None, + *, + include_accept: bool = False, +) -> dict[str, object]: + resolved_api_key: Final = get_gpustack_api_key(api_key) + deduplicated_headers_by_name: Final = { + header_name.lower(): (header_name, header_value) for header_name, header_value in headers.items() + } + deduplicated_headers: Final = { + header_name: header_value for header_name, header_value in deduplicated_headers_by_name.values() + } + header_names: Final = set(deduplicated_headers_by_name) + default_headers: Final = { + **({"Content-Type": "application/json"} if "content-type" not in header_names else {}), + **({"Accept": "application/json"} if include_accept and "accept" not in header_names else {}), + **( + {"Authorization": f"Bearer {resolved_api_key}"} + if resolved_api_key is not None and "authorization" not in header_names + else {} + ), + } + return {**default_headers, **deduplicated_headers} + + +def strip_gpustack_model_prefix(model: str) -> str: + if model.startswith("gpustack/"): + return model.replace("gpustack/", "", 1) + return model diff --git a/litellm/llms/gpustack/embedding/transformation.py b/litellm/llms/gpustack/embedding/transformation.py new file mode 100644 index 00000000000..294a92692fc --- /dev/null +++ b/litellm/llms/gpustack/embedding/transformation.py @@ -0,0 +1,94 @@ +from typing import Final + +import httpx +from pydantic import TypeAdapter + +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.embedding.transformation import BaseEmbeddingConfig +from litellm.llms.gpustack.common_utils import ( + get_gpustack_endpoint, + get_gpustack_headers, + strip_gpustack_model_prefix, +) +from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues +from litellm.types.utils import EmbeddingResponse + + +class GPUStackEmbeddingError(BaseLLMException): + pass + + +class GPUStackEmbeddingConfig(BaseEmbeddingConfig): + def validate_environment( + self, + headers: dict[str, object], + model: str, + messages: list[AllMessageValues], + optional_params: dict[str, object], + litellm_params: dict[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, object]: + return get_gpustack_headers(headers=headers, api_key=api_key) + + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: dict[str, object], + litellm_params: dict[str, object], + stream: bool | None = None, + ) -> str: + return get_gpustack_endpoint(api_base=api_base, endpoint="/embeddings") + + def transform_embedding_request( + self, + model: str, + input: AllEmbeddingInputValues, + optional_params: dict[str, object], + headers: dict[str, object], + ) -> dict[str, object]: + encoding_format: Final[object | None] = optional_params.get("encoding_format") + encoding_format_body: Final = {"encoding_format": encoding_format} if encoding_format not in (None, "") else {} + return { + "model": strip_gpustack_model_prefix(model), + "input": input, + **encoding_format_body, + } + + def transform_embedding_response( + self, + model: str, + raw_response: httpx.Response, + model_response: EmbeddingResponse, + logging_obj: object, + api_key: str | None, + request_data: dict[str, object], + optional_params: dict[str, object], + litellm_params: dict[str, object], + ) -> EmbeddingResponse: + return TypeAdapter(EmbeddingResponse).validate_json(raw_response.content) + + def get_supported_openai_params(self, model: str) -> list[str]: + return ["encoding_format", "timeout"] + + def map_openai_params( + self, + non_default_params: dict[str, object], + optional_params: dict[str, object], + model: str, + drop_params: bool, + ) -> dict[str, object]: + return { + **optional_params, + **{param: value for param, value in non_default_params.items() if param == "encoding_format"}, + } + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, + ) -> BaseLLMException: + return GPUStackEmbeddingError(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/gpustack/rerank/transformation.py b/litellm/llms/gpustack/rerank/transformation.py new file mode 100644 index 00000000000..56f79379278 --- /dev/null +++ b/litellm/llms/gpustack/rerank/transformation.py @@ -0,0 +1,151 @@ +from typing import Final + +import httpx +from pydantic import BaseModel, Field, TypeAdapter + +from litellm._uuid import uuid +from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig +from litellm.llms.gpustack.common_utils import ( + get_gpustack_endpoint, + get_gpustack_headers, + strip_gpustack_model_prefix, +) +from litellm.types.rerank import ( + RerankBilledUnits, + RerankResponse, + RerankResponseDocument, + RerankResponseMeta, + RerankResponseResult, + RerankTokens, +) + + +class GPUStackRerankError(BaseLLMException): + pass + + +class GPUStackRerankDocumentPayload(BaseModel): + text: str | None = None + + +class GPUStackRerankResultPayload(BaseModel): + index: int + relevance_score: float + document: GPUStackRerankDocumentPayload | None = None + + +class GPUStackRerankUsagePayload(BaseModel): + total_tokens: int | None = None + + +class GPUStackRerankResponsePayload(BaseModel): + id: str | None = None + results: list[GPUStackRerankResultPayload] + usage: GPUStackRerankUsagePayload = Field(default_factory=GPUStackRerankUsagePayload) + + +class GPUStackRerankConfig(BaseRerankConfig): + def validate_environment( + self, + headers: dict[str, object], + model: str, + api_key: str | None = None, + optional_params: dict[str, object] | None = None, + ) -> dict[str, object]: + return get_gpustack_headers(headers=headers, api_key=api_key, include_accept=True) + + def get_complete_url( + self, + api_base: str | None, + model: str, + optional_params: dict[str, object] | None = None, + ) -> str: + return get_gpustack_endpoint(api_base=api_base, endpoint="/rerank") + + def get_supported_cohere_rerank_params(self, model: str) -> list[str]: + return ["query", "documents", "top_n", "return_documents"] + + def map_cohere_rerank_params( + self, + non_default_params: dict[str, object], + model: str, + drop_params: bool, + query: str, + documents: list[str | dict[str, object]], + custom_llm_provider: str | None = None, + top_n: int | None = None, + rank_fields: list[str] | None = None, + return_documents: bool | None = True, + max_chunks_per_doc: int | None = None, + max_tokens_per_doc: int | None = None, + instruction: str | None = None, + ) -> dict[str, object]: + top_n_body: Final = {"top_n": top_n} if top_n is not None else {} + return_documents_body: Final = {"return_documents": return_documents} if return_documents is not None else {} + return { + "query": query, + "documents": documents, + **top_n_body, + **return_documents_body, + } + + def transform_rerank_request( + self, + model: str, + optional_rerank_params: dict[str, object], + headers: dict[str, object], + litellm_params: dict[str, object] | None = None, + ) -> dict[str, object]: + return { + "model": strip_gpustack_model_prefix(model), + "query": optional_rerank_params["query"], + "documents": optional_rerank_params["documents"], + **({"top_n": optional_rerank_params["top_n"]} if optional_rerank_params.get("top_n") is not None else {}), + **( + {"return_documents": optional_rerank_params["return_documents"]} + if optional_rerank_params.get("return_documents") is not None + else {} + ), + } + + def transform_rerank_response( + self, + model: str, + raw_response: httpx.Response, + model_response: RerankResponse, + logging_obj: object, + api_key: str | None = None, + request_data: dict[str, object] = {}, + optional_params: dict[str, object] = {}, + litellm_params: dict[str, object] = {}, + ) -> RerankResponse: + response_json: Final = TypeAdapter(GPUStackRerankResponsePayload).validate_json(raw_response.content) + total_tokens: Final = response_json.usage.total_tokens or 0 + return RerankResponse( + id=response_json.id or str(uuid.uuid4()), + results=[ + RerankResponseResult( + index=result.index, + relevance_score=result.relevance_score, + **( + {"document": RerankResponseDocument(text=result.document.text)} + if result.document is not None and result.document.text is not None + else {} + ), + ) + for result in response_json.results + ], + meta=RerankResponseMeta( + billed_units=RerankBilledUnits(total_tokens=total_tokens), + tokens=RerankTokens(input_tokens=total_tokens), + ), + ) + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: dict[str, object] | httpx.Headers, + ) -> BaseLLMException: + return GPUStackRerankError(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/main.py b/litellm/main.py index 2a8ed6c87b6..49359bb2aab 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6285,6 +6285,26 @@ def embedding( if api_key is None: api_key = litellm.api_key or get_secret_str("HOSTED_VLLM_API_KEY") + response = base_llm_http_handler.embedding( + model=model, + input=input, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + logging_obj=logging, + timeout=timeout, + model_response=EmbeddingResponse(), + optional_params=optional_params, + client=client, + aembedding=aembedding, + litellm_params=litellm_params_dict, + headers=headers or {}, + ) + elif custom_llm_provider == "gpustack": + api_base = api_base or litellm.api_base or get_secret_str("GPUSTACK_API_BASE") + if api_key is None: + api_key = litellm.api_key or get_secret_str("GPUSTACK_API_KEY") + response = base_llm_http_handler.embedding( model=model, input=input, diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index 15a6f18a6bb..0b77bf68a8a 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -85,6 +85,7 @@ def rerank( "azure_ai", "infinity", "litellm_proxy", + "gpustack", "hosted_vllm", "deepinfra", "fireworks_ai", @@ -392,6 +393,30 @@ def rerank( litellm_params=rerank_litellm_params, ) + elif _custom_llm_provider == litellm.LlmProviders.GPUSTACK: + api_key = ( + dynamic_api_key or optional_params.api_key or litellm.api_key or get_secret_str("GPUSTACK_API_KEY") + ) + api_base = ( + dynamic_api_base or optional_params.api_base or litellm.api_base or get_secret_str("GPUSTACK_API_BASE") + ) + + response = base_llm_http_handler.rerank( + model=model, + custom_llm_provider=_custom_llm_provider, + provider_config=rerank_provider_config, + optional_rerank_params=optional_rerank_params, + logging_obj=litellm_logging_obj, + timeout=optional_params.timeout, + api_key=api_key, + api_base=api_base, + _is_async=_is_async, + headers=headers or litellm.headers or {}, + client=client, + model_response=model_response, + litellm_params=rerank_litellm_params, + ) + elif _custom_llm_provider == litellm.LlmProviders.DEEPINFRA: api_key = dynamic_api_key or optional_params.api_key or get_secret_str("DEEPINFRA_API_KEY") diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 272fbabf807..8f29e746312 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3668,6 +3668,7 @@ class LlmProviders(str, Enum): DOCKER_MODEL_RUNNER = "docker_model_runner" CUSTOM = "custom" LITELLM_PROXY = "litellm_proxy" + GPUSTACK = "gpustack" HOSTED_VLLM = "hosted_vllm" TENCENT = "tencent" LLAMAFILE = "llamafile" diff --git a/litellm/utils.py b/litellm/utils.py index d91d3092624..24792a45756 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8135,6 +8135,8 @@ class ProviderConfigManager: return litellm.CometAPIEmbeddingConfig() elif litellm.LlmProviders.GITHUB_COPILOT == provider: return litellm.GithubCopilotEmbeddingConfig() + elif litellm.LlmProviders.GPUSTACK == provider: + return litellm.GPUStackEmbeddingConfig() elif litellm.LlmProviders.OPENROUTER == provider: from litellm.llms.openrouter.embedding.transformation import ( OpenrouterEmbeddingConfig, @@ -8179,6 +8181,8 @@ class ProviderConfigManager: return litellm.InfinityRerankConfig() elif litellm.LlmProviders.JINA_AI == provider: return litellm.JinaAIRerankConfig() + elif litellm.LlmProviders.GPUSTACK == provider: + return litellm.GPUStackRerankConfig() elif litellm.LlmProviders.HOSTED_VLLM == provider: return litellm.HostedVLLMRerankConfig() elif litellm.LlmProviders.HUGGINGFACE == provider: diff --git a/tests/test_litellm/llms/gpustack/test_gpustack_embedding_rerank.py b/tests/test_litellm/llms/gpustack/test_gpustack_embedding_rerank.py new file mode 100644 index 00000000000..b468ee6e622 --- /dev/null +++ b/tests/test_litellm/llms/gpustack/test_gpustack_embedding_rerank.py @@ -0,0 +1,252 @@ +import json +from typing import Final + +import httpx +import pytest + +import litellm +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.llms.gpustack.common_utils import get_gpustack_endpoint, get_gpustack_headers + + +def test_gpustack_provider_resolution_preserves_owner_route_id() -> None: + model, provider, _, _ = litellm.get_llm_provider(model="gpustack/owner/bge-m3") + + assert model == "owner/bge-m3" + assert provider == litellm.LlmProviders.GPUSTACK.value + + +@pytest.mark.parametrize( + ("api_base", "endpoint", "expected"), + [ + ( + "https://gpustack.test/v1/embeddings?tenant=a#fragment", + "/embeddings", + "https://gpustack.test/v1/embeddings?tenant=a#fragment", + ), + ( + "https://embeddings/v1?tenant=a", + "/embeddings", + "https://embeddings/v1/embeddings?tenant=a", + ), + ( + "https://rerank/prefix?tenant=a", + "/rerank", + "https://rerank/prefix/v1/rerank?tenant=a", + ), + ( + "https://gpustack.test/v1?tenant=a/", + "/embeddings", + "https://gpustack.test/v1/embeddings?tenant=a/", + ), + ( + "https://gpustack.test/v1#section/", + "/rerank", + "https://gpustack.test/v1/rerank#section/", + ), + ], +) +def test_gpustack_endpoint_normalization_uses_url_path( + api_base: str, + endpoint: str, + expected: str, +) -> None: + assert get_gpustack_endpoint(api_base, endpoint) == expected + + +def test_gpustack_headers_deduplicate_caller_header_casing() -> None: + headers = get_gpustack_headers( + { + "Authorization": "Bearer first", + "authorization": "Bearer second", + "CONTENT-TYPE": "application/custom", + }, + "generated-key", + include_accept=True, + ) + + assert [key for key in headers if key.lower() == "authorization"] == ["authorization"] + assert headers["authorization"] == "Bearer second" + assert [key for key in headers if key.lower() == "content-type"] == ["CONTENT-TYPE"] + + +def test_gpustack_embedding_posts_to_v1_embeddings_with_caller_authorization() -> None: + captured_requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "object": "list", + "model": "owner/bge-m3", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2, 0.3]}], + "usage": {"prompt_tokens": 3, "total_tokens": 3}, + }, + ) + + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + + response = litellm.embedding( + model="gpustack/owner/bge-m3", + input=["hello"], + api_base="https://gpustack.test", + api_key="generated-key", + encoding_format="float", + headers={"authorization": "Bearer caller-key"}, + client=client, + ) + + assert len(captured_requests) == 1 + request: Final = captured_requests[0] + assert str(request.url) == "https://gpustack.test/v1/embeddings" + assert request.headers["Authorization"] == "Bearer caller-key" + assert request.headers.get_list("authorization") == ["Bearer caller-key"] + assert json.loads(request.content) == { + "model": "owner/bge-m3", + "input": ["hello"], + "encoding_format": "float", + } + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + assert response.usage.total_tokens == 3 + + +@pytest.mark.asyncio() +async def test_gpustack_aembedding_uses_env_base_and_key_without_double_v1(monkeypatch: pytest.MonkeyPatch) -> None: + captured_requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "object": "list", + "model": "embeddings", + "data": [{"object": "embedding", "index": 0, "embedding": [0.4, 0.5]}], + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + }, + ) + + monkeypatch.setenv("GPUSTACK_API_BASE", "https://gpustack-env.test/v1") + monkeypatch.setenv("GPUSTACK_API_KEY", "env-key") + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + try: + response = await litellm.aembedding( + model="gpustack/embeddings", + input=["hello"], + client=client, + ) + await GLOBAL_LOGGING_WORKER.clear_queue() + await GLOBAL_LOGGING_WORKER.stop() + finally: + await client.close() + + assert len(captured_requests) == 1 + request: Final = captured_requests[0] + assert str(request.url) == "https://gpustack-env.test/v1/embeddings" + assert request.headers["Authorization"] == "Bearer env-key" + assert json.loads(request.content) == {"model": "embeddings", "input": ["hello"]} + assert response.data[0]["embedding"] == [0.4, 0.5] + + +def test_gpustack_rerank_posts_supported_fields_and_preserves_usage() -> None: + captured_requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "model": "owner/reranker", + "results": [ + { + "index": 1, + "document": {"text": "second"}, + "relevance_score": 0.98, + } + ], + "usage": {"total_tokens": 7}, + }, + ) + + client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(handler))) + + response = litellm.rerank( + model="gpustack/owner/reranker", + query="best doc", + documents=["first", "second"], + top_n=1, + return_documents=True, + rank_fields=["ignored"], + max_tokens_per_doc=128, + api_base="https://gpustack.test/v1/rerank", + api_key="generated-key", + headers={"authorization": "Bearer caller-key"}, + client=client, + ) + + assert len(captured_requests) == 1 + request: Final = captured_requests[0] + assert str(request.url) == "https://gpustack.test/v1/rerank" + assert request.headers["Authorization"] == "Bearer caller-key" + assert request.headers.get_list("authorization") == ["Bearer caller-key"] + assert json.loads(request.content) == { + "model": "owner/reranker", + "query": "best doc", + "documents": ["first", "second"], + "top_n": 1, + "return_documents": True, + } + assert response.results[0]["document"]["text"] == "second" + assert response.results[0]["relevance_score"] == 0.98 + assert response.meta["billed_units"]["total_tokens"] == 7 + assert response.meta["tokens"]["input_tokens"] == 7 + + +@pytest.mark.asyncio() +async def test_gpustack_arerank_uses_env_base_and_key(monkeypatch: pytest.MonkeyPatch) -> None: + captured_requests: list[httpx.Request] = [] + + async def handler(request: httpx.Request) -> httpx.Response: + captured_requests.append(request) + return httpx.Response( + 200, + json={ + "results": [{"index": 0, "document": {"text": "first"}, "relevance_score": 0.8}], + "usage": {"total_tokens": None}, + }, + ) + + monkeypatch.setenv("GPUSTACK_API_BASE", "https://gpustack.test/") + monkeypatch.setenv("GPUSTACK_API_KEY", "generated-key") + client: Final = AsyncHTTPHandler() + client.client = httpx.AsyncClient(transport=httpx.MockTransport(handler)) + + try: + response = await litellm.arerank( + model="gpustack/reranker", + query="best doc", + documents=["first", "second"], + return_documents=True, + client=client, + ) + await GLOBAL_LOGGING_WORKER.clear_queue() + await GLOBAL_LOGGING_WORKER.stop() + finally: + await client.close() + + assert len(captured_requests) == 1 + request: Final = captured_requests[0] + assert str(request.url) == "https://gpustack.test/v1/rerank" + assert request.headers["Authorization"] == "Bearer generated-key" + assert json.loads(request.content) == { + "model": "reranker", + "query": "best doc", + "documents": ["first", "second"], + "return_documents": True, + } + assert response.results[0]["index"] == 0 + assert response.meta["billed_units"]["total_tokens"] == 0