generate_dataset.py

"""
Dataset generator for SwiftContext search-strategy router.

SwiftContext improves on FastContext (Microsoft, arXiv:2606.14066) by adding a
lightweight pre-router that classifies the coding agent's exploration intent
BEFORE the 4B explorer LLM is invoked.  This eliminates wasted first turns
where FastContext's LLM had to "discover" what kind of search was needed.

Router labels (search strategy):
  0 - broad_scan       : wide exploration, file/module locations unknown
  1 - targeted_search  : specific named symbol (function/class/const) to locate
  2 - pinpoint_cite    : exact line-level citation of already-scoped code

Outputs:
  data/train.jsonl   (~210 examples per class)
  data/val.jsonl     (~45  examples per class)
  data/test.jsonl    (~45  examples per class)
"""

import json
import random
import os
import string
from pathlib import Path

random.seed(42)

# ── Broad-scan templates & fillers ──────────────────────────────────────────

BROAD_TEMPLATES = [
    "Where is {concern} handled in this codebase?",
    "Which files deal with {concern}?",
    "Find all code related to {concern}.",
    "What parts of the repo handle {concern}?",
    "Which modules are responsible for {concern}?",
    "Show me everything related to {concern}.",
    "Find all {concern}-related files.",
    "Where in the codebase is {concern} implemented?",
    "Give me an overview of how {concern} is structured.",
    "Which directories contain {concern} logic?",
    "Where are {concern} utilities defined?",
    "Find the {concern} layer of this application.",
    "What handles {concern} across the whole project?",
    "List all files involved in {concern}.",
    "Where does {concern} get initialized?",
    "Which entry points trigger {concern}?",
    "Find every place that touches {concern}.",
    "What is the overall structure of {concern} in this repo?",
    "Walk me through how {concern} works across files.",
    "Which classes/modules collaborate on {concern}?",
]

BROAD_CONCERNS = [
    "authentication", "authorization", "database connections", "error logging",
    "user sessions", "API rate limiting", "caching", "file uploads",
    "email sending", "payment processing", "background jobs", "webhooks",
    "configuration loading", "request validation", "response serialization",
    "middleware", "routing", "dependency injection", "event handling",
    "data migration", "test fixtures", "metrics collection", "audit logging",
    "data encryption", "token refresh", "session expiry", "retry logic",
    "circuit breaking", "connection pooling", "search indexing", "image processing",
    "push notifications", "feature flags", "A/B testing", "internationalization",
    "health checks", "access control lists", "multi-tenancy", "CORS handling",
    "GraphQL resolvers", "OpenAPI schema generation", "database seeding",
    "password hashing", "OAuth flows", "CSRF protection", "SQL query building",
    "file streaming", "PDF generation", "CSV import/export", "job scheduling",
]

# ── Targeted-search templates & symbols ─────────────────────────────────────

TARGETED_TEMPLATES = [
    "Find the `{symbol}` {kind}.",
    "Where is `{symbol}` defined?",
    "Show me the `{symbol}` {kind}.",
    "Locate `{symbol}` in the codebase.",
    "Where is `{symbol}` implemented?",
    "Find the definition of `{symbol}`.",
    "Where can I find the `{symbol}` {kind}?",
    "Which file contains `{symbol}`?",
    "Where is `{symbol}` declared?",
    "Show me where `{symbol}` lives.",
    "Find usages of `{symbol}`.",
    "Where is `{symbol}` called?",
    "Locate all references to `{symbol}`.",
    "Which file exports `{symbol}`?",
    "Where does `{symbol}` get imported from?",
    "Which module defines `{symbol}`?",
    "Find every call site of `{symbol}`.",
    "Where is `{symbol}` first instantiated?",
    "Show me all places that import `{symbol}`.",
    "Where is `{symbol}` registered or configured?",
]

TARGETED_SYMBOLS = [
    ("authenticate_user", "function"),
    ("UserModel", "class"),
    ("MAX_RETRIES", "constant"),
    ("parse_config", "function"),
    ("DatabaseConnection", "class"),
    ("validate_email", "function"),
    ("send_notification", "function"),
    ("PaymentService", "class"),
    ("API_BASE_URL", "constant"),
    ("refresh_token", "function"),
    ("SessionManager", "class"),
    ("retry_with_backoff", "function"),
    ("CacheClient", "class"),
    ("RATE_LIMIT_WINDOW", "constant"),
    ("process_webhook", "function"),
    ("generate_report", "function"),
    ("FileUploader", "class"),
    ("encrypt_payload", "function"),
    ("get_user_by_id", "function"),
    ("EventEmitter", "class"),
    ("DEFAULT_TIMEOUT", "constant"),
    ("validate_token", "function"),
    ("MigrationRunner", "class"),
    ("log_error", "function"),
    ("SearchIndex", "class"),
    ("compress_image", "function"),
    ("ALLOWED_ORIGINS", "constant"),
    ("schedule_job", "function"),
    ("FeatureFlag", "class"),
    ("decode_jwt", "function"),
    ("ConnectionPool", "class"),
    ("sanitize_input", "function"),
    ("BATCH_SIZE", "constant"),
    ("send_transactional_email", "function"),
    ("AuditLogger", "class"),
    ("hash_password", "function"),
    ("RateLimiter", "class"),
    ("parse_request_body", "function"),
    ("ENV_CONFIG", "constant"),
    ("HealthCheck", "class"),
    ("format_api_response", "function"),
    ("MAX_CONNECTIONS", "constant"),
    ("build_query", "function"),
    ("OAuthProvider", "class"),
    ("revoke_token", "function"),
    ("TaskQueue", "class"),
    ("normalize_path", "function"),
    ("APP_VERSION", "constant"),
    ("CircuitBreaker", "class"),
    ("stream_file", "function"),
]

# ── Pinpoint-cite templates & fillers ───────────────────────────────────────

PINPOINT_TEMPLATES = [
    "Show me the full body of `{symbol}`.",
    "What does `{symbol}` return?",
    "What parameters does `{symbol}` accept?",
    "Show me the exact signature of `{symbol}`.",
    "What type hints does `{symbol}` use?",
    "Show me the error handling inside `{symbol}`.",
    "What does `{symbol}` raise when {condition}?",
    "Show me lines {start}-{end} of {filename}.",
    "What is the exact implementation of `{symbol}`?",
    "Show me all arguments to `{symbol}` including defaults.",
    "What does the decorator on `{symbol}` do?",
    "Show me the docstring of `{symbol}`.",
    "What does `{symbol}` do when {condition}?",
    "Show me the `{prop}` method of `{klass}`.",
    "What is the return type annotation of `{symbol}`?",
    "Show me the try/except block inside `{symbol}`.",
    "What are the default values for `{symbol}` parameters?",
    "Show me the class variables declared in `{klass}`.",
    "What does `{symbol}` log at the {level} level?",
    "Show me the SQL query built inside `{symbol}`.",
]

PINPOINT_CONDITIONS = [
    "the token is expired",
    "the user is not found",
    "the connection fails",
    "the input is invalid",
    "the file doesn't exist",
    "the rate limit is exceeded",
    "authentication fails",
    "the cache is empty",
    "the request times out",
    "the database is unreachable",
    "permissions are denied",
    "the payload is too large",
    "the queue is full",
    "the secret is rotated",
]

PINPOINT_FILENAMES = [
    "auth.py", "models/user.py", "services/payment.py", "utils/crypto.py",
    "api/routes.py", "db/connection.py", "middleware/rate_limit.py",
    "handlers/webhook.py", "tasks/email.py", "core/config.py",
    "lib/cache.py", "src/auth/index.ts", "controllers/UserController.ts",
    "services/AuthService.go", "internal/db/pool.go", "src/lib.rs",
    "src/handlers/api.rs", "util/StringHelper.java", "src/models.rs",
]

PINPOINT_PROPS = [
    "save", "delete", "update", "validate", "serialize", "deserialize",
    "connect", "disconnect", "refresh", "reset", "flush", "close",
]

PINPOINT_LOG_LEVELS = ["debug", "info", "warning", "error", "critical"]

PINPOINT_CLASSES = [
    "UserModel", "DatabaseConnection", "SessionManager", "CacheClient",
    "PaymentService", "FileUploader", "AuditLogger", "RateLimiter",
    "ConnectionPool", "TaskQueue", "CircuitBreaker",
]


# ── Generators ───────────────────────────────────────────────────────────────

def gen_broad(n: int) -> list[dict]:
    examples = []
    for tmpl in BROAD_TEMPLATES:
        for concern in BROAD_CONCERNS:
            examples.append({"text": tmpl.format(concern=concern), "label": 0})
    random.shuffle(examples)
    return examples[:n]


def gen_targeted(n: int) -> list[dict]:
    examples = []
    for tmpl in TARGETED_TEMPLATES:
        for symbol, kind in TARGETED_SYMBOLS:
            examples.append({"text": tmpl.format(symbol=symbol, kind=kind), "label": 1})
    random.shuffle(examples)
    return examples[:n]


def gen_pinpoint(n: int) -> list[dict]:
    examples = []
    for tmpl in PINPOINT_TEMPLATES:
        required = set(
            k
            for _, k, _, _ in string.Formatter().parse(tmpl)
            if k is not None
        )
        if required == {"symbol"}:
            for symbol, _ in TARGETED_SYMBOLS:
                examples.append({"text": tmpl.format(symbol=symbol), "label": 2})
        elif required == {"symbol", "condition"}:
            for symbol, _ in TARGETED_SYMBOLS[:10]:
                for cond in PINPOINT_CONDITIONS[:5]:
                    examples.append({"text": tmpl.format(symbol=symbol, condition=cond), "label": 2})
        elif required == {"start", "end", "filename"}:
            for fname in PINPOINT_FILENAMES:
                start = random.randint(1, 200)
                end = start + random.randint(5, 40)
                examples.append({"text": tmpl.format(start=start, end=end, filename=fname), "label": 2})
        elif required == {"prop", "klass"}:
            for klass in PINPOINT_CLASSES:
                for prop in PINPOINT_PROPS:
                    examples.append({"text": tmpl.format(prop=prop, klass=klass), "label": 2})
        elif required == {"symbol", "level"}:
            for symbol, _ in TARGETED_SYMBOLS[:10]:
                for level in PINPOINT_LOG_LEVELS:
                    examples.append({"text": tmpl.format(symbol=symbol, level=level), "label": 2})
        elif required == {"klass"}:
            for klass in PINPOINT_CLASSES:
                examples.append({"text": tmpl.format(klass=klass), "label": 2})
        else:
            # fallback: just use symbol
            for symbol, _ in TARGETED_SYMBOLS:
                try:
                    examples.append({"text": tmpl.format(symbol=symbol), "label": 2})
                except KeyError:
                    pass
    random.shuffle(examples)
    return examples[:n]


# ── Split & write ─────────────────────────────────────────────────────────────

def split_and_write(examples: list[dict], out_dir: Path, train_frac=0.70, val_frac=0.15):
    # Stratified split: divide per-class first so every split has balanced labels.
    # A plain shuffle-then-slice can produce lopsided splits when N is small.
    from collections import defaultdict
    by_label: dict[int, list[dict]] = defaultdict(list)
    for ex in examples:
        by_label[ex["label"]].append(ex)

    splits: dict[str, list[dict]] = {"train": [], "val": [], "test": []}
    for label_examples in by_label.values():
        random.shuffle(label_examples)
        n       = len(label_examples)
        n_train = int(n * train_frac)
        n_val   = int(n * val_frac)
        splits["train"].extend(label_examples[:n_train])
        splits["val"].extend(label_examples[n_train : n_train + n_val])
        splits["test"].extend(label_examples[n_train + n_val :])

    out_dir.mkdir(parents=True, exist_ok=True)
    for split_name, split_data in splits.items():
        random.shuffle(split_data)   # inter-class shuffle within each split
        path = out_dir / f"{split_name}.jsonl"
        with open(path, "w", encoding="utf-8") as f:
            for ex in split_data:
                f.write(json.dumps(ex) + "\n")
        print(f"  Wrote {len(split_data):>4d} examples → {path}")


def main():
    N_PER_CLASS = 300  # balanced classes

    print("Generating broad_scan examples...")
    broad = gen_broad(N_PER_CLASS)

    print("Generating targeted_search examples...")
    targeted = gen_targeted(N_PER_CLASS)

    print("Generating pinpoint_cite examples...")
    pinpoint = gen_pinpoint(N_PER_CLASS)

    # Deduplicate by exact text — template × filler combos can produce collisions
    seen_texts: set[str] = set()
    all_examples: list[dict] = []
    for ex in broad + targeted + pinpoint:
        if ex["text"] not in seen_texts:
            seen_texts.add(ex["text"])
            all_examples.append(ex)
    random.shuffle(all_examples)

    print(f"\nTotal examples: {len(all_examples)}")
    print("Splitting and writing to data/...\n")
    split_and_write(all_examples, Path("data"))

    # Label distribution report
    from collections import Counter
    counts = Counter(ex["label"] for ex in all_examples)
    labels = {0: "broad_scan", 1: "targeted_search", 2: "pinpoint_cite"}
    print("\nLabel distribution:")
    for lbl, name in labels.items():
        print(f"  {name:<20s}: {counts[lbl]}")


if __name__ == "__main__":
    main()