From eb5119d75ea569157eb424e5ae61516811bd3c62 Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Sat, 24 Oct 2026 21:04:55 +0200 Subject: [PATCH 01/11] Add apple_fm integration for Apple's on-device model Track LanguageModelSession.respond() and stream_response() from apple-fm-sdk: duration, time to first token for streams, prompt size, generation options, and the SDK's error names on failure. Sessions share one model per use case, versioned by the macOS release. Token counts stay empty: the SDK reports no usage, and counting would cost an extra model call per request. Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 8 + README.md | 1 + docs/apple_fm.md | 63 +++++ docs/compatibility.md | 4 +- docs/providers.md | 3 + examples/apple_fm_example.py | 52 ++++ scripts/build_llms_txt.py | 1 + scripts/compat_matrix.py | 4 + tests/compat/test_apple_fm_compat.py | 53 +++++ tests/test_integrations_apple_fm.py | 337 ++++++++++++++++++++++++++ tests/test_lazy_loading.py | 1 + wildedge/client.py | 4 + wildedge/integrations/apple_fm.py | 343 +++++++++++++++++++++++++++ wildedge/integrations/registry.py | 2 + 14 files changed, 875 insertions(+), 1 deletion(-) create mode 100644 docs/apple_fm.md create mode 100644 examples/apple_fm_example.py create mode 100644 tests/compat/test_apple_fm_compat.py create mode 100644 tests/test_integrations_apple_fm.py create mode 100644 wildedge/integrations/apple_fm.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4d7b234..59edcb3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,14 @@ under Unreleased and move into a version section at release time. ## Unreleased +### Added + +- `apple_fm` integration tracks Apple's on-device model through apple-fm-sdk. + `LanguageModelSession.respond()` and `stream_response()` record duration, + time to first token for streams, prompt size and generation options, and + failures record the SDK's error name. Token counts are not recorded because + the SDK doesn't report them. See docs/apple_fm.md. + ## 0.3.0 - 2026-10-24 ### Added diff --git a/README.md b/README.md index 03c63e0..ebff969 100644 --- a/README.md +++ b/README.md @@ -180,6 +180,7 @@ explicit registration through `client.load(...)` or `client.register_model(...)` | Transformers | `transformers` | Automatically through `pipeline` / `from_pretrained` | [Example](examples/transformers_example.py) | | llama.cpp (`llama-cpp-python`) | `llama_cpp` | Automatically when constructing `Llama` | [Example](examples/gguf_example.py) | | MLX-LM | `mlx` | Automatically through `mlx_lm.load` | [Example](examples/mlx_example.py) | +| Apple Foundation Models (`apple-fm-sdk`) | `apple_fm` | Automatically for sessions created after initialization ([docs](docs/apple_fm.md)) | [Example](examples/apple_fm_example.py) | | faster-whisper / MLX Whisper | `faster_whisper` / `mlx_whisper` | Automatically through supported model loaders | [Example](examples/whisper_example.py) | | ONNX Runtime | `onnx` | Automatically when constructing `InferenceSession` | [Example](examples/onnx_example.py) | | timm | `timm` | Automatically through `create_model` | [Example](examples/timm_example.py) | diff --git a/docs/apple_fm.md b/docs/apple_fm.md new file mode 100644 index 0000000..4d82930 --- /dev/null +++ b/docs/apple_fm.md @@ -0,0 +1,63 @@ +# Apple Foundation Models + +The `apple_fm` integration tracks the on-device model behind Apple +Intelligence, called through Apple's Python SDK +([apple-fm-sdk](https://github.com/apple/python-apple-fm-sdk)). It needs macOS +26 or later on Apple silicon with Apple Intelligence turned on, plus Xcode to +install the SDK. + +Initialize the SDK before creating sessions: + +```python +import asyncio + +import apple_fm_sdk as fm +import wildedge + +client = wildedge.init(integrations=["apple_fm"]) + + +async def main(): + session = fm.LanguageModelSession(instructions="Answer in one sentence.") + print(await session.respond("How does rain form?")) + # Each snapshot is the full text so far, not just the new part. + text = "" + async for snapshot in session.stream_response("Name three clouds."): + text = snapshot + print(text) + + +asyncio.run(main()) +client.close(timeout=10) +``` + +## What is recorded + +Each `respond()` or `stream_response()` call records one inference event with +its duration, the prompt's character and word counts, and the +`GenerationOptions` you passed (temperature, maximum response tokens, and +sampling top-k, threshold and seed). Streams also record time to first token. + +Sessions share one model per use case: `apple-foundation-model` for the +general model, and `apple-foundation-model-content-tagging` for +`SystemLanguageModelUseCase.CONTENT_TAGGING`. The model version is the macOS +version, since the model ships with the OS. + +Failed calls record an error event named after the SDK's error, such as +`guardrail_violation`, `exceeded_context_window_size`, `assets_unavailable` or +`refusal`. A model that isn't ready because Apple Intelligence is off fails +this way too. + +## Caveats + +- Token counts and tokens per second are empty. The SDK doesn't report token + usage, and counting tokens with `SystemLanguageModel.token_count()` is an + extra call to the model, so WildEdge doesn't make it for you. +- Prompts and responses are not uploaded. + +## `fm serve` + +On macOS 27, `fm serve` exposes the same model through a local +OpenAI-compatible server. Calls made to it with the `openai` client are +tracked by the `openai` integration. Name the provider, because a local host is +otherwise recorded as `127.0.0.1` (see [Providers](providers.md#other-hosts)). diff --git a/docs/compatibility.md b/docs/compatibility.md index e5aec96..d6daedd 100644 --- a/docs/compatibility.md +++ b/docs/compatibility.md @@ -29,7 +29,8 @@ once the bug is fixed and the marker has to come off. On pull requests into `main` or a release branch, on pushes to them, nightly against `devel`, and on manual runs. Every run covers all the combinations below. Python 3.14 rows are canaries and don't fail the run. mlx runs on -macOS. +macOS. apple_fm runs on macOS 26, where Apple Intelligence is off, so its +test checks that a failed call is recorded as an error event. To run the full matrix locally: `python scripts/run_compat_local.py`. `python scripts/compat_matrix.py table` prints the table below. @@ -69,3 +70,4 @@ To run the full matrix locally: `python scripts/run_compat_local.py`. | `mlx` | `current` | `mlx==0.32.4`, `mlx-lm==0.32.2` | 3.12, 3.13 | | `faster_whisper` | `current` | `faster-whisper==1.2.1`, `ctranslate2==4.8.2`, `numpy==2.2.6` | 3.12, 3.13 | | `mlx_whisper` | `current` | `mlx-whisper==0.4.3`, `mlx==0.32.3`, `numpy==2.2.6` | 3.12, 3.13 | +| `apple_fm` | `current` | `apple-fm-sdk==0.2.1` | 3.12, 3.13 | diff --git a/docs/providers.md b/docs/providers.md index 1c29502..4a9b8c9 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -49,6 +49,9 @@ wildedge.llm_api(model="gemma-7b", provider="ollama") wildedge.register_model(client, source="ollama", model_id="gemma-7b") ``` +The same applies to `fm serve`, Apple's local server for its on-device model on +macOS 27. To track that model without a server, see [Apple Foundation Models](apple_fm.md). + ## Caveats - `tokens_in` / `tokens_out` are populated by every provider above. diff --git a/examples/apple_fm_example.py b/examples/apple_fm_example.py new file mode 100644 index 0000000..09261d7 --- /dev/null +++ b/examples/apple_fm_example.py @@ -0,0 +1,52 @@ +# /// script +# requires-python = ">=3.10" +# dependencies = ["wildedge-sdk", "apple-fm-sdk"] +# +# [tool.uv.sources] +# wildedge-sdk = { path = "..", editable = true } +# /// +"""Apple Foundation Models integration example. + +WildEdge patches apple_fm_sdk.LanguageModelSession at instrumentation time, so +every respond() and stream_response() call on Apple's on-device model is +tracked. + +Run with: uv run apple_fm_example.py +Requires: macOS 26+ on Apple silicon with Apple Intelligence on, and Xcode to +install apple-fm-sdk. Set WILDEDGE_DSN to send events. +""" + +import asyncio + +import apple_fm_sdk as fm + +import wildedge + +client = wildedge.init( + app_version="1.0.0", # uses WILDEDGE_DSN if set; otherwise no-op + integrations="apple_fm", +) + + +async def main() -> None: + model = fm.SystemLanguageModel() + available, reason = model.is_available() + if not available: + raise SystemExit(f"Apple's on-device model is unavailable: {reason}") + + session = fm.LanguageModelSession( + instructions="Answer in one sentence.", model=model + ) + print(await session.respond("Explain transformer attention.")) + + # Each snapshot is the full text so far, not just the new part. + text = "" + async for snapshot in session.stream_response("Write a haiku about edge AI."): + text = snapshot + print(text) + + +try: + asyncio.run(main()) +finally: + client.close(timeout=10) diff --git a/scripts/build_llms_txt.py b/scripts/build_llms_txt.py index d4eaa3e..9797cc1 100644 --- a/scripts/build_llms_txt.py +++ b/scripts/build_llms_txt.py @@ -26,6 +26,7 @@ ("docs/providers.md", "providers", "Providers"), ("docs/streaming.md", "streaming", "Streaming lifecycle"), ("docs/whisper.md", "whisper", "Whisper"), + ("docs/apple_fm.md", "apple_fm", "Apple Foundation Models"), ("docs/compatibility.md", "compatibility", "Compatibility matrix"), ] diff --git a/scripts/compat_matrix.py b/scripts/compat_matrix.py index 97f424e..a413b1a 100644 --- a/scripts/compat_matrix.py +++ b/scripts/compat_matrix.py @@ -62,6 +62,7 @@ "current": ["faster-whisper==1.2.1", "ctranslate2==4.8.2", "numpy==2.2.6"] }, "mlx_whisper": {"current": ["mlx-whisper==0.4.3", "mlx==0.32.3", "numpy==2.2.6"]}, + "apple_fm": {"current": ["apple-fm-sdk==0.2.1"]}, } SUPPORTED_PYTHON = { @@ -77,6 +78,7 @@ "mlx": ["3.12", "3.13"], "faster_whisper": ["3.12", "3.13"], "mlx_whisper": ["3.12", "3.13"], + "apple_fm": ["3.12", "3.13"], } # Integrations without a row of their own, and the row that exercises them. @@ -88,6 +90,8 @@ RUNNER = { "mlx": "macos-14", "mlx_whisper": "macos-14", + # apple-fm-sdk builds Swift bindings at install, which needs Xcode 26. + "apple_fm": "macos-26", } # Python versions run as non-blocking canaries, and only with current pins. diff --git a/tests/compat/test_apple_fm_compat.py b/tests/compat/test_apple_fm_compat.py new file mode 100644 index 0000000..8d5489e --- /dev/null +++ b/tests/compat/test_apple_fm_compat.py @@ -0,0 +1,53 @@ +"""Real apple-fm-sdk sessions. + +CI runners can't turn on Apple Intelligence, so a call there fails with the +SDK's own error. Either way the call must produce exactly one event for the +model it ran on. +""" + +import asyncio +import inspect + +import apple_fm_sdk as fm + + +def test_patched_signatures_match_the_sdk(): + """The patches read these arguments by position.""" + model_params = list(inspect.signature(fm.SystemLanguageModel).parameters) + session_params = list(inspect.signature(fm.LanguageModelSession).parameters) + stream_params = list( + inspect.signature(fm.LanguageModelSession.stream_response).parameters + ) + respond = inspect.signature(fm.LanguageModelSession.respond).parameters + assert model_params[0] == "use_case" + assert session_params[1] == "model" + assert stream_params[1:3] == ["prompt", "options"] + assert respond["options"].kind is inspect.Parameter.KEYWORD_ONLY + + +def test_respond_records_one_event(compat_client): + compat_client.instrument("apple_fm") + model = fm.SystemLanguageModel( + use_case=fm.SystemLanguageModelUseCase.CONTENT_TAGGING + ) + available, reason = model.is_available() + session = fm.LanguageModelSession(model=model) + options = fm.GenerationOptions(maximum_response_tokens=8) + + async def respond(): + try: + return await session.respond("Tag this: rain", options=options) + except fm.FoundationModelsError as exc: + return exc + + result = asyncio.run(respond()) + + [event] = compat_client.events + assert event["model_id"] == "apple-foundation-model-content-tagging" + if available: + assert isinstance(result, str), result + assert event["event_type"] == "inference" + assert event["inference"]["generation_config"]["max_tokens"] == 8 + else: + assert isinstance(result, fm.FoundationModelsError), (reason, result) + assert event["event_type"] == "error" diff --git a/tests/test_integrations_apple_fm.py b/tests/test_integrations_apple_fm.py new file mode 100644 index 0000000..f260dce --- /dev/null +++ b/tests/test_integrations_apple_fm.py @@ -0,0 +1,337 @@ +"""Tests for the Apple Foundation Models (apple-fm-sdk) integration.""" + +from __future__ import annotations + +import asyncio +import enum +from types import ModuleType, SimpleNamespace + +import pytest + +import wildedge.integrations.apple_fm as apple_fm_mod +from wildedge import accelerators +from wildedge.integrations.apple_fm import ( + MODEL_ID, + AppleFmExtractor, + build_input_meta, + model_id_for, + provider_error_code_of, +) +from wildedge.integrations.registry import find_extractor +from wildedge.model import ModelHandle + +# --------------------------------------------------------------------------- +# Fake apple_fm_sdk, matching the signatures of apple-fm-sdk 0.2 +# --------------------------------------------------------------------------- + + +class SystemLanguageModelUseCase(enum.IntEnum): + GENERAL = 0 + CONTENT_TAGGING = 1 + + +class FoundationModelsError(Exception): + pass + + +class GenerationError(FoundationModelsError): + pass + + +class GuardrailViolationError(GenerationError): + pass + + +class ExceededContextWindowSizeError(GenerationError): + pass + + +def make_fake_sdk( + *, + reply="Hello there", + snapshots=("Hel", "Hello", "Hello there"), + error=None, + stream_error=None, +) -> ModuleType: + """A fresh module per test, since the integration patches its classes.""" + + class SystemLanguageModel: + def __init__( + self, + use_case=SystemLanguageModelUseCase.GENERAL, + guardrails=None, + _ptr=None, + ): + self._ptr = object() + + class LanguageModelSession: + def __init__(self, instructions=None, model=None, tools=None, _ptr=None): + self.instructions = instructions + + @classmethod + def from_transcript(cls, transcript, model=None, tools=None): + return cls(_ptr=object()) + + async def respond( + self, + prompt, + generating=None, + *, + schema=None, + json_schema=None, + options=None, + ): + await asyncio.sleep(0) + if error is not None: + raise error + return reply + + async def stream_response(self, prompt, options=None): + for snapshot in snapshots: + await asyncio.sleep(0) + yield snapshot + if stream_error is not None: + raise stream_error + + SystemLanguageModel.__module__ = "apple_fm_sdk.core" + LanguageModelSession.__module__ = "apple_fm_sdk.session" + + sdk = ModuleType("apple_fm_sdk") + sdk.SystemLanguageModel = SystemLanguageModel + sdk.LanguageModelSession = LanguageModelSession + sdk.SystemLanguageModelUseCase = SystemLanguageModelUseCase + return sdk + + +class FakeClient: + """Registers through the real extractor so handles carry real ModelInfo.""" + + def __init__(self, closed=False): + self.closed = closed + self.events: list[dict] = [] + self.handles: dict[str, ModelHandle] = {} + + def register_model(self, obj, *, model_id=None, **_kwargs): + extractor = AppleFmExtractor() + model_id, info = extractor.extract_info(obj, {"id": model_id}) + if model_id not in self.handles: + handle = ModelHandle(model_id, info, self.events.append) + extractor.install_hooks(obj, handle) + self.handles[model_id] = handle + return self.handles[model_id] + + def of_type(self, event_type): + return [e for e in self.events if e["event_type"] == event_type] + + +@pytest.fixture +def install(monkeypatch): + """Activate the integration on a fake SDK and return (sdk, client).""" + + def _install(client=None, **sdk_kwargs): + sdk = make_fake_sdk(**sdk_kwargs) + client = client or FakeClient() + monkeypatch.setattr(apple_fm_mod, "_fm", sdk) + monkeypatch.setattr(apple_fm_mod, "_apple_fm_patched", False) + AppleFmExtractor.install_auto_load_patch(lambda: client) + return sdk, client + + return _install + + +def inference_of(event): + return event["inference"] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def test_model_id_names_non_general_use_cases(): + assert model_id_for(None) == MODEL_ID + assert model_id_for(SystemLanguageModelUseCase.GENERAL) == MODEL_ID + assert ( + model_id_for(SystemLanguageModelUseCase.CONTENT_TAGGING) + == "apple-foundation-model-content-tagging" + ) + + +def test_provider_error_code_follows_the_class_hierarchy(): + class AppGuardrailError(GuardrailViolationError): + pass + + assert provider_error_code_of(GuardrailViolationError()) == "guardrail_violation" + assert provider_error_code_of(AppGuardrailError()) == "guardrail_violation" + assert provider_error_code_of(FoundationModelsError()) is None + + +def test_input_meta_counts_text_parts_only(): + modality, meta = build_input_meta(["Describe this", object(), "photo"]) + assert modality == "text" + assert meta.char_count == len("Describe this photo") + assert meta.word_count == 3 + + modality, meta = build_input_meta(object()) + assert (modality, meta) == ("image", None) + + +# --------------------------------------------------------------------------- +# Extractor +# --------------------------------------------------------------------------- + + +def test_extractor_is_found_for_sdk_objects_only(): + sdk = make_fake_sdk() + assert isinstance(find_extractor(sdk.LanguageModelSession()), AppleFmExtractor) + assert isinstance(find_extractor(sdk.SystemLanguageModel()), AppleFmExtractor) + + class LanguageModelSession: # same name, different package + pass + + assert not AppleFmExtractor().can_handle(LanguageModelSession()) + + +def test_extract_info_describes_the_system_model(monkeypatch): + monkeypatch.setattr(apple_fm_mod.platform, "mac_ver", lambda: ("26.6.2", "", "")) + model_id, info = AppleFmExtractor().extract_info( + make_fake_sdk().SystemLanguageModel(), {} + ) + assert model_id == MODEL_ID + assert info.model_version == "26.6.2" + assert info.model_source == "apple" + assert info.model_format == "apple_fm" + + +def test_registering_a_session_through_the_client(client_with_stubbed_runtime): + session = make_fake_sdk().LanguageModelSession() + handle = client_with_stubbed_runtime.register_model(session) + assert handle.model_id == MODEL_ID + assert handle.info.model_format == "apple_fm" + assert handle.detected_accelerator == accelerators.COREML + assert getattr(session, apple_fm_mod.APPLE_FM_HANDLE_ATTR) is handle + client_with_stubbed_runtime.close() + + +# --------------------------------------------------------------------------- +# Auto-load patch +# --------------------------------------------------------------------------- + + +async def test_respond_records_one_inference(install): + sdk, client = install() + session = sdk.LanguageModelSession(instructions="Be brief.") + options = SimpleNamespace( + temperature=0.3, + maximum_response_tokens=64, + sampling=SimpleNamespace(top=40, probability_threshold=None, seed=7), + ) + + assert ( + await session.respond("How does rain form?", options=options) == "Hello there" + ) + + [event] = client.of_type("inference") + inference = inference_of(event) + assert event["model_id"] == MODEL_ID + assert inference["success"] is True + assert inference["input_modality"] == "text" + assert inference["output_modality"] == "generation" + assert inference["input_meta"]["word_count"] == 4 + assert inference["generation_config"] == { + "temperature": 0.3, + "top_k": 40, + "max_tokens": 64, + "seed": 7, + } + assert "tokens_out" not in inference["output_meta"] + + +async def test_sessions_share_a_handle_per_use_case(install): + sdk, client = install() + tagging = sdk.SystemLanguageModel( + use_case=SystemLanguageModelUseCase.CONTENT_TAGGING + ) + await sdk.LanguageModelSession().respond("a") + await sdk.LanguageModelSession(model=sdk.SystemLanguageModel()).respond("b") + await sdk.LanguageModelSession(model=tagging).respond("c") + await sdk.LanguageModelSession.from_transcript(object()).respond("d") + + model_ids = [e["model_id"] for e in client.of_type("inference")] + assert model_ids == [MODEL_ID, MODEL_ID, f"{MODEL_ID}-content-tagging", MODEL_ID] + assert set(client.handles) == {MODEL_ID, f"{MODEL_ID}-content-tagging"} + + +async def test_respond_error_records_the_sdk_error_name(install): + sdk, client = install(error=GuardrailViolationError("Guardrail violation occurred")) + session = sdk.LanguageModelSession() + + with pytest.raises(GuardrailViolationError): + await session.respond("something unsafe") + + [error] = client.of_type("error") + assert error["error"]["error_code"] == "UNKNOWN" + assert error["error"]["provider_error_code"] == "guardrail_violation" + assert "duration_ms" in error["error"] + assert client.of_type("inference") == [] + + +async def test_stream_records_ttft_and_passes_snapshots_through(install): + sdk, client = install() + session = sdk.LanguageModelSession() + + chunks = [c async for c in session.stream_response("Tell me", None)] + + assert chunks == ["Hel", "Hello", "Hello there"] + [event] = client.of_type("inference") + inference = inference_of(event) + assert inference["success"] is True + assert "time_to_first_token_ms" in inference["output_meta"] + + +async def test_stream_closed_early_is_not_a_success(install): + sdk, client = install() + stream = sdk.LanguageModelSession().stream_response("Tell me") + async for _ in stream: + break + await stream.aclose() + + [event] = client.of_type("inference") + assert inference_of(event)["success"] is False + assert inference_of(event)["error_code"] == "STREAM_CLOSED" + + +async def test_stream_error_records_the_sdk_error_name(install): + sdk, client = install(stream_error=ExceededContextWindowSizeError("too long")) + with pytest.raises(ExceededContextWindowSizeError): + async for _ in sdk.LanguageModelSession().stream_response("Tell me"): + pass + + [error] = client.of_type("error") + assert error["error"]["provider_error_code"] == "exceeded_context_window_size" + assert client.of_type("inference") == [] + + +async def test_closed_client_leaves_sessions_untracked(install): + sdk, client = install(client=FakeClient(closed=True)) + session = sdk.LanguageModelSession() + assert await session.respond("hi") == "Hello there" + assert not hasattr(session, apple_fm_mod.APPLE_FM_HANDLE_ATTR) + assert client.events == [] + + +def test_install_is_idempotent(install): + sdk, _client = install() + respond = sdk.LanguageModelSession.respond + AppleFmExtractor.install_auto_load_patch(lambda: None) + apple_fm_mod._apple_fm_patched = False + AppleFmExtractor.install_auto_load_patch(lambda: None) + assert sdk.LanguageModelSession.respond is respond + + +def test_install_skips_when_sdk_missing(monkeypatch): + monkeypatch.setattr(apple_fm_mod, "_fm", None) + monkeypatch.setattr(apple_fm_mod, "_apple_fm_patched", False) + monkeypatch.setattr(apple_fm_mod, "optional_module", lambda _name: None) + AppleFmExtractor.install_auto_load_patch(lambda: None) + assert apple_fm_mod._apple_fm_patched is False diff --git a/tests/test_lazy_loading.py b/tests/test_lazy_loading.py index c06982b..5bce82b 100644 --- a/tests/test_lazy_loading.py +++ b/tests/test_lazy_loading.py @@ -35,6 +35,7 @@ "llama_cpp", "openai", "anthropic", + "apple_fm_sdk", "huggingface_hub", "numpy", ) diff --git a/wildedge/client.py b/wildedge/client.py index 83862ee..0bafd92 100644 --- a/wildedge/client.py +++ b/wildedge/client.py @@ -789,6 +789,10 @@ def instrument( ``"onnx"`` Replaces ``ort.InferenceSession`` with a timed subclass. Requires ``onnxruntime``. + ``"apple_fm"`` + Patches ``apple_fm_sdk.LanguageModelSession`` to track + ``respond()`` and ``stream_response()`` on Apple's on-device + model. Requires ``apple-fm-sdk``. ``"mlx_whisper"`` / ``"faster_whisper"`` Tracks local speech recognition, including lazy segment consumption. ``"timm"`` diff --git a/wildedge/integrations/apple_fm.py b/wildedge/integrations/apple_fm.py new file mode 100644 index 0000000..044c8ed --- /dev/null +++ b/wildedge/integrations/apple_fm.py @@ -0,0 +1,343 @@ +"""Apple Foundation Models integration (``apple-fm-sdk``). + +Tracks the on-device model behind Apple Intelligence through Apple's Python +SDK. Sessions are tied to a model handle when they are created, and +``respond()`` and ``stream_response()`` record one inference event per call. + +The SDK reports no token usage, and counting tokens takes an extra call into +the model, so token counts are left empty. +""" + +from __future__ import annotations + +import functools +import platform +import sys +import threading +import time +from typing import TYPE_CHECKING, Any + +from wildedge import accelerators +from wildedge.events.error import ErrorCode +from wildedge.events.inference import ( + GenerationConfig, + GenerationOutputMeta, + TextInputMeta, +) +from wildedge.integrations.base import BaseExtractor +from wildedge.integrations.common import ( + AsyncStreamWrapper, + debug_failure, + stream_error_code, +) +from wildedge.loading import optional_module +from wildedge.model import ModelInfo +from wildedge.timing import elapsed_ms + +# Imported when the integration is activated: importing the SDK loads its +# native library. +_fm = sys.modules.get("apple_fm_sdk") + +if TYPE_CHECKING: + from wildedge.model import ModelHandle + +_apple_fm_patched = False +_APPLE_FM_PATCH_LOCK = threading.Lock() +APPLE_FM_PATCH_NAME = "apple_fm_auto_load" +APPLE_FM_HANDLE_ATTR = "__wildedge_apple_fm_handle__" +# SystemLanguageModel passes its use case to native code without keeping it. +APPLE_FM_USE_CASE_ATTR = "__wildedge_apple_fm_use_case__" + +MODEL_ID = "apple-foundation-model" +MODEL_SOURCE = "apple" +MODEL_FORMAT = "apple_fm" +APPLE_FM_PACKAGE = "apple_fm_sdk" + +# The SDK's error classes, named as GenerationErrorCode names them. +PROVIDER_ERROR_CODES = { + "ExceededContextWindowSizeError": "exceeded_context_window_size", + "AssetsUnavailableError": "assets_unavailable", + "GuardrailViolationError": "guardrail_violation", + "UnsupportedGuideError": "unsupported_guide", + "UnsupportedLanguageOrLocaleError": "unsupported_language_or_locale", + "DecodingFailureError": "decoding_failure", + "RateLimitedError": "rate_limited", + "ConcurrentRequestsError": "concurrent_requests", + "RefusalError": "refusal", + "InvalidGenerationSchemaError": "invalid_schema", + "ToolCallError": "tool_call", +} + +debug_apple_fm_failure = functools.partial(debug_failure, "apple_fm") + + +def model_id_for(use_case: object) -> str: + """``apple-foundation-model``, plus the use case when it is not GENERAL.""" + name = getattr(use_case, "name", None) + if not name or name == "GENERAL": + return MODEL_ID + return f"{MODEL_ID}-{name.lower().replace('_', '-')}" + + +def os_version() -> str: + """The system model ships with macOS, so its version is the OS version.""" + return platform.mac_ver()[0] or "unknown" + + +def provider_error_code_of(exc: BaseException) -> str | None: + for cls in type(exc).__mro__: + code = PROVIDER_ERROR_CODES.get(cls.__name__) + if code is not None: + return code + return None + + +def track_apple_fm_error( + handle: ModelHandle, exc: BaseException, duration_ms: int +) -> None: + handle.track_error( + error_code=ErrorCode.UNKNOWN, + error_message=str(exc), + provider_error_code=provider_error_code_of(exc), + duration_ms=duration_ms, + ) + + +def build_input_meta(prompt: object) -> tuple[str, TextInputMeta | None]: + """Input modality and text metadata for a prompt. + + A prompt is a string, an attachment, or a list of both. + """ + parts = prompt if isinstance(prompt, list) else [prompt] + text = " ".join(part for part in parts if isinstance(part, str)) + if not text: + return ("image" if parts else "text"), None + return "text", TextInputMeta( + char_count=len(text), + word_count=len(text.split()), + prompt_type="chat", + ) + + +def build_generation_config(options: object) -> GenerationConfig | None: + if options is None: + return None + sampling = getattr(options, "sampling", None) + return GenerationConfig( + temperature=getattr(options, "temperature", None), + max_tokens=getattr(options, "maximum_response_tokens", None), + top_k=getattr(sampling, "top", None), + top_p=getattr(sampling, "probability_threshold", None), + seed=getattr(sampling, "seed", None), + ) + + +def describe_call( + args: tuple, kwargs: dict, *, options_position: int | None +) -> tuple[str, TextInputMeta | None, GenerationConfig | None]: + """Input modality, input metadata and generation config. Never raises.""" + try: + prompt = kwargs.get("prompt", args[0] if args else None) + options = kwargs.get("options") + if ( + options is None + and options_position is not None + and len(args) > options_position + ): + options = args[options_position] + modality, input_meta = build_input_meta(prompt) + return modality, input_meta, build_generation_config(options) + except Exception as exc: + debug_apple_fm_failure("call metadata", exc) + return "text", None, None + + +class AppleFmStream(AsyncStreamWrapper): + """Stream wrapper that reports the SDK's own error names.""" + + def _error(self, exc: Exception) -> None: + if not self._finished: + self._finished = True + self._safe(track_apple_fm_error, self._handle, exc, elapsed_ms(self._t0)) + + +# --------------------------------------------------------------------------- +# Method patches +# --------------------------------------------------------------------------- + + +def build_model_init_patch(original_init): # type: ignore[no-untyped-def] + @functools.wraps(original_init) + def patched_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] + original_init(self_inner, *args, **kwargs) + use_case = kwargs.get("use_case", args[0] if args else None) + if use_case is not None: + try: + setattr(self_inner, APPLE_FM_USE_CASE_ATTR, use_case) + except Exception as exc: + debug_apple_fm_failure("use case capture", exc) + + return patched_init + + +def build_session_init_patch(original_init, client_ref): # type: ignore[no-untyped-def] + @functools.wraps(original_init) + def patched_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] + original_init(self_inner, *args, **kwargs) + c = client_ref() # type: ignore[call-arg] + if c is None or c.closed: + return + try: + model = kwargs.get("model", args[1] if len(args) > 1 else None) + model_id = model_id_for(getattr(model, APPLE_FM_USE_CASE_ATTR, None)) + handle = c.register_model(self_inner, model_id=model_id) + setattr(self_inner, APPLE_FM_HANDLE_ATTR, handle) + except Exception as exc: + debug_apple_fm_failure("session registration", exc) + + return patched_init + + +def build_respond_patch(original_respond): # type: ignore[no-untyped-def] + @functools.wraps(original_respond) + async def patched_respond(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] + handle: ModelHandle | None = getattr(self_inner, APPLE_FM_HANDLE_ATTR, None) + if handle is None: + return await original_respond(self_inner, *args, **kwargs) + # respond(prompt, generating=None, *, schema, json_schema, options) + modality, input_meta, config = describe_call( + args, kwargs, options_position=None + ) + t0 = time.perf_counter() + try: + result = await original_respond(self_inner, *args, **kwargs) + except Exception as exc: + track_apple_fm_error(handle, exc, elapsed_ms(t0)) + raise + handle.track_inference( + duration_ms=elapsed_ms(t0), + input_modality=modality, + output_modality="generation", + success=True, + input_meta=input_meta, + output_meta=GenerationOutputMeta(task="generation"), + generation_config=config, + ) + return result + + return patched_respond + + +def build_stream_patch(original_stream): # type: ignore[no-untyped-def] + @functools.wraps(original_stream) + def patched_stream(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] + stream = original_stream(self_inner, *args, **kwargs) + handle: ModelHandle | None = getattr(self_inner, APPLE_FM_HANDLE_ATTR, None) + if handle is None: + return stream + # stream_response(prompt, options=None) + modality, input_meta, config = describe_call(args, kwargs, options_position=1) + + # Each chunk is the full text so far, so the first one marks the + # first output and nothing needs to be joined. + def on_done( + duration_ms: int, ttft_ms: int | None, *, outcome: str = "completed" + ) -> None: + handle.track_inference( + duration_ms=duration_ms, + input_modality=modality, + output_modality="generation", + success=outcome == "completed", + error_code=stream_error_code(outcome), + input_meta=input_meta, + output_meta=GenerationOutputMeta( + task="generation", time_to_first_token_ms=ttft_ms + ), + generation_config=config, + ) + + return AppleFmStream(stream, handle, time.perf_counter(), None, on_done) + + return patched_stream + + +def patch_method(cls: type, name: str, make_patch, *extra) -> None: # type: ignore[no-untyped-def] + original = getattr(cls, name) + if getattr(original, "__wildedge_patch_name__", None) == APPLE_FM_PATCH_NAME: + return + patched = make_patch(original, *extra) + patched.__wildedge_patch_name__ = APPLE_FM_PATCH_NAME + patched.__wildedge_original_call__ = original + setattr(cls, name, patched) + + +# --------------------------------------------------------------------------- +# Extractor +# --------------------------------------------------------------------------- + + +def is_apple_fm_object(obj: object) -> bool: + return any( + cls.__module__.partition(".")[0] == APPLE_FM_PACKAGE + and cls.__name__ in ("SystemLanguageModel", "LanguageModelSession") + for cls in type(obj).__mro__ + ) + + +class AppleFmExtractor(BaseExtractor): + def can_handle(self, obj: object) -> bool: + return is_apple_fm_object(obj) + + def extract_info( + self, obj: object, overrides: dict + ) -> tuple[str | None, ModelInfo]: + model_id = overrides.pop("id", None) or model_id_for( + getattr(obj, APPLE_FM_USE_CASE_ATTR, None) + ) + info = ModelInfo( + model_name=model_id, + model_version=overrides.pop("version", None) or os_version(), + model_source=overrides.pop("source", MODEL_SOURCE), + model_format=overrides.pop("format", MODEL_FORMAT), + model_family=overrides.pop("family", MODEL_ID), + quantization=overrides.pop("quantization", None), + ) + return model_id, info + + def install_hooks(self, obj: object, handle: ModelHandle) -> None: + # Apple runs the system model on the GPU and Neural Engine; coreml is + # the closest accelerator name. + handle.detected_accelerator = accelerators.COREML + if type(obj).__name__ == "LanguageModelSession": + setattr(obj, APPLE_FM_HANDLE_ATTR, handle) + + @classmethod + def install_auto_load_patch(cls, client_ref: Any) -> None: + """Patch apple_fm_sdk so every new session is tracked. + + - ``SystemLanguageModel.__init__`` keeps the use case, which names the + model. + - ``LanguageModelSession.__init__`` registers the session's model. + - ``respond()`` and ``stream_response()`` record inference events. + """ + global _apple_fm_patched, _fm + if _apple_fm_patched: + return + _fm = _fm or optional_module(APPLE_FM_PACKAGE) + if _fm is None: + return + + with _APPLE_FM_PATCH_LOCK: + if _apple_fm_patched: + return + model_cls = _fm.SystemLanguageModel + session_cls = _fm.LanguageModelSession + patch_method(model_cls, "__init__", build_model_init_patch) + patch_method(session_cls, "__init__", build_session_init_patch, client_ref) + patch_method(session_cls, "respond", build_respond_patch) + patch_method(session_cls, "stream_response", build_stream_patch) + _apple_fm_patched = True + + +# Entry point named by the integration registry. +install_auto_load_patch = AppleFmExtractor.install_auto_load_patch diff --git a/wildedge/integrations/registry.py b/wildedge/integrations/registry.py index 94fc78c..453b2ec 100644 --- a/wildedge/integrations/registry.py +++ b/wildedge/integrations/registry.py @@ -44,6 +44,7 @@ def installer(self) -> str | None: INTEGRATION_SPECS: tuple[IntegrationSpec, ...] = ( IntegrationSpec("anthropic", ("anthropic",), "client_patch"), + IntegrationSpec("apple_fm", ("apple_fm_sdk",), "client_patch"), IntegrationSpec("llama_cpp", ("llama_cpp",), "client_patch", adapter_module="gguf"), IntegrationSpec("onnx", ("onnxruntime",), "client_patch"), IntegrationSpec("openai", ("openai",), "client_patch"), @@ -147,6 +148,7 @@ def load_extractor(target: str): "AnthropicExtractor", class_names=("Anthropic", "AsyncAnthropic", "Messages", "AsyncMessages"), ), + ExtractorSpec("apple_fm", "AppleFmExtractor", ("apple_fm_sdk",)), ExtractorSpec("ultralytics", "UltralyticsExtractor", class_names=("YOLO",)), ExtractorSpec("transformers", "TransformersExtractor", ("transformers",)), # Before MLX, which would otherwise claim MLX Whisper models. From a28c38660927648228206606d19545f8c8e913be Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Sun, 25 Oct 2026 00:41:27 +0200 Subject: [PATCH 02/11] Classify apple_fm errors with specific error codes Apple's on-device model never returns an HTTP status, so the server had nothing to derive a code from and stored every failure as UNKNOWN. The integration now sends a specific code for each SDK error and keeps the SDK's name in provider_error_code. New codes MODEL_UNAVAILABLE, CONTEXT_LENGTH_EXCEEDED and CONTENT_FILTERED cover failures the existing codes don't. RATE_LIMITED and BAD_REQUEST match the names the server already derives for remote APIs. The server keeps a specific code as sent. docs/manual-tracking.md now lists every code. --- CHANGELOG.md | 10 +++++++--- docs/apple_fm.md | 24 ++++++++++++++++++++---- docs/manual-tracking.md | 17 +++++++++++++++++ tests/test_integrations_apple_fm.py | 23 ++++++++++++++++++++++- wildedge/events/error.py | 5 +++++ wildedge/integrations/apple_fm.py | 21 +++++++++++++++++++-- 6 files changed, 90 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 59edcb3..f5b4a5e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,9 +9,13 @@ under Unreleased and move into a version section at release time. - `apple_fm` integration tracks Apple's on-device model through apple-fm-sdk. `LanguageModelSession.respond()` and `stream_response()` record duration, - time to first token for streams, prompt size and generation options, and - failures record the SDK's error name. Token counts are not recorded because - the SDK doesn't report them. See docs/apple_fm.md. + time to first token for streams, prompt size and generation options. + Failures record the SDK's error name and a matching error code. + Token counts are not recorded because the SDK doesn't report them. See + docs/apple_fm.md. +- Error codes `MODEL_UNAVAILABLE`, `CONTEXT_LENGTH_EXCEEDED`, + `CONTENT_FILTERED`, `RATE_LIMITED` and `BAD_REQUEST`, listed in + docs/manual-tracking.md. ## 0.3.0 - 2026-10-24 diff --git a/docs/apple_fm.md b/docs/apple_fm.md index 4d82930..6bb55ad 100644 --- a/docs/apple_fm.md +++ b/docs/apple_fm.md @@ -43,10 +43,26 @@ general model, and `apple-foundation-model-content-tagging` for `SystemLanguageModelUseCase.CONTENT_TAGGING`. The model version is the macOS version, since the model ships with the OS. -Failed calls record an error event named after the SDK's error, such as -`guardrail_violation`, `exceeded_context_window_size`, `assets_unavailable` or -`refusal`. A model that isn't ready because Apple Intelligence is off fails -this way too. +### Errors + +Failed calls record an error event. Its `provider_error_code` is the SDK's +error name, and its error code says what kind of failure it was: + +| SDK error | `provider_error_code` | Error code | +|---|---|---| +| `AssetsUnavailableError` | `assets_unavailable` | `MODEL_UNAVAILABLE` | +| `ExceededContextWindowSizeError` | `exceeded_context_window_size` | `CONTEXT_LENGTH_EXCEEDED` | +| `GuardrailViolationError` | `guardrail_violation` | `CONTENT_FILTERED` | +| `RefusalError` | `refusal` | `CONTENT_FILTERED` | +| `RateLimitedError` | `rate_limited` | `RATE_LIMITED` | +| `ConcurrentRequestsError` | `concurrent_requests` | `RATE_LIMITED` | +| `UnsupportedGuideError` | `unsupported_guide` | `BAD_REQUEST` | +| `UnsupportedLanguageOrLocaleError` | `unsupported_language_or_locale` | `BAD_REQUEST` | +| `InvalidGenerationSchemaError` | `invalid_schema` | `BAD_REQUEST` | +| `DecodingFailureError` | `decoding_failure` | `UNKNOWN` | +| `ToolCallError` | `tool_call` | `UNKNOWN` | + +Calling the model before it's ready fails with `AssetsUnavailableError`. ## Caveats diff --git a/docs/manual-tracking.md b/docs/manual-tracking.md index 0bcc769..5f3485f 100644 --- a/docs/manual-tracking.md +++ b/docs/manual-tracking.md @@ -233,6 +233,23 @@ except Exception as exc: `wildedge.track` (decorator and context manager) captures errors automatically when `capture_errors=True` (the default). +Use the most specific `ErrorCode` that fits, or `UNKNOWN`: + +| Code | Meaning | +|---|---| +| `OOM` | Ran out of memory | +| `CORRUPTED_MODEL` | Model file is damaged or unreadable | +| `INFERENCE_TIMEOUT` | The call timed out | +| `UNSUPPORTED_OP` | The runtime can't run an operation in the model | +| `THERMAL_SHUTDOWN` | The device stopped work because it was too hot | +| `CONNECTION_ERROR` | A remote model API couldn't be reached | +| `MODEL_UNAVAILABLE` | The model or its assets aren't on the device yet | +| `CONTEXT_LENGTH_EXCEEDED` | The prompt and history don't fit the model's context window | +| `CONTENT_FILTERED` | A safety filter blocked the input or output, or the model refused | +| `RATE_LIMITED` | Too many requests, or too many at once | +| `BAD_REQUEST` | The request was invalid, such as an unsupported language or schema | +| `UNKNOWN` | Anything else | + For failed calls to a remote model API, pass `http_status` and, if the provider returned one, `provider_error_code` with `error_code="UNKNOWN"`; Wild Edge derives the error code from them. The openai and anthropic integrations do this automatically, and report timeouts and connection failures as `INFERENCE_TIMEOUT` and `CONNECTION_ERROR`. ## Track feedback diff --git a/tests/test_integrations_apple_fm.py b/tests/test_integrations_apple_fm.py index f260dce..3634053 100644 --- a/tests/test_integrations_apple_fm.py +++ b/tests/test_integrations_apple_fm.py @@ -46,6 +46,10 @@ class ExceededContextWindowSizeError(GenerationError): pass +class DecodingFailureError(GenerationError): + pass + + def make_fake_sdk( *, reply="Hello there", @@ -270,7 +274,7 @@ async def test_respond_error_records_the_sdk_error_name(install): await session.respond("something unsafe") [error] = client.of_type("error") - assert error["error"]["error_code"] == "UNKNOWN" + assert error["error"]["error_code"] == "CONTENT_FILTERED" assert error["error"]["provider_error_code"] == "guardrail_violation" assert "duration_ms" in error["error"] assert client.of_type("inference") == [] @@ -308,10 +312,27 @@ async def test_stream_error_records_the_sdk_error_name(install): pass [error] = client.of_type("error") + assert error["error"]["error_code"] == "CONTEXT_LENGTH_EXCEEDED" assert error["error"]["provider_error_code"] == "exceeded_context_window_size" assert client.of_type("inference") == [] +async def test_error_without_a_specific_code_is_unknown(install): + sdk, client = install(error=DecodingFailureError("bad output")) + with pytest.raises(DecodingFailureError): + await sdk.LanguageModelSession().respond("hi") + + [error] = client.of_type("error") + assert error["error"]["error_code"] == "UNKNOWN" + assert error["error"]["provider_error_code"] == "decoding_failure" + + +def test_every_mapped_provider_code_is_a_known_sdk_error(): + assert set(apple_fm_mod.ERROR_CODES_BY_PROVIDER_CODE) <= set( + apple_fm_mod.PROVIDER_ERROR_CODES.values() + ) + + async def test_closed_client_leaves_sessions_untracked(install): sdk, client = install(client=FakeClient(closed=True)) session = sdk.LanguageModelSession() diff --git a/wildedge/events/error.py b/wildedge/events/error.py index a3d6bd2..79a6807 100644 --- a/wildedge/events/error.py +++ b/wildedge/events/error.py @@ -17,6 +17,11 @@ class ErrorCode(str, Enum): UNSUPPORTED_OP = "UNSUPPORTED_OP" THERMAL_SHUTDOWN = "THERMAL_SHUTDOWN" CONNECTION_ERROR = "CONNECTION_ERROR" + MODEL_UNAVAILABLE = "MODEL_UNAVAILABLE" + CONTEXT_LENGTH_EXCEEDED = "CONTEXT_LENGTH_EXCEEDED" + CONTENT_FILTERED = "CONTENT_FILTERED" + RATE_LIMITED = "RATE_LIMITED" + BAD_REQUEST = "BAD_REQUEST" UNKNOWN = "UNKNOWN" diff --git a/wildedge/integrations/apple_fm.py b/wildedge/integrations/apple_fm.py index 044c8ed..01a553d 100644 --- a/wildedge/integrations/apple_fm.py +++ b/wildedge/integrations/apple_fm.py @@ -68,6 +68,20 @@ "ToolCallError": "tool_call", } +# The error code sent for each SDK error. The server keeps a specific code +# as sent, and none of these failures has an HTTP status to derive one from. +ERROR_CODES_BY_PROVIDER_CODE = { + "assets_unavailable": ErrorCode.MODEL_UNAVAILABLE, + "exceeded_context_window_size": ErrorCode.CONTEXT_LENGTH_EXCEEDED, + "guardrail_violation": ErrorCode.CONTENT_FILTERED, + "refusal": ErrorCode.CONTENT_FILTERED, + "rate_limited": ErrorCode.RATE_LIMITED, + "concurrent_requests": ErrorCode.RATE_LIMITED, + "unsupported_guide": ErrorCode.BAD_REQUEST, + "unsupported_language_or_locale": ErrorCode.BAD_REQUEST, + "invalid_schema": ErrorCode.BAD_REQUEST, +} + debug_apple_fm_failure = functools.partial(debug_failure, "apple_fm") @@ -95,10 +109,13 @@ def provider_error_code_of(exc: BaseException) -> str | None: def track_apple_fm_error( handle: ModelHandle, exc: BaseException, duration_ms: int ) -> None: + provider_error_code = provider_error_code_of(exc) handle.track_error( - error_code=ErrorCode.UNKNOWN, + error_code=ERROR_CODES_BY_PROVIDER_CODE.get( + provider_error_code, ErrorCode.UNKNOWN + ), error_message=str(exc), - provider_error_code=provider_error_code_of(exc), + provider_error_code=provider_error_code, duration_ms=duration_ms, ) From 6d17df45a73965bc0686177e0b73b57a9468fe0b Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Sun, 25 Oct 2026 00:41:36 +0200 Subject: [PATCH 03/11] Record apple_fm model availability as a load event An app that checks SystemLanguageModel.is_available() and falls back when the model can't be used left no trace. is_available() now records a model load. When the model is unavailable the load fails with the SDK's reason as its error code: APPLE_INTELLIGENCE_NOT_ENABLED, DEVICE_NOT_ELIGIBLE or MODEL_NOT_READY. Apps poll is_available() while the model downloads, so an event is sent only when the answer changes, per process and model. --- CHANGELOG.md | 4 +- docs/apple_fm.md | 9 ++++ tests/compat/test_apple_fm_compat.py | 11 +++-- tests/test_integrations_apple_fm.py | 64 ++++++++++++++++++++++++++++ wildedge/integrations/apple_fm.py | 53 +++++++++++++++++++++++ 5 files changed, 137 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f5b4a5e..bd1cb12 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,9 @@ under Unreleased and move into a version section at release time. `LanguageModelSession.respond()` and `stream_response()` record duration, time to first token for streams, prompt size and generation options. Failures record the SDK's error name and a matching error code. - Token counts are not recorded because the SDK doesn't report them. See + `SystemLanguageModel.is_available()` records a model load, failed with + the reason when the model can't be used, such as `MODEL_NOT_READY`. Token + counts are not recorded because the SDK doesn't report them. See docs/apple_fm.md. - Error codes `MODEL_UNAVAILABLE`, `CONTEXT_LENGTH_EXCEEDED`, `CONTENT_FILTERED`, `RATE_LIMITED` and `BAD_REQUEST`, listed in diff --git a/docs/apple_fm.md b/docs/apple_fm.md index 6bb55ad..df425ff 100644 --- a/docs/apple_fm.md +++ b/docs/apple_fm.md @@ -43,6 +43,15 @@ general model, and `apple-foundation-model-content-tagging` for `SystemLanguageModelUseCase.CONTENT_TAGGING`. The model version is the macOS version, since the model ships with the OS. +### Availability + +`SystemLanguageModel.is_available()` records a model load event. When the +model can't be used, the load fails with the SDK's reason as its error code: +`APPLE_INTELLIGENCE_NOT_ENABLED`, `DEVICE_NOT_ELIGIBLE`, `MODEL_NOT_READY` +(still downloading) or `UNKNOWN`. Apps often poll `is_available()` while the +model downloads, so an event is sent only when the answer changes, once per +process and model. + ### Errors Failed calls record an error event. Its `provider_error_code` is the SDK's diff --git a/tests/compat/test_apple_fm_compat.py b/tests/compat/test_apple_fm_compat.py index 8d5489e..659dfad 100644 --- a/tests/compat/test_apple_fm_compat.py +++ b/tests/compat/test_apple_fm_compat.py @@ -1,8 +1,8 @@ """Real apple-fm-sdk sessions. CI runners can't turn on Apple Intelligence, so a call there fails with the -SDK's own error. Either way the call must produce exactly one event for the -model it ran on. +SDK's own error. Either way the availability check must produce one load +event, and the call exactly one event, for the model it ran on. """ import asyncio @@ -42,7 +42,11 @@ async def respond(): result = asyncio.run(respond()) - [event] = compat_client.events + [load, event] = compat_client.events + assert load["event_type"] == "model_load" + assert load["load"]["success"] is available + if not available: + assert load["load"]["error_code"] == reason.name assert event["model_id"] == "apple-foundation-model-content-tagging" if available: assert isinstance(result, str), result @@ -51,3 +55,4 @@ async def respond(): else: assert isinstance(result, fm.FoundationModelsError), (reason, result) assert event["event_type"] == "error" + assert event["error"]["provider_error_code"] is not None diff --git a/tests/test_integrations_apple_fm.py b/tests/test_integrations_apple_fm.py index 3634053..d38820f 100644 --- a/tests/test_integrations_apple_fm.py +++ b/tests/test_integrations_apple_fm.py @@ -25,6 +25,13 @@ # --------------------------------------------------------------------------- +class SystemLanguageModelUnavailableReason(enum.IntEnum): + APPLE_INTELLIGENCE_NOT_ENABLED = 0 + DEVICE_NOT_ELIGIBLE = 1 + MODEL_NOT_READY = 2 + UNKNOWN = 255 + + class SystemLanguageModelUseCase(enum.IntEnum): GENERAL = 0 CONTENT_TAGGING = 1 @@ -56,6 +63,7 @@ def make_fake_sdk( snapshots=("Hel", "Hello", "Hello there"), error=None, stream_error=None, + availability=(True, None), ) -> ModuleType: """A fresh module per test, since the integration patches its classes.""" @@ -68,6 +76,9 @@ def __init__( ): self._ptr = object() + def is_available(self): + return availability() if callable(availability) else availability + class LanguageModelSession: def __init__(self, instructions=None, model=None, tools=None, _ptr=None): self.instructions = instructions @@ -137,6 +148,7 @@ def _install(client=None, **sdk_kwargs): client = client or FakeClient() monkeypatch.setattr(apple_fm_mod, "_fm", sdk) monkeypatch.setattr(apple_fm_mod, "_apple_fm_patched", False) + monkeypatch.setattr(apple_fm_mod, "_reported_availability", {}) AppleFmExtractor.install_auto_load_patch(lambda: client) return sdk, client @@ -333,6 +345,58 @@ def test_every_mapped_provider_code_is_a_known_sdk_error(): ) +# --------------------------------------------------------------------------- +# Availability +# --------------------------------------------------------------------------- + + +def test_available_model_records_one_successful_load(install): + sdk, client = install() + model = sdk.SystemLanguageModel() + + assert model.is_available() == (True, None) + model.is_available() + sdk.SystemLanguageModel().is_available() + + [event] = client.of_type("model_load") + assert event["model_id"] == MODEL_ID + assert event["load"]["success"] is True + assert event["load"]["error_code"] is None + + +def test_unavailable_model_records_the_reason(install): + reason = SystemLanguageModelUnavailableReason.MODEL_NOT_READY + sdk, client = install(availability=(False, reason)) + + assert sdk.SystemLanguageModel().is_available() == (False, reason) + + [event] = client.of_type("model_load") + assert event["load"]["success"] is False + assert event["load"]["error_code"] == "MODEL_NOT_READY" + + +def test_availability_is_recorded_again_when_it_changes(install): + not_ready = (False, SystemLanguageModelUnavailableReason.MODEL_NOT_READY) + answers = iter([not_ready, not_ready, (True, None), (True, None)]) + sdk, client = install(availability=lambda: next(answers)) + model = sdk.SystemLanguageModel() + + for _ in range(4): + model.is_available() + + loads = [e["load"] for e in client.of_type("model_load")] + assert [(load["success"], load["error_code"]) for load in loads] == [ + (False, "MODEL_NOT_READY"), + (True, None), + ] + + +def test_closed_client_leaves_availability_untracked(install): + sdk, client = install(client=FakeClient(closed=True)) + assert sdk.SystemLanguageModel().is_available() == (True, None) + assert client.events == [] + + async def test_closed_client_leaves_sessions_untracked(install): sdk, client = install(client=FakeClient(closed=True)) session = sdk.LanguageModelSession() diff --git a/wildedge/integrations/apple_fm.py b/wildedge/integrations/apple_fm.py index 01a553d..7b7dcf7 100644 --- a/wildedge/integrations/apple_fm.py +++ b/wildedge/integrations/apple_fm.py @@ -84,6 +84,11 @@ debug_apple_fm_failure = functools.partial(debug_failure, "apple_fm") +# The last is_available() answer reported per model id. A load event is sent +# only when the answer changes, since apps poll while the model downloads. +_reported_availability: dict[str, str | None] = {} +_AVAILABILITY_LOCK = threading.Lock() + def model_id_for(use_case: object) -> str: """``apple-foundation-model``, plus the use case when it is not GENERAL.""" @@ -120,6 +125,28 @@ def track_apple_fm_error( ) +def unavailable_error_code(reason: object) -> str: + """``MODEL_NOT_READY`` and the like, from the SDK's unavailable reason.""" + name = getattr(reason, "name", None) + return name if isinstance(name, str) and name else ErrorCode.UNKNOWN.value + + +def track_availability( + handle: ModelHandle, available: bool, reason: object, duration_ms: int +) -> None: + """Send a load event for an is_available() answer that changed.""" + error_code = None if available else unavailable_error_code(reason) + with _AVAILABILITY_LOCK: + model_id = handle.model_id + if ( + model_id in _reported_availability + and _reported_availability[model_id] == error_code + ): + return + _reported_availability[model_id] = error_code + handle.track_load(duration_ms=duration_ms, success=available, error_code=error_code) + + def build_input_meta(prompt: object) -> tuple[str, TextInputMeta | None]: """Input modality and text metadata for a prompt. @@ -197,6 +224,27 @@ def patched_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] return patched_init +def build_availability_patch(original_is_available, client_ref): # type: ignore[no-untyped-def] + @functools.wraps(original_is_available) + def patched_is_available(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] + t0 = time.perf_counter() + result = original_is_available(self_inner, *args, **kwargs) + duration_ms = elapsed_ms(t0) + c = client_ref() # type: ignore[call-arg] + if c is None or c.closed: + return result + try: + available, reason = result + model_id = model_id_for(getattr(self_inner, APPLE_FM_USE_CASE_ATTR, None)) + handle = c.register_model(self_inner, model_id=model_id) + track_availability(handle, bool(available), reason, duration_ms) + except Exception as exc: + debug_apple_fm_failure("availability tracking", exc) + return result + + return patched_is_available + + def build_session_init_patch(original_init, client_ref): # type: ignore[no-untyped-def] @functools.wraps(original_init) def patched_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] @@ -334,6 +382,8 @@ def install_auto_load_patch(cls, client_ref: Any) -> None: - ``SystemLanguageModel.__init__`` keeps the use case, which names the model. + - ``SystemLanguageModel.is_available()`` records a load event when + its answer changes, failed with the reason the model is unavailable. - ``LanguageModelSession.__init__`` registers the session's model. - ``respond()`` and ``stream_response()`` record inference events. """ @@ -350,6 +400,9 @@ def install_auto_load_patch(cls, client_ref: Any) -> None: model_cls = _fm.SystemLanguageModel session_cls = _fm.LanguageModelSession patch_method(model_cls, "__init__", build_model_init_patch) + patch_method( + model_cls, "is_available", build_availability_patch, client_ref + ) patch_method(session_cls, "__init__", build_session_init_patch, client_ref) patch_method(session_cls, "respond", build_respond_patch) patch_method(session_cls, "stream_response", build_stream_patch) From 2d1152415dcb392a3c6a30125159e61c750d203d Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Sun, 25 Oct 2026 00:41:36 +0200 Subject: [PATCH 04/11] Print the unavailable reason's name in the apple_fm example The example printed the bare enum value, such as "unavailable: 2". --- examples/apple_fm_example.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/apple_fm_example.py b/examples/apple_fm_example.py index 09261d7..e15b09e 100644 --- a/examples/apple_fm_example.py +++ b/examples/apple_fm_example.py @@ -32,7 +32,7 @@ async def main() -> None: model = fm.SystemLanguageModel() available, reason = model.is_available() if not available: - raise SystemExit(f"Apple's on-device model is unavailable: {reason}") + raise SystemExit(f"Apple's on-device model is unavailable: {reason.name}") session = fm.LanguageModelSession( instructions="Answer in one sentence.", model=model From d478a1dbe0993c2cd44d84d7a46ab6438b5613e2 Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Sun, 25 Oct 2026 00:45:12 +0200 Subject: [PATCH 05/11] Give apple_fm decoding and tool failures their own error codes Both were recorded as UNKNOWN. DecodingFailureError means the model's output couldn't be parsed into the requested structure, so it is now INVALID_OUTPUT. ToolCallError wraps an exception from the app's own tool, so it is now TOOL_CALL_FAILED and tells app failures apart from model failures. Both codes fit other integrations too. --- CHANGELOG.md | 4 ++-- docs/apple_fm.md | 6 ++++-- docs/manual-tracking.md | 2 ++ tests/test_integrations_apple_fm.py | 30 +++++++++++++++++++++++++++-- wildedge/events/error.py | 2 ++ wildedge/integrations/apple_fm.py | 2 ++ 6 files changed, 40 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bd1cb12..3b1dbd4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,8 +16,8 @@ under Unreleased and move into a version section at release time. counts are not recorded because the SDK doesn't report them. See docs/apple_fm.md. - Error codes `MODEL_UNAVAILABLE`, `CONTEXT_LENGTH_EXCEEDED`, - `CONTENT_FILTERED`, `RATE_LIMITED` and `BAD_REQUEST`, listed in - docs/manual-tracking.md. + `CONTENT_FILTERED`, `RATE_LIMITED`, `BAD_REQUEST`, `INVALID_OUTPUT` and + `TOOL_CALL_FAILED`, listed in docs/manual-tracking.md. ## 0.3.0 - 2026-10-24 diff --git a/docs/apple_fm.md b/docs/apple_fm.md index df425ff..07501e4 100644 --- a/docs/apple_fm.md +++ b/docs/apple_fm.md @@ -68,8 +68,10 @@ error name, and its error code says what kind of failure it was: | `UnsupportedGuideError` | `unsupported_guide` | `BAD_REQUEST` | | `UnsupportedLanguageOrLocaleError` | `unsupported_language_or_locale` | `BAD_REQUEST` | | `InvalidGenerationSchemaError` | `invalid_schema` | `BAD_REQUEST` | -| `DecodingFailureError` | `decoding_failure` | `UNKNOWN` | -| `ToolCallError` | `tool_call` | `UNKNOWN` | +| `DecodingFailureError` | `decoding_failure` | `INVALID_OUTPUT` | +| `ToolCallError` | `tool_call` | `TOOL_CALL_FAILED` | + +Any other failure records `UNKNOWN`. Calling the model before it's ready fails with `AssetsUnavailableError`. diff --git a/docs/manual-tracking.md b/docs/manual-tracking.md index 5f3485f..d16ea65 100644 --- a/docs/manual-tracking.md +++ b/docs/manual-tracking.md @@ -248,6 +248,8 @@ Use the most specific `ErrorCode` that fits, or `UNKNOWN`: | `CONTENT_FILTERED` | A safety filter blocked the input or output, or the model refused | | `RATE_LIMITED` | Too many requests, or too many at once | | `BAD_REQUEST` | The request was invalid, such as an unsupported language or schema | +| `INVALID_OUTPUT` | The model's output couldn't be parsed into the requested structure | +| `TOOL_CALL_FAILED` | A tool the model called raised an error | | `UNKNOWN` | Anything else | For failed calls to a remote model API, pass `http_status` and, if the provider returned one, `provider_error_code` with `error_code="UNKNOWN"`; Wild Edge derives the error code from them. The openai and anthropic integrations do this automatically, and report timeouts and connection failures as `INFERENCE_TIMEOUT` and `CONNECTION_ERROR`. diff --git a/tests/test_integrations_apple_fm.py b/tests/test_integrations_apple_fm.py index d38820f..047af8a 100644 --- a/tests/test_integrations_apple_fm.py +++ b/tests/test_integrations_apple_fm.py @@ -57,6 +57,11 @@ class DecodingFailureError(GenerationError): pass +class ToolCallError(FoundationModelsError): + def __init__(self, tool_name, underlying_error): + super().__init__(f"Tool '{tool_name}' failed: {underlying_error}") + + def make_fake_sdk( *, reply="Hello there", @@ -329,16 +334,37 @@ async def test_stream_error_records_the_sdk_error_name(install): assert client.of_type("inference") == [] -async def test_error_without_a_specific_code_is_unknown(install): +async def test_decoding_failure_is_invalid_output(install): sdk, client = install(error=DecodingFailureError("bad output")) with pytest.raises(DecodingFailureError): await sdk.LanguageModelSession().respond("hi") [error] = client.of_type("error") - assert error["error"]["error_code"] == "UNKNOWN" + assert error["error"]["error_code"] == "INVALID_OUTPUT" assert error["error"]["provider_error_code"] == "decoding_failure" +async def test_tool_failure_is_tool_call_failed(install): + sdk, client = install(error=ToolCallError("weather", ValueError("no city"))) + with pytest.raises(ToolCallError): + await sdk.LanguageModelSession().respond("hi") + + [error] = client.of_type("error") + assert error["error"]["error_code"] == "TOOL_CALL_FAILED" + assert error["error"]["provider_error_code"] == "tool_call" + assert "weather" in error["error"]["error_message"] + + +async def test_error_without_a_specific_code_is_unknown(install): + sdk, client = install(error=FoundationModelsError("something else")) + with pytest.raises(FoundationModelsError): + await sdk.LanguageModelSession().respond("hi") + + [error] = client.of_type("error") + assert error["error"]["error_code"] == "UNKNOWN" + assert "provider_error_code" not in error["error"] + + def test_every_mapped_provider_code_is_a_known_sdk_error(): assert set(apple_fm_mod.ERROR_CODES_BY_PROVIDER_CODE) <= set( apple_fm_mod.PROVIDER_ERROR_CODES.values() diff --git a/wildedge/events/error.py b/wildedge/events/error.py index 79a6807..2bffe9d 100644 --- a/wildedge/events/error.py +++ b/wildedge/events/error.py @@ -22,6 +22,8 @@ class ErrorCode(str, Enum): CONTENT_FILTERED = "CONTENT_FILTERED" RATE_LIMITED = "RATE_LIMITED" BAD_REQUEST = "BAD_REQUEST" + INVALID_OUTPUT = "INVALID_OUTPUT" + TOOL_CALL_FAILED = "TOOL_CALL_FAILED" UNKNOWN = "UNKNOWN" diff --git a/wildedge/integrations/apple_fm.py b/wildedge/integrations/apple_fm.py index 7b7dcf7..60cc389 100644 --- a/wildedge/integrations/apple_fm.py +++ b/wildedge/integrations/apple_fm.py @@ -80,6 +80,8 @@ "unsupported_guide": ErrorCode.BAD_REQUEST, "unsupported_language_or_locale": ErrorCode.BAD_REQUEST, "invalid_schema": ErrorCode.BAD_REQUEST, + "decoding_failure": ErrorCode.INVALID_OUTPUT, + "tool_call": ErrorCode.TOOL_CALL_FAILED, } debug_apple_fm_failure = functools.partial(debug_failure, "apple_fm") From f086b100ed75736820a900885f24c7640012eb52 Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Sun, 25 Oct 2026 01:18:54 +0200 Subject: [PATCH 06/11] Give the apple_fm compat test's reply room to finish The test had only run where the model was unavailable. With the model ready, the content tagging model can't finish a reply in 8 tokens, and the SDK raises DecodingFailureError. 64 tokens is enough. --- tests/compat/test_apple_fm_compat.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/compat/test_apple_fm_compat.py b/tests/compat/test_apple_fm_compat.py index 659dfad..cb10a42 100644 --- a/tests/compat/test_apple_fm_compat.py +++ b/tests/compat/test_apple_fm_compat.py @@ -32,7 +32,7 @@ def test_respond_records_one_event(compat_client): ) available, reason = model.is_available() session = fm.LanguageModelSession(model=model) - options = fm.GenerationOptions(maximum_response_tokens=8) + options = fm.GenerationOptions(maximum_response_tokens=64) async def respond(): try: @@ -51,7 +51,7 @@ async def respond(): if available: assert isinstance(result, str), result assert event["event_type"] == "inference" - assert event["inference"]["generation_config"]["max_tokens"] == 8 + assert event["inference"]["generation_config"]["max_tokens"] == 64 else: assert isinstance(result, fm.FoundationModelsError), (reason, result) assert event["event_type"] == "error" From d294433f81774572ea25ca95b9f1d1d72b98c1f0 Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Sun, 25 Oct 2026 01:18:54 +0200 Subject: [PATCH 07/11] Count apple_fm tokens when WILDEDGE_APPLE_FM_COUNT_TOKENS is set apple-fm-sdk reports no token usage. SystemLanguageModel.token_count() counts tokens without generating, at about 40 ms per call on an M2 Mac. With the variable set, respond() and stream_response() count the prompt and the reply in a background task and send the inference event with tokens in, tokens out and tokens per second. The call returns as soon as the model has answered. macOS queues counts behind generation, so a request that follows right after waits about 80 ms, which is why counting is opt-in. A done callback sends the event whether counting finished, failed or was cancelled before it ran, as when asyncio.run() closes the loop. The setting is an env var, not an init() parameter, so wildedge run can use it too. docs/configuration.md lists it under integration settings. --- CHANGELOG.md | 5 +- docs/apple_fm.md | 22 ++++- docs/configuration.md | 9 ++ tests/test_integrations_apple_fm.py | 108 +++++++++++++++++++++- wildedge/constants.py | 2 + wildedge/integrations/apple_fm.py | 134 ++++++++++++++++++++++++++-- 6 files changed, 266 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b1dbd4..a584983 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,8 +12,9 @@ under Unreleased and move into a version section at release time. time to first token for streams, prompt size and generation options. Failures record the SDK's error name and a matching error code. `SystemLanguageModel.is_available()` records a model load, failed with - the reason when the model can't be used, such as `MODEL_NOT_READY`. Token - counts are not recorded because the SDK doesn't report them. See + the reason when the model can't be used, such as `MODEL_NOT_READY`. The SDK + doesn't report token usage, and `WILDEDGE_APPLE_FM_COUNT_TOKENS=1` counts + it with `SystemLanguageModel.token_count()` in the background. See docs/apple_fm.md. - Error codes `MODEL_UNAVAILABLE`, `CONTEXT_LENGTH_EXCEEDED`, `CONTENT_FILTERED`, `RATE_LIMITED`, `BAD_REQUEST`, `INVALID_OUTPUT` and diff --git a/docs/apple_fm.md b/docs/apple_fm.md index 07501e4..5a34098 100644 --- a/docs/apple_fm.md +++ b/docs/apple_fm.md @@ -43,6 +43,24 @@ general model, and `apple-foundation-model-content-tagging` for `SystemLanguageModelUseCase.CONTENT_TAGGING`. The model version is the macOS version, since the model ships with the OS. +### Token counts + +The SDK doesn't report token usage, so token counts are empty by default. Set +`WILDEDGE_APPLE_FM_COUNT_TOKENS=1` to count them with +`SystemLanguageModel.token_count()`. Each call then records tokens in (the +prompt, not the instructions or earlier turns), tokens out (the reply) and +tokens per second. + +Counting costs two calls to the model service per request, about 40 ms each +on an M2 Mac. They run in the background after the response returns, so your +call is never slower. macOS queues them behind generation, though, so they +can delay a request that follows right after by about 80 ms. The event is +sent once counting finishes, usually soon after the model is idle. + +Structured replies (`generating=` or a schema) count tokens in only. A count +that fails sends the event without token counts. So does the last call before +the event loop closes, since `asyncio.run()` cancels the pending count. + ### Availability `SystemLanguageModel.is_available()` records a model load event. When the @@ -77,9 +95,7 @@ Calling the model before it's ready fails with `AssetsUnavailableError`. ## Caveats -- Token counts and tokens per second are empty. The SDK doesn't report token - usage, and counting tokens with `SystemLanguageModel.token_count()` is an - extra call to the model, so WildEdge doesn't make it for you. +- Token counts are opt-in, see [Token counts](#token-counts). - Prompts and responses are not uploaded. ## `fm serve` diff --git a/docs/configuration.md b/docs/configuration.md index 02f6f46..e17db85 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -47,6 +47,15 @@ inputs or outputs. See [Track attachments](manual-tracking.md#track-attachments) for usage. +## Integration settings + +Settings for a single integration are environment variables only. They are +read when the integration is activated. + +| Env var | Default | Description | +|---|---|---| +| `WILDEDGE_APPLE_FM_COUNT_TOKENS` | `false` | Count tokens for Apple's on-device model, at about 80 ms of background work per call. See [Token counts](apple_fm.md#token-counts) | + ### When ML libraries are imported `import wildedge` does not import any ML framework or API client. The SDK imports a library only when you use it: diff --git a/tests/test_integrations_apple_fm.py b/tests/test_integrations_apple_fm.py index 047af8a..8c908e3 100644 --- a/tests/test_integrations_apple_fm.py +++ b/tests/test_integrations_apple_fm.py @@ -10,6 +10,7 @@ import wildedge.integrations.apple_fm as apple_fm_mod from wildedge import accelerators +from wildedge.events.inference import GenerationOutputMeta from wildedge.integrations.apple_fm import ( MODEL_ID, AppleFmExtractor, @@ -69,6 +70,7 @@ def make_fake_sdk( error=None, stream_error=None, availability=(True, None), + token_count_error=None, ) -> ModuleType: """A fresh module per test, since the integration patches its classes.""" @@ -84,6 +86,12 @@ def __init__( def is_available(self): return availability() if callable(availability) else availability + async def token_count(self, value=None, *, instructions=None): + await asyncio.sleep(0) + if token_count_error is not None: + raise token_count_error + return len(str(value).split()) + class LanguageModelSession: def __init__(self, instructions=None, model=None, tools=None, _ptr=None): self.instructions = instructions @@ -148,9 +156,14 @@ def of_type(self, event_type): def install(monkeypatch): """Activate the integration on a fake SDK and return (sdk, client).""" - def _install(client=None, **sdk_kwargs): + def _install(client=None, count_tokens=False, **sdk_kwargs): sdk = make_fake_sdk(**sdk_kwargs) client = client or FakeClient() + if count_tokens: + monkeypatch.setenv("WILDEDGE_APPLE_FM_COUNT_TOKENS", "1") + else: + monkeypatch.delenv("WILDEDGE_APPLE_FM_COUNT_TOKENS", raising=False) + monkeypatch.setattr(apple_fm_mod, "_default_model", None) monkeypatch.setattr(apple_fm_mod, "_fm", sdk) monkeypatch.setattr(apple_fm_mod, "_apple_fm_patched", False) monkeypatch.setattr(apple_fm_mod, "_reported_availability", {}) @@ -164,6 +177,12 @@ def inference_of(event): return event["inference"] +async def counted(): + """Wait for pending token counts to send their events.""" + await asyncio.gather(*apple_fm_mod._counting_tasks, return_exceptions=True) + await asyncio.sleep(0) # done callbacks run on the next loop iteration + + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -423,6 +442,93 @@ def test_closed_client_leaves_availability_untracked(install): assert client.events == [] +# --------------------------------------------------------------------------- +# Token counting +# --------------------------------------------------------------------------- + + +async def test_tokens_are_not_counted_by_default(install): + sdk, client = install() + await sdk.LanguageModelSession().respond("How does rain form?") + + [event] = client.of_type("inference") + assert "tokens_in" not in inference_of(event)["output_meta"] + assert apple_fm_mod._counting_tasks == set() + + +async def test_respond_counts_prompt_and_reply_tokens(install): + sdk, client = install(count_tokens=True, reply="Water vapour condenses") + session = sdk.LanguageModelSession(model=sdk.SystemLanguageModel()) + + assert await session.respond("How does rain form?") == "Water vapour condenses" + await counted() + + [event] = client.of_type("inference") + inference = inference_of(event) + assert inference["output_meta"]["tokens_in"] == 4 + assert inference["output_meta"]["tokens_out"] == 3 + assert inference["input_meta"]["token_count"] == 4 + + +def test_tokens_per_second_needs_a_duration(): + meta = GenerationOutputMeta() + apple_fm_mod.add_token_counts( + {"output_meta": meta, "duration_ms": 500}, tokens_in=10, tokens_out=20 + ) + assert (meta.tokens_in, meta.tokens_out, meta.tokens_per_second) == (10, 20, 40.0) + + meta = GenerationOutputMeta() + apple_fm_mod.add_token_counts( + {"output_meta": meta, "duration_ms": 0}, tokens_in=10, tokens_out=20 + ) + assert meta.tokens_per_second is None + + +async def test_stream_counts_the_last_snapshot(install): + sdk, client = install(count_tokens=True, snapshots=("Rain", "Rain falls down")) + async for _ in sdk.LanguageModelSession().stream_response("Tell me"): + pass + await counted() + + [event] = client.of_type("inference") + output_meta = inference_of(event)["output_meta"] + assert (output_meta["tokens_in"], output_meta["tokens_out"]) == (2, 3) + assert "time_to_first_token_ms" in output_meta + + +async def test_session_without_a_model_counts_with_the_default_model(install): + sdk, client = install(count_tokens=True) + await sdk.LanguageModelSession().respond("hi there") + await counted() + + assert apple_fm_mod._default_model is not None + [event] = client.of_type("inference") + assert inference_of(event)["output_meta"]["tokens_in"] == 2 + + +async def test_failed_count_still_records_the_inference(install): + sdk, client = install( + count_tokens=True, token_count_error=FoundationModelsError("no assets") + ) + await sdk.LanguageModelSession().respond("hi") + await counted() + + [event] = client.of_type("inference") + assert inference_of(event)["success"] is True + assert "tokens_in" not in inference_of(event)["output_meta"] + + +async def test_cancelled_count_still_records_the_inference_once(install): + sdk, client = install(count_tokens=True) + await sdk.LanguageModelSession().respond("hi") + for task in apple_fm_mod._counting_tasks: + task.cancel() + await counted() + + [event] = client.of_type("inference") + assert "tokens_in" not in inference_of(event)["output_meta"] + + async def test_closed_client_leaves_sessions_untracked(install): sdk, client = install(client=FakeClient(closed=True)) session = sdk.LanguageModelSession() diff --git a/wildedge/constants.py b/wildedge/constants.py index acd0fcc..b8f35e7 100644 --- a/wildedge/constants.py +++ b/wildedge/constants.py @@ -21,6 +21,8 @@ ENV_PRINT_STARTUP_REPORT = "WILDEDGE_PRINT_STARTUP_REPORT" ENV_SAMPLING_INTERVAL = "WILDEDGE_SAMPLING_INTERVAL_S" ENV_ATTACHMENTS_ENABLED = "WILDEDGE_ATTACHMENTS_ENABLED" +# Integration settings +ENV_APPLE_FM_COUNT_TOKENS = "WILDEDGE_APPLE_FM_COUNT_TOKENS" # Defaults DEFAULT_SAMPLING_INTERVAL_S = 30.0 diff --git a/wildedge/integrations/apple_fm.py b/wildedge/integrations/apple_fm.py index 60cc389..579cb1d 100644 --- a/wildedge/integrations/apple_fm.py +++ b/wildedge/integrations/apple_fm.py @@ -4,20 +4,23 @@ SDK. Sessions are tied to a model handle when they are created, and ``respond()`` and ``stream_response()`` record one inference event per call. -The SDK reports no token usage, and counting tokens takes an extra call into -the model, so token counts are left empty. +The SDK reports no token usage. Counting tokens takes two extra calls into +the model service per request, so it is opt-in: set +``WILDEDGE_APPLE_FM_COUNT_TOKENS=1``. """ from __future__ import annotations +import asyncio import functools +import os import platform import sys import threading import time from typing import TYPE_CHECKING, Any -from wildedge import accelerators +from wildedge import accelerators, constants from wildedge.events.error import ErrorCode from wildedge.events.inference import ( GenerationConfig, @@ -32,6 +35,7 @@ ) from wildedge.loading import optional_module from wildedge.model import ModelInfo +from wildedge.settings import parse_bool from wildedge.timing import elapsed_ms # Imported when the integration is activated: importing the SDK loads its @@ -47,6 +51,15 @@ APPLE_FM_HANDLE_ATTR = "__wildedge_apple_fm_handle__" # SystemLanguageModel passes its use case to native code without keeping it. APPLE_FM_USE_CASE_ATTR = "__wildedge_apple_fm_use_case__" +# The model a session was created with, which counts its tokens. +APPLE_FM_MODEL_ATTR = "__wildedge_apple_fm_model__" + +# Read from WILDEDGE_APPLE_FM_COUNT_TOKENS when the integration is activated. +_count_tokens = False +# A default model counts tokens for sessions created without one. +_default_model: Any = None +# Strong references, so pending token counts aren't garbage collected. +_counting_tasks: set[asyncio.Task] = set() MODEL_ID = "apple-foundation-model" MODEL_SOURCE = "apple" @@ -207,6 +220,94 @@ def _error(self, exc: Exception) -> None: self._safe(track_apple_fm_error, self._handle, exc, elapsed_ms(self._t0)) +def prompt_of(args: tuple, kwargs: dict) -> object: + return kwargs.get("prompt", args[0] if args else None) + + +def counting_model(session: object) -> Any: + """The model that counts a session's tokens.""" + global _default_model + model = getattr(session, APPLE_FM_MODEL_ATTR, None) + if model is not None: + return model + if _default_model is None and _fm is not None: + _default_model = _fm.SystemLanguageModel() + return _default_model + + +async def count_tokens( + model: Any, prompt: object, reply: object +) -> tuple[int | None, int | None]: + tokens_in = await model.token_count(prompt) if prompt is not None else None + tokens_out = await model.token_count(reply) if isinstance(reply, str) else None + return tokens_in, tokens_out + + +def add_token_counts( + kwargs: dict, tokens_in: int | None, tokens_out: int | None +) -> None: + output_meta: GenerationOutputMeta = kwargs["output_meta"] + output_meta.tokens_in = tokens_in + output_meta.tokens_out = tokens_out + duration_ms = kwargs["duration_ms"] + if tokens_out and duration_ms > 0: + output_meta.tokens_per_second = round(tokens_out / duration_ms * 1000, 1) + input_meta: TextInputMeta | None = kwargs.get("input_meta") + if input_meta is not None: + input_meta.token_count = tokens_in + + +async def count_into(kwargs: dict, model: Any, prompt: object, reply: object) -> None: + try: + add_token_counts(kwargs, *await count_tokens(model, prompt, reply)) + except Exception as exc: + debug_apple_fm_failure("token count", exc) + + +def send_counted_inference( + handle: ModelHandle, kwargs: dict, task: asyncio.Task +) -> None: + """Send the inference once counting ends, however it ends. + + A failed count, or one cancelled before it ran (as when asyncio.run() + shuts down), sends the event without token counts. + """ + _counting_tasks.discard(task) + try: + handle.track_inference(**kwargs) + except Exception as exc: + debug_apple_fm_failure("inference tracking", exc) + + +def track_generation( + handle: ModelHandle, + session: object, + prompt: object, + reply: object, + **kwargs: Any, +) -> None: + """Record an inference, counting its tokens first when that is enabled. + + Counting runs in a background task, so the call returns as soon as the + model has answered. + """ + if _count_tokens: + try: + loop = asyncio.get_running_loop() + model = counting_model(session) + if model is not None: + task = loop.create_task(count_into(kwargs, model, prompt, reply)) + _counting_tasks.add(task) + task.add_done_callback( + functools.partial(send_counted_inference, handle, kwargs) + ) + return + except Exception as exc: + # No running loop, as when an abandoned stream is collected. + debug_apple_fm_failure("token count scheduling", exc) + handle.track_inference(**kwargs) + + # --------------------------------------------------------------------------- # Method patches # --------------------------------------------------------------------------- @@ -259,6 +360,7 @@ def patched_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] model_id = model_id_for(getattr(model, APPLE_FM_USE_CASE_ATTR, None)) handle = c.register_model(self_inner, model_id=model_id) setattr(self_inner, APPLE_FM_HANDLE_ATTR, handle) + setattr(self_inner, APPLE_FM_MODEL_ATTR, model) except Exception as exc: debug_apple_fm_failure("session registration", exc) @@ -281,7 +383,11 @@ async def patched_respond(self_inner, *args, **kwargs): # type: ignore[no-untyp except Exception as exc: track_apple_fm_error(handle, exc, elapsed_ms(t0)) raise - handle.track_inference( + track_generation( + handle, + self_inner, + prompt_of(args, kwargs), + result, duration_ms=elapsed_ms(t0), input_modality=modality, output_modality="generation", @@ -306,11 +412,20 @@ def patched_stream(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] modality, input_meta, config = describe_call(args, kwargs, options_position=1) # Each chunk is the full text so far, so the first one marks the - # first output and nothing needs to be joined. + # first output, the last one is the reply, and nothing is joined. + last_snapshot: list[object] = [None] + + def on_chunk(chunk: object) -> None: + last_snapshot[0] = chunk + def on_done( duration_ms: int, ttft_ms: int | None, *, outcome: str = "completed" ) -> None: - handle.track_inference( + track_generation( + handle, + self_inner, + prompt_of(args, kwargs), + last_snapshot[0] if outcome == "completed" else None, duration_ms=duration_ms, input_modality=modality, output_modality="generation", @@ -323,7 +438,7 @@ def on_done( generation_config=config, ) - return AppleFmStream(stream, handle, time.perf_counter(), None, on_done) + return AppleFmStream(stream, handle, time.perf_counter(), on_chunk, on_done) return patched_stream @@ -389,7 +504,7 @@ def install_auto_load_patch(cls, client_ref: Any) -> None: - ``LanguageModelSession.__init__`` registers the session's model. - ``respond()`` and ``stream_response()`` record inference events. """ - global _apple_fm_patched, _fm + global _apple_fm_patched, _fm, _count_tokens if _apple_fm_patched: return _fm = _fm or optional_module(APPLE_FM_PACKAGE) @@ -408,6 +523,9 @@ def install_auto_load_patch(cls, client_ref: Any) -> None: patch_method(session_cls, "__init__", build_session_init_patch, client_ref) patch_method(session_cls, "respond", build_respond_patch) patch_method(session_cls, "stream_response", build_stream_patch) + _count_tokens = parse_bool( + os.environ.get(constants.ENV_APPLE_FM_COUNT_TOKENS) + ) _apple_fm_patched = True From 1a68ecf8271ecf55df438f788fdb2f6953143bb7 Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Sun, 25 Oct 2026 01:23:19 +0200 Subject: [PATCH 08/11] Describe apple_fm token counts as opt-in, not missing The docs, changelog and module docstring opened with "the SDK doesn't report token usage". SystemLanguageModel.token_count() counts them, and the integration uses it when WILDEDGE_APPLE_FM_COUNT_TOKENS is set. --- CHANGELOG.md | 7 +++---- docs/apple_fm.md | 5 ++--- wildedge/integrations/apple_fm.py | 4 ++-- 3 files changed, 7 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a584983..3dcb994 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,10 +12,9 @@ under Unreleased and move into a version section at release time. time to first token for streams, prompt size and generation options. Failures record the SDK's error name and a matching error code. `SystemLanguageModel.is_available()` records a model load, failed with - the reason when the model can't be used, such as `MODEL_NOT_READY`. The SDK - doesn't report token usage, and `WILDEDGE_APPLE_FM_COUNT_TOKENS=1` counts - it with `SystemLanguageModel.token_count()` in the background. See - docs/apple_fm.md. + the reason when the model can't be used, such as `MODEL_NOT_READY`. Token + counts, from `SystemLanguageModel.token_count()`, are opt-in with + `WILDEDGE_APPLE_FM_COUNT_TOKENS=1`. See docs/apple_fm.md. - Error codes `MODEL_UNAVAILABLE`, `CONTEXT_LENGTH_EXCEEDED`, `CONTENT_FILTERED`, `RATE_LIMITED`, `BAD_REQUEST`, `INVALID_OUTPUT` and `TOOL_CALL_FAILED`, listed in docs/manual-tracking.md. diff --git a/docs/apple_fm.md b/docs/apple_fm.md index 5a34098..3c58476 100644 --- a/docs/apple_fm.md +++ b/docs/apple_fm.md @@ -45,9 +45,8 @@ version, since the model ships with the OS. ### Token counts -The SDK doesn't report token usage, so token counts are empty by default. Set -`WILDEDGE_APPLE_FM_COUNT_TOKENS=1` to count them with -`SystemLanguageModel.token_count()`. Each call then records tokens in (the +Token counts come from `SystemLanguageModel.token_count()` and are opt-in. +Set `WILDEDGE_APPLE_FM_COUNT_TOKENS=1`, and each call records tokens in (the prompt, not the instructions or earlier turns), tokens out (the reply) and tokens per second. diff --git a/wildedge/integrations/apple_fm.py b/wildedge/integrations/apple_fm.py index 579cb1d..fdce534 100644 --- a/wildedge/integrations/apple_fm.py +++ b/wildedge/integrations/apple_fm.py @@ -4,8 +4,8 @@ SDK. Sessions are tied to a model handle when they are created, and ``respond()`` and ``stream_response()`` record one inference event per call. -The SDK reports no token usage. Counting tokens takes two extra calls into -the model service per request, so it is opt-in: set +Token counts come from ``SystemLanguageModel.token_count()``. That takes two +extra calls into the model service per request, so it is opt-in: set ``WILDEDGE_APPLE_FM_COUNT_TOKENS=1``. """ From e3eb6cdfb7bd7b64ae68d7151ea90db1dec87859 Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Sun, 25 Oct 2026 01:28:22 +0200 Subject: [PATCH 09/11] Report the apple_fm accelerator as npu The integration reported coreml because Apple doesn't say where the system model runs. powermetrics on an M2 settles it: while the model generates, Neural Engine power rises from 0 to about 5 W, GPU power stays near idle (38 to 67 mW) and CPU power rises to 1.5 W for the app and the model service. npu is a new accelerator name. The Swift and Android SDKs already use it. --- CHANGELOG.md | 4 +++- docs/apple_fm.md | 3 ++- tests/test_integrations_apple_fm.py | 2 +- wildedge/accelerators.py | 1 + wildedge/integrations/apple_fm.py | 7 ++++--- 5 files changed, 11 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3dcb994..556420c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,7 +14,9 @@ under Unreleased and move into a version section at release time. `SystemLanguageModel.is_available()` records a model load, failed with the reason when the model can't be used, such as `MODEL_NOT_READY`. Token counts, from `SystemLanguageModel.token_count()`, are opt-in with - `WILDEDGE_APPLE_FM_COUNT_TOKENS=1`. See docs/apple_fm.md. + `WILDEDGE_APPLE_FM_COUNT_TOKENS=1`. Its accelerator is `npu`, a new + accelerator name, since the model runs on the Neural Engine. See + docs/apple_fm.md. - Error codes `MODEL_UNAVAILABLE`, `CONTEXT_LENGTH_EXCEEDED`, `CONTENT_FILTERED`, `RATE_LIMITED`, `BAD_REQUEST`, `INVALID_OUTPUT` and `TOOL_CALL_FAILED`, listed in docs/manual-tracking.md. diff --git a/docs/apple_fm.md b/docs/apple_fm.md index 3c58476..2bf93f8 100644 --- a/docs/apple_fm.md +++ b/docs/apple_fm.md @@ -41,7 +41,8 @@ sampling top-k, threshold and seed). Streams also record time to first token. Sessions share one model per use case: `apple-foundation-model` for the general model, and `apple-foundation-model-content-tagging` for `SystemLanguageModelUseCase.CONTENT_TAGGING`. The model version is the macOS -version, since the model ships with the OS. +version, since the model ships with the OS. The accelerator is `npu`, since +the model runs on the Neural Engine. ### Token counts diff --git a/tests/test_integrations_apple_fm.py b/tests/test_integrations_apple_fm.py index 8c908e3..db6719a 100644 --- a/tests/test_integrations_apple_fm.py +++ b/tests/test_integrations_apple_fm.py @@ -248,7 +248,7 @@ def test_registering_a_session_through_the_client(client_with_stubbed_runtime): handle = client_with_stubbed_runtime.register_model(session) assert handle.model_id == MODEL_ID assert handle.info.model_format == "apple_fm" - assert handle.detected_accelerator == accelerators.COREML + assert handle.detected_accelerator == accelerators.NPU assert getattr(session, apple_fm_mod.APPLE_FM_HANDLE_ATTR) is handle client_with_stubbed_runtime.close() diff --git a/wildedge/accelerators.py b/wildedge/accelerators.py index 617e395..c984ec4 100644 --- a/wildedge/accelerators.py +++ b/wildedge/accelerators.py @@ -9,5 +9,6 @@ ROCM = "rocm" # AMD GPU MPS = "mps" # Apple GPU (Metal) COREML = "coreml" # Apple Core ML, which may use the GPU or Neural Engine +NPU = "npu" # Neural processing unit, such as Apple's Neural Engine DIRECTML = "directml" # Windows DirectML OPENVINO = "openvino" # Intel OpenVINO diff --git a/wildedge/integrations/apple_fm.py b/wildedge/integrations/apple_fm.py index fdce534..681b90d 100644 --- a/wildedge/integrations/apple_fm.py +++ b/wildedge/integrations/apple_fm.py @@ -487,9 +487,10 @@ def extract_info( return model_id, info def install_hooks(self, obj: object, handle: ModelHandle) -> None: - # Apple runs the system model on the GPU and Neural Engine; coreml is - # the closest accelerator name. - handle.detected_accelerator = accelerators.COREML + # The system model runs on the Neural Engine. While it generates, + # Neural Engine power rises from 0 to about 5 W on an M2, and GPU + # power stays near idle. + handle.detected_accelerator = accelerators.NPU if type(obj).__name__ == "LanguageModelSession": setattr(obj, APPLE_FM_HANDLE_ATTR, handle) From 80d648fcf079531ac5159acd01e8c393ae5c22cd Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Sun, 25 Oct 2026 01:29:27 +0200 Subject: [PATCH 10/11] Shorten the apple_fm accelerator comment --- wildedge/integrations/apple_fm.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/wildedge/integrations/apple_fm.py b/wildedge/integrations/apple_fm.py index 681b90d..34b43b5 100644 --- a/wildedge/integrations/apple_fm.py +++ b/wildedge/integrations/apple_fm.py @@ -487,9 +487,7 @@ def extract_info( return model_id, info def install_hooks(self, obj: object, handle: ModelHandle) -> None: - # The system model runs on the Neural Engine. While it generates, - # Neural Engine power rises from 0 to about 5 W on an M2, and GPU - # power stays near idle. + # Apple's system model runs on the Neural Engine. handle.detected_accelerator = accelerators.NPU if type(obj).__name__ == "LanguageModelSession": setattr(obj, APPLE_FM_HANDLE_ATTR, handle) From 7a174dcb4cb5be2953b5b648f2cec3f1fdcdbbff Mon Sep 17 00:00:00 2001 From: Piotr Duda Date: Sun, 25 Oct 2026 01:38:06 +0200 Subject: [PATCH 11/11] Keep the model of apple_fm sessions resumed with from_transcript() LanguageModelSession.from_transcript(transcript, model=...) creates the session as cls(_ptr=...), so the patched __init__ never saw the model. A resumed content tagging session was recorded as the general model, and token counts used the default model. from_transcript() is now patched too, and registers the session for the model passed to it. A compat test resumes a content tagging session on the real SDK, and fails without the fix. --- docs/apple_fm.md | 3 +- tests/compat/test_apple_fm_compat.py | 21 ++++++++++ tests/test_integrations_apple_fm.py | 29 +++++++++++++- wildedge/integrations/apple_fm.py | 60 ++++++++++++++++++++++------ 4 files changed, 99 insertions(+), 14 deletions(-) diff --git a/docs/apple_fm.md b/docs/apple_fm.md index 2bf93f8..1b418b4 100644 --- a/docs/apple_fm.md +++ b/docs/apple_fm.md @@ -40,7 +40,8 @@ sampling top-k, threshold and seed). Streams also record time to first token. Sessions share one model per use case: `apple-foundation-model` for the general model, and `apple-foundation-model-content-tagging` for -`SystemLanguageModelUseCase.CONTENT_TAGGING`. The model version is the macOS +`SystemLanguageModelUseCase.CONTENT_TAGGING`, also for sessions resumed with +`from_transcript()`. The model version is the macOS version, since the model ships with the OS. The accelerator is `npu`, since the model runs on the Neural Engine. diff --git a/tests/compat/test_apple_fm_compat.py b/tests/compat/test_apple_fm_compat.py index cb10a42..9721b46 100644 --- a/tests/compat/test_apple_fm_compat.py +++ b/tests/compat/test_apple_fm_compat.py @@ -56,3 +56,24 @@ async def respond(): assert isinstance(result, fm.FoundationModelsError), (reason, result) assert event["event_type"] == "error" assert event["error"]["provider_error_code"] is not None + + +def test_resumed_session_keeps_its_model(compat_client): + """from_transcript() creates the session without passing the model to __init__.""" + compat_client.instrument("apple_fm") + model = fm.SystemLanguageModel( + use_case=fm.SystemLanguageModelUseCase.CONTENT_TAGGING + ) + original = fm.LanguageModelSession(model=model) + session = fm.LanguageModelSession.from_transcript(original.transcript, model=model) + + async def respond(): + try: + return await session.respond("Tag this: snow") + except fm.FoundationModelsError as exc: + return exc + + asyncio.run(respond()) + + [event] = [e for e in compat_client.events if e["event_type"] != "model_load"] + assert event["model_id"] == "apple-foundation-model-content-tagging" diff --git a/tests/test_integrations_apple_fm.py b/tests/test_integrations_apple_fm.py index db6719a..212b87f 100644 --- a/tests/test_integrations_apple_fm.py +++ b/tests/test_integrations_apple_fm.py @@ -296,9 +296,19 @@ async def test_sessions_share_a_handle_per_use_case(install): await sdk.LanguageModelSession(model=sdk.SystemLanguageModel()).respond("b") await sdk.LanguageModelSession(model=tagging).respond("c") await sdk.LanguageModelSession.from_transcript(object()).respond("d") + await sdk.LanguageModelSession.from_transcript(object(), model=tagging).respond("e") + await sdk.LanguageModelSession.from_transcript(object(), tagging).respond("f") model_ids = [e["model_id"] for e in client.of_type("inference")] - assert model_ids == [MODEL_ID, MODEL_ID, f"{MODEL_ID}-content-tagging", MODEL_ID] + tagging_id = f"{MODEL_ID}-content-tagging" + assert model_ids == [ + MODEL_ID, + MODEL_ID, + tagging_id, + MODEL_ID, + tagging_id, + tagging_id, + ] assert set(client.handles) == {MODEL_ID, f"{MODEL_ID}-content-tagging"} @@ -537,13 +547,30 @@ async def test_closed_client_leaves_sessions_untracked(install): assert client.events == [] +async def test_from_transcript_counts_tokens_with_the_passed_model(install): + sdk, client = install(count_tokens=True) + tagging = sdk.SystemLanguageModel( + use_case=SystemLanguageModelUseCase.CONTENT_TAGGING + ) + session = sdk.LanguageModelSession.from_transcript(object(), model=tagging) + assert getattr(session, apple_fm_mod.APPLE_FM_MODEL_ATTR) is tagging + await session.respond("tag this") + await counted() + + [event] = client.of_type("inference") + assert inference_of(event)["output_meta"]["tokens_in"] == 2 + assert apple_fm_mod._default_model is None + + def test_install_is_idempotent(install): sdk, _client = install() respond = sdk.LanguageModelSession.respond + from_transcript = sdk.LanguageModelSession.__dict__["from_transcript"] AppleFmExtractor.install_auto_load_patch(lambda: None) apple_fm_mod._apple_fm_patched = False AppleFmExtractor.install_auto_load_patch(lambda: None) assert sdk.LanguageModelSession.respond is respond + assert sdk.LanguageModelSession.__dict__["from_transcript"] is from_transcript def test_install_skips_when_sdk_missing(monkeypatch): diff --git a/wildedge/integrations/apple_fm.py b/wildedge/integrations/apple_fm.py index 34b43b5..8e6eda5 100644 --- a/wildedge/integrations/apple_fm.py +++ b/wildedge/integrations/apple_fm.py @@ -348,25 +348,44 @@ def patched_is_available(self_inner, *args, **kwargs): # type: ignore[no-untype return patched_is_available +def register_session(client_ref, session: object, model: object) -> None: # type: ignore[no-untyped-def] + """Tie a session to the handle for its model's use case.""" + c = client_ref() # type: ignore[call-arg] + if c is None or c.closed: + return + try: + model_id = model_id_for(getattr(model, APPLE_FM_USE_CASE_ATTR, None)) + handle = c.register_model(session, model_id=model_id) + setattr(session, APPLE_FM_HANDLE_ATTR, handle) + setattr(session, APPLE_FM_MODEL_ATTR, model) + except Exception as exc: + debug_apple_fm_failure("session registration", exc) + + def build_session_init_patch(original_init, client_ref): # type: ignore[no-untyped-def] @functools.wraps(original_init) def patched_init(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] original_init(self_inner, *args, **kwargs) - c = client_ref() # type: ignore[call-arg] - if c is None or c.closed: - return - try: - model = kwargs.get("model", args[1] if len(args) > 1 else None) - model_id = model_id_for(getattr(model, APPLE_FM_USE_CASE_ATTR, None)) - handle = c.register_model(self_inner, model_id=model_id) - setattr(self_inner, APPLE_FM_HANDLE_ATTR, handle) - setattr(self_inner, APPLE_FM_MODEL_ATTR, model) - except Exception as exc: - debug_apple_fm_failure("session registration", exc) + model = kwargs.get("model", args[1] if len(args) > 1 else None) + register_session(client_ref, self_inner, model) return patched_init +def build_from_transcript_patch(original_from_transcript, client_ref): # type: ignore[no-untyped-def] + # from_transcript(transcript, model=None, tools=None) creates the session + # as cls(_ptr=...), so __init__ never sees the model. + @functools.wraps(original_from_transcript) + def patched_from_transcript(cls_inner, *args, **kwargs): # type: ignore[no-untyped-def] + session = original_from_transcript(cls_inner, *args, **kwargs) + model = kwargs.get("model", args[1] if len(args) > 1 else None) + if model is not None: + register_session(client_ref, session, model) + return session + + return patched_from_transcript + + def build_respond_patch(original_respond): # type: ignore[no-untyped-def] @functools.wraps(original_respond) async def patched_respond(self_inner, *args, **kwargs): # type: ignore[no-untyped-def] @@ -453,6 +472,19 @@ def patch_method(cls: type, name: str, make_patch, *extra) -> None: # type: ign setattr(cls, name, patched) +def patch_classmethod(cls: type, name: str, make_patch, *extra) -> None: # type: ignore[no-untyped-def] + original = cls.__dict__.get(name) + if not isinstance(original, classmethod): + return + func = original.__func__ + if getattr(func, "__wildedge_patch_name__", None) == APPLE_FM_PATCH_NAME: + return + patched = make_patch(func, *extra) + patched.__wildedge_patch_name__ = APPLE_FM_PATCH_NAME + patched.__wildedge_original_call__ = func + setattr(cls, name, classmethod(patched)) + + # --------------------------------------------------------------------------- # Extractor # --------------------------------------------------------------------------- @@ -500,7 +532,8 @@ def install_auto_load_patch(cls, client_ref: Any) -> None: model. - ``SystemLanguageModel.is_available()`` records a load event when its answer changes, failed with the reason the model is unavailable. - - ``LanguageModelSession.__init__`` registers the session's model. + - ``LanguageModelSession.__init__`` registers the session's model, and + ``from_transcript()`` the model passed to it. - ``respond()`` and ``stream_response()`` record inference events. """ global _apple_fm_patched, _fm, _count_tokens @@ -520,6 +553,9 @@ def install_auto_load_patch(cls, client_ref: Any) -> None: model_cls, "is_available", build_availability_patch, client_ref ) patch_method(session_cls, "__init__", build_session_init_patch, client_ref) + patch_classmethod( + session_cls, "from_transcript", build_from_transcript_patch, client_ref + ) patch_method(session_cls, "respond", build_respond_patch) patch_method(session_cls, "stream_response", build_stream_patch) _count_tokens = parse_bool(