scripts/generate_dataset.py

"""
Synthetic Trajectory Generator for Gemma Fine-Tuning.
Generates multi-step tool execution sequences for agent training.
"""

import json
from pathlib import Path

def create_full_trajectory(symbol: str, file_path: str, bug_desc: str) -> dict:
    return {
        "messages": [
            {
                "role": "user",
                "content": f"Investigate and resolve issue in symbol '{symbol}': {bug_desc}"
            },
            {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "type": "function",
                        "function": {
                            "name": "code_graph_search",
                            "parameters": {"symbol_name": symbol}
                        }
                    }
                ]
            },
            {
                "role": "tool",
                "content": json.dumps({
                    "status": "success",
                    "query": symbol,
                    "definitions": [{"file": file_path, "type": "class", "lineno": 10}]
                })
            },
            {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "type": "function",
                        "function": {
                            "name": "read_scope",
                            "parameters": {"file_path": file_path, "start_line": 10, "end_line": 35}
                        }
                    }
                ]
            },
            {
                "role": "tool",
                "content": json.dumps({
                    "status": "success",
                    "file_path": file_path,
                    "content": "def execute(self):\n    pass\n"
                })
            },
            {
                "role": "assistant",
                "content": f"Located `{symbol}` in `{file_path}`. Verified scope and ready for resolution."
            }
        ]
    }

def main():
    output_path = Path("data/train_trajectories.jsonl")
    output_path.parent.mkdir(parents=True, exist_ok=True)

    samples = [
        create_full_trajectory("FileOperations", "src/tools/file_ops.py", "Boundary check validation failure."),
        create_full_trajectory("RepositoryASTParser", "src/tools/ast_parser.py", "AST scope extraction missing line bounds."),
        create_full_trajectory("CodeGraphIndexer", "src/tools/code_graph.py", "Symbol indexing cache invalidation issue."),
        create_full_trajectory("PyTestExecutor", "src/tools/executor.py", "Subprocess execution timeout handling."),
        create_full_trajectory("ContextCompactor", "src/agent/compactor.py", "Token limit calculation drift.")
    ]

    with open(output_path, "w", encoding="utf-8") as f:
        for sample in samples:
            f.write(json.dumps(sample) + "\n")

    print(f"Generated {len(samples)} multi-step training trajectories at {output_path}")

if __name__ == "__main__":
    main()