superlinear 0.1.0__py3-none-any.whl

superlinear/engine/chat_types.py

@@ -0,0 +1,130 @@
+"""Core chat request/streaming event types.
+
+These types are internal to the library and are intentionally decoupled from:
+- HTTP transport (FastAPI / SSE)
+- OpenAI request/response JSON envelopes
+
+The goal is to keep the core engine reusable for future API surfaces.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Literal
+
+
+ToolChoice = Literal["auto", "none", "required"] | dict[str, Any]
+
+
+@dataclass(frozen=True)
+class StreamOptions:
+    """Chunked streaming policy."""
+
+    flush_every_n_tokens: int = 8
+    flush_every_ms: int = 50
+
+
+@dataclass(frozen=True)
+class ToolCall:
+    """A parsed tool call (function name + JSON-serializable arguments)."""
+
+    id: str
+    name: str
+    arguments: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class ChatMessage:
+    """A normalized chat message.
+
+    Notes:
+    - `tool_calls` is only meaningful for assistant messages.
+    - `tool_call_id` is only meaningful for tool messages (tool outputs).
+    """
+
+    role: Literal["system", "user", "assistant", "tool"]
+    content: str | None = None
+    tool_calls: list[ToolCall] = field(default_factory=list)
+    tool_call_id: str | None = None
+
+
+@dataclass(frozen=True)
+class ChatRequest:
+    """Normalized internal chat request."""
+
+    messages: list[ChatMessage]
+    tools: list[dict[str, Any]] = field(default_factory=list)
+    tool_choice: ToolChoice | None = None
+    max_tokens: int = 4096
+    temperature: float = 0.0
+    top_p: float = 1.0
+    stop: list[str] = field(default_factory=list)
+    stream: bool = False
+    stream_options: StreamOptions = field(default_factory=StreamOptions)
+    chat_template_kwargs: dict[str, Any] | None = None  # Additional kwargs for chat template
+    reasoning_budget: int | None = None  # Max tokens for thinking phase (enables thinking when set)
+    discard_thinking: bool | None = None  # If set, discard <think>...</think> from persisted session state
+    stream_thinking: bool | None = None  # If True, stream <think>...</think> content as separate deltas
+    session_id: str | None = None  # Session ID for stateful chat (KV cache reuse)
+    session_append_from_pos: int | None = None  # If set, only append prompt tokens from this position
+    extra: dict[str, Any] = field(default_factory=dict)  # Engine-specific per-request overrides
+
+
+@dataclass(frozen=True)
+class Usage:
+    prompt_tokens: int
+    completion_tokens: int
+
+    @property
+    def total_tokens(self) -> int:
+        return self.prompt_tokens + self.completion_tokens
+
+
+@dataclass(frozen=True)
+class Timing:
+    prefill_s: float | None = None
+    decode_s: float | None = None
+    total_s: float | None = None
+    tok_per_s: float | None = None
+
+
+@dataclass(frozen=True)
+class DeltaEvent:
+    """Chunked text delta."""
+
+    text: str
+
+
+@dataclass(frozen=True)
+class ThinkingDeltaEvent:
+    """Chunked thinking delta (text inside <think>...</think>)."""
+
+    text: str
+
+
+@dataclass(frozen=True)
+class ToolCallEvent:
+    """A completed tool call (or batch of tool calls)."""
+
+    tool_calls: list[ToolCall]
+
+
+@dataclass(frozen=True)
+class FinalEvent:
+    """Terminal event for a generation."""
+
+    finish_reason: Literal["stop", "length", "tool_calls", "cancelled", "error", "repetition"]
+    usage: Usage
+    timing: Timing
+    raw_content: str | None = None  # Unstripped content (includes <think> if any) when discard_thinking=False
+
+
+@dataclass(frozen=True)
+class ErrorEvent:
+    """Non-terminal or terminal error event."""
+
+    message: str
+    retryable: bool = False
+
+
+StreamEvent = DeltaEvent | ThinkingDeltaEvent | ToolCallEvent | FinalEvent | ErrorEvent

superlinear/engine/registry.py

@@ -0,0 +1,51 @@
+"""Model adapter registry.
+
+Maps model identifiers to their corresponding adapter classes.
+"""
+
+from typing import Type
+
+from .adapters.base import BaseAdapter
+from .adapters.superlinear import SuperlinearAdapter
+
+# Registry mapping model family names to adapter classes
+_ADAPTER_REGISTRY: dict[str, Type[BaseAdapter]] = {
+    "superlinear": SuperlinearAdapter,
+}
+
+
+def get_adapter(model_family: str) -> BaseAdapter:
+    """
+    Get an adapter instance for the given model family.
+    
+    Args:
+        model_family: Name of the model family (e.g., "superlinear").
+    
+    Returns:
+        An adapter instance for the model family.
+    
+    Raises:
+        ValueError: If the model family is not registered.
+    """
+    if model_family not in _ADAPTER_REGISTRY:
+        available = ", ".join(_ADAPTER_REGISTRY.keys())
+        raise ValueError(
+            f"Unknown model family: {model_family!r}. Available: {available}"
+        )
+    return _ADAPTER_REGISTRY[model_family]()
+
+
+def register_adapter(model_family: str, adapter_cls: Type[BaseAdapter]) -> None:
+    """
+    Register a new adapter for a model family.
+    
+    Args:
+        model_family: Name of the model family.
+        adapter_cls: Adapter class (must inherit from BaseAdapter).
+    """
+    _ADAPTER_REGISTRY[model_family] = adapter_cls
+
+
+def list_model_families() -> list[str]:
+    """Return list of registered model family names."""
+    return list(_ADAPTER_REGISTRY.keys())

superlinear/engine/repetition.py

@@ -0,0 +1,203 @@
+"""Token-level repetition detection for early stopping.
+
+This module implements a high-precision detector for exact token periodicity
+in the recent tail of generated token IDs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass(frozen=True)
+class RepeatHit:
+    period: int
+    repeats: int
+    checked_tail_len: int
+
+
+@dataclass(frozen=True)
+class RepetitionDetectionConfig:
+    """Configuration for repetition early-stop.
+
+    Notes:
+    - Defaults are tuned for typical long-context Q&A.
+    - `enabled` defaults to True to catch repetition loops early.
+    - Settings are conservative to avoid false positives on legitimate content.
+    """
+
+    enabled: bool = True
+    tail_len: int = 1024
+    check_every: int = 32
+    min_generated_tokens: int = 256
+    min_repeats: int = 3
+    max_period: int = 512
+    min_unique_tokens: int = 5
+
+    def validate(self) -> None:
+        if self.tail_len <= 0:
+            raise ValueError("'repetition_detection.tail_len' must be > 0.")
+        if self.check_every <= 0:
+            raise ValueError("'repetition_detection.check_every' must be > 0.")
+        if self.min_generated_tokens < 0:
+            raise ValueError("'repetition_detection.min_generated_tokens' must be >= 0.")
+        if self.min_repeats < 2:
+            raise ValueError("'repetition_detection.min_repeats' must be >= 2.")
+        if self.max_period <= 0:
+            raise ValueError("'repetition_detection.max_period' must be > 0.")
+        if self.min_unique_tokens <= 0:
+            raise ValueError("'repetition_detection.min_unique_tokens' must be > 0.")
+
+    def merged(self, override: Any | None) -> "RepetitionDetectionConfig":
+        """Merge a request-level override (typically request.extra['repetition_detection'])."""
+        if override is None:
+            return self
+        if isinstance(override, RepetitionDetectionConfig):
+            override.validate()
+            return override
+        if not isinstance(override, dict):
+            raise ValueError("'repetition_detection' must be an object.")
+
+        data: dict[str, Any] = dict(override)
+        if "min_unique_tokens" not in data and "min_unique_tokens_in_period" in data:
+            data["min_unique_tokens"] = data["min_unique_tokens_in_period"]
+
+        enabled = self.enabled
+        if "enabled" in data:
+            raw_enabled = data["enabled"]
+            if not isinstance(raw_enabled, bool):
+                raise ValueError("'repetition_detection.enabled' must be a boolean.")
+            enabled = raw_enabled
+        tail_len = self.tail_len if "tail_len" not in data else _coerce_int(data["tail_len"], "tail_len")
+        check_every = (
+            self.check_every
+            if "check_every" not in data
+            else _coerce_int(data["check_every"], "check_every")
+        )
+        min_generated_tokens = (
+            self.min_generated_tokens
+            if "min_generated_tokens" not in data
+            else _coerce_int(data["min_generated_tokens"], "min_generated_tokens", min_value=0)
+        )
+        min_repeats = (
+            self.min_repeats
+            if "min_repeats" not in data
+            else _coerce_int(data["min_repeats"], "min_repeats", min_value=2)
+        )
+        max_period = (
+            self.max_period
+            if "max_period" not in data
+            else _coerce_int(data["max_period"], "max_period")
+        )
+        min_unique_tokens = (
+            self.min_unique_tokens
+            if "min_unique_tokens" not in data
+            else _coerce_int(data["min_unique_tokens"], "min_unique_tokens", min_value=1)
+        )
+
+        merged = RepetitionDetectionConfig(
+            enabled=enabled,
+            tail_len=tail_len,
+            check_every=check_every,
+            min_generated_tokens=min_generated_tokens,
+            min_repeats=min_repeats,
+            max_period=max_period,
+            min_unique_tokens=min_unique_tokens,
+        )
+        merged.validate()
+        return merged
+
+
+def _coerce_int(value: Any, name: str, *, min_value: int = 1) -> int:
+    if isinstance(value, bool):
+        raise ValueError(f"'repetition_detection.{name}' must be an integer.")
+    try:
+        out = int(value)
+    except Exception as exc:
+        raise ValueError(f"'repetition_detection.{name}' must be an integer.") from exc
+    if out < min_value:
+        raise ValueError(f"'repetition_detection.{name}' must be >= {min_value}.")
+    return out
+
+
+def prefix_function(seq: list[int]) -> list[int]:
+    """Classic KMP prefix-function (pi array) for a sequence of ints.
+
+    pi[i] = length of the longest proper prefix of seq[:i+1]
+            that is also a suffix of seq[:i+1].
+    """
+    n = len(seq)
+    pi = [0] * n
+    j = 0
+    for i in range(1, n):
+        while j > 0 and seq[i] != seq[j]:
+            j = pi[j - 1]
+        if j < n and seq[i] == seq[j]:
+            j += 1
+        pi[i] = j
+    return pi
+
+
+def _nontrivial_period(period_tokens: list[int], *, min_unique_tokens: int) -> bool:
+    # Avoid stopping on junk like a single token or whitespace/punctuation-only loops.
+    return len(set(period_tokens)) >= min_unique_tokens
+
+
+def detect_repetition_kmp_tail(
+    tokens: list[int],
+    *,
+    tail_len: int = 1024,
+    min_generated_tokens: int = 256,
+    min_repeats: int = 3,
+    max_period: int = 512,
+    min_unique_tokens: int = 4,
+) -> RepeatHit | None:
+    """Detect exact periodic repetition using only the last `tail_len` tokens.
+
+    Strategy:
+    - Compute KMP prefix function on the tail.
+    - Walk the border chain to derive candidate periods.
+    - Validate candidate periods by checking the last `min_repeats` blocks match.
+    """
+    if tail_len <= 0:
+        raise ValueError("'tail_len' must be > 0.")
+    if min_generated_tokens < 0:
+        raise ValueError("'min_generated_tokens' must be >= 0.")
+    if min_repeats < 2:
+        raise ValueError("'min_repeats' must be >= 2.")
+    if max_period <= 0:
+        raise ValueError("'max_period' must be > 0.")
+    if min_unique_tokens <= 0:
+        raise ValueError("'min_unique_tokens' must be > 0.")
+
+    if len(tokens) < min_generated_tokens:
+        return None
+
+    tail = tokens[-tail_len:] if len(tokens) > tail_len else tokens
+    L = len(tail)
+    if L < min_repeats:
+        return None
+
+    pi = prefix_function(tail)
+    if not pi:
+        return None
+    b = pi[-1]
+
+    # Walk border chain: b -> pi[b-1] -> ...
+    while b > 0:
+        p = L - b
+        if 1 <= p <= max_period:
+            need = p * min_repeats
+            if L >= need:
+                a = tail[-p:]
+                ok = True
+                for r in range(2, min_repeats + 1):
+                    if tail[-r * p : -(r - 1) * p] != a:
+                        ok = False
+                        break
+                if ok and _nontrivial_period(a, min_unique_tokens=min_unique_tokens):
+                    return RepeatHit(period=p, repeats=min_repeats, checked_tail_len=L)
+        b = pi[b - 1]
+
+    return None