diff --git a/python/src/donkey_kit/integrations/adk.py b/python/src/donkey_kit/integrations/adk.py index 328a958f..a816f1ab 100644 --- a/python/src/donkey_kit/integrations/adk.py +++ b/python/src/donkey_kit/integrations/adk.py @@ -34,7 +34,12 @@ the flow's context, so without it the record stays in that copy and an ``after_model_callback`` reads ``UNOBSERVED``. :func:`last_call_bridged` returns that subclass for a model built from ``connection_kwargs()`` or -``gemini_connection_kwargs()`` (#1218). +``gemini_connection_kwargs()`` (#1218). The factories' classes, +``GovernedLiteLlm``, ``GovernedGemini`` and the ``GovernedLiteLLMClient`` that +``model()`` passes as ``llm_client``, resolve by name on this module, so pickle +finds them. A model the factories build still holds the SDK's http client, so +it does not pickle; nor does a plain ``LiteLlm`` given the same kwargs. A class +:func:`last_call_bridged` builds from your own class has no name here. APPROVALS (BG ยง2.3, #1235): an approval-gated tool (``@donkey.tool(approval="required")`` in an agent's ``tools``) pauses the run @@ -68,7 +73,7 @@ from collections.abc import Iterable, Mapping from contextlib import aclosing from types import MappingProxyType -from typing import TYPE_CHECKING, Any, TypeVar +from typing import TYPE_CHECKING, Any, ClassVar, TypeVar from ..core import _verify from ..core.approvals import Approval, record_decision, record_pending @@ -317,7 +322,9 @@ class LastCallBridged(base): """An ADK model whose ``donkey.last_call`` reaches the context that started the call (#1207).""" - _donkey_last_call_bridged = True + # A ClassVar, so pydantic leaves it a plain class attribute rather than + # turning it into a private attribute. + _donkey_last_call_bridged: ClassVar[bool] = True def generate_content_async(self, *args: Any, **kwargs: Any) -> AsyncGenerator[Any, None]: bridge = open_last_call_bridge() @@ -368,9 +375,47 @@ async def acompletion(self, *args: Any, **kwargs: Any) -> Any: capture.tag(exc) raise + GovernedLiteLLMClient.__qualname__ = GovernedLiteLLMClient.__name__ return GovernedLiteLLMClient +def _governed_lite_llm() -> Any: + from google.adk.models.lite_llm import LiteLlm + + return _last_call_bridged_class(LiteLlm) + + +def _governed_gemini() -> Any: + from google.adk.models import Gemini + + return _last_call_bridged_class(Gemini) + + +def _governed_lite_llm_client() -> Any: + from google.adk.models.lite_llm import LiteLLMClient + + return _governed_llm_client_class(LiteLLMClient) + + +_GOVERNED_CLASSES = { + "GovernedLiteLlm": _governed_lite_llm, + "GovernedGemini": _governed_gemini, + "GovernedLiteLLMClient": _governed_lite_llm_client, +} + +if not TYPE_CHECKING: # mypy would otherwise accept any name on this module + + def __getattr__(name: str) -> Any: + """The factories' classes by name, so pickle finds them (PEP 562).""" + build = _GOVERNED_CLASSES.get(name) + if build is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + try: + return build() + except ModuleNotFoundError as exc: + raise missing_framework_error("adk", exc.name) from exc + + def model(model: str, **kw: Any) -> LiteLlm: """Module-level convenience: a native ``LiteLlm`` at the proxy using a cached default env-configured Donkey. Equivalent to diff --git a/python/tests/unit/test_adapter_async_last_call.py b/python/tests/unit/test_adapter_async_last_call.py index f925c83c..5a45b9cf 100644 --- a/python/tests/unit/test_adapter_async_last_call.py +++ b/python/tests/unit/test_adapter_async_last_call.py @@ -36,6 +36,7 @@ import asyncio import json +import pickle from pathlib import Path from typing import Any @@ -459,6 +460,45 @@ class Mine(Gemini): last_call_bridged(wrong) # type: ignore[arg-type] +def test_the_bridge_marker_is_a_plain_class_attribute() -> None: + """Not a pydantic private attribute, which is a truthy ``ModelPrivateAttr``.""" + pytest.importorskip("google.adk") + pytest.importorskip("litellm") + from google.adk.models import Gemini + from google.adk.models.lite_llm import LiteLlm + + from donkey_kit.integrations.adk import last_call_bridged + + for base in (Gemini, LiteLlm): + bridged = last_call_bridged(base) + assert bridged.__dict__["_donkey_last_call_bridged"] is True + assert "_donkey_last_call_bridged" not in bridged.__private_attributes__ + + +async def test_the_factory_classes_resolve_by_name_so_they_pickle() -> None: + pytest.importorskip("google.adk") + pytest.importorskip("litellm") + from google.adk.models import Gemini + from google.adk.models.lite_llm import LiteLlm + + from donkey_kit.integrations import adk + + donkey = _donkey() + async with donkey: + lite = donkey.adk.model("gpt-4o") + gemini = donkey.adk.gemini("gemini-3.8-flash") + assert type(lite) is adk.GovernedLiteLlm + assert type(lite.llm_client) is adk.GovernedLiteLLMClient + assert type(gemini) is adk.GovernedGemini + for cls in (adk.GovernedLiteLlm, adk.GovernedGemini, adk.GovernedLiteLLMClient): + assert pickle.loads(pickle.dumps(cls)) is cls + plain = adk.last_call_bridged(LiteLlm)(model="openai/gpt-4o") + assert type(pickle.loads(pickle.dumps(plain))) is adk.GovernedLiteLlm + assert adk.last_call_bridged(Gemini) is adk.GovernedGemini + with pytest.raises(AttributeError, match="GovernedNothing"): + adk.GovernedNothing # noqa: B018 + + async def test_adk_factories_return_the_adk_classes() -> None: pytest.importorskip("google.adk") pytest.importorskip("litellm")