synth-ai 0.2.9.dev3__py3-none-any.whl → 0.2.9.dev4__py3-none-any.whl

examples/warming_up_to_rl/configs/rl_from_ft.toml

@@ -0,0 +1,56 @@
+# RL training starting from a finetuned model id (TOML-only model selection)
+
+[services]
+# Task app base URL used by the RL job for rollouts
+# task_url = "https://YOUR-TASK-APP.modal.run"
+
+[compute]
+# Cluster shape for RL pipeline
+gpu_type = "H100"
+gpu_count = 8
+
+[topology]
+# Split GPUs across vLLM, training, and reference
+# Must sum to compute.gpu_count
+#gpus_for_vllm = 4
+#gpus_for_training = 3
+#gpus_for_ref = 1
+
+[vllm]
+# Serving tensor parallel size
+# tensor_parallel_size = 4
+
+[model]
+# Finetuned model id to continue training from (required for this config)
+# source = "ft:YOUR_FT_MODEL_ID"
+label = "crafter-rl-from-ft"
+
+[rollout]
+max_turns = 10
+episodes_per_batch = 64
+
+[evaluation]
+# Run baseline evaluation on the first 100 task seeds every 20 iterations
+instances = 100
+every_n_iters = 20
+seeds = [
+  0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
+  10, 11, 12, 13, 14, 15, 16, 17, 18, 19,
+  20, 21, 22, 23, 24, 25, 26, 27, 28, 29,
+  30, 31, 32, 33, 34, 35, 36, 37, 38, 39,
+  40, 41, 42, 43, 44, 45, 46, 47, 48, 49,
+  50, 51, 52, 53, 54, 55, 56, 57, 58, 59,
+  60, 61, 62, 63, 64, 65, 66, 67, 68, 69,
+  70, 71, 72, 73, 74, 75, 76, 77, 78, 79,
+  80, 81, 82, 83, 84, 85, 86, 87, 88, 89,
+  90, 91, 92, 93, 94, 95, 96, 97, 98, 99,
+]
+
+[training]
+log_interval = 1
+# Additional RL hyperparameters can go here
+
+[training.weight_sync]
+enable  = true
+targets = ["policy"]
+weight_sync_interval = 1

examples/warming_up_to_rl/export_trace_sft.py

@@ -0,0 +1,541 @@
+#!/usr/bin/env python3
+"""Export behavioural-cloning datasets from tracing_v3 SQLite traces with filters."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sqlite3
+import sys
+from collections import Counter, defaultdict
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Set, Tuple
+
+Row = sqlite3.Row
+
+
+def connect(db_path: Path) -> sqlite3.Connection:
+    conn = sqlite3.connect(str(db_path))
+    conn.row_factory = sqlite3.Row
+    return conn
+
+
+def _parse_json(value: Any) -> Any:
+    if value is None:
+        return None
+    if isinstance(value, (dict, list)):
+        return value
+    try:
+        return json.loads(value)
+    except Exception:
+        return None
+
+
+AchievementMap = dict[Tuple[str, int], dict[str, list[str]]]
+
+
+def fetch_achievement_data(
+    conn: sqlite3.Connection,
+) -> tuple[
+    AchievementMap,
+    Counter,
+    Counter,
+    Counter,
+    dict[str, set[str]],
+    dict[str, set[str]],
+]:
+    achievements_map: AchievementMap = defaultdict(lambda: {"unique": [], "all": []})
+    session_unique_sets: dict[str, set[str]] = defaultdict(set)
+    session_final_achievements: dict[str, set[str]] = defaultdict(set)
+    achievement_name_counts: Counter = Counter()
+
+    rows = conn.execute(
+        """
+        SELECT er.session_id, er.reward_value, er.annotation, ev.metadata
+        FROM event_rewards er
+        JOIN events ev ON er.event_id = ev.id
+        WHERE er.reward_type = 'unique_achievement_delta' AND er.reward_value > 0
+        """
+    ).fetchall()
+    for row in rows:
+        session_id = row["session_id"]
+        annotation = _parse_json(row["annotation"]) or {}
+        metadata = _parse_json(row["metadata"]) or {}
+        turn = metadata.get("turn")
+        if turn is None:
+            continue
+        new_unique = annotation.get("new_unique") or []
+        if not isinstance(new_unique, list):
+            continue
+        if new_unique:
+            achievements_map[(session_id, int(turn))]["unique"].extend(new_unique)
+            session_unique_sets[session_id].update(new_unique)
+
+    rows = conn.execute(
+        """
+        SELECT er.session_id, er.reward_value, er.annotation, ev.metadata
+        FROM event_rewards er
+        JOIN events ev ON er.event_id = ev.id
+        WHERE er.reward_type = 'achievement_delta' AND er.reward_value > 0
+        """
+    ).fetchall()
+    for row in rows:
+        session_id = row["session_id"]
+        annotation = _parse_json(row["annotation"]) or {}
+        metadata = _parse_json(row["metadata"]) or {}
+        turn = metadata.get("turn")
+        if turn is None:
+            continue
+        turned_true = annotation.get("turned_true") or []
+        if not isinstance(turned_true, list):
+            continue
+        if turned_true:
+            achievements_map[(session_id, int(turn))]["all"].extend(turned_true)
+
+    rows = conn.execute(
+        """
+        SELECT session_id, reward_metadata
+        FROM outcome_rewards
+        WHERE reward_metadata IS NOT NULL
+        """
+    ).fetchall()
+    for row in rows:
+        session_id = row["session_id"]
+        metadata = _parse_json(row["reward_metadata"])
+        if not isinstance(metadata, dict):
+            continue
+        final_achievements = metadata.get("achievements") or []
+        if isinstance(final_achievements, list):
+            cleaned = [a for a in final_achievements if isinstance(a, str)]
+            session_unique_sets[session_id].update(cleaned)
+            session_final_achievements[session_id].update(cleaned)
+
+    unique_counts_per_session: Counter = Counter()
+    for session_id, achievement_set in session_unique_sets.items():
+        unique_counts_per_session[session_id] = len(achievement_set)
+        achievement_name_counts.update(achievement_set)
+
+    achievement_size_counts: Counter = Counter()
+    for session_id, count in unique_counts_per_session.items():
+        achievement_size_counts[count] += 1
+
+    return (
+        achievements_map,
+        unique_counts_per_session,
+        achievement_name_counts,
+        achievement_size_counts,
+        session_unique_sets,
+        session_final_achievements,
+    )
+
+
+def fetch_session_models(conn: sqlite3.Connection) -> dict[str, tuple[str, str, int]]:
+    rows = conn.execute(
+        """
+        SELECT session_id, model_name, provider, COUNT(*) AS calls
+        FROM events
+        WHERE event_type = 'cais' AND model_name IS NOT NULL
+        GROUP BY session_id, model_name, provider
+        """
+    ).fetchall()
+
+    session_models: dict[str, tuple[str, str, int]] = {}
+    for row in rows:
+        session_id = row["session_id"]
+        calls = int(row["calls"] or 0)
+        current = session_models.get(session_id)
+        if current is None or calls > current[2]:
+            session_models[session_id] = (row["model_name"], row["provider"], calls)
+    return session_models
+
+
+def fetch_outcome_rewards(conn: sqlite3.Connection) -> dict[str, dict[str, Any]]:
+    rows = conn.execute(
+        """
+        SELECT session_id, total_reward, reward_metadata
+        FROM outcome_rewards
+        """
+    ).fetchall()
+
+    outcome_data: dict[str, dict[str, Any]] = {}
+    for row in rows:
+        metadata = _parse_json(row["reward_metadata"])
+        achievements = set()
+        if isinstance(metadata, dict):
+            ach = metadata.get("achievements") or []
+            if isinstance(ach, list):
+                achievements = {a for a in ach if isinstance(a, str)}
+        outcome_data[row["session_id"]] = {
+            "total_reward": float(row["total_reward"] or 0.0),
+            "achievements": achievements,
+        }
+    return outcome_data
+
+
+def fetch_event_reward_totals(conn: sqlite3.Connection) -> dict[str, dict[str, dict[str, float]]]:
+    rows = conn.execute(
+        """
+        SELECT session_id, reward_type, COUNT(*) AS events, COALESCE(SUM(reward_value), 0) AS total_value
+        FROM event_rewards
+        GROUP BY session_id, reward_type
+        """
+    ).fetchall()
+
+    event_totals: dict[str, dict[str, dict[str, float]]] = defaultdict(dict)
+    for row in rows:
+        event_totals[row["session_id"]][row["reward_type"]] = {
+            "events": int(row["events"] or 0),
+            "total": float(row["total_value"] or 0.0),
+        }
+    return event_totals
+
+
+def parse_event_filters(specs: list[str] | None) -> list[tuple[str, float]]:
+    filters: list[tuple[str, float]] = []
+    if not specs:
+        return filters
+    for spec in specs:
+        reward_type, _, min_val_str = spec.partition(":")
+        reward_type = reward_type.strip()
+        if not reward_type:
+            continue
+        min_val = 0.0
+        if min_val_str:
+            try:
+                min_val = float(min_val_str)
+            except ValueError:
+                print(f"Invalid event reward specification '{spec}'", file=sys.stderr)
+                raise SystemExit(1)
+        filters.append((reward_type, min_val))
+    return filters
+
+
+def _collect_text(parts: Iterable[dict[str, Any]] | None) -> str:
+    texts: list[str] = []
+    if not parts:
+        return ""
+    for part in parts:
+        if not isinstance(part, dict):
+            continue
+        if part.get("type") == "text":
+            text = part.get("text")
+            if isinstance(text, str) and text:
+                texts.append(text)
+    return "\n".join(texts)
+
+
+def _normalise_tool_calls(tool_calls: list[dict[str, Any]] | None) -> list[dict[str, Any]]:
+    normalised: list[dict[str, Any]] = []
+    if not tool_calls:
+        return normalised
+    for idx, call in enumerate(tool_calls):
+        if not isinstance(call, dict):
+            continue
+        entry = dict(call)
+
+        func_payload: dict[str, Any] | None = entry.get("function") if isinstance(entry.get("function"), dict) else None
+        name = entry.get("name") or (func_payload.get("name") if func_payload else None) or "tool"
+
+        args = None
+        if func_payload and "arguments" in func_payload:
+            args = func_payload.get("arguments")
+        else:
+            args = entry.get("arguments")
+            if args is None:
+                raw = entry.pop("arguments_json", None)
+                if isinstance(raw, str):
+                    try:
+                        args = json.loads(raw)
+                    except Exception:
+                        args = raw
+
+        if isinstance(args, (dict, list)):
+            args_str = json.dumps(args, ensure_ascii=False)
+        elif isinstance(args, str):
+            args_str = args
+        elif args is None:
+            args_str = "{}"
+        else:
+            args_str = str(args)
+
+        call_id = entry.get("id") or entry.get("call_id") or f"call_{idx}"
+
+        normalised.append(
+            {
+                "id": str(call_id),
+                "type": "function",
+                "function": {
+                    "name": str(name),
+                    "arguments": args_str,
+                },
+            }
+        )
+
+    return normalised
+
+
+def build_sft_dataset(
+    conn: sqlite3.Connection,
+    achievements_map: AchievementMap,
+    sessions_filter: Set[str],
+    *,
+    allowed_models: set[str] | None = None,
+    limit: int | None = None,
+) -> list[dict[str, Any]]:
+    rows = conn.execute(
+        """
+        SELECT id, session_id, metadata, model_name, provider, call_records
+        FROM events
+        WHERE event_type = 'cais' AND call_records IS NOT NULL
+        ORDER BY session_id, id
+        """
+    ).fetchall()
+
+    dataset: list[dict[str, Any]] = []
+    cumulative_unique: dict[str, int] = defaultdict(int)
+    session_turn_counters: dict[str, int] = defaultdict(int)
+
+    for row in rows:
+        session_id = row["session_id"]
+        if session_id not in sessions_filter:
+            continue
+        if allowed_models and row["model_name"] not in allowed_models:
+            continue
+
+        metadata = _parse_json(row["metadata"]) or {}
+        turn = metadata.get("turn")
+        if turn is None:
+            step_id = metadata.get("step_id")
+            if isinstance(step_id, str) and step_id.startswith("turn_"):
+                try:
+                    turn = int(step_id.split("_", 1)[1])
+                except (ValueError, IndexError):
+                    turn = None
+        if turn is None:
+            turn = session_turn_counters[session_id]
+            session_turn_counters[session_id] = turn + 1
+        else:
+            try:
+                turn = int(turn)
+            except (TypeError, ValueError):
+                continue
+            session_turn_counters[session_id] = max(session_turn_counters[session_id], turn + 1)
+
+        call_records = _parse_json(row["call_records"]) or []
+        if not isinstance(call_records, list) or not call_records:
+            continue
+
+        for record in call_records:
+            messages: list[dict[str, Any]] = []
+            for message in record.get("input_messages", []):
+                role = message.get("role", "unknown")
+                content = _collect_text(message.get("parts"))
+                if not content:
+                    continue
+                messages.append({"role": role, "content": content})
+
+            assistant_content = ""
+            assistant_tool_calls: list[dict[str, Any]] = []
+
+            output_text = record.get("output_text")
+            parsed_response: dict[str, Any] | None = None
+            if isinstance(output_text, str) and output_text:
+                try:
+                    parsed_response = json.loads(output_text)
+                except json.JSONDecodeError:
+                    parsed_response = None
+
+            if parsed_response:
+                choices = parsed_response.get("choices") or []
+                if choices:
+                    message = choices[0].get("message") or {}
+                    assistant_content = message.get("content") or ""
+                    assistant_tool_calls = _normalise_tool_calls(message.get("tool_calls"))
+
+            if not assistant_tool_calls:
+                assistant_tool_calls = _normalise_tool_calls(record.get("output_tool_calls"))
+
+            assistant_message: dict[str, Any] = {"role": "assistant", "content": assistant_content or ""}
+            if assistant_tool_calls:
+                assistant_message["tool_calls"] = assistant_tool_calls
+
+            if assistant_message.get("content") == "" and not assistant_message.get("tool_calls"):
+                continue
+
+            messages.append(assistant_message)
+
+            if len(messages) < 2:
+                continue
+
+            achievements = achievements_map.get((session_id, turn), {"unique": [], "all": []})
+            cumulative_unique[session_id] += len(achievements.get("unique", []))
+
+            metadata = {
+                "session_id": session_id,
+                "turn": turn,
+                "model": row["model_name"],
+                "provider": row["provider"] or "unknown",
+                "achievements": {
+                    "new_unique": achievements.get("unique", []),
+                    "turned_true": achievements.get("all", []),
+                    "cumulative_unique": cumulative_unique[session_id],
+                },
+            }
+
+            dataset.append({"messages": messages, "metadata": metadata})
+            if limit is not None and len(dataset) >= limit:
+                return dataset
+
+    return dataset
+
+
+def write_jsonl(path: Path, records: Iterable[dict[str, Any]]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as fh:
+        for record in records:
+            json.dump(record, fh, ensure_ascii=False)
+            fh.write("\n")
+
+
+def _validate_dataset(records: list[dict[str, Any]]) -> None:
+    errors: list[str] = []
+    for idx, record in enumerate(records, start=1):
+        messages = record.get("messages")
+        if not isinstance(messages, list) or not messages:
+            errors.append(f"row {idx}: missing messages list")
+            if len(errors) >= 20:
+                break
+            continue
+        for msg_idx, msg in enumerate(messages):
+            if not isinstance(msg, dict):
+                errors.append(f"row {idx}: message {msg_idx} is not an object")
+                break
+            if "role" not in msg or "content" not in msg:
+                errors.append(f"row {idx}: message {msg_idx} missing role/content")
+                break
+            if not isinstance(msg["role"], str):
+                errors.append(f"row {idx}: message {msg_idx} role not string")
+                break
+            if not isinstance(msg["content"], str):
+                errors.append(f"row {idx}: message {msg_idx} content not string")
+                break
+        if len(errors) >= 20:
+            break
+    if errors:
+        summary = "\n - ".join(errors)
+        raise SystemExit(f"Validation error while exporting dataset:\n - {summary}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--db", type=Path, default=Path("traces/v3/synth_ai.db"), help="Path to tracing_v3 SQLite DB")
+    parser.add_argument("--output", type=Path, required=True, help="Destination JSONL path for the exported dataset")
+    parser.add_argument("--model", action="append", dest="models", help="Restrict to sessions whose dominant model matches (repeatable)")
+    parser.add_argument("--provider", action="append", dest="providers", help="Restrict to sessions whose dominant provider matches (repeatable)")
+    parser.add_argument("--min-unique", type=int, default=None, help="Minimum unique achievements per session")
+    parser.add_argument("--max-unique", type=int, default=None, help="Maximum unique achievements per session")
+    parser.add_argument(
+        "--exclude-achievement",
+        action="append",
+        dest="exclude_achievements",
+        help="Achievements to ignore when evaluating --min-unique/--max-unique (repeatable)",
+    )
+    parser.add_argument("--require-achievement", action="append", dest="required_achievements", help="Require these outcome achievements (repeatable)")
+    parser.add_argument("--min-outcome-reward", type=float, default=None, help="Minimum total outcome reward per session")
+    parser.add_argument("--max-outcome-reward", type=float, default=None, help="Maximum total outcome reward per session")
+    parser.add_argument("--event-reward", action="append", dest="event_reward_filters", help="Require reward_type[:min_total] in event_rewards (repeatable)")
+    parser.add_argument("--limit", type=int, default=None, help="Maximum number of examples to emit")
+    args = parser.parse_args()
+
+    if not args.db.exists():
+        print(f"Database not found: {args.db}", file=sys.stderr)
+        raise SystemExit(1)
+
+    conn = connect(args.db)
+    try:
+        (
+            achievements_map,
+            unique_counts_per_session,
+            _name_counts,
+            _size_counts,
+            session_unique_sets,
+            session_final_achievements,
+        ) = fetch_achievement_data(conn)
+        session_models = fetch_session_models(conn)
+        outcome_data = fetch_outcome_rewards(conn)
+        event_totals = fetch_event_reward_totals(conn)
+        event_filters = parse_event_filters(args.event_reward_filters)
+
+        allowed_models = set(args.models) if args.models else None
+        allowed_providers = set(args.providers) if args.providers else None
+        required_achievements = set(args.required_achievements or [])
+        excluded_achievements = set(args.exclude_achievements or [])
+
+        eligible_sessions: set[str] = set()
+        for session_id, (model_name, provider, _calls) in session_models.items():
+            if allowed_models and model_name not in allowed_models:
+                continue
+            if allowed_providers and (provider or "unknown") not in allowed_providers:
+                continue
+
+            session_uniques = session_unique_sets.get(session_id, set())
+            adjusted_uniques = {a for a in session_uniques if a not in excluded_achievements}
+            unique_count = len(adjusted_uniques)
+            if args.min_unique is not None and unique_count < args.min_unique:
+                continue
+            if args.max_unique is not None and unique_count > args.max_unique:
+                continue
+
+            outcome = outcome_data.get(session_id)
+            total_reward = outcome["total_reward"] if outcome else 0.0
+            final_achievements = outcome["achievements"] if outcome else session_final_achievements.get(session_id, set())
+
+            if args.min_outcome_reward is not None and total_reward < args.min_outcome_reward:
+                continue
+            if args.max_outcome_reward is not None and total_reward > args.max_outcome_reward:
+                continue
+            if required_achievements and not required_achievements.issubset(final_achievements):
+                continue
+
+            session_event_totals = event_totals.get(session_id, {})
+            meets_event_filters = True
+            for reward_type, min_total in event_filters:
+                total = session_event_totals.get(reward_type, {}).get("total", 0.0)
+                if total < min_total:
+                    meets_event_filters = False
+                    break
+            if not meets_event_filters:
+                continue
+
+            eligible_sessions.add(session_id)
+
+        if not eligible_sessions:
+            print("No sessions matched the provided filters.", file=sys.stderr)
+            raise SystemExit(1)
+
+        dataset = build_sft_dataset(
+            conn,
+            achievements_map,
+            eligible_sessions,
+            allowed_models=allowed_models,
+            limit=args.limit,
+        )
+
+        if not dataset:
+            print("No rollout steps matched the filters (after session selection).", file=sys.stderr)
+            raise SystemExit(1)
+
+        _validate_dataset(dataset)
+        write_jsonl(args.output, dataset)
+        session_ids = {item.get("metadata", {}).get("session_id") for item in dataset}
+        session_ids.discard(None)
+        print(
+            f"Wrote {len(dataset)} examples from {len(session_ids)} session(s) -> {args.output}",
+            file=sys.stderr,
+        )
+    finally:
+        conn.close()
+
+
+if __name__ == "__main__":
+    main()

examples/warming_up_to_rl/groq_test.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+"""Quick smoke test that drives a rollout through the Groq proxy-backed Crafter Task App."""
+
+import argparse
+import asyncio
+import os
+from typing import Any
+
+from synth_ai.task import (
+    INTERACT_TOOL_SCHEMA,
+    RolloutEnvSpec,
+    RolloutPolicySpec,
+    RolloutRequest,
+    TaskAppClient,
+    to_jsonable,
+)
+
+
+def _build_policy_payload(seed: int, model: str) -> dict[str, Any]:
+    return {
+        "model": model,
+        "tools": INTERACT_TOOL_SCHEMA,
+        "messages": [
+            {
+                "role": "system",
+                "content": "You control the Crafter agent. Think briefly, then call the interact tool with 3-5 actions to maximize achievements.",
+            },
+            {
+                "role": "user",
+                "content": (
+                    "Environment seed {seed}. Plan initial survival/crafting steps and then call interact with concrete actions."
+                ).format(seed=seed),
+            },
+        ],
+    }
+
+
+async def run(args: argparse.Namespace) -> None:
+    client = TaskAppClient(args.base_url, api_key=args.api_key, timeout=args.timeout)
+
+    health = await client.health()
+    print("/health →", to_jsonable(health))
+
+    info = await client.info()
+    print("/info →", to_jsonable(info))
+
+    inference_url = args.inference_url or f"{args.base_url.rstrip('/')}/proxy/groq"
+
+    request = RolloutRequest(
+        run_id=args.run_id,
+        env=RolloutEnvSpec(env_name="crafter", seed=args.seed, config={"seed": args.seed}),
+        policy=RolloutPolicySpec(
+            policy_name="groq-smoke",
+            config={"model": args.model, "inference_url": inference_url.rstrip("/")},
+        ),
+        ops=[
+            {"type": "policy", "payload": _build_policy_payload(args.seed, args.model)},
+            {"type": "env"},
+        ],
+    )
+
+    response = await client.rollout(request)
+    print("rollout.metrics →", to_jsonable(response.metrics.model_dump()))
+    for idx, step in enumerate(response.trajectories[0].steps, start=1):
+        print(f"step[{idx}] tool_calls={step.tool_calls} reward={step.reward} info={to_jsonable(step.info)}")
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--base-url", default=os.getenv("TASK_APP_BASE_URL", "http://localhost:8000"))
+    parser.add_argument("--api-key", default=os.getenv("TASK_APP_API_KEY"), required=os.getenv("TASK_APP_API_KEY") is None)
+    parser.add_argument("--model", default=os.getenv("GROQ_MODEL", "groq/mixtral-8x7b"))
+    parser.add_argument("--inference-url", default=os.getenv("TASK_APP_INFERENCE_URL"))
+    parser.add_argument("--seed", type=int, default=int(os.getenv("CRAFTER_TEST_SEED", "42")))
+    parser.add_argument("--run-id", default=os.getenv("TASK_APP_RUN_ID", "groq-test"))
+    parser.add_argument("--timeout", type=float, default=float(os.getenv("TASK_APP_TIMEOUT", "60")))
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    asyncio.run(run(args))
+
+
+if __name__ == "__main__":
+    main()
+