diff --git a/README.md b/README.md index 2d5173b..1c311fe 100644 --- a/README.md +++ b/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 diff --git a/microtensor/__init__.py b/microtensor/__init__.py index 8951ba3..f71476e 100644 --- a/microtensor/__init__.py +++ b/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", diff --git a/microtensor/backends/__init__.py b/microtensor/backends/__init__.py index 820b3aa..a8d62a0 100644 --- a/microtensor/backends/__init__.py +++ b/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"] diff --git a/microtensor/backends/base.py b/microtensor/backends/base.py new file mode 100644 index 0000000..19b661c --- /dev/null +++ b/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: ... diff --git a/microtensor/backends/hosted/__init__.py b/microtensor/backends/hosted/__init__.py new file mode 100644 index 0000000..61c40b1 --- /dev/null +++ b/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"] diff --git a/microtensor/backends/hosted.py b/microtensor/backends/hosted/client.py similarity index 82% rename from microtensor/backends/hosted.py rename to microtensor/backends/hosted/client.py index c4ece31..dfb5615 100644 --- a/microtensor/backends/hosted.py +++ b/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() diff --git a/microtensor/backends/hosted/messages.py b/microtensor/backends/hosted/messages.py new file mode 100644 index 0000000..f624724 --- /dev/null +++ b/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 diff --git a/microtensor/backends/local/__init__.py b/microtensor/backends/local/__init__.py new file mode 100644 index 0000000..12623cc --- /dev/null +++ b/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"] diff --git a/microtensor/catalogue/__init__.py b/microtensor/catalogue/__init__.py new file mode 100644 index 0000000..cdcf900 --- /dev/null +++ b/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"] diff --git a/microtensor/catalogue.py b/microtensor/catalogue/catalogue.py similarity index 76% rename from microtensor/catalogue.py rename to microtensor/catalogue/catalogue.py index c563343..bc87712 100644 --- a/microtensor/catalogue.py +++ b/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) diff --git a/microtensor/catalogue/served.py b/microtensor/catalogue/served.py new file mode 100644 index 0000000..3c9d220 --- /dev/null +++ b/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 diff --git a/microtensor/core/__init__.py b/microtensor/core/__init__.py new file mode 100644 index 0000000..45e4ca0 --- /dev/null +++ b/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", +] diff --git a/microtensor/constants.py b/microtensor/core/constants.py similarity index 100% rename from microtensor/constants.py rename to microtensor/core/constants.py diff --git a/microtensor/errors.py b/microtensor/core/errors.py similarity index 100% rename from microtensor/errors.py rename to microtensor/core/errors.py diff --git a/microtensor/result.py b/microtensor/core/result.py similarity index 100% rename from microtensor/result.py rename to microtensor/core/result.py diff --git a/microtensor/engine.py b/microtensor/engine.py index d9b6b0a..4af7e50 100644 --- a/microtensor/engine.py +++ b/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,7 +51,7 @@ class Engine: def __init__( self, *, - hosted: Hosted, + backend: Backend, catalogue: Catalogue, tasks: dict[str, str], quality_floor: float, @@ -62,7 +59,7 @@ def __init__( 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 diff --git a/microtensor/escalation/__init__.py b/microtensor/escalation/__init__.py new file mode 100644 index 0000000..86ffe0e --- /dev/null +++ b/microtensor/escalation/__init__.py @@ -0,0 +1,4 @@ +from microtensor.escalation.budget import Budget +from microtensor.escalation.policy import choose + +__all__ = ["Budget", "choose"] diff --git a/microtensor/budget.py b/microtensor/escalation/budget.py similarity index 100% rename from microtensor/budget.py rename to microtensor/escalation/budget.py diff --git a/microtensor/escalation/policy.py b/microtensor/escalation/policy.py new file mode 100644 index 0000000..5855092 --- /dev/null +++ b/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 diff --git a/tests/backends/test_hosted.py b/tests/backends/test_hosted.py new file mode 100644 index 0000000..c1d2ba3 --- /dev/null +++ b/tests/backends/test_hosted.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +import httpx +import pytest + +from microtensor import EngineError, Hosted +from microtensor.backends import local +from microtensor.backends.hosted import messages + + +def test_the_hosted_backend_needs_a_key_and_reads_the_environment( + monkeypatch: pytest.MonkeyPatch, transport: httpx.MockTransport +) -> None: + monkeypatch.delenv("MICROTENSOR_API_KEY", raising=False) + with pytest.raises(EngineError): + Hosted.connect(None, "http://g/v1", transport) + monkeypatch.setenv("MICROTENSOR_API_KEY", "mt_sk_env") + hosted = Hosted.connect(None, "http://g/v1", transport) + assert [s.task for s in hosted.catalogue().served] == ["classification", "sql_generation"] + hosted.close() + + +def test_messages_put_the_context_in_the_system_message() -> None: + assert messages("q", "ctx") == [ + {"role": "system", "content": "ctx"}, + {"role": "user", "content": "q"}, + ] + assert messages("q", " ") == [{"role": "user", "content": "q"}] + + +def test_the_local_backend_names_the_preview() -> None: + with pytest.raises(EngineError, match="Enterprise preview"): + local.connect(device="cuda") diff --git a/tests/test_catalogue.py b/tests/catalogue/test_catalogue.py similarity index 100% rename from tests/test_catalogue.py rename to tests/catalogue/test_catalogue.py diff --git a/tests/test_engine.py b/tests/engine/test_engine.py similarity index 100% rename from tests/test_engine.py rename to tests/engine/test_engine.py