WIP: Add replay logging mechanism

openadapt/alembic/versions/c84664aeb5ae_add_replay_models.py

@@ -0,0 +1,62 @@
+"""add_replay_models
+
+Revision ID: c84664aeb5ae
+Revises: bb25e889ad71
+Create Date: 2024-06-25 15:05:09.110171
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+import openadapt
+
+# revision identifiers, used by Alembic.
+revision = "c84664aeb5ae"
+down_revision = "bb25e889ad71"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "replay",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column(
+            "timestamp",
+            openadapt.models.ForceFloat(precision=10, scale=2, asdecimal=False),
+            nullable=True,
+        ),
+        sa.Column("strategy_name", sa.String(), nullable=True),
+        sa.Column("strategy_args", sa.JSON(), nullable=True),
+        sa.Column("git_hash", sa.String(), nullable=True),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_replay")),
+    )
+    op.create_table(
+        "replay_log",
+        sa.Column("id", sa.Integer(), nullable=False),
+        sa.Column("replay_id", sa.Integer(), nullable=True),
+        sa.Column("lineno", sa.Integer(), nullable=True),
+        sa.Column("filename", sa.String(), nullable=True),
+        sa.Column("git_hash", sa.String(), nullable=True),
+        sa.Column(
+            "timestamp",
+            openadapt.models.ForceFloat(precision=10, scale=2, asdecimal=False),
+            nullable=True,
+        ),
+        sa.Column("log_level", sa.String(), nullable=True),
+        sa.Column("key", sa.String(), nullable=True),
+        sa.Column("data", sa.LargeBinary(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["replay_id"], ["replay.id"], name=op.f("fk_replay_log_replay_id_replay")
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("pk_replay_log")),
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_table("replay_log")
+    op.drop_table("replay")
+    # ### end Alembic commands ###

openadapt/app/dashboard/api/index.py

@@ -11,6 +11,7 @@
 
 from openadapt.app.dashboard.api.action_events import ActionEventsAPI
 from openadapt.app.dashboard.api.recordings import RecordingsAPI
+from openadapt.app.dashboard.api.replays import ReplaysAPI
 from openadapt.app.dashboard.api.scrubbing import ScrubbingAPI
 from openadapt.app.dashboard.api.settings import SettingsAPI
 from openadapt.build_utils import is_running_from_executable
@@ -23,11 +24,13 @@
 
 action_events_app = ActionEventsAPI().attach_routes()
 recordings_app = RecordingsAPI().attach_routes()
+replays_app = ReplaysAPI().attach_routes()
 scrubbing_app = ScrubbingAPI().attach_routes()
 settings_app = SettingsAPI().attach_routes()
 
 api.include_router(action_events_app, prefix="/action-events")
 api.include_router(recordings_app, prefix="/recordings")
+api.include_router(replays_app, prefix="/replays")
 api.include_router(scrubbing_app, prefix="/scrubbing")
 api.include_router(settings_app, prefix="/settings")
 

openadapt/app/dashboard/api/replays.py

@@ -0,0 +1,70 @@
+"""API endpoints for replays."""
+import pickle
+
+from fastapi import APIRouter, WebSocket
+from loguru import logger
+
+from openadapt.db import crud
+from openadapt.models import Replay
+from openadapt.utils import image2utf8
+
+
+class ReplaysAPI:
+    """API endpoints for replays."""
+
+    def __init__(self) -> None:
+        """Initialize the ReplaysAPI class."""
+        self.app = APIRouter()
+
+    def attach_routes(self) -> APIRouter:
+        """Attach routes to the FastAPI app."""
+        self.app.add_api_route(
+            "", self.get_replays, methods=["GET"], response_model=None
+        )
+        self.replay_logs_route()
+        return self.app
+
+    @staticmethod
+    def get_replays() -> list[Replay]:
+        """Get all replays."""
+        session = crud.get_new_session(read_only=True)
+        return crud.get_replays(session)
+
+    def replay_logs_route(self) -> None:
+        """Add the replay detail route as a websocket."""
+
+        @self.app.websocket("/{replay_id}/logs")
+        async def get_replay_logs(websocket: WebSocket, replay_id: int) -> None:
+            """Get a specific replay and its logs."""
+            await websocket.accept()
+            session = crud.get_new_session(read_only=True)
+
+            replay = crud.get_replay(session, replay_id)
+
+            await websocket.send_json({"type": "replay", "value": replay.asdict()})
+
+            logs = crud.get_replay_logs(session, replay_id)
+
+            await websocket.send_json({"type": "num_logs", "value": len(logs)})
+
+            for log in logs:
+                log.data = pickle.loads(log.data)
+                if log.key == "screenshot":
+                    log.data = image2utf8(log.data)
+                log_dict = log.asdict()
+                if log.key == "window_event":
+                    log_dict["data"] = log_dict["data"].asdict()
+                if log.key == "action_event_dict":
+                    log_dict["data"]["reducer_names"] = list(
+                        log_dict["data"]["reducer_names"]
+                    )
+                if log.key == "segmentation":
+                    log_dict["data"] = log_dict["data"].asdict()
+                try:
+                    await websocket.send_json({"type": "log", "value": log_dict})
+                except Exception as e:
+                    logger.error(f"Error sending log: {e}")
+                    logger.info(log_dict["data"])
+                    logger.info(log_dict["key"])
+
+            await websocket.close()

openadapt/app/dashboard/app/replays/logs/page.tsx

@@ -0,0 +1,87 @@
+'use client';
+
+import { ReplayDetails } from "@/components/ReplayDetails";
+import { ReplayLogs } from "@/components/ReplayLogs";
+import { ReplayLog, Replay as ReplayType } from "@/types/replay";
+import { Box, Loader, Progress } from "@mantine/core";
+import { useSearchParams } from "next/navigation";
+import { Suspense, useEffect, useState } from "react";
+
+function Replay() {
+    const searchParams = useSearchParams();
+    const id = searchParams.get("id");
+    const [replayInfo, setReplayInfo] = useState<{
+        replay: ReplayType,
+        logs: ReplayLog[],
+        num_logs: number,
+    }>();
+    useEffect(() => {
+        if (!id) {
+            return;
+        }
+        const websocket = new WebSocket(`ws://${window.location.host}/api/replays/${id}/logs`);
+        websocket.onmessage = (event) => {
+            const data = JSON.parse(event.data);
+            if (data.type === "replay") {
+                setReplayInfo(prev => {
+                    if (!prev) {
+                        return {
+                            "replay": data.value,
+                            "logs": [],
+                            "num_logs": 0,
+                        }
+                    }
+                    return prev;
+                });
+            } else if (data.type === "log") {
+                setReplayInfo(prev => {
+                    if (!prev) return prev;
+                    return {
+                        ...prev,
+                        "logs": [...prev.logs, data.value],
+                    }
+                });
+            } else if (data.type === "num_logs") {
+                setReplayInfo(prev => {
+                    if (!prev) return prev;
+                    return {
+                        ...prev,
+                        "num_logs": data.value,
+                    }
+                });
+            }
+        }
+
+        return () => {
+            websocket.close();
+        }
+    }, [id]);
+    if (!replayInfo) {
+        return <Loader />;
+    }
+
+    const logs = replayInfo.logs;
+
+    return (
+        <Box>
+            <ReplayDetails replay={replayInfo.replay} />
+            {logs.length && logs.length < replayInfo.num_logs && (
+                <Progress.Root size={30} my={30}>
+                    <Progress.Section value={(logs.length / replayInfo.num_logs) * 100}>
+                        <Progress.Label>Loading events {logs.length}/{replayInfo.num_logs}</Progress.Label>
+                    </Progress.Section>
+                </Progress.Root>
+            )}
+            <ReplayLogs logs={logs} />
+        </Box>
+    )
+}
+
+
+export default function ReplayPage() {
+    return (
+        <Suspense>
+            <Replay />
+        </Suspense>
+    )
+}

openadapt/app/dashboard/app/replays/page.tsx

@@ -0,0 +1,47 @@
+'use client';
+
+import { SimpleTable } from "@/components/SimpleTable";
+import { Replay } from "@/types/replay";
+import { useEffect, useState } from "react";
+import { timeStampToDateString } from "../utils";
+import { Box, Text } from "@mantine/core";
+import { useRouter } from "next/navigation";
+
+export default function Replays() {
+    const [replays, setReplays] = useState<Replay[]>([]);
+    const fetchReplays = () => {
+        fetch('/api/replays').then(res => {
+            if (res.ok) {
+                res.json().then((data) => {
+                    setReplays(data);
+                });
+            }
+        });
+    }
+    const router = useRouter();
+
+    function onClickRow(replay: Replay) {
+        return () => router.push(`/replays/logs?id=${replay.id}`);
+    }
+
+    useEffect(() => {
+        fetchReplays();
+    }, []);
+
+    return (
+        <Box>
+            <Text size="xl">Replay logs</Text>
+            <SimpleTable
+                columns={[
+                    {name: 'ID', accessor: 'id'},
+                    {name: 'Strategy name', accessor: 'strategy_name'},
+                    {name: 'Strategy arguments', accessor: (replay) => JSON.stringify(replay.strategy_args)},
+                    {name: 'Time', accessor: (replay) => timeStampToDateString(replay.timestamp)},
+                ]}
+                data={replays}
+                refreshData={fetchReplays}
+                onClickRow={onClickRow}
+            />
+        </Box>
+    )
+}

openadapt/app/dashboard/app/routes.ts

@@ -19,5 +19,9 @@ export const routes: Route[] = [
     {
         name: 'Onboarding',
         path: '/onboarding',
+    },
+    {
+        name: "Replays",
+        path: "/replays",
     }
 ]

openadapt/app/dashboard/components/ReplayDetails/ReplayDetails.tsx

@@ -0,0 +1,49 @@
+'use client';
+
+import { timeStampToDateString } from '@/app/utils';
+import { Replay } from '@/types/replay'
+import { Code, Table } from '@mantine/core'
+import React from 'react'
+
+type Props = {
+    replay: Replay;
+}
+
+const TableRowWithBorder = ({ children }: { children: React.ReactNode }) => (
+    <Table.Tr className='border-2 border-gray-300 border-solid'>
+        {children}
+    </Table.Tr>
+)
+
+const TableCellWithBorder = ({ children }: { children: React.ReactNode }) => (
+    <Table.Td className='border-2 border-gray-300 border-solid'>
+        {children}
+    </Table.Td>
+)
+
+export const ReplayDetails = ({
+    replay
+}: Props) => {
+  return (
+        <Table withTableBorder withColumnBorders w={600}>
+            <Table.Tbody>
+                <TableRowWithBorder>
+                    <TableCellWithBorder>Replay ID</TableCellWithBorder>
+                    <TableCellWithBorder>{replay.id}</TableCellWithBorder>
+                </TableRowWithBorder>
+                <TableRowWithBorder>
+                    <TableCellWithBorder>timestamp</TableCellWithBorder>
+                    <TableCellWithBorder>{timeStampToDateString(replay.timestamp)}</TableCellWithBorder>
+                </TableRowWithBorder>
+                <TableRowWithBorder>
+                    <TableCellWithBorder>strategy name</TableCellWithBorder>
+                    <TableCellWithBorder>{replay.strategy_name}</TableCellWithBorder>
+                </TableRowWithBorder>
+                <TableRowWithBorder>
+                    <TableCellWithBorder>strategy args</TableCellWithBorder>
+                    <TableCellWithBorder><Code>{JSON.stringify(replay.strategy_args)}</Code></TableCellWithBorder>
+                </TableRowWithBorder>
+            </Table.Tbody>
+        </Table>
+  )
+}

openadapt/app/dashboard/components/ReplayDetails/index.tsx

@@ -0,0 +1 @@
+export { ReplayDetails } from './ReplayDetails';