From d9c460cb4bebaf03852b9845b1f791b47d67fbb9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9bastien=20Han?= Date: Thu, 6 Feb 2025 14:57:19 +0100 Subject: [PATCH] refactor(ollama): model availability check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Moved model availability check logic into a dedicated check_model_availability function. Eliminated redundant code by reusing the helper function in both embedding and non-embedding model registration. Signed-off-by: Sébastien Han --- .../remote/inference/ollama/ollama.py | 22 ++++++++----------- 1 file changed, 9 insertions(+), 13 deletions(-) diff --git a/llama_stack/providers/remote/inference/ollama/ollama.py b/llama_stack/providers/remote/inference/ollama/ollama.py index cff8aa742..ecd195854 100644 --- a/llama_stack/providers/remote/inference/ollama/ollama.py +++ b/llama_stack/providers/remote/inference/ollama/ollama.py @@ -352,24 +352,20 @@ class OllamaInferenceAdapter(Inference, ModelsProtocolPrivate): return EmbeddingsResponse(embeddings=embeddings) async def register_model(self, model: Model) -> Model: - # ollama does not have embedding models running. Check if the model is in list of available models. - if model.model_type == ModelType.embedding: - response = await self.client.list() + async def check_model_availability(model_id: str): + response = await self.client.ps() available_models = [m["model"] for m in response["models"]] - if model.provider_resource_id not in available_models: + if model_id not in available_models: raise ValueError( - f"Model '{model.provider_resource_id}' is not available in Ollama. " - f"Available models: {', '.join(available_models)}" + f"Model '{model_id}' is not available in Ollama. Available models: {', '.join(available_models)}" ) + + if model.model_type == ModelType.embedding: + await check_model_availability(model.provider_resource_id) return model + model = await self.register_helper.register_model(model) - models = await self.client.ps() - available_models = [m["model"] for m in models["models"]] - if model.provider_resource_id not in available_models: - raise ValueError( - f"Model '{model.provider_resource_id}' is not available in Ollama. " - f"Available models: {', '.join(available_models)}" - ) + await check_model_availability(model.provider_resource_id) return model