dot-inference 2.0.2__py3-none-any.whl

dot_inference/_openai_compat.py

@@ -0,0 +1,570 @@
+# SPDX-FileCopyrightText: Kannon For Deep Tech
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+"""Shared OpenAI-compatible client.
+
+Provides ``OpenAICompatClient`` — a class-based LLM client for any
+OpenAI-compatible endpoint (OpenAI, OpenRouter, Azure, local servers).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import logging
+from collections.abc import AsyncGenerator
+from typing import Any
+
+import openai
+from openai import AsyncOpenAI
+from pydantic import BaseModel
+
+from dot_inference._client import EmbeddingClient, LLMClient
+from dot_inference._cost_tracker import compute_cost
+from dot_inference._llm_registry import ResolvedModelConfig
+from dot_inference._retry import compute_wait
+from dot_inference._types import (
+    InputJsonDelta,
+    LLMResponse,
+    Message,
+    ReasoningDelta,
+    StreamDone,
+    StreamEvent,
+    StreamUsage,
+    TextDelta,
+    ToolCall,
+    ToolCallDone,
+    ToolDef,
+)
+
+LOGGER = logging.getLogger(__name__)
+# TODO: include the deepika logger
+
+_MAX_ATTEMPTS = 4
+_RETRY_BASE_WAIT = 3  # seconds
+
+# Client-side errors that should NOT be retried.
+_NON_RETRYABLE = (
+    openai.BadRequestError,
+    openai.AuthenticationError,
+    openai.PermissionDeniedError,
+    openai.NotFoundError,
+    openai.UnprocessableEntityError,
+)
+
+
+def _tools_to_openai(tools: list[ToolDef]) -> list[dict[str, Any]]:
+    """Convert ToolDef list to OpenAI tools format."""
+    return [
+        {
+            "type": "function",
+            "function": {
+                "name": t.name,
+                "description": t.description,
+                "parameters": t.parameters,
+            },
+        }
+        for t in tools
+    ]
+
+
+def _extract_reasoning(delta: Any) -> str | None:
+    """Pull a reasoning/thinking fragment out of a chat-completion delta.
+
+    Reasoning is not part of the official Chat Completions schema, so
+    providers extend it in different ways: vLLM/DeepSeek-style servers
+    use ``delta.reasoning_content``, OpenRouter normalizes it to
+    ``delta.reasoning``. OpenAI itself returns neither — its reasoning
+    stays hidden (summaries exist only on the Responses API).
+    Returns None when the chunk carries no reasoning.
+    """
+    for attr in ("reasoning_content", "reasoning"):
+        value = getattr(delta, attr, None)
+        if isinstance(value, str) and value:
+            return value
+    return None
+
+
+def _extract_cached_tokens(usage: Any) -> int | None:
+    """Pull cached-token count out of a usage object if present.
+
+    OpenAI / OpenRouter expose it as ``usage.prompt_tokens_details.cached_tokens``.
+    Returns None if the provider doesn't report it.
+    """
+    if usage is None:
+        return None
+    details = getattr(usage, "prompt_tokens_details", None)
+    if details is None:
+        return None
+    return getattr(details, "cached_tokens", None)
+
+
+class OpenAICompatClient(LLMClient):
+    """LLM client for any OpenAI-compatible endpoint.
+
+    The underlying ``AsyncOpenAI`` SDK client is stored as ``self._client``
+    and created once at construction time.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        base_url: str,
+        api_key: str | None = None,
+        model_config: ResolvedModelConfig | None = None,
+        client: AsyncOpenAI | None = None,
+        temperature: float | None = None,
+        reasoning_effort: str | None = None,
+        max_tokens: int | None = None,
+        timeout: float = 300,
+    ) -> None:
+        if client is None:
+            client_kwargs: dict[str, Any] = {
+                "base_url": base_url,
+                "timeout": timeout,
+            }
+            if api_key is not None:
+                client_kwargs["api_key"] = api_key
+            client = AsyncOpenAI(**client_kwargs)
+
+        super().__init__()
+        self._client = client
+        self._model = model
+        self._model_config = model_config
+        self._temperature = temperature
+        self._reasoning_effort = reasoning_effort
+        self._max_tokens = max_tokens
+        self._strict = model_config.supports_strict_json if model_config else True
+        self._requires_cache_control = (
+            model_config.requires_cache_control if model_config else False
+        )
+
+        # OpenRouter extra_body (provider routing, fallback models)
+        self._extra_body: dict[str, Any] = {}
+        if model_config:
+            if model_config.fallback_models:
+                self._extra_body["models"] = model_config.fallback_models
+            if model_config.provider_routing:
+                self._extra_body["provider"] = model_config.provider_routing
+
+    def _build_create_kwargs(
+        self,
+        messages: list[Message],
+        response_format: type[BaseModel] | dict[str, Any] | None,
+        tools: list[ToolDef] | None,
+        temperature: float | None,
+        reasoning_effort: str | None,
+        tool_choice: str | None = None,
+    ) -> tuple[dict[str, Any], type[BaseModel] | None]:
+        """Build kwargs for ``chat.completions.create``.
+
+        Returns ``(create_kwargs, response_model)`` where *response_model* is
+        the Pydantic class to validate structured output against (or ``None``).
+        """
+        create_kwargs: dict[str, Any] = {
+            "model": self._model,
+            "messages": [m.model_dump(exclude_none=True) for m in messages],
+        }
+
+        # Reasoning models use reasoning_effort instead of temperature
+        is_reasoning = self._model_config.supports_reasoning if self._model_config else False
+        if is_reasoning and reasoning_effort is not None:
+            create_kwargs["reasoning_effort"] = reasoning_effort
+        elif not is_reasoning and temperature is not None:
+            create_kwargs["temperature"] = temperature
+
+        response_model: type[BaseModel] | None = None
+
+        self._validate_params(response_format, tools)
+
+        # Native tool calling mode
+        if tools:
+            create_kwargs["tools"] = _tools_to_openai(tools)
+            create_kwargs["tool_choice"] = tool_choice if tool_choice is not None else "auto"
+
+        # Structured JSON output mode
+        elif response_format is not None:
+            if isinstance(response_format, type) and issubclass(response_format, BaseModel):
+                response_model = response_format
+                schema = response_format.model_json_schema()
+            else:
+                schema = response_format
+
+            # OpenAI strict mode requires additionalProperties: false
+            if self._strict and "additionalProperties" not in schema:
+                schema["additionalProperties"] = False
+
+            create_kwargs["response_format"] = {
+                "type": "json_schema",
+                "json_schema": {
+                    "name": "Output",
+                    "strict": self._strict,
+                    "schema": schema,
+                },
+            }
+
+        if self._max_tokens is not None:
+            create_kwargs["max_tokens"] = self._max_tokens
+
+        extra_body = dict(self._extra_body) if self._extra_body else {}
+        if self._requires_cache_control:
+            # OpenRouter top-level toggle that enables provider caching.
+            # Per-block cache_control markers on messages control where
+            # breakpoints go; this flag is still required for Haiku/Opus
+            # via OpenRouter to actually engage the Anthropic cache.
+            extra_body["cache_control"] = {"type": "ephemeral"}
+        if extra_body:
+            create_kwargs["extra_body"] = extra_body
+
+        return create_kwargs, response_model
+
+    async def _create_with_retry(self, create_kwargs: dict[str, Any]) -> Any:
+        """Call ``chat.completions.create`` with transient-error retry."""
+        for attempt in range(1, _MAX_ATTEMPTS + 1):
+            try:
+                return await self._client.chat.completions.create(**create_kwargs)
+            except openai.APIError as exc:
+                if isinstance(exc, _NON_RETRYABLE) or attempt == _MAX_ATTEMPTS:
+                    raise
+                resp = getattr(exc, "response", None)
+                retry_after = resp.headers.get("retry-after") if resp is not None else None
+                wait = compute_wait(attempt, _RETRY_BASE_WAIT, retry_after)
+                LOGGER.warning(
+                    "Transient API error (attempt %d/%d): %s — retrying in %.1fs",
+                    attempt,
+                    _MAX_ATTEMPTS,
+                    exc,
+                    wait,
+                )
+                if self.on_retry is not None:
+                    self.on_retry(attempt, _MAX_ATTEMPTS, str(exc))
+                await asyncio.sleep(wait)
+
+    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:
+        temperature = self._resolve_temperature(temperature)
+        reasoning_effort = self._resolve_reasoning_effort(reasoning_effort)
+
+        create_kwargs, response_model = self._build_create_kwargs(
+            messages,
+            response_format,
+            tools,
+            temperature,
+            reasoning_effort,
+            tool_choice,
+        )
+        response = await self._create_with_retry(create_kwargs)
+
+        prompt_tokens = response.usage.prompt_tokens if response.usage else None
+        completion_tokens = response.usage.completion_tokens if response.usage else None
+        cached_tokens = _extract_cached_tokens(response.usage)
+
+        cost = None
+        if self._model_config:
+            cost = compute_cost(
+                response.usage,
+                self._model_config.input_token_cost,
+                self._model_config.output_token_cost,
+            )
+        if response.usage:
+            LOGGER.debug(
+                "tokens: %s in (%s cached) / %s out, cost=$%s",
+                prompt_tokens,
+                cached_tokens if cached_tokens is not None else "?",
+                completion_tokens,
+                cost or "?",
+            )
+
+        msg = response.choices[0].message
+
+        # Parse native tool calls if present
+        tool_calls = None
+        if msg.tool_calls:
+            tool_calls = [
+                ToolCall(
+                    id=tc.id,
+                    name=tc.function.name,
+                    arguments=tc.function.arguments,
+                )
+                for tc in msg.tool_calls
+            ]
+
+        # Parse structured output into Pydantic model if requested
+        structured_output = None
+        if response_model is not None:
+            if not msg.content:
+                finish_reason = response.choices[0].finish_reason
+                refusal = getattr(msg, "refusal", None)
+                raise RuntimeError(
+                    f"LLM returned empty content with response_format set "
+                    f"(finish_reason={finish_reason!r}"
+                    + (f", refusal={refusal!r}" if refusal else "")
+                    + ")"
+                )
+            structured_output = response_model.model_validate_json(msg.content)
+
+        response = LLMResponse(
+            content=msg.content,
+            tool_calls=tool_calls,
+            structured_output=structured_output,
+            cost=cost,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            cached_tokens=cached_tokens,
+        )
+        self.cost_tracker.record(response)
+        return response
+
+    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]:
+        temperature = self._resolve_temperature(temperature)
+        reasoning_effort = self._resolve_reasoning_effort(reasoning_effort)
+
+        create_kwargs, response_model = self._build_create_kwargs(
+            messages,
+            response_format,
+            tools,
+            temperature,
+            reasoning_effort,
+            tool_choice,
+        )
+        create_kwargs["stream"] = True
+        create_kwargs["stream_options"] = {"include_usage": True}
+
+        # Accumulators (reset on each retry attempt)
+        content_parts: list[str] = []
+        tool_acc: dict[int, dict[str, str]] = {}
+        finished_tool_calls: list[ToolCall] = []
+        prompt_tokens: int | None = None
+        completion_tokens: int | None = None
+        cached_tokens: int | None = None
+        last_usage: Any = None
+
+        for attempt in range(1, _MAX_ATTEMPTS + 1):
+            has_yielded = False
+            reasoning_seen = False
+            content_parts.clear()
+            tool_acc.clear()
+            finished_tool_calls.clear()
+            prompt_tokens = None
+            completion_tokens = None
+            cached_tokens = None
+            last_usage = None
+            response = None
+
+            try:
+                response = await self._client.chat.completions.create(**create_kwargs)
+
+                async for chunk in response:
+                    if chunk.usage:
+                        prompt_tokens = chunk.usage.prompt_tokens
+                        completion_tokens = chunk.usage.completion_tokens
+                        cached_tokens = _extract_cached_tokens(chunk.usage)
+                        last_usage = chunk.usage
+
+                    if not chunk.choices:
+                        continue
+
+                    choice = chunk.choices[0]
+                    delta = choice.delta
+
+                    # Reasoning/thinking content (provider extension field)
+                    reasoning = _extract_reasoning(delta)
+                    if reasoning:
+                        if not reasoning_seen:
+                            # Once per stream: whether a model exposes
+                            # reasoning is provider-dependent and hard to
+                            # predict, so make the detection observable.
+                            reasoning_seen = True
+                            field = (
+                                "reasoning_content"
+                                if getattr(delta, "reasoning_content", None)
+                                else "reasoning"
+                            )
+                            LOGGER.info(
+                                "model '%s' streams reasoning via delta.%s — "
+                                "emitting reasoning_delta events",
+                                self._model,
+                                field,
+                            )
+                        has_yielded = True
+                        yield ReasoningDelta(content=reasoning)
+
+                    # Text content
+                    if delta.content:
+                        content_parts.append(delta.content)
+                        has_yielded = True
+                        yield TextDelta(content=delta.content)
+
+                    # Tool call deltas
+                    if delta.tool_calls:
+                        for tc_delta in delta.tool_calls:
+                            idx = tc_delta.index
+                            if idx not in tool_acc:
+                                tool_acc[idx] = {"id": "", "name": "", "arguments": ""}
+
+                            acc = tool_acc[idx]
+                            if tc_delta.id:
+                                acc["id"] = tc_delta.id
+                            if tc_delta.function and tc_delta.function.name:
+                                acc["name"] = tc_delta.function.name
+                            if tc_delta.function and tc_delta.function.arguments:
+                                acc["arguments"] += tc_delta.function.arguments
+                                has_yielded = True
+                                yield InputJsonDelta(
+                                    tool_call_index=idx,
+                                    tool_name=acc["name"],
+                                    partial_json=tc_delta.function.arguments,
+                                )
+
+                    # Emit fully assembled tool calls when the model is done
+                    if choice.finish_reason in ("tool_calls", "stop") and tool_acc:
+                        for idx in sorted(tool_acc):
+                            acc = tool_acc[idx]
+                            tc = ToolCall(
+                                id=acc["id"], name=acc["name"], arguments=acc["arguments"]
+                            )
+                            finished_tool_calls.append(tc)
+                            has_yielded = True
+                            yield ToolCallDone(tool_call=tc)
+                        tool_acc.clear()
+
+                break  # Stream completed successfully
+
+            except openai.APIError as exc:
+                # Can't retry after partial output — consumer already has events
+                if isinstance(exc, _NON_RETRYABLE) or has_yielded or attempt == _MAX_ATTEMPTS:
+                    raise
+                resp = getattr(exc, "response", None)
+                retry_after = resp.headers.get("retry-after") if resp is not None else None
+                wait = compute_wait(attempt, _RETRY_BASE_WAIT, retry_after)
+                LOGGER.warning(
+                    "Transient API error during stream (attempt %d/%d): %s — retrying in %.1fs",
+                    attempt,
+                    _MAX_ATTEMPTS,
+                    exc,
+                    wait,
+                )
+                if self.on_retry is not None:
+                    self.on_retry(attempt, _MAX_ATTEMPTS, str(exc))
+                await asyncio.sleep(wait)
+
+            finally:
+                if response is not None:
+                    aclose = getattr(response, "aclose", None) or getattr(response, "close", None)
+                    if aclose is not None:
+                        result = aclose()
+                        if inspect.isawaitable(result):
+                            await result
+
+        # Cost
+        cost = None
+        if self._model_config:
+            cost = compute_cost(
+                last_usage,
+                self._model_config.input_token_cost,
+                self._model_config.output_token_cost,
+            )
+        if prompt_tokens is not None:
+            LOGGER.debug(
+                "tokens: %s in (%s cached) / %s out, cost=$%s",
+                prompt_tokens,
+                cached_tokens if cached_tokens is not None else "?",
+                completion_tokens,
+                cost or "?",
+            )
+
+        if prompt_tokens is not None or completion_tokens is not None:
+            yield StreamUsage(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                cached_tokens=cached_tokens,
+                cost=cost,
+            )
+
+        # Final assembled response
+        content = "".join(content_parts) or None
+        structured_output = None
+        if response_model is not None:
+            if not content:
+                raise RuntimeError("LLM returned empty streamed content with response_format set")
+            structured_output = response_model.model_validate_json(content)
+
+        final_response = LLMResponse(
+            content=content,
+            tool_calls=finished_tool_calls or None,
+            structured_output=structured_output,
+            cost=cost,
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            cached_tokens=cached_tokens,
+        )
+        self.cost_tracker.record(final_response)
+        yield StreamDone(response=final_response)
+
+
+class OpenAICompatEmbeddingClient(EmbeddingClient):
+    """Embedding client for any OpenAI-compatible endpoint.
+
+    The underlying ``AsyncOpenAI`` SDK client is stored as ``self._client``
+    and created once at construction time.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        base_url: str,
+        api_key: str | None = None,
+        client: AsyncOpenAI | None = None,
+        batch_size: int = 0,
+        timeout: float = 300,
+    ) -> None:
+        if client is None:
+            client_kwargs: dict[str, Any] = {
+                "base_url": base_url,
+                "timeout": timeout,
+            }
+            if api_key is not None:
+                client_kwargs["api_key"] = api_key
+            client = AsyncOpenAI(**client_kwargs)
+
+        self._client = client
+        self._model = model
+        self._batch_size = batch_size
+
+    async def embed(self, texts: list[str]) -> list[list[float]]:
+        """Embed texts. The API natively accepts a list, so batch_size=0
+        sends everything in one call. Set batch_size > 0 only to stay
+        under provider per-request limits (e.g. OpenAI caps at 2048 inputs)."""
+        if not texts:
+            return []
+
+        batches = (
+            [texts[i : i + self._batch_size] for i in range(0, len(texts), self._batch_size)]
+            if self._batch_size > 0
+            else [texts]
+        )
+
+        all_embeddings: list[list[float]] = []
+        for batch in batches:
+            response = await self._client.embeddings.create(model=self._model, input=batch)
+            items = sorted(response.data, key=lambda x: x.index)
+            all_embeddings.extend(item.embedding for item in items)
+
+        return all_embeddings

dot_inference/_rate_limit.py

@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Kannon For Deep Tech
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+"""Async rate-limiting wrapper for ``LLMCallable``."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from collections import deque
+from typing import Any
+
+from pydantic import BaseModel
+
+from dot_inference._types import LLMCallable, LLMResponse, Message, ToolDef
+
+
+def with_rate_limit(
+    callable: LLMCallable,
+    max_calls_per_minute: int,
+) -> LLMCallable:
+    """Wrap an ``LLMCallable`` with a sliding-window rate limiter.
+
+    If *max_calls_per_minute* calls have already been made in the last 60
+    seconds, the wrapper sleeps until a slot frees up before forwarding
+    the call.
+
+    Returns a new ``LLMCallable`` — the original is not modified.
+    """
+    if max_calls_per_minute <= 0:
+        return callable
+
+    window = 60.0  # seconds, #TODO: consider placing it in the arguments for flexicility
+    timestamps: deque[float] = deque()
+    lock = asyncio.Lock()
+
+    async def _rate_limited_call(
+        messages: list[Message],
+        response_format: type[BaseModel] | dict[str, Any] | None = None,
+        tools: list[ToolDef] | None = None,
+        **kwargs: Any,
+    ) -> LLMResponse:
+        async with lock:
+            now = time.monotonic()
+
+            # Evict timestamps older than the window
+            while timestamps and timestamps[0] <= now - window:
+                timestamps.popleft()
+
+            # If at capacity, sleep until the oldest call falls out of the window
+            if len(timestamps) >= max_calls_per_minute:
+                sleep_for = timestamps[0] - (now - window)
+                if sleep_for > 0:
+                    await asyncio.sleep(sleep_for)
+                    # Re-evict after sleeping
+                    now = time.monotonic()
+                    while timestamps and timestamps[0] <= now - window:
+                        timestamps.popleft()
+
+            timestamps.append(time.monotonic())
+
+        return await callable(messages, response_format, tools, **kwargs)
+
+    return _rate_limited_call

dot_inference/_retry.py

@@ -0,0 +1,66 @@
+# SPDX-FileCopyrightText: Kannon For Deep Tech
+# SPDX-License-Identifier: AGPL-3.0-or-later
+
+"""Decides how long to wait before retrying a transient provider error.
+
+Shared by the OpenAI and Mistral clients so both honor ``Retry-After`` the
+same way and fall back to the same exponential-backoff schedule when the
+server gives no hint.
+"""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+from email.utils import parsedate_to_datetime
+
+# Cap on Retry-After to avoid a hostile/misconfigured server stalling us.
+_RETRY_AFTER_MAX = 60.0
+
+
+def _parse_retry_after(value: str | None) -> float | None:
+    """Translate a raw ``Retry-After`` header into seconds-to-wait.
+
+    The HTTP spec allows two forms, both handled here:
+      * delta-seconds — e.g. ``"5"`` or ``"2.5"``
+      * HTTP-date     — e.g. ``"Wed, 21 Oct 2026 07:28:00 GMT"`` (treated as
+        UTC when naïve, per RFC 7231)
+
+    Return contract (important — callers branch on this):
+      * ``None``     — no header, or header unparseable → caller should use
+        its own backoff policy.
+      * ``0.0``      — header was valid but says "don't wait" (past date, or
+        negative delta) → retry immediately.
+      * ``float > 0`` — server's hint, in seconds. Caller is responsible for
+        any upper-bound clamping.
+    """
+    if not value:
+        return None
+    value = value.strip()
+    try:
+        return max(0.0, float(value))
+    except ValueError:
+        pass
+    try:
+        dt = parsedate_to_datetime(value)
+    except (TypeError, ValueError):
+        return None
+    if dt is None:
+        return None
+    if dt.tzinfo is None:
+        dt = dt.replace(tzinfo=UTC)
+    return max(0.0, (dt - datetime.now(UTC)).total_seconds())
+
+
+def compute_wait(attempt: int, base: float, retry_after: str | None = None) -> float:
+    """Seconds to sleep before attempt ``attempt + 1``.
+
+    Policy:
+      1. If the server sent a parseable ``Retry-After``, honor it — clamped
+         to ``_RETRY_AFTER_MAX`` so a misconfigured server can't stall us.
+      2. Otherwise, exponential backoff ``base * 2**(attempt - 1)``
+         (1st retry = base, 2nd = 2·base, 3rd = 4·base, …).
+    """
+    hinted = _parse_retry_after(retry_after)
+    if hinted is not None:
+        return min(hinted, _RETRY_AFTER_MAX)
+    return base * (2 ** (attempt - 1))