ellements 0.2.0__py3-none-any.whl

ellements/core/budgeting/trackers.py

@@ -0,0 +1,177 @@
+"""Concrete :class:`BudgetTrackerProtocol` implementations.
+
+Two strategies are provided out of the box:
+
+* :class:`CallCountBudget` — caps the *number* of LLM calls. Useful for
+  test harnesses where token counts are unpredictable but you want a
+  hard ceiling on traffic.
+
+* :class:`TokenBudget` — accounts for input and output tokens against
+  a per-model price table. The unit of the price is up to the caller
+  (USD, micro-dollars, abstract "credits") as long as it is consistent
+  across the price table and the cap.
+
+Both trackers are thread-/coroutine-safe within a single event loop
+because they perform whole-method updates under no contention assumptions
+(LLM calls are coarse-grained). If you need stricter guarantees, wrap
+them in an :class:`asyncio.Lock` at the call site.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from ..exceptions import BudgetExceededError
+
+
+@dataclass(frozen=True, slots=True)
+class TokenPrice:
+    """Price per token for a given model, split by direction.
+
+    Attributes:
+        input: Cost per input token (units defined by the caller).
+        output: Cost per output token.
+    """
+
+    input: float
+    output: float
+
+
+@dataclass(slots=True)
+class CallCountBudget:
+    """Bound the total number of completed LLM calls.
+
+    Args:
+        limit: Maximum number of charges accepted before subsequent
+            :meth:`charge` invocations raise. Must be positive.
+
+    Each :meth:`charge` increments ``spent`` by one regardless of the
+    ``cost`` argument; the tracker treats every call as identical.
+    """
+
+    limit: int
+    spent: int = 0
+
+    def __post_init__(self) -> None:
+        if self.limit <= 0:
+            raise ValueError("limit must be positive")
+
+    def can_afford(self, estimate: float = 0.0) -> bool:
+        """Return ``True`` while we have room for at least one more call."""
+        return self.spent < self.limit
+
+    def charge(self, cost: float, *, details: dict[str, Any] | None = None) -> None:
+        """Record one call against the budget.
+
+        Args:
+            cost: Ignored — kept for protocol compatibility.
+            details: Ignored — kept for protocol compatibility.
+
+        Raises:
+            BudgetExceededError: When the call count has already hit
+                ``limit`` (i.e. this call would push it over).
+        """
+        if self.spent >= self.limit:
+            raise BudgetExceededError(
+                f"call count budget exhausted ({self.spent}/{self.limit})",
+                limit=float(self.limit),
+                spent=float(self.spent),
+                attempted=1.0,
+            )
+        self.spent += 1
+
+    def remaining(self) -> float:
+        """Return the number of calls left before the cap is hit."""
+        return float(self.limit - self.spent)
+
+
+@dataclass(slots=True)
+class TokenBudget:
+    """Bound cumulative token spend across one or more models.
+
+    Args:
+        limit: Maximum cumulative cost. Caller-defined units.
+        prices: Mapping of model name to :class:`TokenPrice`. Models
+            not present in the table are charged the ``default_price``.
+            If neither is set, calls with unknown models raise.
+        default_price: Fallback price used when a charged ``model`` is
+            not in ``prices``. ``None`` means "fail on unknown models".
+
+    Attributes:
+        spent: Running total of charges accumulated through :meth:`charge`.
+
+    Notes:
+        The tracker assumes ``details`` carries integer ``input_tokens``
+        and ``output_tokens`` plus a ``model`` string. Missing keys
+        default to ``0`` tokens, which means unaccounted calls cost
+        nothing but still count toward provider rate-limits.
+    """
+
+    limit: float
+    prices: dict[str, TokenPrice] = field(default_factory=dict)
+    default_price: TokenPrice | None = None
+    spent: float = 0.0
+
+    def __post_init__(self) -> None:
+        if self.limit <= 0:
+            raise ValueError("limit must be positive")
+
+    def can_afford(self, estimate: float = 0.0) -> bool:
+        """Return ``True`` while ``spent + estimate`` is within the cap."""
+        return self.spent + max(estimate, 0.0) <= self.limit
+
+    def charge(self, cost: float, *, details: dict[str, Any] | None = None) -> None:
+        """Record token-derived spend against the budget.
+
+        The actual charge is computed from ``details`` when it carries
+        ``input_tokens`` (and optionally ``output_tokens`` plus
+        ``model``) so providers' usage data drives the bill. Otherwise
+        the supplied ``cost`` is used as a flat fallback — this is the
+        path taken by :class:`BudgetedLLMClient`, which knows the
+        per-method estimate but not the exact token usage.
+
+        Args:
+            cost: Flat fallback charge applied when ``details`` lacks
+                token counts.
+            details: Optional mapping. When it contains an
+                ``input_tokens`` key the charge is computed as
+                ``input_tokens * price.input + output_tokens * price.output``
+                using ``prices[details["model"]]`` (falling back to
+                ``default_price``); otherwise ``details`` is ignored
+                for the math.
+
+        Raises:
+            BudgetExceededError: When the new charge would push ``spent``
+                past ``limit``.
+            KeyError: When ``details`` carries token counts for a model
+                that is not in ``prices`` and ``default_price`` is unset.
+        """
+        if details is None or "input_tokens" not in details:
+            charge = max(cost, 0.0)
+        else:
+            input_tokens = int(details.get("input_tokens", 0))
+            output_tokens = int(details.get("output_tokens", 0))
+            model = str(details.get("model", ""))
+            price = self.prices.get(model, self.default_price)
+            if price is None:
+                raise KeyError(
+                    f"no price configured for model {model!r}; set default_price "
+                    "or add it to prices"
+                )
+            charge = input_tokens * price.input + output_tokens * price.output
+
+        projected = self.spent + charge
+        if projected > self.limit:
+            raise BudgetExceededError(
+                f"token budget exhausted "
+                f"(spent={self.spent:.4f}, attempted={charge:.4f}, limit={self.limit:.4f})",
+                limit=self.limit,
+                spent=self.spent,
+                attempted=charge,
+            )
+        self.spent = projected
+
+    def remaining(self) -> float:
+        """Return the unspent portion of the budget (clipped at zero)."""
+        return max(self.limit - self.spent, 0.0)

ellements/core/caching/__init__.py

@@ -0,0 +1,51 @@
+"""Optional caching layer for LLM calls.
+
+This subpackage offers three orthogonal pieces:
+
+1. :class:`Cache` — a tiny async Protocol with ``get`` / ``set`` /
+   ``delete`` / ``clear`` semantics. Anything that satisfies it can be
+   plugged into the rest of the package.
+2. Concrete backends: :class:`InMemoryCache` (process-local dict with
+   optional TTL) and :class:`JsonDiskCache` (one JSON file per key in
+   a directory; survives process restarts).
+3. :class:`CachingLLMClient` — a thin wrapper around any
+   :class:`ellements.core.LLMClientProtocol` that memoizes
+   :meth:`complete` and :meth:`complete_structured` responses keyed
+   off the canonical request payload.
+
+Caching is **opt-in**: nothing in ``ellements.core.LLMClient`` itself
+caches. Apps wrap their client when they want repeatable, cheap
+responses (smoke tests, demos, deterministic CI). Streams, tool calls,
+and image generation are deliberately **not** cached — they have
+side-effects or unbounded outputs that don't map cleanly onto a
+key→value model.
+
+Example::
+
+    from ellements.core import LLMClient
+    from ellements.core.caching import CachingLLMClient, InMemoryCache
+
+    client = CachingLLMClient(
+        inner=LLMClient(model="openai/gpt-4o-mini"),
+        cache=InMemoryCache(default_ttl=3600),
+    )
+    answer1 = await client.complete("Hello?")  # hits the network
+    answer2 = await client.complete("Hello?")  # served from cache
+"""
+
+from __future__ import annotations
+
+from .cache import Cache, CacheEntry
+from .client import CachingLLMClient
+from .disk import JsonDiskCache
+from .keys import build_cache_key
+from .memory import InMemoryCache
+
+__all__ = [
+    "Cache",
+    "CacheEntry",
+    "CachingLLMClient",
+    "InMemoryCache",
+    "JsonDiskCache",
+    "build_cache_key",
+]

ellements/core/caching/cache.py

@@ -0,0 +1,73 @@
+"""Cache Protocol and shared value type."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Protocol, runtime_checkable
+
+
+@dataclass(frozen=True, slots=True)
+class CacheEntry:
+    """A single cache hit.
+
+    Backends return :data:`None` on a miss and an :class:`CacheEntry`
+    on a hit so callers never have to disambiguate "cached ``None``"
+    from "absent". The dataclass is frozen + slotted to keep cache
+    handling allocation-light.
+
+    Attributes:
+        value: The cached value. May be any JSON-serializable object
+            (or any object the backend can round-trip — e.g. an
+            in-memory backend accepts anything picklable / addressable
+            by reference).
+        created_at: Wall-clock timestamp (Unix seconds) when the entry
+            was first written. Useful for callers that want to expose
+            cache age in UIs.
+    """
+
+    value: Any
+    created_at: float
+
+
+@runtime_checkable
+class Cache(Protocol):
+    """Async-safe key→value cache.
+
+    Implementations must tolerate concurrent ``get`` / ``set`` calls
+    from many tasks. Backends are free to interpret keys as opaque
+    strings (the in-package helpers always produce stable, hashed
+    keys via :func:`build_cache_key`).
+
+    The TTL on :meth:`set` is **per-entry**: when present it overrides
+    any default TTL the backend may have. When ``ttl`` is ``None`` the
+    backend's default policy applies (which may be "never expire").
+    """
+
+    async def get(self, key: str) -> CacheEntry | None:
+        """Return the entry stored at *key*, or ``None`` on a miss."""
+
+    async def set(
+        self,
+        key: str,
+        value: Any,
+        *,
+        ttl: float | None = None,
+    ) -> None:
+        """Store *value* under *key*.
+
+        Args:
+            key: Opaque key. Treated as a string by every backend.
+            value: Object to store. Backends may impose serialization
+                constraints (e.g. JSON-only).
+            ttl: Optional time-to-live in seconds. When omitted the
+                backend's default is used.
+        """
+
+    async def delete(self, key: str) -> None:
+        """Remove the entry at *key* if it exists. No-op on miss."""
+
+    async def clear(self) -> None:
+        """Drop every entry held by the backend."""
+
+
+__all__ = ["Cache", "CacheEntry"]

ellements/core/caching/client.py

@@ -0,0 +1,300 @@
+"""Caching wrapper around any :class:`LLMClientProtocol`.
+
+Wraps an inner LLM client and memoizes deterministic calls
+(:meth:`complete`, :meth:`complete_structured`). Streams, tool calls,
+``continue_conversation``, ``loglikelihood``, and ``generate_image``
+are **not** cached and pass straight through — they either have
+side-effects (tool calls, conversation mutation) or produce outputs
+that can't be cleanly serialized into a key→value store (streams,
+images).
+
+The wrapper itself satisfies :class:`LLMClientProtocol`, so it slots
+into anything that consumes a client structurally.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Iterable, Mapping
+from typing import Any, TypeVar
+
+from ..llm.images import ImageGenerationResponse
+from ..llm.messages import MessageInput, normalize_message_input
+from ..llm.protocol import LLMClientProtocol
+from ..llm.wrapper import LLMClientWrapper
+from ..tools import ToolDialect, ToolExecutor, ToolRegistry
+from ..tools.records import ToolCallResponse
+from .cache import Cache
+from .keys import build_cache_key
+
+T = TypeVar("T")
+
+
+class CachingLLMClient(LLMClientWrapper):
+    """LLM client wrapper that memoizes deterministic completions.
+
+    Args:
+        inner: Any object satisfying :class:`LLMClientProtocol`. All
+            uncached methods delegate to it unchanged.
+        cache: Backend storing the cached responses. Any
+            :class:`Cache` implementation works.
+        ttl: Optional time-to-live (seconds) applied to every cache
+            write made by this wrapper. ``None`` defers to the
+            backend's policy.
+        only_cache_low_temperature: When True (the default), only
+            cache calls with ``temperature == 0.0`` — i.e. requests
+            that are deterministic by construction. Set to False to
+            cache everything (useful when callers know the model's
+            seed/decoding policy is stable).
+
+    The wrapper exposes ``inner``, ``cache``, and ``model`` as
+    attributes so callers can introspect or replace pieces without
+    rebuilding the wrapper. ``model`` proxies the inner client's
+    default model and satisfies :class:`LLMClientProtocol`.
+    """
+
+    def __init__(
+        self,
+        *,
+        inner: LLMClientProtocol,
+        cache: Cache,
+        ttl: float | None = None,
+        only_cache_low_temperature: bool = True,
+    ) -> None:
+        super().__init__(inner=inner)
+        self.cache = cache
+        self.ttl = ttl
+        self.only_cache_low_temperature = only_cache_low_temperature
+
+    # ── Cached methods ────────────────────────────────────────────────
+
+    async def complete(
+        self,
+        messages: MessageInput,
+        *,
+        model: str | None = None,
+        temperature: float = 0.7,
+        max_tokens: int | None = None,
+        **kwargs: Any,
+    ) -> str:
+        """Return a cached completion when available; otherwise call
+        through and cache the result.
+
+        Bypasses the cache when ``only_cache_low_temperature`` is True
+        and ``temperature > 0`` (non-deterministic). See
+        :meth:`LLMClientProtocol.complete` for the underlying contract.
+        """
+        if not self._should_cache(temperature):
+            return await self._inner.complete(
+                messages,
+                model=model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                **kwargs,
+            )
+        resolved_model = model or self._inner.model
+        key = build_cache_key(
+            method="complete",
+            model=resolved_model,
+            messages=_messages_payload(messages),
+            temperature=temperature,
+            max_tokens=max_tokens,
+            extra=kwargs,
+        )
+        hit = await self.cache.get(key)
+        if hit is not None and isinstance(hit.value, str):
+            return hit.value
+        result = await self._inner.complete(
+            messages,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            **kwargs,
+        )
+        await self.cache.set(key, result, ttl=self.ttl)
+        return result
+
+    async def complete_structured(
+        self,
+        messages: MessageInput,
+        response_model: type[T],
+        *,
+        model: str | None = None,
+        temperature: float = 0.7,
+        max_tokens: int | None = None,
+        **kwargs: Any,
+    ) -> T:
+        """Return a cached structured completion when available; otherwise
+        call through and cache the serialised result.
+
+        Cached values are stored as ``model_dump()`` payloads and
+        re-hydrated through *response_model* on read, so the cache is
+        immune to in-process model identity changes. See
+        :meth:`LLMClientProtocol.complete_structured` for the
+        underlying contract.
+        """
+        if not self._should_cache(temperature):
+            return await self._inner.complete_structured(
+                messages,
+                response_model,
+                model=model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                **kwargs,
+            )
+        resolved_model = model or self._inner.model
+        key = build_cache_key(
+            method="complete_structured",
+            model=resolved_model,
+            messages=_messages_payload(messages),
+            temperature=temperature,
+            max_tokens=max_tokens,
+            extra=kwargs,
+            response_model=response_model,
+        )
+        hit = await self.cache.get(key)
+        if hit is not None:
+            if isinstance(hit.value, response_model):
+                return hit.value
+            if isinstance(hit.value, dict):
+                return response_model(**hit.value)
+        result = await self._inner.complete_structured(
+            messages,
+            response_model,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            **kwargs,
+        )
+        cacheable = result.model_dump() if hasattr(result, "model_dump") else result
+        await self.cache.set(key, cacheable, ttl=self.ttl)
+        return result
+
+    # ── Uncached pass-through methods ─────────────────────────────────
+    #
+    # ``continue_conversation`` is inherited from LLMClientWrapper and
+    # forwarded unchanged. The remaining methods below are also pure
+    # passthroughs (no cache lookups would ever hit for stateful or
+    # non-serializable outputs), but they are explicit here so the
+    # type checker can see the LLMClientProtocol surface in full.
+
+    async def complete_with_tools(
+        self,
+        messages: MessageInput,
+        tools: ToolRegistry | Mapping[str, Any] | Iterable[Any],
+        *,
+        tool_executor: ToolExecutor | None = None,
+        dialect: ToolDialect | None = None,
+        model: str | None = None,
+        temperature: float = 0.7,
+        max_tokens: int | None = None,
+        max_iterations: int = 10,
+        **kwargs: Any,
+    ) -> ToolCallResponse:
+        """Forward tool-using completions to the inner client unchanged.
+
+        Tool executions have side-effects, so the wrapper does not
+        cache the response. See
+        :meth:`LLMClientProtocol.complete_with_tools`.
+        """
+        return await self._inner.complete_with_tools(
+            messages,
+            tools,
+            tool_executor=tool_executor,
+            dialect=dialect,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            max_iterations=max_iterations,
+            **kwargs,
+        )
+
+    def stream(
+        self,
+        messages: MessageInput,
+        *,
+        model: str | None = None,
+        temperature: float = 0.7,
+        max_tokens: int | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[str]:
+        """Forward streamed completions to the inner client unchanged.
+
+        Stream chunks cannot be cleanly re-played from a cache without
+        materialising the full sequence in memory; callers should
+        cache at a higher level if they need that. See
+        :meth:`LLMClientProtocol.stream`.
+        """
+        return self._inner.stream(
+            messages,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            **kwargs,
+        )
+
+    async def loglikelihood(
+        self,
+        context: str,
+        continuation: str,
+        *,
+        model: str | None = None,
+        max_tokens: int | None = None,
+        **kwargs: Any,
+    ) -> tuple[float, bool]:
+        """Forward log-likelihood scoring to the inner client unchanged.
+
+        Scoring is deterministic per (model, context, continuation),
+        but is not a frequent enough call to warrant a cache layer
+        here. See :meth:`LLMClientProtocol.loglikelihood`.
+        """
+        return await self._inner.loglikelihood(
+            context,
+            continuation,
+            model=model,
+            max_tokens=max_tokens,
+            **kwargs,
+        )
+
+    async def generate_image(
+        self,
+        prompt: str,
+        *,
+        model: str | None = None,
+        n: int = 1,
+        size: str | None = None,
+        quality: str | None = None,
+        style: str | None = None,
+        response_format: str = "url",
+        **kwargs: Any,
+    ) -> ImageGenerationResponse:
+        """Forward image generation to the inner client unchanged.
+
+        Image bytes / URLs are out of scope for a generic LLM cache;
+        callers should use a dedicated image cache if needed. See
+        :meth:`LLMClientProtocol.generate_image`.
+        """
+        return await self._inner.generate_image(
+            prompt,
+            model=model,
+            n=n,
+            size=size,
+            quality=quality,
+            style=style,
+            response_format=response_format,
+            **kwargs,
+        )
+
+    # ── Internals ─────────────────────────────────────────────────────
+
+    def _should_cache(self, temperature: float) -> bool:
+        if not self.only_cache_low_temperature:
+            return True
+        return temperature == 0.0
+
+
+def _messages_payload(messages: MessageInput) -> Any:
+    """Normalize *messages* into a stable JSON-friendly payload."""
+    return normalize_message_input(messages)
+
+
+__all__ = ["CachingLLMClient"]