63 lines
1.6 KiB
Python
63 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from storage.artifact_store import ArtifactStore
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TimelineRecord:
|
|
index: int
|
|
scene: dict[str, Any]
|
|
prompt: str
|
|
tool_call: dict[str, Any]
|
|
result: dict[str, Any]
|
|
timestamp: str
|
|
screenshot_path: str | None = None
|
|
|
|
|
|
class Timeline:
|
|
def __init__(self, artifact_store: ArtifactStore | None = None) -> None:
|
|
self.artifact_store = artifact_store or ArtifactStore()
|
|
|
|
def append(
|
|
self,
|
|
*,
|
|
task_id: str,
|
|
scene: Any,
|
|
prompt: str,
|
|
tool_call: dict[str, Any],
|
|
result: dict[str, Any],
|
|
screenshot: bytes | None = None,
|
|
) -> TimelineRecord:
|
|
index = len(self.read(task_id)) + 1
|
|
record = {
|
|
"index": index,
|
|
"scene": scene,
|
|
"prompt": prompt,
|
|
"tool_call": tool_call,
|
|
"result": result,
|
|
"timestamp": datetime.now().astimezone().isoformat(),
|
|
}
|
|
paths = self.artifact_store.write_step(
|
|
task_id=task_id,
|
|
index=index,
|
|
screenshot=screenshot,
|
|
record=record,
|
|
)
|
|
return TimelineRecord(
|
|
index=index,
|
|
scene=record["scene"],
|
|
prompt=prompt,
|
|
tool_call=tool_call,
|
|
result=result,
|
|
timestamp=record["timestamp"],
|
|
screenshot_path=paths["screenshot_path"],
|
|
)
|
|
|
|
def read(self, task_id: str) -> list[dict[str, Any]]:
|
|
return self.artifact_store.read_steps(task_id)
|
|
|