ollama

Ollama model integration nodes for ComfyUI.

17 nodes: client configuration, auth headers, model discovery, load/unload, chat completion, composable inference options, and history utilities.

ADR-004: aiohttp (already in ComfyUI dep tree) for all Ollama HTTP calls. ADR-005: OllamaClient node is the single source of the host URL (and, per US7, of any auth headers — every downstream node reaches the same server). ADR-006: structured_output uses Ollama's OpenAI-compatible tool-calling endpoint + a dynamically-built pydantic model for validation, not pydantic-ai — still plain aiohttp, no new HTTP client.

OllamaChatCompletion

Source code in src/comfydv/ollama.py

class OllamaChatCompletion:
    OUTPUT_NODE = True

    _BASE_RETURN_TYPES = ("STRING", "OLLAMA_HISTORY", "STRING")
    _BASE_RETURN_NAMES = ("response", "updated_history", "model_name")

    # Per-node-instance structured-output config, keyed by unique_id — same
    # pattern as FormatString.node_configs.
    node_configs: dict = {}

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "client": ("OLLAMA_CLIENT",),
                # 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",),
                "timeout_secs": ("INT", {"default": 300, "min": 30, "max": 3600}),
                "structured_output": ("BOOLEAN", {"default": False}),
                "output_schema": (
                    "STRING",
                    {"multiline": True, "default": _DEFAULT_OUTPUT_SCHEMA},
                ),
                "max_retries": ("INT", {"default": 2, "min": 0, "max": 5}),
            },
            "hidden": {"unique_id": "UNIQUE_ID"},
        }

    RETURN_TYPES = _BASE_RETURN_TYPES
    RETURN_NAMES = _BASE_RETURN_NAMES
    FUNCTION = "chat"
    CATEGORY = "dv/ollama"

    @classmethod
    def update_outputs(
        cls, unique_id: str, structured_output: bool, schema: dict | None
    ) -> None:
        """Mutate class-level RETURN_TYPES/RETURN_NAMES for structured_output mode.

        Same class-level-shared-state pattern (and limitation) as
        FormatString.update_widget: RETURN_TYPES/RETURN_NAMES are shared
        across all instances of this node type in a graph, so the very first
        execution after toggling structured_output or editing output_schema
        may show stale downstream socket typing until it runs once.
        """
        cls.node_configs[unique_id] = {
            "structured_output": structured_output,
            "schema": schema,
        }
        if not structured_output or not schema:
            cls.RETURN_TYPES = cls._BASE_RETURN_TYPES
            cls.RETURN_NAMES = cls._BASE_RETURN_NAMES
            return
        names = tuple(schema["properties"].keys())
        cls.RETURN_TYPES = cls._BASE_RETURN_TYPES + _comfy_types_for_schema(schema)
        cls.RETURN_NAMES = cls._BASE_RETURN_NAMES + names

    def chat(
        self,
        client,
        model,
        prompt,
        system="",
        history=None,
        options=None,
        timeout_secs=300,
        structured_output=False,
        output_schema=_DEFAULT_OUTPUT_SCHEMA,
        max_retries=2,
        unique_id="",
    ):
        effective_model = model.strip()
        if not effective_model:
            raise ValueError("model cannot be empty — type a model name or wire one in")

        schema = None
        pydantic_model = None
        if structured_output:
            schema = _parse_output_schema(output_schema)  # fail fast, no network call
            pydantic_model = _build_structured_model(schema)

        if unique_id:
            type(self).update_outputs(unique_id, structured_output, schema)

        if history is None:
            history = []
        messages = list(history)
        if system:
            messages = [{"role": "system", "content": system}] + messages
        messages.append({"role": "user", "content": prompt})
        payload: dict = {
            "model": effective_model,
            "messages": messages,
            "stream": False,
        }
        if options:
            payload["options"] = options
        if structured_output:
            # ADR-006: OpenAI-compatible tool-calling, not native /api/chat
            # "format" — forces a single call whose arguments match schema.
            assert schema is not None  # structured_output implies this was parsed
            tools, tool_choice = _build_tool_call_payload(schema)
            payload["tools"] = tools
            payload["tool_choice"] = tool_choice

        headers = _client_headers(client)
        cache_key = _cache_key(
            "chat",
            client,
            headers or {},
            effective_model,
            messages,
            options or {},
            schema or {},
        )
        cached, hit = _CHAT_RESPONSE_CACHE.get(cache_key)
        url = (
            f"{client}/v1/chat/completions"
            if structured_output
            else f"{client}/api/chat"
        )

        parsed = None
        if not structured_output:
            # Unchanged from pre-structured-output behavior.
            if hit:
                response_text = cached
            else:
                result = _run_async(
                    _post_json(
                        url, payload, timeout=float(timeout_secs), headers=headers
                    )
                )
                response_text = result.get("message", {}).get("content", "")
                _CHAT_RESPONSE_CACHE.set(cache_key, response_text)
        elif hit:
            # schema is part of the cache key, so a hit was necessarily
            # validated against this exact schema when it was written —
            # re-parsing here is guaranteed-successful deserialization, not
            # a fallible re-check.
            assert (
                pydantic_model is not None
            )  # structured_output implies this was built
            response_text = cached
            parsed = pydantic_model.model_validate_json(response_text)
        else:
            assert (
                pydantic_model is not None
            )  # structured_output implies this was built
            total_attempts = max(0, min(int(max_retries), 5)) + 1
            response_text = None
            last_invalid_text = ""
            last_error = None
            for _attempt in range(1, total_attempts + 1):
                result = _run_async(
                    _post_json(
                        url, payload, timeout=float(timeout_secs), headers=headers
                    )
                )
                candidate = _extract_tool_call_arguments(result)
                try:
                    parsed = pydantic_model.model_validate_json(candidate)
                    response_text = candidate
                    _CHAT_RESPONSE_CACHE.set(cache_key, response_text)
                    break
                except ValidationError as exc:
                    last_invalid_text = candidate
                    last_error = exc
            else:
                raise RuntimeError(
                    f"OllamaChatCompletion: structured_output response failed "
                    f"validation against output_schema after {total_attempts} "
                    f"attempt(s) (model={effective_model!r}). Last error: "
                    f"{last_error}. Last response (truncated): "
                    f"{last_invalid_text[:300]!r}"
                )

        updated = list(history)
        updated.append({"role": "user", "content": prompt})
        updated.append({"role": "assistant", "content": response_text})
        n = len(updated)
        ui_text = (
            f"{response_text}\n\n── History: {n} message(s) ──\n{_history_preview(updated)}"
            if n > 2
            else response_text
        )

        result_tuple = (response_text, updated, effective_model)
        if structured_output:
            assert schema is not None  # structured_output implies this was parsed
            comfy_types = _comfy_types_for_schema(schema)
            extra = tuple(
                _coerce_structured_value(getattr(parsed, name), ctype)
                for name, ctype in zip(schema["properties"].keys(), comfy_types)
            )
            result_tuple += extra

        return {
            "ui": {"text": [ui_text]},
            "result": result_tuple,
        }

update_outputs(unique_id, structured_output, schema) classmethod

Mutate class-level RETURN_TYPES/RETURN_NAMES for structured_output mode.

Same class-level-shared-state pattern (and limitation) as FormatString.update_widget: RETURN_TYPES/RETURN_NAMES are shared across all instances of this node type in a graph, so the very first execution after toggling structured_output or editing output_schema may show stale downstream socket typing until it runs once.

Source code in src/comfydv/ollama.py

@classmethod
def update_outputs(
    cls, unique_id: str, structured_output: bool, schema: dict | None
) -> None:
    """Mutate class-level RETURN_TYPES/RETURN_NAMES for structured_output mode.

    Same class-level-shared-state pattern (and limitation) as
    FormatString.update_widget: RETURN_TYPES/RETURN_NAMES are shared
    across all instances of this node type in a graph, so the very first
    execution after toggling structured_output or editing output_schema
    may show stale downstream socket typing until it runs once.
    """
    cls.node_configs[unique_id] = {
        "structured_output": structured_output,
        "schema": schema,
    }
    if not structured_output or not schema:
        cls.RETURN_TYPES = cls._BASE_RETURN_TYPES
        cls.RETURN_NAMES = cls._BASE_RETURN_NAMES
        return
    names = tuple(schema["properties"].keys())
    cls.RETURN_TYPES = cls._BASE_RETURN_TYPES + _comfy_types_for_schema(schema)
    cls.RETURN_NAMES = cls._BASE_RETURN_NAMES + names

OllamaClientType

Bases: str

Typed string carrying the Ollama host URL through the node graph.

Also carries an optional .headers dict (US7 — basic auth / bearer tokens for Ollama servers behind a reverse proxy). Since this is a plain str subclass, every existing f"{client}/api/..." call site keeps working unchanged; only code that wants auth reads .headers.

Source code in src/comfydv/ollama.py

class OllamaClientType(str):
    """Typed string carrying the Ollama host URL through the node graph.

    Also carries an optional ``.headers`` dict (US7 — basic auth / bearer
    tokens for Ollama servers behind a reverse proxy). Since this is a plain
    ``str`` subclass, every existing ``f"{client}/api/..."`` call site keeps
    working unchanged; only code that wants auth reads ``.headers``.
    """

    def __new__(cls, host, headers=None):
        obj = super().__new__(cls, host)
        obj.headers = dict(headers) if headers else {}
        return obj

OllamaUnloadModel

Evict a model from VRAM and pass a value through unchanged.

Wire passthrough from a downstream node (e.g. the response output of OllamaChatCompletion) so ComfyUI executes this node after that node completes. The value is returned unchanged so the rest of the workflow can continue using it.

Source code in src/comfydv/ollama.py

class OllamaUnloadModel:
    """Evict a model from VRAM and pass a value through unchanged.

    Wire ``passthrough`` from a downstream node (e.g. the ``response`` output
    of OllamaChatCompletion) so ComfyUI executes this node *after* that node
    completes.  The value is returned unchanged so the rest of the workflow can
    continue using it.
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "client": ("OLLAMA_CLIENT",),
                "model": ("STRING", {"forceInput": True}),
            },
            "optional": {
                "passthrough": ("STRING", {"forceInput": True}),
            },
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("model_name", "passthrough")
    FUNCTION = "unload_model"
    CATEGORY = "dv/ollama"

    def unload_model(self, client, model: str, passthrough: str = ""):
        if not model.strip():
            raise ValueError("model name cannot be empty")
        _run_async(
            _post_json(
                f"{client}/api/generate",
                {"model": model, "keep_alive": 0, "stream": False},
                timeout=30.0,
                headers=_client_headers(client),
            )
        )
        return (model, passthrough)