mirror of
https://github.com/agent0ai/agent-zero.git
synced 2026-05-21 10:31:45 +00:00
4f95250d0b
Teach document_artifact to create embedded spreadsheet charts through a native create_chart operation, including generic line/bar/column/pie/area/scatter support and stock-style OHLC charts. Parse CSV, TSV, and Markdown table content into real XLSX cells during spreadsheet creation so chart ranges bind to typed data instead of row text blobs. Update the Office artifact skill and tool prompt to prefer native chart creation over Python fallback, and cover the workflow with regression tests.
151 lines
7 KiB
Python
151 lines
7 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from helpers.tool import Response, Tool
|
|
from plugins._office.helpers import artifact_editor, collabora_status, wopi_store
|
|
|
|
|
|
class DocumentArtifact(Tool):
|
|
async def execute(
|
|
self,
|
|
action: str = "",
|
|
kind: str = "document",
|
|
title: str = "Untitled",
|
|
format: str = "docx",
|
|
content: str = "",
|
|
path: str = "",
|
|
file_id: str = "",
|
|
version_id: int | str | None = None,
|
|
operation: str = "",
|
|
find: str = "",
|
|
replace: str = "",
|
|
sheet: str = "",
|
|
cells: Any = None,
|
|
rows: Any = None,
|
|
chart: Any = None,
|
|
slides: Any = None,
|
|
max_chars: int | str = 12000,
|
|
**kwargs: Any,
|
|
) -> Response:
|
|
action = str(action or self.method or "status").strip().lower().replace("-", "_")
|
|
try:
|
|
if action == "create":
|
|
doc = wopi_store.create_document(kind=kind, title=title, fmt=format, content=content, path=path)
|
|
return self._document_response("Created document artifact.", doc, action=action)
|
|
if action == "open":
|
|
doc = self._document_from_input(file_id=file_id, path=path)
|
|
return self._document_response("Opened document artifact.", doc, action=action)
|
|
if action in {"read", "extract"}:
|
|
doc = self._document_from_input(file_id=file_id, path=path)
|
|
payload = {
|
|
"ok": True,
|
|
"action": "read",
|
|
"document": self._public_doc(doc),
|
|
"content": artifact_editor.read_artifact(doc, max_chars=int(max_chars or 12000)),
|
|
}
|
|
return self._json_response(payload, doc=doc, action="read")
|
|
if action in {"edit", "update", "patch"}:
|
|
doc = self._document_from_input(file_id=file_id, path=path)
|
|
updated_doc, payload = artifact_editor.edit_artifact(
|
|
doc,
|
|
operation=operation,
|
|
content=content,
|
|
find=find,
|
|
replace=replace,
|
|
sheet=sheet,
|
|
cells=cells,
|
|
rows=rows,
|
|
chart=chart,
|
|
slides=slides,
|
|
**kwargs,
|
|
)
|
|
payload["document"] = self._public_doc(updated_doc)
|
|
return self._json_response(payload, doc=updated_doc, action="edit")
|
|
if action == "inspect":
|
|
doc = self._document_from_input(file_id=file_id, path=path)
|
|
return self._json_response({"ok": True, "action": action, "document": self._public_doc(doc)}, doc=doc, action=action)
|
|
if action == "version_history":
|
|
doc = self._document_from_input(file_id=file_id, path=path)
|
|
versions = wopi_store.version_history(doc["file_id"])
|
|
return self._json_response({"ok": True, "action": action, "versions": versions}, doc=doc, action=action)
|
|
if action == "restore_version":
|
|
if version_id is None or str(version_id).strip() == "":
|
|
return Response(message="version_id is required for restore_version.", break_loop=False)
|
|
doc = self._document_from_input(file_id=file_id, path=path)
|
|
restored = wopi_store.restore_version(doc["file_id"], int(version_id))
|
|
return self._document_response("Restored document artifact version.", restored, action=action)
|
|
if action == "export":
|
|
doc = self._document_from_input(file_id=file_id, path=path)
|
|
target_format = str(kwargs.get("target_format") or kwargs.get("export_format") or "").lower().lstrip(".")
|
|
if target_format and target_format != doc["extension"]:
|
|
return Response(
|
|
message=f"Export to .{target_format} is not available yet. The source file remains unchanged at {doc['path']}.",
|
|
break_loop=False,
|
|
additional=self._additional(doc, action=action),
|
|
)
|
|
return self._document_response("Document artifact export path is ready.", doc, action=action)
|
|
if action == "status":
|
|
return self._json_response({"ok": True, "action": action, "status": collabora_status.collect_status()}, action=action)
|
|
return Response(message=f"Unknown document_artifact action: {action}", break_loop=False)
|
|
except Exception as exc:
|
|
return Response(message=f"document_artifact {action} failed: {exc}", break_loop=False)
|
|
|
|
def get_log_object(self):
|
|
return self.agent.context.log.log(
|
|
type="tool",
|
|
heading=f"icon://description {self.agent.agent_name}: Using document artifact",
|
|
content="",
|
|
kvps={**self.args, "_tool_name": self.name},
|
|
_tool_name=self.name,
|
|
)
|
|
|
|
def _document_from_input(self, file_id: str = "", path: str = "") -> dict[str, Any]:
|
|
if file_id:
|
|
return wopi_store.get_document(file_id)
|
|
if path:
|
|
return wopi_store.register_document(path)
|
|
raise ValueError("file_id or path is required")
|
|
|
|
def _document_response(self, message: str, doc: dict[str, Any], action: str = "") -> Response:
|
|
payload = {"ok": True, "action": action, "message": message, "document": self._public_doc(doc)}
|
|
return Response(
|
|
message=json.dumps(payload, indent=2, ensure_ascii=False),
|
|
break_loop=False,
|
|
additional=self._additional(doc, action=action),
|
|
)
|
|
|
|
def _json_response(self, payload: dict[str, Any], doc: dict[str, Any] | None = None, action: str = "") -> Response:
|
|
return Response(
|
|
message=json.dumps(payload, indent=2, ensure_ascii=False, default=str),
|
|
break_loop=False,
|
|
additional=self._additional(doc, action=action) if doc else {"_tool_name": self.name, "canvas_surface": "office", "action": action},
|
|
)
|
|
|
|
def _additional(self, doc: dict[str, Any] | None, action: str = "") -> dict[str, Any]:
|
|
if not doc:
|
|
return {"_tool_name": self.name, "canvas_surface": "office", "action": action}
|
|
return {
|
|
"_tool_name": self.name,
|
|
"canvas_surface": "office",
|
|
"action": action,
|
|
"file_id": doc["file_id"],
|
|
"title": doc["basename"],
|
|
"format": doc["extension"],
|
|
"path": doc["path"],
|
|
"version": wopi_store.item_version(doc),
|
|
}
|
|
|
|
def _public_doc(self, doc: dict[str, Any]) -> dict[str, Any]:
|
|
return {
|
|
"file_id": doc["file_id"],
|
|
"path": doc["path"],
|
|
"basename": doc["basename"],
|
|
"extension": doc["extension"],
|
|
"size": doc["size"],
|
|
"version": wopi_store.item_version(doc),
|
|
"last_modified": doc["last_modified"],
|
|
"exists": Path(doc["path"]).exists(),
|
|
}
|