feat(studio): display images from Python tool execution in chat UI (#4778)
* feat(studio): display images from Python tool execution in chat UI
When the model calls the Python tool to create a matplotlib plot or
other image file, the image now displays inline in the chat output
instead of being invisible to the user.
Backend:
- Detect new image files (png/jpg/gif/webp/bmp) after Python subprocess
completes by diffing os.listdir before/after execution
- Append __IMAGES__ sentinel to tool result for frontend consumption
- Strip sentinel before injecting result into LLM context (role: tool)
so the model never sees file paths
- Add GET /sandbox/{session_id}/{filename} endpoint with JWT auth
(header or query param), path traversal protection, extension
allowlist, realpath containment check, and nosniff header
Frontend:
- Parse __IMAGES__ sentinel in tool_end SSE events, create structured
result with text/images/sessionId
- Render <img> tags in Python tool UI pointing at the sandbox endpoint
Also fixes a bug where SyntaxError in user code was misreported as
"unsafe code detected" instead of showing the actual Python traceback.
The _check_code_safety function now lets SyntaxError pass through to
the subprocess for a proper error message.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix(studio): improve SVG detection and strip XML preamble
Handle <?xml ...?> declarations before <svg> tags in code fences,
strip XML declaration from SVGs before data URI rendering, and
update the sloth suggestion prompt to request showing code.
* fix(studio): persist parentId so retries survive reload
The append() handler was destructuring only { message } from
ExportedMessageRepositoryItem and discarding parentId. When loading
a saved thread, load() used ExportedMessageRepository.fromArray()
which chains all messages sequentially, flattening retry branches
into a linear list.
Now append() writes parentId to the MessageRecord, and load()
reconstructs the tree when parentIds are present. Old threads
without parentId fall back to the existing fromArray() behavior.
* fix(studio): address review findings for image display and retry persistence
Image detection:
- Use mtime comparison instead of filename-only diff so overwritten
files (e.g. plt.savefig("chart.png") called twice) are detected
Sentinel parsing:
- Use rsplit/lastIndexOf instead of split/indexOf so user code that
prints __IMAGES__: does not collide with the backend sentinel
Mixed legacy/new threads:
- For old messages without a stored parentId, infer sequential parent
from the previous message instead of null, preventing multiple roots
Sandbox endpoint:
- Change Cache-Control from "public, max-age=3600" to "private,
no-store" since these are authenticated responses
---------
Co-authored-by: Daniel Han <danielhanchen@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han
GitHub
parent
5a5f1a4f34
commit
c8d311a053
9 changed files with 245 additions and 17 deletions
|
|
@ -2905,10 +2905,15 @@ class LlamaCppBackend:
|
|||
_error_prefixes
|
||||
)
|
||||
_tool_call_history.append((_tc_key, _is_error))
|
||||
# Strip image sentinel before feeding result to the LLM
|
||||
# (the full result with sentinel is still yielded via
|
||||
# tool_end so the frontend can extract image paths).
|
||||
_result_content = result
|
||||
if "\n__IMAGES__:" in _result_content:
|
||||
_result_content = _result_content.rsplit("\n__IMAGES__:", 1)[0]
|
||||
if _is_error:
|
||||
_result_content = (
|
||||
result + "\n\nThe tool call encountered an issue. "
|
||||
_result_content + "\n\nThe tool call encountered an issue. "
|
||||
"Please try a different approach or rephrase your request."
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -26,6 +26,10 @@ from loggers import get_logger
|
|||
logger = get_logger(__name__)
|
||||
|
||||
_EXEC_TIMEOUT = 300 # 5 minutes
|
||||
|
||||
# Strict raster-image allowlist for sandbox file serving.
|
||||
# No .svg (XSS risk via embedded scripts), no .html, no .pdf.
|
||||
_IMAGE_EXTS = frozenset({".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"})
|
||||
_MAX_OUTPUT_CHARS = 8000 # truncate long output
|
||||
_BASH_BLOCKED_WORDS = {"rm", "sudo", "dd", "chmod", "mkfs", "shutdown", "reboot"}
|
||||
|
||||
|
|
@ -591,6 +595,12 @@ def _check_code_safety(code: str) -> str | None:
|
|||
"""
|
||||
safe, info = _check_signal_escape_patterns(code)
|
||||
if not safe:
|
||||
# SyntaxError from ast.parse -- let these through so the subprocess
|
||||
# produces a normal Python traceback instead of a misleading
|
||||
# "unsafe code detected" message.
|
||||
if info.get("error"):
|
||||
return None
|
||||
|
||||
reasons = [
|
||||
item.get("description", "") for item in info.get("signal_tampering", [])
|
||||
]
|
||||
|
|
@ -634,6 +644,17 @@ def _python_exec(
|
|||
|
||||
tmp_path = None
|
||||
workdir = _get_workdir(session_id)
|
||||
# Snapshot image mtimes so we detect both new and overwritten files.
|
||||
_before: dict[str, int] = {}
|
||||
if os.path.isdir(workdir):
|
||||
for _name in os.listdir(workdir):
|
||||
if os.path.splitext(_name)[1].lower() in _IMAGE_EXTS:
|
||||
_p = os.path.join(workdir, _name)
|
||||
if os.path.isfile(_p):
|
||||
try:
|
||||
_before[_name] = os.stat(_p).st_mtime_ns
|
||||
except OSError:
|
||||
pass
|
||||
try:
|
||||
fd, tmp_path = tempfile.mkstemp(
|
||||
suffix = ".py", prefix = "studio_exec_", dir = workdir
|
||||
|
|
@ -669,7 +690,29 @@ def _python_exec(
|
|||
result = output or ""
|
||||
if proc.returncode != 0:
|
||||
result = f"Exit code {proc.returncode}:\n{result}"
|
||||
return _truncate(result) if result.strip() else "(no output)"
|
||||
result = _truncate(result) if result.strip() else "(no output)"
|
||||
|
||||
# Detect new or overwritten image files and append sentinel for frontend
|
||||
if session_id and os.path.isdir(workdir):
|
||||
new_images = []
|
||||
for _name in os.listdir(workdir):
|
||||
if os.path.splitext(_name)[1].lower() not in _IMAGE_EXTS:
|
||||
continue
|
||||
_p = os.path.join(workdir, _name)
|
||||
if not os.path.isfile(_p):
|
||||
continue
|
||||
try:
|
||||
_mtime = os.stat(_p).st_mtime_ns
|
||||
except OSError:
|
||||
continue
|
||||
if _name not in _before or _mtime != _before[_name]:
|
||||
new_images.append(_name)
|
||||
if new_images:
|
||||
import json as _json
|
||||
|
||||
result += f"\n__IMAGES__:{_json.dumps(sorted(new_images))}"
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
return f"Execution error: {e}"
|
||||
|
|
|
|||
|
|
@ -5,11 +5,12 @@
|
|||
Inference API routes for model loading and text generation.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import uuid
|
||||
from pathlib import Path
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from fastapi.responses import StreamingResponse, JSONResponse
|
||||
from typing import Optional
|
||||
import json
|
||||
|
|
@ -1680,6 +1681,94 @@ async def openai_chat_completions(
|
|||
raise HTTPException(status_code = 500, detail = str(e))
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# Sandbox file serving (/sandbox/{session_id}/{filename})
|
||||
# =====================================================================
|
||||
|
||||
_SANDBOX_MEDIA_TYPES = {
|
||||
".png": "image/png",
|
||||
".jpg": "image/jpeg",
|
||||
".jpeg": "image/jpeg",
|
||||
".gif": "image/gif",
|
||||
".webp": "image/webp",
|
||||
".bmp": "image/bmp",
|
||||
}
|
||||
|
||||
|
||||
@router.get("/sandbox/{session_id}/{filename}")
|
||||
async def serve_sandbox_file(
|
||||
session_id: str,
|
||||
filename: str,
|
||||
request: Request,
|
||||
token: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Serve image files created by Python tool execution.
|
||||
|
||||
Accepts auth via Authorization header OR ?token= query param
|
||||
(needed because <img src> cannot send custom headers).
|
||||
"""
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
# ── Authentication (header or query param) ──────────────────
|
||||
auth_header = request.headers.get("authorization")
|
||||
if auth_header and auth_header.lower().startswith("bearer "):
|
||||
jwt_token = auth_header[7:]
|
||||
elif token:
|
||||
jwt_token = token
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code = status.HTTP_401_UNAUTHORIZED,
|
||||
detail = "Missing authentication token",
|
||||
)
|
||||
from fastapi.security import HTTPAuthorizationCredentials
|
||||
|
||||
creds = HTTPAuthorizationCredentials(scheme = "Bearer", credentials = jwt_token)
|
||||
await get_current_subject(creds)
|
||||
|
||||
# ── Filename sanitization ───────────────────────────────────
|
||||
safe_filename = os.path.basename(filename)
|
||||
if not safe_filename or safe_filename in (".", ".."):
|
||||
raise HTTPException(status_code = 404, detail = "Not found")
|
||||
|
||||
# ── Extension allowlist ─────────────────────────────────────
|
||||
ext = os.path.splitext(safe_filename)[1].lower()
|
||||
media_type = _SANDBOX_MEDIA_TYPES.get(ext)
|
||||
if not media_type:
|
||||
raise HTTPException(
|
||||
status_code = status.HTTP_403_FORBIDDEN,
|
||||
detail = "File type not allowed",
|
||||
)
|
||||
|
||||
# ── Path containment check ──────────────────────────────────
|
||||
home = os.path.expanduser("~")
|
||||
sandbox_root = os.path.realpath(os.path.join(home, "studio_sandbox"))
|
||||
safe_session = os.path.basename(session_id.replace("..", ""))
|
||||
if not safe_session:
|
||||
raise HTTPException(status_code = 404, detail = "Not found")
|
||||
|
||||
file_path = os.path.realpath(
|
||||
os.path.join(sandbox_root, safe_session, safe_filename)
|
||||
)
|
||||
if not file_path.startswith(sandbox_root + os.sep):
|
||||
raise HTTPException(
|
||||
status_code = status.HTTP_403_FORBIDDEN,
|
||||
detail = "Access denied",
|
||||
)
|
||||
|
||||
if not os.path.isfile(file_path):
|
||||
raise HTTPException(status_code = 404, detail = "Not found")
|
||||
|
||||
return FileResponse(
|
||||
path = file_path,
|
||||
media_type = media_type,
|
||||
headers = {
|
||||
"Cache-Control": "private, no-store",
|
||||
"X-Content-Type-Options": "nosniff",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# =====================================================================
|
||||
# OpenAI-Compatible Models Listing (/models → /v1/models)
|
||||
# =====================================================================
|
||||
|
|
|
|||
|
|
@ -100,20 +100,27 @@ function getCodeFilename(language: string | null) {
|
|||
function isSvgFence(codeFence: CodeFence): boolean {
|
||||
const lang = codeFence.language?.toLowerCase() ?? "";
|
||||
if (lang === "svg") return true;
|
||||
if ((lang === "xml" || lang === "html") && codeFence.source.trimStart().startsWith("<svg")) return true;
|
||||
if (lang === "xml" || lang === "html") {
|
||||
const trimmed = codeFence.source.trimStart();
|
||||
// Match <svg directly or <?xml ...?> followed by <svg
|
||||
if (trimmed.startsWith("<svg")) return true;
|
||||
if (trimmed.startsWith("<?xml") && trimmed.includes("<svg")) return true;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
function isHtmlFence(codeFence: CodeFence): boolean {
|
||||
const lang = codeFence.language?.toLowerCase() ?? "";
|
||||
return lang === "html" && !codeFence.source.trimStart().startsWith("<svg");
|
||||
return lang === "html" && !isSvgFence(codeFence);
|
||||
}
|
||||
|
||||
const UNSAFE_SVG_RE = /<script[\s>]|on\w+\s*=|javascript:|<foreignObject[\s>]|<iframe[\s>]|<embed[\s>]|<object[\s>]/i;
|
||||
|
||||
function sanitizeSvg(source: string): string | null {
|
||||
if (UNSAFE_SVG_RE.test(source)) return null;
|
||||
return source;
|
||||
// Strip XML declaration (<?xml ...?>) -- not needed for data URI
|
||||
// rendering and can cause issues with some renderers.
|
||||
return source.replace(/^\s*<\?xml[^?]*\?>\s*/i, "");
|
||||
}
|
||||
|
||||
function SvgPreview({ source }: { source: string }) {
|
||||
|
|
|
|||
|
|
@ -124,7 +124,7 @@ const SUGGESTION_TOOLS: Record<string, Array<"thinking" | "search" | "code">> =
|
|||
"How do you fine-tune an audio model with Unsloth?": ["thinking", "search"],
|
||||
"Create a live weather dashboard in HTML using no API key. Show me the code": ["thinking", "code", "search"],
|
||||
"Solve the integral of x·sin(x), and verify it step by step": ["thinking", "code"],
|
||||
"Draw an SVG of a cute sloth": ["thinking", "code", "search"],
|
||||
"Draw an SVG of a cute sloth & show the code": ["thinking", "code", "search"],
|
||||
};
|
||||
|
||||
const toolIconMap = {
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@
|
|||
"use client";
|
||||
|
||||
import { copyToClipboard } from "@/lib/copy-to-clipboard";
|
||||
import { getAuthToken } from "@/features/auth/session";
|
||||
import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
|
||||
import { code as codePlugin } from "@streamdown/code";
|
||||
import { CheckIcon, CodeIcon, CopyIcon, LoaderIcon } from "lucide-react";
|
||||
|
|
@ -15,6 +16,12 @@ import {
|
|||
ToolFallbackTrigger,
|
||||
} from "./tool-fallback";
|
||||
|
||||
interface StructuredResult {
|
||||
text: string;
|
||||
images: string[];
|
||||
sessionId: string;
|
||||
}
|
||||
|
||||
const MAX_DISPLAY = 10_000;
|
||||
const COPY_RESET_MS = 2000;
|
||||
const SHIKI_THEME = ["github-light", "github-dark"] as ["github-light", "github-dark"];
|
||||
|
|
@ -84,6 +91,16 @@ function HighlightedCode({ code: source, language }: { code: string; language: s
|
|||
);
|
||||
}
|
||||
|
||||
function isStructuredResult(val: unknown): val is StructuredResult {
|
||||
return (
|
||||
typeof val === "object" &&
|
||||
val !== null &&
|
||||
"text" in val &&
|
||||
"images" in val &&
|
||||
"sessionId" in val
|
||||
);
|
||||
}
|
||||
|
||||
const PythonToolUIImpl: ToolCallMessagePartComponent = ({
|
||||
args,
|
||||
result,
|
||||
|
|
@ -92,12 +109,24 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
const code = (args as { code?: string })?.code ?? "";
|
||||
const firstLine = code.split("\n")[0]?.slice(0, 60) ?? "";
|
||||
const isRunning = status?.type === "running";
|
||||
const output =
|
||||
typeof result === "string"
|
||||
? result
|
||||
: result
|
||||
? JSON.stringify(result, null, 2)
|
||||
: "";
|
||||
|
||||
let output: string;
|
||||
let images: string[] = [];
|
||||
let sessionId = "";
|
||||
|
||||
if (isStructuredResult(result)) {
|
||||
output = result.text;
|
||||
images = result.images;
|
||||
sessionId = result.sessionId;
|
||||
} else if (typeof result === "string") {
|
||||
output = result;
|
||||
} else if (result) {
|
||||
output = JSON.stringify(result, null, 2);
|
||||
} else {
|
||||
output = "";
|
||||
}
|
||||
|
||||
const authToken = getAuthToken();
|
||||
|
||||
return (
|
||||
<ToolFallbackRoot>
|
||||
|
|
@ -133,6 +162,21 @@ const PythonToolUIImpl: ToolCallMessagePartComponent = ({
|
|||
</pre>
|
||||
</div>
|
||||
) : null}
|
||||
|
||||
{/* Images from Python tool execution */}
|
||||
{images.length > 0 && sessionId && (
|
||||
<div className="mt-2 flex flex-col gap-2">
|
||||
{images.map((filename) => (
|
||||
<img
|
||||
key={filename}
|
||||
src={`/api/inference/sandbox/${encodeURIComponent(sessionId)}/${encodeURIComponent(filename)}${authToken ? `?token=${encodeURIComponent(authToken)}` : ""}`}
|
||||
alt={filename}
|
||||
loading="lazy"
|
||||
className="max-w-full rounded border border-border"
|
||||
/>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</ToolFallbackContent>
|
||||
</ToolFallbackRoot>
|
||||
|
|
|
|||
|
|
@ -635,7 +635,23 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
|
|||
toolCallParts[toolCallParts.length - 1]?.toolCallId || "";
|
||||
const idx = toolCallParts.findIndex((p) => p.toolCallId === id);
|
||||
if (idx !== -1) {
|
||||
toolCallParts[idx] = { ...toolCallParts[idx], result: toolEvent.result as string };
|
||||
const rawResult = (toolEvent.result as string) ?? "";
|
||||
const imgMarker = "\n__IMAGES__:";
|
||||
const imgIdx = rawResult.lastIndexOf(imgMarker);
|
||||
let parsedResult: string | { text: string; images: string[]; sessionId: string };
|
||||
if (imgIdx !== -1) {
|
||||
const text = rawResult.slice(0, imgIdx);
|
||||
const sessionId = unstable_threadId || "";
|
||||
try {
|
||||
const images = JSON.parse(rawResult.slice(imgIdx + imgMarker.length)) as string[];
|
||||
parsedResult = { text, images, sessionId };
|
||||
} catch {
|
||||
parsedResult = rawResult;
|
||||
}
|
||||
} else {
|
||||
parsedResult = rawResult;
|
||||
}
|
||||
toolCallParts[idx] = { ...toolCallParts[idx], result: parsedResult };
|
||||
}
|
||||
}
|
||||
// Yield cumulative state so tool UI updates (tools first, text after)
|
||||
|
|
|
|||
|
|
@ -47,9 +47,9 @@ const DEFAULT_SUGGESTIONS = [
|
|||
prompt: "Solve the integral of x·sin(x), and verify it step by step",
|
||||
},
|
||||
{
|
||||
title: "Draw an SVG of a cute sloth",
|
||||
title: "Draw an SVG of a cute sloth & show the code",
|
||||
label: "SVG sloth",
|
||||
prompt: "Draw an SVG of a cute sloth",
|
||||
prompt: "Draw an SVG of a cute sloth & show the code",
|
||||
},
|
||||
];
|
||||
|
||||
|
|
@ -569,10 +569,32 @@ function ThreadHistoryProvider({
|
|||
store.setContextUsage(savedUsage);
|
||||
}
|
||||
|
||||
// If any message has a stored parentId, reconstruct the tree
|
||||
// so retries/regenerations load as branches instead of being
|
||||
// unrolled into a flat list. For mixed legacy/new threads
|
||||
// (old messages without parentId + new messages with), infer
|
||||
// sequential parents for old messages to preserve the chain.
|
||||
// Fall back to fromArray for fully legacy threads.
|
||||
const hasParentIds = msgs.some((m) => "parentId" in m);
|
||||
if (hasParentIds) {
|
||||
let previousId: string | null = null;
|
||||
return {
|
||||
messages: msgs.map((m) => {
|
||||
const parentId = "parentId" in m
|
||||
? (m.parentId ?? null)
|
||||
: previousId;
|
||||
previousId = m.id;
|
||||
return {
|
||||
parentId,
|
||||
message: toThreadMessage(m),
|
||||
};
|
||||
}),
|
||||
};
|
||||
}
|
||||
return ExportedMessageRepository.fromArray(msgs.map(toThreadMessage));
|
||||
},
|
||||
|
||||
async append({ message }: ExportedMessageRepositoryItem) {
|
||||
async append({ parentId, message }: ExportedMessageRepositoryItem) {
|
||||
const { remoteId } = await aui.threadListItem().initialize();
|
||||
const content = cloneContent(message.content);
|
||||
const attachments =
|
||||
|
|
@ -586,6 +608,7 @@ function ThreadHistoryProvider({
|
|||
await db.messages.put({
|
||||
id: message.id,
|
||||
threadId: remoteId,
|
||||
parentId: parentId ?? null,
|
||||
role: message.role,
|
||||
content,
|
||||
...(attachments.length > 0 && { attachments }),
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ export interface ThreadRecord {
|
|||
export interface MessageRecord {
|
||||
id: string;
|
||||
threadId: string;
|
||||
parentId?: string | null;
|
||||
role: import("@assistant-ui/react").ThreadMessage["role"];
|
||||
content: import("@assistant-ui/react").ThreadMessage["content"];
|
||||
attachments?: import("@assistant-ui/react").ThreadMessage["attachments"];
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue