Skip to content

Commit a82b79c

Browse files
fix: Error out when creating vector store with unknown embedding model (#4154)
# What does this PR do? Error out when creating vector store with unknown embedding model Closes #4047 ## Test Plan Added tests Signed-off-by: Francisco Javier Arceo <farceo@redhat.com>
1 parent 2441ca9 commit a82b79c

File tree

2 files changed

+46
-1
lines changed

2 files changed

+46
-1
lines changed

src/llama_stack/core/routers/vector_io.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,9 @@
1414
HealthResponse,
1515
HealthStatus,
1616
InterleavedContent,
17+
ModelNotFoundError,
1718
ModelType,
19+
ModelTypeError,
1820
OpenAICreateVectorStoreFileBatchRequestWithExtraBody,
1921
OpenAICreateVectorStoreRequestWithExtraBody,
2022
QueryChunksResponse,
@@ -124,6 +126,14 @@ async def openai_create_vector_store(
124126
if embedding_model is not None and embedding_dimension is None:
125127
embedding_dimension = await self._get_embedding_model_dimension(embedding_model)
126128

129+
# Validate that embedding model exists and is of the correct type
130+
if embedding_model is not None:
131+
model = await self.routing_table.get_object_by_identifier("model", embedding_model)
132+
if model is None:
133+
raise ModelNotFoundError(embedding_model)
134+
if model.model_type != ModelType.embedding:
135+
raise ModelTypeError(embedding_model, model.model_type, ModelType.embedding)
136+
127137
# Auto-select provider if not specified
128138
if provider_id is None:
129139
num_providers = len(self.routing_table.impls_by_provider_id)

tests/unit/core/routers/test_vector_io.py

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,12 @@
77
from unittest.mock import AsyncMock, Mock
88

99
import pytest
10-
from llama_stack_api import OpenAICreateVectorStoreRequestWithExtraBody
10+
from llama_stack_api import (
11+
ModelNotFoundError,
12+
ModelType,
13+
ModelTypeError,
14+
OpenAICreateVectorStoreRequestWithExtraBody,
15+
)
1116

1217
from llama_stack.core.routers.vector_io import VectorIORouter
1318

@@ -21,6 +26,7 @@ async def test_single_provider_auto_selection():
2126
Mock(identifier="all-MiniLM-L6-v2", model_type="embedding", metadata={"embedding_dimension": 384})
2227
]
2328
)
29+
mock_routing_table.get_object_by_identifier = AsyncMock(return_value=Mock(model_type=ModelType.embedding))
2430
mock_routing_table.register_vector_store = AsyncMock(
2531
return_value=Mock(identifier="vs_123", provider_id="inline::faiss", provider_resource_id="vs_123")
2632
)
@@ -48,6 +54,7 @@ async def test_create_vector_stores_multiple_providers_missing_provider_id_error
4854
Mock(identifier="all-MiniLM-L6-v2", model_type="embedding", metadata={"embedding_dimension": 384})
4955
]
5056
)
57+
mock_routing_table.get_object_by_identifier = AsyncMock(return_value=Mock(model_type=ModelType.embedding))
5158
router = VectorIORouter(mock_routing_table)
5259
request = OpenAICreateVectorStoreRequestWithExtraBody.model_validate(
5360
{"name": "test_store", "embedding_model": "all-MiniLM-L6-v2"}
@@ -117,3 +124,31 @@ async def test_update_vector_store_same_provider_id_succeeds():
117124
provider.openai_update_vector_store.assert_called_once_with(
118125
vector_store_id="vs_123", name="updated_name", expires_after=None, metadata={"provider_id": "inline::faiss"}
119126
)
127+
128+
129+
async def test_create_vector_store_with_unknown_embedding_model_raises_error():
130+
"""Test that creating a vector store with an unknown embedding model raises ModelNotFoundError."""
131+
mock_routing_table = Mock(impls_by_provider_id={"provider": "mock"})
132+
mock_routing_table.get_object_by_identifier = AsyncMock(return_value=None)
133+
134+
router = VectorIORouter(mock_routing_table)
135+
request = OpenAICreateVectorStoreRequestWithExtraBody.model_validate(
136+
{"embedding_model": "unknown-model", "embedding_dimension": 384}
137+
)
138+
139+
with pytest.raises(ModelNotFoundError, match="Model 'unknown-model' not found"):
140+
await router.openai_create_vector_store(request)
141+
142+
143+
async def test_create_vector_store_with_wrong_model_type_raises_error():
144+
"""Test that creating a vector store with a non-embedding model raises ModelTypeError."""
145+
mock_routing_table = Mock(impls_by_provider_id={"provider": "mock"})
146+
mock_routing_table.get_object_by_identifier = AsyncMock(return_value=Mock(model_type=ModelType.llm))
147+
148+
router = VectorIORouter(mock_routing_table)
149+
request = OpenAICreateVectorStoreRequestWithExtraBody.model_validate(
150+
{"embedding_model": "text-model", "embedding_dimension": 384}
151+
)
152+
153+
with pytest.raises(ModelTypeError, match="Model 'text-model' is of type"):
154+
await router.openai_create_vector_store(request)

0 commit comments

Comments
 (0)