mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(bedrock): keep the embeddings endpoint when a batch deployment overrides the model's mode
A deployment with model_info.mode set to batch overwrites the shared cost-map entry's mode, so the job's model no longer read as an embedding model. An entry with an output vector size is now also read as one.
This commit is contained in:
parent
7f51cb1ce7
commit
d54a31f1a8
2 changed files with 20 additions and 2 deletions
|
|
@ -91,10 +91,13 @@ def openai_endpoint_for_bedrock_job_model(
|
|||
if not model_id:
|
||||
return "/v1/chat/completions"
|
||||
try:
|
||||
mode: Final = get_cached_model_info()(model=model_id, custom_llm_provider="bedrock").get("mode")
|
||||
model_info: Final = get_cached_model_info()(model=model_id, custom_llm_provider="bedrock")
|
||||
except ModelNotMappedError:
|
||||
return "/v1/chat/completions"
|
||||
return "/v1/embeddings" if mode == "embedding" else "/v1/chat/completions"
|
||||
is_embedding_model: Final = (
|
||||
model_info.get("mode") == "embedding" or model_info.get("output_vector_size") is not None
|
||||
)
|
||||
return "/v1/embeddings" if is_embedding_model else "/v1/chat/completions"
|
||||
|
||||
|
||||
def titan_embedding_usage_from_batch_output(model_output: Mapping[str, object]) -> Usage | None:
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ from botocore.auth import SigV4Auth
|
|||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.credentials import Credentials
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig
|
||||
from litellm.types.utils import LlmProviders
|
||||
|
||||
|
|
@ -940,6 +941,20 @@ def test_retrieve_response_endpoint_follows_the_job_models_mode(
|
|||
assert out.endpoint == expected_endpoint
|
||||
|
||||
|
||||
def test_retrieve_response_endpoint_survives_a_batch_deployment_overriding_the_models_mode(
|
||||
config: BedrockBatchesConfig, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
job_model_id = "amazon.titan-embed-text-v2:0"
|
||||
monkeypatch.setitem(litellm.model_cost, job_model_id, {**litellm.model_cost[job_model_id], "mode": "batch"})
|
||||
out = config.transform_retrieve_batch_response(
|
||||
model=None,
|
||||
raw_response=_raw({"jobArn": ARN, "status": "InProgress", "modelId": job_model_id}),
|
||||
logging_obj=MagicMock(),
|
||||
litellm_params={},
|
||||
)
|
||||
assert out.endpoint == "/v1/embeddings"
|
||||
|
||||
|
||||
def test_retrieve_response_enriches_metadata(config):
|
||||
body = {
|
||||
"jobArn": ARN,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue