kitkat 0.1.0__py3-none-any.whl

kitkat/_internal/__init__.py

@@ -0,0 +1,6 @@
+"""Private shared utilities for kitkat internals.
+
+Nothing in this package is part of the public API.  Imports from
+``kitkat._internal`` are not covered by stability guarantees and may
+change without notice between any two releases.
+"""

kitkat/_internal/http.py

@@ -0,0 +1,47 @@
+"""Shared httpx async-client configuration.
+
+All providers that use httpx directly should build their clients through
+:func:`build_async_client` so connection-pool settings, timeout defaults,
+and the ``User-Agent`` header are consistent across the library.
+
+Provider SDKs (anthropic, openai, google-genai) manage their own HTTP
+transport internally, so this factory is primarily useful for custom
+providers or future internal HTTP callers.
+"""
+
+from __future__ import annotations
+
+from importlib.metadata import PackageNotFoundError, version
+
+import httpx
+
+try:
+    _LIB_VERSION = version("kitkat")
+except PackageNotFoundError:
+    _LIB_VERSION = "dev"
+
+_USER_AGENT = f"kitkat/{_LIB_VERSION} httpx/{httpx.__version__}"
+
+
+def build_async_client(
+    base_url: str = "",
+    timeout: float = 120.0,
+    **kwargs: object,
+) -> httpx.AsyncClient:
+    """Create a pre-configured :class:`httpx.AsyncClient`.
+
+    Args:
+        base_url: Optional base URL prefix applied to all requests.
+        timeout: Default request timeout in seconds.
+        **kwargs: Additional keyword arguments forwarded to
+            :class:`httpx.AsyncClient`.
+
+    Returns:
+        A ready-to-use async HTTP client with library defaults applied.
+    """
+    return httpx.AsyncClient(
+        base_url=base_url,
+        timeout=httpx.Timeout(timeout),
+        headers={"User-Agent": _USER_AGENT},
+        **kwargs,  # type: ignore[arg-type]
+    )

kitkat/_internal/retry.py

@@ -0,0 +1,102 @@
+"""Shared retry logic for all provider implementations.
+
+Provider subclasses call :func:`execute_with_retry` from within
+``complete_with_retry()``.  This ensures consistent exponential back-off
+behaviour regardless of which provider is in use.
+
+Non-retriable errors (auth, token-limit, content-filter) are re-raised
+immediately without sleeping, so callers never wait on deterministic
+failures.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, TypeVar
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Coroutine
+
+    from ..core.models import RetryPolicy
+
+from ..core.exceptions import (
+    LLMAuthenticationError,
+    LLMContentFilterError,
+    LLMError,
+    LLMRateLimitError,
+    LLMTokenLimitError,
+)
+
+T = TypeVar("T")
+logger = logging.getLogger(__name__)
+
+# Errors that must never be retried
+_NON_RETRIABLE = (LLMAuthenticationError, LLMTokenLimitError, LLMContentFilterError)
+
+
+async def execute_with_retry(
+    func: Callable[[], Coroutine[None, None, T]],
+    policy: RetryPolicy,
+    provider_name: str,
+) -> T:
+    """Execute an async callable with exponential back-off retry.
+
+    Retries on :exc:`~kitkat.core.exceptions.LLMRateLimitError` and
+    generic :exc:`~kitkat.core.exceptions.LLMError`.  Raises immediately
+    on non-retriable errors (authentication, token limit, content filter).
+
+    Args:
+        func: Zero-argument async callable that performs one inference attempt.
+        policy: Retry configuration (attempts, delays, jitter).
+        provider_name: Used in log messages to identify the provider.
+
+    Returns:
+        The return value of *func* on a successful attempt.
+
+    Raises:
+        LLMAuthenticationError: Immediately — credentials are invalid.
+        LLMTokenLimitError: Immediately — prompt is deterministically too long.
+        LLMContentFilterError: Immediately — content policy violation.
+        LLMRateLimitError: After all retry attempts are exhausted.
+        LLMError: After all retry attempts are exhausted for other errors.
+    """
+    last_exc: Exception | None = None
+
+    for attempt in range(policy.max_attempts):
+        try:
+            return await func()
+
+        except _NON_RETRIABLE:
+            raise  # Deterministic failure — skip retries entirely
+
+        except LLMRateLimitError as exc:
+            wait = exc.retry_after_s or policy.delay_for_attempt(attempt)
+            logger.warning(
+                "[%s] Rate limited. Waiting %.1fs (attempt %d/%d).",
+                provider_name,
+                wait,
+                attempt + 1,
+                policy.max_attempts,
+            )
+            last_exc = exc
+            if attempt < policy.max_attempts - 1:
+                await asyncio.sleep(wait)
+
+        except LLMError as exc:
+            wait = policy.delay_for_attempt(attempt)
+            logger.warning(
+                "[%s] Provider error: %s. Waiting %.1fs (attempt %d/%d).",
+                provider_name,
+                exc,
+                wait,
+                attempt + 1,
+                policy.max_attempts,
+            )
+            last_exc = exc
+            if attempt < policy.max_attempts - 1:
+                await asyncio.sleep(wait)
+
+    # All attempts exhausted.
+    assert last_exc is not None, "execute_with_retry exited without an exception set"
+    raise last_exc

kitkat/_internal/tokenizers.py

@@ -0,0 +1,76 @@
+"""Shared token-counting utilities.
+
+Provides a tiktoken-based counter with a conservative character-ratio
+fallback for models not yet supported by tiktoken (e.g. Gemini variants).
+
+All provider ``count_tokens()`` implementations should delegate here so the
+behaviour is consistent across the library.
+"""
+
+from __future__ import annotations
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Conservative approximation: 4 chars ≈ 1 token (valid for most Latin-script
+# text with GPT-style BPE tokenisers).
+_CHARS_PER_TOKEN: float = 4.0
+
+# Sentinel that is stored on first failed tiktoken load so we never try again
+# in the same process (avoids repeated BPE-download attempts in air-gapped envs).
+_TIKTOKEN_UNAVAILABLE = object()
+
+# Cache per encoding name so we only call get_encoding() once.
+_ENCODER_CACHE: dict[str, object] = {}
+
+
+def count_tokens_tiktoken(text: str, encoding_name: str = "cl100k_base") -> int:
+    """Count tokens using tiktoken, with a char-ratio fallback.
+
+    Args:
+        text: The text to tokenise.
+        encoding_name: The tiktoken BPE encoding to use.
+            ``cl100k_base`` covers GPT-4 / GPT-3.5 and approximates
+            Anthropic Claude tokenisation well enough for budgeting.
+
+    Returns:
+        Token count (always ≥ 0; 0 for empty input; ≥ 1 for non-empty).
+    """
+    if not text:
+        return 0
+
+    enc = _ENCODER_CACHE.get(encoding_name)
+    if enc is None:
+        try:
+            import tiktoken
+
+            enc = tiktoken.get_encoding(encoding_name)
+            _ENCODER_CACHE[encoding_name] = enc
+        except Exception as exc:
+            logger.warning(
+                "tiktoken BPE load failed (%s); falling back to "
+                "character-based token estimate (4 chars ≈ 1 token).",
+                exc,
+            )
+            _ENCODER_CACHE[encoding_name] = _TIKTOKEN_UNAVAILABLE
+            enc = _TIKTOKEN_UNAVAILABLE
+
+    if enc is _TIKTOKEN_UNAVAILABLE:
+        return count_tokens_fallback(text)
+
+    return max(1, len(enc.encode(text)))  # type: ignore[union-attr]
+
+
+def count_tokens_fallback(text: str) -> int:
+    """Character-ratio token estimate for models without tiktoken support.
+
+    Args:
+        text: The text to estimate.
+
+    Returns:
+        Estimated token count (0 for empty input; ≥ 1 for non-empty).
+    """
+    if not text:
+        return 0
+    return max(1, round(len(text) / _CHARS_PER_TOKEN))

kitkat/abc/__init__.py

@@ -0,0 +1,11 @@
+"""kitkat.abc — abstract base classes for the library.
+
+The only stable public export is :class:`~kitkat.abc.provider.LLMProvider`.
+Third-party providers should import from here::
+
+    from kitkat.abc import LLMProvider
+"""
+
+from .provider import LLMProvider
+
+__all__ = ["LLMProvider"]

kitkat/abc/provider.py

@@ -0,0 +1,374 @@
+"""The LLMProvider abstract base class.
+
+Every concrete provider (Anthropic, OpenAI, Gemini, or a custom third-party
+provider) must sub-class :class:`LLMProvider` and implement all abstract
+methods.  The library's service layer works exclusively with this ABC — it
+never imports concrete provider classes directly.
+
+Implementing a custom provider::
+
+    from kitkat.abc import LLMProvider
+    from kitkat.core import (
+        LLMRequest, LLMResponse, ProviderCapabilities, ProviderType,
+        RetryPolicy, StreamChunk,
+    )
+    from collections.abc import AsyncIterator
+
+    class MyProvider(LLMProvider):
+        PROVIDER_TYPE = ProviderType.OPENAI        # reuse an existing slot …
+        DEFAULT_MODEL = "my-model-v1"
+        CAPABILITIES = ProviderCapabilities(
+            supports_streaming=True,
+            supports_thinking=False,
+            max_context_tokens=32_768,
+            provider_type=ProviderType.OPENAI,
+        )
+
+        async def initialize(self) -> None:
+            self._client = MySDKClient(api_key=self._config["api_key"])
+            self._initialized = True
+
+        async def shutdown(self) -> None:
+            await self._client.aclose()
+            self._initialized = False
+
+        async def _init_client_only(self) -> None:
+            if self._initialized:
+                return
+            self._client = MySDKClient(api_key=self._config["api_key"])
+            self._initialized = True
+
+        async def complete(self, request: LLMRequest) -> LLMResponse:
+            ...
+
+        async def stream(self, request: LLMRequest) -> AsyncIterator[StreamChunk]:
+            ...
+
+        async def health_check(self) -> bool:
+            ...
+
+        def count_tokens(self, text: str) -> int:
+            from kitkat._internal.tokenizers import count_tokens_tiktoken
+            return count_tokens_tiktoken(text)
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
+
+    from ..core.enums import ProviderType
+
+from .._internal.retry import execute_with_retry
+from ..core.models import (
+    LLMRequest,
+    LLMResponse,
+    Message,
+    ProviderCapabilities,
+    RetryPolicy,
+    StreamChunk,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class LLMProvider(ABC):
+    """Abstract base class for all LLM provider implementations.
+
+    Concrete providers inherit from this class and implement the five
+    abstract methods below.  The shared helpers (:meth:`complete_with_retry`,
+    :meth:`run_sync`, :meth:`_assert_initialized`, …) are provided here so
+    providers don't duplicate boilerplate.
+
+    Lifecycle::
+
+        async with MyProvider(config) as provider:
+            response = await provider.complete(request)
+
+    Or explicitly::
+
+        provider = MyProvider(config)
+        await provider.initialize()
+        try:
+            response = await provider.complete(request)
+        finally:
+            await provider.shutdown()
+    """
+
+    # -- Class-level attributes providers MUST declare --------------------
+
+    PROVIDER_TYPE: ProviderType
+    """Canonical enum value identifying this provider."""
+
+    DEFAULT_MODEL: str
+    """Default model identifier used when :attr:`LLMRequest.model` is empty."""
+
+    CAPABILITIES: ProviderCapabilities
+    """Static feature-flag descriptor queried by the service layer."""
+
+    RETRY_POLICY: RetryPolicy = RetryPolicy()
+    """Default retry policy; concrete providers may override at class level."""
+
+    # -- Constructor -------------------------------------------------------
+
+    def __init__(self, config: dict[str, Any]) -> None:
+        """Create the provider with a raw configuration dictionary.
+
+        Args:
+            config: Provider-specific key/value pairs (API key, model, etc.).
+                Concrete providers typically accept a typed ``*Config``
+                dataclass and call ``super().__init__(config.__dict__)``.
+        """
+        self._config = config
+        self._initialized = False
+        logger.debug("%s provider created.", self.__class__.__name__)
+
+    # -- Lifecycle (abstract) ---------------------------------------------
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        """Initialize the provider: create the HTTP client and probe credentials.
+
+        This is the *full* initialization path. Callers using managed keys
+        should always prefer this over ``_init_client_only``.
+
+        Raises:
+            LLMProviderInitError: If the provider fails to start due to
+                configuration or credential errors.
+        """
+
+    @abstractmethod
+    async def shutdown(self) -> None:
+        """Gracefully release all resources associated with the provider."""
+
+    @abstractmethod
+    async def _init_client_only(self) -> None:
+        """Create the HTTP client *without* running a credential probe.
+
+        This lightweight initialization path is used by
+        :class:`~kitkat.service.byok.BYOKLLMService` for BYOK requests.
+        Auth failures surface on the first inference call rather than a
+        pre-flight probe, avoiding extra latency and billable requests per
+        user key.
+
+        Concrete implementations must:
+
+        1. Guard against double-initialization (idempotent — return early if
+           ``self._initialized`` is already ``True``).
+        2. Instantiate the provider-specific async HTTP client.
+        3. Set ``self._initialized = True`` after successful client creation.
+
+        Raises:
+            LLMProviderInitError: If the underlying client cannot be created.
+        """
+
+    # -- Async context manager support ------------------------------------
+
+    async def __aenter__(self) -> LLMProvider:
+        """Initialize the provider upon context entry."""
+        await self.initialize()
+        return self
+
+    async def __aexit__(self) -> None:
+        """Ensure provider shutdown on context manager exit."""
+        await self.shutdown()
+
+    # -- Core inference (abstract) ----------------------------------------
+
+    @abstractmethod
+    async def complete(self, request: LLMRequest) -> LLMResponse:
+        """Execute a single non-streaming completion attempt.
+
+        This method does **not** apply retry logic. For retry-wrapped
+        completion use :meth:`complete_with_retry`.
+
+        Args:
+            request: The generation request.
+
+        Returns:
+            The completed response from the provider.
+
+        Raises:
+            LLMTimeoutError: If the request exceeds the configured timeout.
+            LLMRateLimitError: On HTTP 429.
+            LLMTokenLimitError: If the prompt exceeds the context window.
+            LLMProviderError: On any other provider-side failure.
+        """
+
+    @abstractmethod
+    async def stream(self, request: LLMRequest) -> AsyncIterator[StreamChunk]:
+        """Yield token deltas as an async stream.
+
+        Args:
+            request: The streaming generation request.
+
+        Yields:
+            :class:`~kitkat.core.models.StreamChunk` objects — one per
+            token delta.  The final chunk has ``is_final=True`` and
+            carries aggregated ``usage``, ``model``, ``provider``,
+            ``finish_reason``, and ``latency_ms``.
+
+        Raises:
+            LLMTimeoutError: If the stream connection times out.
+            LLMRateLimitError: If rate-limited mid-stream.
+            LLMTokenLimitError: If the context window is exceeded.
+            LLMProviderError: On any other streaming error.
+        """
+        # The ``yield`` below satisfies the type-checker's requirement that an
+        # ``@abstractmethod`` decorated as ``AsyncIterator`` is a generator.
+        # Concrete providers should replace the entire body.
+        raise NotImplementedError  # pragma: no cover
+        yield  # type: ignore[misc]  # makes this an async generator
+
+    # -- Health & introspection (abstract) --------------------------------
+
+    @abstractmethod
+    async def health_check(self) -> bool:
+        """Perform a lightweight liveness probe.
+
+        Returns:
+            ``True`` if the provider is reachable and credentials are valid.
+        """
+
+    @abstractmethod
+    def count_tokens(self, text: str) -> int:
+        """Estimate the token count for a piece of text.
+
+        Providers should delegate to
+        :func:`~kitkat._internal.tokenizers.count_tokens_tiktoken`
+        or their SDK's native token counter.
+
+        Args:
+            text: The text to evaluate.
+
+        Returns:
+            Estimated token count (≥ 1 for non-empty input).
+        """
+
+    # -- Shared helpers ---------------------------------------------------
+
+    def count_prompt_tokens(self, messages: list[Message]) -> int:
+        """Estimate total token count for a list of messages.
+
+        Concatenates all message contents with a single-space separator and
+        delegates to :meth:`count_tokens`.
+
+        Args:
+            messages: The conversation messages to estimate.
+
+        Returns:
+            Estimated token count, or 0 for an empty list.
+        """
+        if not messages:
+            return 0
+        return self.count_tokens(" ".join(m.content for m in messages))
+
+    def _assert_initialized(self) -> None:
+        """Raise if the provider has not been initialized.
+
+        Raises:
+            RuntimeError: If :meth:`initialize` (or :meth:`_init_client_only`)
+                has not been successfully called.
+        """
+        if not self._initialized:
+            raise RuntimeError(
+                f"{self.__class__.__name__}.initialize() must be called "
+                "before making inference requests. Use the async context manager."
+            )
+
+    def _build_base_response_kwargs(
+        self,
+        request: LLMRequest,  # noqa: ARG002  (kept for API compatibility)
+        start_time: float,
+    ) -> dict[str, Any]:
+        """Build common tracing fields for every response.
+
+        Args:
+            request: The original :class:`~kitkat.core.models.LLMRequest`.
+            start_time: Monotonic clock value recorded before the API call.
+
+        Returns:
+            Dict with ``provider`` and ``latency_ms`` keys ready to unpack
+            into :class:`~kitkat.core.models.LLMResponse`.
+        """
+        return {
+            "provider": self.PROVIDER_TYPE,
+            "latency_ms": (time.monotonic() - start_time) * 1_000,
+        }
+
+    async def complete_with_retry(
+        self,
+        request: LLMRequest,
+        *,
+        policy: RetryPolicy | None = None,
+    ) -> LLMResponse:
+        """Execute a completion request with exponential back-off retry.
+
+        Delegates to :func:`~kitkat._internal.retry.execute_with_retry`,
+        which handles non-retriable errors (auth, token limit, content
+        filter) by re-raising immediately.
+
+        Args:
+            request: The completion request.
+            policy: Override the provider's class-level ``RETRY_POLICY``.
+
+        Returns:
+            The completed response after a successful attempt.
+
+        Raises:
+            LLMTimeoutError: If all retries time out.
+            LLMRateLimitError: If all rate-limit retries are exhausted.
+            LLMProviderError: On unrecoverable provider errors.
+        """
+        p = policy or getattr(self, "RETRY_POLICY", RetryPolicy())
+        return await execute_with_retry(
+            func=lambda: self.complete(request),
+            policy=p,
+            provider_name=self.__class__.__name__,
+        )
+
+    def run_sync(self, request: LLMRequest) -> LLMResponse:
+        """Execute a completion synchronously (blocks the calling thread).
+
+        Useful for scripts and tests that do not run inside an asyncio event
+        loop.  **Do not call from within a running loop** — use
+        ``await provider.complete(request)`` instead.
+
+        Args:
+            request: The request to send to the provider.
+
+        Returns:
+            The provider response.
+
+        Raises:
+            RuntimeError: If called from within a running asyncio event loop.
+        """
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            pass  # No running loop — safe to proceed
+        else:
+            raise RuntimeError(
+                "run_sync() cannot be called from within a running event loop. "
+                "Use 'await provider.complete(request)' instead."
+            )
+        return asyncio.run(self.complete(request))
+
+    # -- Representation ---------------------------------------------------
+
+    def __repr__(self) -> str:
+        status = "ready" if self._initialized else "uninitialised"
+        provider_type = getattr(self, "PROVIDER_TYPE", "unknown")
+        model = getattr(self, "DEFAULT_MODEL", "unknown")
+        return (
+            f"<{self.__class__.__name__} "
+            f"provider={getattr(provider_type, 'value', provider_type)!r} "
+            f"model={model!r} "
+            f"status={status}>"
+        )

kitkat/core/__init__.py

@@ -0,0 +1,61 @@
+"""Core layer public API.
+
+Zero-dependency foundation — no provider SDK imports, no optional extras.
+Every other module in the library imports from here.
+
+Usage::
+
+    from kitkat.core import LLMRequest, LLMResponse, Role
+    from kitkat.core import LLMAuthenticationError, LLMRateLimitError
+"""
+
+from .enums import FinishReason, ProviderType, Role
+from .exceptions import (
+    KitkatError,
+    LLMAuthenticationError,
+    LLMContentFilterError,
+    LLMError,
+    LLMProviderError,
+    LLMProviderInitError,
+    LLMRateLimitError,
+    LLMTimeoutError,
+    LLMTokenLimitError,
+)
+from .models import (
+    LLMRequest,
+    LLMResponse,
+    Message,
+    ProviderCapabilities,
+    ProviderCapabilitiesModel,
+    RetryPolicy,
+    StreamChunk,
+    ThinkingConfig,
+    TokenUsage,
+)
+
+__all__ = [
+    # Enums
+    "Role",
+    "FinishReason",
+    "ProviderType",
+    # Models
+    "Message",
+    "ThinkingConfig",
+    "LLMRequest",
+    "LLMResponse",
+    "StreamChunk",
+    "TokenUsage",
+    "RetryPolicy",
+    "ProviderCapabilities",
+    "ProviderCapabilitiesModel",
+    # Exceptions
+    "KitkatError",
+    "LLMError",
+    "LLMProviderError",
+    "LLMProviderInitError",
+    "LLMAuthenticationError",
+    "LLMRateLimitError",
+    "LLMTokenLimitError",
+    "LLMTimeoutError",
+    "LLMContentFilterError",
+]