"""Worker half of the Pydantic AI bridge.
:class:`BaseAgentWorker` hosts one or more :class:`~edri.integrations.pydantic_ai.agent.AgentBase`
agents on an EDRI :class:`~edri.abstract.worker.WorkerThread`, bakes runtime deps
into each, and exposes ``dispatch`` so agent tools can call other services as
plain Events.
"""
from __future__ import annotations
from abc import ABC
from queue import Empty, Queue
from threading import Lock, Thread
from typing import Any, ClassVar, TypeVar
from uuid import UUID
from edri.abstract.worker import WorkerThread
from edri.dataclass.event import Event
from edri.integrations.pydantic_ai.agent import AgentBase, AgentDeps, ToolResponse
from edri.integrations.pydantic_ai.observability import (
AgentActionPolicy,
_ObservabilityContext,
_describe_missing_queue,
_resolve_effective_policy,
resolve_audit_mode,
)
__all__ = ["BaseAgentWorker"]
class _AgentAccessor:
"""Maps agent classes to their built instances (``worker.agent[SomeAgent]``)."""
def __init__(self, agents: dict[str, AgentBase[Any, Any]]) -> None:
self._agents = agents
def __getitem__[A: AgentBase[Any, Any]](self, agent_cls: type[A]) -> A:
try:
return self._agents[agent_cls.default_name()]
except KeyError:
raise KeyError(
f"Agent '{agent_cls.__name__}' is not configured on this worker"
)
# Old-style TypeVar (not PEP 695): WorkerThread is already ``Generic[T], ABC``,
# and a PEP 695 ``[T: Event]`` here adds a second implicit ``Generic`` base that
# cannot be linearized against it ("inconsistent MRO").
_EventT = TypeVar("_EventT", bound=Event)
[docs]
class BaseAgentWorker(WorkerThread[_EventT], ABC):
"""A Worker that hosts one or more :class:`AgentBase` agents.
The worker is an orchestrator, not an agent. It declares which agents it
uses (:attr:`agents`), constructs them, bakes runtime deps into each, then
runs domain logic in :meth:`do`. Agents are reachable via
``self.agent[SomeAgentClass]`` after setup. The EDRI ``dispatch`` and shared
resources are injected via deps, so agent methods never expose them.
Subclasses must set :attr:`agents`, supply each agent's constructor
arguments — at least ``model_name`` — via :attr:`agent_configs` or
:meth:`agent_kwargs`, and implement :meth:`do`.
The agents are built in :meth:`setup`, which the base
:class:`~edri.abstract.worker.Worker` runs before :meth:`do` inside its own
error-handling scope — so this class does not override ``run()``.
"""
#: Agent classes used by this worker. Identity is the class itself; the
#: logical name comes from :meth:`AgentBase.default_name`.
agents: ClassVar[list[type[AgentBase[Any, Any]]]] = []
#: Per-agent constructor kwargs for the common case — each agent class in
#: :attr:`agents` maps to the dict passed to its constructor (e.g.
#: ``{"model_name": "gpt-4o"}``), so a worker hosting several agents on
#: different models is a flat declaration instead of a dispatch you write
#: yourself. See :meth:`agent_kwargs` for kwargs that must be computed at
#: runtime instead (env lookups, secrets, ...).
agent_configs: ClassVar[dict[type[AgentBase[Any, Any]], dict[str, Any]]] = {}
#: Set in :meth:`setup`; access agents through it.
agent: _AgentAccessor
[docs]
def setup(self) -> None:
"""Construct all declared agents and bake runtime deps into each.
Overrides the :meth:`~edri.abstract.worker.Worker.setup` hook, so it runs
once before :meth:`do`. Override further for worker-specific setup, but
call ``super().setup()`` so agents get built.
The audit mode is resolved once here, not per run: a running worker is
not a runtime toggle, and a worker must be restarted for a new mode to
take effect.
Raises:
ValueError: If two agent classes resolve to the same logical name,
or if ``EDRI_PYDANTIC_AI_AUDIT`` holds an unsupported value.
RuntimeError: If observability is enabled but no router queue is
bound, which would leave the actions with nowhere to go.
"""
super().setup()
mode = resolve_audit_mode()
if mode.enabled and self._router_queue is None:
raise RuntimeError(_describe_missing_queue(type(self).__name__))
event_id = self.event._id if self.event is not None else None
instances: dict[str, AgentBase[Any, Any]] = {}
for agent_cls in self.agents:
name = agent_cls.default_name()
if name in instances:
raise ValueError(
f"Duplicate agent name '{name}' on {type(self).__name__} "
f"(from {agent_cls.__name__})."
)
agent = agent_cls(**self.agent_kwargs(agent_cls))
agent.bind_deps(self.agent_deps(agent_cls))
agent._bind_observability(
_ObservabilityContext(
policy=_resolve_effective_policy(
mode, self.agent_action_policy(agent_cls)
),
router_queue=self._router_queue if mode.enabled else None,
edri_event_id=event_id,
agent=agent.name,
model=agent.model_name,
provider=agent.provider,
output_mode=agent.output_mode,
model_version=agent.model_version,
prompt_version=agent.prompt_version,
)
)
instances[name] = agent
self.agent = _AgentAccessor(instances)
# `_dispatch_event` needs concurrent dispatches from this worker to not
# race on the shared pipe: `_send_lock` serializes writes (per the
# `multiprocessing.Connection` docs, concurrent sends *or* receives on
# one end can corrupt it); `_reply_pump` becomes the pipe's sole reader
# and fans replies out to the right caller via `_pending`, keyed on the
# reply's `event._id` — so multiple in-flight dispatches run truly
# concurrently instead of queuing behind each other.
self._send_lock = Lock()
self._pending_lock = Lock()
self._pending: dict[UUID, Queue[Event]] = {}
self._reply_pump = Thread(
target=self._pump_replies, name=f"{self._name}-ReplyPump", daemon=True
)
self._reply_pump.start()
def _pump_replies(self) -> None:
"""Sole reader of the manager pipe once dispatching starts.
Routes each incoming reply to the :class:`queue.Queue` registered for
its ``event._id`` in :meth:`_dispatch_event`. Runs until the pipe
closes (worker shutdown), at which point any still-pending dispatches
are woken with a ``None`` sentinel instead of hanging past their own
timeout on a pipe that will never produce anything more.
"""
while True:
try:
reply = self._manager_pipe.recv()
except (EOFError, OSError):
break
reply._timing.stamp(self.__class__.__name__, "Received")
with self._pending_lock:
result_queue = self._pending.pop(reply._id, None)
if result_queue is None:
self.logger.warning("Unmatched dispatch reply id=%s: %s", reply._id, reply)
continue
result_queue.put(reply)
with self._pending_lock:
stale = list(self._pending.values())
self._pending.clear()
for result_queue in stale:
result_queue.put(None)
[docs]
def agent_kwargs(self, agent_cls: type[AgentBase[Any, Any]]) -> dict[str, Any]:
"""Constructor keyword arguments for ``agent_cls``.
Defaults to looking ``agent_cls`` up in :attr:`agent_configs` — set
that class-level mapping for the common case (each agent just needs a
different ``model_name`` or similar static config). Override this
method instead when kwargs must be computed at runtime (env lookups,
secrets, ...); an agent needs at least ``model_name``, and only the
worker knows the application's model/provider configuration.
Raises:
KeyError: If ``agent_cls`` isn't in :attr:`agent_configs` and this
wasn't overridden.
"""
try:
return self.agent_configs[agent_cls]
except KeyError:
raise KeyError(
f"No agent_kwargs for {agent_cls.__name__}: add it to "
f"{type(self).__name__}.agent_configs, or override agent_kwargs()."
) from None
[docs]
def agent_deps(self, agent_cls: type[AgentBase[Any, Any]]) -> AgentDeps[Any]:
"""Runtime dependencies for ``agent_cls``.
Defaults to sharing the worker's :meth:`dispatch <_dispatch_event>` and
exposing the worker itself as ``resources``. Override to scope resources
per agent.
"""
return AgentDeps(dispatch=self._dispatch_event, resources=self)
[docs]
def agent_action_policy(
self, agent_cls: type[AgentBase[Any, Any]]
) -> AgentActionPolicy:
"""Content policy for ``agent_cls``, resolved once during :meth:`setup`.
Defaults to capturing no payload bodies at all. Override to describe
what would be *useful* to observe; the deployment environment decides
separately what is *permitted*, and the effective policy is their
intersection, so an override here can never widen what
``EDRI_PYDANTIC_AI_AUDIT`` allows.
"""
return AgentActionPolicy()
def _dispatch_event(self, event: Event, *, timeout: float = 30.0) -> Any:
"""Send a request Event to the router and return the tool result.
Blocking (for the calling tool), but *not* serialized against other
concurrent dispatches from this worker: pydantic-ai can run a turn's
sync tool calls from multiple threads (a thread-pool executor per
call), so more than one dispatch can legitimately be in flight on
this worker's single pipe at once. Correlation is by ``event._id``
(universal on every Event) via :meth:`_pump_replies`, which is the
pipe's sole reader — sends are still serialized (``_send_lock``,
since concurrent multi-thread sends on one `Connection` end can also
corrupt it), but each call then waits on its own private
:class:`queue.Queue`, so unrelated dispatches never block each other.
A reply that crosses a *second* worker-owning Manager (the
tool-target Manager forwarding its own worker's response) still finds
its way back to this worker without any per-type wiring: the
dispatching Manager tracks it generically, keyed on ``event._id``
(see ``ManagerBase._track_worker_response``), and forwards it to this
worker's pipe once it arrives — that's also why ``event`` needs no
special base class here, just an ``Event``.
Raises:
TypeError: If ``event`` declares no response at all — it could
never get a reply, so this fails fast instead of hanging for
the full ``timeout``.
TimeoutError: If no reply arrives within ``timeout`` seconds —
turns a misrouted/lost reply into a legible error instead of
hanging the worker forever.
RuntimeError: If the reply's response is not a :class:`ToolResponse`,
or if the manager pipe closed before a reply arrived.
"""
if event.response is None:
raise TypeError(
f"{type(event).__name__} has no declared response and can "
f"never get a reply — a tool's dispatch target must declare "
f"one (and it must be a ToolResponse)."
)
self.logger.debug("Tool dispatch: %s", event)
result_queue: Queue[Event] = Queue(maxsize=1)
with self._pending_lock:
self._pending[event._id] = result_queue
with self._send_lock:
self.event_send(event)
try:
reply = result_queue.get(timeout=timeout)
except Empty:
with self._pending_lock:
self._pending.pop(event._id, None)
raise TimeoutError(f"No reply for {type(event).__name__} within {timeout}s.")
if reply is None:
raise RuntimeError(
"Manager pipe closed while waiting for a dispatch reply."
)
result = reply.response
if not isinstance(result, ToolResponse):
raise RuntimeError(
f"{type(event).__name__} dispatched by a tool must reply with a "
f"ToolResponse, got {type(result).__name__}. Make its response "
f"subclass ToolResponse and implement tool_result."
)
return result.tool_result