feat: add new event methods and update generation metadata fields

Author: JasonLovesDoggo

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "litmus-python-sdk"
-version = "0.4.0"
+version = "0.5.0"
 description = "Litmus Python SDK - implicit evals for AI products"
 readme = "README.md"
 requires-python = ">=3.12"

src/litmus/client.py

@@ -53,6 +53,7 @@
     "$switch_model",
     "$retry_context",
     "$post_accept_edit",
+    "$sessionend",
     "$blur",
     "$return",
     "$scroll_regression",
@@ -106,6 +107,32 @@ def event(self, event_type: EventType, **metadata: object) -> None:
             metadata=merged if merged else None,
         )
 
+    def edit(self, before: str, after: str, **metadata: object) -> None:
+        """Record that the user modified the output before using it.
+
+        Send the raw before/after text. The backend computes edit distance,
+        diff classification, and derived metrics.
+        """
+        self.event("$edit", before=before, after=after, **metadata)
+
+    def accept(self, **metadata: object) -> None:
+        """User used the output as-is."""
+        self.event("$accept", **metadata)
+
+    def copy(self, **metadata: object) -> None:
+        """User copied the output."""
+        self.event("$copy", **metadata)
+
+    def regenerate(self, **metadata: object) -> None:
+        """User requested a new output. Fire BEFORE creating the next generation."""
+        self.event("$regenerate", **metadata)
+
+    def share(self, channel: str | None = None, **metadata: object) -> None:
+        """User shared the output."""
+        if channel:
+            metadata["channel"] = channel
+        self.event("$share", **metadata)
+
 
 class Feature:
     """Scoped handle for an AI feature. Carries defaults so you don't
@@ -271,6 +298,14 @@ def track(
         generation_id: str | None = None,
         metadata: dict | None = None,
         timestamp: datetime | None = None,
+        model: str | None = None,
+        provider: str | None = None,
+        input_tokens: int | None = None,
+        output_tokens: int | None = None,
+        total_tokens: int | None = None,
+        duration_ms: int | None = None,
+        ttft_ms: int | None = None,
+        cost: float | None = None,
     ) -> str | None:
         """Enqueue a single event. Returns the event UUID or None if dropped."""
         if self.disabled:
@@ -294,6 +329,23 @@ def track(
         if generation_id:
             msg["generation_id"] = generation_id
 
+        if model:
+            msg["model"] = model
+        if provider:
+            msg["provider"] = provider
+        if input_tokens is not None:
+            msg["input_tokens"] = input_tokens
+        if output_tokens is not None:
+            msg["output_tokens"] = output_tokens
+        if total_tokens is not None:
+            msg["total_tokens"] = total_tokens
+        if duration_ms is not None:
+            msg["duration_ms"] = duration_ms
+        if ttft_ms is not None:
+            msg["ttft_ms"] = ttft_ms
+        if cost is not None:
+            msg["cost"] = cost
+
         props = {"$lib": "litmus-python", "$lib_version": VERSION}
         if metadata:
             props.update(metadata)
@@ -328,10 +380,25 @@ def generation(
         prompt_id: str | None = None,
         prompt_version: str | None = None,
         model: str | None = None,
+        provider: str | None = None,
+        input_tokens: int | None = None,
+        output_tokens: int | None = None,
+        total_tokens: int | None = None,
+        duration_ms: int | None = None,
+        ttft_ms: int | None = None,
+        cost: float | None = None,
         metadata: dict | None = None,
+        generation_id: str | None = None,
     ) -> Generation:
-        """Create a generation and return a handle for recording signals."""
-        generation_id = str(uuid4())
+        """Create a generation and return a handle for recording signals.
+
+        Pass ``generation_id`` when the caller has already minted an id
+        (e.g. to correlate with an OTel/OpenRouter trace broadcast). When
+        omitted, a fresh UUID4 is used. The id is emitted on the
+        ``$generation`` event, so downstream joins work either way.
+        """
+        if generation_id is None:
+            generation_id = str(uuid4())
         defaults = {
             "user_id": user_id,
             "prompt_id": prompt_id,
@@ -348,6 +415,14 @@ def generation(
             prompt_id=prompt_id,
             prompt_version=prompt_version,
             metadata=metadata,
+            model=model,
+            provider=provider,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            total_tokens=total_tokens,
+            duration_ms=duration_ms,
+            ttft_ms=ttft_ms,
+            cost=cost,
         )
 
         return Generation(self, session_id, generation_id, defaults)

src/litmus/consumer.py

@@ -36,7 +36,7 @@ def __init__(
         api_key: str,
         host: str | None = None,
         on_error: Callable[[Exception, list[dict]], None] | None = None,
-        flush_at: int = 100,
+        flush_at: int = 10,
         flush_interval: float = 0.5,
         use_gzip: bool = False,
         retries: int = 10,

tests/test_client.py

@@ -260,6 +260,95 @@ def test_feature_track(self, mock_transport: CallLog) -> None:
         assert events[0]["user_id"] == "u1"
 
 
+class TestGenerationIdOverride:
+    """generation() accepts an externally-minted generation_id so callers
+    can correlate the $generation event with ids they already emit elsewhere
+    (e.g. OTel span metadata broadcast by OpenRouter)."""
+
+    def test_uses_provided_generation_id(self, mock_transport: CallLog) -> None:
+        client = LitmusClient(api_key="ltm_pk_test_abc", sync_mode=True)
+        gen = client.generation("s1", prompt_id="chat", generation_id="external-uuid")
+        assert gen.id == "external-uuid"
+
+        gen.event("$accept")
+        events = mock_transport.events()
+        assert all(e["generation_id"] == "external-uuid" for e in events)
+
+    def test_default_generates_uuid_when_absent(self, mock_transport: CallLog) -> None:
+        client = LitmusClient(api_key="ltm_pk_test_abc", sync_mode=True)
+        gen = client.generation("s1")
+        # UUID4 string: 36 chars, 4 hyphens
+        assert len(gen.id) == 36
+        assert gen.id.count("-") == 4
+
+
+class TestGenerationTopLevelFields:
+    """generation() forwards model, provider, tokens, cost, latency as
+    TOP-LEVEL wire fields on the $generation event. The ingest server
+    stores these in dedicated columns (model TEXT, provider TEXT, etc.)
+    so dashboards can query them without JSONB lookups.
+
+    Regression guard for the v0.4.0 bug where these were accepted as
+    kwargs but silently dropped before send."""
+
+    def test_model_is_top_level(self, mock_transport: CallLog) -> None:
+        client = LitmusClient(api_key="ltm_pk_test_abc", sync_mode=True)
+        client.generation("s1", prompt_id="chat", model="claude-sonnet-4-20250514")
+
+        event = next(e for e in mock_transport.events() if e["type"] == "$generation")
+        assert event["model"] == "claude-sonnet-4-20250514"
+        # Must NOT be in metadata — that was the old path we're moving off.
+        assert event["metadata"].get("model") is None
+
+    def test_provider_and_usage_are_top_level(self, mock_transport: CallLog) -> None:
+        client = LitmusClient(api_key="ltm_pk_test_abc", sync_mode=True)
+        client.generation(
+            "s1",
+            prompt_id="chat",
+            model="gpt-4o",
+            provider="openai",
+            input_tokens=120,
+            output_tokens=340,
+            total_tokens=460,
+            duration_ms=1850,
+            ttft_ms=240,
+            cost=0.0042,
+        )
+
+        event = next(e for e in mock_transport.events() if e["type"] == "$generation")
+        assert event["model"] == "gpt-4o"
+        assert event["provider"] == "openai"
+        assert event["input_tokens"] == 120
+        assert event["output_tokens"] == 340
+        assert event["total_tokens"] == 460
+        assert event["duration_ms"] == 1850
+        assert event["ttft_ms"] == 240
+        assert event["cost"] == 0.0042
+
+    def test_optional_fields_omitted_when_none(self, mock_transport: CallLog) -> None:
+        """None-valued fields are dropped from the wire payload entirely —
+        the ingest server treats missing and null the same, but we keep
+        payloads lean."""
+        client = LitmusClient(api_key="ltm_pk_test_abc", sync_mode=True)
+        client.generation("s1", prompt_id="chat")
+
+        event = next(e for e in mock_transport.events() if e["type"] == "$generation")
+        for key in ("model", "provider", "input_tokens", "output_tokens",
+                    "total_tokens", "duration_ms", "ttft_ms", "cost"):
+            assert key not in event, f"{key} should be omitted when None"
+
+    def test_feature_forwards_model_top_level(self, mock_transport: CallLog) -> None:
+        """Feature-scoped defaults also end up as top-level wire fields so
+        dashboards using feature() get the same model column as callers
+        using generation() directly."""
+        client = LitmusClient(api_key="ltm_pk_test_abc", sync_mode=True)
+        feat = client.feature("summarizer", model="claude-haiku-4")
+        feat.generation("s1")
+
+        event = next(e for e in mock_transport.events() if e["type"] == "$generation")
+        assert event["model"] == "claude-haiku-4"
+
+
 class TestAttach:
     """attach() returns a Generation handle without emitting $generation."""