From 4fe10aa0a7510f1426adfca236faca97361548c1 Mon Sep 17 00:00:00 2001 From: BANADDA MUBARAKA <83466862+BANADDA@users.noreply.github.com> Date: Fri, 16 Oct 2026 04:16:35 +0300 Subject: [PATCH] package: catalogue, hosted backend, result, budget and errors as their own modules --- README.md | 14 ++ microtensor/__init__.py | 23 ++- microtensor/backends/__init__.py | 3 + microtensor/backends/hosted.py | 82 +++++++++ microtensor/budget.py | 29 ++++ microtensor/catalogue.py | 89 ++++++++++ microtensor/constants.py | 6 + microtensor/engine.py | 289 ++++--------------------------- microtensor/errors.py | 6 + microtensor/result.py | 36 ++++ tests/conftest.py | 72 ++++++++ tests/test_catalogue.py | 41 +++++ tests/test_engine.py | 114 ++++-------- 13 files changed, 462 insertions(+), 342 deletions(-) create mode 100644 microtensor/backends/__init__.py create mode 100644 microtensor/backends/hosted.py create mode 100644 microtensor/budget.py create mode 100644 microtensor/catalogue.py create mode 100644 microtensor/constants.py create mode 100644 microtensor/errors.py create mode 100644 microtensor/result.py create mode 100644 tests/conftest.py create mode 100644 tests/test_catalogue.py diff --git a/README.md b/README.md index 69d355a..2d5173b 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,20 @@ Certified routers can be served with CUDA graphs and fused kernels for RMSNorm, The package installs as `microtensor`, the customer facing name. The network tooling installs as `microtensor_subnet` (distribution `microtensor-subnet`), so both can live in one environment. +## 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. + ## Development pip install -e ".[dev]" diff --git a/microtensor/__init__.py b/microtensor/__init__.py index 18fb766..8951ba3 100644 --- a/microtensor/__init__.py +++ b/microtensor/__init__.py @@ -1,21 +1,20 @@ -from microtensor.engine import ( - Budget, - Catalogue, - Engine, - EngineError, - Result, - Served, - TaskUnavailable, - catalogue, - frontier_models, - served_tasks, -) +from microtensor.backends.hosted import Hosted +from microtensor.budget import Budget +from microtensor.catalogue import Catalogue, Served +from microtensor.constants import AUTO, DEFAULT_BASE_URL, DEFAULT_FLOOR +from microtensor.engine import Engine, catalogue, frontier_models, served_tasks +from microtensor.errors import EngineError, TaskUnavailable +from microtensor.result import Result __all__ = [ + "AUTO", + "DEFAULT_BASE_URL", + "DEFAULT_FLOOR", "Budget", "Catalogue", "Engine", "EngineError", + "Hosted", "Result", "Served", "TaskUnavailable", diff --git a/microtensor/backends/__init__.py b/microtensor/backends/__init__.py new file mode 100644 index 0000000..820b3aa --- /dev/null +++ b/microtensor/backends/__init__.py @@ -0,0 +1,3 @@ +from microtensor.backends.hosted import Hosted + +__all__ = ["Hosted"] diff --git a/microtensor/backends/hosted.py b/microtensor/backends/hosted.py new file mode 100644 index 0000000..c4ece31 --- /dev/null +++ b/microtensor/backends/hosted.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import os +from collections.abc import Mapping, Sequence +from typing import Any, Final + +import httpx + +from microtensor.catalogue import Catalogue +from microtensor.constants import TIMEOUT_SECONDS +from microtensor.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 + + +class Hosted: + def __init__( + self, api_key: str, base_url: str, *, transport: httpx.BaseTransport | None = None + ) -> None: + if not api_key: + raise EngineError(NO_KEY) + self._http = httpx.Client( + base_url=base_url.rstrip("/"), + headers={"authorization": f"Bearer {api_key}", "accept": "application/json"}, + timeout=TIMEOUT_SECONDS, + transport=transport, + ) + + @classmethod + def connect( + cls, + api_key: str | None, + base_url: str, + transport: httpx.BaseTransport | None = None, + ) -> Hosted: + return cls( + api_key or os.environ.get("MICROTENSOR_API_KEY", ""), base_url, transport=transport + ) + + def close(self) -> None: + self._http.close() + + def _body(self, response: httpx.Response) -> dict[str, Any]: + try: + body = response.json() + except ValueError as exc: + raise EngineError("the API answered with something that is not JSON") from exc + if response.status_code >= 400: + error = body.get("error") if isinstance(body, dict) else {} + error = error if isinstance(error, dict) else {} + raise EngineError( + f"{error.get('code', 'error')}: {error.get('message', response.text[:200])}" + ) + if not isinstance(body, dict): + raise EngineError("unexpected answer from the API") + return body + + def models(self) -> list[dict[str, Any]]: + body = self._body(self._http.get("/models")) + return [dict(m) for m in body.get("data", [])] + + def catalogue(self) -> Catalogue: + return Catalogue.from_models(self.models()) + + def chat( + self, *, model: str, messages: Sequence[Mapping[str, str]], quality_floor: float + ) -> dict[str, Any]: + payload = { + "model": model, + "messages": [dict(m) for m in messages], + "quality_floor": quality_floor, + "stream": False, + } + return self._body(self._http.post("/chat/completions", json=payload)) diff --git a/microtensor/budget.py b/microtensor/budget.py new file mode 100644 index 0000000..464a067 --- /dev/null +++ b/microtensor/budget.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +import datetime as dt +from dataclasses import dataclass, field + + +def today() -> dt.date: + return dt.datetime.now(dt.timezone.utc).date() + + +@dataclass(slots=True) +class Budget: + usd_per_day: float | None = None + spent: float = 0.0 + day: dt.date = field(default_factory=today) + + def allows(self, cost: float) -> bool: + self._roll() + return self.usd_per_day is None or self.spent + cost <= self.usd_per_day + + def charge(self, cost: float) -> None: + self._roll() + self.spent += cost + + def _roll(self) -> None: + now = today() + if now != self.day: + self.day = now + self.spent = 0.0 diff --git a/microtensor/catalogue.py b/microtensor/catalogue.py new file mode 100644 index 0000000..c563343 --- /dev/null +++ b/microtensor/catalogue.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from collections.abc import Mapping, Sequence +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 + + +@dataclass(frozen=True, slots=True) +class Catalogue: + served: tuple[Served, ...] + frontier: tuple[str, ...] + + @classmethod + def from_models(cls, models: Sequence[Mapping[str, Any]]) -> Catalogue: + entries = [dict(m) for m in models] + picked: dict[str, str] = {} + live: dict[str, bool] = {} + for entry in entries: + if entry.get("id") != AUTO: + continue + for item in entry.get("tasks") or []: + task = str(item.get("task", "")) + if task: + picked[task] = str(item.get("model", "")) + live[task] = bool(item.get("live", False)) + for entry in entries: + task = str(entry.get("task", "")) + if task and entry.get("id") != AUTO and task not in picked: + picked[task] = str(entry.get("id", "")) + live[task] = bool((entry.get("capacity") or {}).get("live", False)) + by_id = {str(e.get("id", "")): e for e in entries} + served = tuple( + Served( + task=task, + model=model, + title=str(by_id.get(model, {}).get("title", "")), + aliases=tuple(str(a) for a in by_id.get(model, {}).get("aliases") or []), + live=live.get(task, False), + ) + for task, model in sorted(picked.items()) + ) + frontier = tuple( + str(e["id"]) for e in entries if e.get("kind") == "frontier" and e.get("id") + ) + return cls(served=served, frontier=frontier) + + def resolve(self, name: str) -> Served | None: + wanted = normal(name) + for entry in self.served: + if wanted in entry.names: + return entry + return None + + def describe(self) -> str: + parts: list[str] = [] + for entry in self.served: + extra = f" ({', '.join(entry.aliases)})" if entry.aliases else "" + parts.append(f"{entry.task}{extra}") + return ", ".join(parts) or "none" diff --git a/microtensor/constants.py b/microtensor/constants.py new file mode 100644 index 0000000..5ce2646 --- /dev/null +++ b/microtensor/constants.py @@ -0,0 +1,6 @@ +from typing import Final + +DEFAULT_BASE_URL: Final[str] = "https://api.microtensor.cloud/v1" +AUTO: Final[str] = "mt/auto" +DEFAULT_FLOOR: Final[float] = 0.9 +TIMEOUT_SECONDS: Final[float] = 180.0 diff --git a/microtensor/engine.py b/microtensor/engine.py index 96dcac1..d9b6b0a 100644 --- a/microtensor/engine.py +++ b/microtensor/engine.py @@ -1,220 +1,31 @@ from __future__ import annotations -import datetime as dt -import os from collections.abc import Mapping, Sequence -from dataclasses import dataclass, field +from dataclasses import replace from typing import Any, Final import httpx -DEFAULT_BASE_URL: Final[str] = "https://api.microtensor.cloud/v1" -AUTO: Final[str] = "mt/auto" -DEFAULT_FLOOR: Final[float] = 0.9 -TIMEOUT_SECONDS: Final[float] = 180.0 -NO_KEY: Final[str] = "set MICROTENSOR_API_KEY or pass api_key" +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'" ) -class EngineError(RuntimeError): - pass - - -class TaskUnavailable(EngineError): - pass - - -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 - - -@dataclass(frozen=True, slots=True) -class Catalogue: - served: tuple[Served, ...] - frontier: tuple[str, ...] - - @classmethod - def from_models(cls, models: Sequence[Mapping[str, Any]]) -> Catalogue: - entries = [dict(m) for m in models] - picked: dict[str, str] = {} - live: dict[str, bool] = {} - for entry in entries: - if entry.get("id") != AUTO: - continue - for item in entry.get("tasks") or []: - task = str(item.get("task", "")) - if task: - picked[task] = str(item.get("model", "")) - live[task] = bool(item.get("live", False)) - for entry in entries: - task = str(entry.get("task", "")) - if task and entry.get("id") != AUTO and task not in picked: - picked[task] = str(entry.get("id", "")) - live[task] = bool((entry.get("capacity") or {}).get("live", False)) - by_id = {str(e.get("id", "")): e for e in entries} - served = tuple( - Served( - task=task, - model=model, - title=str(by_id.get(model, {}).get("title", "")), - aliases=tuple(str(a) for a in by_id.get(model, {}).get("aliases") or []), - live=live.get(task, False), - ) - for task, model in sorted(picked.items()) - ) - frontier = tuple( - str(e["id"]) for e in entries if e.get("kind") == "frontier" and e.get("id") - ) - return cls(served=served, frontier=frontier) - - def resolve(self, name: str) -> Served | None: - wanted = _normal(name) - for entry in self.served: - if wanted in entry.names: - return entry - return None - - def describe(self) -> str: - parts = [] - for entry in self.served: - extra = f" ({', '.join(entry.aliases)})" if entry.aliases else "" - parts.append(f"{entry.task}{extra}") - return ", ".join(parts) or "none" - - -@dataclass(frozen=True, slots=True) -class Result: - answer: str - task: str - specialist: str - confidence: float - below_floor: bool - escalated: bool - escalated_to: str | None - cost: float - request_id: str - raw: dict[str, Any] = field(default_factory=dict, repr=False) - - -@dataclass(slots=True) -class Budget: - usd_per_day: float | None = None - spent: float = 0.0 - day: dt.date = field(default_factory=lambda: dt.datetime.now(dt.timezone.utc).date()) - - def allows(self, cost: float) -> bool: - self._roll() - return self.usd_per_day is None or self.spent + cost <= self.usd_per_day - - def charge(self, cost: float) -> None: - self._roll() - self.spent += cost - - def _roll(self) -> None: - today = dt.datetime.now(dt.timezone.utc).date() - if today != self.day: - self.day = today - self.spent = 0.0 - - -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 - - -class Hosted: - def __init__( - self, api_key: str, base_url: str, *, transport: httpx.BaseTransport | None = None - ) -> None: - if not api_key: - raise EngineError(NO_KEY) - self._http = httpx.Client( - base_url=base_url.rstrip("/"), - headers={"authorization": f"Bearer {api_key}", "accept": "application/json"}, - timeout=TIMEOUT_SECONDS, - transport=transport, - ) - - def close(self) -> None: - self._http.close() - - def _body(self, response: httpx.Response) -> dict[str, Any]: - try: - body = response.json() - except ValueError as exc: - raise EngineError("the API answered with something that is not JSON") from exc - if response.status_code >= 400: - error = body.get("error") if isinstance(body, dict) else {} - error = error if isinstance(error, dict) else {} - raise EngineError( - f"{error.get('code', 'error')}: {error.get('message', response.text[:200])}" - ) - if not isinstance(body, dict): - raise EngineError("unexpected answer from the API") - return body - - def models(self) -> list[dict[str, Any]]: - body = self._body(self._http.get("/models")) - return [dict(m) for m in body.get("data", [])] - - def catalogue(self) -> Catalogue: - return Catalogue.from_models(self.models()) - - def chat( - self, *, model: str, messages: Sequence[Mapping[str, str]], quality_floor: float - ) -> dict[str, Any]: - payload = { - "model": model, - "messages": [dict(m) for m in messages], - "quality_floor": quality_floor, - "stream": False, - } - return self._body(self._http.post("/chat/completions", json=payload)) - - -def _hosted(api_key: str | None, base_url: str, transport: httpx.BaseTransport | None) -> Hosted: - return Hosted( - api_key or os.environ.get("MICROTENSOR_API_KEY", ""), base_url, transport=transport - ) - - def catalogue( *, api_key: str | None = None, base_url: str = DEFAULT_BASE_URL, transport: httpx.BaseTransport | None = None, ) -> Catalogue: - hosted = _hosted(api_key, base_url, transport) + hosted = Hosted.connect(api_key, base_url, transport) try: return hosted.catalogue() finally: @@ -288,7 +99,7 @@ def load( ) -> Engine: if backend != "hosted": raise EngineError(LOCAL_PREVIEW) - hosted = _hosted(api_key, base_url, transport) + hosted = Hosted.connect(api_key, base_url, transport) found = hosted.catalogue() wanted: dict[str, str] = {} for name in tasks: @@ -329,12 +140,37 @@ def __enter__(self) -> Engine: def __exit__(self, *_: object) -> None: self.close() + def run(self, question: str, *, context: str = "", task: str | None = None) -> Result: + 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) + ) + if not result.below_floor or not self.escalation: + return result + target = self._escalation_for(result.task) + 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.budget.charge(escalated.cost) + return replace(escalated, task=result.task, escalated=True, escalated_to=target) + + def _model_for(self, task: str) -> str: + entry = self._catalogue.resolve(task) + model = self._tasks.get(entry.task, "") if entry else "" + if not model: + 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") + 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 @@ -344,54 +180,3 @@ def _escalation_for(self, task: str) -> str | None: if name not in frontier: raise EngineError(f"{name!r} is not a frontier model the network serves") return name - - def run(self, question: str, *, context: str = "", task: str | None = None) -> Result: - model = AUTO - if task is not None: - entry = self._catalogue.resolve(task) - model = self._tasks.get(entry.task, "") if entry else "" - if not model: - raise TaskUnavailable(f"{task!r} was not loaded; loaded: {self._loaded()}") - messages = _messages(question, context) - found = self._hosted.chat(model=model, messages=messages, quality_floor=self.quality_floor) - result = self._result(found) - if not result.below_floor or not self.escalation: - return result - target = self._escalation_for(result.task) - if target is None or not self.budget.allows(0.0): - return result - again = self._hosted.chat(model=target, messages=messages, quality_floor=self.quality_floor) - escalated = self._result(again) - self.budget.charge(escalated.cost) - return Result( - answer=escalated.answer, - task=result.task, - specialist=escalated.specialist, - confidence=escalated.confidence, - below_floor=escalated.below_floor, - escalated=True, - escalated_to=target, - cost=escalated.cost, - request_id=escalated.request_id, - raw=again, - ) - - def _loaded(self) -> str: - return ", ".join(sorted(self._tasks)) or "none" - - def _result(self, found: dict[str, Any]) -> Result: - choices = found.get("choices") or [{}] - message = choices[0].get("message") or {} - routing = found.get("routing") or {} - return Result( - answer=str(message.get("content", "")), - task=str(routing.get("task", "")), - specialist=str(routing.get("specialist", found.get("model", ""))), - confidence=float(routing.get("answer_confidence") or 0.0), - below_floor=not bool(routing.get("floor_met", False)), - escalated=bool(routing.get("escalated", False)), - escalated_to=None, - cost=float((found.get("usage") or {}).get("cost") or 0.0), - request_id=str(found.get("id", "")), - raw=found, - ) diff --git a/microtensor/errors.py b/microtensor/errors.py new file mode 100644 index 0000000..bc14a53 --- /dev/null +++ b/microtensor/errors.py @@ -0,0 +1,6 @@ +class EngineError(RuntimeError): + pass + + +class TaskUnavailable(EngineError): + pass diff --git a/microtensor/result.py b/microtensor/result.py new file mode 100644 index 0000000..141e4e4 --- /dev/null +++ b/microtensor/result.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + + +@dataclass(frozen=True, slots=True) +class Result: + answer: str + task: str + specialist: str + confidence: float + below_floor: bool + escalated: bool + escalated_to: str | None + cost: float + request_id: str + raw: dict[str, Any] = field(default_factory=dict, repr=False) + + @classmethod + def from_response(cls, found: dict[str, Any]) -> Result: + choices = found.get("choices") or [{}] + message = choices[0].get("message") or {} + routing = found.get("routing") or {} + return cls( + answer=str(message.get("content", "")), + task=str(routing.get("task", "")), + specialist=str(routing.get("specialist", found.get("model", ""))), + confidence=float(routing.get("answer_confidence") or 0.0), + below_floor=not bool(routing.get("floor_met", False)), + escalated=bool(routing.get("escalated", False)), + escalated_to=None, + cost=float((found.get("usage") or {}).get("cost") or 0.0), + request_id=str(found.get("id", "")), + raw=found, + ) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..33091d0 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,72 @@ +from __future__ import annotations + +import json +from collections.abc import Callable +from typing import Any + +import httpx +import pytest + +SQL = { + "id": "mt/sql_gen_s1", + "task": "sql_generation", + "title": "SQL generation, season 1", + "aliases": ["text-to-sql", "sql"], + "quality": 0.73, + "capacity": {"operators": 2, "live": True}, +} +CLASSIFY = { + "id": "mt/classify_s1", + "task": "classification", + "title": "Classification, season 1", + "capacity": {"operators": 0, "live": False}, +} +AUTO = { + "id": "mt/auto", + "tasks": [{"task": "sql_generation", "model": "mt/sql_gen_s1", "live": True}], +} +FRONTIER = {"id": "deepseek-v4-pro", "kind": "frontier", "object": "model"} + +MODELS: list[dict[str, Any]] = [AUTO, SQL, CLASSIFY] + + +def make_gateway(models: list[dict[str, Any]]) -> Callable[[httpx.Request], httpx.Response]: + def gateway(request: httpx.Request) -> httpx.Response: + if request.url.path.endswith("/models"): + return httpx.Response(200, json={"data": models}) + body = json.loads(request.content) + frontier = body["model"] == FRONTIER["id"] + confident = frontier or "2025" in body["messages"][-1]["content"] + return httpx.Response( + 200, + json={ + "id": "cmpl_1", + "model": body["model"], + "choices": [{"message": {"content": "SELECT count(*) FROM orders"}}], + "routing": { + "task": "sql_generation", + "specialist": body["model"], + "answer_confidence": 0.96 if confident else 0.4, + "floor_met": confident, + "quality_floor": body["quality_floor"], + }, + "usage": {"total_tokens": 49, "cost": 0.002 if frontier else 0.0}, + }, + ) + + return gateway + + +@pytest.fixture +def models() -> list[dict[str, Any]]: + return MODELS + + +@pytest.fixture +def transport() -> httpx.MockTransport: + return httpx.MockTransport(make_gateway(MODELS)) + + +@pytest.fixture +def frontier_transport() -> httpx.MockTransport: + return httpx.MockTransport(make_gateway([*MODELS, FRONTIER])) diff --git a/tests/test_catalogue.py b/tests/test_catalogue.py new file mode 100644 index 0000000..c86da62 --- /dev/null +++ b/tests/test_catalogue.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from typing import Any + +import httpx +import pytest + +from microtensor import Catalogue, Engine, TaskUnavailable, frontier_models, served_tasks + +KEY = {"api_key": "mt_sk_x", "base_url": "http://g/v1"} + + +def test_the_catalogue_reads_tasks_aliases_and_liveness(models: list[dict[str, Any]]) -> None: + found = Catalogue.from_models(models) + assert [s.task for s in found.served] == ["classification", "sql_generation"] + sql = found.served[1] + assert sql.aliases == ("text-to-sql", "sql") and sql.live and sql.label == "text-to-sql" + assert not found.served[0].live and found.served[0].label == "classification" + assert found.frontier == () + for name in ("text-to-sql", "SQL", "sql_generation", "sql-generation", "sql_gen_s1"): + assert found.resolve(name) is sql + assert found.resolve("mt/sql_gen_s1") is sql and found.resolve("summarisation") is None + assert found.describe() == "classification, sql_generation (text-to-sql, sql)" + + +def test_served_tasks_and_frontier_models_read_the_live_list( + transport: httpx.MockTransport, frontier_transport: httpx.MockTransport +) -> None: + served = served_tasks(**KEY, transport=transport) + assert [s.task for s in served] == ["classification", "sql_generation"] + assert frontier_models(**KEY, transport=transport) == [] + assert frontier_models(**KEY, transport=frontier_transport) == ["deepseek-v4-pro"] + + +def test_load_resolves_names_against_the_catalogue(transport: httpx.MockTransport) -> None: + for name in ("text-to-sql", "SQL", "sql_generation"): + engine = Engine.load([name], **KEY, transport=transport) + assert engine.tasks == {"sql_generation": "mt/sql_gen_s1"} + engine.close() + with pytest.raises(TaskUnavailable, match="text-to-sql"): + Engine.load(["summarisation"], **KEY, transport=transport) diff --git a/tests/test_engine.py b/tests/test_engine.py index c00998c..ffdfe5d 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -1,95 +1,53 @@ from __future__ import annotations -import json - import httpx import pytest -from microtensor import Engine, EngineError, TaskUnavailable, frontier_models, served_tasks - -MODELS = { - "data": [ - { - "id": "mt/auto", - "tasks": [{"task": "sql_generation", "model": "mt/sql_gen_s1", "live": True}], - }, - { - "id": "mt/sql_gen_s1", - "task": "sql_generation", - "title": "SQL generation, season 1", - "aliases": ["text-to-sql", "sql"], - "quality": 0.73, - "capacity": {"operators": 2, "live": True}, - }, - { - "id": "mt/classify_s1", - "task": "classification", - "title": "Classification, season 1", - "capacity": {"operators": 0, "live": False}, - }, - ] -} +from microtensor import Engine, EngineError, TaskUnavailable +KEY = {"api_key": "mt_sk_x", "base_url": "http://g/v1"} +SCHEMA = "CREATE TABLE orders(x)" -def gateway(request: httpx.Request) -> httpx.Response: - if request.url.path.endswith("/models"): - return httpx.Response(200, json=MODELS) - body = json.loads(request.content) - confident = "2025" in body["messages"][-1]["content"] - return httpx.Response( - 200, - json={ - "id": "cmpl_1", - "model": body["model"], - "choices": [{"message": {"content": "SELECT count(*) FROM orders"}}], - "routing": { - "task": "sql_generation", - "specialist": "mt/sql_gen_s1", - "answer_confidence": 0.96 if confident else 0.4, - "floor_met": confident, - "quality_floor": body["quality_floor"], - }, - "usage": {"total_tokens": 49}, - }, - ) +def test_hosted_engine_reports_the_floor_and_pins_a_loaded_task( + transport: httpx.MockTransport, +) -> None: + with Engine.load(["text-to-sql"], **KEY, transport=transport) as engine: + sure = engine.run("How many orders in 2025?", context=SCHEMA) + assert sure.answer.startswith("SELECT") and sure.task == "sql_generation" + assert sure.confidence == 0.96 and not sure.below_floor and not sure.escalated -def test_task_names_come_from_the_live_catalogue() -> None: - transport = httpx.MockTransport(gateway) - served = served_tasks(api_key="mt_sk_x", base_url="http://g/v1", transport=transport) - assert [s.task for s in served] == ["classification", "sql_generation"] - assert served[1].aliases == ("text-to-sql", "sql") and served[1].live - assert not served[0].live and served[0].label == "classification" + unsure = engine.run("How many orders?", context=SCHEMA) + assert unsure.below_floor and not unsure.escalated and unsure.escalated_to is None - for name in ("text-to-sql", "SQL", "sql_generation", "sql-generation", "sql_gen_s1"): - engine = Engine.load([name], api_key="mt_sk_x", base_url="http://g/v1", transport=transport) - assert engine.tasks == {"sql_generation": "mt/sql_gen_s1"} - engine.close() + pinned = engine.run("How many orders in 2025?", context=SCHEMA, task="text-to-sql") + assert pinned.raw["model"] == "mt/sql_gen_s1" + with pytest.raises(TaskUnavailable): + engine.run("Is this spam?", task="classification") - with pytest.raises(TaskUnavailable, match="text-to-sql"): - Engine.load( - ["summarisation"], api_key="mt_sk_x", base_url="http://g/v1", transport=transport - ) - assert frontier_models(api_key="mt_sk_x", base_url="http://g/v1", transport=transport) == [] +def test_escalation_uses_a_served_frontier_model_within_budget( + frontier_transport: httpx.MockTransport, +) -> None: + with Engine.load( + ["text-to-sql"], + escalation_model={"text-to-sql": "deepseek-v4-pro"}, + escalation_budget={"usd_per_day": 1}, + **KEY, + transport=frontier_transport, + ) as engine: + result = engine.run("How many orders?", context=SCHEMA) + assert result.escalated and result.escalated_to == "deepseek-v4-pro" + assert result.task == "sql_generation" and not result.below_floor + assert result.cost == 0.002 and engine.budget.spent == 0.002 -def test_hosted_engine_reports_the_floor_and_pins_a_loaded_task() -> None: - transport = httpx.MockTransport(gateway) - engine = Engine.load( - ["text-to-sql"], api_key="mt_sk_x", base_url="http://g/v1", transport=transport - ) - sure = engine.run("How many orders in 2025?", context="CREATE TABLE orders(x)") - assert sure.answer.startswith("SELECT") and sure.task == "sql_generation" - assert sure.confidence == 0.96 and not sure.below_floor and not sure.escalated + with Engine.load(["text-to-sql"], escalation=False, **KEY, transport=frontier_transport) as off: + assert off.run("How many orders?", context=SCHEMA).below_floor - unsure = engine.run("How many orders?", context="CREATE TABLE orders(x)") - assert unsure.below_floor and not unsure.escalated and unsure.escalated_to is None - - pinned = engine.run("How many orders in 2025?", context="x", task="text-to-sql") - assert pinned.raw["model"] == "mt/sql_gen_s1" - with pytest.raises(TaskUnavailable): - engine.run("Is this spam?", task="classification") - engine.close() + with pytest.raises(EngineError, match="frontier"): + Engine.load( + ["text-to-sql"], escalation_model="gpt-5", **KEY, transport=frontier_transport + ).run("How many orders?", context=SCHEMA) def test_the_local_backend_is_named_as_a_preview_and_a_key_is_required() -> None: