dbx-tools-model 0.6.78__tar.gz

dbx_tools_model-0.6.78/PKG-INFO

@@ -0,0 +1,60 @@
+Metadata-Version: 2.3
+Name: dbx-tools-model
+Version: 0.6.78
+Summary: Databricks Model Serving invocation, classification, and endpoint resolution
+Requires-Dist: databricks-sdk>=0.63.0,<1
+Requires-Dist: pydantic>=2.9,<3
+Requires-Python: >=3.10
+Project-URL: Source, https://github.com/reggie-db/dbx-tools/tree/main/packages/py/model
+Description-Content-Type: text/markdown
+
+# `dbx-tools-model`
+
+Python contracts and runtime helpers for Databricks Model Serving. This package
+mirrors the reusable parts of `@dbx-tools/shared-model` and `@dbx-tools/model`
+without AppKit cache or Mastra dependencies.
+
+Install directly from this monorepo:
+
+```bash
+pip install "dbx-tools-model @ git+https://github.com/reggie-db/dbx-tools.git@main#subdirectory=packages/py/model"
+```
+
+Key features:
+
+- stable Pydantic endpoint, profile, query, and ranked-result models;
+- live endpoint listing through a structural `WorkspaceClient` protocol;
+- score-driven model classification with family fallbacks;
+- exact and fuzzy endpoint resolution with deterministic class ordering;
+- Databricks invocation URL and per-request authentication helpers;
+- OpenAI chat request sanitization and content extraction;
+- embedding vector extraction with optional dimension validation.
+
+```python
+from databricks.sdk import WorkspaceClient
+from dbx_tools.model import ModelClass, list_serving_endpoints, resolve_model
+
+endpoints = list_serving_endpoints(WorkspaceClient())
+selection = resolve_model(endpoints, model_class=ModelClass.CHAT_BALANCED)
+print(selection.model_id)
+```
+
+The Python port intentionally omits AppKit `CacheManager` integration,
+Mastra-specific adapters, and browser-only schemas. Callers can cache the plain
+Pydantic results with their preferred Python cache.
+
+## Relationship to the Databricks SDK
+
+Use the native SDK directly when an endpoint name is already known and its typed
+query method fits the request. Use this package when endpoint choice, stable
+cross-runtime models, OpenAI-shaped HTTP invocation, or provider-neutral chat
+and embedding normalization is the repetitive part.
+
+## Module map
+
+- `models` — Pydantic wire contracts;
+- `classify`, `classes`, `fallback` — model taxonomy and ordering;
+- `resolve` — exact/fuzzy ranking and single-model selection;
+- `serving` — structural `WorkspaceClient` endpoint listing;
+- `invoke` — URLs, SDK authentication headers, and JSON POST helpers;
+- `chat`, `embedding` — request sanitization and response extraction.

dbx_tools_model-0.6.78/README.md

@@ -0,0 +1,50 @@
+# `dbx-tools-model`
+
+Python contracts and runtime helpers for Databricks Model Serving. This package
+mirrors the reusable parts of `@dbx-tools/shared-model` and `@dbx-tools/model`
+without AppKit cache or Mastra dependencies.
+
+Install directly from this monorepo:
+
+```bash
+pip install "dbx-tools-model @ git+https://github.com/reggie-db/dbx-tools.git@main#subdirectory=packages/py/model"
+```
+
+Key features:
+
+- stable Pydantic endpoint, profile, query, and ranked-result models;
+- live endpoint listing through a structural `WorkspaceClient` protocol;
+- score-driven model classification with family fallbacks;
+- exact and fuzzy endpoint resolution with deterministic class ordering;
+- Databricks invocation URL and per-request authentication helpers;
+- OpenAI chat request sanitization and content extraction;
+- embedding vector extraction with optional dimension validation.
+
+```python
+from databricks.sdk import WorkspaceClient
+from dbx_tools.model import ModelClass, list_serving_endpoints, resolve_model
+
+endpoints = list_serving_endpoints(WorkspaceClient())
+selection = resolve_model(endpoints, model_class=ModelClass.CHAT_BALANCED)
+print(selection.model_id)
+```
+
+The Python port intentionally omits AppKit `CacheManager` integration,
+Mastra-specific adapters, and browser-only schemas. Callers can cache the plain
+Pydantic results with their preferred Python cache.
+
+## Relationship to the Databricks SDK
+
+Use the native SDK directly when an endpoint name is already known and its typed
+query method fits the request. Use this package when endpoint choice, stable
+cross-runtime models, OpenAI-shaped HTTP invocation, or provider-neutral chat
+and embedding normalization is the repetitive part.
+
+## Module map
+
+- `models` — Pydantic wire contracts;
+- `classify`, `classes`, `fallback` — model taxonomy and ordering;
+- `resolve` — exact/fuzzy ranking and single-model selection;
+- `serving` — structural `WorkspaceClient` endpoint listing;
+- `invoke` — URLs, SDK authentication headers, and JSON POST helpers;
+- `chat`, `embedding` — request sanitization and response extraction.

dbx_tools_model-0.6.78/pyproject.toml

@@ -0,0 +1,22 @@
+[project]
+name = "dbx-tools-model"
+version = "0.6.78"
+description = "Databricks Model Serving invocation, classification, and endpoint resolution"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "databricks-sdk>=0.63.0,<1",
+    "pydantic>=2.9,<3",
+]
+
+[project.urls]
+Source = "https://github.com/reggie-db/dbx-tools/tree/main/packages/py/model"
+
+[build-system]
+requires = ["uv_build>=0.11.28,<0.12.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "dbx_tools.model"
+module-root = "src"
+namespace = true

dbx_tools_model-0.6.78/pyproject.toml.orig

@@ -0,0 +1,21 @@
+# ~~ Generated by projen. To modify, edit .projenrc.js and run "bunx projen".
+
+[project]
+name = "dbx-tools-model"
+version = "0.6.78"
+description = "Databricks Model Serving invocation, classification, and endpoint resolution"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [ "databricks-sdk>=0.63.0,<1", "pydantic>=2.9,<3" ]
+
+  [project.urls]
+  Source = "https://github.com/reggie-db/dbx-tools/tree/main/packages/py/model"
+
+[build-system]
+requires = [ "uv_build>=0.11.28,<0.12.0" ]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "dbx_tools.model"
+module-root = "src"
+namespace = true

dbx_tools_model-0.6.78/src/dbx_tools/model/__init__.py

@@ -0,0 +1,113 @@
+from .chat import (
+    UNSUPPORTED_CHAT_FIELDS,
+    chat_content_to_text,
+    sanitize_chat_request,
+    strip_unsupported_chat_fields,
+)
+from .classes import (
+    CHAT_CLASS_ORDER,
+    MODEL_CLASS_ORDER,
+    classes_at_or_below,
+    is_chat_class,
+    parse_model_class,
+)
+from .classify import (
+    CHAT_TASK,
+    EMBEDDING_TASK,
+    classify_by_family,
+    classify_endpoints,
+    endpoint_capabilities,
+    supports_tools_by_family,
+    version_tuple,
+)
+from .embedding import embedding_dimension, extract_embedding, extract_embeddings
+from .fallback import FALLBACK_MODEL_IDS, model_for_class, models_for_class
+from .invoke import (
+    CHAT_COMPLETIONS_PATH,
+    INVOCATIONS_SUFFIX,
+    OPEN_RESPONSES_PATH,
+    RESPONSES_PATH,
+    auth_headers,
+    chat_completions_url,
+    invocations_url,
+    invoke_json,
+    is_responses_only,
+    open_responses_url,
+    post_json,
+    responses_upstream_url,
+    responses_url,
+)
+from .models import (
+    EndpointCapabilities,
+    ModelClass,
+    ModelProfile,
+    ModelQuery,
+    RankedModel,
+    ResolvedModel,
+    ResolvedModelSelection,
+    ServingEndpointSummary,
+)
+from .resolve import (
+    DEFAULT_FUZZY_THRESHOLD,
+    rank_model_id,
+    rank_models,
+    resolve_model,
+    resolve_model_id,
+    search_serving_endpoints,
+)
+from .serving import list_serving_endpoints, list_serving_endpoints_uncached, to_model_display_name
+
+__all__ = [
+    "CHAT_CLASS_ORDER",
+    "CHAT_COMPLETIONS_PATH",
+    "CHAT_TASK",
+    "DEFAULT_FUZZY_THRESHOLD",
+    "EMBEDDING_TASK",
+    "FALLBACK_MODEL_IDS",
+    "INVOCATIONS_SUFFIX",
+    "MODEL_CLASS_ORDER",
+    "OPEN_RESPONSES_PATH",
+    "RESPONSES_PATH",
+    "UNSUPPORTED_CHAT_FIELDS",
+    "EndpointCapabilities",
+    "ModelClass",
+    "ModelProfile",
+    "ModelQuery",
+    "RankedModel",
+    "ResolvedModel",
+    "ResolvedModelSelection",
+    "ServingEndpointSummary",
+    "auth_headers",
+    "chat_completions_url",
+    "chat_content_to_text",
+    "classes_at_or_below",
+    "classify_by_family",
+    "classify_endpoints",
+    "embedding_dimension",
+    "endpoint_capabilities",
+    "extract_embedding",
+    "extract_embeddings",
+    "invocations_url",
+    "invoke_json",
+    "is_chat_class",
+    "is_responses_only",
+    "list_serving_endpoints",
+    "list_serving_endpoints_uncached",
+    "model_for_class",
+    "models_for_class",
+    "open_responses_url",
+    "parse_model_class",
+    "post_json",
+    "rank_model_id",
+    "rank_models",
+    "resolve_model",
+    "resolve_model_id",
+    "responses_upstream_url",
+    "responses_url",
+    "sanitize_chat_request",
+    "search_serving_endpoints",
+    "strip_unsupported_chat_fields",
+    "supports_tools_by_family",
+    "to_model_display_name",
+    "version_tuple",
+]

dbx_tools_model-0.6.78/src/dbx_tools/model/chat.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+UNSUPPORTED_CHAT_FIELDS = (
+    "parallel_tool_calls",
+    "store",
+    "metadata",
+    "service_tier",
+    "prompt_cache_key",
+    "safety_identifier",
+)
+
+
+def strip_unsupported_chat_fields(body: dict[str, object], extra: Iterable[str] = ()) -> list[str]:
+    dropped = []
+    for field in (*UNSUPPORTED_CHAT_FIELDS, *extra):
+        if field in body:
+            del body[field]
+            dropped.append(field)
+    return dropped
+
+
+def sanitize_chat_request(body: dict[str, object], extra: Iterable[str] = ()) -> dict[str, object]:
+    sanitized = dict(body)
+    strip_unsupported_chat_fields(sanitized, extra)
+    return sanitized
+
+
+def chat_content_to_text(
+    content: object,
+    options: dict[str, object] | None = None,
+    *,
+    separator: str = "",
+    types: Iterable[str] | None = None,
+) -> str:
+    if options:
+        separator = str(options.get("separator", separator))
+        option_types = options.get("types")
+        if isinstance(option_types, list) and all(isinstance(value, str) for value in option_types):
+            types = option_types
+    if isinstance(content, str):
+        return content
+    if not isinstance(content, list):
+        return ""
+    accepted = set(types) if types is not None else None
+    parts = []
+    for part in content:
+        if not isinstance(part, dict):
+            continue
+        if accepted is not None and part.get("type") not in accepted:
+            continue
+        text = part.get("text")
+        if isinstance(text, str):
+            parts.append(text)
+    return separator.join(parts)

dbx_tools_model-0.6.78/src/dbx_tools/model/classes.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from .models import ModelClass
+
+CHAT_CLASS_ORDER = (
+    ModelClass.CHAT_THINKING,
+    ModelClass.CHAT_BALANCED,
+    ModelClass.CHAT_FAST,
+)
+MODEL_CLASS_ORDER = (*CHAT_CLASS_ORDER, ModelClass.EMBEDDING)
+
+
+def is_chat_class(model_class: ModelClass) -> bool:
+    return model_class in CHAT_CLASS_ORDER
+
+
+def parse_model_class(value: object) -> ModelClass | None:
+    try:
+        return ModelClass(value)
+    except (TypeError, ValueError):
+        try:
+            return ModelClass(f"chat-{value}")
+        except (TypeError, ValueError):
+            return None
+
+
+def classes_at_or_below(model_class: ModelClass) -> list[ModelClass]:
+    if model_class == ModelClass.EMBEDDING:
+        return [ModelClass.EMBEDDING]
+    try:
+        return list(CHAT_CLASS_ORDER[CHAT_CLASS_ORDER.index(model_class) :])
+    except ValueError:
+        return list(CHAT_CLASS_ORDER)

dbx_tools_model-0.6.78/src/dbx_tools/model/classify.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import math
+import re
+from collections.abc import Iterable
+
+from .models import EndpointCapabilities, ModelClass, ServingEndpointSummary
+
+CHAT_TASK = "llm/v1/chat"
+EMBEDDING_TASK = "llm/v1/embeddings"
+
+
+def supports_tools_by_family(name: str) -> bool:
+    normalized = name.lower()
+    if "gemini" in normalized or "gpt-oss" in normalized:
+        return False
+    return any(family in normalized for family in ("claude", "gpt", "qwen", "glm", "llama"))
+
+
+def endpoint_capabilities(
+    endpoint: ServingEndpointSummary | dict[str, object],
+) -> EndpointCapabilities:
+    summary = _summary(endpoint)
+    embedding = summary.task == EMBEDDING_TASK or summary.model_class == ModelClass.EMBEDDING
+    chat = not embedding and (summary.task == CHAT_TASK or summary.model_class is not None)
+    tools = chat and (
+        summary.supports_tools
+        if summary.supports_tools is not None
+        else supports_tools_by_family(summary.name)
+    )
+    return EndpointCapabilities(chat=chat, embedding=embedding, tools=tools)
+
+
+def version_tuple(name: str) -> list[int]:
+    match = re.search(r"\d", name)
+    if match is None:
+        return [0, 0, 0]
+    numbers = []
+    for chunk in re.split(r"[^a-z0-9]+", name[match.start() :], flags=re.IGNORECASE):
+        digits = re.match(r"^\d+", chunk)
+        if digits:
+            numbers.append(int(digits.group(0)))
+    return (numbers + [0, 0, 0])[:3]
+
+
+def classify_by_family(name: str) -> dict[str, object] | None:
+    normalized = name.lower()
+
+    def result(model_class: ModelClass) -> dict[str, object]:
+        major, minor, patch = version_tuple(normalized)
+        return {"class": model_class.value, "rank": major * 1_000_000 + minor * 1_000 + patch}
+
+    if "opus" in normalized:
+        return result(ModelClass.CHAT_THINKING)
+    if "sonnet" in normalized:
+        return result(ModelClass.CHAT_BALANCED)
+    if "haiku" in normalized:
+        return result(ModelClass.CHAT_FAST)
+    if "gpt-oss" in normalized:
+        return result(ModelClass.CHAT_BALANCED if "120b" in normalized else ModelClass.CHAT_FAST)
+    if "gpt" in normalized:
+        if "pro" in normalized:
+            return result(ModelClass.CHAT_THINKING)
+        if "mini" in normalized or "nano" in normalized:
+            return result(ModelClass.CHAT_FAST)
+        return result(ModelClass.CHAT_BALANCED)
+    if "gemini" in normalized:
+        if "flash-lite" in normalized:
+            return result(ModelClass.CHAT_FAST)
+        if "pro" in normalized:
+            return result(ModelClass.CHAT_THINKING)
+        return result(ModelClass.CHAT_BALANCED)
+    if "gemma" in normalized:
+        return result(ModelClass.CHAT_FAST)
+    if "llama" in normalized:
+        if "maverick" in normalized or "405b" in normalized:
+            return result(ModelClass.CHAT_THINKING)
+        if "70b" in normalized:
+            return result(ModelClass.CHAT_BALANCED)
+        if "8b" in normalized or "1b" in normalized:
+            return result(ModelClass.CHAT_FAST)
+        return result(ModelClass.CHAT_BALANCED)
+    if "qwen" in normalized:
+        return result(ModelClass.CHAT_BALANCED)
+    return None
+
+
+def classify_endpoints(
+    endpoints: Iterable[ServingEndpointSummary | dict[str, object]],
+) -> dict[str, list[dict[str, object]]]:
+    summaries = [_summary(endpoint) for endpoint in endpoints]
+    chat = [endpoint for endpoint in summaries if endpoint.task == CHAT_TASK]
+    qualities = sorted(
+        endpoint.profile.quality
+        for endpoint in chat
+        if endpoint.profile is not None
+        and endpoint.profile.quality is not None
+        and math.isfinite(endpoint.profile.quality)
+    )
+    low = _quantile(qualities, 1 / 3)
+    high = _quantile(qualities, 2 / 3)
+    buckets: dict[ModelClass, list[tuple[tuple[object, ...], ServingEndpointSummary]]] = {
+        model_class: [] for model_class in ModelClass
+    }
+
+    for endpoint in chat:
+        quality = endpoint.profile.quality if endpoint.profile else None
+        if quality is not None and math.isfinite(quality):
+            model_class = (
+                ModelClass.CHAT_THINKING
+                if quality >= high
+                else ModelClass.CHAT_FAST
+                if quality <= low
+                else ModelClass.CHAT_BALANCED
+            )
+            cost = (
+                endpoint.profile.cost
+                if endpoint.profile and endpoint.profile.cost is not None
+                else math.inf
+            )
+            speed = (
+                endpoint.profile.speed
+                if endpoint.profile and endpoint.profile.speed is not None
+                else 0
+            )
+            key = (0, -quality, cost, -speed, *[-part for part in version_tuple(endpoint.name)])
+        else:
+            family = classify_by_family(endpoint.name)
+            if family is None:
+                continue
+            model_class = ModelClass(str(family["class"]))
+            key = (
+                1,
+                -int(family["rank"]),
+                math.inf,
+                0,
+                *[-part for part in version_tuple(endpoint.name)],
+            )
+        buckets[model_class].append((key, endpoint))
+
+    embeddings = [endpoint for endpoint in summaries if endpoint.task == EMBEDDING_TASK]
+    return {
+        ModelClass.CHAT_THINKING.value: _dump_sorted(buckets[ModelClass.CHAT_THINKING]),
+        ModelClass.CHAT_BALANCED.value: _dump_sorted(buckets[ModelClass.CHAT_BALANCED]),
+        ModelClass.CHAT_FAST.value: _dump_sorted(buckets[ModelClass.CHAT_FAST]),
+        ModelClass.EMBEDDING.value: [endpoint.as_dict() for endpoint in embeddings],
+    }
+
+
+def classified_summaries(
+    endpoints: Iterable[ServingEndpointSummary | dict[str, object]],
+) -> dict[ModelClass, list[ServingEndpointSummary]]:
+    return {
+        ModelClass(key): [ServingEndpointSummary.model_validate(value) for value in values]
+        for key, values in classify_endpoints(endpoints).items()
+    }
+
+
+def _summary(value: ServingEndpointSummary | dict[str, object]) -> ServingEndpointSummary:
+    return (
+        value
+        if isinstance(value, ServingEndpointSummary)
+        else ServingEndpointSummary.model_validate(value)
+    )
+
+
+def _quantile(values: list[float], probability: float) -> float:
+    if not values:
+        return math.nan
+    index = (len(values) - 1) * probability
+    lower = math.floor(index)
+    upper = math.ceil(index)
+    if lower == upper:
+        return values[lower]
+    return values[lower] + (values[upper] - values[lower]) * (index - lower)
+
+
+def _dump_sorted(
+    values: list[tuple[tuple[object, ...], ServingEndpointSummary]],
+) -> list[dict[str, object]]:
+    return [endpoint.as_dict() for _, endpoint in sorted(values, key=lambda value: value[0])]

dbx_tools_model-0.6.78/src/dbx_tools/model/embedding.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+
+def extract_embedding(response: object, expected_dimension: int | None = None) -> list[float]:
+    if not isinstance(response, dict):
+        raise TypeError("Embedding response must be an object")
+    data = response.get("data")
+    if not isinstance(data, list) or not data or not isinstance(data[0], dict):
+        raise ValueError("Embedding response contains no data")
+    vector = data[0].get("embedding")
+    if not isinstance(vector, list) or not all(isinstance(value, (int, float)) for value in vector):
+        raise ValueError("Embedding response contains no numeric vector")
+    result = [float(value) for value in vector]
+    if expected_dimension is not None and len(result) != expected_dimension:
+        raise ValueError(
+            f"Expected embedding dimension {expected_dimension}, received {len(result)}"
+        )
+    return result
+
+
+def extract_embeddings(
+    response: object, expected_dimension: int | None = None
+) -> list[list[float]]:
+    if not isinstance(response, dict) or not isinstance(response.get("data"), list):
+        raise TypeError("Embedding response contains no data")
+    return [
+        extract_embedding({"data": [item]}, expected_dimension)
+        for item in response["data"]
+        if isinstance(item, dict)
+    ]
+
+
+def embedding_dimension(vector: Iterable[float]) -> int:
+    return len(list(vector))

dbx_tools_model-0.6.78/src/dbx_tools/model/fallback.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from .classify import classify_by_family
+from .models import ModelClass
+
+_FALLBACK_MODEL_NAMES = (
+    "databricks-claude-opus-4-8",
+    "databricks-gpt-5-5-pro",
+    "databricks-gemini-3-1-pro",
+    "databricks-claude-sonnet-4-6",
+    "databricks-gpt-5-5",
+    "databricks-meta-llama-3-3-70b-instruct",
+    "databricks-claude-haiku-4-5",
+    "databricks-gpt-5-nano",
+    "databricks-meta-llama-3-1-8b-instruct",
+)
+
+
+def models_for_class(model_class: ModelClass) -> list[str]:
+    classified = []
+    for name in _FALLBACK_MODEL_NAMES:
+        family = classify_by_family(name)
+        if family is not None and family["class"] == model_class.value:
+            classified.append((int(family["rank"]), name))
+    return [name for _, name in sorted(classified, reverse=True)]
+
+
+def model_for_class(model_class: ModelClass) -> str:
+    return models_for_class(model_class)[0]
+
+
+FALLBACK_MODEL_IDS = tuple(
+    model
+    for model_class in (
+        ModelClass.CHAT_THINKING,
+        ModelClass.CHAT_BALANCED,
+        ModelClass.CHAT_FAST,
+    )
+    for model in models_for_class(model_class)
+)

dbx_tools_model-0.6.78/src/dbx_tools/model/invoke.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import json
+import re
+from collections.abc import Callable, Mapping
+from typing import Any, Protocol
+from urllib.parse import quote, urljoin
+from urllib.request import Request, urlopen
+
+INVOCATIONS_SUFFIX = "invocations"
+RESPONSES_PATH = "serving-endpoints/responses"
+OPEN_RESPONSES_PATH = "serving-endpoints/open-responses"
+CHAT_COMPLETIONS_PATH = "serving-endpoints/chat/completions"
+
+
+class AuthenticatingConfig(Protocol):
+    def authenticate(self) -> Mapping[str, str]: ...
+
+
+class AuthenticatingClientLike(Protocol):
+    config: AuthenticatingConfig
+
+
+def invocations_url(host: str, endpoint: str) -> str:
+    return urljoin(host, f"serving-endpoints/{quote(endpoint, safe='')}/{INVOCATIONS_SUFFIX}")
+
+
+def responses_url(host: str) -> str:
+    return urljoin(host, RESPONSES_PATH)
+
+
+def open_responses_url(host: str) -> str:
+    return urljoin(host, OPEN_RESPONSES_PATH)
+
+
+def chat_completions_url(host: str) -> str:
+    return urljoin(host, CHAT_COMPLETIONS_PATH)
+
+
+def is_responses_only(endpoint: str) -> bool:
+    return "codex" in endpoint.lower()
+
+
+def responses_upstream_url(host: str, endpoint: str) -> str:
+    normalized = endpoint.lower()
+    openai_family = (
+        "gpt" in normalized
+        or "codex" in normalized
+        or re.search(r"(^|[^a-z])o[1-9]([^a-z]|$)", normalized) is not None
+        or "openai" in normalized
+    )
+    return responses_url(host) if openai_family else open_responses_url(host)
+
+
+def auth_headers(client: AuthenticatingClientLike) -> dict[str, str]:
+    return dict(client.config.authenticate())
+
+
+def post_json(
+    client: AuthenticatingClientLike,
+    url: str,
+    body: Mapping[str, object],
+    *,
+    timeout: float | None = None,
+    opener: Callable[..., Any] = urlopen,
+) -> object:
+    headers = {"Accept": "application/json", "Content-Type": "application/json"}
+    headers.update(auth_headers(client))
+    request = Request(url, data=json.dumps(body).encode(), headers=headers, method="POST")
+    with opener(request, timeout=timeout) as response:
+        return json.loads(response.read())
+
+
+def invoke_json(
+    client: AuthenticatingClientLike,
+    host: str,
+    endpoint: str,
+    body: Mapping[str, object],
+    **options: object,
+) -> object:
+    return post_json(client, invocations_url(host, endpoint), body, **options)

dbx_tools_model-0.6.78/src/dbx_tools/model/models.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+from enum import Enum
+
+from pydantic import BaseModel, ConfigDict, Field
+
+
+class ModelClass(str, Enum):
+    CHAT_THINKING = "chat-thinking"
+    CHAT_BALANCED = "chat-balanced"
+    CHAT_FAST = "chat-fast"
+    EMBEDDING = "embedding"
+
+
+class WireModel(BaseModel):
+    model_config = ConfigDict(populate_by_name=True, extra="ignore")
+
+    def as_dict(self) -> dict[str, object]:
+        return self.model_dump(by_alias=True, exclude_none=True, mode="json")
+
+
+class ModelProfile(WireModel):
+    quality: float | None = None
+    speed: float | None = None
+    cost: float | None = None
+
+
+class ServingEndpointSummary(WireModel):
+    name: str
+    display_name: str | None = Field(default=None, alias="displayName")
+    task: str | None = None
+    state: str | None = None
+    description: str | None = None
+    supports_tools: bool | None = Field(default=None, alias="supportsTools")
+    profile: ModelProfile | None = None
+    model_class: ModelClass | None = Field(default=None, alias="class")
+    dimension: int | None = None
+
+
+class ModelQuery(WireModel):
+    search: str | None = None
+    model_class: ModelClass | None = Field(default=None, alias="modelClass")
+    requires_tools: bool | None = Field(default=None, alias="requiresTools")
+    limit: int | None = None
+    threshold: float | None = None
+
+
+class RankedModel(WireModel):
+    endpoint: ServingEndpointSummary
+    model_class: ModelClass = Field(alias="modelClass")
+    score: float | None = None
+
+
+class ResolvedModel(WireModel):
+    model_id: str = Field(alias="modelId")
+    matched: bool
+    score: float | None = None
+
+
+class ResolvedModelSelection(WireModel):
+    model_id: str = Field(alias="modelId")
+    source: str
+
+
+class EndpointCapabilities(WireModel):
+    chat: bool
+    embedding: bool
+    tools: bool

dbx_tools_model-0.6.78/src/dbx_tools/model/resolve.py

@@ -0,0 +1,218 @@
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable, Sequence
+from difflib import SequenceMatcher
+
+from .classes import CHAT_CLASS_ORDER, MODEL_CLASS_ORDER, classes_at_or_below
+from .classify import classified_summaries, endpoint_capabilities
+from .fallback import FALLBACK_MODEL_IDS, models_for_class
+from .models import (
+    ModelClass,
+    ModelQuery,
+    RankedModel,
+    ResolvedModel,
+    ResolvedModelSelection,
+    ServingEndpointSummary,
+)
+
+DEFAULT_FUZZY_THRESHOLD = 0.4
+
+
+def search_serving_endpoints(
+    input_value: str,
+    endpoints: Iterable[ServingEndpointSummary | dict[str, object]],
+    threshold: float = DEFAULT_FUZZY_THRESHOLD,
+) -> list[dict[str, object]]:
+    summaries = [_summary(endpoint) for endpoint in endpoints]
+    for endpoint in summaries:
+        if endpoint.name == input_value:
+            return [{"endpoint": endpoint.as_dict(), "score": 0}]
+    tokens = re.findall(r"[a-zA-Z0-9]+", input_value.lower())
+    if not tokens:
+        return []
+    matches = []
+    for endpoint in summaries:
+        name = endpoint.name.lower()
+        token_scores = [_token_distance(token, name) for token in tokens]
+        score = sum(token_scores) / len(token_scores)
+        if score <= threshold:
+            matches.append({"endpoint": endpoint.as_dict(), "score": score})
+    return sorted(matches, key=lambda match: (float(match["score"]), match["endpoint"]["name"]))
+
+
+def resolve_model_id(
+    input_value: str,
+    endpoints: Iterable[ServingEndpointSummary | dict[str, object]],
+    *,
+    threshold: float = DEFAULT_FUZZY_THRESHOLD,
+    requires_tools: bool = False,
+) -> ResolvedModel:
+    candidates = [
+        _summary(endpoint)
+        for endpoint in endpoints
+        if not requires_tools or endpoint_capabilities(_summary(endpoint)).tools
+    ]
+    matches = search_serving_endpoints(input_value, candidates, threshold)
+    if not matches:
+        return ResolvedModel(modelId=input_value, matched=False)
+    top = matches[0]
+    return ResolvedModel(
+        modelId=str(top["endpoint"]["name"]), matched=True, score=float(top["score"])
+    )
+
+
+def rank_models(
+    endpoints: Iterable[ServingEndpointSummary | dict[str, object]],
+    query: ModelQuery | dict[str, object] | None = None,
+) -> list[dict[str, object]]:
+    summaries = [_summary(endpoint) for endpoint in endpoints]
+    request = query if isinstance(query, ModelQuery) else ModelQuery.model_validate(query or {})
+    classified = classified_summaries(summaries)
+    eligible = (
+        classes_at_or_below(request.model_class)
+        if request.model_class is not None
+        else list(CHAT_CLASS_ORDER)
+    )
+    candidates: list[RankedModel] = []
+    for model_class in eligible:
+        for endpoint in classified[model_class]:
+            if request.requires_tools and not endpoint_capabilities(endpoint).tools:
+                continue
+            candidates.append(RankedModel(endpoint=endpoint, modelClass=model_class))
+
+    search = request.search.strip() if request.search else ""
+    if search:
+        scores = {
+            str(match["endpoint"]["name"]): float(match["score"])
+            for match in search_serving_endpoints(
+                search,
+                [candidate.endpoint for candidate in candidates],
+                request.threshold if request.threshold is not None else DEFAULT_FUZZY_THRESHOLD,
+            )
+        }
+        candidates = [candidate for candidate in candidates if candidate.endpoint.name in scores]
+        for candidate in candidates:
+            candidate.score = scores[candidate.endpoint.name]
+        candidates.sort(
+            key=lambda candidate: (
+                round((candidate.score or 0) * 1000),
+                MODEL_CLASS_ORDER.index(candidate.model_class),
+            )
+        )
+    if request.limit is not None:
+        candidates = candidates[: max(0, request.limit)]
+    return [candidate.as_dict() for candidate in candidates]
+
+
+def rank_model_id(
+    endpoints: Iterable[ServingEndpointSummary | dict[str, object]],
+    search: str,
+    *,
+    threshold: float = DEFAULT_FUZZY_THRESHOLD,
+    requires_tools: bool = False,
+) -> ResolvedModel:
+    ranked = rank_models(
+        endpoints,
+        ModelQuery(
+            search=search,
+            limit=1,
+            threshold=threshold,
+            requiresTools=requires_tools,
+        ),
+    )
+    if not ranked:
+        return ResolvedModel(modelId=search, matched=False)
+    return ResolvedModel(
+        modelId=str(ranked[0]["endpoint"]["name"]),
+        matched=True,
+        score=float(ranked[0]["score"]),
+    )
+
+
+def resolve_model(
+    endpoints: Iterable[ServingEndpointSummary | dict[str, object]],
+    *,
+    explicit: str | None = None,
+    fuzzy: bool = True,
+    threshold: float = DEFAULT_FUZZY_THRESHOLD,
+    requires_tools: bool = False,
+    model_class: ModelClass | None = None,
+    fallbacks: Sequence[str] = (),
+) -> ResolvedModelSelection:
+    summaries = [_summary(endpoint) for endpoint in endpoints]
+    if explicit is not None:
+        if not fuzzy:
+            if requires_tools:
+                _assert_tool_support(summaries, explicit)
+            return ResolvedModelSelection(modelId=explicit, source="explicit")
+        ranked = rank_models(
+            summaries,
+            ModelQuery(
+                search=explicit,
+                modelClass=model_class,
+                requiresTools=requires_tools,
+                threshold=threshold,
+                limit=1,
+            ),
+        )
+        if requires_tools and not ranked:
+            raise ValueError(f'No tool-capable model matches "{explicit}"')
+        model_id = str(ranked[0]["endpoint"]["name"]) if ranked else explicit
+        return ResolvedModelSelection(modelId=model_id, source="fuzzy-match")
+
+    if model_class is None and fallbacks:
+        present = {
+            endpoint.name
+            for endpoint in summaries
+            if not requires_tools or endpoint_capabilities(endpoint).tools
+        }
+        for fallback in fallbacks:
+            if fallback in present:
+                return ResolvedModelSelection(modelId=fallback, source="fallback")
+
+    source = "class" if model_class is not None else "fallback"
+    ranked = rank_models(
+        summaries,
+        ModelQuery(modelClass=model_class, requiresTools=requires_tools, limit=1),
+    )
+    if ranked:
+        return ResolvedModelSelection(modelId=str(ranked[0]["endpoint"]["name"]), source=source)
+
+    floor_source = models_for_class(model_class) if model_class is not None else list(fallbacks)
+    floor = list(dict.fromkeys([*floor_source, *FALLBACK_MODEL_IDS]))
+    if requires_tools:
+        available = {
+            endpoint.name for endpoint in summaries if endpoint_capabilities(endpoint).tools
+        }
+        selected = next((model_id for model_id in floor if model_id in available), None)
+        if selected is None:
+            raise ValueError("No tool-capable model is available")
+        return ResolvedModelSelection(modelId=selected, source=source)
+    present = {endpoint.name for endpoint in summaries}
+    selected = next((model_id for model_id in floor if model_id in present), floor[0])
+    return ResolvedModelSelection(modelId=selected, source=source)
+
+
+def _token_distance(token: str, name: str) -> float:
+    if token in name:
+        return 0.0
+    segments = re.findall(r"[a-z0-9]+", name)
+    similarity = max(
+        (SequenceMatcher(None, token, segment).ratio() for segment in segments), default=0
+    )
+    return 1 - similarity
+
+
+def _assert_tool_support(endpoints: list[ServingEndpointSummary], model_id: str) -> None:
+    endpoint = next((candidate for candidate in endpoints if candidate.name == model_id), None)
+    if endpoint is None or not endpoint_capabilities(endpoint).tools:
+        raise ValueError(f'Model "{model_id}" does not support function tools')
+
+
+def _summary(value: ServingEndpointSummary | dict[str, object]) -> ServingEndpointSummary:
+    return (
+        value
+        if isinstance(value, ServingEndpointSummary)
+        else ServingEndpointSummary.model_validate(value)
+    )

dbx_tools_model-0.6.78/src/dbx_tools/model/serving.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+import math
+import re
+from collections.abc import Iterable, Mapping
+from typing import Any, Protocol
+
+from .classes import MODEL_CLASS_ORDER
+from .classify import classified_summaries, supports_tools_by_family
+from .models import ModelProfile, ServingEndpointSummary
+
+
+class ServingEndpointsApi(Protocol):
+    def list(self) -> Iterable[object]: ...
+
+
+class WorkspaceClientLike(Protocol):
+    serving_endpoints: ServingEndpointsApi
+
+
+def list_serving_endpoints(client: WorkspaceClientLike) -> list[ServingEndpointSummary]:
+    summaries = list_serving_endpoints_uncached(client)
+    buckets = classified_summaries(summaries)
+    classes = {
+        endpoint.name: model_class
+        for model_class in MODEL_CLASS_ORDER
+        for endpoint in buckets[model_class]
+    }
+    for summary in summaries:
+        summary.model_class = classes.get(summary.name)
+    return summaries
+
+
+def list_serving_endpoints_uncached(client: WorkspaceClientLike) -> list[ServingEndpointSummary]:
+    summaries = []
+    for endpoint in client.serving_endpoints.list():
+        name = _value(endpoint, "name")
+        if not isinstance(name, str) or not name:
+            continue
+        summaries.append(
+            ServingEndpointSummary(
+                name=name,
+                displayName=to_model_display_name(name, _provided_display_name(endpoint)),
+                task=_string(_value(endpoint, "task")),
+                state=_string(_value(_value(endpoint, "state"), "ready")),
+                description=_string(_value(endpoint, "description")),
+                supportsTools=supports_tools_by_family(name),
+                profile=_extract_profile(endpoint),
+            )
+        )
+    return summaries
+
+
+def to_model_display_name(name: str, provided: str | None = None) -> str:
+    if provided is not None and provided.strip():
+        return provided.strip()
+    segments = [segment for segment in re.split(r"[-_.\s/]+", name) if segment]
+    start = 0
+    while start < len(segments) and segments[start].lower() in {"databricks", "system", "dbx"}:
+        if (
+            segments[start].lower() == "system"
+            and start + 1 < len(segments)
+            and segments[start + 1].lower() == "ai"
+        ):
+            start += 1
+        start += 1
+    source = segments[start:] or segments
+    acronyms = {"gpt", "gte", "bge", "dbrx", "oss", "llm", "moe", "ai"}
+    pieces = []
+    numeric = []
+    for segment in source:
+        if segment.isdigit():
+            numeric.append(segment)
+            continue
+        if numeric:
+            pieces.append(".".join(numeric))
+            numeric = []
+        size = re.fullmatch(r"(\d+)([bmk])", segment, flags=re.IGNORECASE)
+        if size:
+            pieces.append(f"{size.group(1)}{size.group(2).upper()}")
+        else:
+            lower = segment.lower()
+            pieces.append(lower.upper() if lower in acronyms else lower.capitalize())
+    if numeric:
+        pieces.append(".".join(numeric))
+    return " ".join(pieces) or name.strip()
+
+
+def _extract_profile(endpoint: object) -> ModelProfile | None:
+    config = _value(endpoint, "config")
+    entities = _value(config, "served_entities") or []
+    for entity in entities:
+        foundation = _value(entity, "foundation_model")
+        raw = _value(foundation, "ai_gateway_model_profile")
+        values = {
+            key: value
+            for key in ("quality", "speed", "cost")
+            if isinstance((value := _value(raw, key)), (int, float)) and math.isfinite(value)
+        }
+        if values:
+            return ModelProfile.model_validate(values)
+    return None
+
+
+def _provided_display_name(endpoint: object) -> str | None:
+    for tag in _value(endpoint, "tags") or []:
+        if _value(tag, "key") in {"display_name", "displayName", "name"}:
+            value = _string(_value(tag, "value"))
+            if value and value.strip():
+                return value.strip()
+    config = _value(endpoint, "config")
+    for entity in _value(config, "served_entities") or []:
+        value = _string(_value(_value(entity, "external_model"), "name"))
+        if value and value.strip():
+            return value.strip()
+    return None
+
+
+def _value(value: object, key: str) -> Any:
+    if isinstance(value, Mapping):
+        return value.get(key)
+    return getattr(value, key, None)
+
+
+def _string(value: object) -> str | None:
+    return None if value is None else str(value)