Skip to content

Commit a485f93

Browse files
feat(sdk): Get Agent Controls from SDK Init (#15)
### Summary agent_control.get_server_controls() returns the cached list of controls that were fetched from the Agent Control server. #### Purpose: Returns controls fetched during init() or the last refresh_controls() call Used internally by the @control() decorator to determine which controls to apply Useful for debugging/inspecting what controls are active for your agent #### Returns: list[dict[str, Any]] — List of control dictionaries if controls are cached None — If init() hasn't been called yet or no controls were fetched #### Usage ``` import agent_control # Initialize the SDK (fetches controls from server) agent_control.init( agent_name="my-agent", agent_id="agent-123", server_url="http://localhost:8000" ) # Inspect the cached controls controls = agent_control.get_server_controls() if controls: print(f"Loaded {len(controls)} controls") for c in controls: print(f" - {c['name']} (execution: {c['control'].get('execution', 'server')})") ```
1 parent d5ae157 commit a485f93

4 files changed

Lines changed: 288 additions & 16 deletions

File tree

‎examples/crewai/pyproject.toml‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ version = "0.1.0"
44
description = "CrewAI content moderation example with Agent Control"
55
requires-python = ">=3.12"
66
dependencies = [
7-
"agent-control-sdk>=2.1.0,<3.0.0",
7+
"agent-control-sdk>=3.0.0,<4.0.0",
88
"crewai>=0.80.0",
99
"crewai-tools>=0.12.0",
1010
"openai>=1.0.0",
@@ -26,4 +26,4 @@ build-backend = "hatchling.build"
2626
include = [
2727
"*.py",
2828
".env.example",
29-
]
29+
]

‎sdks/python/src/agent_control/__init__.py‎

Lines changed: 123 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ async def process(input: str) -> str:
3737
import os
3838
from collections.abc import Callable
3939
from datetime import UTC, datetime
40+
from importlib.metadata import PackageNotFoundError
41+
from importlib.metadata import version as get_version
4042
from typing import TYPE_CHECKING, Any, Literal, TypeVar
4143
from uuid import UUID
4244

@@ -180,11 +182,119 @@ def __init__(
180182
_current_agent: Agent | None = None
181183
_control_engine = None
182184
_client: AgentControlClient | None = None
183-
_server_controls: list | None = None
185+
_server_controls: list[dict[str, Any]] | None = None
186+
_server_url: str | None = None
187+
_api_key: str | None = None
184188

185189
F = TypeVar("F", bound=Callable[..., Any])
186190

187191

192+
def get_server_controls() -> list[dict[str, Any]] | None:
193+
"""Get the cached server controls.
194+
195+
Returns the controls that were fetched during init() or the last
196+
refresh_controls() call. Returns None if no controls are cached.
197+
198+
Returns:
199+
List of control dicts or None if not initialized.
200+
201+
Example:
202+
controls = agent_control.get_server_controls()
203+
if controls:
204+
print(f"Loaded {len(controls)} controls")
205+
for c in controls:
206+
print(f" - {c['name']} (execution: {c['control'].get('execution', 'server')})")
207+
"""
208+
return _server_controls
209+
210+
211+
async def refresh_controls_async() -> list[dict[str, Any]] | None:
212+
"""Refresh controls from the server asynchronously.
213+
214+
Fetches the latest controls from the server and updates the cache.
215+
Use this when you've made changes to controls in the UI or API
216+
and want the SDK to pick them up without restarting.
217+
218+
Returns:
219+
List of control dicts or None if fetch failed.
220+
221+
Example:
222+
# After updating a control in the UI
223+
controls = await agent_control.refresh_controls_async()
224+
print(f"Refreshed {len(controls)} controls")
225+
"""
226+
global _server_controls
227+
228+
if _current_agent is None:
229+
raise RuntimeError("Agent not initialized. Call agent_control.init() first.")
230+
231+
if _server_url is None:
232+
raise RuntimeError("Server URL not set. Call agent_control.init() first.")
233+
234+
async with AgentControlClient(base_url=_server_url, api_key=_api_key) as client:
235+
response = await agents.register_agent(
236+
client,
237+
_current_agent,
238+
steps=[]
239+
)
240+
_server_controls = response.get('controls', [])
241+
logger.info("Refreshed %d control(s) from server", len(_server_controls or []))
242+
return _server_controls
243+
244+
245+
def refresh_controls() -> list[dict[str, Any]] | None:
246+
"""Refresh controls from the server synchronously.
247+
248+
Fetches the latest controls from the server and updates the cache.
249+
Use this when you've made changes to controls in the UI or API
250+
and want the SDK to pick them up without restarting.
251+
252+
Returns:
253+
List of control dicts or None if fetch failed.
254+
255+
Example:
256+
# After updating a control in the UI
257+
controls = agent_control.refresh_controls()
258+
print(f"Refreshed {len(controls)} controls")
259+
"""
260+
import asyncio
261+
262+
try:
263+
loop = asyncio.get_running_loop()
264+
# We're in an async context - run in thread
265+
import threading
266+
267+
result_container: list[list[dict[str, Any]] | None] = [None]
268+
exception_container: list[Exception | None] = [None]
269+
270+
def run_in_thread() -> None:
271+
new_loop = asyncio.new_event_loop()
272+
asyncio.set_event_loop(new_loop)
273+
try:
274+
result_container[0] = new_loop.run_until_complete(refresh_controls_async())
275+
except Exception as e:
276+
exception_container[0] = e
277+
finally:
278+
new_loop.close()
279+
280+
thread = threading.Thread(target=run_in_thread)
281+
thread.start()
282+
thread.join(timeout=10)
283+
284+
if exception_container[0]:
285+
raise exception_container[0]
286+
return result_container[0]
287+
288+
except RuntimeError:
289+
# No running event loop - we're in a sync context
290+
loop = asyncio.new_event_loop()
291+
asyncio.set_event_loop(loop)
292+
try:
293+
return loop.run_until_complete(refresh_controls_async())
294+
finally:
295+
loop.close()
296+
297+
188298
# ============================================================================
189299
# Public API Functions
190300
# ============================================================================
@@ -260,7 +370,7 @@ async def handle(message: str):
260370
Environment Variables:
261371
AGENT_CONTROL_URL: Server URL (default: http://localhost:8000)
262372
"""
263-
global _current_agent, _control_engine, _client, _server_controls
373+
global _current_agent, _control_engine, _client, _server_controls, _server_url, _api_key
264374

265375
if not agent_id:
266376
raise ValueError(
@@ -290,6 +400,7 @@ async def handle(message: str):
290400

291401
# Get server URL (ensure it's always a string)
292402
_server_url = server_url or os.getenv('AGENT_CONTROL_URL') or 'http://localhost:8000'
403+
_api_key = api_key
293404

294405
# Register with server and fetch controls
295406
server_controls = None
@@ -915,6 +1026,11 @@ async def main():
9151026
"init",
9161027
"current_agent",
9171028

1029+
# Control sync
1030+
"get_server_controls",
1031+
"refresh_controls",
1032+
"refresh_controls_async",
1033+
9181034
# SDK Logging
9191035
"get_logger",
9201036

@@ -990,4 +1106,8 @@ async def main():
9901106
"EvaluatorConfig",
9911107
]
9921108

993-
__version__ = "0.1.0"
1109+
try:
1110+
__version__ = get_version("agent-control-sdk")
1111+
except PackageNotFoundError:
1112+
# Package not installed (e.g., running from source without install)
1113+
__version__ = "0.0.0.dev"

‎sdks/python/src/agent_control/control_decorators.py‎

Lines changed: 109 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -189,8 +189,16 @@ async def _evaluate(
189189
server_url: str,
190190
trace_id: str | None = None,
191191
span_id: str | None = None,
192+
controls: list[dict[str, Any]] | None = None,
192193
) -> dict[str, Any]:
193-
"""Call server evaluation endpoint asynchronously."""
194+
"""Call evaluation with support for local (SDK) and server execution.
195+
196+
If controls are provided, uses check_evaluation_with_local() which:
197+
- Evaluates execution="sdk" controls locally in the SDK
198+
- Sends execution="server" controls to the server
199+
200+
If no controls provided, falls back to server-only evaluation.
201+
"""
194202
# Build headers with trace/span IDs for distributed tracing
195203
headers = {}
196204
if trace_id:
@@ -199,6 +207,80 @@ async def _evaluate(
199207
headers["X-Span-Id"] = span_id
200208

201209
async with AgentControlClient(base_url=server_url) as client:
210+
# If we have controls, use local evaluation which handles both SDK and server controls
211+
if controls is not None:
212+
try:
213+
from uuid import UUID
214+
215+
from agent_control.evaluation import check_evaluation_with_local
216+
217+
# Build Step object for evaluation
218+
try:
219+
from agent_control_models import Step
220+
step_obj = Step(**step)
221+
except ImportError:
222+
step_obj = step # type: ignore
223+
224+
result = await check_evaluation_with_local(
225+
client=client,
226+
agent_uuid=UUID(agent_uuid),
227+
step=step_obj,
228+
stage=stage, # type: ignore
229+
controls=controls,
230+
)
231+
232+
# Convert result to dict format expected by process_result
233+
return {
234+
"is_safe": result.is_safe,
235+
"confidence": result.confidence,
236+
"reason": result.reason,
237+
"matches": [
238+
{
239+
"control_id": m.control_id,
240+
"control_name": m.control_name,
241+
"action": m.action,
242+
"result": {
243+
"matched": m.result.matched,
244+
"confidence": m.result.confidence,
245+
"message": m.result.message,
246+
"error": m.result.error,
247+
"metadata": m.result.metadata,
248+
},
249+
"control_execution_id": m.control_execution_id,
250+
}
251+
for m in (result.matches or [])
252+
] if result.matches else None,
253+
"errors": [
254+
{
255+
"control_id": e.control_id,
256+
"control_name": e.control_name,
257+
"action": e.action,
258+
"result": {
259+
"matched": e.result.matched,
260+
"confidence": e.result.confidence,
261+
"message": e.result.message,
262+
"error": e.result.error,
263+
"metadata": e.result.metadata,
264+
},
265+
"control_execution_id": e.control_execution_id,
266+
}
267+
for e in (result.errors or [])
268+
] if result.errors else None,
269+
"non_matches": None, # check_evaluation_with_local doesn't return non_matches
270+
}
271+
except ImportError:
272+
logger.warning(
273+
"Local evaluation not available (missing agent_control_engine). "
274+
"Falling back to server-only evaluation. "
275+
"Controls with execution='sdk' will be skipped."
276+
)
277+
except Exception as e:
278+
logger.warning(
279+
"Local evaluation failed: %s. Falling back to server-only evaluation.",
280+
e,
281+
)
282+
283+
# Fallback: server-only evaluation
202284
response = await client.http_client.post(
203285
"/api/v1/evaluation",
204286
json={
@@ -209,8 +291,8 @@ async def _evaluate(
209291
headers=headers,
210292
)
211293
response.raise_for_status()
212-
result: dict[str, Any] = response.json()
213-
return result
294+
result_dict: dict[str, Any] = response.json()
295+
return result_dict
214296

215297

216298
def _extract_input_from_args(func: Callable, args: tuple, kwargs: dict) -> str:
@@ -379,6 +461,19 @@ def _log_single_control(
379461
)
380462

381463

464+
def _get_server_controls() -> list[dict[str, Any]] | None:
465+
"""Get the cached server controls from agent_control module."""
466+
try:
467+
import agent_control
468+
return agent_control.get_server_controls()
469+
except Exception as exc:
470+
logger.debug(
471+
"Unable to access cached server controls; proceeding without local cache.",
472+
exc_info=exc,
473+
)
474+
return None
475+
476+
382477
async def _execute_with_control(
383478
func: Callable,
384479
args: tuple,
@@ -392,6 +487,10 @@ async def _execute_with_control(
392487
for async functions, or inside asyncio.run() for sync functions), so it can
393488
always use await _evaluate() directly.
394489
490+
Uses cached controls from init() to support both SDK-side and server-side
491+
evaluation. Controls with execution="sdk" are evaluated locally, while
492+
execution="server" controls are sent to the server.
493+
395494
Args:
396495
func: The wrapped function to execute
397496
args: Positional arguments for the function
@@ -414,6 +513,9 @@ async def _execute_with_control(
414513
return await func(*args, **kwargs)
415514
return func(*args, **kwargs)
416515

516+
# Get cached controls for local evaluation support
517+
controls = _get_server_controls()
518+
417519
# Get trace context: inherit trace_id if set, always generate new span_id
418520
# This allows multiple @control() calls to share the same trace but have unique spans
419521
existing_trace_id = get_current_trace_id()
@@ -441,7 +543,8 @@ async def _execute_with_control(
441543
try:
442544
result = await _evaluate(
443545
ctx.agent_uuid, ctx.pre_payload(), "pre",
444-
ctx.server_url, ctx.trace_id, ctx.span_id
546+
ctx.server_url, ctx.trace_id, ctx.span_id,
547+
controls=controls,
445548
)
446549
ctx.process_result(result, "pre")
447550
except ControlViolationError:
@@ -463,7 +566,8 @@ async def _execute_with_control(
463566
try:
464567
result = await _evaluate(
465568
ctx.agent_uuid, ctx.post_payload(output), "post",
466-
ctx.server_url, ctx.trace_id, ctx.span_id
569+
ctx.server_url, ctx.trace_id, ctx.span_id,
570+
controls=controls,
467571
)
468572
ctx.process_result(result, "post")
469573
except ControlViolationError:

0 commit comments

Comments
 (0)