mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(voyage): add mutable-ok pragmas and replace type: ignore with pyright: ignore (LIT001, LIT009)
This commit is contained in:
fzowl
parent
bb6db24d58
commit
eddb2225dc
1 changed files with 43 additions and 41 deletions
|
|
@ -22,11 +22,14 @@ class VoyageError(BaseLLMException):
|
|||
self,
|
||||
status_code: int,
|
||||
message: str,
|
||||
headers: dict | httpx.Headers = {},
|
||||
headers: dict | httpx.Headers = {}, # mutable-ok: base class signature
|
||||
):
|
||||
self.status_code = status_code
|
||||
self.message = message
|
||||
self.request = httpx.Request(method="POST", url="https://api.voyageai.com/v1/contextualizedembeddings")
|
||||
self.request = httpx.Request(
|
||||
method="POST",
|
||||
url="https://api.voyageai.com/v1/contextualizedembeddings",
|
||||
)
|
||||
self.response = httpx.Response(status_code=status_code, request=self.request)
|
||||
super().__init__(
|
||||
status_code=status_code,
|
||||
|
|
@ -48,8 +51,8 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
|
|||
api_base: str | None,
|
||||
api_key: str | None,
|
||||
model: str,
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
optional_params: dict, # mutable-ok: base class signature
|
||||
litellm_params: dict, # mutable-ok: base class signature
|
||||
stream: bool | None = None,
|
||||
) -> str:
|
||||
if api_base:
|
||||
|
|
@ -58,16 +61,16 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
|
|||
return api_base
|
||||
return "https://api.voyageai.com/v1/contextualizedembeddings"
|
||||
|
||||
def get_supported_openai_params(self, model: str) -> list:
|
||||
def get_supported_openai_params(self, model: str) -> list: # mutable-ok: base class signature
|
||||
return ["encoding_format", "dimensions"]
|
||||
|
||||
def map_openai_params(
|
||||
self,
|
||||
non_default_params: dict,
|
||||
optional_params: dict,
|
||||
non_default_params: dict, # mutable-ok: base class signature
|
||||
optional_params: dict, # mutable-ok: base class signature
|
||||
model: str,
|
||||
drop_params: bool,
|
||||
) -> dict:
|
||||
) -> dict: # mutable-ok: base class signature
|
||||
"""
|
||||
Map OpenAI params to Voyage params
|
||||
|
||||
|
|
@ -81,14 +84,14 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
|
|||
|
||||
def validate_environment(
|
||||
self,
|
||||
headers: dict,
|
||||
headers: dict, # mutable-ok: base class signature
|
||||
model: str,
|
||||
messages: list[AllMessageValues],
|
||||
optional_params: dict,
|
||||
litellm_params: dict,
|
||||
messages: list[AllMessageValues], # mutable-ok: base class signature
|
||||
optional_params: dict, # mutable-ok: base class signature
|
||||
litellm_params: dict, # mutable-ok: base class signature
|
||||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> dict:
|
||||
) -> dict: # mutable-ok: base class signature
|
||||
if api_key is None:
|
||||
api_key = (
|
||||
get_secret_str("VOYAGE_API_KEY")
|
||||
|
|
@ -99,18 +102,15 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
|
|||
"Authorization": f"Bearer {api_key}",
|
||||
}
|
||||
|
||||
# Chunk size (in tokens) used when the API auto-chunks a flat ``list[str]``.
|
||||
# Matches the voyage-context-4 context window so each string stays a single
|
||||
# chunk instead of being split.
|
||||
AUTO_CHUNK_SIZE = 32000
|
||||
|
||||
def transform_embedding_request(
|
||||
self,
|
||||
model: str,
|
||||
input: AllEmbeddingInputValues | list[list[str]],
|
||||
optional_params: dict,
|
||||
headers: dict,
|
||||
) -> dict:
|
||||
input: AllEmbeddingInputValues | list[list[str]], # mutable-ok: base class signature
|
||||
optional_params: dict, # mutable-ok: base class signature
|
||||
headers: dict, # mutable-ok: base class signature
|
||||
) -> dict: # mutable-ok: base class signature
|
||||
inputs, extra_params = self._transform_contextual_inputs(input, optional_params)
|
||||
return {
|
||||
"inputs": inputs,
|
||||
|
|
@ -122,9 +122,9 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
|
|||
@classmethod
|
||||
def _transform_contextual_inputs(
|
||||
cls,
|
||||
input: AllEmbeddingInputValues | list[list[str]],
|
||||
optional_params: dict,
|
||||
) -> tuple[list[str] | list[list[str]], dict]:
|
||||
input: AllEmbeddingInputValues | list[list[str]], # mutable-ok: union with AllEmbeddingInputValues
|
||||
optional_params: dict, # mutable-ok: matches public API
|
||||
) -> tuple[list[str] | list[list[str]], dict]: # mutable-ok: returned to caller who owns it
|
||||
"""
|
||||
Normalize ``input`` for Voyage's contextualized embeddings API and
|
||||
return ``(inputs, extra_params)`` where ``extra_params`` carries any
|
||||
|
|
@ -149,32 +149,27 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
|
|||
|
||||
Reference: https://docs.voyageai.com/reference/contextualized-embeddings-api
|
||||
"""
|
||||
# Single string -> a one-element flat list, auto-chunked.
|
||||
if isinstance(input, str):
|
||||
if optional_params.get("input_type") == "query":
|
||||
return [input], {}
|
||||
return [input], cls._auto_chunk_params(optional_params)
|
||||
return [input], {} # mutable-ok: fresh list returned to caller
|
||||
return [input], cls._auto_chunk_params(optional_params) # mutable-ok: fresh list returned to caller
|
||||
|
||||
# Flat list[str].
|
||||
if isinstance(input, list) and all(isinstance(i, str) for i in input):
|
||||
if optional_params.get("input_type") == "query":
|
||||
# The API accepts a flat query list as-is.
|
||||
return input, {} # type: ignore[return-value]
|
||||
# Otherwise let the API auto-chunk the flat list.
|
||||
return input, cls._auto_chunk_params(optional_params) # type: ignore[return-value]
|
||||
return input, {} # pyright: ignore[reportReturnType] # narrowed to list[str] by isinstance+all check
|
||||
return input, cls._auto_chunk_params(optional_params) # pyright: ignore[reportReturnType] # narrowed to list[str]
|
||||
|
||||
# Already list[list[str]] (or another shape) -> pass through unchanged.
|
||||
return input, {} # type: ignore[return-value]
|
||||
return input, {} # pyright: ignore[reportReturnType] # list[list[str]] branch
|
||||
|
||||
@classmethod
|
||||
def _auto_chunk_params(cls, optional_params: dict) -> dict:
|
||||
def _auto_chunk_params(cls, optional_params: dict) -> dict: # mutable-ok: matches public API
|
||||
"""
|
||||
Params required to send a flat ``list[str]`` to the contextualized API.
|
||||
|
||||
``enable_auto_chunking=True`` requires ``input_type="document"``, so set
|
||||
it unless the caller already provided an ``input_type``.
|
||||
"""
|
||||
params: dict[str, Any] = {
|
||||
params: dict[str, Any] = { # mutable-ok: building return value
|
||||
"enable_auto_chunking": True,
|
||||
"chunk_size": cls.AUTO_CHUNK_SIZE,
|
||||
}
|
||||
|
|
@ -189,16 +184,18 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
|
|||
model_response: EmbeddingResponse,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
api_key: str | None = None,
|
||||
request_data: dict = {},
|
||||
optional_params: dict = {},
|
||||
litellm_params: dict = {},
|
||||
request_data: dict = {}, # mutable-ok: base class signature
|
||||
optional_params: dict = {}, # mutable-ok: base class signature
|
||||
litellm_params: dict = {}, # mutable-ok: base class signature
|
||||
) -> EmbeddingResponse:
|
||||
try:
|
||||
raw_response_json = raw_response.json()
|
||||
except Exception:
|
||||
raise VoyageError(message=raw_response.text, status_code=raw_response.status_code)
|
||||
raise VoyageError(
|
||||
message=raw_response.text,
|
||||
status_code=raw_response.status_code,
|
||||
)
|
||||
|
||||
# model_response.usage
|
||||
model_response.model = raw_response_json.get("model")
|
||||
model_response.data = raw_response_json.get("data")
|
||||
model_response.object = raw_response_json.get("object")
|
||||
|
|
@ -210,7 +207,12 @@ class VoyageContextualEmbeddingConfig(BaseEmbeddingConfig):
|
|||
model_response.usage = usage
|
||||
return model_response
|
||||
|
||||
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
|
||||
def get_error_class(
|
||||
self,
|
||||
error_message: str,
|
||||
status_code: int,
|
||||
headers: dict | httpx.Headers, # mutable-ok: base class signature
|
||||
) -> BaseLLMException:
|
||||
return VoyageError(message=error_message, status_code=status_code, headers=headers)
|
||||
|
||||
@staticmethod
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue