diff --git a/CHANGELOG.md b/CHANGELOG.md index 4d7b234..556420c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,22 @@ 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. + 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, from `SystemLanguageModel.token_count()`, are opt-in with + `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. + ## 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..1b418b4 --- /dev/null +++ b/docs/apple_fm.md @@ -0,0 +1,107 @@ +# 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`, 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. + +### Token counts + +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. + +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 +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 +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` | `INVALID_OUTPUT` | +| `ToolCallError` | `tool_call` | `TOOL_CALL_FAILED` | + +Any other failure records `UNKNOWN`. + +Calling the model before it's ready fails with `AssetsUnavailableError`. + +## Caveats + +- Token counts are opt-in, see [Token counts](#token-counts). +- 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/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/docs/manual-tracking.md b/docs/manual-tracking.md index 0bcc769..d16ea65 100644 --- a/docs/manual-tracking.md +++ b/docs/manual-tracking.md @@ -233,6 +233,25 @@ 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 | +| `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`. ## Track feedback 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..e15b09e --- /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.name}") + + 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..9721b46 --- /dev/null +++ b/tests/compat/test_apple_fm_compat.py @@ -0,0 +1,79 @@ +"""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 availability check must produce one load +event, and the call 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=64) + + async def respond(): + try: + return await session.respond("Tag this: rain", options=options) + except fm.FoundationModelsError as exc: + return exc + + result = asyncio.run(respond()) + + [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 + assert event["event_type"] == "inference" + assert event["inference"]["generation_config"]["max_tokens"] == 64 + else: + 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 new file mode 100644 index 0000000..212b87f --- /dev/null +++ b/tests/test_integrations_apple_fm.py @@ -0,0 +1,581 @@ +"""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.events.inference import GenerationOutputMeta +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 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 + + +class FoundationModelsError(Exception): + pass + + +class GenerationError(FoundationModelsError): + pass + + +class GuardrailViolationError(GenerationError): + pass + + +class ExceededContextWindowSizeError(GenerationError): + pass + + +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", + snapshots=("Hel", "Hello", "Hello there"), + error=None, + stream_error=None, + availability=(True, None), + token_count_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() + + 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 + + @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, 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", {}) + AppleFmExtractor.install_auto_load_patch(lambda: client) + return sdk, client + + return _install + + +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 +# --------------------------------------------------------------------------- + + +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.NPU + 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") + 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")] + 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"} + + +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"] == "CONTENT_FILTERED" + 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"]["error_code"] == "CONTEXT_LENGTH_EXCEEDED" + assert error["error"]["provider_error_code"] == "exceeded_context_window_size" + assert client.of_type("inference") == [] + + +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"] == "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() + ) + + +# --------------------------------------------------------------------------- +# 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 == [] + + +# --------------------------------------------------------------------------- +# 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() + assert await session.respond("hi") == "Hello there" + assert not hasattr(session, apple_fm_mod.APPLE_FM_HANDLE_ATTR) + 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): + 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/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/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/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/events/error.py b/wildedge/events/error.py index a3d6bd2..2bffe9d 100644 --- a/wildedge/events/error.py +++ b/wildedge/events/error.py @@ -17,6 +17,13 @@ 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" + 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 new file mode 100644 index 0000000..8e6eda5 --- /dev/null +++ b/wildedge/integrations/apple_fm.py @@ -0,0 +1,568 @@ +"""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. + +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 __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, constants +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.settings import parse_bool +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__" +# 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" +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", +} + +# 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, + "decoding_failure": ErrorCode.INVALID_OUTPUT, + "tool_call": ErrorCode.TOOL_CALL_FAILED, +} + +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.""" + 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: + provider_error_code = provider_error_code_of(exc) + handle.track_error( + error_code=ERROR_CODES_BY_PROVIDER_CODE.get( + provider_error_code, ErrorCode.UNKNOWN + ), + error_message=str(exc), + provider_error_code=provider_error_code, + duration_ms=duration_ms, + ) + + +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. + + 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)) + + +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 +# --------------------------------------------------------------------------- + + +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_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 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) + 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] + 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 + track_generation( + handle, + self_inner, + prompt_of(args, kwargs), + result, + 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, 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: + 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", + 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(), on_chunk, 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) + + +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 +# --------------------------------------------------------------------------- + + +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'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) + + @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. + - ``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, and + ``from_transcript()`` the model passed to it. + - ``respond()`` and ``stream_response()`` record inference events. + """ + global _apple_fm_patched, _fm, _count_tokens + 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( + 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( + os.environ.get(constants.ENV_APPLE_FM_COUNT_TOKENS) + ) + _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.