From 2f603183f02e48af9a54b73a0fc0d7a0a45ddc68 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=9B=90=E7=B2=92=20Yanli?= Date: Wed, 29 Apr 2026 21:55:02 +0800 Subject: [PATCH] add the transformer design for compositor Co-authored-by: Copilot --- dify-agent/Makefile | 2 +- dify-agent/examples/agenton_basics.py | 17 +-- dify-agent/examples/agenton_pydantic_ai.py | 113 +++++++++++++++++ dify-agent/src/agenton/compositor/__init__.py | 97 ++++++++++---- dify-agent/src/agenton/compositor/helpers.py | 110 ---------------- dify-agent/src/agenton/layers/__init__.py | 12 ++ dify-agent/src/agenton/layers/base.py | 20 ++- dify-agent/src/agenton/layers/types.py | 87 ++++++++++++- .../src/agenton_collections/__init__.py | 14 +- .../layers/pydantic_ai/__init__.py | 2 - .../layers/pydantic_ai/bridge.py | 42 ++++-- .../transformers/__init__.py | 11 ++ .../transformers/pydantic_ai.py | 65 ++++++++++ .../unit/agenton/compositor/test_enter.py | 8 +- .../agenton/compositor/test_transformers.py | 120 ++++++++++++++++++ .../unit/agenton/layers/test_layer_deps.py | 2 +- .../layers/pydantic_ai/test_bridge.py | 53 ++++++++ .../transformers/test_pydantic_ai.py | 75 +++++++++++ 18 files changed, 679 insertions(+), 171 deletions(-) create mode 100644 dify-agent/examples/agenton_pydantic_ai.py delete mode 100644 dify-agent/src/agenton/compositor/helpers.py create mode 100644 dify-agent/src/agenton_collections/transformers/__init__.py create mode 100644 dify-agent/src/agenton_collections/transformers/pydantic_ai.py create mode 100644 dify-agent/tests/unit/agenton/compositor/test_transformers.py create mode 100644 dify-agent/tests/unit/agenton_collections/layers/pydantic_ai/test_bridge.py create mode 100644 dify-agent/tests/unit/agenton_collections/transformers/test_pydantic_ai.py diff --git a/dify-agent/Makefile b/dify-agent/Makefile index 28ad989537..0ae3f4bae8 100644 --- a/dify-agent/Makefile +++ b/dify-agent/Makefile @@ -8,7 +8,7 @@ help: @echo " make test - Run dify-agent pytest suite" typecheck: - @uv run --project . python -m basedpyright src tests + @uv run --project . python -m basedpyright --level error src tests test: @uv run --project . python -m pytest tests diff --git a/dify-agent/examples/agenton_basics.py b/dify-agent/examples/agenton_basics.py index 2a2015c4bd..981c789e5c 100644 --- a/dify-agent/examples/agenton_basics.py +++ b/dify-agent/examples/agenton_basics.py @@ -10,7 +10,8 @@ from typing_extensions import override from agenton.compositor import Compositor, CompositorLayerConfig from agenton.layers import LayerControl, LayerDeps, NoLayerDeps, PlainLayer -from agenton_collections.plain import DynamicToolsLayer, ObjectLayer, ToolsLayer, with_object +from agenton.layers.types import PlainPromptType, PlainToolType +from agenton_collections.layers.plain import DynamicToolsLayer, ObjectLayer, ToolsLayer, with_object @dataclass(frozen=True, slots=True) @@ -74,13 +75,13 @@ async def main() -> None: ) trace = TraceLayer() - compositor = Compositor.from_config( + compositor = Compositor[PlainPromptType, PlainToolType].from_config( { "layers": [ { "name": "base_prompt", "layer": { - "import_path": "agenton_collections.plain.basic:PromptLayer", + "import_path": "agenton_collections.layers.plain:PromptLayer", "config": { "prefix": "Use config dicts for serializable layers.", "suffix": "Before finalizing, make the result easy to scan.", @@ -90,7 +91,7 @@ async def main() -> None: { "name": "extra_prompt", "layer": { - "import_path": "agenton_collections.plain.basic:PromptLayer", + "import_path": "agenton_collections.layers.plain:PromptLayer", "config": { "prefix": "Use constructed instances for objects, local code, and callables.", }, @@ -118,17 +119,17 @@ async def main() -> None: ), CompositorLayerConfig(name="trace", layer=trace), ] - } + }, ) print("Prompts:") for prompt in compositor.prompts: - print(f"- {prompt}") + print(f"- {prompt.value}") print("\nTools:") for tool in compositor.tools: - print(f"- {tool.__name__}{signature(tool)}") - print([tool("layer composition") for tool in compositor.tools]) + print(f"- {tool.value.__name__}{signature(tool.value)}") + print([tool.value("layer composition") for tool in compositor.tools]) async with compositor.enter() as lifecycle_control: lifecycle_control.tmp_leave = True diff --git a/dify-agent/examples/agenton_pydantic_ai.py b/dify-agent/examples/agenton_pydantic_ai.py new file mode 100644 index 0000000000..c5fdfffe3c --- /dev/null +++ b/dify-agent/examples/agenton_pydantic_ai.py @@ -0,0 +1,113 @@ +"""Run with: uv run --project dify-agent python examples/agenton_pydantic_ai.py.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import cast + +from pydantic_ai import Agent, RunContext, Tool +from pydantic_ai.messages import ToolCallPart, BuiltinToolCallPart +from pydantic_ai.models.test import TestModel + +from agenton.compositor import Compositor, CompositorLayerConfig +from agenton.layers.types import AllPromptTypes, AllToolTypes, PydanticAIPrompt, PydanticAITool +from agenton_collections.layers.plain import ObjectLayer, ToolsLayer +from agenton_collections.layers.pydantic_ai import PydanticAIBridgeLayer +from agenton_collections.transformers import PYDANTIC_AI_TRANSFORMERS + + +import json +@dataclass(frozen=True, slots=True) +class AgentProfile: + name: str + audience: str + tone: str + + +def count_words(text: str) -> int: + return len(text.split()) + + +def profile_prompt(ctx: RunContext[AgentProfile]) -> str: + profile = ctx.deps + return f"You are {profile.name}, helping {profile.audience}." + + +def tone_prompt(ctx: RunContext[AgentProfile]) -> str: + return f"Keep responses {ctx.deps.tone}." + + +def write_tagline(ctx: RunContext[AgentProfile], topic: str) -> str: + profile = ctx.deps + return f"{profile.name}: {topic} for {profile.audience}, in a {profile.tone} voice." + + +async def main() -> None: + profile = AgentProfile( + name="Agenton Assistant", + audience="engineers composing agent capabilities", + tone="precise and friendly", + ) + pydantic_ai_bridge = PydanticAIBridgeLayer[AgentProfile]( + prefix=(profile_prompt, tone_prompt), + tool_entries=(Tool(write_tagline),), + ) + + compositor = Compositor[ + PydanticAIPrompt[object], + PydanticAITool[object], + AllPromptTypes, + AllToolTypes, + ].from_config( + { + "layers": [ + { + "name": "base_prompt", + "layer": { + "import_path": "agenton_collections.layers.plain:PromptLayer", + "config": { + "prefix": "Use the available tools before answering.", + "suffix": "Return concise, inspectable output.", + }, + }, + }, + CompositorLayerConfig( + name="profile", + layer=ObjectLayer[AgentProfile](profile), + ), + CompositorLayerConfig( + name="plain_tools", + layer=ToolsLayer(tool_entries=(count_words,)), + ), + CompositorLayerConfig( + name="pydantic_ai_bridge", + deps={"object_layer": "profile"}, + layer=pydantic_ai_bridge, + ), + ] + }, + **PYDANTIC_AI_TRANSFORMERS, + ) + + async with compositor.enter(): + agent = Agent[AgentProfile]( + model=TestModel(call_tools=["count_words", "write_tagline"]), + deps_type=AgentProfile, + tools=compositor.tools, + ) + for prompt in compositor.prompts: + agent.system_prompt(prompt) + + result = await agent.run( + "Use the tools for 'layer composition'.", + deps=pydantic_ai_bridge.run_deps, + ) + + for message in result.all_messages(): + for part in message.parts: + print(f"{type(part).__name__}: {part.content if not isinstance(part, (ToolCallPart, BuiltinToolCallPart)) else part.tool_name + '(' + json.dumps(part.args, ensure_ascii=False) + ')'}") + + +if __name__ == "__main__": + import asyncio + asyncio.run(main()) diff --git a/dify-agent/src/agenton/compositor/__init__.py b/dify-agent/src/agenton/compositor/__init__.py index 5f0ac9201e..b5d7caf548 100644 --- a/dify-agent/src/agenton/compositor/__init__.py +++ b/dify-agent/src/agenton/compositor/__init__.py @@ -1,11 +1,10 @@ """Layer composition primitives. The compositor owns a named, ordered set of layers. ``Compositor[PromptT, -ToolT]`` is framework-neutral; callers choose prompt/tool item types by -annotating construction or assignment sites. Use -``agenton.compositor.helpers.make_compositor`` when type inference from layer -arguments is useful; it lives in a child module so the core compositor does not -depend on its helper overloads. +ToolT, LayerPromptT, LayerToolT]`` is framework-neutral; callers choose layer and +exposed prompt/tool item types by annotating construction or assignment sites. +When only the first two type arguments are supplied, ``LayerPromptT`` and +``LayerToolT`` default to the corresponding exposed item types. Dependency mappings use layer-local dependency names as keys and compositor layer names as values. Prompt aggregation depends on insertion order: prefix @@ -18,19 +17,31 @@ reverse order through ``AsyncExitStack``. It accepts an optional omitted, one is created from the compositor's layer names. Reuse the same ``CompositorControl`` after setting ``tmp_leave`` to reenter those layer contexts. + +Optional prompt and tool transformers run after layer aggregation. The +compositor asks each layer to ``wrap_prompt`` and ``wrap_tool`` its native +values, so typed layer families can tag prompt/tool values without changing +their authoring contracts. When transformers are omitted, the compositor +returns those wrapped items unchanged. """ from collections import OrderedDict -from collections.abc import AsyncIterator, Iterable +from collections.abc import AsyncIterator, Callable, Iterable, Sequence from contextlib import AsyncExitStack, asynccontextmanager from dataclasses import dataclass, field from importlib import import_module -from typing import TYPE_CHECKING, Annotated, Any, Mapping, cast +from typing import TYPE_CHECKING, Annotated, Any, Generic, Mapping, TypedDict, cast from pydantic import AfterValidator, BaseModel, ConfigDict, Field, JsonValue -from typing_extensions import Self +from typing_extensions import Self, TypeVar from agenton.layers.base import Layer, LayerControl +from agenton.layers.types import AllPromptTypes, AllToolTypes + +PromptT = TypeVar("PromptT", default=AllPromptTypes) +ToolT = TypeVar("ToolT", default=AllToolTypes) +LayerPromptT = TypeVar("LayerPromptT", default=AllPromptTypes) +LayerToolT = TypeVar("LayerToolT", default=AllToolTypes) class ImportedLayerConfig(BaseModel): @@ -58,6 +69,16 @@ class ImportedLayerConfig(BaseModel): LayerSpec = Layer[Any, Any, Any] | ImportedLayerConfig +type CompositorTransformer[InputT, OutputT] = Callable[[Sequence[InputT]], Sequence[OutputT]] + + +class CompositorTransformerKwargs[PromptT, ToolT, LayerPromptT, LayerToolT](TypedDict): + """Keyword arguments that install prompt and tool transformers together.""" + + prompt_transformer: CompositorTransformer[LayerPromptT, PromptT] + tool_transformer: CompositorTransformer[LayerToolT, ToolT] + + type _ConfigModelValue[ModelT: BaseModel] = ModelT | JsonValue | str | bytes @@ -153,26 +174,46 @@ class CompositorControl: @dataclass(kw_only=True) -class Compositor[PromptT, ToolT]: - """Framework-neutral ordered layer graph with lifecycle and aggregation.""" +class Compositor(Generic[PromptT, ToolT, LayerPromptT, LayerToolT]): + """Framework-neutral ordered layer graph with lifecycle and aggregation. - layers: OrderedDict[str, Layer[Any, PromptT, ToolT]] + ``prompt_transformer`` and ``tool_transformer`` are post-aggregation hooks: + they run whenever ``prompts`` or ``tools`` is read, after layer + contributions have been collected in compositor order. Use two type + arguments for identity aggregation, or all four when layer item types differ + from exposed item types. + """ + + layers: OrderedDict[str, Layer[Any, Any, Any]] deps_name_mapping: Mapping[str, Mapping[str, str]] = field(default_factory=dict) + prompt_transformer: CompositorTransformer[LayerPromptT, PromptT] | None = None + tool_transformer: CompositorTransformer[LayerToolT, ToolT] | None = None _deps_bound: bool = field(default=False, init=False) def __post_init__(self) -> None: self._bind_deps(self.deps_name_mapping) @classmethod - def from_config(cls, conf: CompositorConfigValue) -> Self: + def from_config( + cls, + conf: CompositorConfigValue, + *, + prompt_transformer: CompositorTransformer[LayerPromptT, PromptT] | None = None, + tool_transformer: CompositorTransformer[LayerToolT, ToolT] | None = None, + ) -> Self: """Create layers from config-like input and bind named dependencies.""" conf = _validate_compositor_config_input(conf) - layers: OrderedDict[str, Layer[Any, PromptT, ToolT]] = OrderedDict() + layers: OrderedDict[str, Layer[Any, Any, Any]] = OrderedDict() for layer_conf in conf.layers: - layers[layer_conf.name] = cast(Layer[Any, PromptT, ToolT], layer_conf.create_layer()) + layers[layer_conf.name] = layer_conf.create_layer() deps_name_mapping = {layer_conf.name: layer_conf.deps for layer_conf in conf.layers} - return cls(layers=layers, deps_name_mapping=deps_name_mapping) + return cls( + layers=layers, + deps_name_mapping=deps_name_mapping, + prompt_transformer=prompt_transformer, + tool_transformer=tool_transformer, + ) def _bind_deps(self, deps_name_mapping: Mapping[str, Mapping[str, str]]) -> None: """Resolve dependency-name mappings and bind dependencies on each layer. @@ -230,19 +271,29 @@ class Compositor[PromptT, ToolT]: @property def prompts(self) -> list[PromptT]: - result: list[PromptT] = [] + result: list[LayerPromptT] = [] for layer in self.layers.values(): - result.extend(layer.prefix_prompts) + result.extend( + cast(LayerPromptT, layer.wrap_prompt(prompt)) + for prompt in layer.prefix_prompts + ) for layer in reversed(self.layers.values()): - result.extend(layer.suffix_prompts) - return result + result.extend( + cast(LayerPromptT, layer.wrap_prompt(prompt)) + for prompt in layer.suffix_prompts + ) + if self.prompt_transformer is None: + return cast(list[PromptT], result) + return list(self.prompt_transformer(result)) @property def tools(self) -> list[ToolT]: - result: list[ToolT] = [] + result: list[LayerToolT] = [] for layer in self.layers.values(): - result.extend(layer.tools) - return result + result.extend(cast(LayerToolT, layer.wrap_tool(tool)) for tool in layer.tools) + if self.tool_transformer is None: + return cast(list[ToolT], result) + return list(self.tool_transformer(result)) __all__ = [ @@ -251,6 +302,8 @@ __all__ = [ "CompositorConfigValue", "CompositorLayerConfigInput", "CompositorControl", + "CompositorTransformer", + "CompositorTransformerKwargs", "CompositorLayerConfig", "CompositorLayerConfigValue", "ImportedLayerConfig", diff --git a/dify-agent/src/agenton/compositor/helpers.py b/dify-agent/src/agenton/compositor/helpers.py deleted file mode 100644 index 42122c6c90..0000000000 --- a/dify-agent/src/agenton/compositor/helpers.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Type-inference helpers for compositor construction. - -The core ``Compositor`` stays framework-neutral and usually needs explicit -prompt/tool type parameters. ``make_compositor`` is a small runtime factory -whose overloads let type checkers infer prompt and tool unions from the layer -arguments without introducing annotation-only compositor aliases. -""" - -from __future__ import annotations - -from collections import OrderedDict -from typing import TYPE_CHECKING, Any, Mapping, overload - -from agenton.layers.base import Layer - -if TYPE_CHECKING: - from . import Compositor - -type NamedLayer[PromptT, ToolT] = tuple[str, Layer[Any, PromptT, ToolT]] - - -@overload -def make_compositor[PromptT1, ToolT1]( - layer1: NamedLayer[PromptT1, ToolT1], - /, - *, - deps_name_mapping: Mapping[str, Mapping[str, str]] | None = None, -) -> Compositor[PromptT1, ToolT1]: ... - - -@overload -def make_compositor[PromptT1, ToolT1, PromptT2, ToolT2]( - layer1: NamedLayer[PromptT1, ToolT1], - layer2: NamedLayer[PromptT2, ToolT2], - /, - *, - deps_name_mapping: Mapping[str, Mapping[str, str]] | None = None, -) -> Compositor[PromptT1 | PromptT2, ToolT1 | ToolT2]: ... - - -@overload -def make_compositor[PromptT1, ToolT1, PromptT2, ToolT2, PromptT3, ToolT3]( - layer1: NamedLayer[PromptT1, ToolT1], - layer2: NamedLayer[PromptT2, ToolT2], - layer3: NamedLayer[PromptT3, ToolT3], - /, - *, - deps_name_mapping: Mapping[str, Mapping[str, str]] | None = None, -) -> Compositor[PromptT1 | PromptT2 | PromptT3, ToolT1 | ToolT2 | ToolT3]: ... - - -@overload -def make_compositor[ - PromptT1, - ToolT1, - PromptT2, - ToolT2, - PromptT3, - ToolT3, - PromptT4, - ToolT4, -]( - layer1: NamedLayer[PromptT1, ToolT1], - layer2: NamedLayer[PromptT2, ToolT2], - layer3: NamedLayer[PromptT3, ToolT3], - layer4: NamedLayer[PromptT4, ToolT4], - /, - *, - deps_name_mapping: Mapping[str, Mapping[str, str]] | None = None, -) -> Compositor[PromptT1 | PromptT2 | PromptT3 | PromptT4, ToolT1 | ToolT2 | ToolT3 | ToolT4]: ... - - -@overload -def make_compositor[ - PromptT1, - ToolT1, - PromptT2, - ToolT2, - PromptT3, - ToolT3, - PromptT4, - ToolT4, - PromptT5, - ToolT5, -]( - layer1: NamedLayer[PromptT1, ToolT1], - layer2: NamedLayer[PromptT2, ToolT2], - layer3: NamedLayer[PromptT3, ToolT3], - layer4: NamedLayer[PromptT4, ToolT4], - layer5: NamedLayer[PromptT5, ToolT5], - /, - *, - deps_name_mapping: Mapping[str, Mapping[str, str]] | None = None, -) -> Compositor[ - PromptT1 | PromptT2 | PromptT3 | PromptT4 | PromptT5, - ToolT1 | ToolT2 | ToolT3 | ToolT4 | ToolT5, -]: ... - - -def make_compositor( - *layers: NamedLayer[Any, Any], - deps_name_mapping: Mapping[str, Mapping[str, str]] | None = None, -) -> Compositor[Any, Any]: - """Create a compositor while letting type checkers infer layer item unions.""" - from . import Compositor - - return Compositor( - layers=OrderedDict(layers), - deps_name_mapping=deps_name_mapping or {}, - ) diff --git a/dify-agent/src/agenton/layers/__init__.py b/dify-agent/src/agenton/layers/__init__.py index 47c92108d3..d7f2365fc6 100644 --- a/dify-agent/src/agenton/layers/__init__.py +++ b/dify-agent/src/agenton/layers/__init__.py @@ -7,23 +7,35 @@ families while keeping concrete reusable layers in ``agenton_collections``. from agenton.layers.base import Layer, LayerControl, LayerDeps, NoLayerDeps from agenton.layers.types import ( + AllPromptTypes, + AllToolTypes, PlainLayer, PlainPrompt, + PlainPromptType, PlainTool, + PlainToolType, PydanticAILayer, PydanticAIPrompt, + PydanticAIPromptType, PydanticAITool, + PydanticAIToolType, ) __all__ = [ + "AllPromptTypes", + "AllToolTypes", "Layer", "LayerDeps", "LayerControl", "NoLayerDeps", "PlainLayer", "PlainPrompt", + "PlainPromptType", "PlainTool", + "PlainToolType", "PydanticAILayer", "PydanticAIPrompt", + "PydanticAIPromptType", "PydanticAITool", + "PydanticAIToolType", ] diff --git a/dify-agent/src/agenton/layers/base.py b/dify-agent/src/agenton/layers/base.py index 48023ff337..4ab5bb12eb 100644 --- a/dify-agent/src/agenton/layers/base.py +++ b/dify-agent/src/agenton/layers/base.py @@ -16,12 +16,14 @@ exits from temporary exits. The control is also the external lifecycle state: reuse a ``tmp_leave`` control to reenter, or pass a fresh control to start from create logic. -``Layer`` is framework-neutral over prompt and tool item types. Typed families -such as ``agenton.layers.types.PlainLayer`` bind those generic slots to a -specific contract without pushing framework types into this base module. +``Layer`` is framework-neutral over prompt and tool item types. The native +``prefix_prompts``, ``suffix_prompts``, and ``tools`` properties are the layer +authoring surface. ``wrap_prompt`` and ``wrap_tool`` are the compositor +aggregation surface; typed families such as ``agenton.layers.types.PlainLayer`` +implement them to tag native values without changing layer implementations. """ -from abc import ABC +from abc import ABC, abstractmethod from collections.abc import AsyncIterator from contextlib import AbstractAsyncContextManager, asynccontextmanager from dataclasses import dataclass @@ -202,6 +204,16 @@ class Layer[DepsT: LayerDeps, PromptT, ToolT](ABC): def tools(self) -> Sequence[ToolT]: return [] + @abstractmethod + def wrap_prompt(self, prompt: PromptT) -> object: + """Wrap a native prompt item for compositor aggregation.""" + raise NotImplementedError + + @abstractmethod + def wrap_tool(self, tool: ToolT) -> object: + """Wrap a native tool item for compositor aggregation.""" + raise NotImplementedError + def _get_dep_specs(deps_type: type[LayerDeps]) -> dict[str, LayerDepSpec]: dep_specs: dict[str, LayerDepSpec] = {} diff --git a/dify-agent/src/agenton/layers/types.py b/dify-agent/src/agenton/layers/types.py index 1572b1bc3d..3a705964a6 100644 --- a/dify-agent/src/agenton/layers/types.py +++ b/dify-agent/src/agenton/layers/types.py @@ -3,14 +3,24 @@ ``Layer`` itself is framework-neutral. This module defines typed layer families that bind its prompt/tool generic slots to concrete contracts, such as ordinary string prompts with plain callable tools or pydantic-ai prompt/tool shapes. +Tagged aggregate aliases cover code paths that can accept any supported +prompt/tool family without changing the plain and pydantic-ai layer contracts. +Pydantic-ai names are imported for static analysis only, so ``agenton`` can be +imported without loading that optional integration at runtime. Concrete reusable layers live under ``agenton_collections``. """ -from collections.abc import Callable -from typing import Any +from __future__ import annotations -from pydantic_ai import Tool -from pydantic_ai.tools import SystemPromptFunc, ToolFuncEither +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal + +from typing_extensions import final, override + +if TYPE_CHECKING: + from pydantic_ai import Tool + from pydantic_ai.tools import SystemPromptFunc from agenton.layers.base import Layer, LayerDeps @@ -18,12 +28,58 @@ type PlainPrompt = str type PlainTool = Callable[..., Any] +type PydanticAIPrompt[AgentDepsT] = SystemPromptFunc[AgentDepsT] +type PydanticAITool[AgentDepsT] = Tool[AgentDepsT] + + +@dataclass(frozen=True, slots=True) +class PlainPromptType: + """Tagged plain prompt item for aggregate prompt transformations.""" + + value: PlainPrompt + kind: Literal["plain"] = field(default="plain", init=False) + + +@dataclass(frozen=True, slots=True) +class PlainToolType: + """Tagged plain tool item for aggregate tool transformations.""" + + value: PlainTool + kind: Literal["plain"] = field(default="plain", init=False) + + +@dataclass(frozen=True, slots=True) +class PydanticAIPromptType[AgentDepsT]: + """Tagged pydantic-ai prompt item for aggregate prompt transformations.""" + + value: PydanticAIPrompt[AgentDepsT] + kind: Literal["pydantic_ai"] = field(default="pydantic_ai", init=False) + + +@dataclass(frozen=True, slots=True) +class PydanticAIToolType[AgentDepsT]: + """Tagged pydantic-ai tool item for aggregate tool transformations.""" + + value: PydanticAITool[AgentDepsT] + kind: Literal["pydantic_ai"] = field(default="pydantic_ai", init=False) + + +type AllPromptTypes = PlainPromptType | PydanticAIPromptType[Any] +type AllToolTypes = PlainToolType | PydanticAIToolType[Any] + + class PlainLayer[DepsT: LayerDeps](Layer[DepsT, PlainPrompt, PlainTool]): """Layer base for ordinary string prompts and plain-callable tools.""" + @final + @override + def wrap_prompt(self, prompt: PlainPrompt) -> PlainPromptType: + return PlainPromptType(prompt) -type PydanticAIPrompt[AgentDepsT] = str | SystemPromptFunc[AgentDepsT] -type PydanticAITool[AgentDepsT] = Tool[AgentDepsT] | ToolFuncEither[AgentDepsT, ...] + @final + @override + def wrap_tool(self, tool: PlainTool) -> PlainToolType: + return PlainToolType(tool) class PydanticAILayer[DepsT: LayerDeps, AgentDepsT]( @@ -31,12 +87,31 @@ class PydanticAILayer[DepsT: LayerDeps, AgentDepsT]( ): """Layer base for pydantic-ai prompt and tool adapters.""" + @final + @override + def wrap_prompt( + self, + prompt: PydanticAIPrompt[AgentDepsT], + ) -> PydanticAIPromptType[AgentDepsT]: + return PydanticAIPromptType(prompt) + + @final + @override + def wrap_tool(self, tool: PydanticAITool[AgentDepsT]) -> PydanticAIToolType[AgentDepsT]: + return PydanticAIToolType(tool) + __all__ = [ + "AllPromptTypes", + "AllToolTypes", "PlainLayer", "PlainPrompt", + "PlainPromptType", "PlainTool", + "PlainToolType", "PydanticAILayer", "PydanticAIPrompt", + "PydanticAIPromptType", "PydanticAITool", + "PydanticAIToolType", ] diff --git a/dify-agent/src/agenton_collections/__init__.py b/dify-agent/src/agenton_collections/__init__.py index 3bd8775f64..6e0b3750f0 100644 --- a/dify-agent/src/agenton_collections/__init__.py +++ b/dify-agent/src/agenton_collections/__init__.py @@ -7,17 +7,22 @@ implementation code in ``__init__``. """ from agenton.layers.types import ( + AllPromptTypes, + AllToolTypes, PlainLayer, PlainPrompt, + PlainPromptType, PlainTool, + PlainToolType, PydanticAILayer, PydanticAIPrompt, + PydanticAIPromptType, PydanticAITool, + PydanticAIToolType, ) from agenton_collections.layers.pydantic_ai import ( PydanticAIBridgeLayer, PydanticAIBridgeLayerDeps, - PydanticAIPrompts, ) from agenton_collections.layers.plain import ( DynamicToolsLayer, @@ -29,19 +34,24 @@ from agenton_collections.layers.plain import ( ) __all__ = [ + "AllPromptTypes", + "AllToolTypes", "DynamicToolsLayer", "DynamicToolsLayerDeps", "ObjectLayer", "PlainLayer", "PlainPrompt", + "PlainPromptType", "PlainTool", + "PlainToolType", "PromptLayer", "PydanticAIBridgeLayer", "PydanticAIBridgeLayerDeps", "PydanticAILayer", "PydanticAIPrompt", - "PydanticAIPrompts", + "PydanticAIPromptType", "PydanticAITool", + "PydanticAIToolType", "ToolsLayer", "with_object", ] diff --git a/dify-agent/src/agenton_collections/layers/pydantic_ai/__init__.py b/dify-agent/src/agenton_collections/layers/pydantic_ai/__init__.py index 7399966d18..c8b657cfdb 100644 --- a/dify-agent/src/agenton_collections/layers/pydantic_ai/__init__.py +++ b/dify-agent/src/agenton_collections/layers/pydantic_ai/__init__.py @@ -3,11 +3,9 @@ from agenton_collections.layers.pydantic_ai.bridge import ( PydanticAIBridgeLayer, PydanticAIBridgeLayerDeps, - PydanticAIPrompts, ) __all__ = [ "PydanticAIBridgeLayer", "PydanticAIBridgeLayerDeps", - "PydanticAIPrompts", ] diff --git a/dify-agent/src/agenton_collections/layers/pydantic_ai/bridge.py b/dify-agent/src/agenton_collections/layers/pydantic_ai/bridge.py index 2f8573a60e..89a0ae9023 100644 --- a/dify-agent/src/agenton_collections/layers/pydantic_ai/bridge.py +++ b/dify-agent/src/agenton_collections/layers/pydantic_ai/bridge.py @@ -4,19 +4,22 @@ This module keeps pydantic-ai's callable shapes intact through ``PydanticAILayer``. The bridge layer depends on ``ObjectLayer`` so callers have one explicit graph node that provides the object used as ``RunContext[ObjectT].deps`` in pydantic-ai prompt and tool callables. +Bridge construction accepts pydantic-ai's ergonomic input forms and normalizes +them at the layer boundary: string prompts become zero-arg system prompt +functions, and bare tool functions become ``Tool`` instances. """ from collections.abc import Sequence from dataclasses import dataclass +from pydantic_ai import Tool +from pydantic_ai.tools import ToolFuncEither from typing_extensions import override from agenton.layers.base import LayerDeps from agenton.layers.types import PydanticAILayer, PydanticAIPrompt, PydanticAITool from agenton_collections.layers.plain.basic import ObjectLayer -type PydanticAIPrompts[ObjectT] = PydanticAIPrompt[ObjectT] | Sequence[PydanticAIPrompt[ObjectT]] - class PydanticAIBridgeLayerDeps[ObjectT](LayerDeps): """Dependencies required by ``PydanticAIBridgeLayer``.""" @@ -30,9 +33,9 @@ class PydanticAIBridgeLayer[ObjectT]( ): """Bridge layer for pydantic-ai prompts and tools using one object deps.""" - prefix: PydanticAIPrompts[ObjectT] = () - suffix: PydanticAIPrompts[ObjectT] = () - tool_entries: Sequence[PydanticAITool[ObjectT]] = () + prefix: str | PydanticAIPrompt[ObjectT] | Sequence[str | PydanticAIPrompt[ObjectT]] = () + suffix: str | PydanticAIPrompt[ObjectT] | Sequence[str | PydanticAIPrompt[ObjectT]] = () + tool_entries: Sequence[PydanticAITool[ObjectT] | ToolFuncEither[ObjectT, ...]] = () @property def run_deps(self) -> ObjectT: @@ -52,19 +55,36 @@ class PydanticAIBridgeLayer[ObjectT]( @property @override def tools(self) -> list[PydanticAITool[ObjectT]]: - return list(self.tool_entries) + return [_normalize_tool(tool_entry) for tool_entry in self.tool_entries] def _normalize_prompts[ObjectT]( - prompts: PydanticAIPrompts[ObjectT], + prompts: str | PydanticAIPrompt[ObjectT] | Sequence[str | PydanticAIPrompt[ObjectT]], ) -> list[PydanticAIPrompt[ObjectT]]: - if isinstance(prompts, str) or callable(prompts): - return [prompts] - return list(prompts) + if isinstance(prompts, str): + return [_normalize_prompt(prompts)] + if isinstance(prompts, Sequence): + return [_normalize_prompt(prompt) for prompt in prompts] + return [prompts] + + +def _normalize_prompt[ObjectT]( + prompt: str | PydanticAIPrompt[ObjectT], +) -> PydanticAIPrompt[ObjectT]: + if isinstance(prompt, str): + return (lambda value: lambda: value)(prompt) + return prompt + + +def _normalize_tool[ObjectT]( + tool_entry: PydanticAITool[ObjectT] | ToolFuncEither[ObjectT, ...], +) -> PydanticAITool[ObjectT]: + if isinstance(tool_entry, Tool): + return tool_entry + return Tool(tool_entry) __all__ = [ "PydanticAIBridgeLayer", "PydanticAIBridgeLayerDeps", - "PydanticAIPrompts", ] diff --git a/dify-agent/src/agenton_collections/transformers/__init__.py b/dify-agent/src/agenton_collections/transformers/__init__.py new file mode 100644 index 0000000000..dd46275944 --- /dev/null +++ b/dify-agent/src/agenton_collections/transformers/__init__.py @@ -0,0 +1,11 @@ +"""Reusable compositor transformers for collection integrations.""" + +from agenton_collections.transformers.pydantic_ai import ( + PYDANTIC_AI_TRANSFORMERS, + PydanticAICompositorTransformerKwargs, +) + +__all__ = [ + "PYDANTIC_AI_TRANSFORMERS", + "PydanticAICompositorTransformerKwargs", +] diff --git a/dify-agent/src/agenton_collections/transformers/pydantic_ai.py b/dify-agent/src/agenton_collections/transformers/pydantic_ai.py new file mode 100644 index 0000000000..12e067063f --- /dev/null +++ b/dify-agent/src/agenton_collections/transformers/pydantic_ai.py @@ -0,0 +1,65 @@ +"""Pydantic AI compositor transformer presets. + +This module owns the pydantic-ai runtime dependency for transforming tagged +agenton prompt/tool items into pydantic-ai-compatible items. +""" + +from collections.abc import Sequence +from typing import Final + +from pydantic_ai import Tool + +from agenton.compositor import CompositorTransformerKwargs +from agenton.layers.types import ( + AllPromptTypes, + AllToolTypes, + PydanticAIPrompt, + PydanticAITool, +) + +type PydanticAICompositorTransformerKwargs = CompositorTransformerKwargs[ + PydanticAIPrompt[object], + PydanticAITool[object], + AllPromptTypes, + AllToolTypes, +] + + +def _pydantic_ai_prompt_transformer( + prompts: Sequence[AllPromptTypes], +) -> list[PydanticAIPrompt[object]]: + result: list[PydanticAIPrompt[object]] = [] + for prompt in prompts: + if prompt.kind == "plain": + result.append((lambda value: lambda: value)(prompt.value)) + elif prompt.kind == "pydantic_ai": + result.append(prompt.value) + else: + raise NotImplementedError(f"Unsupported prompt type: {type(prompt).__qualname__}.") + return result + + +def _pydantic_ai_tool_transformer( + tools: Sequence[AllToolTypes], +) -> list[PydanticAITool[object]]: + result: list[PydanticAITool[object]] = [] + for tool in tools: + if tool.kind == "plain": + result.append(Tool(tool.value)) + elif tool.kind == "pydantic_ai": + result.append(tool.value) + else: + raise NotImplementedError(f"Unsupported tool type: {type(tool).__qualname__}.") + return result + + +PYDANTIC_AI_TRANSFORMERS: Final[PydanticAICompositorTransformerKwargs] = { + "prompt_transformer": _pydantic_ai_prompt_transformer, + "tool_transformer": _pydantic_ai_tool_transformer, +} + + +__all__ = [ + "PYDANTIC_AI_TRANSFORMERS", + "PydanticAICompositorTransformerKwargs", +] diff --git a/dify-agent/tests/unit/agenton/compositor/test_enter.py b/dify-agent/tests/unit/agenton/compositor/test_enter.py index 208c541b0c..59b077bf1b 100644 --- a/dify-agent/tests/unit/agenton/compositor/test_enter.py +++ b/dify-agent/tests/unit/agenton/compositor/test_enter.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, field from typing_extensions import override from agenton.compositor import Compositor, CompositorControl -from agenton.layers import LayerControl, NoLayerDeps, PlainLayer +from agenton.layers import LayerControl, NoLayerDeps, PlainLayer, PlainPromptType, PlainToolType @dataclass(slots=True) @@ -34,7 +34,7 @@ class TraceLayer(PlainLayer[NoLayerDeps]): def test_compositor_enter_creates_control_and_applies_tmp_leave_to_all_layers() -> None: first_layer = TraceLayer() second_layer = TraceLayer() - compositor: Compositor[str, object] = Compositor( + compositor: Compositor[PlainPromptType, PlainToolType] = Compositor( layers=OrderedDict( [ ("first", first_layer), @@ -61,7 +61,7 @@ def test_compositor_enter_creates_control_and_applies_tmp_leave_to_all_layers() def test_compositor_enter_does_not_store_tmp_leave_on_layer() -> None: layer = TraceLayer() - compositor: Compositor[str, object] = Compositor( + compositor: Compositor[PlainPromptType, PlainToolType] = Compositor( layers=OrderedDict([("trace", layer)]) ) @@ -79,7 +79,7 @@ def test_compositor_enter_does_not_store_tmp_leave_on_layer() -> None: def test_compositor_enter_rejects_control_with_mismatched_layer_names() -> None: layer = TraceLayer() - compositor: Compositor[str, object] = Compositor( + compositor: Compositor[PlainPromptType, PlainToolType] = Compositor( layers=OrderedDict([("trace", layer)]) ) compositor_control = CompositorControl(["other"]) diff --git a/dify-agent/tests/unit/agenton/compositor/test_transformers.py b/dify-agent/tests/unit/agenton/compositor/test_transformers.py new file mode 100644 index 0000000000..34f7ecccf3 --- /dev/null +++ b/dify-agent/tests/unit/agenton/compositor/test_transformers.py @@ -0,0 +1,120 @@ +from collections import OrderedDict +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from inspect import Parameter, signature + +from typing_extensions import override + +from agenton.compositor import Compositor, CompositorTransformerKwargs +from agenton.layers import NoLayerDeps, PlainLayer, PlainPromptType, PlainToolType + +type ToolCallable = Callable[..., object] +type WrappedPrompt = tuple[str, str] + + +@dataclass(slots=True) +class PromptAndToolLayer(PlainLayer[NoLayerDeps]): + prefix: list[str] + suffix: list[str] + tool_entries: list[ToolCallable] + + @property + @override + def prefix_prompts(self) -> list[str]: + return self.prefix + + @property + @override + def suffix_prompts(self) -> list[str]: + return self.suffix + + @property + @override + def tools(self) -> list[ToolCallable]: + return self.tool_entries + + +def base_tool() -> str: + return "base" + + +def wrapped_tool() -> str: + return "wrapped" + + +def wrap_prompts(prompts: Sequence[PlainPromptType]) -> list[WrappedPrompt]: + return [("wrapped", prompt.value) for prompt in prompts] + + +def describe_tools(tools: Sequence[PlainToolType]) -> list[str]: + return [tool.value.__name__ for tool in tools] + + +def test_compositor_transformer_kwargs_keys_match_constructor_parameters() -> None: + transformer_kwargs = set(CompositorTransformerKwargs.__required_keys__) + parameters = signature(Compositor).parameters + + assert CompositorTransformerKwargs.__optional_keys__ == frozenset() + assert transformer_kwargs == { + name for name in parameters if name.endswith("_transformer") + } + assert all(parameters[name].kind is Parameter.KEYWORD_ONLY for name in transformer_kwargs) + + +def test_compositor_transformer_kwargs_keys_match_from_config_parameters() -> None: + transformer_kwargs = set(CompositorTransformerKwargs.__required_keys__) + parameters = signature(Compositor.from_config).parameters + + assert transformer_kwargs == { + name for name in parameters if name.endswith("_transformer") + } + assert all(parameters[name].kind is Parameter.KEYWORD_ONLY for name in transformer_kwargs) + + +def test_compositor_transforms_prompts_to_another_type_after_layer_ordering() -> None: + compositor: Compositor[WrappedPrompt, PlainToolType, PlainPromptType, PlainToolType] = Compositor( + layers=OrderedDict( + [ + ( + "first", + PromptAndToolLayer( + prefix=["first-prefix"], + suffix=["first-suffix"], + tool_entries=[], + ), + ), + ( + "second", + PromptAndToolLayer( + prefix=["second-prefix"], + suffix=["second-suffix"], + tool_entries=[], + ), + ), + ] + ), + prompt_transformer=wrap_prompts, + ) + + assert compositor.prompts == [ + ("wrapped", "first-prefix"), + ("wrapped", "second-prefix"), + ("wrapped", "second-suffix"), + ("wrapped", "first-suffix"), + ] + + +def test_compositor_transforms_tools_to_another_type_after_layer_aggregation() -> None: + compositor: Compositor[PlainPromptType, str, PlainPromptType, PlainToolType] = Compositor( + layers=OrderedDict( + [ + ( + "tools", + PromptAndToolLayer(prefix=[], suffix=[], tool_entries=[base_tool, wrapped_tool]), + ) + ] + ), + tool_transformer=describe_tools, + ) + + assert compositor.tools == ["base_tool", "wrapped_tool"] diff --git a/dify-agent/tests/unit/agenton/layers/test_layer_deps.py b/dify-agent/tests/unit/agenton/layers/test_layer_deps.py index 950cad03aa..1ab0c543e2 100644 --- a/dify-agent/tests/unit/agenton/layers/test_layer_deps.py +++ b/dify-agent/tests/unit/agenton/layers/test_layer_deps.py @@ -1,7 +1,7 @@ import pytest from agenton.layers import LayerDeps -from agenton_collections.plain import ObjectLayer, PromptLayer +from agenton_collections.layers.plain import ObjectLayer, PromptLayer class ObjectLayerDeps(LayerDeps): diff --git a/dify-agent/tests/unit/agenton_collections/layers/pydantic_ai/test_bridge.py b/dify-agent/tests/unit/agenton_collections/layers/pydantic_ai/test_bridge.py new file mode 100644 index 0000000000..cfa93bd407 --- /dev/null +++ b/dify-agent/tests/unit/agenton_collections/layers/pydantic_ai/test_bridge.py @@ -0,0 +1,53 @@ +from collections.abc import Callable +from dataclasses import dataclass +from typing import cast + +from pydantic_ai import RunContext, Tool + +from agenton_collections.layers.pydantic_ai import PydanticAIBridgeLayer + + +@dataclass(frozen=True, slots=True) +class Profile: + name: str + + +def profile_prompt(ctx: RunContext[Profile]) -> str: + return f"Profile: {ctx.deps.name}" + + +def existing_tool(ctx: RunContext[Profile]) -> str: + return ctx.deps.name + + +def raw_tool(ctx: RunContext[Profile], topic: str) -> str: + return f"{ctx.deps.name}: {topic}" + + +def test_pydantic_ai_bridge_layer_accepts_mixed_string_and_function_prompts() -> None: + layer = PydanticAIBridgeLayer[Profile]( + prefix=("plain prefix", profile_prompt), + suffix="plain suffix", + ) + + prefix_prompts = layer.prefix_prompts + suffix_prompts = layer.suffix_prompts + + plain_prefix = cast(Callable[[], str], prefix_prompts[0]) + plain_suffix = cast(Callable[[], str], suffix_prompts[0]) + assert plain_prefix() == "plain prefix" + assert prefix_prompts[1] is profile_prompt + assert plain_suffix() == "plain suffix" + + +def test_pydantic_ai_bridge_layer_accepts_mixed_tool_and_tool_function_entries() -> None: + pydantic_ai_tool = Tool(existing_tool) + layer = PydanticAIBridgeLayer[Profile]( + tool_entries=(pydantic_ai_tool, raw_tool), + ) + + tools = layer.tools + + assert tools[0] is pydantic_ai_tool + assert isinstance(tools[1], Tool) + assert tools[1].function is raw_tool diff --git a/dify-agent/tests/unit/agenton_collections/transformers/test_pydantic_ai.py b/dify-agent/tests/unit/agenton_collections/transformers/test_pydantic_ai.py new file mode 100644 index 0000000000..cc0dbdc16f --- /dev/null +++ b/dify-agent/tests/unit/agenton_collections/transformers/test_pydantic_ai.py @@ -0,0 +1,75 @@ +from collections.abc import Callable +from typing import cast + +from pydantic_ai import Tool + +from agenton.layers.types import ( + PlainPromptType, + PlainToolType, + PydanticAIPromptType, + PydanticAIToolType, +) +from agenton_collections.transformers.pydantic_ai import PYDANTIC_AI_TRANSFORMERS + + +def plain_tool(name: str) -> str: + return f"hello {name}" + + +def dynamic_prompt() -> str: + return "dynamic prompt" + + +def test_pydantic_ai_transformers_wrap_tagged_plain_prompts() -> None: + prompts = [PlainPromptType("plain prompt")] + + result = PYDANTIC_AI_TRANSFORMERS["prompt_transformer"](prompts) + + assert len(result) == 1 + prompt_func = cast(Callable[[], str], result[0]) + assert prompt_func() == "plain prompt" + + +def test_pydantic_ai_transformers_preserve_tagged_existing_prompt_functions() -> None: + result = PYDANTIC_AI_TRANSFORMERS["prompt_transformer"]([PydanticAIPromptType(dynamic_prompt)]) + + assert result == [dynamic_prompt] + + +def test_pydantic_ai_transformers_accept_mixed_tagged_prompt_types() -> None: + result = PYDANTIC_AI_TRANSFORMERS["prompt_transformer"]( + [PlainPromptType("plain prompt"), PydanticAIPromptType(dynamic_prompt)] + ) + + plain_prompt = cast(Callable[[], str], result[0]) + assert plain_prompt() == "plain prompt" + assert result[1] is dynamic_prompt + + +def test_pydantic_ai_transformers_wrap_tagged_plain_tools() -> None: + result = PYDANTIC_AI_TRANSFORMERS["tool_transformer"]([PlainToolType(plain_tool)]) + + assert len(result) == 1 + tool = result[0] + assert isinstance(tool, Tool) + assert tool.function is plain_tool + + +def test_pydantic_ai_transformers_preserve_tagged_existing_tools() -> None: + pydantic_ai_tool = Tool(plain_tool) + + result = PYDANTIC_AI_TRANSFORMERS["tool_transformer"]([PydanticAIToolType(pydantic_ai_tool)]) + + assert result == [pydantic_ai_tool] + + +def test_pydantic_ai_transformers_accept_tagged_tool_types() -> None: + pydantic_ai_tool = Tool(plain_tool) + + result = PYDANTIC_AI_TRANSFORMERS["tool_transformer"]( + [PlainToolType(plain_tool), PydanticAIToolType(pydantic_ai_tool)] + ) + + assert isinstance(result[0], Tool) + assert result[0].function is plain_tool + assert result[1] is pydantic_ai_tool