Fix mypy errors

This commit is contained in:
Sameer Kankute 2026-04-02 20:49:24 +05:30
parent 1169fbffd0
commit 6633a18efe
No known key found for this signature in database
6 changed files with 19 additions and 16 deletions

View file

@ -1272,15 +1272,15 @@ class BedrockLLM(BaseAWSLLM):
return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider)
return None
@staticmethod
def get_bedrock_model_id(
self,
optional_params: dict,
provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL],
model: str,
) -> str:
modelId = optional_params.pop("model_id", None)
if modelId is not None:
modelId = self.encode_model_id(model_id=modelId)
modelId = BedrockLLM.encode_model_id(model_id=modelId)
else:
modelId = model
@ -1288,19 +1288,19 @@ class BedrockLLM(BaseAWSLLM):
modelId, optional_params
)
if provider == "llama" and "llama/" in modelId:
modelId = self._get_model_id_for_llama_like_model(modelId)
modelId = BedrockLLM._get_model_id_for_llama_like_model(modelId)
return modelId
@staticmethod
def _get_model_id_for_llama_like_model(
self,
model: str,
) -> str:
"""
Remove `llama` from modelID since `llama` is simply a spec to follow for custom bedrock models
"""
model_id = model.replace("llama/", "")
return self.encode_model_id(model_id=model_id)
return BedrockLLM.encode_model_id(model_id=model_id)
def get_response_stream_shape():

View file

@ -568,15 +568,15 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
return cast(litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL, provider)
return None
@staticmethod
def get_bedrock_model_id(
self,
optional_params: dict,
provider: Optional[litellm.BEDROCK_INVOKE_PROVIDERS_LITERAL],
model: str,
) -> str:
modelId = optional_params.pop("model_id", None)
if modelId is not None:
modelId = self.encode_model_id(model_id=modelId)
modelId = AmazonInvokeConfig.encode_model_id(model_id=modelId)
else:
modelId = model
@ -585,15 +585,17 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
modelId, optional_params
)
if provider == "llama" and "llama/" in modelId:
modelId = self._get_model_id_from_model_with_spec(modelId, spec="llama")
modelId = AmazonInvokeConfig._get_model_id_from_model_with_spec(
modelId, spec="llama"
)
elif provider == "deepseek_r1" and "deepseek_r1/" in modelId:
modelId = self._get_model_id_from_model_with_spec(
modelId = AmazonInvokeConfig._get_model_id_from_model_with_spec(
modelId, spec="deepseek_r1"
)
return modelId
@staticmethod
def _get_model_id_from_model_with_spec(
self,
model: str,
spec: str,
) -> str:
@ -601,9 +603,10 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
Remove `llama` from modelID since `llama` is simply a spec to follow for custom bedrock models
"""
model_id = model.replace(spec + "/", "")
return self.encode_model_id(model_id=model_id)
return AmazonInvokeConfig.encode_model_id(model_id=model_id)
def encode_model_id(self, model_id: str) -> str:
@staticmethod
def encode_model_id(model_id: str) -> str:
"""
Double encode the model ID to ensure it matches the expected double-encoded format.
Args:

View file

@ -343,7 +343,7 @@ class GeminiVideoConfig(BaseVideoConfig):
model=model,
)
usage_data = {}
usage_data: Dict[str, Any] = {}
if request_data:
parameters = request_data.get("parameters", {})
duration = (

View file

@ -132,7 +132,7 @@ class VertexAITextEmbeddingConfig(BaseModel):
vertex_request["labels"] = labels
return vertex_request
vertex_request: VertexEmbeddingRequest = VertexEmbeddingRequest()
vertex_request = VertexEmbeddingRequest()
vertex_text_embedding_input_list: List[TextEmbeddingInput] = []
task_type: Optional[TaskType] = optional_params.get("task_type")
title = optional_params.get("title")

View file

@ -363,7 +363,7 @@ class VertexAIVideoConfig(BaseVideoConfig, VertexBase):
id=video_id, object="video", status="processing", model=model
)
usage_data = {}
usage_data: Dict[str, Any] = {}
if request_data:
parameters = request_data.get("parameters", {})
duration = (

View file

@ -181,7 +181,7 @@ class LiteLLMCompletionResponsesConfig:
)
# Extract reasoning_effort from reasoning parameter
reasoning_effort = None
reasoning_effort: Optional[Any] = None
reasoning_param = responses_api_request.get("reasoning")
if reasoning_param:
if isinstance(reasoning_param, dict):