fluxmeter 1.0.0__tar.gz

fluxmeter-1.0.0/.gitignore

@@ -0,0 +1,7 @@
+dist/
+*.egg-info/
+.eggs/
+build/
+.pytest_cache/
+.ruff_cache/
+__pycache__/

fluxmeter-1.0.0/PKG-INFO

@@ -0,0 +1,133 @@
+Metadata-Version: 2.4
+Name: fluxmeter
+Version: 1.0.0
+Summary: Python SDK for FluxMeter — streaming metering for AI token billing
+Project-URL: Homepage, https://github.com/10kshuaizhang/fluxmeter
+Project-URL: Repository, https://github.com/10kshuaizhang/fluxmeter
+Author-email: FluxMeter <hello@fluxmeter.dev>
+License-Expression: Apache-2.0
+Keywords: ai,anthropic,billing,llm,metering,openai,streaming,tokens
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Software Development :: Libraries
+Requires-Python: >=3.9
+Requires-Dist: confluent-kafka>=2.3.0
+Provides-Extra: anthropic
+Requires-Dist: anthropic>=0.20; extra == 'anthropic'
+Provides-Extra: dev
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: pytest-asyncio; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Provides-Extra: openai
+Requires-Dist: openai>=1.0; extra == 'openai'
+Description-Content-Type: text/markdown
+
+# FluxMeter Python SDK
+
+Send AI token usage events to FluxMeter for real-time aggregation and billing.
+
+## Install
+
+```bash
+pip install fluxmeter
+```
+
+## Quick Start (3 lines)
+
+```python
+from fluxmeter import FluxMeter
+
+meter = FluxMeter(kafka_brokers="localhost:9094")
+meter.track("cust_123", "gpt-4o", input_tokens=500, output_tokens=150)
+```
+
+## OpenAI Integration
+
+```python
+import time
+from openai import OpenAI
+from fluxmeter import FluxMeter
+
+client = OpenAI()
+meter = FluxMeter(kafka_brokers="localhost:9094", environment="production")
+
+start = time.time()
+response = client.chat.completions.create(
+    model="gpt-4o",
+    messages=[{"role": "user", "content": "Hello!"}],
+)
+latency = int((time.time() - start) * 1000)
+
+# One line to meter the usage
+meter.track_openai("cust_123", response, latency_ms=latency)
+```
+
+## Anthropic Integration
+
+```python
+import anthropic
+from fluxmeter import FluxMeter
+
+client = anthropic.Anthropic()
+meter = FluxMeter(kafka_brokers="localhost:9094")
+
+response = client.messages.create(
+    model="claude-sonnet-4-20250514",
+    max_tokens=1024,
+    messages=[{"role": "user", "content": "Hello!"}],
+)
+
+meter.track_anthropic("cust_123", response)
+```
+
+## Manual Tracking (any provider)
+
+```python
+meter.track(
+    customer_id="cust_123",
+    model_id="gemini-1.5-pro",
+    provider="google",
+    input_tokens=2000,
+    output_tokens=500,
+    request_id="req_abc123",
+    span_id="span_7f3a",          # link to your tracing
+    session_id="sess_456",        # group by conversation
+    latency_ms=890,
+    environment="production",
+    metadata={"feature": "code-review", "team": "platform"},
+)
+```
+
+## Configuration
+
+```python
+meter = FluxMeter(
+    kafka_brokers="kafka1:9092,kafka2:9092",  # Kafka cluster
+    topic="token-events",                       # Topic name (default)
+    environment="production",                   # Applied to all events
+    producer_config={                           # Extra Kafka producer config
+        "security.protocol": "SASL_SSL",
+        "sasl.mechanisms": "PLAIN",
+        "sasl.username": "...",
+        "sasl.password": "...",
+    },
+)
+```
+
+## How It Works
+
+```
+Your App  →  meter.track(...)  →  Kafka  →  Flink (real-time aggregation)  →  Redis
+                                                                                 ↓
+                                                                           Grafana / API
+```
+
+Events are batched and compressed (lz4) before sending. The SDK flushes automatically on process exit.
+
+## Requirements
+
+- Python 3.9+
+- `confluent-kafka` (librdkafka-based, high performance)
+- FluxMeter infrastructure running (Kafka + Flink + Redis)

fluxmeter-1.0.0/README.md

@@ -0,0 +1,107 @@
+# FluxMeter Python SDK
+
+Send AI token usage events to FluxMeter for real-time aggregation and billing.
+
+## Install
+
+```bash
+pip install fluxmeter
+```
+
+## Quick Start (3 lines)
+
+```python
+from fluxmeter import FluxMeter
+
+meter = FluxMeter(kafka_brokers="localhost:9094")
+meter.track("cust_123", "gpt-4o", input_tokens=500, output_tokens=150)
+```
+
+## OpenAI Integration
+
+```python
+import time
+from openai import OpenAI
+from fluxmeter import FluxMeter
+
+client = OpenAI()
+meter = FluxMeter(kafka_brokers="localhost:9094", environment="production")
+
+start = time.time()
+response = client.chat.completions.create(
+    model="gpt-4o",
+    messages=[{"role": "user", "content": "Hello!"}],
+)
+latency = int((time.time() - start) * 1000)
+
+# One line to meter the usage
+meter.track_openai("cust_123", response, latency_ms=latency)
+```
+
+## Anthropic Integration
+
+```python
+import anthropic
+from fluxmeter import FluxMeter
+
+client = anthropic.Anthropic()
+meter = FluxMeter(kafka_brokers="localhost:9094")
+
+response = client.messages.create(
+    model="claude-sonnet-4-20250514",
+    max_tokens=1024,
+    messages=[{"role": "user", "content": "Hello!"}],
+)
+
+meter.track_anthropic("cust_123", response)
+```
+
+## Manual Tracking (any provider)
+
+```python
+meter.track(
+    customer_id="cust_123",
+    model_id="gemini-1.5-pro",
+    provider="google",
+    input_tokens=2000,
+    output_tokens=500,
+    request_id="req_abc123",
+    span_id="span_7f3a",          # link to your tracing
+    session_id="sess_456",        # group by conversation
+    latency_ms=890,
+    environment="production",
+    metadata={"feature": "code-review", "team": "platform"},
+)
+```
+
+## Configuration
+
+```python
+meter = FluxMeter(
+    kafka_brokers="kafka1:9092,kafka2:9092",  # Kafka cluster
+    topic="token-events",                       # Topic name (default)
+    environment="production",                   # Applied to all events
+    producer_config={                           # Extra Kafka producer config
+        "security.protocol": "SASL_SSL",
+        "sasl.mechanisms": "PLAIN",
+        "sasl.username": "...",
+        "sasl.password": "...",
+    },
+)
+```
+
+## How It Works
+
+```
+Your App  →  meter.track(...)  →  Kafka  →  Flink (real-time aggregation)  →  Redis
+                                                                                 ↓
+                                                                           Grafana / API
+```
+
+Events are batched and compressed (lz4) before sending. The SDK flushes automatically on process exit.
+
+## Requirements
+
+- Python 3.9+
+- `confluent-kafka` (librdkafka-based, high performance)
+- FluxMeter infrastructure running (Kafka + Flink + Redis)

fluxmeter-1.0.0/fluxmeter/__init__.py

@@ -0,0 +1,8 @@
+"""FluxMeter — streaming metering SDK for AI token billing."""
+
+from fluxmeter.client import FluxMeter
+from fluxmeter.event import TokenEvent
+from fluxmeter.streaming import StreamingWrapper
+
+__version__ = "1.0.0"
+__all__ = ["FluxMeter", "TokenEvent", "StreamingWrapper"]

fluxmeter-1.0.0/fluxmeter/client.py

@@ -0,0 +1,398 @@
+"""FluxMeter client — sends token usage events to Kafka."""
+
+from __future__ import annotations
+
+import json
+import logging
+import atexit
+import threading
+import time
+from typing import Optional
+
+from confluent_kafka import Producer
+
+from fluxmeter.event import TokenEvent
+from fluxmeter.streaming import StreamingWrapper
+from fluxmeter.wal import WriteAheadLog
+
+logger = logging.getLogger(__name__)
+
+
+class FluxMeter:
+    """Main FluxMeter client. Sends token events to Kafka for real-time aggregation.
+
+    Events are persisted to a local WAL (write-ahead log) BEFORE sending to Kafka.
+    If Kafka is unavailable, events accumulate on disk and flush when it recovers.
+    This guarantees zero event loss regardless of Kafka availability.
+
+    Usage:
+        from fluxmeter import FluxMeter
+
+        meter = FluxMeter(kafka_brokers="localhost:9094")
+        meter.track(customer_id="cust_123", model_id="gpt-4o", input_tokens=500, output_tokens=150)
+    """
+
+    def __init__(
+        self,
+        kafka_brokers: str = "localhost:9094",
+        topic: str = "token-events",
+        environment: Optional[str] = None,
+        producer_config: Optional[dict] = None,
+        wal_enabled: bool = True,
+        wal_path: str = "~/.fluxmeter/wal",
+    ):
+        self._topic = topic
+        self._environment = environment
+        self._delivery_errors = 0
+        self._events_sent = 0
+        self._wal_enabled = wal_enabled
+
+        config = {
+            "bootstrap.servers": kafka_brokers,
+            "linger.ms": 5,
+            "batch.num.messages": 10000,
+            "compression.type": "lz4",
+            "acks": "all",  # Wait for all replicas (no data loss on broker crash)
+        }
+        if producer_config:
+            config.update(producer_config)
+
+        self._producer = Producer(config)
+
+        # Local WAL: events persisted to disk before Kafka send
+        if wal_enabled:
+            self._wal = WriteAheadLog(path=wal_path)
+            self._flush_thread = threading.Thread(target=self._wal_flush_loop, daemon=True)
+            self._flush_thread.start()
+        else:
+            self._wal = None
+
+        atexit.register(self.flush)
+
+    def track(
+        self,
+        customer_id: str,
+        model_id: str,
+        *,
+        provider: str = "openai",
+        input_tokens: int = 0,
+        output_tokens: int = 0,
+        cache_read_tokens: int = 0,
+        cache_write_tokens: int = 0,
+        reasoning_tokens: int = 0,
+        embedding_tokens: int = 0,
+        request_id: Optional[str] = None,
+        span_id: Optional[str] = None,
+        parent_span_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        latency_ms: int = 0,
+        environment: Optional[str] = None,
+        metadata: Optional[dict[str, str]] = None,
+    ) -> TokenEvent:
+        """Track a single LLM API call's token usage.
+
+        Args:
+            customer_id: Your customer/tenant identifier.
+            model_id: Model name (e.g. "gpt-4o", "claude-sonnet-4").
+            provider: Provider name ("openai", "anthropic", "google").
+            input_tokens: Prompt/input token count.
+            output_tokens: Completion/output token count.
+            cache_read_tokens: Cached prompt tokens read.
+            cache_write_tokens: Tokens written to prompt cache.
+            reasoning_tokens: Internal reasoning tokens (o1/o3).
+            embedding_tokens: Embedding tokens.
+            request_id: Provider's request ID.
+            span_id: Observability span ID.
+            session_id: Conversation/session identifier.
+            latency_ms: Provider response time in milliseconds.
+            environment: Override instance-level environment.
+            metadata: Arbitrary key-value pairs.
+
+        Returns:
+            The TokenEvent that was sent.
+        """
+        event = TokenEvent(
+            customer_id=customer_id,
+            model_id=model_id,
+            provider=provider,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cache_read_tokens=cache_read_tokens,
+            cache_write_tokens=cache_write_tokens,
+            reasoning_tokens=reasoning_tokens,
+            embedding_tokens=embedding_tokens,
+            request_id=request_id,
+            span_id=span_id,
+            parent_span_id=parent_span_id,
+            session_id=session_id,
+            latency_ms=latency_ms,
+            environment=environment or self._environment,
+            metadata=metadata,
+        )
+        self._send(event)
+        return event
+
+    def track_openai(
+        self,
+        customer_id: str,
+        response,
+        *,
+        session_id: Optional[str] = None,
+        span_id: Optional[str] = None,
+        latency_ms: int = 0,
+        environment: Optional[str] = None,
+    ) -> TokenEvent:
+        """Track usage from an OpenAI ChatCompletion response object.
+
+        Args:
+            customer_id: Your customer/tenant identifier.
+            response: OpenAI ChatCompletion response (or dict).
+            session_id: Optional conversation session ID.
+            span_id: Optional observability span ID.
+            latency_ms: Request latency in ms.
+            environment: Override instance-level environment.
+
+        Returns:
+            The TokenEvent that was sent.
+        """
+        # Handle both object and dict responses
+        if hasattr(response, "model"):
+            model = response.model
+            usage = response.usage
+            request_id = response.id
+        else:
+            model = response["model"]
+            usage = response["usage"]
+            request_id = response.get("id")
+
+        # Extract token counts from usage
+        if hasattr(usage, "prompt_tokens"):
+            input_tokens = usage.prompt_tokens or 0
+            output_tokens = usage.completion_tokens or 0
+            cache_read = getattr(usage, "prompt_tokens_details", None)
+            cache_read_tokens = (
+                getattr(cache_read, "cached_tokens", 0) if cache_read else 0
+            )
+            reasoning = getattr(usage, "completion_tokens_details", None)
+            reasoning_tokens = (
+                getattr(reasoning, "reasoning_tokens", 0) if reasoning else 0
+            )
+        else:
+            input_tokens = usage.get("prompt_tokens", 0)
+            output_tokens = usage.get("completion_tokens", 0)
+            details = usage.get("prompt_tokens_details", {}) or {}
+            cache_read_tokens = details.get("cached_tokens", 0)
+            comp_details = usage.get("completion_tokens_details", {}) or {}
+            reasoning_tokens = comp_details.get("reasoning_tokens", 0)
+
+        return self.track(
+            customer_id=customer_id,
+            model_id=model,
+            provider="openai",
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cache_read_tokens=cache_read_tokens,
+            reasoning_tokens=reasoning_tokens,
+            request_id=request_id,
+            span_id=span_id,
+            session_id=session_id,
+            latency_ms=latency_ms,
+            environment=environment,
+        )
+
+    def track_anthropic(
+        self,
+        customer_id: str,
+        response,
+        *,
+        session_id: Optional[str] = None,
+        span_id: Optional[str] = None,
+        latency_ms: int = 0,
+        environment: Optional[str] = None,
+    ) -> TokenEvent:
+        """Track usage from an Anthropic Message response object.
+
+        Args:
+            customer_id: Your customer/tenant identifier.
+            response: Anthropic Message response (or dict).
+            session_id: Optional conversation session ID.
+            span_id: Optional observability span ID.
+            latency_ms: Request latency in ms.
+            environment: Override instance-level environment.
+
+        Returns:
+            The TokenEvent that was sent.
+        """
+        if hasattr(response, "model"):
+            model = response.model
+            usage = response.usage
+            request_id = response.id
+        else:
+            model = response["model"]
+            usage = response["usage"]
+            request_id = response.get("id")
+
+        if hasattr(usage, "input_tokens"):
+            input_tokens = usage.input_tokens or 0
+            output_tokens = usage.output_tokens or 0
+            cache_read_tokens = getattr(usage, "cache_read_input_tokens", 0) or 0
+            cache_write_tokens = getattr(usage, "cache_creation_input_tokens", 0) or 0
+        else:
+            input_tokens = usage.get("input_tokens", 0)
+            output_tokens = usage.get("output_tokens", 0)
+            cache_read_tokens = usage.get("cache_read_input_tokens", 0)
+            cache_write_tokens = usage.get("cache_creation_input_tokens", 0)
+
+        return self.track(
+            customer_id=customer_id,
+            model_id=model,
+            provider="anthropic",
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cache_read_tokens=cache_read_tokens,
+            cache_write_tokens=cache_write_tokens,
+            request_id=request_id,
+            span_id=span_id,
+            session_id=session_id,
+            latency_ms=latency_ms,
+            environment=environment,
+        )
+
+    def wrap_stream(
+        self,
+        stream,
+        customer_id: str,
+        model_id: str,
+        *,
+        provider: str = "openai",
+        input_tokens: int = 0,
+        heartbeat_interval_sec: float = 2.0,
+        parent_span_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        environment: Optional[str] = None,
+    ) -> StreamingWrapper:
+        """Wrap a streaming LLM response for near-real-time usage tracking.
+
+        Emits heartbeat events every heartbeat_interval_sec during the stream,
+        then a final accurate event when the stream completes.
+
+        Usage:
+            stream = client.chat.completions.create(..., stream=True)
+            for chunk in meter.wrap_stream(stream, "cust_1", "gpt-4o"):
+                process(chunk)
+            # Final event emitted automatically
+        """
+        return StreamingWrapper(
+            stream=stream,
+            meter=self,
+            customer_id=customer_id,
+            model_id=model_id,
+            provider=provider,
+            input_tokens=input_tokens,
+            heartbeat_interval_sec=heartbeat_interval_sec,
+            parent_span_id=parent_span_id,
+            session_id=session_id,
+            environment=environment or self._environment,
+        )
+
+    def _send(self, event: TokenEvent) -> None:
+        """Persist event to WAL, then send to Kafka. Zero data loss."""
+        event_dict = event.to_dict()
+
+        if self._wal:
+            self._wal.append(event_dict)
+            return  # WAL flush thread is the sole Kafka sender (no duplicate replay)
+
+        try:
+            value = json.dumps(event_dict, separators=(",", ":")).encode("utf-8")
+            self._producer.produce(
+                topic=self._topic,
+                key=event.customer_id.encode("utf-8"),
+                value=value,
+                on_delivery=self._on_delivery,
+            )
+            self._events_sent += 1
+            self._producer.poll(0)
+        except (BufferError, Exception) as e:
+            self._delivery_errors += 1
+            logger.debug("Kafka send failed: %s", e)
+
+    def _produce_event(self, evt: dict) -> bool:
+        """Send one event to Kafka and wait for broker ack. Returns False on failure."""
+        value = json.dumps(evt, separators=(",", ":")).encode("utf-8")
+        customer_id = evt.get("customerId", "unknown")
+        for _ in range(2):
+            try:
+                self._producer.produce(
+                    topic=self._topic,
+                    key=customer_id.encode("utf-8"),
+                    value=value,
+                    on_delivery=self._on_delivery,
+                )
+                self._producer.flush(timeout=10)
+                self._events_sent += 1
+                return True
+            except BufferError:
+                self._producer.flush(timeout=10)
+            except Exception as e:
+                self._delivery_errors += 1
+                logger.debug("Kafka send failed: %s", e)
+                return False
+        self._delivery_errors += 1
+        return False
+
+    def _flush_wal_once(self) -> bool:
+        """Send at most one pending WAL event across all files. Returns True if one was sent."""
+        if not self._wal:
+            return False
+        for f in self._wal.pending_files():
+            offset = self._wal.get_send_offset(f)
+            evt, new_offset = self._wal.read_next_event_from_offset(f, offset)
+            if evt is None:
+                if f != self._wal._current_file and self._wal.is_fully_sent(f):
+                    self._wal.mark_flushed(f, 0)
+                continue
+            if not self._produce_event(evt):
+                return False
+            self._wal.advance_send_offset(f, new_offset)
+            if f != self._wal._current_file and self._wal.is_fully_sent(f):
+                self._wal.mark_flushed(f, 1)
+            return True
+        return False
+
+    def _wal_flush_loop(self) -> None:
+        """Background thread: sends pending WAL events to Kafka one at a time."""
+        while True:
+            time.sleep(1)
+            if not self._wal:
+                break
+            try:
+                while self._flush_wal_once():
+                    pass
+            except Exception as e:
+                logger.debug("WAL flush error: %s", e)
+
+    def _on_delivery(self, err, msg):
+        if err:
+            self._delivery_errors += 1
+            logger.debug("FluxMeter delivery failed: %s", err)
+
+    def flush(self, timeout: float = 10.0) -> None:
+        """Flush pending events. Drains WAL before closing."""
+        if self._wal:
+            deadline = time.time() + timeout
+            while time.time() < deadline and self._flush_wal_once():
+                pass
+        self._producer.flush(timeout=timeout)
+        if self._wal:
+            self._wal.close()
+
+    @property
+    def events_sent(self) -> int:
+        """Total events sent (including buffered)."""
+        return self._events_sent
+
+    @property
+    def delivery_errors(self) -> int:
+        """Total delivery failures."""
+        return self._delivery_errors

fluxmeter-1.0.0/fluxmeter/event.py

@@ -0,0 +1,85 @@
+"""Token usage event model."""
+
+from __future__ import annotations
+
+import time
+import uuid
+from dataclasses import dataclass, field, asdict
+from typing import Optional
+
+
+@dataclass
+class TokenEvent:
+    """Represents one LLM API call's token usage.
+
+    Supports OpenAI, Anthropic, Google, and custom providers.
+    All token fields are optional — set what's available from your provider response.
+    """
+
+    customer_id: str
+    model_id: str
+    provider: str = "openai"
+
+    # Token counts
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_read_tokens: int = 0
+    cache_write_tokens: int = 0
+    reasoning_tokens: int = 0
+    embedding_tokens: int = 0
+
+    # Identity & tracing
+    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    request_id: Optional[str] = None
+    span_id: Optional[str] = None
+    parent_span_id: Optional[str] = None  # Links child LLM calls to parent agent run
+    session_id: Optional[str] = None
+
+    # Timing
+    timestamp: int = field(default_factory=lambda: int(time.time() * 1000))
+    latency_ms: int = 0
+
+    # Context
+    environment: Optional[str] = None
+    metadata: Optional[dict[str, str]] = None
+
+    def to_dict(self) -> dict:
+        """Serialize to dict with camelCase keys (matches Java consumer)."""
+        d = {
+            "eventId": self.event_id,
+            "customerId": self.customer_id,
+            "provider": self.provider,
+            "modelId": self.model_id,
+            "inputTokens": self.input_tokens,
+            "outputTokens": self.output_tokens,
+            "cacheReadTokens": self.cache_read_tokens,
+            "cacheWriteTokens": self.cache_write_tokens,
+            "reasoningTokens": self.reasoning_tokens,
+            "embeddingTokens": self.embedding_tokens,
+            "timestamp": self.timestamp,
+            "latencyMs": self.latency_ms,
+        }
+        if self.request_id:
+            d["requestId"] = self.request_id
+        if self.span_id:
+            d["spanId"] = self.span_id
+        if self.parent_span_id:
+            d["parentSpanId"] = self.parent_span_id
+        if self.session_id:
+            d["sessionId"] = self.session_id
+        if self.environment:
+            d["environment"] = self.environment
+        if self.metadata:
+            d["metadata"] = self.metadata
+        return d
+
+    @property
+    def total_tokens(self) -> int:
+        return (
+            self.input_tokens
+            + self.output_tokens
+            + self.cache_read_tokens
+            + self.cache_write_tokens
+            + self.reasoning_tokens
+            + self.embedding_tokens
+        )

fluxmeter-1.0.0/fluxmeter/streaming.py

@@ -0,0 +1,157 @@
+"""Streaming response wrapper for FluxMeter.
+
+Wraps OpenAI/Anthropic streaming responses to emit partial usage events
+during the stream (heartbeat every N chunks or every interval_sec).
+Provides near-real-time visibility into long-running LLM calls.
+
+Usage:
+    from fluxmeter import FluxMeter
+
+    meter = FluxMeter(kafka_brokers="localhost:9094")
+
+    # OpenAI streaming
+    stream = client.chat.completions.create(model="gpt-4o", messages=[...], stream=True)
+    for chunk in meter.wrap_stream(stream, customer_id="cust_1", model_id="gpt-4o"):
+        process(chunk)
+    # Final usage event emitted automatically on stream end
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Iterator, Optional, Any
+
+from fluxmeter.event import TokenEvent
+
+
+class StreamingWrapper:
+    """Wraps a streaming LLM response iterator with usage tracking.
+
+    Counts output tokens (approximated from chunks) and emits partial
+    usage events at regular intervals. Emits a final event on stream end.
+    """
+
+    def __init__(
+        self,
+        stream: Iterator[Any],
+        meter,  # FluxMeter instance
+        customer_id: str,
+        model_id: str,
+        provider: str = "openai",
+        input_tokens: int = 0,
+        heartbeat_interval_sec: float = 2.0,
+        parent_span_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        environment: Optional[str] = None,
+    ):
+        self._stream = stream
+        self._meter = meter
+        self._customer_id = customer_id
+        self._model_id = model_id
+        self._provider = provider
+        self._input_tokens = input_tokens
+        self._heartbeat_interval = heartbeat_interval_sec
+        self._parent_span_id = parent_span_id
+        self._session_id = session_id
+        self._environment = environment
+
+        self._output_chunks = 0
+        self._estimated_output_tokens = 0
+        self._last_emitted_output_tokens = 0
+        self._last_heartbeat = time.time()
+        self._start_time = time.time()
+        self._finished = False
+        self._request_id: Optional[str] = None
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        try:
+            chunk = next(self._stream)
+            self._process_chunk(chunk)
+
+            # Emit heartbeat if interval elapsed
+            now = time.time()
+            if now - self._last_heartbeat >= self._heartbeat_interval:
+                self._emit_heartbeat()
+                self._last_heartbeat = now
+
+            return chunk
+        except StopIteration:
+            self._emit_final()
+            raise
+
+    def _process_chunk(self, chunk) -> None:
+        """Extract token info from a streaming chunk."""
+        self._output_chunks += 1
+
+        # OpenAI: chunk.choices[0].delta.content
+        if hasattr(chunk, "choices") and chunk.choices:
+            delta = getattr(chunk.choices[0], "delta", None)
+            if delta and getattr(delta, "content", None):
+                # Approximate: ~0.75 tokens per character for English
+                self._estimated_output_tokens += max(1, len(delta.content) // 4)
+            if not self._request_id and hasattr(chunk, "id"):
+                self._request_id = chunk.id
+
+        # Anthropic: chunk.type == "content_block_delta", chunk.delta.text
+        elif hasattr(chunk, "type") and chunk.type == "content_block_delta":
+            text = getattr(getattr(chunk, "delta", None), "text", "")
+            if text:
+                self._estimated_output_tokens += max(1, len(text) // 4)
+
+        # OpenAI final chunk with usage
+        if hasattr(chunk, "usage") and chunk.usage:
+            usage = chunk.usage
+            if hasattr(usage, "completion_tokens") and usage.completion_tokens:
+                self._estimated_output_tokens = usage.completion_tokens
+            if hasattr(usage, "prompt_tokens") and usage.prompt_tokens:
+                self._input_tokens = usage.prompt_tokens
+
+    def _emit_heartbeat(self) -> None:
+        """Emit a partial usage event (heartbeat) during streaming."""
+        delta = self._estimated_output_tokens - self._last_emitted_output_tokens
+        if delta <= 0:
+            return
+        self._last_emitted_output_tokens = self._estimated_output_tokens
+        self._meter.track(
+            customer_id=self._customer_id,
+            model_id=self._model_id,
+            provider=self._provider,
+            input_tokens=0,  # Only count input once in final event
+            output_tokens=delta,
+            parent_span_id=self._parent_span_id,
+            session_id=self._session_id,
+            environment=self._environment,
+            metadata={"_heartbeat": "true", "_chunks": str(self._output_chunks)},
+        )
+
+    def _emit_final(self) -> None:
+        """Emit the final usage event with accurate totals on stream end."""
+        if self._finished:
+            return
+        self._finished = True
+
+        latency_ms = int((time.time() - self._start_time) * 1000)
+        self._meter.track(
+            customer_id=self._customer_id,
+            model_id=self._model_id,
+            provider=self._provider,
+            input_tokens=self._input_tokens,
+            output_tokens=self._estimated_output_tokens,
+            request_id=self._request_id,
+            parent_span_id=self._parent_span_id,
+            session_id=self._session_id,
+            latency_ms=latency_ms,
+            environment=self._environment,
+            metadata={"_stream_chunks": str(self._output_chunks)},
+        )
+
+    @property
+    def estimated_output_tokens(self) -> int:
+        return self._estimated_output_tokens
+
+    @property
+    def elapsed_ms(self) -> int:
+        return int((time.time() - self._start_time) * 1000)

fluxmeter-1.0.0/fluxmeter/wal.py

@@ -0,0 +1,174 @@
+"""Write-Ahead Log for FluxMeter SDK.
+
+Events are persisted to a local append-only file BEFORE sending to Kafka.
+If Kafka is unavailable, events accumulate on disk and flush when it recovers.
+This guarantees no event loss regardless of Kafka availability.
+
+File format: one JSON object per line (newline-delimited JSON / NDJSON).
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+import threading
+import time
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+class WriteAheadLog:
+    """Append-only local event buffer with background flush to Kafka."""
+
+    def __init__(
+        self,
+        path: str = "~/.fluxmeter/wal",
+        max_file_size_mb: int = 100,
+        flush_interval_sec: float = 1.0,
+    ):
+        self._dir = Path(os.path.expanduser(path))
+        self._dir.mkdir(parents=True, exist_ok=True)
+        self._max_file_size = max_file_size_mb * 1024 * 1024
+        self._flush_interval = flush_interval_sec
+
+        self._current_file: Optional[Path] = None
+        self._file_handle = None
+        self._lock = threading.Lock()
+        self._pending_count = 0
+        self._flushed_count = 0
+        # Byte offset successfully sent to Kafka per file (avoids duplicate replay)
+        self._send_offsets: dict[str, int] = {}
+
+        self._rotate_if_needed()
+
+    def append(self, event_dict: dict) -> None:
+        """Append event to WAL. Returns immediately. Thread-safe.
+        Batch fsync: every 100 events or 500ms, whichever comes first."""
+        line = json.dumps(event_dict, separators=(",", ":")) + "\n"
+        with self._lock:
+            self._rotate_if_needed()
+            self._file_handle.write(line)
+            self._file_handle.flush()
+            self._pending_count += 1
+            if self._pending_count % 100 == 0:
+                os.fsync(self._file_handle.fileno())
+
+    def pending_files(self) -> list[Path]:
+        """List WAL files that may have unsent events (oldest first)."""
+        files = sorted(self._dir.glob("wal-*.jsonl"))
+        return files
+
+    def get_send_offset(self, file_path: Path) -> int:
+        """Return byte offset of last successfully sent event in this file."""
+        return self._send_offsets.get(str(file_path), 0)
+
+    def advance_send_offset(self, file_path: Path, new_offset: int) -> None:
+        """Record how many bytes have been successfully sent from file_path."""
+        with self._lock:
+            self._send_offsets[str(file_path)] = new_offset
+
+    def read_next_event_from_offset(
+        self, file_path: Path, byte_offset: int
+    ) -> tuple[dict | None, int]:
+        """Read at most one event from byte_offset. Returns (event, new_offset)."""
+        try:
+            with open(file_path, "r") as f:
+                f.seek(byte_offset)
+                line = f.readline()
+                if not line:
+                    return None, byte_offset
+                stripped = line.strip()
+                if not stripped:
+                    return None, f.tell()
+                try:
+                    return json.loads(stripped), f.tell()
+                except json.JSONDecodeError:
+                    return None, f.tell()
+        except FileNotFoundError:
+            return None, byte_offset
+
+    def read_events_from_offset(self, file_path: Path, byte_offset: int) -> tuple[list[dict], int]:
+        """Read events starting at byte_offset. Returns (events, new_byte_offset)."""
+        events: list[dict] = []
+        new_offset = byte_offset
+        try:
+            with open(file_path, "r") as f:
+                f.seek(byte_offset)
+                while True:
+                    line_start = f.tell()
+                    line = f.readline()
+                    if not line:
+                        break
+                    stripped = line.strip()
+                    if stripped:
+                        try:
+                            events.append(json.loads(stripped))
+                            new_offset = f.tell()
+                        except json.JSONDecodeError:
+                            new_offset = f.tell()
+                            continue
+                    else:
+                        new_offset = line_start + len(line)
+        except FileNotFoundError:
+            pass
+        return events, new_offset
+
+    def is_fully_sent(self, file_path: Path) -> bool:
+        """True if all bytes in file have been sent to Kafka."""
+        try:
+            size = file_path.stat().st_size
+        except FileNotFoundError:
+            return True
+        return self.get_send_offset(file_path) >= size and size > 0
+
+    def mark_flushed(self, file_path: Path, count: int) -> None:
+        """Mark a WAL file as fully flushed to Kafka. Deletes it."""
+        with self._lock:
+            if file_path == self._current_file:
+                return
+            key = str(file_path)
+            self._send_offsets.pop(key, None)
+            try:
+                file_path.unlink()
+                self._flushed_count += count
+            except FileNotFoundError:
+                pass
+
+    def read_events(self, file_path: Path) -> list[dict]:
+        """Read all events from a WAL file."""
+        events, _ = self.read_events_from_offset(file_path, 0)
+        return events
+
+    def _rotate_if_needed(self) -> None:
+        """Create a new WAL file if current is too large or doesn't exist."""
+        if self._file_handle and self._current_file:
+            try:
+                size = self._current_file.stat().st_size
+                if size < self._max_file_size:
+                    return
+            except FileNotFoundError:
+                pass
+
+        if self._file_handle:
+            self._file_handle.close()
+
+        ts = int(time.time() * 1000)
+        self._current_file = self._dir / f"wal-{ts}.jsonl"
+        self._file_handle = open(self._current_file, "a")
+
+    @property
+    def pending_count(self) -> int:
+        return self._pending_count
+
+    @property
+    def flushed_count(self) -> int:
+        return self._flushed_count
+
+    def close(self) -> None:
+        with self._lock:
+            if self._file_handle:
+                self._file_handle.close()
+                self._file_handle = None

fluxmeter-1.0.0/pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "fluxmeter"
+version = "1.0.0"
+description = "Python SDK for FluxMeter — streaming metering for AI token billing"
+readme = "README.md"
+license = "Apache-2.0"
+requires-python = ">=3.9"
+authors = [{ name = "FluxMeter", email = "hello@fluxmeter.dev" }]
+keywords = ["ai", "llm", "metering", "billing", "tokens", "openai", "anthropic", "streaming"]
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3",
+    "Topic :: Software Development :: Libraries",
+]
+dependencies = [
+    "confluent-kafka>=2.3.0",
+]
+
+[project.optional-dependencies]
+openai = ["openai>=1.0"]
+anthropic = ["anthropic>=0.20"]
+dev = ["pytest", "pytest-asyncio", "ruff"]
+
+[project.urls]
+Homepage = "https://github.com/10kshuaizhang/fluxmeter"
+Repository = "https://github.com/10kshuaizhang/fluxmeter"
+
+[tool.hatch.build.targets.wheel]
+packages = ["fluxmeter"]
+
+[tool.ruff]
+target-version = "py39"
+line-length = 100

fluxmeter-1.0.0/tests/test_client.py

@@ -0,0 +1,69 @@
+"""Tests for FluxMeter client (provider response parsing)."""
+
+from unittest.mock import patch, MagicMock
+from fluxmeter.client import FluxMeter
+
+
+def _mock_meter():
+    """Create a FluxMeter with mocked Kafka producer."""
+    with patch("fluxmeter.client.Producer") as mock_producer_cls:
+        mock_producer = MagicMock()
+        mock_producer_cls.return_value = mock_producer
+        meter = FluxMeter(kafka_brokers="localhost:9094", wal_enabled=False)
+        return meter, mock_producer
+
+
+def test_track_basic():
+    meter, producer = _mock_meter()
+    event = meter.track("cust_1", "gpt-4o", input_tokens=100, output_tokens=50)
+    assert event.customer_id == "cust_1"
+    assert event.model_id == "gpt-4o"
+    assert event.input_tokens == 100
+    assert event.output_tokens == 50
+    assert producer.produce.called
+
+
+def test_track_openai_dict_response():
+    meter, producer = _mock_meter()
+    response = {
+        "id": "chatcmpl-abc123",
+        "model": "gpt-4o-2024-08-06",
+        "usage": {
+            "prompt_tokens": 1200,
+            "completion_tokens": 350,
+            "prompt_tokens_details": {"cached_tokens": 200},
+            "completion_tokens_details": {"reasoning_tokens": 0},
+        },
+    }
+    event = meter.track_openai("cust_42", response, latency_ms=1200)
+    assert event.customer_id == "cust_42"
+    assert event.model_id == "gpt-4o-2024-08-06"
+    assert event.provider == "openai"
+    assert event.input_tokens == 1200
+    assert event.output_tokens == 350
+    assert event.cache_read_tokens == 200
+    assert event.request_id == "chatcmpl-abc123"
+    assert event.latency_ms == 1200
+
+
+def test_track_anthropic_dict_response():
+    meter, producer = _mock_meter()
+    response = {
+        "id": "msg_abc123",
+        "model": "claude-sonnet-4-20250514",
+        "usage": {
+            "input_tokens": 800,
+            "output_tokens": 200,
+            "cache_read_input_tokens": 150,
+            "cache_creation_input_tokens": 50,
+        },
+    }
+    event = meter.track_anthropic("cust_99", response)
+    assert event.customer_id == "cust_99"
+    assert event.model_id == "claude-sonnet-4-20250514"
+    assert event.provider == "anthropic"
+    assert event.input_tokens == 800
+    assert event.output_tokens == 200
+    assert event.cache_read_tokens == 150
+    assert event.cache_write_tokens == 50
+    assert event.request_id == "msg_abc123"

fluxmeter-1.0.0/tests/test_event.py

@@ -0,0 +1,61 @@
+"""Tests for TokenEvent serialization."""
+
+from fluxmeter.event import TokenEvent
+
+
+def test_to_dict_camel_case():
+    event = TokenEvent(
+        customer_id="cust_1",
+        model_id="gpt-4o",
+        provider="openai",
+        input_tokens=100,
+        output_tokens=50,
+    )
+    d = event.to_dict()
+    assert d["customerId"] == "cust_1"
+    assert d["modelId"] == "gpt-4o"
+    assert d["provider"] == "openai"
+    assert d["inputTokens"] == 100
+    assert d["outputTokens"] == 50
+    assert "eventId" in d
+    assert "timestamp" in d
+
+
+def test_total_tokens():
+    event = TokenEvent(
+        customer_id="cust_1",
+        model_id="o1",
+        input_tokens=1000,
+        output_tokens=500,
+        reasoning_tokens=3000,
+        cache_read_tokens=200,
+    )
+    assert event.total_tokens == 4700
+
+
+def test_optional_fields_excluded():
+    event = TokenEvent(customer_id="cust_1", model_id="gpt-4o")
+    d = event.to_dict()
+    assert "requestId" not in d
+    assert "spanId" not in d
+    assert "sessionId" not in d
+    assert "environment" not in d
+    assert "metadata" not in d
+
+
+def test_optional_fields_included():
+    event = TokenEvent(
+        customer_id="cust_1",
+        model_id="gpt-4o",
+        request_id="chatcmpl-abc",
+        span_id="span_123",
+        session_id="sess_456",
+        environment="production",
+        metadata={"feature": "chat"},
+    )
+    d = event.to_dict()
+    assert d["requestId"] == "chatcmpl-abc"
+    assert d["spanId"] == "span_123"
+    assert d["sessionId"] == "sess_456"
+    assert d["environment"] == "production"
+    assert d["metadata"] == {"feature": "chat"}

fluxmeter-1.0.0/tests/test_wal.py

@@ -0,0 +1,71 @@
+"""Tests for WAL duplicate-send prevention and flush semantics."""
+
+from unittest.mock import MagicMock, patch
+import tempfile
+
+from fluxmeter.client import FluxMeter
+from fluxmeter.wal import WriteAheadLog
+
+
+def test_wal_single_send_path():
+    """Reading from offset after advance returns no duplicate events."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        wal = WriteAheadLog(path=tmpdir, flush_interval_sec=0.1)
+        wal.append({"eventId": "evt-1", "customerId": "c1", "modelId": "gpt-4o"})
+        wal.append({"eventId": "evt-2", "customerId": "c1", "modelId": "gpt-4o"})
+
+        f = wal.pending_files()[0]
+        evt1, off1 = wal.read_next_event_from_offset(f, 0)
+        assert evt1["eventId"] == "evt-1"
+        wal.advance_send_offset(f, off1)
+        evt2, off2 = wal.read_next_event_from_offset(f, off1)
+        assert evt2["eventId"] == "evt-2"
+        wal.advance_send_offset(f, off2)
+        evt3, _ = wal.read_next_event_from_offset(f, off2)
+        assert evt3 is None
+
+
+def test_wal_enabled_no_immediate_kafka():
+    """With WAL on, _send does not call produce directly."""
+    with patch("fluxmeter.client.Producer") as mock_cls:
+        mock_producer = MagicMock()
+        mock_cls.return_value = mock_producer
+        with tempfile.TemporaryDirectory() as tmpdir:
+            meter = FluxMeter(
+                kafka_brokers="localhost:9094",
+                wal_enabled=True,
+                wal_path=tmpdir,
+            )
+            meter.track("cust_1", "gpt-4o", input_tokens=10, output_tokens=5)
+            assert not mock_producer.produce.called
+
+
+def test_flush_drains_wal_before_close():
+    """flush() sends WAL events synchronously before closing."""
+    with patch("fluxmeter.client.Producer") as mock_cls:
+        mock_producer = MagicMock()
+        mock_cls.return_value = mock_producer
+        with tempfile.TemporaryDirectory() as tmpdir:
+            meter = FluxMeter(
+                kafka_brokers="localhost:9094",
+                wal_enabled=True,
+                wal_path=tmpdir,
+            )
+            meter.track("cust_1", "gpt-4o", input_tokens=1, output_tokens=1)
+            meter.flush(timeout=5.0)
+            assert mock_producer.produce.call_count == 1
+            mock_producer.flush.assert_called()
+
+
+def test_partial_send_advances_one_event_at_a_time():
+    """First event ack advances offset; second event remains for retry."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        wal = WriteAheadLog(path=tmpdir)
+        f = wal._current_file
+        wal.append({"eventId": "a", "customerId": "c1"})
+        wal.append({"eventId": "b", "customerId": "c1"})
+
+        _, off1 = wal.read_next_event_from_offset(f, 0)
+        wal.advance_send_offset(f, off1)
+        evt_b, _ = wal.read_next_event_from_offset(f, off1)
+        assert evt_b["eventId"] == "b"