Add wildedge.llm_api(): track LLM calls made over any HTTP API (#43)

Author: piotrekno1

CHANGELOG.md

@@ -7,6 +7,8 @@ under Unreleased and move into a version section at release time.
 
 ### Added
 
+- `wildedge.llm_api()`: provider-agnostic tracking for LLM calls made with any HTTP client (OpenRouter, vLLM, Ollama, OpenAI/Anthropic-compatible endpoints). Times the block, normalizes usage payloads from either provider shape via `call.usage()` / `call.response()`, supports TTFT marks and async use, and records exceptions as error events. See `docs/llm_api.md`.
+- `register_model()` without a matching extractor defaults the model name to the explicit `model_id` instead of the placeholder object's type name.
 - Process-wide default client: `wildedge.get_client()`, `wildedge.set_default_client()`, and module-level `trace`, `span`, `track_span`, `register_model`, `flush` delegating to it (#41)
 - `init()` reuses the default client installed by `wildedge run` unless `dsn` is passed, so CLI and in-process init share one client (#41)
 - `SpanContextManager.set_attributes()` and `fail()` for recording outcomes on an open span (#41)

docs/llm_api.md

@@ -0,0 +1,79 @@
+# Tracking LLM calls without a client library
+
+`wildedge.llm_api()` records LLM API calls made with any HTTP client: httpx or
+requests against OpenRouter, vLLM, Ollama, or any OpenAI-compatible or
+Anthropic-compatible endpoint. Use it when auto-instrumentation does not
+apply because your app does not use the `openai` or `anthropic` packages. If
+it does use them, prefer the integrations; they capture the same events with
+zero code.
+
+The name is deliberate: this tracks calls to an LLM behind an *API boundary*.
+An LLM running inside your own process (llama.cpp, transformers, MLX) is
+covered by the [framework integrations](manual-tracking.md), which also
+capture what no API boundary can expose: quantization from the model
+artifact, memory footprint, load/unload timing, and hardware linkage.
+
+The block is timed automatically, events correlate with surrounding
+`wildedge.trace()` / `wildedge.span()` blocks, and everything is a silent
+no-op without a DSN.
+
+## Quickstart
+
+```python
+import httpx
+import wildedge
+
+async def generate(prompt: str, model: str) -> dict:
+    with wildedge.llm_api(model=model, provider="openrouter", prompt=prompt) as call:
+        async with httpx.AsyncClient(timeout=300) as http:
+            response = await http.post(
+                "https://openrouter.ai/api/v1/chat/completions",
+                headers={"Authorization": f"Bearer {API_KEY}"},
+                json={"model": model, "messages": [{"role": "user", "content": prompt}]},
+            )
+        data = response.json()
+        call.response(data)
+    return data
+```
+
+`call.response()` pulls token usage, stop reason, and API metadata from the
+full response payload, dict or SDK object, OpenAI shape
+(`usage.prompt_tokens`, `choices[0].finish_reason`) or Anthropic shape
+(`usage.input_tokens`, top-level `stop_reason`).
+
+## Recording pieces individually
+
+When you do not have a full response payload, set what you know:
+
+```python
+with wildedge.llm_api(model="gemma-7b", base_url="http://localhost:11434") as call:
+    result = post_to_ollama(...)
+    call.usage(result["usage"])            # payload in either provider shape
+    call.usage(tokens_in=52, tokens_out=209)  # or explicit fields; these win
+    call.stop_reason = "stop"
+    call.first_token()                     # TTFT mark for streaming
+    call.success = False                   # delivered but unusable output
+```
+
+An exception escaping the block records an error event with the exception
+class as the error code, and no inference event.
+
+## Model identity
+
+Events register under `model` as the model id. `provider` names the source
+directly; alternatively `base_url` derives it (`openrouter.ai` becomes
+`openrouter`, unknown hosts record the hostname). Pass `prompt` or `messages`
+(chat format) to attach input metadata such as prompt length.
+
+## Agentic pipelines
+
+Combine with traces for multi-step visibility:
+
+```python
+with wildedge.trace(run_id=run_id, agent_id="skin-generator"):
+    for attempt in range(2):
+        with wildedge.span(kind="agent_step", name="generate", step_index=attempt):
+            with wildedge.llm_api(model=model, provider="openrouter", prompt=prompt) as call:
+                data = await post_chat_completion(prompt)
+                call.response(data)
+```

docs/manual-tracking.md

@@ -5,7 +5,8 @@ Use manual tracking when your framework is not covered by a WildEdge integration
 ## When to use it
 
 - Your model class is custom (e.g. a `torch.nn.Module` subclass not loaded via `timm` or `transformers`)
-- You are calling a remote API not yet covered by an integration
+- You are calling a remote API not yet covered by an integration (for LLM
+  APIs called over raw HTTP, [`wildedge.llm_api()`](llm_api.md) is the shortcut)
 - You want to attach input/output metadata (token counts, image dimensions, confidence scores, etc.)
 - You want to record user feedback tied to a specific inference
 

tests/test_llm_api.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+import asyncio
+from types import SimpleNamespace
+
+import pytest
+
+import wildedge
+from wildedge import constants
+from wildedge.defaults import set_default_client
+
+
+class FakeHandle:
+    def __init__(self):
+        self.inferences: list[dict] = []
+        self.errors: list[dict] = []
+
+    def track_inference(self, **kwargs):
+        self.inferences.append(kwargs)
+
+    def track_error(self, error_code=None, *, error_message=None, **kwargs):
+        self.errors.append({"error_code": error_code, "error_message": error_message})
+
+
+class FakeClient:
+    def __init__(self):
+        self.handle = FakeHandle()
+        self.registered: list[dict] = []
+
+    def register_model(self, model_obj, **kwargs):
+        self.registered.append({"model_obj": model_obj, **kwargs})
+        return self.handle
+
+
+@pytest.fixture
+def client():
+    fake = FakeClient()
+    set_default_client(fake)
+    return fake
+
+
+def test_openai_shape_usage(client):
+    with wildedge.llm_api(model="openai/gpt-4o-mini", provider="openrouter") as call:
+        call.usage(
+            {
+                "prompt_tokens": 100,
+                "completion_tokens": 50,
+                "prompt_tokens_details": {"cached_tokens": 25},
+                "completion_tokens_details": {"reasoning_tokens": 10},
+            }
+        )
+        call.stop_reason = "stop"
+
+    assert client.registered[0]["model_id"] == "openai/gpt-4o-mini"
+    assert client.registered[0]["source"] == "openrouter"
+    event = client.handle.inferences[0]
+    assert event["input_modality"] == "text"
+    assert event["output_modality"] == "generation"
+    assert event["success"] is True
+    meta = event["output_meta"]
+    assert meta.tokens_in == 100
+    assert meta.tokens_out == 50
+    assert meta.cached_input_tokens == 25
+    assert meta.reasoning_tokens_out == 10
+    assert meta.stop_reason == "stop"
+
+
+def test_anthropic_shape_usage(client):
+    with wildedge.llm_api(model="claude-sonnet-4-5", provider="anthropic") as call:
+        call.usage(
+            {"input_tokens": 10, "output_tokens": 5, "cache_read_input_tokens": 3}
+        )
+
+    meta = client.handle.inferences[0]["output_meta"]
+    assert meta.tokens_in == 10
+    assert meta.tokens_out == 5
+    assert meta.cached_input_tokens == 3
+
+
+def test_attribute_object_usage(client):
+    usage = SimpleNamespace(
+        prompt_tokens=7,
+        completion_tokens=3,
+        prompt_tokens_details=SimpleNamespace(cached_tokens=2),
+        completion_tokens_details=None,
+    )
+    with wildedge.llm_api(model="m", provider="p") as call:
+        call.usage(usage)
+
+    meta = client.handle.inferences[0]["output_meta"]
+    assert meta.tokens_in == 7
+    assert meta.tokens_out == 3
+    assert meta.cached_input_tokens == 2
+
+
+def test_explicit_field_usage_wins_over_payload(client):
+    with wildedge.llm_api(model="m", provider="p") as call:
+        call.usage({"prompt_tokens": 1}, tokens_in=11, tokens_out=22)
+
+    meta = client.handle.inferences[0]["output_meta"]
+    assert meta.tokens_in == 11
+    assert meta.tokens_out == 22
+
+
+def test_response_openai_shape(client):
+    with wildedge.llm_api(model="m", provider="p") as call:
+        call.response(
+            {
+                "model": "gpt-4o-mini-2024-07-18",
+                "system_fingerprint": "fp_1",
+                "usage": {"prompt_tokens": 4, "completion_tokens": 2},
+                "choices": [{"finish_reason": "length", "message": {}}],
+            }
+        )
+
+    event = client.handle.inferences[0]
+    assert event["output_meta"].stop_reason == "length"
+    assert event["output_meta"].tokens_in == 4
+    assert event["api_meta"].resolved_model_id == "gpt-4o-mini-2024-07-18"
+    assert event["api_meta"].system_fingerprint == "fp_1"
+
+
+def test_response_anthropic_shape(client):
+    with wildedge.llm_api(model="m", provider="anthropic") as call:
+        call.response(
+            {
+                "model": "claude-sonnet-4-5",
+                "stop_reason": "end_turn",
+                "usage": {"input_tokens": 9, "output_tokens": 1},
+            }
+        )
+
+    meta = client.handle.inferences[0]["output_meta"]
+    assert meta.stop_reason == "end_turn"
+    assert meta.tokens_in == 9
+
+
+def test_exception_records_error_not_inference(client):
+    with pytest.raises(ValueError):
+        with wildedge.llm_api(model="m", provider="p"):
+            raise ValueError("bad payload")
+
+    assert client.handle.inferences == []
+    assert client.handle.errors == [
+        {"error_code": "ValueError", "error_message": "bad payload"}
+    ]
+
+
+def test_prompt_input_meta(client):
+    with wildedge.llm_api(
+        model="m", provider="p", prompt="a knight with two swords"
+    ) as call:
+        call.usage(tokens_in=6)
+
+    meta = client.handle.inferences[0]["input_meta"]
+    assert meta.char_count == len("a knight with two swords")
+    assert meta.word_count == 5
+    assert meta.token_count == 6
+    assert meta.prompt_type == "chat"
+
+
+def test_messages_input_meta(client):
+    messages = [
+        {"role": "system", "content": "be brief"},
+        {"role": "user", "content": "hello there"},
+    ]
+    with wildedge.llm_api(model="m", provider="p", messages=messages) as call:
+        call.usage(tokens_in=3)
+
+    meta = client.handle.inferences[0]["input_meta"]
+    assert meta.char_count == len("hello there")
+    assert meta.word_count == 2
+
+
+def test_first_token_sets_ttft(client):
+    with wildedge.llm_api(model="m", provider="p") as call:
+        call.first_token()
+        call.usage(tokens_out=1)
+
+    assert client.handle.inferences[0]["output_meta"].time_to_first_token_ms is not None
+
+
+def test_async_context_manager(client):
+    async def run():
+        async with wildedge.llm_api(model="m", provider="p") as call:
+            call.usage(tokens_in=1)
+
+    asyncio.run(run())
+    assert len(client.handle.inferences) == 1
+
+
+def test_source_derived_from_base_url(client):
+    with wildedge.llm_api(model="m", base_url="https://openrouter.ai/api/v1"):
+        pass
+
+    assert client.registered[0]["source"] == "openrouter"
+
+
+def test_success_flag_recorded(client):
+    with wildedge.llm_api(model="m", provider="p") as call:
+        call.usage(tokens_out=1)
+        call.success = False
+
+    assert client.handle.inferences[0]["success"] is False
+
+
+def test_no_output_meta_when_nothing_known(client):
+    with wildedge.llm_api(model="m", provider="p"):
+        pass
+
+    event = client.handle.inferences[0]
+    assert event["output_meta"] is None
+    assert event["api_meta"] is None
+
+
+def test_noop_without_dsn_end_to_end(monkeypatch):
+    monkeypatch.delenv(constants.ENV_DSN, raising=False)
+    monkeypatch.delenv(constants.ENV_INTEGRATIONS, raising=False)
+
+    with wildedge.llm_api(model="m", provider="p") as call:
+        call.usage(tokens_in=1, tokens_out=2)
+
+    assert wildedge.get_client().noop is True

tests/test_offline_replay.py

@@ -37,6 +37,12 @@ def test_offline_replay_restores_model_registry_for_pending_events(tmp_path):
         ),
         patch("wildedge.client.Transmitter"),
         patch("wildedge.client.Consumer", _DummyConsumer),
+        # The registry path has no constructor override; without this patch the
+        # test reads and writes machine-global state across runs.
+        patch(
+            "wildedge.client.default_model_registry_path",
+            return_value=tmp_path / "model_registry.json",
+        ),
     ):
         client_a = WildEdge(
             dsn="https://secret@ingest.wildedge.dev/proj",
@@ -69,4 +75,5 @@ def test_offline_replay_restores_model_registry_for_pending_events(tmp_path):
     assert client_b.queue.length() == 1
     models = client_b.registry.snapshot()
     assert "ResNet" in models
-    assert models["ResNet"]["model_name"] == "_Model"
+    # register_model with no matching extractor names the model after model_id
+    assert models["ResNet"]["model_name"] == "ResNet"

wildedge/__init__.py

@@ -29,6 +29,7 @@
     TextInputMeta,
 )
 from wildedge.events.span import SpanKind, SpanStatus
+from wildedge.llm_api import LLMCall, llm_api
 from wildedge.platforms import capture_hardware
 from wildedge.platforms.device_info import DeviceInfo
 from wildedge.platforms.hardware import HardwareContext, ThermalContext
@@ -52,6 +53,8 @@
     "track_span",
     "register_model",
     "flush",
+    "llm_api",
+    "LLMCall",
     "Attachment",
     "capture_hardware",
     "HardwareContext",

wildedge/client.py

@@ -529,8 +529,10 @@ def register_model(
         else:
             # No extractor matched - require explicit id
             model_id = overrides.pop("id", None)
-            model_name = overrides.pop("model_name", None) or (
-                str(type(model_obj).__name__)
+            model_name = (
+                overrides.pop("model_name", None)
+                or model_id
+                or str(type(model_obj).__name__)
             )
             info = ModelInfo(
                 model_name=model_name,