synth-ai 0.2.14__py3-none-any.whl → 0.2.16__py3-none-any.whl

synth_ai/cli/setup.py

@@ -0,0 +1,266 @@
+from __future__ import annotations
+
+import contextlib
+import os
+import time
+import webbrowser
+from pathlib import Path
+from typing import Any, cast
+from urllib.parse import urljoin, urlsplit, urlunsplit
+
+import requests
+from click.exceptions import Exit
+from synth_ai.demos import core as demo_core
+from synth_ai.utils.cli import print_next_step
+from synth_ai.utils.env import mask_str
+from synth_ai.utils.modal import is_modal_public_url
+from synth_ai.utils.process import popen_capture
+from synth_ai.utils.user_config import USER_CONFIG_PATH, update_user_config
+
+
+class HandshakeError(Exception):
+    pass
+
+
+def _get_canonical_origin() -> str:
+    """Resolve the dashboard origin for the browser handshake.
+
+    Priority order:
+      1. Explicit ``SYNTH_CANONICAL_ORIGIN`` override.
+      2. Development flag ``SYNTH_CANONICAL_DEV`` (case-insensitive truthy) → localhost.
+      3. Production dashboard at ``https://www.usesynth.ai/dashboard``.
+    """
+
+    override = (os.getenv("SYNTH_CANONICAL_ORIGIN") or "").strip()
+    if override:
+        return override.rstrip("/")
+
+    dev_flag = (os.getenv("SYNTH_CANONICAL_DEV") or "").strip().lower()
+    if dev_flag in { "1", "true", "yes", "on" }:
+        print("USING DEV ORIGIN")
+        return "http://localhost:3000"
+
+    return "https://www.usesynth.ai/dashboard"
+
+
+def _split_origin(origin: str) -> tuple[str, str]:
+    parsed = urlsplit(origin)
+    bare = cast(str, urlunsplit((parsed.scheme, parsed.netloc, "", "", "")))
+    path = parsed.path.rstrip("/")
+    return bare, path
+
+
+def _ensure_verification_uri(data: dict[str, Any], base_with_path: str) -> None:
+    uri = data.get("verification_uri")
+    if not isinstance(uri, str) or not uri:
+        return
+    if uri.startswith("http://") or uri.startswith("https://"):
+        return
+    data["verification_uri"] = urljoin(base_with_path.rstrip("/") + "/", uri.lstrip("/"))
+
+
+def _start_handshake_session(origin: str | None = None) -> tuple[str, str, int, int]:
+    base = (origin or _get_canonical_origin()).rstrip("/")
+    api_origin, _ = _split_origin(base)
+    url = urljoin(api_origin.rstrip("/") + "/", "api/sdk/handshake/init")
+    r = requests.post(url, timeout=10)
+    if r.status_code != 200:
+        raise HandshakeError(f"init failed: {r.status_code} {r.text}")
+    try:
+        data = r.json()
+    except ValueError as exc:  # pragma: no cover - network dependent
+        raise HandshakeError(f"init returned malformed JSON: {exc}") from exc
+    _ensure_verification_uri(data, base)
+    return (
+        str(data.get("device_code")),
+        str(data.get("verification_uri")),
+        int(data.get("expires_in", 600)),
+        int(data.get("interval", 3)),
+    )
+
+
+def _poll_handshake_token(
+    device_code: str, origin: str | None = None, *, timeout_s: int | None = None
+) -> dict[str, Any]:
+    base = (origin or _get_canonical_origin()).rstrip("/")
+    api_origin, _ = _split_origin(base)
+    url = urljoin(api_origin.rstrip("/") + "/", "api/sdk/handshake/token")
+    deadline = time.time() + (timeout_s or 600)
+    while True:
+        if time.time() > deadline:
+            raise HandshakeError("handshake timed out")
+        try:
+            r = requests.post(url, json={"device_code": device_code}, timeout=10)
+        except Exception:
+            time.sleep(2)
+            continue
+        if r.status_code == 200:
+            try:
+                data = r.json()
+            except ValueError as exc:  # pragma: no cover - network dependent
+                raise HandshakeError(f"token returned malformed JSON: {exc}") from exc
+            _ensure_verification_uri(data, base)
+            return data
+        elif r.status_code in (404, 410):
+            raise HandshakeError(f"handshake failed: {r.status_code}")
+        # 428 authorization_pending or others → wait and retry
+        time.sleep(2)
+
+
+def _run_handshake(origin: str | None = None) -> dict[str, Any]:
+    device_code, verification_uri, expires_in, interval = _start_handshake_session(origin)
+    with contextlib.suppress(Exception):
+        webbrowser.open(verification_uri)
+    return _poll_handshake_token(device_code, origin, timeout_s=expires_in)
+
+
+
+def setup() -> int:
+    # Prefer the demo directory provided in the current shell session, then fall back to persisted state
+    demo_dir_env = (os.environ.get("DEMO_DIR") or "").strip()
+    demo_dir: str | None = None
+    if demo_dir_env:
+        candidate = Path(demo_dir_env).expanduser()
+        if candidate.is_dir():
+            demo_dir = str(candidate.resolve())
+        else:
+            print(f"Warning: DEMO_DIR={demo_dir_env} does not exist; falling back to stored demo directory.")
+
+    if demo_dir is None:
+        loaded = demo_core.load_demo_dir()
+        if loaded:
+            demo_dir = loaded
+
+    if demo_dir and os.path.isdir(demo_dir):
+        os.chdir(demo_dir)
+        print(f"Using demo directory: {demo_dir}")
+
+    synth_key = ""
+    rl_env_key = ""
+    org_name = ""
+
+    try:
+        print("\n⏳ Connecting to your browser session…")
+        res = _run_handshake()
+        org = res.get("org") or {}
+        keys = res.get("keys") or {}
+        synth_key = str(keys.get("synth") or "").strip()
+        rl_env_key = str(keys.get("rl_env") or "").strip()
+        org_name = org.get("name") or "Unamed Organization ™️"
+        print(f"✅ Connected to {org_name}!")
+    except (HandshakeError, Exception) as exc:
+        print(f"⚠️  Failed to fetch keys from frontend: {exc}")
+        print("Falling back to manual entry...")
+
+    if not synth_key:
+        try:
+            synth_key = input(
+                "Failed to fetch your Synth API key. Please enter your Synth API key here:\n> "
+            ).strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nSetup cancelled.")
+            return 1
+        if not synth_key:
+            print("Synth API key is required.")
+            return 1
+
+    if not rl_env_key:
+        try:
+            rl_env_key = input(
+                "Failed to fetch your Environment API key. Please enter your Environment API key here:\n> "
+            ).strip()
+        except (EOFError, KeyboardInterrupt):
+            print("\nSetup cancelled.")
+            return 1
+        if not rl_env_key:
+            print("Environment API key is required.")
+            return 1
+
+    # Persist keys to user config
+    config_updates = {
+        "SYNTH_API_KEY": synth_key,
+        "ENVIRONMENT_API_KEY": rl_env_key,
+    }
+    update_user_config(config_updates)
+
+    os.environ["SYNTH_API_KEY"] = synth_key
+    os.environ["ENVIRONMENT_API_KEY"] = rl_env_key
+
+    env = demo_core.load_env()
+
+    def _refresh_env() -> None:
+        nonlocal env
+        env = demo_core.load_env()
+
+    def _maybe_fix_task_url() -> None:
+        if not env.task_app_name:
+            return
+        current = env.task_app_base_url
+        needs_lookup = not current or not is_modal_public_url(current)
+        if not needs_lookup:
+            return
+        code, out = popen_capture(
+            [
+                "uv",
+                "run",
+                "python",
+                "-m",
+                "modal",
+                "app",
+                "url",
+                env.task_app_name,
+            ]
+        )
+        if code != 0 or not out:
+            return
+        new_url = ""
+        for token in out.split():
+            if is_modal_public_url(token):
+                new_url = token.strip().rstrip("/")
+                break
+        if new_url and new_url != current:
+            print(f"Updating TASK_APP_BASE_URL from Modal CLI → {new_url}")
+            persist_path = demo_dir or os.getcwd()
+            demo_core.persist_task_url(new_url, name=env.task_app_name, path=persist_path)
+            os.environ["TASK_APP_BASE_URL"] = new_url
+            _refresh_env()
+
+    modal_ok, modal_msg = demo_core.modal_auth_status()
+    if modal_ok:
+        print(f"✓ Modal authenticated: {modal_msg}")
+    else:
+        print(f"[setup] Modal authentication status: {modal_msg}")
+
+    _maybe_fix_task_url()
+
+    if env.dev_backend_url:
+        api = env.dev_backend_url.rstrip("/") + (
+            "" if env.dev_backend_url.endswith("/api") else "/api"
+        )
+        demo_core.assert_http_ok(api + "/health", method="GET")
+    if env.task_app_base_url:
+        base = env.task_app_base_url.rstrip("/")
+        demo_core.assert_http_ok(
+            base + "/health", method="GET"
+        ) or demo_core.assert_http_ok(
+            base, method="GET"
+        )
+    print("\nSaved keys:")
+    print(f"  SYNTH_API_KEY={mask_str(synth_key)}")
+    print(f"  ENVIRONMENT_API_KEY={mask_str(rl_env_key)}")
+    if env.task_app_base_url:
+        print(f"  TASK_APP_BASE_URL={env.task_app_base_url}")
+    print(f"Configuration persisted to: {USER_CONFIG_PATH}")
+
+    demo_core.persist_demo_dir(os.getcwd())
+
+    print_next_step("deploy our task app", ["uvx synth-ai deploy"])
+    return 0
+
+
+def register(group):
+    @group.command("setup")
+    def demo_setup():
+        code = setup()
+        if code:
+            raise Exit(code)

synth_ai/cli/status.py

@@ -12,7 +12,7 @@ from rich.console import Console
 from rich.panel import Panel
 from rich.table import Table
 
-from synth_ai.cli._storage import load_storage
+from ._storage import load_storage
 
 
 async def _db_stats(db_url: str) -> dict:

synth_ai/cli/task_app_deploy.py

@@ -0,0 +1,16 @@
+"""Compatibility wrapper for task-app deploy command."""
+
+from __future__ import annotations
+
+import click
+
+from .task_apps import task_app_group
+
+_deploy = task_app_group.commands.get("deploy")
+
+if _deploy is None:
+    raise RuntimeError("task_app_group does not define a 'deploy' command")
+
+deploy_command: click.Command = _deploy
+
+__all__ = ["deploy_command"]

synth_ai/cli/task_app_list.py

@@ -0,0 +1,25 @@
+"""Task app list command."""
+
+from __future__ import annotations
+
+import click
+from synth_ai.task.apps import registry
+
+from .task_apps import task_app_group
+
+
+@task_app_group.command("list")
+def list_apps() -> None:
+    """List registered task apps."""
+
+    entries = registry.list()
+    if not entries:
+        click.echo("No task apps registered.")
+        return
+
+    for entry in entries:
+        aliases = f" (aliases: {', '.join(entry.aliases)})" if entry.aliases else ""
+        click.echo(f"- {entry.app_id}{aliases}: {entry.description}")
+
+
+__all__ = ["list_apps"]

synth_ai/cli/task_app_modal_serve.py

@@ -0,0 +1,16 @@
+"""Compatibility wrapper for task-app modal-serve command."""
+
+from __future__ import annotations
+
+import click
+
+from .task_apps import task_app_group
+
+_modal_serve = task_app_group.commands.get("modal-serve")
+
+if _modal_serve is None:
+    raise RuntimeError("task_app_group does not define a 'modal-serve' command")
+
+modal_serve_command: click.Command = _modal_serve
+
+__all__ = ["modal_serve_command"]

synth_ai/cli/task_app_serve.py

@@ -0,0 +1,18 @@
+"""Compatibility wrapper for task-app serve command."""
+
+from __future__ import annotations
+
+import click
+
+from .task_apps import serve_command as task_app_serve_command
+from .task_apps import task_app_group
+
+serve_command = task_app_serve_command
+
+_group_serve = task_app_group.commands.get("serve")
+if _group_serve is None:
+    raise RuntimeError("task_app_group does not define a 'serve' command")
+
+serve_task_group: click.Command = _group_serve
+
+__all__ = ["serve_command", "serve_task_group"]

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 datetime, timezone
+from datetime import UTC, datetime
 from pathlib import Path
-from typing import Any, Optional, cast
+from typing import Any, cast
 
 try:  # Python 3.11+
     import tomllib as _toml
@@ -92,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() -> Optional[Path]:
+def _load_demo_directory() -> Path | None:
     """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[[], Optional[str | Path]], module.load_demo_dir)
+        loader = cast(Callable[[], str | Path | None], module.load_demo_dir)
         demo_dir = loader()
         if isinstance(demo_dir, str | Path):
             demo_path = Path(demo_dir)
@@ -139,7 +139,7 @@ DEFAULT_SEARCH_RELATIVE = (
 )
 
 
-def _pearson(xs: Sequence[float], ys: Sequence[float]) -> Optional[float]:
+def _pearson(xs: Sequence[float], ys: Sequence[float]) -> float | None:
     if len(xs) != len(ys) or len(xs) < 2:
         return None
     mean_x = sum(xs) / len(xs)
@@ -164,7 +164,7 @@ class AppChoice:
     label: str
     path: Path
     source: str
-    description: Optional[str] = None
+    description: str | None = None
     aliases: tuple[str, ...] = ()
     entry: TaskAppEntryType | None = None
     entry_loader: Callable[[], TaskAppEntryType] | None = None
@@ -188,21 +188,21 @@ class JudgeSpec:
     kwargs: dict[str, Any]
 
 
-def _parse_datetime_for_trace(value: Any) -> Optional[datetime]:
+def _parse_datetime_for_trace(value: Any) -> datetime | None:
     if isinstance(value, datetime):
-        return value if value.tzinfo else value.replace(tzinfo=timezone.utc)
+        return value if value.tzinfo else value.replace(tzinfo=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=timezone.utc)
+                dt = datetime.fromtimestamp(float(value), tz=UTC)
             except Exception:
                 return None
-        return dt if dt.tzinfo else dt.replace(tzinfo=timezone.utc)
+        return dt if dt.tzinfo else dt.replace(tzinfo=UTC)
     if isinstance(value, int | float):
-        return datetime.fromtimestamp(float(value), tz=timezone.utc)
+        return datetime.fromtimestamp(float(value), tz=UTC)
     return None
 
 
@@ -241,6 +241,24 @@ def _event_from_dict(payload: dict[str, Any]) -> BaseEvent:
             system_state_after=payload.get("system_state_after"),
             **base_kwargs,
         )
+    # Check for LM CAIS event fields
+    if any(key in payload for key in ("model_name", "provider", "call_records")):
+        from synth_ai.tracing_v3.abstractions import LMCAISEvent
+        # Note: call_records are left as dicts - the storage layer will handle serialization
+        call_records = payload.get("call_records") or []
+        return LMCAISEvent(
+            model_name=payload.get("model_name", ""),
+            provider=payload.get("provider", ""),
+            input_tokens=payload.get("input_tokens"),
+            output_tokens=payload.get("output_tokens"),
+            total_tokens=payload.get("total_tokens"),
+            cost_usd=payload.get("cost_usd"),
+            latency_ms=payload.get("latency_ms"),
+            span_id=payload.get("span_id"),
+            trace_id=payload.get("trace_id"),
+            call_records=call_records,
+            **base_kwargs,
+        )
     return BaseEvent(**base_kwargs)
 
 
@@ -279,7 +297,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(timezone.utc)
+    timestamp = _parse_datetime_for_trace(payload.get("timestamp")) or datetime.now(UTC)
     completed_at = _parse_datetime_for_trace(payload.get("completed_at"))
     return SessionTimeStep(
         step_id=payload.get("step_id", ""),
@@ -293,7 +311,7 @@ def _step_from_dict(payload: dict[str, Any]) -> SessionTimeStep:
     )
 
 
-def _session_trace_from_dict(payload: dict[str, Any]) -> Optional[V3SessionTrace]:
+def _session_trace_from_dict(payload: dict[str, Any]) -> V3SessionTrace | None:
     if not isinstance(payload, dict):
         return None
     steps = [
@@ -311,7 +329,7 @@ def _session_trace_from_dict(payload: dict[str, Any]) -> Optional[V3SessionTrace
         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(timezone.utc)
+    created_at = _parse_datetime_for_trace(payload.get("created_at")) or datetime.now(UTC)
     metadata = payload.get("metadata") or {}
     session_metadata = payload.get("session_metadata")
     return V3SessionTrace(
@@ -341,10 +359,18 @@ async def _store_trace(
     
     _logger.info(f"[STORE_TRACE_DEBUG] trace_namespace keys: {list(trace_namespace.keys())}")
     
+    # Handle both formats:
+    # - With session_trace key: {"session_trace": {...}}
+    # - Without session_trace key (trace itself is the session): {"session_id": ..., "markov_blanket_message_history": ...}
     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
+        # If no session_trace key, assume "full" format where trace itself is the session_trace
+        if "session_id" in trace_namespace:
+            session_payload = trace_namespace
+            _logger.info("[STORE_TRACE_DEBUG] Using trace_namespace directly as session_payload (no session_trace key)")
+        else:
+            _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", []))
@@ -352,7 +378,7 @@ async def _store_trace(
     
     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")
+        _logger.warning("[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")
@@ -366,7 +392,7 @@ async def _store_trace(
     
     _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")
+    _logger.info("[STORE_TRACE_DEBUG] Successfully inserted trace")
 
 def _temporary_sys_path(paths: Sequence[Path]):
     """Context manager to prepend entries to sys.path temporarily."""
@@ -913,43 +939,43 @@ def _build_modal_config_from_ast(modal_call: ast.Call) -> ModalDeploymentConfigT
         for kw in modal_call.keywords:
             if kw.arg and isinstance(kw.value, ast.Constant):
                 kwargs[kw.arg] = kw.value.value
-            elif kw.arg == "pip_packages" and isinstance(kw.value, (ast.List, ast.Tuple)):
+            elif kw.arg == "pip_packages" and isinstance(kw.value, ast.List | ast.Tuple):
                 # Handle pip_packages list/tuple
                 packages: list[str] = []
                 value_node = kw.value
-                if isinstance(value_node, (ast.List, ast.Tuple)):
+                if isinstance(value_node, ast.List | ast.Tuple):
                     for elt in value_node.elts:
                         if isinstance(elt, ast.Constant):
                             packages.append(elt.value)
                 kwargs[kw.arg] = tuple(packages)
-            elif kw.arg == "extra_local_dirs" and isinstance(kw.value, (ast.List, ast.Tuple)):
+            elif kw.arg == "extra_local_dirs" and isinstance(kw.value, ast.List | ast.Tuple):
                 # Handle extra_local_dirs list/tuple of tuples
                 dirs = []
                 value_node = kw.value
-                if isinstance(value_node, (ast.List, ast.Tuple)):
+                if isinstance(value_node, ast.List | ast.Tuple):
                     for elt in value_node.elts:
-                        if isinstance(elt, (ast.List, ast.Tuple)) and len(elt.elts) == 2:
+                        if isinstance(elt, ast.List | ast.Tuple) and len(elt.elts) == 2:
                             src = elt.elts[0].value if isinstance(elt.elts[0], ast.Constant) else None
                             dst = elt.elts[1].value if isinstance(elt.elts[1], ast.Constant) else None
                             if src and dst:
                                 dirs.append((src, dst))
                 kwargs[kw.arg] = tuple(dirs)
-            elif kw.arg == "secret_names" and isinstance(kw.value, (ast.List, ast.Tuple)):
+            elif kw.arg == "secret_names" and isinstance(kw.value, ast.List | ast.Tuple):
                 # Handle secret_names list/tuple
                 secrets = []
                 value_node = kw.value
-                if isinstance(value_node, (ast.List, ast.Tuple)):
+                if isinstance(value_node, ast.List | ast.Tuple):
                     for elt in value_node.elts:
                         if isinstance(elt, ast.Constant):
                             secrets.append(elt.value)
                 kwargs[kw.arg] = tuple(secrets)
-            elif kw.arg == "volume_mounts" and isinstance(kw.value, (ast.List, ast.Tuple)):
+            elif kw.arg == "volume_mounts" and isinstance(kw.value, ast.List | ast.Tuple):
                 # Handle volume_mounts list/tuple of tuples
                 mounts = []
                 value_node = kw.value
-                if isinstance(value_node, (ast.List, ast.Tuple)):
+                if isinstance(value_node, ast.List | ast.Tuple):
                     for elt in value_node.elts:
-                        if isinstance(elt, (ast.List, ast.Tuple)) and len(elt.elts) == 2:
+                        if isinstance(elt, ast.List | ast.Tuple) and len(elt.elts) == 2:
                             name = elt.elts[0].value if isinstance(elt.elts[0], ast.Constant) else None
                             mount = elt.elts[1].value if isinstance(elt.elts[1], ast.Constant) else None
                             if name and mount:
@@ -2245,7 +2271,7 @@ def validate_task_app_cmd(
     import time
     
     # Import the validate_task_app function defined in this module
-    from synth_ai.cli._validate_task_app import validate_task_app  # type: ignore[attr-defined]
+    from ._validate_task_app import validate_task_app  # type: ignore[attr-defined]
     
     proc = None
     task_app_url = url
@@ -4442,6 +4468,10 @@ def filter_command(config_path: str) -> None:
                     except Exception:
                         user_content = content_raw
                     
+                    # If user_content is a message dict with a 'content' key, extract it
+                    if isinstance(user_content, dict) and "content" in user_content:
+                        user_content = user_content["content"]
+                    
                     # Extract text from structured content
                     def extract_text(content: Any) -> str:
                         if isinstance(content, str):
@@ -4472,21 +4502,31 @@ def filter_command(config_path: str) -> None:
                     
                     # For assistant, we might not have it recorded, so use tool calls as completion
                     assistant_text = ""
+                    assistant_content = None
                     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
+                        
+                        # If assistant_content is a message dict with a 'content' key, extract it
+                        if isinstance(assistant_content, dict) and "content" in assistant_content:
+                            assistant_content = assistant_content["content"]
+                        
                         assistant_text = extract_text(assistant_content)
                     
                     if not user_text:
                         continue
 
+                    # Use full multimodal content if it's a list (contains images), otherwise use text
+                    user_content_for_message = user_content if isinstance(user_content, list) else user_text
+                    assistant_content_for_message = assistant_content if isinstance(assistant_content, list) else (assistant_text if assistant_text else "[no response recorded]")
+
                     record = {
                         "messages": [
-                            {"role": "user", "content": user_text},
-                            {"role": "assistant", "content": assistant_text if assistant_text else "[no response recorded]"},
+                            {"role": "user", "content": user_content_for_message},
+                            {"role": "assistant", "content": assistant_content_for_message},
                         ],
                         "metadata": {
                             "session_id": session_id,

synth_ai/cli/traces.py

@@ -11,7 +11,7 @@ from rich import box
 from rich.console import Console
 from rich.table import Table
 
-from synth_ai.cli._storage import load_storage
+from ._storage import load_storage
 
 
 def register(cli):

synth_ai/cli/train.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import Any
+
+from synth_ai.api.train.cli import register as _register
+from synth_ai.api.train.cli import train_command as _train_command
+
+__all__ = ["register", "train_command"]
+
+
+def register(cli: Any) -> None:
+    """Compatibility wrapper for the legacy train CLI location."""
+
+    _register(cli)
+
+
+def train_command(*args: Any, **kwargs: Any) -> Any:
+    return _train_command(*args, **kwargs)