From 7cbb9fb7c32ddd60507aa08c66c29924240c054d Mon Sep 17 00:00:00 2001 From: aanokh Date: Thu, 3 Sep 2026 09:48:59 -0700 Subject: [PATCH 1/9] feat: alora adapters for ollama Signed-off-by: aanokh --- docs/docs/advanced/intrinsics.md | 9 +- docs/docs/advanced/lora-and-alora-adapters.md | 4 +- .../examples/intrinsics/uncertainty_ollama.py | 35 ++ mellea/backends/ollama.py | 346 +++++++++++++++- test/backends/test_ollama_intrinsics_unit.py | 387 ++++++++++++++++++ 5 files changed, 767 insertions(+), 14 deletions(-) create mode 100644 docs/examples/intrinsics/uncertainty_ollama.py create mode 100644 test/backends/test_ollama_intrinsics_unit.py diff --git a/docs/docs/advanced/intrinsics.md b/docs/docs/advanced/intrinsics.md index 5147123653..3ebf1af55f 100644 --- a/docs/docs/advanced/intrinsics.md +++ b/docs/docs/advanced/intrinsics.md @@ -14,7 +14,7 @@ Adapter functions are adapter-accelerated operations for RAG quality checks. The LoRA/aLoRA adapters loaded directly into the Hugging Face backend — faster and more reliable than prompting a general-purpose model for these specialized micro-tasks. -> **Backend note:** Adapter functions work with two backends: +> **Backend note:** Adapter functions work with three backends: > > - **LocalHFBackend** — loads LoRA/aLoRA adapters from the catalog at runtime. > A local Granite Switch checkpoint can instead use @@ -25,8 +25,13 @@ reliable than prompting a general-purpose model for these specialized micro-task > `load_embedded_adapters=True`. Only adapter functions embedded in the model are > available — check the model's `adapter_index.json` for the list. > See `docs/docs/examples/granite-switch/README.md` +> - **OllamaModelBackend**: uses an Ollama model that bundles the adapter, +> for example `gabegoodhart/granite4.1-uncertainty:3b`, which is `granite4.1:3b` +> plus the uncertainty aLoRA. Ollama bundles one adapter per model, so pass +> `adapter_models={"uncertainty": "", ...}` to route each adapter function +> to its model. Install `mellea[switch]` to download the adapter's `io.yaml`. > -> Adapter functions do not work with Ollama or other remote backends. +> Adapter functions do not work with other remote backends. Set up the backend once and reuse it across adapter function calls: diff --git a/docs/docs/advanced/lora-and-alora-adapters.md b/docs/docs/advanced/lora-and-alora-adapters.md index f52b69fefd..2da57a2ff5 100644 --- a/docs/docs/advanced/lora-and-alora-adapters.md +++ b/docs/docs/advanced/lora-and-alora-adapters.md @@ -16,7 +16,9 @@ Apple Silicon Mac with sufficient VRAM for the chosen base model. Uploading requ Hugging Face account. > **Backend note:** Custom-trained adapters can only be loaded into `LocalHFBackend`. -> They do not work with Ollama, OpenAI, or other remote backends. +> They do not work with OpenAI or other remote backends. To use one with Ollama, +> convert it to GGUF and bundle it into an Ollama model with a Modelfile `ADAPTER` +> line, then pass that model tag via `OllamaModelBackend(adapter_models=...)`. > > Granite Switch models ship with pre-trained adapter functions embedded in the > model weights. Use them through `OpenAIBackend` with a served checkpoint, or diff --git a/docs/examples/intrinsics/uncertainty_ollama.py b/docs/examples/intrinsics/uncertainty_ollama.py new file mode 100644 index 0000000000..c00cf8180e --- /dev/null +++ b/docs/examples/intrinsics/uncertainty_ollama.py @@ -0,0 +1,35 @@ +# pytest: e2e, ollama, qualitative + +"""Example usage of the uncertainty/certainty intrinsic with Ollama. + +Evaluates how certain the model is about its response to a user question. +The context should contain a user question followed by an assistant answer. + +Ollama bundles one adapter per model, so the uncertainty adapter is served by +its own model tag (`granite4.1:3b` plus the uncertainty aLoRA). Pass that tag +via `adapter_models`; normal chat still uses the base model. + +Requires `mellea[switch]` to download the adapter's `io.yaml`. + +To run this script from the root of the Mellea source tree, use the command: +``` +uv run python docs/examples/intrinsics/uncertainty_ollama.py +``` +""" + +from mellea import model_ids, start_backend +from mellea.stdlib import functional as mfuncs +from mellea.stdlib.components.intrinsic import core + +ctx, backend = start_backend( + "ollama", + model_id=model_ids.IBM_GRANITE_4_1_3B, + context_type="chat", + adapter_models={"uncertainty": "gabegoodhart/granite4.1-uncertainty:3b"}, +) + +response, ctx = mfuncs.chat("What is 2 + 2?", ctx, backend) # type: ignore +print(f"Response: {response.content}") + +result = core.check_certainty(ctx, backend) # type: ignore +print(f"Certainty score: {result}") diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index d555706ca1..8cac1b5214 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -30,18 +30,21 @@ RawProviderResponse, ) from ..core.base import AbstractMelleaTool -from ..formatters import ChatFormatter, TemplateFormatter +from ..formatters import ChatFormatter, TemplateFormatter, granite as granite_formatters from ..helpers import ( DEFAULT_CHUNK_TIMEOUT, ClientCache, get_current_event_loop, merge_provider_fields, + message_to_openai_message, + messages_to_docs, send_to_queue, should_replay_reasoning, ) -from ..stdlib.components import Message -from ..stdlib.requirements import ALoraRequirement +from ..stdlib.components import Intrinsic, Message +from ..stdlib.requirements import ALoraRequirement, LLMaJRequirement, Requirement from ..telemetry.context import generate_request_id, with_context +from .adapters.adapter import AdapterInput, AdapterMixin, IntrinsicAdapter from .backend import FormatterBackend from .model_options import ModelOption from .tools import add_tools_from_context_actions, add_tools_from_model_options @@ -49,6 +52,46 @@ format: None = None # typing this variable in order to shadow the global format function and ensure mypy checks for errors +def _to_chat_completion_dict(response: ollama.ChatResponse) -> dict: + """Convert an Ollama chat response into the OpenAI-shaped dict the intrinsic result processor reads. + + Args: + response: A non-streaming Ollama chat response. + + Returns: + dict: A chat completion dict with one choice, carrying the message content + and, when present, per-token logprobs with their top alternatives. + """ + logprobs = None + if response.logprobs: + logprobs = { + "content": [ + { + "token": lp.token, + "logprob": lp.logprob, + "top_logprobs": [ + {"token": t.token, "logprob": t.logprob} + for t in lp.top_logprobs or [] + ], + } + for lp in response.logprobs + ] + } + return { + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": response.message.content or "", + }, + "logprobs": logprobs, + "finish_reason": response.done_reason, + } + ] + } + + def _strip_data_uri_prefix(images: list[str]) -> list[str]: """Strip data URI prefix from base64 image strings for Ollama. @@ -138,7 +181,7 @@ def _to_ollama_tool_calls(openai_tool_calls: list[dict[str, Any]]) -> list[dict] return translated -class OllamaModelBackend(FormatterBackend): +class OllamaModelBackend(FormatterBackend, AdapterMixin): """A model that uses the Ollama Python SDK for local inference. Args: @@ -160,6 +203,11 @@ class OllamaModelBackend(FormatterBackend): bounds the wait between consecutive chunks; for non-streaming requests it bounds total time-to-response. Pass `None` to use the upstream `ollama` SDK default (no timeout). + adapter_models (dict[str, str] | None): Mapping from adapter function name + (e.g. `"uncertainty"`) to the Ollama model tag that bundles that + adapter (e.g. `"gabegoodhart/granite4.1-uncertainty:3b"`). Ollama + bundles one adapter per model, so each adapter function is served by + its own tag. Adapter functions not listed here run against `model_id`. Attributes: to_mellea_model_opts_map (dict): Mapping from Ollama-specific option names @@ -180,6 +228,7 @@ def __init__( base_url: str | None = None, model_options: dict | None = None, timeout: float | None = 300.0, + adapter_models: dict[str, str] | None = None, ): """Initialize an Ollama backend, connecting to the server and pulling the model if needed.""" super().__init__( @@ -204,6 +253,9 @@ def __init__( self._model_id: str = ollama_model_id self._provider: str = "ollama" + self._added_adapters: dict[str, IntrinsicAdapter] = {} + self._adapter_models: dict[str, str] = adapter_models or {} + # Setup the client and ensure that we have the model available. self._base_url = base_url self._timeout = timeout @@ -256,6 +308,55 @@ def __init__( ModelOption.STOP_SEQUENCES: "stop", } + @property + def base_model_name(self) -> str: + """Return the short base model name used for adapter config lookup. + + Adapter configs are laid out by Hugging Face model name, so a known + Ollama tag is mapped back to it (e.g. `"granite4.1:3b"` to + `"granite-4.1-3b"`). Unknown tags are returned unchanged. + + Returns: + str: The short base model name. + """ + for ident in vars(model_ids).values(): + if ( + isinstance(ident, ModelIdentifier) + and ident.ollama_name == self._model_id + and ident.hf_model_name + ): + return ident.hf_model_name.split("/")[-1] + return self._model_id + + def add_adapter(self, adapter: AdapterInput) -> None: + """Register an adapter with this backend. + + Ollama serves adapter weights bundled into a model, so only the + adapter's I/O config is used here; no weights are loaded. + + Args: + adapter (AdapterInput): The adapter to register. Must be an + `IntrinsicAdapter`. + + Raises: + TypeError: If `adapter` is not an `IntrinsicAdapter`. + """ + if not isinstance(adapter, IntrinsicAdapter): + raise TypeError( + f"OllamaModelBackend currently only supports IntrinsicAdapter. " + f"Got: {type(adapter).__name__}" + ) + adapter.backend = self + self._added_adapters[adapter.qualified_name] = adapter + + def list_adapters(self) -> list[str]: + """Return qualified names of all registered adapters. + + Returns: + list[str]: Qualified adapter names. + """ + return list(self._added_adapters.keys()) + def _check_ollama_server(self) -> bool: """Requests generic info about the Ollama server to ensure it's running.""" try: @@ -394,6 +495,191 @@ def _make_backend_specific_and_remove( ) return ModelOption.remove_special_keys(backend_specific) + async def _generate_from_intrinsic( + self, + action: Intrinsic, + ctx: Context, + *, + model_options: dict[str, Any], + tool_calls: bool = False, + ) -> ModelOutputThunk: + """Generate a completion for an intrinsic action via the Ollama chat API. + + Applies the intrinsic's I/O rewriter to transform the conversation, sends + it to the Ollama model that bundles the adapter, and post-processes the + model output through the intrinsic's result processor. The adapter is + activated by the instruction text the rewriter appends to the + conversation, so no extra request field is needed. + + Intrinsics default to options provided by `io.yaml`. Model options + override these defaults. All model options besides streaming are + respected. + + Args: + action (Intrinsic): The intrinsic component to execute. + ctx (Context): The current generation context (must be a chat context). + model_options (dict[str, Any]): Merged model options for this call. + tool_calls (bool): If `True`, expose available tools to the model + and parse tool-call responses. + + Returns: + ModelOutputThunk: A thunk that lazily resolves to the processed + intrinsic output. + + Raises: + NotImplementedError: If the context isn't a chat context, or if + streaming is requested (intrinsic post-processing requires + the complete response). + ValueError: If no adapter is registered for the requested intrinsic. + TypeError: If the adapter isn't an `IntrinsicAdapter`. + """ + if not ctx.is_chat_context: + raise NotImplementedError("Intrinsics require a chat context.") + + # Intrinsics don't support streaming because of their post-processing step. + if model_options.get(ModelOption.STREAM, False): + raise NotImplementedError( + "Intrinsics do not support streaming due to structured output parsing." + ) + + allowed_types = tuple(at.value for at in action.adapter_types) + adapter = self._find_adapter(action.intrinsic_name, allowed_types) + if adapter is None: + raise ValueError( + f"backend ({self}) has no adapter for processing adapter function: " + f"{action.intrinsic_name}" + ) + if not isinstance(adapter, IntrinsicAdapter): + raise TypeError( + f"OllamaModelBackend only supports IntrinsicAdapter, got: {type(adapter).__name__}" + ) + + intrinsic_config = adapter.config + assert intrinsic_config is not None + + rewriter = granite_formatters.IntrinsicsRewriter( + config_dict=intrinsic_config, model_name=adapter.name + ) + result_processor = granite_formatters.IntrinsicsResultProcessor( + config_dict=intrinsic_config + ) + + linearized_context = ctx.view_for_generation() + assert linearized_context is not None, ( + "If ctx.is_chat_context, then the context should be linearizable." + ) + + # NOTE: Explicitly do not add the action to the context here. + # Intrinsics modify the context through their rewriters. + messages: list[Message] = self.formatter.to_chat_messages(linearized_context) + + system_prompt = model_options.get(ModelOption.SYSTEM_PROMPT, "") + conversation: list[dict] = [] + if system_prompt != "": + conversation.append({"role": "system", "content": system_prompt}) + conversation.extend( + [message_to_openai_message(m, provider=self._provider) for m in messages] + ) + + docs = messages_to_docs(messages) + + request_json: dict = { + "messages": conversation, + "extra_body": {"documents": docs}, + } + + rewritten = rewriter.transform(request_json, **action.intrinsic_kwargs) + + tools: dict[str, AbstractMelleaTool] = dict() + if tool_calls: + add_tools_from_model_options(tools, model_options) + add_tools_from_context_actions(tools, ctx.actions_for_available_tools()) + MelleaLogger.get_logger().info(f"Tools for call: {tools.keys()}") + + # io.yaml parameters are defaults, user model options override them + params = dict(rewriter.parameters) + params.pop("model", None) + if "max_completion_tokens" in params: + params[ModelOption.MAX_NEW_TOKENS] = params.pop("max_completion_tokens") + model_opts = ModelOption.merge_model_options(params, model_options) + logprobs = model_opts.pop("logprobs", None) + top_logprobs = model_opts.pop("top_logprobs", None) + + # each adapter function is served by its own ollama model tag + model = self._adapter_models.get(action.intrinsic_name, self._model_id) + + messages_dicts = [] + for m in rewritten.messages: + d = m.model_dump(exclude_unset=True) + if "role" not in d: + d["role"] = m.role + messages_dicts.append(d) + + chat_response: Coroutine[Any, Any, ollama.ChatResponse] = ( + self._async_client.chat( + model=model, + messages=messages_dicts, + tools=[t.as_json_tool for t in tools.values()], + think=model_opts.get(ModelOption.THINKING, None), + stream=False, + options=self._make_backend_specific_and_remove(model_opts), + format=rewriter.config["response_format"], + logprobs=logprobs, + top_logprobs=top_logprobs, + ) + ) # type: ignore + + output = ModelOutputThunk(None) + output._gen.start = datetime.datetime.now() + output._call.context = linearized_context + output._call.action = action + output._call.model_options = model_opts + + async def granite_formatters_processing( + mot: ModelOutputThunk, + chunk: ollama.ChatResponse, + rewritten: granite_formatters.ChatCompletion, + result_processor: granite_formatters.IntrinsicsResultProcessor, + ): + """Accumulate content and apply intrinsic result processing.""" + await self.processing(mot, chunk, tools=tools) + + try: + res = result_processor.transform( + _to_chat_completion_dict(chunk), rewritten + ) + except json.JSONDecodeError as e: + raise Exception( + f"Intrinsic did not return a JSON: {chunk.message.content}" + ) from e + + mot._underlying_value = res.choices[0].message.content + + output._gen.process = functools.partial( + granite_formatters_processing, + rewritten=rewritten, + result_processor=result_processor, + ) + output._gen.post_process = functools.partial( + self.post_processing, conversation=messages_dicts, tools=tools, _format=None + ) + + output.generation.model = model + output.generation.provider = self._provider + + output._gen.generate = asyncio.create_task( + send_to_queue( + chat_response, + output._gen.queue, + chunk_timeout=model_opts.get( + ModelOption.STREAM_TIMEOUT, DEFAULT_CHUNK_TIMEOUT + ), + ) + ) + output._gen.generate_type = GenerateType.ASYNC + + return output + async def _generate_from_context( self, action: Component[C] | CBlock | ModelOutputThunk, @@ -428,6 +714,50 @@ async def _generate_from_context( _model_id_str = str(getattr(self, "model_id", "unknown")) with with_context(request_id=generate_request_id(), model_id=_model_id_str): + model_opts = self._simplify_and_merge(model_options) + + # Requirements can be automatically rerouted to a requirement adapter. + if isinstance(action, Requirement): + reroute_to_alora = isinstance(action, ALoraRequirement) + adapter_name = "requirement-check" + + if isinstance(action, ALoraRequirement): + adapter_name = action.intrinsic_name + alora_action = action + else: + assert action.description is not None, ( + "must have a description when generating from a requirement" + ) + alora_action = ALoraRequirement(action.description, adapter_name) + + alora_req_adapter = self._find_adapter(adapter_name, ("alora",)) + if alora_req_adapter is None: + if reroute_to_alora: + MelleaLogger.get_logger().warning( + f"attempted to use an AloraRequirement but backend {self} " + f"doesn't have the specified adapter added {adapter_name}; " + f"defaulting to regular generation" + ) + reroute_to_alora = False + + if issubclass(type(action), LLMaJRequirement): + reroute_to_alora = False + + if reroute_to_alora: + mot = await self._generate_from_intrinsic( + alora_action, + ctx, + model_options=model_opts, + tool_calls=tool_calls, + ) + return mot, ctx.add(alora_action).add(mot) + + elif isinstance(action, Intrinsic): + mot = await self._generate_from_intrinsic( + action, ctx, model_options=model_opts, tool_calls=tool_calls + ) + return mot, ctx.add(action).add(mot) + mot = await self.generate_from_chat_context( action, ctx, @@ -487,13 +817,7 @@ async def generate_from_chat_context( # Convert our linearized context into a sequence of chat messages. Template formatters have a standard way of doing this. messages: list[Message] = self.formatter.to_chat_messages(linearized_context) # Add the final message. - match action: - case ALoraRequirement(): - raise Exception( - "The ollama backend does not currently support aLoRA adapters." - ) - case _: - messages.extend(self.formatter.to_chat_messages([action])) + messages.extend(self.formatter.to_chat_messages([action])) # construct the conversation from our messages, adding a system prompt at the first message if one was provided. conversation: list[dict] = [] # We use system prompt None/empty-string semantics in a way that is consistent with Hugging Face and other libraries. diff --git a/test/backends/test_ollama_intrinsics_unit.py b/test/backends/test_ollama_intrinsics_unit.py new file mode 100644 index 0000000000..cda3ea2d7c --- /dev/null +++ b/test/backends/test_ollama_intrinsics_unit.py @@ -0,0 +1,387 @@ +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for the Ollama backend intrinsic generation path. No server required. + +Mocks the Ollama async client to verify that `_generate_from_intrinsic` correctly: +- appends the io.yaml instruction (the adapter's activation text) as the last message +- passes the io.yaml response schema as `format` and requests logprobs +- routes the call to the Ollama model tag registered for the adapter function +- applies the `IntrinsicsResultProcessor` to the raw response +- user-provided model options override io.yaml parameter defaults +- raises when no adapter is registered or streaming is requested +""" + +import json +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch + +import ollama +import pytest + +from mellea.backends import ModelOption +from mellea.backends.adapters.adapter import AdapterType, IntrinsicAdapter +from mellea.backends.ollama import OllamaModelBackend, _to_chat_completion_dict +from mellea.stdlib import functional as mfuncs +from mellea.stdlib.components import Intrinsic, Message +from mellea.stdlib.context import ChatContext + +# --------------------------------------------------------------------------- +# Configs +# --------------------------------------------------------------------------- + +_SCORE_SCHEMA = { + "type": "object", + "properties": {"score": {"type": "string", "enum": [str(i) for i in range(10)]}}, + "required": ["score"], + "additionalProperties": False, +} + +# Minimal config: no transformations, no logprobs. Good enough for tests that +# only inspect the API call. +_SIMPLE_CONFIG = { + "model": None, + "response_format": _SCORE_SCHEMA, + "transformations": None, + "instruction": "", + "parameters": {"max_completion_tokens": 64, "temperature": 1.0}, + "sentence_boundaries": None, +} + +# Mirrors the real uncertainty io.yaml: likelihood + project transformations. +_UNCERTAINTY_CONFIG = { + "model": None, + "response_format": _SCORE_SCHEMA, + "transformations": [ + { + "type": "likelihood", + "categories_to_values": {str(i): 0.1 * i + 0.05 for i in range(10)}, + "input_path": ["score"], + }, + { + "type": "project", + "input_path": [], + "retained_fields": {"score": "certainty"}, + }, + ], + "instruction": "", + "parameters": {"max_completion_tokens": 15, "temperature": 0.0}, + "sentence_boundaries": None, +} + +_ADAPTER_TAG = "gabegoodhart/granite4.1-uncertainty:3b" + +# --------------------------------------------------------------------------- +# Canned responses +# --------------------------------------------------------------------------- + + +def _simple_chat_response(content: str = '{"score": "9"}') -> ollama.ChatResponse: + """Build a minimal ChatResponse with no logprobs.""" + return ollama.ChatResponse.model_validate( + { + "model": _ADAPTER_TAG, + "message": {"role": "assistant", "content": content}, + "done": True, + "done_reason": "stop", + } + ) + + +def _uncertainty_chat_response() -> ollama.ChatResponse: + """Build a ChatResponse that the uncertainty result processor can parse. + + The likelihood transformation reads top_logprobs to compute an expected value. + """ + return ollama.ChatResponse.model_validate( + { + "model": _ADAPTER_TAG, + "message": {"role": "assistant", "content": '{"score": "9"}'}, + "done": True, + "done_reason": "stop", + "logprobs": [ + { + "token": '{"', + "logprob": 0.0, + "top_logprobs": [{"token": '{"', "logprob": 0.0}], + }, + { + "token": "score", + "logprob": 0.0, + "top_logprobs": [{"token": "score", "logprob": 0.0}], + }, + { + "token": '":', + "logprob": 0.0, + "top_logprobs": [{"token": '":', "logprob": 0.0}], + }, + { + "token": ' "', + "logprob": 0.0, + "top_logprobs": [{"token": ' "', "logprob": 0.0}], + }, + { + "token": "9", + "logprob": -0.05, + "top_logprobs": [ + {"token": "9", "logprob": -0.05}, + {"token": "4", "logprob": -3.0}, + ], + }, + { + "token": '"}', + "logprob": 0.0, + "top_logprobs": [{"token": '"}', "logprob": 0.0}], + }, + ], + } + ) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +def _make_backend( + *, model_options: dict | None = None, adapter_models: dict | None = None +) -> OllamaModelBackend: + """Return an OllamaModelBackend with all network calls patched out.""" + with ( + patch.object(OllamaModelBackend, "_check_ollama_server", return_value=True), + patch.object(OllamaModelBackend, "_pull_ollama_model", return_value=True), + patch("mellea.backends.ollama.ollama.Client", return_value=MagicMock()), + patch("mellea.backends.ollama.ollama.AsyncClient", return_value=MagicMock()), + ): + return OllamaModelBackend( + model_id="granite4.1:3b", + model_options=model_options, + adapter_models=adapter_models, + ) + + +def _make_backend_with_adapter( + config: dict, + *, + model_options: dict | None = None, + adapter_models: dict | None = None, +) -> OllamaModelBackend: + """Return an OllamaModelBackend with a registered uncertainty adapter.""" + backend = _make_backend(model_options=model_options, adapter_models=adapter_models) + adapter = IntrinsicAdapter( + "uncertainty", adapter_type=AdapterType.LORA, config_dict=config + ) + backend.add_adapter(adapter) + return backend + + +def _make_context() -> ChatContext: + """Return a simple two-turn chat context.""" + return ( + ChatContext() + .add(Message("user", "What is the square root of 4?")) + .add(Message("assistant", "The square root of 4 is 2.")) + ) + + +async def _run_intrinsic( + backend: OllamaModelBackend, response: ollama.ChatResponse, **kwargs +): + """Run the uncertainty intrinsic against a mocked client; return (mot, mock_chat).""" + mock_chat = AsyncMock(return_value=response) + mock_client = MagicMock() + mock_client.chat = mock_chat + + with patch.object( + OllamaModelBackend, + "_async_client", + new_callable=PropertyMock, + return_value=mock_client, + ): + mot, _ = await mfuncs.aact( + Intrinsic("uncertainty"), _make_context(), backend, strategy=None, **kwargs + ) + await mot.avalue() + return mot, mock_chat + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +async def test_instruction_appended_as_last_message(): + """The io.yaml instruction (the adapter activation text) is the final user message.""" + backend = _make_backend_with_adapter(_SIMPLE_CONFIG) + _, mock_chat = await _run_intrinsic(backend, _simple_chat_response()) + + mock_chat.assert_called_once() + messages = mock_chat.call_args.kwargs["messages"] + assert messages[-1]["role"] == "user" + assert messages[-1]["content"] == "" + assert messages[0]["content"] == "What is the square root of 4?" + + +async def test_format_and_logprobs_requested(): + """The io.yaml response schema is passed as `format`; likelihood rules request logprobs.""" + backend = _make_backend_with_adapter(_UNCERTAINTY_CONFIG) + _, mock_chat = await _run_intrinsic(backend, _uncertainty_chat_response()) + + call_kwargs = mock_chat.call_args.kwargs + assert call_kwargs["format"] == _SCORE_SCHEMA + assert call_kwargs["logprobs"] is True + assert call_kwargs["top_logprobs"] == 10 + assert call_kwargs["stream"] is False + + +async def test_adapter_model_tag_used(): + """The call goes to the Ollama model registered for the adapter function.""" + backend = _make_backend_with_adapter( + _SIMPLE_CONFIG, adapter_models={"uncertainty": _ADAPTER_TAG} + ) + _, mock_chat = await _run_intrinsic(backend, _simple_chat_response()) + + assert mock_chat.call_args.kwargs["model"] == _ADAPTER_TAG + + +async def test_adapter_model_tag_defaults_to_model_id(): + """Adapter functions without a registered tag run against the backend's model.""" + backend = _make_backend_with_adapter(_SIMPLE_CONFIG) + _, mock_chat = await _run_intrinsic(backend, _simple_chat_response()) + + assert mock_chat.call_args.kwargs["model"] == "granite4.1:3b" + + +async def test_result_processor_applied(): + """Full uncertainty config: likelihood + project transforms produce the expected JSON.""" + backend = _make_backend_with_adapter(_UNCERTAINTY_CONFIG) + mot, _ = await _run_intrinsic(backend, _uncertainty_chat_response()) + + parsed = json.loads(mot.value) + assert list(parsed.keys()) == ["certainty"] + score = parsed["certainty"] + assert isinstance(score, float) + # Expected value over {9: 0.95, 4: 0.45} weighted by exp(logprob); 9 dominates. + assert 0.9 < score < 0.95 + + +async def test_io_yaml_parameters_forwarded(): + """io.yaml max_completion_tokens and temperature reach Ollama's options.""" + backend = _make_backend_with_adapter(_SIMPLE_CONFIG) + _, mock_chat = await _run_intrinsic(backend, _simple_chat_response()) + + options = mock_chat.call_args.kwargs["options"] + assert options["num_predict"] == 64 + assert options["temperature"] == 1.0 + + +async def test_model_options_override_io_yaml_defaults(): + """User-provided temperature overrides the io.yaml default; other defaults remain.""" + backend = _make_backend_with_adapter(_SIMPLE_CONFIG) + _, mock_chat = await _run_intrinsic( + backend, + _simple_chat_response(), + model_options={ModelOption.TEMPERATURE: 0.5, ModelOption.SEED: 42}, + ) + + options = mock_chat.call_args.kwargs["options"] + assert options["temperature"] == 0.5 + assert options["seed"] == 42 + assert options["num_predict"] == 64 + + +async def test_no_adapter_raises_valueerror(): + """Calling an intrinsic with no registered adapter raises ValueError.""" + backend = _make_backend() + + with pytest.raises(ValueError, match="has no adapter"): + await mfuncs.aact( + Intrinsic("uncertainty"), _make_context(), backend, strategy=None + ) + + +async def test_streaming_raises(): + """Intrinsics do not support streaming, so this raises NotImplementedError.""" + backend = _make_backend_with_adapter(_SIMPLE_CONFIG) + + with pytest.raises(NotImplementedError, match="do not support streaming"): + await mfuncs.aact( + Intrinsic("uncertainty"), + _make_context(), + backend, + strategy=None, + model_options={ModelOption.STREAM: True}, + ) + + +async def test_tools_passed_to_api(): + """Tools are forwarded to the chat call when tool_calls=True.""" + from mellea.backends.tools import MelleaTool + + def get_temperature(location: str) -> int: + """Returns the temperature of a city. + + Args: + location: A city name. + """ + return 21 + + backend = _make_backend_with_adapter(_SIMPLE_CONFIG) + _, mock_chat = await _run_intrinsic( + backend, + _simple_chat_response(), + tool_calls=True, + model_options={ModelOption.TOOLS: [MelleaTool.from_callable(get_temperature)]}, + ) + + tools = mock_chat.call_args.kwargs["tools"] + assert len(tools) == 1 + assert tools[0]["function"]["name"] == "get_temperature" + + +# --------------------------------------------------------------------------- +# Adapter registration and response conversion +# --------------------------------------------------------------------------- + + +def test_add_adapter_registers_intrinsic_adapter(): + backend = _make_backend_with_adapter(_SIMPLE_CONFIG) + assert backend.list_adapters() == ["uncertainty_lora"] + + +def test_add_adapter_rejects_other_adapter_types(): + backend = _make_backend() + with pytest.raises(TypeError, match="only supports IntrinsicAdapter"): + backend.add_adapter(object()) # type: ignore[arg-type] + + +def test_base_model_name_maps_ollama_tag_to_hf_name(): + backend = _make_backend() + assert backend.base_model_name == "granite-4.1-3b" + + +def test_base_model_name_unknown_tag_unchanged(): + with ( + patch.object(OllamaModelBackend, "_check_ollama_server", return_value=True), + patch.object(OllamaModelBackend, "_pull_ollama_model", return_value=True), + patch("mellea.backends.ollama.ollama.Client", return_value=MagicMock()), + patch("mellea.backends.ollama.ollama.AsyncClient", return_value=MagicMock()), + ): + backend = OllamaModelBackend(model_id="someone/custom:3b") + assert backend.base_model_name == "someone/custom:3b" + + +def test_to_chat_completion_dict_with_logprobs(): + result = _to_chat_completion_dict(_uncertainty_chat_response()) + + choice = result["choices"][0] + assert choice["message"]["content"] == '{"score": "9"}' + assert choice["finish_reason"] == "stop" + digit = choice["logprobs"]["content"][4] + assert digit["token"] == "9" + assert digit["top_logprobs"][1] == {"token": "4", "logprob": -3.0} + + +def test_to_chat_completion_dict_without_logprobs(): + result = _to_chat_completion_dict(_simple_chat_response()) + + assert result["choices"][0]["logprobs"] is None From 9427b4f50586f8fdf6538cfb1ffccbccb6de6925 Mon Sep 17 00:00:00 2001 From: Nigel Jones Date: Tue, 8 Sep 2026 08:48:12 +0100 Subject: [PATCH 2/9] feat(backends): support Ollama adapter functions Assisted-by: Codex Signed-off-by: Nigel Jones --- .github/workflows/quality.yml | 2 +- AGENTS.md | 8 +- docs/docs/advanced/intrinsics.md | 66 +++++++- docs/docs/advanced/lora-and-alora-adapters.md | 11 +- .../tutorials/04-making-agents-reliable.md | 12 +- .../examples/intrinsics/uncertainty_ollama.py | 10 +- mellea/backends/adapters/adapter.py | 22 +-- mellea/backends/ollama.py | 153 +++++++++++++++--- test/backends/test_ollama.py | 92 ++++++++++- test/backends/test_ollama_intrinsics_unit.py | 106 +++++++++++- .../build_ollama_uncertainty_adapter.sh | 106 ++++++++++++ .../scripts/run_tests_with_ollama_and_vllm.sh | 6 + 12 files changed, 529 insertions(+), 65 deletions(-) create mode 100755 test/scripts/build_ollama_uncertainty_adapter.sh diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml index 1ec9eec6a9..a1d116260a 100644 --- a/.github/workflows/quality.yml +++ b/.github/workflows/quality.yml @@ -136,7 +136,7 @@ jobs: # projector Ollama needs for image input). Drop the hf.co prefix if the # model is ever published to the Ollama library directly. run: | - for model in granite4.2:3b granite4:micro-h hf.co/ibm-granite/granite-vision-4.1-4b-GGUF:Q4_K_M; do + for model in granite4.2:3b granite4:micro-h granite4.1:3b hf.co/ibm-granite/granite-vision-4.1-4b-GGUF:Q4_K_M; do pulled=false for i in 1 2 3 4 5; do ollama pull "$model" && { pulled=true; break; } diff --git a/AGENTS.md b/AGENTS.md index f46016c90e..aa9b88147f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -226,15 +226,17 @@ For lower-level control (custom adapters, model options), use `mfuncs.act()` wit ### Weights binding shapes -`Adapter.weights` normalizes each deployment's activation mechanism behind one of -two shapes — a `WeightsBinding` lifecycle for weights you stage yourself, or -`EmbeddedBinding.apply_activation` for weights already in the served model. The +`Adapter.weights` normalizes each deployment's activation mechanism behind three +shapes — a `WeightsBinding` lifecycle for weights you stage yourself, +`EmbeddedBinding.apply_activation` for weights already in the served model, or +`ServerMediatedBinding` for a model tag selected by the provider. The post-activation shape each produces: | Binding | Reality | Lifecycle verbs | Caller invokes | Normalized post-activation state | |---------|---------|------------------|-----------------|-----------------------------------| | `LocalFileBinding` | LocalFile/PEFT | `prepare` / `activate` / `deactivate` / `release` | `activate()` / `deactivate()`, via `adapter_scope` | Backend-internal PEFT adapter state toggled; the outgoing request is untouched | | `EmbeddedBinding` | Embedded/Granite Switch | none — weights are already in the served model | `apply_activation(request, identity)` | `request.extra_body["chat_template_kwargs"]["adapter_name"]` set; `request.api_params["model"]` removed if present | +| `ServerMediatedBinding` | Ollama bundled adapter model | none for the current Ollama path | select the configured model tag during intrinsic generation | Ollama request's `model` is the bundled adapter tag; full lifecycle telemetry remains follow-up work | ### Project Resources diff --git a/docs/docs/advanced/intrinsics.md b/docs/docs/advanced/intrinsics.md index 3ebf1af55f..518610ae29 100644 --- a/docs/docs/advanced/intrinsics.md +++ b/docs/docs/advanced/intrinsics.md @@ -10,9 +10,10 @@ checkpoints. Both local paths require a GPU or Apple Silicon Mac. An OpenAIBackend using a Granite Switch model served via vLLM uses `uv sync --extra switch` when it downloads embedded adapter metadata. -Adapter functions are adapter-accelerated operations for RAG quality checks. They use -LoRA/aLoRA adapters loaded directly into the Hugging Face backend — faster and more -reliable than prompting a general-purpose model for these specialized micro-tasks. +Adapter functions are adapter-accelerated operations for RAG quality checks. Their +LoRA/aLoRA weights are loaded locally, selected from an embedded checkpoint, or +bundled into an Ollama model — faster and more reliable than prompting a +general-purpose model for these specialised micro-tasks. > **Backend note:** Adapter functions work with three backends: > @@ -25,9 +26,8 @@ reliable than prompting a general-purpose model for these specialized micro-task > `load_embedded_adapters=True`. Only adapter functions embedded in the model are > available — check the model's `adapter_index.json` for the list. > See `docs/docs/examples/granite-switch/README.md` -> - **OllamaModelBackend**: uses an Ollama model that bundles the adapter, -> for example `gabegoodhart/granite4.1-uncertainty:3b`, which is `granite4.1:3b` -> plus the uncertainty aLoRA. Ollama bundles one adapter per model, so pass +> - **OllamaModelBackend** — uses an Ollama model that bundles the adapter. +> Ollama bundles one adapter per model, so pass > `adapter_models={"uncertainty": "", ...}` to route each adapter function > to its model. Install `mellea[switch]` to download the adapter's `io.yaml`. > @@ -43,6 +43,56 @@ from mellea.backends.huggingface import LocalHFBackend backend = LocalHFBackend(model_id="ibm-granite/granite-4.1-3b") ``` +## Use an adapter bundled in an Ollama model + +Ollama serves adapter weights as part of a model tag; Mellea does not load the +weights separately. For local development, build a bundled uncertainty model +from the pinned official Granite base and adapter artefacts: + +```bash +export MELLEA_OLLAMA_UNCERTAINTY_MODEL="$( + ./test/scripts/build_ollama_uncertainty_adapter.sh +)" +``` + +Install the lightweight Hugging Face Hub dependency that retrieves the +catalogued `io.yaml`: + +```bash +uv sync --extra switch +``` + +Pass the bundled model tag for each adapter function. The usual helper API +stays unchanged: + +```python +import os + +from mellea.backends import ModelOption +from mellea.backends.ollama import OllamaModelBackend +from mellea.stdlib.components import Message +from mellea.stdlib.components.intrinsic import core +from mellea.stdlib.context import ChatContext + +backend = OllamaModelBackend( + model_id="granite4.1:3b", + model_options={ModelOption.CONTEXT_WINDOW: 4096}, + adapter_models={ + "uncertainty": os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"], + }, +) +context = ( + ChatContext() + .add(Message("user", "What is the square root of 4?")) + .add(Message("assistant", "The square root of 4 is 2.")) +) + +print(core.check_certainty(context, backend)) +``` + +See `docs/examples/intrinsics/uncertainty_ollama.py` for the complete +executable example. + ## Use a local Granite Switch checkpoint Granite Switch checkpoints contain their adapter functions already. Pass the @@ -347,8 +397,10 @@ Weights-binding support by backend today: | --- | --- | --- | --- | | `LocalHFBackend` | ✅ shipping — `add_adapter` accepts a composed `Adapter` or a bare `LocalFileBinding` directly | ✅ shipping — `load_embedded_adapters=True`, or `add_adapter(adapter, config=...)`/`register_embedded_adapter_model` with a composed `Adapter` | — | | `OpenAIBackend` | — | ✅ shipping — `load_embedded_adapters=True`, or `add_adapter(adapter, config=...)`/`register_embedded_adapter_model` with a composed `Adapter` | — | +| `OllamaModelBackend` | — | — | ✅ model selection through `adapter_models` for catalogued adapter functions; lifecycle telemetry is tracked separately | -`ServerMediatedBinding` has no backend implementation yet — see discussion #1486. +`ServerMediatedBinding` currently supports Ollama's bundled-model path. A full +server-mediated lifecycle and telemetry contract is tracked separately. Discovering *multiple* embedded adapters from a Granite Switch checkpoint or Hub repo (rather than one already-known name) still goes through `register_embedded_adapter_model`, which builds the composed `Adapter` diff --git a/docs/docs/advanced/lora-and-alora-adapters.md b/docs/docs/advanced/lora-and-alora-adapters.md index 2da57a2ff5..2fde7a5d6b 100644 --- a/docs/docs/advanced/lora-and-alora-adapters.md +++ b/docs/docs/advanced/lora-and-alora-adapters.md @@ -15,10 +15,13 @@ and use it as a requirement validator in any Mellea program. Apple Silicon Mac with sufficient VRAM for the chosen base model. Uploading requires a Hugging Face account. -> **Backend note:** Custom-trained adapters can only be loaded into `LocalHFBackend`. -> They do not work with OpenAI or other remote backends. To use one with Ollama, -> convert it to GGUF and bundle it into an Ollama model with a Modelfile `ADAPTER` -> line, then pass that model tag via `OllamaModelBackend(adapter_models=...)`. +> **Backend note:** Custom-trained adapters can only be loaded directly into +> `LocalHFBackend`. Ollama can use a custom adapter bundled into a model with a +> Modelfile `ADAPTER` line, but Mellea does not discover custom Ollama adapter +> functions from `adapter_models` alone. Register a composed adapter with its +> `io.yaml` explicitly, then map its name to the bundled model tag. See +> [Adapter functions](./intrinsics.md) for the supported catalogue-adapter +> workflow. > > Granite Switch models ship with pre-trained adapter functions embedded in the > model weights. Use them through `OpenAIBackend` with a served checkpoint, or diff --git a/docs/docs/tutorials/04-making-agents-reliable.md b/docs/docs/tutorials/04-making-agents-reliable.md index c1ddcb4999..f1f51a92e3 100644 --- a/docs/docs/tutorials/04-making-agents-reliable.md +++ b/docs/docs/tutorials/04-making-agents-reliable.md @@ -365,8 +365,8 @@ response = m.instruct( output_text = str(response) -# Guardian adapter functions require a LocalHFBackend — they load LoRA adapters -# that are not supported by OllamaModelBackend. +# This tutorial uses LocalHFBackend. Ollama requires a compatible bundled model +# for each Guardian adapter function; see Adapter functions for that setup. guardian_backend = LocalHFBackend(model_id="ibm-granite/granite-4.1-3b") # Build a context containing the exchange to check. @@ -397,9 +397,11 @@ and dynamic applications with ease. The word "Mellea" consists of 6 characters. ``` -> **Note:** Guardian adapter functions load LoRA adapters and require `LocalHFBackend`. -> They cannot run against `OllamaModelBackend`. The main agent and the Guardian -> checks can use different backends — only the Guardian calls need `LocalHFBackend`. +> **Note:** This tutorial uses `LocalHFBackend` because it loads Guardian LoRA +> adapters directly. `OllamaModelBackend` can run a Guardian adapter function +> only when you provide an Ollama model tag that bundles that adapter through +> `adapter_models`. The main agent and the Guardian checks can use different +> backends. Scores are floats between 0.0 (safe) and 1.0 (risk detected); 0.5 is the threshold. The available criteria are: `"harm"`, `"jailbreak"`, `"social_bias"`, diff --git a/docs/examples/intrinsics/uncertainty_ollama.py b/docs/examples/intrinsics/uncertainty_ollama.py index c00cf8180e..fae1e1929d 100644 --- a/docs/examples/intrinsics/uncertainty_ollama.py +++ b/docs/examples/intrinsics/uncertainty_ollama.py @@ -1,4 +1,4 @@ -# pytest: e2e, ollama, qualitative +# pytest: e2e, ollama """Example usage of the uncertainty/certainty intrinsic with Ollama. @@ -17,19 +17,25 @@ ``` """ +import os + from mellea import model_ids, start_backend from mellea.stdlib import functional as mfuncs from mellea.stdlib.components.intrinsic import core ctx, backend = start_backend( "ollama", + # Regular chat uses this base model. model_id=model_ids.IBM_GRANITE_4_1_3B, context_type="chat", - adapter_models={"uncertainty": "gabegoodhart/granite4.1-uncertainty:3b"}, + # The certainty helper routes only its adapter call to this bundled aLoRA model. + adapter_models={"uncertainty": os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"]}, ) +# Add the exchange whose answer the adapter will score. response, ctx = mfuncs.chat("What is 2 + 2?", ctx, backend) # type: ignore print(f"Response: {response.content}") +# This call uses the mapped bundled adapter model, not the normal chat model. result = core.check_certainty(ctx, backend) # type: ignore print(f"Certainty score: {result}") diff --git a/mellea/backends/adapters/adapter.py b/mellea/backends/adapters/adapter.py index 935daf1467..4ad92c8bad 100644 --- a/mellea/backends/adapters/adapter.py +++ b/mellea/backends/adapters/adapter.py @@ -565,24 +565,26 @@ def add_adapter(self, adapter: AdapterInput, *, config: dict | None = None) -> N still be rejected at runtime. `config` is the raw io.yaml mapping for a composed `Adapter`/ - `_AdapterCore` whose `weights` is an `EmbeddedBinding` — that reality's - config cannot be cheaply re-derived later, so it must be supplied here - (see `_discover_embedded_adapters`/`resolve_adapter`), rather than - being fetched lazily the way a `LocalFileBinding`'s io.yaml is. + `_AdapterCore` whose `weights` is an `EmbeddedBinding` or + `ServerMediatedBinding`. Those realities do not retain the raw + configuration required by the legacy rewriter, so it must be supplied + at registration rather than being fetched lazily like a + `LocalFileBinding`'s io.yaml. Args: adapter (AdapterInput): The adapter to register with this backend. config (dict | None): Raw io.yaml config for a composed - `EmbeddedBinding` adapter. Ignored (and rejected) for every - other adapter reality. + `EmbeddedBinding` or `ServerMediatedBinding` adapter. Ignored + (and rejected) for every other adapter reality. Raises: TypeError: If `adapter` belongs to a reality this backend does not support, or `config` is given for a reality other than a - composed `EmbeddedBinding` adapter. - ValueError: If `adapter.weights` is an `EmbeddedBinding` and - `config` is not given — registering it without a config would - make it discoverable but permanently unable to generate. + composed `EmbeddedBinding` or `ServerMediatedBinding` adapter. + ValueError: If `adapter.weights` is an `EmbeddedBinding` or + `ServerMediatedBinding` and `config` is not given — + registering it without a config would make it discoverable + but permanently unable to generate. """ @abc.abstractmethod diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index 8cac1b5214..a6d057051c 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -12,6 +12,7 @@ import httpx import ollama +import yaml from tqdm import tqdm from ..backends import ModelIdentifier, model_ids @@ -44,7 +45,10 @@ from ..stdlib.components import Intrinsic, Message from ..stdlib.requirements import ALoraRequirement, LLMaJRequirement, Requirement from ..telemetry.context import generate_request_id, with_context -from .adapters.adapter import AdapterInput, AdapterMixin, IntrinsicAdapter +from .adapters._core import Adapter as _AdapterCore, Identity, ServerMediatedBinding +from .adapters.adapter import AdapterInput, AdapterMixin, _composed_adapter_key +from .adapters.catalog import AdapterType, fetch_intrinsic_metadata +from .adapters.io_contracts import get_io_contract from .backend import FormatterBackend from .model_options import ModelOption from .tools import add_tools_from_context_actions, add_tools_from_model_options @@ -205,7 +209,7 @@ class OllamaModelBackend(FormatterBackend, AdapterMixin): `ollama` SDK default (no timeout). adapter_models (dict[str, str] | None): Mapping from adapter function name (e.g. `"uncertainty"`) to the Ollama model tag that bundles that - adapter (e.g. `"gabegoodhart/granite4.1-uncertainty:3b"`). Ollama + adapter (e.g. `"mellea-test/uncertainty-alora:latest"`). Ollama bundles one adapter per model, so each adapter function is served by its own tag. Adapter functions not listed here run against `model_id`. @@ -221,6 +225,8 @@ class OllamaModelBackend(FormatterBackend, AdapterMixin): OSError: If the model cannot be pulled from the Ollama library. """ + _supports_composed_adapters = True + def __init__( self, model_id: str | ModelIdentifier = model_ids.IBM_GRANITE_4_2_3B, @@ -253,7 +259,8 @@ def __init__( self._model_id: str = ollama_model_id self._provider: str = "ollama" - self._added_adapters: dict[str, IntrinsicAdapter] = {} + self._added_adapters: dict[str, _AdapterCore] = {} + self._composed_adapter_configs: dict[str, dict] = {} self._adapter_models: dict[str, str] = adapter_models or {} # Setup the client and ensure that we have the model available. @@ -328,26 +335,48 @@ def base_model_name(self) -> str: return ident.hf_model_name.split("/")[-1] return self._model_id - def add_adapter(self, adapter: AdapterInput) -> None: + def add_adapter(self, adapter: AdapterInput, *, config: dict | None = None) -> None: """Register an adapter with this backend. - Ollama serves adapter weights bundled into a model, so only the - adapter's I/O config is used here; no weights are loaded. + Ollama serves adapter weights bundled into a model, so Mellea only + registers the adapter's I/O contract and raw `io.yaml` configuration. + The configured Ollama model tag activates the adapter at generation + time. Args: - adapter (AdapterInput): The adapter to register. Must be an - `IntrinsicAdapter`. + adapter (AdapterInput): A composed adapter with a + `ServerMediatedBinding`. + config (dict | None): Parsed `io.yaml` configuration for the + adapter. Required because a composed adapter does not retain + the raw configuration that powers the legacy rewriter. Raises: - TypeError: If `adapter` is not an `IntrinsicAdapter`. + TypeError: If `adapter` does not use `ServerMediatedBinding`. + ValueError: If `config` is omitted for a composed adapter. """ - if not isinstance(adapter, IntrinsicAdapter): + if not isinstance(adapter, _AdapterCore) or not isinstance( + adapter.weights, ServerMediatedBinding + ): raise TypeError( - f"OllamaModelBackend currently only supports IntrinsicAdapter. " + "OllamaModelBackend only supports composed Adapters with a " + "ServerMediatedBinding. " f"Got: {type(adapter).__name__}" ) - adapter.backend = self - self._added_adapters[adapter.qualified_name] = adapter + if config is None: + raise ValueError( + f"No io.yaml config given for server-mediated adapter " + f"{adapter.identity.name!r}; registering it without one would " + "leave it discoverable but unable to generate." + ) + key = _composed_adapter_key(adapter) + if key in self._added_adapters: + MelleaLogger.get_logger().warning( + f"attempted to add adapter {key!r} but it is already registered; " + "refusing to overwrite it." + ) + return + self._added_adapters[key] = adapter + self._composed_adapter_configs[key] = config def list_adapters(self) -> list[str]: """Return qualified names of all registered adapters. @@ -357,6 +386,68 @@ def list_adapters(self) -> list[str]: """ return list(self._added_adapters.keys()) + def resolve_adapter(self, name: str) -> _AdapterCore: + """Find or register an Ollama model-mediated adapter by capability name. + + The model tag selected through `adapter_models` contains the adapter + weights. Mellea downloads only the catalogued `io.yaml`, which + rewrites requests and processes the structured response. + + Args: + name (str): Catalogued adapter function name. + + Returns: + _AdapterCore: The registered server-mediated adapter. + + Raises: + ValueError: If the catalogued `io.yaml` is invalid. + KeyError: If registration does not yield a discoverable adapter. + """ + found = self._find_adapter(name) + if found is not None: + return found + + metadata = fetch_intrinsic_metadata(name) + try: + config_path = granite_formatters.intrinsics.obtain_io_yaml( + name, + self.base_model_name, + metadata.repo_id, + revision=metadata.revision, + alora=True, + ) + except ModuleNotFoundError as e: + if e.name != "huggingface_hub": + raise + raise ImportError( + "Ollama adapter functions require Hugging Face Hub support. " + "Install it with `uv sync --extra switch`." + ) from e + with config_path.open(encoding="utf-8") as config_file: + config = yaml.safe_load(config_file) + if not isinstance(config, dict): + raise ValueError( + f"Adapter configuration at {config_path} must be a mapping, " + f"got {type(config).__name__}." + ) + + self.add_adapter( + _AdapterCore( + identity=Identity( + name=name, + adapter_type=AdapterType.ALORA.value, + capability=metadata.effective_capability, + ), + io_contract=get_io_contract(name), + weights=ServerMediatedBinding(), + ), + config=config, + ) + found = self._find_adapter(name) + if found is None: + raise KeyError(f"Adapter {name!r} not found after registration") + return found + def _check_ollama_server(self) -> bool: """Requests generic info about the Ollama server to ensure it's running.""" try: @@ -531,7 +622,7 @@ async def _generate_from_intrinsic( streaming is requested (intrinsic post-processing requires the complete response). ValueError: If no adapter is registered for the requested intrinsic. - TypeError: If the adapter isn't an `IntrinsicAdapter`. + ValueError: If the registered adapter has no cached `io.yaml`. """ if not ctx.is_chat_context: raise NotImplementedError("Intrinsics require a chat context.") @@ -549,16 +640,16 @@ async def _generate_from_intrinsic( f"backend ({self}) has no adapter for processing adapter function: " f"{action.intrinsic_name}" ) - if not isinstance(adapter, IntrinsicAdapter): - raise TypeError( - f"OllamaModelBackend only supports IntrinsicAdapter, got: {type(adapter).__name__}" + key = _composed_adapter_key(adapter) + intrinsic_config = self._composed_adapter_configs.get(key) + if intrinsic_config is None: + raise ValueError( + f"No io.yaml config cached for server-mediated adapter {key!r}; " + "register it via add_adapter() or resolve_adapter()." ) - intrinsic_config = adapter.config - assert intrinsic_config is not None - rewriter = granite_formatters.IntrinsicsRewriter( - config_dict=intrinsic_config, model_name=adapter.name + config_dict=intrinsic_config, model_name=adapter.identity.name ) result_processor = granite_formatters.IntrinsicsResultProcessor( config_dict=intrinsic_config @@ -730,7 +821,20 @@ async def _generate_from_context( ) alora_action = ALoraRequirement(action.description, adapter_name) - alora_req_adapter = self._find_adapter(adapter_name, ("alora",)) + explicit_types = getattr(alora_action, "_adapter_types", None) + search_types = ( + tuple(adapter_type.value for adapter_type in explicit_types) + if explicit_types + else ("alora",) + ) + alora_req_adapter = self._find_adapter(adapter_name, search_types) + if ( + alora_req_adapter is None + and adapter_name in self._adapter_models + and not explicit_types + ): + await asyncio.to_thread(self.resolve_adapter, adapter_name) + alora_req_adapter = self._find_adapter(adapter_name, search_types) if alora_req_adapter is None: if reroute_to_alora: MelleaLogger.get_logger().warning( @@ -1217,9 +1321,10 @@ async def post_processing( ) # Generate the log for this ModelOutputThunk. + selected_model = mot.generation.model or self._model_id generate_log = GenerateLog() generate_log.prompt = conversation - generate_log.backend = f"ollama::{self._model_id}" + generate_log.backend = f"ollama::{selected_model}" generate_log.model_options = mot._call.model_options generate_log.date = datetime.datetime.now() generate_log.model_output = mot.raw.response @@ -1252,7 +1357,7 @@ async def post_processing( } # Populate model and provider metadata - mot.generation.model = self._model_id + mot.generation.model = selected_model mot.generation.provider = self._provider mot.raw.provider = self._provider diff --git a/test/backends/test_ollama.py b/test/backends/test_ollama.py index d96eaeceea..d565b2453c 100644 --- a/test/backends/test_ollama.py +++ b/test/backends/test_ollama.py @@ -4,7 +4,10 @@ import asyncio import json import os +import subprocess +from pathlib import Path from typing import Annotated +from unittest.mock import patch import ollama as _ollama import pydantic @@ -15,8 +18,12 @@ from mellea.backends.model_ids import IBM_GRANITE_4_2_3B from mellea.backends.ollama import OllamaModelBackend from mellea.core import CBlock, Requirement -from mellea.stdlib.context import SimpleContext +from mellea.stdlib import functional as mfuncs +from mellea.stdlib.components import Intrinsic, Message +from mellea.stdlib.components.intrinsic import core +from mellea.stdlib.context import ChatContext, SimpleContext from mellea.stdlib.requirements import simple_validate +from test.conftest import hf_skip # Mark all tests in this module as requiring Ollama pytestmark = [pytest.mark.ollama, pytest.mark.e2e] @@ -24,6 +31,9 @@ # Match granite4.2:3b's constrained default (Modelfile num_ctx: 8192) so the # runner is loaded once and never reloaded for a context-size mismatch. TEST_CONTEXT_WINDOW = 8192 +_UNCERTAINTY_ADAPTER_BUILDER = ( + Path(__file__).parents[2] / "test/scripts/build_ollama_uncertainty_adapter.sh" +) def _ollama_model_for_eval() -> str: @@ -90,6 +100,17 @@ def session(): session.reset() +@pytest.fixture(scope="session") +def uncertainty_adapter_model() -> str: + """Build and return the official uncertainty aLoRA Ollama model tag.""" + if configured_model := os.environ.get("MELLEA_OLLAMA_UNCERTAINTY_MODEL"): + return configured_model + completed = subprocess.run( + [_UNCERTAINTY_ADAPTER_BUILDER], check=True, stdout=subprocess.PIPE, text=True + ) + return completed.stdout.strip() + + @pytest.mark.qualitative def test_simple_instruct(session) -> None: result = session.instruct( @@ -275,6 +296,75 @@ async def test_async_avalue(session) -> None: assert mot1.generation.ttfb_ms is None +@pytest.mark.timeout(120) +def test_uncertainty_adapter_function(uncertainty_adapter_model: str) -> None: + """A bundled aLoRA model serves the intrinsic and public certainty helper.""" + backend = OllamaModelBackend( + model_id="granite4.1:3b", + model_options={ModelOption.CONTEXT_WINDOW: 4096}, + adapter_models={"uncertainty": uncertainty_adapter_model}, + ) + context = ( + ChatContext() + .add(Message("user", "What is the square root of 4?")) + .add(Message("assistant", "The square root of 4 is 2.")) + ) + with hf_skip(): + adapter = backend.resolve_adapter("uncertainty") + requested_models: list[str] = [] + original_chat = _ollama.AsyncClient.chat + + async def record_model_selection(client, *args, **kwargs): + model = kwargs.get("model") + assert isinstance(model, str) + requested_models.append(model) + return await original_chat(client, *args, **kwargs) + + with patch.object(_ollama.AsyncClient, "chat", new=record_model_selection): + output, _ = mfuncs.act( + Intrinsic("uncertainty"), context, backend, strategy=None + ) + score = core.check_certainty(context, backend) + + assert requested_models == [uncertainty_adapter_model, uncertainty_adapter_model] + assert output.generation.model == uncertainty_adapter_model + parsed = json.loads(output.value) + assert 0.0 <= parsed["certainty"] <= 1.0 + + assert 0.0 <= score <= 1.0 + assert adapter.identity.adapter_type == "alora" + + +@pytest.mark.qualitative +@pytest.mark.timeout(120) +def test_uncertainty_adapter_changes_base_model_score( + uncertainty_adapter_model: str, +) -> None: + """The bundled aLoRA materially changes the base model's certainty score.""" + context = ( + ChatContext() + .add(Message("user", "What is the square root of 4?")) + .add(Message("assistant", "The square root of 4 is 2.")) + ) + model_options = {ModelOption.CONTEXT_WINDOW: 4096, ModelOption.TEMPERATURE: 0.0} + base_backend = OllamaModelBackend( + model_id="granite4.1:3b", + model_options=model_options, + adapter_models={"uncertainty": "granite4.1:3b"}, + ) + adapter_backend = OllamaModelBackend( + model_id="granite4.1:3b", + model_options=model_options, + adapter_models={"uncertainty": uncertainty_adapter_model}, + ) + + with hf_skip(): + base_score = core.check_certainty(context, base_backend) + adapter_score = core.check_certainty(context, adapter_backend) + + assert abs(adapter_score - base_score) >= 0.2 + + def test_multiple_asyncio_runs(session) -> None: async def test(): result = await session.achat("hello") diff --git a/test/backends/test_ollama_intrinsics_unit.py b/test/backends/test_ollama_intrinsics_unit.py index cda3ea2d7c..d20cb352e0 100644 --- a/test/backends/test_ollama_intrinsics_unit.py +++ b/test/backends/test_ollama_intrinsics_unit.py @@ -13,17 +13,24 @@ """ import json +from pathlib import Path from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import ollama import pytest from mellea.backends import ModelOption -from mellea.backends.adapters.adapter import AdapterType, IntrinsicAdapter +from mellea.backends.adapters import ( + Adapter, + Identity, + ServerMediatedBinding, + get_io_contract, +) from mellea.backends.ollama import OllamaModelBackend, _to_chat_completion_dict from mellea.stdlib import functional as mfuncs from mellea.stdlib.components import Intrinsic, Message from mellea.stdlib.context import ChatContext +from mellea.stdlib.requirements import ALoraRequirement # --------------------------------------------------------------------------- # Configs @@ -68,7 +75,7 @@ "sentence_boundaries": None, } -_ADAPTER_TAG = "gabegoodhart/granite4.1-uncertainty:3b" +_ADAPTER_TAG = "mellea-test/uncertainty-alora:latest" # --------------------------------------------------------------------------- # Canned responses @@ -167,10 +174,12 @@ def _make_backend_with_adapter( ) -> OllamaModelBackend: """Return an OllamaModelBackend with a registered uncertainty adapter.""" backend = _make_backend(model_options=model_options, adapter_models=adapter_models) - adapter = IntrinsicAdapter( - "uncertainty", adapter_type=AdapterType.LORA, config_dict=config + adapter = Adapter( + identity=Identity(name="uncertainty", adapter_type="alora"), + io_contract=get_io_contract("uncertainty"), + weights=ServerMediatedBinding(), ) - backend.add_adapter(adapter) + backend.add_adapter(adapter, config=config) return backend @@ -251,6 +260,39 @@ async def test_adapter_model_tag_defaults_to_model_id(): assert mock_chat.call_args.kwargs["model"] == "granite4.1:3b" +async def test_alora_requirement_resolves_mapped_adapter( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + """A configured catalogue adapter is resolved before requirement routing.""" + config_path = tmp_path / "io.yaml" + config_path.write_text(json.dumps(_SIMPLE_CONFIG), encoding="utf-8") + monkeypatch.setattr( + "mellea.backends.ollama.granite_formatters.intrinsics.obtain_io_yaml", + lambda *_args, **_kwargs: config_path, + ) + backend = _make_backend(adapter_models={"requirement-check": _ADAPTER_TAG}) + mock_chat = AsyncMock(return_value=_simple_chat_response()) + mock_client = MagicMock() + mock_client.chat = mock_chat + + with patch.object( + OllamaModelBackend, + "_async_client", + new_callable=PropertyMock, + return_value=mock_client, + ): + mot, _ = await mfuncs.aact( + ALoraRequirement("The response is correct."), + _make_context(), + backend, + strategy=None, + ) + await mot.avalue() + + assert mock_chat.call_args.kwargs["model"] == _ADAPTER_TAG + assert backend.list_adapters() == ["requirement-check_alora"] + + async def test_result_processor_applied(): """Full uncertainty config: likelihood + project transforms produce the expected JSON.""" backend = _make_backend_with_adapter(_UNCERTAINTY_CONFIG) @@ -343,17 +385,65 @@ def get_temperature(location: str) -> int: # --------------------------------------------------------------------------- -def test_add_adapter_registers_intrinsic_adapter(): +def test_add_adapter_registers_server_mediated_adapter(): backend = _make_backend_with_adapter(_SIMPLE_CONFIG) - assert backend.list_adapters() == ["uncertainty_lora"] + assert backend.list_adapters() == ["uncertainty_alora"] def test_add_adapter_rejects_other_adapter_types(): backend = _make_backend() - with pytest.raises(TypeError, match="only supports IntrinsicAdapter"): + with pytest.raises(TypeError, match="ServerMediatedBinding"): backend.add_adapter(object()) # type: ignore[arg-type] +def test_add_adapter_requires_io_yaml_config(): + backend = _make_backend() + adapter = Adapter( + identity=Identity(name="uncertainty", adapter_type="alora"), + io_contract=get_io_contract("uncertainty"), + weights=ServerMediatedBinding(), + ) + + with pytest.raises(ValueError, match=r"No io\.yaml config"): + backend.add_adapter(adapter) + + +def test_resolve_adapter_registers_server_mediated_adapter( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + """Resolving a capability creates the modern adapter shape without a server call.""" + config_path = tmp_path / "io.yaml" + config_path.write_text(json.dumps(_SIMPLE_CONFIG), encoding="utf-8") + backend = _make_backend() + monkeypatch.setattr( + "mellea.backends.ollama.granite_formatters.intrinsics.obtain_io_yaml", + lambda *_args, **_kwargs: config_path, + ) + + adapter = backend.resolve_adapter("uncertainty") + + assert adapter.identity.name == "uncertainty" + assert adapter.identity.adapter_type == "alora" + assert isinstance(adapter.weights, ServerMediatedBinding) + assert backend.list_adapters() == ["uncertainty_alora"] + + +def test_resolve_adapter_explains_missing_huggingface_extra( + monkeypatch: pytest.MonkeyPatch, +): + """The optional dependency failure names the extra users need to install.""" + backend = _make_backend() + monkeypatch.setattr( + "mellea.backends.ollama.granite_formatters.intrinsics.obtain_io_yaml", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + ModuleNotFoundError(name="huggingface_hub") + ), + ) + + with pytest.raises(ImportError, match=r"uv sync --extra switch"): + backend.resolve_adapter("uncertainty") + + def test_base_model_name_maps_ollama_tag_to_hf_name(): backend = _make_backend() assert backend.base_model_name == "granite-4.1-3b" diff --git a/test/scripts/build_ollama_uncertainty_adapter.sh b/test/scripts/build_ollama_uncertainty_adapter.sh new file mode 100755 index 0000000000..204b1106ef --- /dev/null +++ b/test/scripts/build_ollama_uncertainty_adapter.sh @@ -0,0 +1,106 @@ +#!/bin/bash +# Copyright IBM Corp. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Build a local Ollama model from the official pinned Granite uncertainty aLoRA. +# +# Prints only the resulting Ollama tag to stdout so callers can assign it: +# export MELLEA_OLLAMA_UNCERTAINTY_MODEL="$(./test/scripts/build_ollama_uncertainty_adapter.sh)" + +set -euo pipefail + +log() { echo "[ollama-adapter] $*" >&2; } +die() { log "ERROR: $*"; exit 1; } + +OLLAMA_BIN="${OLLAMA_BIN:-$(command -v ollama)}" +BASE_MODEL="${OLLAMA_BASE_MODEL:-granite4.1:3b}" +OUTPUT_MODEL="${MELLEA_OLLAMA_UNCERTAINTY_MODEL:-mellea-test/uncertainty-alora:latest}" +ADAPTER_REPO="ibm-granite/granitelib-core-r1.0" +ADAPTER_REVISION="d0a2a96a4cd07e96f0fe7ca29a42bfe088299d43" +ADAPTER_PATH="uncertainty/granite-4.1-3b/alora" +BASE_REPO="ibm-granite/granite-4.1-3b" +LLAMA_CPP_REVISION="e71b80510c848c00175924ecf3c40333ccae8eb5" + +if [[ -n "${MELLEA_OLLAMA_ADAPTER_CACHE_DIR:-}" ]]; then + BUILD_CACHE="$MELLEA_OLLAMA_ADAPTER_CACHE_DIR" +elif [[ -n "${CACHE_DIR:-}" ]]; then + BUILD_CACHE="${CACHE_DIR}/ollama-adapter-build" +else + BUILD_CACHE="${XDG_CACHE_HOME:-$HOME/.cache}/mellea/ollama-adapter-build" +fi + +BASE_DIR="${BUILD_CACHE}/base" +ADAPTER_DIR="${BUILD_CACHE}/adapter" +HF_HOME_DIR="${BUILD_CACHE}/huggingface" +LLAMA_CPP_DIR="${BUILD_CACHE}/llama.cpp" +CONVERTER_VENV="${BUILD_CACHE}/converter-venv" +ADAPTER_GGUF="${BUILD_CACHE}/uncertainty-alora-f16.gguf" +MODELFILE="${BUILD_CACHE}/Modelfile" + +mkdir -p "$BUILD_CACHE" + +if [[ ! -f "$ADAPTER_GGUF" ]]; then + log "Downloading official Granite base and uncertainty aLoRA..." + HF_HOME="$HF_HOME_DIR" uv run --quiet --frozen --all-extras --all-groups \ + python - "$BASE_DIR" "$ADAPTER_DIR" <<'PY' +import shutil +import sys +from pathlib import Path + +from huggingface_hub import hf_hub_download, snapshot_download + +base_dir = Path(sys.argv[1]) +adapter_dir = Path(sys.argv[2]) +snapshot_download( + repo_id="ibm-granite/granite-4.1-3b", + local_dir=base_dir, + ignore_patterns=["*.gguf", "*.onnx", "*.tflite"], +) +adapter_dir.mkdir(parents=True, exist_ok=True) +for filename in ( + "adapter_config.json", + "adapter_model.safetensors", + "io.yaml", + "model.sig", +): + source = hf_hub_download( + repo_id="ibm-granite/granitelib-core-r1.0", + filename=f"uncertainty/granite-4.1-3b/alora/{filename}", + revision="d0a2a96a4cd07e96f0fe7ca29a42bfe088299d43", + ) + shutil.copy2(source, adapter_dir / filename) +PY + + if [[ ! -d "$LLAMA_CPP_DIR/.git" ]]; then + log "Cloning llama.cpp at ${LLAMA_CPP_REVISION}..." + git clone --quiet https://github.com/ggml-org/llama.cpp.git "$LLAMA_CPP_DIR" + fi + git -C "$LLAMA_CPP_DIR" fetch --quiet --depth 1 origin "$LLAMA_CPP_REVISION" + git -C "$LLAMA_CPP_DIR" checkout --quiet --detach "$LLAMA_CPP_REVISION" + + if [[ ! -x "$CONVERTER_VENV/bin/python" ]]; then + log "Installing pinned converter dependencies..." + uv venv "$CONVERTER_VENV" + uv pip install --python "$CONVERTER_VENV/bin/python" \ + -r "$LLAMA_CPP_DIR/requirements/requirements-convert_lora_to_gguf.txt" + fi + + log "Converting official uncertainty aLoRA to GGUF..." + "$CONVERTER_VENV/bin/python" "$LLAMA_CPP_DIR/convert_lora_to_gguf.py" \ + "$ADAPTER_DIR" \ + --base "$BASE_DIR" \ + --outtype f16 \ + --outfile "$ADAPTER_GGUF" + + if ! "$CONVERTER_VENV/bin/python" \ + "$LLAMA_CPP_DIR/gguf-py/gguf/scripts/gguf_dump.py" "$ADAPTER_GGUF" \ + | grep -q "adapter.alora.invocation_tokens"; then + die "Converted adapter is missing aLoRA invocation-token metadata." + fi +fi + +printf 'FROM %s\nADAPTER %s\n' "$BASE_MODEL" "$ADAPTER_GGUF" > "$MODELFILE" +log "Creating Ollama model ${OUTPUT_MODEL}..." +"$OLLAMA_BIN" create "$OUTPUT_MODEL" -f "$MODELFILE" >/dev/null + +printf '%s\n' "$OUTPUT_MODEL" diff --git a/test/scripts/run_tests_with_ollama_and_vllm.sh b/test/scripts/run_tests_with_ollama_and_vllm.sh index 34965e7e5c..65106b22cb 100755 --- a/test/scripts/run_tests_with_ollama_and_vllm.sh +++ b/test/scripts/run_tests_with_ollama_and_vllm.sh @@ -46,6 +46,7 @@ OLLAMA_CONTEXT_LENGTH="${OLLAMA_CONTEXT_LENGTH:-2048}" OLLAMA_MODEL_LIST=( "granite4.2:3b" "granite4:micro-h" + "granite4.1:3b" "hf.co/ibm-granite/granite-vision-4.1-4b-GGUF:Q4_K_M" "llama3.2" "qwen2.5vl:7b" @@ -194,6 +195,11 @@ done log "All ollama models ready." +export MELLEA_OLLAMA_UNCERTAINTY_MODEL="$( + ./test/scripts/build_ollama_uncertainty_adapter.sh +)" +OLLAMA_MODEL_LIST+=("$MELLEA_OLLAMA_UNCERTAINTY_MODEL") + # --- Warm up models (first load into memory is slow) --- # Disable with SKIP_WARMUP=1 (covers all backends) or OLLAMA_SKIP_WARMUP=1 (ollama only). # Note: vLLM has no warmup step — it serves immediately after the readiness check. From aec57f9fffa4cf9b9998f1e199784ce8ec40cbef Mon Sep 17 00:00:00 2001 From: Nigel Jones Date: Tue, 8 Sep 2026 09:00:55 +0100 Subject: [PATCH 3/9] docs: explain Ollama adapter cache trade-off Assisted-by: Codex Signed-off-by: Nigel Jones --- docs/docs/advanced/intrinsics.md | 14 +++++++++++--- docs/examples/intrinsics/uncertainty_ollama.py | 14 +++++++------- 2 files changed, 18 insertions(+), 10 deletions(-) diff --git a/docs/docs/advanced/intrinsics.md b/docs/docs/advanced/intrinsics.md index 518610ae29..21b162b16a 100644 --- a/docs/docs/advanced/intrinsics.md +++ b/docs/docs/advanced/intrinsics.md @@ -62,8 +62,10 @@ catalogued `io.yaml`: uv sync --extra switch ``` -Pass the bundled model tag for each adapter function. The usual helper API -stays unchanged: +For one adapter function, use the bundled tag for both ordinary chat and the +adapter route. Before the invocation tokens appear, the aLoRA model behaves as +the base model; keeping one model identity lets Ollama reuse its own prefix +cache where available. ```python import os @@ -75,7 +77,7 @@ from mellea.stdlib.components.intrinsic import core from mellea.stdlib.context import ChatContext backend = OllamaModelBackend( - model_id="granite4.1:3b", + model_id=os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"], model_options={ModelOption.CONTEXT_WINDOW: 4096}, adapter_models={ "uncertainty": os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"], @@ -90,6 +92,12 @@ context = ( print(core.check_certainty(context, backend)) ``` +When an application needs several adapter functions, keep the base model as +`model_id` and map each function to its bundled model tag. Current Ollama +model packaging exposes one adapter per tag, so calls across those tags do not +share a KV cache. Use a Granite Switch checkpoint when multi-adapter, +single-model serving is important. + See `docs/examples/intrinsics/uncertainty_ollama.py` for the complete executable example. diff --git a/docs/examples/intrinsics/uncertainty_ollama.py b/docs/examples/intrinsics/uncertainty_ollama.py index fae1e1929d..2a305a9dc7 100644 --- a/docs/examples/intrinsics/uncertainty_ollama.py +++ b/docs/examples/intrinsics/uncertainty_ollama.py @@ -5,9 +5,9 @@ Evaluates how certain the model is about its response to a user question. The context should contain a user question followed by an assistant answer. -Ollama bundles one adapter per model, so the uncertainty adapter is served by -its own model tag (`granite4.1:3b` plus the uncertainty aLoRA). Pass that tag -via `adapter_models`; normal chat still uses the base model. +Ollama bundles one adapter per model. For this single-adapter example, the +uncertainty aLoRA tag is used for normal chat and the certainty helper, so +Ollama can retain one model identity. Requires `mellea[switch]` to download the adapter's `io.yaml`. @@ -19,16 +19,16 @@ import os -from mellea import model_ids, start_backend +from mellea import start_backend from mellea.stdlib import functional as mfuncs from mellea.stdlib.components.intrinsic import core ctx, backend = start_backend( "ollama", - # Regular chat uses this base model. - model_id=model_ids.IBM_GRANITE_4_1_3B, + # Before its invocation tokens, this bundled aLoRA behaves as the base model. + model_id=os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"], context_type="chat", - # The certainty helper routes only its adapter call to this bundled aLoRA model. + # The certainty helper uses the same model identity. adapter_models={"uncertainty": os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"]}, ) From cd2a2b0c3dfce99d035f1593317a9f3cf3beb212 Mon Sep 17 00:00:00 2001 From: Nigel Jones Date: Tue, 8 Sep 2026 09:04:29 +0100 Subject: [PATCH 4/9] fix(backends): support bundled Ollama adapter models Assisted-by: Codex Signed-off-by: Nigel Jones --- docs/docs/advanced/intrinsics.md | 1 + docs/examples/intrinsics/uncertainty_ollama.py | 2 ++ mellea/backends/ollama.py | 8 ++++++++ test/backends/test_ollama.py | 3 ++- 4 files changed, 13 insertions(+), 1 deletion(-) diff --git a/docs/docs/advanced/intrinsics.md b/docs/docs/advanced/intrinsics.md index 21b162b16a..da1aa705ac 100644 --- a/docs/docs/advanced/intrinsics.md +++ b/docs/docs/advanced/intrinsics.md @@ -78,6 +78,7 @@ from mellea.stdlib.context import ChatContext backend = OllamaModelBackend( model_id=os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"], + adapter_base_model_name="granite-4.1-3b", model_options={ModelOption.CONTEXT_WINDOW: 4096}, adapter_models={ "uncertainty": os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"], diff --git a/docs/examples/intrinsics/uncertainty_ollama.py b/docs/examples/intrinsics/uncertainty_ollama.py index 2a305a9dc7..439def30c1 100644 --- a/docs/examples/intrinsics/uncertainty_ollama.py +++ b/docs/examples/intrinsics/uncertainty_ollama.py @@ -27,6 +27,8 @@ "ollama", # Before its invocation tokens, this bundled aLoRA behaves as the base model. model_id=os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"], + # The bundled tag cannot identify the Hugging Face adapter directory itself. + adapter_base_model_name="granite-4.1-3b", context_type="chat", # The certainty helper uses the same model identity. adapter_models={"uncertainty": os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"]}, diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index a6d057051c..f3e0911ab3 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -212,6 +212,10 @@ class OllamaModelBackend(FormatterBackend, AdapterMixin): adapter (e.g. `"mellea-test/uncertainty-alora:latest"`). Ollama bundles one adapter per model, so each adapter function is served by its own tag. Adapter functions not listed here run against `model_id`. + adapter_base_model_name (str | None): Hugging Face base-model directory + name used to find an adapter's `io.yaml` (for example, + `"granite-4.1-3b"`). Required when `model_id` is itself a bundled + Ollama adapter tag rather than a known base-model tag. Attributes: to_mellea_model_opts_map (dict): Mapping from Ollama-specific option names @@ -235,6 +239,7 @@ def __init__( model_options: dict | None = None, timeout: float | None = 300.0, adapter_models: dict[str, str] | None = None, + adapter_base_model_name: str | None = None, ): """Initialize an Ollama backend, connecting to the server and pulling the model if needed.""" super().__init__( @@ -262,6 +267,7 @@ def __init__( self._added_adapters: dict[str, _AdapterCore] = {} self._composed_adapter_configs: dict[str, dict] = {} self._adapter_models: dict[str, str] = adapter_models or {} + self._adapter_base_model_name = adapter_base_model_name # Setup the client and ensure that we have the model available. self._base_url = base_url @@ -326,6 +332,8 @@ def base_model_name(self) -> str: Returns: str: The short base model name. """ + if self._adapter_base_model_name is not None: + return self._adapter_base_model_name for ident in vars(model_ids).values(): if ( isinstance(ident, ModelIdentifier) diff --git a/test/backends/test_ollama.py b/test/backends/test_ollama.py index d565b2453c..25866e4dc2 100644 --- a/test/backends/test_ollama.py +++ b/test/backends/test_ollama.py @@ -300,7 +300,8 @@ async def test_async_avalue(session) -> None: def test_uncertainty_adapter_function(uncertainty_adapter_model: str) -> None: """A bundled aLoRA model serves the intrinsic and public certainty helper.""" backend = OllamaModelBackend( - model_id="granite4.1:3b", + model_id=uncertainty_adapter_model, + adapter_base_model_name="granite-4.1-3b", model_options={ModelOption.CONTEXT_WINDOW: 4096}, adapter_models={"uncertainty": uncertainty_adapter_model}, ) From 8d2d5728078db6ad48319c66e173361e390d0f7c Mon Sep 17 00:00:00 2001 From: Nigel Jones Date: Tue, 8 Sep 2026 10:41:38 +0100 Subject: [PATCH 5/9] fix: harden Ollama adapter provisioning Assisted-by: Codex Signed-off-by: Nigel Jones --- .../examples/intrinsics/uncertainty_ollama.py | 15 ++++++++-- mellea/backends/ollama.py | 1 + .../build_ollama_uncertainty_adapter.sh | 29 ++++++++++++++----- 3 files changed, 35 insertions(+), 10 deletions(-) diff --git a/docs/examples/intrinsics/uncertainty_ollama.py b/docs/examples/intrinsics/uncertainty_ollama.py index 439def30c1..ffcbecb78d 100644 --- a/docs/examples/intrinsics/uncertainty_ollama.py +++ b/docs/examples/intrinsics/uncertainty_ollama.py @@ -18,20 +18,31 @@ """ import os +import subprocess +from pathlib import Path from mellea import start_backend from mellea.stdlib import functional as mfuncs from mellea.stdlib.components.intrinsic import core +adapter_model = os.environ.get("MELLEA_OLLAMA_UNCERTAINTY_MODEL") +if adapter_model is None: + builder = ( + Path(__file__).parents[2] / "test/scripts/build_ollama_uncertainty_adapter.sh" + ) + adapter_model = subprocess.run( + [builder], check=True, stdout=subprocess.PIPE, text=True + ).stdout.strip() + ctx, backend = start_backend( "ollama", # Before its invocation tokens, this bundled aLoRA behaves as the base model. - model_id=os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"], + model_id=adapter_model, # The bundled tag cannot identify the Hugging Face adapter directory itself. adapter_base_model_name="granite-4.1-3b", context_type="chat", # The certainty helper uses the same model identity. - adapter_models={"uncertainty": os.environ["MELLEA_OLLAMA_UNCERTAINTY_MODEL"]}, + adapter_models={"uncertainty": adapter_model}, ) # Add the exchange whose answer the adapter will score. diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index f3e0911ab3..7de82ec6e7 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -838,6 +838,7 @@ async def _generate_from_context( alora_req_adapter = self._find_adapter(adapter_name, search_types) if ( alora_req_adapter is None + and reroute_to_alora and adapter_name in self._adapter_models and not explicit_types ): diff --git a/test/scripts/build_ollama_uncertainty_adapter.sh b/test/scripts/build_ollama_uncertainty_adapter.sh index 204b1106ef..7c3b05b991 100755 --- a/test/scripts/build_ollama_uncertainty_adapter.sh +++ b/test/scripts/build_ollama_uncertainty_adapter.sh @@ -12,13 +12,18 @@ set -euo pipefail log() { echo "[ollama-adapter] $*" >&2; } die() { log "ERROR: $*"; exit 1; } -OLLAMA_BIN="${OLLAMA_BIN:-$(command -v ollama)}" +OLLAMA_BIN="${OLLAMA_BIN:-}" +if [[ -z "$OLLAMA_BIN" ]]; then + OLLAMA_BIN="$(command -v ollama || true)" +fi +[[ -n "$OLLAMA_BIN" ]] || die "ollama is not installed or not on PATH." BASE_MODEL="${OLLAMA_BASE_MODEL:-granite4.1:3b}" OUTPUT_MODEL="${MELLEA_OLLAMA_UNCERTAINTY_MODEL:-mellea-test/uncertainty-alora:latest}" ADAPTER_REPO="ibm-granite/granitelib-core-r1.0" ADAPTER_REVISION="d0a2a96a4cd07e96f0fe7ca29a42bfe088299d43" ADAPTER_PATH="uncertainty/granite-4.1-3b/alora" BASE_REPO="ibm-granite/granite-4.1-3b" +BASE_REVISION="c0650403e44e78ec0262dab1c90914c65b196c4e" LLAMA_CPP_REVISION="e71b80510c848c00175924ecf3c40333ccae8eb5" if [[ -n "${MELLEA_OLLAMA_ADAPTER_CACHE_DIR:-}" ]]; then @@ -42,7 +47,7 @@ mkdir -p "$BUILD_CACHE" if [[ ! -f "$ADAPTER_GGUF" ]]; then log "Downloading official Granite base and uncertainty aLoRA..." HF_HOME="$HF_HOME_DIR" uv run --quiet --frozen --all-extras --all-groups \ - python - "$BASE_DIR" "$ADAPTER_DIR" <<'PY' + python - "$BASE_DIR" "$ADAPTER_DIR" "$BASE_REPO" "$BASE_REVISION" "$ADAPTER_REPO" "$ADAPTER_REVISION" "$ADAPTER_PATH" <<'PY' import shutil import sys from pathlib import Path @@ -51,8 +56,14 @@ from huggingface_hub import hf_hub_download, snapshot_download base_dir = Path(sys.argv[1]) adapter_dir = Path(sys.argv[2]) +base_repo = sys.argv[3] +base_revision = sys.argv[4] +adapter_repo = sys.argv[5] +adapter_revision = sys.argv[6] +adapter_path = sys.argv[7] snapshot_download( - repo_id="ibm-granite/granite-4.1-3b", + repo_id=base_repo, + revision=base_revision, local_dir=base_dir, ignore_patterns=["*.gguf", "*.onnx", "*.tflite"], ) @@ -64,9 +75,9 @@ for filename in ( "model.sig", ): source = hf_hub_download( - repo_id="ibm-granite/granitelib-core-r1.0", - filename=f"uncertainty/granite-4.1-3b/alora/{filename}", - revision="d0a2a96a4cd07e96f0fe7ca29a42bfe088299d43", + repo_id=adapter_repo, + filename=f"{adapter_path}/{filename}", + revision=adapter_revision, ) shutil.copy2(source, adapter_dir / filename) PY @@ -92,9 +103,11 @@ PY --outtype f16 \ --outfile "$ADAPTER_GGUF" - if ! "$CONVERTER_VENV/bin/python" \ + METADATA_DUMP="${BUILD_CACHE}/uncertainty-alora-metadata.txt" + "$CONVERTER_VENV/bin/python" \ "$LLAMA_CPP_DIR/gguf-py/gguf/scripts/gguf_dump.py" "$ADAPTER_GGUF" \ - | grep -q "adapter.alora.invocation_tokens"; then + > "$METADATA_DUMP" + if ! grep -q "adapter.alora.invocation_tokens" "$METADATA_DUMP"; then die "Converted adapter is missing aLoRA invocation-token metadata." fi fi From 7c3d4949ccf671cb065c2664bd6021afe69a0fba Mon Sep 17 00:00:00 2001 From: Nigel Jones Date: Tue, 8 Sep 2026 11:21:44 +0100 Subject: [PATCH 6/9] test: provision Ollama adapter before pytest Assisted-by: Codex Signed-off-by: Nigel Jones --- .github/workflows/quality.yml | 8 ++++++++ test/backends/test_ollama.py | 2 -- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml index a1d116260a..500e929c14 100644 --- a/.github/workflows/quality.yml +++ b/.github/workflows/quality.yml @@ -153,6 +153,14 @@ jobs: # The names below are what the tests match against; print them so a # tag-name mismatch is diagnosable from the log alone. ollama list + - name: Build the official uncertainty aLoRA Ollama model + env: + # read-only public-repo HF token; used only to make the pinned build + # reliable under Hub rate limits. + HF_TOKEN: ${{ secrets.HF_TOKEN_READ_PUBLIC_ONLY }} # zizmor: ignore[secrets-outside-env] + run: | + model=$(./test/scripts/build_ollama_uncertainty_adapter.sh) + echo "MELLEA_OLLAMA_UNCERTAINTY_MODEL=$model" >> "$GITHUB_ENV" - name: Constrain the granite4.2:3b context for the CI runner # The published tag ships num_ctx=131072, a ~6 GB KV cache at load — # too much for the 16 GB runner and a cause of the CI stalls. Re-point diff --git a/test/backends/test_ollama.py b/test/backends/test_ollama.py index 25866e4dc2..d3d63f5579 100644 --- a/test/backends/test_ollama.py +++ b/test/backends/test_ollama.py @@ -296,7 +296,6 @@ async def test_async_avalue(session) -> None: assert mot1.generation.ttfb_ms is None -@pytest.mark.timeout(120) def test_uncertainty_adapter_function(uncertainty_adapter_model: str) -> None: """A bundled aLoRA model serves the intrinsic and public certainty helper.""" backend = OllamaModelBackend( @@ -337,7 +336,6 @@ async def record_model_selection(client, *args, **kwargs): @pytest.mark.qualitative -@pytest.mark.timeout(120) def test_uncertainty_adapter_changes_base_model_score( uncertainty_adapter_model: str, ) -> None: From 6fae19a81901b4a6fbe94605effe47e4c8d60ecb Mon Sep 17 00:00:00 2001 From: Nigel Jones Date: Wed, 9 Sep 2026 08:21:09 +0100 Subject: [PATCH 7/9] fix(backends): address Ollama adapter review findings (PR #1634) Fixes 8 findings from psschwei's review plus AngeloDanducci's suggestion: - Fold extra_body.documents into a message in _generate_from_intrinsic; Ollama's chat SDK has no extra_body passthrough, so answerability, citations, hallucination_detection, clarify_query, and find_context_attributions were sending zero documents. - resolve_adapter now raises if name has no adapter_models entry, instead of registering successfully and letting generation silently fall back to the base model. - resolve_adapter picks LoRA vs aLoRA from the catalog's adapter_types instead of hardcoding aLoRA; restricted context-attribution, citations, and hallucination_detection to LoRA-only (verified against the Hub). - reroute_to_alora now follows a new default_to_constraint_checking_alora flag (default True), matching OpenAIBackend/LocalHFBackend, instead of only firing for ALoraRequirement. - _generate_from_context now awaits do_generate_walk before the Requirement/Intrinsic dispatch, matching the other backends. - Fix Path(__file__).parents[2] -> parents[3] in the example script. - Constrain granite4.1:3b's context in CI before the adapter build step; only granite4.2:3b was constrained, so the new bundled tag inherited the unconstrained default and reintroduced the CI-stall risk. - Fix export VAR="$(cmd)" masking command-substitution failures under set -e in the test runner script, the build script's own usage comment, and docs/docs/advanced/intrinsics.md. - test_ollama.py's adapter-build fixture now skips on build failure instead of erroring. Added regression tests for the resolve_adapter and document-forwarding fixes; updated two existing unit tests that had encoded the old (buggy) unconditional-resolve behaviour. Assisted-by: Claude Code Signed-off-by: Nigel Jones --- .github/workflows/quality.yml | 11 +++ docs/docs/advanced/intrinsics.md | 3 +- .../examples/intrinsics/uncertainty_ollama.py | 2 +- mellea/backends/adapters/catalog.py | 16 ++++- mellea/backends/ollama.py | 41 +++++++++-- test/backends/test_ollama.py | 12 +++- test/backends/test_ollama_intrinsics_unit.py | 71 ++++++++++++++++++- .../build_ollama_uncertainty_adapter.sh | 3 +- .../scripts/run_tests_with_ollama_and_vllm.sh | 3 +- 9 files changed, 146 insertions(+), 16 deletions(-) diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml index 500e929c14..bba043b5db 100644 --- a/.github/workflows/quality.yml +++ b/.github/workflows/quality.yml @@ -153,6 +153,17 @@ jobs: # The names below are what the tests match against; print them so a # tag-name mismatch is diagnosable from the log alone. ollama list + - name: Constrain the granite4.1:3b context for the CI runner + # granite4.1:3b ships the same oversized num_ctx default as + # granite4.2:3b below. The uncertainty aLoRA Ollama model is built + # `FROM` this tag, so it must be constrained *before* that build step + # or the bundled adapter model inherits the unconstrained context. + run: | + ollama cp granite4.1:3b granite4.1:3b-128k + printf 'FROM granite4.1:3b-128k\nPARAMETER num_ctx 8192\n' \ + > /tmp/MODELFILE-granite41-ci + ollama create granite4.1:3b -f /tmp/MODELFILE-granite41-ci + ollama show granite4.1:3b --modelfile | grep -qx 'PARAMETER num_ctx 8192' - name: Build the official uncertainty aLoRA Ollama model env: # read-only public-repo HF token; used only to make the pinned build diff --git a/docs/docs/advanced/intrinsics.md b/docs/docs/advanced/intrinsics.md index da1aa705ac..2572599160 100644 --- a/docs/docs/advanced/intrinsics.md +++ b/docs/docs/advanced/intrinsics.md @@ -50,9 +50,10 @@ weights separately. For local development, build a bundled uncertainty model from the pinned official Granite base and adapter artefacts: ```bash -export MELLEA_OLLAMA_UNCERTAINTY_MODEL="$( +MELLEA_OLLAMA_UNCERTAINTY_MODEL="$( ./test/scripts/build_ollama_uncertainty_adapter.sh )" +export MELLEA_OLLAMA_UNCERTAINTY_MODEL ``` Install the lightweight Hugging Face Hub dependency that retrieves the diff --git a/docs/examples/intrinsics/uncertainty_ollama.py b/docs/examples/intrinsics/uncertainty_ollama.py index ffcbecb78d..c93f0f4f66 100644 --- a/docs/examples/intrinsics/uncertainty_ollama.py +++ b/docs/examples/intrinsics/uncertainty_ollama.py @@ -28,7 +28,7 @@ adapter_model = os.environ.get("MELLEA_OLLAMA_UNCERTAINTY_MODEL") if adapter_model is None: builder = ( - Path(__file__).parents[2] / "test/scripts/build_ollama_uncertainty_adapter.sh" + Path(__file__).parents[3] / "test/scripts/build_ollama_uncertainty_adapter.sh" ) adapter_model = subprocess.run( [builder], check=True, stdout=subprocess.PIPE, text=True diff --git a/mellea/backends/adapters/catalog.py b/mellea/backends/adapters/catalog.py index b01cd38898..426d06d4c1 100644 --- a/mellea/backends/adapters/catalog.py +++ b/mellea/backends/adapters/catalog.py @@ -162,11 +162,15 @@ def effective_capability(self) -> str: ############################################ # Core adapter functions ############################################ + # context-attribution, citations, and hallucination_detection publish only + # a `lora/` subdirectory on the Hub for granite-4.1-3b (no `alora/`), unlike + # every other entry below. IntrinsicsCatalogEntry( name="context-attribution", capability="context_attribution", repo_id=_CORE_R1_REPO, revision=_CORE_R1_SHA, + adapter_types=(AdapterType.LORA,), ), IntrinsicsCatalogEntry( name="requirement-check", @@ -181,9 +185,17 @@ def effective_capability(self) -> str: # RAG adapter functions ############################################ IntrinsicsCatalogEntry(name="answerability", repo_id=_RAG_REPO, revision=_RAG_SHA), - IntrinsicsCatalogEntry(name="citations", repo_id=_RAG_REPO, revision=_RAG_SHA), IntrinsicsCatalogEntry( - name="hallucination_detection", repo_id=_RAG_REPO, revision=_RAG_SHA + name="citations", + repo_id=_RAG_REPO, + revision=_RAG_SHA, + adapter_types=(AdapterType.LORA,), + ), + IntrinsicsCatalogEntry( + name="hallucination_detection", + repo_id=_RAG_REPO, + revision=_RAG_SHA, + adapter_types=(AdapterType.LORA,), ), IntrinsicsCatalogEntry( name="query_clarification", repo_id=_RAG_REPO, revision=_RAG_SHA diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index 7de82ec6e7..e4725ceec5 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -32,6 +32,7 @@ ) from ..core.base import AbstractMelleaTool from ..formatters import ChatFormatter, TemplateFormatter, granite as granite_formatters +from ..formatters.granite.intrinsics.input import move_documents_to_message from ..helpers import ( DEFAULT_CHUNK_TIMEOUT, ClientCache, @@ -216,6 +217,10 @@ class OllamaModelBackend(FormatterBackend, AdapterMixin): name used to find an adapter's `io.yaml` (for example, `"granite-4.1-3b"`). Required when `model_id` is itself a bundled Ollama adapter tag rather than a known base-model tag. + default_to_constraint_checking_alora (bool): If `False`, deactivates + automatic rerouting of a plain `Requirement` to the + `requirement-check` adapter. Matches `OpenAIBackend` and + `LocalHFBackend`. Attributes: to_mellea_model_opts_map (dict): Mapping from Ollama-specific option names @@ -240,6 +245,7 @@ def __init__( timeout: float | None = 300.0, adapter_models: dict[str, str] | None = None, adapter_base_model_name: str | None = None, + default_to_constraint_checking_alora: bool = True, ): """Initialize an Ollama backend, connecting to the server and pulling the model if needed.""" super().__init__( @@ -268,6 +274,7 @@ def __init__( self._composed_adapter_configs: dict[str, dict] = {} self._adapter_models: dict[str, str] = adapter_models or {} self._adapter_base_model_name = adapter_base_model_name + self.default_to_constraint_checking_alora = default_to_constraint_checking_alora # Setup the client and ensure that we have the model available. self._base_url = base_url @@ -408,21 +415,33 @@ def resolve_adapter(self, name: str) -> _AdapterCore: _AdapterCore: The registered server-mediated adapter. Raises: - ValueError: If the catalogued `io.yaml` is invalid. + ValueError: If `name` has no entry in `adapter_models`, or if the + catalogued `io.yaml` is invalid. KeyError: If registration does not yield a discoverable adapter. """ found = self._find_adapter(name) if found is not None: return found + if name not in self._adapter_models: + raise ValueError( + f"No Ollama model tag configured for adapter function {name!r}; " + "add one to `adapter_models` before resolving it." + ) + metadata = fetch_intrinsic_metadata(name) + # Prefer aLoRA (shares the base model's KV cache) where the catalog + # says it's published; fall back to LoRA otherwise (e.g. citations, + # hallucination_detection, context-attribution ship LoRA only). + use_alora = AdapterType.ALORA in metadata.adapter_types + adapter_type = AdapterType.ALORA if use_alora else AdapterType.LORA try: config_path = granite_formatters.intrinsics.obtain_io_yaml( name, self.base_model_name, metadata.repo_id, revision=metadata.revision, - alora=True, + alora=use_alora, ) except ModuleNotFoundError as e: if e.name != "huggingface_hub": @@ -443,7 +462,7 @@ def resolve_adapter(self, name: str) -> _AdapterCore: _AdapterCore( identity=Identity( name=name, - adapter_type=AdapterType.ALORA.value, + adapter_type=adapter_type.value, capability=metadata.effective_capability, ), io_contract=get_io_contract(name), @@ -689,6 +708,15 @@ async def _generate_from_intrinsic( rewritten = rewriter.transform(request_json, **action.intrinsic_kwargs) + # Ollama's chat API has no extra_body passthrough (unlike the OpenAI- + # compatible endpoints this rewriter otherwise targets), so any + # documents the io.yaml didn't already fold into a message via + # `docs_as_message` must be folded in here or they're silently dropped. + if rewritten.extra_body is not None and rewritten.extra_body.documents: + rewritten = move_documents_to_message( # type: ignore[assignment] + rewritten, "string" + ) + tools: dict[str, AbstractMelleaTool] = dict() if tool_calls: add_tools_from_model_options(tools, model_options) @@ -813,14 +841,17 @@ async def _generate_from_context( _model_id_str = str(getattr(self, "model_id", "unknown")) with with_context(request_id=generate_request_id(), model_id=_model_id_str): + await self.do_generate_walk(action) + model_opts = self._simplify_and_merge(model_options) # Requirements can be automatically rerouted to a requirement adapter. if isinstance(action, Requirement): - reroute_to_alora = isinstance(action, ALoraRequirement) + reroute_to_alora = self.default_to_constraint_checking_alora adapter_name = "requirement-check" if isinstance(action, ALoraRequirement): + reroute_to_alora = True adapter_name = action.intrinsic_name alora_action = action else: @@ -845,7 +876,7 @@ async def _generate_from_context( await asyncio.to_thread(self.resolve_adapter, adapter_name) alora_req_adapter = self._find_adapter(adapter_name, search_types) if alora_req_adapter is None: - if reroute_to_alora: + if reroute_to_alora and isinstance(action, ALoraRequirement): MelleaLogger.get_logger().warning( f"attempted to use an AloraRequirement but backend {self} " f"doesn't have the specified adapter added {adapter_name}; " diff --git a/test/backends/test_ollama.py b/test/backends/test_ollama.py index d3d63f5579..2794dcdd62 100644 --- a/test/backends/test_ollama.py +++ b/test/backends/test_ollama.py @@ -105,9 +105,15 @@ def uncertainty_adapter_model() -> str: """Build and return the official uncertainty aLoRA Ollama model tag.""" if configured_model := os.environ.get("MELLEA_OLLAMA_UNCERTAINTY_MODEL"): return configured_model - completed = subprocess.run( - [_UNCERTAINTY_ADAPTER_BUILDER], check=True, stdout=subprocess.PIPE, text=True - ) + try: + completed = subprocess.run( + [_UNCERTAINTY_ADAPTER_BUILDER], + check=True, + stdout=subprocess.PIPE, + text=True, + ) + except subprocess.CalledProcessError as e: + pytest.skip(f"uncertainty adapter build failed: {e}") return completed.stdout.strip() diff --git a/test/backends/test_ollama_intrinsics_unit.py b/test/backends/test_ollama_intrinsics_unit.py index d20cb352e0..f936b68265 100644 --- a/test/backends/test_ollama_intrinsics_unit.py +++ b/test/backends/test_ollama_intrinsics_unit.py @@ -218,6 +218,36 @@ async def _run_intrinsic( # --------------------------------------------------------------------------- +async def test_documents_folded_into_message_when_not_docs_as_message(): + """Documents extra_body has no Ollama transport, so they must land in a message. + + `_SIMPLE_CONFIG` sets no `docs_as_message`, mirroring the shipped + answerability/citations/etc. io.yaml configs. Without folding, Ollama + receives zero documents and still returns a schema-valid, meaningless score. + """ + backend = _make_backend_with_adapter(_SIMPLE_CONFIG) + context = ChatContext().add( + Message("user", "What is the square root of 4?", documents=["4 is 2 squared."]) + ) + mock_chat = AsyncMock(return_value=_simple_chat_response()) + mock_client = MagicMock() + mock_client.chat = mock_chat + + with patch.object( + OllamaModelBackend, + "_async_client", + new_callable=PropertyMock, + return_value=mock_client, + ): + mot, _ = await mfuncs.aact( + Intrinsic("uncertainty"), context, backend, strategy=None + ) + await mot.avalue() + + messages = mock_chat.call_args.kwargs["messages"] + assert any("4 is 2 squared." in m["content"] for m in messages) + + async def test_instruction_appended_as_last_message(): """The io.yaml instruction (the adapter activation text) is the final user message.""" backend = _make_backend_with_adapter(_SIMPLE_CONFIG) @@ -408,13 +438,50 @@ def test_add_adapter_requires_io_yaml_config(): backend.add_adapter(adapter) +def test_resolve_adapter_requires_configured_model_tag(): + """Resolving without a configured tag raises instead of silently succeeding. + + A cold resolve with no `adapter_models` entry used to register the io.yaml + anyway; generation would then fall back to the base `model_id`, producing a + schema-valid score from a model with no adapter weights. + """ + backend = _make_backend() + + with pytest.raises(ValueError, match=r"No Ollama model tag configured"): + backend.resolve_adapter("uncertainty") + + +def test_resolve_adapter_falls_back_to_lora_when_alora_unavailable( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + """citations/hallucination_detection/context-attribution ship LoRA only.""" + config_path = tmp_path / "io.yaml" + config_path.write_text(json.dumps(_SIMPLE_CONFIG), encoding="utf-8") + backend = _make_backend(adapter_models={"citations": _ADAPTER_TAG}) + recorded_alora: list[bool] = [] + + def _fake_obtain_io_yaml(*_args, alora: bool, **_kwargs) -> Path: + recorded_alora.append(alora) + return config_path + + monkeypatch.setattr( + "mellea.backends.ollama.granite_formatters.intrinsics.obtain_io_yaml", + _fake_obtain_io_yaml, + ) + + adapter = backend.resolve_adapter("citations") + + assert recorded_alora == [False] + assert adapter.identity.adapter_type == "lora" + + def test_resolve_adapter_registers_server_mediated_adapter( monkeypatch: pytest.MonkeyPatch, tmp_path: Path ): """Resolving a capability creates the modern adapter shape without a server call.""" config_path = tmp_path / "io.yaml" config_path.write_text(json.dumps(_SIMPLE_CONFIG), encoding="utf-8") - backend = _make_backend() + backend = _make_backend(adapter_models={"uncertainty": _ADAPTER_TAG}) monkeypatch.setattr( "mellea.backends.ollama.granite_formatters.intrinsics.obtain_io_yaml", lambda *_args, **_kwargs: config_path, @@ -432,7 +499,7 @@ def test_resolve_adapter_explains_missing_huggingface_extra( monkeypatch: pytest.MonkeyPatch, ): """The optional dependency failure names the extra users need to install.""" - backend = _make_backend() + backend = _make_backend(adapter_models={"uncertainty": _ADAPTER_TAG}) monkeypatch.setattr( "mellea.backends.ollama.granite_formatters.intrinsics.obtain_io_yaml", lambda *_args, **_kwargs: (_ for _ in ()).throw( diff --git a/test/scripts/build_ollama_uncertainty_adapter.sh b/test/scripts/build_ollama_uncertainty_adapter.sh index 7c3b05b991..cb2e26e66d 100755 --- a/test/scripts/build_ollama_uncertainty_adapter.sh +++ b/test/scripts/build_ollama_uncertainty_adapter.sh @@ -5,7 +5,8 @@ # Build a local Ollama model from the official pinned Granite uncertainty aLoRA. # # Prints only the resulting Ollama tag to stdout so callers can assign it: -# export MELLEA_OLLAMA_UNCERTAINTY_MODEL="$(./test/scripts/build_ollama_uncertainty_adapter.sh)" +# MELLEA_OLLAMA_UNCERTAINTY_MODEL="$(./test/scripts/build_ollama_uncertainty_adapter.sh)" +# export MELLEA_OLLAMA_UNCERTAINTY_MODEL set -euo pipefail diff --git a/test/scripts/run_tests_with_ollama_and_vllm.sh b/test/scripts/run_tests_with_ollama_and_vllm.sh index 65106b22cb..d678b1c116 100755 --- a/test/scripts/run_tests_with_ollama_and_vllm.sh +++ b/test/scripts/run_tests_with_ollama_and_vllm.sh @@ -195,9 +195,10 @@ done log "All ollama models ready." -export MELLEA_OLLAMA_UNCERTAINTY_MODEL="$( +MELLEA_OLLAMA_UNCERTAINTY_MODEL="$( ./test/scripts/build_ollama_uncertainty_adapter.sh )" +export MELLEA_OLLAMA_UNCERTAINTY_MODEL OLLAMA_MODEL_LIST+=("$MELLEA_OLLAMA_UNCERTAINTY_MODEL") # --- Warm up models (first load into memory is slow) --- From 203da8c906d8bf5f23d5414554511eaba1db4258 Mon Sep 17 00:00:00 2001 From: Nigel Jones Date: Tue, 15 Sep 2026 17:03:03 +0000 Subject: [PATCH 8/9] fix(backends): address second review round on Ollama adapter functions Fix three issues from jakelorocco's review of the Ollama adapter-function support: - _generate_from_intrinsic() looked up the adapter's model tag via self._adapter_models.get(name, self._model_id), silently falling back to the plain base model for an adapter registered directly via add_adapter() (bypassing adapter_models). The rewriter still built the adapter's activation prompt and enforced its response schema, so the base model returned a schema-valid but meaningless answer with no adapter weights behind it. Now raises instead, matching the guard resolve_adapter() already has for the same failure class. - The opportunistic resolve_adapter() call in the automatic requirement-check reroute path was unguarded; a network or config error during resolution killed the whole generate call instead of falling back to regular generation as intended. Wrapped in try/except. - That same reroute path skipped the opportunistic resolve whenever any explicit adapter_types override was given, even one that included aLoRA and would have been satisfied by the resolve. Now only skips when the override excludes aLoRA. - Investigating the docs_as_message question surfaced an unrelated bug: the result processor was constructed from the adapter's unmodified io.yaml config while the rewriter separately folded documents into the message, so the two disagreed about where documents lived. Citations and hallucination flags decoded as empty on every call, silently. Fixed by forcing docs_as_message onto the config before constructing both the rewriter and the result processor, and removed the now-redundant manual fold. Updated test/backends/test_ollama_intrinsics_unit.py to match: replaced test_adapter_model_tag_defaults_to_model_id (asserted the old silent fallback) with test_generation_without_configured_tag_raises, defaulted the adapter fixture to a configured tag so unrelated tests aren't affected, and added coverage for the explicit-adapter-types resolve behaviour. Assisted-by: Claude Code Co-Authored-By: Claude Sonnet 5 Signed-off-by: Nigel Jones --- mellea/backends/ollama.py | 60 +++++++--- test/backends/test_ollama_intrinsics_unit.py | 111 +++++++++++++++++-- 2 files changed, 149 insertions(+), 22 deletions(-) diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index e4725ceec5..dad4a1fe63 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -32,7 +32,6 @@ ) from ..core.base import AbstractMelleaTool from ..formatters import ChatFormatter, TemplateFormatter, granite as granite_formatters -from ..formatters.granite.intrinsics.input import move_documents_to_message from ..helpers import ( DEFAULT_CHUNK_TIMEOUT, ClientCache, @@ -675,6 +674,17 @@ async def _generate_from_intrinsic( "register it via add_adapter() or resolve_adapter()." ) + # Ollama's chat API has no extra_body passthrough (unlike the OpenAI- + # compatible endpoints this rewriter otherwise targets). If the + # io.yaml leaves `docs_as_message` unset, force "roles" here so the + # rewriter (encode) and result processor (decode) agree on where + # documents live. Encoding them into the message without updating + # this config would leave the result processor still looking for + # documents in `extra_body`, where they no longer are — it would + # silently decode zero document sentences instead of raising. + if not intrinsic_config.get("docs_as_message"): + intrinsic_config = {**intrinsic_config, "docs_as_message": "roles"} + rewriter = granite_formatters.IntrinsicsRewriter( config_dict=intrinsic_config, model_name=adapter.identity.name ) @@ -708,15 +718,6 @@ async def _generate_from_intrinsic( rewritten = rewriter.transform(request_json, **action.intrinsic_kwargs) - # Ollama's chat API has no extra_body passthrough (unlike the OpenAI- - # compatible endpoints this rewriter otherwise targets), so any - # documents the io.yaml didn't already fold into a message via - # `docs_as_message` must be folded in here or they're silently dropped. - if rewritten.extra_body is not None and rewritten.extra_body.documents: - rewritten = move_documents_to_message( # type: ignore[assignment] - rewritten, "string" - ) - tools: dict[str, AbstractMelleaTool] = dict() if tool_calls: add_tools_from_model_options(tools, model_options) @@ -732,8 +733,20 @@ async def _generate_from_intrinsic( logprobs = model_opts.pop("logprobs", None) top_logprobs = model_opts.pop("top_logprobs", None) - # each adapter function is served by its own ollama model tag - model = self._adapter_models.get(action.intrinsic_name, self._model_id) + # Each adapter function is served by its own ollama model tag. + # `ServerMediatedBinding` doesn't yet carry its own server target + # (issue #1633), so it's not enough that `_find_adapter` above found + # a registered adapter — that only proves it's discoverable, not that + # we know which tag serves it. An adapter added directly via + # add_adapter() (bypassing adapter_models) would otherwise silently + # fall back to the plain base model here. + if action.intrinsic_name not in self._adapter_models: + raise ValueError( + f"No Ollama model tag configured for adapter function " + f"{action.intrinsic_name!r}; add one to `adapter_models` " + "before generating with it." + ) + model = self._adapter_models[action.intrinsic_name] messages_dicts = [] for m in rewritten.messages: @@ -867,14 +880,31 @@ async def _generate_from_context( else ("alora",) ) alora_req_adapter = self._find_adapter(adapter_name, search_types) + # resolve_adapter() has no way to request a specific adapter + # type — it picks aLoRA over LoRA whenever the catalog says + # aLoRA is published. Skip the opportunistic resolve only + # when an explicit override would rule out whatever it + # produces (i.e. the override excludes aLoRA); an explicit + # override that includes aLoRA is exactly what a cold resolve + # would satisfy, so it shouldn't block the attempt. if ( alora_req_adapter is None and reroute_to_alora and adapter_name in self._adapter_models - and not explicit_types + and (not explicit_types or AdapterType.ALORA in explicit_types) ): - await asyncio.to_thread(self.resolve_adapter, adapter_name) - alora_req_adapter = self._find_adapter(adapter_name, search_types) + try: + await asyncio.to_thread(self.resolve_adapter, adapter_name) + except Exception as e: + MelleaLogger.get_logger().warning( + f"failed to resolve adapter {adapter_name!r} while " + f"attempting automatic requirement-check rerouting; " + f"defaulting to regular generation: {e}" + ) + else: + alora_req_adapter = self._find_adapter( + adapter_name, search_types + ) if alora_req_adapter is None: if reroute_to_alora and isinstance(action, ALoraRequirement): MelleaLogger.get_logger().warning( diff --git a/test/backends/test_ollama_intrinsics_unit.py b/test/backends/test_ollama_intrinsics_unit.py index f936b68265..3a8b508e4a 100644 --- a/test/backends/test_ollama_intrinsics_unit.py +++ b/test/backends/test_ollama_intrinsics_unit.py @@ -9,7 +9,8 @@ - routes the call to the Ollama model tag registered for the adapter function - applies the `IntrinsicsResultProcessor` to the raw response - user-provided model options override io.yaml parameter defaults -- raises when no adapter is registered or streaming is requested +- raises when no adapter is registered, no model tag is configured for it, or + streaming is requested """ import json @@ -22,6 +23,7 @@ from mellea.backends import ModelOption from mellea.backends.adapters import ( Adapter, + AdapterType, Identity, ServerMediatedBinding, get_io_contract, @@ -172,7 +174,16 @@ def _make_backend_with_adapter( model_options: dict | None = None, adapter_models: dict | None = None, ) -> OllamaModelBackend: - """Return an OllamaModelBackend with a registered uncertainty adapter.""" + """Return an OllamaModelBackend with a registered uncertainty adapter. + + Defaults `adapter_models` to a tag for "uncertainty" so tests that don't + care about model-tag routing aren't affected by it: generating against an + adapter with no configured tag now raises (it would otherwise silently + run against the plain base model with none of the adapter's weights). + Pass `adapter_models={}` explicitly to opt into that unconfigured case. + """ + if adapter_models is None: + adapter_models = {"uncertainty": _ADAPTER_TAG} backend = _make_backend(model_options=model_options, adapter_models=adapter_models) adapter = Adapter( identity=Identity(name="uncertainty", adapter_type="alora"), @@ -282,12 +293,20 @@ async def test_adapter_model_tag_used(): assert mock_chat.call_args.kwargs["model"] == _ADAPTER_TAG -async def test_adapter_model_tag_defaults_to_model_id(): - """Adapter functions without a registered tag run against the backend's model.""" - backend = _make_backend_with_adapter(_SIMPLE_CONFIG) - _, mock_chat = await _run_intrinsic(backend, _simple_chat_response()) +async def test_generation_without_configured_tag_raises(): + """An adapter with no `adapter_models` entry raises rather than silently + running against the plain base model. - assert mock_chat.call_args.kwargs["model"] == "granite4.1:3b" + The rewriter still builds the adapter's activation prompt and enforces + its response schema either way, so a silent fallback would return a + schema-valid, meaningless answer from a model that never saw the + adapter's weights — the same failure class already guarded against in + `resolve_adapter()`. + """ + backend = _make_backend_with_adapter(_SIMPLE_CONFIG, adapter_models={}) + + with pytest.raises(ValueError, match="No Ollama model tag configured"): + await _run_intrinsic(backend, _simple_chat_response()) async def test_alora_requirement_resolves_mapped_adapter( @@ -323,6 +342,84 @@ async def test_alora_requirement_resolves_mapped_adapter( assert backend.list_adapters() == ["requirement-check_alora"] +async def test_alora_requirement_with_explicit_alora_type_still_resolves( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + """An explicit `adapter_types=(ALORA,)` override doesn't block the resolve. + + `resolve_adapter()` has no way to request a specific type — it always + prefers aLoRA when the catalog publishes it. An override that asks for + exactly that is satisfied by a cold resolve, so it shouldn't be treated + like an override the resolve can't satisfy. + """ + config_path = tmp_path / "io.yaml" + config_path.write_text(json.dumps(_SIMPLE_CONFIG), encoding="utf-8") + monkeypatch.setattr( + "mellea.backends.ollama.granite_formatters.intrinsics.obtain_io_yaml", + lambda *_args, **_kwargs: config_path, + ) + backend = _make_backend(adapter_models={"requirement-check": _ADAPTER_TAG}) + mock_chat = AsyncMock(return_value=_simple_chat_response()) + mock_client = MagicMock() + mock_client.chat = mock_chat + + with patch.object( + OllamaModelBackend, + "_async_client", + new_callable=PropertyMock, + return_value=mock_client, + ): + mot, _ = await mfuncs.aact( + ALoraRequirement( + "The response is correct.", adapter_types=(AdapterType.ALORA,) + ), + _make_context(), + backend, + strategy=None, + ) + await mot.avalue() + + assert mock_chat.call_args.kwargs["model"] == _ADAPTER_TAG + + +async def test_alora_requirement_with_explicit_lora_type_skips_resolve( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path +): + """An explicit override that excludes aLoRA does not trigger a cold resolve. + + `resolve_adapter()` can't be steered to a specific type, so resolving + here could register the wrong one; falling back to regular generation is + correct instead. + """ + config_path = tmp_path / "io.yaml" + config_path.write_text(json.dumps(_SIMPLE_CONFIG), encoding="utf-8") + resolve_calls: list[str] = [] + + def _fake_obtain_io_yaml(name, *_args, **_kwargs) -> Path: + resolve_calls.append(name) + return config_path + + monkeypatch.setattr( + "mellea.backends.ollama.granite_formatters.intrinsics.obtain_io_yaml", + _fake_obtain_io_yaml, + ) + backend = _make_backend(adapter_models={"requirement-check": _ADAPTER_TAG}) + action = ALoraRequirement( + "The response is correct.", adapter_types=(AdapterType.LORA,) + ) + ctx = _make_context() + + with patch.object( + OllamaModelBackend, "generate_from_chat_context", new_callable=AsyncMock + ) as mock_standard: + mock_standard.return_value = MagicMock() + await backend._generate_from_context(action, ctx, model_options={}) + + assert resolve_calls == [] + assert backend.list_adapters() == [] + mock_standard.assert_awaited_once() + + async def test_result_processor_applied(): """Full uncertainty config: likelihood + project transforms produce the expected JSON.""" backend = _make_backend_with_adapter(_UNCERTAINTY_CONFIG) From 17cf8a01b45399e9e892b49358170d3c06c11a71 Mon Sep 17 00:00:00 2001 From: Nigel Jones Date: Wed, 16 Sep 2026 07:41:47 +0000 Subject: [PATCH 9/9] fix(backends): pass ModelOutputThunk to send_to_queue in Ollama adapter path Upstream's send_to_queue() (#1631) now takes the ModelOutputThunk directly instead of a bare queue, to stamp TTFB at provider receipt. The adapter function code path added by this PR still passed output._gen.queue, which broke type-checking after merging upstream/main. Assisted-by: Claude Code Co-Authored-By: Claude Sonnet 5 Signed-off-by: Nigel Jones --- mellea/backends/ollama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mellea/backends/ollama.py b/mellea/backends/ollama.py index af45bc4b02..99f3894616 100644 --- a/mellea/backends/ollama.py +++ b/mellea/backends/ollama.py @@ -810,7 +810,7 @@ async def granite_formatters_processing( output._gen.generate = asyncio.create_task( send_to_queue( chat_response, - output._gen.queue, + output, chunk_timeout=model_opts.get( ModelOption.STREAM_TIMEOUT, DEFAULT_CHUNK_TIMEOUT ),