Reliable, observable agents

budget.py# An agent loop with HARD CAPS. The 'model' is a deterministic fake so this
# runs offline; a real LLM drops into the same slot. We simulate tokens and
# tool calls with fixed counts (no wall-clock, no randomness).

def fake_model(step):
    # Turns 0..2 call a tool; turn 3 would answer -- but caps may stop us first.
    if step < 3:
        return {'type': 'tool_call', 'name': 'search', 'sim_tokens': 120}
    return {'type': 'final', 'text': 'done: found the answer', 'sim_tokens': 90}

def run_agent(question, max_steps=8, max_tokens=1000, max_tool_calls=5):
    used_tokens = 0
    tool_calls = 0
    for step in range(max_steps + 1):
        if step >= max_steps:
            return {'stop': 'STEP_CAP', 'step': step, 'tokens': used_tokens, 'tool_calls': tool_calls}
        decision = fake_model(step)
        used_tokens += decision['sim_tokens']
        if used_tokens > max_tokens:
            return {'stop': 'TOKEN_BUDGET', 'step': step, 'tokens': used_tokens, 'tool_calls': tool_calls}
        if decision['type'] == 'final':
            return {'stop': 'FINAL_ANSWER', 'step': step, 'tokens': used_tokens,
                    'tool_calls': tool_calls, 'text': decision['text']}
        tool_calls += 1
        if tool_calls > max_tool_calls:
            return {'stop': 'TOOL_CALL_CAP', 'step': step, 'tokens': used_tokens, 'tool_calls': tool_calls}
    return {'stop': 'STEP_CAP', 'step': max_steps, 'tokens': used_tokens, 'tool_calls': tool_calls}

# 1. Normal run: reaches a final answer within every budget.
print('normal :', run_agent('find the answer'))
# 2. A tight token budget trips first -- the loop stops cleanly and says why.
print('budget :', run_agent('find the answer', max_tokens=200))
# 3. A tight step cap trips before the model would ever answer.
print('stepcap:', run_agent('find the answer', max_steps=2, max_tokens=10000))

guardrails.py# A guardrail layer with THREE gates around the agent:
#   input  -> block a disallowed request before the model ever runs,
#   tool   -> block a dangerous tool call the model proposes,
#   output -> validate/filter the final answer before it reaches the user.
# All checks are plain rules; a real system may add a model-based classifier.

class Blocked(Exception):
    pass

BANNED_INPUT = ('password', 'ssn', 'wire money')
DANGEROUS_TOOLS = {'delete_database', 'transfer_funds', 'run_shell'}

def input_guard(request):
    low = request.lower()
    for term in BANNED_INPUT:
        if term in low:
            raise Blocked(f'input refused: contains disallowed term {term!r}')
    return request

def tool_guard(name, args):
    if name in DANGEROUS_TOOLS:
        raise Blocked(f'tool refused: {name!r} is dangerous (needs human approval)')
    return name, args

def output_guard(answer):
    # Redact anything that looks like a secret before it leaves the system.
    import re
    redacted = re.sub(r'\bsk-[A-Za-z0-9]+', '[REDACTED_KEY]', answer)
    if len(redacted) > 200:
        redacted = redacted[:200] + ' ...'
    return redacted

def handle(request, proposed_tool, answer):
    try:
        input_guard(request)
        tool_guard(*proposed_tool)
    except Blocked as e:
        return {'status': 'blocked', 'reason': str(e)}
    return {'status': 'ok', 'answer': output_guard(answer)}

# 1. Clean request, safe tool, secret leaked in the answer -> runs, output redacted.
print(handle('what is our refund policy?', ('search_docs', {'q': 'refund'}),
             'Our policy is 30 days. (debug key sk-abc123XYZ)'))
# 2. Disallowed input -> blocked before the model runs at all.
print(handle('email me the admin password', ('search_docs', {}), 'irrelevant'))
# 3. Model proposes a dangerous tool -> tool guard intercepts it.
print(handle('clean up old records', ('delete_database', {'name': 'prod'}), 'irrelevant'))

tracer.py# A tiny TRACER: record every step of a run (thought, tool, args, result,
# simulated tokens + latency), then print a readable trace and a summary.
# Numbers are simulated fixed values -- no wall-clock, no randomness.

class Tracer:
    def __init__(self):
        self.spans = []
    def record(self, kind, name, detail, sim_tokens, sim_ms):
        self.spans.append({'kind': kind, 'name': name, 'detail': detail,
                           'tokens': sim_tokens, 'ms': sim_ms})
    def summary(self):
        steps = len(self.spans)
        tool_calls = sum(1 for s in self.spans if s['kind'] == 'tool')
        tokens = sum(s['tokens'] for s in self.spans)
        ms = sum(s['ms'] for s in self.spans)
        return {'steps': steps, 'tool_calls': tool_calls, 'sim_tokens': tokens, 'sim_ms': ms}
    def show(self):
        for i, s in enumerate(self.spans):
            print(f"  [{i}] {s['kind']:7} {s['name']:14} {s['detail']:32} "
                  f"tok={s['tokens']:>3} ms={s['ms']:>3}")

# Simulate one agent run over the tracer. A real agent calls tracer.record()
# at each step of its loop; here we replay a fixed, deterministic trajectory.
def run_traced():
    t = Tracer()
    t.record('thought', 'plan',   'need city then its population',   40, 20)
    t.record('tool',    'geocode', "args={'city':'Paris'}",          25, 60)
    t.record('thought', 'reflect', 'have coords, now population',     35, 20)
    t.record('tool',    'lookup',  "args={'id':'FR-75'}",            25, 55)
    t.record('answer',  'final',   'Paris population is 2,100,000',   50, 25)
    return t

tracer = run_traced()
print('TRACE:')
tracer.show()
print('SUMMARY:', tracer.summary())

trajectory.py# Evaluate a whole TRAJECTORY, not just the final answer. We score three
# things: did it use the expected tools, did it reach the goal, did it stay
# in budget? Each check is a deterministic rule over a recorded trace.

def evaluate(trace, expected_tools, goal_marker, max_steps, max_tool_calls):
    used_tools = [s['tool'] for s in trace if s['kind'] == 'tool']
    tool_calls = len(used_tools)
    steps = len(trace)
    reached = any(goal_marker in s.get('text', '') for s in trace if s['kind'] == 'answer')

    reasons = []
    # 1. right tool set (order-independent: did it use exactly what was expected?)
    tools_ok = set(used_tools) == set(expected_tools)
    reasons.append(('expected_tools', tools_ok,
                    f'used {used_tools}, expected {sorted(expected_tools)}'))
    # 2. reached the goal
    reasons.append(('goal_reached', reached, f'goal_marker={goal_marker!r}'))
    # 3. within budget
    budget_ok = steps <= max_steps and tool_calls <= max_tool_calls
    reasons.append(('within_budget', budget_ok,
                    f'{steps} steps / {tool_calls} tool calls '
                    f'(caps {max_steps}/{max_tool_calls})'))

    passed = all(ok for _, ok, _ in reasons)
    return passed, reasons

# A GOOD trajectory: uses both expected tools, reaches the goal, in budget.
good = [
    {'kind': 'tool',   'tool': 'geocode'},
    {'kind': 'tool',   'tool': 'lookup'},
    {'kind': 'answer', 'text': 'Paris population is 2,100,000 [done]'},
]
# A BAD trajectory: wrong tool, never reached the goal, over the tool cap.
bad = [
    {'kind': 'tool',   'tool': 'geocode'},
    {'kind': 'tool',   'tool': 'guess'},
    {'kind': 'tool',   'tool': 'guess'},
    {'kind': 'answer', 'text': 'I think it is somewhere around a lot'},
]

for label, tr in [('GOOD', good), ('BAD', bad)]:
    passed, reasons = evaluate(tr, expected_tools={'geocode', 'lookup'},
                               goal_marker='[done]', max_steps=5, max_tool_calls=2)
    print(f'{label}: {"PASS" if passed else "FAIL"}')
    for name, ok, detail in reasons:
        print(f'   {"ok " if ok else "X  "} {name:15} {detail}')