fix: refactor Ollama model fetching to use async and filter capabilities

src/lfx/src/lfx/components/models/language_model.py

@@ -1,6 +1,8 @@
+import asyncio
 from typing import Any
 from urllib.parse import urljoin
 
+import httpx
 import requests
 from langchain_anthropic import ChatAnthropic
 from langchain_ibm import ChatWatsonx
@@ -32,6 +34,14 @@
     "https://ca-tor.ml.cloud.ibm.com",
 ]
 
+# Ollama API constants
+HTTP_STATUS_OK = 200
+JSON_MODELS_KEY = "models"
+JSON_NAME_KEY = "name"
+JSON_CAPABILITIES_KEY = "capabilities"
+DESIRED_CAPABILITY = "completion"
+DEFAULT_OLLAMA_URL = "http://localhost:11434"
+
 
 class LanguageModelComponent(LCModelComponent):
     display_name = "Language Model"
@@ -56,27 +66,85 @@
             logger.exception("Error fetching IBM watsonx models. Using default models.")
             return IBM_WATSONX_DEFAULT_MODELS
 
-    @staticmethod
-    def fetch_ollama_models(base_url: str) -> list[str]:
-        """Fetch available models from the Ollama API."""
+    async def is_valid_ollama_url(self, url: str) -> bool:
+        """Check if the provided URL is a valid Ollama API endpoint."""
+        try:
+            async with httpx.AsyncClient() as client:
+                url = transform_localhost_url(url)
+                if not url:
+                    return False
+                # Strip /v1 suffix if present, as Ollama API endpoints are at root level
+                url = url.rstrip("/").removesuffix("/v1")
+                if not url.endswith("/"):
+                    url = url + "/"
+                return (await client.get(url=urljoin(url, "api/tags"))).status_code == HTTP_STATUS_OK
+        except httpx.RequestError:
+            logger.debug(f"Invalid Ollama URL: {url}")
+            return False
+
+    async def get_ollama_models(self, base_url_value: str) -> list[str]:
+        """Fetch available completion models from the Ollama API.
+
+        Filters out embedding models and only returns models with completion capability.
+
+        Args:
+            base_url_value (str): The base URL of the Ollama API.
+
+        Returns:
+            list[str]: A sorted list of model names that support completion.
+
+        Raises:
+            ValueError: If there is an issue with the API request or response.
+        """
         try:
             # Strip /v1 suffix if present, as Ollama API endpoints are at root level
-            base_url = base_url.rstrip("/").removesuffix("/v1")
+            base_url = base_url_value.rstrip("/").removesuffix("/v1")
             if not base_url.endswith("/"):
                 base_url = base_url + "/"
             base_url = transform_localhost_url(base_url)
 
             # Ollama REST API to return models
             tags_url = urljoin(base_url, "api/tags")
 
-            response = requests.get(tags_url, timeout=10)
-            response.raise_for_status()
-            data = response.json()
-            models = [model["name"] for model in data.get("models", [])]
-            return sorted(models)
-        except Exception:  # noqa: BLE001
-            logger.exception("Error fetching Ollama models. Returning empty list.")
-            return []
+            # Ollama REST API to return model capabilities
+            show_url = urljoin(base_url, "api/show")
+
+            async with httpx.AsyncClient() as client:
+                # Fetch available models
+                tags_response = await client.get(url=tags_url)
+                tags_response.raise_for_status()
+                models = tags_response.json()
+                if asyncio.iscoroutine(models):
+                    models = await models
+                await logger.adebug(f"Available models: {models}")
+
+                # Filter models that are NOT embedding models
+                model_ids = []
+                for model in models.get(JSON_MODELS_KEY, []):
+                    model_name = model.get(JSON_NAME_KEY)
+                    if not model_name:
+                        continue
+                    await logger.adebug(f"Checking model: {model_name}")
+
+                    payload = {"model": model_name}
+                    show_response = await client.post(url=show_url, json=payload)
+                    show_response.raise_for_status()
+                    json_data = show_response.json()
+                    if asyncio.iscoroutine(json_data):
+                        json_data = await json_data
+
+                    capabilities = json_data.get(JSON_CAPABILITIES_KEY, [])
+                    await logger.adebug(f"Model: {model_name}, Capabilities: {capabilities}")
+
+                    if DESIRED_CAPABILITY in capabilities:
+                        model_ids.append(model_name)
+
+                return sorted(model_ids)
+
+        except (httpx.RequestError, ValueError) as e:
+            msg = "Could not get model names from Ollama."
+            await logger.aexception(msg)
+            raise ValueError(msg) from e
 
     inputs = [
         DropdownInput(
@@ -101,6 +169,7 @@
             value=OPENAI_CHAT_MODEL_NAMES[0],
             info="Select the model to use",
             real_time_refresh=True,
+            refresh_button=True,
         ),
         SecretStrInput(
             name="api_key",
@@ -129,10 +198,11 @@
         MessageTextInput(
             name="ollama_base_url",
             display_name="Ollama API URL",
-            info="Endpoint of the Ollama API (Ollama only). Defaults to http://localhost:11434",
-            value="http://localhost:11434",
+            info=f"Endpoint of the Ollama API (Ollama only). Defaults to {DEFAULT_OLLAMA_URL}",
+            value=DEFAULT_OLLAMA_URL,
             show=False,
             real_time_refresh=True,
+            load_from_db=True,
         ),
         MessageInput(
             name="input_value",
@@ -252,7 +322,9 @@
         msg = f"Unknown provider: {provider}"
         raise ValueError(msg)
 
-    def update_build_config(self, build_config: dotdict, field_value: Any, field_name: str | None = None) -> dotdict:
+    async def update_build_config(
+        self, build_config: dotdict, field_value: Any, field_name: str | None = None
+    ) -> dotdict:
         if field_name == "provider":
             if field_value == "OpenAI":
                 build_config["model_name"]["options"] = OPENAI_CHAT_MODEL_NAMES + OPENAI_REASONING_MODEL_NAMES
@@ -288,14 +360,48 @@
                 build_config["ollama_base_url"]["show"] = False
             elif field_value == "Ollama":
                 # Fetch Ollama models from the API
-                ollama_url = build_config["ollama_base_url"].get("value", "http://localhost:11434")
-                models = self.fetch_ollama_models(base_url=ollama_url)
-                build_config["model_name"]["options"] = models
-                build_config["model_name"]["value"] = models[0] if models else ""
                 build_config["api_key"]["show"] = False
                 build_config["base_url_ibm_watsonx"]["show"] = False
                 build_config["project_id"]["show"] = False
                 build_config["ollama_base_url"]["show"] = True
+                build_config["ollama_base_url"]["load_from_db"] = True
+
+                # Try multiple sources to get the URL (in order of preference):
+                # 1. Instance attribute (already resolved from global/db)
+                # 2. Build config value (may be a global variable reference)
+                # 3. Default value
+                ollama_url = getattr(self, "ollama_base_url", None)
+                if not ollama_url:
+                    config_value = build_config["ollama_base_url"].get("value", DEFAULT_OLLAMA_URL)
+                    # If config_value looks like a variable name (all caps with underscores), use default
+                    is_variable_ref = (
+                        config_value
+                        and isinstance(config_value, str)
+                        and config_value.isupper()
+                        and "_" in config_value
+                    )
+                    if is_variable_ref:
+                        await logger.adebug(
+                            f"Config value appears to be a variable reference: {config_value}, using default"
+                        )
+                        ollama_url = DEFAULT_OLLAMA_URL
+                    else:
+                        ollama_url = config_value
+
+                await logger.adebug(f"Fetching Ollama models for provider switch. URL: {ollama_url}")
+                if await self.is_valid_ollama_url(ollama_url):
+                    try:
+                        models = await self.get_ollama_models(base_url_value=ollama_url)
+                        build_config["model_name"]["options"] = models
+                        build_config["model_name"]["value"] = models[0] if models else ""
+                    except ValueError:
+                        await logger.awarning("Failed to fetch Ollama models. Setting empty options.")
+                        build_config["model_name"]["options"] = []
+                        build_config["model_name"]["value"] = ""
+                else:
+                    await logger.awarning(f"Invalid Ollama URL: {ollama_url}")
+                    build_config["model_name"]["options"] = []
+                    build_config["model_name"]["value"] = ""
         elif (
             field_name == "base_url_ibm_watsonx"
             and field_value
@@ -311,22 +417,46 @@
                 logger.info(info_message)
             except Exception:  # noqa: BLE001
                 logger.exception("Error updating IBM model options.")
-        elif (
-            field_name == "ollama_base_url" and field_value and hasattr(self, "provider") and self.provider == "Ollama"
-        ):
+        elif field_name == "ollama_base_url":
             # Fetch Ollama models when ollama_base_url changes
-            try:
-                models = self.fetch_ollama_models(base_url=field_value)
-                build_config["model_name"]["options"] = models
-                build_config["model_name"]["value"] = models[0] if models else ""
-                info_message = f"Updated model options: {len(models)} models found in {field_value}"
-                logger.info(info_message)
-            except Exception:  # noqa: BLE001
-                logger.exception("Error updating Ollama model options.")
-        elif field_name == "model_name" and field_value.startswith("o1") and self.provider == "OpenAI":
+            # Use the field_value directly since this is triggered when the field changes
+            logger.debug(
+                f"Fetching Ollama models from updated URL: {build_config['ollama_base_url']} and value {self.ollama_base_url}"
+            )
+            await logger.adebug(f"Fetching Ollama models from updated URL: {self.ollama_base_url}")
+            if await self.is_valid_ollama_url(self.ollama_base_url):
+                try:
+                    models = await self.get_ollama_models(base_url_value=self.ollama_base_url)
+                    build_config["model_name"]["options"] = models
+                    build_config["model_name"]["value"] = models[0] if models else ""
+                    info_message = f"Updated model options: {len(models)} models found in {self.ollama_base_url}"
+                    await logger.ainfo(info_message)
+                except ValueError:
+                    await logger.awarning("Error updating Ollama model options.")
+                    build_config["model_name"]["options"] = []
+                    build_config["model_name"]["value"] = ""
+            else:
+                await logger.awarning(f"Invalid Ollama URL: {self.ollama_base_url}")
+                build_config["model_name"]["options"] = []
+                build_config["model_name"]["value"] = ""
+        elif field_name == "model_name":
+            # Refresh Ollama models when model_name field is accessed
+            if hasattr(self, "provider") and self.provider == "Ollama":
+                ollama_url = getattr(self, "ollama_base_url", DEFAULT_OLLAMA_URL)
+                if await self.is_valid_ollama_url(ollama_url):
+                    try:
+                        models = await self.get_ollama_models(base_url_value=ollama_url)
+                        build_config["model_name"]["options"] = models
+                    except ValueError:
+                        await logger.awarning("Failed to refresh Ollama models.")
+                        build_config["model_name"]["options"] = []
+                else:
+                    build_config["model_name"]["options"] = []
+
             # Hide system_message for o1 models - currently unsupported
-            if "system_message" in build_config:
-                build_config["system_message"]["show"] = False
-        elif field_name == "model_name" and not field_value.startswith("o1") and "system_message" in build_config:
-            build_config["system_message"]["show"] = True
+            if field_value and field_value.startswith("o1") and hasattr(self, "provider") and self.provider == "OpenAI":
+                if "system_message" in build_config:
+                    build_config["system_message"]["show"] = False
+            elif "system_message" in build_config:
+                build_config["system_message"]["show"] = True
         return build_config

src/lfx/src/lfx/components/ollama/ollama.py

@@ -350,6 +350,8 @@ async def update_build_config(self, build_config: dict, field_value: Any, field_
                     build_config["mirostat_tau"]["value"] = 5
 
         if field_name in {"model_name", "base_url", "tool_model_enabled"}:
+            logger.warning(f"Fetching Ollama models from updated URL: {build_config['base_url']}")
+
             if await self.is_valid_ollama_url(self.base_url):
                 tool_model_enabled = build_config["tool_model_enabled"].get("value", False) or self.tool_model_enabled
                 build_config["model_name"]["options"] = await self.get_models(