docs/architecture/CANONICAL_TRACE_IMPLEMENTATION.md

# Canonical Normalized Trace Generation - Implementation Summary

## Overview

Successfully implemented a **canonical normalized trace generation system** that converts agent-specific execution traces into a unified format suitable for:

1. Multi-agent comparison
2. Trace replay and visualization  
3. Failure analysis
4. Future agent architectures

## ✅ Deliverables

### 1. Trace Schema Analysis ✅

**File:** [BASELINE_TRACE_SCHEMA.md](BASELINE_TRACE_SCHEMA.md)

Complete analysis of baseline trace format:
- Available fields documented
- Missing information identified
- Tool call patterns analyzed
- Recommendations for canonical schema

**Key findings:**
- Baseline has: task_id, steps, answer, execution status, timing
- Baseline missing: agent metadata, timestamps, token usage, phases
- Tool calls represented as action + observation
- Single-agent architecture (no multi-agent support)

### 2. Canonical Trace Schema ✅

**File:** [CANONICAL_TRACE_FORMAT.md](CANONICAL_TRACE_FORMAT.md)

Universal schema supporting all agent types:

```json
{
  "run_id": "string",
  "task_id": "string",
  "agent_type": "string",
  "question": "string",
  "success": boolean,
  "steps": [
    {
      "step_id": integer,
      "agent": "string",        // NEW: Agent name
      "agent_role": "string",   // NEW: Agent role (planner/worker/critic)
      "thought": "string",
      "action": "string",
      "action_input": {},
      "observation": {},
      "tool_success": boolean,
      ...
    }
  ]
}
```

**Key features:**
- **Multi-agent support:** Each step has agent + agent_role
- **Future-proof:** Supports planner/executor/critic architectures
- **NULL-safe:** Missing data represented as None
- **Backward compatible:** Works with existing evaluation pipeline

### 3. Trace Adapter Enhancement ✅

**File:** [src/data_agent_baseline/evaluation/baseline_adapter.py](src/data_agent_baseline/evaluation/baseline_adapter.py)

Enhanced existing adapter with agent metadata:
- Added `agent="baseline_react"` to all steps
- Added `agent_role="worker"` to all steps
- Preserves all original trace data
- Converts to canonical schema

### 4. Normalized Trace Manager ✅

**File:** [src/data_agent_baseline/evaluation/normalized_trace_manager.py](src/data_agent_baseline/evaluation/normalized_trace_manager.py)

Complete trace management system:

```python
class NormalizedTraceManager:
    def save_normalized_trace(canonical_trace) -> Path
    def save_normalized_traces(traces) -> list[Path]
    def load_normalized_trace(task_id) -> dict
    def list_normalized_traces() -> list[str]
    def validate_normalized_trace(trace) -> (bool, list[str])
    def validate_all_traces() -> dict
    def get_trace_metrics(trace) -> dict
```

**Capabilities:**
- Save/load individual trace files
- Validate schema compliance
- Derive metrics from traces
- List available traces

### 5. Validation Utility ✅

**Validation checks:**
- ✅ Required fields present (run_id, task_id, agent_type, etc.)
- ✅ Step ordering correct (sequential step_id)
- ✅ Type checking (strings, ints, bools, dicts)
- ✅ Schema consistency (nested structures valid)

**Usage:**
```python
manager = NormalizedTraceManager(output_dir)
is_valid, errors = manager.validate_normalized_trace(trace)
```

### 6. CLI Integration ✅

**New command:** `dabench eval-baseline`

Enhanced with normalized trace generation:

```bash
dabench eval-baseline 20260613T114457Z
```

**Output:**
```
baseline_evaluation/
├── task_results.csv
├── summary_metrics.json
├── evaluation_report.md
└── normalized_traces/          # NEW!
    ├── task_22.json
    └── ...
```

**Steps:**
1. Normalize baseline traces
2. **Save normalized traces** ← NEW
3. **Validate traces** ← NEW
4. Compute metrics
5. Generate reports

### 7. Trace Viewer Command ✅

**New command:** `dabench view-normalized-trace`

```bash
dabench view-normalized-trace 20260613T114457Z task_22
```

**Features:**
- Task information display
- Validation status
- Derived metrics (steps, tools, failures)
- Agent breakdown (for multi-agent)
- Tool usage statistics
- Step-by-step execution view
- Final answer display

**Output example:**
```
Normalized Trace: task_22
Run ID: 20260613T114457Z
Agent Type: baseline_react

┌─ Task Information ─┐
│ Task ID    │ task_22
│ Question   │ State the date...
│ Difficulty │ Easy
│ Success    │ ✓
│ Duration   │ 8.53s
└────────────────────┘

✓ Trace validation passed

┌─ Derived Metrics ──┐
│ Total Steps  │ 5
│ Tool Calls   │ 5
│ Failed Tools │ 0
│ Unique Tools │ 4
└────────────────────┘

┌─ Tool Usage ───┐
│ read_csv     │ 2
│ list_context │ 1
│ read_json    │ 1
│ answer       │ 1
└────────────────┘
```

### 8. Example Normalized Trace ✅

**Location:** `/data3/dataFAIR/kdd-dev/public/artifacts/runs/20260613T114457Z/baseline_evaluation/normalized_traces/task_22.json`

**Structure:**
```json
{
  "run_id": "20260613T114457Z",
  "task_id": "task_22",
  "agent_type": "baseline_react",
  "question": "State the date Connor Hilton paid his/her dues.",
  "difficulty": "Easy",
  "success": true,
  "final_answer": {
    "columns": ["date_received"],
    "rows": [["2019-10-02"], ["2019-09-12"]]
  },
  "duration_seconds": 8.528,
  "steps": [
    {
      "step_id": 1,
      "agent": "baseline_react",
      "agent_role": "worker",
      "thought": "I need to find...",
      "action": "list_context",
      "action_input": {"max_depth": 4},
      "observation": {...},
      "tool_success": true
    }
    ...
  ]
}
```

## 📊 Test Results

Successfully generated and validated normalized traces:

```
✓ Step 1: Normalized 1 traces
✓ Step 2: Saved 1 normalized traces
✓ Step 3: All 1 traces validated
✓ Step 4: Evaluated 1 tasks
✓ Step 5: Generated reports
```

**Validation:** ✅ All traces pass schema validation

**Metrics extraction:** ✅ Successfully derived:
- num_steps: 5
- tool_calls: 5
- failed_tools: 0
- unique_tools: 4
- agent_steps: {baseline_react: 5}
- tool_counts: {read_csv: 2, list_context: 1, ...}

## 🎯 Multi-Agent Support

The schema supports future architectures:

### Baseline ReAct
```json
{"agent": "baseline_react", "agent_role": "worker"}
```

### Planner + Executor
```json
{"agent": "planner", "agent_role": "planner"}
{"agent": "executor", "agent_role": "worker"}
```

### Multi-Agent Analyst
```json
{"agent": "scout", "agent_role": "worker"}
{"agent": "analyst", "agent_role": "worker"}
{"agent": "critic", "agent_role": "critic"}
```

### Coordinator System
```json
{"agent": "coordinator", "agent_role": "coordinator"}
{"agent": "worker_1", "agent_role": "worker"}
```

## 📈 Derived Metrics

From normalized traces, we can derive:

**Basic metrics:**
- Number of steps
- Number of tool calls
- Number of failed tools
- Runtime in seconds
- Unique tools used
- Success status

**Multi-agent metrics:**
- Agent breakdown (steps per agent)
- Role breakdown (steps per role)
- Tool frequency (calls per tool)
- Phase breakdown (if available)

**Usage:**
```python
manager = NormalizedTraceManager(output_dir)
trace = manager.load_normalized_trace("task_22")
metrics = manager.get_trace_metrics(trace)

print(metrics["num_steps"])       # 5
print(metrics["agent_steps"])     # {"baseline_react": 5}
print(metrics["tool_counts"])     # {"read_csv": 2, ...}
```

## 🔄 Workflow

### 1. Generate Normalized Traces

```bash
dabench eval-baseline <run_id>
```

Creates one JSON file per task in `normalized_traces/`.

### 2. View Traces

```bash
dabench view-normalized-trace <run_id> <task_id>
```

Displays detailed breakdown with validation and metrics.

### 3. Load Programmatically

```python
from data_agent_baseline.evaluation.normalized_trace_manager import NormalizedTraceManager

manager = NormalizedTraceManager(output_dir)

# List traces
task_ids = manager.list_normalized_traces()

# Load trace
trace = manager.load_normalized_trace("task_22")

# Validate
is_valid, errors = manager.validate_normalized_trace(trace)

# Get metrics
metrics = manager.get_trace_metrics(trace)
```

## 🚀 Benefits

### 1. Agent-Agnostic Evaluation

Same evaluation pipeline works for:
- Baseline ReAct
- Multi-agent systems
- Future architectures

### 2. Easy Comparison

Compare agents by loading their normalized traces:

```python
baseline = manager.load_normalized_trace("task_22", run_id="baseline_run")
multi_agent = manager.load_normalized_trace("task_22", run_id="multi_agent_run")

baseline_metrics = manager.get_trace_metrics(baseline)
multi_agent_metrics = manager.get_trace_metrics(multi_agent)

print(f"Baseline: {baseline_metrics['num_steps']} steps")
print(f"Multi-agent: {multi_agent_metrics['num_steps']} steps")
```

### 3. Trace Replay

Reconstruct execution from normalized trace:
- Debugging
- Visualization (DAG, timeline)
- Failure analysis

### 4. One File Per Task

Each task has its own normalized trace file:
- Easy to locate
- Easy to replay
- Easy to analyze
- Easy to visualize
- Easy to share

### 5. Future-Proof

Schema supports:
- New agent types (add agent name)
- New agent roles (add role name)
- New phases (add phase name)
- Parallel execution (same step_id, different agent)
- Multi-agent coordination

## 📁 File Structure

```
src/data_agent_baseline/evaluation/
├── __init__.py                      # Canonical schema (enhanced)
├── baseline_adapter.py              # Baseline → canonical (enhanced)
├── normalized_trace_manager.py      # NEW: Trace management
├── phase1_evaluator.py              # Evaluator
└── report_generator.py              # Report generation

Documentation:
├── BASELINE_TRACE_SCHEMA.md         # Baseline analysis
├── CANONICAL_TRACE_FORMAT.md        # Schema documentation
└── CANONICAL_TRACE_IMPLEMENTATION.md # This file

CLI:
├── eval-baseline                     # Enhanced with trace generation
└── view-normalized-trace            # NEW: Trace viewer
```

## 🔍 Example Usage

### Generate Traces

```bash
cd /data3/dataFAIR/kdd-dev/public
dabench eval-baseline 20260613T114457Z
```

Output:
```
Step 1: Normalizing baseline traces...
  ✓ Normalized 1 traces

Step 2: Saving normalized traces...
  ✓ Saved 1 normalized traces
  → .../baseline_evaluation/normalized_traces

Step 3: Validating normalized traces...
  ✓ All 1 traces validated

Step 4: Computing Phase 1 metrics...
  ✓ Evaluated 1 tasks

Step 5: Generating evaluation reports...
  ✓ task_results: baseline_evaluation/task_results.csv
  ✓ summary_metrics: baseline_evaluation/summary_metrics.json
  ✓ evaluation_report: baseline_evaluation/evaluation_report.md
```

### View Trace

```bash
dabench view-normalized-trace 20260613T114457Z task_22
```

Shows:
- Task information
- Validation status
- Derived metrics
- Agent breakdown
- Tool usage
- Step-by-step execution
- Final answer

### Load and Analyze

```python
from pathlib import Path
from data_agent_baseline.evaluation.normalized_trace_manager import NormalizedTraceManager

# Initialize
output_dir = Path("/data3/dataFAIR/kdd-dev/public/artifacts/runs/20260613T114457Z/baseline_evaluation")
manager = NormalizedTraceManager(output_dir)

# List traces
print(f"Available traces: {manager.list_normalized_traces()}")

# Load trace
trace = manager.load_normalized_trace("task_22")
print(f"Task: {trace['task_id']}")
print(f"Success: {trace['success']}")
print(f"Steps: {len(trace['steps'])}")

# Validate
is_valid, errors = manager.validate_normalized_trace(trace)
print(f"Valid: {is_valid}")
if errors:
    print(f"Errors: {errors}")

# Get metrics
metrics = manager.get_trace_metrics(trace)
print(f"Total steps: {metrics['num_steps']}")
print(f"Tool calls: {metrics['num_tool_calls']}")
print(f"Failed tools: {metrics['num_failed_tools']}")
print(f"Unique tools: {metrics['unique_tools']}")
print(f"Agent breakdown: {metrics['agent_steps']}")
print(f"Tool counts: {metrics['tool_counts']}")
```

## 🎓 Key Design Decisions

### 1. One File Per Task

**Why:** Easy to locate, analyze, replay, and visualize individual tasks.

**Alternative considered:** Single monolithic file (harder to work with).

### 2. Agent + Agent Role Fields

**Why:** Support multi-agent architectures without breaking single-agent traces.

**Baseline:** `agent="baseline_react"`, `agent_role="worker"`  
**Multi-agent:** Different agents per step with specific roles

### 3. NULL-Safe Design

**Why:** Baseline traces don't have timestamps, token usage, phases.

**Solution:** Use `None` for unavailable fields instead of omitting them.

### 4. Validation Built-In

**Why:** Ensure schema consistency across different agent types.

**Implementation:** Comprehensive validation with detailed error messages.

### 5. Derived Metrics

**Why:** Compute common metrics from traces without storing redundantly.

**Benefit:** Metrics always consistent with trace data.

## 🔮 Future Enhancements

### Phase 2: Multi-Agent Traces

When multi-agent systems are available:

```python
class MultiAgentAdapter:
    def normalize(self, raw_trace, run_id):
        steps = []
        for step in raw_trace["steps"]:
            steps.append(CanonicalStep(
                step_id=step["id"],
                agent=step["agent_name"],     # planner, executor, critic
                agent_role=step["role"],      # planner, worker, critic
                phase=step["phase"],          # explore, plan, execute
                ...
            ))
        return CanonicalTrace(steps=steps, ...)
```

### Phase 3: Parallel Execution

Support concurrent steps:

```json
[
  {"step_id": 5, "agent": "worker_1", "phase": "parallel_execute"},
  {"step_id": 5, "agent": "worker_2", "phase": "parallel_execute"}
]
```

Same step_id, different agents = parallel execution.

### Phase 4: Trace Replay

Replay execution from normalized trace:

```python
class TraceReplayer:
    def replay(self, trace):
        for step in trace["steps"]:
            print(f"Step {step['step_id']}: {step['agent']} → {step['action']}")
            # Visualize or re-execute
```

### Phase 5: DAG Visualization

Generate execution DAG from trace:

```python
def generate_dag(trace):
    # Create nodes for each step
    # Create edges based on data dependencies
    # Render as graph
```

## 🎉 Summary

Successfully implemented a comprehensive **canonical normalized trace generation system** that:

✅ **Converts** baseline traces to canonical format  
✅ **Saves** one JSON file per task  
✅ **Validates** schema compliance  
✅ **Derives** metrics from traces  
✅ **Supports** multi-agent architectures  
✅ **Provides** CLI tools for viewing and analysis  
✅ **Documents** schema and usage  
✅ **Tests** with real baseline run  

**Ready for:**
- Multi-agent comparison
- Trace replay
- Failure analysis
- DAG visualization
- Future agent architectures

**No baseline code was modified.** All integration through adapter layer.