aimai-kit 1.0.0__py3-none-any.whl

aimai_kit/__init__.py

@@ -0,0 +1 @@
+"""aimai-kit — a framework-free LLM engineering toolkit."""

aimai_kit/agent/__init__.py

@@ -0,0 +1,28 @@
+"""Agent layer: a bounded loop over the tool layer.
+
+The loop owns very little on purpose. Tool execution, request construction,
+telemetry and loop detection each live behind their own seam, so each can be
+tested without starting a run — and so the loop stays short enough to read in
+one sitting.
+"""
+
+from .budget import Budgets, Spend, StopReason
+from .loop import Agent, RunResult
+from .loopdetect import LoopDetector
+from .metrics import RunMetrics, run_metrics
+from .scripted import ScriptedClient
+from .thread import Thread, TraceEvent
+
+__all__ = [
+    "Agent",
+    "Budgets",
+    "LoopDetector",
+    "RunMetrics",
+    "RunResult",
+    "ScriptedClient",
+    "Spend",
+    "StopReason",
+    "Thread",
+    "TraceEvent",
+    "run_metrics",
+]

aimai_kit/agent/budget.py

@@ -0,0 +1,118 @@
+"""Four budgets and one stop reason.
+
+An agent that runs until it decides it is finished is not a system, it is a
+bet. Four budgets bound it, and they bound different failure modes:
+
+    steps    a loop that makes progress but never converges
+    tokens   a loop that grows its own context until the window bursts
+    cost     the one the finance team asks about
+    seconds  the one the user experiences
+
+Three of them are cheap to check and one is not: cost requires a pricing
+catalog, so it is optional. A missing catalog does not silently disable the
+budget — `Spend.usd` simply stays at zero and the cost budget never fires,
+which is visible in the metrics.
+
+`StopReason` is a closed set because it is what the caller branches on. Half
+of these values mean "the work is not done"; only `finished` means it is.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from decimal import Decimal
+from enum import StrEnum
+
+__all__ = ["Budgets", "Spend", "StopReason"]
+
+
+class StopReason(StrEnum):
+    """Why a run ended. Exactly one value per run."""
+
+    FINISHED = "finished"
+    STEP_BUDGET = "step_budget"
+    TOKEN_BUDGET = "token_budget"
+    COST_BUDGET = "cost_budget"
+    TIME_BUDGET = "time_budget"
+    LOOP_DETECTED = "loop_detected"
+    TOOL_ERROR_CEILING = "tool_error_ceiling"
+    NEEDS_APPROVAL = "needs_approval"
+    CANCELLED = "cancelled"
+
+    @property
+    def is_success(self) -> bool:
+        return self is StopReason.FINISHED
+
+    @property
+    def is_budget(self) -> bool:
+        return self in {
+            StopReason.STEP_BUDGET,
+            StopReason.TOKEN_BUDGET,
+            StopReason.COST_BUDGET,
+            StopReason.TIME_BUDGET,
+        }
+
+
+@dataclass(frozen=True, slots=True)
+class Budgets:
+    """The ceilings for one run.
+
+    `max_seconds` is wall-clock, not CPU: what bounds a user's patience is
+    the clock, and most of the time here is spent waiting on a network call
+    anyway.
+
+    Leaving a budget at `None` disables it. That is allowed but should be
+    deliberate — an agent with no step budget will eventually find an input
+    that makes it run forever.
+    """
+
+    max_steps: int | None = 12
+    max_tokens: int | None = 120_000
+    max_usd: Decimal | None = None
+    max_seconds: float | None = 120.0
+
+    def exceeded_by(self, spend: Spend) -> StopReason | None:
+        """Which budget, if any, this spend has already crossed.
+
+        Checked BEFORE a step, not after: knowing you are over budget after
+        spending the money is an audit, not a budget.
+        """
+        if self.max_steps is not None and spend.steps >= self.max_steps:
+            return StopReason.STEP_BUDGET
+        if self.max_tokens is not None and spend.tokens >= self.max_tokens:
+            return StopReason.TOKEN_BUDGET
+        if self.max_usd is not None and spend.usd >= self.max_usd:
+            return StopReason.COST_BUDGET
+        if self.max_seconds is not None and spend.seconds >= self.max_seconds:
+            return StopReason.TIME_BUDGET
+        return None
+
+
+@dataclass
+class Spend:
+    """What a run has consumed so far."""
+
+    steps: int = 0
+    tokens: int = 0
+    usd: Decimal = field(default_factory=lambda: Decimal(0))
+    seconds: float = 0.0
+    tool_errors: int = 0
+
+    def as_dict(self) -> dict[str, object]:
+        return {
+            "steps": self.steps,
+            "tokens": self.tokens,
+            "usd": str(self.usd),
+            "seconds": round(self.seconds, 3),
+            "tool_errors": self.tool_errors,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> Spend:
+        return cls(
+            steps=int(data.get("steps", 0)),
+            tokens=int(data.get("tokens", 0)),
+            usd=Decimal(str(data.get("usd", "0"))),
+            seconds=float(data.get("seconds", 0.0)),
+            tool_errors=int(data.get("tool_errors", 0)),
+        )

aimai_kit/agent/loop.py

@@ -0,0 +1,296 @@
+"""The agent loop — framework-free, four budgets, one stop reason.
+
+The loop itself is short, and that is the design. Everything it could have
+absorbed lives elsewhere:
+
+    tool execution      the tool layer's executor (five gates, timeouts)
+    request building    the prompt layer (registry, budget, cache prefix)
+    telemetry           the provider layer's resilient client
+    loop detection      a separate object with its own tests
+
+What is left here is the part that is genuinely about the loop: check the
+budget, call the model, hand any tool calls to the executor, feed the results
+back, repeat.
+
+Two decisions are worth reading the code for.
+
+**Every tool call gets a result.** Even a refused one, even a timed-out one.
+An assistant turn that requested three calls and received two results leaves
+the conversation in a shape most providers reject, and the ones that do not
+reject it produce confused output. The loop makes half-turns impossible.
+
+**A budget stop gets one last turn without tools.** Cutting the run dead at
+the ceiling leaves the user with nothing. Giving the model one final call with
+tools disabled produces a partial answer that says what is missing, which is
+almost always more useful — and the cost of that last turn has to be inside
+the ceiling you chose, not on top of it.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+
+from ..provider.client import LLMClient
+from ..provider.counters import COUNTERS
+from ..provider.pricing import ModelPricing
+from ..provider.types import ChatRequest, Role, ToolCall
+from ..tools.executor import CallContext, ToolExecutor
+from ..tools.export import export_for
+from ..tools.idempotency import call_signature
+from .budget import Budgets, StopReason
+from .loopdetect import REPEAT_WARNING
+from .thread import Thread, TraceEvent
+
+__all__ = ["Agent", "RunResult"]
+
+
+@dataclass
+class RunResult:
+    """What a finished run produced. The thread carries the details."""
+
+    thread: Thread
+    stop_reason: StopReason
+    answer: str
+
+    @property
+    def ok(self) -> bool:
+        return self.stop_reason is StopReason.FINISHED
+
+
+class Agent:
+    """A stateless agent. All state lives in the `Thread` passed to `run`.
+
+    Stateless means one instance can serve concurrent requests, and it means
+    a run can be reconstructed from its thread alone rather than from an
+    object graph that no longer exists.
+    """
+
+    def __init__(
+        self,
+        client: LLMClient,
+        executor: ToolExecutor,
+        *,
+        system: str = "",
+        budgets: Budgets | None = None,
+        pricing: dict[str, ModelPricing] | None = None,
+        allowlist: Sequence[str] | None = None,
+        max_tool_errors: int = 4,
+        max_output_tokens: int = 2048,
+        on_event: Callable[[TraceEvent], None] | None = None,
+        now: Callable[[], float] = time.monotonic,
+    ) -> None:
+        self.client = client
+        self.executor = executor
+        self.system = system
+        self.budgets = budgets or Budgets()
+        self.pricing = pricing or {}
+        self.allowlist = allowlist
+        self.max_tool_errors = max_tool_errors
+        self.max_output_tokens = max_output_tokens
+        self.on_event = on_event
+        self._now = now
+
+    # --- internals --------------------------------------------------------
+
+    def _emit(
+        self, thread: Thread, kind: str, detail: str = "", *, ok: bool = True
+    ) -> None:
+        event = thread.record(kind, detail, ok=ok)
+        if self.on_event:
+            self.on_event(event)
+
+    def _tool_declarations(self) -> list[dict]:
+        visible = self.executor.registry.visible(self.allowlist)
+        return export_for(self.client.provider, visible)
+
+    def _charge(self, thread: Thread, result) -> None:
+        """Accounting happens in ONE place in the loop, not per call site."""
+        thread.spend.tokens += result.usage.total_tokens
+        price = self.pricing.get(result.model)
+        if price is not None:
+            thread.spend.usd += price.cost(
+                result.usage.input_tokens,
+                result.usage.output_tokens,
+                result.usage.cached_input_tokens,
+            )
+
+    def _request(self, thread: Thread, *, with_tools: bool) -> ChatRequest:
+        return ChatRequest(
+            messages=list(thread.messages),
+            system=self.system or None,
+            max_output_tokens=self.max_output_tokens,
+            tools=self._tool_declarations() if with_tools else (),
+            tool_choice=None if with_tools else "none",
+            operation="agent",
+        )
+
+    def _final_turn(self, thread: Thread, reason: StopReason) -> StopReason:
+        """One last call with tools disabled, so the user gets something.
+
+        The model is told what happened and asked to answer with what it has.
+        A partial answer that names the gap beats an empty response, and it
+        beats an answer that pretends the gap is not there.
+        """
+        thread.add(
+            Role.USER,
+            "You have reached the limit for this task "
+            f"({reason.value}). Do not call any more tools. Answer with what "
+            "you already have, and state explicitly what is still missing.",
+        )
+        try:
+            result = self.client.complete(self._request(thread, with_tools=False))
+        except Exception as error:  # noqa: BLE001 - the run is already ending
+            self._emit(thread, "final_turn", str(error), ok=False)
+            return reason
+        self._charge(thread, result)
+        thread.add(Role.ASSISTANT, result.text)
+        thread.answer = result.text
+        # `tool_choice="none"` should make this impossible. If a provider
+        # ignores it, the run ends with an empty answer and that has to be
+        # visible in the trace rather than look like a normal stop.
+        self._emit(
+            thread,
+            "final_turn",
+            reason.value,
+            ok=bool(result.text) and not result.wants_tools,
+        )
+        return reason
+
+    def _run_tool_calls(
+        self, thread: Thread, calls: Sequence[ToolCall], ctx: CallContext
+    ) -> StopReason | None:
+        """Execute a turn's tool calls and feed every result back.
+
+        Returns a stop reason when the turn itself ends the run (a loop, an
+        approval gate, or too many tool errors); otherwise None.
+        """
+        pairs = [(c.name, c.arguments) for c in calls]
+        results = self.executor.call_many(pairs, ctx, allowlist=self.allowlist)
+
+        stop: StopReason | None = None
+        for call, result in zip(calls, results, strict=True):
+            signature = call_signature(call.name, call.arguments)
+            occurrences = thread.detector.observe(signature)
+
+            content = result.content
+            if thread.detector.should_warn(occurrences):
+                # The warning goes INTO the tool result, where the model is
+                # already looking, rather than into a separate system message
+                # it may not weigh the same way.
+                content += REPEAT_WARNING
+                self._emit(thread, "repeat_warning", call.name)
+            elif thread.detector.should_stop(occurrences):
+                stop = stop or StopReason.LOOP_DETECTED
+                self._emit(thread, "loop_detected", call.name, ok=False)
+
+            if not result.ok:
+                thread.spend.tool_errors += 1
+                if result.error_code == "needs_approval":
+                    thread.pending_approval = signature
+                    stop = stop or StopReason.NEEDS_APPROVAL
+
+            if result.ok and self.executor.registry.get(call.name) is not None:
+                spec = self.executor.registry.get(call.name)
+                if spec and spec.side_effect:
+                    thread.tool_results[signature] = result.content
+
+            # Every call gets a result, including refused and failed ones.
+            thread.add(
+                Role.TOOL,
+                json.dumps(
+                    {
+                        "tool_call_id": call.id,
+                        "tool": call.name,
+                        "ok": result.ok,
+                        "error_code": result.error_code,
+                        "content": content,
+                    },
+                    ensure_ascii=False,
+                ),
+            )
+            self._emit(
+                thread,
+                "tool_result",
+                f"{call.name}:{result.error_code or 'ok'}",
+                ok=result.ok,
+            )
+
+        if thread.spend.tool_errors >= self.max_tool_errors:
+            stop = stop or StopReason.TOOL_ERROR_CEILING
+        return stop
+
+    # --- public API -------------------------------------------------------
+
+    def run(
+        self, thread: Thread, task: str | None = None, *, ctx: CallContext | None = None
+    ) -> RunResult:
+        """Run until the task is done or a budget stops it.
+
+        Passing an existing thread resumes it; that is what makes checkpoint
+        recovery ordinary rather than special.
+        """
+        ctx = ctx or CallContext()
+        started = self._now()
+        if task:
+            thread.goal = thread.goal or task
+            thread.add(Role.USER, task)
+
+        stop: StopReason | None = None
+        while True:
+            thread.spend.seconds = self._now() - started
+            over = self.budgets.exceeded_by(thread.spend)
+            if over is not None:
+                COUNTERS.increment("agent_budget_stops_total", reason=over.value)
+                stop = self._final_turn(thread, over)
+                break
+
+            thread.step += 1
+            thread.spend.steps += 1
+            self._emit(thread, "step", f"step {thread.step}")
+
+            try:
+                result = self.client.complete(self._request(thread, with_tools=True))
+            except Exception as error:  # noqa: BLE001 - surface, do not crash
+                self._emit(thread, "model_error", str(error), ok=False)
+                stop = StopReason.CANCELLED
+                break
+
+            self._charge(thread, result)
+
+            if not result.wants_tools:
+                thread.add(Role.ASSISTANT, result.text)
+                thread.answer = result.text
+                stop = StopReason.FINISHED
+                self._emit(thread, "finished", "no further tool calls")
+                break
+
+            thread.add(
+                Role.ASSISTANT,
+                result.text
+                or json.dumps(
+                    [
+                        {"tool": c.name, "arguments": c.arguments}
+                        for c in result.tool_calls
+                    ],
+                    ensure_ascii=False,
+                ),
+            )
+            turn_stop = self._run_tool_calls(thread, result.tool_calls, ctx)
+            if turn_stop is not None:
+                if turn_stop is StopReason.NEEDS_APPROVAL:
+                    stop = turn_stop
+                else:
+                    stop = self._final_turn(thread, turn_stop)
+                break
+
+        thread.spend.seconds = self._now() - started
+        thread.stop_reason = (stop or StopReason.FINISHED).value
+        COUNTERS.increment("agent_runs_total", reason=thread.stop_reason)
+        return RunResult(
+            thread=thread,
+            stop_reason=StopReason(thread.stop_reason),
+            answer=thread.answer,
+        )

aimai_kit/agent/loopdetect.py

@@ -0,0 +1,70 @@
+"""Signature-based loop detection.
+
+The failure mode this catches is specific and common: an agent calls
+`get_order(id=42)`, does not like the answer, and calls `get_order(id=42)`
+again. Nothing errors. Every step looks healthy. The budget drains.
+
+Detection needs a stable identity for "the same call", which the tool layer
+already provides — the call signature is the tool name plus normalized
+arguments, so argument reordering does not disguise a repeat.
+
+The response is graduated, and that matters:
+
+    2nd occurrence  a warning is injected INTO the tool result
+    3rd occurrence  the run stops with `loop_detected`
+
+Warning before stopping is the useful part. A model that is told "this call
+returned the same result as before; try a different approach" frequently does
+try a different approach, and the run completes. Stopping at the second
+occurrence would kill runs that were about to recover; never stopping turns a
+stuck run into a full budget burn.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+__all__ = ["LoopDetector", "REPEAT_WARNING"]
+
+REPEAT_WARNING = (
+    "\n\n[note: this exact call was already made earlier in this run and "
+    "returned the same result. Repeating it will not produce new information. "
+    "Use a different tool, different arguments, or answer with what you have.]"
+)
+
+
+@dataclass
+class LoopDetector:
+    """Counts repeated call signatures within one run."""
+
+    warn_at: int = 2
+    stop_at: int = 3
+    counts: dict[str, int] = field(default_factory=dict)
+
+    def observe(self, signature: str) -> int:
+        """Record an occurrence and return how many times it has been seen."""
+        self.counts[signature] = self.counts.get(signature, 0) + 1
+        return self.counts[signature]
+
+    def should_warn(self, occurrences: int) -> bool:
+        return self.warn_at <= occurrences < self.stop_at
+
+    def should_stop(self, occurrences: int) -> bool:
+        return occurrences >= self.stop_at
+
+    @property
+    def repeated_signatures(self) -> list[str]:
+        return [sig for sig, count in self.counts.items() if count > 1]
+
+    def as_dict(self) -> dict[str, int]:
+        return dict(self.counts)
+
+    @classmethod
+    def from_dict(
+        cls, data: dict, *, warn_at: int = 2, stop_at: int = 3
+    ) -> LoopDetector:
+        return cls(
+            warn_at=warn_at,
+            stop_at=stop_at,
+            counts={k: int(v) for k, v in data.items()},
+        )

aimai_kit/agent/metrics.py

@@ -0,0 +1,112 @@
+"""Run metrics, and the order in which to read them.
+
+Seven numbers, and reading them in the wrong order wastes days. The order is:
+
+    1. loop_rate          are runs getting stuck?
+    2. budget_stop_rate   are they being cut off?
+    3. recovery_rate      do they survive a tool failure?
+
+Loop rate comes first because a loop inflates everything downstream: steps,
+cost, and the budget stop rate. Tuning budgets while the loop rate is high
+means raising the ceiling on wasted work.
+
+Budget stop rate comes second. A high rate with a low loop rate means the
+budget is genuinely too tight; the same rate with a high loop rate means it is
+doing its job.
+
+Recovery rate comes last and is the most interesting: it is the share of runs
+that hit a tool error and still finished. A low recovery rate says the error
+messages the tool layer produces are not actionable — which is a prompt and
+tool-description problem, not a budget one.
+"""
+
+from __future__ import annotations
+
+import statistics
+from collections.abc import Sequence
+from dataclasses import dataclass
+from decimal import Decimal
+
+from .budget import StopReason
+from .thread import Thread
+
+__all__ = ["RunMetrics", "run_metrics"]
+
+
+@dataclass
+class RunMetrics:
+    runs: int = 0
+    completion_rate: float = 0.0
+    steps_p50: float = 0.0
+    steps_p95: float = 0.0
+    tool_error_rate: float = 0.0
+    loop_rate: float = 0.0
+    budget_stop_rate: float = 0.0
+    recovery_rate: float = 0.0
+    mean_usd: str = "0"
+
+    def as_dict(self) -> dict[str, object]:
+        return {
+            "runs": self.runs,
+            "completion_rate": self.completion_rate,
+            "steps_p50": self.steps_p50,
+            "steps_p95": self.steps_p95,
+            "tool_error_rate": self.tool_error_rate,
+            "loop_rate": self.loop_rate,
+            "budget_stop_rate": self.budget_stop_rate,
+            "recovery_rate": self.recovery_rate,
+            "mean_usd": self.mean_usd,
+        }
+
+
+def _percentile(values: Sequence[float], p: int) -> float:
+    if not values:
+        return 0.0
+    ordered = sorted(values)
+    index = max(0, min(len(ordered) - 1, round(p / 100 * len(ordered)) - 1))
+    return ordered[index]
+
+
+def run_metrics(threads: Sequence[Thread]) -> RunMetrics:
+    """Aggregate a batch of finished runs."""
+    if not threads:
+        return RunMetrics()
+
+    total = len(threads)
+    steps = [float(t.step) for t in threads]
+    finished = sum(t.stop_reason == StopReason.FINISHED for t in threads)
+    looped = sum(t.stop_reason == StopReason.LOOP_DETECTED for t in threads)
+    budget_stopped = sum(
+        StopReason(t.stop_reason).is_budget for t in threads if t.stop_reason
+    )
+
+    tool_events = [e for t in threads for e in t.trace if e.kind == "tool_result"]
+    tool_errors = sum(not e.ok for e in tool_events)
+
+    # A run "recovered" if it saw a tool error and still finished. Runs that
+    # never hit an error are excluded from the denominator — including them
+    # would make the rate a measure of how rarely tools fail.
+    had_errors = [
+        t for t in threads if any(e.kind == "tool_result" and not e.ok for e in t.trace)
+    ]
+    recovered = sum(t.stop_reason == StopReason.FINISHED for t in had_errors)
+
+    mean_usd = (
+        sum((t.spend.usd for t in threads), Decimal(0)) / Decimal(total)
+        if total
+        else Decimal(0)
+    )
+
+    return RunMetrics(
+        runs=total,
+        completion_rate=round(finished / total, 3),
+        steps_p50=round(statistics.median(steps), 2),
+        steps_p95=round(_percentile(steps, 95), 2),
+        tool_error_rate=(
+            round(tool_errors / len(tool_events), 3) if tool_events else 0.0
+        ),
+        loop_rate=round(looped / total, 3),
+        budget_stop_rate=round(budget_stopped / total, 3),
+        recovery_rate=round(recovered / len(had_errors), 3) if had_errors else 1.0,
+        mean_usd=f"{mean_usd:.6f}",
+    )