synth-ai 0.2.13.dev2__py3-none-any.whl → 0.2.14__py3-none-any.whl

synth_ai/cli/task_apps.py

@@ -24,9 +24,9 @@ import types
 import uuid
 from collections.abc import Callable, Iterable, Iterator, Sequence
 from dataclasses import dataclass
-from datetime import UTC, datetime
+from datetime import datetime, timezone
 from pathlib import Path
-from typing import Any, cast
+from typing import Any, Optional, cast
 
 try:  # Python 3.11+
     import tomllib as _toml
@@ -36,19 +36,29 @@ except Exception:  # pragma: no cover - fallback
 import click
 from click.exceptions import Abort
 
-from synth_ai.tracing_v3 import (  # type: ignore[import-untyped]
-    BaseEvent,
-    EnvironmentEvent,
-    RuntimeEvent,
-    SessionEventMarkovBlanketMessage,
-    SessionMessageContent,
-    SessionTimeStep,
-    SessionTracer,
-    TimeRecord,
-)
-from synth_ai.tracing_v3 import (  # type: ignore[import-untyped]
-    SessionTrace as V3SessionTrace,
-)
+# Tracing imports - make conditional for optional dependencies
+try:
+    from synth_ai.tracing_v3 import (  # type: ignore[import-untyped]
+        BaseEvent,
+        EnvironmentEvent,
+        RuntimeEvent,
+        SessionEventMarkovBlanketMessage,
+        SessionMessageContent,
+        SessionTimeStep,
+        SessionTracer,
+        TimeRecord,
+    )
+    from synth_ai.tracing_v3 import (  # type: ignore[import-untyped]
+        SessionTrace as V3SessionTrace,
+    )
+    _TRACING_AVAILABLE = True
+except (ImportError, ModuleNotFoundError, TypeError):
+    # Tracing system not available (missing optional dependencies)
+    BaseEvent = EnvironmentEvent = RuntimeEvent = None  # type: ignore
+    SessionEventMarkovBlanketMessage = SessionMessageContent = None  # type: ignore
+    SessionTimeStep = SessionTracer = TimeRecord = None  # type: ignore
+    V3SessionTrace = None  # type: ignore
+    _TRACING_AVAILABLE = False
 
 # ---------------------------------------------------------------------------
 # Dynamic imports to avoid hard dependencies during type checking.
@@ -82,14 +92,14 @@ except Exception as exc:  # pragma: no cover - critical dependency
     raise RuntimeError("Unable to load task app server utilities") from exc
 
 
-def _load_demo_directory() -> Path | None:
+def _load_demo_directory() -> Optional[Path]:
     """Return the demo task apps directory if available."""
 
     try:
         module = cast(
             Any, importlib.import_module("synth_ai.demos.demo_task_apps.core")
         )
-        loader = cast(Callable[[], str | Path | None], module.load_demo_dir)
+        loader = cast(Callable[[], Optional[str | Path]], module.load_demo_dir)
         demo_dir = loader()
         if isinstance(demo_dir, str | Path):
             demo_path = Path(demo_dir)
@@ -129,7 +139,7 @@ DEFAULT_SEARCH_RELATIVE = (
 )
 
 
-def _pearson(xs: Sequence[float], ys: Sequence[float]) -> float | None:
+def _pearson(xs: Sequence[float], ys: Sequence[float]) -> Optional[float]:
     if len(xs) != len(ys) or len(xs) < 2:
         return None
     mean_x = sum(xs) / len(xs)
@@ -154,7 +164,7 @@ class AppChoice:
     label: str
     path: Path
     source: str
-    description: str | None = None
+    description: Optional[str] = None
     aliases: tuple[str, ...] = ()
     entry: TaskAppEntryType | None = None
     entry_loader: Callable[[], TaskAppEntryType] | None = None
@@ -178,21 +188,21 @@ class JudgeSpec:
     kwargs: dict[str, Any]
 
 
-def _parse_datetime_for_trace(value: Any) -> datetime | None:
+def _parse_datetime_for_trace(value: Any) -> Optional[datetime]:
     if isinstance(value, datetime):
-        return value if value.tzinfo else value.replace(tzinfo=UTC)
+        return value if value.tzinfo else value.replace(tzinfo=timezone.utc)
     if isinstance(value, str):
         value = value.replace("Z", "+00:00")
         try:
             dt = datetime.fromisoformat(value)
         except ValueError:
             try:
-                dt = datetime.fromtimestamp(float(value), tz=UTC)
+                dt = datetime.fromtimestamp(float(value), tz=timezone.utc)
             except Exception:
                 return None
-        return dt if dt.tzinfo else dt.replace(tzinfo=UTC)
+        return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)
     if isinstance(value, int | float):
-        return datetime.fromtimestamp(float(value), tz=UTC)
+        return datetime.fromtimestamp(float(value), tz=timezone.utc)
     return None
 
 
@@ -269,7 +279,7 @@ def _step_from_dict(payload: dict[str, Any]) -> SessionTimeStep:
         for msg in payload.get("markov_blanket_messages", [])
         if isinstance(msg, dict)
     ]
-    timestamp = _parse_datetime_for_trace(payload.get("timestamp")) or datetime.now(UTC)
+    timestamp = _parse_datetime_for_trace(payload.get("timestamp")) or datetime.now(timezone.utc)
     completed_at = _parse_datetime_for_trace(payload.get("completed_at"))
     return SessionTimeStep(
         step_id=payload.get("step_id", ""),
@@ -283,7 +293,7 @@ def _step_from_dict(payload: dict[str, Any]) -> SessionTimeStep:
     )
 
 
-def _session_trace_from_dict(payload: dict[str, Any]) -> V3SessionTrace | None:
+def _session_trace_from_dict(payload: dict[str, Any]) -> Optional[V3SessionTrace]:
     if not isinstance(payload, dict):
         return None
     steps = [
@@ -301,7 +311,7 @@ def _session_trace_from_dict(payload: dict[str, Any]) -> V3SessionTrace | None:
         for msg in payload.get("markov_blanket_message_history", [])
         if isinstance(msg, dict)
     ]
-    created_at = _parse_datetime_for_trace(payload.get("created_at")) or datetime.now(UTC)
+    created_at = _parse_datetime_for_trace(payload.get("created_at")) or datetime.now(timezone.utc)
     metadata = payload.get("metadata") or {}
     session_metadata = payload.get("session_metadata")
     return V3SessionTrace(
@@ -320,21 +330,43 @@ async def _store_trace(
     trace_namespace: dict[str, Any] | None,
     extra_metadata: dict[str, Any] | None = None,
 ):
+    import logging
+    _logger = logging.getLogger(__name__)
+    
+    _logger.info(f"[STORE_TRACE_DEBUG] Called with tracer={tracer is not None}, trace_namespace={trace_namespace is not None}")
+    
     if tracer is None or not isinstance(trace_namespace, dict):
+        _logger.warning(f"[STORE_TRACE_DEBUG] Early return: tracer={tracer is not None}, trace_namespace type={type(trace_namespace)}")
         return
+    
+    _logger.info(f"[STORE_TRACE_DEBUG] trace_namespace keys: {list(trace_namespace.keys())}")
+    
     session_payload = trace_namespace.get("session_trace")
     if not isinstance(session_payload, dict):
+        _logger.warning(f"[STORE_TRACE_DEBUG] No session_trace found or wrong type: {type(session_payload)}")
         return
+    
+    _logger.info(f"[STORE_TRACE_DEBUG] session_payload keys: {list(session_payload.keys())}")
+    msg_count = len(session_payload.get("markov_blanket_message_history", []))
+    _logger.info(f"[STORE_TRACE_DEBUG] Found {msg_count} messages in session_payload")
+    
     trace_obj = _session_trace_from_dict(session_payload)
     if trace_obj is None:
+        _logger.warning(f"[STORE_TRACE_DEBUG] _session_trace_from_dict returned None")
         return
+    
+    _logger.info(f"[STORE_TRACE_DEBUG] Created SessionTrace object with {len(trace_obj.markov_blanket_message_history)} messages")
+    
     if tracer.db is None:
         await tracer.initialize()
     meta = dict(trace_obj.metadata or {})
     if extra_metadata:
         meta.update(extra_metadata)
     trace_obj.metadata = meta
+    
+    _logger.info(f"[STORE_TRACE_DEBUG] Calling insert_session_trace for session_id={trace_obj.session_id}")
     await tracer.db.insert_session_trace(trace_obj)
+    _logger.info(f"[STORE_TRACE_DEBUG] Successfully inserted trace")
 
 def _temporary_sys_path(paths: Sequence[Path]):
     """Context manager to prepend entries to sys.path temporarily."""
@@ -3044,6 +3076,11 @@ def _write_modal_entrypoint(
     if not any(str(p).startswith("synth-ai") for p in pip_packages):
         pip_packages.insert(0, synth_pkg)
 
+    apt_packages = list(modal_cfg.apt_packages)
+    click.echo(f"[DEBUG] modal_cfg.apt_packages type: {type(modal_cfg.apt_packages)}")
+    click.echo(f"[DEBUG] modal_cfg.apt_packages value: {modal_cfg.apt_packages}")
+    click.echo(f"[DEBUG] apt_packages after list(): {apt_packages}")
+    
     local_dirs = [(str(Path(src)), dst) for src, dst in modal_cfg.extra_local_dirs]
     # Also mount the host synth_ai source if available to ensure latest code is used
     if host_synth is not None:
@@ -3090,6 +3127,15 @@ INLINE_SECRET_VALUES = {inline_secret_values!r}
 
 image = Image.debian_slim(python_version={modal_cfg.python_version!r})
 
+# CRITICAL: Install iverilog for Verilog task app (hardcoded to prevent config issues)
+if {entry.app_id!r} == "grpo-verilog":
+    image = image.apt_install("iverilog")
+
+# Install apt packages first (before pip)
+apt_packages = {apt_packages!r}
+if apt_packages:
+    image = image.apt_install(*apt_packages)
+
 pip_packages = {pip_packages!r}
 if pip_packages:
     image = image.pip_install(*pip_packages)
@@ -3251,7 +3297,7 @@ def register(cli: click.Group) -> None:
 )
 @click.option(
     "--trace-db",
-    default="traces/v3/eval_traces.db",
+    default="traces/v3/synth_ai.db",
     show_default=True,
     help="SQLite/Turso URL for storing rollout traces set to 'none' to disable persistence.",
 )
@@ -3284,8 +3330,13 @@ def eval_command(
     pointing at a remote `--url`, supply matching `--env-file` values so the CLI can
     forward authentication headers to the running service.
     """
+    # Parse and validate TOML config
+    from synth_ai.task.config import EvalConfig
+    
     cfg: dict[str, Any] = {}
+    eval_cfg: EvalConfig | None = None
     config_path: Path | None = None
+    
     if config:
         config_path = Path(config)
     else:
@@ -3307,21 +3358,37 @@ def eval_command(
             if isinstance(parsed, dict):
                 section = parsed.get("eval")
                 cfg = dict(section) if isinstance(section, dict) else dict(parsed)
+            
+            # Validate config with dataclass
+            try:
+                eval_cfg = EvalConfig.from_dict(cfg)
+                click.echo(f"✓ Config validated: {len(eval_cfg.seeds)} seeds, model={eval_cfg.model}")
+            except (ValueError, TypeError) as validation_error:
+                raise click.ClickException(f"Invalid eval config: {validation_error}") from validation_error
+        except click.ClickException:
+            raise
         except Exception as exc:
             raise click.ClickException(f"Failed to parse TOML '{config_path}': {exc}") from exc
 
-    app_id = app_id or (cfg.get("app_id") if isinstance(cfg.get("app_id"), str) else None)  # type: ignore
+    # CLI args override config
+    if eval_cfg:
+        app_id = app_id or eval_cfg.app_id
+    else:
+        app_id = app_id or (cfg.get("app_id") if isinstance(cfg.get("app_id"), str) else None)  # type: ignore
 
     metadata_filters: dict[str, str] = {}
-    cfg_metadata = cfg.get("metadata")
-    if isinstance(cfg_metadata, dict):
-        for key, value in cfg_metadata.items():
-            metadata_filters[str(key)] = str(value)
-    elif isinstance(cfg_metadata, list):
-        for item in cfg_metadata:
-            if isinstance(item, str) and "=" in item:
-                key, value = item.split("=", 1)
-                metadata_filters[key.strip()] = value.strip()
+    if eval_cfg:
+        metadata_filters.update(eval_cfg.metadata)
+    else:
+        cfg_metadata = cfg.get("metadata")
+        if isinstance(cfg_metadata, dict):
+            for key, value in cfg_metadata.items():
+                metadata_filters[str(key)] = str(value)
+        elif isinstance(cfg_metadata, list):
+            for item in cfg_metadata:
+                if isinstance(item, str) and "=" in item:
+                    key, value = item.split("=", 1)
+                    metadata_filters[key.strip()] = value.strip()
 
     for item in metadata or ():
         if "=" not in item:
@@ -3334,11 +3401,14 @@ def eval_command(
         metadata_filters[key] = value
 
     metadata_sql_query: str | None = None
-    cfg_metadata_sql = cfg.get("metadata_sql")
-    if isinstance(cfg_metadata_sql, dict):
-        metadata_sql_query = cfg_metadata_sql.get("query") or cfg_metadata_sql.get("sql")
-    elif isinstance(cfg_metadata_sql, str):
-        metadata_sql_query = cfg_metadata_sql
+    if eval_cfg and eval_cfg.metadata_sql:
+        metadata_sql_query = eval_cfg.metadata_sql
+    else:
+        cfg_metadata_sql = cfg.get("metadata_sql")
+        if isinstance(cfg_metadata_sql, dict):
+            metadata_sql_query = cfg_metadata_sql.get("query") or cfg_metadata_sql.get("sql")
+        elif isinstance(cfg_metadata_sql, str):
+            metadata_sql_query = cfg_metadata_sql
 
     if metadata_sql:
         metadata_sql_query = metadata_sql
@@ -3780,18 +3850,52 @@ def eval_command(
 
             async def _run_seed(seed_val: int) -> None:
                 nonlocal successes, failures, outcome_sum, outcome_count, outcome_correct, records
+                # Read env_name and policy_name from config if available
+                env_name = cfg.get("env_name") or (cfg.get("env", {}).get("env_name") if isinstance(cfg.get("env"), dict) else None)
+                policy_name = cfg.get("policy_name") or (cfg.get("policy", {}).get("policy_name") if isinstance(cfg.get("policy"), dict) else None)
+                env_config_overrides = cfg.get("env_config", {}) if isinstance(cfg.get("env_config"), dict) else {}
+                policy_config_overrides = cfg.get("policy_config", {}) if isinstance(cfg.get("policy_config"), dict) else {}
+                
+                # Debug: print config parsing
+                if seed_val == 0:
+                    click.echo(f"[DEBUG] env_name from config: {env_name}")
+                    click.echo(f"[DEBUG] policy_name from config: {policy_name}")
+                
+                # Generate default ops sequence if not provided
+                max_llm_calls = policy_config_overrides.get("max_llm_calls", 10)
+                ops_list = cfg.get("ops", [])
+                if not ops_list:
+                    # Generate default "agent, env" pairs for max_llm_calls
+                    ops_list = ["agent", "env"] * int(max_llm_calls)
+                
                 body = {
                     "run_id": str(uuid.uuid4()),
-                    "env": {"config": {"split": split, "index": seed_val}, "seed": seed_val},
+                    "env": {"config": {"split": split, "index": seed_val, **env_config_overrides}, "seed": seed_val},
                     "policy": {
-                        "policy_name": selected_model,
-                        "config": {"model": selected_model, **policy_overrides},
+                        "policy_name": policy_name or selected_model,
+                        "config": {"model": selected_model, **policy_overrides, **policy_config_overrides},
+                    },
+                    "ops": ops_list,
+                    "record": {
+                        "return_trace": cfg.get("return_trace", True),
+                        "trace_format": cfg.get("trace_format", "structured"),
                     },
-                    "ops": [],
+                    "mode": "eval",  # RolloutMode.EVAL: use inference URLs as-is, no transformations
                 }
+                if env_name:
+                    body["env"]["env_name"] = env_name
+                
+                # Debug: print the body being sent
+                if seed_val == 0:
+                    click.echo(f"[DEBUG] rollout body env: {body['env']}")
+                    click.echo(f"[DEBUG] rollout body policy: {body['policy']}")
+                    click.echo(f"[DEBUG] rollout body mode: {body.get('mode', 'NOT SET')}")
                 rollout_elapsed: float | None = None
                 rollout_start = time.perf_counter()
                 try:
+                    import logging
+                    _log = logging.getLogger(__name__)
+                    _log.info(f"[EVAL_BODY_DEBUG] Sending body with mode={body.get('mode')}")
                     async with semaphore:
                         response = await async_client.post("/rollout", json=body)
                     rollout_elapsed = time.perf_counter() - rollout_start
@@ -3812,6 +3916,10 @@ def eval_command(
                     data = response.json()
                 except Exception:
                     data = None
+                
+                # Debug: print validation errors
+                if response.status_code == 422 and data:
+                    click.echo(f"[DEBUG] 422 Validation Error: {data}")
 
                 metrics: dict[str, Any] | None = None
                 completion: str | None = None
@@ -3825,16 +3933,33 @@ def eval_command(
                 session_trace_dict: dict[str, Any] | None = None
 
                 if isinstance(data, dict):
+                    import logging
+                    _logger = logging.getLogger(__name__)
+                    _logger.info(f"[EVAL_DEBUG] Response data keys: {list(data.keys())}")
+                    if "detail" in data:
+                        _logger.error(f"[EVAL_DEBUG] Task app returned error: {data['detail']}")
                     trace_namespace = data.get("trace")
+                    _logger.info(f"[EVAL_DEBUG] trace_namespace type: {type(trace_namespace)}, value: {trace_namespace if not isinstance(trace_namespace, dict) else 'dict with keys: ' + str(list(trace_namespace.keys()) if trace_namespace else 'None')}")
                     if not isinstance(trace_namespace, dict):
                         raise RuntimeError(
-                            "rollout response missing trace payload; task app must return tracing_v3 data"
+                            "The 'synth-ai eval' command requires trace payloads in rollout responses. "
+                            "Ensure the rollout request includes 'trace_format': 'structured' and 'return_trace': true, "
+                            "and that task app tracing is enabled (TASKAPP_TRACING_ENABLED=1). "
+                            "Note: This is specific to the eval command - general rollout endpoints don't require traces."
                         )
+                    # Handle both "compact" and "full" trace formats:
+                    # - compact: trace_namespace contains {session_id, metadata, ...}
+                    # - full: trace_namespace IS the full session_trace dict
                     session_trace_dict = trace_namespace.get("session_trace")
                     if not isinstance(session_trace_dict, dict):
-                        raise RuntimeError(
-                            "rollout response trace missing 'session_trace'; ensure the task app is serving the tracing_v3 build"
-                        )
+                        # If no session_trace key, assume "full" format where trace itself is the session_trace
+                        if "session_id" in trace_namespace:
+                            session_trace_dict = trace_namespace
+                        else:
+                            raise RuntimeError(
+                                "The 'synth-ai eval' command requires 'session_trace' in the trace payload or a valid full trace format. "
+                                "Ensure the task app is using tracing_v3 and returning structured trace data."
+                            )
                     metrics = data.get("metrics") if isinstance(data.get("metrics"), dict) else None
                     if metrics:
                         mean_return = metrics.get("mean_return") or metrics.get("total_reward")
@@ -3956,26 +4081,27 @@ def eval_command(
                     for spec in judge_specs:
                         score_value: float | None = None
                         judge_elapsed: float | None = None
-                        if completion is not None:
-                            judge_payload = {
-                                "seed": seed_val,
-                                "prompt_index": prompt_index,
-                                "prompt": prompt_text,
-                                "completion": completion,
-                                "metrics": metrics,
-                                "response": data,
-                                "trace": trace_namespace,
-                            }
-                            try:
-                                judge_start = time.perf_counter()
-                                result = spec.fn(judge_payload, **spec.kwargs)
+                        # Run judges for all tasks (text-based and trajectory-based)
+                        # Text-based tasks have completion, trajectory-based tasks use response
+                        judge_payload = {
+                            "seed": seed_val,
+                            "prompt_index": prompt_index,
+                            "prompt": prompt_text,
+                            "completion": completion,
+                            "metrics": metrics,
+                            "response": data,
+                            "trace": trace_namespace,
+                        }
+                        try:
+                            judge_start = time.perf_counter()
+                            result = spec.fn(judge_payload, **spec.kwargs)
+                            judge_elapsed = time.perf_counter() - judge_start
+                            if isinstance(result, int | float):
+                                score_value = float(result)
+                        except Exception as exc:
+                            if judge_elapsed is None:
                                 judge_elapsed = time.perf_counter() - judge_start
-                                if isinstance(result, int | float):
-                                    score_value = float(result)
-                            except Exception as exc:
-                                if judge_elapsed is None:
-                                    judge_elapsed = time.perf_counter() - judge_start
-                                click.echo(f"seed={seed_val} judge[{spec.name}]_error={exc}")
+                            click.echo(f"seed={seed_val} judge[{spec.name}]_error={exc}")
                         judges_timings[spec.name] = judge_elapsed
                         judge_scores[spec.name] = score_value
 
@@ -4129,6 +4255,9 @@ def filter_command(config_path: str) -> None:
     high-quality traces. See `customers/agora_single_file/configs/filter_local.toml`
     for a working example.
     """
+    # Parse and validate TOML config
+    from synth_ai.task.config import FilterConfig
+    
     if _toml is None:
         raise click.ClickException("TOML parser not available; install tomli or use Python 3.11+")
 
@@ -4141,58 +4270,37 @@ def filter_command(config_path: str) -> None:
     except Exception as exc:
         raise click.ClickException(f"Failed to parse TOML '{cfg_path}': {exc}") from exc
 
-    filter_cfg = config_data.get("filter") if isinstance(config_data, dict) else None
-    if not isinstance(filter_cfg, dict):
+    filter_cfg_dict = config_data.get("filter") if isinstance(config_data, dict) else None
+    if not isinstance(filter_cfg_dict, dict):
         raise click.ClickException("Config must contain a [filter] table")
 
-    db_value = str(filter_cfg.get("db", "traces/v3/eval_traces.db")).strip()
-    if not db_value:
-        raise click.ClickException("filter.db must be provided")
-    if "://" in db_value:
-        db_url = db_value
-    else:
-        db_path = Path(db_value).expanduser()
-        db_path.parent.mkdir(parents=True, exist_ok=True)
-        db_url = f"sqlite+aiosqlite:///{db_path}"
-
-    output_value = filter_cfg.get("output")
-    if not output_value:
-        raise click.ClickException("filter.output must be provided")
-    output_path = Path(str(output_value)).expanduser()
-
-    splits = set(filter_cfg.get("splits", []) or [])
-    task_ids = set(filter_cfg.get("task_ids", []) or [])
-    models = set(filter_cfg.get("models", []) or [])
-    min_official = filter_cfg.get("min_official_score")
-    max_official = filter_cfg.get("max_official_score")
-    if min_official is not None:
-        try:
-            min_official = float(min_official)
-        except Exception as err:
-            raise click.ClickException("filter.min_official_score must be numeric") from err
-    if max_official is not None:
-        try:
-            max_official = float(max_official)
-        except Exception as err:
-            raise click.ClickException("filter.max_official_score must be numeric") from err
-    min_judge_scores = filter_cfg.get("min_judge_scores", {}) or {}
-    max_judge_scores = filter_cfg.get("max_judge_scores", {}) or {}
+    # Validate config with dataclass
     try:
-        min_judge_scores = {k: float(v) for k, v in min_judge_scores.items()}
-    except Exception as err:
-        raise click.ClickException("filter.min_judge_scores values must be numeric") from err
-    try:
-        max_judge_scores = {k: float(v) for k, v in max_judge_scores.items()}
-    except Exception as err:
-        raise click.ClickException("filter.max_judge_scores values must be numeric") from err
-    min_created = _parse_datetime_for_trace(filter_cfg.get("min_created_at"))
-    max_created = _parse_datetime_for_trace(filter_cfg.get("max_created_at"))
-    limit = filter_cfg.get("limit")
-    if limit is not None:
-        try:
-            limit = int(limit)
-        except Exception as err:
-            raise click.ClickException("filter.limit must be an integer") from err
+        filter_cfg = FilterConfig.from_dict(filter_cfg_dict)
+        click.echo(f"✓ Config validated: db={filter_cfg.db}, output={filter_cfg.output}")
+        if filter_cfg.min_official_score is not None:
+            click.echo(f"  → Filtering for official score >= {filter_cfg.min_official_score}")
+        if filter_cfg.limit:
+            click.echo(f"  → Limiting to {filter_cfg.limit} examples")
+    except (ValueError, TypeError) as validation_error:
+        raise click.ClickException(f"Invalid filter config: {validation_error}") from validation_error
+
+    # Use validated config
+    db_url = filter_cfg.get_db_url()
+    output_path = filter_cfg.get_output_path()
+
+    # Extract validated fields from dataclass
+    splits = set(filter_cfg.splits)
+    task_ids = set(filter_cfg.task_ids)
+    models = set(filter_cfg.models)
+    min_official = filter_cfg.min_official_score
+    max_official = filter_cfg.max_official_score
+    min_judge_scores = filter_cfg.min_judge_scores
+    max_judge_scores = filter_cfg.max_judge_scores
+    # Note: min_created_at and max_created_at not yet in FilterConfig dataclass
+    min_created = _parse_datetime_for_trace(filter_cfg_dict.get("min_created_at"))
+    max_created = _parse_datetime_for_trace(filter_cfg_dict.get("max_created_at"))
+    limit = filter_cfg.limit
 
     def _score_ok(value: Any, min_val: Any, max_val: Any) -> bool:
         try:
@@ -4247,8 +4355,21 @@ def filter_command(config_path: str) -> None:
             if max_created and (created_at_dt is None or created_at_dt > max_created):
                 continue
 
-            if not _score_ok(metadata.get("official_score"), min_official, max_official):
-                continue
+            # Check against outcome_rewards if score filter is set
+            total_reward = None
+            achievements_count = None
+            if min_official is not None or max_official is not None:
+                reward_query = "SELECT total_reward, achievements_count FROM outcome_rewards WHERE session_id = :session_id"
+                reward_rows = await tracer.db.query_traces(reward_query, {"session_id": session_id})
+                reward_records = reward_rows.to_dict("records") if hasattr(reward_rows, "to_dict") else []
+                if reward_records:
+                    total_reward = reward_records[0].get("total_reward")
+                    achievements_count = reward_records[0].get("achievements_count")
+                    if not _score_ok(total_reward, min_official, max_official):
+                        continue
+                elif min_official is not None:
+                    # No reward found, but score filter requires it
+                    continue
 
             judge_scores = metadata.get("judge_scores") or {}
             include = True
@@ -4265,30 +4386,120 @@ def filter_command(config_path: str) -> None:
             if not include:
                 continue
 
-            prompt = metadata.get("prompt") or ""
-            completion = metadata.get("completion") or ""
-            if not prompt or not completion:
+            # Query messages for this session
+            messages_query = """
+                SELECT message_type, content, timestamp 
+                FROM messages 
+                WHERE session_id = :session_id
+                ORDER BY timestamp ASC, id ASC
+            """
+            msg_df = await tracer.db.query_traces(messages_query, {"session_id": session_id})
+            message_rows = msg_df.to_dict("records") if hasattr(msg_df, "to_dict") else []
+            
+            if not message_rows:
+                # Fallback: check if prompt/completion in metadata (old format)
+                prompt = metadata.get("prompt") or ""
+                completion = metadata.get("completion") or ""
+                if prompt and completion:
+                    record = {
+                        "messages": [
+                            {"role": "user", "content": str(prompt)},
+                            {"role": "assistant", "content": str(completion)},
+                        ],
+                        "metadata": {
+                            "session_id": session_id,
+                            "env_name": metadata.get("env_name"),
+                            "policy_name": metadata.get("policy_name"),
+                            "seed": metadata.get("seed"),
+                            "total_reward": total_reward,
+                            "achievements_count": achievements_count,
+                            "model": metadata.get("model"),
+                            "created_at": created_at_dt.isoformat() if created_at_dt else created_at_raw,
+                        },
+                    }
+                    accepted.append(record)
                 continue
 
-            record = {
-                "messages": [
-                    {"role": "user", "content": str(prompt)},
-                    {"role": "assistant", "content": str(completion)},
-                ],
-                "metadata": {
-                    "session_id": session_id,
-                    "task_id": metadata.get("task_id"),
-                    "task_split": metadata.get("task_split"),
-                    "task_rubric_id": metadata.get("task_rubric_id"),
-                    "official_score": metadata.get("official_score"),
-                    "judge_scores": judge_scores,
-                    "model": metadata.get("model"),
-                    "created_at": created_at_dt.isoformat() if created_at_dt else created_at_raw,
-                    "prompt": prompt,
-                    "completion": completion,
-                },
-            }
-            accepted.append(record)
+            # Extract user/assistant pairs from messages
+            for i, msg_row in enumerate(message_rows):
+                msg_type = msg_row.get("message_type")
+                content_raw = msg_row.get("content")
+                
+                # Look for user message
+                if msg_type in ("user", "policy_user_prompt"):
+                    # Find next policy_system_prompt or assistant
+                    assistant_msg = None
+                    for j in range(i + 1, len(message_rows)):
+                        next_type = message_rows[j].get("message_type")
+                        if next_type in ("assistant", "policy_system_prompt"):
+                            if next_type == "assistant":
+                                assistant_msg = message_rows[j]
+                            break
+                    
+                    # Parse content
+                    try:
+                        user_content = json.loads(content_raw) if isinstance(content_raw, str) else content_raw
+                    except Exception:
+                        user_content = content_raw
+                    
+                    # Extract text from structured content
+                    def extract_text(content: Any) -> str:
+                        if isinstance(content, str):
+                            return content
+                        if isinstance(content, dict):
+                            # Try payload.content for user prompts
+                            if "payload" in content and isinstance(content["payload"], dict):
+                                payload = content["payload"]
+                                if "content" in payload:
+                                    return extract_text(payload["content"])
+                            # Try common keys
+                            for key in ["text", "content", "content_text"]:
+                                if key in content:
+                                    val = content[key]
+                                    if isinstance(val, str):
+                                        return val
+                            return json.dumps(content)
+                        if isinstance(content, list):
+                            # Multimodal content - concatenate text parts
+                            parts = []
+                            for item in content:
+                                if isinstance(item, dict) and item.get("type") == "text":
+                                    parts.append(item.get("text", ""))
+                            return " ".join(parts) if parts else str(content)
+                        return str(content)
+                    
+                    user_text = extract_text(user_content)
+                    
+                    # For assistant, we might not have it recorded, so use tool calls as completion
+                    assistant_text = ""
+                    if assistant_msg:
+                        assistant_content_raw = assistant_msg.get("content")
+                        try:
+                            assistant_content = json.loads(assistant_content_raw) if isinstance(assistant_content_raw, str) else assistant_content_raw
+                        except Exception:
+                            assistant_content = assistant_content_raw
+                        assistant_text = extract_text(assistant_content)
+                    
+                    if not user_text:
+                        continue
+
+                    record = {
+                        "messages": [
+                            {"role": "user", "content": user_text},
+                            {"role": "assistant", "content": assistant_text if assistant_text else "[no response recorded]"},
+                        ],
+                        "metadata": {
+                            "session_id": session_id,
+                            "env_name": metadata.get("env_name"),
+                            "policy_name": metadata.get("policy_name"),
+                            "seed": metadata.get("seed"),
+                            "total_reward": total_reward,
+                            "achievements_count": achievements_count,
+                            "model": metadata.get("model"),
+                            "created_at": created_at_dt.isoformat() if created_at_dt else created_at_raw,
+                        },
+                    }
+                    accepted.append(record)
 
         if not accepted:
             raise click.ClickException("No sessions matched the provided filters")