dify/dify-agent/tests/local/agenton/layers/test_schema_inference.py
盐粒 Yanli 55f95dbc36
feat(agent): init agent server (#36087)
Co-authored-by: Copilot <copilot@github.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
2026-05-14 06:04:44 +00:00

93 lines
2.7 KiB
Python

from dataclasses import dataclass
from pydantic import BaseModel, ConfigDict
from agenton.compositor import LayerProvider
from agenton.layers import (
EmptyLayerConfig,
EmptyRuntimeState,
LayerConfig,
NoLayerDeps,
PlainLayer,
)
class InferredConfig(LayerConfig):
value: str = "configured"
model_config = ConfigDict(extra="forbid")
class InferredState(BaseModel):
count: int = 0
model_config = ConfigDict(extra="forbid", validate_assignment=True)
@dataclass(slots=True)
class GenericSchemaLayer(PlainLayer[NoLayerDeps, InferredConfig, InferredState]):
type_id = "test.generic-schema"
@classmethod
def from_config(cls, config: InferredConfig) -> "GenericSchemaLayer":
return cls()
async def on_context_create(self) -> None:
self.runtime_state.count += 1
@dataclass(slots=True)
class DefaultSchemaLayer(PlainLayer[NoLayerDeps]):
type_id = "test.default-schema"
def test_layer_infers_config_and_runtime_state_from_generics() -> None:
layer = GenericSchemaLayer()
layer.runtime_state = InferredState(count=3)
assert GenericSchemaLayer.config_type is InferredConfig
assert GenericSchemaLayer.runtime_state_type is InferredState
assert isinstance(layer.runtime_state, InferredState)
assert layer.runtime_state.count == 3
def test_layer_uses_empty_schema_defaults_when_omitted() -> None:
layer = DefaultSchemaLayer()
assert DefaultSchemaLayer.config_type is EmptyLayerConfig
assert DefaultSchemaLayer.runtime_state_type is EmptyRuntimeState
assert isinstance(layer.runtime_state, EmptyRuntimeState)
def test_invalid_declared_schema_type_is_rejected_clearly() -> None:
try:
class InvalidSchemaLayer(PlainLayer[NoLayerDeps]):
config_type = dict # pyright: ignore[reportAssignmentType]
except TypeError as e:
assert str(e) == "InvalidSchemaLayer.config_type must be a LayerConfig subclass."
else:
raise AssertionError("Expected TypeError.")
try:
class InvalidGenericSchemaLayer(PlainLayer[NoLayerDeps, dict[str, object]]): # pyright: ignore[reportInvalidTypeArguments]
pass
except TypeError as e:
assert str(e) == "InvalidGenericSchemaLayer.config_type must be a LayerConfig subclass."
else:
raise AssertionError("Expected TypeError.")
def test_layer_provider_uses_inferred_schema_types() -> None:
provider = LayerProvider.from_layer_type(GenericSchemaLayer)
layer = provider.create_layer({"value": "configured"})
assert provider.type_id == "test.generic-schema"
assert provider.layer_type.config_type is InferredConfig
assert provider.layer_type.runtime_state_type is InferredState
assert isinstance(layer.config, InferredConfig)