diff --git a/.github/workflows/quality.yml b/.github/workflows/quality.yml index 1ec9eec6a9..bba043b5db 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; } @@ -153,6 +153,25 @@ 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 + # 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/AGENTS.md b/AGENTS.md index 94e8b20aa2..b37d943fe6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -232,15 +232,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 89726d9f69..715ae55fc2 100644 --- a/docs/docs/advanced/intrinsics.md +++ b/docs/docs/advanced/intrinsics.md @@ -10,11 +10,12 @@ 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 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 +26,12 @@ 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. +> 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: @@ -38,6 +43,66 @@ 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 +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 +catalogued `io.yaml`: + +```bash +uv sync --extra switch +``` + +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 + +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=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"], + }, +) +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)) +``` + +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. + ## Use a local Granite Switch checkpoint Granite Switch checkpoints contain their adapter functions already. Pass the @@ -342,8 +407,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 c2f27f0d45..dc772df7df 100644 --- a/docs/docs/advanced/lora-and-alora-adapters.md +++ b/docs/docs/advanced/lora-and-alora-adapters.md @@ -15,8 +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 Ollama, OpenAI, or other remote backends. +> **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 48bac52f6b..bfc52bb5e4 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 new file mode 100644 index 0000000000..c93f0f4f66 --- /dev/null +++ b/docs/examples/intrinsics/uncertainty_ollama.py @@ -0,0 +1,54 @@ +# pytest: e2e, ollama + +"""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. 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`. + +To run this script from the root of the Mellea source tree, use the command: +``` +uv run python docs/examples/intrinsics/uncertainty_ollama.py +``` +""" + +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[3] / "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=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": adapter_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/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 864d05dbe8..99f3894616 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 @@ -30,18 +31,24 @@ 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._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 @@ -49,6 +56,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 +185,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 +207,19 @@ 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. `"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. + 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 @@ -173,6 +233,8 @@ class OllamaModelBackend(FormatterBackend): 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, @@ -180,6 +242,9 @@ def __init__( base_url: str | None = None, model_options: dict | None = None, 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__( @@ -204,6 +269,12 @@ def __init__( self._model_id: str = ollama_model_id self._provider: str = "ollama" + 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 + 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 self._timeout = timeout @@ -256,6 +327,153 @@ 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. + """ + 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) + 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, *, config: dict | None = None) -> None: + """Register an adapter with this backend. + + 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): 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` does not use `ServerMediatedBinding`. + ValueError: If `config` is omitted for a composed adapter. + """ + if not isinstance(adapter, _AdapterCore) or not isinstance( + adapter.weights, ServerMediatedBinding + ): + raise TypeError( + "OllamaModelBackend only supports composed Adapters with a " + "ServerMediatedBinding. " + f"Got: {type(adapter).__name__}" + ) + 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. + + Returns: + list[str]: Qualified adapter names. + """ + 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 `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=use_alora, + ) + 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=adapter_type.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: @@ -394,6 +612,214 @@ 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. + ValueError: If the registered adapter has no cached `io.yaml`. + """ + 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}" + ) + 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()." + ) + + # 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 + ) + 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. + # `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: + 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, + 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 +854,84 @@ 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 = 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: + assert action.description is not None, ( + "must have a description when generating from a requirement" + ) + alora_action = ALoraRequirement(action.description, adapter_name) + + 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) + # 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 or AdapterType.ALORA in explicit_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( + 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 +991,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. @@ -893,9 +1391,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 @@ -928,7 +1427,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 29cd593422..497e961eb3 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: @@ -95,6 +105,23 @@ 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 + 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() + + @pytest.fixture(scope="function") def thinking_session(): """Fresh Ollama session with THINKING on at construction time. @@ -293,6 +320,74 @@ async def test_async_avalue(session) -> None: assert mot1.generation.ttfb_ms is 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=uncertainty_adapter_model, + adapter_base_model_name="granite-4.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 +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 new file mode 100644 index 0000000000..3a8b508e4a --- /dev/null +++ b/test/backends/test_ollama_intrinsics_unit.py @@ -0,0 +1,641 @@ +# 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, no model tag is configured for it, or + streaming is requested +""" + +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 import ( + Adapter, + AdapterType, + 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 +# --------------------------------------------------------------------------- + +_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 = "mellea-test/uncertainty-alora:latest" + +# --------------------------------------------------------------------------- +# 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. + + 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"), + io_contract=get_io_contract("uncertainty"), + weights=ServerMediatedBinding(), + ) + backend.add_adapter(adapter, config=config) + 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_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) + _, 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_generation_without_configured_tag_raises(): + """An adapter with no `adapter_models` entry raises rather than silently + running against the plain base model. + + 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( + 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_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) + 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_server_mediated_adapter(): + backend = _make_backend_with_adapter(_SIMPLE_CONFIG) + assert backend.list_adapters() == ["uncertainty_alora"] + + +def test_add_adapter_rejects_other_adapter_types(): + backend = _make_backend() + 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_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(adapter_models={"uncertainty": _ADAPTER_TAG}) + 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(adapter_models={"uncertainty": _ADAPTER_TAG}) + 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" + + +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 diff --git a/test/scripts/build_ollama_uncertainty_adapter.sh b/test/scripts/build_ollama_uncertainty_adapter.sh new file mode 100755 index 0000000000..cb2e26e66d --- /dev/null +++ b/test/scripts/build_ollama_uncertainty_adapter.sh @@ -0,0 +1,120 @@ +#!/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: +# MELLEA_OLLAMA_UNCERTAINTY_MODEL="$(./test/scripts/build_ollama_uncertainty_adapter.sh)" +# export MELLEA_OLLAMA_UNCERTAINTY_MODEL + +set -euo pipefail + +log() { echo "[ollama-adapter] $*" >&2; } +die() { log "ERROR: $*"; exit 1; } + +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 + 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" "$BASE_REPO" "$BASE_REVISION" "$ADAPTER_REPO" "$ADAPTER_REVISION" "$ADAPTER_PATH" <<'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]) +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=base_repo, + revision=base_revision, + 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=adapter_repo, + filename=f"{adapter_path}/{filename}", + revision=adapter_revision, + ) + 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" + + METADATA_DUMP="${BUILD_CACHE}/uncertainty-alora-metadata.txt" + "$CONVERTER_VENV/bin/python" \ + "$LLAMA_CPP_DIR/gguf-py/gguf/scripts/gguf_dump.py" "$ADAPTER_GGUF" \ + > "$METADATA_DUMP" + if ! grep -q "adapter.alora.invocation_tokens" "$METADATA_DUMP"; 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..d678b1c116 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,12 @@ done log "All ollama models ready." +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) --- # 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.