ellements 0.2.0__py3-none-any.whl

ellements/core/rate_limit/__init__.py

@@ -0,0 +1,44 @@
+"""Optional client-side rate limiting for LLM calls.
+
+Two pieces:
+
+1. :class:`RateLimiterProtocol` — a tiny async Protocol with a single
+   :meth:`acquire` method. Anything implementing it is a valid
+   limiter.
+2. :class:`TokenBucketRateLimiter` — a classic token-bucket
+   implementation: refills at a configurable rate, accepts bursts up
+   to a configurable capacity. Bounded contention via a single
+   :class:`asyncio.Lock`.
+3. :class:`RateLimitedLLMClient` — a wrapper around any
+   :class:`ellements.core.LLMClientProtocol` that calls
+   ``await limiter.acquire(...)`` before each request.
+
+The wrapper itself satisfies :class:`LLMClientProtocol` so it slots
+into any consumer structurally.
+
+Example::
+
+    from ellements.core import LLMClient
+    from ellements.core.rate_limit import (
+        RateLimitedLLMClient,
+        TokenBucketRateLimiter,
+    )
+
+    limiter = TokenBucketRateLimiter(rate_per_second=2.0, capacity=4)
+    client = RateLimitedLLMClient(
+        inner=LLMClient(model="openai/gpt-4o-mini"),
+        limiter=limiter,
+    )
+"""
+
+from __future__ import annotations
+
+from .bucket import TokenBucketRateLimiter
+from .client import RateLimitedLLMClient
+from .protocol import RateLimiterProtocol
+
+__all__ = [
+    "RateLimitedLLMClient",
+    "RateLimiterProtocol",
+    "TokenBucketRateLimiter",
+]

ellements/core/rate_limit/bucket.py

@@ -0,0 +1,85 @@
+"""Token-bucket :class:`RateLimiterProtocol` implementation."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+
+
+class TokenBucketRateLimiter:
+    """Classic token-bucket limiter satisfying :class:`RateLimiterProtocol`.
+
+    The bucket is refilled continuously at ``rate_per_second`` units
+    per second, up to a maximum of ``capacity`` units. Each
+    :meth:`acquire` call removes ``cost`` units; if the bucket lacks
+    enough, the caller sleeps until refill catches up.
+
+    Args:
+        rate_per_second: Refill rate. Must be positive. A value of
+            ``2`` means up to two single-unit calls per second
+            (steady state).
+        capacity: Maximum bucket size. Bursts up to this size are
+            allowed when the bucket starts full / has been idle.
+            Defaults to ``rate_per_second`` (no extra burst budget).
+
+    Raises:
+        ValueError: If *rate_per_second* or *capacity* is not
+            positive, or if a single :meth:`acquire` call requests a
+            *cost* larger than the bucket's *capacity* (such a call
+            could never be satisfied).
+    """
+
+    def __init__(
+        self,
+        rate_per_second: float,
+        *,
+        capacity: float | None = None,
+    ) -> None:
+        if rate_per_second <= 0:
+            raise ValueError(
+                f"rate_per_second must be positive, got {rate_per_second}."
+            )
+        if capacity is None:
+            capacity = rate_per_second
+        if capacity <= 0:
+            raise ValueError(f"capacity must be positive, got {capacity}.")
+        self._rate = rate_per_second
+        self._capacity = float(capacity)
+        self._tokens = float(capacity)
+        self._last_refill = time.monotonic()
+        self._lock = asyncio.Lock()
+
+    async def acquire(self, cost: float = 1.0) -> None:
+        """Block until *cost* tokens are available, then deduct them.
+
+        Returns immediately if *cost* is zero or negative. Raises
+        :class:`ValueError` if *cost* exceeds the bucket's capacity
+        (such a request could never be satisfied).
+        """
+        if cost <= 0:
+            return
+        if cost > self._capacity:
+            raise ValueError(
+                f"acquire(cost={cost}) exceeds capacity={self._capacity}; "
+                f"this call could never be satisfied."
+            )
+        while True:
+            async with self._lock:
+                self._refill_locked()
+                if self._tokens >= cost:
+                    self._tokens -= cost
+                    return
+                deficit = cost - self._tokens
+                wait_seconds = deficit / self._rate
+            await asyncio.sleep(wait_seconds)
+
+    def _refill_locked(self) -> None:
+        now = time.monotonic()
+        elapsed = now - self._last_refill
+        if elapsed <= 0:
+            return
+        self._tokens = min(self._capacity, self._tokens + elapsed * self._rate)
+        self._last_refill = now
+
+
+__all__ = ["TokenBucketRateLimiter"]

ellements/core/rate_limit/client.py

@@ -0,0 +1,216 @@
+"""Rate-limited :class:`LLMClientProtocol` wrapper."""
+
+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 Conversation, MessageInput
+from ..llm.protocol import LLMClientProtocol
+from ..llm.wrapper import LLMClientWrapper
+from ..tools import ToolDialect, ToolExecutor, ToolRegistry
+from ..tools.records import ToolCallResponse
+from .protocol import RateLimiterProtocol
+
+T = TypeVar("T")
+
+
+class RateLimitedLLMClient(LLMClientWrapper):
+    """LLM client wrapper that acquires a permit before each call.
+
+    Every public method calls ``await limiter.acquire(cost)`` before
+    delegating to the inner client. The default ``cost`` is ``1.0``
+    per call, but callers can supply a custom per-method cost map
+    (e.g. ``{"generate_image": 5.0}``) to reflect the relative
+    expense of each operation.
+
+    Args:
+        inner: Any object satisfying :class:`LLMClientProtocol`.
+        limiter: Any object satisfying :class:`RateLimiterProtocol`. Often a
+            :class:`TokenBucketRateLimiter`, but could be backed by
+            Redis or an external service.
+        cost_per_method: Optional ``{method_name: cost}`` overrides.
+            Missing methods fall back to ``default_cost``.
+        default_cost: Cost used for any method not present in
+            *cost_per_method*. Defaults to ``1.0``.
+    """
+
+    def __init__(
+        self,
+        *,
+        inner: LLMClientProtocol,
+        limiter: RateLimiterProtocol,
+        cost_per_method: Mapping[str, float] | None = None,
+        default_cost: float = 1.0,
+    ) -> None:
+        super().__init__(inner=inner)
+        self.limiter = limiter
+        self._cost_per_method = dict(cost_per_method or {})
+        self._default_cost = default_cost
+
+    def _cost(self, method_name: str) -> float:
+        return self._cost_per_method.get(method_name, self._default_cost)
+
+    async def complete(
+        self,
+        messages: MessageInput,
+        *,
+        model: str | None = None,
+        temperature: float = 0.7,
+        max_tokens: int | None = None,
+        **kwargs: Any,
+    ) -> str:
+        """Acquire a permit for ``complete`` then forward to the inner client."""
+        await self.limiter.acquire(self._cost("complete"))
+        return await self._inner.complete(
+            messages,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            **kwargs,
+        )
+
+    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:
+        """Acquire a permit for ``complete_structured`` then forward."""
+        await self.limiter.acquire(self._cost("complete_structured"))
+        return await self._inner.complete_structured(
+            messages,
+            response_model,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            **kwargs,
+        )
+
+    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:
+        """Acquire a single permit for the whole tool-using turn.
+
+        Note: tool follow-up calls made by the inner client are **not**
+        re-throttled here. If you want per-iteration throttling, wrap
+        the inner client's tool-execution path instead.
+        """
+        await self.limiter.acquire(self._cost("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,
+        )
+
+    async def stream(
+        self,
+        messages: MessageInput,
+        *,
+        model: str | None = None,
+        temperature: float = 0.7,
+        max_tokens: int | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[str]:
+        """Acquire a permit for ``stream`` then yield from the inner stream."""
+        await self.limiter.acquire(self._cost("stream"))
+        async for chunk in self._inner.stream(
+            messages,
+            model=model,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            **kwargs,
+        ):
+            yield chunk
+
+    async def continue_conversation(
+        self,
+        conversation: Conversation,
+        user_message: str,
+        *,
+        temperature: float = 0.7,
+        max_tokens: int | None = None,
+        **kwargs: Any,
+    ) -> str:
+        """Acquire a permit for ``continue_conversation`` then forward."""
+        await self.limiter.acquire(self._cost("continue_conversation"))
+        return await self._inner.continue_conversation(
+            conversation,
+            user_message,
+            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]:
+        """Acquire a permit for ``loglikelihood`` then forward."""
+        await self.limiter.acquire(self._cost("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:
+        """Acquire a permit for ``generate_image`` then forward.
+
+        Image generation typically costs more than text completions;
+        pass a higher cost via ``cost_per_method={"generate_image": …}``
+        to reflect that.
+        """
+        await self.limiter.acquire(self._cost("generate_image"))
+        return await self._inner.generate_image(
+            prompt,
+            model=model,
+            n=n,
+            size=size,
+            quality=quality,
+            style=style,
+            response_format=response_format,
+            **kwargs,
+        )
+
+
+__all__ = ["RateLimitedLLMClient"]

ellements/core/rate_limit/protocol.py

@@ -0,0 +1,27 @@
+"""Structural protocol for rate limiters."""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class RateLimiterProtocol(Protocol):
+    """Async-safe permit acquisition.
+
+    A limiter exposes a single coroutine: :meth:`acquire`. Callers
+    ``await`` it before performing the rate-limited action and the
+    limiter blocks until enough capacity is available.
+
+    Implementations are free to model "capacity" however they want
+    (requests-per-second, tokens-per-minute, bytes-per-hour, ...).
+    The ``cost`` argument lets callers indicate that a single
+    operation consumes multiple units of capacity (e.g. a long
+    prompt counts as 50 tokens).
+    """
+
+    async def acquire(self, cost: float = 1.0) -> None:
+        """Block until *cost* units of capacity become available."""
+
+
+__all__ = ["RateLimiterProtocol"]

ellements/core/templating.py

@@ -0,0 +1,126 @@
+"""Template loading and rendering helpers.
+
+Templates can be loaded either from a filesystem directory or from
+package resources via ``importlib.resources``. Mustache-style
+substitution is delegated to :mod:`chevron`.
+
+The renderer validates its source at construction so misconfiguration
+fails fast (rather than at the first call).
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Mapping
+from importlib.resources import files
+from importlib.resources.abc import Traversable
+from pathlib import Path
+from typing import Any
+
+import chevron
+
+_logger = logging.getLogger(__name__)
+
+
+class TemplateRenderer:
+    """Load and render mustache templates from a filesystem dir or a package.
+
+    Exactly one of ``template_dir`` or ``package`` must be provided.
+
+    Args:
+        template_dir: Directory containing template files.
+        package: Importable package name whose resources hold templates.
+        resource_root: Sub-path within *package* (or *template_dir*) to
+            resolve template paths against. Default ``"prompts"``.
+    """
+
+    def __init__(
+        self,
+        template_dir: Path | str | None = None,
+        *,
+        package: str | None = None,
+        resource_root: str = "prompts",
+    ) -> None:
+        if (template_dir is None) == (package is None):
+            raise ValueError(
+                "TemplateRenderer requires exactly one of template_dir or package."
+            )
+
+        self.resource_root = resource_root.strip("/")
+
+        if template_dir is not None:
+            self.template_dir: Path | None = Path(template_dir)
+            if not self.template_dir.exists():
+                raise FileNotFoundError(
+                    f"Template directory not found: {self.template_dir}"
+                )
+            self.package: str | None = None
+        else:
+            assert package is not None
+            self.template_dir = None
+            self.package = package
+            # Probe the package resource root to fail-fast on bad config.
+            try:
+                files(package)
+            except ModuleNotFoundError as exc:
+                raise ValueError(
+                    f"TemplateRenderer package {package!r} cannot be imported: {exc}"
+                ) from exc
+
+    def _resource_traversable(self, template_path: str) -> Traversable:
+        assert self.package is not None
+        resource = files(self.package)
+        if self.resource_root:
+            resource = resource.joinpath(*self.resource_root.split("/"))
+        return resource.joinpath(*template_path.split("/"))
+
+    def load_template(self, template_path: str) -> str:
+        """Read the raw template text at *template_path*."""
+        if self.template_dir is not None:
+            return (self.template_dir / template_path).read_text(encoding="utf-8")
+        return self._resource_traversable(template_path).read_text(encoding="utf-8")
+
+    @staticmethod
+    def render_template(template_content: str, context: Mapping[str, Any]) -> str:
+        """Render *template_content* against *context* using mustache syntax."""
+        rendered: str = chevron.render(template_content, dict(context))
+        return rendered
+
+    def load_and_render(self, template_path: str, context: Mapping[str, Any]) -> str:
+        """Convenience: load *template_path* and render it against *context*."""
+        return self.render_template(self.load_template(template_path), context)
+
+    def render_with_fallback(
+        self,
+        template_path: str,
+        context: Mapping[str, Any],
+        fallback: str,
+        *,
+        silent: bool = False,
+    ) -> str:
+        """Render a template, returning *fallback* on any error."""
+        try:
+            return self.load_and_render(template_path, context)
+        except Exception as exc:
+            if not silent:
+                _logger.warning(
+                    "Error loading template %r: %s. Using fallback content.",
+                    template_path,
+                    exc,
+                )
+            return fallback
+
+
+def render_prompt_template(
+    template_path: str,
+    *,
+    package: str,
+    resource_root: str = "prompts",
+    **kwargs: Any,
+) -> str:
+    """One-shot helper to load+render a packaged template with kwargs as context."""
+    renderer = TemplateRenderer(package=package, resource_root=resource_root)
+    return renderer.load_and_render(template_path, kwargs)
+
+
+__all__ = ["TemplateRenderer", "render_prompt_template"]

ellements/core/tools/__init__.py

@@ -0,0 +1,51 @@
+"""Provider-neutral tool architecture.
+
+This subpackage defines the canonical tool surface every LLM call and
+every agent backend speaks to:
+
+- :class:`Tool` — the duck-typed Protocol.
+- :class:`ToolSpec` — the canonical in-memory record.
+- :class:`SimpleTool` — convenience builder from a Python callable.
+- :class:`ToolRegistry` — composable typed collection.
+- :class:`ToolDialect` and built-in dialects — provider-specific wire
+  formats (:class:`OpenAIChatDialect`, :class:`OpenAIResponsesDialect`,
+  :class:`AnthropicDialect`, :class:`GeminiDialect`).
+- :class:`ToolCallRecord` / :class:`ToolCallResponse` — result records.
+"""
+
+from __future__ import annotations
+
+from .dialects import (
+    AnthropicDialect,
+    GeminiDialect,
+    OpenAIChatDialect,
+    OpenAIResponsesDialect,
+    ToolDialect,
+    default_dialect_for_model,
+)
+from .executor import ToolExecutor, stringify_tool_result
+from .protocol import Tool
+from .records import ToolCallRecord, ToolCallResponse
+from .registry import ToolExecutorFn, ToolInput, ToolRegistry
+from .simple import SimpleTool, bind_json_invoker
+from .spec import ToolSpec
+
+__all__ = [
+    "AnthropicDialect",
+    "GeminiDialect",
+    "OpenAIChatDialect",
+    "OpenAIResponsesDialect",
+    "SimpleTool",
+    "Tool",
+    "ToolCallRecord",
+    "ToolCallResponse",
+    "ToolDialect",
+    "ToolExecutor",
+    "ToolExecutorFn",
+    "ToolInput",
+    "ToolRegistry",
+    "ToolSpec",
+    "bind_json_invoker",
+    "default_dialect_for_model",
+    "stringify_tool_result",
+]

ellements/core/tools/dialects.py

@@ -0,0 +1,136 @@
+"""Wire-format dialects for LLM-provider tool definitions.
+
+A :class:`ToolDialect` translates a provider-neutral
+:class:`~ellements.core.tools.ToolSpec` into the JSON shape a given
+provider's API expects. Adding support for a new provider is one new
+dialect class with no churn elsewhere.
+
+The built-in dialects are:
+
+- :class:`OpenAIChatDialect` — OpenAI Chat Completions / LiteLLM default.
+- :class:`OpenAIResponsesDialect` — OpenAI Responses API (flatter shape).
+- :class:`AnthropicDialect` — Anthropic Messages API direct.
+- :class:`GeminiDialect` — Google Gemini direct API.
+
+:func:`default_dialect_for_model` picks a sensible dialect from a model
+name; callers can override explicitly when needed.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, ClassVar, Protocol, runtime_checkable
+
+from .spec import ToolSpec
+
+
+@runtime_checkable
+class ToolDialect(Protocol):
+    """Emits a provider-specific tool definition from a :class:`ToolSpec`."""
+
+    name: ClassVar[str]
+    """Stable dialect identifier (e.g. ``"openai_chat"``)."""
+
+    def emit(self, spec: ToolSpec) -> dict[str, Any]:
+        """Return the provider-specific dictionary for *spec*."""
+
+
+@dataclass(slots=True, frozen=True)
+class OpenAIChatDialect:
+    """OpenAI Chat Completions wire format (also used by LiteLLM)."""
+
+    name: ClassVar[str] = "openai_chat"
+
+    def emit(self, spec: ToolSpec) -> dict[str, Any]:
+        return {
+            "type": "function",
+            "function": {
+                "name": spec.name,
+                "description": spec.description,
+                "parameters": spec.params_json_schema,
+            },
+        }
+
+
+@dataclass(slots=True, frozen=True)
+class OpenAIResponsesDialect:
+    """OpenAI Responses API wire format (flat, no ``function`` envelope)."""
+
+    name: ClassVar[str] = "openai_responses"
+
+    def emit(self, spec: ToolSpec) -> dict[str, Any]:
+        return {
+            "type": "function",
+            "name": spec.name,
+            "description": spec.description,
+            "parameters": spec.params_json_schema,
+        }
+
+
+@dataclass(slots=True, frozen=True)
+class AnthropicDialect:
+    """Anthropic Messages API wire format (``input_schema`` key)."""
+
+    name: ClassVar[str] = "anthropic"
+
+    def emit(self, spec: ToolSpec) -> dict[str, Any]:
+        return {
+            "name": spec.name,
+            "description": spec.description,
+            "input_schema": spec.params_json_schema,
+        }
+
+
+@dataclass(slots=True, frozen=True)
+class GeminiDialect:
+    """Google Gemini direct API wire format."""
+
+    name: ClassVar[str] = "gemini"
+
+    def emit(self, spec: ToolSpec) -> dict[str, Any]:
+        return {
+            "name": spec.name,
+            "description": spec.description,
+            "parameters": spec.params_json_schema,
+        }
+
+
+def default_dialect_for_model(
+    model: str,
+    *,
+    use_responses_api: bool = False,
+) -> ToolDialect:
+    """Pick the appropriate :class:`ToolDialect` for *model*.
+
+    LiteLLM normalizes every provider onto the OpenAI Chat Completions
+    tool wire format, so the default for litellm-routed calls is
+    :class:`OpenAIChatDialect`. For direct (non-litellm) clients,
+    Anthropic/Gemini-specific dialects apply.
+
+    Args:
+        model: The model identifier, possibly prefixed with a provider
+            slug (e.g. ``"anthropic/claude-3-5-sonnet"``).
+        use_responses_api: If true, prefer
+            :class:`OpenAIResponsesDialect` for OpenAI models.
+
+    Returns:
+        A :class:`ToolDialect` instance.
+    """
+    model_lc = model.lower()
+    if model_lc.startswith("anthropic/") or model_lc.startswith("claude-"):
+        return AnthropicDialect()
+    if model_lc.startswith("gemini/") or model_lc.startswith("gemini-"):
+        return GeminiDialect()
+    if use_responses_api:
+        return OpenAIResponsesDialect()
+    return OpenAIChatDialect()
+
+
+__all__ = [
+    "AnthropicDialect",
+    "GeminiDialect",
+    "OpenAIChatDialect",
+    "OpenAIResponsesDialect",
+    "ToolDialect",
+    "default_dialect_for_model",
+]