Merge branch 'main' into allow-dynamic-models-ollama

llama_stack/providers/remote/inference/ollama/ollama.py

@@ -5,6 +5,7 @@
 # the root directory of this source tree.
 
 
+import asyncio
 import base64
 import uuid
 from collections.abc import AsyncGenerator, AsyncIterator
@@ -89,23 +90,88 @@ class OllamaInferenceAdapter(
     InferenceProvider,
     ModelRegistryHelper,
 ):
+    # automatically set by the resolver when instantiating the provider
+    __provider_id__: str
+
     def __init__(self, config: OllamaImplConfig) -> None:
         ModelRegistryHelper.__init__(self, MODEL_ENTRIES)
-        self.url = config.url
+        self.config = config
+        self._client = None
+        self._openai_client = None
 
     @property
     def client(self) -> AsyncClient:
-        return AsyncClient(host=self.url)
+        if self._client is None:
+            self._client = AsyncClient(host=self.config.url)
+        return self._client
 
     @property
     def openai_client(self) -> AsyncOpenAI:
-        return AsyncOpenAI(base_url=f"{self.url}/v1", api_key="ollama")
+        if self._openai_client is None:
+            self._openai_client = AsyncOpenAI(base_url=f"{self.config.url}/v1", api_key="ollama")
+        return self._openai_client
 
     async def initialize(self) -> None:
-        logger.debug(f"checking connectivity to Ollama at `{self.url}`...")
+        logger.info(f"checking connectivity to Ollama at `{self.config.url}`...")
         health_response = await self.health()
         if health_response["status"] == HealthStatus.ERROR:
-            raise RuntimeError("Ollama Server is not running, start it using `ollama serve` in a separate terminal")
+            logger.warning(
+                "Ollama Server is not running, make sure to start it using `ollama serve` in a separate terminal"
+            )
+
+        if self.config.refresh_models:
+            logger.debug("ollama starting background model refresh task")
+            self._refresh_task = asyncio.create_task(self._refresh_models())
+
+            def cb(task):
+                if task.cancelled():
+                    import traceback
+
+                    logger.error(f"ollama background refresh task canceled:\n{''.join(traceback.format_stack())}")
+                elif task.exception():
+                    logger.error(f"ollama background refresh task died: {task.exception()}")
+                else:
+                    logger.error("ollama background refresh task completed unexpectedly")
+
+            self._refresh_task.add_done_callback(cb)
+
+    async def _refresh_models(self) -> None:
+        # Wait for model store to be available (with timeout)
+        waited_time = 0
+        while not self.model_store and waited_time < 60:
+            await asyncio.sleep(1)
+            waited_time += 1
+
+        if not self.model_store:
+            raise ValueError("Model store not set after waiting 60 seconds")
+
+        provider_id = self.__provider_id__
+        while True:
+            try:
+                response = await self.client.list()
+            except Exception as e:
+                logger.warning(f"Failed to list models: {str(e)}")
+                await asyncio.sleep(self.config.refresh_models_interval)
+                continue
+
+            models = []
+            for m in response.models:
+                model_type = ModelType.embedding if m.details.family in ["bert"] else ModelType.llm
+                if model_type == ModelType.embedding:
+                    continue
+                models.append(
+                    Model(
+                        identifier=m.model,
+                        provider_resource_id=m.model,
+                        provider_id=provider_id,
+                        metadata={},
+                        model_type=model_type,
+                    )
+                )
+            await self.model_store.update_registered_llm_models(provider_id, models)
+            logger.debug(f"ollama refreshed model list ({len(models)} models)")
+
+            await asyncio.sleep(self.config.refresh_models_interval)
 
     async def health(self) -> HealthResponse:
         """
@@ -157,7 +223,12 @@ class OllamaInferenceAdapter(
         return available_models
 
     async def shutdown(self) -> None:
-        pass
+        if hasattr(self, "_refresh_task") and not self._refresh_task.done():
+            logger.debug("ollama cancelling background refresh task")
+            self._refresh_task.cancel()
+
+        self._client = None
+        self._openai_client = None
 
     async def unregister_model(self, model_id: str) -> None:
         pass