Identifying Slow and Expensive Steps

Performance Profiling for Agents

Timing Each Step

import time
from contextlib import contextmanager

@contextmanager
def timer(step_name: str, timings: dict):
    start = time.perf_counter()
    try:
        yield
    finally:
        end = time.perf_counter()
        duration_ms = (end - start) * 1000
        timings[step_name] = duration_ms
        print(f'{step_name}: {duration_ms:.1f}ms')

# Usage
timings = {}

with timer('entity_extraction', timings):
    time.sleep(0.05)  # Simulate work

with timer('vector_search', timings):
    time.sleep(0.12)  # Simulate work

with timer('llm_call', timings):
    time.sleep(0.80)  # Simulate LLM latency

print('\nTimings:', timings)
print('Slowest step:', max(timings, key=timings.get))

Building a Step Profiler

import time
from dataclasses import dataclass, field
from typing import List, Optional

@dataclass
class StepProfile:
    name: str
    start_ms: float
    end_ms: float
    duration_ms: float
    prompt_tokens: int = 0
    completion_tokens: int = 0
    cost_usd: float = 0.0
    error: Optional[str] = None

class StepProfiler:
    def __init__(self):
        self.steps: List[StepProfile] = []
        self.run_start = time.perf_counter()
    
    def start_step(self, name: str) -> float:
        return time.perf_counter()
    
    def end_step(self, name: str, start_time: float, tokens: dict = None, error: str = None):
        end = time.perf_counter()
        run_elapsed = (start_time - self.run_start) * 1000
        duration = (end - start_time) * 1000
        
        profile = StepProfile(
            name=name,
            start_ms=run_elapsed,
            end_ms=run_elapsed + duration,
            duration_ms=duration,
            error=error
        )
        if tokens:
            profile.prompt_tokens = tokens.get('prompt', 0)
            profile.completion_tokens = tokens.get('completion', 0)
        self.steps.append(profile)
        return profile

profiler = StepProfiler()
t = profiler.start_step('entity_extraction')
time.sleep(0.05)
profiler.end_step('entity_extraction', t, {'prompt': 200, 'completion': 50})
print('Step recorded:', profiler.steps[0].duration_ms)

Waterfall Chart in Terminal

class Step:
    def __init__(self, name, start_ms, end_ms, error=False):
        self.name, self.start_ms, self.end_ms, self.error = name, start_ms, end_ms, error
    @property
    def duration_ms(self):
        return self.end_ms - self.start_ms

class StepProfiler:
    def __init__(self, steps):
        self.steps = steps

def print_waterfall(profiler):
    if not profiler.steps:
        print('No steps recorded')
        return

    total_ms = max(s.end_ms for s in profiler.steps)
    bar_width = 50

    print('\n=== Agent Step Waterfall ===')
    for step in profiler.steps:
        start_pos = int(step.start_ms / total_ms * bar_width)
        end_pos = int(step.end_ms / total_ms * bar_width)
        bar = ' ' * start_pos + '#' * max(1, end_pos - start_pos) + ' ' * (bar_width - end_pos)
        status = 'ERR' if step.error else '   '
        print(f'{status} {step.name:<20} {step.duration_ms:>6.0f}ms  |{bar}|')

    print(f'\nTotal run: {total_ms:.0f}ms')

profiler = StepProfiler([Step('plan', 0, 120), Step('search', 120, 480), Step('generate', 480, 900, error=True)])
print_waterfall(profiler)

Collecting P95 Latency Across Runs

import statistics
from collections import defaultdict

class LatencyCollector:
    def __init__(self):
        self.step_durations = defaultdict(list)
    
    def record(self, step_name: str, duration_ms: float):
        self.step_durations[step_name].append(duration_ms)
    
    def percentile(self, data: list, pct: float) -> float:
        sorted_data = sorted(data)
        index = int(len(sorted_data) * pct / 100)
        return sorted_data[min(index, len(sorted_data) - 1)]
    
    def report(self):
        print('=== Latency Report (ms) ===')
        print(f'{"Step":<30} {"Count":>6} {"P50":>8} {"P95":>8} {"P99":>8} {"Max":>8}')
        print('-' * 75)
        for step_name, durations in sorted(self.step_durations.items()):
            p50 = self.percentile(durations, 50)
            p95 = self.percentile(durations, 95)
            p99 = self.percentile(durations, 99)
            max_d = max(durations)
            print(f'{step_name:<30} {len(durations):>6} {p50:>8.0f} {p95:>8.0f} {p99:>8.0f} {max_d:>8.0f}')

collector = LatencyCollector()
import random
for _ in range(100):
    collector.record('vector_search', random.gauss(120, 30))
    collector.record('llm_call', random.gauss(800, 150))
collector.report()

Identifying P95 Slow Steps

class LatencyCollector:
    def __init__(self, step_durations):
        self.step_durations = step_durations
    def percentile(self, durations, p):
        s = sorted(durations)
        k = int(len(s) * p / 100)
        return s[min(k, len(s) - 1)]

def find_optimization_targets(collector, p95_threshold_ms=500):
    targets = []
    for step_name, durations in collector.step_durations.items():
        p95 = collector.percentile(durations, 95)
        avg = sum(durations) / len(durations)
        p95_to_avg_ratio = p95 / avg if avg > 0 else 0

        target = {
            'step': step_name,
            'p95_ms': round(p95),
            'avg_ms': round(avg),
            'p95_to_avg_ratio': round(p95_to_avg_ratio, 2),
            'call_count': len(durations),
            'needs_optimization': p95 > p95_threshold_ms
        }

        if target['needs_optimization']:
            if p95_to_avg_ratio > 2.0:
                target['suggestion'] = 'High variance: consider timeout and retry or caching'
            else:
                target['suggestion'] = 'Consistently slow: consider parallel execution or faster model'

        targets.append(target)

    targets.sort(key=lambda x: x['p95_ms'], reverse=True)
    return targets

collector = LatencyCollector({'search': [100, 150, 900], 'generate': [400, 420, 410]})
targets = find_optimization_targets(collector, p95_threshold_ms=500)
for t in targets:
    print(f"{t['step']}: P95={t['p95_ms']}ms, Avg={t['avg_ms']}ms, Action: {t.get('suggestion', 'OK')}")

Caching Expensive Tool Results

import hashlib
import json
import time
from typing import Callable, Any

class ToolResultCache:
    def __init__(self, ttl_seconds: int = 300):
        self.cache = {}
        self.ttl = ttl_seconds
    
    def _make_key(self, tool_name: str, args: dict) -> str:
        content = json.dumps({'tool': tool_name, 'args': args}, sort_keys=True)
        return hashlib.sha256(content.encode()).hexdigest()[:16]
    
    def get_or_compute(self, tool_name: str, args: dict, compute_fn: Callable) -> Any:
        key = self._make_key(tool_name, args)
        now = time.time()
        
        if key in self.cache:
            entry = self.cache[key]
            if now - entry['ts'] < self.ttl:
                print(f'Cache HIT for {tool_name}')
                return entry['result']
        
        print(f'Cache MISS for {tool_name} - computing...')
        start = time.perf_counter()
        result = compute_fn(**args)
        elapsed = (time.perf_counter() - start) * 1000
        print(f'{tool_name} computed in {elapsed:.0f}ms')
        
        self.cache[key] = {'result': result, 'ts': now}
        return result

cache = ToolResultCache(ttl_seconds=60)

def expensive_web_search(query: str) -> list:
    time.sleep(0.2)  # Simulate slow API call
    return [f'Result for: {query}']

result1 = cache.get_or_compute('web_search', {'query': 'AI news'}, expensive_web_search)
result2 = cache.get_or_compute('web_search', {'query': 'AI news'}, expensive_web_search)  # Cache hit

Detecting Token-Heavy Steps

def find_token_heavy_steps(tracker: 'CostTracker', token_threshold: int = 2000) -> list:
    heavy_steps = []
    for step in tracker.steps:
        total_tokens = step.prompt_tokens + step.completion_tokens
        if total_tokens > token_threshold:
            heavy_steps.append({
                'step': step.step_name,
                'total_tokens': total_tokens,
                'prompt_tokens': step.prompt_tokens,
                'completion_tokens': step.completion_tokens,
                'cost_usd': step.cost_usd,
                'suggestions': []
            })
            entry = heavy_steps[-1]
            if step.prompt_tokens > token_threshold * 0.9:
                entry['suggestions'].append('Prompt is very large: truncate context documents or summarize')
            if step.completion_tokens > 1000:
                entry['suggestions'].append('Large output: use max_tokens limit if full response not needed')
    
    heavy_steps.sort(key=lambda x: x['total_tokens'], reverse=True)
    return heavy_steps

tracker = CostTracker()
tracker.record('answer_generation', 'gpt-4o-mini', 4500, 1200)
tracker.record('entity_extraction', 'gpt-4o-mini', 200, 50)
heavy = find_token_heavy_steps(tracker)
for s in heavy:
    print(f"{s['step']}: {s['total_tokens']} tokens, suggestions: {s['suggestions']}")

Model Downgrade Analysis

MODEL_TIERS = {
    'gpt-4o': {'tier': 'premium', 'capabilities': ['complex reasoning', 'nuanced writing']},
    'gpt-4o-mini': {'tier': 'standard', 'capabilities': ['extraction', 'classification', 'summarization']},
    'claude-3-haiku-20240307': {'tier': 'fast', 'capabilities': ['simple tasks', 'routing']}
}

STEP_MODEL_RECOMMENDATIONS = {
    'entity_extraction': 'gpt-4o-mini',
    'intent_classification': 'gpt-4o-mini',
    'simple_summarization': 'gpt-4o-mini',
    'complex_reasoning': 'gpt-4o',
    'final_answer_generation': 'gpt-4o-mini'
}

def audit_model_usage(tracker: 'CostTracker') -> list:
    recommendations = []
    for step in tracker.steps:
        recommended = STEP_MODEL_RECOMMENDATIONS.get(step.step_name)
        if recommended and recommended != step.model:
            current_cost = step.cost_usd
            # Estimate cost with recommended model (rough calculation)
            recommendations.append({
                'step': step.step_name,
                'current_model': step.model,
                'recommended_model': recommended,
                'potential_savings': 'significant' if step.model == 'gpt-4o' else 'moderate'
            })
    return recommendations

print('Model downgrade analysis function defined')

Profiling in Production

import random

class SampledProfiler:
    def __init__(self, sample_rate: float = 0.1):
        self.sample_rate = sample_rate
        self.full_profiles = []
        self.aggregate_timings = defaultdict(list)
    
    def should_profile_full(self) -> bool:
        return random.random() < self.sample_rate
    
    def record_run(self, profiler: 'StepProfiler', full_profile: bool):
        # Always record aggregate timing
        for step in profiler.steps:
            self.aggregate_timings[step.name].append(step.duration_ms)
        
        # Only store full profiles for sampled runs
        if full_profile:
            self.full_profiles.append(profiler.steps)
    
    def get_summary(self) -> dict:
        return {
            'full_profiles_stored': len(self.full_profiles),
            'steps_tracked': {k: len(v) for k, v in self.aggregate_timings.items()}
        }

sampled = SampledProfiler(sample_rate=0.1)
print('Sampled profiler: recording 10% of runs in full detail')
print('Summary:', sampled.get_summary())

Combining Latency and Cost Signals

def combined_optimization_priority(profiler: 'StepProfiler', tracker: 'CostTracker') -> list:
    # Build combined step data
    cost_by_step = {s.step_name: s.cost_usd for s in tracker.steps}
    
    combined = []
    for step in profiler.steps:
        cost = cost_by_step.get(step.name, 0)
        # Priority score: normalize and combine
        # High latency + High cost = top priority
        priority = (step.duration_ms / 1000) + (cost * 1000)  # Rough normalization
        combined.append({
            'step': step.name,
            'duration_ms': round(step.duration_ms),
            'cost_usd': round(cost, 6),
            'priority_score': round(priority, 3)
        })
    
    combined.sort(key=lambda x: x['priority_score'], reverse=True)
    return combined

print('Combined optimization priority function defined')
print('High priority = slow + expensive')

Knowledge Check: Performance Profiling

Performance Profiling Summary