This commit is contained in:
Ashwin Bharambe 2025-10-16 06:43:17 -07:00
parent 54354257dc
commit bbd40d8cc4
3 changed files with 3 additions and 6 deletions

View file

@ -102,9 +102,6 @@ class DiskDistributionRegistry(DistributionRegistry):
"Unregister it first if you want to replace it." "Unregister it first if you want to replace it."
) )
if "sentence-transformers/sentence-transformers" in obj.identifier:
raise Exception("OMG")
await self.kvstore.set( await self.kvstore.set(
KEY_FORMAT.format(type=obj.type, identifier=obj.identifier), KEY_FORMAT.format(type=obj.type, identifier=obj.identifier),
obj.model_dump_json(), obj.model_dump_json(),

View file

@ -435,7 +435,7 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
""" """
# First check if the model is pre-registered in the model store # First check if the model is pre-registered in the model store
if hasattr(self, "model_store") and self.model_store: if hasattr(self, "model_store") and self.model_store:
qualified_model = f"{self.__provider_id__}/{model}" qualified_model = f"{self.__provider_id__}/{model}" # type: ignore[attr-defined]
if await self.model_store.has_model(qualified_model): if await self.model_store.has_model(qualified_model):
return True return True

View file

@ -205,7 +205,7 @@ class TestOpenAIMixinCheckModelAvailability:
assert await mixin.check_model_availability("pre-registered-model") assert await mixin.check_model_availability("pre-registered-model")
# Should not call the provider's list_models since model was found in store # Should not call the provider's list_models since model was found in store
mock_client_with_models.models.list.assert_not_called() mock_client_with_models.models.list.assert_not_called()
mock_model_store.has_model.assert_called_once_with("pre-registered-model") mock_model_store.has_model.assert_called_once_with("test-provider/pre-registered-model")
async def test_check_model_availability_fallback_to_provider_when_not_in_store( async def test_check_model_availability_fallback_to_provider_when_not_in_store(
self, mixin, mock_client_with_models, mock_client_context self, mixin, mock_client_with_models, mock_client_context
@ -222,7 +222,7 @@ class TestOpenAIMixinCheckModelAvailability:
assert await mixin.check_model_availability("some-mock-model-id") assert await mixin.check_model_availability("some-mock-model-id")
# Should call the provider's list_models since model was not found in store # Should call the provider's list_models since model was not found in store
mock_client_with_models.models.list.assert_called_once() mock_client_with_models.models.list.assert_called_once()
mock_model_store.has_model.assert_called_once_with("some-mock-model-id") mock_model_store.has_model.assert_called_once_with("test-provider/some-mock-model-id")
class TestOpenAIMixinCacheBehavior: class TestOpenAIMixinCacheBehavior: