Nested package layout

README.md

@@ -53,17 +53,19 @@ The package installs as `microtensor`, the customer facing name. The network too
 
 ## Layout
 
-| Module | Holds |
-|---|---|
-| `microtensor/engine.py` | `Engine`, plus `catalogue()`, `served_tasks()` and `frontier_models()` |
-| `microtensor/catalogue.py` | `Catalogue` and `Served`: the live task list and name resolution |
-| `microtensor/backends/hosted.py` | The hosted backend: the HTTP client for the API |
-| `microtensor/result.py` | `Result`, built from a chat completion |
-| `microtensor/budget.py` | `Budget`: the daily escalation allowance |
-| `microtensor/errors.py` | `EngineError` and `TaskUnavailable` |
-| `microtensor/constants.py` | Base URL, `mt/auto`, default floor and timeout |
-
-The local backend joins `backends/` with the Enterprise preview.
+    microtensor/
+      __init__.py          the public API
+      engine.py            Engine, plus catalogue(), served_tasks() and frontier_models()
+      core/                constants, errors and Result
+      catalogue/           Served and Catalogue: the live task list and name resolution
+      backends/
+        base.py            the Backend protocol every backend implements
+        hosted/            the hosted backend: the API client and message shaping
+        local/             the local backend; names the Enterprise preview until it ships
+      escalation/          Budget and the escalation policy
+    tests/
+      conftest.py          a fake gateway and the shared transports
+      catalogue/ backends/ engine/
 
 ## Development
 

microtensor/__init__.py

@@ -1,15 +1,21 @@
-from microtensor.backends.hosted import Hosted
-from microtensor.budget import Budget
+from microtensor.backends import Backend, Hosted
 from microtensor.catalogue import Catalogue, Served
-from microtensor.constants import AUTO, DEFAULT_BASE_URL, DEFAULT_FLOOR
+from microtensor.core import (
+    AUTO,
+    DEFAULT_BASE_URL,
+    DEFAULT_FLOOR,
+    EngineError,
+    Result,
+    TaskUnavailable,
+)
 from microtensor.engine import Engine, catalogue, frontier_models, served_tasks
-from microtensor.errors import EngineError, TaskUnavailable
-from microtensor.result import Result
+from microtensor.escalation import Budget
 
 __all__ = [
     "AUTO",
     "DEFAULT_BASE_URL",
     "DEFAULT_FLOOR",
+    "Backend",
     "Budget",
     "Catalogue",
     "Engine",

microtensor/backends/__init__.py

@@ -1,3 +1,4 @@
+from microtensor.backends.base import Backend
 from microtensor.backends.hosted import Hosted
 
-__all__ = ["Hosted"]
+__all__ = ["Backend", "Hosted"]

microtensor/backends/base.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any, Protocol
+
+from microtensor.catalogue import Catalogue
+
+
+class Backend(Protocol):
+    def models(self) -> list[dict[str, Any]]: ...
+
+    def catalogue(self) -> Catalogue: ...
+
+    def chat(
+        self, *, model: str, messages: Sequence[Mapping[str, str]], quality_floor: float
+    ) -> dict[str, Any]: ...
+
+    def close(self) -> None: ...

microtensor/backends/hosted/__init__.py

@@ -0,0 +1,4 @@
+from microtensor.backends.hosted.client import ENV_KEY, NO_KEY, Hosted
+from microtensor.backends.hosted.messages import messages
+
+__all__ = ["ENV_KEY", "NO_KEY", "Hosted", "messages"]

microtensor/backends/hosted.py → microtensor/backends/hosted/client.py

@@ -7,18 +7,11 @@
 import httpx
 
 from microtensor.catalogue import Catalogue
-from microtensor.constants import TIMEOUT_SECONDS
-from microtensor.errors import EngineError
+from microtensor.core.constants import TIMEOUT_SECONDS
+from microtensor.core.errors import EngineError
 
 NO_KEY: Final[str] = "set MICROTENSOR_API_KEY or pass api_key"
-
-
-def messages(question: str, context: str) -> list[dict[str, str]]:
-    found: list[dict[str, str]] = []
-    if context.strip():
-        found.append({"role": "system", "content": context})
-    found.append({"role": "user", "content": question})
-    return found
+ENV_KEY: Final[str] = "MICROTENSOR_API_KEY"
 
 
 class Hosted:
@@ -41,9 +34,7 @@ def connect(
         base_url: str,
         transport: httpx.BaseTransport | None = None,
     ) -> Hosted:
-        return cls(
-            api_key or os.environ.get("MICROTENSOR_API_KEY", ""), base_url, transport=transport
-        )
+        return cls(api_key or os.environ.get(ENV_KEY, ""), base_url, transport=transport)
 
     def close(self) -> None:
         self._http.close()

microtensor/backends/hosted/messages.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+
+def messages(question: str, context: str = "") -> list[dict[str, str]]:
+    found: list[dict[str, str]] = []
+    if context.strip():
+        found.append({"role": "system", "content": context})
+    found.append({"role": "user", "content": question})
+    return found

microtensor/backends/local/__init__.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from typing import Final, NoReturn
+
+from microtensor.core.errors import EngineError
+
+LOCAL_PREVIEW: Final[str] = (
+    "the local backend ships with the Enterprise preview; install microtensor-engine[local] "
+    "with a licence token, or use backend='hosted'"
+)
+
+
+def connect(*_: object, **__: object) -> NoReturn:
+    raise EngineError(LOCAL_PREVIEW)
+
+
+__all__ = ["LOCAL_PREVIEW", "connect"]

microtensor/catalogue/__init__.py

@@ -0,0 +1,4 @@
+from microtensor.catalogue.catalogue import Catalogue
+from microtensor.catalogue.served import Served, normal
+
+__all__ = ["Catalogue", "Served", "normal"]

microtensor/catalogue.py → microtensor/catalogue/catalogue.py

@@ -4,35 +4,8 @@
 from dataclasses import dataclass
 from typing import Any
 
-from microtensor.constants import AUTO
-
-
-def normal(name: str) -> str:
-    return name.strip().lower().replace(" ", "-")
-
-
-@dataclass(frozen=True, slots=True)
-class Served:
-    task: str
-    model: str
-    title: str
-    aliases: tuple[str, ...]
-    live: bool
-
-    @property
-    def names(self) -> tuple[str, ...]:
-        raw = [
-            self.task,
-            self.task.replace("_", "-"),
-            self.model,
-            self.model.removeprefix("mt/"),
-            *self.aliases,
-        ]
-        return tuple(dict.fromkeys(normal(n) for n in raw if n))
-
-    @property
-    def label(self) -> str:
-        return self.aliases[0] if self.aliases else self.task
+from microtensor.catalogue.served import Served, normal
+from microtensor.core.constants import AUTO
 
 
 @dataclass(frozen=True, slots=True)

microtensor/catalogue/served.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+def normal(name: str) -> str:
+    return name.strip().lower().replace(" ", "-")
+
+
+@dataclass(frozen=True, slots=True)
+class Served:
+    task: str
+    model: str
+    title: str
+    aliases: tuple[str, ...]
+    live: bool
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        raw = [
+            self.task,
+            self.task.replace("_", "-"),
+            self.model,
+            self.model.removeprefix("mt/"),
+            *self.aliases,
+        ]
+        return tuple(dict.fromkeys(normal(n) for n in raw if n))
+
+    @property
+    def label(self) -> str:
+        return self.aliases[0] if self.aliases else self.task

microtensor/core/__init__.py

@@ -0,0 +1,13 @@
+from microtensor.core.constants import AUTO, DEFAULT_BASE_URL, DEFAULT_FLOOR, TIMEOUT_SECONDS
+from microtensor.core.errors import EngineError, TaskUnavailable
+from microtensor.core.result import Result
+
+__all__ = [
+    "AUTO",
+    "DEFAULT_BASE_URL",
+    "DEFAULT_FLOOR",
+    "TIMEOUT_SECONDS",
+    "EngineError",
+    "Result",
+    "TaskUnavailable",
+]

microtensor/engine.py

@@ -2,21 +2,18 @@
 
 from collections.abc import Mapping, Sequence
 from dataclasses import replace
-from typing import Any, Final
+from typing import Any
 
 import httpx
 
-from microtensor.backends.hosted import Hosted, messages
-from microtensor.budget import Budget
-from microtensor.catalogue import Catalogue, Served, normal
-from microtensor.constants import AUTO, DEFAULT_BASE_URL, DEFAULT_FLOOR
-from microtensor.errors import EngineError, TaskUnavailable
-from microtensor.result import Result
-
-LOCAL_PREVIEW: Final[str] = (
-    "the local backend ships with the Enterprise preview; install microtensor-engine[local] "
-    "with a licence token, or use backend='hosted'"
-)
+from microtensor.backends import Backend, Hosted
+from microtensor.backends import local as local_backend
+from microtensor.backends.hosted import messages
+from microtensor.catalogue import Catalogue, Served
+from microtensor.core.constants import AUTO, DEFAULT_BASE_URL, DEFAULT_FLOOR
+from microtensor.core.errors import TaskUnavailable
+from microtensor.core.result import Result
+from microtensor.escalation import Budget, choose
 
 
 def catalogue(
@@ -54,15 +51,15 @@ class Engine:
     def __init__(
         self,
         *,
-        hosted: Hosted,
+        backend: Backend,
         catalogue: Catalogue,
         tasks: dict[str, str],
         quality_floor: float,
         escalation: bool,
         escalation_model: str | Mapping[str, Any],
         budget: Budget,
     ) -> None:
-        self._hosted = hosted
+        self._backend = backend
         self._catalogue = catalogue
         self._tasks = tasks
         self.quality_floor = quality_floor
@@ -98,7 +95,7 @@ def load(
         transport: httpx.BaseTransport | None = None,
     ) -> Engine:
         if backend != "hosted":
-            raise EngineError(LOCAL_PREVIEW)
+            local_backend.connect(device=device)
         hosted = Hosted.connect(api_key, base_url, transport)
         found = hosted.catalogue()
         wanted: dict[str, str] = {}
@@ -114,7 +111,7 @@ def load(
             usd_per_day=float((escalation_budget or {}).get("usd_per_day", 0) or 0) or None
         )
         return cls(
-            hosted=hosted,
+            backend=hosted,
             catalogue=found,
             tasks=wanted,
             quality_floor=quality_floor,
@@ -132,7 +129,7 @@ def catalogue(self) -> Catalogue:
         return self._catalogue
 
     def close(self) -> None:
-        self._hosted.close()
+        self._backend.close()
 
     def __enter__(self) -> Engine:
         return self
@@ -144,15 +141,15 @@ def run(self, question: str, *, context: str = "", task: str | None = None) -> R
         model = AUTO if task is None else self._model_for(task)
         sent = messages(question, context)
         result = Result.from_response(
-            self._hosted.chat(model=model, messages=sent, quality_floor=self.quality_floor)
+            self._backend.chat(model=model, messages=sent, quality_floor=self.quality_floor)
         )
         if not result.below_floor or not self.escalation:
             return result
-        target = self._escalation_for(result.task)
+        target = choose(self._escalation_model, task=result.task, catalogue=self._catalogue)
         if target is None or not self.budget.allows(0.0):
             return result
         escalated = Result.from_response(
-            self._hosted.chat(model=target, messages=sent, quality_floor=self.quality_floor)
+            self._backend.chat(model=target, messages=sent, quality_floor=self.quality_floor)
         )
         self.budget.charge(escalated.cost)
         return replace(escalated, task=result.task, escalated=True, escalated_to=target)
@@ -164,19 +161,3 @@ def _model_for(self, task: str) -> str:
             loaded = ", ".join(sorted(self._tasks)) or "none"
             raise TaskUnavailable(f"{task!r} was not loaded; loaded: {loaded}")
         return model
-
-    def _escalation_for(self, task: str) -> str | None:
-        chosen: Any = self._escalation_model
-        if isinstance(chosen, Mapping):
-            entry = self._catalogue.resolve(task)
-            names = entry.names if entry else (normal(task),)
-            chosen = next((chosen[k] for k in chosen if normal(str(k)) in names), "auto")
-        if chosen is False or chosen is None:
-            return None
-        frontier = self._catalogue.frontier
-        if chosen == "auto":
-            return frontier[0] if frontier else None
-        name = str(chosen)
-        if name not in frontier:
-            raise EngineError(f"{name!r} is not a frontier model the network serves")
-        return name

microtensor/escalation/__init__.py

@@ -0,0 +1,4 @@
+from microtensor.escalation.budget import Budget
+from microtensor.escalation.policy import choose
+
+__all__ = ["Budget", "choose"]

microtensor/escalation/policy.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from microtensor.catalogue import Catalogue, normal
+from microtensor.core.errors import EngineError
+
+NOT_FRONTIER = "{name!r} is not a frontier model the network serves"
+
+
+def choose(target: str | Mapping[str, Any], *, task: str, catalogue: Catalogue) -> str | None:
+    chosen: Any = target
+    if isinstance(chosen, Mapping):
+        entry = catalogue.resolve(task)
+        names = entry.names if entry else (normal(task),)
+        chosen = next((chosen[k] for k in chosen if normal(str(k)) in names), "auto")
+    if chosen is False or chosen is None:
+        return None
+    if chosen == "auto":
+        return catalogue.frontier[0] if catalogue.frontier else None
+    name = str(chosen)
+    if name not in catalogue.frontier:
+        raise EngineError(NOT_FRONTIER.format(name=name))
+    return name