83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass, field
|
|
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]
|
|
# The prompt actually sent to the LLM for this step. For non-LLM planners
|
|
# (or older records persisted before D9), this falls back to the task goal.
|
|
prompt: str
|
|
tool_call: dict[str, Any]
|
|
result: dict[str, Any]
|
|
timestamp: str
|
|
ocr_results: list[dict[str, Any]] = field(default_factory=list)
|
|
before_screenshot_path: str | None = None
|
|
after_screenshot_path: str | None = None
|
|
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,
|
|
before_screenshot: bytes | None = None,
|
|
after_screenshot: bytes | None = None,
|
|
ocr_results: list[dict[str, Any]] | None = None,
|
|
) -> TimelineRecord:
|
|
index = len(self.read(task_id)) + 1
|
|
resolved_after_screenshot = (
|
|
after_screenshot if after_screenshot is not None else screenshot
|
|
)
|
|
resolved_ocr_results = list(ocr_results or [])
|
|
record = {
|
|
"index": index,
|
|
"scene": scene,
|
|
"prompt": prompt,
|
|
"tool_call": tool_call,
|
|
"result": result,
|
|
"ocr_results": resolved_ocr_results,
|
|
"timestamp": datetime.now().astimezone().isoformat(),
|
|
}
|
|
paths = self.artifact_store.write_step(
|
|
task_id=task_id,
|
|
index=index,
|
|
before_screenshot=before_screenshot,
|
|
after_screenshot=resolved_after_screenshot,
|
|
record=record,
|
|
)
|
|
return TimelineRecord(
|
|
index=index,
|
|
scene=record["scene"],
|
|
prompt=prompt,
|
|
tool_call=tool_call,
|
|
result=result,
|
|
timestamp=record["timestamp"],
|
|
ocr_results=resolved_ocr_results,
|
|
before_screenshot_path=paths["before_screenshot_path"],
|
|
after_screenshot_path=paths["after_screenshot_path"],
|
|
screenshot_path=paths["screenshot_path"],
|
|
)
|
|
|
|
def read(self, task_id: str) -> list[dict[str, Any]]:
|
|
return self.artifact_store.read_steps(task_id)
|
|
|
|
def delete_task(self, task_id: str) -> None:
|
|
"""Delete all timeline records and screenshots for a task."""
|
|
self.artifact_store.delete_task(task_id)
|