phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-12-12 12:06:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add rerank models and rerank API change

Browse source
This commit is contained in:
Jiayi 2025-10-16 17:27:38 -07:00
parent f675fdda0f
commit 51c923f096
12 changed files with 215 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

11
llama_stack/providers/utils/inference/openai_mixin.py
View file

@ -78,6 +78,10 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
# Format: {"model_id": {"embedding_dimension": 1536, "context_length": 8192}}
embedding_model_metadata: dict[str, dict[str, int]] = {}
# List of rerank model IDs for this provider
# Can be set by subclasses or instances to provide rerank models
rerank_model_list: list[str] = []
# Cache of available models keyed by model ID
# This is set in list_models() and used in check_model_availability()
_model_cache: dict[str, Model] = {}
@ -424,6 +428,13 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
model_type=ModelType.embedding,
metadata=metadata,
)
elif provider_model_id in self.rerank_model_list:
model = Model(
provider_id=self.__provider_id__, # type: ignore[attr-defined]
provider_resource_id=provider_model_id,
identifier=provider_model_id,
model_type=ModelType.rerank,
)
else:
model = Model(
provider_id=self.__provider_id__, # type: ignore[attr-defined]