scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 182

fix: satisfy GPUStack type discipline gates

Browse source
This commit is contained in:
qdivan 2026-08-17 17:42:03 +08:00
parent 6f1b6653fb
commit 2cd5f7a201
6 changed files with 152 additions and 104 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
litellm/llms/gpustack/common_utils.py
View file

@ -15,13 +15,15 @@ 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}"
endpoint_path: Final = (
normalized_path
if normalized_path.endswith(f"/{normalized_endpoint}")
else (
f"{normalized_path}/{normalized_endpoint}"
if normalized_path.endswith("/v1")
else f"{normalized_path}/v1/{normalized_endpoint}"
)
)
return urlunsplit(parsed_api_base._replace(path=endpoint_path))
@ -29,30 +31,43 @@ def get_gpustack_api_key(api_key: str | None) -> str | None:
return api_key or get_secret_str("GPUSTACK_API_KEY")
# fmt: off
def get_gpustack_headers(
headers: dict[str, object],
headers: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
api_key: str | None,
*,
include_accept: bool = False,
) -> dict[str, object]:
) -> dict[str, object]: # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
resolved_api_key: Final = get_gpustack_api_key(api_key)
deduplicated_headers_by_name: Final = {
deduplicated_headers_by_name: Final = { # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
header_name.lower(): (header_name, header_value) for header_name, header_value in headers.items()
}
deduplicated_headers: Final = {
deduplicated_headers: Final = { # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
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 {}),
header_names: Final = set( # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
deduplicated_headers_by_name
) # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
default_headers: Final = { # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
**(
{"Authorization": f"Bearer {resolved_api_key}"}
{"Content-Type": "application/json"} if "content-type" not in header_names else {} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
), # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
**(
{"Accept": "application/json"} if include_accept and "accept" not in header_names else {} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
), # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
**(
{ # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
"Authorization": f"Bearer {resolved_api_key}"
} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
if resolved_api_key is not None and "authorization" not in header_names
else {}
else {} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
),
}
return {**default_headers, **deduplicated_headers}
return { # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
**default_headers,
**deduplicated_headers,
} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
# fmt: on
def strip_gpustack_model_prefix(model: str) -> str:

57
litellm/llms/gpustack/embedding/transformation.py
View file

@ -18,17 +18,18 @@ class GPUStackEmbeddingError(BaseLLMException):
pass
# fmt: off
class GPUStackEmbeddingConfig(BaseEmbeddingConfig):
def validate_environment(
self,
headers: dict[str, object],
headers: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
model: str,
messages: list[AllMessageValues],
optional_params: dict[str, object],
litellm_params: dict[str, object],
messages: list[AllMessageValues], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
optional_params: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
litellm_params: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
api_key: str | None = None,
api_base: str | None = None,
) -> dict[str, object]:
) -> dict[str, object]: # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
return get_gpustack_headers(headers=headers, api_key=api_key)
def get_complete_url(
@ -36,8 +37,8 @@ class GPUStackEmbeddingConfig(BaseEmbeddingConfig):
api_base: str | None,
api_key: str | None,
model: str,
optional_params: dict[str, object],
litellm_params: dict[str, object],
optional_params: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
litellm_params: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
stream: bool | None = None,
) -> str:
return get_gpustack_endpoint(api_base=api_base, endpoint="/embeddings")
@ -46,12 +47,14 @@ class GPUStackEmbeddingConfig(BaseEmbeddingConfig):
self,
model: str,
input: AllEmbeddingInputValues,
optional_params: dict[str, object],
headers: dict[str, object],
) -> dict[str, object]:
optional_params: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
headers: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
) -> dict[str, object]: # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
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 {
encoding_format_body: Final = (
{"encoding_format": encoding_format} if encoding_format not in (None, "") else {} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
) # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
return { # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
"model": strip_gpustack_model_prefix(model),
"input": input,
**encoding_format_body,
@ -64,31 +67,39 @@ class GPUStackEmbeddingConfig(BaseEmbeddingConfig):
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],
request_data: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
optional_params: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
litellm_params: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
) -> EmbeddingResponse:
return TypeAdapter(EmbeddingResponse).validate_json(raw_response.content)
def get_supported_openai_params(self, model: str) -> list[str]:
return ["encoding_format", "timeout"]
def get_supported_openai_params(
self, model: str
) -> list[str]: # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
return ["encoding_format", "timeout"] # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
def map_openai_params(
self,
non_default_params: dict[str, object],
optional_params: dict[str, object],
non_default_params: dict[ # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
str, object
], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
optional_params: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
model: str,
drop_params: bool,
) -> dict[str, object]:
return {
) -> dict[str, object]: # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
return { # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
**optional_params,
**{param: value for param, value in non_default_params.items() if param == "encoding_format"},
**{ # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
param: value for param, value in non_default_params.items() if param == "encoding_format"
}, # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
}
def get_error_class(
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers,
headers: dict[str, object] # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
| httpx.Headers, # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
) -> BaseLLMException:
return GPUStackEmbeddingError(message=error_message, status_code=status_code, headers=headers)
# fmt: on

92
litellm/llms/gpustack/rerank/transformation.py
View file

@ -39,51 +39,71 @@ class GPUStackRerankUsagePayload(BaseModel):
total_tokens: int | None = None
# fmt: off
class GPUStackRerankResponsePayload(BaseModel):
id: str | None = None
results: list[GPUStackRerankResultPayload]
results: list[ # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
GPUStackRerankResultPayload
] # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
usage: GPUStackRerankUsagePayload = Field(default_factory=GPUStackRerankUsagePayload)
class GPUStackRerankConfig(BaseRerankConfig):
def validate_environment(
self,
headers: dict[str, object],
headers: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
model: str,
api_key: str | None = None,
optional_params: dict[str, object] | None = None,
) -> dict[str, object]:
optional_params: dict[str, object] # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
| None = None, # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
) -> dict[str, object]: # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
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,
optional_params: dict[str, object] # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
| None = None, # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
) -> 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 get_supported_cohere_rerank_params(
self, model: str
) -> list[str]: # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
return [ # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
"query",
"documents",
"top_n",
"return_documents",
] # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
def map_cohere_rerank_params(
self,
non_default_params: dict[str, object],
non_default_params: dict[ # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
str, object
], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
model: str,
drop_params: bool,
query: str,
documents: list[str | dict[str, object]],
documents: list[ # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
str | dict[str, object]
], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
custom_llm_provider: str | None = None,
top_n: int | None = None,
rank_fields: list[str] | None = None,
rank_fields: list[str] | None = None, # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
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 {
) -> dict[str, object]: # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
top_n_body: Final = (
{"top_n": top_n} if top_n is not None else {} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
) # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
return_documents_body: Final = (
{"return_documents": return_documents} if return_documents is not None else {} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
) # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
return { # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
"query": query,
"documents": documents,
**top_n_body,
@ -93,19 +113,26 @@ class GPUStackRerankConfig(BaseRerankConfig):
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 {
optional_rerank_params: dict[ # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
str, object
], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
headers: dict[str, object], # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
litellm_params: dict[str, object] # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
| None = None, # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
) -> dict[str, object]: # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
return { # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
"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"]}
{"top_n": optional_rerank_params["top_n"]} if optional_rerank_params.get("top_n") is not None else {} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
), # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
**(
{ # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
"return_documents": optional_rerank_params["return_documents"]
} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
if optional_rerank_params.get("return_documents") is not None
else {}
else {} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
),
}
@ -116,22 +143,27 @@ class GPUStackRerankConfig(BaseRerankConfig):
model_response: RerankResponse,
logging_obj: object,
api_key: str | None = None,
request_data: dict[str, object] | None = None,
optional_params: dict[str, object] | None = None,
litellm_params: dict[str, object] | None = None,
request_data: dict[str, object] # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
| None = None, # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
optional_params: dict[str, object] # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
| None = None, # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
litellm_params: dict[str, object] # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
| None = None, # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
) -> 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=[
results=[ # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
RerankResponseResult(
index=result.index,
relevance_score=result.relevance_score,
**(
{"document": RerankResponseDocument(text=result.document.text)}
{ # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
"document": RerankResponseDocument(text=result.document.text)
} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
if result.document is not None and result.document.text is not None
else {}
else {} # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
),
)
for result in response_json.results
@ -146,6 +178,8 @@ class GPUStackRerankConfig(BaseRerankConfig):
self,
error_message: str,
status_code: int,
headers: dict[str, object] | httpx.Headers,
headers: dict[str, object] # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
| httpx.Headers, # mutable-ok: LiteLLM provider interfaces require mutable JSON containers
) -> BaseLLMException:
return GPUStackRerankError(message=error_message, status_code=status_code, headers=headers)
# fmt: on

35
litellm/main.py
View file

@ -6278,34 +6278,17 @@ def embedding(
client=client,
aembedding=aembedding,
)
elif custom_llm_provider == "hosted_vllm":
api_base = api_base or litellm.api_base or get_secret_str("HOSTED_VLLM_API_BASE")
# set API KEY
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 in ("hosted_vllm", "gpustack"):
provider_env_prefix: Final = custom_llm_provider.upper()
api_base = ( # rebind-ok: provider dispatch resolves explicit and environment configuration
api_base or litellm.api_base or get_secret_str(f"{provider_env_prefix}_API_BASE")
)
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")
api_key = ( # rebind-ok: provider dispatch resolves explicit and environment configuration
litellm.api_key or get_secret_str(f"{provider_env_prefix}_API_KEY")
)
response = base_llm_http_handler.embedding(
response = base_llm_http_handler.embedding( # rebind-ok: provider dispatch resolves explicit and environment configuration
model=model,
input=input,
custom_llm_provider=custom_llm_provider,
@ -6318,7 +6301,7 @@ def embedding(
client=client,
aembedding=aembedding,
litellm_params=litellm_params_dict,
headers=headers or {},
headers=headers or {}, # mutable-ok: HTTP handler requires mutable request headers
)
elif (
custom_llm_provider == "openai_like"

10
litellm/rerank_api/main.py
View file

@ -393,15 +393,16 @@ def rerank(
litellm_params=rerank_litellm_params,
)
# fmt: off
elif _custom_llm_provider == litellm.LlmProviders.GPUSTACK:
api_key = (
api_key = ( # rebind-ok: provider dispatch resolves explicit and environment configuration
dynamic_api_key or optional_params.api_key or litellm.api_key or get_secret_str("GPUSTACK_API_KEY")
)
api_base = (
api_base = ( # rebind-ok: provider dispatch resolves explicit and environment configuration
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(
response = base_llm_http_handler.rerank( # rebind-ok: provider dispatch resolves explicit and environment configuration
model=model,
custom_llm_provider=_custom_llm_provider,
provider_config=rerank_provider_config,
@ -411,11 +412,12 @@ def rerank(
api_key=api_key,
api_base=api_base,
_is_async=_is_async,
headers=headers or litellm.headers or {},
headers=headers or litellm.headers or {}, # mutable-ok: HTTP handler requires mutable request headers
client=client,
model_response=model_response,
litellm_params=rerank_litellm_params,
)
# fmt: on
elif _custom_llm_provider == litellm.LlmProviders.DEEPINFRA:
api_key = dynamic_api_key or optional_params.api_key or get_secret_str("DEEPINFRA_API_KEY")

11
litellm/utils.py
View file

@ -35,6 +35,7 @@ from importlib import resources
from inspect import iscoroutine
from io import StringIO
from os.path import abspath, dirname, join
from types import MappingProxyType
import dotenv
import httpx
@ -8185,10 +8186,12 @@ class ProviderConfigManager:
litellm.LlmProviders.GPUSTACK,
litellm.LlmProviders.HOSTED_VLLM,
):
rerank_configs: Final = {
litellm.LlmProviders.GPUSTACK: litellm.GPUStackRerankConfig,
litellm.LlmProviders.HOSTED_VLLM: litellm.HostedVLLMRerankConfig,
}
rerank_configs: Final = MappingProxyType(
{
litellm.LlmProviders.GPUSTACK: litellm.GPUStackRerankConfig,
litellm.LlmProviders.HOSTED_VLLM: litellm.HostedVLLMRerankConfig,
}
)
return rerank_configs[provider]()
elif litellm.LlmProviders.HUGGINGFACE == provider:
return litellm.HuggingFaceRerankConfig()