"""Synthetic Trace Generator. Generates 10,000+ agent traces with: - task type - model used - tool calls - context size - cost - latency - failure mode - final outcome - optimal cheaper alternative - recovery action - verifier need - escalation decision Includes traces with: - cheap model success - cheap model failure - frontier model unnecessary - tool overuse - tool underuse - retrieval overuse - verifier overuse - retry loops - cache breaks - false-DONE - successful meta-tool reuse - bad meta-tool reuse """ import uuid import random import json from datetime import datetime, timedelta from typing import List, Dict, Any, Optional from dataclasses import asdict from aco.trace_schema import ( AgentTrace, TraceStep, ModelCall, ToolCall, VerifierCall, TaskType, Outcome, FailureTag, ) class SyntheticTraceGenerator: """Generates diverse synthetic agent traces for training and benchmarking.""" MODEL_CONFIGS = { "tiny_local": {"tier": 1, "cost_input": 0.0001, "cost_output": 0.0002, "latency": 200, "strength": 0.3}, "cheap_cloud": {"tier": 2, "cost_input": 0.0005, "cost_output": 0.001, "latency": 500, "strength": 0.5}, "medium": {"tier": 3, "cost_input": 0.003, "cost_output": 0.006, "latency": 800, "strength": 0.75}, "frontier": {"tier": 4, "cost_input": 0.01, "cost_output": 0.03, "latency": 1500, "strength": 0.95}, "specialist": {"tier": 5, "cost_input": 0.015, "cost_output": 0.045, "latency": 2000, "strength": 0.98}, } TOOL_COSTS = { "search": 0.002, "retrieve": 0.001, "fetch": 0.003, "code_execution": 0.005, "linter": 0.001, "test_runner": 0.003, "file_read": 0.0005, "file_write": 0.0005, "calculator": 0.0001, "database_query": 0.004, "compliance_check": 0.01, "summarize": 0.002, "task_planner": 0.001, "progress_tracker": 0.0005, } TASK_TYPE_DISTRIBUTION = { TaskType.QUICK_ANSWER: 0.20, TaskType.CODING: 0.20, TaskType.RESEARCH: 0.15, TaskType.DOCUMENT_DRAFTING: 0.10, TaskType.LEGAL_REGULATED: 0.05, TaskType.TOOL_HEAVY: 0.10, TaskType.RETRIEVAL_HEAVY: 0.10, TaskType.LONG_HORIZON: 0.08, TaskType.UNKNOWN_AMBIGUOUS: 0.02, } # Scenario templates for generating realistic traces SCENARIOS = [ # cheap model success {"name": "cheap_success", "prob": 0.15, "tier": [1, 2], "outcome": Outcome.SUCCESS, "failure_tags": []}, # cheap model failure -> should have escalated {"name": "cheap_failure", "prob": 0.10, "tier": [1, 2], "outcome": Outcome.FAILURE, "failure_tags": [FailureTag.MODEL_TOO_WEAK]}, # frontier model unnecessary -> overpaid {"name": "frontier_unnecessary", "prob": 0.08, "tier": [4], "outcome": Outcome.SUCCESS, "failure_tags": [], "optimal_tier": [1, 2]}, # tool overuse {"name": "tool_overuse", "prob": 0.07, "tier": [3, 4], "outcome": Outcome.PARTIAL_SUCCESS, "failure_tags": [FailureTag.TOOL_UNNECESSARY], "extra_tools": 3}, # tool underuse {"name": "tool_underuse", "prob": 0.05, "tier": [3, 4], "outcome": Outcome.FAILURE, "failure_tags": [FailureTag.TOOL_MISSED], "missing_tools": 2}, # retrieval overuse {"name": "retrieval_overuse", "prob": 0.04, "tier": [3, 4], "outcome": Outcome.SUCCESS, "failure_tags": [], "extra_retrievals": 5}, # verifier overuse {"name": "verifier_overuse", "prob": 0.03, "tier": [3, 4], "outcome": Outcome.SUCCESS, "failure_tags": [], "extra_verifiers": 2}, # retry loops {"name": "retry_loop", "prob": 0.05, "tier": [3, 4], "outcome": Outcome.FAILURE, "failure_tags": [FailureTag.RETRY_LOOP], "retries": 5}, # cache breaks {"name": "cache_break", "prob": 0.04, "tier": [3, 4], "outcome": Outcome.PARTIAL_SUCCESS, "failure_tags": [FailureTag.CACHE_BREAK]}, # false DONE {"name": "false_done", "prob": 0.05, "tier": [3, 4], "outcome": Outcome.FALSE_DONE, "failure_tags": [FailureTag.VERIFIER_FALSE_PASS]}, # meta-tool success {"name": "meta_tool_success", "prob": 0.06, "tier": [2, 3], "outcome": Outcome.SUCCESS, "failure_tags": [], "uses_meta_tool": True}, # meta-tool bad reuse {"name": "meta_tool_bad", "prob": 0.02, "tier": [2, 3], "outcome": Outcome.FAILURE, "failure_tags": [FailureTag.MODEL_TOO_WEAK], "uses_meta_tool": True}, # normal success {"name": "normal_success", "prob": 0.20, "tier": [3, 4], "outcome": Outcome.SUCCESS, "failure_tags": []}, # blocked {"name": "blocked", "prob": 0.03, "tier": [4], "outcome": Outcome.BLOCKED, "failure_tags": [FailureTag.MISSED_ESCALATION]}, # human escalation {"name": "human_escalation", "prob": 0.02, "tier": [4, 5], "outcome": Outcome.ESCALATED_HUMAN, "failure_tags": [FailureTag.MISSED_ESCALATION]}, # stopped doom {"name": "stopped_doom", "prob": 0.03, "tier": [3, 4], "outcome": Outcome.STOPPED_DOOM, "failure_tags": [FailureTag.COST_EXCEEDED]}, ] def __init__(self, seed: int = 42): self.rng = random.Random(seed) def generate(self, n: int = 10000) -> List[AgentTrace]: """Generate n synthetic traces.""" traces = [] for i in range(n): trace = self._generate_trace(i) traces.append(trace) return traces def _generate_trace(self, idx: int) -> AgentTrace: trace_id = f"synth_{idx}_{uuid.uuid4().hex[:8]}" # Pick task type task_type = self.rng.choices( list(self.TASK_TYPE_DISTRIBUTION.keys()), weights=list(self.TASK_TYPE_DISTRIBUTION.values()), )[0] # Pick scenario scenario = self._pick_scenario() # Generate user request based on task type user_request = self._generate_request(task_type, scenario["name"]) # Number of steps base_steps = self.rng.randint(1, 8) if scenario["name"] in ("retry_loop", "false_done"): base_steps = self.rng.randint(5, 12) if scenario.get("uses_meta_tool"): base_steps = max(2, base_steps // 2) # meta-tools compress steps steps = [] tier = self.rng.choice(scenario["tier"]) model_key = self._tier_to_model(tier) model_cfg = self.MODEL_CONFIGS[model_key] total_cost = 0.0 total_latency = 0.0 for step_idx in range(base_steps): step = self._generate_step( trace_id, step_idx, task_type, model_key, model_cfg, scenario, step_idx == base_steps - 1, ) steps.append(step) total_cost += step.step_cost total_latency += step.step_latency_ms # Generate final outcome outcome = scenario["outcome"] failure_tags = list(scenario["failure_tags"]) # Optimal cheaper alternative optimal_tier = scenario.get("optimal_tier") if optimal_tier: optimal_model = self._tier_to_model(self.rng.choice(optimal_tier)) optimal_cost = self.MODEL_CONFIGS[optimal_model]["cost_input"] * 2000 # rough estimate else: optimal_cost = total_cost * 0.6 # Compute cost saved vs always frontier frontier_cost = self.MODEL_CONFIGS["frontier"]["cost_input"] * 2000 * base_steps cost_saved = frontier_cost - total_cost return AgentTrace( trace_id=trace_id, user_request=user_request, task_type=task_type, steps=steps, final_outcome=outcome, failure_tags=failure_tags, total_cost=total_cost, total_cost_saved_vs_frontier=cost_saved, optimal_cost=optimal_cost, metadata={ "scenario": scenario["name"], "synthetic": True, "generation_timestamp": datetime.utcnow().isoformat(), "optimal_tier": optimal_tier[0] if optimal_tier else tier, }, ) def _pick_scenario(self) -> Dict: names = [s["name"] for s in self.SCENARIOS] probs = [s["prob"] for s in self.SCENARIOS] return self.rng.choices(self.SCENARIOS, weights=probs)[0] def _generate_request(self, task_type: TaskType, scenario: str) -> str: templates = { TaskType.QUICK_ANSWER: [ "What is the capital of France?", "Briefly explain quantum computing.", "Summarize the key points of article X.", "What is 237 * 452?", ], TaskType.CODING: [ "Write a Python function to reverse a linked list.", "Fix the bug in this React component.", "Refactor the authentication module to use JWT.", "Implement a LRU cache in Go.", ], TaskType.RESEARCH: [ "Research the latest advancements in transformer architectures.", "Find sources comparing LoRA and full fine-tuning.", "Investigate the climate impact of data centers.", "What does the literature say about speculative decoding?", ], TaskType.DOCUMENT_DRAFTING: [ "Draft a project proposal for the ML pipeline.", "Write an email to the team about the deployment schedule.", "Create a technical report on system performance.", ], TaskType.LEGAL_REGULATED: [ "Review this contract for liability clauses.", "Check compliance with GDPR for this data processing pipeline.", "Draft a privacy policy section for user data.", ], TaskType.TOOL_HEAVY: [ "Search for open issues in the repo and create a summary.", "Fetch the latest API documentation and generate client code.", "Query the database for Q3 sales and produce a chart.", ], TaskType.RETRIEVAL_HEAVY: [ "Answer based on the attached 50-page document.", "Find all mentions of 'payment processing' in my files.", "Retrieve relevant cases for this legal query.", ], TaskType.LONG_HORIZON: [ "Plan a 3-month roadmap for the agent framework.", "Orchestrate the deployment of the multi-region system.", "Project: redesign the data architecture end-to-end.", ], TaskType.UNKNOWN_AMBIGUOUS: [ "Help me with this thing.", "I need something done about the server.", "Can you look into that issue we discussed?", ], } options = templates.get(task_type, ["Generic request"]) return self.rng.choice(options) def _tier_to_model(self, tier: int) -> str: mapping = {1: "tiny_local", 2: "cheap_cloud", 3: "medium", 4: "frontier", 5: "specialist"} return mapping.get(tier, "medium") def _generate_step( self, trace_id: str, step_idx: int, task_type: TaskType, model_key: str, model_cfg: Dict, scenario: Dict, is_last: bool, ) -> TraceStep: step_id = f"{trace_id}_step_{step_idx}" # Model call input_tokens = self.rng.randint(500, 8000) output_tokens = self.rng.randint(100, 4000) cache_hit = self.rng.random() < 0.3 cache_hit_tokens = int(input_tokens * self.rng.random() * 0.5) if cache_hit else 0 model_call = ModelCall( model_id=model_key, provider="synthetic", input_tokens=input_tokens, output_tokens=output_tokens, reasoning_tokens=output_tokens // 5 if model_key == "frontier" else 0, cost_per_1k_input=model_cfg["cost_input"], cost_per_1k_output=model_cfg["cost_output"], cache_hit_input_tokens=cache_hit_tokens, latency_ms=model_cfg["latency"] * self.rng.uniform(0.8, 1.5), ) # Tool calls tool_calls = [] base_tools = self._get_tools_for_task(task_type) num_tools = self.rng.randint(0, len(base_tools)) if scenario.get("extra_tools"): num_tools += scenario["extra_tools"] if scenario.get("missing_tools"): num_tools = max(0, num_tools - scenario["missing_tools"]) for t in range(min(num_tools, len(base_tools))): tool_name = base_tools[t] tool_cost = self.TOOL_COSTS.get(tool_name, 0.001) repeated = self.rng.random() < 0.1 ignored = self.rng.random() < 0.05 failed = self.rng.random() < (0.2 if scenario["name"] in ("retry_loop", "tool_underuse") else 0.05) tool_calls.append(ToolCall( tool_name=tool_name, tool_input={"query": f"auto_{tool_name}"}, tool_cost=tool_cost, tool_latency_ms=self.rng.uniform(100, 1000), cache_hit=self.rng.random() < 0.2, repeated=repeated, ignored_result=ignored, failed=failed, )) # Verifier calls verifier_calls = [] num_verifiers = 0 if task_type in (TaskType.LEGAL_REGULATED, TaskType.CODING, TaskType.RESEARCH): num_verifiers = 1 if self.rng.random() < 0.5 else 0 if scenario.get("extra_verifiers"): num_verifiers += scenario["extra_verifiers"] for v in range(num_verifiers): passed = self.rng.random() < 0.8 verifier_calls.append(VerifierCall( verifier_model_id="verifier_medium", target_step_id=step_id, passed=passed, confidence=self.rng.uniform(0.6, 0.99), cost=0.005, latency_ms=500, )) # Context size context_size = self.rng.randint(1000, 15000) if scenario["name"] == "cache_break": context_size += self.rng.randint(5000, 20000) # excessive context # Retry count retries = 0 if scenario.get("retries"): retries = self.rng.randint(scenario["retries"] - 1, scenario["retries"] + 1) elif self.rng.random() < 0.15: retries = self.rng.randint(1, 2) # Recovery action recovery = None if retries > 0: recovery = self.rng.choice([ "retry_same", "retry_changed_prompt", "repair_tool", "retrieve_more_context", "switch_model", "ask_clarification", ]) # Outcome per step step_outcome = Outcome.SUCCESS if is_last: step_outcome = scenario["outcome"] elif scenario["name"] == "retry_loop" and step_idx >= 2: step_outcome = Outcome.FAILURE elif scenario["name"] == "false_done" and is_last: step_outcome = Outcome.FALSE_DONE return TraceStep( step_id=step_id, timestamp=datetime.utcnow() + timedelta(seconds=step_idx * 30), task_type=task_type, model_call=model_call, tool_calls=tool_calls, verifier_calls=verifier_calls, context_size_tokens=context_size, context_sources=["system_rules", "tool_descriptions", "user_preferences", "recent_messages"], retry_count=retries, recovery_action=recovery, artifacts_created=[f"artifact_{step_idx}"] if self.rng.random() < 0.3 else [], step_outcome=step_outcome, ) def _get_tools_for_task(self, task_type: TaskType) -> List[str]: mapping = { TaskType.QUICK_ANSWER: ["calculator", "search"], TaskType.CODING: ["file_read", "file_write", "code_execution", "linter", "test_runner"], TaskType.RESEARCH: ["search", "retrieve", "fetch", "summarize"], TaskType.DOCUMENT_DRAFTING: ["file_read", "summarize"], TaskType.LEGAL_REGULATED: ["document_retrieval", "compliance_check", "search"], TaskType.TOOL_HEAVY: ["search", "fetch", "api_call", "database_query"], TaskType.RETRIEVAL_HEAVY: ["retrieve", "search", "fetch"], TaskType.LONG_HORIZON: ["task_planner", "progress_tracker", "file_read"], TaskType.UNKNOWN_AMBIGUOUS: ["search"], } return mapping.get(task_type, ["search"]) def to_dicts(self, traces: List[AgentTrace]) -> List[Dict[str, Any]]: return [t.to_dict() for t in traces] def save(self, traces: List[AgentTrace], path: str) -> None: with open(path, "w") as f: for trace in traces: f.write(json.dumps(trace.to_dict()) + "\n") def load(self, path: str) -> List[Dict[str, Any]]: traces = [] with open(path, "r") as f: for line in f: traces.append(json.loads(line)) return traces