token-runtime 0.1.0a2__py3-none-any.whl

token_runtime/contracts.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Mapping, Protocol, runtime_checkable
+
+from .model import RequestEnvelope
+
+if TYPE_CHECKING:
+    from .capabilities import CapabilityKey, CapabilityProfile
+
+
+@runtime_checkable
+class ClientAdapterContract(Protocol):
+    @property
+    def client_id(self) -> str: ...
+
+    @property
+    def protocol_ids(self) -> tuple[str, ...]: ...
+
+
+@runtime_checkable
+class ProtocolAdapterContract(Protocol):
+    @property
+    def protocol_id(self) -> str: ...
+
+    def parse(self, payload: Mapping[str, Any]) -> RequestEnvelope: ...
+
+    def serialize(self, envelope: RequestEnvelope) -> dict[str, Any]: ...
+
+
+@runtime_checkable
+class CapabilityProviderContract(Protocol):
+    @property
+    def provider_id(self) -> str: ...
+
+    def profile_for(self, key: "CapabilityKey") -> "CapabilityProfile | None": ...
+
+
+@runtime_checkable
+class AgentIntegrationContract(Protocol):
+    @property
+    def agent_id(self) -> str: ...
+
+    @property
+    def capability_key(self) -> "CapabilityKey": ...
+
+    @property
+    def endpoint(self) -> str: ...
+
+
+@runtime_checkable
+class TokenizerAdapterContract(Protocol):
+    @property
+    def tokenizer_id(self) -> str: ...
+
+    def estimate(self, text: str) -> int: ...
+
+
+@runtime_checkable
+class StrategyContract(Protocol):
+    @property
+    def strategy_id(self) -> str: ...
+
+
+@runtime_checkable
+class BenchmarkAdapterContract(Protocol):
+    @property
+    def benchmark_id(self) -> str: ...
+
+    @property
+    def assertion_ids(self) -> tuple[str, ...]: ...

token_runtime/engine.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+from dataclasses import replace
+import re
+from typing import Iterable, Protocol
+
+from .model import OptimizationDecision, OptimizationResult, RequestEnvelope
+from .planner import ContextPlanner
+from .reducers import JsonToolOutputReducer, RepeatedLineReducer, RetrievedDuplicateReducer, ToolOutputReducer
+from .store import RecoveryStore
+
+
+_TOKEN_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE)
+
+
+def estimate_tokens(envelope: RequestEnvelope) -> int:
+    units = sum(len(_TOKEN_RE.findall(block.text)) for block in envelope.blocks)
+    return int(round(units * 1.33))
+
+
+class ReducerLike(Protocol):
+    def reduce(self, block, store: RecoveryStore, *, protected: bool): ...
+
+
+class OptimizationEngine:
+    def __init__(
+        self,
+        *,
+        planner: ContextPlanner,
+        store: RecoveryStore,
+        reducers: Iterable[ReducerLike] | None = None,
+    ):
+        self.planner = planner
+        self.store = store
+        self.reducers = tuple(reducers) if reducers is not None else (
+            JsonToolOutputReducer(),
+            RepeatedLineReducer(),
+            RetrievedDuplicateReducer(),
+            ToolOutputReducer(),
+        )
+
+    def optimize(self, envelope: RequestEnvelope) -> OptimizationResult:
+        before = estimate_tokens(envelope)
+        plan = self.planner.plan(envelope)
+        if plan.decision is OptimizationDecision.BYPASS:
+            return OptimizationResult(
+                decision=OptimizationDecision.BYPASS,
+                envelope=envelope,
+                reasons=plan.reasons,
+                before_estimated_tokens=before,
+                after_estimated_tokens=before,
+            )
+
+        blocks = []
+        changed_ids: list[str] = []
+        recovery_refs: list[str] = []
+        try:
+            for original in envelope.blocks:
+                current = original
+                for reducer in self.reducers:
+                    reduction = reducer.reduce(
+                        current,
+                        self.store,
+                        protected=original.id in plan.protected_ids,
+                    )
+                    current = reduction.block
+                    if reduction.changed:
+                        if original.id not in changed_ids:
+                            changed_ids.append(original.id)
+                        if reduction.recovery_ref and reduction.recovery_ref not in recovery_refs:
+                            recovery_refs.append(reduction.recovery_ref)
+                blocks.append(current)
+        except Exception:
+            return self._fallback(envelope, before, "fallback_internal_error")
+
+        blocks, dedup_changed, dedup_refs = self._dedup_exact_evidence(
+            envelope, blocks, plan.protected_ids
+        )
+        for block_id in dedup_changed:
+            if block_id not in changed_ids:
+                changed_ids.append(block_id)
+        for ref in dedup_refs:
+            if ref not in recovery_refs:
+                recovery_refs.append(ref)
+
+        optimized = replace(envelope, blocks=tuple(blocks))
+        after = estimate_tokens(optimized)
+        if not changed_ids or after >= before:
+            return self._fallback(envelope, before, "no_eligible_reduction")
+        return OptimizationResult(
+            decision=OptimizationDecision.OPTIMIZE,
+            envelope=optimized,
+            reasons=("optimized",),
+            before_estimated_tokens=before,
+            after_estimated_tokens=after,
+            changed_block_ids=tuple(changed_ids),
+            recovery_refs=tuple(recovery_refs),
+        )
+
+    def _dedup_exact_evidence(self, envelope, blocks, protected_ids):
+        evidence_kinds = {"retrieved", "tool_output"}
+        canonical: dict[str, tuple[int, str]] = {}
+        for index, (original, current) in enumerate(zip(envelope.blocks, blocks)):
+            if original.kind in evidence_kinds:
+                canonical[current.text] = (index, original.id)
+
+        out = []
+        changed: list[str] = []
+        refs: list[str] = []
+        for index, (original, current) in enumerate(zip(envelope.blocks, blocks)):
+            if original.kind not in evidence_kinds or original.id in protected_ids:
+                out.append(current)
+                continue
+            canonical_index, canonical_id = canonical[current.text]
+            if index == canonical_index:
+                out.append(current)
+                continue
+            ref = self.store.put(original.text.encode())
+            marker = f"[TOKEN_DUPLICATE_OF:{canonical_id};TOKEN_REF:{ref}]"
+            if len(marker) >= len(current.text):
+                out.append(current)
+                continue
+            out.append(replace(current, text=marker))
+            changed.append(original.id)
+            refs.append(ref)
+        return out, tuple(changed), tuple(refs)
+
+    @staticmethod
+    def _fallback(envelope: RequestEnvelope, before: int, reason: str) -> OptimizationResult:
+        return OptimizationResult(
+            decision=OptimizationDecision.BYPASS,
+            envelope=envelope,
+            reasons=(reason,),
+            before_estimated_tokens=before,
+            after_estimated_tokens=before,
+        )

token_runtime/feature_flags.py

@@ -0,0 +1,78 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+
+from .compatibility import CompatibilityState
+
+
+class FeatureFlagState(str, Enum):
+    DISABLED = "disabled"
+    SHADOW = "shadow"
+    CANARY = "canary"
+    ENABLED = "enabled"
+
+
+class FeatureEffect(str, Enum):
+    NONE = "none"
+    OBSERVE = "observe"
+    EXECUTE = "execute"
+
+
+@dataclass(frozen=True, slots=True)
+class FeatureFlag:
+    name: str
+    state: FeatureFlagState
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.state, FeatureFlagState):
+            raise TypeError("state must be FeatureFlagState")
+
+
+class DuplicateFeatureFlagError(ValueError):
+    pass
+
+
+def permitted_effect(
+    flag_state: FeatureFlagState,
+    compatibility_state: CompatibilityState,
+) -> FeatureEffect:
+    if flag_state is FeatureFlagState.DISABLED:
+        return FeatureEffect.NONE
+    if compatibility_state in {
+        CompatibilityState.BLOCKED,
+        CompatibilityState.UNSUPPORTED,
+    }:
+        return FeatureEffect.NONE
+    if flag_state is FeatureFlagState.SHADOW:
+        return FeatureEffect.OBSERVE
+    if flag_state is FeatureFlagState.CANARY:
+        if compatibility_state in {
+            CompatibilityState.CANARY,
+            CompatibilityState.CERTIFIED,
+        }:
+            return FeatureEffect.EXECUTE
+        return FeatureEffect.OBSERVE
+    if compatibility_state is CompatibilityState.CERTIFIED:
+        return FeatureEffect.EXECUTE
+    return FeatureEffect.OBSERVE
+
+
+class FeatureFlagRegistry:
+    def __init__(self) -> None:
+        self._flags: dict[str, FeatureFlag] = {}
+
+    def register(self, flag: FeatureFlag) -> FeatureFlag:
+        existing = self._flags.get(flag.name)
+        if existing is None:
+            self._flags[flag.name] = flag
+            return flag
+        if existing == flag:
+            return existing
+        raise DuplicateFeatureFlagError(f"conflicting feature flag: {flag.name}")
+
+    def get(self, name: str) -> FeatureFlag:
+        return self._flags.get(name, FeatureFlag(name, FeatureFlagState.DISABLED))
+
+    def snapshot(self) -> tuple[FeatureFlag, ...]:
+        return tuple(self._flags[name] for name in sorted(self._flags))

token_runtime/gateway.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import json
+import re
+import time
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from typing import Any
+from urllib.error import HTTPError
+from urllib.request import Request, urlopen
+
+from .adapters import ChatCompletionsAdapter, ResponsesAdapter
+from .metrics import MetricsStore
+from .model import OptimizationDecision
+
+
+_SAFE_LABEL_RE = re.compile(r"^[A-Za-z0-9._:/-]{1,128}$")
+_HOP_HEADERS = {"host", "content-length", "connection", "transfer-encoding"}
+
+
+@dataclass(frozen=True, slots=True)
+class PreparedRequest:
+    body: bytes
+    optimized: bool
+    reasons: tuple[str, ...]
+    adapter: str | None = None
+
+
+class GatewayCore:
+    def __init__(self, *, engine, metrics: MetricsStore):
+        self.engine = engine
+        self.metrics = metrics
+
+    def prepare(self, path: str, raw_body: bytes) -> PreparedRequest:
+        started = time.perf_counter()
+        adapter = self._adapter(path)
+        if adapter is None:
+            prepared = PreparedRequest(raw_body, False, ("unsupported_endpoint",), None)
+            self._record(prepared, None, None, None, started)
+            return prepared
+
+        try:
+            payload = json.loads(raw_body.decode("utf-8"))
+            if not isinstance(payload, dict):
+                raise ValueError("JSON request must be an object")
+        except (UnicodeDecodeError, json.JSONDecodeError, ValueError):
+            prepared = PreparedRequest(raw_body, False, ("invalid_json",), adapter.name)
+            self._record(prepared, None, None, None, started)
+            return prepared
+
+        try:
+            envelope = adapter.parse(payload)
+            result = self.engine.optimize(envelope)
+            if result.decision is OptimizationDecision.OPTIMIZE:
+                serialized = adapter.serialize(result.envelope)
+                body = json.dumps(
+                    serialized, ensure_ascii=False, separators=(",", ":")
+                ).encode("utf-8")
+                prepared = PreparedRequest(body, True, result.reasons, adapter.name)
+            else:
+                prepared = PreparedRequest(raw_body, False, result.reasons, adapter.name)
+            self._record(prepared, result, payload, adapter.name, started)
+            return prepared
+        except Exception:
+            prepared = PreparedRequest(
+                raw_body, False, ("fallback_internal_error",), adapter.name
+            )
+            self._record(prepared, None, payload, adapter.name, started)
+            return prepared
+
+    @staticmethod
+    def _adapter(path: str):
+        clean_path = path.split("?", 1)[0]
+        if clean_path == "/v1/responses":
+            return ResponsesAdapter()
+        if clean_path == "/v1/chat/completions":
+            return ChatCompletionsAdapter()
+        return None
+
+    def _record(self, prepared, result, payload, adapter_name, started) -> None:
+        before = getattr(result, "before_estimated_tokens", None)
+        after = getattr(result, "after_estimated_tokens", None)
+        model = payload.get("model") if isinstance(payload, dict) else None
+        if not isinstance(model, str) or not _SAFE_LABEL_RE.fullmatch(model):
+            model = None
+        event = {
+            "decision": "optimize" if prepared.optimized else "bypass",
+            "adapter": adapter_name,
+            "model": model,
+            "before_input_tokens": before,
+            "after_input_tokens": after,
+            "latency_ms": round((time.perf_counter() - started) * 1000, 3),
+            "reason_codes": prepared.reasons,
+            "reducer_ids": (),
+        }
+        self.metrics.record(event)
+
+
+def _request_headers(handler: BaseHTTPRequestHandler) -> dict[str, str]:
+    return {
+        key: value
+        for key, value in handler.headers.items()
+        if key.lower() not in _HOP_HEADERS
+    }
+
+
+def _write_upstream_response(handler, response) -> None:
+    handler.send_response(response.status)
+    for key, value in response.headers.items():
+        if key.lower() not in _HOP_HEADERS:
+            handler.send_header(key, value)
+    handler.end_headers()
+    while True:
+        chunk = response.read(64 * 1024)
+        if not chunk:
+            break
+        handler.wfile.write(chunk)
+
+
+def build_server(host: str, port: int, *, upstream: str, core: GatewayCore):
+    upstream_base = upstream.rstrip("/")
+
+    class Handler(BaseHTTPRequestHandler):
+        def do_POST(self):
+            length = int(self.headers.get("Content-Length", "0"))
+            raw_body = self.rfile.read(length)
+            prepared = core.prepare(self.path, raw_body)
+            request = Request(
+                upstream_base + self.path,
+                data=prepared.body,
+                headers=_request_headers(self),
+                method="POST",
+            )
+            try:
+                with urlopen(request, timeout=60) as response:
+                    _write_upstream_response(self, response)
+            except HTTPError as response:
+                _write_upstream_response(self, response)
+
+        def log_message(self, *args):
+            return
+
+    return ThreadingHTTPServer((host, port), Handler)
+
+
+def serve(host: str, port: int, *, upstream: str, core: GatewayCore) -> None:
+    server = build_server(host, port, upstream=upstream, core=core)
+    try:
+        server.serve_forever()
+    finally:
+        server.server_close()

token_runtime/gemini_adapter.py

@@ -0,0 +1,260 @@
+from __future__ import annotations
+
+from copy import deepcopy
+import json
+from typing import Any, Mapping, Sequence
+
+from .model import ContextBlock, RequestEnvelope
+
+
+def _set_path(root: Any, path: Sequence[Any], value: str) -> None:
+    target = root
+    for key in path[:-1]:
+        target = target[key]
+    target[path[-1]] = value
+
+
+def _block(
+    block_id: str,
+    kind: str,
+    text: str,
+    *,
+    role: str | None = None,
+    turn: int = 0,
+    path: Sequence[Any] | None = None,
+) -> ContextBlock:
+    metadata = {} if path is None else {"path": tuple(path)}
+    return ContextBlock(
+        block_id,
+        kind,
+        text,
+        role=role,
+        turn_index=turn,
+        metadata=metadata,
+    )
+
+
+def _encoded(value: Mapping[str, Any]) -> str:
+    return json.dumps(dict(value), sort_keys=True, separators=(",", ":"))
+
+
+class GeminiGenerateContentAdapter:
+    protocol_id = "gemini_generate_content"
+
+    def serialize(self, envelope: RequestEnvelope) -> dict[str, Any]:
+        payload = deepcopy(envelope.opaque["original_payload"])
+        for block in envelope.blocks:
+            path = block.metadata.get("path")
+            if path:
+                _set_path(payload, path, block.text)
+        return payload
+
+    def parse(self, payload: Mapping[str, Any]) -> RequestEnvelope:
+        original = deepcopy(dict(payload))
+        blocks: list[ContextBlock] = []
+        safe = True
+
+        if "cachedContent" in payload:
+            cached_content = payload.get("cachedContent")
+            safe = (
+                isinstance(cached_content, str)
+                and cached_content.startswith("cachedContents/")
+                and len(cached_content) > len("cachedContents/")
+            )
+
+        system_blocks, system_safe = self._parse_system(payload.get("systemInstruction"))
+        blocks.extend(system_blocks)
+        safe = safe and system_safe
+
+        tool_blocks, tools_safe = self._parse_tools(payload.get("tools"))
+        blocks.extend(tool_blocks)
+        safe = safe and tools_safe
+
+        contents = payload.get("contents")
+        if not isinstance(contents, list) or not contents:
+            safe = False
+            contents = []
+
+        for index, content in enumerate(contents):
+            content_blocks, content_safe = self._parse_content(index, content)
+            blocks.extend(content_blocks)
+            safe = safe and content_safe
+
+        return RequestEnvelope(
+            blocks=tuple(blocks),
+            opaque={"original_payload": original},
+            adapter=self.protocol_id,
+            wire_safe=safe,
+        )
+
+    @staticmethod
+    def _parse_system(system: Any) -> tuple[list[ContextBlock], bool]:
+        if system is None:
+            return [], True
+        if not isinstance(system, Mapping):
+            return [], False
+        parts = system.get("parts")
+        if not isinstance(parts, list) or not parts:
+            return [], False
+
+        blocks: list[ContextBlock] = []
+        safe = True
+        for index, part in enumerate(parts):
+            if not isinstance(part, Mapping) or not isinstance(part.get("text"), str):
+                safe = False
+                continue
+            if set(part) != {"text"}:
+                safe = False
+                continue
+            blocks.append(
+                _block(
+                    f"system-{index}",
+                    "system",
+                    part["text"],
+                    role="system",
+                    path=("systemInstruction", "parts", index, "text"),
+                )
+            )
+        return blocks, safe
+
+    @staticmethod
+    def _parse_tools(tools: Any) -> tuple[list[ContextBlock], bool]:
+        if tools is None:
+            return [], True
+        if not isinstance(tools, list):
+            return [], False
+        blocks: list[ContextBlock] = []
+        safe = True
+        for index, tool in enumerate(tools):
+            if not isinstance(tool, Mapping):
+                safe = False
+                continue
+            declarations = tool.get("functionDeclarations")
+            valid = (
+                set(tool) == {"functionDeclarations"}
+                and isinstance(declarations, list)
+                and bool(declarations)
+                and all(
+                    isinstance(declaration, Mapping)
+                    and isinstance(declaration.get("name"), str)
+                    for declaration in declarations
+                )
+            )
+            blocks.append(_block(f"tool-schema-{index}", "tool_schema", _encoded(tool)))
+            safe = safe and valid
+        return blocks, safe
+
+    def _parse_content(
+        self,
+        index: int,
+        content: Any,
+    ) -> tuple[list[ContextBlock], bool]:
+        if not isinstance(content, Mapping):
+            return [], False
+        role = content.get("role")
+        parts = content.get("parts")
+        if not isinstance(parts, list) or not parts:
+            return [], False
+        if role is None:
+            roleless_text = all(
+                isinstance(part, Mapping)
+                and set(part) == {"text"}
+                and isinstance(part.get("text"), str)
+                for part in parts
+            )
+            if not roleless_text:
+                return [], False
+            role = "user"
+        if role not in {"user", "model", "function"}:
+            return [], False
+
+        blocks: list[ContextBlock] = []
+        safe = True
+        for part_index, part in enumerate(parts):
+            part_blocks, part_safe = self._parse_part(index, part_index, role, part)
+            blocks.extend(part_blocks)
+            safe = safe and part_safe
+        return blocks, safe
+
+    @staticmethod
+    def _parse_part(
+        content_index: int,
+        part_index: int,
+        role: str,
+        part: Any,
+    ) -> tuple[list[ContextBlock], bool]:
+        if not isinstance(part, Mapping):
+            return [], False
+
+        if "thoughtSignature" in part or "thought" in part:
+            return [
+                _block(
+                    f"content-{content_index}-protocol-{part_index}",
+                    "protocol_state",
+                    _encoded(part),
+                    role="assistant" if role == "model" else role,
+                    turn=content_index,
+                )
+            ], False
+
+        if "text" in part:
+            text = part.get("text")
+            if role not in {"user", "model"} or not isinstance(text, str):
+                return [], False
+            if set(part) != {"text"}:
+                return [], False
+            kind = "assistant" if role == "model" else "user"
+            return [
+                _block(
+                    f"content-{content_index}-part-{part_index}",
+                    kind,
+                    text,
+                    role=kind,
+                    turn=content_index,
+                    path=("contents", content_index, "parts", part_index, "text"),
+                )
+            ], True
+
+        if "functionCall" in part:
+            call = part.get("functionCall")
+            valid = (
+                set(part) == {"functionCall"}
+                and role == "model"
+                and isinstance(call, Mapping)
+                and isinstance(call.get("name"), str)
+                and ("id" not in call or isinstance(call.get("id"), str))
+                and ("args" not in call or isinstance(call.get("args"), Mapping))
+            )
+            return [
+                _block(
+                    f"content-{content_index}-function-call-{part_index}",
+                    "tool_call",
+                    _encoded(part),
+                    role="assistant",
+                    turn=content_index,
+                )
+            ], valid
+
+        if "functionResponse" in part:
+            response = part.get("functionResponse")
+            valid = (
+                set(part) == {"functionResponse"}
+                and role in {"user", "function"}
+                and isinstance(response, Mapping)
+                and isinstance(response.get("name"), str)
+                and isinstance(response.get("response"), Mapping)
+                and ("id" not in response or isinstance(response.get("id"), str))
+                and not response.get("parts")
+            )
+            kind = "tool_output" if valid else "protocol_state"
+            return [
+                _block(
+                    f"content-{content_index}-function-response-{part_index}",
+                    kind,
+                    _encoded(part),
+                    role="tool",
+                    turn=content_index,
+                )
+            ], False
+
+        return [], False