dot-inference 2.0.2__py3-none-any.whl

dot_inference/README.md

@@ -0,0 +1,83 @@
+# dot_inference
+
+A thin, async-first abstraction over multiple LLM providers. It exposes class-based clients for generation, streaming, and embedding behind a unified interface, so application code never couples to a specific SDK.
+
+## Core concepts
+
+### Client classes (`_client.py`)
+
+The package defines two abstract base classes that every provider must implement:
+
+| Class | Key methods | Purpose |
+| --- | --- | --- |
+| `LLMClient` | `call()`, `stream()`, `stream_text()`, `call_sync()` | Chat completion with optional structured output or tool calling |
+| `EmbeddingClient` | `embed()`, `embed_sync()` | Batched text embeddings |
+
+`LLMResponse` carries `content`, `tool_calls`, `structured_output`, `cost`, and token counts.
+
+### Type aliases (`_types.py`)
+
+Functional type aliases are also available for interoperability:
+
+| Type | Signature | Purpose |
+| --- | --- | --- |
+| `LLMCallable` | `(messages, json_schema?, tools?) -> LLMResponse` | Async chat completion |
+| `LLMStreamCallable` | `(messages, json_schema?, tools?) -> AsyncGenerator[StreamEvent]` | Streaming events (text, tool calls, usage, done) |
+| `EmbeddingCallable` | `(texts) -> list[list[float]]` | Batched text embeddings |
+
+### Settings (`settings.py`)
+
+Configuration is driven by `pydantic-settings` and environment variables prefixed with `DOTI_`:
+
+```bash
+DOTI_GENERATION__PROVIDER=openrouter
+DOTI_GENERATION__OPENROUTER__MODEL_NAME=openai/gpt-5-mini
+DOTI_GENERATION__OPENROUTER__API_KEY=sk-or-...
+DOTI_GENERATION__MISTRAL__MODEL_NAME=mistral-small-latest
+DOTI_GENERATION__MISTRAL__API_KEY=...
+```
+
+Each provider has its own nested section for `model_name`, `api_key`, and `base_url`. Switching providers only requires changing the `PROVIDER` value — all credentials and model names are pre-configured per provider.
+
+### Factories (`_client.py`)
+
+`get_llm_client()` and `get_embedding_client()` read the active provider from an `InferenceSettings` instance and return the appropriate concrete client. This is the recommended entry point for application code.
+
+### Provider implementations
+
+| Module | Purpose |
+| --- | --- |
+| `_openai_compat.py` | OpenAI-compatible client used by OpenAI, OpenRouter, and local providers |
+| `_mistral_compat.py` | Native Mistral SDK client |
+
+## Other modules
+
+| Module | Purpose |
+| --- | --- |
+| `_rate_limit.py` | Sliding-window rate limiter wrapping any `LLMClient.call()` |
+| `_cost_tracker.py` | Accumulates token usage and cost across multiple calls |
+| `_llm_registry.py` | Loads model metadata (pricing, capabilities, routing) from `models.yaml` |
+| `mock.py` | Mock clients (`MockLLMClient`, `MockEmbeddingClient`) for testing without API calls |
+| `models.yaml` | Model registry with per-model pricing and feature flags |
+
+## Quick start
+
+```python
+from dot_inference import InferenceSettings, get_llm_client
+
+settings = InferenceSettings()  # reads from DOTI_* env vars (or a .env file)
+client = get_llm_client(settings)
+
+response = await client.call([{"role": "user", "content": "Hello!"}])
+print(response.content)
+```
+
+For explicit provider selection:
+
+```python
+from dot_inference import InferenceSettings, Provider, get_llm_client
+
+settings = InferenceSettings()
+client = get_llm_client(settings, provider=Provider.OPENROUTER, model="openai/gpt-5-mini")
+response = await client.call([{"role": "user", "content": "Hello!"}])
+```

dot_inference/__init__.py

@@ -0,0 +1,86 @@
+# SPDX-FileCopyrightText: Kannon For Deep Tech
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+from dot_inference._client import (
+    EmbeddingClient,
+    LLMClient,
+    get_embedding_client,
+    get_llm_client,
+)
+from dot_inference._cost_tracker import CostTracker, compute_cost
+from dot_inference._llm_registry import (
+    ModelConfig,
+    ResolvedModelConfig,
+    allowed_reasoning_efforts,
+    default_reasoning_effort,
+    get_context_window,
+    get_model_config,
+    list_models,
+    register_model,
+    resolve_model_config,
+)
+from dot_inference._local_check import LocalEndpointCheck, check_local_endpoint
+from dot_inference._rate_limit import with_rate_limit
+from dot_inference._types import (
+    EmbeddingCallable,
+    InputJsonDelta,
+    LLMCallable,
+    LLMResponse,
+    LLMStreamCallable,
+    Message,
+    ReasoningDelta,
+    Role,
+    StreamDone,
+    StreamEvent,
+    StreamUsage,
+    SyncLLMCallable,
+    TextDelta,
+    ToolCall,
+    ToolCallDone,
+    ToolDef,
+    make_sync,
+    text_stream,
+)
+from dot_inference.settings import InferenceSettings, Provider, ReasoningEffort
+
+__all__ = [
+    "CostTracker",
+    "EmbeddingCallable",
+    "EmbeddingClient",
+    "InferenceSettings",
+    "InputJsonDelta",
+    "LLMCallable",
+    "LLMClient",
+    "LLMResponse",
+    "LLMStreamCallable",
+    "LocalEndpointCheck",
+    "Message",
+    "ModelConfig",
+    "Provider",
+    "ReasoningDelta",
+    "ReasoningEffort",
+    "ResolvedModelConfig",
+    "Role",
+    "StreamDone",
+    "StreamEvent",
+    "StreamUsage",
+    "SyncLLMCallable",
+    "TextDelta",
+    "ToolCall",
+    "ToolCallDone",
+    "ToolDef",
+    "allowed_reasoning_efforts",
+    "check_local_endpoint",
+    "compute_cost",
+    "default_reasoning_effort",
+    "get_context_window",
+    "get_embedding_client",
+    "get_llm_client",
+    "get_model_config",
+    "list_models",
+    "make_sync",
+    "register_model",
+    "resolve_model_config",
+    "text_stream",
+    "with_rate_limit",
+]

dot_inference/_client.py

@@ -0,0 +1,520 @@
+# SPDX-FileCopyrightText: Kannon For Deep Tech
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+"""Class-based LLM client interface.
+
+Provides ``LLMClient`` and ``EmbeddingClient`` — abstract base classes with
+explicit ``self._client`` attributes for readability — and ``get_llm_client()``
+/ ``get_embedding_client()`` settings-based factories that resolve API keys
+and provider configuration automatically.
+"""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator, Callable
+from typing import Any
+
+from pydantic import BaseModel, SecretStr
+
+from dot_inference._cost_tracker import CostTracker
+from dot_inference._llm_registry import (
+    ResolvedModelConfig,
+    allowed_reasoning_efforts,
+    default_reasoning_effort,
+    get_default_provider,
+    resolve_model_config,
+)
+from dot_inference._rate_limit import with_rate_limit
+from dot_inference._types import (
+    LLMResponse,
+    Message,
+    StreamDone,
+    StreamEvent,
+    TextDelta,
+    ToolDef,
+)
+from dot_inference.settings import InferenceSettings, Provider, ReasoningEffort
+
+logger = logging.getLogger(__name__)
+
+
+class LLMClient(ABC):
+    """Abstract base class for LLM clients.
+
+    Subclasses hold the provider SDK client as ``self._client`` and expose
+    ``call()``, ``stream()``, and ``call_sync()`` methods.
+    """
+
+    # for readability:
+    _client: Any
+    _model: str
+    _model_config: ResolvedModelConfig | None
+    _temperature: float | None
+    _reasoning_effort: str | None
+    _max_tokens: int | None
+
+    def __init__(self) -> None:
+        self.cost_tracker = CostTracker()
+        self.on_retry: Callable[[int, int, str], None] | None = None
+        """Optional callback invoked after each failed attempt, before the retry sleep.
+
+        Signature: ``(attempt, max_attempts, error_message) -> None``
+        where *attempt* is the 1-indexed number of the attempt that just failed
+        (1 = first attempt failed / first retry upcoming) and *error_message*
+        is a human-readable description of the failure.
+        """
+
+    @abstractmethod
+    async def call(
+        self,
+        messages: list[Message],
+        response_format: type[BaseModel] | dict[str, Any] | None = None,
+        tools: list[ToolDef] | None = None,
+        *,
+        temperature: float | None = None,
+        reasoning_effort: str | None = None,
+        tool_choice: str | None = None,
+    ) -> LLMResponse:
+        """Send messages to the LLM and return a response.
+
+        Per-call *temperature* and *reasoning_effort* override the instance
+        defaults when provided.  *tool_choice* controls whether the model
+        must call a tool (``"auto"``, ``"required"``, or ``"none"``).
+        Defaults to ``"auto"`` when *tools* is provided, ignored otherwise.
+        """
+
+    @abstractmethod
+    async def stream(
+        self,
+        messages: list[Message],
+        response_format: type[BaseModel] | dict[str, Any] | None = None,
+        tools: list[ToolDef] | None = None,
+        *,
+        temperature: float | None = None,
+        reasoning_effort: str | None = None,
+        tool_choice: str | None = None,
+    ) -> AsyncGenerator[StreamEvent, None]:
+        """Stream events from the LLM.
+
+        Yields ``StreamEvent`` objects: ``TextDelta`` for text,
+        ``ToolCallDelta``/``ToolCallDone`` for tool calls, ``StreamUsage``
+        for token counts, and ``StreamDone`` with the final assembled
+        ``LLMResponse``.
+
+        *tool_choice* controls whether the model must call a tool
+        (``"auto"``, ``"required"``, or ``"none"``).  Defaults to
+        ``"auto"`` when *tools* is provided, ignored otherwise.
+        """
+        yield StreamDone(response=LLMResponse())  # pragma: no cover
+
+    async def stream_text(
+        self,
+        messages: list[Message],
+        response_format: type[BaseModel] | dict[str, Any] | None = None,
+        tools: list[ToolDef] | None = None,
+        *,
+        temperature: float | None = None,
+        reasoning_effort: str | None = None,
+        tool_choice: str | None = None,
+    ) -> AsyncGenerator[str, None]:
+        """Stream only text content from the LLM.
+
+        Convenience wrapper around :meth:`stream` that yields plain ``str``
+        chunks, filtering out tool-call, usage, and done events.
+        """
+        async for event in self.stream(
+            messages,
+            response_format,
+            tools,
+            temperature=temperature,
+            reasoning_effort=reasoning_effort,
+            tool_choice=tool_choice,
+        ):
+            if isinstance(event, TextDelta):
+                yield event.content
+
+    def call_sync(
+        self,
+        messages: list[Message],
+        response_format: type[BaseModel] | dict[str, Any] | None = None,
+        tools: list[ToolDef] | None = None,
+        *,
+        temperature: float | None = None,
+        reasoning_effort: str | None = None,
+        tool_choice: str | None = None,
+    ) -> LLMResponse:
+        """Synchronous wrapper around :meth:`call`.
+
+        Uses a persistent background event loop so async SDK clients
+        can clean up properly.
+        """
+        from dot_inference._sync import run_sync
+
+        return run_sync(
+            self.call(
+                messages,
+                response_format,
+                tools,
+                temperature=temperature,
+                reasoning_effort=reasoning_effort,
+                tool_choice=tool_choice,
+            )
+        )
+
+    @staticmethod
+    def _validate_params(
+        response_format: type[BaseModel] | dict[str, Any] | None,
+        tools: list[ToolDef] | None,
+    ) -> None:
+        """Raise if incompatible parameters are provided together."""
+        if tools and response_format is not None:
+            raise ValueError(
+                "Cannot use both 'tools' and 'response_format' simultaneously. "
+                "Most LLM providers do not support combining tool calling with "
+                "structured JSON output. Use one or the other."
+            )
+
+    def _resolve_temperature(self, temperature: float | None) -> float | None:
+        """Return per-call override or instance default."""
+        return self._temperature if temperature is None else temperature
+
+    def _resolve_reasoning_effort(self, reasoning_effort: str | None) -> str | None:
+        """Return per-call override or instance default, validated for the model.
+
+        Covers the per-call path (``call``/``stream`` with ``reasoning_effort=``)
+        that bypasses the construction-time check in ``_warn_compatibility``.
+        """
+        resolved = self._reasoning_effort if reasoning_effort is None else reasoning_effort
+        if self._model_config is not None:
+            _validate_reasoning_effort_for_model(
+                self._model_config.model_id,
+                resolved,
+                self._model_config.supports_reasoning,
+            )
+        return resolved
+
+
+class EmbeddingClient(ABC):
+    """Abstract base class for embedding clients.
+
+    Subclasses hold the provider SDK client as ``self._client`` and expose
+    an ``embed()`` method.
+    """
+
+    _client: Any
+    _model: str
+    _batch_size: int
+
+    @abstractmethod
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        """Embed a list of texts and return their vectors."""
+
+    def embed_sync(self, texts: list[str]) -> list[list[float]]:
+        """Synchronous wrapper around :meth:`embed`.
+
+        Uses a persistent background event loop so async SDK clients
+        can clean up properly.
+        """
+        from dot_inference._sync import run_sync
+
+        return run_sync(self.embed(texts))
+
+
+def _reveal(secret: SecretStr | None) -> str | None:
+    """Unwrap a ``SecretStr`` right at the point the SDK client is built.
+
+    Keeping the unwrap here means the key exists in clear text only as an
+    argument to the provider SDK, never on a settings object that could be
+    printed or serialized.
+    """
+    return secret.get_secret_value() if secret is not None else None
+
+
+def _validate_reasoning_effort_for_model(
+    model_id: str,
+    reasoning_effort: str | None,
+    supports_reasoning: bool,
+) -> None:
+    """Raise ValueError if *reasoning_effort* isn't accepted by *model_id*.
+
+    A no-op for non-reasoning models (the value is ignored downstream), for
+    ``None`` (means "use instance default"), and for models that declare no
+    restricted option set.
+    """
+    if reasoning_effort is None or not supports_reasoning:
+        return
+    allowed = allowed_reasoning_efforts(model_id)
+    if allowed and str(reasoning_effort) not in allowed:
+        raise ValueError(
+            f"Model '{model_id}' does not accept reasoning_effort "
+            f"'{reasoning_effort}'. Valid values for this model: "
+            f"{', '.join(allowed)}."
+        )
+
+
+def _warn_compatibility(
+    *,
+    model: str,
+    provider: Provider,
+    temperature: float | None,
+    reasoning_effort: str | None,
+) -> None:
+    """Emit warnings when user settings are incompatible with the chosen model."""
+
+    _valid_efforts = {e.value for e in ReasoningEffort}
+    if reasoning_effort is not None and reasoning_effort not in _valid_efforts:
+        raise ValueError(
+            f"Invalid reasoning_effort '{reasoning_effort}'. "
+            f"Valid values: {', '.join(sorted(_valid_efforts))}."
+        )
+
+    config = resolve_model_config(model, provider.value)
+
+    if config is None:
+        logger.warning(
+            "Model '%s' is not registered for provider '%s'. "
+            "The raw model name will be sent to the API — if this is intentional "
+            "you can silence this warning by adding the model to a custom models "
+            "file (DOTI_CUSTOM_MODELS).",
+            model,
+            provider.value,
+        )
+        return  # Can't check further without registry info.
+
+    # Per-model effort validation: a reasoning model may accept only a subset
+    # of the global ladder (Mistral Medium 3.5 takes only [none, high]). Reject
+    # an out-of-range value with the choices valid for *this* model rather than
+    # letting the provider API fail opaquely. The per-call path is covered
+    # separately in LLMClient._resolve_reasoning_effort.
+    _validate_reasoning_effort_for_model(model, reasoning_effort, config.supports_reasoning)
+
+    if config.supports_reasoning and temperature:
+        logger.warning(
+            "Model '%s' is a reasoning model — the temperature setting (%.2f) "
+            "will be ignored in favor of reasoning_effort.",
+            model,
+            temperature,
+        )
+
+
+def get_llm_client(
+    settings: InferenceSettings,
+    *,
+    provider: Provider | None = None,
+    model: str | None = None,
+    temperature: float | None = None,
+    reasoning_effort: str | None = None,
+) -> LLMClient:
+    """Construct an :class:`LLMClient` from settings.
+
+    API keys are resolved from *settings* — nothing needs to be passed
+    explicitly.  Optional keyword arguments override the corresponding
+    values in *settings* without mutating the original object.
+    """
+
+    assert isinstance(settings, InferenceSettings)
+
+    gen_settings = settings.generation
+    if provider is not None:
+        resolved_provider = provider
+    elif model is not None:
+        default = get_default_provider(model)
+        resolved_provider = Provider(default) if default else gen_settings.provider
+    else:
+        resolved_provider = gen_settings.provider
+    resolved_provider_cfg = getattr(gen_settings, resolved_provider.value)
+    resolved_model = model or resolved_provider_cfg.model_name
+    if not resolved_model:
+        raise ValueError(
+            f"No model configured for provider '{resolved_provider.value}'. "
+            f"Set DOTI_GENERATION__{resolved_provider.value.upper()}__MODEL_NAME "
+            f"or pass model= explicitly."
+        )
+    resolved_temperature = gen_settings.temperature if temperature is None else temperature
+    # Precedence: explicit call arg > the model's declared default > global
+    # default. The per-model default lets a restricted model (e.g. Mistral
+    # Medium 3.5, which rejects the global "medium") pick a value it accepts.
+    if reasoning_effort is not None:
+        resolved_reasoning_effort = reasoning_effort
+    else:
+        resolved_reasoning_effort = (
+            default_reasoning_effort(resolved_model) or gen_settings.reasoning_effort
+        )
+    resolved_max_tokens = gen_settings.max_tokens
+    resolved_timeout = gen_settings.timeout
+
+    _warn_compatibility(
+        model=resolved_model,
+        provider=resolved_provider,
+        temperature=resolved_temperature,
+        reasoning_effort=resolved_reasoning_effort,
+    )
+
+    match resolved_provider:
+        case Provider.OPENAI:
+            from dot_inference._openai_compat import OpenAICompatClient
+
+            config = resolve_model_config(resolved_model, "openai")
+            client = OpenAICompatClient(
+                model=config.api_model_name if config else resolved_model,
+                base_url="https://api.openai.com/v1",
+                api_key=_reveal(resolved_provider_cfg.api_key),
+                model_config=config,
+                temperature=resolved_temperature,
+                reasoning_effort=resolved_reasoning_effort,
+                max_tokens=resolved_max_tokens,
+                timeout=resolved_timeout,
+            )
+
+        case Provider.OPENROUTER:
+            from dot_inference._openai_compat import OpenAICompatClient
+
+            config = resolve_model_config(resolved_model, "openrouter")
+            client = OpenAICompatClient(
+                model=config.api_model_name if config else resolved_model,
+                base_url="https://openrouter.ai/api/v1",
+                api_key=_reveal(resolved_provider_cfg.api_key),
+                model_config=config,
+                temperature=resolved_temperature,
+                reasoning_effort=resolved_reasoning_effort,
+                max_tokens=resolved_max_tokens,
+                timeout=resolved_timeout,
+            )
+
+        case Provider.MISTRAL:
+            from dot_inference._mistral_compat import MistralClient
+
+            config = resolve_model_config(resolved_model, "mistral")
+            client = MistralClient(
+                model=config.api_model_name if config else resolved_model,
+                api_key=_reveal(resolved_provider_cfg.api_key),
+                model_config=config,
+                temperature=resolved_temperature,
+                reasoning_effort=resolved_reasoning_effort,
+                max_tokens=resolved_max_tokens,
+                timeout=resolved_timeout,
+            )
+
+        case Provider.LOCAL:
+            from dot_inference._openai_compat import OpenAICompatClient
+
+            if not resolved_provider_cfg.base_url:
+                raise ValueError(
+                    "base_url is required for the local provider. "
+                    "Set DOTI_GENERATION__LOCAL__BASE_URL or pass it in settings."
+                )
+            config = resolve_model_config(resolved_model, "local")
+            client = OpenAICompatClient(
+                model=config.api_model_name if config else resolved_model,
+                base_url=resolved_provider_cfg.base_url,
+                api_key="not-needed",
+                model_config=config,
+                temperature=resolved_temperature,
+                reasoning_effort=resolved_reasoning_effort,
+                max_tokens=resolved_max_tokens,
+                timeout=resolved_timeout,
+            )
+
+        case Provider.AZURE:
+            raise NotImplementedError("Azure provider is not yet supported.")
+
+    # Apply rate limiting if configured.
+    if gen_settings.calls_per_minute_limit > 0:
+        # Deliberate per-instance override: the rate-limited wrapper keeps the
+        # exact LLMCallable signature, but a bound method is not a writable
+        # attribute as far as the type checker is concerned.
+        client.call = with_rate_limit(  # ty: ignore[invalid-assignment]
+            client.call, gen_settings.calls_per_minute_limit
+        )
+
+    return client
+
+
+def get_embedding_client(
+    settings: InferenceSettings,
+    *,
+    provider: Provider | None = None,
+    model: str | None = None,
+    batch_size: int | None = None,
+) -> EmbeddingClient:
+    """Construct an :class:`EmbeddingClient` from settings.
+
+    API keys are resolved from ``settings.embedding`` — nothing needs to be
+    passed explicitly.  Optional keyword arguments override the corresponding
+    values in *settings* without mutating the original object.
+    """
+
+    assert isinstance(settings, InferenceSettings)
+
+    emb_settings = settings.embedding
+    if provider is not None:
+        resolved_provider = provider
+    elif model is not None:
+        default = get_default_provider(model)
+        resolved_provider = Provider(default) if default else emb_settings.provider
+    else:
+        resolved_provider = emb_settings.provider
+    resolved_provider_cfg = getattr(emb_settings, resolved_provider.value)
+    resolved_model = model or resolved_provider_cfg.model_name
+    if not resolved_model:
+        raise ValueError(
+            f"No model configured for provider '{resolved_provider.value}'. "
+            f"Set DOTI_EMBEDDING__{resolved_provider.value.upper()}__MODEL_NAME "
+            f"or pass model= explicitly."
+        )
+    resolved_batch_size = emb_settings.batch_size if batch_size is None else batch_size
+    resolved_timeout = emb_settings.timeout
+
+    match resolved_provider:
+        case Provider.OPENAI:
+            from dot_inference._openai_compat import OpenAICompatEmbeddingClient
+
+            return OpenAICompatEmbeddingClient(
+                model=resolved_model,
+                base_url="https://api.openai.com/v1",
+                api_key=_reveal(resolved_provider_cfg.api_key),
+                batch_size=resolved_batch_size,
+                timeout=resolved_timeout,
+            )
+
+        case Provider.OPENROUTER:
+            from dot_inference._openai_compat import OpenAICompatEmbeddingClient
+
+            return OpenAICompatEmbeddingClient(
+                model=resolved_model,
+                base_url="https://openrouter.ai/api/v1",
+                api_key=_reveal(resolved_provider_cfg.api_key),
+                batch_size=resolved_batch_size,
+                timeout=resolved_timeout,
+            )
+
+        case Provider.MISTRAL:
+            from dot_inference._mistral_compat import MistralEmbeddingClient
+
+            return MistralEmbeddingClient(
+                model=resolved_model,
+                api_key=_reveal(resolved_provider_cfg.api_key),
+                batch_size=resolved_batch_size,
+                timeout=resolved_timeout,
+            )
+
+        case Provider.LOCAL:
+            from dot_inference._openai_compat import OpenAICompatEmbeddingClient
+
+            if not resolved_provider_cfg.base_url:
+                raise ValueError(
+                    "base_url is required for the local provider. "
+                    "Set DOTI_EMBEDDING__LOCAL__BASE_URL or pass it in settings."
+                )
+            return OpenAICompatEmbeddingClient(
+                model=resolved_model,
+                base_url=resolved_provider_cfg.base_url,
+                api_key="not-needed",
+                batch_size=resolved_batch_size,
+                timeout=resolved_timeout,
+            )
+
+        case Provider.AZURE:
+            raise NotImplementedError("Azure provider is not yet supported.")