334 lines
11 KiB
Python
334 lines
11 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Export Timmy session logs as LoRA training data (ChatML JSONL).
|
||
|
|
|
||
|
|
Reads session JSONL files written by ``SessionLogger`` and converts them into
|
||
|
|
conversation pairs suitable for fine-tuning with ``mlx_lm.lora``.
|
||
|
|
|
||
|
|
Output format — one JSON object per line::
|
||
|
|
|
||
|
|
{"messages": [
|
||
|
|
{"role": "system", "content": "<Timmy system prompt>"},
|
||
|
|
{"role": "user", "content": "<user turn>"},
|
||
|
|
{"role": "assistant", "content": "<timmy response, with tool calls embedded>"}
|
||
|
|
]}
|
||
|
|
|
||
|
|
Tool calls that appear between a user turn and the next assistant message are
|
||
|
|
embedded in the assistant content using the Hermes 4 ``<tool_call>`` XML format
|
||
|
|
so the fine-tuned model learns both when to call tools and what JSON to emit.
|
||
|
|
|
||
|
|
Usage::
|
||
|
|
|
||
|
|
# Export all session logs (default paths)
|
||
|
|
python scripts/export_trajectories.py
|
||
|
|
|
||
|
|
# Custom source / destination
|
||
|
|
python scripts/export_trajectories.py \\
|
||
|
|
--logs-dir ~/custom-logs \\
|
||
|
|
--output ~/timmy-training-data.jsonl \\
|
||
|
|
--min-turns 2 \\
|
||
|
|
--verbose
|
||
|
|
|
||
|
|
Epic: #1091 Project Bannerlord — AutoLoRA Sovereignty Loop (Step 3 of 7)
|
||
|
|
Refs: #1103
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import logging
|
||
|
|
import sys
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
# ── Constants ─────────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
TIMMY_SYSTEM_PROMPT = (
|
||
|
|
"You are Timmy, Alexander's personal AI agent running on a local Mac. "
|
||
|
|
"You are concise, direct, and action-oriented. "
|
||
|
|
"You have access to a broad set of tools — use them proactively. "
|
||
|
|
"When you need to call a tool, output it in this format:\n"
|
||
|
|
"<tool_call>\n"
|
||
|
|
'{"name": "function_name", "arguments": {"param": "value"}}\n'
|
||
|
|
"</tool_call>\n\n"
|
||
|
|
"Always provide structured, accurate responses."
|
||
|
|
)
|
||
|
|
|
||
|
|
# ── Entry grouping ─────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
|
||
|
|
def _load_entries(logs_dir: Path) -> list[dict[str, Any]]:
|
||
|
|
"""Load all session log entries, sorted chronologically."""
|
||
|
|
entries: list[dict[str, Any]] = []
|
||
|
|
log_files = sorted(logs_dir.glob("session_*.jsonl"))
|
||
|
|
for log_file in log_files:
|
||
|
|
try:
|
||
|
|
with open(log_file) as f:
|
||
|
|
for line in f:
|
||
|
|
line = line.strip()
|
||
|
|
if not line:
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
entries.append(json.loads(line))
|
||
|
|
except json.JSONDecodeError:
|
||
|
|
logger.warning("Skipping malformed line in %s", log_file.name)
|
||
|
|
except OSError as exc:
|
||
|
|
logger.warning("Cannot read %s: %s", log_file, exc)
|
||
|
|
return entries
|
||
|
|
|
||
|
|
|
||
|
|
def _format_tool_call(entry: dict[str, Any]) -> str:
|
||
|
|
"""Render a tool_call entry as a Hermes 4 <tool_call> XML block."""
|
||
|
|
payload = {"name": entry.get("tool", "unknown"), "arguments": entry.get("args", {})}
|
||
|
|
return f"<tool_call>\n{json.dumps(payload)}\n</tool_call>"
|
||
|
|
|
||
|
|
|
||
|
|
def _format_tool_result(entry: dict[str, Any]) -> str:
|
||
|
|
"""Render a tool result observation."""
|
||
|
|
result = entry.get("result", "")
|
||
|
|
tool = entry.get("tool", "unknown")
|
||
|
|
return f"<tool_response>\n{{\"name\": \"{tool}\", \"result\": {json.dumps(result)}}}\n</tool_response>"
|
||
|
|
|
||
|
|
|
||
|
|
def _group_into_turns(entries: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||
|
|
"""Group raw session entries into (user_text, assistant_parts) turn pairs.
|
||
|
|
|
||
|
|
Returns a list of dicts with keys:
|
||
|
|
``user`` - user message content
|
||
|
|
``assistant`` - assembled assistant content (responses + tool calls)
|
||
|
|
"""
|
||
|
|
turns: list[dict[str, Any]] = []
|
||
|
|
pending_user: str | None = None
|
||
|
|
assistant_parts: list[str] = []
|
||
|
|
|
||
|
|
for entry in entries:
|
||
|
|
etype = entry.get("type", "")
|
||
|
|
role = entry.get("role", "")
|
||
|
|
|
||
|
|
if etype == "message" and role == "user":
|
||
|
|
# Flush any open turn
|
||
|
|
if pending_user is not None and assistant_parts:
|
||
|
|
turns.append(
|
||
|
|
{
|
||
|
|
"user": pending_user,
|
||
|
|
"assistant": "\n".join(assistant_parts).strip(),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
elif pending_user is not None:
|
||
|
|
# User message with no assistant response — discard
|
||
|
|
pass
|
||
|
|
pending_user = entry.get("content", "").strip()
|
||
|
|
assistant_parts = []
|
||
|
|
|
||
|
|
elif etype == "message" and role == "timmy":
|
||
|
|
if pending_user is not None:
|
||
|
|
content = entry.get("content", "").strip()
|
||
|
|
if content:
|
||
|
|
assistant_parts.append(content)
|
||
|
|
|
||
|
|
elif etype == "tool_call":
|
||
|
|
if pending_user is not None:
|
||
|
|
assistant_parts.append(_format_tool_call(entry))
|
||
|
|
# Also append tool result as context so model learns the full loop
|
||
|
|
if entry.get("result"):
|
||
|
|
assistant_parts.append(_format_tool_result(entry))
|
||
|
|
|
||
|
|
# decision / error entries are skipped — they are meta-data, not conversation
|
||
|
|
|
||
|
|
# Flush final open turn
|
||
|
|
if pending_user is not None and assistant_parts:
|
||
|
|
turns.append(
|
||
|
|
{
|
||
|
|
"user": pending_user,
|
||
|
|
"assistant": "\n".join(assistant_parts).strip(),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
return turns
|
||
|
|
|
||
|
|
|
||
|
|
# ── Conversion ────────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
|
||
|
|
def turns_to_training_examples(
|
||
|
|
turns: list[dict[str, Any]],
|
||
|
|
system_prompt: str = TIMMY_SYSTEM_PROMPT,
|
||
|
|
min_assistant_len: int = 10,
|
||
|
|
) -> list[dict[str, Any]]:
|
||
|
|
"""Convert grouped turns into mlx-lm training examples.
|
||
|
|
|
||
|
|
Each example has a ``messages`` list in ChatML order:
|
||
|
|
``[system, user, assistant]``.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
turns: Output of ``_group_into_turns``.
|
||
|
|
system_prompt: System prompt prepended to every example.
|
||
|
|
min_assistant_len: Skip examples where the assistant turn is shorter
|
||
|
|
than this many characters (filters out empty/trivial turns).
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
List of training example dicts.
|
||
|
|
"""
|
||
|
|
examples: list[dict[str, Any]] = []
|
||
|
|
for turn in turns:
|
||
|
|
assistant_text = turn.get("assistant", "").strip()
|
||
|
|
user_text = turn.get("user", "").strip()
|
||
|
|
if not user_text or len(assistant_text) < min_assistant_len:
|
||
|
|
continue
|
||
|
|
examples.append(
|
||
|
|
{
|
||
|
|
"messages": [
|
||
|
|
{"role": "system", "content": system_prompt},
|
||
|
|
{"role": "user", "content": user_text},
|
||
|
|
{"role": "assistant", "content": assistant_text},
|
||
|
|
]
|
||
|
|
}
|
||
|
|
)
|
||
|
|
return examples
|
||
|
|
|
||
|
|
|
||
|
|
def export_training_data(
|
||
|
|
logs_dir: Path,
|
||
|
|
output_path: Path,
|
||
|
|
min_turns: int = 1,
|
||
|
|
min_assistant_len: int = 10,
|
||
|
|
verbose: bool = False,
|
||
|
|
) -> int:
|
||
|
|
"""Full export pipeline: load → group → convert → write.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
logs_dir: Directory containing ``session_*.jsonl`` files.
|
||
|
|
output_path: Destination ``.jsonl`` file for training data.
|
||
|
|
min_turns: Minimum number of turns required (used for logging only).
|
||
|
|
min_assistant_len: Minimum assistant response length to include.
|
||
|
|
verbose: Print progress to stdout.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
Number of training examples written.
|
||
|
|
"""
|
||
|
|
if verbose:
|
||
|
|
print(f"Loading session logs from: {logs_dir}")
|
||
|
|
|
||
|
|
entries = _load_entries(logs_dir)
|
||
|
|
if verbose:
|
||
|
|
print(f" Loaded {len(entries)} raw entries")
|
||
|
|
|
||
|
|
turns = _group_into_turns(entries)
|
||
|
|
if verbose:
|
||
|
|
print(f" Grouped into {len(turns)} conversation turns")
|
||
|
|
|
||
|
|
examples = turns_to_training_examples(
|
||
|
|
turns, min_assistant_len=min_assistant_len
|
||
|
|
)
|
||
|
|
if verbose:
|
||
|
|
print(f" Generated {len(examples)} training examples")
|
||
|
|
|
||
|
|
if not examples:
|
||
|
|
print("WARNING: No training examples generated. Check that session logs exist.")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
with open(output_path, "w") as f:
|
||
|
|
for ex in examples:
|
||
|
|
f.write(json.dumps(ex) + "\n")
|
||
|
|
|
||
|
|
if verbose:
|
||
|
|
print(f" Wrote {len(examples)} examples → {output_path}")
|
||
|
|
|
||
|
|
return len(examples)
|
||
|
|
|
||
|
|
|
||
|
|
# ── CLI ───────────────────────────────────────────────────────────────────────
|
||
|
|
|
||
|
|
|
||
|
|
def _default_logs_dir() -> Path:
|
||
|
|
"""Return default logs directory (repo root / logs)."""
|
||
|
|
# Walk up from this script to find repo root (contains pyproject.toml)
|
||
|
|
candidate = Path(__file__).resolve().parent
|
||
|
|
for _ in range(5):
|
||
|
|
candidate = candidate.parent
|
||
|
|
if (candidate / "pyproject.toml").exists():
|
||
|
|
return candidate / "logs"
|
||
|
|
return Path.home() / "logs"
|
||
|
|
|
||
|
|
|
||
|
|
def _default_output_path() -> Path:
|
||
|
|
return Path.home() / "timmy-training-data.jsonl"
|
||
|
|
|
||
|
|
|
||
|
|
def main(argv: list[str] | None = None) -> int:
|
||
|
|
parser = argparse.ArgumentParser(
|
||
|
|
description="Export Timmy session logs as LoRA training data (ChatML JSONL)",
|
||
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||
|
|
epilog=__doc__,
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--logs-dir",
|
||
|
|
type=Path,
|
||
|
|
default=_default_logs_dir(),
|
||
|
|
help="Directory containing session_*.jsonl files (default: <repo>/logs)",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--output",
|
||
|
|
type=Path,
|
||
|
|
default=_default_output_path(),
|
||
|
|
help="Output JSONL path (default: ~/timmy-training-data.jsonl)",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--min-turns",
|
||
|
|
type=int,
|
||
|
|
default=1,
|
||
|
|
help="Minimum turns to process (informational, default: 1)",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--min-assistant-len",
|
||
|
|
type=int,
|
||
|
|
default=10,
|
||
|
|
help="Minimum assistant response length in chars (default: 10)",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--verbose",
|
||
|
|
"-v",
|
||
|
|
action="store_true",
|
||
|
|
help="Print progress information",
|
||
|
|
)
|
||
|
|
|
||
|
|
args = parser.parse_args(argv)
|
||
|
|
|
||
|
|
logging.basicConfig(
|
||
|
|
level=logging.DEBUG if args.verbose else logging.WARNING,
|
||
|
|
format="%(levelname)s: %(message)s",
|
||
|
|
)
|
||
|
|
|
||
|
|
if not args.logs_dir.exists():
|
||
|
|
print(f"ERROR: Logs directory not found: {args.logs_dir}")
|
||
|
|
print("Run the Timmy dashboard first to generate session logs.")
|
||
|
|
return 1
|
||
|
|
|
||
|
|
count = export_training_data(
|
||
|
|
logs_dir=args.logs_dir,
|
||
|
|
output_path=args.output,
|
||
|
|
min_turns=args.min_turns,
|
||
|
|
min_assistant_len=args.min_assistant_len,
|
||
|
|
verbose=args.verbose,
|
||
|
|
)
|
||
|
|
|
||
|
|
if count > 0:
|
||
|
|
print(f"Exported {count} training examples to: {args.output}")
|
||
|
|
print()
|
||
|
|
print("Next steps:")
|
||
|
|
print(f" mkdir -p ~/timmy-lora-training")
|
||
|
|
print(f" cp {args.output} ~/timmy-lora-training/train.jsonl")
|
||
|
|
print(f" python scripts/lora_finetune.py --data ~/timmy-lora-training")
|
||
|
|
else:
|
||
|
|
print("No training examples exported.")
|
||
|
|
return 1
|
||
|
|
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.exit(main())
|