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/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 } diff --git a/sdks/python/src/agent_control/__init__.py b/sdks/python/src/agent_control/__init__.py index 06a16069..71e9a519 100644 --- a/sdks/python/src/agent_control/__init__.py +++ b/sdks/python/src/agent_control/__init__.py @@ -66,6 +66,14 @@ async def process(input: str) -> str: ) from . import agents, controls, evaluation, evaluators, policies +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 @@ -340,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 @@ -435,6 +443,29 @@ 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 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: + if merge_result.overridden_keys: + formatted = ", ".join( + 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(merge_result.overridden_keys), + formatted, + ) + + logger.debug( + "Auto-discovered %d step(s) from @control() decorators (%d after merge)", + len(auto_steps), + len(registration_steps), + ) + # Register with server and fetch controls server_controls = None try: @@ -456,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', []) @@ -466,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: @@ -1054,6 +1085,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..5b37c47c --- /dev/null +++ b/sdks/python/src/agent_control/_control_registry.py @@ -0,0 +1,189 @@ +"""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 logging +from collections.abc import Callable +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] + + +@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, _RegisteredControl] = {} +"""Maps ``(type, name)`` -> registration metadata. Keyed by type+name to deduplicate.""" + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def register(func: Callable[..., Any], policy: str | None = None) -> None: + """Register a decorated function's step schema in the registry. + + 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. + 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 + + metadata: dict[str, Any] = {} + if policy is not None: + metadata["policy"] = policy + + 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: + logger.debug( + "Overwriting previously registered step '%s' (type=%s)", + step_name, + step_type, + ) + _registered_steps[key] = registered + logger.debug("Registered step schema: %s (type=%s)", step_name, step_type) + + +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 + by ``init(steps=...)``. + """ + 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( + 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 new file mode 100644 index 00000000..f3d97a0b --- /dev/null +++ b/sdks/python/src/agent_control/_schema_derivation.py @@ -0,0 +1,256 @@ +"""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_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, unwrapped, hints) + + output_schema = _infer_output_schema(func, hints) + 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], + unwrapped: Callable[..., Any], + hints: dict[str, Any], +) -> dict[str, Any]: + """Infer input schema from signature + type hints using dynamic Pydantic model.""" + + 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], hints: dict[str, Any]) -> dict[str, Any]: + """Infer output schema from return type annotation using ``TypeAdapter``.""" + 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/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..3ba2d320 --- /dev/null +++ b/sdks/python/tests/test_control_registry.py @@ -0,0 +1,440 @@ +"""Tests for control step registry behavior.""" + +from __future__ import annotations + +import functools +import logging +from collections.abc import Generator +from typing import Any + +import pytest +from agent_control._control_registry import ( + clear, + get_registered_steps, + merge_explicit_and_auto_steps, + register, +) +from pydantic import BaseModel + + +@pytest.fixture(autouse=True) +def _clean_registry() -> Generator[None, None, None]: + """Ensure each test starts with an empty registry.""" + clear() + yield + clear() + + +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" + assert step["type"] == "llm" + assert step["description"] == "Chat with the agent." + 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.""" + + # Given a function that advertises tool-style metadata via attributes. + 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] + + # 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.""" + ... + + def chat_v2(message: str) -> int: # noqa: ARG001 + """Second version.""" + ... + + 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_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: + ... + + # 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" + 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().""" + + def test_clear_empties_registry(self) -> None: + # Given multiple registered functions in the global registry. + def f1(x: str) -> str: + ... + + def f2(x: str) -> str: + ... + + register(f1) + register(f2) + assert len(get_registered_steps()) == 2 + + # When clear() is called. + clear() + + # Then the registry becomes empty. + assert len(get_registered_steps()) == 0 + + +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() + 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: + """Look up a record.""" + return query + + _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 + 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.""" + + # Given a function wrapped by two @control() decorators with different policies. + from agent_control.control_decorators import control + + @control(policy="p1") + @control(policy="p2") + 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"} + + 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().""" + + 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() + 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: + ... + + register(my_step) + + explicit_steps: list[dict[str, Any]] = [ + {"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() + merge_result = merge_explicit_and_auto_steps(explicit_steps, auto_steps) + merged = merge_result.steps + + # 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} + 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_explicit_and_auto_steps(explicit, get_registered_steps()).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_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} + 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_init_step_merge.py b/sdks/python/tests/test_init_step_merge.py new file mode 100644 index 00000000..e34c5155 --- /dev/null +++ b/sdks/python/tests/test_init_step_merge.py @@ -0,0 +1,162 @@ +"""Tests for init() step merge wiring into register_agent.""" + +from __future__ import annotations + +import logging +from collections.abc import Generator +from typing import TYPE_CHECKING +from unittest.mock import AsyncMock, patch +from uuid import uuid4 + +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]: + """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 + + +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 new file mode 100644 index 00000000..edefc69c --- /dev/null +++ b/sdks/python/tests/test_schema_derivation.py @@ -0,0 +1,862 @@ +"""Tests for isolated schema derivation logic.""" + +from __future__ import annotations + +import functools +import logging +from dataclasses import dataclass +from enum import Enum +from typing import TYPE_CHECKING, Annotated, Any, Literal +from unittest.mock import MagicMock + +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): + query: str + limit: int = 5 + + +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 + + +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 + + +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.""" + + 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) + and "null" in conversation_schema["type"] + ) + 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"} + + 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.""" + + 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.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: + # 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 + ) + assert schemas.output_schema["type"] == "object" + + +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", + "properties": {"q": {"type": "string"}}, + "required": ["q"], + } + + 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"}}, + "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: + # 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") + + 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 + + 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.""" + + 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 + + +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"} + + +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. + 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"}