From 2586925528ad1a1ff73a29c6e74e4d6ca3a5f2ef Mon Sep 17 00:00:00 2001 From: Nachiket Paranjape Date: Fri, 6 Feb 2026 14:42:02 -0800 Subject: [PATCH 1/9] control wrapped --- sdks/python/src/agent_control/__init__.py | 20 + .../src/agent_control/_control_registry.py | 191 +++++++++ .../src/agent_control/control_decorators.py | 5 + sdks/python/tests/test_control_registry.py | 396 ++++++++++++++++++ 4 files changed, 612 insertions(+) create mode 100644 sdks/python/src/agent_control/_control_registry.py create mode 100644 sdks/python/tests/test_control_registry.py diff --git a/sdks/python/src/agent_control/__init__.py b/sdks/python/src/agent_control/__init__.py index 06a16069..7b2f6f1b 100644 --- a/sdks/python/src/agent_control/__init__.py +++ b/sdks/python/src/agent_control/__init__.py @@ -66,6 +66,8 @@ async def process(input: str) -> str: ) from . import agents, controls, evaluation, evaluators, policies +from ._control_registry import clear as clear_step_registry +from ._control_registry import get_registered_steps # Import client and operations modules from .client import AgentControlClient @@ -435,6 +437,21 @@ async def handle(message: str): _server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000' _api_key = api_key + # Merge auto-discovered steps from @control() decorators with explicit steps. + # Explicit steps take precedence (by name) over auto-discovered ones. + from agent_control._control_registry import get_registered_steps + + auto_steps = get_registered_steps() + if auto_steps: + explicit_names = {s["name"] for s in (steps or [])} + merged = list(steps or []) + [s for s in auto_steps if s["name"] not in explicit_names] + steps = merged + logger.debug( + "Auto-discovered %d step(s) from @control() decorators (%d after merge)", + len(auto_steps), + len(steps), + ) + # Register with server and fetch controls server_controls = None try: @@ -1054,6 +1071,9 @@ async def main(): "get_server_controls", "refresh_controls", "refresh_controls_async", + # Step registry (auto-discovered from @control decorators) + "get_registered_steps", + "clear_step_registry", # SDK Logging "get_logger", diff --git a/sdks/python/src/agent_control/_control_registry.py b/sdks/python/src/agent_control/_control_registry.py new file mode 100644 index 00000000..6824c7ff --- /dev/null +++ b/sdks/python/src/agent_control/_control_registry.py @@ -0,0 +1,191 @@ +"""Registry for @control()-decorated functions. + +Tracks step schemas (name, type, input/output schema) from decorated functions +so they can be auto-populated into ``init(steps=...)`` without the user having +to specify them manually. + +Registration happens at **decoration time** (import time), so all decorated +functions are captured before ``init()`` is called -- as long as ``init()`` +is called after the module containing the decorated functions has been imported. +""" + +from __future__ import annotations + +import inspect +import logging +import typing +from collections.abc import Callable +from typing import Any + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Internal registry +# --------------------------------------------------------------------------- + +_registered_steps: dict[str, dict[str, Any]] = {} +"""Maps step name -> step schema dict. Keyed by name to deduplicate.""" + + +# --------------------------------------------------------------------------- +# Schema extraction helpers +# --------------------------------------------------------------------------- + +# Mapping from Python primitive types to JSON Schema type strings. +_PRIMITIVE_TYPE_MAP: dict[type, str] = { + str: "string", + int: "integer", + float: "number", + bool: "boolean", +} + + +def _type_to_json_schema(annotation: Any) -> dict[str, Any] | None: + """Convert a Python type annotation to a JSON Schema dict. + + Supports: + - Pydantic models (via ``model_json_schema()``) + - Primitive types (str, int, float, bool) + - Framework-specific ``args_schema`` on the function (checked separately) + + Returns ``None`` for complex or unrecognised types so callers can + gracefully degrade. + """ + if annotation is None or annotation is inspect.Parameter.empty: + return None + + # Pydantic v2 models expose model_json_schema() + if hasattr(annotation, "model_json_schema"): + try: + result: dict[str, Any] = annotation.model_json_schema() + return result + except Exception: + logger.debug("Failed to extract JSON schema from Pydantic model %s", annotation) + return None + + # Primitive types + type_str = _PRIMITIVE_TYPE_MAP.get(annotation) # type: ignore[arg-type] + if type_str is not None: + return {"type": type_str} + + return None + + +def _extract_input_schema(func: Callable[..., Any]) -> dict[str, Any] | None: + """Build a JSON Schema ``object`` from the function's parameter type hints. + + Skips ``self`` and ``cls`` parameters. Returns ``None`` when no useful + schema can be derived (e.g. no type hints at all). + """ + # Framework tools (e.g. LangChain) may expose a Pydantic args_schema + args_schema = getattr(func, "args_schema", None) + if args_schema is not None and hasattr(args_schema, "model_json_schema"): + try: + result: dict[str, Any] = args_schema.model_json_schema() + return result + except Exception: + logger.debug("Failed to extract args_schema from %s", func) + + try: + hints = typing.get_type_hints(func) + except Exception: + # get_type_hints can fail on some decorated / wrapped functions + return None + + sig = inspect.signature(func) + properties: dict[str, Any] = {} + for name, _param in sig.parameters.items(): + if name in ("self", "cls"): + continue + hint = hints.get(name) + if hint is None: + continue + schema = _type_to_json_schema(hint) + if schema is not None: + properties[name] = schema + + if not properties: + return None + return {"type": "object", "properties": properties} + + +def _extract_output_schema(func: Callable[..., Any]) -> dict[str, Any] | None: + """Derive output schema from the function's return type annotation.""" + try: + hints = typing.get_type_hints(func) + except Exception: + return None + + return_hint = hints.get("return") + if return_hint is None: + return None + return _type_to_json_schema(return_hint) + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def register(func: Callable[..., Any], policy: str | None = None) -> None: + """Register a decorated function's step schema in the registry. + + Extracts name, type (tool vs llm), description, and input/output schemas + from the function and stores them for later retrieval via + ``get_registered_steps()``. + + Args: + func: The original (unwrapped) function being decorated. + policy: Optional policy name (stored as metadata). + """ + # Determine step name -- tools typically have .name or .tool_name + tool_name = getattr(func, "name", None) or getattr(func, "tool_name", None) + step_name: str = tool_name if isinstance(tool_name, str) else func.__name__ + step_type: str = "tool" if isinstance(tool_name, str) else "llm" + + # Extract description from docstring (first line only) + description: str | None = None + if func.__doc__: + first_line = func.__doc__.strip().split("\n")[0].strip() + if first_line: + description = first_line + + input_schema = _extract_input_schema(func) + output_schema = _extract_output_schema(func) + + step: dict[str, Any] = { + "type": step_type, + "name": step_name, + } + if description is not None: + step["description"] = description + if input_schema is not None: + step["input_schema"] = input_schema + if output_schema is not None: + step["output_schema"] = output_schema + + metadata: dict[str, Any] = {} + if policy is not None: + metadata["policy"] = policy + if metadata: + step["metadata"] = metadata + + # Store (last-write-wins for duplicate names) + if step_name in _registered_steps: + logger.debug("Overwriting previously registered step '%s'", step_name) + _registered_steps[step_name] = step + logger.debug("Registered step schema: %s (type=%s)", step_name, step_type) + + +def get_registered_steps() -> list[dict[str, Any]]: + """Return all registered step schemas as a list of dicts. + + The returned dicts conform to the ``StepSchema`` model format expected + by ``init(steps=...)``. + """ + return list(_registered_steps.values()) + + +def clear() -> None: + """Clear all registered steps. Useful for testing.""" + _registered_steps.clear() diff --git a/sdks/python/src/agent_control/control_decorators.py b/sdks/python/src/agent_control/control_decorators.py index 7f6ae7bc..d9214d67 100644 --- a/sdks/python/src/agent_control/control_decorators.py +++ b/sdks/python/src/agent_control/control_decorators.py @@ -641,6 +641,11 @@ async def process(input: str) -> str: _ = policy def decorator(func: F) -> F: + # Register this function's step schema for auto-discovery by init() + from agent_control._control_registry import register + + register(func, policy) + @functools.wraps(func) async def async_wrapper(*args: Any, **kwargs: Any) -> Any: return await _execute_with_control(func, args, kwargs, is_async=True) diff --git a/sdks/python/tests/test_control_registry.py b/sdks/python/tests/test_control_registry.py new file mode 100644 index 00000000..2b5571c9 --- /dev/null +++ b/sdks/python/tests/test_control_registry.py @@ -0,0 +1,396 @@ +"""Tests for the control step registry (_control_registry).""" + +from __future__ import annotations + +from typing import Any +from unittest.mock import MagicMock + +import pytest + +from agent_control._control_registry import ( + _extract_input_schema, + _extract_output_schema, + _type_to_json_schema, + clear, + get_registered_steps, + register, +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _clean_registry() -> None: # noqa: PT004 + """Ensure each test starts with an empty registry.""" + clear() + yield # type: ignore[misc] + clear() + + +# =========================================================================== +# Schema extraction helpers +# =========================================================================== + + +class TestTypeToJsonSchema: + """Tests for _type_to_json_schema.""" + + def test_primitive_str(self) -> None: + assert _type_to_json_schema(str) == {"type": "string"} + + def test_primitive_int(self) -> None: + assert _type_to_json_schema(int) == {"type": "integer"} + + def test_primitive_float(self) -> None: + assert _type_to_json_schema(float) == {"type": "number"} + + def test_primitive_bool(self) -> None: + assert _type_to_json_schema(bool) == {"type": "boolean"} + + def test_none_returns_none(self) -> None: + assert _type_to_json_schema(None) is None + + def test_complex_type_returns_none(self) -> None: + """Complex / unrecognised types should return None.""" + assert _type_to_json_schema(list) is None + assert _type_to_json_schema(dict) is None + + def test_pydantic_model(self) -> None: + """Pydantic v2 models should delegate to model_json_schema().""" + mock_model = MagicMock() + mock_model.model_json_schema.return_value = { + "type": "object", + "properties": {"query": {"type": "string"}}, + } + result = _type_to_json_schema(mock_model) + assert result == {"type": "object", "properties": {"query": {"type": "string"}}} + mock_model.model_json_schema.assert_called_once() + + +class TestExtractInputSchema: + """Tests for _extract_input_schema.""" + + def test_simple_function(self) -> None: + def my_func(query: str, limit: int = 10) -> list: + ... + + schema = _extract_input_schema(my_func) + assert schema is not None + assert schema["type"] == "object" + assert schema["properties"]["query"] == {"type": "string"} + assert schema["properties"]["limit"] == {"type": "integer"} + + def test_no_type_hints(self) -> None: + def my_func(x, y): + ... + + assert _extract_input_schema(my_func) is None + + def test_skips_self_and_cls(self) -> None: + def my_method(self, query: str) -> str: # noqa: ANN001 + ... + + schema = _extract_input_schema(my_method) + assert schema is not None + assert "self" not in schema["properties"] + assert "query" in schema["properties"] + + def test_framework_args_schema(self) -> None: + """If func has .args_schema with model_json_schema(), use that.""" + mock_schema = MagicMock() + mock_schema.model_json_schema.return_value = { + "type": "object", + "properties": {"q": {"type": "string"}}, + } + + def my_func(q: str) -> str: + ... + + my_func.args_schema = mock_schema # type: ignore[attr-defined] + + result = _extract_input_schema(my_func) + assert result == {"type": "object", "properties": {"q": {"type": "string"}}} + mock_schema.model_json_schema.assert_called_once() + + +class TestExtractOutputSchema: + """Tests for _extract_output_schema.""" + + def test_str_return(self) -> None: + def my_func() -> str: + ... + + assert _extract_output_schema(my_func) == {"type": "string"} + + def test_int_return(self) -> None: + def my_func() -> int: + ... + + assert _extract_output_schema(my_func) == {"type": "integer"} + + def test_no_return_annotation(self) -> None: + def my_func(): + ... + + assert _extract_output_schema(my_func) is None + + def test_complex_return_type(self) -> None: + def my_func() -> list[str]: + ... + + # list[str] is not a supported primitive, should return None + assert _extract_output_schema(my_func) is None + + +# =========================================================================== +# register() and get_registered_steps() +# =========================================================================== + + +class TestRegister: + """Tests for register() and get_registered_steps().""" + + def test_register_simple_function(self) -> None: + def chat(message: str) -> str: + """Chat with the agent.""" + ... + + register(chat) + steps = get_registered_steps() + assert len(steps) == 1 + step = steps[0] + assert step["name"] == "chat" + assert step["type"] == "llm" + assert step["description"] == "Chat with the agent." + assert step["input_schema"] is not None + assert step["output_schema"] == {"type": "string"} + + def test_register_tool_function(self) -> None: + """Functions with .name or .tool_name should be registered as tools.""" + def search_db(query: str, limit: int = 10) -> str: + """Search the database.""" + ... + + search_db.name = "search_db" # type: ignore[attr-defined] + search_db.tool_name = "search_db" # type: ignore[attr-defined] + + register(search_db) + steps = get_registered_steps() + assert len(steps) == 1 + assert steps[0]["type"] == "tool" + assert steps[0]["name"] == "search_db" + + def test_register_with_policy(self) -> None: + def my_func(x: str) -> str: + ... + + register(my_func, policy="safety-policy") + steps = get_registered_steps() + assert steps[0]["metadata"] == {"policy": "safety-policy"} + + def test_register_no_policy_no_metadata(self) -> None: + def my_func(x: str) -> str: + ... + + register(my_func) + steps = get_registered_steps() + assert "metadata" not in steps[0] + + def test_deduplicate_by_name(self) -> None: + """Registering two functions with the same name should keep the last one.""" + def chat(message: str) -> str: + """First version.""" + ... + + def chat_v2(message: str) -> int: # noqa: ARG001 + """Second version.""" + ... + + chat_v2.__name__ = "chat" # simulate same name + + register(chat) + register(chat_v2) + steps = get_registered_steps() + assert len(steps) == 1 + assert steps[0]["description"] == "Second version." + + def test_no_docstring(self) -> None: + def my_func(x: str) -> str: + ... + + register(my_func) + steps = get_registered_steps() + assert "description" not in steps[0] + + def test_no_type_hints(self) -> None: + """Functions with no type hints should still register with None schemas.""" + def my_func(x, y): + ... + + register(my_func) + steps = get_registered_steps() + assert len(steps) == 1 + assert steps[0]["name"] == "my_func" + assert "input_schema" not in steps[0] + assert "output_schema" not in steps[0] + + +class TestClear: + """Tests for clear().""" + + def test_clear_empties_registry(self) -> None: + def f1(x: str) -> str: + ... + + def f2(x: str) -> str: + ... + + register(f1) + register(f2) + assert len(get_registered_steps()) == 2 + + clear() + assert len(get_registered_steps()) == 0 + + +# =========================================================================== +# Decorator integration +# =========================================================================== + + +class TestDecoratorRegistration: + """Tests that @control() decorator registers functions in the registry.""" + + def test_decorator_registers_async_function(self) -> None: + from agent_control.control_decorators import control + + @control() + async def my_chat(message: str) -> str: + """Handle a chat message.""" + return message + + steps = get_registered_steps() + assert len(steps) == 1 + assert steps[0]["name"] == "my_chat" + assert steps[0]["type"] == "llm" + assert steps[0]["description"] == "Handle a chat message." + + def test_decorator_registers_sync_function(self) -> None: + from agent_control.control_decorators import control + + @control() + def my_process(input: str) -> str: + return input.upper() + + steps = get_registered_steps() + assert len(steps) == 1 + assert steps[0]["name"] == "my_process" + + def test_decorator_registers_with_policy(self) -> None: + from agent_control.control_decorators import control + + @control(policy="my-policy") + async def guarded(msg: str) -> str: + return msg + + steps = get_registered_steps() + assert steps[0]["metadata"] == {"policy": "my-policy"} + + def test_decorator_registers_tool(self) -> None: + """Tool-like functions (with .name attribute) should register as type=tool.""" + from agent_control.control_decorators import control + + def _lookup(query: str) -> str: + """Look up a record.""" + return query + + _lookup.name = "lookup_tool" # type: ignore[attr-defined] + _lookup.tool_name = "lookup_tool" # type: ignore[attr-defined] + control()(_lookup) + + steps = get_registered_steps() + assert len(steps) == 1 + assert steps[0]["type"] == "tool" + assert steps[0]["name"] == "lookup_tool" + + def test_stacked_decorators_deduplicate(self) -> None: + """Stacking @control() twice on the same function deduplicates by name.""" + from agent_control.control_decorators import control + + @control(policy="p1") + @control(policy="p2") + async def stacked(msg: str) -> str: + return msg + + steps = get_registered_steps() + # The inner decorator runs first (p2), then the outer (p1). + # Both use func.__name__ == "stacked", so last-write-wins -> p1. + assert len(steps) == 1 + assert steps[0]["name"] == "stacked" + assert steps[0]["metadata"] == {"policy": "p1"} + + +# =========================================================================== +# init() merge behaviour +# =========================================================================== + + +class TestInitMerge: + """Tests that init() merges auto-discovered steps with explicit steps.""" + + def test_auto_steps_merged_into_init(self) -> None: + """Steps from @control() decorators should be sent to register_agent.""" + from agent_control.control_decorators import control + + @control() + async def auto_tool(query: str) -> str: + """Automatically discovered tool.""" + return query + + # Verify the step is registered + steps = get_registered_steps() + assert any(s["name"] == "auto_tool" for s in steps) + + def test_explicit_steps_take_precedence(self) -> None: + """Explicit steps override auto-discovered steps with the same name.""" + # Register via decorator + register(lambda x: x) # name will be "" + + def my_tool(query: str) -> str: + ... + + register(my_tool) + + explicit_steps: list[dict[str, Any]] = [ + {"type": "tool", "name": "my_tool", "input_schema": {"custom": True}} + ] + + # Simulate what init() does: merge logic + auto_steps = get_registered_steps() + explicit_names = {s["name"] for s in explicit_steps} + merged = list(explicit_steps) + [ + s for s in auto_steps if s["name"] not in explicit_names + ] + + # Explicit "my_tool" should win over the auto-registered one + my_tool_entries = [s for s in merged if s["name"] == "my_tool"] + assert len(my_tool_entries) == 1 + assert my_tool_entries[0]["input_schema"] == {"custom": True} + + # The lambda step should still be present + assert any(s["name"] == "" for s in merged) + + def test_no_auto_steps_leaves_explicit_unchanged(self) -> None: + """When no decorators are used, explicit steps pass through unchanged.""" + # Registry is empty (autouse fixture cleared it) + explicit: list[dict[str, Any]] = [ + {"type": "tool", "name": "manual_tool"} + ] + auto = get_registered_steps() + assert auto == [] + + # Simulate merge + merged = list(explicit) # no auto steps to add + assert merged == explicit From e1f07b3b4f253a505faebd69280fb7726f64a6e4 Mon Sep 17 00:00:00 2001 From: Lev Neiman Date: Thu, 12 Feb 2026 16:13:42 -0800 Subject: [PATCH 2/9] refactor(sdk): isolate step schema derivation and merge by type+name --- sdks/python/src/agent_control/__init__.py | 30 ++- .../src/agent_control/_control_registry.py | 108 +------- .../src/agent_control/_schema_derivation.py | 245 ++++++++++++++++++ sdks/python/tests/test_control_registry.py | 231 ++++------------- sdks/python/tests/test_schema_derivation.py | 165 ++++++++++++ 5 files changed, 493 insertions(+), 286 deletions(-) create mode 100644 sdks/python/src/agent_control/_schema_derivation.py create mode 100644 sdks/python/tests/test_schema_derivation.py diff --git a/sdks/python/src/agent_control/__init__.py b/sdks/python/src/agent_control/__init__.py index 7b2f6f1b..a7933a60 100644 --- a/sdks/python/src/agent_control/__init__.py +++ b/sdks/python/src/agent_control/__init__.py @@ -438,14 +438,32 @@ async def handle(message: str): _api_key = api_key # Merge auto-discovered steps from @control() decorators with explicit steps. - # Explicit steps take precedence (by name) over auto-discovered ones. - from agent_control._control_registry import get_registered_steps - + # Explicit steps take precedence when (type, name) collides. auto_steps = get_registered_steps() if auto_steps: - explicit_names = {s["name"] for s in (steps or [])} - merged = list(steps or []) + [s for s in auto_steps if s["name"] not in explicit_names] - steps = merged + explicit_steps = list(steps or []) + explicit_keys = {(step["type"], step["name"]) for step in explicit_steps} + merged_auto_steps: list[dict[str, Any]] = [] + overridden_keys: list[tuple[str, str]] = [] + + for auto_step in auto_steps: + key = (auto_step["type"], auto_step["name"]) + if key in explicit_keys: + overridden_keys.append(key) + continue + merged_auto_steps.append(auto_step) + + if overridden_keys: + formatted = ", ".join( + f"{step_type}:{step_name}" for step_type, step_name in overridden_keys + ) + logger.warning( + "Skipping %d auto-discovered step(s) overridden by explicit steps: %s", + len(overridden_keys), + formatted, + ) + + steps = explicit_steps + merged_auto_steps logger.debug( "Auto-discovered %d step(s) from @control() decorators (%d after merge)", len(auto_steps), diff --git a/sdks/python/src/agent_control/_control_registry.py b/sdks/python/src/agent_control/_control_registry.py index 6824c7ff..55c43c52 100644 --- a/sdks/python/src/agent_control/_control_registry.py +++ b/sdks/python/src/agent_control/_control_registry.py @@ -11,12 +11,12 @@ from __future__ import annotations -import inspect import logging -import typing from collections.abc import Callable from typing import Any +from ._schema_derivation import derive_schemas + logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- @@ -27,101 +27,6 @@ """Maps step name -> step schema dict. Keyed by name to deduplicate.""" -# --------------------------------------------------------------------------- -# Schema extraction helpers -# --------------------------------------------------------------------------- - -# Mapping from Python primitive types to JSON Schema type strings. -_PRIMITIVE_TYPE_MAP: dict[type, str] = { - str: "string", - int: "integer", - float: "number", - bool: "boolean", -} - - -def _type_to_json_schema(annotation: Any) -> dict[str, Any] | None: - """Convert a Python type annotation to a JSON Schema dict. - - Supports: - - Pydantic models (via ``model_json_schema()``) - - Primitive types (str, int, float, bool) - - Framework-specific ``args_schema`` on the function (checked separately) - - Returns ``None`` for complex or unrecognised types so callers can - gracefully degrade. - """ - if annotation is None or annotation is inspect.Parameter.empty: - return None - - # Pydantic v2 models expose model_json_schema() - if hasattr(annotation, "model_json_schema"): - try: - result: dict[str, Any] = annotation.model_json_schema() - return result - except Exception: - logger.debug("Failed to extract JSON schema from Pydantic model %s", annotation) - return None - - # Primitive types - type_str = _PRIMITIVE_TYPE_MAP.get(annotation) # type: ignore[arg-type] - if type_str is not None: - return {"type": type_str} - - return None - - -def _extract_input_schema(func: Callable[..., Any]) -> dict[str, Any] | None: - """Build a JSON Schema ``object`` from the function's parameter type hints. - - Skips ``self`` and ``cls`` parameters. Returns ``None`` when no useful - schema can be derived (e.g. no type hints at all). - """ - # Framework tools (e.g. LangChain) may expose a Pydantic args_schema - args_schema = getattr(func, "args_schema", None) - if args_schema is not None and hasattr(args_schema, "model_json_schema"): - try: - result: dict[str, Any] = args_schema.model_json_schema() - return result - except Exception: - logger.debug("Failed to extract args_schema from %s", func) - - try: - hints = typing.get_type_hints(func) - except Exception: - # get_type_hints can fail on some decorated / wrapped functions - return None - - sig = inspect.signature(func) - properties: dict[str, Any] = {} - for name, _param in sig.parameters.items(): - if name in ("self", "cls"): - continue - hint = hints.get(name) - if hint is None: - continue - schema = _type_to_json_schema(hint) - if schema is not None: - properties[name] = schema - - if not properties: - return None - return {"type": "object", "properties": properties} - - -def _extract_output_schema(func: Callable[..., Any]) -> dict[str, Any] | None: - """Derive output schema from the function's return type annotation.""" - try: - hints = typing.get_type_hints(func) - except Exception: - return None - - return_hint = hints.get("return") - if return_hint is None: - return None - return _type_to_json_schema(return_hint) - - # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- @@ -150,19 +55,16 @@ def register(func: Callable[..., Any], policy: str | None = None) -> None: if first_line: description = first_line - input_schema = _extract_input_schema(func) - output_schema = _extract_output_schema(func) + schemas = derive_schemas(func) step: dict[str, Any] = { "type": step_type, "name": step_name, + "input_schema": schemas.input_schema, + "output_schema": schemas.output_schema, } if description is not None: step["description"] = description - if input_schema is not None: - step["input_schema"] = input_schema - if output_schema is not None: - step["output_schema"] = output_schema metadata: dict[str, Any] = {} if policy is not None: diff --git a/sdks/python/src/agent_control/_schema_derivation.py b/sdks/python/src/agent_control/_schema_derivation.py new file mode 100644 index 00000000..e197283c --- /dev/null +++ b/sdks/python/src/agent_control/_schema_derivation.py @@ -0,0 +1,245 @@ +"""Schema derivation helpers for agent step registration. + +This module centralizes JSON schema derivation so registry code can focus on +bookkeeping (naming, metadata, and deduplication) rather than inference logic. + +Derivation is best-effort: +- Prefer framework-provided ``args_schema`` for input if available. +- Otherwise infer schemas from function type hints using Pydantic. +- On inference failures, log a warning and return permissive fallback schemas. +""" + +from __future__ import annotations + +import inspect +import logging +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any, get_type_hints + +from pydantic import TypeAdapter, create_model + +logger = logging.getLogger(__name__) + +_INPUT_FALLBACK_SCHEMA: dict[str, Any] = {"type": "object", "additionalProperties": True} +_OUTPUT_FALLBACK_SCHEMA: dict[str, Any] = {} + + +@dataclass(frozen=True) +class DerivedSchemas: + """Container for derived input/output JSON schemas.""" + + input_schema: dict[str, Any] + output_schema: dict[str, Any] + + +def derive_schemas(func: Callable[..., Any]) -> DerivedSchemas: + """Derive input/output JSON schemas for a function. + + Args: + func: Function for schema derivation. + + Returns: + DerivedSchemas containing input and output JSON schemas. + """ + input_schema = _extract_args_schema_override(func) + if input_schema is None: + input_schema = _infer_input_schema(func) + + output_schema = _infer_output_schema(func) + return DerivedSchemas(input_schema=input_schema, output_schema=output_schema) + + +def _extract_args_schema_override(func: Callable[..., Any]) -> dict[str, Any] | None: + """Extract framework-provided input schema from ``func.args_schema``.""" + args_schema = getattr(func, "args_schema", None) + if args_schema is None: + return None + if not hasattr(args_schema, "model_json_schema"): + return None + + try: + schema = args_schema.model_json_schema() + except Exception as exc: # pragma: no cover - exercised via tests with caplog + _warn( + func, + phase="input", + reason="args_schema.model_json_schema() failed; using signature inference", + exc=exc, + fallback=_INPUT_FALLBACK_SCHEMA, + ) + return None + + if not isinstance(schema, dict): + _warn( + func, + phase="input", + reason="args_schema.model_json_schema() returned non-dict; using signature inference", + fallback=_INPUT_FALLBACK_SCHEMA, + ) + return None + + return schema + + +def _infer_input_schema(func: Callable[..., Any]) -> dict[str, Any]: + """Infer input schema from signature + type hints using dynamic Pydantic model.""" + unwrapped = inspect.unwrap(func) + + try: + hints = get_type_hints(unwrapped, include_extras=True) + except Exception as exc: + _warn( + func, + phase="input", + reason="failed to resolve type hints", + exc=exc, + fallback=_INPUT_FALLBACK_SCHEMA, + ) + return _fallback_input_schema() + + try: + signature = inspect.signature(unwrapped) + except Exception as exc: + _warn( + func, + phase="input", + reason="failed to inspect function signature", + exc=exc, + fallback=_INPUT_FALLBACK_SCHEMA, + ) + return _fallback_input_schema() + + fields: dict[str, tuple[Any, Any]] = {} + for name, param in signature.parameters.items(): + if name in {"self", "cls"}: + continue + if param.kind in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}: + continue + + annotation = hints.get(name, Any) + default = ... if param.default is inspect.Parameter.empty else param.default + fields[name] = (annotation, default) + + model_name = _build_model_name(unwrapped, suffix="Input") + try: + model = create_model(model_name, **fields) # type: ignore[call-overload] + schema = model.model_json_schema() + except Exception as exc: + _warn( + func, + phase="input", + reason="failed to infer schema from signature", + exc=exc, + fallback=_INPUT_FALLBACK_SCHEMA, + ) + return _fallback_input_schema() + + if not isinstance(schema, dict): + _warn( + func, + phase="input", + reason="inferred input schema is not a dict", + fallback=_INPUT_FALLBACK_SCHEMA, + ) + return _fallback_input_schema() + + return schema + + +def _infer_output_schema(func: Callable[..., Any]) -> dict[str, Any]: + """Infer output schema from return type annotation using ``TypeAdapter``.""" + unwrapped = inspect.unwrap(func) + + try: + hints = get_type_hints(unwrapped, include_extras=True) + except Exception as exc: + _warn( + func, + phase="output", + reason="failed to resolve type hints", + exc=exc, + fallback=_OUTPUT_FALLBACK_SCHEMA, + ) + return _fallback_output_schema() + + if "return" not in hints: + _warn( + func, + phase="output", + reason="missing return type annotation", + fallback=_OUTPUT_FALLBACK_SCHEMA, + ) + return _fallback_output_schema() + + return_annotation = hints["return"] + try: + schema = TypeAdapter(return_annotation).json_schema() + except Exception as exc: + _warn( + func, + phase="output", + reason="failed to infer output schema from return annotation", + exc=exc, + fallback=_OUTPUT_FALLBACK_SCHEMA, + ) + return _fallback_output_schema() + + if not isinstance(schema, dict): + _warn( + func, + phase="output", + reason="inferred output schema is not a dict", + fallback=_OUTPUT_FALLBACK_SCHEMA, + ) + return _fallback_output_schema() + + return schema + + +def _build_model_name(func: Callable[..., Any], suffix: str) -> str: + """Build a safe dynamic model name from function metadata.""" + raw = f"{func.__module__}_{func.__qualname__}_{suffix}" + safe = "".join(ch if ch.isalnum() else "_" for ch in raw) + return safe or f"Derived_{suffix}" + + +def _fallback_input_schema() -> dict[str, Any]: + """Return a permissive fallback input schema.""" + return dict(_INPUT_FALLBACK_SCHEMA) + + +def _fallback_output_schema() -> dict[str, Any]: + """Return a permissive fallback output schema.""" + return dict(_OUTPUT_FALLBACK_SCHEMA) + + +def _warn( + func: Callable[..., Any], + *, + phase: str, + reason: str, + fallback: dict[str, Any], + exc: Exception | None = None, +) -> None: + """Emit a structured warning for schema fallback paths.""" + function_name = f"{func.__module__}.{func.__qualname__}" + if exc is None: + logger.warning( + "Using fallback %s schema for %s: %s. fallback=%s", + phase, + function_name, + reason, + fallback, + ) + return + + logger.warning( + "Using fallback %s schema for %s: %s (%s: %s). fallback=%s", + phase, + function_name, + reason, + exc.__class__.__name__, + exc, + fallback, + ) diff --git a/sdks/python/tests/test_control_registry.py b/sdks/python/tests/test_control_registry.py index 2b5571c9..ffda3067 100644 --- a/sdks/python/tests/test_control_registry.py +++ b/sdks/python/tests/test_control_registry.py @@ -1,24 +1,12 @@ -"""Tests for the control step registry (_control_registry).""" +"""Tests for control step registry behavior.""" from __future__ import annotations from typing import Any -from unittest.mock import MagicMock import pytest -from agent_control._control_registry import ( - _extract_input_schema, - _extract_output_schema, - _type_to_json_schema, - clear, - get_registered_steps, - register, -) - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- +from agent_control._control_registry import clear, get_registered_steps, register @pytest.fixture(autouse=True) @@ -29,124 +17,14 @@ def _clean_registry() -> None: # noqa: PT004 clear() -# =========================================================================== -# Schema extraction helpers -# =========================================================================== - - -class TestTypeToJsonSchema: - """Tests for _type_to_json_schema.""" - - def test_primitive_str(self) -> None: - assert _type_to_json_schema(str) == {"type": "string"} - - def test_primitive_int(self) -> None: - assert _type_to_json_schema(int) == {"type": "integer"} - - def test_primitive_float(self) -> None: - assert _type_to_json_schema(float) == {"type": "number"} - - def test_primitive_bool(self) -> None: - assert _type_to_json_schema(bool) == {"type": "boolean"} - - def test_none_returns_none(self) -> None: - assert _type_to_json_schema(None) is None - - def test_complex_type_returns_none(self) -> None: - """Complex / unrecognised types should return None.""" - assert _type_to_json_schema(list) is None - assert _type_to_json_schema(dict) is None - - def test_pydantic_model(self) -> None: - """Pydantic v2 models should delegate to model_json_schema().""" - mock_model = MagicMock() - mock_model.model_json_schema.return_value = { - "type": "object", - "properties": {"query": {"type": "string"}}, - } - result = _type_to_json_schema(mock_model) - assert result == {"type": "object", "properties": {"query": {"type": "string"}}} - mock_model.model_json_schema.assert_called_once() - - -class TestExtractInputSchema: - """Tests for _extract_input_schema.""" - - def test_simple_function(self) -> None: - def my_func(query: str, limit: int = 10) -> list: - ... - - schema = _extract_input_schema(my_func) - assert schema is not None - assert schema["type"] == "object" - assert schema["properties"]["query"] == {"type": "string"} - assert schema["properties"]["limit"] == {"type": "integer"} - - def test_no_type_hints(self) -> None: - def my_func(x, y): - ... - - assert _extract_input_schema(my_func) is None - - def test_skips_self_and_cls(self) -> None: - def my_method(self, query: str) -> str: # noqa: ANN001 - ... - - schema = _extract_input_schema(my_method) - assert schema is not None - assert "self" not in schema["properties"] - assert "query" in schema["properties"] - - def test_framework_args_schema(self) -> None: - """If func has .args_schema with model_json_schema(), use that.""" - mock_schema = MagicMock() - mock_schema.model_json_schema.return_value = { - "type": "object", - "properties": {"q": {"type": "string"}}, - } - - def my_func(q: str) -> str: - ... - - my_func.args_schema = mock_schema # type: ignore[attr-defined] - - result = _extract_input_schema(my_func) - assert result == {"type": "object", "properties": {"q": {"type": "string"}}} - mock_schema.model_json_schema.assert_called_once() - - -class TestExtractOutputSchema: - """Tests for _extract_output_schema.""" - - def test_str_return(self) -> None: - def my_func() -> str: - ... - - assert _extract_output_schema(my_func) == {"type": "string"} - - def test_int_return(self) -> None: - def my_func() -> int: - ... - - assert _extract_output_schema(my_func) == {"type": "integer"} - - def test_no_return_annotation(self) -> None: - def my_func(): - ... - - assert _extract_output_schema(my_func) is None - - def test_complex_return_type(self) -> None: - def my_func() -> list[str]: - ... - - # list[str] is not a supported primitive, should return None - assert _extract_output_schema(my_func) is None - - -# =========================================================================== -# register() and get_registered_steps() -# =========================================================================== +def _merge_steps_by_key( + explicit_steps: list[dict[str, Any]], auto_steps: list[dict[str, Any]] +) -> list[dict[str, Any]]: + """Mirror init() merge behavior (explicit wins by type+name).""" + explicit_keys = {(s["type"], s["name"]) for s in explicit_steps} + return list(explicit_steps) + [ + s for s in auto_steps if (s["type"], s["name"]) not in explicit_keys + ] class TestRegister: @@ -159,16 +37,18 @@ def chat(message: str) -> str: register(chat) steps = get_registered_steps() + assert len(steps) == 1 step = steps[0] assert step["name"] == "chat" assert step["type"] == "llm" assert step["description"] == "Chat with the agent." - assert step["input_schema"] is not None - assert step["output_schema"] == {"type": "string"} + assert step["input_schema"]["type"] == "object" + assert step["output_schema"]["type"] == "string" def test_register_tool_function(self) -> None: """Functions with .name or .tool_name should be registered as tools.""" + def search_db(query: str, limit: int = 10) -> str: """Search the database.""" ... @@ -178,6 +58,7 @@ def search_db(query: str, limit: int = 10) -> str: register(search_db) steps = get_registered_steps() + assert len(steps) == 1 assert steps[0]["type"] == "tool" assert steps[0]["name"] == "search_db" @@ -188,6 +69,7 @@ def my_func(x: str) -> str: register(my_func, policy="safety-policy") steps = get_registered_steps() + assert steps[0]["metadata"] == {"policy": "safety-policy"} def test_register_no_policy_no_metadata(self) -> None: @@ -196,10 +78,12 @@ def my_func(x: str) -> str: register(my_func) steps = get_registered_steps() + assert "metadata" not in steps[0] def test_deduplicate_by_name(self) -> None: """Registering two functions with the same name should keep the last one.""" + def chat(message: str) -> str: """First version.""" ... @@ -213,6 +97,7 @@ def chat_v2(message: str) -> int: # noqa: ARG001 register(chat) register(chat_v2) steps = get_registered_steps() + assert len(steps) == 1 assert steps[0]["description"] == "Second version." @@ -222,19 +107,23 @@ def my_func(x: str) -> str: register(my_func) steps = get_registered_steps() + assert "description" not in steps[0] def test_no_type_hints(self) -> None: - """Functions with no type hints should still register with None schemas.""" + """Untyped functions still register with permissive schemas.""" + def my_func(x, y): ... register(my_func) steps = get_registered_steps() + assert len(steps) == 1 assert steps[0]["name"] == "my_func" - assert "input_schema" not in steps[0] - assert "output_schema" not in steps[0] + assert steps[0]["input_schema"]["type"] == "object" + assert set(steps[0]["input_schema"]["properties"]) == {"x", "y"} + assert steps[0]["output_schema"] == {} class TestClear: @@ -255,11 +144,6 @@ def f2(x: str) -> str: assert len(get_registered_steps()) == 0 -# =========================================================================== -# Decorator integration -# =========================================================================== - - class TestDecoratorRegistration: """Tests that @control() decorator registers functions in the registry.""" @@ -272,6 +156,7 @@ async def my_chat(message: str) -> str: return message steps = get_registered_steps() + assert len(steps) == 1 assert steps[0]["name"] == "my_chat" assert steps[0]["type"] == "llm" @@ -285,6 +170,7 @@ def my_process(input: str) -> str: return input.upper() steps = get_registered_steps() + assert len(steps) == 1 assert steps[0]["name"] == "my_process" @@ -296,6 +182,7 @@ async def guarded(msg: str) -> str: return msg steps = get_registered_steps() + assert steps[0]["metadata"] == {"policy": "my-policy"} def test_decorator_registers_tool(self) -> None: @@ -311,6 +198,7 @@ def _lookup(query: str) -> str: control()(_lookup) steps = get_registered_steps() + assert len(steps) == 1 assert steps[0]["type"] == "tool" assert steps[0]["name"] == "lookup_tool" @@ -325,23 +213,17 @@ async def stacked(msg: str) -> str: return msg steps = get_registered_steps() - # The inner decorator runs first (p2), then the outer (p1). - # Both use func.__name__ == "stacked", so last-write-wins -> p1. + assert len(steps) == 1 assert steps[0]["name"] == "stacked" assert steps[0]["metadata"] == {"policy": "p1"} -# =========================================================================== -# init() merge behaviour -# =========================================================================== - - class TestInitMerge: - """Tests that init() merges auto-discovered steps with explicit steps.""" + """Tests the explicit+auto merge behavior used by init().""" def test_auto_steps_merged_into_init(self) -> None: - """Steps from @control() decorators should be sent to register_agent.""" + """Steps from @control() decorators should be available for init merge.""" from agent_control.control_decorators import control @control() @@ -349,48 +231,43 @@ async def auto_tool(query: str) -> str: """Automatically discovered tool.""" return query - # Verify the step is registered steps = get_registered_steps() + assert any(s["name"] == "auto_tool" for s in steps) def test_explicit_steps_take_precedence(self) -> None: - """Explicit steps override auto-discovered steps with the same name.""" - # Register via decorator + """Explicit steps override auto-discovered steps on exact type+name key.""" register(lambda x: x) # name will be "" - def my_tool(query: str) -> str: + def my_step(query: str) -> str: ... - register(my_tool) + register(my_step) explicit_steps: list[dict[str, Any]] = [ - {"type": "tool", "name": "my_tool", "input_schema": {"custom": True}} + {"type": "llm", "name": "my_step", "input_schema": {"custom": True}} ] - # Simulate what init() does: merge logic auto_steps = get_registered_steps() - explicit_names = {s["name"] for s in explicit_steps} - merged = list(explicit_steps) + [ - s for s in auto_steps if s["name"] not in explicit_names - ] + merged = _merge_steps_by_key(explicit_steps, auto_steps) - # Explicit "my_tool" should win over the auto-registered one - my_tool_entries = [s for s in merged if s["name"] == "my_tool"] - assert len(my_tool_entries) == 1 - assert my_tool_entries[0]["input_schema"] == {"custom": True} - - # The lambda step should still be present + my_step_entries = [s for s in merged if (s["type"], s["name"]) == ("llm", "my_step")] + assert len(my_step_entries) == 1 + assert my_step_entries[0]["input_schema"] == {"custom": True} assert any(s["name"] == "" for s in merged) def test_no_auto_steps_leaves_explicit_unchanged(self) -> None: - """When no decorators are used, explicit steps pass through unchanged.""" - # Registry is empty (autouse fixture cleared it) - explicit: list[dict[str, Any]] = [ - {"type": "tool", "name": "manual_tool"} - ] - auto = get_registered_steps() - assert auto == [] + explicit: list[dict[str, Any]] = [{"type": "tool", "name": "manual_tool"}] + + merged = _merge_steps_by_key(explicit, get_registered_steps()) - # Simulate merge - merged = list(explicit) # no auto steps to add assert merged == explicit + + def test_merge_keeps_same_name_steps_with_different_types(self) -> None: + auto_steps: list[dict[str, Any]] = [{"type": "llm", "name": "shared"}] + explicit_steps: list[dict[str, Any]] = [{"type": "tool", "name": "shared"}] + + merged = _merge_steps_by_key(explicit_steps, auto_steps) + + merged_keys = {(s["type"], s["name"]) for s in merged} + assert merged_keys == {("tool", "shared"), ("llm", "shared")} diff --git a/sdks/python/tests/test_schema_derivation.py b/sdks/python/tests/test_schema_derivation.py new file mode 100644 index 00000000..0130911a --- /dev/null +++ b/sdks/python/tests/test_schema_derivation.py @@ -0,0 +1,165 @@ +"""Tests for isolated schema derivation logic.""" + +from __future__ import annotations + +import logging +from typing import Any +from unittest.mock import MagicMock + +import pytest +from pydantic import BaseModel + +from agent_control._schema_derivation import derive_schemas + + +class _InputModel(BaseModel): + query: str + limit: int = 5 + + +class _OutputModel(BaseModel): + answer: str + + +class TestInputInference: + """Input schema derivation tests.""" + + def test_required_and_default_parameters(self) -> None: + def my_func(query: str, limit: int = 10) -> str: + ... + + schemas = derive_schemas(my_func) + + assert schemas.input_schema["type"] == "object" + assert set(schemas.input_schema.get("required", [])) == {"query"} + assert schemas.input_schema["properties"]["query"]["type"] == "string" + assert schemas.input_schema["properties"]["limit"]["type"] == "integer" + + def test_optional_union_parameter_is_preserved(self) -> None: + def my_func(query: str, conversation_id: str | None = None) -> str: + ... + + schemas = derive_schemas(my_func) + + conversation_schema = schemas.input_schema["properties"]["conversation_id"] + has_null = "anyOf" in conversation_schema or ( + isinstance(conversation_schema.get("type"), list) + and "null" in conversation_schema["type"] + ) + assert has_null + + def test_collection_types_are_represented(self) -> None: + def my_func(tags: list[str], metadata: dict[str, int]) -> list[str]: + ... + + schemas = derive_schemas(my_func) + + assert schemas.input_schema["properties"]["tags"]["type"] == "array" + assert schemas.input_schema["properties"]["metadata"]["type"] == "object" + assert schemas.output_schema["type"] == "array" + + def test_untyped_parameters_fall_back_to_any_fields(self) -> None: + def my_func(x, y): + ... + + schemas = derive_schemas(my_func) + + assert schemas.input_schema["type"] == "object" + assert set(schemas.input_schema["properties"]) == {"x", "y"} + + +class TestOutputInference: + """Output schema derivation tests.""" + + def test_primitive_output(self) -> None: + def my_func() -> str: + ... + + schemas = derive_schemas(my_func) + + assert schemas.output_schema["type"] == "string" + + def test_pydantic_input_and_output(self) -> None: + def my_func(payload: _InputModel) -> _OutputModel: + ... + + schemas = derive_schemas(my_func) + + payload_schema = schemas.input_schema["properties"]["payload"] + assert ("type" in payload_schema and payload_schema["type"] == "object") or ( + "$ref" in payload_schema + ) + assert schemas.output_schema["type"] == "object" + + +class TestArgsSchemaOverride: + """args_schema precedence and fallback behavior.""" + + def test_args_schema_precedence(self) -> None: + mock_schema = MagicMock() + mock_schema.model_json_schema.return_value = { + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + } + + def my_func(query: str) -> str: + ... + + my_func.args_schema = mock_schema # type: ignore[attr-defined] + + schemas = derive_schemas(my_func) + + assert schemas.input_schema == { + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + } + mock_schema.model_json_schema.assert_called_once() + + def test_args_schema_failure_falls_back_to_signature_inference( + self, caplog: pytest.LogCaptureFixture + ) -> None: + class BrokenArgsSchema: + def model_json_schema(self) -> dict[str, Any]: + raise RuntimeError("broken args schema") + + def my_func(query: str) -> str: + ... + + my_func.args_schema = BrokenArgsSchema() # type: ignore[attr-defined] + + with caplog.at_level(logging.WARNING): + schemas = derive_schemas(my_func) + + assert schemas.input_schema["properties"]["query"]["type"] == "string" + assert "args_schema.model_json_schema() failed" in caplog.text + + +class TestFallbackWarnings: + """Warnings and fallback behavior for unresolved/incomplete typing.""" + + def test_missing_return_annotation_warns_and_falls_back( + self, caplog: pytest.LogCaptureFixture + ) -> None: + def my_func(query: str): + ... + + with caplog.at_level(logging.WARNING): + schemas = derive_schemas(my_func) + + assert schemas.output_schema == {} + assert "missing return type annotation" in caplog.text + + def test_unresolved_type_hints_warn_and_fall_back( + self, caplog: pytest.LogCaptureFixture + ) -> None: + def my_func(query: "DoesNotExist") -> str: + ... + + with caplog.at_level(logging.WARNING): + schemas = derive_schemas(my_func) + + assert schemas.input_schema == {"type": "object", "additionalProperties": True} + assert schemas.output_schema == {} + assert "failed to resolve type hints" in caplog.text From b9c6fb1a79de648c19e29a8f3b5aa571b557f9d8 Mon Sep 17 00:00:00 2001 From: Lev Neiman Date: Thu, 12 Feb 2026 16:17:43 -0800 Subject: [PATCH 3/9] test: add Given/When/Then comments to SDK schema tests --- sdks/python/tests/test_control_registry.py | 56 +++++++++++++++++++++ sdks/python/tests/test_schema_derivation.py | 30 +++++++++++ 2 files changed, 86 insertions(+) diff --git a/sdks/python/tests/test_control_registry.py b/sdks/python/tests/test_control_registry.py index ffda3067..a6e89650 100644 --- a/sdks/python/tests/test_control_registry.py +++ b/sdks/python/tests/test_control_registry.py @@ -31,13 +31,16 @@ class TestRegister: """Tests for register() and get_registered_steps().""" def test_register_simple_function(self) -> None: + # Given a basic typed control function with a docstring. def chat(message: str) -> str: """Chat with the agent.""" ... + # When the function is registered and steps are retrieved. register(chat) steps = get_registered_steps() + # Then one llm step is created with the expected name, description, and schemas. assert len(steps) == 1 step = steps[0] assert step["name"] == "chat" @@ -49,6 +52,7 @@ def chat(message: str) -> str: def test_register_tool_function(self) -> None: """Functions with .name or .tool_name should be registered as tools.""" + # Given a function that advertises tool-style metadata via attributes. def search_db(query: str, limit: int = 10) -> str: """Search the database.""" ... @@ -56,34 +60,43 @@ def search_db(query: str, limit: int = 10) -> str: search_db.name = "search_db" # type: ignore[attr-defined] search_db.tool_name = "search_db" # type: ignore[attr-defined] + # When the function is registered. register(search_db) steps = get_registered_steps() + # Then the registered step is classified as a tool using the declared tool name. assert len(steps) == 1 assert steps[0]["type"] == "tool" assert steps[0]["name"] == "search_db" def test_register_with_policy(self) -> None: + # Given a typed function and an explicit policy value at registration time. def my_func(x: str) -> str: ... + # When the function is registered with that policy. register(my_func, policy="safety-policy") steps = get_registered_steps() + # Then metadata contains the policy exactly once. assert steps[0]["metadata"] == {"policy": "safety-policy"} def test_register_no_policy_no_metadata(self) -> None: + # Given a typed function registered without policy metadata. def my_func(x: str) -> str: ... + # When registration completes. register(my_func) steps = get_registered_steps() + # Then no metadata field is emitted for the step. assert "metadata" not in steps[0] def test_deduplicate_by_name(self) -> None: """Registering two functions with the same name should keep the last one.""" + # Given two functions that resolve to the same step name. def chat(message: str) -> str: """First version.""" ... @@ -94,31 +107,39 @@ def chat_v2(message: str) -> int: # noqa: ARG001 chat_v2.__name__ = "chat" # simulate same name + # When both are registered in order. register(chat) register(chat_v2) steps = get_registered_steps() + # Then only the second registration remains for that name. assert len(steps) == 1 assert steps[0]["description"] == "Second version." def test_no_docstring(self) -> None: + # Given a typed function without a docstring. def my_func(x: str) -> str: ... + # When the function is registered. register(my_func) steps = get_registered_steps() + # Then description is omitted from the resulting step. assert "description" not in steps[0] def test_no_type_hints(self) -> None: """Untyped functions still register with permissive schemas.""" + # Given an untyped function. def my_func(x, y): ... + # When the function is registered. register(my_func) steps = get_registered_steps() + # Then registration succeeds with a permissive input schema and fallback output schema. assert len(steps) == 1 assert steps[0]["name"] == "my_func" assert steps[0]["input_schema"]["type"] == "object" @@ -130,6 +151,7 @@ class TestClear: """Tests for clear().""" def test_clear_empties_registry(self) -> None: + # Given multiple registered functions in the global registry. def f1(x: str) -> str: ... @@ -140,7 +162,10 @@ def f2(x: str) -> str: register(f2) assert len(get_registered_steps()) == 2 + # When clear() is called. clear() + + # Then the registry becomes empty. assert len(get_registered_steps()) == 0 @@ -148,6 +173,7 @@ class TestDecoratorRegistration: """Tests that @control() decorator registers functions in the registry.""" def test_decorator_registers_async_function(self) -> None: + # Given an async function decorated with @control(). from agent_control.control_decorators import control @control() @@ -155,38 +181,48 @@ async def my_chat(message: str) -> str: """Handle a chat message.""" return message + # When registered steps are queried after decoration. steps = get_registered_steps() + # Then the decorated async function appears as a single llm step with its docstring. assert len(steps) == 1 assert steps[0]["name"] == "my_chat" assert steps[0]["type"] == "llm" assert steps[0]["description"] == "Handle a chat message." def test_decorator_registers_sync_function(self) -> None: + # Given a sync function decorated with @control(). from agent_control.control_decorators import control @control() def my_process(input: str) -> str: return input.upper() + # When registered steps are retrieved. steps = get_registered_steps() + # Then the function is registered with the expected name. assert len(steps) == 1 assert steps[0]["name"] == "my_process" def test_decorator_registers_with_policy(self) -> None: + # Given a decorated function that includes a policy in the decorator arguments. from agent_control.control_decorators import control @control(policy="my-policy") async def guarded(msg: str) -> str: return msg + # When registered steps are fetched. steps = get_registered_steps() + # Then the policy is persisted in step metadata. assert steps[0]["metadata"] == {"policy": "my-policy"} def test_decorator_registers_tool(self) -> None: """Tool-like functions (with .name attribute) should register as type=tool.""" + + # Given a function decorated via @control() and marked with tool metadata attributes. from agent_control.control_decorators import control def _lookup(query: str) -> str: @@ -195,8 +231,11 @@ def _lookup(query: str) -> str: _lookup.name = "lookup_tool" # type: ignore[attr-defined] _lookup.tool_name = "lookup_tool" # type: ignore[attr-defined] + + # When the decorator is applied to the tool-like function. control()(_lookup) + # Then the registry contains a tool step with the tool's declared name. steps = get_registered_steps() assert len(steps) == 1 @@ -205,6 +244,8 @@ def _lookup(query: str) -> str: def test_stacked_decorators_deduplicate(self) -> None: """Stacking @control() twice on the same function deduplicates by name.""" + + # Given a function wrapped by two @control() decorators with different policies. from agent_control.control_decorators import control @control(policy="p1") @@ -212,8 +253,10 @@ def test_stacked_decorators_deduplicate(self) -> None: async def stacked(msg: str) -> str: return msg + # When the stacked decorators register the same function name. steps = get_registered_steps() + # Then deduplication keeps one step, preserving the outermost decorator metadata. assert len(steps) == 1 assert steps[0]["name"] == "stacked" assert steps[0]["metadata"] == {"policy": "p1"} @@ -224,6 +267,8 @@ class TestInitMerge: def test_auto_steps_merged_into_init(self) -> None: """Steps from @control() decorators should be available for init merge.""" + + # Given an auto-registered control via decorator discovery. from agent_control.control_decorators import control @control() @@ -231,12 +276,15 @@ async def auto_tool(query: str) -> str: """Automatically discovered tool.""" return query + # When registered steps are collected for init-time merge. steps = get_registered_steps() + # Then the auto-registered step is present by name. assert any(s["name"] == "auto_tool" for s in steps) def test_explicit_steps_take_precedence(self) -> None: """Explicit steps override auto-discovered steps on exact type+name key.""" + # Given auto-discovered steps and an explicit step that shares the same (type, name) key. register(lambda x: x) # name will be "" def my_step(query: str) -> str: @@ -248,26 +296,34 @@ def my_step(query: str) -> str: {"type": "llm", "name": "my_step", "input_schema": {"custom": True}} ] + # When explicit and auto steps are merged with explicit-first precedence. auto_steps = get_registered_steps() merged = _merge_steps_by_key(explicit_steps, auto_steps) + # Then the explicit version wins for the duplicate key while unrelated auto steps are retained. my_step_entries = [s for s in merged if (s["type"], s["name"]) == ("llm", "my_step")] assert len(my_step_entries) == 1 assert my_step_entries[0]["input_schema"] == {"custom": True} assert any(s["name"] == "" for s in merged) def test_no_auto_steps_leaves_explicit_unchanged(self) -> None: + # Given only explicit steps and no auto-registered steps. explicit: list[dict[str, Any]] = [{"type": "tool", "name": "manual_tool"}] + # When merge logic runs against an empty auto-step set. merged = _merge_steps_by_key(explicit, get_registered_steps()) + # Then the output matches the explicit list unchanged. assert merged == explicit def test_merge_keeps_same_name_steps_with_different_types(self) -> None: + # Given explicit and auto steps that share a name but have different step types. auto_steps: list[dict[str, Any]] = [{"type": "llm", "name": "shared"}] explicit_steps: list[dict[str, Any]] = [{"type": "tool", "name": "shared"}] + # When merge logic deduplicates by (type, name) rather than by name only. merged = _merge_steps_by_key(explicit_steps, auto_steps) + # Then both entries are preserved because their type dimension differs. merged_keys = {(s["type"], s["name"]) for s in merged} assert merged_keys == {("tool", "shared"), ("llm", "shared")} diff --git a/sdks/python/tests/test_schema_derivation.py b/sdks/python/tests/test_schema_derivation.py index 0130911a..a9c5519c 100644 --- a/sdks/python/tests/test_schema_derivation.py +++ b/sdks/python/tests/test_schema_derivation.py @@ -25,22 +25,28 @@ class TestInputInference: """Input schema derivation tests.""" def test_required_and_default_parameters(self) -> None: + # Given a callable with one required parameter and one parameter with a default. def my_func(query: str, limit: int = 10) -> str: ... + # When JSON schemas are derived from the function signature. schemas = derive_schemas(my_func) + # Then the input schema marks only the required field as required and preserves types. assert schemas.input_schema["type"] == "object" assert set(schemas.input_schema.get("required", [])) == {"query"} assert schemas.input_schema["properties"]["query"]["type"] == "string" assert schemas.input_schema["properties"]["limit"]["type"] == "integer" def test_optional_union_parameter_is_preserved(self) -> None: + # Given a callable with an optional union parameter (`str | None`). def my_func(query: str, conversation_id: str | None = None) -> str: ... + # When JSON schemas are derived from that signature. schemas = derive_schemas(my_func) + # Then the derived property schema includes a nullable representation. conversation_schema = schemas.input_schema["properties"]["conversation_id"] has_null = "anyOf" in conversation_schema or ( isinstance(conversation_schema.get("type"), list) @@ -49,21 +55,27 @@ def my_func(query: str, conversation_id: str | None = None) -> str: assert has_null def test_collection_types_are_represented(self) -> None: + # Given collection-typed inputs and a collection-typed return annotation. def my_func(tags: list[str], metadata: dict[str, int]) -> list[str]: ... + # When JSON schemas are derived. schemas = derive_schemas(my_func) + # Then list/dict/return collection types are preserved in the emitted schemas. assert schemas.input_schema["properties"]["tags"]["type"] == "array" assert schemas.input_schema["properties"]["metadata"]["type"] == "object" assert schemas.output_schema["type"] == "array" def test_untyped_parameters_fall_back_to_any_fields(self) -> None: + # Given a callable with untyped parameters. def my_func(x, y): ... + # When JSON schemas are derived. schemas = derive_schemas(my_func) + # Then schema derivation still exposes both fields under a permissive object schema. assert schemas.input_schema["type"] == "object" assert set(schemas.input_schema["properties"]) == {"x", "y"} @@ -72,19 +84,25 @@ class TestOutputInference: """Output schema derivation tests.""" def test_primitive_output(self) -> None: + # Given a callable with a primitive return annotation. def my_func() -> str: ... + # When JSON schemas are derived. schemas = derive_schemas(my_func) + # Then the output schema is emitted as a string type. assert schemas.output_schema["type"] == "string" def test_pydantic_input_and_output(self) -> None: + # Given a callable that uses Pydantic models for input and output. def my_func(payload: _InputModel) -> _OutputModel: ... + # When JSON schemas are derived. schemas = derive_schemas(my_func) + # Then the payload is represented as an object/$ref and the output resolves to object. payload_schema = schemas.input_schema["properties"]["payload"] assert ("type" in payload_schema and payload_schema["type"] == "object") or ( "$ref" in payload_schema @@ -96,6 +114,7 @@ class TestArgsSchemaOverride: """args_schema precedence and fallback behavior.""" def test_args_schema_precedence(self) -> None: + # Given a callable that provides a working args_schema override. mock_schema = MagicMock() mock_schema.model_json_schema.return_value = { "type": "object", @@ -108,8 +127,10 @@ def my_func(query: str) -> str: my_func.args_schema = mock_schema # type: ignore[attr-defined] + # When schemas are derived for the callable. schemas = derive_schemas(my_func) + # Then args_schema is used as the authoritative input schema source. assert schemas.input_schema == { "type": "object", "properties": {"q": {"type": "string"}}, @@ -120,6 +141,7 @@ def my_func(query: str) -> str: def test_args_schema_failure_falls_back_to_signature_inference( self, caplog: pytest.LogCaptureFixture ) -> None: + # Given a callable whose args_schema override raises at schema generation time. class BrokenArgsSchema: def model_json_schema(self) -> dict[str, Any]: raise RuntimeError("broken args schema") @@ -129,9 +151,11 @@ def my_func(query: str) -> str: my_func.args_schema = BrokenArgsSchema() # type: ignore[attr-defined] + # When schema derivation runs with warning capture enabled. with caplog.at_level(logging.WARNING): schemas = derive_schemas(my_func) + # Then derivation falls back to signature inference and emits a warning. assert schemas.input_schema["properties"]["query"]["type"] == "string" assert "args_schema.model_json_schema() failed" in caplog.text @@ -142,24 +166,30 @@ class TestFallbackWarnings: def test_missing_return_annotation_warns_and_falls_back( self, caplog: pytest.LogCaptureFixture ) -> None: + # Given a callable without an explicit return type annotation. def my_func(query: str): ... + # When schemas are derived while warnings are captured. with caplog.at_level(logging.WARNING): schemas = derive_schemas(my_func) + # Then output falls back to a permissive schema and a warning is emitted. assert schemas.output_schema == {} assert "missing return type annotation" in caplog.text def test_unresolved_type_hints_warn_and_fall_back( self, caplog: pytest.LogCaptureFixture ) -> None: + # Given a callable that references an unresolved forward type hint. def my_func(query: "DoesNotExist") -> str: ... + # When schema derivation attempts to resolve type hints. with caplog.at_level(logging.WARNING): schemas = derive_schemas(my_func) + # Then both schemas fall back to permissive defaults and a warning is emitted. assert schemas.input_schema == {"type": "object", "additionalProperties": True} assert schemas.output_schema == {} assert "failed to resolve type hints" in caplog.text From 1e1e634d4e732a8f6ddf77d8bc2674f71c42cd21 Mon Sep 17 00:00:00 2001 From: Lev Neiman Date: Thu, 12 Feb 2026 16:26:41 -0800 Subject: [PATCH 4/9] refactor: strengthen step schema derivation and merge flow --- sdks/python/src/agent_control/__init__.py | 42 ++++----- .../src/agent_control/_control_registry.py | 92 +++++++++++++++++-- .../src/agent_control/_schema_derivation.py | 75 ++++++++------- sdks/python/tests/test_control_registry.py | 48 +++++++--- 4 files changed, 179 insertions(+), 78 deletions(-) diff --git a/sdks/python/src/agent_control/__init__.py b/sdks/python/src/agent_control/__init__.py index a7933a60..71e9a519 100644 --- a/sdks/python/src/agent_control/__init__.py +++ b/sdks/python/src/agent_control/__init__.py @@ -66,8 +66,14 @@ async def process(input: str) -> str: ) from . import agents, controls, evaluation, evaluators, policies -from ._control_registry import clear as clear_step_registry -from ._control_registry import get_registered_steps +from ._control_registry import ( + StepSchemaDict, + get_registered_steps, + merge_explicit_and_auto_steps, +) +from ._control_registry import ( + clear as clear_step_registry, +) # Import client and operations modules from .client import AgentControlClient @@ -342,7 +348,7 @@ def init( server_url: str | None = None, api_key: str | None = None, controls_file: str | None = None, - steps: list[dict[str, Any]] | None = None, + steps: list[StepSchemaDict] | None = None, observability_enabled: bool | None = None, log_config: dict[str, Any] | None = None, **kwargs: object @@ -440,34 +446,24 @@ async def handle(message: str): # Merge auto-discovered steps from @control() decorators with explicit steps. # Explicit steps take precedence when (type, name) collides. auto_steps = get_registered_steps() + merge_result = merge_explicit_and_auto_steps(steps, auto_steps) + registration_steps: list[dict[str, Any]] = [dict(step) for step in merge_result.steps] + if auto_steps: - explicit_steps = list(steps or []) - explicit_keys = {(step["type"], step["name"]) for step in explicit_steps} - merged_auto_steps: list[dict[str, Any]] = [] - overridden_keys: list[tuple[str, str]] = [] - - for auto_step in auto_steps: - key = (auto_step["type"], auto_step["name"]) - if key in explicit_keys: - overridden_keys.append(key) - continue - merged_auto_steps.append(auto_step) - - if overridden_keys: + if merge_result.overridden_keys: formatted = ", ".join( - f"{step_type}:{step_name}" for step_type, step_name in overridden_keys + f"{step_type}:{step_name}" for step_type, step_name in merge_result.overridden_keys ) logger.warning( "Skipping %d auto-discovered step(s) overridden by explicit steps: %s", - len(overridden_keys), + len(merge_result.overridden_keys), formatted, ) - steps = explicit_steps + merged_auto_steps logger.debug( "Auto-discovered %d step(s) from @control() decorators (%d after merge)", len(auto_steps), - len(steps), + len(registration_steps), ) # Register with server and fetch controls @@ -491,7 +487,7 @@ async def register() -> list[dict[str, Any]] | None: response = await agents.register_agent( client, _current_agent, - steps=steps or [] + steps=registration_steps ) created = response.get('created', False) controls: list[dict[str, Any]] = response.get('controls', []) @@ -501,8 +497,8 @@ async def register() -> list[dict[str, Any]] | None: else: logger.info("Agent updated: %s (ID: %s)", agent_name, _agent_uuid) - if steps: - logger.debug("Registered %d step(s)", len(steps)) + if registration_steps: + logger.debug("Registered %d step(s)", len(registration_steps)) return controls except httpx.HTTPStatusError: diff --git a/sdks/python/src/agent_control/_control_registry.py b/sdks/python/src/agent_control/_control_registry.py index 55c43c52..04e23b74 100644 --- a/sdks/python/src/agent_control/_control_registry.py +++ b/sdks/python/src/agent_control/_control_registry.py @@ -13,18 +13,45 @@ import logging from collections.abc import Callable -from typing import Any +from dataclasses import dataclass +from typing import Any, NotRequired, TypedDict from ._schema_derivation import derive_schemas logger = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# Step schema types +# --------------------------------------------------------------------------- + +type StepKey = tuple[str, str] + + +class StepSchemaDict(TypedDict): + """Runtime representation of a step schema payload.""" + + type: str + name: str + description: NotRequired[str] + input_schema: NotRequired[dict[str, Any] | None] + output_schema: NotRequired[dict[str, Any] | None] + metadata: NotRequired[dict[str, Any]] + + +@dataclass(frozen=True) +class StepMergeResult: + """Result of merging explicit and auto-discovered step schemas.""" + + steps: list[StepSchemaDict] + overridden_keys: list[StepKey] + + # --------------------------------------------------------------------------- # Internal registry # --------------------------------------------------------------------------- -_registered_steps: dict[str, dict[str, Any]] = {} -"""Maps step name -> step schema dict. Keyed by name to deduplicate.""" +_registered_steps: dict[StepKey, StepSchemaDict] = {} +"""Maps ``(type, name)`` -> step schema dict. Keyed by type+name to deduplicate.""" # --------------------------------------------------------------------------- @@ -57,7 +84,7 @@ def register(func: Callable[..., Any], policy: str | None = None) -> None: schemas = derive_schemas(func) - step: dict[str, Any] = { + step: StepSchemaDict = { "type": step_type, "name": step_name, "input_schema": schemas.input_schema, @@ -72,14 +99,20 @@ def register(func: Callable[..., Any], policy: str | None = None) -> None: if metadata: step["metadata"] = metadata - # Store (last-write-wins for duplicate names) - if step_name in _registered_steps: - logger.debug("Overwriting previously registered step '%s'", step_name) - _registered_steps[step_name] = step + key = _step_key(step_type, step_name) + + # Store (last-write-wins for duplicate type+name pairs). + if key in _registered_steps: + logger.debug( + "Overwriting previously registered step '%s' (type=%s)", + step_name, + step_type, + ) + _registered_steps[key] = step logger.debug("Registered step schema: %s (type=%s)", step_name, step_type) -def get_registered_steps() -> list[dict[str, Any]]: +def get_registered_steps() -> list[StepSchemaDict]: """Return all registered step schemas as a list of dicts. The returned dicts conform to the ``StepSchema`` model format expected @@ -88,6 +121,47 @@ def get_registered_steps() -> list[dict[str, Any]]: return list(_registered_steps.values()) +def merge_explicit_and_auto_steps( + explicit_steps: list[StepSchemaDict] | None, + auto_steps: list[StepSchemaDict], +) -> StepMergeResult: + """Merge explicit and auto-discovered steps. + + Explicit steps win on exact ``(type, name)`` collisions. + + Args: + explicit_steps: Steps provided explicitly to ``init(steps=...)``. + auto_steps: Steps auto-discovered from ``@control()`` registration. + + Returns: + Merge result containing final merged steps and the overridden auto keys. + """ + explicit = list(explicit_steps or []) + if not auto_steps: + return StepMergeResult(steps=explicit, overridden_keys=[]) + + explicit_keys = {_step_key(step["type"], step["name"]) for step in explicit} + merged_auto_steps: list[StepSchemaDict] = [] + overridden_keys: list[StepKey] = [] + + for auto_step in auto_steps: + key = _step_key(auto_step["type"], auto_step["name"]) + if key in explicit_keys: + overridden_keys.append(key) + continue + merged_auto_steps.append(auto_step) + + return StepMergeResult( + steps=explicit + merged_auto_steps, + overridden_keys=overridden_keys, + ) + + def clear() -> None: """Clear all registered steps. Useful for testing.""" _registered_steps.clear() + + +def _step_key(step_type: str, step_name: str) -> StepKey: + """Create a canonical deduplication key for step schemas.""" + return (step_type, step_name) diff --git a/sdks/python/src/agent_control/_schema_derivation.py b/sdks/python/src/agent_control/_schema_derivation.py index e197283c..f3d97a0b 100644 --- a/sdks/python/src/agent_control/_schema_derivation.py +++ b/sdks/python/src/agent_control/_schema_derivation.py @@ -42,11 +42,45 @@ def derive_schemas(func: Callable[..., Any]) -> DerivedSchemas: Returns: DerivedSchemas containing input and output JSON schemas. """ - input_schema = _extract_args_schema_override(func) + input_schema_override = _extract_args_schema_override(func) + unwrapped = inspect.unwrap(func) + + try: + hints = get_type_hints(unwrapped, include_extras=True) + except Exception as exc: + if input_schema_override is not None: + _warn( + func, + phase="output", + reason="failed to resolve type hints", + exc=exc, + fallback=_OUTPUT_FALLBACK_SCHEMA, + ) + return DerivedSchemas( + input_schema=input_schema_override, + output_schema=_fallback_output_schema(), + ) + + _warn( + func, + phase="input/output", + reason="failed to resolve type hints", + exc=exc, + fallback={ + "input": _INPUT_FALLBACK_SCHEMA, + "output": _OUTPUT_FALLBACK_SCHEMA, + }, + ) + return DerivedSchemas( + input_schema=_fallback_input_schema(), + output_schema=_fallback_output_schema(), + ) + + input_schema = input_schema_override if input_schema is None: - input_schema = _infer_input_schema(func) + input_schema = _infer_input_schema(func, unwrapped, hints) - output_schema = _infer_output_schema(func) + output_schema = _infer_output_schema(func, hints) return DerivedSchemas(input_schema=input_schema, output_schema=output_schema) @@ -82,21 +116,12 @@ def _extract_args_schema_override(func: Callable[..., Any]) -> dict[str, Any] | return schema -def _infer_input_schema(func: Callable[..., Any]) -> dict[str, Any]: +def _infer_input_schema( + func: Callable[..., Any], + unwrapped: Callable[..., Any], + hints: dict[str, Any], +) -> dict[str, Any]: """Infer input schema from signature + type hints using dynamic Pydantic model.""" - unwrapped = inspect.unwrap(func) - - try: - hints = get_type_hints(unwrapped, include_extras=True) - except Exception as exc: - _warn( - func, - phase="input", - reason="failed to resolve type hints", - exc=exc, - fallback=_INPUT_FALLBACK_SCHEMA, - ) - return _fallback_input_schema() try: signature = inspect.signature(unwrapped) @@ -147,22 +172,8 @@ def _infer_input_schema(func: Callable[..., Any]) -> dict[str, Any]: return schema -def _infer_output_schema(func: Callable[..., Any]) -> dict[str, Any]: +def _infer_output_schema(func: Callable[..., Any], hints: dict[str, Any]) -> dict[str, Any]: """Infer output schema from return type annotation using ``TypeAdapter``.""" - unwrapped = inspect.unwrap(func) - - try: - hints = get_type_hints(unwrapped, include_extras=True) - except Exception as exc: - _warn( - func, - phase="output", - reason="failed to resolve type hints", - exc=exc, - fallback=_OUTPUT_FALLBACK_SCHEMA, - ) - return _fallback_output_schema() - if "return" not in hints: _warn( func, diff --git a/sdks/python/tests/test_control_registry.py b/sdks/python/tests/test_control_registry.py index a6e89650..292d70fe 100644 --- a/sdks/python/tests/test_control_registry.py +++ b/sdks/python/tests/test_control_registry.py @@ -6,7 +6,12 @@ import pytest -from agent_control._control_registry import clear, get_registered_steps, register +from agent_control._control_registry import ( + clear, + get_registered_steps, + merge_explicit_and_auto_steps, + register, +) @pytest.fixture(autouse=True) @@ -17,16 +22,6 @@ def _clean_registry() -> None: # noqa: PT004 clear() -def _merge_steps_by_key( - explicit_steps: list[dict[str, Any]], auto_steps: list[dict[str, Any]] -) -> list[dict[str, Any]]: - """Mirror init() merge behavior (explicit wins by type+name).""" - explicit_keys = {(s["type"], s["name"]) for s in explicit_steps} - return list(explicit_steps) + [ - s for s in auto_steps if (s["type"], s["name"]) not in explicit_keys - ] - - class TestRegister: """Tests for register() and get_registered_steps().""" @@ -116,6 +111,29 @@ def chat_v2(message: str) -> int: # noqa: ARG001 assert len(steps) == 1 assert steps[0]["description"] == "Second version." + def test_same_name_different_types_are_kept_distinct(self) -> None: + # Given two controls that share a name but represent different step types. + def llm_step(message: str) -> str: + """LLM variant.""" + ... + + def tool_step(query: str) -> str: + """Tool variant.""" + ... + + llm_step.__name__ = "shared_name" + tool_step.name = "shared_name" # type: ignore[attr-defined] + tool_step.tool_name = "shared_name" # type: ignore[attr-defined] + + # When both controls are registered. + register(llm_step) + register(tool_step) + steps = get_registered_steps() + + # Then both entries are retained because deduplication key is (type, name). + keys = {(step["type"], step["name"]) for step in steps} + assert keys == {("llm", "shared_name"), ("tool", "shared_name")} + def test_no_docstring(self) -> None: # Given a typed function without a docstring. def my_func(x: str) -> str: @@ -298,20 +316,22 @@ def my_step(query: str) -> str: # When explicit and auto steps are merged with explicit-first precedence. auto_steps = get_registered_steps() - merged = _merge_steps_by_key(explicit_steps, auto_steps) + merge_result = merge_explicit_and_auto_steps(explicit_steps, auto_steps) + merged = merge_result.steps # Then the explicit version wins for the duplicate key while unrelated auto steps are retained. my_step_entries = [s for s in merged if (s["type"], s["name"]) == ("llm", "my_step")] assert len(my_step_entries) == 1 assert my_step_entries[0]["input_schema"] == {"custom": True} assert any(s["name"] == "" for s in merged) + assert merge_result.overridden_keys == [("llm", "my_step")] def test_no_auto_steps_leaves_explicit_unchanged(self) -> None: # Given only explicit steps and no auto-registered steps. explicit: list[dict[str, Any]] = [{"type": "tool", "name": "manual_tool"}] # When merge logic runs against an empty auto-step set. - merged = _merge_steps_by_key(explicit, get_registered_steps()) + merged = merge_explicit_and_auto_steps(explicit, get_registered_steps()).steps # Then the output matches the explicit list unchanged. assert merged == explicit @@ -322,7 +342,7 @@ def test_merge_keeps_same_name_steps_with_different_types(self) -> None: explicit_steps: list[dict[str, Any]] = [{"type": "tool", "name": "shared"}] # When merge logic deduplicates by (type, name) rather than by name only. - merged = _merge_steps_by_key(explicit_steps, auto_steps) + merged = merge_explicit_and_auto_steps(explicit_steps, auto_steps).steps # Then both entries are preserved because their type dimension differs. merged_keys = {(s["type"], s["name"]) for s in merged} From c6bafc20f53f0aa8998708969e590fc3f424fc9d Mon Sep 17 00:00:00 2001 From: Lev Neiman Date: Thu, 12 Feb 2026 16:54:47 -0800 Subject: [PATCH 5/9] test: add golden snapshots and fallback branch coverage --- sdks/python/tests/test_control_registry.py | 37 +++ sdks/python/tests/test_schema_derivation.py | 338 ++++++++++++++++++++ 2 files changed, 375 insertions(+) diff --git a/sdks/python/tests/test_control_registry.py b/sdks/python/tests/test_control_registry.py index 292d70fe..152cd5dd 100644 --- a/sdks/python/tests/test_control_registry.py +++ b/sdks/python/tests/test_control_registry.py @@ -347,3 +347,40 @@ def test_merge_keeps_same_name_steps_with_different_types(self) -> None: # Then both entries are preserved because their type dimension differs. merged_keys = {(s["type"], s["name"]) for s in merged} assert merged_keys == {("tool", "shared"), ("llm", "shared")} + + +class TestStepSchemaContract: + """Contract tests that merged registry payloads satisfy StepSchema model.""" + + def test_merged_steps_validate_against_stepschema_model(self) -> None: + # Given one auto-discovered step and one explicit override for the same (type, name). + from agent_control_models import StepSchema + + def auto_llm_step(query: str) -> str: + """Auto-discovered llm step.""" + ... + + register(auto_llm_step) + + explicit_steps: list[dict[str, Any]] = [ + { + "type": "llm", + "name": "auto_llm_step", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, + "output_schema": {"type": "string"}, + }, + { + "type": "tool", + "name": "manual_tool", + "input_schema": {"type": "object"}, + "output_schema": {"type": "object"}, + }, + ] + + # When explicit and auto steps are merged by the registry helper. + merge_result = merge_explicit_and_auto_steps(explicit_steps, get_registered_steps()) + + # Then every merged payload validates against the shared StepSchema contract. + validated_steps = [StepSchema.model_validate(step) for step in merge_result.steps] + assert len(validated_steps) == len(merge_result.steps) + assert merge_result.overridden_keys == [("llm", "auto_llm_step")] diff --git a/sdks/python/tests/test_schema_derivation.py b/sdks/python/tests/test_schema_derivation.py index a9c5519c..4f87331e 100644 --- a/sdks/python/tests/test_schema_derivation.py +++ b/sdks/python/tests/test_schema_derivation.py @@ -9,6 +9,7 @@ import pytest from pydantic import BaseModel +import agent_control._schema_derivation as schema_derivation from agent_control._schema_derivation import derive_schemas @@ -21,6 +22,26 @@ class _OutputModel(BaseModel): answer: str +def golden_primitive_defaults(query: str, limit: int = 10) -> str: + """Golden-case primitive/default function.""" + raise NotImplementedError + + +def golden_optional_union(conversation_id: str | None = None) -> None: + """Golden-case optional union function.""" + raise NotImplementedError + + +def golden_collections(tags: list[str], metadata: dict[str, int]) -> list[str]: + """Golden-case collection function.""" + raise NotImplementedError + + +def golden_nested_models(payload: _InputModel) -> _OutputModel: + """Golden-case nested Pydantic model function.""" + raise NotImplementedError + + class TestInputInference: """Input schema derivation tests.""" @@ -193,3 +214,320 @@ def my_func(query: "DoesNotExist") -> str: assert schemas.input_schema == {"type": "object", "additionalProperties": True} assert schemas.output_schema == {} assert "failed to resolve type hints" in caplog.text + + +class TestAdditionalFallbackBranches: + """Additional branch coverage for defensive schema fallback paths.""" + + def test_args_schema_without_model_json_schema_is_ignored(self) -> None: + # Given a callable with an args_schema object that is missing model_json_schema(). + class _MissingArgsSchemaMethod: + pass + + def my_func(query: str) -> str: + ... + + my_func.args_schema = _MissingArgsSchemaMethod() # type: ignore[attr-defined] + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then derivation ignores args_schema and falls back to signature inference. + assert schemas.input_schema["properties"]["query"]["type"] == "string" + assert schemas.output_schema == {"type": "string"} + + def test_args_schema_non_dict_warns_and_falls_back( + self, caplog: pytest.LogCaptureFixture + ) -> None: + # Given a callable whose args_schema returns a non-dict payload. + class _NonDictArgsSchema: + @staticmethod + def model_json_schema() -> list[str]: + return ["not-a-dict"] + + def my_func(query: str) -> str: + ... + + my_func.args_schema = _NonDictArgsSchema() # type: ignore[attr-defined] + + # When schemas are derived with warning capture. + with caplog.at_level(logging.WARNING): + schemas = derive_schemas(my_func) + + # Then derivation warns and uses signature-based input inference. + assert schemas.input_schema["properties"]["query"]["type"] == "string" + assert "returned non-dict" in caplog.text + + def test_args_schema_override_kept_when_type_hints_resolution_fails( + self, caplog: pytest.LogCaptureFixture + ) -> None: + # Given a callable with args_schema override but unresolved type hints. + class _GoodArgsSchema: + @staticmethod + def model_json_schema() -> dict[str, Any]: + return {"type": "object", "properties": {"q": {"type": "string"}}} + + def my_func(query: "DoesNotExist") -> str: + ... + + my_func.args_schema = _GoodArgsSchema() # type: ignore[attr-defined] + + # When derivation attempts to resolve type hints. + with caplog.at_level(logging.WARNING): + schemas = derive_schemas(my_func) + + # Then input stays overridden while output falls back with warning. + assert schemas.input_schema == {"type": "object", "properties": {"q": {"type": "string"}}} + assert schemas.output_schema == {} + assert "failed to resolve type hints" in caplog.text + + def test_signature_inspection_failure_uses_input_fallback( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + # Given inspect.signature raises while deriving input fields. + def _broken_signature(_func: Any) -> Any: + raise RuntimeError("signature failed") + + def my_func(query: str) -> str: + ... + + monkeypatch.setattr(schema_derivation.inspect, "signature", _broken_signature) + + # When schemas are derived. + with caplog.at_level(logging.WARNING): + schemas = derive_schemas(my_func) + + # Then input falls back to permissive schema and warning is emitted. + assert schemas.input_schema == {"type": "object", "additionalProperties": True} + assert "failed to inspect function signature" in caplog.text + + def test_self_cls_varargs_and_kwargs_are_excluded_from_input_schema(self) -> None: + # Given a callable containing self/cls placeholders and variadic parameters. + def my_func(self, cls, query: str, *args: Any, **kwargs: Any) -> str: # noqa: ANN001 + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then only concrete named fields remain in the inferred input schema. + assert set(schemas.input_schema["properties"]) == {"query"} + assert schemas.input_schema.get("required") == ["query"] + + def test_input_model_creation_failure_uses_fallback( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + # Given dynamic input model creation raises unexpectedly. + def _broken_create_model(*_args: Any, **_kwargs: Any) -> Any: + raise RuntimeError("create_model failed") + + def my_func(query: str) -> str: + ... + + monkeypatch.setattr(schema_derivation, "create_model", _broken_create_model) + + # When schemas are derived. + with caplog.at_level(logging.WARNING): + schemas = derive_schemas(my_func) + + # Then derivation emits warning and returns permissive input fallback. + assert schemas.input_schema == {"type": "object", "additionalProperties": True} + assert "failed to infer schema from signature" in caplog.text + + def test_input_non_dict_schema_from_model_uses_fallback( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + # Given create_model returns a model whose model_json_schema() is non-dict. + class _NonDictModel: + @staticmethod + def model_json_schema() -> list[str]: + return ["not-a-dict"] + + def _fake_create_model(*_args: Any, **_kwargs: Any) -> _NonDictModel: + return _NonDictModel() + + def my_func(query: str) -> str: + ... + + monkeypatch.setattr(schema_derivation, "create_model", _fake_create_model) + + # When schemas are derived. + with caplog.at_level(logging.WARNING): + schemas = derive_schemas(my_func) + + # Then input inference falls back after warning about non-dict schema output. + assert schemas.input_schema == {"type": "object", "additionalProperties": True} + assert "inferred input schema is not a dict" in caplog.text + + def test_output_type_adapter_failure_uses_fallback( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + # Given output TypeAdapter json_schema() raises unexpectedly. + class _BrokenTypeAdapter: + def __init__(self, *_args: Any, **_kwargs: Any) -> None: + pass + + @staticmethod + def json_schema() -> dict[str, Any]: + raise RuntimeError("adapter failed") + + def my_func(query: str) -> str: + ... + + monkeypatch.setattr(schema_derivation, "TypeAdapter", _BrokenTypeAdapter) + + # When schemas are derived. + with caplog.at_level(logging.WARNING): + schemas = derive_schemas(my_func) + + # Then output derivation falls back and emits a warning. + assert schemas.output_schema == {} + assert "failed to infer output schema from return annotation" in caplog.text + + def test_output_non_dict_schema_from_type_adapter_uses_fallback( + self, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture + ) -> None: + # Given output TypeAdapter returns a non-dict JSON schema. + class _NonDictTypeAdapter: + def __init__(self, *_args: Any, **_kwargs: Any) -> None: + pass + + @staticmethod + def json_schema() -> list[str]: + return ["not-a-dict"] + + def my_func(query: str) -> str: + ... + + monkeypatch.setattr(schema_derivation, "TypeAdapter", _NonDictTypeAdapter) + + # When schemas are derived. + with caplog.at_level(logging.WARNING): + schemas = derive_schemas(my_func) + + # Then output derivation falls back after warning about non-dict schema output. + assert schemas.output_schema == {} + assert "inferred output schema is not a dict" in caplog.text + + +class TestGoldenSchemas: + """Golden-style schema snapshots for representative function signatures.""" + + def test_golden_primitive_defaults_snapshot(self) -> None: + # Given a primitive+default function with stable module-level identity. + # When schemas are derived. + schemas = derive_schemas(golden_primitive_defaults) + + # Then the full input/output schema snapshots match exactly. + assert schemas.input_schema == { + "properties": { + "query": {"title": "Query", "type": "string"}, + "limit": {"default": 10, "title": "Limit", "type": "integer"}, + }, + "required": ["query"], + "title": "tests_test_schema_derivation_golden_primitive_defaults_Input", + "type": "object", + } + assert schemas.output_schema == {"type": "string"} + + def test_golden_optional_union_snapshot(self) -> None: + # Given an optional-union input and explicit `None` return annotation. + # When schemas are derived. + schemas = derive_schemas(golden_optional_union) + + # Then the full schema snapshots preserve nullable input and null output types. + assert schemas.input_schema == { + "properties": { + "conversation_id": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "default": None, + "title": "Conversation Id", + } + }, + "title": "tests_test_schema_derivation_golden_optional_union_Input", + "type": "object", + } + assert schemas.output_schema == {"type": "null"} + + def test_golden_collection_snapshot(self) -> None: + # Given list/dict input collections and a list return annotation. + # When schemas are derived. + schemas = derive_schemas(golden_collections) + + # Then the full schema snapshots preserve array/object structures exactly. + assert schemas.input_schema == { + "properties": { + "tags": {"items": {"type": "string"}, "title": "Tags", "type": "array"}, + "metadata": { + "additionalProperties": {"type": "integer"}, + "title": "Metadata", + "type": "object", + }, + }, + "required": ["tags", "metadata"], + "title": "tests_test_schema_derivation_golden_collections_Input", + "type": "object", + } + assert schemas.output_schema == {"items": {"type": "string"}, "type": "array"} + + def test_golden_nested_pydantic_model_snapshot(self) -> None: + # Given nested Pydantic input/output models. + # When schemas are derived. + schemas = derive_schemas(golden_nested_models) + + # Then `$defs` and `$ref` are preserved in the exact schema snapshot. + assert schemas.input_schema == { + "$defs": { + "_InputModel": { + "properties": { + "query": {"title": "Query", "type": "string"}, + "limit": {"default": 5, "title": "Limit", "type": "integer"}, + }, + "required": ["query"], + "title": "_InputModel", + "type": "object", + } + }, + "properties": {"payload": {"$ref": "#/$defs/_InputModel"}}, + "required": ["payload"], + "title": "tests_test_schema_derivation_golden_nested_models_Input", + "type": "object", + } + assert schemas.output_schema == { + "properties": {"answer": {"title": "Answer", "type": "string"}}, + "required": ["answer"], + "title": "_OutputModel", + "type": "object", + } + + def test_golden_args_schema_override_snapshot(self) -> None: + # Given a callable with an explicit framework-style args_schema override. + class _GoldenArgsSchema: + @staticmethod + def model_json_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": { + "q": {"type": "string"}, + "limit": {"type": "integer", "default": 5}, + }, + "required": ["q"], + } + + def golden_args_schema_override(query: str) -> str: + raise NotImplementedError + + golden_args_schema_override.args_schema = _GoldenArgsSchema() # type: ignore[attr-defined] + + # When schemas are derived. + schemas = derive_schemas(golden_args_schema_override) + + # Then the full input snapshot is sourced from args_schema and output remains inferred. + assert schemas.input_schema == { + "type": "object", + "properties": { + "q": {"type": "string"}, + "limit": {"type": "integer", "default": 5}, + }, + "required": ["q"], + } + assert schemas.output_schema == {"type": "string"} From 2c4d7d41679e923db8e8fa6c76221b8dc0d9d2c1 Mon Sep 17 00:00:00 2001 From: Lev Neiman Date: Thu, 12 Feb 2026 17:08:31 -0800 Subject: [PATCH 6/9] test: cover init step merge wiring and deferred schema derivation --- .../src/agent_control/_control_registry.py | 62 ++++++++++----- sdks/python/tests/test_control_registry.py | 30 +++++++- sdks/python/tests/test_init_step_merge.py | 76 +++++++++++++++++++ 3 files changed, 146 insertions(+), 22 deletions(-) create mode 100644 sdks/python/tests/test_init_step_merge.py diff --git a/sdks/python/src/agent_control/_control_registry.py b/sdks/python/src/agent_control/_control_registry.py index 04e23b74..5b37c47c 100644 --- a/sdks/python/src/agent_control/_control_registry.py +++ b/sdks/python/src/agent_control/_control_registry.py @@ -46,12 +46,23 @@ class StepMergeResult: overridden_keys: list[StepKey] +@dataclass(frozen=True) +class _RegisteredControl: + """Internal metadata stored at decorator-registration time.""" + + func: Callable[..., Any] + step_type: str + step_name: str + description: str | None + metadata: dict[str, Any] + + # --------------------------------------------------------------------------- # Internal registry # --------------------------------------------------------------------------- -_registered_steps: dict[StepKey, StepSchemaDict] = {} -"""Maps ``(type, name)`` -> step schema dict. Keyed by type+name to deduplicate.""" +_registered_steps: dict[StepKey, _RegisteredControl] = {} +"""Maps ``(type, name)`` -> registration metadata. Keyed by type+name to deduplicate.""" # --------------------------------------------------------------------------- @@ -62,9 +73,12 @@ class StepMergeResult: def register(func: Callable[..., Any], policy: str | None = None) -> None: """Register a decorated function's step schema in the registry. - Extracts name, type (tool vs llm), description, and input/output schemas - from the function and stores them for later retrieval via - ``get_registered_steps()``. + Extracts step metadata from the function and stores it for later retrieval + via ``get_registered_steps()``. + + Input/output schema derivation is intentionally deferred to retrieval time + so forward references are more likely to be resolvable by the time + ``init()`` runs. Args: func: The original (unwrapped) function being decorated. @@ -82,24 +96,18 @@ def register(func: Callable[..., Any], policy: str | None = None) -> None: if first_line: description = first_line - schemas = derive_schemas(func) - - step: StepSchemaDict = { - "type": step_type, - "name": step_name, - "input_schema": schemas.input_schema, - "output_schema": schemas.output_schema, - } - if description is not None: - step["description"] = description - metadata: dict[str, Any] = {} if policy is not None: metadata["policy"] = policy - if metadata: - step["metadata"] = metadata key = _step_key(step_type, step_name) + registered = _RegisteredControl( + func=func, + step_type=step_type, + step_name=step_name, + description=description, + metadata=metadata, + ) # Store (last-write-wins for duplicate type+name pairs). if key in _registered_steps: @@ -108,7 +116,7 @@ def register(func: Callable[..., Any], policy: str | None = None) -> None: step_name, step_type, ) - _registered_steps[key] = step + _registered_steps[key] = registered logger.debug("Registered step schema: %s (type=%s)", step_name, step_type) @@ -118,7 +126,21 @@ def get_registered_steps() -> list[StepSchemaDict]: The returned dicts conform to the ``StepSchema`` model format expected by ``init(steps=...)``. """ - return list(_registered_steps.values()) + steps: list[StepSchemaDict] = [] + for registered in _registered_steps.values(): + schemas = derive_schemas(registered.func) + step: StepSchemaDict = { + "type": registered.step_type, + "name": registered.step_name, + "input_schema": schemas.input_schema, + "output_schema": schemas.output_schema, + } + if registered.description is not None: + step["description"] = registered.description + if registered.metadata: + step["metadata"] = dict(registered.metadata) + steps.append(step) + return steps def merge_explicit_and_auto_steps( diff --git a/sdks/python/tests/test_control_registry.py b/sdks/python/tests/test_control_registry.py index 152cd5dd..b5fae211 100644 --- a/sdks/python/tests/test_control_registry.py +++ b/sdks/python/tests/test_control_registry.py @@ -2,9 +2,12 @@ from __future__ import annotations +from collections.abc import Generator +import logging from typing import Any import pytest +from pydantic import BaseModel from agent_control._control_registry import ( clear, @@ -15,10 +18,10 @@ @pytest.fixture(autouse=True) -def _clean_registry() -> None: # noqa: PT004 +def _clean_registry() -> Generator[None, None, None]: """Ensure each test starts with an empty registry.""" clear() - yield # type: ignore[misc] + yield clear() @@ -164,6 +167,29 @@ def my_func(x, y): assert set(steps[0]["input_schema"]["properties"]) == {"x", "y"} assert steps[0]["output_schema"] == {} + def test_forward_reference_can_resolve_at_retrieval_time( + self, caplog: pytest.LogCaptureFixture + ) -> None: + # Given a function registered before its forward-referenced model is available in globals. + def my_func(payload: "LaterModel") -> str: + ... + + register(my_func) + + class LaterModel(BaseModel): + value: str + + my_func.__globals__["LaterModel"] = LaterModel + + # When registered steps are materialized (schema derivation occurs at retrieval time). + with caplog.at_level(logging.WARNING): + steps = get_registered_steps() + + # Then the forward reference resolves successfully without fallback warnings. + payload_schema = steps[0]["input_schema"]["properties"]["payload"] + assert ("$ref" in payload_schema) or (payload_schema.get("type") == "object") + assert "failed to resolve type hints" not in caplog.text + class TestClear: """Tests for clear().""" diff --git a/sdks/python/tests/test_init_step_merge.py b/sdks/python/tests/test_init_step_merge.py new file mode 100644 index 00000000..b21b11a1 --- /dev/null +++ b/sdks/python/tests/test_init_step_merge.py @@ -0,0 +1,76 @@ +"""Tests for init() step merge wiring into register_agent.""" + +from __future__ import annotations + +from collections.abc import Generator +import logging +from unittest.mock import AsyncMock, patch +from uuid import uuid4 + +import pytest + +import agent_control +from agent_control._control_registry import clear, register + + +@pytest.fixture(autouse=True) +def _clean_registry() -> Generator[None, None, None]: + """Ensure each test starts with an empty step registry.""" + clear() + yield + clear() + + +def test_init_passes_merged_steps_to_register_agent( + caplog: pytest.LogCaptureFixture, +) -> None: + # Given one auto-discovered step and explicit steps including a conflicting override. + def auto_llm(query: str) -> str: + """Auto-discovered step.""" + ... + + register(auto_llm) + explicit_steps = [ + { + "type": "llm", + "name": "auto_llm", + "input_schema": {"type": "object", "properties": {"query": {"type": "string"}}}, + "output_schema": {"type": "string"}, + "description": "Explicit override for auto_llm.", + }, + { + "type": "tool", + "name": "manual_tool", + "input_schema": {"type": "object", "properties": {"q": {"type": "string"}}}, + "output_schema": {"type": "string"}, + }, + ] + + register_agent_mock = AsyncMock(return_value={"created": True, "controls": []}) + health_check_mock = AsyncMock(return_value={"status": "healthy"}) + + # When init() performs registration with patched network-facing calls. + with patch( + "agent_control.__init__.AgentControlClient.health_check", + new=health_check_mock, + ), patch( + "agent_control.__init__.agents.register_agent", + new=register_agent_mock, + ): + with caplog.at_level(logging.WARNING): + agent_control.init( + agent_name="Init Merge Agent", + agent_id=str(uuid4()), + steps=explicit_steps, + ) + + # Then register_agent() receives merged steps with explicit precedence on conflicts. + assert register_agent_mock.await_count == 1 + assert register_agent_mock.await_args is not None + merged_steps = register_agent_mock.await_args.kwargs["steps"] + + llm_entries = [s for s in merged_steps if (s["type"], s["name"]) == ("llm", "auto_llm")] + assert len(llm_entries) == 1 + assert llm_entries[0]["description"] == "Explicit override for auto_llm." + assert any((s["type"], s["name"]) == ("tool", "manual_tool") for s in merged_steps) + assert "Skipping 1 auto-discovered step(s) overridden by explicit steps" in caplog.text From b45c0ed7000c7eabe79e8d7e3074b6dce7411c6f Mon Sep 17 00:00:00 2001 From: Lev Neiman Date: Thu, 12 Feb 2026 17:14:09 -0800 Subject: [PATCH 7/9] feat(examples): add langgraph auto-schema control demo --- examples/README.md | 17 +- examples/langchain/README.md | 17 ++ .../langchain/langgraph_auto_schema_agent.py | 278 ++++++++++++++++++ examples/langchain/pyproject.toml | 3 +- 4 files changed, 302 insertions(+), 13 deletions(-) create mode 100644 examples/langchain/langgraph_auto_schema_agent.py diff --git a/examples/README.md b/examples/README.md index 6fa2fe6f..1149325a 100644 --- a/examples/README.md +++ b/examples/README.md @@ -57,22 +57,17 @@ Programmatic control setup using SDK models: uv run python examples/demo_setup_controls.py ``` -### 🤖 LangGraph Integration (`langgraph/my_agent/`) +### 🤖 LangGraph Integration (`langchain/`) -LangGraph agent with built-in safety checks: +LangGraph examples are available in `examples/langchain`: ```bash -cd examples/langgraph/my_agent -pip install -e . -cp env.example .env -python cli.py +cd examples/langchain +uv run langgraph_auto_schema_agent.py ``` -**Files:** -- `agent.py` - LangGraph agent with safety check node -- `simple_example.py` - Simplified protection engine usage -- `decorator_example.py` - Visual demonstration of data extraction -- `protect_engine.py` - Local YAML-based protection engine +This specific example demonstrates auto-derived step schemas from `@control()` +decorators, so no explicit `steps=...` list is required in `agent_control.init(...)`. ### 🛡️ Luna-2 Demo (`luna2_demo.py`) diff --git a/examples/langchain/README.md b/examples/langchain/README.md index d9169459..9a32f49f 100644 --- a/examples/langchain/README.md +++ b/examples/langchain/README.md @@ -55,6 +55,23 @@ cd examples/langchain uv run sql_agent_protection.py ``` +## Auto-Derived Step Schema Example (LangGraph) + +This repository now also includes a LangGraph example that demonstrates +automatic step schema derivation from `@control()`-decorated functions. + +Key behavior: +- `agent_control.init(...)` is called **without** explicit `steps=...` +- Tool step schemas are auto-discovered from decorated functions +- Input/output JSON schema is inferred from Python type hints (via Pydantic) + +Run: + +```bash +cd examples/langchain +uv run langgraph_auto_schema_agent.py +``` + ### Local vs Remote Control Execution **Remote (server-side) controls**: diff --git a/examples/langchain/langgraph_auto_schema_agent.py b/examples/langchain/langgraph_auto_schema_agent.py new file mode 100644 index 00000000..7a6e1c3e --- /dev/null +++ b/examples/langchain/langgraph_auto_schema_agent.py @@ -0,0 +1,278 @@ +"""LangGraph agent that relies on auto-derived step schemas from @control tools. + +This example demonstrates the SDK flow we want: +1. Define tool-like functions with Python type hints. +2. Decorate them with ``@control()``. +3. Call ``agent_control.init(...)`` without explicit ``steps=...``. +4. Let the SDK auto-discover decorated functions and derive JSON Schemas. + +Run: + cd examples/langchain + uv run langgraph_auto_schema_agent.py + +Prerequisite: + Start the Agent Control server (`cd server && make run`) so @control() + evaluations can execute successfully at runtime. +""" + +from __future__ import annotations + +import asyncio +import json +import os +import re +from typing import Annotated, Literal, TypedDict +from uuid import UUID + +import agent_control +from agent_control import ControlViolationError, control, get_registered_steps +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage +from langchain_core.tools import tool +from langgraph.graph import END, START, StateGraph +from langgraph.graph.message import add_messages +from langgraph.prebuilt import ToolNode +from pydantic import BaseModel, Field + +AGENT_ID = UUID("736dc6fa-6f6d-4464-a655-c3fe2f5d2e6e") +AGENT_NAME = "LangGraph Auto Schema Demo" +AGENT_DESCRIPTION = "LangGraph tool routing with @control auto step schema derivation" + + +class AgentState(TypedDict): + """LangGraph state object.""" + + messages: Annotated[list[BaseMessage], add_messages] + + +class OrderStatus(BaseModel): + """Structured result for order status lookups.""" + + order_id: str = Field(description="External order identifier") + status: Literal["processing", "shipped", "delivered"] + estimated_delivery_days: int | None = Field( + default=None, + ge=0, + description="Days until delivery; null when already delivered", + ) + history: list[str] = Field(default_factory=list) + + +class RefundDecision(BaseModel): + """Structured result for refund checks.""" + + order_id: str + reason: Literal["damaged", "late", "cancelled"] + approved: bool + approved_amount: float = Field(ge=0) + + +class OrderStatusPayload(TypedDict): + """Tool payload returned to LangGraph.""" + + order_id: str + status: Literal["processing", "shipped", "delivered"] + estimated_delivery_days: int | None + history: list[str] + + +class RefundDecisionPayload(TypedDict): + """Tool payload returned to LangGraph.""" + + order_id: str + reason: Literal["damaged", "late", "cancelled"] + approved: bool + approved_amount: float + + +def _parse_order_id(user_text: str) -> str: + """Extract a stable order identifier from user text.""" + match = re.search(r"(?:order\s*)?(\d{4,8})", user_text) + if match: + return f"ORD-{match.group(1)}" + return "ORD-1001" + + +async def _lookup_order_status(order_id: str, include_history: bool = False) -> OrderStatus: + """Fetch fulfillment status for an order.""" + base = OrderStatus(order_id=order_id, status="shipped", estimated_delivery_days=2) + if include_history: + base.history = [ + "Label created", + "Picked up by carrier", + "Arrived at regional hub", + ] + return base + + +setattr(_lookup_order_status, "name", "lookup_order_status") +setattr(_lookup_order_status, "tool_name", "lookup_order_status") +_lookup_order_status_checked = control()(_lookup_order_status) + + +async def _issue_refund( + order_id: str, + reason: Literal["damaged", "late", "cancelled"], + requested_amount: float | None = None, +) -> RefundDecision: + """Evaluate refund eligibility for an order.""" + approved_amount = requested_amount if requested_amount is not None else 25.0 + is_approved = approved_amount <= 100.0 + return RefundDecision( + order_id=order_id, + reason=reason, + approved=is_approved, + approved_amount=approved_amount if is_approved else 0.0, + ) + + +setattr(_issue_refund, "name", "issue_refund") +setattr(_issue_refund, "tool_name", "issue_refund") +_issue_refund_checked = control()(_issue_refund) + + +@tool("lookup_order_status") +async def lookup_order_status(order_id: str, include_history: bool = False) -> OrderStatusPayload: + """LangGraph tool wrapper for order status.""" + result = await _lookup_order_status_checked( + order_id=order_id, + include_history=include_history, + ) + return { + "order_id": result.order_id, + "status": result.status, + "estimated_delivery_days": result.estimated_delivery_days, + "history": result.history, + } + + +@tool("issue_refund") +async def issue_refund( + order_id: str, + reason: Literal["damaged", "late", "cancelled"], + requested_amount: float | None = None, +) -> RefundDecisionPayload: + """LangGraph tool wrapper for refund decisions.""" + result = await _issue_refund_checked( + order_id=order_id, + reason=reason, + requested_amount=requested_amount, + ) + return { + "order_id": result.order_id, + "reason": result.reason, + "approved": result.approved, + "approved_amount": result.approved_amount, + } + + +def _build_graph(): + """Build a simple deterministic LangGraph flow with ToolNode.""" + tool_node = ToolNode([lookup_order_status, issue_refund]) + + def planner(state: AgentState) -> dict[str, list[AIMessage]]: + last_message = state["messages"][-1] + user_text = str(last_message.content) + lower = user_text.lower() + order_id = _parse_order_id(user_text) + + if "refund" in lower: + reason: Literal["damaged", "late", "cancelled"] = "late" + if "damaged" in lower: + reason = "damaged" + elif "cancel" in lower: + reason = "cancelled" + + requested_amount: float | None = 49.0 if "49" in lower else None + tool_call = { + "name": "issue_refund", + "args": { + "order_id": order_id, + "reason": reason, + "requested_amount": requested_amount, + }, + "id": "call-refund-1", + "type": "tool_call", + } + else: + include_history = "history" in lower or "timeline" in lower + tool_call = { + "name": "lookup_order_status", + "args": { + "order_id": order_id, + "include_history": include_history, + }, + "id": "call-status-1", + "type": "tool_call", + } + + return {"messages": [AIMessage(content="", tool_calls=[tool_call])]} # type: ignore[arg-type] + + def finalize(state: AgentState) -> dict[str, list[AIMessage]]: + tool_message = next( + message for message in reversed(state["messages"]) if isinstance(message, ToolMessage) + ) + return { + "messages": [ + AIMessage(content=f"Tool `{tool_message.name}` returned: {tool_message.content}") + ] + } + + graph = StateGraph(AgentState) + graph.add_node("planner", planner) + graph.add_node("tools", tool_node) + graph.add_node("finalize", finalize) + + graph.add_edge(START, "planner") + graph.add_edge("planner", "tools") + graph.add_edge("tools", "finalize") + graph.add_edge("finalize", END) + return graph.compile() + + +def _print_auto_derived_steps() -> None: + """Show the step schemas auto-derived from @control-decorated functions.""" + print("\nAuto-derived step schemas from @control():") + for step in get_registered_steps(): + print("-" * 80) + print(json.dumps(step, indent=2, sort_keys=True)) + + +async def main() -> None: + """Run the demo end-to-end.""" + print("Initializing Agent Control (no explicit steps passed)...") + agent_control.init( + agent_name=AGENT_NAME, + agent_id=AGENT_ID, + agent_description=AGENT_DESCRIPTION, + server_url=os.getenv("AGENT_CONTROL_URL"), + ) + + _print_auto_derived_steps() + + app = _build_graph() + + scenarios = [ + "Track order 1001 and include its history", + "Issue a refund for order 2048 because it was late (49 dollars)", + ] + + print("\nRunning LangGraph scenarios...") + for prompt in scenarios: + print("=" * 80) + print(f"User: {prompt}") + try: + result = await app.ainvoke({"messages": [HumanMessage(content=prompt)]}) + final_message = result["messages"][-1] + print(f"Assistant: {final_message.content}") + except ControlViolationError as exc: + print(f"Assistant: Request blocked by control policy: {exc.message}") + except RuntimeError as exc: + print( + "Assistant: Control evaluation is unavailable. " + f"Start the Agent Control server and retry. Details: {exc}" + ) + break + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/langchain/pyproject.toml b/examples/langchain/pyproject.toml index 19ade2a3..bf249abf 100644 --- a/examples/langchain/pyproject.toml +++ b/examples/langchain/pyproject.toml @@ -42,5 +42,4 @@ include = [ agent-control-sdk = { path = "../../sdks/python", editable = true } agent-control-models = { path = "../../models", editable = true } agent-control-engine = { path = "../../engine", editable = true } -agent-control-evaluators = { path = "../../evaluators", editable = true } - +agent-control-evaluators = { path = "../../evaluators/builtin", editable = true } From d42f7d29ccb9c89265d15058ae49859ff6fcd658 Mon Sep 17 00:00:00 2001 From: Lev Neiman Date: Thu, 12 Feb 2026 17:17:59 -0800 Subject: [PATCH 8/9] docs(examples): remove auto-schema section heading --- examples/langchain/README.md | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/examples/langchain/README.md b/examples/langchain/README.md index 9a32f49f..d9169459 100644 --- a/examples/langchain/README.md +++ b/examples/langchain/README.md @@ -55,23 +55,6 @@ cd examples/langchain uv run sql_agent_protection.py ``` -## Auto-Derived Step Schema Example (LangGraph) - -This repository now also includes a LangGraph example that demonstrates -automatic step schema derivation from `@control()`-decorated functions. - -Key behavior: -- `agent_control.init(...)` is called **without** explicit `steps=...` -- Tool step schemas are auto-discovered from decorated functions -- Input/output JSON schema is inferred from Python type hints (via Pydantic) - -Run: - -```bash -cd examples/langchain -uv run langgraph_auto_schema_agent.py -``` - ### Local vs Remote Control Execution **Remote (server-side) controls**: From 222c34f9f647e6ad07cd6e73c941f67a43c7d50a Mon Sep 17 00:00:00 2001 From: Lev Neiman Date: Thu, 12 Feb 2026 17:29:21 -0800 Subject: [PATCH 9/9] test(sdk): expand schema derivation and init auto-step coverage --- sdks/python/tests/test_control_registry.py | 38 ++- sdks/python/tests/test_init_step_merge.py | 92 +++++- sdks/python/tests/test_schema_derivation.py | 341 +++++++++++++++++++- 3 files changed, 457 insertions(+), 14 deletions(-) diff --git a/sdks/python/tests/test_control_registry.py b/sdks/python/tests/test_control_registry.py index b5fae211..3ba2d320 100644 --- a/sdks/python/tests/test_control_registry.py +++ b/sdks/python/tests/test_control_registry.py @@ -2,19 +2,19 @@ from __future__ import annotations -from collections.abc import Generator +import functools import logging +from collections.abc import Generator from typing import Any import pytest -from pydantic import BaseModel - from agent_control._control_registry import ( clear, get_registered_steps, merge_explicit_and_auto_steps, register, ) +from pydantic import BaseModel @pytest.fixture(autouse=True) @@ -171,7 +171,7 @@ def test_forward_reference_can_resolve_at_retrieval_time( self, caplog: pytest.LogCaptureFixture ) -> None: # Given a function registered before its forward-referenced model is available in globals. - def my_func(payload: "LaterModel") -> str: + def my_func(payload: LaterModel) -> str: ... register(my_func) @@ -305,6 +305,34 @@ async def stacked(msg: str) -> str: assert steps[0]["name"] == "stacked" assert steps[0]["metadata"] == {"policy": "p1"} + def test_control_with_prior_wrapper_preserves_schema_fields(self) -> None: + # Given a user wrapper applied before @control() using functools.wraps. + from agent_control.control_decorators import control + + def with_tracing(func: Any) -> Any: + @functools.wraps(func) + async def wrapped(*args: Any, **kwargs: Any) -> Any: + return await func(*args, **kwargs) + + return wrapped + + @control() + @with_tracing + async def wrapped_chat(message: str, include_context: bool = False) -> str: + return message + + # When registered steps are materialized. + steps = get_registered_steps() + + # Then registration keeps the original step identity and inferred schema fields. + assert len(steps) == 1 + step = steps[0] + assert step["type"] == "llm" + assert step["name"] == "wrapped_chat" + assert step["input_schema"]["properties"]["message"]["type"] == "string" + assert step["input_schema"]["properties"]["include_context"]["type"] == "boolean" + assert step["output_schema"]["type"] == "string" + class TestInitMerge: """Tests the explicit+auto merge behavior used by init().""" @@ -345,7 +373,7 @@ def my_step(query: str) -> str: merge_result = merge_explicit_and_auto_steps(explicit_steps, auto_steps) merged = merge_result.steps - # Then the explicit version wins for the duplicate key while unrelated auto steps are retained. + # Then explicit wins for duplicates while unrelated auto steps are retained. my_step_entries = [s for s in merged if (s["type"], s["name"]) == ("llm", "my_step")] assert len(my_step_entries) == 1 assert my_step_entries[0]["input_schema"] == {"custom": True} diff --git a/sdks/python/tests/test_init_step_merge.py b/sdks/python/tests/test_init_step_merge.py index b21b11a1..e34c5155 100644 --- a/sdks/python/tests/test_init_step_merge.py +++ b/sdks/python/tests/test_init_step_merge.py @@ -2,16 +2,20 @@ from __future__ import annotations -from collections.abc import Generator import logging +from collections.abc import Generator +from typing import TYPE_CHECKING from unittest.mock import AsyncMock, patch from uuid import uuid4 -import pytest - import agent_control +import pytest from agent_control._control_registry import clear, register +if TYPE_CHECKING: + # Intentionally unavailable at runtime to trigger unresolved forward-ref fallback. + class DoesNotExist: ... + @pytest.fixture(autouse=True) def _clean_registry() -> Generator[None, None, None]: @@ -74,3 +78,85 @@ def auto_llm(query: str) -> str: assert llm_entries[0]["description"] == "Explicit override for auto_llm." assert any((s["type"], s["name"]) == ("tool", "manual_tool") for s in merged_steps) assert "Skipping 1 auto-discovered step(s) overridden by explicit steps" in caplog.text + + +def test_init_uses_auto_discovered_steps_from_control_decorator() -> None: + # Given a real @control()-decorated async function and no explicit steps passed to init(). + from agent_control.control_decorators import control + + @control() + async def auto_chat(message: str, temperature: float = 0.2) -> str: + """Auto-discovered chat step.""" + return message + + register_agent_mock = AsyncMock(return_value={"created": True, "controls": []}) + health_check_mock = AsyncMock(return_value={"status": "healthy"}) + + # When init() performs registration. + with patch( + "agent_control.__init__.AgentControlClient.health_check", + new=health_check_mock, + ), patch( + "agent_control.__init__.agents.register_agent", + new=register_agent_mock, + ): + agent_control.init( + agent_name="Auto Discovery Agent", + agent_id=str(uuid4()), + ) + + # Then register_agent() receives the auto-derived step schema payload. + assert register_agent_mock.await_count == 1 + assert register_agent_mock.await_args is not None + merged_steps = register_agent_mock.await_args.kwargs["steps"] + + auto_entries = [s for s in merged_steps if (s["type"], s["name"]) == ("llm", "auto_chat")] + assert len(auto_entries) == 1 + auto_step = auto_entries[0] + assert auto_step["description"] == "Auto-discovered chat step." + assert auto_step["input_schema"]["properties"]["message"]["type"] == "string" + assert auto_step["input_schema"]["properties"]["temperature"]["type"] == "number" + assert auto_step["output_schema"]["type"] == "string" + + +def test_init_logs_fallback_warning_for_unresolved_type_hints( + caplog: pytest.LogCaptureFixture, +) -> None: + # Given a decorated function whose forward reference cannot be resolved at runtime. + from agent_control.control_decorators import control + + @control() + async def unresolved(payload: DoesNotExist) -> str: + """Function with unresolved forward reference.""" + return "ok" + + register_agent_mock = AsyncMock(return_value={"created": True, "controls": []}) + health_check_mock = AsyncMock(return_value={"status": "healthy"}) + + # When init() materializes auto-discovered steps. + with patch( + "agent_control.__init__.AgentControlClient.health_check", + new=health_check_mock, + ), patch( + "agent_control.__init__.agents.register_agent", + new=register_agent_mock, + ): + with caplog.at_level(logging.WARNING): + agent_control.init( + agent_name="Fallback Warning Agent", + agent_id=str(uuid4()), + ) + + # Then initialization continues, using fallback schemas and emitting a warning. + assert register_agent_mock.await_count == 1 + assert register_agent_mock.await_args is not None + merged_steps = register_agent_mock.await_args.kwargs["steps"] + + unresolved_entries = [ + s for s in merged_steps if (s["type"], s["name"]) == ("llm", "unresolved") + ] + assert len(unresolved_entries) == 1 + unresolved_step = unresolved_entries[0] + assert unresolved_step["input_schema"] == {"type": "object", "additionalProperties": True} + assert unresolved_step["output_schema"] == {} + assert "failed to resolve type hints" in caplog.text diff --git a/sdks/python/tests/test_schema_derivation.py b/sdks/python/tests/test_schema_derivation.py index 4f87331e..edefc69c 100644 --- a/sdks/python/tests/test_schema_derivation.py +++ b/sdks/python/tests/test_schema_derivation.py @@ -2,15 +2,21 @@ from __future__ import annotations +import functools import logging -from typing import Any +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING, Annotated, Any, Literal from unittest.mock import MagicMock -import pytest -from pydantic import BaseModel - import agent_control._schema_derivation as schema_derivation +import pytest from agent_control._schema_derivation import derive_schemas +from pydantic import BaseModel, Field + +if TYPE_CHECKING: + # Intentionally not defined at runtime; used to test unresolved forward-ref fallback behavior. + class DoesNotExist: ... class _InputModel(BaseModel): @@ -22,6 +28,18 @@ class _OutputModel(BaseModel): answer: str +class _OrderState(str, Enum): + PENDING = "pending" + APPROVED = "approved" + REJECTED = "rejected" + + +@dataclass +class _DataPayload: + query: str + limit: int = 5 + + def golden_primitive_defaults(query: str, limit: int = 10) -> str: """Golden-case primitive/default function.""" raise NotImplementedError @@ -42,6 +60,16 @@ def golden_nested_models(payload: _InputModel) -> _OutputModel: raise NotImplementedError +def _resolve_local_ref(container: dict[str, Any], schema: dict[str, Any]) -> dict[str, Any]: + """Resolve local ``#/$defs/...`` refs for assertions in tests.""" + ref = schema.get("$ref") + if not isinstance(ref, str): + return schema + assert ref.startswith("#/$defs/") + def_name = ref.split("/")[-1] + return container["$defs"][def_name] + + class TestInputInference: """Input schema derivation tests.""" @@ -100,6 +128,139 @@ def my_func(x, y): assert schemas.input_schema["type"] == "object" assert set(schemas.input_schema["properties"]) == {"x", "y"} + def test_keyword_only_parameters_are_included(self) -> None: + # Given a callable with keyword-only parameters. + def my_func(*, key: str, verbose: bool = False) -> str: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then keyword-only fields appear with expected required/default behavior. + assert schemas.input_schema["properties"]["key"]["type"] == "string" + assert schemas.input_schema["properties"]["verbose"]["type"] == "boolean" + assert schemas.input_schema["properties"]["verbose"]["default"] is False + assert set(schemas.input_schema.get("required", [])) == {"key"} + + def test_non_nullable_multi_union_input_is_preserved(self) -> None: + # Given a callable with a non-nullable multi-member union input. + def my_func(value: str | int) -> str: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then the input union preserves both primitive branches without adding null. + value_schema = schemas.input_schema["properties"]["value"] + any_of = value_schema.get("anyOf") + assert isinstance(any_of, list) + any_of_types = {item["type"] for item in any_of} + assert any_of_types == {"string", "integer"} + + def test_literal_input_is_emitted_as_enum(self) -> None: + # Given a callable with a Literal-constrained input parameter. + def my_func(mode: Literal["fast", "accurate"]) -> str: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then the parameter schema is emitted as an enum. + assert schemas.input_schema["properties"]["mode"]["enum"] == ["fast", "accurate"] + + def test_annotated_input_preserves_field_metadata(self) -> None: + # Given a callable using Annotated with Field metadata. + def my_func( + query: Annotated[str, Field(description="Natural language query", min_length=3)] + ) -> str: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then Annotated metadata is preserved in the emitted input schema. + query_schema = schemas.input_schema["properties"]["query"] + assert query_schema["description"] == "Natural language query" + assert query_schema["minLength"] == 3 + + def test_set_input_is_array_with_unique_items(self) -> None: + # Given a callable with a set-typed input parameter. + def my_func(tags: set[str]) -> str: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then the set is represented as an array with uniqueItems. + tags_schema = schemas.input_schema["properties"]["tags"] + assert tags_schema["type"] == "array" + assert tags_schema["uniqueItems"] is True + assert tags_schema["items"]["type"] == "string" + + def test_tuple_input_uses_prefix_items(self) -> None: + # Given a callable with a fixed-length tuple input parameter. + def my_func(pair: tuple[str, int]) -> str: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then tuple structure is represented via prefixItems and tuple bounds. + pair_schema = schemas.input_schema["properties"]["pair"] + assert pair_schema["type"] == "array" + assert pair_schema["minItems"] == 2 + assert pair_schema["maxItems"] == 2 + assert len(pair_schema["prefixItems"]) == 2 + assert pair_schema["prefixItems"][0]["type"] == "string" + assert pair_schema["prefixItems"][1]["type"] == "integer" + + def test_default_none_without_optional_annotation(self) -> None: + # Given a callable with `str` annotation but None default value. + def my_func(query: str = None) -> str: # type: ignore[assignment] + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then field is optional in requirements and carries a None default. + query_schema = schemas.input_schema["properties"]["query"] + assert query_schema["default"] is None + assert "query" not in schemas.input_schema.get("required", []) + assert query_schema["type"] == "string" + + def test_enum_input_schema_smoke(self) -> None: + # Given a callable with an Enum-constrained input parameter. + def my_func(state: _OrderState) -> str: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then the enum values are preserved in the input property schema. + state_schema = _resolve_local_ref( + schemas.input_schema, + schemas.input_schema["properties"]["state"], + ) + assert state_schema["type"] == "string" + assert state_schema["enum"] == ["pending", "approved", "rejected"] + + def test_dataclass_input_schema_smoke(self) -> None: + # Given a callable that accepts a standard-library dataclass input. + def my_func(payload: _DataPayload) -> str: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then the dataclass shape is reflected in the input schema. + payload_schema = _resolve_local_ref( + schemas.input_schema, + schemas.input_schema["properties"]["payload"], + ) + assert payload_schema["type"] == "object" + assert payload_schema["properties"]["query"]["type"] == "string" + assert payload_schema["properties"]["limit"]["type"] == "integer" + class TestOutputInference: """Output schema derivation tests.""" @@ -113,6 +274,103 @@ def my_func() -> str: schemas = derive_schemas(my_func) # Then the output schema is emitted as a string type. + assert schemas.input_schema["type"] == "object" + assert schemas.input_schema.get("properties") == {} + assert schemas.input_schema.get("required", []) == [] + assert schemas.output_schema["type"] == "string" + + def test_async_function_output_derivation(self) -> None: + # Given an async callable with typed input and output annotations. + async def my_func(query: str) -> str: + return query + + # When schemas are derived directly from the async function. + schemas = derive_schemas(my_func) + + # Then input and output schemas are inferred from the annotated signature. + assert schemas.input_schema["properties"]["query"]["type"] == "string" + assert schemas.output_schema["type"] == "string" + + def test_literal_output_is_emitted_as_enum(self) -> None: + # Given a callable returning a Literal-constrained value. + def my_func() -> Literal["ok", "retry"]: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then the output schema is emitted as an enum. + assert schemas.output_schema["enum"] == ["ok", "retry"] + + def test_any_output_is_permissive(self) -> None: + # Given a callable explicitly annotated with Any output. + def my_func(query: str) -> Any: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then output schema stays permissive (empty object schema). + assert schemas.output_schema == {} + + def test_annotated_output_preserves_field_metadata(self) -> None: + # Given a callable with Annotated return metadata. + def my_func() -> Annotated[str, Field(description="Normalized answer")]: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then output metadata from Annotated is preserved. + assert schemas.output_schema["type"] == "string" + assert schemas.output_schema["description"] == "Normalized answer" + + def test_enum_output_schema_smoke(self) -> None: + # Given a callable returning an Enum value. + def my_func() -> _OrderState: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then output schema preserves enum value constraints. + output_schema = _resolve_local_ref(schemas.output_schema, schemas.output_schema) + assert output_schema["type"] == "string" + assert output_schema["enum"] == ["pending", "approved", "rejected"] + + def test_dataclass_output_schema_smoke(self) -> None: + # Given a callable returning a standard-library dataclass. + def my_func(query: str) -> _DataPayload: + ... + + # When schemas are derived. + schemas = derive_schemas(my_func) + + # Then output schema reflects the dataclass field structure. + output_schema = _resolve_local_ref(schemas.output_schema, schemas.output_schema) + assert output_schema["type"] == "object" + assert output_schema["properties"]["query"]["type"] == "string" + assert output_schema["properties"]["limit"]["type"] == "integer" + + +class TestFunctionUnwrapBehavior: + """unwrap() behavior for decorated callables.""" + + def test_wrapped_function_uses_unwrapped_signature(self) -> None: + # Given a wrapped function where only the unwrapped callable has useful type hints. + def base(query: str, limit: int = 3) -> str: + ... + + @functools.wraps(base) + def wrapped(*args: Any, **kwargs: Any) -> Any: + return base(*args, **kwargs) + + # When schemas are derived from the wrapped callable. + schemas = derive_schemas(wrapped) + + # Then derive_schemas() uses inspect.unwrap() and reflects the base signature. + assert set(schemas.input_schema["properties"]) == {"query", "limit"} + assert set(schemas.input_schema.get("required", [])) == {"query"} assert schemas.output_schema["type"] == "string" def test_pydantic_input_and_output(self) -> None: @@ -180,6 +438,37 @@ def my_func(query: str) -> str: assert schemas.input_schema["properties"]["query"]["type"] == "string" assert "args_schema.model_json_schema() failed" in caplog.text + def test_args_schema_override_wins_for_wrapped_function(self) -> None: + # Given a wrapped function with args_schema on the wrapper and typed signature on the base. + class _WrapperArgsSchema: + @staticmethod + def model_json_schema() -> dict[str, Any]: + return { + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + } + + def base(query: str, limit: int = 3) -> str: + ... + + @functools.wraps(base) + def wrapped(*args: Any, **kwargs: Any) -> Any: + return base(*args, **kwargs) + + wrapped.args_schema = _WrapperArgsSchema() # type: ignore[attr-defined] + + # When schemas are derived from the wrapped callable. + schemas = derive_schemas(wrapped) + + # Then wrapper args_schema is used for input while output is inferred from unwrapped return. + assert schemas.input_schema == { + "type": "object", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + } + assert schemas.output_schema["type"] == "string" + class TestFallbackWarnings: """Warnings and fallback behavior for unresolved/incomplete typing.""" @@ -203,7 +492,7 @@ def test_unresolved_type_hints_warn_and_fall_back( self, caplog: pytest.LogCaptureFixture ) -> None: # Given a callable that references an unresolved forward type hint. - def my_func(query: "DoesNotExist") -> str: + def my_func(query: DoesNotExist) -> str: ... # When schema derivation attempts to resolve type hints. @@ -267,7 +556,7 @@ class _GoodArgsSchema: def model_json_schema() -> dict[str, Any]: return {"type": "object", "properties": {"q": {"type": "string"}}} - def my_func(query: "DoesNotExist") -> str: + def my_func(query: DoesNotExist) -> str: ... my_func.args_schema = _GoodArgsSchema() # type: ignore[attr-defined] @@ -429,6 +718,46 @@ def test_golden_primitive_defaults_snapshot(self) -> None: } assert schemas.output_schema == {"type": "string"} + +class TestJsonSchemaContract: + """Validate derived schemas are valid JSON Schemas.""" + + def test_derived_schemas_are_json_schema_valid(self) -> None: + # Given a representative set of callables and their derived schemas. + jsonschema = pytest.importorskip("jsonschema") + draft202012_validator = jsonschema.Draft202012Validator + + def enum_case(state: _OrderState) -> _OrderState: + ... + + def tuple_case(pair: tuple[str, int]) -> tuple[str, int]: + ... + + def dataclass_case(payload: _DataPayload) -> _DataPayload: + ... + + async def async_case(message: str) -> str: + return message + + cases = [ + golden_primitive_defaults, + golden_optional_union, + golden_collections, + golden_nested_models, + enum_case, + tuple_case, + dataclass_case, + async_case, + ] + + # When each callable is passed through schema derivation. + for func in cases: + schemas = derive_schemas(func) + + # Then both input and output schemas pass Draft 2020-12 structural validation. + draft202012_validator.check_schema(schemas.input_schema) + draft202012_validator.check_schema(schemas.output_schema) + def test_golden_optional_union_snapshot(self) -> None: # Given an optional-union input and explicit `None` return annotation. # When schemas are derived.