mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-04 14:12:58 +00:00
Add document generation tool with editable spec-backed artifacts
Add an artifact generator tool with create, edit, and rewrite actions that render presentations, documents, spreadsheets, and PDFs from a validated JSON spec by running a fixed program in the sandbox (the spec travels as data, so its contents can never execute). The spec is stored as the source of truth on each artifact version; edits apply a JSON merge-patch and re-render, appending a new version while earlier versions stay intact.
This commit is contained in:
Alex committed
1 parent
1ac1f793b7
commit
381fb3335d
5 files changed
+1006
-2
No files matched your search
@@ -0,0 +1,545 @@
|
||||
"""Artifact Generator tool: render editable documents from a JSON spec and version them append-only.
|
||||
|
||||
The ``artifact_versions.spec`` JSONB is the source of truth; the rendered
|
||||
``.pptx``/``.docx``/``.xlsx``/``.pdf`` is derived. ``create_artifact`` stores
|
||||
v1, ``edit_artifact`` applies an RFC 7386 merge-patch to the current spec and
|
||||
appends a version, ``rewrite_artifact`` replaces the spec wholesale and appends
|
||||
a version. Rendering runs a FIXED program in the sandbox that reads the spec as
|
||||
DATA (``json.loads``) — spec values are never interpolated into the program, so
|
||||
a spec string containing code/quotes is rendered as literal text, not executed.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from application.agents.tools.base import Tool
|
||||
from application.core.settings import settings
|
||||
from application.sandbox.artifacts_capture import append_artifact_version, persist_new_artifact
|
||||
from application.sandbox.sandbox_creator import SandboxCreator
|
||||
from application.storage.db.repositories.artifacts import ArtifactsRepository
|
||||
from application.storage.db.session import db_readonly
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
import jsonschema
|
||||
except Exception: # pragma: no cover - jsonschema is a declared dependency
|
||||
jsonschema = None # type: ignore[assignment]
|
||||
|
||||
# Per-kind output metadata: artifact ``kind`` + produced file extension + mime.
|
||||
_KIND_INFO: Dict[str, Dict[str, str]] = {
|
||||
"presentation": {
|
||||
"kind": "presentation",
|
||||
"ext": "pptx",
|
||||
"mime": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||
},
|
||||
"document": {
|
||||
"kind": "document",
|
||||
"ext": "docx",
|
||||
"mime": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
},
|
||||
"spreadsheet": {
|
||||
"kind": "spreadsheet",
|
||||
"ext": "xlsx",
|
||||
"mime": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
},
|
||||
"pdf": {
|
||||
"kind": "document",
|
||||
"ext": "pdf",
|
||||
"mime": "application/pdf",
|
||||
},
|
||||
}
|
||||
|
||||
# Tight per-kind JSON schemas. ``additionalProperties: false`` keeps specs minimal
|
||||
# and rejects stray keys before any rendering happens.
|
||||
_SCHEMAS: Dict[str, Dict[str, Any]] = {
|
||||
"presentation": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"title": {"type": "string"},
|
||||
"slides": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"title": {"type": "string"},
|
||||
"bullets": {"type": "array", "items": {"type": "string"}},
|
||||
"notes": {"type": "string"},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["slides"],
|
||||
},
|
||||
"document": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"title": {"type": "string"},
|
||||
"sections": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"heading": {"type": "string"},
|
||||
"paragraphs": {"type": "array", "items": {"type": "string"}},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["sections"],
|
||||
},
|
||||
"spreadsheet": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"sheets": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"name": {"type": "string"},
|
||||
"rows": {
|
||||
"type": "array",
|
||||
"items": {"type": "array", "items": {}},
|
||||
},
|
||||
},
|
||||
"required": ["rows"],
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["sheets"],
|
||||
},
|
||||
"pdf": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"title": {"type": "string"},
|
||||
"blocks": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"type": {"type": "string", "enum": ["heading", "paragraph"]},
|
||||
"text": {"type": "string"},
|
||||
},
|
||||
"required": ["type", "text"],
|
||||
},
|
||||
},
|
||||
},
|
||||
"required": ["blocks"],
|
||||
},
|
||||
}
|
||||
|
||||
# FIXED renderer programs. Each reads ``spec.json`` from the workspace as DATA
|
||||
# and writes ``out.<ext>``. The spec is NEVER string-interpolated into the
|
||||
# program; ``{spec_path}``/``{out_path}`` are server-controlled path literals.
|
||||
_RENDERERS: Dict[str, str] = {
|
||||
"presentation": (
|
||||
"import json\n"
|
||||
"from pptx import Presentation\n"
|
||||
"from pptx.util import Pt\n"
|
||||
"spec = json.load(open({spec_path!r}))\n"
|
||||
"prs = Presentation()\n"
|
||||
"blank = prs.slide_layouts[6]\n"
|
||||
"title_only = prs.slide_layouts[5]\n"
|
||||
"for s in spec.get('slides', []):\n"
|
||||
" slide = prs.slides.add_slide(title_only)\n"
|
||||
" slide.shapes.title.text = str(s.get('title', '') or '')\n"
|
||||
" bullets = s.get('bullets') or []\n"
|
||||
" if bullets:\n"
|
||||
" left = top = Pt(72)\n"
|
||||
" width = prs.slide_width - Pt(144)\n"
|
||||
" height = prs.slide_height - Pt(216)\n"
|
||||
" box = slide.shapes.add_textbox(left, Pt(150), width, height)\n"
|
||||
" tf = box.text_frame\n"
|
||||
" tf.word_wrap = True\n"
|
||||
" for i, b in enumerate(bullets):\n"
|
||||
" para = tf.paragraphs[0] if i == 0 else tf.add_paragraph()\n"
|
||||
" para.text = str(b)\n"
|
||||
" notes = s.get('notes')\n"
|
||||
" if notes:\n"
|
||||
" slide.notes_slide.notes_text_frame.text = str(notes)\n"
|
||||
"prs.save({out_path!r})\n"
|
||||
),
|
||||
"document": (
|
||||
"import json\n"
|
||||
"from docx import Document\n"
|
||||
"spec = json.load(open({spec_path!r}))\n"
|
||||
"doc = Document()\n"
|
||||
"title = spec.get('title')\n"
|
||||
"if title:\n"
|
||||
" doc.add_heading(str(title), level=0)\n"
|
||||
"for sec in spec.get('sections', []):\n"
|
||||
" heading = sec.get('heading')\n"
|
||||
" if heading:\n"
|
||||
" doc.add_heading(str(heading), level=1)\n"
|
||||
" for p in (sec.get('paragraphs') or []):\n"
|
||||
" doc.add_paragraph(str(p))\n"
|
||||
"doc.save({out_path!r})\n"
|
||||
),
|
||||
"spreadsheet": (
|
||||
"import json\n"
|
||||
"from openpyxl import Workbook\n"
|
||||
"spec = json.load(open({spec_path!r}))\n"
|
||||
"wb = Workbook()\n"
|
||||
"wb.remove(wb.active)\n"
|
||||
"for idx, sheet in enumerate(spec.get('sheets', [])):\n"
|
||||
" name = str(sheet.get('name') or ('Sheet%d' % (idx + 1)))[:31]\n"
|
||||
" ws = wb.create_sheet(title=name)\n"
|
||||
" for row in (sheet.get('rows') or []):\n"
|
||||
" ws.append([('' if c is None else c) for c in row])\n"
|
||||
"if not wb.sheetnames:\n"
|
||||
" wb.create_sheet(title='Sheet1')\n"
|
||||
"wb.save({out_path!r})\n"
|
||||
),
|
||||
"pdf": (
|
||||
"import json\n"
|
||||
"from reportlab.lib.pagesizes import letter\n"
|
||||
"from reportlab.lib.styles import getSampleStyleSheet\n"
|
||||
"from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer\n"
|
||||
"from xml.sax.saxutils import escape\n"
|
||||
"spec = json.load(open({spec_path!r}))\n"
|
||||
"styles = getSampleStyleSheet()\n"
|
||||
"story = []\n"
|
||||
"title = spec.get('title')\n"
|
||||
"if title:\n"
|
||||
" story.append(Paragraph(escape(str(title)), styles['Title']))\n"
|
||||
" story.append(Spacer(1, 12))\n"
|
||||
"for block in spec.get('blocks', []):\n"
|
||||
" style = styles['Heading1'] if block.get('type') == 'heading' else styles['BodyText']\n"
|
||||
" story.append(Paragraph(escape(str(block.get('text', ''))), style))\n"
|
||||
" story.append(Spacer(1, 6))\n"
|
||||
"SimpleDocTemplate({out_path!r}, pagesize=letter).build(story)\n"
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def merge_patch(target: Any, patch: Any) -> Any:
|
||||
"""Apply an RFC 7386 JSON Merge Patch to ``target`` and return the result."""
|
||||
if not isinstance(patch, dict):
|
||||
return copy.deepcopy(patch)
|
||||
if not isinstance(target, dict):
|
||||
target = {}
|
||||
result = copy.deepcopy(target)
|
||||
for key, value in patch.items():
|
||||
if value is None:
|
||||
result.pop(key, None)
|
||||
else:
|
||||
result[key] = merge_patch(result.get(key), value)
|
||||
return result
|
||||
|
||||
|
||||
class ArtifactGeneratorTool(Tool):
|
||||
"""Artifact Generator
|
||||
Render editable documents (presentation/document/spreadsheet/pdf) from a JSON spec and version them.
|
||||
"""
|
||||
|
||||
def __init__(self, tool_config: Optional[Dict[str, Any]] = None, user_id: Optional[str] = None) -> None:
|
||||
"""Bind the tool to the invoker and its conversation/run-scoped sandbox session."""
|
||||
self.config: Dict[str, Any] = tool_config or {}
|
||||
self.user_id: Optional[str] = user_id
|
||||
self.tool_id: Optional[str] = self.config.get("tool_id")
|
||||
self.conversation_id: Optional[str] = self.config.get("conversation_id")
|
||||
self.workflow_run_id: Optional[str] = self.config.get("workflow_run_id")
|
||||
self._last_artifact_id: Optional[str] = None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Tool ABC
|
||||
# ------------------------------------------------------------------
|
||||
def get_actions_metadata(self) -> List[Dict[str, Any]]:
|
||||
"""Return JSON metadata describing the create/edit/rewrite actions for tool schemas."""
|
||||
kinds = sorted(_KIND_INFO.keys())
|
||||
return [
|
||||
{
|
||||
"name": "create_artifact",
|
||||
"description": (
|
||||
"Render a new editable document from a JSON spec and store it as version 1. "
|
||||
"The spec is the source of truth; the rendered file is derived."
|
||||
),
|
||||
"active": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"kind": {"type": "string", "enum": kinds, "description": "Document kind to render."},
|
||||
"title": {"type": "string", "description": "Optional artifact title."},
|
||||
"spec": {"type": "object", "description": "Document spec matching the kind's schema."},
|
||||
},
|
||||
"required": ["kind", "spec"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "edit_artifact",
|
||||
"description": (
|
||||
"Apply a JSON merge-patch (RFC 7386) to the current spec, re-render, and append a "
|
||||
"new version. Preferred for small, targeted changes."
|
||||
),
|
||||
"active": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": {"type": "string", "description": "Artifact id to edit."},
|
||||
"spec_patch": {
|
||||
"type": "object",
|
||||
"description": "RFC 7386 merge-patch; null values delete keys.",
|
||||
},
|
||||
},
|
||||
"required": ["id", "spec_patch"],
|
||||
},
|
||||
},
|
||||
{
|
||||
"name": "rewrite_artifact",
|
||||
"description": "Replace the spec wholesale, re-render, and append a new version.",
|
||||
"active": True,
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": {"type": "string", "description": "Artifact id to rewrite."},
|
||||
"spec": {"type": "object", "description": "Replacement spec matching the kind's schema."},
|
||||
},
|
||||
"required": ["id", "spec"],
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
def get_config_requirements(self) -> Dict[str, Any]:
|
||||
"""Return configuration requirements (none beyond the deployment sandbox backend)."""
|
||||
return {}
|
||||
|
||||
def get_artifact_id(self, action_name: str, **kwargs: Any) -> Optional[str]:
|
||||
"""Return the produced artifact id so the UI artifact rail lights up."""
|
||||
return self._last_artifact_id
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Dispatch
|
||||
# ------------------------------------------------------------------
|
||||
def execute_action(self, action_name: str, **kwargs: Any) -> Dict[str, Any]:
|
||||
"""Dispatch a create/edit/rewrite action."""
|
||||
self._last_artifact_id = None
|
||||
if not self.user_id:
|
||||
return {"status": "error", "error": "artifact_generator requires a valid user_id."}
|
||||
if self.conversation_id is None and self.workflow_run_id is None:
|
||||
return {"status": "error", "error": "artifact_generator requires a conversation_id or workflow_run_id."}
|
||||
if jsonschema is None:
|
||||
return {"status": "error", "error": "jsonschema is required for spec validation."}
|
||||
if action_name == "create_artifact":
|
||||
return self._create(**kwargs)
|
||||
if action_name == "edit_artifact":
|
||||
return self._edit(**kwargs)
|
||||
if action_name == "rewrite_artifact":
|
||||
return self._rewrite(**kwargs)
|
||||
return {"status": "error", "error": f"unknown action: {action_name}"}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Actions
|
||||
# ------------------------------------------------------------------
|
||||
def _create(self, **kwargs: Any) -> Dict[str, Any]:
|
||||
"""Validate, render, and persist a new artifact at version 1."""
|
||||
kind = kwargs.get("kind")
|
||||
spec = kwargs.get("spec")
|
||||
title = kwargs.get("title")
|
||||
if kind not in _KIND_INFO:
|
||||
return {"status": "error", "error": f"unsupported kind: {kind!r}; expected one of {sorted(_KIND_INFO)}."}
|
||||
valid = self._validate(kind, spec)
|
||||
if valid is not None:
|
||||
return valid
|
||||
|
||||
rendered = self._render(kind, spec)
|
||||
if rendered.get("error"):
|
||||
return {"status": "error", "error": rendered["error"]}
|
||||
|
||||
info = _KIND_INFO[kind]
|
||||
filename = self._filename(title, info["ext"])
|
||||
ref = persist_new_artifact(
|
||||
user_id=self.user_id,
|
||||
kind=info["kind"],
|
||||
data=rendered["data"],
|
||||
filename=filename,
|
||||
mime_type=info["mime"],
|
||||
title=title,
|
||||
conversation_id=self.conversation_id,
|
||||
workflow_run_id=self.workflow_run_id,
|
||||
spec=spec,
|
||||
produced_by=self._produced_by("create_artifact", kind),
|
||||
)
|
||||
if ref is None:
|
||||
return {"status": "error", "error": "failed to persist artifact."}
|
||||
self._last_artifact_id = ref["artifact_id"]
|
||||
return {"status": "ok", **ref}
|
||||
|
||||
def _edit(self, **kwargs: Any) -> Dict[str, Any]:
|
||||
"""Merge-patch the current spec, re-render, and append a version."""
|
||||
artifact_id = kwargs.get("id")
|
||||
spec_patch = kwargs.get("spec_patch")
|
||||
if not isinstance(spec_patch, dict):
|
||||
return {"status": "error", "error": "spec_patch must be a JSON object (merge-patch)."}
|
||||
loaded = self._load_current(artifact_id)
|
||||
if loaded.get("error"):
|
||||
return {"status": "error", "error": loaded["error"]}
|
||||
kind = loaded["kind"]
|
||||
new_spec = merge_patch(loaded["spec"], spec_patch)
|
||||
return self._reversion(artifact_id, kind, new_spec, "edit_artifact")
|
||||
|
||||
def _rewrite(self, **kwargs: Any) -> Dict[str, Any]:
|
||||
"""Replace the spec wholesale, re-render, and append a version."""
|
||||
artifact_id = kwargs.get("id")
|
||||
spec = kwargs.get("spec")
|
||||
loaded = self._load_current(artifact_id)
|
||||
if loaded.get("error"):
|
||||
return {"status": "error", "error": loaded["error"]}
|
||||
return self._reversion(artifact_id, loaded["kind"], spec, "rewrite_artifact")
|
||||
|
||||
def _reversion(self, artifact_id: str, kind: str, spec: Any, action: str) -> Dict[str, Any]:
|
||||
"""Validate the new spec, re-render, and append the next version of an existing artifact."""
|
||||
valid = self._validate(kind, spec)
|
||||
if valid is not None:
|
||||
return valid
|
||||
rendered = self._render(kind, spec)
|
||||
if rendered.get("error"):
|
||||
return {"status": "error", "error": rendered["error"]}
|
||||
info = _KIND_INFO[kind]
|
||||
filename = self._filename(None, info["ext"])
|
||||
ref = append_artifact_version(
|
||||
user_id=self.user_id,
|
||||
artifact_id=artifact_id,
|
||||
data=rendered["data"],
|
||||
filename=filename,
|
||||
mime_type=info["mime"],
|
||||
spec=spec,
|
||||
produced_by=self._produced_by(action, kind),
|
||||
)
|
||||
if ref is None:
|
||||
return {"status": "error", "error": "failed to persist artifact version."}
|
||||
self._last_artifact_id = ref["artifact_id"]
|
||||
return {"status": "ok", **ref}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Spec / render helpers
|
||||
# ------------------------------------------------------------------
|
||||
def _validate(self, kind: str, spec: Any) -> Optional[Dict[str, Any]]:
|
||||
"""Return an error payload when ``spec`` is invalid for ``kind``, else None."""
|
||||
if not isinstance(spec, dict):
|
||||
return {"status": "error", "error": "spec must be a JSON object."}
|
||||
try:
|
||||
jsonschema.validate(spec, _SCHEMAS[kind])
|
||||
except jsonschema.ValidationError as exc:
|
||||
return {"status": "error", "error": f"invalid {kind} spec: {exc.message}"}
|
||||
return None
|
||||
|
||||
def _load_current(self, artifact_id: Any) -> Dict[str, Any]:
|
||||
"""Fetch the parent-scoped artifact and its current-version spec for edit/rewrite."""
|
||||
if not isinstance(artifact_id, str) or not artifact_id.strip():
|
||||
return {"error": "id is required."}
|
||||
try:
|
||||
with db_readonly() as conn:
|
||||
repo = ArtifactsRepository(conn)
|
||||
artifact = repo.get_artifact_in_parent(
|
||||
artifact_id,
|
||||
conversation_id=self.conversation_id,
|
||||
workflow_run_id=self.workflow_run_id,
|
||||
)
|
||||
if artifact is None:
|
||||
return {"error": f"artifact {artifact_id} not found in this conversation/run."}
|
||||
version = repo.get_version(artifact_id, artifact["current_version"])
|
||||
except Exception:
|
||||
logger.exception("artifact_generator: failed to load artifact")
|
||||
return {"error": f"failed to load artifact {artifact_id}."}
|
||||
if not version or version.get("spec") is None:
|
||||
return {"error": f"artifact {artifact_id} has no editable spec."}
|
||||
kind = self._kind_for(artifact, version)
|
||||
if kind is None:
|
||||
return {"error": f"artifact {artifact_id} is not a spec-rendered document."}
|
||||
return {"kind": kind, "spec": version["spec"]}
|
||||
|
||||
@staticmethod
|
||||
def _kind_for(artifact: Dict[str, Any], version: Dict[str, Any]) -> Optional[str]:
|
||||
"""Resolve the spec kind from ``produced_by`` (preferred) or the version mime type."""
|
||||
produced = version.get("produced_by")
|
||||
if isinstance(produced, dict):
|
||||
spec_kind = produced.get("spec_kind")
|
||||
if spec_kind in _KIND_INFO:
|
||||
return spec_kind
|
||||
mime = version.get("mime_type") or ""
|
||||
for spec_kind, info in _KIND_INFO.items():
|
||||
if info["mime"] == mime:
|
||||
return spec_kind
|
||||
return None
|
||||
|
||||
def _render(self, kind: str, spec: Any) -> Dict[str, Any]:
|
||||
"""Run the fixed renderer in the sandbox and return the produced file bytes."""
|
||||
session_id = self._resolve_session_id()
|
||||
if session_id is None:
|
||||
return {"error": "artifact_generator requires a conversation_id or workflow_run_id."}
|
||||
|
||||
token = uuid.uuid4().hex
|
||||
spec_path = f"artifacts/{token}/spec.json"
|
||||
out_path = f"artifacts/{token}/out.{_KIND_INFO[kind]['ext']}"
|
||||
program = _RENDERERS[kind].format(spec_path=spec_path, out_path=out_path)
|
||||
timeout = float(getattr(settings, "SANDBOX_EXEC_TIMEOUT", 60))
|
||||
|
||||
manager = SandboxCreator.get_manager()
|
||||
try:
|
||||
manager.open(session_id, ttl=timeout)
|
||||
except Exception as exc:
|
||||
logger.exception("artifact_generator: failed to open sandbox session")
|
||||
return {"error": f"sandbox unavailable: {type(exc).__name__}: {exc}"}
|
||||
try:
|
||||
# The spec rides in as a JSON file the program ``json.load``s; it is
|
||||
# never interpolated into the program, so its contents stay data.
|
||||
manager.put_file(session_id, spec_path, json.dumps(spec).encode("utf-8"))
|
||||
result = manager.exec(session_id, program, timeout=timeout)
|
||||
if not result.ok:
|
||||
detail = (
|
||||
f"{result.error_name}: {result.error_value}"
|
||||
if result.error_name
|
||||
else (result.error_value or "render failed")
|
||||
)
|
||||
return {"error": f"render failed: {detail}"}
|
||||
data = manager.get_file(session_id, out_path)
|
||||
except Exception as exc:
|
||||
logger.exception("artifact_generator: render failed")
|
||||
return {"error": f"render failed: {type(exc).__name__}: {exc}"}
|
||||
finally:
|
||||
try:
|
||||
manager.close(session_id)
|
||||
except Exception:
|
||||
logger.exception("artifact_generator: session close failed")
|
||||
if not data:
|
||||
return {"error": "renderer produced an empty file."}
|
||||
return {"data": data}
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Misc helpers
|
||||
# ------------------------------------------------------------------
|
||||
def _produced_by(self, action: str, kind: str) -> Dict[str, Any]:
|
||||
"""Build the ``produced_by`` provenance record, carrying the spec kind for re-editing."""
|
||||
return {
|
||||
"tool": "artifact_generator",
|
||||
"action": action,
|
||||
"spec_kind": kind,
|
||||
"tool_id": self.tool_id,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _filename(title: Optional[str], ext: str) -> str:
|
||||
"""Derive a download filename from a title (or a generic stem) plus the kind extension."""
|
||||
stem = (title or "artifact").strip() or "artifact"
|
||||
return f"{stem}.{ext}"
|
||||
|
||||
def _resolve_session_id(self) -> Optional[str]:
|
||||
"""Derive the sandbox session id from the bound conversation/run; sanitize to the gateway charset."""
|
||||
raw = self.conversation_id or self.workflow_run_id
|
||||
if not raw:
|
||||
return None
|
||||
sanitized = "".join(c if c.isalnum() or c in "-_" else "-" for c in str(raw))
|
||||
return sanitized or None
|
||||
@@ -30,7 +30,16 @@ class ToolManager:
|
||||
if issubclass(obj, Tool) and obj is not Tool:
|
||||
if (
|
||||
tool_name
|
||||
in {"mcp_tool", "notes", "memory", "todo_list", "scheduler", "remote_device", "code_executor"}
|
||||
in {
|
||||
"mcp_tool",
|
||||
"notes",
|
||||
"memory",
|
||||
"todo_list",
|
||||
"scheduler",
|
||||
"remote_device",
|
||||
"code_executor",
|
||||
"artifact_generator",
|
||||
}
|
||||
and user_id
|
||||
):
|
||||
return obj(tool_config, user_id)
|
||||
@@ -42,7 +51,16 @@ class ToolManager:
|
||||
raise ValueError(f"Tool '{tool_name}' not loaded")
|
||||
if (
|
||||
tool_name
|
||||
in {"mcp_tool", "memory", "todo_list", "notes", "scheduler", "remote_device", "code_executor"}
|
||||
in {
|
||||
"mcp_tool",
|
||||
"memory",
|
||||
"todo_list",
|
||||
"notes",
|
||||
"scheduler",
|
||||
"remote_device",
|
||||
"code_executor",
|
||||
"artifact_generator",
|
||||
}
|
||||
and user_id
|
||||
):
|
||||
tool_config = self.config.get(tool_name, {})
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
"""Unit tests for ArtifactGeneratorTool: spec validation, merge-patch, and renderer injection safety.
|
||||
|
||||
No DB or sandbox: these exercise the pure logic (schema gate + RFC 7386 merge)
|
||||
and the FIXED renderer programs directly (executed in-process against a temp dir)
|
||||
to prove a spec value containing Python/quotes is rendered as literal text.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
|
||||
from application.agents.tools.artifact_generator import (
|
||||
_KIND_INFO,
|
||||
_RENDERERS,
|
||||
ArtifactGeneratorTool,
|
||||
merge_patch,
|
||||
)
|
||||
|
||||
pytest.importorskip("pptx")
|
||||
pytest.importorskip("docx")
|
||||
pytest.importorskip("openpyxl")
|
||||
pytest.importorskip("reportlab")
|
||||
|
||||
|
||||
def _tool():
|
||||
return ArtifactGeneratorTool(
|
||||
tool_config={"conversation_id": "conv-1", "tool_id": "t-1"}, user_id="u-1"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Spec validation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_validate_accepts_minimal_presentation():
|
||||
assert _tool()._validate("presentation", {"slides": [{"title": "x"}]}) is None
|
||||
|
||||
|
||||
def test_validate_rejects_missing_required_key():
|
||||
err = _tool()._validate("presentation", {"title": "no slides"})
|
||||
assert err["status"] == "error"
|
||||
assert "invalid presentation spec" in err["error"]
|
||||
|
||||
|
||||
def test_validate_rejects_unknown_key():
|
||||
err = _tool()._validate("document", {"sections": [], "bogus": 1})
|
||||
assert err["status"] == "error"
|
||||
|
||||
|
||||
def test_validate_rejects_wrong_type():
|
||||
err = _tool()._validate("spreadsheet", {"sheets": "not-a-list"})
|
||||
assert err["status"] == "error"
|
||||
|
||||
|
||||
def test_validate_rejects_non_object_spec():
|
||||
err = _tool()._validate("pdf", "just a string")
|
||||
assert err["status"] == "error"
|
||||
assert "spec must be a JSON object" in err["error"]
|
||||
|
||||
|
||||
def test_create_rejects_unknown_kind():
|
||||
out = _tool()._create(kind="hologram", spec={})
|
||||
assert out["status"] == "error"
|
||||
assert "unsupported kind" in out["error"]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RFC 7386 JSON merge-patch
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_merge_patch_adds_and_overwrites():
|
||||
assert merge_patch({"a": 1, "b": 2}, {"b": 3, "c": 4}) == {"a": 1, "b": 3, "c": 4}
|
||||
|
||||
|
||||
def test_merge_patch_null_deletes_key():
|
||||
assert merge_patch({"a": 1, "b": 2}, {"b": None}) == {"a": 1}
|
||||
|
||||
|
||||
def test_merge_patch_recurses_into_objects():
|
||||
assert merge_patch({"x": {"a": 1, "b": 2}}, {"x": {"b": None, "c": 3}}) == {"x": {"a": 1, "c": 3}}
|
||||
|
||||
|
||||
def test_merge_patch_replaces_array_wholesale():
|
||||
# RFC 7386: arrays are replaced, not merged element-wise.
|
||||
assert merge_patch({"l": [1, 2, 3]}, {"l": [9]}) == {"l": [9]}
|
||||
|
||||
|
||||
def test_merge_patch_non_object_patch_replaces_target():
|
||||
assert merge_patch({"a": 1}, "scalar") == "scalar"
|
||||
|
||||
|
||||
def test_merge_patch_does_not_mutate_inputs():
|
||||
target = {"a": {"b": 1}}
|
||||
patch = {"a": {"c": 2}}
|
||||
merge_patch(target, patch)
|
||||
assert target == {"a": {"b": 1}}
|
||||
assert patch == {"a": {"c": 2}}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Renderer injection safety
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _render_in_process(kind: str, spec: dict) -> str:
|
||||
"""Execute the FIXED renderer program against a temp dir; return the output path."""
|
||||
workdir = tempfile.mkdtemp()
|
||||
spec_path = os.path.join(workdir, "spec.json")
|
||||
out_path = os.path.join(workdir, f"out.{_KIND_INFO[kind]['ext']}")
|
||||
with open(spec_path, "w") as handle:
|
||||
json.dump(spec, handle)
|
||||
program = _RENDERERS[kind].format(spec_path=spec_path, out_path=out_path)
|
||||
namespace: dict = {}
|
||||
exec(compile(program, "<renderer>", "exec"), namespace, namespace) # noqa: S102
|
||||
return out_path
|
||||
|
||||
|
||||
def test_renderer_does_not_execute_spec_code(capfd, tmp_path):
|
||||
# A spec whose values are Python source / shell payloads must be treated as
|
||||
# literal text. If the renderer string-interpolated the spec it would run
|
||||
# this; instead it json.loads the spec as data, so nothing executes.
|
||||
sentinel = tmp_path / "pwned.txt"
|
||||
payload = (
|
||||
f"'''__import__('os').system('echo PWNED > {sentinel}')'''; "
|
||||
"print('SHOULD_NOT_PRINT')"
|
||||
)
|
||||
spec = {
|
||||
"title": payload,
|
||||
"slides": [{"title": payload, "bullets": [payload], "notes": payload}],
|
||||
}
|
||||
out_path = _render_in_process("presentation", spec)
|
||||
|
||||
captured = capfd.readouterr()
|
||||
assert "SHOULD_NOT_PRINT" not in captured.out
|
||||
assert "PWNED" not in captured.out
|
||||
assert not sentinel.exists()
|
||||
assert os.path.getsize(out_path) > 0
|
||||
|
||||
|
||||
def test_renderer_keeps_injection_text_as_literal_content():
|
||||
from pptx import Presentation
|
||||
|
||||
payload = "'''; import os; os.system('echo HACK'); x = '''"
|
||||
spec = {"slides": [{"title": payload, "bullets": [payload]}]}
|
||||
prs = Presentation(_render_in_process("presentation", spec))
|
||||
assert len(prs.slides) == 1
|
||||
assert prs.slides[0].shapes.title.text == payload
|
||||
|
||||
|
||||
def test_pdf_renderer_escapes_markup_and_does_not_execute():
|
||||
payload = "<b>not bold</b> & \"</para>\" '''os.system('x')'''"
|
||||
spec = {"title": payload, "blocks": [{"type": "paragraph", "text": payload}]}
|
||||
out_path = _render_in_process("pdf", spec)
|
||||
assert os.path.getsize(out_path) > 0
|
||||
@@ -0,0 +1,274 @@
|
||||
"""End-to-end ArtifactGeneratorTool: live Jupyter gateway + ephemeral Postgres + local storage.
|
||||
|
||||
Launches a real ``jupyter kernelgateway`` (no Docker — the credential helper
|
||||
hangs on dev machines), wires the tool to the ephemeral pytest-postgresql DB and
|
||||
a temp-dir ``LocalStorage``, and drives create/edit/rewrite through the full
|
||||
path: a spec is rendered in the kernel by the FIXED renderer, the produced file
|
||||
is persisted as an artifact version (spec = source of truth), and the rendered
|
||||
bytes are re-opened with the real library to assert structure.
|
||||
|
||||
Skips gracefully when the gateway binary, websocket-client, or a renderer
|
||||
library is unavailable.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import shutil
|
||||
import socket
|
||||
import subprocess
|
||||
import time
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
|
||||
requests = pytest.importorskip("requests")
|
||||
pytest.importorskip("websocket") # websocket-client
|
||||
pytest.importorskip("pptx")
|
||||
pytest.importorskip("docx")
|
||||
pytest.importorskip("openpyxl")
|
||||
|
||||
from application.agents.tools.artifact_generator import ArtifactGeneratorTool # noqa: E402
|
||||
from application.sandbox.jupyter_gateway import JupyterKernelGatewaySandbox # noqa: E402
|
||||
from application.sandbox.manager import SandboxManager # noqa: E402
|
||||
from application.sandbox.sandbox_creator import SandboxCreator # noqa: E402
|
||||
from application.storage.db.repositories.artifacts import ArtifactsRepository # noqa: E402
|
||||
from application.storage.local import LocalStorage # noqa: E402
|
||||
from application.storage.storage_creator import StorageCreator # noqa: E402
|
||||
|
||||
_GATEWAY_BIN = shutil.which("jupyter-kernelgateway") or shutil.which("jupyter")
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.integration,
|
||||
pytest.mark.skipif(
|
||||
_GATEWAY_BIN is None,
|
||||
reason="jupyter kernel gateway not installed (pip install jupyter-kernel-gateway)",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def _free_port() -> int:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
||||
sock.bind(("127.0.0.1", 0))
|
||||
return sock.getsockname()[1]
|
||||
|
||||
|
||||
def _gateway_cmd(port: int) -> list:
|
||||
if _GATEWAY_BIN.endswith("jupyter-kernelgateway"):
|
||||
base = [_GATEWAY_BIN]
|
||||
else:
|
||||
base = [_GATEWAY_BIN, "kernelgateway"]
|
||||
return base + [
|
||||
"--KernelGatewayApp.ip=127.0.0.1",
|
||||
f"--KernelGatewayApp.port={port}",
|
||||
"--ZMQChannelsWebsocketConnection.limit_rate=False",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def gateway_url():
|
||||
port = _free_port()
|
||||
proc = subprocess.Popen(_gateway_cmd(port), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
|
||||
url = f"http://127.0.0.1:{port}"
|
||||
deadline = time.time() + 30
|
||||
ready = False
|
||||
try:
|
||||
while time.time() < deadline:
|
||||
if proc.poll() is not None:
|
||||
pytest.skip("jupyter kernelgateway process exited during startup")
|
||||
try:
|
||||
if requests.get(f"{url}/api", timeout=1).status_code == 200:
|
||||
ready = True
|
||||
break
|
||||
except requests.RequestException:
|
||||
time.sleep(0.3)
|
||||
if not ready:
|
||||
pytest.skip("jupyter kernelgateway did not become ready in time")
|
||||
yield url
|
||||
finally:
|
||||
proc.terminate()
|
||||
try:
|
||||
proc.wait(timeout=10)
|
||||
except subprocess.TimeoutExpired:
|
||||
proc.kill()
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def wired_tool(gateway_url, pg_engine, tmp_path, monkeypatch):
|
||||
"""An ArtifactGeneratorTool wired to the live gateway, ephemeral PG, and a temp local store."""
|
||||
backend = JupyterKernelGatewaySandbox(gateway_url=gateway_url, default_timeout=60.0)
|
||||
SandboxCreator._instance = SandboxManager(backend=backend, max_ttl=1200.0)
|
||||
|
||||
storage = LocalStorage(base_dir=str(tmp_path))
|
||||
monkeypatch.setattr(StorageCreator, "_instance", storage, raising=False)
|
||||
|
||||
monkeypatch.setattr("application.storage.db.session.get_engine", lambda: pg_engine)
|
||||
|
||||
conversation_id = str(uuid.uuid4())
|
||||
tool = ArtifactGeneratorTool(
|
||||
tool_config={"conversation_id": conversation_id, "tool_id": str(uuid.uuid4())},
|
||||
user_id="user-e2e",
|
||||
)
|
||||
try:
|
||||
yield tool, conversation_id, pg_engine, storage
|
||||
finally:
|
||||
SandboxCreator.reset()
|
||||
|
||||
|
||||
def test_create_edit_rewrite_presentation_versions(wired_tool):
|
||||
tool, conversation_id, pg_engine, storage = wired_tool
|
||||
from pptx import Presentation
|
||||
|
||||
# create -> v1
|
||||
spec = {"title": "Q3", "slides": [{"title": "Intro", "bullets": ["one", "two"]}]}
|
||||
created = tool.execute_action("create_artifact", kind="presentation", title="Deck", spec=spec)
|
||||
assert created["status"] == "ok", created
|
||||
assert created["version"] == 1
|
||||
assert created["mime_type"].endswith("presentationml.presentation")
|
||||
assert created["filename"] == "Deck.pptx"
|
||||
artifact_id = created["artifact_id"]
|
||||
assert tool.get_artifact_id("create_artifact") == artifact_id
|
||||
|
||||
# v1 spec is the source of truth; the rendered .pptx re-opens with one slide.
|
||||
with pg_engine.connect() as conn:
|
||||
repo = ArtifactsRepository(conn)
|
||||
artifact = repo.get_artifact_in_parent(artifact_id, conversation_id=conversation_id)
|
||||
v1 = repo.get_version(artifact_id, 1)
|
||||
assert artifact["current_version"] == 1
|
||||
assert v1["spec"] == spec
|
||||
assert v1["produced_by"]["spec_kind"] == "presentation"
|
||||
prs = Presentation(storage.get_file(v1["storage_path"]))
|
||||
assert len(prs.slides) == 1
|
||||
assert prs.slides[0].shapes.title.text == "Intro"
|
||||
|
||||
# edit (merge-patch) -> v2: add a second slide, keep title.
|
||||
patch = {"slides": [{"title": "Intro", "bullets": ["one", "two"]}, {"title": "Outro"}]}
|
||||
edited = tool.execute_action("edit_artifact", id=artifact_id, spec_patch=patch)
|
||||
assert edited["status"] == "ok", edited
|
||||
assert edited["version"] == 2
|
||||
assert edited["artifact_id"] == artifact_id
|
||||
|
||||
with pg_engine.connect() as conn:
|
||||
repo = ArtifactsRepository(conn)
|
||||
artifact = repo.get_artifact_in_parent(artifact_id, conversation_id=conversation_id)
|
||||
v1_again = repo.get_version(artifact_id, 1)
|
||||
v2 = repo.get_version(artifact_id, 2)
|
||||
assert artifact["current_version"] == 2
|
||||
# Append-only: v1 unchanged and still present.
|
||||
assert v1_again["spec"] == spec
|
||||
assert v2["spec"]["title"] == "Q3"
|
||||
assert len(v2["spec"]["slides"]) == 2
|
||||
prs2 = Presentation(storage.get_file(v2["storage_path"]))
|
||||
assert len(prs2.slides) == 2
|
||||
|
||||
# rewrite -> v3: wholesale new spec.
|
||||
new_spec = {"title": "Fresh", "slides": [{"title": "A"}, {"title": "B"}, {"title": "C"}]}
|
||||
rewritten = tool.execute_action("rewrite_artifact", id=artifact_id, spec=new_spec)
|
||||
assert rewritten["status"] == "ok", rewritten
|
||||
assert rewritten["version"] == 3
|
||||
|
||||
with pg_engine.connect() as conn:
|
||||
repo = ArtifactsRepository(conn)
|
||||
artifact = repo.get_artifact_in_parent(artifact_id, conversation_id=conversation_id)
|
||||
versions = repo.list_versions(artifact_id)
|
||||
v3 = repo.get_version(artifact_id, 3)
|
||||
assert artifact["current_version"] == 3
|
||||
assert [v["version"] for v in versions] == [1, 2, 3]
|
||||
assert v3["spec"] == new_spec
|
||||
assert len(Presentation(storage.get_file(v3["storage_path"])).slides) == 3
|
||||
|
||||
|
||||
def test_create_document_docx(wired_tool):
|
||||
tool, conversation_id, pg_engine, storage = wired_tool
|
||||
from docx import Document
|
||||
|
||||
spec = {"title": "Report", "sections": [{"heading": "Intro", "paragraphs": ["Hello.", "World."]}]}
|
||||
created = tool.execute_action("create_artifact", kind="document", spec=spec)
|
||||
assert created["status"] == "ok", created
|
||||
assert created["mime_type"].endswith("wordprocessingml.document")
|
||||
|
||||
with pg_engine.connect() as conn:
|
||||
version = ArtifactsRepository(conn).get_version(created["artifact_id"], 1)
|
||||
doc = Document(storage.get_file(version["storage_path"]))
|
||||
texts = [p.text for p in doc.paragraphs]
|
||||
assert "Hello." in texts and "World." in texts
|
||||
|
||||
|
||||
def test_create_spreadsheet_xlsx(wired_tool):
|
||||
tool, conversation_id, pg_engine, storage = wired_tool
|
||||
from openpyxl import load_workbook
|
||||
|
||||
spec = {"sheets": [{"name": "Data", "rows": [["a", "b"], [1, 2]]}]}
|
||||
created = tool.execute_action("create_artifact", kind="spreadsheet", spec=spec)
|
||||
assert created["status"] == "ok", created
|
||||
assert created["mime_type"].endswith("spreadsheetml.sheet")
|
||||
|
||||
with pg_engine.connect() as conn:
|
||||
version = ArtifactsRepository(conn).get_version(created["artifact_id"], 1)
|
||||
wb = load_workbook(storage.get_file(version["storage_path"]))
|
||||
ws = wb["Data"]
|
||||
assert ws.cell(row=1, column=1).value == "a"
|
||||
assert ws.cell(row=2, column=2).value == 2
|
||||
|
||||
|
||||
def test_invalid_spec_creates_no_artifact(wired_tool):
|
||||
tool, conversation_id, pg_engine, _storage = wired_tool
|
||||
|
||||
out = tool.execute_action("create_artifact", kind="presentation", spec={"title": "no slides"})
|
||||
assert out["status"] == "error"
|
||||
assert "invalid presentation spec" in out["error"]
|
||||
|
||||
with pg_engine.connect() as conn:
|
||||
rows = ArtifactsRepository(conn).list_artifacts(conversation_id=conversation_id)
|
||||
assert rows == []
|
||||
|
||||
|
||||
def test_edit_cross_tenant_denied(wired_tool):
|
||||
tool, conversation_id, pg_engine, storage = wired_tool
|
||||
|
||||
# An artifact owned by a DIFFERENT conversation must not be editable here.
|
||||
other_conversation = str(uuid.uuid4())
|
||||
foreign_id = _seed_presentation(pg_engine, storage, other_conversation)
|
||||
|
||||
out = tool.execute_action("edit_artifact", id=foreign_id, spec_patch={"title": "x"})
|
||||
assert out["status"] == "error"
|
||||
assert "not found in this conversation/run" in out["error"]
|
||||
|
||||
# The foreign artifact stayed at v1 (no version appended).
|
||||
with pg_engine.connect() as conn:
|
||||
artifact = ArtifactsRepository(conn).get_artifact(foreign_id)
|
||||
assert artifact["current_version"] == 1
|
||||
|
||||
|
||||
def _seed_presentation(pg_engine, storage, conversation_id) -> str:
|
||||
"""Create a presentation artifact (row + version + bytes) in a given conversation; return its id."""
|
||||
spec = {"slides": [{"title": "seed"}]}
|
||||
data = b"%PDF-stub-bytes" # content is irrelevant to the authz check
|
||||
with pg_engine.begin() as conn:
|
||||
repo = ArtifactsRepository(conn)
|
||||
artifact = repo.create_artifact(
|
||||
"user-e2e",
|
||||
"presentation",
|
||||
conversation_id=conversation_id,
|
||||
title="seed",
|
||||
mime_type="application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||
filename="seed.pptx",
|
||||
storage_path=None,
|
||||
size=len(data),
|
||||
sha256="0" * 64,
|
||||
spec=spec,
|
||||
produced_by={"tool": "artifact_generator", "action": "create_artifact", "spec_kind": "presentation"},
|
||||
)
|
||||
artifact_id = str(artifact["id"])
|
||||
from sqlalchemy import text
|
||||
|
||||
storage_path = f"inputs/user-e2e/artifacts/{artifact_id}/v1/seed.pptx"
|
||||
storage.save_file(io.BytesIO(data), storage_path)
|
||||
conn.execute(
|
||||
text(
|
||||
"UPDATE artifact_versions SET storage_path = :p "
|
||||
"WHERE artifact_id = CAST(:aid AS uuid) AND version = 1"
|
||||
),
|
||||
{"p": storage_path, "aid": artifact_id},
|
||||
)
|
||||
return artifact_id
|
||||
@@ -11,3 +11,10 @@ cryptography>=46.0.0
|
||||
jupyter-kernel-gateway>=2.5.0
|
||||
ipykernel>=6.0.0
|
||||
websocket-client>=1.6.0
|
||||
# Document renderers the artifact_generator runs inside the sandbox kernel.
|
||||
# In production these are pre-baked into the sandbox image; locally the kernel
|
||||
# uses the app .venv, so the integration test needs them installed. python-pptx
|
||||
# and openpyxl already ship in application/requirements.txt; only the rest are
|
||||
# test-only here (kept out of application/requirements.txt — sandbox-image libs).
|
||||
python-docx>=1.1
|
||||
reportlab>=4.0,<5
|
||||
Reference in new issue
Block a user