secondshift 0.1.0b12__py3-none-any.whl

secondshift/gateway/metering.py

@@ -0,0 +1,171 @@
+"""Token usage read from provider responses, normalized across providers."""
+
+import json
+from dataclasses import asdict, dataclass, field
+from typing import Any
+
+from secondshift.config.constants import ModelFamily
+from secondshift.gateway.dao import aadd_used_tokens, acreate_model_usage
+from secondshift.gateway.models import RunToken
+from secondshift.gateway.pricing import request_cost
+from secondshift.runner.agent_image.outcomes import budget_count
+
+MAX_SSE_LINE_BYTES = 1024 * 1024
+
+
+@dataclass
+class Usage:
+    input_tokens: int = 0
+    output_tokens: int = 0
+    cache_read_tokens: int = 0
+    cache_write_tokens: int = 0
+
+    @property
+    def total(self) -> int:
+        return (
+            self.input_tokens
+            + self.output_tokens
+            + self.cache_read_tokens
+            + self.cache_write_tokens
+        )
+
+    @property
+    def budget_tokens(self) -> int:
+        """What counts toward the run's token budget."""
+        return budget_count(
+            self.input_tokens + self.output_tokens + self.cache_write_tokens,
+            self.cache_read_tokens,
+        )
+
+
+def token_count(value: Any) -> int | None:
+    return value if isinstance(value, int) and value >= 0 else None
+
+
+def apply_anthropic_usage(usage: Usage, data: Any) -> None:
+    # message_delta repeats cumulative counts, so later values replace earlier.
+    if not isinstance(data, dict):
+        return
+    fields = {
+        "input_tokens": "input_tokens",
+        "output_tokens": "output_tokens",
+        "cache_read_input_tokens": "cache_read_tokens",
+        "cache_creation_input_tokens": "cache_write_tokens",
+    }
+    for source, target in fields.items():
+        if (value := token_count(data.get(source))) is not None:
+            setattr(usage, target, value)
+
+
+def apply_openai_usage(usage: Usage, data: Any) -> None:
+    if not isinstance(data, dict):
+        return
+    prompt = token_count(data.get("prompt_tokens"))
+    if prompt is None:
+        prompt = token_count(data.get("input_tokens"))
+    completion = token_count(data.get("completion_tokens"))
+    if completion is None:
+        completion = token_count(data.get("output_tokens"))
+    details = data.get("prompt_tokens_details") or data.get("input_tokens_details")
+    cached = (
+        token_count(details.get("cached_tokens")) if isinstance(details, dict) else None
+    )
+    cached = cached or 0
+    # OpenAI counts cached tokens inside the prompt. Store them separately so
+    # totals match Anthropic's accounting.
+    if prompt is not None:
+        usage.input_tokens = max(prompt - cached, 0)
+        usage.cache_read_tokens = cached
+    if completion is not None:
+        usage.output_tokens = completion
+
+
+def usage_from_json(family: ModelFamily, body: Any) -> Usage:
+    usage = Usage()
+    if not isinstance(body, dict):
+        return usage
+    if family == ModelFamily.ANTHROPIC:
+        apply_anthropic_usage(usage, body.get("usage"))
+    else:
+        apply_openai_usage(usage, body.get("usage"))
+    return usage
+
+
+@dataclass
+class StreamMeter:
+    """Reads usage out of a server-sent event stream without altering it."""
+
+    family: ModelFamily
+    usage: Usage = field(default_factory=Usage)
+    # The partial line waiting for its newline. A bytearray grows in place, so
+    # a long line read in small chunks costs time linear in its length.
+    buffer: bytearray = field(default_factory=bytearray, init=False, repr=False)
+
+    def feed(self, chunk: bytes) -> None:
+        self.buffer += chunk
+        # Only the new chunk can hold the newline that ends the partial line.
+        if b"\n" in chunk:
+            *lines, rest = self.buffer.split(b"\n")
+            self.buffer = rest
+            for line in lines:
+                self.read_line(bytes(line).strip())
+        if len(self.buffer) > MAX_SSE_LINE_BYTES:
+            self.buffer.clear()
+
+    def read_line(self, line: bytes) -> None:
+        event = sse_event(line)
+        if event is None:
+            return
+        if self.family == ModelFamily.ANTHROPIC:
+            apply_anthropic_event(self.usage, event)
+        else:
+            apply_openai_event(self.usage, event)
+
+
+def sse_event(line: bytes) -> dict[str, Any] | None:
+    """The JSON object in a server-sent `data:` line, or None for anything else."""
+    if not line.startswith(b"data:"):
+        return None
+    data = line[5:].strip()
+    if not data or data == b"[DONE]":
+        return None
+    try:
+        event = json.loads(data)
+    except ValueError:
+        return None
+    return event if isinstance(event, dict) else None
+
+
+def apply_anthropic_event(usage: Usage, event: dict[str, Any]) -> None:
+    """Usage arrives at the message's start and in its final delta."""
+    if event.get("type") == "message_start":
+        message = event.get("message")
+        if isinstance(message, dict):
+            apply_anthropic_usage(usage, message.get("usage"))
+    elif event.get("type") == "message_delta":
+        apply_anthropic_usage(usage, event.get("usage"))
+
+
+def apply_openai_event(usage: Usage, event: dict[str, Any]) -> None:
+    """Responses report usage when complete. Chat completions in a last chunk."""
+    if event.get("type") == "response.completed":
+        response = event.get("response")
+        if isinstance(response, dict):
+            apply_openai_usage(usage, response.get("usage"))
+    elif event.get("usage"):
+        apply_openai_usage(usage, event["usage"])
+
+
+async def record_usage(
+    run_token: RunToken, family: ModelFamily, model: str, status_code: int, usage: Usage
+) -> None:
+    await acreate_model_usage(
+        run_token,
+        family,
+        model[:200],
+        status_code,
+        asdict(usage),
+        request_cost(run_token.connection, model, usage),
+    )
+    if usage.total:
+        await aadd_used_tokens(run_token, usage.budget_tokens)

secondshift/gateway/migrations/0001_initial.py

@@ -0,0 +1,54 @@
+# Generated by Django 6.0.9 on 2026-10-26 14:22
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    initial = True
+
+    dependencies = [
+        ('connections', '0001_initial'),
+        ('runs', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='ModelUsage',
+            fields=[
+                ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('provider', models.CharField(max_length=24)),
+                ('model', models.CharField(max_length=200)),
+                ('status_code', models.PositiveSmallIntegerField()),
+                ('input_tokens', models.PositiveBigIntegerField(default=0)),
+                ('output_tokens', models.PositiveBigIntegerField(default=0)),
+                ('cache_read_tokens', models.PositiveBigIntegerField(default=0)),
+                ('cache_write_tokens', models.PositiveBigIntegerField(default=0)),
+                ('cost_usd', models.DecimalField(blank=True, decimal_places=8, max_digits=14, null=True)),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('connection', models.ForeignKey(null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='model_usage', to='connections.connection')),
+                ('run', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='model_usage', to='runs.run')),
+            ],
+            options={
+                'ordering': ['created_at'],
+            },
+        ),
+        migrations.CreateModel(
+            name='RunToken',
+            fields=[
+                ('id', models.BigAutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')),
+                ('token_hash', models.CharField(max_length=64, unique=True)),
+                ('expires_at', models.DateTimeField()),
+                ('revoked_at', models.DateTimeField(blank=True, null=True)),
+                ('budget_tokens', models.PositiveBigIntegerField()),
+                ('used_tokens', models.PositiveBigIntegerField(default=0)),
+                ('created_at', models.DateTimeField(auto_now_add=True)),
+                ('connection', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='run_tokens', to='connections.connection')),
+                ('run', models.ForeignKey(on_delete=django.db.models.deletion.CASCADE, related_name='tokens', to='runs.run')),
+            ],
+            options={
+                'ordering': ['-created_at'],
+            },
+        ),
+    ]

secondshift/gateway/migrations/0002_modelusage_gateway_usage_created_at_idx.py

@@ -0,0 +1,19 @@
+# Generated by Django 6.0.9 on 2026-11-07 23:22
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('connections', '0008_agentcredentialcleanup'),
+        ('gateway', '0001_initial'),
+        ('runs', '0012_alter_rundelivery_status'),
+    ]
+
+    operations = [
+        migrations.AddIndex(
+            model_name='modelusage',
+            index=models.Index(fields=['created_at'], name='gateway_usage_created_at_idx'),
+        ),
+    ]

secondshift/gateway/models.py

@@ -0,0 +1,61 @@
+from django.db import models
+
+from secondshift.connections.models import Connection
+
+
+class RunToken(models.Model):
+    """A per-run credential for the model gateway. Only its SHA-256 is stored."""
+
+    run = models.ForeignKey("runs.Run", on_delete=models.CASCADE, related_name="tokens")
+    connection = models.ForeignKey(
+        Connection,
+        on_delete=models.CASCADE,
+        related_name="run_tokens",
+    )
+    token_hash = models.CharField(max_length=64, unique=True)
+    expires_at = models.DateTimeField()
+    revoked_at = models.DateTimeField(null=True, blank=True)
+    budget_tokens = models.PositiveBigIntegerField()
+    used_tokens = models.PositiveBigIntegerField(default=0)
+    created_at = models.DateTimeField(auto_now_add=True)
+
+    class Meta:
+        ordering = ["-created_at"]
+
+    def __str__(self) -> str:
+        return f"Run {self.run_id} gateway token"
+
+
+class ModelUsage(models.Model):
+    """One gateway request: what was called and what it cost in tokens."""
+
+    run = models.ForeignKey(
+        "runs.Run", on_delete=models.CASCADE, related_name="model_usage"
+    )
+    connection = models.ForeignKey(
+        Connection,
+        on_delete=models.SET_NULL,
+        related_name="model_usage",
+        null=True,
+    )
+    provider = models.CharField(max_length=24)
+    model = models.CharField(max_length=200)
+    status_code = models.PositiveSmallIntegerField()
+    input_tokens = models.PositiveBigIntegerField(default=0)
+    output_tokens = models.PositiveBigIntegerField(default=0)
+    cache_read_tokens = models.PositiveBigIntegerField(default=0)
+    cache_write_tokens = models.PositiveBigIntegerField(default=0)
+    # USD at list price. Null when the model has no known price.
+    cost_usd = models.DecimalField(
+        max_digits=14, decimal_places=8, null=True, blank=True
+    )
+    created_at = models.DateTimeField(auto_now_add=True)
+
+    class Meta:
+        ordering = ["created_at"]
+        indexes = [
+            models.Index(fields=["created_at"], name="gateway_usage_created_at_idx")
+        ]
+
+    def __str__(self) -> str:
+        return f"Run {self.run_id}: {self.model} ({self.status_code})"

secondshift/gateway/pricing.py

@@ -0,0 +1,39 @@
+"""Request cost from genai-prices' bundled price tables (no network lookups)."""
+
+import logging
+from decimal import Decimal
+from typing import TYPE_CHECKING
+
+import genai_prices
+
+from secondshift.connections.model_providers import base_url, model_provider
+from secondshift.connections.models import Connection
+
+if TYPE_CHECKING:
+    from secondshift.gateway.metering import Usage
+
+logger = logging.getLogger(__name__)
+
+
+def request_cost(connection: Connection, model: str, usage: "Usage") -> Decimal | None:
+    """Price one request in USD, or None when the model has no known price."""
+    if not usage.total:
+        return Decimal(0)
+    # genai-prices counts cache reads and writes inside input_tokens.
+    priced = genai_prices.Usage(
+        input_tokens=usage.input_tokens
+        + usage.cache_read_tokens
+        + usage.cache_write_tokens,
+        output_tokens=usage.output_tokens,
+        cache_read_tokens=usage.cache_read_tokens,
+        cache_write_tokens=usage.cache_write_tokens,
+    )
+    if model_provider(connection.provider).router:
+        lookup: dict[str, str] = {"provider_api_url": base_url(connection)}
+    else:
+        lookup = {"provider_id": connection.provider}
+    try:
+        return genai_prices.calc_price(priced, model, **lookup).total_price
+    except (LookupError, ValueError):
+        logger.info("no price for model %s on connection %s", model, connection.pk)
+        return None