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

ensure embedding models have dimensions

Browse source
This commit is contained in:
melonkernel 2025-09-04 14:01:51 +03:00
parent 2a478fb1d5
commit 27a6918834
2 changed files with 24 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

23
tests/unit/providers/inference/test_model_registry.py
View file

@ -15,9 +15,9 @@ class TestProviderDataValidator(BaseModel):
test_api_key: str | None = Field(default=None)
MODEL_ENTRIES_WITHOUT_ALIASES = [
ProviderModelEntry(model_type=ModelType.llm, provider_model_id="test-llm-model", aliases=[]),
ProviderModelEntry(model_type=ModelType.embedding, provider_model_id="test-text-embedding-model", aliases=[], metadata={"embedding_dimension": 1536, "context_length": 8192}),
]
ProviderModelEntry(model_type=ModelType.llm, provider_model_id="test-llm-model", aliases=[]),
ProviderModelEntry(model_type=ModelType.embedding, provider_model_id="test-text-embedding-model", aliases=[], metadata={"embedding_dimension": 1536, "context_length": 8192}),
]
class TestLiteLLMAdapterWithModelEntries(LiteLLMOpenAIMixin):
def __init__(self, config: TestConfig):
@ -38,7 +38,6 @@ def adapter_with_model_entries():
return adapter
async def test_model_types_are_correct(adapter_with_model_entries):
"""Test that model types are correct"""
model_entries = adapter_with_model_entries.model_entries
@ -54,4 +53,20 @@ async def test_model_types_are_correct(adapter_with_model_entries):
embedding_models = [m for m in models if m.model_type == ModelType.embedding]
assert len(embedding_models) == len(embedding_model_entries)
def test_embedding_metadata_is_required():
with pytest.raises(ValueError):
entry1 = ProviderModelEntry(
model_type=ModelType.embedding,
provider_model_id="test-text-embedding-model",
aliases=[],
metadata={}
)
entry2 = ProviderModelEntry(
model_type=ModelType.embedding,
provider_model_id="test-text-embedding-model",
aliases=[],
metadata={"embedding_dimension": 1536}
)
assert entry2.metadata["embedding_dimension"] == 1536