bedrock-kit 0.1.0__py3-none-any.whl

bedrock_kit/__init__.py

@@ -0,0 +1,35 @@
+"""bedrock-kit - small, opinionated AWS Bedrock client wrapper.
+
+Adds the things every Bedrock production team rebuilds: adaptive throttle,
+per-call cost tracking with cache-aware accounting, and structured-output
+parse-and-repair. Single-cloud, single-purpose, < 3000 LOC.
+"""
+
+from bedrock_kit.client import BedrockClient, BedrockResponse, Usage
+from bedrock_kit.cost import CostEntry, CostLedger, Pricing
+from bedrock_kit.exceptions import (
+    BedrockKitError,
+    JsonParseError,
+    PricingNotFoundError,
+    ThrottleExhausted,
+)
+from bedrock_kit.retry import AdaptiveThrottle
+from bedrock_kit.schema import JsonSchema
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "AdaptiveThrottle",
+    "BedrockClient",
+    "BedrockKitError",
+    "BedrockResponse",
+    "CostEntry",
+    "CostLedger",
+    "JsonParseError",
+    "JsonSchema",
+    "Pricing",
+    "PricingNotFoundError",
+    "ThrottleExhausted",
+    "Usage",
+    "__version__",
+]

bedrock_kit/client.py

@@ -0,0 +1,254 @@
+"""BedrockClient - the user-facing wrapper.
+
+Wraps a boto3 bedrock-runtime client (or anything quacking like one) and
+adds adaptive throttle, cost tracking, and JSON-schema parsing-and-repair.
+Uses the Bedrock Converse API (`converse`) which normalizes Anthropic /
+Mistral / Cohere / Llama into one shape.
+
+We keep boto3 as a runtime dependency and let users inject their own client
+for testing.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from typing import Any, TypeVar
+
+from bedrock_kit.cost import CostEntry, CostLedger
+from bedrock_kit.exceptions import JsonParseError
+from bedrock_kit.retry import AdaptiveThrottle
+from bedrock_kit.schema import JsonSchema
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True)
+class Usage:
+    input_tokens: int
+    output_tokens: int
+    cache_read_tokens: int
+    cache_write_tokens: int
+
+    @classmethod
+    def from_converse(cls, raw: dict[str, Any]) -> Usage:
+        usage = raw.get("usage", {}) or {}
+        return cls(
+            input_tokens=int(usage.get("inputTokens", 0)),
+            output_tokens=int(usage.get("outputTokens", 0)),
+            # Bedrock Converse uses cacheReadInputTokens/cacheWriteInputTokens
+            cache_read_tokens=int(usage.get("cacheReadInputTokens", 0)),
+            cache_write_tokens=int(usage.get("cacheWriteInputTokens", 0)),
+        )
+
+
+@dataclass
+class BedrockResponse:
+    text: str
+    usage: Usage
+    model_id: str
+    stop_reason: str
+    cost_usd: float
+    cache_hit: bool
+    parsed: Any = None
+    raw: dict[str, Any] = field(default_factory=dict)
+    cost_entry: CostEntry | None = None
+
+
+def _extract_text_from_converse(raw: dict[str, Any]) -> str:
+    output = raw.get("output", {}) or {}
+    message = output.get("message", {}) or {}
+    parts = []
+    for block in message.get("content", []) or []:
+        text = block.get("text")
+        if text:
+            parts.append(text)
+    return "".join(parts)
+
+
+class _BotoClientFactory:
+    @staticmethod
+    def make(region: str, **client_kwargs: Any) -> Any:
+        try:
+            import boto3  # type: ignore
+        except ImportError as e:
+            raise RuntimeError(
+                "boto3 is required for the default BedrockClient. "
+                "Install with `pip install bedrock-kit[boto]` or pass `client=` explicitly."
+            ) from e
+        return boto3.client("bedrock-runtime", region_name=region, **client_kwargs)
+
+
+class BedrockClient:
+    """Wrapper around bedrock-runtime.converse with cost + retry + schema parsing.
+
+        client = BedrockClient(region="us-east-1")
+        resp = client.invoke(
+            model_id="anthropic.claude-sonnet-4-5",
+            messages=[{"role": "user", "content": "hello"}],
+            system="You are helpful",
+            max_tokens=512,
+        )
+        resp.text
+        resp.cost_usd
+        resp.cache_hit
+
+    Inject a fake for testing:
+
+        client = BedrockClient(client=fake_client)
+    """
+
+    def __init__(
+        self,
+        region: str | None = None,
+        *,
+        client: Any | None = None,
+        retry: AdaptiveThrottle | None = None,
+        cost_ledger: CostLedger | None = None,
+        boto_kwargs: dict[str, Any] | None = None,
+    ) -> None:
+        if client is None:
+            if region is None:
+                raise ValueError("either `region` or `client` must be provided")
+            client = _BotoClientFactory.make(region, **(boto_kwargs or {}))
+        self._client = client
+        self._retry = retry or AdaptiveThrottle()
+        self._ledger = cost_ledger or CostLedger()
+
+    @property
+    def cost_ledger(self) -> CostLedger:
+        return self._ledger
+
+    def invoke(
+        self,
+        model_id: str,
+        messages: Sequence[dict[str, Any]],
+        *,
+        system: str | list[dict[str, Any]] | None = None,
+        max_tokens: int = 1024,
+        temperature: float | None = None,
+        top_p: float | None = None,
+        stop_sequences: list[str] | None = None,
+        response_schema: JsonSchema[T] | None = None,
+        additional_model_fields: dict[str, Any] | None = None,
+    ) -> BedrockResponse:
+        body = self._build_body(
+            messages=messages,
+            system=system,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            stop_sequences=stop_sequences,
+            additional_model_fields=additional_model_fields,
+        )
+
+        raw = self._retry.call(self._client.converse, modelId=model_id, **body)
+        response = self._build_response(model_id, raw)
+
+        if response_schema is not None:
+            response = self._apply_schema(
+                response=response,
+                model_id=model_id,
+                schema=response_schema,
+                body=body,
+            )
+
+        return response
+
+    def _apply_schema(
+        self,
+        response: BedrockResponse,
+        model_id: str,
+        schema: JsonSchema[T],
+        body: dict[str, Any],
+    ) -> BedrockResponse:
+        last_err: Exception | None = None
+        for attempt in range(schema.max_repair_attempts + 1):
+            try:
+                response.parsed = schema.parse(response.text)
+                return response
+            except JsonParseError as e:
+                last_err = e
+                if attempt == schema.max_repair_attempts:
+                    raise
+                # Append a follow-up turn asking the model to fix the output
+                followup_messages = list(body["messages"]) + [
+                    {"role": "assistant", "content": [{"text": response.text}]},
+                    {
+                        "role": "user",
+                        "content": [{"text": schema.repair_prompt(response.text, e)}],
+                    },
+                ]
+                retry_body = dict(body)
+                retry_body["messages"] = followup_messages
+                raw = self._retry.call(self._client.converse, modelId=model_id, **retry_body)
+                response = self._build_response(model_id, raw)
+        # unreachable: the loop either returns or raises
+        raise JsonParseError(
+            f"schema retries exhausted: {last_err}",
+            raw_text=response.text,
+            attempts=schema.max_repair_attempts + 1,
+        )
+
+    def _build_body(
+        self,
+        messages: Sequence[dict[str, Any]],
+        system: str | list[dict[str, Any]] | None,
+        max_tokens: int,
+        temperature: float | None,
+        top_p: float | None,
+        stop_sequences: list[str] | None,
+        additional_model_fields: dict[str, Any] | None,
+    ) -> dict[str, Any]:
+        # Normalize messages: callers may pass {"role":..., "content": "string"}
+        # but Converse expects content as a list of content blocks.
+        norm_messages: list[dict[str, Any]] = []
+        for m in messages:
+            content = m.get("content")
+            if isinstance(content, str):
+                norm_messages.append({"role": m["role"], "content": [{"text": content}]})
+            else:
+                norm_messages.append(dict(m))
+
+        inference_config: dict[str, Any] = {"maxTokens": max_tokens}
+        if temperature is not None:
+            inference_config["temperature"] = temperature
+        if top_p is not None:
+            inference_config["topP"] = top_p
+        if stop_sequences:
+            inference_config["stopSequences"] = list(stop_sequences)
+
+        body: dict[str, Any] = {
+            "messages": norm_messages,
+            "inferenceConfig": inference_config,
+        }
+        if system is not None:
+            if isinstance(system, str):
+                body["system"] = [{"text": system}]
+            else:
+                body["system"] = list(system)
+        if additional_model_fields:
+            body["additionalModelRequestFields"] = additional_model_fields
+        return body
+
+    def _build_response(self, model_id: str, raw: dict[str, Any]) -> BedrockResponse:
+        usage = Usage.from_converse(raw)
+        cost_entry = self._ledger.record(
+            model_id=model_id,
+            input_tokens=usage.input_tokens,
+            output_tokens=usage.output_tokens,
+            cache_read_tokens=usage.cache_read_tokens,
+            cache_write_tokens=usage.cache_write_tokens,
+        )
+        text = _extract_text_from_converse(raw)
+        return BedrockResponse(
+            text=text,
+            usage=usage,
+            model_id=model_id,
+            stop_reason=str(raw.get("stopReason", "")),
+            cost_usd=cost_entry.cost_usd,
+            cache_hit=cost_entry.cache_hit,
+            parsed=None,
+            raw=raw,
+            cost_entry=cost_entry,
+        )

bedrock_kit/cost.py

@@ -0,0 +1,236 @@
+"""CostLedger - per-call cost accounting with cache-aware token math.
+
+Bedrock returns four token counts on Anthropic models:
+  * input_tokens
+  * output_tokens
+  * cache_creation_input_tokens (when cache_control is set on a block)
+  * cache_read_input_tokens (when a previous request hit the cache)
+
+Each has its own price. cache_read is typically 10% of input; cache_creation
+is 1.25x input. We compute total cost from all four. Pricing is in USD per
+1M tokens, matching how AWS publishes the rates.
+
+The default pricing table covers the popular Anthropic models on Bedrock as
+of 2026-Q2. Verify against AWS docs - rates change. You can pass `pricing=`
+to override per-model.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any
+
+from bedrock_kit.exceptions import PricingNotFoundError
+
+
+@dataclass(frozen=True)
+class Pricing:
+    """USD per 1,000,000 tokens. cache_* default to standard Anthropic ratios."""
+
+    input: float
+    output: float
+    cache_read: float | None = None
+    cache_write: float | None = None
+
+    def cost_for(
+        self,
+        input_tokens: int,
+        output_tokens: int,
+        cache_read_tokens: int = 0,
+        cache_write_tokens: int = 0,
+    ) -> float:
+        cache_read_rate = self.cache_read if self.cache_read is not None else self.input * 0.1
+        cache_write_rate = self.cache_write if self.cache_write is not None else self.input * 1.25
+        total = (
+            input_tokens * self.input
+            + output_tokens * self.output
+            + cache_read_tokens * cache_read_rate
+            + cache_write_tokens * cache_write_rate
+        )
+        return total / 1_000_000.0
+
+
+# Default pricing for popular Bedrock models (USD per 1M tokens).
+# Source: aws.amazon.com/bedrock/pricing as of 2026-Q2. VERIFY before billing on these.
+DEFAULT_PRICING: dict[str, Pricing] = {
+    "anthropic.claude-sonnet-4-5": Pricing(
+        input=3.0, output=15.0, cache_read=0.3, cache_write=3.75
+    ),
+    "anthropic.claude-sonnet-4-5-v1:0": Pricing(
+        input=3.0, output=15.0, cache_read=0.3, cache_write=3.75
+    ),
+    "anthropic.claude-opus-4-7": Pricing(
+        input=15.0, output=75.0, cache_read=1.5, cache_write=18.75
+    ),
+    "anthropic.claude-opus-4-7-v1:0": Pricing(
+        input=15.0, output=75.0, cache_read=1.5, cache_write=18.75
+    ),
+    "anthropic.claude-haiku-4-5": Pricing(
+        input=1.0, output=5.0, cache_read=0.1, cache_write=1.25
+    ),
+    "anthropic.claude-3-5-sonnet-20241022-v2:0": Pricing(
+        input=3.0, output=15.0, cache_read=0.3, cache_write=3.75
+    ),
+    "anthropic.claude-3-5-haiku-20241022-v1:0": Pricing(
+        input=0.8, output=4.0, cache_read=0.08, cache_write=1.0
+    ),
+}
+
+
+def _normalize_model_id(model_id: str) -> str:
+    """Strip arn:aws:bedrock:region::foundation-model/ prefix and inference-profile prefix."""
+    if model_id.startswith("arn:aws:bedrock:"):
+        # arn:...:foundation-model/anthropic.claude-...   or .../inference-profile/...
+        tail = model_id.rsplit("/", 1)[-1]
+        return tail
+    if model_id.startswith("us.") or model_id.startswith("eu.") or model_id.startswith("apac."):
+        # cross-region inference profile prefix
+        return model_id.split(".", 1)[1]
+    return model_id
+
+
+@dataclass(frozen=True)
+class CostEntry:
+    timestamp: datetime
+    model_id: str
+    input_tokens: int
+    output_tokens: int
+    cache_read_tokens: int
+    cache_write_tokens: int
+    cost_usd: float
+    cache_hit: bool
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "timestamp": self.timestamp.isoformat(),
+            "model_id": self.model_id,
+            "input_tokens": self.input_tokens,
+            "output_tokens": self.output_tokens,
+            "cache_read_tokens": self.cache_read_tokens,
+            "cache_write_tokens": self.cache_write_tokens,
+            "cost_usd": self.cost_usd,
+            "cache_hit": self.cache_hit,
+            "metadata": self.metadata,
+        }
+
+
+class CostLedger:
+    """Accumulates per-call cost entries.
+
+        ledger = CostLedger()
+        client = BedrockClient(..., cost_ledger=ledger)
+        # ... after some invokes:
+        ledger.total_usd
+        ledger.by_model       # {"anthropic.claude-sonnet-4-5": 0.0234, ...}
+        ledger.entries        # list[CostEntry]
+        ledger.to_dict()
+        ledger.to_pandas()    # requires `pandas` extra
+    """
+
+    def __init__(
+        self,
+        pricing: dict[str, Pricing] | None = None,
+        strict: bool = False,
+    ) -> None:
+        self._pricing = {**DEFAULT_PRICING, **(pricing or {})}
+        self.strict = strict
+        self._entries: list[CostEntry] = []
+
+    def record(
+        self,
+        model_id: str,
+        input_tokens: int,
+        output_tokens: int,
+        cache_read_tokens: int = 0,
+        cache_write_tokens: int = 0,
+        metadata: dict[str, Any] | None = None,
+    ) -> CostEntry:
+        cost = self.compute_cost(
+            model_id, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens
+        )
+        entry = CostEntry(
+            timestamp=datetime.now(timezone.utc),
+            model_id=model_id,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cache_read_tokens=cache_read_tokens,
+            cache_write_tokens=cache_write_tokens,
+            cost_usd=cost,
+            cache_hit=cache_read_tokens > 0,
+            metadata=metadata or {},
+        )
+        self._entries.append(entry)
+        return entry
+
+    def compute_cost(
+        self,
+        model_id: str,
+        input_tokens: int,
+        output_tokens: int,
+        cache_read_tokens: int = 0,
+        cache_write_tokens: int = 0,
+    ) -> float:
+        pricing = self._lookup(model_id)
+        if pricing is None:
+            if self.strict:
+                raise PricingNotFoundError(f"no pricing for {model_id!r}")
+            return 0.0
+        return pricing.cost_for(
+            input_tokens, output_tokens, cache_read_tokens, cache_write_tokens
+        )
+
+    def _lookup(self, model_id: str) -> Pricing | None:
+        if model_id in self._pricing:
+            return self._pricing[model_id]
+        normalized = _normalize_model_id(model_id)
+        return self._pricing.get(normalized)
+
+    @property
+    def entries(self) -> list[CostEntry]:
+        return list(self._entries)
+
+    @property
+    def total_usd(self) -> float:
+        return sum(e.cost_usd for e in self._entries)
+
+    @property
+    def by_model(self) -> dict[str, float]:
+        out: dict[str, float] = {}
+        for e in self._entries:
+            out[e.model_id] = out.get(e.model_id, 0.0) + e.cost_usd
+        return out
+
+    def reset(self) -> None:
+        self._entries.clear()
+
+    def extend(self, entries: Iterable[CostEntry]) -> None:
+        self._entries.extend(entries)
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "total_usd": self.total_usd,
+            "by_model": self.by_model,
+            "n_entries": len(self._entries),
+            "entries": [e.to_dict() for e in self._entries],
+        }
+
+    def to_pandas(self):
+        import pandas as pd
+
+        if not self._entries:
+            return pd.DataFrame(
+                columns=[
+                    "timestamp",
+                    "model_id",
+                    "input_tokens",
+                    "output_tokens",
+                    "cache_read_tokens",
+                    "cache_write_tokens",
+                    "cost_usd",
+                    "cache_hit",
+                ]
+            )
+        return pd.DataFrame([e.to_dict() for e in self._entries])

bedrock_kit/exceptions.py

@@ -0,0 +1,27 @@
+"""bedrock-kit exception hierarchy."""
+
+
+class BedrockKitError(Exception):
+    """Base for all bedrock-kit errors."""
+
+
+class ThrottleExhausted(BedrockKitError):
+    """Raised when AdaptiveThrottle gives up after max_attempts."""
+
+    def __init__(self, attempts: int, last_exc: BaseException):
+        self.attempts = attempts
+        self.last_exc = last_exc
+        super().__init__(f"throttle exhausted after {attempts} attempts: {last_exc}")
+
+
+class PricingNotFoundError(BedrockKitError):
+    """Raised when CostLedger has no pricing for a model and strict=True."""
+
+
+class JsonParseError(BedrockKitError):
+    """Raised when JsonSchema cannot parse/validate a response after retries."""
+
+    def __init__(self, message: str, raw_text: str, attempts: int):
+        self.raw_text = raw_text
+        self.attempts = attempts
+        super().__init__(f"{message} (after {attempts} attempts)")

bedrock_kit/retry.py

@@ -0,0 +1,120 @@
+"""AdaptiveThrottle - exponential backoff with full jitter for Bedrock throttling.
+
+AWS recommends "full jitter" for throttle retries (delay = uniform(0, base * 2^n)).
+This avoids the thundering-herd that "decorrelated jitter" still allows when
+many workers retry against the same throttled region.
+
+We retry on:
+  * boto3 ClientError with ErrorCode in THROTTLE_CODES
+  * any exception in `also_retry` (caller-supplied)
+
+We do NOT retry on validation errors, auth errors, or model-not-found - those
+are bugs to fix, not transient failures.
+"""
+
+from __future__ import annotations
+
+import logging
+import random
+import time
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+from bedrock_kit.exceptions import ThrottleExhausted
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+# Bedrock + STS error codes that mean "back off and retry"
+THROTTLE_CODES: frozenset[str] = frozenset({
+    "ThrottlingException",
+    "Throttling",
+    "TooManyRequestsException",
+    "RequestLimitExceeded",
+    "ServiceUnavailableException",
+    "ServiceUnavailable",
+    "ProvisionedThroughputExceededException",
+    "ModelTimeoutException",
+})
+
+
+def _error_code(exc: BaseException) -> str | None:
+    """Extract AWS error code from a botocore ClientError without importing botocore."""
+    response = getattr(exc, "response", None)
+    if isinstance(response, dict):
+        code = response.get("Error", {}).get("Code")
+        if isinstance(code, str):
+            return code
+    # botocore.exceptions.ClientError also carries operation_name, response.Error.Code
+    return None
+
+
+class AdaptiveThrottle:
+    """Retry a callable with exponential backoff + full jitter.
+
+    Construct once, reuse across many calls:
+
+        throttle = AdaptiveThrottle(max_attempts=8, base_delay=0.5, max_delay=30.0)
+        result = throttle.call(client.invoke_model, body=...)
+    """
+
+    def __init__(
+        self,
+        max_attempts: int = 6,
+        base_delay: float = 1.0,
+        max_delay: float = 30.0,
+        jitter: bool = True,
+        also_retry: tuple[type[BaseException], ...] = (),
+        sleep_fn: Callable[[float], None] | None = None,
+        rng: random.Random | None = None,
+    ) -> None:
+        if max_attempts < 1:
+            raise ValueError("max_attempts must be >= 1")
+        if base_delay <= 0:
+            raise ValueError("base_delay must be > 0")
+        if max_delay < base_delay:
+            raise ValueError("max_delay must be >= base_delay")
+        self.max_attempts = max_attempts
+        self.base_delay = base_delay
+        self.max_delay = max_delay
+        self.jitter = jitter
+        self.also_retry = also_retry
+        self._sleep = sleep_fn or time.sleep
+        self._rng = rng or random.Random()
+
+    def _should_retry(self, exc: BaseException) -> bool:
+        if isinstance(exc, self.also_retry):
+            return True
+        code = _error_code(exc)
+        return code in THROTTLE_CODES
+
+    def _delay_for(self, attempt_index: int) -> float:
+        """attempt_index is 0-based: 0 is the first retry sleep, 1 is the second, ..."""
+        capped = min(self.base_delay * (2**attempt_index), self.max_delay)
+        if self.jitter:
+            return self._rng.uniform(0.0, capped)
+        return capped
+
+    def call(self, fn: Callable[..., T], *args: Any, **kwargs: Any) -> T:
+        last_exc: BaseException | None = None
+        for attempt in range(self.max_attempts):
+            try:
+                return fn(*args, **kwargs)
+            except Exception as exc:  # noqa: BLE001 - we re-raise selectively
+                if not self._should_retry(exc):
+                    raise
+                last_exc = exc
+                if attempt == self.max_attempts - 1:
+                    break
+                delay = self._delay_for(attempt)
+                logger.debug(
+                    "throttle: attempt %d/%d failed (%s); sleeping %.2fs",
+                    attempt + 1,
+                    self.max_attempts,
+                    _error_code(exc) or type(exc).__name__,
+                    delay,
+                )
+                self._sleep(delay)
+        assert last_exc is not None
+        raise ThrottleExhausted(self.max_attempts, last_exc) from last_exc

bedrock_kit/schema.py

@@ -0,0 +1,156 @@
+"""JsonSchema - parse + validate + repair JSON output from Bedrock.
+
+Models still emit invalid JSON sometimes: trailing commas, smart quotes,
+markdown fences around the JSON. We do three passes before giving up:
+
+  1. json.loads on the raw text
+  2. lightweight repair: strip ```json fences, strip leading/trailing prose,
+     trim trailing comma before } or ]
+  3. ask the LLM again with the parse error as context (if `repair=True`
+     and a retry callback is provided)
+
+Each pass is opt-in. The retry-with-LLM step requires the caller to wire it
+in via BedrockClient.invoke; standalone JsonSchema only does passes 1 and 2.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from collections.abc import Callable
+from typing import Any, Generic, TypeVar
+
+from bedrock_kit.exceptions import JsonParseError
+
+T = TypeVar("T")
+
+
+_FENCE_RE = re.compile(r"^\s*```(?:json|JSON)?\s*\n?|\n?```\s*$", re.MULTILINE)
+_TRAILING_COMMA_RE = re.compile(r",(\s*[}\]])")
+
+
+def _strip_fences(text: str) -> str:
+    return _FENCE_RE.sub("", text).strip()
+
+
+def _extract_first_json_object(text: str) -> str:
+    """Find the first balanced {...} or [...] in the text. Useful when the model
+    wraps JSON in prose like 'Here you go: { ... } let me know if...'."""
+    if not text:
+        return text
+    open_chars = {"{": "}", "[": "]"}
+    for i, ch in enumerate(text):
+        if ch in open_chars:
+            close = open_chars[ch]
+            depth = 0
+            in_str = False
+            escape = False
+            for j in range(i, len(text)):
+                c = text[j]
+                if escape:
+                    escape = False
+                    continue
+                if c == "\\":
+                    escape = True
+                    continue
+                if c == '"':
+                    in_str = not in_str
+                    continue
+                if in_str:
+                    continue
+                if c == ch:
+                    depth += 1
+                elif c == close:
+                    depth -= 1
+                    if depth == 0:
+                        return text[i : j + 1]
+            break
+    return text
+
+
+def _light_repair(text: str) -> str:
+    text = _strip_fences(text)
+    text = _extract_first_json_object(text)
+    text = _TRAILING_COMMA_RE.sub(r"\1", text)
+    return text.strip()
+
+
+def parse_json(text: str, *, repair: bool = True) -> Any:
+    """Parse JSON from `text`. With repair=True, try a light repair pass on failure."""
+    try:
+        return json.loads(text)
+    except json.JSONDecodeError:
+        if not repair:
+            raise
+        repaired = _light_repair(text)
+        return json.loads(repaired)
+
+
+class JsonSchema(Generic[T]):
+    """Validate-and-repair wrapper around a pydantic model.
+
+        from pydantic import BaseModel
+
+        class Answer(BaseModel):
+            text: str
+            confidence: float
+
+        schema = JsonSchema(Answer, repair=True, max_repair_attempts=2)
+        parsed: Answer = schema.parse(response_text)
+
+    `max_repair_attempts` is for the LLM-retry path used by BedrockClient.invoke.
+    Standalone .parse() does the JSON repair pass once; if you want the LLM
+    retry behavior, pass the schema to client.invoke(..., response_schema=schema).
+    """
+
+    def __init__(
+        self,
+        model: type,
+        *,
+        repair: bool = True,
+        max_repair_attempts: int = 2,
+        validator: Callable[[Any], T] | None = None,
+    ) -> None:
+        self.model = model
+        self.repair = repair
+        self.max_repair_attempts = max_repair_attempts
+        self._validator = validator or self._default_validator(model)
+
+    @staticmethod
+    def _default_validator(model: type) -> Callable[[Any], Any]:
+        # pydantic v2: model.model_validate
+        if hasattr(model, "model_validate"):
+            return model.model_validate  # type: ignore[no-any-return]
+        # dataclass-style or callable
+        return lambda obj: model(**obj) if isinstance(obj, dict) else model(obj)
+
+    def parse(self, text: str) -> T:
+        """Parse one response. Does NOT call the LLM again - that's BedrockClient's job."""
+        attempts = 0
+        last_err: Exception | None = None
+        candidates = [text]
+        if self.repair:
+            candidates.append(_light_repair(text))
+        for candidate in candidates:
+            attempts += 1
+            try:
+                obj = json.loads(candidate)
+                return self._validator(obj)
+            except Exception as e:  # noqa: BLE001 - we collect and re-raise
+                last_err = e
+                continue
+        raise JsonParseError(
+            f"could not parse response into {self.model.__name__}: {last_err}",
+            raw_text=text,
+            attempts=attempts,
+        )
+
+    def repair_prompt(self, raw: str, error: Exception) -> str:
+        """Build a follow-up prompt for the LLM to fix its output."""
+        return (
+            "Your previous response could not be parsed as valid JSON for the "
+            f"{self.model.__name__} schema. The error was:\n\n{error}\n\n"
+            f"Your previous output was:\n\n{raw}\n\n"
+            "Return ONLY a valid JSON object that matches the schema. No prose, "
+            "no markdown fences."
+        )

bedrock_kit-0.1.0.dist-info/METADATA

@@ -0,0 +1,188 @@
+Metadata-Version: 2.4
+Name: bedrock-kit
+Version: 0.1.0
+Summary: Small, opinionated AWS Bedrock client wrapper: adaptive throttle, cache-aware cost tracking, and structured-output parse-and-repair. Single-cloud, single-purpose.
+Project-URL: Homepage, https://github.com/MukundaKatta/bedrock-kit
+Project-URL: Issues, https://github.com/MukundaKatta/bedrock-kit/issues
+Project-URL: Source, https://github.com/MukundaKatta/bedrock-kit
+Author-email: Mukunda Rao Katta <mukunda.vjcs6@gmail.com>
+License: MIT
+License-File: LICENSE
+Keywords: anthropic,aws,bedrock,claude,cost-tracking,llm,mlops,rag,retry,structured-output
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.10
+Requires-Dist: pydantic>=2.0
+Provides-Extra: boto
+Requires-Dist: boto3>=1.28; extra == 'boto'
+Provides-Extra: dev
+Requires-Dist: pandas>=2.0; extra == 'dev'
+Requires-Dist: pytest>=8.0; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'
+Provides-Extra: pandas
+Requires-Dist: pandas>=2.0; extra == 'pandas'
+Description-Content-Type: text/markdown
+
+# bedrock-kit
+
+[![CI](https://github.com/MukundaKatta/bedrock-kit/actions/workflows/ci.yml/badge.svg)](https://github.com/MukundaKatta/bedrock-kit/actions/workflows/ci.yml)
+[![PyPI](https://img.shields.io/pypi/v/bedrock-kit.svg)](https://pypi.org/project/bedrock-kit/)
+[![Python](https://img.shields.io/pypi/pyversions/bedrock-kit.svg)](https://pypi.org/project/bedrock-kit/)
+[![License: MIT](https://img.shields.io/badge/license-MIT-green.svg)](LICENSE)
+
+**Small, opinionated AWS Bedrock client wrapper.**
+
+Every Bedrock production team rebuilds the same three things: adaptive
+throttle for `ThrottlingException`, per-call cost tracking that handles
+cache-read tokens, and structured-output parsing-with-repair. `bedrock-kit`
+ships those, and nothing else. Single-cloud, single-purpose. No proxy
+server. Wraps the official boto3 client; you can inject a fake for testing.
+
+## Install
+
+```bash
+pip install "bedrock-kit[boto]"
+# optional
+pip install "bedrock-kit[boto,pandas]"  # adds CostLedger.to_pandas()
+```
+
+## Quickstart
+
+```python
+from bedrock_kit import BedrockClient, AdaptiveThrottle, CostLedger
+
+ledger = CostLedger()
+client = BedrockClient(
+    region="us-east-1",
+    retry=AdaptiveThrottle(max_attempts=8, base_delay=1.0, max_delay=30.0),
+    cost_ledger=ledger,
+)
+
+resp = client.invoke(
+    model_id="anthropic.claude-sonnet-4-5",
+    messages=[{"role": "user", "content": "Hello"}],
+    system="You are concise.",
+    max_tokens=512,
+)
+
+resp.text                       # the model's reply
+resp.usage.input_tokens         # 12
+resp.usage.cache_read_tokens    # 0
+resp.cost_usd                   # 0.000176
+resp.cache_hit                  # False
+ledger.total_usd                # accumulates across all calls
+ledger.by_model                 # {"anthropic.claude-sonnet-4-5": 0.000176}
+```
+
+## Structured output with repair
+
+```python
+from pydantic import BaseModel
+from bedrock_kit import JsonSchema
+
+class Sentiment(BaseModel):
+    label: str
+    confidence: float
+
+resp = client.invoke(
+    model_id="anthropic.claude-sonnet-4-5",
+    messages=[{"role": "user", "content": "Classify: 'this is great!'"}],
+    response_schema=JsonSchema(Sentiment, max_repair_attempts=2),
+)
+resp.parsed                  # Sentiment(label="positive", confidence=0.95)
+```
+
+If the model returns invalid JSON, `bedrock-kit` first does a light local
+repair pass (strip markdown fences, trailing commas, surrounding prose).
+If that still fails, it asks the model to fix its own output, up to
+`max_repair_attempts` times.
+
+## Adaptive throttle
+
+```python
+throttle = AdaptiveThrottle(
+    max_attempts=8,    # total attempts (incl. first)
+    base_delay=1.0,    # seconds
+    max_delay=30.0,
+    jitter=True,       # full-jitter (uniform(0, capped_delay))
+)
+```
+
+Retries on Bedrock throttle codes: `ThrottlingException`,
+`TooManyRequestsException`, `ServiceUnavailableException`,
+`ProvisionedThroughputExceededException`, `ModelTimeoutException`. Does
+**not** retry validation, auth, or model-not-found errors - those are
+your bugs to fix, not transient.
+
+## Cost tracking
+
+`CostLedger` ships pricing for popular Anthropic Bedrock models. Override
+or extend with `pricing=`:
+
+```python
+from bedrock_kit import CostLedger, Pricing
+
+ledger = CostLedger(
+    pricing={
+        "amazon.nova-pro-v1:0": Pricing(
+            input=0.8, output=3.2, cache_read=0.2, cache_write=1.0
+        ),
+    },
+    strict=True,  # raise PricingNotFoundError on unknown models
+)
+```
+
+Default pricing is best-effort and dated; verify against
+[aws.amazon.com/bedrock/pricing](https://aws.amazon.com/bedrock/pricing)
+before using these numbers for billing.
+
+## Why not LiteLLM?
+
+LiteLLM is great if you need cross-provider routing. `bedrock-kit` is for
+the case where you've already decided on Bedrock, you don't want a 46k-LOC
+multi-provider abstraction, and you want a small surface a security team
+can audit. We don't proxy, don't include a server, don't ship a UI. We're
+< 1000 LOC of Python.
+
+## What it explicitly does NOT do
+
+- No multi-provider routing
+- No proxy server, no UI
+- No prompt management
+- No agent loop
+- No image generation
+- No SageMaker, Bedrock Agents, or Knowledge Bases SDK wrapping
+- No streaming or cancellation yet (planned for v0.2)
+- No OpenTelemetry emission yet (planned for v0.2)
+
+## Testing without AWS
+
+The default `BedrockClient` makes a real boto3 client. For tests, inject a
+fake that quacks like one:
+
+```python
+from bedrock_kit import BedrockClient, AdaptiveThrottle
+
+class FakeClient:
+    def converse(self, **kwargs):
+        return {"output": {"message": {"content": [{"text": "stub"}]}},
+                "stopReason": "end_turn",
+                "usage": {"inputTokens": 1, "outputTokens": 1}}
+
+client = BedrockClient(client=FakeClient(), retry=AdaptiveThrottle(sleep_fn=lambda _: None))
+```
+
+## Status
+
+v0.1 - alpha. Public API may change before v1.0. Issues and PRs welcome.

bedrock_kit-0.1.0.dist-info/RECORD

@@ -0,0 +1,10 @@
+bedrock_kit/__init__.py,sha256=JmrMV-0j1RxQiJpY4iA2dg7pVhHwTulqiiu9E8Ra2c0,930
+bedrock_kit/client.py,sha256=3_1wjyd6hMiBjR5BzzdNIPL-EttuENMw2O5eL1HsKEc,8718
+bedrock_kit/cost.py,sha256=2dgzajQQFXPwWHXIS7GAJW3B-2U9p4zYd5FnCEQfajY,7772
+bedrock_kit/exceptions.py,sha256=35vZlOvMOvSOyBrxYYeJ8JfQBuUzdNHtRBe0zIGgTU0,895
+bedrock_kit/retry.py,sha256=cyJhqZUhq0Sj7eYequg1KpbFEA3m94S5Nk8Q7rfyvM0,4228
+bedrock_kit/schema.py,sha256=TfDpEyS0oiDeXL5jrUgKSU5uGcNCqji4Ed4O8UAlvPQ,5273
+bedrock_kit-0.1.0.dist-info/METADATA,sha256=jAlkIfLLvMslu5RlCPhvf5zGrz0U67eHYhllebKpfeo,6571
+bedrock_kit-0.1.0.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+bedrock_kit-0.1.0.dist-info/licenses/LICENSE,sha256=p1GujHnprYaKo-fuZc9Tpy9i711QOy8PeYBhNM0VOdw,1074
+bedrock_kit-0.1.0.dist-info/RECORD,,

bedrock_kit-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.29.0
+Root-Is-Purelib: true
+Tag: py3-none-any

bedrock_kit-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Mukunda Rao Katta
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.