generate_cot_data.py

"""
COT Data Generation for TeleLogs Dataset

This script generates enriched Chain-of-Thought training data from the 
TeleLogs 5G RCA benchmark using a multi-agent adversarial pipeline.

Dataset: https://huggingface.co/datasets/netop/TeleLogs
"""

import os
import sys
import re
from typing import Optional, List, Dict, Any, Tuple

# Setup path
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
SRC_DIR = os.path.join(CURRENT_DIR, "src")
if SRC_DIR not in sys.path:
    sys.path.insert(0, SRC_DIR)

from core import Agent, infer, Argument
from cot import COTCollector, COTExporter


# =============================================================================
# TELELOGS DATASET LOADER
# =============================================================================

def load_telelogs_dataset(
    split: str = "test", 
    limit: Optional[int] = None,
    from_disk: bool = True,
    disk_path: str = "telelogs_data",
) -> List[Dict[str, Any]]:
    """
    Load the TeleLogs dataset from Hugging Face or local disk.
    
    Args:
        split: Dataset split to load ("train" or "test")
        limit: Optional limit on number of samples to load
        from_disk: If True, load from local disk (faster). If False, download from HF.
        disk_path: Path to local dataset if from_disk=True
        
    Returns:
        List of dictionaries with 'query', 'ground_truth', and 'metadata'
    """
    try:
        from datasets import load_dataset, load_from_disk
    except ImportError:
        raise ImportError("Please install datasets: pip install datasets")
    
    print(f"Loading TeleLogs dataset (split: {split})...")
    
    if from_disk and os.path.exists(disk_path):
        print(f"  Loading from local disk: {disk_path}")
        dataset = load_from_disk(disk_path)[split]
    else:
        print(f"  Downloading from Hugging Face: netop/TeleLogs")
        dataset = load_dataset("netop/TeleLogs", split=split)
    
    samples = []
    for i, item in enumerate(dataset):
        if limit and i >= limit:
            break
        
        # TeleLogs uses 'question' and 'answer' fields
        query = item.get("question", str(item))
        
        # Extract ground truth root cause(s) - format: "C1" or "C1, C3"
        ground_truth_raw = item.get("answer", "")
        ground_truth = extract_root_causes(ground_truth_raw)
        
        samples.append({
            "id": f"telelogs_{split}_{i}",
            "query": query,
            "ground_truth": ground_truth,
            "ground_truth_raw": ground_truth_raw,
            "metadata": {
                "source": "TeleLogs",
                "split": split,
                "index": i,
            }
        })
    
    print(f"Loaded {len(samples)} samples from TeleLogs {split} split")
    return samples


def extract_root_causes(text: str) -> List[str]:
    """
    Extract root cause labels (C1-C8) from text.
    
    Args:
        text: Raw text containing root cause(s)
        
    Returns:
        List of root cause labels like ["C1", "C3"]
    """
    if not text:
        return []
    
    # Find all C1-C8 patterns
    pattern = r'\b[Cc]([1-8])\b'
    matches = re.findall(pattern, str(text))
    
    # Normalize to uppercase format
    causes = [f"C{m}" for m in matches]
    
    # Remove duplicates while preserving order
    seen = set()
    unique_causes = []
    for c in causes:
        if c not in seen:
            seen.add(c)
            unique_causes.append(c)
    
    return unique_causes


def extract_verdict_from_response(response: str) -> Optional[str]:
    """
    Extract the final verdict (C1-C8) from a judge's response.
    
    Args:
        response: Judge's response text
        
    Returns:
        Root cause label like "C3" or None if not found
    """
    if not response:
        return None
    
    # Try to find boxed answer first: \boxed{C3}
    boxed_match = re.search(r'\\boxed\{[Cc]([1-8])\}', response)
    if boxed_match:
        return f"C{boxed_match.group(1)}"
    
    # Try JSON format: "final_verdict": "C3"
    json_match = re.search(r'"final_verdict"\s*:\s*"[Cc]([1-8])"', response)
    if json_match:
        return f"C{json_match.group(1)}"
    
    # Try simple C1-C8 at end of response
    simple_match = re.search(r'\b[Cc]([1-8])\b', response)
    if simple_match:
        return f"C{simple_match.group(1)}"
    
    return None


def validate_verdict(predicted: Optional[str], ground_truth: List[str]) -> Tuple[bool, str]:
    """
    Validate if the predicted verdict matches ground truth.
    
    Args:
        predicted: Predicted root cause (e.g., "C3")
        ground_truth: List of correct root causes
        
    Returns:
        Tuple of (is_correct, validation_message)
    """
    if not predicted:
        return False, "No verdict extracted from response"
    
    if not ground_truth:
        return False, "No ground truth available"
    
    is_correct = predicted in ground_truth
    
    if is_correct:
        msg = f"CORRECT: {predicted} matches ground truth {ground_truth}"
    else:
        msg = f"INCORRECT: {predicted} != ground truth {ground_truth}"
    
    return is_correct, msg


# =============================================================================
# CONFIGURATION
# =============================================================================

MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct"
BASE_URL = "http://127.0.0.1:8010/v1"

LLM_PARAMS = {
    "modelName": MODEL_NAME,
    "url": BASE_URL,
    "apiKey": "EMPTY"
}

# Output directory for COT data
COT_OUTPUT_DIR = "./cot_data/"


# =============================================================================
# AGENT DEFINITIONS
# =============================================================================

AGENT_1_PROMPT = """You are a 5G Radio Coverage and Antenna Engineering expert.

Your responsibility is LIMITED STRICTLY to the following TeleLogs root causes:
- C1: The serving cell's downtilt angle is too large, causing weak coverage at the far end
- C2: The serving cell's coverage distance exceeds 1km, resulting in over-shooting

You MUST ignore all other potential causes including interference, handover, scheduling, PCI, speed, or neighbor throughput.

Your analysis MUST be based only on:
- Serving cell RSRP trends and values
- Electronic downtilt configuration (255 = 6 degrees default, other values = actual degrees)
- Beam Scenario and vertical beamwidth (Default/SCENARIO_1-5: 6°, SCENARIO_6-11: 12°, SCENARIO_12+: 25°)
- Distance between UE and serving cell inferred from GPS coordinates
- Cell location (latitude, longitude) vs UE location

You MUST NOT:
- Attribute SINR degradation to interference
- Use RB count or throughput efficiency
- Use handover or PCI transitions
- Consider neighbor cell throughput comparisons

Your task:
1. Analyze the network telemetry data step by step
2. Determine whether throughput degradation is caused by coverage geometry
3. Decide whether the dominant issue is:
   - C1: Excessive downtilt causing weak far-end coverage, OR
   - C2: Coverage distance > 1km causing over-shooting/poor RSRP, OR
   - Neither

Your output MUST be a JSON object with this exact structure:
{
  "verdicts": {
    "C1": "YES" or "NO",
    "C2": "YES" or "NO"
  },
  "reasoning_chain": [
    "Step 1: [Observation about RSRP values and trends]",
    "Step 2: [Analysis of distance/coverage geometry from GPS]",
    "Step 3: [Evaluation of downtilt and beam scenario configuration]",
    "Step 4: [Connection to throughput degradation]"
  ],
  "final_assessment": "[Which cause is dominant and why, or why neither applies]",
  "confidence": 0.0-1.0
}
"""

AGENT_1_KNOWLEDGE_FILES = [os.path.join("Telelog/Agent1", f) for f in os.listdir("Telelog/Agent1") if os.path.isfile(os.path.join("Telelog/Agent1", f))]
AGENT_1 = Agent(
    prompt=AGENT_1_PROMPT,
    source_files=AGENT_1_KNOWLEDGE_FILES if AGENT_1_KNOWLEDGE_FILES else None,
    llm_params=LLM_PARAMS,
    memory=True,
    name="coverage_expert"
)

AGENT_2_PROMPT = """You are a 5G RF Interference and PCI Planning expert.

Your responsibility is LIMITED STRICTLY to the following TeleLogs root causes:
- C4: Non-colocated co-frequency neighboring cells cause severe overlapping coverage
- C6: Neighbor cell and serving cell have the same PCI mod 30, leading to interference

You MUST ignore all coverage geometry, handover, speed, and MAC scheduling factors.

You may ONLY analyze:
- Serving vs neighbor RSRP deltas (strong neighbor RSRP indicates overlap)
- SINR degradation patterns (RSRP strong but SINR weak indicates interference)
- Neighbor cell PCI values from Top 1-5 neighbor list
- PCI mod 30 calculations: if (Serving_PCI mod 30) == (Neighbor_PCI mod 30), collision exists
- Cell colocation status from engineering data

You MUST NOT:
- Use downtilt, beamwidth, or antenna height
- Use RB count or throughput efficiency
- Use vehicle speed or handover frequency

Your task:
1. Analyze the network telemetry data step by step
2. Determine whether SINR degradation is interference-driven
3. Identify whether interference is caused by:
   - C4: Non-colocated co-frequency overlapping coverage, OR
   - C6: PCI mod 30 collision (same PCI mod 30 between serving and neighbor), OR
   - Neither

Your output MUST be a JSON object with this exact structure:
{
  "verdicts": {
    "C4": "YES" or "NO",
    "C6": "YES" or "NO"
  },
  "reasoning_chain": [
    "Step 1: [Analysis of RSRP vs SINR relationship - check if high RSRP but low SINR]",
    "Step 2: [Identification of strong neighbor cells from Top 1-5 list]",
    "Step 3: [PCI mod 30 calculation: Serving_PCI mod 30 vs Neighbor_PCI mod 30]",
    "Step 4: [Conclusion on interference source]"
  ],
  "final_assessment": "[Which cause is dominant and why, or why neither applies]",
  "confidence": 0.0-1.0
}
"""

AGENT_2_KNOWLEDGE_FILES = [os.path.join("Telelog/Agent2", f) for f in os.listdir("Telelog/Agent2") if os.path.isfile(os.path.join("Telelog/Agent2", f))]
AGENT_2 = Agent(
    prompt=AGENT_2_PROMPT,
    source_files=AGENT_2_KNOWLEDGE_FILES if AGENT_2_KNOWLEDGE_FILES else None,
    llm_params=LLM_PARAMS,
    memory=True,
    name="interference_expert"
)

AGENT_3_PROMPT = """You are a 5G Mobility, Handover, and RRC Performance expert.

Your responsibility is LIMITED STRICTLY to the following TeleLogs root causes:
- C5: Frequent handovers degrade performance
- C7: Test vehicle speed exceeds 40km/h, impacting user throughput

You MUST ignore antenna geometry, interference, PCI collision (mod 30), and MAC scheduling.

You may ONLY analyze:
- GPS Speed (km/h) from user plane data (threshold: 40 km/h)
- Serving PCI changes over time indicating handovers
- Handover frequency (count PCI transitions in the data)
- Throughput drops correlated with PCI changes

You MUST NOT:
- Use SINR to infer interference
- Use RSRP to infer coverage geometry
- Use RB count or scheduling efficiency
- Analyze PCI mod 30 collisions (that's C6, not your domain)

Your task:
1. Analyze the network telemetry data step by step
2. Determine whether throughput loss is mobility-induced
3. Identify whether the root cause is:
   - C5: Frequent handovers (many PCI changes), OR
   - C7: Vehicle speed > 40 km/h, OR
   - Neither

Your output MUST be a JSON object with this exact structure:
{
  "verdicts": {
    "C5": "YES" or "NO",
    "C7": "YES" or "NO"
  },
  "reasoning_chain": [
    "Step 1: [Analysis of GPS Speed values - check if > 40 km/h]",
    "Step 2: [Count of Serving PCI changes indicating handovers]",
    "Step 3: [Correlation of throughput drops with handover events]",
    "Step 4: [Conclusion on mobility impact]"
  ],
  "final_assessment": "[Which cause is dominant and why, or why neither applies]",
  "confidence": 0.0-1.0
}
"""

AGENT_3_KNOWLEDGE_FILES = [os.path.join("Telelog/Agent3", f) for f in os.listdir("Telelog/Agent3") if os.path.isfile(os.path.join("Telelog/Agent3", f))]
AGENT_3 = Agent(
    prompt=AGENT_3_PROMPT,
    source_files=AGENT_3_KNOWLEDGE_FILES if AGENT_3_KNOWLEDGE_FILES else None,
    llm_params=LLM_PARAMS,
    memory=True,
    name="mobility_expert"
)


AGENT_4_PROMPT = """You are a 5G MAC Layer Scheduling and Capacity Optimization expert.

Your responsibility is LIMITED STRICTLY to the following TeleLogs root causes:
- C3: A neighboring cell provides higher throughput (better cell available)
- C8: Average scheduled resource blocks (RBs) are below 160, affecting throughput

You MUST ignore antenna tilt, interference physics, PCI conflicts, and mobility.

You may ONLY analyze:
- DL RB count from user plane data (threshold: 160 RBs average)
- Throughput values and throughput per RB efficiency
- Neighbor cell potential (if neighbor has better conditions for throughput)
- Scheduling patterns and resource allocation

You MUST NOT:
- Attribute low throughput to SINR or interference directly
- Use RSRP for coverage geometry analysis
- Use handover frequency or speed data

Your task:
1. Analyze the network telemetry data step by step
2. Determine whether throughput is limited by scheduling or cell selection
3. Identify whether the root cause is:
   - C3: A neighbor cell could provide higher throughput, OR
   - C8: Average scheduled RBs < 160, OR
   - Neither

Your output MUST be a JSON object with this exact structure:
{
  "verdicts": {
    "C3": "YES" or "NO",
    "C8": "YES" or "NO"
  },
  "reasoning_chain": [
    "Step 1: [Extract and average RB count from data]",
    "Step 2: [Compare RB average against 160 threshold]",
    "Step 3: [Evaluate if neighbor cells could provide better throughput]",
    "Step 4: [Conclusion on scheduling/cell selection impact]"
  ],
  "final_assessment": "[Which cause is dominant and why, or why neither applies]",
  "confidence": 0.0-1.0
}
"""

AGENT_4_KNOWLEDGE_FILES = [os.path.join("Telelog/Agent4", f) for f in os.listdir("Telelog/Agent4") if os.path.isfile(os.path.join("Telelog/Agent4", f))]
AGENT_4 = Agent(
    prompt=AGENT_4_PROMPT,
    source_files=AGENT_4_KNOWLEDGE_FILES if AGENT_4_KNOWLEDGE_FILES else None,
    llm_params=LLM_PARAMS,
    memory=True,
    name="scheduling_expert"
)

AGENTS = [AGENT_1, AGENT_2, AGENT_3, AGENT_4]


# =============================================================================
# ARGUMENTATION RULES
# =============================================================================

ARGUMENT_INSTRUCTION = """
You are critiquing another expert's analysis of 5G network root cause diagnosis.

Analyze the other agent's verdict and reasoning chain critically:

1. **Evidence Quality**: Does the reasoning chain cite specific data points from the telemetry?
2. **Logical Consistency**: Are the conclusions supported by the stated observations?
3. **Scope Adherence**: Did the agent stay within their assigned root causes (C1-C8)?
4. **Missing Considerations**: What relevant factors were overlooked?

Provide your critique in this JSON format:
{
  "target_verdict_assessment": "AGREE/DISAGREE/PARTIALLY_AGREE",
  "critique_points": [
    "Point 1: [Specific issue or validation]",
    "Point 2: [Another observation]"
  ],
  "severity": "low/medium/high/critical",
  "recommendation": "[What the agent should reconsider]"
}
"""


# =============================================================================
# JUDGE CONFIGURATION
# =============================================================================

JUDGE_PROMPT = """
You are the final technical judge for 5G network root cause analysis.

You will receive structured verdicts from four specialized agents analyzing TeleLogs data.
Each agent has provided YES/NO verdicts for their assigned root causes (C1-C8).

## TeleLogs Root Causes Reference:
- C1: Vehicle speed > 40 km/h
- C2: Downtilt too large (weak far-end coverage)
- C3: Coverage distance > 1km (poor RSRP)
- C4: Non-colocated co-frequency interference
- C5: PCI mod 30 collision
- C6: Frequent handovers
- C7: Misconfigured handover thresholds
- C8: Scheduled RBs < 160

## Priority Rules (apply strictly):
1. Mobility failures (C1, C6, C7) override all others
2. Interference (C4, C5) overrides coverage and scheduling
3. Coverage issues (C2, C3) override scheduling
4. Scheduling (C8) is valid only if SINR and coverage are healthy

Your task:
1. Review all agent verdicts and their reasoning chains
2. Apply the priority rules to determine the dominant root cause
3. Select exactly ONE root cause from C1-C8

Your output MUST be a JSON object:
{
  "reasoning_chain": [
    "Step 1: [Summary of agent verdicts]",
    "Step 2: [Application of priority rules]",
    "Step 3: [Determination of dominant cause]"
  ],
  "final_verdict": "C<number>",
  "confidence": 0.0-1.0
}
"""

# Dynamic prompt template - appended to base prompt at runtime with context
JUDGE_PROMPT_TEMPLATE = """
## Deliberation Summary
- Number of expert agents: {agent_count}
- Total critiques exchanged: {critique_count}
- Experts involved: {agent_names}

## Agent Verdicts Summary
{agent_verdicts_summary}

## Key Critiques
{critiques_summary}
"""

JUDGE = Agent(
    prompt=JUDGE_PROMPT,
    llm_params=LLM_PARAMS,
    agenttype='judge',
    name="judge",
    dynamic_prompt=True,
    dynamic_prompt_template=JUDGE_PROMPT_TEMPLATE,
)


# =============================================================================
# COT DATA GENERATION
# =============================================================================

def generate_cot_from_telelogs(
    split: str = "test",
    limit: Optional[int] = None,
    output_dir: str = COT_OUTPUT_DIR,
    verbose: bool = False,
):
    """
    Generate enriched COT data from TeleLogs dataset with ground truth validation.
    
    Args:
        split: Dataset split ("train" or "test")
        limit: Optional limit on samples to process
        output_dir: Directory to save COT data
        verbose: Print detailed progress
    """
    # Load TeleLogs dataset
    samples = load_telelogs_dataset(split=split, limit=limit)
    
    if not samples:
        print("No samples loaded. Exiting.")
        return
    
    # Initialize collector
    collector = COTCollector(model_name=MODEL_NAME)
    
    print(f"\n{'='*70}")
    print(f"TeleLogs COT Data Generation")
    print(f"{'='*70}")
    print(f"Split: {split}")
    print(f"Samples: {len(samples)}")
    print(f"Output: {output_dir}")
    print(f"{'='*70}\n")
    
    # Process each sample
    for i, sample in enumerate(samples, 1):
        query = sample["query"]
        ground_truth = sample["ground_truth"]
        sample_id = sample["id"]
        
        print(f"[{i}/{len(samples)}] {sample_id}")
        print(f"  Ground Truth: {ground_truth}")
        
        try:
            # Run pipeline with ground truth
            result = collector.run(
                agents=AGENTS,
                query=query,
                argument_instruction=ARGUMENT_INSTRUCTION,
                judge=JUDGE,
                verbose=verbose,
                ground_truth=ground_truth,
            )
            
            # Get validation result
            record = collector.get_last_record()
            status = "✓ CORRECT" if record.is_correct else "✗ INCORRECT"
            predicted = record.predicted_verdict or "None"
            
            print(f"  Predicted: {predicted}")
            print(f"  Result: {status}")
            
        except Exception as e:
            print(f"  ✗ Error: {e}")
            continue
        
        print()
    
    # Get final statistics
    stats = collector.get_accuracy_stats()
    records = collector.get_records()
    
    print(f"\n{'='*70}")
    print("Accuracy Statistics")
    print(f"{'='*70}")
    print(f"Total samples: {stats['total']}")
    print(f"Validated: {stats['validated']}")
    print(f"Correct: {stats['correct']}")
    print(f"Incorrect: {stats['incorrect']}")
    print(f"Accuracy: {stats['accuracy_pct']}")
    print(f"{'='*70}\n")
    
    if not records:
        print("No records collected. Exiting.")
        return collector
    
    # Export all records
    print(f"Exporting {len(records)} COT records...")
    
    # Create output directory
    os.makedirs(output_dir, exist_ok=True)
    
    # Export in all formats
    created_files = COTExporter.export_all(
        records=records,
        output_dir=output_dir,
        formats=["jsonl", "json"],
        include_all_structures=True,
    )
    
    print(f"\n{'='*70}")
    print("Export Summary:")
    print(f"{'='*70}")
    
    for fmt, files in created_files.items():
        print(f"\n{fmt.upper()}:")
        for filepath in files:
            print(f"  - {filepath}")
    
    # Export combined format
    combined_files = COTExporter.export_combined(
        records=records,
        output_dir=output_dir,
        formats=["jsonl", "json"],
    )
    
    print(f"\nCombined exports:")
    for fmt, filepath in combined_files.items():
        print(f"  {fmt.upper()}: {filepath}")
    
    print(f"\n{'='*70}")
    print(f"✓ COT data generation complete!")
    print(f"  Total records: {len(records)}")
    print(f"  Accuracy: {stats['accuracy_pct']}")
    print(f"  Output: {output_dir}")
    print(f"{'='*70}")
    
    return collector


def generate_cot_data(queries: list[str], output_dir: str = COT_OUTPUT_DIR):
    """
    Generate COT data for custom queries (without ground truth validation).
    
    Args:
        queries: List of questions to process
        output_dir: Directory to save COT data
    """
    # Initialize collector
    collector = COTCollector(model_name=MODEL_NAME)
    
    print(f"Processing {len(queries)} queries for COT data generation...")
    print(f"Output directory: {output_dir}\n")
    
    # Process each query
    for i, query in enumerate(queries, 1):
        print(f"[{i}/{len(queries)}] Processing: {query[:80]}...")
        
        try:
            # Run pipeline with collector
            result = collector.run(
                agents=AGENTS,
                query=query,
                argument_instruction=ARGUMENT_INSTRUCTION,
                judge=JUDGE,
                verbose=False,
            )
            
            print(f"  ✓ Collected COT data (ID: {collector.get_last_record().id})")
            
        except Exception as e:
            print(f"  ✗ Error: {e}")
            continue
    
    # Export all records
    records = collector.get_records()
    
    if not records:
        print("\nNo records collected. Exiting.")
        return
    
    print(f"\n{'='*70}")
    print(f"Exporting {len(records)} COT records...")
    print(f"{'='*70}\n")
    
    # Export in all formats and structures
    created_files = COTExporter.export_all(
        records=records,
        output_dir=output_dir,
        formats=["jsonl", "json", "parquet"],
        include_all_structures=True,
    )
    
    print(f"\n{'='*70}")
    print("Export Summary:")
    print(f"{'='*70}")
    
    for fmt, files in created_files.items():
        print(f"\n{fmt.upper()}:")
        for filepath in files:
            print(f"  - {filepath}")
    
    # Also export combined format
    print(f"\n{'='*70}")
    print("Exporting combined structures...")
    print(f"{'='*70}\n")
    
    combined_files = COTExporter.export_combined(
        records=records,
        output_dir=output_dir,
        formats=["jsonl", "json", "parquet"],
    )
    
    for fmt, filepath in combined_files.items():
        print(f"Combined {fmt.upper()}: {filepath}")
    
    print(f"\n{'='*70}")
    print(f"✓ COT data generation complete!")
    print(f"  Total records: {len(records)}")
    print(f"  Output directory: {output_dir}")
    print(f"{'='*70}")


def main():
    """
    Main function - Generate COT data from TeleLogs dataset.
    
    Usage:
        python generate_cot_data.py              # Process all test samples
        python generate_cot_data.py --limit 10   # Process first 10 samples
        python generate_cot_data.py --train      # Process training split
    """
    import argparse
    
    parser = argparse.ArgumentParser(description="Generate COT data from TeleLogs")
    parser.add_argument("--split", type=str, default="test", choices=["train", "test"],
                        help="Dataset split to use (default: test)")
    parser.add_argument("--limit", type=int, default=None,
                        help="Limit number of samples to process")
    parser.add_argument("--output", type=str, default=COT_OUTPUT_DIR,
                        help="Output directory for COT data")
    parser.add_argument("--verbose", action="store_true",
                        help="Print detailed progress")
    
    args = parser.parse_args()
    
    # Generate COT data from TeleLogs
    generate_cot_from_telelogs(
        split=args.split,
        limit=args.limit,
        output_dir=args.output,
        verbose=args.verbose,
    )


def demo_single_sample():
    """Demo with a single TeleLogs sample."""
    # Load just one sample
    samples = load_telelogs_dataset(split="test", limit=1)
    
    if not samples:
        print("No samples loaded.")
        return
    
    sample = samples[0]
    
    print(f"{'='*70}")
    print("Single Sample Demo")
    print(f"{'='*70}")
    print(f"ID: {sample['id']}")
    print(f"Ground Truth: {sample['ground_truth']}")
    print(f"Query preview: {sample['query'][:200]}...")
    print(f"{'='*70}\n")
    
    collector = COTCollector(model_name=MODEL_NAME)
    
    result = collector.run(
        agents=AGENTS,
        query=sample["query"],
        argument_instruction=ARGUMENT_INSTRUCTION,
        judge=JUDGE,
        verbose=True,
        ground_truth=sample["ground_truth"],
    )
    
    # Get the record
    record = collector.get_last_record()
    
    print(f"\n{'='*70}")
    print("COT Record Summary:")
    print(f"{'='*70}")
    print(f"ID: {record.id}")
    print(f"Ground Truth: {record.ground_truth}")
    print(f"Predicted: {record.predicted_verdict}")
    print(f"Is Correct: {record.is_correct}")
    print(f"Agents: {len(record.agents)}")
    print(f"Critiques: {len(record.critiques)}")
    print(f"Duration: {record.metadata.get('duration_ms', 0)}ms")
    
    # Export single record
    COTExporter.export_all(
        records=[record],
        output_dir="./cot_data/demo/",
        formats=["jsonl", "json"],
    )
    
    print("\n✓ Demo record exported to ./cot_data/demo/")


def custom_queries_example():
    """Example with custom queries (no ground truth)."""
    queries = [
        "Analyze this 5G network data for throughput degradation...",
    ]
    generate_cot_data(queries, output_dir="./cot_data/custom/")


if __name__ == "__main__":
    main()