modelab 0.1.0__py3-none-any.whl

modelab/__init__.py

@@ -0,0 +1,126 @@
+"""modelab — provider-agnostic A/B testing for LLM systems."""
+
+from __future__ import annotations
+
+import json
+import logging
+import urllib.request
+from typing import Any, Sequence
+
+from modelab._assignment import Assignment
+from modelab._engine import assign_variant
+from modelab._errors import FlagNotFoundError, NotInitializedError
+from modelab._server_storage import ServerStorage
+from modelab._state import _global_state
+from modelab._types import (
+    AssignmentRecord,
+    EvalContext,
+    Flag,
+    Variant,
+)
+
+__all__ = [
+    "init",
+    "assign",
+    "evaluate",
+    "Flag",
+    "Variant",
+    "EvalContext",
+    "Assignment",
+]
+
+
+def init(
+    server: str,
+    flags: Sequence[Flag] = (),
+    api_key: str = "",
+) -> None:
+    """Initialize modelab with a server URL and flag definitions.
+
+    Args:
+        server: The modelab server URL (e.g. "http://localhost:8100").
+        flags: The experiment flags to register.
+        api_key: Optional API key for server authentication.
+    """
+    storage = ServerStorage(server, api_key=api_key)
+    _global_state.configure(storage, list(flags), server_url=server)
+
+
+def assign(flag_name: str, ctx: EvalContext) -> Assignment | None:
+    """Assign a variant for the given flag and context.
+
+    Returns None if the user is outside the rollout percentage.
+    Raises NotInitializedError if init() hasn't been called.
+    Raises FlagNotFoundError if the flag name isn't registered.
+    """
+    if not _global_state.initialized:
+        raise NotInitializedError()
+
+    flag = _global_state.flags.get(flag_name)
+    if flag is None:
+        raise FlagNotFoundError(flag_name)
+
+    variant = assign_variant(flag, ctx)
+    if variant is None:
+        return None
+
+    storage = _global_state.storage
+    assert storage is not None
+
+    record = AssignmentRecord(
+        flag_name=flag_name,
+        variant_name=variant.name,
+        user_id=ctx.user_id,
+        session_id=ctx.session_id,
+        config_json=dict(variant.config),
+    )
+
+    try:
+        storage.save_assignment(record)
+    except Exception:
+        logging.getLogger("modelab").warning(
+            "Failed to save assignment for %s", flag_name, exc_info=True
+        )
+
+    return Assignment(
+        flag_name=flag_name,
+        variant_name=variant.name,
+        config=dict(variant.config),
+        context=ctx,
+        storage=storage,
+        assignment_id=record.assignment_id,
+    )
+
+
+def evaluate(flag_name: str) -> dict[str, Any]:
+    """Fetch per-variant metrics for a flag from the server.
+
+    Flushes any buffered data, then queries the server's
+    GET /api/v1/flags/{flag_name} endpoint.
+
+    Raises NotInitializedError if init() hasn't been called.
+    Raises FlagNotFoundError if the flag name isn't registered.
+    """
+    if not _global_state.initialized:
+        raise NotInitializedError()
+
+    if flag_name not in _global_state.flags:
+        raise FlagNotFoundError(flag_name)
+
+    storage = _global_state.storage
+    assert storage is not None
+    storage.flush()
+
+    url = f"{_global_state.server_url.rstrip('/')}/api/v1/flags/{flag_name}"
+    req = urllib.request.Request(url, method="GET")
+    with urllib.request.urlopen(req, timeout=10) as resp:
+        return json.loads(resp.read())
+
+
+def reset() -> None:
+    """Reset global state to uninitialized.
+
+    Clears all flags, storage, and server URL. Primarily used for testing
+    to ensure clean state between test runs.
+    """
+    _global_state.reset()

modelab/_assignment.py

@@ -0,0 +1,163 @@
+"""Assignment class — the main object returned by modelab.assign()."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from modelab._types import (
+    EvalContext,
+    EventRecord,
+    ExecutionRecord,
+)
+
+if TYPE_CHECKING:
+    from modelab._server_storage import ServerStorage
+
+logger = logging.getLogger("modelab")
+
+
+def _extract_usage_from_response(response: Any) -> tuple[int | None, int | None]:
+    """Duck-type token usage from a provider response object.
+
+    Tries OpenAI-style attrs first (usage.prompt_tokens / completion_tokens),
+    then Anthropic-style (usage.input_tokens / output_tokens).
+    Returns (input_tokens, output_tokens) or (None, None) if not found.
+    """
+    usage = getattr(response, "usage", None)
+    if usage is None:
+        return None, None
+
+    # OpenAI: usage.prompt_tokens / usage.completion_tokens
+    prompt = getattr(usage, "prompt_tokens", None)
+    completion = getattr(usage, "completion_tokens", None)
+    if prompt is not None or completion is not None:
+        return prompt, completion
+
+    # Anthropic: usage.input_tokens / usage.output_tokens
+    inp = getattr(usage, "input_tokens", None)
+    out = getattr(usage, "output_tokens", None)
+    if inp is not None or out is not None:
+        return inp, out
+
+    return None, None
+
+
+class Assignment:
+    """Represents a variant assignment for a specific evaluation context."""
+
+    def __init__(
+        self,
+        flag_name: str,
+        variant_name: str,
+        config: dict[str, Any],
+        context: EvalContext,
+        storage: ServerStorage,
+        assignment_id: str,
+    ) -> None:
+        self._flag_name = flag_name
+        self._variant_name = variant_name
+        self._config = config
+        self._context = context
+        self._storage = storage
+        self._assignment_id = assignment_id
+
+    @property
+    def flag_name(self) -> str:
+        return self._flag_name
+
+    @property
+    def variant_name(self) -> str:
+        return self._variant_name
+
+    @property
+    def config(self) -> dict[str, Any]:
+        return self._config
+
+    @property
+    def context(self) -> EvalContext:
+        return self._context
+
+    @property
+    def assignment_id(self) -> str:
+        return self._assignment_id
+
+    def record(
+        self,
+        response: Any = None,
+        *,
+        latency_ms: float | None = None,
+        input_tokens: int | None = None,
+        output_tokens: int | None = None,
+        cost: float | None = None,
+        error: str | None = None,
+        **metadata: Any,
+    ) -> None:
+        """Record execution metrics, optionally extracting tokens from a provider response.
+
+        If ``response`` is passed, token counts are duck-typed from
+        ``response.usage`` (OpenAI and Anthropic formats).
+        Explicit keyword arguments always override extracted values.
+        """
+        if response is not None:
+            extracted_in, extracted_out = _extract_usage_from_response(response)
+            if input_tokens is None:
+                input_tokens = extracted_in
+            if output_tokens is None:
+                output_tokens = extracted_out
+
+        self._save_execution(
+            ExecutionRecord(
+                assignment_id=self._assignment_id,
+                latency_ms=latency_ms,
+                input_tokens=input_tokens,
+                output_tokens=output_tokens,
+                cost=cost,
+                error=error,
+                metadata_json=metadata,
+            )
+        )
+
+    def mark_success(self, payload: dict[str, Any] | None = None) -> None:
+        """Mark this assignment as a success event.
+
+        Args:
+            payload: Optional additional data to attach to the event.
+        """
+        self._save_event("success", "", payload or {})
+
+    def mark_failure(self, payload: dict[str, Any] | None = None) -> None:
+        """Mark this assignment as a failure event.
+
+        Args:
+            payload: Optional additional data to attach to the event.
+        """
+        self._save_event("failure", "", payload or {})
+
+    def mark_custom_event(self, name: str, payload: dict[str, Any] | None = None) -> None:
+        """Mark a custom event for this assignment.
+
+        Args:
+            name: The name of the custom event (e.g., "copied", "dismissed").
+            payload: Optional additional data to attach to the event.
+        """
+        self._save_event("custom", name, payload or {})
+
+    def _save_execution(self, record: ExecutionRecord) -> None:
+        try:
+            self._storage.save_execution(record)
+        except Exception:
+            logger.warning("Failed to save execution for %s", self._assignment_id, exc_info=True)
+
+    def _save_event(self, event_type: str, event_name: str, payload: dict[str, Any]) -> None:
+        try:
+            self._storage.save_event(
+                EventRecord(
+                    assignment_id=self._assignment_id,
+                    event_type=event_type,
+                    event_name=event_name,
+                    payload_json=payload,
+                )
+            )
+        except Exception:
+            logger.warning("Failed to save event for %s", self._assignment_id, exc_info=True)

modelab/_engine.py

@@ -0,0 +1,40 @@
+"""Deterministic assignment engine: hash → bucket → rollout gate → variant."""
+
+from __future__ import annotations
+
+import hashlib
+
+from modelab._types import EvalContext, Flag, Variant
+
+BUCKET_COUNT = 10_000
+
+
+def _bucket(flag_name: str, user_id: str) -> int:
+    """Deterministic bucket in [0, BUCKET_COUNT) from flag + user."""
+    key = f"{flag_name}:{user_id}"
+    digest = hashlib.md5(key.encode()).hexdigest()
+    return int(digest, 16) % BUCKET_COUNT
+
+
+def assign_variant(flag: Flag, ctx: EvalContext) -> Variant | None:
+    """Return the assigned Variant, or None if outside rollout."""
+    bucket = _bucket(flag.name, ctx.user_id)
+
+    # Rollout gate: rollout_pct of 100 means buckets 0–9999 pass
+    rollout_threshold = int(flag.rollout_pct / 100.0 * BUCKET_COUNT)
+    if bucket >= rollout_threshold:
+        return None
+
+    # Weighted variant selection within the rollout population
+    total_weight = sum(v.weight for v in flag.variants)
+    if total_weight == 0:
+        return flag.variants[0] if flag.variants else None
+    point = bucket % total_weight
+    cumulative = 0
+    for variant in flag.variants:
+        cumulative += variant.weight
+        if point < cumulative:
+            return variant
+
+    # Fallback (shouldn't happen if weights > 0)
+    return flag.variants[-1] if flag.variants else None

modelab/_errors.py

@@ -0,0 +1,26 @@
+"""Exception hierarchy for modelab."""
+
+from __future__ import annotations
+
+
+class ModelabError(Exception):
+    """Base exception for all modelab errors."""
+
+
+class NotInitializedError(ModelabError):
+    """Raised when modelab.assign() is called before modelab.init()."""
+
+    def __init__(self) -> None:
+        super().__init__("modelab.init() must be called before assign()")
+
+
+class FlagNotFoundError(ModelabError):
+    """Raised when a flag name is not in the registry."""
+
+    def __init__(self, name: str) -> None:
+        super().__init__(f"Flag not found: {name!r}")
+        self.name = name
+
+
+class InvalidFlagError(ModelabError):
+    """Raised when a Flag definition is invalid."""

modelab/_server_storage.py

@@ -0,0 +1,102 @@
+"""ServerStorage — HTTP storage backend that buffers and flushes to modelab-server."""
+
+from __future__ import annotations
+
+import atexit
+import json
+import logging
+import threading
+import urllib.request
+from dataclasses import asdict
+from datetime import datetime, timezone
+from typing import Any
+
+from modelab._types import AssignmentRecord, EventRecord, ExecutionRecord
+
+logger = logging.getLogger("modelab")
+
+_FLUSH_SIZE = 50
+_FLUSH_INTERVAL = 5.0  # seconds
+
+
+def _default_serializer(obj: Any) -> str:
+    if isinstance(obj, datetime):
+        return obj.isoformat()
+    raise TypeError(f"Not serializable: {type(obj)}")
+
+
+class ServerStorage:
+    """HTTP storage that buffers records and flushes to the modelab server."""
+
+    def __init__(self, base_url: str, api_key: str = "") -> None:
+        self._base_url = base_url.rstrip("/")
+        self._api_key = api_key
+        self._lock = threading.Lock()
+        self._assignments: list[dict[str, Any]] = []
+        self._executions: list[dict[str, Any]] = []
+        self._events: list[dict[str, Any]] = []
+
+        # Background flush timer
+        self._timer: threading.Timer | None = None
+        self._start_timer()
+        atexit.register(self.flush)
+
+    def _start_timer(self) -> None:
+        self._timer = threading.Timer(_FLUSH_INTERVAL, self._timer_flush)
+        self._timer.daemon = True
+        self._timer.start()
+
+    def _timer_flush(self) -> None:
+        self.flush()
+        self._start_timer()
+
+    def save_assignment(self, record: AssignmentRecord) -> None:
+        with self._lock:
+            self._assignments.append(asdict(record))
+            if len(self._assignments) >= _FLUSH_SIZE:
+                self._flush_locked("assignments", self._assignments)
+                self._assignments = []
+
+    def save_execution(self, record: ExecutionRecord) -> None:
+        with self._lock:
+            self._executions.append(asdict(record))
+            if len(self._executions) >= _FLUSH_SIZE:
+                self._flush_locked("executions", self._executions)
+                self._executions = []
+
+    def save_event(self, record: EventRecord) -> None:
+        with self._lock:
+            self._events.append(asdict(record))
+            if len(self._events) >= _FLUSH_SIZE:
+                self._flush_locked("events", self._events)
+                self._events = []
+
+    def flush(self) -> None:
+        with self._lock:
+            if self._assignments:
+                self._flush_locked("assignments", self._assignments)
+                self._assignments = []
+            if self._executions:
+                self._flush_locked("executions", self._executions)
+                self._executions = []
+            if self._events:
+                self._flush_locked("events", self._events)
+                self._events = []
+
+    def _flush_locked(self, endpoint: str, records: list[dict[str, Any]]) -> None:
+        url = f"{self._base_url}/api/v1/ingest/{endpoint}"
+        data = json.dumps(records, default=_default_serializer).encode()
+        req = urllib.request.Request(
+            url,
+            data=data,
+            headers={
+                "Content-Type": "application/json",
+                **({"X-API-Key": self._api_key} if self._api_key else {}),
+            },
+            method="POST",
+        )
+        try:
+            with urllib.request.urlopen(req, timeout=10) as resp:
+                resp.read()
+        except Exception:
+            logger.warning("Failed to flush %d %s to %s", len(records), endpoint, url, exc_info=True)

modelab/_state.py

@@ -0,0 +1,42 @@
+"""Module-level singleton state for modelab."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from modelab._errors import InvalidFlagError
+from modelab._types import Flag
+
+if TYPE_CHECKING:
+    from modelab._server_storage import ServerStorage
+
+
+class _State:
+    def __init__(self) -> None:
+        self.storage: ServerStorage | None = None
+        self.flags: dict[str, Flag] = {}
+        self.server_url: str = ""
+
+    def configure(self, storage: ServerStorage, flags: list[Flag], server_url: str = "") -> None:
+        for flag in flags:
+            if not flag.variants:
+                raise InvalidFlagError(f"Flag {flag.name!r} has no variants")
+            if not (0 <= flag.rollout_pct <= 100):
+                raise InvalidFlagError(
+                    f"Flag {flag.name!r} rollout_pct must be 0-100, got {flag.rollout_pct}"
+                )
+        self.storage = storage
+        self.flags = {f.name: f for f in flags}
+        self.server_url = server_url
+
+    def reset(self) -> None:
+        self.storage = None
+        self.flags = {}
+        self.server_url = ""
+
+    @property
+    def initialized(self) -> bool:
+        return self.storage is not None
+
+
+_global_state = _State()

modelab/_types.py

@@ -0,0 +1,70 @@
+"""Core types for modelab: Flag, Variant, EvalContext, records, Storage protocol."""
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+
+@dataclass(frozen=True)
+class Variant:
+    """A single variant within a flag."""
+
+    name: str
+    weight: int = 50
+    config: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class Flag:
+    """An experiment flag with one or more variants."""
+
+    name: str
+    variants: list[Variant] = field(default_factory=list)
+    rollout_pct: float = 100.0
+
+
+@dataclass(frozen=True)
+class EvalContext:
+    """Context for assignment — identifies who is being assigned."""
+
+    user_id: str
+    session_id: str = ""
+
+
+# ── Records (persisted to storage) ──────────────────────────────────
+
+
+@dataclass
+class AssignmentRecord:
+    assignment_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    flag_name: str = ""
+    variant_name: str = ""
+    user_id: str = ""
+    session_id: str = ""
+    config_json: dict[str, Any] = field(default_factory=dict)
+    assigned_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+
+@dataclass
+class ExecutionRecord:
+    assignment_id: str = ""
+    latency_ms: float | None = None
+    input_tokens: int | None = None
+    output_tokens: int | None = None
+    cost: float | None = None
+    error: str | None = None
+    metadata_json: dict[str, Any] = field(default_factory=dict)
+    recorded_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+
+@dataclass
+class EventRecord:
+    event_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    assignment_id: str = ""
+    event_type: str = ""  # success / failure / custom
+    event_name: str = ""  # for custom events
+    payload_json: dict[str, Any] = field(default_factory=dict)
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

modelab-0.1.0.dist-info/METADATA

@@ -0,0 +1,163 @@
+Metadata-Version: 2.4
+Name: modelab
+Version: 0.1.0
+Summary: Provider-agnostic A/B testing for LLM systems
+License-Expression: MIT
+Keywords: ab-testing,experiments,feature-flags,llm
+Requires-Python: >=3.10
+Provides-Extra: dev
+Requires-Dist: httpx>=0.27; extra == 'dev'
+Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Provides-Extra: server
+Requires-Dist: fastapi>=0.110; extra == 'server'
+Requires-Dist: psycopg[binary,pool]>=3.1; extra == 'server'
+Requires-Dist: uvicorn[standard]>=0.29; extra == 'server'
+Description-Content-Type: text/markdown
+
+# modelab
+
+Provider-agnostic A/B testing for LLM systems in production.
+
+**Two components:**
+1. **Python SDK** — zero-dependency library for assignment, tracking, and evaluation
+2. **Server + Dashboard** — self-hosted FastAPI + React app for visualization (Docker Compose)
+
+## Quick Start
+
+### SDK (local development)
+
+```bash
+pip install modelab
+```
+
+```python
+import modelab
+from modelab import Flag, Variant, EvalContext
+
+# Initialize — point to the modelab server
+modelab.init(
+    server="http://localhost:8100",
+    flags=[
+        Flag(
+            name="summarizer_v2",
+            variants=[
+                Variant("control", weight=50, config={"model": "gpt-3.5-turbo", "prompt": "Summarize: {input}"}),
+                Variant("treatment", weight=50, config={"model": "gpt-4", "prompt": "Concisely summarize: {input}"}),
+            ],
+            rollout_pct=100,
+        ),
+    ],
+)
+
+# Assign a variant
+ctx = EvalContext(user_id="user_123", session_id="abc")
+assignment = modelab.assign("summarizer_v2", ctx)
+
+if assignment is None:
+    # Outside rollout — use default behavior
+    response = call_llm(model="gpt-3.5-turbo", prompt=text)
+else:
+    # In experiment — use assigned variant config
+    response = call_llm(
+        model=assignment.config["model"],
+        prompt=assignment.config["prompt"].format(input=text),
+    )
+    assignment.record(response, cost=0.013)
+    assignment.mark_success()
+
+# Evaluate results
+results = modelab.evaluate("summarizer_v2")
+print(results)
+```
+
+### Self-Hosted Server + Dashboard
+
+```bash
+docker compose up
+```
+
+This starts:
+- **PostgreSQL** on port 5432
+- **modelab server + dashboard** on port 8100
+
+## Concepts
+
+### Flags
+An experiment with one or more variants and a rollout percentage (0-100%).
+
+### Variants
+Each variant has a name, weight (for traffic splitting), and a config dict you use to parameterize your LLM calls.
+
+### Assignment
+Deterministic — the same `(flag_name, user_id)` always maps to the same variant. Uses MD5 hashing into 10,000 buckets for 0.01% rollout granularity.
+
+### Recording
+
+Use `assignment.record(response)` to capture execution metrics. Token counts are automatically extracted from the response object via duck-typing (supports OpenAI and Anthropic response formats). Cost, latency, error, and arbitrary metadata can be passed as keyword arguments:
+
+```python
+assignment.record(response, cost=0.013, latency_ms=250.0, model="gpt-4o")
+```
+
+You can also record without a response object:
+
+```python
+assignment.record(input_tokens=50, output_tokens=100, cost=0.01)
+```
+
+### Events
+Mark assignments as success/failure or record custom events (e.g., "copied", "thumbs_up").
+
+### Evaluation
+`modelab.evaluate(flag_name)` returns per-variant metrics: success rate, avg latency, avg cost, token usage, and custom event counts.
+
+## Server API
+
+### Ingestion (from SDK)
+```
+POST /api/v1/ingest/assignments    (batch)
+POST /api/v1/ingest/executions     (batch)
+POST /api/v1/ingest/events         (batch)
+```
+
+### Dashboard API
+```
+GET /api/v1/flags                  — list flags with summary stats
+GET /api/v1/flags/{name}           — detailed per-variant evaluation
+GET /api/v1/flags/{name}/timeline  — time-series metrics
+```
+
+## Development
+
+```bash
+# Install in dev mode
+pip install -e ".[dev]"
+
+# Run tests
+pytest
+
+# Run dashboard dev server
+cd dashboard && npm install && npm run dev
+
+# Run API server (requires Postgres)
+uvicorn server.app:app --reload --port 8100
+```
+
+## Architecture
+
+```
+Developer's App
+│
+├── modelab SDK (pip install modelab)
+│   └── ServerStorage ──HTTP POST──▶ modelab-server
+│
+modelab-server (docker compose up)
+├── FastAPI backend
+├── React dashboard (served as static files)
+└── PostgreSQL
+```
+
+## License
+
+MIT

modelab-0.1.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+modelab/__init__.py,sha256=ToaPXEjWZYi8poObLXTHXOkTz8DugUHLyKcZsqQ_jLg,3463
+modelab/_assignment.py,sha256=XXH9ZNMftobh6m2rj62T-pqpjwxw8wKG41JLmXz4Mvo,5211
+modelab/_engine.py,sha256=rPhlSXIRQLvhHWxnLs6YayJ1HJX5iKwZbVFrkmIB49A,1333
+modelab/_errors.py,sha256=y6aqP28CihLrlM2MTxm0v3ISdK-qbkW4aOetOZRh_oM,693
+modelab/_server_storage.py,sha256=pmlIRwIY9mmNKEUemVeF2f41YnJ5RYcr_2YZlmCSksA,3537
+modelab/_state.py,sha256=pERDfuCx5Ylh7YzqTdpbgZ90aTHz9bO5HyqCcdbk0kY,1225
+modelab/_types.py,sha256=3zc_be5PTmlANlLrY1ZhYoX9vaTyWO4PdBd5NK6zh_s,2033
+modelab/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+modelab-0.1.0.dist-info/METADATA,sha256=lPiLzc6Hs0Ni5Qk4V1TvHxSHgVLloG0bteexe3ea5-g,4386
+modelab-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+modelab-0.1.0.dist-info/RECORD,,

modelab-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any