mirror of
https://github.com/HKUDS/OpenSpace.git
synced 2026-08-28 05:15:00 +00:00
487 lines
18 KiB
Python
487 lines
18 KiB
Python
"""Build redacted openspace_task_trace_v2 artifact archives."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import tempfile
|
|
import zipfile
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Iterable, Mapping
|
|
|
|
from openspace.cloud.local_mapping import CloudLocalMappingStore
|
|
from openspace.cloud.redaction import (
|
|
REDACTION_POLICY_VERSION,
|
|
redact_task_trace_value,
|
|
redaction_report_for_payload,
|
|
sanitize_upload_path,
|
|
)
|
|
from openspace.cloud.task_trace_schema import (
|
|
TASK_TRACE_ARTIFACT_FORMAT,
|
|
TASK_TRACE_REQUIRED_FILES,
|
|
TASK_TRACE_SCHEMA_VERSION,
|
|
TaskTraceArtifact,
|
|
classify_cloud_task_involvement,
|
|
)
|
|
from openspace.cloud.telemetry_payloads import short_cloud_request_id
|
|
|
|
|
|
class TaskTraceExportError(RuntimeError):
|
|
"""Raised when a task trace artifact cannot be safely exported."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class TaskTraceExportInput:
|
|
task_id: str
|
|
session_id: str
|
|
collection_scope: str
|
|
collection_reason: str
|
|
status: str = "unknown"
|
|
entrypoint: str = "openspace"
|
|
workspace_ref: str = ""
|
|
cloud_skill_ids: tuple[str, ...] = ()
|
|
package_ids: tuple[str, ...] = ()
|
|
messages: tuple[Mapping[str, Any], ...] = ()
|
|
llm_calls: tuple[Mapping[str, Any], ...] = ()
|
|
tool_calls: tuple[Mapping[str, Any], ...] = ()
|
|
skill_invocations: tuple[Mapping[str, Any], ...] = ()
|
|
package_skill_provenance: tuple[Mapping[str, Any], ...] = ()
|
|
evolution_events: tuple[Mapping[str, Any], ...] = ()
|
|
task_metadata: Mapping[str, Any] | None = None
|
|
model_inventory: Mapping[str, Any] | None = None
|
|
|
|
|
|
class TaskTraceExporter:
|
|
"""Write a redacted task trace directory and zip archive."""
|
|
|
|
def __init__(
|
|
self,
|
|
output_dir: str | Path,
|
|
*,
|
|
workspace_root: str | Path | None = None,
|
|
) -> None:
|
|
self.output_dir = Path(output_dir)
|
|
self.output_dir.mkdir(parents=True, exist_ok=True)
|
|
self.workspace_root = Path(workspace_root).resolve() if workspace_root else None
|
|
|
|
def export(self, data: TaskTraceExportInput) -> TaskTraceArtifact:
|
|
if not data.task_id or not data.session_id:
|
|
raise TaskTraceExportError("task_id and session_id are required")
|
|
staging_parent = Path(
|
|
tempfile.mkdtemp(prefix="openspace-task-trace-", dir=str(self.output_dir))
|
|
)
|
|
root = staging_parent / "task_trace_artifact"
|
|
root.mkdir(parents=True)
|
|
created_at = utc_now_iso()
|
|
|
|
task_metadata = self._task_metadata(data, created_at=created_at)
|
|
rows = {
|
|
"messages.jsonl": self._message_rows(data.messages),
|
|
"llm_calls.jsonl": self._generic_rows(data.llm_calls),
|
|
"tool_calls.jsonl": self._generic_rows(data.tool_calls),
|
|
"skill_invocations.jsonl": self._generic_rows(data.skill_invocations),
|
|
"package_skill_provenance.jsonl": self._generic_rows(
|
|
data.package_skill_provenance
|
|
),
|
|
"evolution_events.jsonl": self._generic_rows(data.evolution_events),
|
|
}
|
|
_write_json(root / "task_metadata.json", task_metadata)
|
|
for rel_path, file_rows in rows.items():
|
|
_write_jsonl(root / rel_path, file_rows)
|
|
|
|
derived_summary = {
|
|
"task_id": data.task_id,
|
|
"session_id": data.session_id,
|
|
"status": data.status,
|
|
"collection_scope": data.collection_scope,
|
|
"collection_reason": data.collection_reason,
|
|
"cloud_skill_count": len(data.cloud_skill_ids),
|
|
"package_count": len(data.package_ids),
|
|
"message_count": len(rows["messages.jsonl"]),
|
|
"llm_call_count": len(rows["llm_calls.jsonl"]),
|
|
"tool_call_count": len(rows["tool_calls.jsonl"]),
|
|
"skill_invocation_count": len(rows["skill_invocations.jsonl"]),
|
|
"model_inventory": redact_task_trace_value(
|
|
dict(data.model_inventory or {}),
|
|
workspace_root=self.workspace_root,
|
|
),
|
|
}
|
|
_write_json(root / "derived_summary.json", derived_summary)
|
|
|
|
payload_for_report = {
|
|
"task_metadata": task_metadata,
|
|
"rows": rows,
|
|
"derived_summary": derived_summary,
|
|
}
|
|
redaction_report = redaction_report_for_payload(payload_for_report)
|
|
_write_json(root / "redaction_report.json", redaction_report)
|
|
|
|
file_entries = self._file_entries(root)
|
|
manifest = {
|
|
"schema_version": TASK_TRACE_SCHEMA_VERSION,
|
|
"artifact_format": TASK_TRACE_ARTIFACT_FORMAT,
|
|
"task_id": data.task_id,
|
|
"session_id": data.session_id,
|
|
"collection_scope": data.collection_scope,
|
|
"collection_reason": data.collection_reason,
|
|
"created_at": created_at,
|
|
"redaction_level": "complete_redacted",
|
|
"redaction_policy_version": REDACTION_POLICY_VERSION,
|
|
"compression": "zip",
|
|
"cloud_skill_ids": list(data.cloud_skill_ids),
|
|
"package_ids": list(data.package_ids),
|
|
"files": file_entries,
|
|
}
|
|
_write_json(root / "manifest.json", manifest)
|
|
_ensure_required_files(root)
|
|
|
|
archive_stem = short_cloud_request_id(
|
|
"task-trace-archive",
|
|
data.task_id,
|
|
data.session_id,
|
|
created_at,
|
|
).replace(":", "-")
|
|
archive_path = self.output_dir / f"{archive_stem}.zip"
|
|
archive_path = _unique_path(archive_path)
|
|
_zip_dir(root, archive_path)
|
|
digest = _sha256_file(archive_path)
|
|
request_id = short_cloud_request_id(
|
|
"task-trace",
|
|
data.task_id,
|
|
data.session_id,
|
|
digest,
|
|
)
|
|
return TaskTraceArtifact(
|
|
archive_path=archive_path,
|
|
request_id=request_id,
|
|
task_id=data.task_id,
|
|
session_id=data.session_id,
|
|
manifest=manifest,
|
|
sha256=digest,
|
|
size_bytes=archive_path.stat().st_size,
|
|
compression="zip",
|
|
collection_scope=data.collection_scope,
|
|
collection_reason=data.collection_reason,
|
|
cloud_skill_ids=tuple(data.cloud_skill_ids),
|
|
package_ids=tuple(data.package_ids),
|
|
upload_allowed=bool(redaction_report.get("upload_allowed")),
|
|
deny_findings=tuple(redaction_report.get("deny_findings") or ()),
|
|
)
|
|
|
|
def from_execution_result(
|
|
self,
|
|
final_result: Mapping[str, Any],
|
|
*,
|
|
task_id: str,
|
|
session_id: str,
|
|
mapping_store: CloudLocalMappingStore | None = None,
|
|
model_inventory: Mapping[str, Any] | None = None,
|
|
) -> TaskTraceArtifact | None:
|
|
"""Create an artifact only when cloud involvement requires one."""
|
|
|
|
local_skill_ids = _extract_local_skill_ids(final_result)
|
|
cloud_bindings = _cloud_bindings_for_local_ids(mapping_store, local_skill_ids)
|
|
cloud_skill_ids = tuple(
|
|
binding.cloud_skill_id
|
|
for binding in cloud_bindings
|
|
if binding.cloud_skill_id
|
|
)
|
|
package_ids = tuple(
|
|
item
|
|
for binding in cloud_bindings
|
|
for item in (binding.current_package_id, binding.package_id_at_pull)
|
|
if item
|
|
)
|
|
cloud_evolve = any(
|
|
_mapping_get(item, "source_cloud_skill_id")
|
|
or _mapping_get(item, "cloud_skill_id")
|
|
for item in final_result.get("evolved_skills") or []
|
|
if isinstance(item, Mapping)
|
|
)
|
|
status = str(final_result.get("status") or "unknown")
|
|
involvement = classify_cloud_task_involvement(
|
|
cloud_skill_ids=list(cloud_skill_ids),
|
|
local_skill_ids=local_skill_ids,
|
|
package_ids=list(package_ids),
|
|
cloud_evolve=cloud_evolve,
|
|
cloud_attempted=status not in {"success", "completed", "ok"},
|
|
)
|
|
if not involvement.requires_task_trace_artifact:
|
|
return None
|
|
|
|
data = TaskTraceExportInput(
|
|
task_id=task_id,
|
|
session_id=session_id,
|
|
collection_scope=involvement.scope,
|
|
collection_reason=involvement.reason,
|
|
status=status,
|
|
workspace_ref=(
|
|
sanitize_upload_path(self.workspace_root, workspace_root=self.workspace_root)
|
|
if self.workspace_root
|
|
else ""
|
|
),
|
|
cloud_skill_ids=involvement.cloud_skill_ids,
|
|
package_ids=involvement.package_ids,
|
|
messages=tuple(_message_inputs(final_result)),
|
|
tool_calls=tuple(_tool_call_inputs(final_result)),
|
|
skill_invocations=tuple(
|
|
_skill_invocation_rows(cloud_bindings, final_result=final_result)
|
|
),
|
|
package_skill_provenance=tuple(
|
|
_provenance_rows(cloud_bindings, workspace_root=self.workspace_root)
|
|
),
|
|
evolution_events=tuple(_evolution_event_rows(final_result)),
|
|
task_metadata={
|
|
"iterations": final_result.get("iterations"),
|
|
"stop_reason": final_result.get("stop_reason"),
|
|
"execution_time": final_result.get("execution_time"),
|
|
},
|
|
model_inventory=model_inventory or {},
|
|
)
|
|
return self.export(data)
|
|
|
|
def _task_metadata(
|
|
self,
|
|
data: TaskTraceExportInput,
|
|
*,
|
|
created_at: str,
|
|
) -> dict[str, Any]:
|
|
metadata = {
|
|
"task_id": data.task_id,
|
|
"session_id": data.session_id,
|
|
"agent_id": "",
|
|
"started_at": "",
|
|
"ended_at": created_at,
|
|
"status": data.status,
|
|
"entrypoint": data.entrypoint,
|
|
"workspace_ref": data.workspace_ref,
|
|
"collection_scope": data.collection_scope,
|
|
"collection_reason": data.collection_reason,
|
|
"model_inventory": dict(data.model_inventory or {}),
|
|
}
|
|
metadata.update(dict(data.task_metadata or {}))
|
|
return redact_task_trace_value(metadata, workspace_root=self.workspace_root)
|
|
|
|
def _message_rows(self, messages: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]:
|
|
rows: list[dict[str, Any]] = []
|
|
for index, message in enumerate(messages):
|
|
content = _extract_content_text(message)
|
|
redacted = redact_task_trace_value(content, workspace_root=self.workspace_root)
|
|
rows.append({
|
|
"message_id": str(message.get("message_id") or f"msg_{index + 1:06d}"),
|
|
"turn_id": str(message.get("turn_id") or ""),
|
|
"role": str(message.get("role") or "unknown"),
|
|
"created_at": str(message.get("created_at") or ""),
|
|
"content_type": "text",
|
|
"content": redacted,
|
|
"content_sha256": hashlib.sha256(str(redacted).encode("utf-8")).hexdigest(),
|
|
"redaction_applied": True,
|
|
"source": str(message.get("source") or "execution_result"),
|
|
})
|
|
return rows
|
|
|
|
def _generic_rows(self, rows: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]:
|
|
return [
|
|
redact_task_trace_value(dict(row), workspace_root=self.workspace_root)
|
|
for row in rows
|
|
if isinstance(row, Mapping)
|
|
]
|
|
|
|
def _file_entries(self, root: Path) -> list[dict[str, Any]]:
|
|
entries: list[dict[str, Any]] = []
|
|
for path in sorted(root.rglob("*")):
|
|
if not path.is_file() or path.name == "manifest.json":
|
|
continue
|
|
rel = path.relative_to(root).as_posix()
|
|
entries.append({
|
|
"path": rel,
|
|
"size_bytes": path.stat().st_size,
|
|
"sha256": _sha256_file(path),
|
|
})
|
|
return entries
|
|
|
|
|
|
def utc_now_iso() -> str:
|
|
return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z")
|
|
|
|
|
|
def _write_json(path: Path, payload: Mapping[str, Any]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(
|
|
json.dumps(payload, ensure_ascii=False, sort_keys=True, indent=2) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def _write_jsonl(path: Path, rows: Iterable[Mapping[str, Any]]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w", encoding="utf-8") as handle:
|
|
for row in rows:
|
|
handle.write(json.dumps(dict(row), ensure_ascii=False, sort_keys=True) + "\n")
|
|
|
|
|
|
def _ensure_required_files(root: Path) -> None:
|
|
missing = [name for name in TASK_TRACE_REQUIRED_FILES if not (root / name).exists()]
|
|
if missing:
|
|
raise TaskTraceExportError(f"task trace missing required files: {missing}")
|
|
|
|
|
|
def _zip_dir(root: Path, archive_path: Path) -> None:
|
|
with zipfile.ZipFile(archive_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
|
|
for path in sorted(root.rglob("*")):
|
|
if path.is_file():
|
|
zf.write(path, Path("task_trace_artifact") / path.relative_to(root))
|
|
|
|
|
|
def _sha256_file(path: Path) -> str:
|
|
hasher = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
|
|
hasher.update(chunk)
|
|
return hasher.hexdigest()
|
|
|
|
|
|
def _unique_path(path: Path) -> Path:
|
|
if not path.exists():
|
|
return path
|
|
stem = path.stem
|
|
suffix = path.suffix
|
|
for index in range(1, 1000):
|
|
candidate = path.with_name(f"{stem}-{index}{suffix}")
|
|
if not candidate.exists():
|
|
return candidate
|
|
raise TaskTraceExportError(f"could not allocate artifact path under {path.parent}")
|
|
|
|
|
|
def _extract_content_text(message: Mapping[str, Any]) -> str:
|
|
content = message.get("content")
|
|
if isinstance(content, str):
|
|
return content
|
|
if isinstance(content, list):
|
|
parts: list[str] = []
|
|
for item in content:
|
|
if isinstance(item, Mapping):
|
|
parts.append(str(item.get("text") or item.get("content") or ""))
|
|
else:
|
|
parts.append(str(item))
|
|
return "\n".join(part for part in parts if part)
|
|
return json.dumps(content, ensure_ascii=False, default=str)
|
|
|
|
|
|
def _message_inputs(final_result: Mapping[str, Any]) -> list[Mapping[str, Any]]:
|
|
raw = final_result.get("messages")
|
|
if isinstance(raw, list):
|
|
return [item for item in raw if isinstance(item, Mapping)]
|
|
response = final_result.get("response")
|
|
if response:
|
|
return [{"role": "assistant", "content": str(response), "source": "final_response"}]
|
|
return []
|
|
|
|
|
|
def _tool_call_inputs(final_result: Mapping[str, Any]) -> list[Mapping[str, Any]]:
|
|
rows: list[Mapping[str, Any]] = []
|
|
for index, item in enumerate(final_result.get("tool_executions") or []):
|
|
if not isinstance(item, Mapping):
|
|
continue
|
|
rows.append({
|
|
"tool_call_id": str(item.get("tool_use_id") or item.get("id") or f"tool_{index + 1:06d}"),
|
|
"tool_name": str(item.get("tool_name") or item.get("name") or item.get("tool") or ""),
|
|
"status": str(item.get("status") or ""),
|
|
"duration_ms": item.get("duration_ms"),
|
|
"input": item.get("input") or item.get("args") or {},
|
|
"output": item.get("output") or item.get("result") or "",
|
|
})
|
|
return rows
|
|
|
|
|
|
def _extract_local_skill_ids(final_result: Mapping[str, Any]) -> list[str]:
|
|
raw = (
|
|
final_result.get("active_skill_ids")
|
|
or final_result.get("active_skills")
|
|
or final_result.get("skills_used")
|
|
or []
|
|
)
|
|
if isinstance(raw, str):
|
|
raw = [raw]
|
|
return [str(item) for item in raw if str(item)]
|
|
|
|
|
|
def _cloud_bindings_for_local_ids(
|
|
store: CloudLocalMappingStore | None,
|
|
local_skill_ids: Iterable[str],
|
|
):
|
|
if store is None:
|
|
return []
|
|
bindings = []
|
|
for local_skill_id in local_skill_ids:
|
|
try:
|
|
binding = store.get_binding_by_local(local_skill_id)
|
|
except Exception:
|
|
binding = None
|
|
if binding is not None and binding.cloud_skill_id:
|
|
bindings.append(binding)
|
|
return bindings
|
|
|
|
|
|
def _skill_invocation_rows(
|
|
bindings: Iterable[Any],
|
|
*,
|
|
final_result: Mapping[str, Any],
|
|
) -> list[dict[str, Any]]:
|
|
status = str(final_result.get("status") or "")
|
|
rows: list[dict[str, Any]] = []
|
|
for index, binding in enumerate(bindings):
|
|
rows.append({
|
|
"skill_invocation_id": f"skill_use_{index + 1:06d}",
|
|
"cloud_skill_id": binding.cloud_skill_id,
|
|
"local_skill_id": binding.local_skill_id,
|
|
"package_id": binding.current_package_id or binding.package_id_at_pull,
|
|
"skill_name": "",
|
|
"status": status or "unknown",
|
|
"used_cloud_origin": True,
|
|
})
|
|
return rows
|
|
|
|
|
|
def _provenance_rows(
|
|
bindings: Iterable[Any],
|
|
*,
|
|
workspace_root: Path | None,
|
|
) -> list[dict[str, Any]]:
|
|
rows: list[dict[str, Any]] = []
|
|
for binding in bindings:
|
|
rows.append({
|
|
"cloud_skill_id": binding.cloud_skill_id,
|
|
"local_skill_id": binding.local_skill_id,
|
|
"package_id": binding.current_package_id or binding.package_id_at_pull,
|
|
"package_path_at_pull": binding.package_path_at_pull,
|
|
"local_path": sanitize_upload_path(binding.local_path, workspace_root=workspace_root),
|
|
"manifest_hash": binding.manifest_hash,
|
|
"absorbed_at": binding.last_pulled_at,
|
|
})
|
|
return rows
|
|
|
|
|
|
def _evolution_event_rows(final_result: Mapping[str, Any]) -> list[Mapping[str, Any]]:
|
|
rows: list[Mapping[str, Any]] = []
|
|
for index, item in enumerate(final_result.get("evolved_skills") or []):
|
|
if not isinstance(item, Mapping):
|
|
continue
|
|
rows.append({
|
|
"evolve_event_id": str(item.get("evolve_event_id") or f"evolve_{index + 1:06d}"),
|
|
"event_type": str(item.get("origin") or item.get("type") or "evolve"),
|
|
"source_cloud_skill_id": item.get("source_cloud_skill_id"),
|
|
"source_local_skill_id": item.get("source_local_skill_id"),
|
|
"result_local_skill_id": item.get("local_skill_id") or item.get("skill_id"),
|
|
"result_cloud_skill_id": item.get("cloud_skill_id"),
|
|
"status": str(item.get("status") or "success"),
|
|
"summary": item.get("summary") or item.get("change_summary") or "",
|
|
"created_at": item.get("created_at") or "",
|
|
})
|
|
return rows
|
|
|
|
|
|
def _mapping_get(value: Any, key: str) -> Any:
|
|
return value.get(key) if isinstance(value, Mapping) else None
|