feat(sdk): Auto-populate init() steps from @control() decorators

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`)
 

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())

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 }

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",