scripts/build_features.py

"""
This script extracts handcrafted features from agent execution traces for the classical
ML pipeline. These features capture structural, behavioral, and linguistic
signals of trace anomalies (circular reasoning, goal drift, unnecessary
tool calls, failure patterns).

Feature groups:
  1. Structural     : turn counts, trace length, conversation depth
  2. Tool-usage     : call counts, diversity, repetition, density
  3. Behavioral     : circular patterns, consecutive repeats, call-response ratio
  4. Linguistic     : error keywords, sentiment signals, response length stats
  5. Sequence       : positional features, where tool calls appear in the trace
"""

import argparse
import json
import os
import re
from collections import Counter

import numpy as np
import pandas as pd
from tqdm import tqdm


# list of keywords/phrases indicative of errors, failures, or issues in tool responses
ERROR_KEYWORDS = [
    "error", "failed", "failure", "exception", "timeout",
    "unauthorized", "forbidden", "not found", "invalid",
    "rate limit", "500", "502", "503", "404", "400",
]

APOLOGY_KEYWORDS = [
    "sorry", "apologize", "unfortunately", "regret",
    "unable", "cannot", "can't", "couldn't",
]

HEDGING_KEYWORDS = [
    "might", "perhaps", "possibly", "try again",
    "not sure", "uncertain", "unclear",
]

GIVE_UP_KEYWORDS = [
    "give up", "stop here", "will not", "won't be able",
    "no further", "end here", "conclude",
]

# regex to detect Action/Action Input inside gpt turns (ReAct format)
ACTION_RE = re.compile(r"Action\s*:\s*(.+)", re.IGNORECASE)


def normalize_message(msg):
    """
    Ensure a conversation message is a dict with 'from' and 'value' keys.
    ToolBench stores some turns as serialized JSON strings instead of dicts.
    """
    if isinstance(msg, str):
        try:
            msg = json.loads(msg)
        except (json.JSONDecodeError, TypeError):
            msg = {"from": "unknown", "value": msg}

    if not isinstance(msg, dict):
        msg = {"from": "unknown", "value": str(msg)}

    return msg


def extract_actions_from_gpt_turn(text: str) -> list[str]:
    """
    Extract tool call names from a gpt/assistant turn that uses ReAct format.
    Returns list of action names found in the text.
    """
    actions = ACTION_RE.findall(text)
    cleaned = []
    for a in actions:
        a = a.strip()
        if a.lower() not in ("finish", "none", ""):
            cleaned.append(a)
    return cleaned


def count_keyword_hits(text: str, keywords: list[str]) -> int:
    """Count how many keyword patterns appear in text."""
    text_lower = text.lower()
    return sum(1 for kw in keywords if kw in text_lower)


def extract_tool_names(tool_calls: list[str]) -> list[str]:
    """
    Extract function/API names from tool call strings.
    ToolBench tool calls are typically JSON with a 'name' or 'api_name' field,
    OR plain action names extracted from ReAct-format gpt turns.
    """
    names = []
    for tc in tool_calls:
        # first try parsing as JSON (explicit function_call format)
        try:
            parsed = json.loads(tc)
            name = parsed.get("name", parsed.get("api_name", parsed.get("function", "")))
            if name:
                names.append(name)
                continue
            else:
                for key in parsed:
                    if "name" in key.lower():
                        names.append(str(parsed[key]))
                        break
                continue
        except (json.JSONDecodeError, TypeError, AttributeError):
            pass

        # fallback: regex for common patterns
        match = re.search(r'"(?:name|api_name|function)"\s*:\s*"([^"]+)"', tc)
        if match:
            names.append(match.group(1))
        else:
            # already a plain action name from ReAct extraction
            names.append(tc.strip()[:80])
    return names


def compute_repetition_features(tool_names: list[str]) -> dict:
    """
    Compute features related to repeated/circular tool usage.
    """
    if not tool_names:
        return {
            "unique_tools": 0,
            "tool_diversity_ratio": 0.0,
            "max_single_tool_freq": 0,
            "max_consecutive_same_tool": 0,
            "num_repeated_exact_calls": 0,
        }

    counts = Counter(tool_names)
    unique = len(counts)
    total = len(tool_names)

    # longest streak of the same tool called consecutively
    max_consec = 1
    current_consec = 1
    for i in range(1, len(tool_names)):
        if tool_names[i] == tool_names[i - 1]:
            current_consec += 1
            max_consec = max(max_consec, current_consec)
        else:
            current_consec = 1

    return {
        "unique_tools": unique,
        "tool_diversity_ratio": unique / total if total > 0 else 0.0,
        "max_single_tool_freq": max(counts.values()),
        "max_consecutive_same_tool": max_consec,
        "num_repeated_exact_calls": total - unique,
    }


def compute_response_length_features(assistant_turns: list[str]) -> dict:
    """Stats on assistant response lengths."""
    if not assistant_turns:
        return {
            "mean_response_len": 0.0,
            "std_response_len": 0.0,
            "max_response_len": 0,
            "min_response_len": 0,
            "total_response_chars": 0,
        }

    lengths = [len(t) for t in assistant_turns]
    return {
        "mean_response_len": float(np.mean(lengths)),
        "std_response_len": float(np.std(lengths)),
        "max_response_len": max(lengths),
        "min_response_len": min(lengths),
        "total_response_chars": sum(lengths),
    }


def compute_observation_features(observations: list[str]) -> dict:
    """Features from tool responses/observations."""
    if not observations:
        return {
            "mean_obs_len": 0.0,
            "num_error_observations": 0,
            "num_empty_observations": 0,
        }

    lengths = [len(o) for o in observations]
    n_errors = sum(
        1 for o in observations
        if count_keyword_hits(o, ERROR_KEYWORDS) > 0
    )
    n_empty = sum(1 for o in observations if len(o.strip()) < 10)

    return {
        "mean_obs_len": float(np.mean(lengths)),
        "num_error_observations": n_errors,
        "num_empty_observations": n_empty,
    }


def compute_positional_features(turns: list[tuple], tool_call_turn_indices: list[int]) -> dict:
    """
    Where in the trace do tool calls appear?
    Normalized position (0 = start, 1 = end).
    """
    total = len(turns) if turns else 1

    if not tool_call_turn_indices:
        return {
            "first_tool_position": 1.0,
            "last_tool_position": 0.0,
            "tool_position_spread": 0.0,
        }

    positions = [i / total for i in tool_call_turn_indices]

    return {
        "first_tool_position": positions[0],
        "last_tool_position": positions[-1],
        "tool_position_spread": positions[-1] - positions[0],
    }


def extract_features_from_row(row: pd.Series) -> dict:
    """
    Extract all features from a single processed dataset row.
    """
    conv = json.loads(row["conversations_json"])

    # rebuild parsed trace
    turns = []
    tool_calls = []
    observations = []
    assistant_turns = []
    tool_call_turn_indices = []  # which turn indices contain tool calls

    for raw_msg in conv:
        msg = normalize_message(raw_msg)
        role = msg.get("from", msg.get("role", "unknown"))
        content = msg.get("value", msg.get("content", ""))
        if content is None:
            content = ""

        turn_idx = len(turns)
        turns.append((role, content))

        if role in ("gpt", "assistant", "chatgpt"):
            assistant_turns.append(content)
            # extract embedded tool calls from ReAct-style gpt turns
            actions = extract_actions_from_gpt_turn(content)
            if actions:
                tool_calls.extend(actions)
                tool_call_turn_indices.append(turn_idx)
        elif role in ("function_call", "tool_call"):
            tool_calls.append(content)
            tool_call_turn_indices.append(turn_idx)
        elif role in ("observation", "tool_response", "function"):
            observations.append(content)

    # extract tool names for repetition analysis
    tool_names = extract_tool_names(tool_calls)

    # structural features
    features = {
        "num_turns": len(turns),
        "num_tool_calls": len(tool_calls),
        "num_observations": len(observations),
        "num_assistant_turns": len(assistant_turns),
        "trace_length_chars": len(row["raw_trace"]),
    }

    # tool usage features
    features["tool_call_density"] = (
        len(tool_calls) / len(turns) if len(turns) > 0 else 0.0
    )
    features["tool_obs_ratio"] = (
        len(observations) / len(tool_calls) if len(tool_calls) > 0 else 0.0
    )

    # repetition and circularity features
    features.update(compute_repetition_features(tool_names))

    # response length features
    features.update(compute_response_length_features(assistant_turns))

    # observation features
    features.update(compute_observation_features(observations))

    # positional features
    features.update(compute_positional_features(turns, tool_call_turn_indices))

    # linguistic features
    full_text = " ".join(content for _, content in turns)
    last_assistant = assistant_turns[-1] if assistant_turns else ""

    features["error_keyword_count"] = count_keyword_hits(full_text, ERROR_KEYWORDS)
    features["apology_keyword_count"] = count_keyword_hits(full_text, APOLOGY_KEYWORDS)
    features["hedging_keyword_count"] = count_keyword_hits(full_text, HEDGING_KEYWORDS)
    features["give_up_keyword_count"] = count_keyword_hits(full_text, GIVE_UP_KEYWORDS)
    features["last_turn_error_keywords"] = count_keyword_hits(last_assistant, ERROR_KEYWORDS)
    features["last_turn_apology_keywords"] = count_keyword_hits(last_assistant, APOLOGY_KEYWORDS)

    return features


def build_feature_matrix(df: pd.DataFrame) -> pd.DataFrame:
    """
    Apply feature extraction to every row. Returns a DataFrame of features + label.
    """
    print(f"[INFO] Extracting features from {len(df)} traces...")
    feature_rows = []
    for _, row in tqdm(df.iterrows(), total=len(df), desc="Extracting features"):
        feats = extract_features_from_row(row)
        feats["id"] = row["id"]
        feats["label"] = row["label"]
        feature_rows.append(feats)

    feat_df = pd.DataFrame(feature_rows)
    print(f"[INFO] Extracted {len(feat_df.columns) - 2} features.")  # minus id and label
    return feat_df


def get_feature_columns(df: pd.DataFrame) -> list[str]:
    """Return the list of feature column names (everything except id and label)."""
    return [c for c in df.columns if c not in ("id", "label")]


def main():
    parser = argparse.ArgumentParser(description="Build features from processed traces")
    parser.add_argument("--input_dir", type=str, default="data/processed")
    parser.add_argument("--output_dir", type=str, default="data/processed")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    for split_name in ["train", "val", "test"]:
        input_path = os.path.join(args.input_dir, f"{split_name}.parquet")
        if not os.path.exists(input_path):
            print(f"[WARN] {input_path} not found, skipping.")
            continue

        df = pd.read_parquet(input_path)
        feat_df = build_feature_matrix(df)

        output_path = os.path.join(args.output_dir, f"{split_name}_features.parquet")
        feat_df.to_parquet(output_path, index=False)
        print(f"[INFO] Saved features → {output_path}")

    print("[DONE] Feature extraction complete.")


if __name__ == "__main__":
    main()