diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index 692f80c..c9098bd 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -9,6 +9,10 @@ services: # On macOS/Windows Docker Desktop: host.docker.internal is built-in. extra_hosts: - "host.docker.internal:host-gateway" + environment: + # Tells comfydv where to reach Ollama on the host so model dropdowns are + # pre-populated at server start without needing a manual Refresh click. + - OLLAMA_HOST=http://host.docker.internal:11434 command: - sh - -c diff --git a/docs/assets/ollama_chat.png b/docs/assets/ollama_chat.png index 72ae97c..de55783 100644 Binary files a/docs/assets/ollama_chat.png and b/docs/assets/ollama_chat.png differ diff --git a/docs/assets/ollama_lifecycle.png b/docs/assets/ollama_lifecycle.png index 9579107..a8d8c97 100644 Binary files a/docs/assets/ollama_lifecycle.png and b/docs/assets/ollama_lifecycle.png differ diff --git a/docs/assets/ollama_workflow.png b/docs/assets/ollama_workflow.png index de17548..298f828 100644 Binary files a/docs/assets/ollama_workflow.png and b/docs/assets/ollama_workflow.png differ diff --git a/src/comfydv/ollama.py b/src/comfydv/ollama.py index 985f611..a8f97a6 100644 --- a/src/comfydv/ollama.py +++ b/src/comfydv/ollama.py @@ -11,6 +11,7 @@ ADR-005: OllamaClient node is the single source of the host URL. import asyncio import json import logging +import os import sys logger = logging.getLogger(__name__) @@ -82,10 +83,36 @@ async def _post_json(url: str, payload: dict, *, timeout: float = 120.0) -> dict raise RuntimeError(f"Cannot reach Ollama at {url}: {exc}") from exc -# Static placeholder — avoids a network call at import time (which always fails -# in CI and slows cold-start). The JS refresh button calls /dv/ollama/models -# at runtime to populate the live list. -_DEFAULT_MODELS: list[str] = ["(⟳ click Refresh models)"] +def _load_default_models() -> list[str]: + """Fetch the installed model list at server start-up. + + Tries OLLAMA_HOST env var first (set in docker-compose for host.docker.internal), + then falls back to localhost. Returns a one-element placeholder list only when + Ollama is genuinely unreachable so that COMBO validation doesn't reject saved + workflow values. + """ + candidates = [] + env_host = os.environ.get("OLLAMA_HOST", "").strip() + if env_host: + candidates.append(env_host) + candidates.append("http://host.docker.internal:11434") + candidates.append("http://localhost:11434") + + for host in candidates: + models = _run_async(_fetch_models(host)) + if models: + logger.info("Ollama models loaded from %s: %s", host, models) + return models + + logger.warning( + "Could not reach Ollama at any candidate host %s — " + "model dropdowns will be empty until a Refresh is clicked.", + candidates, + ) + return ["(start Ollama — click ⟳ Refresh)"] + + +_DEFAULT_MODELS: list[str] = _load_default_models() # --------------------------------------------------------------------------- @@ -255,16 +282,15 @@ class OllamaChatCompletion: return { "required": { "client": ("OLLAMA_CLIENT",), - "model": (_DEFAULT_MODELS, {}), + # Plain STRING so it can receive a wired value from OllamaLoadModel + # (or OllamaModelSelector) without needing a separate model_name socket. + "model": ("STRING", {"default": ""}), "prompt": ("STRING", {"multiline": True, "default": ""}), }, "optional": { "system": ("STRING", {"multiline": True, "default": ""}), "history": ("OLLAMA_HISTORY",), "options": ("OLLAMA_OPTIONS",), - # Wire OllamaLoadModel.model_name here to guarantee load runs - # before chat and to override the dropdown with the wired value. - "model_name": ("STRING", {"forceInput": True}), "timeout_secs": ("INT", {"default": 300, "min": 30, "max": 3600}), }, } @@ -282,12 +308,11 @@ class OllamaChatCompletion: system="", history=None, options=None, - model_name=None, timeout_secs=300, ): - effective_model = ( - model_name.strip() if model_name and model_name.strip() else model - ) + effective_model = model.strip() + if not effective_model: + raise ValueError("model cannot be empty — type a model name or wire one in") if history is None: history = [] messages = list(history) diff --git a/src/js/ollama.js b/src/js/ollama.js index c840535..bfe1248 100644 --- a/src/js/ollama.js +++ b/src/js/ollama.js @@ -1,24 +1,33 @@ /** * ollama.js — ComfyUI frontend extension for comfydv Ollama nodes. * - * Populates model COMBO widgets on OllamaModelSelector, OllamaLoadModel, and - * OllamaChatCompletion from a live call to GET /dv/ollama/models?host=. + * Populates the model widget on Ollama nodes from a live call to + * GET /dv/ollama/models?host=. + * + * OllamaModelSelector and OllamaLoadModel use a COMBO widget (dropdown). + * OllamaChatCompletion uses a plain STRING widget (accepts wired values). + * The Refresh button works the same way for both: it fetches the live list + * and sets the widget value / updates COMBO options as appropriate. */ import { app } from "../../scripts/app.js"; -const OLLAMA_COMBO_NODES = new Set([ - "OllamaModelSelector", - "OllamaLoadModel", - "OllamaChatCompletion", -]); +/** Nodes whose model widget is a COMBO dropdown. */ +const OLLAMA_COMBO_NODES = new Set(["OllamaModelSelector", "OllamaLoadModel"]); + +/** Nodes whose model widget is a plain STRING (accepts wired input). */ +const OLLAMA_STRING_MODEL_NODES = new Set(["OllamaChatCompletion"]); + +const OLLAMA_ALL_NODES = new Set([...OLLAMA_COMBO_NODES, ...OLLAMA_STRING_MODEL_NODES]); /** - * Fetch model list from the backend and repopulate the COMBO widget. - * @param {LGraphNode} node - * @param {string} host - Ollama host URL, e.g. "http://localhost:11434" + * Fetch model list and update the node's model widget. + * Works for both COMBO and STRING widgets: + * - COMBO: updates options.values + preserves selection if model still exists. + * - STRING: sets the value to the first model; keeps existing value if it + * still appears in the live list (user may have typed a valid name). */ -async function refreshModelDropdown(node, host) { +async function refreshModelWidget(node, host) { try { const resp = await fetch(`/dv/ollama/models?host=${encodeURIComponent(host)}`); if (!resp.ok) return; @@ -29,12 +38,22 @@ async function refreshModelDropdown(node, host) { const modelWidget = node.widgets?.find(w => w.name === "model"); if (!modelWidget) return; - const current = modelWidget.value; - modelWidget.options.values = models; - modelWidget.value = models.includes(current) ? current : models[0]; + const current = (modelWidget.value ?? "").trim(); + + if (Array.isArray(modelWidget.options?.values)) { + // COMBO widget + modelWidget.options.values = models; + modelWidget.value = models.includes(current) ? current : models[0]; + } else { + // STRING widget — keep current value if it is a known model + if (!current || !models.includes(current)) { + modelWidget.value = models[0]; + } + } + node.setDirtyCanvas(true, false); } catch (_) { - // Ollama unreachable — leave COMBO with server-side defaults + // Ollama unreachable — leave widget unchanged } } @@ -43,8 +62,7 @@ async function refreshModelDropdown(node, host) { * * First checks the node's own widgets (OllamaClient has a "host" widget). * Otherwise traverses graph links to find a connected OllamaClient node and - * reads its "host" widget — this is the common case for downstream nodes - * (ModelSelector, LoadModel, ChatCompletion) that receive the client socket. + * reads its "host" widget — this is the common case for downstream nodes. */ function getHostFromNode(node) { const ownHostWidget = node.widgets?.find(w => w.name === "host"); @@ -68,21 +86,21 @@ app.registerExtension({ name: "comfydv.ollama", async beforeRegisterNodeDef(nodeType, nodeData) { - if (!OLLAMA_COMBO_NODES.has(nodeData.name)) return; + if (!OLLAMA_ALL_NODES.has(nodeData.name)) return; const onNodeCreated = nodeType.prototype.onNodeCreated; nodeType.prototype.onNodeCreated = function () { const result = onNodeCreated?.apply(this, arguments); - // Add a refresh button below the model widget + // Add a Refresh button below the model widget this.addWidget("button", "⟳ Refresh models", null, () => { const host = getHostFromNode(this); - refreshModelDropdown(this, host); + refreshModelWidget(this, host); }); - // Initial population + // Initial population on node creation const host = getHostFromNode(this); - refreshModelDropdown(this, host); + refreshModelWidget(this, host); return result; }; diff --git a/tests/conftest.py b/tests/conftest.py index 117949b..0b57861 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -89,6 +89,25 @@ def skip_if_no_ollama(ollama_available): ) +@pytest.fixture(scope="session") +def first_generative_model(ollama_host, ollama_available): + """Return the first model available from Ollama, skipping embedding-only models. + + Used by lifecycle tests that call /api/generate — embedding models like + embeddinggemma reject that endpoint with HTTP 400. + """ + if not ollama_available: + pytest.skip("Ollama not reachable at localhost:11434") + import asyncio + + from comfydv.ollama import _fetch_models + + models = asyncio.run(_fetch_models(ollama_host)) + if not models: + pytest.skip("No models installed in Ollama") + return models[0] + + # --------------------------------------------------------------------------- # Existing ComfyUI node fixtures # --------------------------------------------------------------------------- diff --git a/tests/test_ollama.py b/tests/test_ollama.py index 7514403..2372692 100644 --- a/tests/test_ollama.py +++ b/tests/test_ollama.py @@ -374,12 +374,14 @@ class TestUS3ModelLifecycle: ) @pytest.mark.integration - def test_load_model_returns_name(self, ollama_host, skip_if_no_ollama): + def test_load_model_returns_name( + self, ollama_host, skip_if_no_ollama, first_generative_model + ): """Scenario: Load Model loads model into Ollama memory.""" (result,) = OllamaLoadModel().load_model( - client=ollama_host, model="embeddinggemma:latest" + client=ollama_host, model=first_generative_model ) - assert result == "embeddinggemma:latest" + assert result == first_generative_model def test_unload_returns_two_values(self): """OllamaUnloadModel returns (model_name, passthrough) tuple.""" @@ -396,15 +398,15 @@ class TestUS3ModelLifecycle: @pytest.mark.integration def test_unload_model_returns_name_and_passthrough( - self, ollama_host, skip_if_no_ollama + self, ollama_host, skip_if_no_ollama, first_generative_model ): """Scenario: Unload evicts model; passthrough flows through unchanged.""" model_name, passthrough = OllamaUnloadModel().unload_model( client=ollama_host, - model="embeddinggemma:latest", + model=first_generative_model, passthrough="sentinel", ) - assert model_name == "embeddinggemma:latest" + assert model_name == first_generative_model assert passthrough == "sentinel" @@ -521,7 +523,6 @@ class TestUS4ChatCompletion: assert isinstance(history, list) assert model_name == "m" - # ---- Issue 6: Chat timeout widget ----------------------------------------- def test_chat_has_timeout_secs_input(self): @@ -586,18 +587,25 @@ class TestUS4ChatCompletion: @pytest.mark.integration def test_multi_turn_receives_context(self, ollama_host, skip_if_no_ollama): - """Scenario: Multi-turn completion receives full conversation context.""" + """Scenario: Multi-turn completion receives full conversation context. + + Passes think=False to prevent Qwen3-family models from returning all + output as thinking tokens with an empty content field. + """ + no_think = {"think": False} _, history, _ = OllamaChatCompletion().chat( client=ollama_host, model=_CHAT_MODEL, prompt="My name is Alice. Remember it.", history=[], + options=no_think, )["result"] response, updated, _ = OllamaChatCompletion().chat( client=ollama_host, model=_CHAT_MODEL, prompt="What is my name?", history=history, + options=no_think, )["result"] assert "Alice" in response assert len(updated) == 4