Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
79 changes: 79 additions & 0 deletions docs/llm_api.md
Original file line number Diff line number Diff line change
@@ -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)
```
3 changes: 2 additions & 1 deletion docs/manual-tracking.md
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
223 changes: 223 additions & 0 deletions tests/test_llm_api.py
Original file line number Diff line number Diff line change
@@ -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
9 changes: 8 additions & 1 deletion tests/test_offline_replay.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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"
3 changes: 3 additions & 0 deletions wildedge/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -52,6 +53,8 @@
"track_span",
"register_model",
"flush",
"llm_api",
"LLMCall",
"Attachment",
"capture_hardware",
"HardwareContext",
Expand Down
6 changes: 4 additions & 2 deletions wildedge/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Loading
Loading