diff --git a/docs/docs/advanced/intrinsics.md b/docs/docs/advanced/intrinsics.md index 298beb27dc..a2329ba5e3 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 b6198cdb23..e48d08b840 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