aid-runtime-python 0.1.39__py3-none-any.whl

aid_runtime/tools.py

@@ -0,0 +1,271 @@
+"""Tool definitions and the executor that runs them.
+
+A tool handler is business logic -- in a generated app it lives under
+``business/`` and is never overwritten by regeneration. The executor only owns
+argument validation, dispatch, and turning a handler crash into a tool error the
+model can read.
+
+Validation covers the JSON Schema subset the model is actually given (object
+``required`` and ``properties`` with scalar/array types). It is deliberately not a
+general JSON Schema implementation: a fuller validator belongs to the schema
+library, not to a runtime that must stay small.
+"""
+
+from __future__ import annotations
+
+import inspect
+import json
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from typing import Any
+
+from aid_runtime.concurrency import ConcurrencyLimiter
+from aid_runtime.errors import ToolArgumentError, ToolExecutionError, ToolNotFoundError
+from aid_runtime.types import ToolCall, ToolDefinition, ToolResult
+
+__all__ = ["Tool", "ToolExecutor"]
+
+#: A tool handler: sync or async, taking keyword arguments, returning anything JSON-ish.
+Handler = Callable[..., Any]
+
+_SCALAR_TYPES: dict[str, tuple[type, ...]] = {
+    "string": (str,),
+    "integer": (int,),
+    "number": (int, float),
+    "boolean": (bool,),
+    "object": (dict,),
+    "array": (list, tuple),
+    "null": (type(None),),
+}
+
+
+class Tool:
+    """A definition plus the callable that implements it."""
+
+    __slots__ = ("definition", "handler")
+
+    def __init__(self, definition: ToolDefinition, handler: Handler) -> None:
+        self.definition = definition
+        self.handler = handler
+
+    @property
+    def name(self) -> str:
+        return self.definition.name
+
+    @classmethod
+    def of(
+        cls,
+        name: str,
+        handler: Handler,
+        *,
+        description: str = "",
+        parameters: Mapping[str, Any] | None = None,
+    ) -> Tool:
+        """Declare a tool, inferring the parameter schema from the handler signature.
+
+        Inference covers names and required-ness only. Types must be declared
+        explicitly: guessing a type from a default is how a schema quietly
+        disagrees with the implementation.
+
+        An explicit ``parameters`` schema is checked against the handler at
+        registration, so a generated tool wired to the wrong function fails at
+        import time rather than the first time a model calls it.
+        """
+        declared = parameters if parameters is not None else _infer_parameters(handler)
+        _check_declaration(name, handler, declared)
+        return cls(
+            ToolDefinition(
+                name=name,
+                description=description or (inspect.getdoc(handler) or ""),
+                parameters=dict(declared),
+            ),
+            handler,
+        )
+
+
+class ToolExecutor:
+    """Validates and dispatches tool calls.
+
+    An executor may share a :class:`~aid_runtime.concurrency.ConcurrencyLimiter`
+    with the gateway; when one is supplied, every tool call acquires a slot so tool
+    calls and model calls draw from the same budget of concurrent work.
+    """
+
+    def __init__(
+        self,
+        tools: Iterable[Tool] = (),
+        *,
+        strict: bool = False,
+        limiter: ConcurrencyLimiter | None = None,
+    ) -> None:
+        self._tools: dict[str, Tool] = {}
+        self._strict = strict
+        self._limiter = limiter
+        for tool in tools:
+            self.register(tool)
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        return tuple(sorted(self._tools))
+
+    def register(self, tool: Tool) -> None:
+        self._tools[tool.name] = tool
+
+    def definition(self, name: str) -> ToolDefinition:
+        return self._get(name).definition
+
+    def definitions(self) -> tuple[ToolDefinition, ...]:
+        return tuple(self._tools[name].definition for name in self.names)
+
+    async def execute(self, call: ToolCall) -> ToolResult:
+        """Run a tool call, never raising for a bad argument or a handler crash.
+
+        Both are *tool* errors: the model composed the call, so the model is the
+        party that can fix it, and it needs the message rather than a stack trace
+        to do so. A call naming an unregistered tool is a programming error and
+        does raise -- no prompt repair makes a missing tool exist.
+
+        A ``strict`` executor is for development, where a handler traceback is
+        what you actually want: it raises
+        :class:`~aid_runtime.errors.ToolExecutionError` instead of returning an
+        error the model would silently try to talk its way around.
+
+        When a limiter is shared, the call holds one slot for its whole duration,
+        released on success, handler error, or cancellation.
+        """
+        if self._limiter is None:
+            return await self._execute(call)
+        async with self._limiter.run():
+            return await self._execute(call)
+
+    async def _execute(self, call: ToolCall) -> ToolResult:
+        tool = self._get(call.name)
+        problems = validate_arguments(tool.definition.parameters, call.arguments)
+        if problems:
+            return ToolResult.from_error(
+                call, f"invalid arguments: {'; '.join(problems)} -- {json.dumps(call.arguments)}"
+            )
+        try:
+            value = await _invoke(tool.handler, call.arguments)
+        except Exception as exc:
+            if self._strict:
+                raise ToolExecutionError(
+                    f"tool {call.name!r} raised {type(exc).__name__}: {exc}"
+                ) from exc
+            return ToolResult.from_error(call, f"{type(exc).__name__}: {exc}")
+        return ToolResult.from_value(call, value)
+
+    async def execute_all(self, calls: Iterable[ToolCall]) -> tuple[ToolResult, ...]:
+        """Run calls in order, each independent of the others' outcome."""
+        return tuple([await self.execute(call) for call in calls])
+
+    def _get(self, name: str) -> Tool:
+        try:
+            return self._tools[name]
+        except KeyError:
+            raise ToolNotFoundError(name, self.names) from None
+
+
+def validate_arguments(schema: Mapping[str, Any], arguments: Mapping[str, Any]) -> list[str]:
+    """Return every problem with ``arguments`` against ``schema``, in one pass."""
+    problems: list[str] = []
+
+    properties = schema.get("properties")
+    if not isinstance(properties, Mapping):
+        return problems
+
+    for name in schema.get("required") or []:
+        if name not in arguments:
+            problems.append(f"missing required argument {name!r}")
+
+    if schema.get("additionalProperties") is False:
+        for name in arguments:
+            if name not in properties:
+                problems.append(f"unexpected argument {name!r}")
+
+    for name, value in arguments.items():
+        spec = properties.get(name)
+        if not isinstance(spec, Mapping):
+            continue
+        problems.extend(_check_value(name, value, spec))
+
+    return problems
+
+
+def _check_value(name: str, value: Any, spec: Mapping[str, Any]) -> list[str]:
+    declared = spec.get("type")
+    if declared is None or value is None:
+        return []
+
+    allowed = [declared] if isinstance(declared, str) else list(declared)
+    if any(_matches(value, kind) for kind in allowed):
+        return []
+
+    return [f"argument {name!r} must be {_join(allowed)}, got {type(value).__name__}"]
+
+
+def _matches(value: Any, kind: str) -> bool:
+    expected = _SCALAR_TYPES.get(kind)
+    if expected is None:
+        # An unknown type name cannot be checked, so it cannot be reported as a
+        # mismatch either; rejecting here would fail every call.
+        return True
+    if kind in {"integer", "number"} and isinstance(value, bool):
+        # bool is a subclass of int in Python; JSON Schema does not consider
+        # `true` a valid number.
+        return False
+    return isinstance(value, expected)
+
+
+def _join(kinds: list[str]) -> str:
+    return " or ".join(kinds)
+
+
+def _check_declaration(name: str, handler: Handler, parameters: Mapping[str, Any]) -> None:
+    """Reject a schema that requires an argument the handler cannot receive."""
+    if _accepts_keywords(handler):
+        return
+    accepted = set(inspect.signature(handler).parameters)
+    required = parameters.get("required")
+    if not isinstance(required, Sequence) or isinstance(required, str):
+        return
+    undeclared = sorted(str(item) for item in required if str(item) not in accepted)
+    if undeclared:
+        raise ToolArgumentError(
+            name,
+            [f"required argument(s) {', '.join(undeclared)} not accepted by the handler"],
+        )
+
+
+def _infer_parameters(handler: Handler) -> dict[str, Any]:
+    signature = inspect.signature(handler)
+    names = [
+        name
+        for name, parameter in signature.parameters.items()
+        if parameter.kind in {parameter.POSITIONAL_OR_KEYWORD, parameter.KEYWORD_ONLY}
+    ]
+    return {"type": "object", "properties": {name: {} for name in names}, "required": names}
+
+
+def _accepts_keywords(handler: Handler) -> bool:
+    try:
+        signature = inspect.signature(handler)
+    except (TypeError, ValueError):
+        return False
+    return any(
+        parameter.kind is parameter.VAR_KEYWORD for parameter in signature.parameters.values()
+    )
+
+
+async def _invoke(handler: Handler, arguments: Mapping[str, Any]) -> Any:
+    """Call a sync or async handler with ``arguments`` as keyword arguments."""
+    kwargs = dict(arguments) if _accepts_keywords(handler) else _bind(handler, arguments)
+    result = handler(**kwargs)
+    if inspect.isawaitable(result):
+        return await result
+    return result
+
+
+def _bind(handler: Handler, arguments: Mapping[str, Any]) -> dict[str, Any]:
+    """Drop arguments the handler does not name, so an over-supplying model is tolerated."""
+    accepted = set(inspect.signature(handler).parameters)
+    return {key: value for key, value in arguments.items() if key in accepted}

aid_runtime/tracing.py

@@ -0,0 +1,241 @@
+"""OpenTelemetry tracing using the GenAI semantic conventions.
+
+Attribute names come from the OpenTelemetry GenAI conventions (``gen_ai.*``) so
+spans land in any OTLP backend already understood as model telemetry. Attributes
+the conventions do not cover -- which pinned alias produced the call, and what it
+cost -- are namespaced under ``aid.*`` rather than smuggled into a ``gen_ai.*``
+key with the wrong meaning.
+
+Only the OpenTelemetry *API* is a runtime dependency. With no tracer provider
+installed every call here is a no-op, which is the correct behaviour for a library
+that must not decide a host application's telemetry setup for it.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Iterator
+from typing import TYPE_CHECKING, Any
+
+from opentelemetry import trace
+from opentelemetry.trace import Span, SpanKind, Status, StatusCode
+
+from aid_runtime.types import TokenUsage
+
+if TYPE_CHECKING:
+    from aid_runtime.agents.guardrails import GuardrailRecord
+
+__all__ = [
+    "ATTRIBUTES",
+    "GEN_AI_OPERATION_CHAT",
+    "agent_span",
+    "aid_attributes",
+    "get_tracer",
+    "llm_span",
+    "record_attributes",
+    "record_cost",
+    "record_exception",
+    "record_guardrail",
+    "record_usage",
+]
+
+GEN_AI_OPERATION_CHAT = "chat"
+
+#: The GenAI semantic-convention attribute keys this runtime writes.
+ATTRIBUTES: dict[str, str] = {
+    "system": "gen_ai.system",
+    "operation": "gen_ai.operation.name",
+    "request_model": "gen_ai.request.model",
+    "request_temperature": "gen_ai.request.temperature",
+    "request_max_tokens": "gen_ai.request.max_tokens",
+    "response_model": "gen_ai.response.model",
+    "response_finish_reasons": "gen_ai.response.finish_reasons",
+    "input_tokens": "gen_ai.usage.input_tokens",
+    "output_tokens": "gen_ai.usage.output_tokens",
+    # Framework-specific, deliberately outside the gen_ai namespace.
+    "alias": "aid.model.alias",
+    "cost_usd": "aid.cost.usd",
+    "attempts": "aid.gateway.attempts",
+    "tool_names": "aid.tool.names",
+    # Agent loop. ``agent_status`` is the run's terminal state, recorded on the
+    # agent span so a paused run is distinguishable from a completed one without
+    # reading the loop's return value.
+    "agent_id": "aid.agent.id",
+    "agent_engine": "aid.agent.engine",
+    "agent_steps": "aid.agent.steps",
+    "agent_status": "aid.agent.status",
+    "agent_transfer_to": "aid.agent.transfer.to",
+    "agent_delegate_to": "aid.agent.delegate.to",
+    "agent_budget_degraded": "aid.agent.budget.degraded",
+    "agent_budget_from": "aid.agent.budget.from",
+    "agent_budget_to": "aid.agent.budget.to",
+    # Human-in-the-loop: who was asked, on what channel, at which step, and what
+    # they decided. A pause is an ordinary outcome, not an error.
+    "hitl_channel": "aid.hitl.channel",
+    "hitl_when": "aid.hitl.when",
+    "hitl_decision": "aid.hitl.decision",
+    "hitl_step": "aid.hitl.step",
+    "hitl_phase": "aid.hitl.phase",
+    "memory_scope": "aid.memory.scope",
+    "memory_entries": "aid.memory.entries",
+    "memory_tokens": "aid.memory.tokens",
+    # Guardrails. An event per hit rather than one attribute, for the same reason
+    # HITL uses events: several guardrails can fire at one stage, and a single
+    # attribute would keep only the last one while still looking complete.
+    "guardrail_id": "aid.guardrail.id",
+    "guardrail_stage": "aid.guardrail.stage",
+    "guardrail_action": "aid.guardrail.action",
+    "guardrail_policies": "aid.guardrail.policies",
+}
+
+_TRACER_NAME = "aid_runtime"
+
+
+def get_tracer(name: str = _TRACER_NAME) -> trace.Tracer:
+    return trace.get_tracer(name)
+
+
+@contextlib.contextmanager
+def llm_span(
+    *,
+    model: str,
+    provider: str,
+    operation: str = GEN_AI_OPERATION_CHAT,
+    alias: str | None = None,
+    temperature: float | None = None,
+    max_tokens: int | None = None,
+    tracer: trace.Tracer | None = None,
+) -> Iterator[Span]:
+    """Open a CLIENT span for one model call and close it with an error status."""
+    active = tracer or get_tracer()
+    attributes: dict[str, Any] = {
+        ATTRIBUTES["system"]: provider,
+        ATTRIBUTES["operation"]: operation,
+        ATTRIBUTES["request_model"]: model,
+    }
+    if temperature is not None:
+        attributes[ATTRIBUTES["request_temperature"]] = temperature
+    if max_tokens is not None:
+        attributes[ATTRIBUTES["request_max_tokens"]] = max_tokens
+    if alias is not None:
+        attributes[ATTRIBUTES["alias"]] = alias
+
+    with active.start_as_current_span(
+        f"{operation} {model}",
+        kind=SpanKind.CLIENT,
+        attributes=attributes,
+        # The exception is recorded once, explicitly, from the ``except`` below.
+        # Left at their defaults these would record it a second time and overwrite
+        # the status description with the exception's repr.
+        record_exception=False,
+        set_status_on_exception=False,
+    ) as span:
+        try:
+            yield span
+        except BaseException as exc:
+            record_exception(span, exc)
+            raise
+
+
+def record_usage(span: Span, usage: TokenUsage) -> None:
+    span.set_attribute(ATTRIBUTES["input_tokens"], usage.input_tokens)
+    span.set_attribute(ATTRIBUTES["output_tokens"], usage.output_tokens)
+
+
+def record_attributes(span: Span, **values: Any) -> None:
+    """Attach ``aid.*`` attributes by short name, skipping ``None`` values.
+
+    A recorded ``None`` would look like a decided empty value, so an attribute
+    that does not apply is omitted rather than nulled.
+    """
+    for key, value in aid_attributes(**values).items():
+        if value is not None:
+            span.set_attribute(key, value)
+
+
+@contextlib.contextmanager
+def agent_span(
+    *,
+    agent_id: str,
+    engine: str,
+    tracer: trace.Tracer | None = None,
+) -> Iterator[Span]:
+    """Open an INTERNAL span for one agent's turn.
+
+    INTERNAL, not CLIENT: an agent is not a remote service, and the model calls it
+    makes open their own CLIENT child spans.
+    """
+    active = tracer or get_tracer()
+    with active.start_as_current_span(
+        f"invoke_agent {agent_id}",
+        kind=SpanKind.INTERNAL,
+        attributes=aid_attributes(agent_id=agent_id, agent_engine=engine),
+        record_exception=False,
+        set_status_on_exception=False,
+    ) as span:
+        try:
+            yield span
+        except BaseException as exc:
+            record_exception(span, exc)
+            raise
+
+
+def record_hitl(
+    span: Span,
+    *,
+    phase: str,
+    channel: str = "",
+    decision: str | None = None,
+    when: str | None = None,
+    step: int | None = None,
+) -> None:
+    """Record one pause or resume of a human-in-the-loop gate.
+
+    An event, not only attributes: a run can pause and resume more than once, and a
+    span attribute would silently keep the last decision while the span still looked
+    like it paused exactly once. ``phase`` is ``pause`` or ``resume``.
+    """
+    attributes: dict[str, Any] = {
+        ATTRIBUTES["hitl_phase"]: phase,
+        ATTRIBUTES["hitl_channel"]: channel,
+    }
+    if when is not None:
+        attributes[ATTRIBUTES["hitl_when"]] = when
+    if decision is not None:
+        attributes[ATTRIBUTES["hitl_decision"]] = decision
+    if step is not None:
+        attributes[ATTRIBUTES["hitl_step"]] = step
+    span.add_event("aid.hitl", attributes)
+
+
+def record_guardrail(span: Span, record: GuardrailRecord) -> None:
+    """Record one guardrail hit as an ``aid.guardrail`` span event."""
+    span.add_event(
+        "aid.guardrail",
+        {
+            ATTRIBUTES["guardrail_id"]: record.guardrail,
+            ATTRIBUTES["guardrail_stage"]: record.stage,
+            ATTRIBUTES["guardrail_action"]: record.action,
+            ATTRIBUTES["guardrail_policies"]: ",".join(record.policies),
+        },
+    )
+
+
+def record_cost(span: Span, cost_usd: float, *, alias: str | None = None) -> None:
+    span.set_attribute(ATTRIBUTES["cost_usd"], cost_usd)
+    if alias is not None:
+        span.set_attribute(ATTRIBUTES["alias"], alias)
+
+
+def record_exception(span: Span, exc: BaseException) -> None:
+    span.record_exception(exc)
+    span.set_status(Status(StatusCode.ERROR, str(exc)))
+
+
+def aid_attributes(**values: Any) -> dict[str, Any]:
+    """Map short names from :data:`ATTRIBUTES` to their namespaced keys.
+
+    Unmapped names pass through, so a caller can attach a bespoke ``aid.*``
+    attribute without editing this module.
+    """
+    return {ATTRIBUTES.get(key, key): value for key, value in values.items()}