dot-inference 2.0.2__py3-none-any.whl

dot_inference/_cost_tracker.py

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: Kannon For Deep Tech
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+"""Cost and token usage tracker for aggregating across multiple LLM calls."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from dot_inference._types import LLMResponse
+
+
+# Per-token costs are sourced from models.yaml via ResolvedModelConfig
+# and passed in by the provider clients (OpenAI, Mistral, etc.).
+def compute_cost(
+    usage: Any,
+    input_cost_per_token: float | None,
+    output_cost_per_token: float | None,
+) -> float | None:
+    """Compute USD cost from token usage and per-token prices.
+
+    Returns ``None`` when any input is missing so callers can treat
+    the result as "unknown cost" without special-casing.
+    """
+    if usage is None or input_cost_per_token is None or output_cost_per_token is None:
+        return None
+    return (
+        usage.prompt_tokens * input_cost_per_token + usage.completion_tokens * output_cost_per_token
+    )
+
+
+class CostTracker:
+    """Accumulates token usage and cost across multiple ``LLMResponse`` instances.
+
+    Usage::
+
+        tracker = CostTracker()
+        response = await llm(messages)
+        tracker.record(response)
+        # ... more calls ...
+        print(tracker.total_cost, tracker.total_tokens)
+    """
+
+    def __init__(self) -> None:
+        self.total_prompt_tokens: int = 0
+        self.total_completion_tokens: int = 0
+        self.total_cost: float = 0.0
+        self.call_count: int = 0
+
+    @property
+    def total_tokens(self) -> int:
+        """Total tokens (prompt + completion) across all tracked calls."""
+        return self.total_prompt_tokens + self.total_completion_tokens
+
+    def record(self, response: LLMResponse) -> LLMResponse:
+        """Record usage from a response. Returns the response unchanged for chaining."""
+        if response.prompt_tokens is not None:
+            self.total_prompt_tokens += response.prompt_tokens
+        if response.completion_tokens is not None:
+            self.total_completion_tokens += response.completion_tokens
+        if response.cost is not None:
+            self.total_cost += response.cost
+        self.call_count += 1
+        return response
+
+    def reset(self) -> None:
+        """Reset all counters to zero."""
+        self.total_prompt_tokens = 0
+        self.total_completion_tokens = 0
+        self.total_cost = 0.0
+        self.call_count = 0

dot_inference/_llm_registry.py

@@ -0,0 +1,302 @@
+# SPDX-FileCopyrightText: Kannon For Deep Tech
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+"""Model registry — loads model configurations from models.yaml.
+
+The built-in ``models.yaml`` ships with the package.  Users can extend or
+override it by pointing the ``DOTI_CUSTOM_MODELS`` environment variable at
+their own YAML file (same format).  Entries in the custom file are merged on
+top of the built-in registry — new model IDs are added, and existing ones
+are fully replaced.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass, field, replace
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+@dataclass(frozen=True)
+class ProviderModelConfig:
+    """Provider-specific configuration for a model."""
+
+    api_model_name: str
+    input_token_cost: float | None = None
+    output_token_cost: float | None = None
+    # Per-provider override for reasoning support (None = inherit from model).
+    # Useful when reasoning works through one provider (e.g. OpenRouter) but
+    # not another (e.g. native SDK doesn't expose the parameter yet).
+    supports_reasoning: bool | None = None
+    # OpenRouter-specific features (ignored by other providers)
+    response_healing: bool = False
+    fallback_models: list[str] = field(default_factory=list)
+    provider_routing: dict[str, Any] | None = None
+
+
+@dataclass(frozen=True)
+class ModelConfig:
+    """Configuration for a single model from the registry."""
+
+    model_id: str
+    supports_strict_json: bool = True
+    supports_reasoning: bool = False
+    requires_cache_control: bool = False
+    input_token_cost: float | None = None
+    output_token_cost: float | None = None
+    # Maximum tokens (input + output) the model accepts in one request.
+    # Model-level (not per-provider) — public spec value. ``None`` when
+    # the figure is not declared.
+    context_window: int | None = None
+    default_provider: str | None = None
+    # The reasoning_effort values this model accepts. Empty means "not declared" — callers
+    # fall back to the standard ladder via allowed_reasoning_efforts().
+    reasoning_effort_options: list[str] = field(default_factory=list)
+    # The effort used when the caller specifies none; None falls back to the
+    # global GenerationSettings default.
+    default_reasoning_effort: str | None = None
+    providers: dict[str, ProviderModelConfig] = field(default_factory=dict)
+
+    def for_provider(self, provider: str) -> ResolvedModelConfig:
+        """Resolve a flattened config for a specific provider.
+
+        Merges model-level defaults with provider-level overrides.
+        Raises KeyError if the model is not available through this provider.
+        """
+        pmc = self.providers[provider]
+        return ResolvedModelConfig(
+            model_id=self.model_id,
+            provider=provider,
+            api_model_name=pmc.api_model_name,
+            supports_strict_json=self.supports_strict_json,
+            supports_reasoning=pmc.supports_reasoning
+            if pmc.supports_reasoning is not None
+            else self.supports_reasoning,
+            requires_cache_control=self.requires_cache_control,
+            input_token_cost=pmc.input_token_cost
+            if pmc.input_token_cost is not None
+            else self.input_token_cost,
+            output_token_cost=pmc.output_token_cost
+            if pmc.output_token_cost is not None
+            else self.output_token_cost,
+            context_window=self.context_window,
+            response_healing=pmc.response_healing,
+            fallback_models=pmc.fallback_models,
+            provider_routing=pmc.provider_routing,
+        )
+
+
+@dataclass(frozen=True)
+class ResolvedModelConfig:
+    """Fully resolved config for a model+provider pair. Ready for use by callables."""
+
+    model_id: str
+    provider: str
+    api_model_name: str
+    supports_strict_json: bool = True
+    supports_reasoning: bool = False
+    requires_cache_control: bool = False
+    input_token_cost: float | None = None
+    output_token_cost: float | None = None
+    context_window: int | None = None
+    # OpenRouter-specific (no-ops for other providers)
+    response_healing: bool = False
+    fallback_models: list[str] = field(default_factory=list)
+    provider_routing: dict[str, Any] | None = None
+
+
+_REGISTRY: dict[str, ModelConfig] | None = None
+
+
+def _to_float(val: Any) -> float | None:
+    """Coerce cost values to float (YAML may parse e.g. ``5e-7`` as str)."""
+    if val is None:
+        return None
+    return float(val)
+
+
+def _parse_models_yaml(raw: dict[str, Any]) -> dict[str, ModelConfig]:
+    """Parse a raw YAML dict into ModelConfig entries."""
+    registry: dict[str, ModelConfig] = {}
+    for model_id, cfg in raw.items():
+        providers: dict[str, ProviderModelConfig] = {}
+        for prov_name, prov_cfg in cfg.get("providers", {}).items():
+            features = prov_cfg.get("features", {})
+            providers[prov_name] = ProviderModelConfig(
+                api_model_name=prov_cfg["api_model_name"],
+                input_token_cost=_to_float(prov_cfg.get("input_token_cost")),
+                output_token_cost=_to_float(prov_cfg.get("output_token_cost")),
+                supports_reasoning=prov_cfg.get("supports_reasoning"),
+                response_healing=features.get("response_healing", False),
+                fallback_models=features.get("fallback_models", []),
+                provider_routing=features.get("provider_routing"),
+            )
+        ctx_window = cfg.get("context_window")
+        registry[model_id] = ModelConfig(
+            model_id=model_id,
+            supports_strict_json=cfg.get("supports_strict_json", True),
+            supports_reasoning=cfg.get("supports_reasoning", False),
+            requires_cache_control=cfg.get("requires_cache_control", False),
+            input_token_cost=_to_float(cfg.get("input_token_cost")),
+            output_token_cost=_to_float(cfg.get("output_token_cost")),
+            context_window=int(ctx_window) if ctx_window is not None else None,
+            default_provider=cfg.get("default_provider"),
+            reasoning_effort_options=list(cfg.get("reasoning_effort_options") or []),
+            default_reasoning_effort=cfg.get("default_reasoning_effort"),
+            providers=providers,
+        )
+    return registry
+
+
+def _load_registry() -> dict[str, ModelConfig]:
+    global _REGISTRY
+    if _REGISTRY is not None:
+        return _REGISTRY
+
+    # 1. Load built-in models
+    yaml_path = Path(__file__).parent / "models.yaml"
+    with yaml_path.open() as f:
+        raw = yaml.safe_load(f)
+    registry = _parse_models_yaml(raw)
+
+    # 2. Merge user-defined models (DOTI_CUSTOM_MODELS env var)
+    custom_path = os.environ.get("DOTI_CUSTOM_MODELS")
+    if custom_path:
+        path = Path(custom_path).expanduser()
+        with path.open() as f:
+            custom_raw = yaml.safe_load(f)
+        if custom_raw:
+            registry.update(_parse_models_yaml(custom_raw))
+
+    _REGISTRY = registry
+    return registry
+
+
+def register_model(config: ModelConfig) -> None:
+    """Register (or replace) a model configuration at runtime.
+
+    Adds *config* to the in-memory registry under its ``model_id``, replacing
+    any existing entry with that ID. This is the programmatic counterpart to
+    the ``DOTI_CUSTOM_MODELS`` YAML merge: use it to declare a custom model's
+    costs and capabilities in code so cost tracking and feature detection work
+    for models that are not in the built-in registry::
+
+        register_model(ModelConfig(model_id="my-org/llm", input_token_cost=...))
+
+    The registry is loaded on first use if it has not been already, so this is
+    safe to call as the very first registry interaction in a process.
+    """
+    _load_registry()[config.model_id] = config
+
+
+def resolve_model_config(
+    model: str,
+    provider: str,
+    *,
+    input_token_cost: float | None = None,
+    output_token_cost: float | None = None,
+) -> ResolvedModelConfig | None:
+    """Build a ResolvedModelConfig from the registry + optional cost overrides.
+
+    If the model is in the registry, returns the resolved config with any
+    user-provided costs taking precedence.  If the model is NOT in the
+    registry but costs are provided, returns a minimal config with those
+    costs (so cost tracking still works).  Returns None only when the model
+    is unknown AND no costs were provided.
+    """
+    config = get_model_config(model, provider)
+    has_overrides = input_token_cost is not None or output_token_cost is not None
+
+    if config is not None and has_overrides:
+        overrides: dict[str, float] = {}
+        if input_token_cost is not None:
+            overrides["input_token_cost"] = input_token_cost
+        if output_token_cost is not None:
+            overrides["output_token_cost"] = output_token_cost
+        return replace(config, **overrides)
+
+    if config is None and has_overrides:
+        return ResolvedModelConfig(
+            model_id=model,
+            provider=provider,
+            api_model_name=model,
+            input_token_cost=input_token_cost,
+            output_token_cost=output_token_cost,
+        )
+
+    return config
+
+
+def get_model_config(model_id: str, provider: str) -> ResolvedModelConfig | None:
+    """Look up a model by its canonical ID, resolved for a specific provider.
+
+    Returns None if the model is not in the registry or is not available
+    through the specified provider.
+    """
+    model = _load_registry().get(model_id)
+    if model is None:
+        return None
+    try:
+        return model.for_provider(provider)
+    except KeyError:
+        return None
+
+
+def get_default_provider(model_id: str) -> str | None:
+    """Return the declared default provider for a model, if any."""
+    model = _load_registry().get(model_id)
+    return model.default_provider if model is not None else None
+
+
+def get_context_window(model_id: str) -> int | None:
+    """Return the advertised context window for *model_id*, or None if unknown."""
+    model = _load_registry().get(model_id)
+    return model.context_window if model is not None else None
+
+
+def list_models(provider: str | None = None) -> list[str]:
+    """Return registered model IDs, optionally filtered by provider availability."""
+    registry = _load_registry()
+    if provider is None:
+        return list(registry.keys())
+    return [mid for mid, mc in registry.items() if provider in mc.providers]
+
+
+# Fallback effort ladder for reasoning models that don't declare an explicit
+# option set (the OpenAI/Anthropic convention). "none" is intentionally
+# excluded — it's opt-in via reasoning_effort_options (Mistral models).
+# TODO: this will be made better later.
+_DEFAULT_REASONING_LADDER = ["low", "medium", "high"]
+
+
+def allowed_reasoning_efforts(model_id: str) -> list[str]:
+    """Return the reasoning_effort values valid for *model_id*.
+
+    Empty list when the model is unknown or doesn't support reasoning. When
+    the model supports reasoning but declares no explicit
+    ``reasoning_effort_options``, returns the standard ``[low, medium,
+    high]`` ladder. Lets a caller (e.g. a client UI) offer only the choices a
+    given model accepts — Mistral Medium 3.5 returns ``[none, high]``.
+    """
+    model = _load_registry().get(model_id)
+    if model is None or not model.supports_reasoning:
+        return []
+    if model.reasoning_effort_options:
+        return list(model.reasoning_effort_options)
+    return list(_DEFAULT_REASONING_LADDER)
+
+
+def default_reasoning_effort(model_id: str) -> str | None:
+    """Return the model's declared default reasoning_effort, if any.
+
+    ``None`` when the model is unknown, doesn't support reasoning, or simply
+    doesn't pin a default (in which case the caller's global default
+    applies).
+    """
+    model = _load_registry().get(model_id)
+    if model is None or not model.supports_reasoning:
+        return None
+    return model.default_reasoning_effort

dot_inference/_local_check.py

@@ -0,0 +1,76 @@
+# SPDX-FileCopyrightText: Kannon For Deep Tech
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+"""Sanity check for the ``local`` provider endpoint.
+
+Single-model OpenAI-compatible servers (SGLang, vLLM in single-model mode)
+typically ignore the ``model`` field of a request and answer with whatever
+checkpoint is loaded — some even echo the requested name back in the
+response. Selecting a model the server does not actually serve therefore
+runs the WRONG model with zero failure signal. This module lets callers
+detect that situation up front by comparing the registry's expected
+``api_model_name`` against the endpoint's ``GET /models`` listing.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import httpx
+
+from dot_inference._llm_registry import get_default_provider, get_model_config
+from dot_inference.settings import InferenceSettings
+
+
+@dataclass(frozen=True)
+class LocalEndpointCheck:
+    """Result of probing the local endpoint for a given model."""
+
+    base_url: str
+    # The api_model_name the registry maps the model to on the local provider.
+    expected_model: str
+    # Model ids the endpoint reports via GET /models.
+    served_models: list[str]
+
+    @property
+    def matches(self) -> bool:
+        return self.expected_model in self.served_models
+
+
+async def check_local_endpoint(
+    model_id: str,
+    *,
+    settings: InferenceSettings | None = None,
+    timeout: float = 3.0,
+    transport: httpx.AsyncBaseTransport | None = None,
+) -> LocalEndpointCheck | None:
+    """Probe the ``local`` provider endpoint for *model_id*.
+
+    Returns ``None`` when there is nothing to check: the model is unknown to
+    the registry, its default provider is not ``local``, or no local
+    ``base_url`` is configured. Otherwise returns a
+    :class:`LocalEndpointCheck` — inspect ``.matches`` to know whether the
+    endpoint actually serves the expected model. Network failures propagate
+    as ``httpx`` exceptions; callers decide whether they are fatal.
+
+    *transport* is a test seam (``httpx.MockTransport``).
+    """
+    if get_default_provider(model_id) != "local":
+        return None
+    resolved = get_model_config(model_id, "local")
+    if resolved is None:
+        return None
+    settings = settings or InferenceSettings()
+    base_url = settings.generation.local.base_url
+    if not base_url:
+        return None
+    async with httpx.AsyncClient(timeout=timeout, transport=transport) as client:
+        resp = await client.get(f"{base_url.rstrip('/')}/models")
+        resp.raise_for_status()
+        data = resp.json().get("data", [])
+    served = [m["id"] for m in data if isinstance(m, dict) and "id" in m]
+    return LocalEndpointCheck(
+        base_url=base_url,
+        expected_model=resolved.api_model_name,
+        served_models=served,
+    )