"""Agent-facing half of the Pydantic AI bridge.
:class:`AgentDeps` and :class:`ToolResponse` are the contract between an agent's
tools and the worker hosting them; :class:`AgentBase` is what an application
subclasses to own domain behaviour without touching runtime concerns. The worker
side lives in :mod:`edri.integrations.pydantic_ai.worker`.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from dataclasses import dataclass, replace
from logging import getLogger
from time import monotonic
from typing import Any, ClassVar, Protocol
from uuid import uuid4
from edri.dataclass.event import Event
from edri.dataclass.response import Response
from edri.utility.function import camel2snake
from edri.integrations.pydantic_ai._optional import Agent
from edri.integrations.pydantic_ai.lifecycle import (
audit_hooks,
_current_tool_action,
_emit_failure,
)
from edri.integrations.pydantic_ai.observability import (
AgentAction,
AgentActionNamespace,
AgentActionOperation,
AgentActionOutcome,
AgentActionPhase,
_ActionEmitter,
_Capture,
_EffectivePolicy,
_ObservabilityContext,
_RunState,
_emit_safely,
_unbound_observability,
)
__all__ = ["AgentBase", "AgentDeps", "DispatchFn", "ToolResponse"]
[docs]
class DispatchFn(Protocol):
"""Callable shape of :attr:`AgentDeps.dispatch`.
``timeout`` is keyword-only so each call site (i.e. each tool) can override
the default wait for its own dispatched event — a tool known to hit a slow
downstream Worker can afford a longer wait; one that should fail fast can
ask for less than the 30s default.
"""
def __call__(self, event: Event, *, timeout: float = ...) -> Any: ...
[docs]
@dataclass(frozen=True)
class AgentDeps[ResourcesT]:
"""Runtime dependencies injected into an agent by its worker.
Pydantic AI tools receive this as ``RunContext[AgentDeps].deps``. Agents
never construct or inspect it — the worker builds one and bakes it in via
:meth:`AgentBase.bind_deps`.
Attributes:
dispatch: Sends a request Event to the EDRI router and returns the tool
result extracted from the response. Implemented by the worker
(:meth:`BaseAgentWorker._dispatch_event`); blocking for the caller,
but concurrent dispatches from the same worker don't block each
other. This is the only path from a tool back into the rest of
the system. Pass ``timeout=`` to override the 30s default for a
given call, e.g. ``ctx.deps.dispatch(SlowLookup(...), timeout=5.0)``.
resources: Optional worker-level shared state (in-memory stores, DB
handles, caches). Parametrize ``AgentDeps[MyResources]`` for typing.
"""
dispatch: DispatchFn
resources: ResourcesT | None = None
def _event_payload(event: Event) -> Any:
"""The dispatched event's own fields, without framework transport state.
``_id`` and ``_timing`` are EDRI's own plumbing and are reported as
metadata, not as request content; ``response`` is the empty reply container
on the way out. What remains is the payload the tool actually sent.
"""
try:
return {
name: value
for name, value in vars(event).items()
if not name.startswith("_") and name != "response"
}
except Exception:
return None
def _audited_dispatch(
dispatch: DispatchFn,
state: _RunState,
emitter: _ActionEmitter,
policy: _EffectivePolicy,
) -> DispatchFn:
"""Wrap the dispatch callable this run's tools will actually call.
The *actual* callable, not :meth:`BaseAgentWorker._dispatch_event`:
``agent_deps()`` is public and applications do supply their own dispatch, so
instrumenting the worker method would observe some deployments and silently
miss others.
This is also the bridge from Pydantic AI run correlation to EDRI transport
correlation. It closes over this run's state, so a dispatch carries the run
that issued it even while other runs execute on the same built agent -- and
tool bodies run on other threads, so reading a "current run" attribute
instead would attribute dispatches to whichever run happened to be latest.
"""
def audited(event: Event, **kwargs: Any) -> Any:
# ``timeout`` is deliberately absent from this signature. Declaring it
# would bind a default here, and a caller that omitted the argument
# would then get *this* default instead of the wrapped dispatch's own.
if isinstance(event, AgentAction):
# No auditing the audit: an application that dispatches an action
# event would otherwise start an endless chain.
return dispatch(event, **kwargs)
shared = {
"action_id": uuid4(),
"namespace": AgentActionNamespace.EDRI,
"operation": AgentActionOperation.DISPATCH,
# The tool that made this dispatch when one is in flight in this
# context, so the inspector shows the dispatch under the tool
# rather than as a sibling of it. A dispatch from anywhere else in
# the run attaches to the run itself.
"parent_action_id": (
_current_tool_action(state.root_action_id) or state.root_action_id
),
"root_action_id": state.root_action_id,
"run_id": state.run_id,
"conversation_id": state.conversation_id,
"name": type(event).__name__,
"metadata": {"dispatched_event_id": getattr(event, "_id", None)},
}
_emit_safely(
emitter,
phase=AgentActionPhase.START,
request=(
_Capture(_event_payload(event), True)
if policy.include_dispatch_requests
else None
),
**shared,
)
started = monotonic()
try:
result = dispatch(event, **kwargs)
except BaseException as error:
# The one operation here that owns a deadline of its own: a
# dispatch waits ``timeout`` seconds for its reply, and the
# TimeoutError it raises when that expires *is* the expiry.
_emit_failure(
emitter,
error,
started=started,
shared=shared,
owns_deadline=True,
)
raise
_emit_safely(
emitter,
phase=AgentActionPhase.END,
outcome=AgentActionOutcome.SUCCESS,
duration=monotonic() - started,
response_data=(
_Capture(result, True)
if policy.include_dispatch_responses
else None
),
**shared,
)
return result
return audited
[docs]
class AgentBase[DepsT, OutputT](ABC):
"""Base class for EDRI-hosted Pydantic AI agents.
An agent owns only its domain behavior: model/provider construction
(:meth:`build_agent`), prompt, output configuration and domain methods.
Runtime concerns (the EDRI pipe, shared resources) arrive as
:class:`AgentDeps`, which the worker bakes in via :meth:`bind_deps` after
construction. Tools reach them through ``RunContext[AgentDeps]``, so domain
methods stay clean and never expose deps in their signatures.
"""
provider: str = "unknown"
framework: str = "pydantic_ai"
output_mode: str = "unknown"
version: str | None = None # version of the agent implementation
model_version: str | None = None # specific model revision/checkpoint, if known
prompt_version: str | None = None # version of the system prompt / instructions
#: Stateless, shareable capability/tool bundles. Tools read runtime values
#: from ``RunContext[AgentDeps]``; capabilities themselves hold no per-run state.
capabilities: ClassVar[list[Any]] = []
def __init__(self, *, model_name: str, name: str | None = None) -> None:
"""Initialize the agent.
Args:
model_name: Model identifier used by the underlying provider.
name: Optional logical agent name. Defaults to a snake_case form of
the class name (``ClaimExtractorAgent`` -> ``claim_extractor``).
"""
self.model_name = self._validate_model_name(model_name)
self.name = self._resolve_name(name)
self.logger = getLogger(self.name)
self._deps: AgentDeps[Any] | None = None # baked in by the worker
# Framework-internal, replaced by the worker in setup(). An agent built
# outside a worker is simply unobserved rather than broken, so the
# default is a real disabled context and not None.
self._observability = _unbound_observability()
self.agent = self.build_agent()
[docs]
@abstractmethod
def build_agent(self) -> Agent[DepsT, OutputT]:
"""Build the underlying Pydantic AI agent.
Implementations construct the model/provider and return a configured
``Agent``. Pass ``deps_type=AgentDeps`` (and your capabilities) so tools
can resolve runtime dependencies.
"""
...
[docs]
def bind_deps(self, deps: AgentDeps[Any]) -> AgentBase[DepsT, OutputT]:
"""Bake runtime dependencies into the agent (called once by the worker).
Returns ``self`` to allow fluent chaining.
"""
self._deps = deps
return self
def _bind_observability(self, context: _ObservabilityContext) -> None:
"""Bind framework-internal audit state (called once by the worker).
Kept off :class:`AgentDeps` on purpose: deps are application-facing and
``agent_deps()`` is routinely overridden, so observability must not be
reachable through -- or dependent on -- ``ctx.deps.resources``.
"""
self._observability = context
def _run(self, *args: Any, **kwargs: Any) -> Any:
"""Run the underlying agent synchronously with baked-in deps.
Domain methods call this instead of ``self.agent.run_sync(...)`` so deps
never leak into their signatures.
With observability enabled the run is additionally wired for audit; see
:meth:`_wire_observability`. The guard below deliberately still consults only
:attr:`capabilities`, which describes *domain* capabilities: the audit
capability is attached per invocation and never added there, so turning
observability on cannot make a previously valid deps-less agent fail.
Raises:
RuntimeError: If the agent exposes capabilities but no deps were
bound (a misconfiguration in the worker setup).
"""
if self._deps is None and self.capabilities:
raise RuntimeError(
f"Agent '{self.name}' exposes capabilities but has no deps bound. "
f"The worker must call bind_deps() during setup."
)
deps = self._deps
observability = self._observability
if observability.enabled:
deps = self._wire_observability(observability, kwargs)
return self.agent.run_sync(*args, deps=deps, **kwargs)
def _wire_observability(
self, observability: _ObservabilityContext, kwargs: dict[str, Any]
) -> AgentDeps[Any] | None:
"""Wire this one invocation for audit and return the deps it should use.
Everything the observation needs lives in one :class:`_RunState` closed
over here, so nothing is written to ``self``: concurrent runs on one
built agent are routine (a worker drives several, and tool bodies run on
other threads), and a "current run" attribute would silently attribute
one run's dispatches to another.
The audit capability is appended to the caller's capabilities rather
than replacing them, and the caller's own list object is left alone.
"""
emitter = _ActionEmitter(observability)
state = _RunState(root_action_id=uuid4())
kwargs["capabilities"] = [
*(kwargs.get("capabilities") or ()),
audit_hooks(state, emitter, observability.policy),
]
deps = self._deps
if deps is None:
return None
# ``replace`` on the deps the agent actually received, not a rebuilt
# AgentDeps: ``agent_deps()`` is public, so an application may supply
# its own dispatch, extra fields or a subclass, and the audited run must
# carry all of it unchanged but for ``dispatch``.
return replace(
deps,
dispatch=_audited_dispatch(
deps.dispatch, state, emitter, observability.policy
),
)
[docs]
@classmethod
def default_name(cls) -> str:
"""Derive a stable snake_case agent name from the class name.
A trailing ``Agent`` is stripped: ``ClaimDebateJudgeAgent`` ->
``claim_debate_judge``.
"""
class_name = cls.__name__
if class_name.endswith("Agent"):
class_name = class_name[: -len("Agent")]
return camel2snake(class_name)
@classmethod
def _resolve_name(cls, name: str | None) -> str:
resolved = cls.default_name() if name is None else name.strip()
if not resolved:
raise ValueError("Agent name cannot be empty.")
if not resolved.isidentifier():
raise ValueError(
f"Agent name must be a valid Python identifier, got: {resolved!r}"
)
return resolved
@staticmethod
def _validate_model_name(model_name: str) -> str:
resolved = model_name.strip()
if not resolved:
raise ValueError("model_name cannot be empty.")
return resolved