src/gold_traces.py

"""
Gold-standard search agent traces — hand-crafted examples using real code
from the corpus. These define the quality bar for the synthetic trace generator.

Each trace follows this format:
  1. System prompt (agent role definition)
  2. User query (from the bigger model)
  3. Agent response:
     a. <|reasoning|> — decompose the query, plan searches
     b. <|search|> — issue a search query
     c. <|result|> — retrieved code chunks (injected by the harness)
     d. <|reasoning|> — analyze results, decide if more search needed
     e. (repeat b-d as needed)
     f. <|evidence|> — curated evidence package
     g. <|finish|> — signal completion

The agent learns to:
  - Break complex queries into subqueries
  - Search with specific terms (function names, types, patterns)
  - Judge which results are relevant vs noise
  - Extract key facts concisely (not dump entire code)
  - Synthesize across multiple results
  - Know when to stop searching
"""

import json
import os

PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
GOLD_PATH = os.path.join(PROJECT_DIR, "data", "gold_traces.jsonl")

SYSTEM_PROMPT = (
    "You are a code search agent. Given a query from a reasoning model, "
    "decompose it into subqueries, search the codebase, inspect results, "
    "and return curated evidence. Use <|search|> to issue searches, "
    "<|reasoning|> to analyze, and <|evidence|> to return findings. "
    "Be concise. Extract only the relevant facts. End with <|finish|>."
)

GOLD_TRACES = [
    # ─── Trace 1: Simple function lookup (C) ───────────────────────────────
    {
        "query": "How does nginx handle reusable connections?",
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "How does nginx handle reusable connections?"},
            {"role": "assistant", "content": (
                "<|reasoning|>The query asks about reusable connections in nginx. "
                "I should search for functions related to reusable connections.<|end|>\n"
                "<|search|>reusable connection<|end|>"
            )},
            {"role": "result", "content": (
                "ngx_reusable_connection(ngx_connection_t *c, ngx_uint_t reusable)\n"
                "{\n"
                "    ngx_log_debug1(NGX_LOG_DEBUG_CORE, c->log, 0,\n"
                "                   \"reusable connection: %ui\", reusable);\n"
                "\n"
                "    if (c->reusable) {\n"
                "        ngx_queue_remove(&c->queue);\n"
                "        ngx_cycle->reusable_connections_n--;\n"
                "    }\n"
                "\n"
                "    c->reusable = reusable;\n"
                "\n"
                "    if (reusable) {\n"
                "        ngx_queue_insert_head(\n"
                "            (ngx_queue_t *) &ngx_cycle->reusable_connections_queue, &c->queue);\n"
                "        ngx_cycle->reusable_connections_n++;\n"
                "    }\n"
                "}"
            )},
            {"role": "assistant", "content": (
                "<|reasoning|>This function manages the reusable connection queue. "
                "When a connection is marked reusable, it's added to "
                "reusable_connections_queue and the counter is incremented. "
                "When unreusable, it's removed from the queue and counter decremented. "
                "I have enough to answer.<|end|>\n"
                "<|evidence|>nginx manages reusable connections via `ngx_reusable_connection()`:\n"
                "- Marks connection as reusable/unreusable via the `reusable` flag\n"
                "- When reusable: inserts into `reusable_connections_queue` (head), increments `reusable_connections_n`\n"
                "- When unreusable: removes from queue, decrements counter\n"
                "- Uses `ngx_queue_insert_head` and `ngx_queue_remove` for queue management\n"
                "Source: ngx_connection.c<|end|>\n"
                "<|finish|>"
            )},
        ],
    },
    # ─── Trace 2: Multi-step search (Python) ──────────────────────────────
    {
        "query": "What does the with_params_help decorator do and how is it used?",
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "What does the with_params_help decorator do and how is it used?"},
            {"role": "assistant", "content": (
                "<|reasoning|>This asks about a specific decorator. I need to find its definition first, then find usages.<|end|>\n"
                "<|search|>with_params_help<|end|>"
            )},
            {"role": "result", "content": (
                "def with_params_help(params_cls: type, title: str = \"Parameter defaults\"):\n"
                "    \"\"\"\n"
                "    Class decorator that appends a help table generated from another dataclass\n"
                "    (e.g., VllmParameters) to the decorated class's docstring.\n"
                "    \"\"\"\n"
                "    if not is_dataclass(params_cls):\n"
                "        raise TypeError(f\"{params_cls} must be a dataclass\")\n"
                "\n"
                "    def _decorator(cls: type) -> type:\n"
                "        block = generate_dataclass_help(params_cls)\n"
                "        cls.__doc__ = (cls.__doc__ or \"\") + f\"\\n\\n{title}:\\n{block}\"\n"
                "        return cls\n"
                "\n"
                "    return _decorator"
            )},
            {"role": "assistant", "content": (
                "<|reasoning|>Found the definition. It's a decorator factory that takes a dataclass and appends a help table to another class's docstring. Now I need to find where it's used to show a concrete example.<|end|>\n"
                "<|search|>@with_params_help<|end|>"
            )},
            {"role": "result", "content": (
                "@with_params_help(VllmBuildParameters)\n"
                "class VllmBuildRunner(BaseRunner):\n"
                "    \"\"\"\n"
                "    Build vLLM using docker buildx.\n"
                "    \"\"\""
            )},
            {"role": "assistant", "content": (
                "<|reasoning|>Found a usage. The decorator is applied to VllmBuildRunner with VllmBuildParameters as the params class. This appends parameter defaults from VllmBuildParameters to VllmBuildRunner's docstring. I have enough to answer.<|end|>\n"
                "<|evidence|>`with_params_help` is a class decorator factory:\n"
                "- Takes a dataclass (`params_cls`) and optional title (default \"Parameter defaults\")\n"
                "- Generates a help table from the dataclass fields via `generate_dataclass_help()`\n"
                "- Appends the table to the decorated class's `__doc__` docstring\n"
                "- Usage example: `@with_params_help(VllmBuildParameters)` on `VllmBuildRunner`\n"
                "- Requires `params_cls` to be a dataclass (raises TypeError otherwise)<|end|>\n"
                "<|finish|>"
            )},
        ],
    },
    # ─── Trace 3: Type/struct lookup with noise (C) ───────────────────────
    {
        "query": "What fields does the nginx listening socket structure have?",
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "What fields does the nginx listening socket structure have?"},
            {"role": "assistant", "content": (
                "<|reasoning|>Looking for the listening socket struct definition in nginx. Should search for ngx_listening struct.<|end|>\n"
                "<|search|>ngx_listening struct fields<|end|>"
            )},
            {"role": "result", "content": (
                "typedef struct ngx_listening_s  ngx_listening_t;\n"
                "\n"
                "struct ngx_listening_s {\n"
                "    ngx_socket_t        fd;\n"
                "    struct sockaddr    *sockaddr;\n"
                "    socklen_t           socklen;\n"
                "    size_t              addr_text_max_len;\n"
                "    ngx_str_t           addr_text;\n"
                "    int                 type;\n"
                "    int                 protocol;\n"
                "    int                 backlog;\n"
                "    int                 rcvbuf;\n"
                "    int                 sndbuf;\n"
                "};"
            )},
            {"role": "result", "content": (
                "ngx_create_listening(ngx_conf_t *cf, struct sockaddr *sockaddr,\n"
                "    socklen_t socklen)\n"
                "{\n"
                "    ls = ngx_array_push(&cf->cycle->listening);\n"
                "    ngx_memzero(ls, sizeof(ngx_listening_t));\n"
                "    ls->sockaddr = sa;\n"
                "    ls->socklen = socklen;\n"
                "}"
            )},
            {"role": "assistant", "content": (
                "<|reasoning|>The first result is the struct definition — exactly what I need. The second result is the creation function, which is related but not asked for. I'll focus on the struct fields.<|end|>\n"
                "<|evidence|>`ngx_listening_s` (aka `ngx_listening_t`) structure fields:\n"
                "- `fd` (ngx_socket_t): file descriptor\n"
                "- `sockaddr` (struct sockaddr*): socket address pointer\n"
                "- `socklen` (socklen_t): socket address length\n"
                "- `addr_text_max_len` (size_t): max address text length\n"
                "- `addr_text` (ngx_str_t): address text representation\n"
                "- `type` (int): socket type\n"
                "- `protocol` (int): protocol\n"
                "- `backlog` (int): listen backlog\n"
                "- `rcvbuf` (int): receive buffer size\n"
                "- `sndbuf` (int): send buffer size\n"
                "Source: ngx_connection.h<|end|>\n"
                "<|finish|>"
            )},
        ],
    },
    # ─── Trace 4: Query decomposition (JS/TS) ─────────────────────────────
    {
        "query": "How does the concatenate function work and what types does it accept?",
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "How does the concatenate function work and what types does it accept?"},
            {"role": "assistant", "content": (
                "<|reasoning|>This has two parts: (1) how concatenate works, (2) what types it accepts. Let me search for the function definition.<|end|>\n"
                "<|search|>concatenate function array<|end|>"
            )},
            {"role": "result", "content": (
                "export function concatenate<T>(array1: T[], array2: T[] | undefined): T[];\n"
                "// eslint-disable-line @typescript-eslint/unified-signatures"
            )},
            {"role": "result", "content": (
                "function concatenate<T>(array1: T[], array2: T[] | undefined): T[] {\n"
                "    if (array2 === undefined) return array1;\n"
                "    return array1.concat(array2);\n"
                "}"
            )},
            {"role": "assistant", "content": (
                "<|reasoning|>Found both the declaration and implementation. It's a generic TypeScript function that concatenates two arrays. The second parameter is optional (can be undefined). If undefined, it returns the first array unchanged. Otherwise uses Array.concat. I have everything needed.<|end|>\n"
                "<|evidence|>`concatenate<T>` is a TypeScript utility function:\n"
                "- Signature: `concatenate<T>(array1: T[], array2: T[] | undefined): T[]`\n"
                "- Generic over type T — accepts arrays of any type\n"
                "- If `array2` is `undefined`, returns `array1` unchanged (no-op)\n"
                "- Otherwise returns `array1.concat(array2)` (standard Array.concat)\n"
                "- Exported as a public function<|end|>\n"
                "<|finish|>"
            )},
        ],
    },
    # ─── Trace 5: Search with irrelevant results (Rust) ───────────────────
    {
        "query": "How is error handling done in the Rust codebase?",
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "How is error handling done in the Rust codebase?"},
            {"role": "assistant", "content": (
                "<|reasoning|>Error handling in Rust typically uses Result types. Let me search for Result usage and error types.<|end|>\n"
                "<|search|>Result error handling<|end|>"
            )},
            {"role": "result", "content": (
                "pub fn parse_config(input: &str) -> Result<Config, ParseError> {\n"
                "    let parts: Vec<&str> = input.split_whitespace().collect();\n"
                "    if parts.len() < 2 {\n"
                "        return Err(ParseError::TooFewArguments);\n"
                "    }\n"
                "    Ok(Config { query: parts[0].to_string(), filename: parts[1].to_string() })\n"
                "}"
            )},
            {"role": "result", "content": (
                "use std::fmt;\n"
                "\n"
                "pub enum ParseError {\n"
                "    TooFewArguments,\n"
                "    InvalidSyntax,\n"
                "}\n"
                "\n"
                "impl fmt::Display for ParseError {\n"
                "    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {\n"
                "        match self {\n"
                "            ParseError::TooFewArguments => write!(f, \"too few arguments\"),\n"
                "            ParseError::InvalidSyntax => write!(f, \"invalid syntax\"),\n"
                "        }\n"
                "    }\n"
                "}"
            )},
            {"role": "result", "content": (
                "mod tests {\n"
                "    #[test]\n"
                "    fn test_basic() {\n"
                "        assert_eq!(2 + 2, 4);\n"
                "    }\n"
                "}"
            )},
            {"role": "assistant", "content": (
                "<|reasoning|>Result 1 shows a function returning Result<Config, ParseError> with early return on error. Result 2 shows the ParseError enum with Display impl. Result 3 is an unrelated test — not relevant. I have enough to answer.<|end|>\n"
                "<|evidence|>Error handling pattern in this Rust codebase:\n"
                "- Uses `Result<T, E>` for fallible operations\n"
                "- Custom error type: `ParseError` enum with variants (TooFewArguments, InvalidSyntax)\n"
                "- Implements `std::fmt::Display` for human-readable errors\n"
                "- Early return pattern: `return Err(ParseError::...)` on validation failure\n"
                "- Success path: `Ok(Config { ... })`\n"
                "- No use of `unwrap()` or `expect()` in the found code — proper error propagation<|end|>\n"
                "<|finish|>"
            )},
        ],
    },
    # ─── Trace 6: No results found ────────────────────────────────────────
    {
        "query": "Where is the database connection pool implemented?",
        "trace": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": "Where is the database connection pool implemented?"},
            {"role": "assistant", "content": (
                "<|reasoning|>Looking for database connection pool implementation. Let me search for relevant terms.<|end|>\n"
                "<|search|>database connection pool<|end|>"
            )},
            {"role": "result", "content": ""},
            {"role": "assistant", "content": (
                "<|reasoning|>No results found for 'database connection pool'. Let me try broader terms.<|end|>\n"
                "<|search|>pool db<|end|>"
            )},
            {"role": "result", "content": ""},
            {"role": "assistant", "content": (
                "<|reasoning|>Still no results. The codebase does not appear to contain a database connection pool implementation.<|end|>\n"
                "<|evidence|>No database connection pool implementation was found in the codebase. Searched for: 'database connection pool', 'pool db'. The indexed code may not include database layer code.<|end|>\n"
                "<|finish|>"
            )},
        ],
    },
]


def main():
    print(f"Writing {len(GOLD_TRACES)} gold-standard traces to {GOLD_PATH}...")
    with open(GOLD_PATH, "w", encoding="utf-8") as f:
        for trace in GOLD_TRACES:
            f.write(json.dumps(trace) + "\n")

    print(f"Done. {len(GOLD_TRACES)} traces written.")

    # Print summary
    for i, trace in enumerate(GOLD_TRACES):
        query = trace["query"]
        n_steps = len([m for m in trace["trace"] if m["role"] == "assistant"])
        n_searches = len([m for m in trace["trace"] if m["role"] == "assistant" and "<|search|>" in m["content"]])
        n_results = len([m for m in trace["trace"] if m["role"] == "result"])
        print(f"  Trace {i+1}: '{query[:50]}...' | {n_steps} agent turns | {n_searches} searches | {n_results} results")


if __name__ == "__main__":
    main()